From de77d16a5ac36494d85964f846b50ad36f7acc21 Mon Sep 17 00:00:00 2001 From: dejay2 <218806300+dejay2@users.noreply.github.com> Date: Wed, 2 Sep 2026 05:03:53 +0100 Subject: [PATCH] perf(ple): fuse the n-gram row-id hash into one Triton kernel 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 Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK --- python/freetoken/kernel/triton/ple_hash.py | 145 +++++++++++++++++++ python/freetoken/models/qwen4_exp/ple.py | 98 ++++++++++++- tests/models/qwen4_exp/test_ple.py | 153 +++++++++++++++++++++ 3 files changed, 394 insertions(+), 2 deletions(-) create mode 100644 python/freetoken/kernel/triton/ple_hash.py diff --git a/python/freetoken/kernel/triton/ple_hash.py b/python/freetoken/kernel/triton/ple_hash.py new file mode 100644 index 000000000..c8c5e6302 --- /dev/null +++ b/python/freetoken/kernel/triton/ple_hash.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Fused n-gram hash -> PLE table row ids. + +The eager form of this hash (``NGramEmbedding._window`` + ``_shift_ignore_eos`` + the per-ngram +XOR/multiply/remainder/offset loop) is 39 tiny CUDA kernels per PLE layer per step -- ~400 us of +launch wall for a few us of GPU work. This is the same arithmetic as one kernel, one program per +token; ``NGramEmbedding.row_ids_reference`` keeps the torch-op form as the oracle and CPU path. + +The key observation that collapses the ``cummax``-over-the-whole-window boundary scan into a +``ngram_size-1`` step walk: ``_shift_ignore_eos`` marks shift ``s`` valid at position ``p`` iff +``p-s >= 0`` and no boundary token sits anywhere in ``[p-s, p-1]``. Only shifts ``< ngram_size`` +are ever used, so the scan never needs to look further back than that, and the predicate is +built incrementally as the walk goes. + +Window addressing avoids materializing the ``[B, ctx+max_len]`` packed window entirely. Token +``t`` of the forward 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 as the eos-filled packed window was. + +Capture-safe: fixed shapes, every input on device, no host reads. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _ple_row_ids_kernel( + ids_ptr, # [T] int64 -- this forward's tokens, ragged, in request order + ctx_ptr, # [B, CTX_LEN] int64 -- the tokens immediately before each request's first + req_ptr, # [T] int32 -- request index of each token + local_ptr, # [T] int32 -- intra-request offset of each token + mult_ptr, # [NGRAM] int64 + vocab_ptr, # [NUM_HEADS] int64 + off_ptr, # [NUM_HEADS] int64 + out_ptr, # [T, NUM_HEADS] int64 + EOS: tl.constexpr, + CTX_LEN: tl.constexpr, + NGRAM: tl.constexpr, + HEADS_PER: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_H: tl.constexpr, +): + token = tl.program_id(0).to(tl.int64) + req = tl.load(req_ptr + token).to(tl.int64) + local = tl.load(local_ptr + token).to(tl.int64) + + head = tl.arange(0, BLOCK_H) + head_ok = head < NUM_HEADS + + # shift 0 is always the token itself; the eos-crossing rule never masks it + mixed = tl.load(ids_ptr + token).to(tl.int64) * tl.load(mult_ptr).to(tl.int64) + acc = tl.zeros([BLOCK_H], dtype=tl.int64) + + valid = 1 + for shift in tl.static_range(1, NGRAM): + column = CTX_LEN + local - shift # column in the (virtual) packed window + from_ids = column >= CTX_LEN + from_ctx = (column >= 0) & (column < CTX_LEN) + token_ids = tl.load(ids_ptr + (token - shift), mask=from_ids, other=0) + token_ctx = tl.load(ctx_ptr + req * CTX_LEN + column, mask=from_ctx, other=EOS) + raw = tl.where(from_ids, token_ids, token_ctx).to(tl.int64) + # the window may not cross a boundary token, and a boundary token is itself the wall + valid = valid * tl.where((column >= 0) & (raw != EOS), 1, 0) + mixed = mixed ^ ( + tl.where(valid == 1, raw, EOS) * tl.load(mult_ptr + shift).to(tl.int64) + ) + # after ``shift`` taps the (shift+1)-gram mix is complete; it owns one head block + ngram = shift + 1 + block = (head >= (ngram - 2) * HEADS_PER) & (head < (ngram - 1) * HEADS_PER) + acc = tl.where(block, mixed, acc) + + vocab = tl.load(vocab_ptr + head, mask=head_ok, other=1).to(tl.int64) + offset = tl.load(off_ptr + head, mask=head_ok, other=0).to(tl.int64) + # torch.remainder is floored, triton's % is truncated; the divisor is always positive + rem = acc % vocab + rem = tl.where(rem < 0, rem + vocab, rem) + tl.store(out_ptr + token * NUM_HEADS + head, rem + offset, mask=head_ok) + + +def ple_row_ids( + input_ids: torch.Tensor, + ngram_context: torch.Tensor, + req_index: torch.Tensor, + local_index: torch.Tensor, + multipliers: torch.Tensor, + vocab_sizes: torch.Tensor, + offsets: torch.Tensor, + *, + eos_token_id: int, + heads_per_ngram: int, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """``[T, num_heads]`` int64 global table rows for this forward's tokens. + + ``input_ids`` [T] and ``ngram_context`` [B, ngram_size-1] are int64 device tensors; + ``req_index`` / ``local_index`` are [T] int32 device tensors naming each token's request and + its offset within that request. The three hash constant tensors are int64 and on the same + device. ``out``, when given, is the destination (a CUDA graph replays into a fixed buffer). + """ + tokens = input_ids.numel() + ngram_size = int(multipliers.numel()) + num_heads = int(vocab_sizes.numel()) + ctx_len = int(ngram_context.shape[-1]) + # Checked as raises, not asserts: the kernel addresses the context row and the head + # blocks by these, and ``python -O`` must not turn a geometry mismatch into an OOB read. + if ctx_len != ngram_size - 1: + raise ValueError( + f"PLE hash: ngram_context has {ctx_len} context ids but ngram_size {ngram_size} " + f"needs {ngram_size - 1}" + ) + if num_heads != heads_per_ngram * (ngram_size - 1): + raise ValueError( + f"PLE hash: {num_heads} heads is not heads_per_ngram {heads_per_ngram} x " + f"{ngram_size - 1} n-gram orders" + ) + if out is None: + out = torch.empty((tokens, num_heads), dtype=torch.int64, device=input_ids.device) + if tokens == 0: + return out + _ple_row_ids_kernel[(tokens,)]( + input_ids, + ngram_context, + req_index, + local_index, + multipliers, + vocab_sizes, + offsets, + out, + EOS=int(eos_token_id), + CTX_LEN=ctx_len, + NGRAM=ngram_size, + HEADS_PER=int(heads_per_ngram), + NUM_HEADS=num_heads, + BLOCK_H=triton.next_power_of_2(num_heads), + num_warps=1, + ) + return out + + +__all__ = ["ple_row_ids"] diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index 5100229a9..2d1d66637 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/python/freetoken/models/qwen4_exp/ple.py @@ -19,6 +19,7 @@ from __future__ import annotations import math +import os from dataclasses import dataclass from typing import TYPE_CHECKING, List, Protocol, Sequence, Tuple @@ -42,6 +43,14 @@ _SPLITMIX_M1 = 0xBF58476D1CE4E5B9 _SPLITMIX_M2 = 0x94D049BB133111EB _PLE_LAYER_PRIME = 10007 +# Distinct (is_decode, shape, device) keys the n-gram row-id token index is memoized for. +_TOKEN_INDEX_CACHE_SIZE = 64 +_FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH" + + +def _fused_row_ids_enabled() -> bool: + """The fused hash kernel, on unless ``FREETOKEN_PLE_FUSED_HASH=0`` takes it back to torch ops.""" + return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in ("0", "false", "False") class PLETableBackend(Protocol): @@ -416,6 +425,7 @@ def __init__(self, args: Qwen4ExpArgs, table: PLETableBackend | None = None) -> self.ngram_heads_vocab_sizes = torch.empty(self.num_heads, dtype=torch.int64) self.ngram_heads_offsets = torch.empty(self.num_heads, dtype=torch.int64) self._table = table + self._token_index_cache: dict[tuple, Tuple[torch.Tensor, torch.Tensor]] = {} def attach_table(self, table: PLETableBackend) -> None: self._table = table @@ -462,8 +472,92 @@ def _shift_ignore_eos(self, packed: torch.Tensor) -> List[torch.Tensor]: shifted.append(torch.where(valid, gathered, packed.new_full((), self.eos_token_id))) return shifted - def row_ids(self, meta: PLEMetadata) -> torch.Tensor: - """Global table row per (token, hash head): ``[T, num_ngram_heads]`` int64.""" + def _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]: + """``(req[T], local[T])`` int32: each token's request, and its offset inside it. + + The fused kernel addresses the hash window through these instead of materializing the + ``[B, ctx+max_len]`` packed window. Memoized on the shape (which is all they depend on) + so a captured replay reads a stable address instead of re-running the build; a build + that happens DURING capture is not cached, since its buffers live in the graph pool. + """ + device = meta.input_ids.device + num_tokens = meta.input_ids.numel() + capturing = device.type == "cuda" and torch.cuda.is_current_stream_capturing() + # is_decode is part of the key, not just the shape: a decode of B requests and a + # prefill of ONE B-token request are the same [T] and mean opposite things (one token + # per request at offset 0 vs B offsets inside one request). + key = ( + meta.is_decode, + (num_tokens,) if meta.is_decode else tuple(meta.seq_lens), + str(device), + ) + cached = self._token_index_cache.get(key) + if cached is not None: + return cached + if meta.is_decode: # one token per request, each at offset 0 + index = ( + torch.arange(num_tokens, dtype=torch.int32, device=device), + torch.zeros(num_tokens, dtype=torch.int32, device=device), + ) + else: + cu = meta.cu_seqlens.long() + flat_pos = torch.arange(num_tokens, device=device) + req = (torch.searchsorted(cu, flat_pos, right=True) - 1).clamp_( + max=len(meta.seq_lens) - 1 + ) + index = ((req).to(torch.int32), (flat_pos - cu[req]).to(torch.int32)) + if not capturing: + if len(self._token_index_cache) >= _TOKEN_INDEX_CACHE_SIZE: + self._token_index_cache.pop(next(iter(self._token_index_cache))) + self._token_index_cache[key] = index + return index + + def _use_fused_row_ids(self, meta: PLEMetadata) -> bool: + device = meta.input_ids.device + if device.type != "cuda": + return False + if not _fused_row_ids_enabled(): + return False + return all( + t.device == device + for t in ( + meta.ngram_context, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + ) + ) + + def row_ids(self, meta: PLEMetadata, out: torch.Tensor | None = None) -> torch.Tensor: + """Global table row per (token, hash head): ``[T, num_ngram_heads]`` int64. + + One Triton program per token on CUDA; ``row_ids_reference`` is the same arithmetic in + torch ops and stays the oracle (and the CPU path). + """ + if self._use_fused_row_ids(meta): + from freetoken.kernel.triton.ple_hash import ple_row_ids + + req, local = self._token_index(meta) + return ple_row_ids( + meta.input_ids.long(), + meta.ngram_context, + req, + local, + self.layer_multipliers, + self.ngram_heads_vocab_sizes, + self.ngram_heads_offsets, + eos_token_id=self.eos_token_id, + heads_per_ngram=self.heads_per_ngram, + out=out, + ) + rows = self.row_ids_reference(meta) + if out is None: + return rows + out.copy_(rows) + return out + + def row_ids_reference(self, meta: PLEMetadata) -> torch.Tensor: + """Torch-op transcription of the hash; the oracle the fused kernel is diffed against.""" packed, select = self._window(meta) tokens = [select(s) for s in self._shift_ignore_eos(packed)] blocks = [] diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 0d2461427..74d4fd614 100644 --- a/tests/models/qwen4_exp/test_ple.py +++ b/tests/models/qwen4_exp/test_ple.py @@ -17,6 +17,7 @@ import pytest import torch +import freetoken.models.qwen4_exp.ple as ple_module from freetoken.models.config import ModelConfig from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.ple import ( @@ -173,6 +174,158 @@ def test_decode_hash_matches_prefill_hash(): assert torch.equal(got[i], prefill[i * len(sequences[0]) + step]) +_FUSED_HASH_SEQS = [[3, 4, EOS, 5, 6, 8], [2, EOS, 11, 12, 13, 14], [7]] +_FUSED_HASH_CTX = [[EOS, EOS], [21, 22], [31, EOS]] + + +def test_token_index_addresses_the_same_window_as_the_packed_build(): + """``_token_index`` names the (request, column) the packed window would have used.""" + layer = _make_layer(_config()) + embedding = layer.ple_embedding + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX) + packed, select = embedding._window(meta) + req, local = embedding._token_index(meta) + ctx_len = embedding.ngram_size - 1 + assert torch.equal(packed[req.long(), local.long() + ctx_len], meta.input_ids) + assert torch.equal(select(packed), meta.input_ids) + + +def test_token_index_cache_is_keyed_by_shape(): + """A repeat of the same shape reuses the buffers; a different one does not.""" + layer = _make_layer(_config()) + embedding = layer.ple_embedding + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX) + first = embedding._token_index(meta) + assert embedding._token_index(meta)[0] is first[0] + other = _meta([[3, 4]], [[EOS, EOS]]) + assert embedding._token_index(other)[0] is not first[0] + + +def test_token_index_does_not_confuse_a_decode_with_a_one_request_prefill(): + """Same [T], opposite meaning: B decode rows at offset 0 vs B offsets in one request.""" + layer = _make_layer(_config()) + embedding = layer.ple_embedding + tokens = [5, 6, 7] + context = [[EOS, EOS], [21, 22], [31, EOS]] + decode_req, decode_local = embedding._token_index( + _meta([[t] for t in tokens], context, decode=True) + ) + prefill_req, prefill_local = embedding._token_index(_meta([tokens], context[:1])) + assert decode_req.tolist() == [0, 1, 2] and decode_local.tolist() == [0, 0, 0] + assert prefill_req.tolist() == [0, 0, 0] and prefill_local.tolist() == [0, 1, 2] + + +def test_token_index_cache_is_bounded(): + """The memo is an LRU-ish ring, not an unbounded map keyed by every prefill shape.""" + embedding = _make_layer(_config()).ple_embedding + for length in range(4 * ple_module._TOKEN_INDEX_CACHE_SIZE): + embedding._token_index(_meta([list(range(length + 1))], [[EOS, EOS]])) + assert len(embedding._token_index_cache) == ple_module._TOKEN_INDEX_CACHE_SIZE + + +def test_fused_hash_is_off_on_cpu(): + """The Triton path needs CUDA; a CPU meta stays on the torch reference.""" + layer = _make_layer(_config()) + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX) + assert not layer.ple_embedding._use_fused_row_ids(meta) + assert torch.equal( + layer.ple_embedding.row_ids(meta), layer.ple_embedding.row_ids_reference(meta) + ) + + +@requires_cuda +@pytest.mark.parametrize("decode", [False, True]) +def test_fused_hash_matches_the_torch_reference(decode): + """The fused kernel is byte-identical to the torch-op hash, prefill and decode.""" + layer = _make_layer(_config(), device="cuda") + embedding = layer.ple_embedding + if decode: + meta = _meta( + [[s[0]] for s in _FUSED_HASH_SEQS], _FUSED_HASH_CTX, device="cuda", decode=True + ) + else: + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX, device="cuda") + assert embedding._use_fused_row_ids(meta) + assert torch.equal(embedding.row_ids(meta), embedding.row_ids_reference(meta)) + + +@requires_cuda +def test_fused_hash_captures_and_replays_in_a_cuda_graph(): + """Fixed shapes, no host reads: the hash is part of the captured decode step.""" + layer = _make_layer(_config(), device="cuda") + embedding = layer.ple_embedding + ids = torch.randint(0, VOCAB, (4,), dtype=torch.int64, device="cuda") + context = torch.randint(0, VOCAB, (4, embedding.ngram_size - 1), dtype=torch.int64, + device="cuda") + meta = PLEMetadata( + input_ids=ids, + cu_seqlens=torch.arange(5, dtype=torch.int32, device="cuda"), + seq_lens=(1,) * 4, + ngram_context=context, + state_slots=torch.arange(4, dtype=torch.int64, device="cuda"), + fresh_slots=None, + is_decode=True, + ) + out = torch.zeros(4, embedding.num_heads, dtype=torch.int64, device="cuda") + embedding.row_ids(meta, out) # prime the index cache and the JIT before capture + + graph = torch.cuda.CUDAGraph() + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + embedding.row_ids(meta, out) + torch.cuda.current_stream().wait_stream(side) + with torch.cuda.graph(graph): + embedding.row_ids(meta, out) + + for seed in range(3): + torch.manual_seed(seed) + ids.copy_(torch.randint(0, VOCAB, (4,), dtype=torch.int64, device="cuda")) + context.copy_( + torch.randint(0, VOCAB, context.shape, dtype=torch.int64, device="cuda") + ) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, embedding.row_ids_reference(meta)), seed + + +@requires_cuda +def test_fused_hash_env_switch_restores_the_torch_path(monkeypatch): + """FREETOKEN_PLE_FUSED_HASH=0 takes the hash back to the reference ops.""" + layer = _make_layer(_config(), device="cuda") + meta = _meta(_FUSED_HASH_SEQS, _FUSED_HASH_CTX, device="cuda") + monkeypatch.setenv("FREETOKEN_PLE_FUSED_HASH", "0") + assert not layer.ple_embedding._use_fused_row_ids(meta) + assert torch.equal( + layer.ple_embedding.row_ids(meta), layer.ple_embedding.row_ids_reference(meta) + ) + + +@pytest.mark.parametrize( + "ctx_len, heads_per_ngram, message", + [(1, 8, "context"), (2, 5, "heads")], +) +def test_the_fused_hash_refuses_a_geometry_it_cannot_address(ctx_len, heads_per_ngram, message): + """The kernel's block layout (``heads_per_ngram`` heads per n-gram order, ``ngram_size-1`` + context ids) is checked BEFORE anything is launched, as a sentence -- not as an ``assert`` + that ``python -O`` strips, leaving the kernel to read past the context row.""" + from freetoken.kernel.triton.ple_hash import ple_row_ids + + tokens = 3 + with pytest.raises(ValueError, match=message): + ple_row_ids( + torch.arange(tokens, dtype=torch.int64), + torch.zeros(1, ctx_len, dtype=torch.int64), + torch.zeros(tokens, dtype=torch.int32), + torch.arange(tokens, dtype=torch.int32), + torch.tensor([3, 5, 7], dtype=torch.int64), # ngram_size 3 + torch.full((16,), 11, dtype=torch.int64), # 16 heads + torch.arange(16, dtype=torch.int64) * 11, + eos_token_id=EOS, + heads_per_ngram=heads_per_ngram, + ) + + # -------------------------------------------------------------------------------------- # table backends # --------------------------------------------------------------------------------------