diff --git a/exllamav3/generator/sampler/__init__.py b/exllamav3/generator/sampler/__init__.py index 10bf464e..0c38ada0 100644 --- a/exllamav3/generator/sampler/__init__.py +++ b/exllamav3/generator/sampler/__init__.py @@ -16,6 +16,9 @@ SS_RepP, SS_PresFreqP, SS_AdaptiveP, + SS_BanTokens, + SS_XTC, + xtc_default_protected_token_ids, ) from .presets import ( DefaultSampler, diff --git a/exllamav3/generator/sampler/custom.py b/exllamav3/generator/sampler/custom.py index 91de53ba..32175e82 100644 --- a/exllamav3/generator/sampler/custom.py +++ b/exllamav3/generator/sampler/custom.py @@ -10,6 +10,7 @@ import random from dataclasses import dataclass from enum import Enum +from functools import lru_cache from ...util import profile_opt import torch.nn.functional as F @@ -100,10 +101,10 @@ def run(self, state: SamplingState): state.sample = torch.argmax(state.probs, dim = -1) case SS.LOGITS_S: temp = torch.argmax(state.logits, dim = -1) - state.sample = state.indices[temp] + state.sample = state.indices[buffered_arange(state.bsz, state.in_logits.device), temp] case SS.PROBS_S | SS.PROBS_N_S: temp = torch.argmax(state.probs, dim = -1) - state.sample = state.indices[temp] + state.sample = state.indices[buffered_arange(state.bsz, state.in_logits.device), temp] state.state = SS.DONE @@ -453,7 +454,7 @@ def alt(self): class SS_TopP(SS_Base): """ Identify the smallest set of top tokens with a cumulative probability greater than P, mask out all - remainig tokens + remaining tokens """ def __init__(self, top_p: float): self.top_p = top_p @@ -522,6 +523,161 @@ def alt(self): return None +class TokenMask: + """ + Boolean mask over the vocabulary selecting a fixed set of token IDs with per device and + vocabulary size caching. + """ + def __init__(self, token_ids): + self.token_ids = sorted({int(t) for t in token_ids}) + self.masks = {} + + def get(self, dim: int, device: torch.device) -> torch.Tensor: + key = (dim, device) + mask = self.masks.get(key) + if mask is None: + mask = torch.zeros((dim,), dtype = torch.bool, device = device) + # IDs outside the vocabulary would be out of bounds for the mask + ids = [t for t in self.token_ids if 0 <= t < dim] + if ids: + mask[torch.tensor(ids, dtype = torch.long, device = device)] = True + self.masks[key] = mask + return mask + + +class SS_BanTokens(SS_Base): + """ + Mask out a fixed set of token IDs + """ + def __init__(self, token_ids: list[int]): + self.mask = TokenMask(token_ids) + + def run(self, state: SamplingState): + match state.state: + case SS.INIT: + state.logits = state.in_logits.to(torch.float, copy = True) + state.logits.masked_fill_(self.mask.get(state.dim, state.logits.device), -float("inf")) + state.state = SS.LOGITS + case SS.LOGITS: + state.logits.masked_fill_(self.mask.get(state.dim, state.logits.device), -float("inf")) + case SS.LOGITS_S: + mask = self.mask.get(state.dim, state.logits.device) + state.logits.masked_fill_(mask[state.indices], -float("inf")) + case SS.PROBS | SS.PROBS_N: + state.probs.masked_fill_(self.mask.get(state.dim, state.probs.device), 0.0) + state.state = SS.PROBS + case SS.PROBS_S | SS.PROBS_N_S: + mask = self.mask.get(state.dim, state.probs.device) + state.probs.masked_fill_(mask[state.indices], 0.0) + state.state = SS.PROBS_S + + def alt(self): + if not self.mask.token_ids: + return SS_NoOp() + return None + + +@lru_cache(10) +def xtc_default_protected_token_ids(tokenizer: Tokenizer) -> frozenset[int]: + """ + Default set of token IDs for SS_XTC to leave in place, being every piece containing a newline + plus every extended token. Matches ExLlamaV2Sampler.get_default_xtc_mask_tokens in exllamav2. + """ + pieces = tokenizer.get_id_to_piece_list(include_special_tokens = True) + protected = {t for t, piece in enumerate(pieces) if "\n" in piece} + protected.update(tokenizer.extended_id_to_piece.keys()) + return frozenset(protected) + + +class SS_XTC(SS_Base): + """ + Exclude Top Choices. Of the tokens reaching the threshold probability, all but the least likely + one are excluded with probability p. Protected token IDs are held out of consideration. + + The step averages the two outcomes of the exclusion, weighted by p. A categorical final step + then samples from exactly the distribution a random exclusion would give, while SS_Argmax + returns a fixed token where a random exclusion would alternate between two. + + Matches xtc_cpu in exllamav2. The original XTC sampler and llama.cpp draw the outcome once per + sampler call and truncate the distribution. + """ + def __init__( + self, + probability: float, + threshold: float, + protected_token_ids: frozenset[int] | list[int] | None = None, + tokenizer: Tokenizer | None = None, + ): + """ + :param probability: + Probability of the exclusion happening. 0.0 disables the step, 1.0 makes it always happen + :param threshold: + Minimum probability for a token to be considered for exclusion. Values above 0.5 make the step + a no-op, since at most one token can exceed half the total probability + :param protected_token_ids: + Token IDs removed from the candidate set before the exclusion is chosen + :param tokenizer: + Used to derive a default set of protected IDs when protected_token_ids is None. Ignored if + protected_token_ids is given. The default set comes from xtc_default_protected_token_ids and + covers every token whose piece contains a newline plus every token in Tokenizer.extended_id_to_piece + """ + self.probability = probability + self.threshold = threshold + assert 0.0 <= probability <= 1.0 + assert 0.0 <= threshold <= 1.0 + if protected_token_ids is None and tokenizer is not None: + protected_token_ids = xtc_default_protected_token_ids(tokenizer) + self.protected = TokenMask(protected_token_ids) if protected_token_ids else None + + def run(self, state: SamplingState): + match state.state: + case SS.PROBS_N_S: + pass + case _: + raise ValueError("Sampling logic error") + + # The probabilities are normalized and sorted descending, so at most 1/threshold of them + # can reach the threshold and those are the first ones + k = state.dim if self.threshold <= 0.0 else min(state.dim, int(1.0 / self.threshold) + 1) + probs = state.probs[:, :k] + + qualifies = probs >= self.threshold + if self.protected is not None: + protected = self.protected.get(state.dim, probs.device) + qualifies &= ~protected[state.indices[:, :k]] + + # Sorted descending, so the last qualifying position in a row is its least likely one + counts = qualifies.sum(dim = -1, keepdim = True) + excluded = qualifies & (qualifies.cumsum(dim = -1) < counts) + + # Averaged over the two outcomes of the roll, an excluded token is scaled by (1 - p) and + # every other token by (1 + p * m / (1 - m)), for excluded mass m. Multiplying only the + # excluded tokens, by the ratio of those two factors, is proportional to that average and + # leaves the result unnormalized + x_mass = (probs * excluded).sum(dim = -1, keepdim = True) + scale = (1.0 - self.probability) / (1.0 + self.probability * x_mass / (1.0 - x_mass)) + probs *= torch.where(excluded, scale, torch.ones_like(scale)) + state.state = SS.PROBS_S + + def prep(self, in_state: SS): + match in_state: + case SS.PROBS_N: + return [SS_Sort] + case SS.INIT | SS.LOGITS | SS.PROBS: + return [SS_Normalize, SS_Sort] + case SS.LOGITS_S | SS.PROBS_S: + return [SS_Normalize] + case _: + return None + + def alt(self): + # Two tokens cannot both exceed half the total probability, so above that threshold there + # is never more than one token to choose between + if self.probability == 0.0 or self.threshold > 0.5: + return SS_NoOp() + return None + + class SS_RepP(SS_Base): """ Apply Transformers style repetition penalties based on past token IDs. Must be the first step in sampler @@ -733,8 +889,10 @@ def __init__( super().__init__() # Simplify the stack (identity steps become no-ops), then collapse an eligible tail - # into the fused kernel step. Leading penalty steps are kept as-is; they feed fp32 - # logits to the fused step. Ineligible stacks fall through to the step-by-step path. + # into the fused kernel step. Leading penalty and ban steps are kept as-is; they feed + # fp32 logits to the fused step, which already accepts -inf entries from the logit mask + # and the padded vocabulary region. Ineligible stacks fall through to the step-by-step + # path. simplified = [] for step in steps: self.reqs_past_ids = self.reqs_past_ids or step.reqs_past_ids() @@ -749,7 +907,7 @@ def __init__( fused_tail = None if fused_sampler_enable: i = 0 - while i < len(simplified) and type(simplified[i]) in (SS_RepP, SS_PresFreqP): + while i < len(simplified) and type(simplified[i]) in (SS_RepP, SS_PresFreqP, SS_BanTokens): i += 1 fused_tail = _match_fused_tail(simplified[i:]) if fused_tail is not None: diff --git a/exllamav3/generator/sampler/presets.py b/exllamav3/generator/sampler/presets.py index 5d839df7..3db487de 100644 --- a/exllamav3/generator/sampler/presets.py +++ b/exllamav3/generator/sampler/presets.py @@ -81,7 +81,7 @@ def __init__(self, top_p: float, temperature: float = 1.0, temperature_last = Fa class ComboSampler(CustomSampler): """ - Single class with an argument for each sampling step + Single class with an argument for common sampling steps """ def __init__( self, diff --git a/tests/test_sampler.py b/tests/test_sampler.py index 8fa20f07..4eca8c06 100644 --- a/tests/test_sampler.py +++ b/tests/test_sampler.py @@ -127,6 +127,100 @@ "expect_logits": [[3.0, 2.9, 2.8, 2.7, 2.6]] * 3, "expect_indices": [[0, 9, 8, 7, 6]] * 3, }, + { + "name": "ban_tokens", + "sampler": CustomSampler([ + SS_BanTokens([1, 3]) + ]), + "input": [[1.0, 2.0, 3.0, 4.0, 5.0]] * 2, + "expect_logits": [[1.0, ni, 3.0, ni, 5.0]] * 2, + }, + { + "name": "ban_tokens, sorted", + "sampler": CustomSampler([ + SS_Sort(), + SS_BanTokens([0, 2]) + ]), + "input": [[5.0, 3.0, 4.0, 1.0]], + "expect_indices": [[0, 2, 1, 3]], + "expect_logits": [[ni, ni, 3.0, 1.0]], + }, + { + "name": "ban_tokens, after temperature", + "sampler": CustomSampler([ + SS_Temperature(2.0), + SS_BanTokens([1]) + ]), + "input": [[2.0, 4.0, 6.0, 8.0]], + "expect_logits": [[1.0, ni, 3.0, 4.0]], + }, + { + "name": "ban_tokens, normalized", + "sampler": CustomSampler([ + SS_Normalize(), + SS_BanTokens([0]) + ]), + "input": [[2.0, 1.0, 0.0, -1.0]], + "expect_probs": [[0.0, 0.236883, 0.087144, 0.032059]], + }, + { + # Token IDs past the end of the vocabulary are ignored + "name": "ban_tokens, normalized and sorted", + "sampler": CustomSampler([ + SS_Normalize(), + SS_Sort(), + SS_BanTokens([1, 999]) + ]), + "input": [[2.0, 1.0, 0.0, -1.0]], + "expect_indices": [[0, 1, 2, 3]], + "expect_probs": [[0.643914, 0.0, 0.087144, 0.032059]], + }, + { + # Two tokens over the threshold, so only the more likely one is excluded + "name": "xtc", + "sampler": CustomSampler([ + SS_XTC(1.0, 0.1) + ]), + "input": [[2.0, 1.0, 0.0, -1.0]] * 2, + "expect_probs": [[0.0, 0.236883, 0.087144, 0.032059]] * 2, + }, + { + # Only one token over the threshold, leaving nothing to choose between + "name": "xtc, single candidate", + "sampler": CustomSampler([ + SS_XTC(1.0, 0.3) + ]), + "input": [[2.0, 1.0, 0.0, -1.0]], + "expect_probs": [[0.643914, 0.236883, 0.087144, 0.032059]], + }, + { + # Three tokens over the threshold, of which the two more likely are excluded + "name": "xtc, unprotected", + "sampler": CustomSampler([ + SS_XTC(1.0, 0.1) + ]), + "input": [[2.0, 1.0, 0.5, -1.0]], + "expect_probs": [[0.0, 0.0, 0.135989, 0.030343]], + }, + { + # Same distribution with the least likely of the three protected, so it neither is excluded + # nor counts as the one to keep, and the token above it survives in its place + "name": "xtc, protected token", + "sampler": CustomSampler([ + SS_XTC(1.0, 0.1, protected_token_ids = [2]) + ]), + "input": [[2.0, 1.0, 0.5, -1.0]], + "expect_probs": [[0.0, 0.224208, 0.135989, 0.030343]], + }, + { + # Reweighting for a probability below 1, leaving the result unnormalized + "name": "xtc, partial probability", + "sampler": CustomSampler([ + SS_XTC(0.5, 0.1) + ]), + "input": [[2.0, 1.0, 0.0, -1.0]], + "expect_probs": [[0.169081, 0.236883, 0.087144, 0.032059]], + }, ] @@ -276,9 +370,25 @@ def test_fused_collapse(): assert fused_mode(ComboSampler(temperature = 0.8, min_p = 0.05, top_k = 50, top_p = 0.9)) == SS_Fused.MODE_SAMPLE_FILTERS # Leading penalties keep a fused tail assert fused_mode(ComboSampler(rep_p = 1.2, temperature = 0.8, min_p = 0.05)) == SS_Fused.MODE_SAMPLE_MINP + # A leading ban step also keeps a fused tail + assert fused_mode(CustomSampler([ + SS_BanTokens([1, 2]), SS_Temperature(0.8), SS_MinP(0.05), SS_Sample() + ])) == SS_Fused.MODE_SAMPLE_MINP # Reference/multinomial nodes and non-canonical filter orders stay on the eager path assert fused_mode(CustomSampler([SS_MinP(0.1), SS_Sample_mn()])) is None assert fused_mode(CustomSampler([SS_TopP(0.9), SS_MinP(0.1), SS_Temperature(0.8), SS_Sample()])) is None + # XTC needs the sorted distribution, so it has no fused form + assert fused_mode(CustomSampler([SS_Temperature(0.8), SS_XTC(0.5, 0.1), SS_Sample()])) is None + # An empty ban list, a zero XTC probability and an XTC threshold above half are no-ops, and a + # no-op does not keep the tail off the fused path + for neutral in [ + SS_BanTokens([]), + SS_XTC(0.0, 0.1), + SS_XTC(0.5, 0.6), + ]: + assert fused_mode(CustomSampler([ + neutral, SS_Temperature(0.8), SS_MinP(0.05), SS_Sample() + ])) == SS_Fused.MODE_SAMPLE_MINP @pytest.mark.parametrize("dim", dims) @@ -392,3 +502,148 @@ def build(): if not torch.equal(a, b): mismatches += 1 assert mismatches == 0 + + +@torch.inference_mode() +def test_argmax_sorted(): + """ + Stacks ending in Argmax normally collapse to the fused step. An explicit sort, XTC, or a + disabled fused sampler reaches the eager path, where Argmax has to map a sorted position back + to a token ID per row. + """ + logits = torch.tensor( + [[2.0, 1.0, 0.0, -1.0], + [0.0, 3.0, 1.0, 2.0]], + dtype = torch.float, + device = device + ) + + enabled = sampler_custom.fused_sampler_enable + try: + sampler_custom.fused_sampler_enable = False + cases = [ + ([SS_Argmax()], [0, 1]), + ([SS_Sort(), SS_Argmax()], [0, 1]), + ([SS_TopK(3), SS_Argmax()], [0, 1]), + ([SS_MinP(0.1), SS_Argmax()], [0, 1]), + # Two tokens per row clear the threshold, and scaling the more likely one drops it + # below the one XTC keeps + ([SS_XTC(1.0, 0.1), SS_Argmax()], [1, 3]), + # The same scaling at a low probability is too small to reorder them + ([SS_XTC(0.05, 0.1), SS_Argmax()], [0, 1]), + ] + for steps, expected in cases: + sample = CustomSampler(steps).forward(logits, rand_u32 = 0) + assert sample.tolist() == expected, (steps, sample.tolist(), expected) + finally: + sampler_custom.fused_sampler_enable = enabled + + +@pytest.mark.parametrize("dim", dims) +@torch.inference_mode() +def test_ban_tokens(dim: tuple): + """ + The ban runs before the collapsed tail, so the fused kernel has to honor the -inf logits it + writes. + """ + torch.manual_seed(0) + random.seed(0) + logits = torch.randn(dim, dtype = torch.half, device = device) * 2 + + banned = sorted({0, 1, 2, dim[-1] // 2, dim[-1] - 1}) + sampler = CustomSampler([ + SS_BanTokens(banned), + SS_Temperature(0.8), + SS_MinP(0.02), + SS_Sample() + ]) + if sampler_custom.fused_sampler_enable: + assert fused_mode(sampler) == SS_Fused.MODE_SAMPLE_MINP + + banned_t = torch.tensor(banned, dtype = torch.long, device = device) + for seed in range(200): + sample = sampler.forward(logits, rand_u32 = seed) + assert not torch.isin(sample, banned_t).any() + + +@pytest.mark.parametrize("dim", dims) +@pytest.mark.parametrize("probability", [0.5, 1.0]) +@torch.inference_mode() +def test_xtc(dim: tuple, probability): + """ + SS_XTC reweights instead of drawing an outcome per token. The reference below is the + distribution that draw produces on average, so the sampled tokens have to follow it. + """ + torch.manual_seed(0) + random.seed(0) + threshold = 0.1 + + logits = torch.randn(dim, dtype = torch.half, device = device) * 2 + probs_ref = torch.softmax(logits.float(), dim = -1) + + # Reference for the excluded set, which is every token above the threshold except the least + # likely of them + sorted_probs, sorted_indices = torch.sort(probs_ref, dim = -1, descending = True) + qualifies = sorted_probs >= threshold + counts = qualifies.sum(dim = -1, keepdim = True) + excluded = torch.zeros_like(qualifies).scatter_( + -1, sorted_indices, qualifies & (qualifies.cumsum(dim = -1) < counts) + ) + + truncated = probs_ref.masked_fill(excluded, 0.0) + truncated /= truncated.sum(dim = -1, keepdim = True) + probs_ref = (1.0 - probability) * probs_ref + probability * truncated + + sampler = CustomSampler([SS_XTC(probability, threshold), SS_Sample()]) + assert fused_mode(sampler) is None + + num_samples = min(dim[-1] * 200, 10000) + samples = torch.empty((dim[0], 0), dtype = torch.long, device = device) + for _ in range(num_samples): + sample = sampler.forward(logits).unsqueeze(-1) + samples = torch.cat((samples, sample), dim = -1) + + hb = [torch.bincount(samples[b], minlength = dim[1]) for b in range(dim[0])] + histogram = torch.stack(hb).float() + histogram /= num_samples + + # Untruncated sampling has a large effective support on big vocabularies; see test_minp + k_eff = (probs_ref > 1e-5).sum(dim = -1).max().item() + chisq = compare(histogram, probs_ref) + assert chisq < max(0.01, 3.0 * k_eff / num_samples) + + +@pytest.mark.parametrize("dim", dims) +@pytest.mark.parametrize("threshold", [0.02, 0.1, 0.4]) +@pytest.mark.parametrize("probability", [0.35, 1.0]) +@torch.inference_mode() +def test_xtc_bound(dim: tuple, threshold, probability): + """ + SS_XTC reads only the first 1/threshold positions. The reference below scans the whole + vocabulary, so a wrong bound shows up as a wrong distribution. Both are normalized before + comparing, the step leaving its result unnormalized. + """ + torch.manual_seed(0) + random.seed(0) + + logits = torch.randn(dim, dtype = torch.half, device = device) * 2 + probs = torch.softmax(logits.float(), dim = -1) + + sorted_probs, sorted_indices = torch.sort(probs, dim = -1, descending = True) + qualifies = sorted_probs >= threshold + counts = qualifies.sum(dim = -1, keepdim = True) + excluded = qualifies & (qualifies.cumsum(dim = -1) < counts) + x_mass = (sorted_probs * excluded).sum(dim = -1, keepdim = True) + scale = (1.0 - probability) / (1.0 + probability * x_mass / (1.0 - x_mass)) + reference = torch.zeros_like(probs).scatter_( + -1, + sorted_indices, + sorted_probs * torch.where(excluded, scale, torch.ones_like(scale)) + ) + reference /= reference.sum(dim = -1, keepdim = True) + + state = CustomSampler([SS_XTC(probability, threshold)]).forward(logits, return_state = True) + result = torch.zeros_like(probs).scatter_(-1, state.indices, state.probs) + result /= result.sum(dim = -1, keepdim = True) + + torch.testing.assert_close(result, reference)