From aae638d877f23f3e1f954cc8c9b7b2703bf85478 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Wed, 30 Sep 2026 14:07:09 +0300 Subject: [PATCH 01/18] test(moe): a short prefill chunk loads its experts through the slot cache, bit-identical to the stream --- tests/moe/test_nvfp4_backends.py | 35 ++++++++++++++++++++++++++ tests/moe/test_offload.py | 43 ++++++++++++++++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 0546fbcd6f..91abeee1a8 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -325,6 +325,41 @@ def test_triton_overlap_prefill_matches_dequant_reference(): cache.release_prefill_layer(layer_id) +@cuda +def test_triton_demand_prefill_matches_the_streamed_layer_bit_for_bit(): + """A short prefill chunk loads only its routed experts, through the slot cache, and runs the + same grouped GEMM over the slots: its output must equal the whole-layer stream's exactly, and + its admissions must not count as decode misses.""" + from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4 + + device = torch.device("cuda") + torch.manual_seed(15) + M = 12 + hidden = torch.randn(M, H, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.rand(M, E, device=device).argsort(dim=1)[:, :TOPK].to(torch.int32).contiguous() + topk_weights = torch.rand(M, TOPK, dtype=torch.float32, device=device) + + streamed_cache, _ = _triton_cache(device) + streamed_cache.materialize_layer(1) + streamed_cache.copy_missing() + streamed = fused_experts_nvfp4( + hidden, *streamed_cache.bank_views(E), topk_weights, topk_ids.clone(), E, "silu", False + ) + + demand_cache, _ = _triton_cache(device, cache_size=2 * S) + demand_cache.collect_stats = True + slots = topk_ids.clone() + demand_cache.ensure_experts(1, slots, counts_as_decode=False) + demand_cache.copy_missing() + demanded = fused_experts_nvfp4( + hidden, *demand_cache.bank_views(), topk_weights, slots, demand_cache.cache_size, "silu", False + ) + + assert not torch.equal(slots, topk_ids), "the admission must have rewritten expert ids to slots" + assert torch.equal(demanded, streamed) + assert int(demand_cache.lru_stats.sum()) == 0 + + @cuda def test_triton_swigluoai_matches_dequant_reference(): """MiniMax-M3's swigluoai routed experts through the Triton prefill grouped GEMM diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index cd5f5dc773..5e3a20c8fb 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -125,6 +125,49 @@ def fake_fused( assert calls["topk_ids"].tolist() == [[2, 1]] +def test_offload_moe_layer_short_prefill_loads_its_experts_through_the_slot_cache(monkeypatch): + layer, cache = _make_layer_and_cache() + cache.prefill_demand_tokens = 1 + topk_weights = torch.tensor([[0.7, 0.3]], dtype=torch.float32) + hidden_states = torch.randn(1, 8) + calls = {} + + monkeypatch.setattr( + "freetoken.layers.moe.fused_topk", + lambda *, hidden_states, gating_output, topk, renormalize: (topk_weights, torch.tensor([[2, 1]], dtype=torch.int32)), + ) + monkeypatch.setattr(cache, "materialize_layer", lambda layer_id: calls.setdefault("materialized", layer_id)) + + def ensure(layer_id, expert_ids, *, counts_as_decode=True): + calls["ensured"] = (layer_id, expert_ids.tolist(), counts_as_decode) + expert_ids.copy_(torch.tensor([[5, 3]], dtype=torch.int32)) + + monkeypatch.setattr(cache, "ensure_experts", ensure) + monkeypatch.setattr(cache, "copy_missing", lambda: calls.setdefault("copied", True)) + + def fake_fused(hidden_states, w1, w2, got_topk_weights, got_topk_ids, activation, apply_router_weight_on_input, + act_alpha=1.0, act_limit=float("inf")): + calls["w1"] = w1 + calls["topk_ids"] = got_topk_ids.clone() + return hidden_states + + monkeypatch.setattr("freetoken.moe.fused.fused_experts_impl", fake_fused) + + out = layer.prefill_forward(hidden_states, torch.randn(1, 4)) + + assert out is hidden_states + assert "materialized" not in calls + assert calls["ensured"] == (0, [[2, 1]], False) + assert calls["copied"] is True + # the grouped GEMM reads the whole slot cache at the slots the admission wrote + assert calls["w1"].shape[0] == cache.cache_size + assert calls["topk_ids"].tolist() == [[5, 3]] + + calls.clear() + layer.prefill_forward(torch.randn(2, 8), torch.randn(2, 4)) + assert calls["materialized"] == 0 and "ensured" not in calls + + def test_offload_moe_layer_prefill_overlap_prefetches_layers_into_two_buffers(monkeypatch): from freetoken.moe.offload_cache import OffloadMoeCache From 689f3d137e23aefbb1eecb961e59f4afd72e2a0c Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Wed, 30 Sep 2026 14:08:19 +0300 Subject: [PATCH 02/18] feat(moe): load a short prefill chunk's experts on demand through the slot cache --- docs/cli.md | 1 + python/freetoken/engine/config.py | 1 + python/freetoken/engine/engine.py | 2 ++ python/freetoken/layers/moe.py | 27 +++++++++++++++++++------ python/freetoken/moe/offload_cache.py | 16 ++++++++++++--- python/freetoken/moe/offload_kernels.py | 4 ++-- python/freetoken/server/args.py | 11 ++++++++++ 7 files changed, 51 insertions(+), 11 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 2166e9e86e..b4c318e424 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -89,6 +89,7 @@ See [models.md](models.md#moe-strategies) for what each strategy does. | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, a fraction, or `auto`). `auto` is for Windows/WSL only, where CUDA pinned memory is capped; every value needs an expert format the CPU executor serves (bf16, fp8_block, nvfp4, mxfp4) | | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU | +| `--moe-prefill-demand-tokens` | 0 | Prefill: a chunk of at most this many tokens loads only the experts it routes to, through the expert cache as decode does (GPU decode only); 0 streams whole layers | | `--moe-prefill-hit-d2d` | off | Prefill: copy cache-hit experts device-side, stream only misses (CUDA >= 13) | | `--disable-moe-prefill-overlap` | overlap on | Disable the two-buffer prefill copy overlap | diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index b6fa8b8949..4b30fe7fda 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -47,6 +47,7 @@ class EngineConfig: moe_cache_auto: bool = False kv_reserve_tokens: int = 8192 # KV floor for --moe-cache-auto; small by design (MoE-priority) moe_cache_policy: str = "lru" + moe_prefill_demand_tokens: int = 0 moe_prefill_overlap: bool = True # Prefill hit/miss split: serve cache-resident experts D2D during prefill # prefetch instead of re-streaming the full layer over PCIe. Needs CUDA >= 12.8 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3988caf602..5e02fb2ed8 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -753,6 +753,7 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache: cache_size=config.moe_cache_size, device=self.device, cache_policy=config.moe_cache_policy, + prefill_demand_tokens=config.moe_prefill_demand_tokens, prefill_overlap=config.moe_prefill_overlap, prefill_hit_d2d=config.moe_prefill_hit_d2d, quant_format=banks.quant_format, @@ -1459,6 +1460,7 @@ def _auto_cpu_layers(config: EngineConfig, num_moe_layers: int, *, reserved: int "moe_cpu_layers": None, "moe_cpu_threads": 0, "moe_hybrid_max_fetch": -1, + "moe_prefill_demand_tokens": 0, "moe_prefill_overlap": True, "moe_prefill_hit_d2d": False, "expert_load": "auto", diff --git a/python/freetoken/layers/moe.py b/python/freetoken/layers/moe.py index 4918a032fe..710a44c79a 100644 --- a/python/freetoken/layers/moe.py +++ b/python/freetoken/layers/moe.py @@ -253,8 +253,9 @@ def prefill_forward( # ------------------------------------------------------------------ # Data movement -- one decision tree for every quant format (the banks # registry makes the cache machinery bank-count agnostic). Decode loads - # on demand; prefill streams whole layers, double-buffered when overlap - # is enabled. The kernels only ever see bank views plus row indices; + # on demand, and so does a short enough prefill chunk; a longer one + # streams whole layers, double-buffered when overlap is enabled. The + # kernels only ever see bank views plus row indices; # which kernel runs is decided afterwards, in ``_expert_gemm``. # ------------------------------------------------------------------ @@ -355,12 +356,26 @@ def _prefill_routed( topk_weights: torch.Tensor, topk_ids: torch.Tensor, ) -> torch.Tensor: - """Prefill movement: stream whole layers -- double-buffered behind the - previous layer's GEMMs when ``prefill_overlap`` is on, else a synchronous - ``materialize_layer``. In both, position == expert id, so the routing ids - pass through unmapped.""" + """Prefill movement. A chunk within ``prefill_demand_tokens`` loads its routed experts + into the slot cache as decode does and runs the same grouped GEMM over the slots. A + longer one streams whole layers -- double-buffered behind the previous layer's GEMMs + when ``prefill_overlap`` is on, else a synchronous ``materialize_layer`` -- where + position == expert id, so the routing ids pass through unmapped.""" cache = self.offload_cache assert cache is not None + if cache.loads_prefill_on_demand(hidden_states.shape[0]): + cache.ensure_experts(self.layer_id, topk_ids, counts_as_decode=False) + cache.copy_missing() + return self._expert_gemm( + cache, + hidden_states, + topk_weights, + topk_ids, + views=cache.bank_views(), + n=cache.cache_size, + alphas=cache.alphas_for_slots(self.layer_id), + is_prefill=True, + ) if cache.prefill_overlap: views = self._wait_prefill_overlap(cache) out = self._expert_gemm( diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py index 0969f73829..2afbb1a79f 100644 --- a/python/freetoken/moe/offload_cache.py +++ b/python/freetoken/moe/offload_cache.py @@ -108,6 +108,9 @@ class OffloadMoeCache: cache_size: int device: torch.device cache_policy: str = "lru" + # A prefill chunk of at most this many tokens admits the experts it routes to into the slot + # cache, as decode does, instead of streaming every expert of each layer; 0 streams them all. + prefill_demand_tokens: int = 0 prefill_overlap: bool = False # Prefill hit/miss split: experts already resident in the slot cache (slots # >= 2 * num_experts) are gathered device-side into the double buffer instead @@ -150,6 +153,10 @@ def __post_init__(self) -> None: policy_ids = {"lru": 0} assert self.cache_policy in policy_ids assert self.decode_target in ("gpu", "cpu", "hybrid"), self.decode_target + if self.prefill_demand_tokens and self.decode_target != "gpu": + raise ValueError( + f"--moe-prefill-demand-tokens needs GPU decode; this cache decodes on {self.decode_target!r}" + ) if self.layout is None: assert self.quant_format in _BANK_SCHEMAS, f"unknown quant_format {self.quant_format!r}" # Attached by the engine for decode_target == "cpu" (CpuMoeExecutor); None @@ -843,17 +850,20 @@ def release_prefill_layer(self, layer_id: int) -> None: self._prefill_buffer_has_release_event[buffer_id] = True self._prefill_buffer_released[buffer_id] = True - def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor) -> None: + def loads_prefill_on_demand(self, num_tokens: int) -> bool: + return num_tokens <= self.prefill_demand_tokens + + def ensure_experts(self, layer_id: int, expert_ids: torch.Tensor, *, counts_as_decode: bool = True) -> None: from freetoken.moe.offload_kernels import ensure_experts - if self.collect_decode_freq: + if self.collect_decode_freq and counts_as_decode: # ``expert_ids`` still holds raw expert ids here (the kernel rewrites them to # slot ids in place), so snapshot the routing histogram before that happens. ids = expert_ids.reshape(-1).long() self.decode_freq[layer_id].scatter_add_(0, ids, torch.ones_like(ids)) self._pending_src_layer = layer_id self._pending_whole_layer = False - ensure_experts(self, layer_id, expert_ids) + ensure_experts(self, layer_id, expert_ids, record_stats=counts_as_decode) def ensure_experts_hybrid(self, layer_id: int, expert_ids: torch.Tensor) -> None: """Capped-fetch LRU for the hybrid backend. diff --git a/python/freetoken/moe/offload_kernels.py b/python/freetoken/moe/offload_kernels.py index cf513f52d1..5df2ffec68 100644 --- a/python/freetoken/moe/offload_kernels.py +++ b/python/freetoken/moe/offload_kernels.py @@ -16,7 +16,7 @@ ) -def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: +def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor, *, record_stats: bool = True) -> None: """Make this layer's routed experts resident; rewrite ``expert_ids`` to slot ids. Delegates to flashlib's slot cache. ``id_base`` maps this layer's expert ids into the @@ -35,7 +35,7 @@ def ensure_experts(cache, layer_id: int, expert_ids: torch.Tensor) -> None: cache.src_indices, cache.evict_slots, cache.num_indices, - stats=cache.lru_stats[layer_id] if cache.collect_stats else None, + stats=cache.lru_stats[layer_id] if cache.collect_stats and record_stats else None, id_base=layer_id * cache.num_experts, ) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 7c0ed9b047..519447081b 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -781,6 +781,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--moe-prefill-demand-tokens", + type=int, + default=ServerArgs.moe_prefill_demand_tokens, + help=( + "A prefill chunk of at most this many tokens loads only the experts it routes to, " + "through the expert cache as decode does, instead of every expert of each layer. " + "0 streams whole layers for every chunk." + ), + ) + parser.add_argument( "--enable-special-token-ckpt", action="store_true", From 96c731085aca69b74f5fae066bcdc806df540948 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Wed, 30 Sep 2026 17:38:27 +0300 Subject: [PATCH 03/18] test(moe): a dense model drops the demand-prefill threshold --- tests/moe/test_offload.py | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 5e3a20c8fb..665cd7876e 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -498,6 +498,37 @@ def test_adjust_config_converts_moe_cache_rate_to_cache_size(monkeypatch): assert is_offload_moe_strategy(config.moe_strategy) +def test_adjust_config_drops_the_demand_prefill_threshold_for_a_dense_model(): + from types import SimpleNamespace + + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + from freetoken.engine.engine import _adjust_config + + config = EngineConfig( + model_path="/tmp/freetoken-test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.float16, + attention_backend="triton", + moe_prefill_demand_tokens=512, + ) + object.__setattr__( + config, + "model_config", + SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + is_moe=False, + num_layers=10, + expert_quant="none", + ), + ) + + _adjust_config(config) + + assert config.moe_prefill_demand_tokens == 0 + + def test_graph_capture_reuses_warm_offload_cache_before_capture(monkeypatch): import freetoken.core as core from freetoken.core import Context, Req, get_global_ctx From a8f24b1764ccc54a7a1509be7f9a923ca23f1a24 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 06:27:12 +0300 Subject: [PATCH 04/18] test(moe): the demand-prefill threshold is refused off GPU decode OffloadMoeCache raises when --moe-prefill-demand-tokens is set for cpu or hybrid decode, whose experts are not read from the slot cache a demand-loaded chunk fills; the test drives both targets. --- tests/moe/test_offload.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 665cd7876e..2cbd29eb29 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -283,6 +283,22 @@ def test_offload_moe_cache_marlin_rejects_slot_count_beyond_kernel_limit(): ) +@pytest.mark.parametrize("decode_target", ["cpu", "hybrid"]) +def test_offload_moe_cache_rejects_demand_prefill_without_gpu_decode(decode_target): + # a demand-loaded chunk admits its experts into GPU decode's slot cache, which cpu and hybrid decode do not read + from freetoken.moe.offload_cache import OffloadMoeCache + + with pytest.raises(ValueError, match="needs GPU decode"): + OffloadMoeCache( + num_layers=1, + num_experts=4, + cache_size=6, + device=torch.device("cpu"), + decode_target=decode_target, + prefill_demand_tokens=64, + ) + + def test_prefill_overlap_prefetch_invalidates_borrowed_unified_cache_slots(): from freetoken.moe.offload_cache import OffloadMoeCache From b880e7cfaf48904478a63bca7e68fc363a829aab Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 05:16:35 -0400 Subject: [PATCH 05/18] fix(kernel): size triton's autotune flush buffer to the device l2 A kernel autotunes on its first launch of each key, which can be the first request after the cache planner has handed out nearly all VRAM; triton's flat 256 MB flush buffer was an OOM there on an 11 GB card. Four times the L2 still evicts it between timed runs. Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit ffbf001acec04f4f882d6cccbc1d0d4cbbec633f) (cherry picked from commit e097977e525302a6f773ace34ca553666e593876) --- python/freetoken/engine/engine.py | 3 +++ .../freetoken/kernel/triton/autotune_cache.py | 26 ++++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index dbdbc36674..a8e0aaff8a 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -494,6 +494,9 @@ def __init__(self, config: EngineConfig): logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}") # ======================= Model initialization ======================== + from freetoken.kernel.triton.autotune_cache import bound_autotune_flush_buffer + + bound_autotune_flush_buffer(self.device.index) set_rope_device(self.device) with torch.device("meta"), torch_dtype(config.dtype): self.model = create_model(config.model_config) diff --git a/python/freetoken/kernel/triton/autotune_cache.py b/python/freetoken/kernel/triton/autotune_cache.py index ea95768368..2e57fadde9 100644 --- a/python/freetoken/kernel/triton/autotune_cache.py +++ b/python/freetoken/kernel/triton/autotune_cache.py @@ -34,4 +34,28 @@ {"cache_results": _CACHE_RESULTS} if _SUPPORTS_AUTOTUNE_CACHE else {} ) -__all__ = ["autotune_cache_kwargs"] + + +def bound_autotune_flush_buffer(device_index: int) -> None: + """Size triton's L2-flush buffer for autotune benchmarks to the device instead of a flat 256 MB. + + A kernel autotunes on its first launch of each key, which can be the first request, after the + cache planner has handed out nearly all VRAM; a 256 MB transient there is an OOM on a small + card. Four times the L2 still evicts everything between timed runs. + """ + import torch + + try: + from triton.backends.nvidia.driver import CudaDriver + except ImportError: + return + l2 = int(getattr(torch.cuda.get_device_properties(device_index), "L2_cache_size", 0) or 0) + nbytes = min(256 << 20, max(8 << 20, 4 * l2)) + + def get_empty_cache_for_benchmark(self): + return torch.empty(nbytes // 4, dtype=torch.int, device="cuda") + + CudaDriver.get_empty_cache_for_benchmark = get_empty_cache_for_benchmark + + +__all__ = ["autotune_cache_kwargs", "bound_autotune_flush_buffer"] From 47afdaf1e3036c08f6099b4e98aa332e3d12f928 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 10:02:35 -0400 Subject: [PATCH 06/18] test(utils): drive the async coalesced-pull test with asyncio.run Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 32007ecb48aee2dc619fc9ceb89ad9b5dd70ef14) --- tests/utils/test_mp_coalesced.py | 31 +++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/tests/utils/test_mp_coalesced.py b/tests/utils/test_mp_coalesced.py index e599db7265..c64ee1c5c3 100644 --- a/tests/utils/test_mp_coalesced.py +++ b/tests/utils/test_mp_coalesced.py @@ -92,7 +92,7 @@ def test_sub_queue_coalesced(self): sub.stop() -async def test_async_pull_coalesced(): +def test_async_pull_coalesced(): import sys from freetoken.utils.mp import ZmqAsyncPullQueue as _AsyncPull @@ -101,16 +101,19 @@ async def test_async_pull_coalesced(): # zmq.asyncio needs a Selector loop on Windows; the project runs on Linux CI. if sys.platform == "win32": pytest.skip("zmq.asyncio requires SelectorEventLoop on Windows (prod runs on Linux)") - push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) - pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) - try: - await push.put({"n": 1}) - await push.put({"n": 2}) - await asyncio.sleep(0.1) - - # Even if both messages arrive in a single frame, two gets must succeed. - assert await pull.get() == {"n": 1} - assert await pull.get() == {"n": 2} - finally: - push.stop() - pull.stop() \ No newline at end of file + async def scenario(): + push = ZmqAsyncPushQueue(_addr(5604), create=True, encoder=lambda o: o) + pull = ZmqAsyncPullQueue(_addr(5604), create=False, decoder=_identity_decoder) + try: + await push.put({"n": 1}) + await push.put({"n": 2}) + await asyncio.sleep(0.1) + + # Even if both messages arrive in a single frame, two gets must succeed. + assert await pull.get() == {"n": 1} + assert await pull.get() == {"n": 2} + finally: + push.stop() + pull.stop() + + asyncio.run(scenario()) From c0190aa51a8f8f2f3e8206fe3982bb12913ac70b Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 20:35:40 -0400 Subject: [PATCH 07/18] fix(engine): keep multimodal placeholder ids out of the repetition-penalty prompt mask Fork #47's sampler plan scattered the raw prompt ids into a [rows, vocab+1] mask; image tokens carry pseudo-ids >= 1,000,000, so repetition_penalty with an image indexed out of bounds -- a device-side assert that takes the server down. They now land in the scratch column. Also keeps presence/frequency out of the LogitsPlan: with #527 device counts and #393 plan both applying them, a request with a penalty plus any plan processor was penalized twice. (cherry picked from commit 66053ded) Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit b9883902735cadb2232e3ae4d9dd18e02481a189) --- python/freetoken/engine/sample.py | 23 ++++++----------------- tests/engine/test_logits_processors.py | 23 ++++++++++++++++++++--- tests/engine/test_sample_logprobs.py | 4 ++-- 3 files changed, 28 insertions(+), 22 deletions(-) diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 89d2a6fb88..8e05f80694 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -172,35 +172,24 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: ) pad = self.vocab_size - want_hist = any( - p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 - for _, p in picked - ) - if want_hist: + # presence/frequency are not plan work: they ride the device counts (Sampler.prepare's penalties) + if any(p.repetition_penalty != 1.0 for _, p in picked): hist: list[list[int]] = [] prompts: list[list[int]] = [] - want_prompt = False for req, p in picked: prompt_len = req.max_device_len - req.output_len ids = req.input_ids - penalized = p.presence_penalty != 0.0 or p.frequency_penalty != 0.0 or p.repetition_penalty != 1.0 - hist.append(ids[prompt_len:].tolist() if penalized else []) if p.repetition_penalty != 1.0: - want_prompt = True + hist.append(ids[prompt_len:].tolist()) # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); # send them to the scratch column, not past the scatter target prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) else: + hist.append([]) prompts.append([]) plan.hist_ids = _padded_rows(hist, pad, self.device) - if want_prompt: - plan.prompt_ids = _padded_rows(prompts, pad, self.device) - if any(p.presence_penalty != 0.0 for _, p in picked): - plan.presence = make_device_tensor([[p.presence_penalty] for _, p in picked], torch.float32, self.device) - if any(p.frequency_penalty != 0.0 for _, p in picked): - plan.frequency = make_device_tensor([[p.frequency_penalty] for _, p in picked], torch.float32, self.device) - if any(p.repetition_penalty != 1.0 for _, p in picked): - plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) + plan.prompt_ids = _padded_rows(prompts, pad, self.device) + plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) bias_rows: list[int] = [] bias_ids: list[int] = [] diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index abd2e1f939..237ce9f602 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -172,9 +172,8 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): logits = torch.zeros((4, V)) out = apply_logits_processors(logits, plan, V) - assert ( - out[1, 3].item() == -1.0 - ) # presence on the repeated 3 (repetition on 0 is a no-op) + # presence is the sampler's device counts, not the plan; repetition on 0 is a no-op + assert plan.presence is None and out[1, 3].item() == 0.0 assert out[1, 5].item() == 2.0 assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") assert torch.isfinite(out[3]).all() @@ -182,6 +181,24 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert sampler.sample(logits, args)[1].item() == 5 +def test_presence_and_frequency_ride_the_device_counts_not_the_plan(): + # the plan must not apply them a second time on top of the sampler's own counts + sampler = Sampler(CPU, V) + only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)]) + args = sampler.prepare(only) + assert args.plan is None and [row for row, *_ in args.penalties] == [0] + + req = _req([1], [], 4, presence_penalty=1.0, logit_bias=[[5, 0.25]]) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + assert args.plan is not None and [row for row, *_ in args.penalties] == [0] + req.output_token_counts[3] = 1 + logits = torch.zeros((1, V)) + logits[0, 3] = 1.5 + # once: 3 -> 0.5 beats 5 -> 0.25; twice would drop 3 to -0.5 and pick 5 + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts[3].item() == 2 + + def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): # image tokens carry pseudo-ids >= MM_PAD_SHIFT_VALUE; scattering them into the # [rows, V+1] prompt mask indexed out of bounds (a device-side assert on CUDA) diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index 6bd204604d..fae5b7c25b 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -137,8 +137,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=True, top_logprobs=17, is_greedy=True)), + SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=False, top_logprobs=0, is_greedy=True)), ] ) args = sampler.prepare(batch) From 3af6e25146bf3f2b8ca83b0993b741d212a8a97c Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:43:17 -0400 Subject: [PATCH 08/18] fix(engine): apply the repetition penalty from device state, not a per-step prompt upload (up to 262K ids) and a synchronous pageable copy, which stalled the overlap loop and saw the output one token late. The prompt mask is now built once per request, the output side reuses #527's device counts, and the order is vLLM's (repetition, then presence and frequency). Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 343d7421dfebf72244b9e9d16261dba073b48cbc) (cherry picked from commit 3df811dee9bcd4af287e3cc94abb689d531e58bd) --- python/freetoken/core.py | 11 +- python/freetoken/engine/sample.py | 188 ++++++++++++++----------- tests/engine/test_logits_processors.py | 83 +++++++---- tests/engine/test_sample_logprobs.py | 5 +- 4 files changed, 172 insertions(+), 115 deletions(-) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ab5b81b0af..a7811ed712 100644 --- a/python/freetoken/core.py +++ b/python/freetoken/core.py @@ -30,7 +30,8 @@ class SamplingParams: # min_p: drop tokens whose probability is below min_p x the top probability. min_p: float = 0.0 # OpenAI presence/frequency penalties over the generated tokens; repetition_penalty is the - # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). + # HF/vLLM multiplicative penalty over prompt + generated tokens (1.0 = off). The sampler + # keeps their state on the device (Req.output_token_counts, Req.prompt_token_mask). presence_penalty: float = 0.0 frequency_penalty: float = 0.0 repetition_penalty: float = 1.0 @@ -60,11 +61,10 @@ def is_greedy(self) -> bool: @property def needs_logits_processing(self) -> bool: + # the penalties are not LogitsPlan work: the sampler applies them from its own + # device-side state (Sampler.prepare's penalties). return bool( self.min_p > 0.0 - or self.presence_penalty != 0.0 - or self.frequency_penalty != 0.0 - or self.repetition_penalty != 1.0 or self.logit_bias or self.min_tokens > 0 ) @@ -106,7 +106,10 @@ class Req: # handler must not free resources under an in-flight forward; it sets this flag and # _process_last_data frees the request when the batch drains (after copy_done.synchronize). aborted: bool = False + output_token_counts: torch.Tensor | None = field(default=None, init=False, repr=False) stop_decode_status: DecodeStatus | None = None + # [vocab] bool on the device: the prompt's ids, for repetition_penalty. Built once. + prompt_token_mask: torch.Tensor | None = field(default=None, init=False, repr=False) def __post_init__(self) -> None: assert self.input_ids.is_cpu diff --git a/python/freetoken/engine/sample.py b/python/freetoken/engine/sample.py index 8e05f80694..9d12b5f571 100644 --- a/python/freetoken/engine/sample.py +++ b/python/freetoken/engine/sample.py @@ -1,7 +1,7 @@ from __future__ import annotations -from dataclasses import dataclass -from typing import TYPE_CHECKING, List +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, List, NamedTuple import torch from freetoken.utils import is_sm90_supported, nvtx_annotate @@ -13,41 +13,49 @@ @dataclass class LogitsPlan: """Inputs of the logits processors for the rows of one batch that asked for any - (presence / frequency / repetition penalties, logit_bias, min_tokens, min_p). - - Built on the host in ``Sampler.prepare`` from the requests' SamplingParams and - token histories, applied on the device by ``apply_logits_processors`` before the - sampling kernel. ``None`` when no request in the batch uses a processor, so the - default path stays exactly as it was. Row indices inside the plan are LOCAL (0..r-1, - in the order of ``rows``); ``rows`` maps them back to the batch. + (logit_bias, min_tokens, min_p). The penalties are not here: the sampler applies + them from per-request device state (``RowPenalty``). + + Built on the host in ``Sampler.prepare`` from the requests' SamplingParams, applied + on the device by ``apply_logits_processors`` before the sampling kernel. ``None`` + when no request in the batch uses a processor, so the default path stays exactly + as it was. Row indices inside the plan are LOCAL (0..r-1, in the order of ``rows``); + ``rows`` maps them back to the batch. """ rows: torch.Tensor # [r] int64 batch rows temps: torch.Tensor # [r, 1] float32 temperatures (min_p works on probabilities) - presence: torch.Tensor | None = None # [r, 1] float32 - frequency: torch.Tensor | None = None # [r, 1] float32 - repetition: torch.Tensor | None = None # [r, 1] float32 (1.0 = off) - # Generated ids per row, right-padded with vocab_size (a scratch column that is - # dropped), for the presence / frequency counts and the repetition mask. - hist_ids: torch.Tensor | None = None # [r, L] int64 - # Prompt ids per row, same padding: repetition_penalty also covers the prompt. - prompt_ids: torch.Tensor | None = None # [r, P] int64 bias_rows: torch.Tensor | None = None # [b] int64 local rows bias_ids: torch.Tensor | None = None # [b] int64 bias_vals: torch.Tensor | None = None # [b] float32 - # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded) they must - # not sample yet. + # Rows still under min_tokens and the ids (EOS + stop_token_ids, padded with + # vocab_size, a scratch column that is dropped) they must not sample yet. min_rows: torch.Tensor | None = None # [m] int64 local rows min_ids: torch.Tensor | None = None # [m, S] int64 min_p: torch.Tensor | None = None # [r, 1] float32 (0 = off for that row) +class RowPenalty(NamedTuple): + """One batch row's presence / frequency / repetition penalty. ``counts`` are the + request's generated-token counts ([V] int32), updated on the device after each draw so + overlap scheduling never sees them late; ``prompt_mask`` ([V] bool) marks the prompt's + ids for the repetition penalty and is None when that is off.""" + + row: int + counts: torch.Tensor + presence: float + frequency: float + repetition: float = 1.0 + prompt_mask: torch.Tensor | None = None + + @dataclass class BatchSamplingArgs: temperatures: torch.Tensor | None top_k: torch.Tensor | None = None top_p: torch.Tensor | None = None greedy_mask: torch.Tensor | None = None + penalties: list[RowPenalty] = field(default_factory=list) plan: LogitsPlan | None = None logprob_rows: torch.Tensor | None = None max_top_logprobs: int = 0 @@ -60,47 +68,39 @@ def make_device_tensor(data: List, dtype: torch.dtype, device: torch.device) -> def _padded_rows(rows: list[list[int]], pad: int, device: torch.device) -> torch.Tensor: - """[len(rows), max_len] int64 with ``pad`` on the right. A synchronous copy: the - histories can be long (prompt + output) and a pinned buffer must not be freed under - an in-flight non_blocking copy.""" + """[len(rows), max_len] int64 with ``pad`` on the right.""" width = max(1, max((len(r) for r in rows), default=1)) - out = torch.full((len(rows), width), pad, dtype=torch.int64) - for i, r in enumerate(rows): - if r: - out[i, : len(r)] = torch.as_tensor(r, dtype=torch.int64) - return out.to(device) + return make_device_tensor([r + [pad] * (width - len(r)) for r in rows], torch.int64, device) + + +def apply_penalties(logits: torch.Tensor, penalties: list[RowPenalty]) -> torch.Tensor: + """A float32 copy of ``logits`` with each row's penalties applied in vLLM's order: + repetition (over prompt + generated ids), then frequency and presence (generated).""" + logits = logits.float().clone() + for row, counts, presence, frequency, repetition, prompt_mask in penalties: + if repetition != 1.0: + x = logits[row] + seen = counts > 0 + if prompt_mask is not None: + seen = seen | prompt_mask + # HF / vLLM semantics: divide positive logits, multiply negative ones. + penalized = torch.where(x > 0, x / repetition, x * repetition) + logits[row] = torch.where(seen, penalized, x) + if presence or frequency: + logits[row] -= frequency * counts + presence * (counts > 0) + return logits def apply_logits_processors( logits: torch.Tensor, plan: LogitsPlan, vocab_size: int ) -> torch.Tensor: - """A float32 copy of ``logits`` with the plan applied to its rows, in vLLM's order: - repetition -> presence -> frequency penalties, logit_bias, the min_tokens mask, then - min_p. Rows outside the plan are copied unchanged.""" + """A float32 copy of ``logits`` with the plan applied to its rows: logit_bias, the + min_tokens mask, then min_p. Rows outside the plan are copied unchanged.""" out = logits.to(torch.float32, copy=True) sub = out[plan.rows] # [r, V] (advanced indexing copies) - r = sub.shape[0] device = sub.device neg_inf = float("-inf") - if plan.hist_ids is not None: - counts = torch.zeros((r, vocab_size + 1), dtype=torch.float32, device=device) - counts.scatter_add_(1, plan.hist_ids, torch.ones_like(plan.hist_ids, dtype=torch.float32)) - counts = counts[:, :vocab_size] - if plan.repetition is not None: - seen = counts > 0 - if plan.prompt_ids is not None: - in_prompt = torch.zeros((r, vocab_size + 1), dtype=torch.bool, device=device) - in_prompt.scatter_(1, plan.prompt_ids, torch.ones_like(plan.prompt_ids, dtype=torch.bool)) - seen = seen | in_prompt[:, :vocab_size] - # HF / vLLM semantics: divide positive logits, multiply negative ones. - penalized = torch.where(sub > 0, sub / plan.repetition, sub * plan.repetition) - sub = torch.where(seen, penalized, sub) - if plan.presence is not None: - sub = sub - plan.presence * (counts > 0).to(sub.dtype) - if plan.frequency is not None: - sub = sub - plan.frequency * counts - if plan.bias_ids is not None: sub.index_put_((plan.bias_rows, plan.bias_ids), plan.bias_vals, accumulate=True) @@ -155,11 +155,7 @@ class Sampler: vocab_size: int def _plan(self, batch: Batch, params) -> LogitsPlan | None: - """The LogitsPlan for this batch, or None when no request asks for a processor. - Token histories come from the host-side ``req.input_ids``; under overlap - scheduling the previous step's token may not be appended yet, so a penalty sees - the history one token late. That is the accepted cost of keeping the counts off - the device.""" + """The LogitsPlan for this batch, or None when no request asks for a processor.""" MIN_T = 1e-6 rows = [i for i, p in enumerate(params) if p.needs_logits_processing] if not rows: @@ -172,25 +168,6 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: ) pad = self.vocab_size - # presence/frequency are not plan work: they ride the device counts (Sampler.prepare's penalties) - if any(p.repetition_penalty != 1.0 for _, p in picked): - hist: list[list[int]] = [] - prompts: list[list[int]] = [] - for req, p in picked: - prompt_len = req.max_device_len - req.output_len - ids = req.input_ids - if p.repetition_penalty != 1.0: - hist.append(ids[prompt_len:].tolist()) - # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); - # send them to the scratch column, not past the scatter target - prompts.append(ids[:prompt_len].clamp(max=pad).tolist()) - else: - hist.append([]) - prompts.append([]) - plan.hist_ids = _padded_rows(hist, pad, self.device) - plan.prompt_ids = _padded_rows(prompts, pad, self.device) - plan.repetition = make_device_tensor([[p.repetition_penalty] for _, p in picked], torch.float32, self.device) - bias_rows: list[int] = [] bias_ids: list[int] = [] bias_vals: list[float] = [] @@ -223,9 +200,48 @@ def _plan(self, batch: Batch, params) -> LogitsPlan | None: plan.min_p = make_device_tensor([[p.min_p] for _, p in picked], torch.float32, self.device) return plan + def _prompt_mask(self, req) -> torch.Tensor: + """[V] bool marking the request's prompt ids, uploaded once per request.""" + prompt_len = req.max_device_len - req.output_len + # multimodal placeholders are pseudo-ids >= vocab_size (MM_PAD_SHIFT_VALUE); + # send them to the scratch column, not past the scatter target + ids = req.input_ids[:prompt_len].to(torch.int64).clamp(max=self.vocab_size) + if self.device.type == "cuda": + ids = ids.pin_memory().to(self.device, non_blocking=True) + mask = torch.zeros(self.vocab_size + 1, dtype=torch.bool, device=self.device) + mask[ids] = True + return mask[: self.vocab_size] + def prepare(self, batch: Batch) -> BatchSamplingArgs: params = [r.sampling_params for r in batch.reqs] is_greedy = [p.is_greedy for p in params] + penalties = [] + for row, req in enumerate(batch.reqs): + p = req.sampling_params + repetition = p.repetition_penalty + if not (p.presence_penalty or p.frequency_penalty or repetition != 1.0): + continue + if not req.can_decode: + continue + if req.output_token_counts is None: + req.output_token_counts = torch.zeros( + self.vocab_size, dtype=torch.int32, device=self.device + ) + prompt_mask = None + if repetition != 1.0: + if req.prompt_token_mask is None: + req.prompt_token_mask = self._prompt_mask(req) + prompt_mask = req.prompt_token_mask + penalties.append( + RowPenalty( + row, + req.output_token_counts, + p.presence_penalty, + p.frequency_penalty, + repetition, + prompt_mask, + ) + ) plan = self._plan(batch, params) want_logprobs = [p.logprobs for p in params] logprob_rows = ( @@ -239,6 +255,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: if all(is_greedy): return BatchSamplingArgs( temperatures=None, + penalties=penalties, plan=plan, logprob_rows=logprob_rows, max_top_logprobs=max_top_logprobs, @@ -270,6 +287,7 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: top_k=top_k, top_p=top_p, greedy_mask=greedy_mask, + penalties=penalties, plan=plan, logprob_rows=logprob_rows, max_top_logprobs=max_top_logprobs, @@ -278,16 +296,26 @@ def prepare(self, batch: Batch) -> BatchSamplingArgs: @nvtx_annotate("Sampler") def sample(self, logits: torch.Tensor, args: BatchSamplingArgs) -> torch.Tensor: with torch.cuda.nvtx.range("Sampler"): + if args.penalties: + logits = apply_penalties(logits, args.penalties) + # after the penalties: min_p must see the final distribution if args.plan is not None: logits = apply_logits_processors(logits, args.plan, self.vocab_size) if args.temperatures is None: # greedy sampling - return torch.argmax(logits, dim=-1) - tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) - if args.greedy_mask is not None: - # Mixed batches still run probability sampling for all rows, but - # greedy rows must follow argmax's deterministic tie-breaking. - greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) - tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + tokens = torch.argmax(logits, dim=-1) + else: + tokens = sample_impl(logits.float(), args.temperatures, args.top_k, args.top_p) + if args.greedy_mask is not None: + # Mixed batches still run probability sampling for all rows, but + # greedy rows must follow argmax's deterministic tie-breaking. + greedy_tokens = torch.argmax(logits, dim=-1).to(tokens.dtype) + tokens = torch.where(args.greedy_mask, greedy_tokens, tokens) + # Update on the sampling stream: overlapped scheduling can prepare the next + # batch before the previous token reaches Req.input_ids on the CPU. + for row, counts, *_ in args.penalties: + counts.scatter_add_( + 0, tokens[row : row + 1].long(), counts.new_ones(1) + ) return tokens def compute_logprobs( diff --git a/tests/engine/test_logits_processors.py b/tests/engine/test_logits_processors.py index 237ce9f602..1829f35f32 100644 --- a/tests/engine/test_logits_processors.py +++ b/tests/engine/test_logits_processors.py @@ -8,7 +8,13 @@ import pytest import torch from freetoken.core import SamplingParams -from freetoken.engine.sample import LogitsPlan, Sampler, apply_logits_processors +from freetoken.engine.sample import ( + LogitsPlan, + RowPenalty, + Sampler, + apply_logits_processors, + apply_penalties, +) V = 8 CPU = torch.device("cpu") @@ -32,33 +38,17 @@ def _ids(rows: list[list[int]]) -> torch.Tensor: return out -def test_presence_and_frequency_penalties_count_generated_tokens(): - logits = torch.zeros((2, V)) - plan = _plan( - [1], - hist_ids=_ids([[3, 3, 5]]), - presence=torch.tensor([[0.5]]), - frequency=torch.tensor([[0.25]]), - ) - out = apply_logits_processors(logits, plan, V) - assert torch.equal(out[0], logits[0]) # row outside the plan untouched - assert out[1, 3].item() == -(0.5 + 2 * 0.25) - assert out[1, 5].item() == -(0.5 + 0.25) - assert out[1, 0].item() == 0.0 - - def test_repetition_penalty_covers_prompt_and_output_with_hf_semantics(): logits = torch.tensor([[2.0, -2.0, 2.0, -2.0, 1.0, 1.0, 1.0, 1.0]]) - plan = _plan( - [0], - hist_ids=_ids([[0, 1]]), # generated: 0 (positive), 1 (negative) - prompt_ids=_ids([[2, 3]]), # prompt: 2 (positive), 3 (negative) - repetition=torch.tensor([[2.0]]), - ) - out = apply_logits_processors(logits, plan, V) + counts = torch.zeros(V, dtype=torch.int32) + counts[[0, 1]] = 1 # generated: 0 (positive), 1 (negative) + prompt = torch.zeros(V, dtype=torch.bool) + prompt[[2, 3]] = True # prompt: 2 (positive), 3 (negative) + out = apply_penalties(logits, [RowPenalty(0, counts, 0.0, 0.0, 2.0, prompt)]) assert out[0, 0].item() == 1.0 and out[0, 2].item() == 1.0 # divided assert out[0, 1].item() == -4.0 and out[0, 3].item() == -4.0 # multiplied assert out[0, 4].item() == 1.0 # unseen + assert logits[0, 0].item() == 2.0 # the caller's logits stay untouched def test_logit_bias_adds_per_row(): @@ -134,6 +124,9 @@ def _req( output_len=max_tokens, max_device_len=len(prompt) + max_tokens, sampling_params=SamplingParams(**sp), + can_decode=True, + output_token_counts=None, + prompt_token_mask=None, ) @@ -164,16 +157,15 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): plan = args.plan assert plan is not None assert plan.rows.tolist() == [1, 2, 3] - assert plan.hist_ids[0].tolist()[:2] == [3, 3] # row 1's generated tokens - assert plan.prompt_ids[0].tolist()[:2] == [1, 2] + assert [pen.row for pen in args.penalties] == [1] + assert batch.reqs[1].prompt_token_mask.nonzero().flatten().tolist() == [1, 2] assert plan.bias_rows.tolist() == [0] and plan.bias_ids.tolist() == [5] assert plan.min_rows.tolist() == [1] # local row of the third request only assert plan.min_ids[0].tolist() == [0, 7] logits = torch.zeros((4, V)) out = apply_logits_processors(logits, plan, V) - # presence is the sampler's device counts, not the plan; repetition on 0 is a no-op - assert plan.presence is None and out[1, 3].item() == 0.0 + assert out[1, 3].item() == 0.0 # the penalties are not plan work assert out[1, 5].item() == 2.0 assert out[2, 0].item() == float("-inf") and out[2, 7].item() == float("-inf") assert torch.isfinite(out[3]).all() @@ -181,7 +173,7 @@ def test_prepare_builds_no_plan_without_processors_and_a_plan_with_them(): assert sampler.sample(logits, args)[1].item() == 5 -def test_presence_and_frequency_ride_the_device_counts_not_the_plan(): +def test_penalties_ride_the_device_state_not_the_plan(): # the plan must not apply them a second time on top of the sampler's own counts sampler = Sampler(CPU, V) only = SimpleNamespace(reqs=[_req([1], [], 4, presence_penalty=1.0, frequency_penalty=0.5)]) @@ -205,5 +197,38 @@ def test_repetition_penalty_ignores_multimodal_placeholder_ids_in_the_prompt(): sampler = Sampler(CPU, V) batch = SimpleNamespace(reqs=[_req([1, 1_000_000, 2], [3], 4, repetition_penalty=2.0)]) args = sampler.prepare(batch) - out = apply_logits_processors(torch.ones(1, V), args.plan, V) + out = apply_penalties(torch.ones(1, V), args.penalties) assert out[0, 1] == 0.5 and out[0, 2] == 0.5 and out[0, 4] == 1.0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_min_p_with_temperature_draws_only_the_kept_tokens(): + # min_p leaves -inf over nearly the whole row; the sampling softmax must not turn that + # into NaN (which the Pascal fallback then samples as a uniform row) + V = 32768 + sampler = Sampler(torch.device("cuda"), V) + req = _req([1, 2], [], 4, temperature=0.8, min_p=0.1) + args = sampler.prepare(SimpleNamespace(reqs=[req])) + logits = torch.full((1, V), -10.0, device="cuda") + kept = [100, 20000, 31000] + logits[0, kept] = torch.tensor([5.0, 4.5, 4.0], device="cuda") + for _ in range(32): + assert sampler.sample(logits, args).item() in kept + + +def test_repetition_penalty_reads_the_prompt_once_and_counts_output_on_the_device(): + # the prompt (up to the whole context) is uploaded once per request, not every step, + # and the generated side comes from the counts sample() keeps, not the host history + sampler = Sampler(CPU, V) + req = _req([1, 2], [], 4, repetition_penalty=2.0) + batch = SimpleNamespace(reqs=[req]) + args = sampler.prepare(batch) + mask = req.prompt_token_mask + logits = torch.tensor([[0.0, 1.0, 1.0, 4.0, 0.5, 0.0, 0.0, 0.0]]) + assert sampler.sample(logits, args).item() == 3 + assert req.output_token_counts.tolist() == [0, 0, 0, 1, 0, 0, 0, 0] + req.input_ids = torch.tensor([9, 9, 3], dtype=torch.int32) # a stale host view is not read + args = sampler.prepare(batch) + assert req.prompt_token_mask is mask + out = apply_penalties(logits, args.penalties) + assert out[0].tolist() == [0.0, 0.5, 0.5, 2.0, 0.5, 0.0, 0.0, 0.0] diff --git a/tests/engine/test_sample_logprobs.py b/tests/engine/test_sample_logprobs.py index fae5b7c25b..88fffcc2ed 100644 --- a/tests/engine/test_sample_logprobs.py +++ b/tests/engine/test_sample_logprobs.py @@ -2,6 +2,7 @@ import torch +from freetoken.core import SamplingParams from freetoken.engine.sample import BatchSamplingArgs, Sampler @@ -137,8 +138,8 @@ def test_compute_logprobs_with_prepare_keeps_max_top_logprobs_clamped(monkeypatc batch = SimpleNamespace( reqs=[ - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=True, top_logprobs=17, is_greedy=True)), - SimpleNamespace(sampling_params=SimpleNamespace(needs_logits_processing=False, presence_penalty=0.0, frequency_penalty=0.0, logprobs=False, top_logprobs=0, is_greedy=True)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=True, top_logprobs=17)), + SimpleNamespace(sampling_params=SamplingParams(logprobs=False, top_logprobs=0)), ] ) args = sampler.prepare(batch) From 1e32ddbc8a958427f647b7f200c99b5fdbdb7bad Mon Sep 17 00:00:00 2001 From: Siarhei Berdachuk Date: Mon, 5 Oct 2026 19:11:56 +0300 Subject: [PATCH 09/18] fix(server): clarify unsupported response_format error text After soft-accepting json_object/json_schema, the old "no constrained decoding" message was misleading for remaining unknown types. (cherry picked from commit 51e17e6c3380af25d34a301859a8991224f30c2f) (cherry picked from commit e2df92f2ade5cd0ed4f5f1bd421eef96ded793cb) --- python/freetoken/server/openai_api.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 7f94efc250..79dff72963 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -305,8 +305,10 @@ async def handle_chat_completion( if logprobs_conflict is not None: return create_error_response(logprobs_conflict, param="logprobs") if _response_format_unsupported(req.response_format): + rtype = (req.response_format or {}).get("type") return create_error_response( - "response_format json_object/json_schema is not supported (no constrained decoding)", + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)", param="response_format", ) if not 1 <= req.n <= MAX_N: @@ -1072,7 +1074,11 @@ def _completion_unsupported_reason(req: CompletionRequest) -> str | None: if logprobs_error is not None: return logprobs_error if _response_format_unsupported(req.response_format): - return "response_format json_object/json_schema is not supported (no constrained decoding)" + rtype = (req.response_format or {}).get("type") + return ( + f"response_format.type {rtype!r} is not supported " + "(supported: text, json_object, json_schema)" + ) return None From 958688538b6ea03c2f69365fef26133b5b9831af Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?L=C3=A9on=20Avic=20Simmons?= Date: Sat, 3 Oct 2026 09:45:33 -0400 Subject: [PATCH 10/18] fix(mm): classify image fetch errors and re-check redirect targets (cherry picked from commit 8952c053603397040c0bdf58c9d3e020b82e319e) (cherry picked from commit 5eeb4802f3dde07fdc1b4b1c74a9b33d1726e264) --- docs/cli.md | 2 +- python/freetoken/mm/media.py | 124 ++++++++++++++++++++++++++-------- tests/server/test_mm_input.py | 121 +++++++++++++++++++++++++++++++++ 3 files changed, 219 insertions(+), 28 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 420c329351..06077d8a27 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -203,7 +203,7 @@ so a client can gate its attachment controls without reading the checkpoint conf | `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up | | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins | | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back | -| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain | +| `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain. Redirects are re-checked against the list and may not land on a private or loopback address | | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files | ## ft shell diff --git a/python/freetoken/mm/media.py b/python/freetoken/mm/media.py index 5c3c4ef789..a51e5bc98b 100644 --- a/python/freetoken/mm/media.py +++ b/python/freetoken/mm/media.py @@ -1,6 +1,8 @@ """Client-supplied image handling: ref collection, the accept gate, and byte fetching. Raises plain ValueError on bad input; the server layer maps it to its wire error type. +The message names the class of failure only: no local path, allowlist root, errno or +library text reaches the client (that detail goes to the server log). """ from __future__ import annotations @@ -8,10 +10,15 @@ import base64 from typing import TYPE_CHECKING, Any +from freetoken.utils import init_logger + if TYPE_CHECKING: from freetoken.server.args import ServerArgs +logger = init_logger(__name__) + _MAX_IMAGE_BYTES = 32 << 20 +_MAX_REDIRECTS = 5 def collect_image_refs(messages: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -56,6 +63,57 @@ def _check_media_domain(url: str, config: ServerArgs) -> None: ) +def _check_redirect_target(url: str, config: ServerArgs) -> None: + """A redirect hop is the remote server's choice, not the client's URL: it must pass the + allowlist like the original, and may never land on a loopback / private / link-local + address (an allowed public host 302-ing into the operator's network is the SSRF shape).""" + import ipaddress + from urllib.parse import urlparse + + if not url.startswith(("http://", "https://")): + raise ValueError("the URL redirects to a non-http(s) location") + try: + _check_media_domain(url, config) + except ValueError: + raise ValueError("the URL redirects outside the allowed domains") from None + host = (urlparse(url).hostname or "").rstrip(".").lower() + if host == "localhost" or host.endswith(".localhost"): + raise ValueError("the URL redirects to a private or loopback address") + try: + ip = ipaddress.ip_address(host) + except ValueError: + return # a name; the allowlist above is the gate for names (no DNS resolution here) + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + if not ip.is_global or ip.is_multicast: # covers loopback, private, link-local, CGNAT, unspecified + raise ValueError("the URL redirects to a private or loopback address") + + +async def _fetch_remote_media(url: str, config: ServerArgs) -> bytes: + """GET an http(s) ref, following redirects one hop at a time so every target is re-checked.""" + import httpx + + _check_media_domain(url, config) + async with httpx.AsyncClient(timeout=30.0, follow_redirects=False) as client: + for _ in range(_MAX_REDIRECTS + 1): + async with client.stream("GET", url) as resp: + if resp.next_request is not None: # a 3xx with a Location: httpx built the hop + url = str(resp.next_request.url) + _check_redirect_target(url, config) + continue + resp.raise_for_status() + declared = resp.headers.get("content-length", "") + if declared.isdigit() and int(declared) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + buf = bytearray() + async for chunk in resp.aiter_bytes(): + buf += chunk + if len(buf) > _MAX_IMAGE_BYTES: + raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") + return bytes(buf) + raise ValueError(f"the URL redirects more than {_MAX_REDIRECTS} times") + + def _load_local_media(url: str, config: ServerArgs) -> bytes: """Read a file:// ref; the resolved path must be a strict subpath of --allowed-local-media-path.""" from pathlib import Path @@ -70,41 +128,53 @@ def _load_local_media(url: str, config: ServerArgs) -> bytes: # resolve() follows symlinks, so a link inside the root escaping it is rejected too resolved = filepath.resolve() if Path(root).resolve() not in resolved.parents: - raise ValueError( - f"the file path {filepath} must be a subpath of " - f"--allowed-local-media-path {root}" - ) + raise ValueError("the file path must be a subpath of --allowed-local-media-path") return resolved.read_bytes() +async def _fetch_one(ref: dict[str, Any], config: ServerArgs) -> bytes: + """One ref to bytes. Our own ValueErrors are already client-safe and pass through; anything + a library raises is reduced to its failure class and logged with the detail.""" + data = ref.get("data") or "" + try: + if ref.get("kind") == "b64": + return base64.b64decode(data) + if data.startswith("data:"): + return base64.b64decode(data.split(",", 1)[1]) + except Exception as exc: # noqa: BLE001 -- binascii.Error / IndexError, input-driven + raise ValueError("image data is not valid base64") from exc + if data.startswith(("http://", "https://")): + import httpx + + try: + return await _fetch_remote_media(data, config) + except ValueError: + raise + except httpx.HTTPStatusError as exc: + raise ValueError(f"the image URL returned HTTP {exc.response.status_code}") from exc + except httpx.TimeoutException as exc: + raise ValueError("timed out fetching the image URL") from exc + except Exception as exc: # noqa: BLE001 -- transport / protocol / bad URL + logger.warning("image URL fetch failed: %r", exc) + raise ValueError("the image URL is unreachable") from exc + if data.startswith("file://"): + try: + return _load_local_media(data, config) + except ValueError: + raise + except OSError as exc: + logger.warning("local image read failed: %r", exc) + raise ValueError("the local file is missing or unreadable") from exc + raise ValueError("unsupported image source (expect http(s)/file url or base64)") + + async def fetch_image_bytes(refs: list[dict[str, Any]], config: ServerArgs) -> list[bytes]: """Resolve collected image refs (URLs / base64) to raw bytes, in order.""" out: list[bytes] = [] for ref in refs: - data = ref.get("data") or "" try: - if ref.get("kind") == "b64": - out.append(base64.b64decode(data)) - elif data.startswith("data:"): - out.append(base64.b64decode(data.split(",", 1)[1])) - elif data.startswith(("http://", "https://")): - import httpx - - _check_media_domain(data, config) - # redirect targets are not re-checked against the allowlist - async with httpx.AsyncClient(timeout=30.0, follow_redirects=True) as client: - resp = await client.get(data) - resp.raise_for_status() - if len(resp.content) > _MAX_IMAGE_BYTES: - raise ValueError(f"image exceeds {_MAX_IMAGE_BYTES} bytes") - out.append(resp.content) - elif data.startswith("file://"): - out.append(_load_local_media(data, config)) - else: - raise ValueError( - "unsupported image source (expect http(s)/file url or base64)" - ) - except Exception as exc: # noqa: BLE001 -- input-driven, client-classifiable + out.append(await _fetch_one(ref, config)) + except ValueError as exc: raise ValueError(f"could not load image: {exc}") from exc return out diff --git a/tests/server/test_mm_input.py b/tests/server/test_mm_input.py index a694a50a84..33561acdcd 100644 --- a/tests/server/test_mm_input.py +++ b/tests/server/test_mm_input.py @@ -10,6 +10,7 @@ import pytest +from freetoken.mm.media import _MAX_IMAGE_BYTES as _MAX from freetoken.mm.media import collect_image_refs, fetch_image_bytes, image_reject_reason from freetoken.server.generation import GenerationError, render_messages from freetoken.server.stats import derive_model_card @@ -205,3 +206,123 @@ def test_render_messages_rejects_an_image_in_a_merged_system_message(): ]}, {"role": "user", "content": "hi"}, ]) +# ---- error bodies and redirect policy (no network: httpx.MockTransport stands in for the remote) ---- + + +def _mock_remote(monkeypatch, routes): + """Serve ``routes`` ({url: (status, headers, body) | Exception}) through the client media.py builds; + returns the list of URLs the transport actually saw.""" + import httpx + + seen = [] + + def handler(request): + seen.append(str(request.url)) + route = routes[str(request.url)] + if isinstance(route, Exception): + raise route + status, headers, body = route + return httpx.Response(status, headers=headers, content=body, request=request) + + real = httpx.AsyncClient + monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: real(transport=httpx.MockTransport(handler), **kw)) + return seen + + +def _fetch(url, config): + return asyncio.run(fetch_image_bytes([{"kind": "url", "data": url}], config)) + + +def test_local_media_error_bodies_carry_no_path_root_or_errno(tmp_path): + root = tmp_path / "media" + root.mkdir() + config = _config(allowed_local_media_path=str(root)) + outside = tmp_path / "secret.png" + outside.write_bytes(b"x") + (root / "dir").mkdir() + + cases = { + f"file://{outside}": "subpath", # escapes the root + f"file://{root / 'missing.png'}": "missing or unreadable", # ENOENT + f"file://{root / 'dir'}": "missing or unreadable", # EISDIR + } + for url, expect in cases.items(): + with pytest.raises(ValueError, match=expect) as ei: + _fetch(url, config) + body = str(ei.value) + assert str(tmp_path) not in body and "Errno" not in body, body + assert "secret.png" not in body and "missing.png" not in body, body + + +def test_remote_fetch_failures_are_classified_without_library_text(monkeypatch): + import httpx + + big = _MAX + 1 + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/404.png": (404, {}, b""), + "https://cdn.example.com/down.png": httpx.ConnectError("[Errno -2] Name or service not known"), + "https://cdn.example.com/slow.png": httpx.ReadTimeout("timed out"), + "https://cdn.example.com/declared.png": (200, {"content-length": str(big)}, b""), + "https://cdn.example.com/streamed.png": (200, {}, b"\0" * big), + }) + config = _config(allowed_media_domains="cdn.example.com") + for name, expect in [("404", "HTTP 404"), ("down", "unreachable"), ("slow", "timed out"), + ("declared", "exceeds"), ("streamed", "exceeds")]: + with pytest.raises(ValueError, match=expect) as ei: + _fetch(f"https://cdn.example.com/{name}.png", config) + assert "Errno" not in str(ei.value) + assert len(seen) == 5 + + +def test_redirects_are_rechecked_against_the_allowlist(monkeypatch): + seen = _mock_remote(monkeypatch, { + "https://cdn.example.com/a.png": (302, {"location": "https://evil.com/a.png"}, b""), + "https://cdn.example.com/b.png": (302, {"location": "/real.png"}, b""), + "https://cdn.example.com/real.png": (200, {}, b"fakepng"), + "https://evil.com/a.png": (200, {}, b"pwned"), + }) + config = _config(allowed_media_domains="cdn.example.com") + with pytest.raises(ValueError, match="redirects outside the allowed domains") as ei: + _fetch("https://cdn.example.com/a.png", config) + assert "evil.com" not in str(ei.value) + # a same-host relative redirect is followed and served + assert _fetch("https://cdn.example.com/b.png", config) == [b"fakepng"] + # the disallowed hop was never requested + assert seen == ["https://cdn.example.com/a.png", "https://cdn.example.com/b.png", "https://cdn.example.com/real.png"] + + +@pytest.mark.parametrize("target", [ + "http://127.0.0.1/x", "http://10.0.0.1/x", "http://169.254.169.254/latest/meta-data", + "http://[::1]/x", "http://[::ffff:10.0.0.1]/x", "http://localhost/x", "http://0.0.0.0/x", + "http://100.64.0.1/x", +]) +def test_redirects_to_private_or_loopback_addresses_are_refused(monkeypatch, target): + seen = _mock_remote(monkeypatch, { + "https://public.example.com/a.png": (302, {"location": target}, b""), + target: (200, {}, b"internal"), + }) + # even with an empty allowlist (any domain admitted) a redirect may not reach inside + with pytest.raises(ValueError, match="private or loopback"): + _fetch("https://public.example.com/a.png", _config()) + assert seen == ["https://public.example.com/a.png"] + + +def test_redirect_hop_count_and_scheme_are_bounded(monkeypatch): + chain = {f"https://h.example.com/{i}.png": (302, {"location": f"https://h.example.com/{i + 1}.png"}, b"") + for i in range(7)} + chain["https://h.example.com/7.png"] = (200, {}, b"fakepng") + chain["https://h.example.com/file.png"] = (302, {"location": "file:///etc/passwd"}, b"") + seen = _mock_remote(monkeypatch, chain) + with pytest.raises(ValueError, match="redirects more than"): + _fetch("https://h.example.com/0.png", _config()) + assert len(seen) == 6 # the original plus five hops; the sixth hop is never requested + with pytest.raises(ValueError, match="non-http"): + _fetch("https://h.example.com/file.png", _config()) + # a short chain that stays public is fine + assert _fetch("https://h.example.com/5.png", _config()) == [b"fakepng"] + + +def test_bad_base64_is_named_without_decoder_text(): + for ref in ({"kind": "b64", "data": "not base64!!"}, {"kind": "url", "data": "data:image/png;base64"}): + with pytest.raises(ValueError, match="not valid base64"): + asyncio.run(fetch_image_bytes([ref], _config())) From cbe217b2a0f77cca05200251dd931ec61532ce75 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Fri, 25 Sep 2026 06:06:12 -0400 Subject: [PATCH 11/18] test(engine): cover presence/frequency penalties on the greedy and sampled paths Assisted-by: Claude Opus 5.5 (1M context) (cherry picked from commit 30b818114ce353da17bf6c3c9933e6c752eb7c11) (cherry picked from commit dc7cfc13fa5eaef06e105530b55e752529c9d0c6) --- tests/engine/test_sample.py | 84 +++++++++++++++++++++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 tests/engine/test_sample.py diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py new file mode 100644 index 0000000000..c1cc925f00 --- /dev/null +++ b/tests/engine/test_sample.py @@ -0,0 +1,84 @@ +"""Presence/frequency penalties in engine.sample.Sampler (greedy and sampled rows).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.core import SamplingParams +from freetoken.engine import sample as S + + +def _batch(*params: SamplingParams, can_decode: bool = True): + reqs = [ + SimpleNamespace(sampling_params=p, can_decode=can_decode, output_token_counts=None) + for p in params + ] + return SimpleNamespace(reqs=reqs) + + +def test_no_penalty_allocates_nothing(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0)) + args = sampler.prepare(batch) + assert args.penalties == [] and batch.reqs[0].output_token_counts is None + + +def test_frequency_penalty_counts_generated_tokens_on_the_greedy_path(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, frequency_penalty=10.0)) + logits = torch.tensor([[0.0, 5.0, 1.0, 0.0, 0.0]]) + + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] + # Token 1 now carries -10, so the runner-up wins; the caller's logits stay untouched. + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] + assert batch.reqs[0].output_token_counts.tolist() == [0, 1, 1, 0, 0] + assert logits[0, 1] == 5.0 + + +def test_intermediate_prefill_chunks_are_not_counted(): + sampler = S.Sampler(torch.device("cpu"), 5) + batch = _batch(SamplingParams(temperature=0.0, presence_penalty=1.0), can_decode=False) + assert sampler.prepare(batch).penalties == [] + + +def test_penalized_logits_reach_the_probability_sampler(monkeypatch): + """Sampled rows (the torch top-k/top-p fallback on Pascal) see the penalized logits.""" + seen = {} + + def fake_sample_impl(logits, temperatures, top_k, top_p): + seen["logits"] = logits.clone() + return torch.argmax(logits, dim=-1) + + monkeypatch.setattr(S, "sample_impl", fake_sample_impl) + monkeypatch.setattr(S, "make_device_tensor", lambda data, dtype, device: torch.tensor(data, dtype=dtype)) + sampler = S.Sampler(torch.device("cpu"), 4) + batch = _batch( + SamplingParams(temperature=1.0, top_k=2, presence_penalty=3.0), + SamplingParams(temperature=1.0, top_k=2), + ) + batch.reqs[0].output_token_counts = torch.tensor([0, 2, 0, 0], dtype=torch.int32) + logits = torch.tensor([[0.0, 4.0, 2.0, 0.0], [0.0, 4.0, 2.0, 0.0]]) + + tokens = sampler.sample(logits, sampler.prepare(batch)) + + assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0] # presence: -3 once, not per count + assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0] + assert tokens.tolist() == [2, 1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_penalty_through_the_pre_volta_torch_sampling_fallback(monkeypatch): + import freetoken.kernel.backend as backend + import freetoken.utils.arch as arch + + monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False) + monkeypatch.setattr(arch, "is_sm70_supported", lambda: False) + sampler = S.Sampler(torch.device("cuda"), 4) + batch = _batch(SamplingParams(temperature=1.0, top_k=1, frequency_penalty=5.0)) + logits = torch.tensor([[0.0, 4.0, 2.0, 0.0]], device="cuda") + + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] + assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] From d8d261b36e865c4d09ff0af18b2dde27aa5db9e8 Mon Sep 17 00:00:00 2001 From: "Eva (port)" Date: Sat, 3 Oct 2026 01:21:44 +0300 Subject: [PATCH 12/18] perf(model): fuse hc combine with the next block's RMSNorm via hc_combine_norm the vendored combine+norm kernel was exported but unwired; mix() takes a pre-computed rn, the model loop carries the fused pair across the layer boundary, skipped before PLE layers where ple(hidden) lands between. (cherry picked from commit 7636dd76cbe64989cbdec09e5c09a4c67be89542) (cherry picked from commit 0e847f8cd849e4a75cab781a7846af972afbbc58) --- python/freetoken/models/qwen4_exp/hc.py | 30 +++++++++++++++++----- python/freetoken/models/qwen4_exp/model.py | 30 ++++++++++++++++------ 2 files changed, 45 insertions(+), 15 deletions(-) diff --git a/python/freetoken/models/qwen4_exp/hc.py b/python/freetoken/models/qwen4_exp/hc.py index f13808dccd..1574cd473f 100644 --- a/python/freetoken/models/qwen4_exp/hc.py +++ b/python/freetoken/models/qwen4_exp/hc.py @@ -17,6 +17,7 @@ from freetoken.kernel.triton.hc import ( grouped_gemma_rmsnorm, hc_combine, + hc_combine_norm, hc_gate_mix, hc_silu, ) @@ -119,14 +120,16 @@ def _down(self, rn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: # both slices keep unit inner stride, so the kernels read them without a copy return down[:, : self.lowrank], down[:, self.lowrank : self.lowrank + self.hc_count] - def _mix_kernel(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_kernel(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_gemma_rmsnorm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) gate = self.input_mix_weight_up.forward(hc_silu(lora, self.hc_count)) return hc_gate_mix(rn, gate, self.hc_count), s - def _mix_torch(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) + def _mix_torch(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + if rn is None: + rn = grouped_plus_one_rms_norm(R, self.hc_norm.weight, self.hc_norm.eps, self.hc_count) lora, s = self._down(rn) lora = F.silu(lora.float() / self.hc_count) gate = self.input_mix_weight_up.forward(lora.to(R.dtype)) @@ -140,9 +143,12 @@ def _combine_torch(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> t out = out + y.float().unsqueeze(-2) * inject.unsqueeze(-1) return out.flatten(-2).to(R.dtype) - def mix(self, R: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor | None]: - """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine).""" - return self._mix_kernel(R) if R.is_cuda else self._mix_torch(R) + def mix(self, R: torch.Tensor, rn: torch.Tensor | None = None) -> Tuple[torch.Tensor, torch.Tensor | None]: + """Return the block input ``x [T, hidden]`` and the inject logits ``s [T, hc_count]`` (None if no combine). + + ``rn`` is the pre-normalized R from a preceding :meth:`combine_norm` (same norm weight as + this block's ``hc_norm``), letting the caller fuse the two kernels; None recomputes it.""" + return self._mix_kernel(R, rn) if R.is_cuda else self._mix_torch(R, rn) def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Tensor: """Inject the block output ``y [T, hidden]`` back into every stream of ``R``.""" @@ -150,5 +156,15 @@ def combine(self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor) -> torch.Te return hc_combine(R, y, s, self.hc_count) return self._combine_torch(R, y, s) + def combine_norm( + self, R: torch.Tensor, y: torch.Tensor, s: torch.Tensor, next_block: "GatedResidual" + ) -> Tuple[torch.Tensor, torch.Tensor | None]: + """combine() fused with ``next_block``'s hc_norm; returns (R', Rn'). CPU keeps the split path (Rn' = None).""" + if not R.is_cuda: + return self.combine(R, y, s), None + return hc_combine_norm( + R, y, s, next_block.hc_norm.weight, next_block.hc_norm.eps, self.hc_count + ) + __all__ = ["GatedResidual", "GroupedPlusOneRMSNorm", "grouped_plus_one_rms_norm"] diff --git a/python/freetoken/models/qwen4_exp/model.py b/python/freetoken/models/qwen4_exp/model.py index 7f4f826807..6ce2da985b 100644 --- a/python/freetoken/models/qwen4_exp/model.py +++ b/python/freetoken/models/qwen4_exp/model.py @@ -75,17 +75,19 @@ def __init__(self, config: ModelConfig, layer_id: int, *, prefix: str = "") -> N ) @nvtx_annotate("Layer_{}", layer_id_field="_layer_id") - def forward(self, hidden: torch.Tensor, batch: Batch) -> torch.Tensor: + def forward(self, hidden: torch.Tensor, batch: Batch, rn: torch.Tensor | None = None): + """Returns (R, mlp_out, inject) un-combined: the caller fuses the mlp combine with the + next consumer's norm (next layer's attn mix, or the final mixer).""" if self.ple is not None: hidden = hidden + self.ple.forward(hidden, batch) - block_input, inject = self.attn_hyper_connection.mix(hidden) + block_input, inject = self.attn_hyper_connection.mix(hidden, rn=rn) if self._is_linear: block_output = self.linear_attn.forward(block_input) else: block_output = self.self_attn.forward(block_input, batch) - hidden = self.attn_hyper_connection.combine(hidden, block_output, inject) - block_input, inject = self.mlp_hyper_connection.mix(hidden) - return self.mlp_hyper_connection.combine(hidden, self.mlp.forward(block_input), inject) + hidden, rn_mlp = self.attn_hyper_connection.combine_norm(hidden, block_output, inject, self.mlp_hyper_connection) + block_input, inject = self.mlp_hyper_connection.mix(hidden, rn=rn_mlp) + return hidden, self.mlp.forward(block_input), inject class Qwen4ExpModel(BaseOP): @@ -121,13 +123,25 @@ def forward(self, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor: meta = build_ple_metadata(batch, self._ple[0].args, input_ids.device) for ple in self._ple: # gather the pinned-host PLE rows while the early layers run ple.start_prefetch(batch, meta) - for layer in self.layers.op_list: - hidden = layer.forward(hidden, batch) + # The last block of each layer combines through the layer boundary, fusing the + # combine with the next consumer's hc_norm -- except before a PLE layer, whose + # `hidden + ple(hidden)` happens after the boundary and changes what the norm sees. + layers = self.layers.op_list + rn = None + for i, layer in enumerate(layers): + hidden, mlp_out, inject = layer.forward(hidden, batch, rn=rn) + nxt = layers[i + 1] if i + 1 < len(layers) else None + if nxt is not None and nxt.ple is not None: + hidden = layer.mlp_hyper_connection.combine(hidden, mlp_out, inject) + rn = None + else: + consumer = nxt.attn_hyper_connection if nxt is not None else self.hyper_connection_mixer + hidden, rn = layer.mlp_hyper_connection.combine_norm(hidden, mlp_out, inject, consumer) if meta is not None: # single writer: the layers only read the context, so a second PLE layer's # prefetch sees the un-rolled window commit_ngram_context(meta, getattr(batch, "fla_metadata", None)) - return self.hyper_connection_mixer.mix(hidden)[0] + return self.hyper_connection_mixer.mix(hidden, rn=rn)[0] class Qwen4ExpForCausalLM(BaseLLMModel): From 7ec5a115bb5728df0b59c67e82803460ee2eeacf Mon Sep 17 00:00:00 2001 From: "Eva (port)" Date: Sat, 3 Oct 2026 01:21:44 +0300 Subject: [PATCH 13/18] fix(kernel): bit-exact rmsnorm reduction tree in hc_combine_norm the fused kernel summed squares over a two-axis padded tile; the standalone grouped RMSNorm reduces one flat next_pow2 vector - different fp32 tree, ulp-level drift caught by the greedy gate. Reshape to the identical tree. (cherry picked from commit c1c5e76f5d98e08f0ff525fe71415bac028af3bc) (cherry picked from commit 8f8d7fbd7f0940e6ea4f0fa1b5ddfee4892efe0f) --- python/freetoken/kernel/triton/hc.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/python/freetoken/kernel/triton/hc.py b/python/freetoken/kernel/triton/hc.py index 45d71bd731..acae534c45 100644 --- a/python/freetoken/kernel/triton/hc.py +++ b/python/freetoken/kernel/triton/hc.py @@ -326,9 +326,11 @@ def _hc_combine_norm_kernel( tl.store(out_ptr + row * stride_out + offs, out, mask=mask_inner) out = out.to(tl.float32) - # Keep the two-axis reduction: flattening the padded tile is ~40% slower - # at decode sizes. - sum_sq = tl.sum(tl.sum(out * out, axis=1), axis=0) + # The standalone grouped RMSNorm reduces one padded [next_pow2(HC_DIM)] + # vector; reshape keeps the identical tree (same element + # order, same masked zeros), so the fused statistic is bit-exact vs the split. + flat = tl.reshape(out * out, [NUM_TILES_PAD * BLOCK_SIZE]) + sum_sq = tl.sum(flat) rrms = tl.rsqrt(sum_sq / HC_DIM + EPS) if launch_pdl: From c039107877b183765fd0cc60628150c15eac8ec1 Mon Sep 17 00:00:00 2001 From: "Eva (port)" Date: Sat, 3 Oct 2026 01:22:10 +0300 Subject: [PATCH 14/18] test(kernel): bitwise pin for hc_combine_norm against the split path (cherry picked from commit 7cba55ea8ba0d3576523c75ef0fcacb3d3e817bb) (cherry picked from commit 081534bb5574c7074ff37a2ecb9dca5a645073f1) --- tests/kernels/test_hc_combine_norm.py | 48 +++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) create mode 100644 tests/kernels/test_hc_combine_norm.py diff --git a/tests/kernels/test_hc_combine_norm.py b/tests/kernels/test_hc_combine_norm.py new file mode 100644 index 0000000000..7f46755479 --- /dev/null +++ b/tests/kernels/test_hc_combine_norm.py @@ -0,0 +1,48 @@ +"""Pins for the fused hc combine + grouped RMSNorm path (hc_combine_norm). + +hc_combine_norm replaces the standalone pair (hc_combine -> grouped_gemma_rmsnorm) +in the Qwen4-exp decoder loop. The swap is only sound if it is BITWISE identical: +the hyper-connection streams feed every later layer, and a bf16 flip here compounds +into different greedy tokens by the end of the stack. We assert torch.equal, not +assert_close: an approx-tolerance test would not have caught the regression this +path shipped with at first attempt -- the fused kernel reduced squares over the +[8, block] tile view (a two-axis tree) while the standalone kernel reduces the +flat row tree; fp32 ulps that differ between trees flip bf16 rounding. Both paths +on the same data must return the same bytes: fails before the reduction-tree fix +in _hc_combine_norm_kernel, passes after. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.hc import hc_combine, grouped_gemma_rmsnorm, hc_combine_norm + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + +HC, H = 4, 2560 # model geometry: the reduction tree specializes on the block layout +DIM = HC * H +EPS = 1e-6 + + +def _inputs(seed: int, T: int): + g = torch.Generator(device="cuda").manual_seed(seed) + bf = lambda shape: torch.randn(shape, generator=g, device="cuda", dtype=torch.float32).to(torch.bfloat16) + return ( + bf((T, DIM)), # residual streams R + bf((T, H)), # block output y + bf((T, HC)), # inject scores s + (torch.randn(DIM, generator=g, device="cuda", dtype=torch.float32) * 0.1).to(torch.bfloat16), # norm weight + ) + + +@pytest.mark.parametrize("T", [1, 8]) +@pytest.mark.parametrize("seed", range(6)) +def test_combine_norm_matches_split_bitwise(seed: int, T: int) -> None: + R, y, s, w = _inputs(seed, T) + r2_split = hc_combine(R, y, s, HC) + rn_split = grouped_gemma_rmsnorm(r2_split, w, EPS, HC) + r2_fused, rn_fused = hc_combine_norm(R, y, s, w, EPS, HC) + assert torch.equal(r2_fused, r2_split) + assert torch.equal(rn_fused, rn_split) From 0aa00e4592279c0b0b3e60800abbf606b24e076f Mon Sep 17 00:00:00 2001 From: Siarhei Berdachuk Date: Mon, 5 Oct 2026 19:11:42 +0300 Subject: [PATCH 15/18] fix(server): soft-accept OpenAI response_format json_object/json_schema Rejecting response_format with "no constrained decoding" breaks OpenAI-compatible clients that rely on json_object/json_schema even when the model can emit valid JSON from a prompt. Accept those wire types and prepend a best-effort system hint; schema is not token-level enforced. Unknown response_format types still return 400. (cherry picked from commit ca7d704bb84e255d8677b788f295e9cb6ebb06f0) (cherry picked from commit db64d8eff1a7778b6bc093634303d31e1324af76) --- python/freetoken/server/openai_api.py | 44 +++++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 79dff72963..7a345cadfa 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -71,6 +71,43 @@ def _thinking_type(req: Any) -> str | None: +def _json_mode_system_message(response_format: dict[str, Any] | None) -> dict[str, str] | None: + """Best-effort JSON mode hint for OpenAI response_format. + + FreeToken has no constrained/guided decoding yet. OpenAI-compatible clients + still send response_format=json_object|json_schema. Rejecting them with 400 + breaks those clients even though the model can emit valid JSON from a + prompt. Accept the wire shape and prepend a system reminder so the model is + steered toward JSON; the schema is NOT token-level enforced. + """ + if not response_format: + return None + rtype = response_format.get("type") + if rtype == "json_object": + return { + "role": "system", + "content": ( + "Respond with a single valid JSON object only. " + "Do not wrap it in markdown fences or add commentary." + ), + } + if rtype == "json_schema": + wrapper = response_format.get("json_schema") or {} + name = wrapper.get("name") or "response" + schema = wrapper.get("schema", wrapper) + return { + "role": "system", + "content": ( + f"Respond with a single JSON object that validates against the " + f"JSON Schema below (schema name: '{name}'). Put required fields " + "at the top level of the object — do not wrap them under an extra " + "key named after the schema. Do not wrap the answer in markdown " + f"fences or add commentary.\nJSON Schema:\n{json.dumps(schema, ensure_ascii=False)}" + ), + } + return None + + def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -1063,8 +1100,11 @@ def _usage( def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # We have no constrained/guided decoding; only plain text ('text' or unset) is honored. - return response_format is not None and response_format.get("type") not in (None, "text") + # Soft-accept OpenAI json_object/json_schema (prompt-steered; no constrained decoding). + # Unknown types still 400. See _json_mode_system_message. + if response_format is None: + return False + return response_format.get("type") not in (None, "text", "json_object", "json_schema") def _completion_unsupported_reason(req: CompletionRequest) -> str | None: From 97758b6cac9413cad4f497344b71b767453ccaab Mon Sep 17 00:00:00 2001 From: Siarhei Berdachuk Date: Mon, 5 Oct 2026 19:18:06 +0300 Subject: [PATCH 16/18] refactor(server): drop response_format system-hint injection Soft-accepting json_object/json_schema is enough when clients already request JSON in the prompt. The injected schema wrapper was redundant and could distort the output shape. (cherry picked from commit 65907c8bc2c1c11a4af11dff328f7cc293806433) (cherry picked from commit c82625f6eed239a9edcc6e76e935d3b51359f425) --- python/freetoken/server/openai_api.py | 45 ++------------------------- 1 file changed, 3 insertions(+), 42 deletions(-) diff --git a/python/freetoken/server/openai_api.py b/python/freetoken/server/openai_api.py index 7a345cadfa..0fc8113bc8 100644 --- a/python/freetoken/server/openai_api.py +++ b/python/freetoken/server/openai_api.py @@ -68,46 +68,6 @@ def _thinking_type(req: Any) -> str | None: if value in ("enabled", "disabled"): return value return None - - - -def _json_mode_system_message(response_format: dict[str, Any] | None) -> dict[str, str] | None: - """Best-effort JSON mode hint for OpenAI response_format. - - FreeToken has no constrained/guided decoding yet. OpenAI-compatible clients - still send response_format=json_object|json_schema. Rejecting them with 400 - breaks those clients even though the model can emit valid JSON from a - prompt. Accept the wire shape and prepend a system reminder so the model is - steered toward JSON; the schema is NOT token-level enforced. - """ - if not response_format: - return None - rtype = response_format.get("type") - if rtype == "json_object": - return { - "role": "system", - "content": ( - "Respond with a single valid JSON object only. " - "Do not wrap it in markdown fences or add commentary." - ), - } - if rtype == "json_schema": - wrapper = response_format.get("json_schema") or {} - name = wrapper.get("name") or "response" - schema = wrapper.get("schema", wrapper) - return { - "role": "system", - "content": ( - f"Respond with a single JSON object that validates against the " - f"JSON Schema below (schema name: '{name}'). Put required fields " - "at the top level of the object — do not wrap them under an extra " - "key named after the schema. Do not wrap the answer in markdown " - f"fences or add commentary.\nJSON Schema:\n{json.dumps(schema, ensure_ascii=False)}" - ), - } - return None - - def chat_request_to_genspec( req: ChatCompletionRequest, model_sampling: dict[str, Any], @@ -1100,8 +1060,9 @@ def _usage( def _response_format_unsupported(response_format: dict[str, Any] | None) -> bool: - # Soft-accept OpenAI json_object/json_schema (prompt-steered; no constrained decoding). - # Unknown types still 400. See _json_mode_system_message. + # Soft-accept OpenAI json_object/json_schema. FreeToken has no constrained + # decoding; clients that already ask for JSON in the prompt still need the + # wire field accepted (HTTP 200) instead of a hard 400. Unknown types 400. if response_format is None: return False return response_format.get("type") not in (None, "text", "json_object", "json_schema") From 941fe3be06e8d9b2bc8b7dda2a5bb410d808031c Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 13:14:39 -0400 Subject: [PATCH 17/18] test(engine): drop the pre-volta sampling penalty test, whose torch fallback this branch does not carry Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- tests/engine/test_sample.py | 15 --------------- 1 file changed, 15 deletions(-) diff --git a/tests/engine/test_sample.py b/tests/engine/test_sample.py index c1cc925f00..28e7195b58 100644 --- a/tests/engine/test_sample.py +++ b/tests/engine/test_sample.py @@ -67,18 +67,3 @@ def fake_sample_impl(logits, temperatures, top_k, top_p): assert seen["logits"][0].tolist() == [0.0, 1.0, 2.0, 0.0] # presence: -3 once, not per count assert seen["logits"][1].tolist() == [0.0, 4.0, 2.0, 0.0] assert tokens.tolist() == [2, 1] - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") -def test_penalty_through_the_pre_volta_torch_sampling_fallback(monkeypatch): - import freetoken.kernel.backend as backend - import freetoken.utils.arch as arch - - monkeypatch.setattr(backend, "is_flashinfer_installed", lambda: False) - monkeypatch.setattr(arch, "is_sm70_supported", lambda: False) - sampler = S.Sampler(torch.device("cuda"), 4) - batch = _batch(SamplingParams(temperature=1.0, top_k=1, frequency_penalty=5.0)) - logits = torch.tensor([[0.0, 4.0, 2.0, 0.0]], device="cuda") - - assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [1] - assert sampler.sample(logits, sampler.prepare(batch)).tolist() == [2] From 8eb353b30a3a1047f79f281cabe7d5aba838df14 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Tue, 13 Oct 2026 13:14:39 -0400 Subject: [PATCH 18/18] fix(moe): align on triton past sgl_kernel's 1024-expert limit sgl_kernel's moe_align_block_size (0.4.5 and 0.4.9) reads out of bounds once num_experts + 1 exceeds 1024: an illegal memory access that poisons the CUDA context. A prefill over the offload slot cache (FlashML-org/FreeToken#676) aligns against the slot count, 2100 here, so the first short prefill crashed the server; the bf16 and NVFP4 prefill paths both go through it. The Triton align has no such limit and matches sgl_kernel's output below it. Assisted-by: Claude Opus 5.5 Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01Bu6LgoxLR4wETqb7RPR2vt --- python/freetoken/moe/fused.py | 8 +++++++- tests/moe/test_nvfp4_backends.py | 19 +++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/python/freetoken/moe/fused.py b/python/freetoken/moe/fused.py index a90aa4c792..b0b9c01c71 100644 --- a/python/freetoken/moe/fused.py +++ b/python/freetoken/moe/fused.py @@ -44,6 +44,9 @@ def fused_topk( return fused_topk_softmax(gating_output, topk, renormalize, num_token_non_padded) +_SGL_ALIGN_MAX_EXPERTS = 1024 + + def moe_align_block_size( topk_ids: torch.Tensor, block_size: int, num_experts: int ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: @@ -86,7 +89,10 @@ def moe_align_block_size( """ from freetoken.kernel.backend import is_sgl_kernel_installed - if not is_sgl_kernel_installed(): + # sgl_kernel's align kernel takes num_experts + 1 counters in one block and reads out of + # bounds past 1024 of them (0.4.5 and 0.4.9). A prefill over an offload slot cache passes the + # slot count here (#676), which is thousands; the Triton kernel has no such limit. + if not is_sgl_kernel_installed() or num_experts + 1 > _SGL_ALIGN_MAX_EXPERTS: from freetoken.kernel.triton.moe_align import ( moe_align_block_size as triton_moe_align_block_size, ) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 0a2c68ecb5..a31d6ff763 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -639,3 +639,22 @@ def test_b12x_pack_keeps_per_layer_banks_and_flat_alphas(): assert len(cache.bank_sources["gate_up"]) == L assert sum(t.shape[0] for t in cache.bank_sources["gate_up"]) == total assert cache.gate_up_alpha.shape == (total,) + + +@cuda +@pytest.mark.parametrize("num_experts", [1023, 2100]) +def test_moe_align_handles_a_slot_cache_wider_than_1024(num_experts): + """A prefill over the offload slot cache aligns against its slot count (#676): sgl_kernel's + align reads out of bounds past 1024 counters, so a 2100-slot cache must still align.""" + from freetoken.moe.fused import moe_align_block_size + + torch.manual_seed(0) + ids = torch.randint(0, num_experts, (300, 10), device="cuda", dtype=torch.int32) + sorted_ids, expert_ids, padded = moe_align_block_size(ids, 16, num_experts) + n = int(padded) + flat = ids.flatten() + real = sorted_ids[:n][sorted_ids[:n] < flat.numel()] + assert sorted(real.tolist()) == list(range(flat.numel())) + block_of = torch.arange(n, device="cuda") // 16 + rows = sorted_ids[:n] < flat.numel() + assert torch.equal(flat[sorted_ids[:n][rows].long()], expert_ids[block_of[rows]])