perf(ple): fuse the n-gram row-id hash into one Triton kernel - #338
perf(ple): fuse the n-gram row-id hash into one Triton kernel#338dejay2 wants to merge 1 commit into
Conversation
The PLE hash builds its row ids from a packed ``[B, ctx+max_len]`` window, a
cummax boundary scan and a per-ngram XOR/multiply/remainder/offset loop. Every
one of those is a tiny elementwise op, so a single PLE layer spends 39 CUDA
launches (~400 us of launch wall) on a few microseconds of GPU work, on the
critical path of every step.
``freetoken.kernel.triton.ple_hash.ple_row_ids`` is the same arithmetic as one
kernel, one program per token:
- the cummax over the whole window collapses to an ``ngram_size-1`` step walk,
because ``_shift_ignore_eos`` only ever needs shifts below ``ngram_size`` and
the "no boundary token in between" predicate can be carried incrementally;
- the packed window is never materialized: token ``t`` at intra-request offset
``local[t]`` reads ``input_ids[t-s]`` when ``local[t] >= s`` and
``ngram_context[req[t], ...]`` otherwise, out of range on the left being the
boundary token exactly as the eos-filled window was.
``NGramEmbedding.row_ids_reference`` is the old torch-op transcription, kept as
the oracle the kernel is diffed against and as the CPU path; ``row_ids`` picks
the kernel only when every input is on the same CUDA device.
``FREETOKEN_PLE_FUSED_HASH=0`` puts the hash back on the reference.
Fixed shapes, all inputs on device, no host reads, an optional ``out`` buffer:
the hash is capture-safe and replays inside a captured decode step. The
``(request, offset)`` index the kernel addresses through is memoized per
(is_decode, shape, device) so a replay reads a stable address; a build that
happens during capture is not cached, since its buffers live in the graph pool.
``is_decode`` is part of that key because a decode of B requests and a prefill
of one B-token request are the same ``[T]`` and mean opposite things.
Measured on an RTX 5090 (toy config, ngram_size 3, 4 heads), median of 7 x 300
iterations, torch profiler for the launch count:
launches per call 39 -> 1
decode B=1 351 us -> 17 us
decode B=4 365 us -> 16 us
decode B=8 361 us -> 16 us
prefill 512 tokens 482 us -> 16 us
Byte-identical to the reference: the new tests assert ``torch.equal`` on both
the prefill and decode shapes, and after three CUDA-graph replays.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK
…cle stat) Ports from upstream FreeToken, adapted to the tier: FlashML-org#342 lm_head on sampled rows only (already generalised here via select_lm_head_rows); FlashML-org#339 the varlen GDN/KDA prefill conv takes max_seq_len from the scheduler on the Triton fallback (inert when sgl_kernel is installed, which every install path pins, so no node-4 change); FlashML-org#338 the n-gram PLE row-id hash as one Triton kernel with a bounded memo that is bypassed during CUDA graph capture (consumed by the pinned and cached PLE backends; the disk backend stages from its host hash); FlashML-org#231 the routing-oracle hit rate on the stats line next to the realised hot-pair rate, with the baseline reset on a live cache rebuild so the oracle can never read below realised. FlashML-org#89 (route-density tile selection) is skipped: its ds_fp4 tile table does not match the NVFP4 kernel's, which needs its own sm_89 sweep. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01A88MCbnLtwsFSHmqwuJezY
|
Nice one — I pulled this into my daily branch and it has been serving since. Two things I can add Setup. RTX 4090 24 GB (sm_89), CUDA toolkit 13.3, The measurement, and it confirms your diagnosisThe
The row I would point at is not the speedup, it is that the torch path costs the same ~143 µs at What it buys end to end here: 0.14%, and I would rather say so
Noise is ~23× the effect, so I did not run a serving A/B — it could not have resolved this, and a None of that is an argument against the change. It is 12× on its own terms, it is free, and the CorrectnessI checked the kernel's arithmetic against Empirically the equivalence tests pass here in both directions One failing test, and it is not yours
Before and after over those 84 files: 7 failures either way, nothing gained, nothing broken. Happy to run anything else against this — the A/B harness needs no model load, so re-measuring on a Written with AI assistance; every number above was measured on my hardware (RTX 4090, sm_89) |
|
Good to have a second set of numbers on a different card, and on the real 16-head geometry. The flat ~143 µs from bs=1 up to a 2048-token prefill is the clearest version of the launch-bound argument I have seen, clearer than my own table. The end-to-end point is fair and I should have spelled it out in the description: On Thanks for the hand check of the valid-shift predicate. That was the part I most wanted someone else to read. |
|
Tested this on a Qwen3.8-Flash-Next deployment; effect on decode is within noise here, no regression.
The pinned PLE backend was used, so the fused Measured on a 2x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4, no NVLink) / 2x Xeon Gold 6526Y / 503 GiB box, CUDA 13.3, torch 2.11+cu130, sgl_kernel 0.4.5, model Single-stream = median of three 64-vs-256-token completion pairs, aggregate = 8 concurrent 256-token completions, TTFT on a ~1k-token prompt. Run-to-run spread of the baseline on this box is about +-4% single-stream. |
|
Borrowing the nsys approach from your comment and running the same trace on this box. It
The correction is to the per-launch figure in my earlier comment. "143 us over 39 launches is As a second data point on where it would show: the step here is launch-bound and the change One note on the tooling, since it cost some time here. With CUDA graphs enabled, nsys records Measured on an RTX 4090 (24 GiB, sm_89) / i9-14900KF / 61 GiB box, driver 595.84, CUDA 13.3, Launches and GPU-busy are ten iterations inside one NVTX range, 30 warm-ups outside it; wall The reason is a measurement now rather than an inference, which is what I wanted. Thank you. Assisted-by: Claude Opus 5 |
3 participants
What
The Qwen3.8-Flash-Next PLE layer turns each token's n-gram window into
[T, num_ngram_heads]table row ids. Today that is done in torch ops:NGramEmbedding._windowpacks the ragged tokens into a[B, ctx+max_len]buffer,
_shift_ignore_eosruns acummaxboundary scan over the whole windowto mark the shifts that would cross a boundary token, and then a per-n-gram
loop does multiply / XOR /
remainder/ offset.Every one of those is a tiny elementwise op. A single PLE layer therefore
issues 39 CUDA kernel launches for a few microseconds of actual GPU work,
on the critical path of every decode step — the wall is launch overhead, not
arithmetic.
This PR adds
freetoken/kernel/triton/ple_hash.py(ple_row_ids), which is thesame arithmetic in one kernel, one program per token, and dispatches
NGramEmbedding.row_idsto it on CUDA.Two observations make the one-pass form possible:
cummax-over-the-whole-window boundary scan collapses to anngram_size - 1step walk._shift_ignore_eosmarks shiftsvalid atposition
piffp - s >= 0and no boundary token sits in[p-s, p-1].Only shifts below
ngram_sizeare ever used, so the scan never looks furtherback than that and the predicate can be carried incrementally as the walk
goes.
tbelongs to requestreq[t]at intra-request offsetlocal[t], so the tokensplaces to itsleft is
input_ids[t - s]whenlocal[t] >= sandngram_context[req[t], ctx_len + local[t] - s]otherwise. Out of range onthe left is the boundary token — exactly what the eos-filled packed window
gave.
Why it is safe
NGramEmbedding.row_ids_referenceand is still the CPU path and the oraclethe kernel is diffed against in tests. Nothing about it changed.
row_idstakes thefused path only when
input_idsis on CUDA andngram_context,layer_multipliers,ngram_heads_vocab_sizesandngram_heads_offsetsareall on that same device. CPU stays on the reference.
FREETOKEN_PLE_FUSED_HASH=0puts the hash back on the reference opswithout a rebuild.
torch.equalagainst thereference on prefill and decode shapes, and after CUDA-graph replays.
torch.remainderis floored and Triton's%is truncated, so the kernelfixes the sign up explicitly.
optional
out=destination so a replay writes into a fixed buffer. The(request, offset)index the kernel addresses through is memoized per(is_decode, shape, device)so a replay reads a stable address; a build thathappens during capture is deliberately not cached, since its buffers live in
the graph pool.
is_decodeis part of that key because a decode ofBrequests and a prefill of one
B-token request have the same[T]and meanopposite things (one token per request at offset 0, vs
Boffsets inside onerequest). The memo is bounded at
_TOKEN_INDEX_CACHE_SIZEentries.ngram_size - 1context ids, andheads_per_ngramheads per n-gram order) are raised asValueError, notasserted — the kernel addresses the context row and the head blocks by them,
and
python -Omust not turn a geometry mismatch into an out-of-bounds read.Measured
RTX 5090, this repo's toy PLE config (
ngram_size3, 4 hash heads), median of7 x 300 iterations per point; launch count from
torch.profilerover one call.That is per PLE layer per step. (The same change measured on the full
Qwen3.8-Flash-Next serving path on the same box: 415 us -> 17 us.)
Note the multiplier: on
Qwen3.8-Flash-Next-NVFP4ple_layer_idsis[2], so thisis one layer out of 48 and the per-step saving is one layer's worth. Against a
~95 ms decode step on a 4090 with experts streaming from NVMe that is ~0.14 %,
inside run-to-run noise (measured by @MT-z below). This PR does not claim a
tok/s win on that checkpoint; the win is per-layer, and it compounds on any
checkpoint with PLE on several layers or in graph-off small-batch decode.
How it was tested
Windows 11, RTX 5090, CUDA available.
origin/main, a80b4d3)The 4 skips are unchanged and pre-existing (3 need
FREETOKEN_QWEN4_HF_PYTHON,1 needs
FREETOKEN_QWEN4EXP_MODEL).Also run, unchanged by this PR:
The same 6 failures on both: they are all
tests/models/qwen4_exp/test_weight.pyand pre-existing on this platform — they need
os.O_DIRECT, which Windows doesnot have.
New tests (11, in
tests/models/qwen4_exp/test_ple.py):test_token_index_addresses_the_same_window_as_the_packed_build— the(request, offset)pair names the very cell the packed window would have read.test_token_index_cache_is_keyed_by_shapetest_token_index_does_not_confuse_a_decode_with_a_one_request_prefilltest_token_index_cache_is_boundedtest_fused_hash_is_off_on_cputest_fused_hash_matches_the_torch_reference[prefill|decode](CUDA)test_fused_hash_captures_and_replays_in_a_cuda_graph(CUDA) — capture, thenthree replays with fresh ids and context, each checked against the reference.
test_fused_hash_env_switch_restores_the_torch_path(CUDA)test_the_fused_hash_refuses_a_geometry_it_cannot_address[ctx|heads]Not included
This is one topic split out of a larger fork branch. Deliberately left out:
that shared a commit with this work;
No serving-path behaviour outside
NGramEmbedding.row_idschanges; the onlysignature change is the new optional
out=argument onrow_ids, which allexisting callers ignore.
🤖 Generated with Claude Code
https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK