feat: port upstream fixes (lm_head rows, GDN conv sync, PLE hash, oracle 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

Author: jomcgi

python/freetoken/attention/linear.py

@@ -26,12 +26,17 @@ class FLAMetadata:
       fresh_state_indices prefill only: the state-pool slots whose sequence is fresh
                           (cached_len == 0) and must be zeroed before the chunk kernel
                           reads them in place. None if there are none / for decode.
+      max_seq_len         prefill only: the longest extend_len as a host int. It reaches
+                          only the Triton fallback, where it sizes the varlen convolution
+                          launch without a device-to-host sync, and is ignored when
+                          sgl_kernel is installed.
     """
 
     cu_seqlens: torch.Tensor
     cache_indices: torch.Tensor
     has_initial_state: torch.Tensor | None = None
     fresh_state_indices: torch.Tensor | None = None
+    max_seq_len: int | None = None
 
     # --- hybrid-radix track-checkpoint (extra_buffer) fields; all None when not caching ---
     # For each request crossing a chunk-aligned (×CHUNK) boundary this forward, snapshot its
@@ -87,6 +92,7 @@ def gdn_slot(r):
         fresh_state_indices=(
             fresh_host.to(device, non_blocking=True) if fresh_host is not None else None
         ),
+        max_seq_len=max(lens),
         **track,
     )
 

python/freetoken/kernel/backend.py

@@ -31,6 +31,12 @@ def is_sgl_kernel_installed() -> bool:
     return _importable("sgl_kernel")
 
 
+@functools.cache
+def is_triton_installed() -> bool:
+    """Whether the Triton runtime used by ``freetoken.kernel.triton`` is available."""
+    return _importable("triton")
+
+
 @functools.cache
 def is_triton_kernels_installed() -> bool:
     """OpenAI's ``triton_kernels`` (the fused MoE router used by ``moe.fused.fused_topk``).

python/freetoken/kernel/causal_conv1d.py

@@ -22,9 +22,14 @@ def causal_conv1d_varlen(
     cu_seqlens: torch.Tensor,   # [batch+1] int32 prefix sums of per-request lengths
     cache_indices: torch.Tensor,    # [batch] int32 slot id per request
     has_initial_state: torch.Tensor,  # [batch] bool (carry conv state across chunks)
+    max_seq_len: int | None = None,  # host-known longest extend_len
 ) -> torch.Tensor:
     """Varlen (prefill) depthwise causal conv with silu; writes silu(conv) into ``x``
-    in place and refreshes ``conv_states[cache_indices]`` with each request's tail."""
+    in place and refreshes ``conv_states[cache_indices]`` with each request's tail.
+
+    ``max_seq_len`` only sizes the Triton launch grid. Passing the scheduler's host value
+    avoids the device-to-host sync otherwise needed to derive it from ``cu_seqlens``.
+    """
     from freetoken.kernel.backend import is_sgl_kernel_installed
 
     if not is_sgl_kernel_installed():
@@ -33,7 +38,13 @@ def causal_conv1d_varlen(
         )
 
         return triton_causal_conv1d_varlen(
-            x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state
+            x,
+            weight,
+            conv_states,
+            cu_seqlens,
+            cache_indices,
+            has_initial_state,
+            max_seq_len=max_seq_len,
         )
 
     from sgl_kernel import causal_conv1d_fwd

python/freetoken/kernel/triton/ple_hash.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Fused n-gram hash to PLE table row ids.
+
+The eager hash builds a packed ragged window, scans boundaries, and launches many
+small elementwise operations. This kernel performs the same signed int64 arithmetic
+with one Triton program per token and without materializing the packed window.
+"""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _ple_row_ids_kernel(
+    ids_ptr,
+    ctx_ptr,
+    req_ptr,
+    local_ptr,
+    mult_ptr,
+    vocab_ptr,
+    off_ptr,
+    out_ptr,
+    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
+    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
+        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)
+        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)
+        )
+        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)
+    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,
+) -> torch.Tensor:
+    """Return ``[tokens, heads]`` int64 global PLE table row ids."""
+    tokens = input_ids.numel()
+    ngram_size = int(multipliers.numel())
+    num_heads = int(vocab_sizes.numel())
+    ctx_len = int(ngram_context.shape[-1])
+    if ctx_len != ngram_size - 1:
+        raise ValueError(
+            f"PLE hash context has {ctx_len} 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 has {num_heads} heads, expected heads_per_ngram "
+            f"{heads_per_ngram} x {ngram_size - 1} n-gram orders"
+        )
+    out = torch.empty(
+        (tokens, num_heads), dtype=torch.int64, device=input_ids.device
+    )
+    if tokens == 0:
+        return out
+    assert ngram_context.is_contiguous()
+    assert out.is_contiguous()
+    _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/glm5_next/linear.py

@@ -102,6 +102,7 @@ def _conv_prefill(self, conv_in, pool, fla):
             fla.cu_seqlens,
             fla.cache_indices,
             fla.has_initial_state,
+            max_seq_len=fla.max_seq_len,
         )
         return out.transpose(0, 1)
 

python/freetoken/models/qwen3_5_moe/gdn.py

@@ -113,14 +113,19 @@ def _gate_params(self, a: torch.Tensor, b: torch.Tensor):
     def _conv_weight(self) -> torch.Tensor:
         return self.conv1d.weight.squeeze(1)  # [conv_dim, kernel] for the fused kernel
 
-    def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor:
+    def _conv_prefill(
+        self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state, max_seq_len=None
+    ) -> torch.Tensor:
         """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's
         conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim].
-        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata."""
+        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from
+        ``FLAMetadata``. Its host-known ``max_seq_len`` lets the Triton fallback size
+        its launch without a sync."""
         li = pool.local_index(self.layer_id)
         x = conv_in.transpose(0, 1).contiguous()  # [conv_dim, total]
         out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li],
-                                   cu_seqlens, cache_indices, has_initial_state)
+                                   cu_seqlens, cache_indices, has_initial_state,
+                                   max_seq_len=max_seq_len)
         return out.transpose(0, 1)  # [total, conv_dim]
 
     def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor:
@@ -189,7 +194,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             )
         else:
             mixed = self._conv_prefill(
-                conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
+                conv_in, pool, fla.cu_seqlens, fla.cache_indices,
+                fla.has_initial_state, fla.max_seq_len)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
             qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
             q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)

python/freetoken/models/qwen4_exp/gdn.py

@@ -122,14 +122,19 @@ def _gate_params(self, a: torch.Tensor, b: torch.Tensor):
     def _conv_weight(self) -> torch.Tensor:
         return self.conv1d.weight.squeeze(1)  # [conv_dim, kernel] for the fused kernel
 
-    def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor:
+    def _conv_prefill(
+        self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state, max_seq_len=None
+    ) -> torch.Tensor:
         """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's
         conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim].
-        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata."""
+        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from
+        ``FLAMetadata``. Its host-known ``max_seq_len`` lets the Triton fallback size
+        its launch without a sync."""
         li = pool.local_index(self.layer_id)
         x = conv_in.transpose(0, 1).contiguous()  # [conv_dim, total]
         out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li],
-                                   cu_seqlens, cache_indices, has_initial_state)
+                                   cu_seqlens, cache_indices, has_initial_state,
+                                   max_seq_len=max_seq_len)
         return out.transpose(0, 1)  # [total, conv_dim]
 
     def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor:
@@ -228,7 +233,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
                 )
         else:
             mixed = self._conv_prefill(
-                conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
+                conv_in, pool, fla.cu_seqlens, fla.cache_indices,
+                fla.has_initial_state, fla.max_seq_len)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
             qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
             q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)

python/freetoken/models/qwen4_exp/ple.py

@@ -53,10 +53,21 @@
 _SPLITMIX_M2 = 0x94D049BB133111EB
 _PLE_LAYER_PRIME = 10007
 _PLE_CACHE_SLAB_ROWS = 1 << 16
+_TOKEN_INDEX_CACHE_SIZE = 64
+_FUSED_HASH_ENV = "FREETOKEN_PLE_FUSED_HASH"
 
 logger = init_logger(__name__)
 
 
+def _fused_row_ids_enabled() -> bool:
+    """Enable the fused hash unless the environment explicitly disables it."""
+    return (os.getenv(_FUSED_HASH_ENV) or "1").strip() not in (
+        "0",
+        "false",
+        "False",
+    )
+
+
 class _IOVec(ctypes.Structure):
     _fields_ = [("base", ctypes.c_void_p), ("length", ctypes.c_size_t)]
 
@@ -2644,6 +2655,7 @@ def __init__(self, args: Qwen4ExpArgs, table: PLETableBackend | None = None) ->
         self._table = table
         self._host_hash_constants: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
         self._host_hash_buffers: tuple[torch.Tensor, ...] | None = None
+        self._token_index_cache: dict[tuple, Tuple[torch.Tensor, torch.Tensor]] = {}
 
     def attach_table(self, table: PLETableBackend) -> None:
         self._table = table
@@ -2775,8 +2787,84 @@ 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 _token_index(self, meta: PLEMetadata) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return each token's request index and request-local offset."""
+        device = meta.input_ids.device
+        num_tokens = meta.input_ids.numel()
+        capturing = device.type == "cuda" and torch.cuda.is_current_stream_capturing()
+        key = (
+            meta.is_decode,
+            (num_tokens,) if meta.is_decode else tuple(meta.seq_lens),
+            str(device),
+        )
+        # A graph must own its index allocations. Reusing an eager memo entry here
+        # lets normal cache eviction free storage whose pointers were baked into the
+        # captured Triton launch.
+        cached = None if capturing else self._token_index_cache.get(key)
+        if cached is not None:
+            return cached
+        if meta.is_decode:
+            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:
+        from freetoken.kernel.backend import is_triton_installed
+
+        device = meta.input_ids.device
+        if (
+            device.type != "cuda"
+            or not is_triton_installed()
+            or not _fused_row_ids_enabled()
+        ):
+            return False
+        return all(
+            tensor.device == device
+            for tensor in (
+                meta.ngram_context,
+                self.layer_multipliers,
+                self.ngram_heads_vocab_sizes,
+                self.ngram_heads_offsets,
+            )
+        )
+
     def row_ids(self, meta: PLEMetadata) -> torch.Tensor:
         """Global table row per (token, hash head): ``[T, num_ngram_heads]`` int64."""
+        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,
+            )
+        return self.row_ids_reference(meta)
+
+    def row_ids_reference(self, meta: PLEMetadata) -> torch.Tensor:
+        """Torch implementation retained as the CPU path and fused-kernel oracle."""
         packed, select = self._window(meta)
         tokens = [select(s) for s in self._shift_ignore_eos(packed)]
         blocks = []