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 |
|
The fusion is clearly right — 39 elementwise kernels to compute some row ids is work the launcher should not be doing. Two things from a deployment that may temper the headline, and one offer. The "39 launches per decode step" figure is per PLE layer, and the layer count is a checkpoint property. "ple_layer_ids": [2]— a single PLE layer out of 48. So this checkpoint pays 39 launches per step, not 39 × N. Worth saying explicitly in the PR, because a reader sizing the win will otherwise assume it scales with model depth. On the decode path those launches are already inside a CUDA graph. with get_global_ctx().forward_batch(batch):
self.buffer.logits[:bs] = model.forward()and Where the win should survive intact is prefill, which is not graph-captured. That is also where it matters most for us — we serve 262k contexts, so a wide prefill walks many chunks. None of that argues against the PR. It argues that the benchmark to quote is prefill/TTFT rather than decode tok/s, and that the decode figure should be qualified with "uncaptured" or measured with graphs on, or a reader with graphs enabled will not reproduce it. The offer: we run this exact checkpoint at TP=2 on 2 × RTX 6000 Ada with graphs on, and I have a TTFT harness across prompt widths (1k/4k/8k/16k) already pointed at it. If you want the prefill number on a real 48-layer Qwen3.8-Flash-Next rather than a synthetic, say the word and I will run your branch and post TTFT before/after. It merges cleanly onto our deploy line as of today. 🤖 Generated with Claude Code |
|
You're right on both, and I went and read the code for the second one before replying rather than take it on trust. On main, start_prefetch calls row_ids from inside the model forward, and the runner captures the whole forward, so with graphs on the 39 host launches get recorded once and never issued again on decode. I'd been picturing my own box, where the Windows mmap table backend does the hash outside the graph before each replay, so decode there really does pay the uncaptured cost every step. That backend isn't in this PR though, so it doesn't count. What's left per replay with graphs on is the GPU-side work, and MT-z's nsys row already has it: 44.7 µs GPU-busy per call on the torch path vs 1.3 µs fused. So the fair decode number for a graphs-on server is about 43 µs saved per step per PLE layer, not 131, which against your ~20 ms step is ~0.2% and well inside your spread. The 131 µs / 12× figures belong to the uncaptured case, which is prefill or a server running graphs off like MT-z's. The one-layer note is already in the description from the earlier exchange, but I'll go through the Measured section and mark the decode rows as uncaptured, and add the 44.7 → 1.3 µs row for the graphs-on case, so nobody with graphs enabled tries to reproduce a number they can't get. And yes please to the TTFT run, that's the measurement this thread is missing. Two small asks: flip FREETOKEN_PLE_FUSED_HASH=0/1 on the same build rather than main vs branch so the hash path is the only thing that moves, and say what the prefill chunk size is. Going by MT-z's 2048-token row the saving is roughly 175 µs per chunk per PLE layer, so a 16k prompt at 2k chunks is about 1.4 ms off a TTFT measured in seconds. I expect it to disappear into your noise, and I'd rather that be stated from a measurement than left implied. If anything shows it'll be at the 1k end where the prompt is a single chunk. Either way the table goes in the description. |
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.The decode rows are wall time with the forward uncaptured. On
mainthe hash runsinside
Qwen4ExpModel.forward, which the runner records into the CUDA graph, so aserver with
--cuda-graph-max-bsset pays the 39 host launches once at capture andnever again; what a replay still pays is the last row, the GPU-side work of 39 tiny
kernels against one. So the graphs-on decode saving is ~43 us per step per PLE layer,
not ~131 us (@gdevenyi's point below). Prefill is not captured and keeps the full
uncaptured saving per chunk.
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, in prefill, 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.