Add XTC and token ban sampling steps

[ethanj801]

Summary

This PR introduces sampler steps that enable XTC sampling and the banning of tokens. Previous versions such as EXL2 had these sampling features, and they are present (but not hooked up) downstream in TabbyAPI (for the EXL3 backend, they work for the EXL2 backend).

This PR also includes a bugfix for an issue I encountered during testing.

Design Notes

• SS_BanTokens: Very straightforward class. We cache the mask to avoid needing to recompute each token.
• SS_XTC:
  • We implement this in the same way that EXL2 does, avoiding explicit randomness and instead creating the distribution mixture directly. This leads to the same sampling results as the probabilistic approach, except in cases where argmax sampling is used. This is reasonable since greedy sampling would seem to defeat the purpose of using XTC in the first place.
  • We use the same token exclusions as EXL2: all new line containing tokens and all special tokens.

Bugfix

During testing, SS_Argmax appeared to fail on unsorted states. It appears that this is due to the following.

The sorted branches indexed state.indices with the argmax positions alone, which selects
rows of the index tensor instead of one token per row:

case SS.PROBS_S | SS.PROBS_N_S:
    temp = torch.argmax(state.probs, dim = -1)
    state.sample = state.indices[temp]      # PR NOTE: issue is on this line

This appears to have been unreachable before due to the fusion meaning this code was basically never reached. I corrected it by applying the pattern from the SS.LOGITS_S state.

state.sample = state.indices[buffered_arange(state.bsz, state.in_logits.device), temp]

Testing

I added some tests to tests/test_sampler.py following the existing pattern. The tests check a few various edge cases and also show that our implementation has the same distributional result as the random version of XTC.

137 passed in 53.80s                          # fused sampler enabled
136 passed, 1 skipped in 119.60s              # EXL3_FUSED_SAMPLER=0

Reverting only custom.py to the first commit reproduces the SS_Argmax failure (RuntimeError: shape '[2]' is invalid for input of size 8), and restoring it passes.

Limitations

I did not implement either token banning or XTC into ComboSampler. My understanding is ComboSampler is designed to provide an easy interface for sampling, and so it would be ideal to add these samplers there. However, both of these steps require a tokenizer, which ComboSampler currently does not take. If desired we could add a tokenizer parameter to ComboSampler and add xtc parameter, token bans and an eos token ban option.

[turboderp]

Thanks. 👍