From 278267e06389f493ffec53d977ca24287d162259 Mon Sep 17 00:00:00 2001 From: wenlive <33243841+wenlive@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:23:37 +0800 Subject: [PATCH 1/2] fix(sampling): bound top-k before int32 conversion --- python/freetoken/engine/sample.py | 4 ++- tests/engine/test_sampling.py | 59 +++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) create mode 100644 tests/engine/test_sampling.py diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index c6fb2f40ee..bfff814d4b 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -79,8 +79,10 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: # Greedy outputs are selected explicitly in sample(); use neutral sampling # parameters for those rows instead of approximating argmax at low temperature. ts = [1.0 if g else max(p.temperature, MIN_T) for p, g in zip(params, is_greedy)] + # Keeping the full vocabulary is unfiltered. Bound Python integers before + # converting them to the device's int32 tensor. top_ks = [ - p.top_k if not g and p.top_k >= 1 else self.vocab_size + min(p.top_k, self.vocab_size) if not g and p.top_k >= 1 else self.vocab_size for p, g in zip(params, is_greedy) ] top_ps = [ diff --git a/tests/engine/test_sampling.py b/tests/engine/test_sampling.py new file mode 100644 index 0000000000..eb738771a7 --- /dev/null +++ b/tests/engine/test_sampling.py @@ -0,0 +1,59 @@ +"""Sampling bounds must preserve full-vocabulary and mixed-batch semantics.""" + +import pytest +import torch + +from freetoken.core import Batch, Req, SamplingParams +from freetoken.engine.sample import Sampler + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="a CUDA or ROCm GPU is required" +) + + +def _batch(params): + return Batch(reqs=[ + Req(input_ids=torch.tensor([0], dtype=torch.int32), table_idx=i, + cached_len=0, output_len=4, uid=i, sampling_params=p, cache_handle=None) + for i, p in enumerate(params) + ], phase="decode") + + +@pytest.mark.parametrize("top_k", [-1, 0, 17, 18, 2**31 - 1, 2**31, 2**80]) +@pytest.mark.parametrize("top_p", [1.0, 0.8]) +def test_full_vocabulary_top_k_disables_filter(top_k, top_p): + """Keeping every entry disables top-k without changing temperature or top-p.""" + sampler = Sampler(torch.device("cuda"), vocab_size=17) + args = sampler.prepare(_batch([SamplingParams(temperature=0.7, top_k=top_k, top_p=top_p)])) + assert args.top_k is None + assert args.temperatures.tolist() == pytest.approx([0.7]) + if top_p == 1: + assert args.top_p is None + else: + assert args.top_p.tolist() == pytest.approx([top_p]) + # One finite logit gives a deterministic oracle without depending on RNG state. + logits = torch.full((1, 17), -float("inf"), device="cuda") + logits[0, 8] = 0 + assert sampler.sample(logits, args).item() == 8 + + +@pytest.mark.parametrize("top_k", [2**31, 2**80]) +def test_large_top_k_preserves_other_rows_in_mixed_batch(top_k): + """An unfiltered row must not change another row's top-k, top-p or greedy choice.""" + sampler = Sampler(torch.device("cuda"), vocab_size=17) + args = sampler.prepare(_batch([ + SamplingParams(temperature=0.7, top_k=top_k), + SamplingParams(temperature=0.7, top_k=3, top_p=0.8), + SamplingParams(temperature=0, top_k=top_k), + SamplingParams(temperature=0.7, top_k=1), + ])) + assert args.top_k.tolist() == [17, 3, 17, 17] + assert args.top_p.tolist() == pytest.approx([1.0, 0.8, 1.0, 1.0]) + assert args.greedy_mask.tolist() == [False, False, True, True] + logits = torch.arange(68, device="cuda", dtype=torch.float32).reshape(4, 17) + logits[2:, 0] = logits[2:, -1] # Greedy ties must retain argmax's first-index rule. + tokens = sampler.sample(logits, args) + assert 0 <= tokens[0].item() < 17 + assert tokens[1].item() in torch.topk(logits[1], 3).indices.tolist() + torch.testing.assert_close(tokens[2:].long(), logits[2:].argmax(-1)) From 0a08db860f3ed5470430fec2dfa039ff1492d8ab Mon Sep 17 00:00:00 2001 From: wenlive <33243841+wenlive@users.noreply.github.com> Date: Wed, 4 Nov 2026 10:33:09 +0800 Subject: [PATCH 2/2] chore(sampling): remove added tests and comments --- python/freetoken/engine/sample.py | 2 -- tests/engine/test_sampling.py | 59 ------------------------------- 2 files changed, 61 deletions(-) delete mode 100644 tests/engine/test_sampling.py diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index bfff814d4b..5bc12ca941 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -79,8 +79,6 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: # Greedy outputs are selected explicitly in sample(); use neutral sampling # parameters for those rows instead of approximating argmax at low temperature. ts = [1.0 if g else max(p.temperature, MIN_T) for p, g in zip(params, is_greedy)] - # Keeping the full vocabulary is unfiltered. Bound Python integers before - # converting them to the device's int32 tensor. top_ks = [ min(p.top_k, self.vocab_size) if not g and p.top_k >= 1 else self.vocab_size for p, g in zip(params, is_greedy) diff --git a/tests/engine/test_sampling.py b/tests/engine/test_sampling.py deleted file mode 100644 index eb738771a7..0000000000 --- a/tests/engine/test_sampling.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Sampling bounds must preserve full-vocabulary and mixed-batch semantics.""" - -import pytest -import torch - -from freetoken.core import Batch, Req, SamplingParams -from freetoken.engine.sample import Sampler - - -pytestmark = pytest.mark.skipif( - not torch.cuda.is_available(), reason="a CUDA or ROCm GPU is required" -) - - -def _batch(params): - return Batch(reqs=[ - Req(input_ids=torch.tensor([0], dtype=torch.int32), table_idx=i, - cached_len=0, output_len=4, uid=i, sampling_params=p, cache_handle=None) - for i, p in enumerate(params) - ], phase="decode") - - -@pytest.mark.parametrize("top_k", [-1, 0, 17, 18, 2**31 - 1, 2**31, 2**80]) -@pytest.mark.parametrize("top_p", [1.0, 0.8]) -def test_full_vocabulary_top_k_disables_filter(top_k, top_p): - """Keeping every entry disables top-k without changing temperature or top-p.""" - sampler = Sampler(torch.device("cuda"), vocab_size=17) - args = sampler.prepare(_batch([SamplingParams(temperature=0.7, top_k=top_k, top_p=top_p)])) - assert args.top_k is None - assert args.temperatures.tolist() == pytest.approx([0.7]) - if top_p == 1: - assert args.top_p is None - else: - assert args.top_p.tolist() == pytest.approx([top_p]) - # One finite logit gives a deterministic oracle without depending on RNG state. - logits = torch.full((1, 17), -float("inf"), device="cuda") - logits[0, 8] = 0 - assert sampler.sample(logits, args).item() == 8 - - -@pytest.mark.parametrize("top_k", [2**31, 2**80]) -def test_large_top_k_preserves_other_rows_in_mixed_batch(top_k): - """An unfiltered row must not change another row's top-k, top-p or greedy choice.""" - sampler = Sampler(torch.device("cuda"), vocab_size=17) - args = sampler.prepare(_batch([ - SamplingParams(temperature=0.7, top_k=top_k), - SamplingParams(temperature=0.7, top_k=3, top_p=0.8), - SamplingParams(temperature=0, top_k=top_k), - SamplingParams(temperature=0.7, top_k=1), - ])) - assert args.top_k.tolist() == [17, 3, 17, 17] - assert args.top_p.tolist() == pytest.approx([1.0, 0.8, 1.0, 1.0]) - assert args.greedy_mask.tolist() == [False, False, True, True] - logits = torch.arange(68, device="cuda", dtype=torch.float32).reshape(4, 17) - logits[2:, 0] = logits[2:, -1] # Greedy ties must retain argmax's first-index rule. - tokens = sampler.sample(logits, args) - assert 0 <= tokens[0].item() < 17 - assert tokens[1].item() in torch.topk(logits[1], 3).indices.tolist() - torch.testing.assert_close(tokens[2:].long(), logits[2:].argmax(-1))