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

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"]

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 = []