perf(ple): fuse the n-gram row-id hash into one Triton kernel

[dejay2]

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._window packs the ragged tokens into a [B, ctx+max_len]
buffer, _shift_ignore_eos runs a cummax boundary scan over the whole window
to 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 the
same arithmetic in one kernel, one program per token, and dispatches
NGramEmbedding.row_ids to it on CUDA.

Two observations make the one-pass form possible:

• The cummax-over-the-whole-window boundary scan collapses to an
  ngram_size - 1 step walk. _shift_ignore_eos marks shift s valid at
  position p iff p - s >= 0 and no boundary token sits in [p-s, p-1].
  Only shifts below ngram_size are ever used, so the scan never looks further
  back than that and the predicate can be carried incrementally as the walk
  goes.
• The packed window never has to be materialized. Token t belongs to request
  req[t] at intra-request offset local[t], so the token s places to its
  left is input_ids[t - s] when local[t] >= s and
  ngram_context[req[t], ctx_len + local[t] - s] otherwise. Out of range on
  the left is the boundary token — exactly what the eos-filled packed window
  gave.

Why it is safe

• The torch path stays. The old implementation is renamed
  NGramEmbedding.row_ids_reference and is still the CPU path and the oracle
  the kernel is diffed against in tests. Nothing about it changed.
• Default on only where the reference already ran. row_ids takes the
  fused path only when input_ids is on CUDA and ngram_context,
  layer_multipliers, ngram_heads_vocab_sizes and ngram_heads_offsets are
  all on that same device. CPU stays on the reference.
• FREETOKEN_PLE_FUSED_HASH=0 puts the hash back on the reference ops
  without a rebuild.
• Byte-identical, not "close": the tests assert torch.equal against the
  reference on prefill and decode shapes, and after CUDA-graph replays.
  torch.remainder is floored and Triton's % is truncated, so the kernel
  fixes the sign up explicitly.
• Capture-safe: fixed shapes, every input on device, no host reads, and an
  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 that
  happens during capture is deliberately 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 have the same [T] and mean
  opposite things (one token per request at offset 0, vs B offsets inside one
  request). The memo is bounded at _TOKEN_INDEX_CACHE_SIZE entries.
• The kernel's two geometry preconditions (ngram_size - 1 context ids, and
  heads_per_ngram heads per n-gram order) are raised as ValueError, not
  asserted — the kernel addresses the context row and the head blocks by them,
  and python -O must not turn a geometry mismatch into an out-of-bounds read.

Measured

RTX 5090, this repo's toy PLE config (ngram_size 3, 4 hash heads), median of
7 x 300 iterations per point; launch count from torch.profiler over one call.

	torch reference	fused kernel
CUDA 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

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.)

How it was tested

Windows 11, RTX 5090, CUDA available.

$env:PYTHONPATH = "python;<pytest-lib>"
python -m pytest -q tests/models/qwen4_exp/test_ple.py -p no:cacheprovider -W ignore

	result
base (origin/main, a80b4d3)	15 passed, 4 skipped
this branch	26 passed, 4 skipped

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:

python -m pytest -q tests/models/qwen4_exp -p no:cacheprovider -W ignore
# base:        86 passed, 52 skipped, 6 failed
# this branch: 97 passed, 52 skipped, 6 failed

The same 6 failures on both: they are all tests/models/qwen4_exp/test_weight.py
and pre-existing on this platform — they need os.O_DIRECT, which Windows does
not 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_shape
• test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill
• test_token_index_cache_is_bounded
• test_fused_hash_is_off_on_cpu
• test_fused_hash_matches_the_torch_reference[prefill|decode] (CUDA)
• test_fused_hash_captures_and_replays_in_a_cuda_graph (CUDA) — capture, then
  three 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:

• the layer-ahead MoE expert prefetch and the small-prefill decode movement
  that shared a commit with this work;
• the mmap/disk PLE table backend and its pinned staging-buffer ring;
• anything MTP / speculative-decoding related.

No serving-path behaviour outside NGramEmbedding.row_ids changes; the only
signature change is the new optional out= argument on row_ids, which all
existing callers ignore.

🤖 Generated with Claude Code

https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK