refactor(attention): move linear-attention metadata into a LinearAttnBackend

python/freetoken/attention/linear.py

@@ -1,9 +1,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List
 
 import torch
+from freetoken.attention.base import AttentionSpec, BaseAttnBackend
 
 if TYPE_CHECKING:
     from freetoken.core import Batch
@@ -20,7 +21,8 @@ class FLAMetadata:
     Fields:
       cu_seqlens          query indptr; decode = arange(bs+1) (1 token/req), prefill =
                           cumsum of extend_len. int32 on device.
-      cache_indices       per-request recurrent/conv state slot (= Req.table_idx). int32.
+      cache_indices       per-request recurrent/conv state slot (Req.linear_slot_idx, else
+                          Req.table_idx). int32.
       has_initial_state   prefill only: whether each request continues a cached prefix
                           (cached_len > 0). None for decode (state always present).
       fresh_state_indices prefill only: the state-pool slots whose sequence is fresh
@@ -49,11 +51,9 @@ def build_fla_metadata(batch: "Batch", device: torch.device) -> FLAMetadata:
     (the input_ids/attn-metadata pattern), so the copies overlap the forward instead of
     stalling it.
 
-    Decode is one token per request, so ``cu_seqlens`` is a plain ``arange(bs+1)`` and
-    ``cache_indices`` is ``batch.linear_table_idx`` (already int32) -- reused as-is. Under
-    CUDA graph the decode ``FLAMetadata`` is instead built directly in
-    ``GraphCaptureBuffer.set_batch`` against the persistent buffers (stable addresses); this
-    builder serves the eager scheduler path and direct-op test callers.
+    Decode is one token per request, so ``cu_seqlens`` is a plain ``arange(bs+1)``. Under
+    CUDA graph ``LinearAttnBackend`` copies this metadata into its persistent buffers
+    (stable addresses) before each replay.
     """
     reqs = batch.padded_reqs
     pin = {"device": "cpu", "pin_memory": torch.cuda.is_available()}
@@ -66,9 +66,8 @@ def gdn_slot(r):
     if batch.is_decode:
         bs = len(reqs)
         cu_seqlens = torch.arange(bs + 1, dtype=torch.int32, device=device)
-        # the scheduler stages linear_table_idx from gdn_slot (decode), reused as-is here
-        assert batch.linear_table_idx is not None
-        return FLAMetadata(cu_seqlens=cu_seqlens, cache_indices=batch.linear_table_idx)
+        idx_host = torch.tensor([gdn_slot(r) for r in reqs], dtype=torch.int32, **pin)
+        return FLAMetadata(cu_seqlens=cu_seqlens, cache_indices=idx_host.to(device, non_blocking=True))
 
     # prefill: cumsum of query (extend) lengths, per-request slot + continuation flags.
     lens = [r.extend_len for r in reqs]
@@ -137,4 +136,66 @@ def _build_track_metadata(reqs, cu_host, device, pin):
     )
 
 
-__all__ = ["FLAMetadata", "build_fla_metadata"]
+@dataclass
+class FLACaptureData:
+    """The static buffers the linear-attention layers' CUDA graphs read, sized for the largest captured batch."""
+
+    cu_seqlens: torch.Tensor
+    cache_indices: torch.Tensor
+
+    @classmethod
+    def create(cls, max_bs: int, device: torch.device, rows_per_req: int = 1) -> FLACaptureData:
+        return cls(
+            # constant under a graph: every request has exactly `rows_per_req` query rows
+            cu_seqlens=torch.arange(max_bs + 1, dtype=torch.int32, device=device) * rows_per_req,
+            cache_indices=torch.zeros(max_bs, dtype=torch.int32, device=device),
+        )
+
+
+class LinearAttnBackend(BaseAttnBackend):
+    """Metadata of the linear-attention layers (GatedDeltaNet, KDA), next to the attention backend: ``FLAMetadata`` per forward, ``FLACaptureData`` for CUDA graphs.
+
+    The layers run their own kernels and read ``batch.fla_metadata``, so ``forward`` is never called.
+    """
+
+    def __init__(self, device: torch.device) -> None:
+        self.device = device
+        self.capture: FLACaptureData | None = None
+        self.capture_bs: List[int] = []
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: "Batch",
+        attn_spec: AttentionSpec | None = None,
+    ) -> torch.Tensor:
+        # the GDN and KDA layers call their fla kernels themselves on batch.fla_metadata
+        raise NotImplementedError("linear-attention layers run their own kernels")
+
+    def prepare_metadata(self, batch: "Batch") -> None:
+        batch.fla_metadata = build_fla_metadata(batch, self.device)
+
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        assert self.capture is None, "Capture already initialized."
+        self.capture = FLACaptureData.create(max(bs_list), self.device, rows_per_req=rows_per_req)
+        self.capture_bs = sorted(bs_list)
+
+    def prepare_for_capture(self, batch: "Batch") -> None:
+        assert (bs := batch.padded_size) in self.capture_bs and self.capture
+        capture = self.capture
+        # capture on the padding request's state slot, so the captured gather/scatter touches scratch; replay writes the real slots
+        dummy = batch.padded_reqs[0]
+        capture.cache_indices[:bs].fill_(dummy.linear_slot_idx if dummy.linear_slot_idx is not None else dummy.table_idx)
+        batch.fla_metadata = FLAMetadata(cu_seqlens=capture.cu_seqlens[: bs + 1], cache_indices=capture.cache_indices[:bs])
+
+    def prepare_for_replay(self, batch: "Batch") -> None:
+        metadata, bs = batch.fla_metadata, batch.padded_size
+        assert self.capture is not None and bs in self.capture_bs
+        # cu_seqlens is the constant arange * rows_per_req under a graph (i.e. no-op)
+        self.capture.cache_indices[:bs].copy_(metadata.cache_indices)
+
+
+__all__ = ["FLACaptureData", "FLAMetadata", "LinearAttnBackend", "build_fla_metadata"]

python/freetoken/core.py

@@ -136,10 +136,6 @@ class Batch:
     # [3, n] t/h/w rope positions on mrope models; positions keeps its sequence-index meaning for token_pool / page_table
     mrope_positions: torch.Tensor | None = field(default=None, init=False)
     out_loc: torch.Tensor | None = field(init=False)
-    # Per-(padded-)request table_idx as a GPU int64 tensor, used by GatedDeltaNet
-    # decode to gather/scatter recurrent+conv state without host-side loops (so the
-    # decode step is CUDA-graph capturable). Set by the scheduler / graph buffer.
-    linear_table_idx: torch.Tensor | None = field(default=None, init=False)
     # Per-forward GatedDeltaNet metadata (cu_seqlens / cache_indices / continuation
     # flags), built once and shared by all GDN layers. Lazily built by the GDN op if
     # the scheduler/graph didn't set it.

python/freetoken/engine/engine.py

@@ -11,6 +11,7 @@
 
 import torch
 from freetoken.attention import AttnType, attention_backend_info, create_attention_backend
+from freetoken.attention.linear import LinearAttnBackend
 from freetoken.core import Batch, Context, Req, set_global_ctx
 from freetoken.distributed import destroy_distributed, enable_pynccl_distributed, set_tp_info
 from freetoken.gpu_select import gpu_identity
@@ -546,6 +547,7 @@ def __init__(self, config: EngineConfig):
         self.ctx.attn_backend = self.attn_backend = create_attention_backend(
             config.attention_backend, config.model_config
         )
+        self.linear_backend = LinearAttnBackend(self.device) if self.linear_state_pool is not None else None
 
         # ======================= Sampler initialization ========================
         self.sampler = Sampler(self.device, config.model_config.vocab_size)
@@ -564,7 +566,8 @@ def __init__(self, config: EngineConfig):
             cache_handle=None,  # type: ignore
         )
         # padded/dummy rows index the GDN padding slot (0) so gather/scatter hits scratch.
-        if self.linear_state_pool is not None:
+        # Naive GDN keys state by table_idx, where slot 0 is a live request and the dummy's own row is the scratch.
+        if self.linear_state_pool is not None and config.cache_type == "hybrid_radix":
             self.dummy_req.linear_slot_idx = self.linear_state_pool.padding_slot
         self.page_table[self.dummy_req.table_idx].fill_(num_tokens)  # point to dummy page
         self.graph_runner = GraphRunner(
@@ -580,6 +583,7 @@ def __init__(self, config: EngineConfig):
             moe_offload_cache=self.moe_offload_cache,
             mrope=config.model_config.model_is_mrope,
             out_dtype=torch.float32,
+            linear_backend=self.linear_backend,
         )
         if config.attention_backend.split(",")[0] == "triton":
             # Prefill runs on the first comma part; warm its autotune cache.
@@ -1223,6 +1227,8 @@ def rebuild_runtime_cache(
         self.rebuild_teardown_started = True
         # 1. Tear down CUDA graphs + backend capture scratch (free-before-alloc).
         self.attn_backend.reset_capture()
+        if self.linear_backend is not None:
+            self.linear_backend.reset_capture()
         self.graph_runner.destroy_cuda_graphs()
         # 2. Resize caches in place (each frees its old GPU tensors before allocating).
         # Pin the new window first (validated above) so any KV-pool rebuild below sizes the window
@@ -1267,6 +1273,7 @@ def rebuild_runtime_cache(
             moe_offload_cache=self.moe_offload_cache,
             mrope=config.model_config.model_is_mrope,
             out_dtype=torch.float32,
+            linear_backend=self.linear_backend,
         )
 
     def forward_logits(self, batch: Batch) -> torch.Tensor:

python/freetoken/engine/graph.py

@@ -16,6 +16,7 @@
 
 if TYPE_CHECKING:
     from freetoken.attention import BaseAttnBackend
+    from freetoken.attention.linear import LinearAttnBackend
     from freetoken.moe.offload_cache import OffloadMoeCache
 
 logger = init_logger(__name__)
@@ -28,9 +29,6 @@ class GraphCaptureBuffer:
     positions: torch.Tensor
     # [3, bs * rows_per_req] t/h/w rope positions; allocated only for mrope models (else None).
     mrope_positions: torch.Tensor | None
-    table_idx: torch.Tensor  # per-request slot id for GatedDeltaNet state gather/scatter
-    # GDN query indptr = arange(bs+1) * rows_per_req; a constant per captured bs, filled once.
-    fla_cu_seqlens: torch.Tensor
     rows_per_req: int = 1
 
     @classmethod
@@ -43,26 +41,17 @@ def init(cls, bs: int, device: torch.device, mrope: bool = False, rows_per_req:
             mrope_positions=(
                 torch.zeros(3, n, dtype=torch.int32, device=device) if mrope else None
             ),
-            table_idx=torch.zeros(bs, dtype=torch.int32, device=device),
-            fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device) * rows_per_req,
             rows_per_req=rows_per_req,
         )
 
     def set_batch(self, batch: Batch) -> None:
-        from freetoken.attention.linear import FLAMetadata
-
         bs = batch.padded_size
         _rows = slice(bs * self.rows_per_req)
         batch.input_ids = self.input_ids[_rows]
         batch.out_loc = self.out_loc[_rows]
         batch.positions = self.positions[_rows]
         if self.mrope_positions is not None:
             batch.mrope_positions = self.mrope_positions[:, _rows]
-        batch.linear_table_idx = self.table_idx[:bs]
-        # GDN metadata reads the persistent cu_seqlens (constant arange * rows_per_req) and table_idx slot map, so the captured kernels see stable addresses.
-        batch.fla_metadata = FLAMetadata(
-            cu_seqlens=self.fla_cu_seqlens[: bs + 1], cache_indices=self.table_idx[:bs]
-        )
 
     def copy_from(self, batch: Batch) -> None:
         _rows = slice(batch.padded_size * self.rows_per_req)
@@ -72,8 +61,6 @@ def copy_from(self, batch: Batch) -> None:
         self.positions[_rows] = batch.positions
         if self.mrope_positions is not None:
             self.mrope_positions[:, _rows] = batch.mrope_positions
-        if batch.linear_table_idx is not None:
-            self.table_idx[: batch.padded_size] = batch.linear_table_idx
 
 
 def _determine_cuda_graph_bs(
@@ -149,6 +136,7 @@ def __init__(
         batch_hook: Callable[[Batch], None] | None = None,
         pool: Any = None,
         out_dtype: torch.dtype | None = None,
+        linear_backend: LinearAttnBackend | None = None,
     ) -> None:
         cuda_graph_bs = _determine_cuda_graph_bs(
             cuda_graph_bs=cuda_graph_bs,
@@ -159,6 +147,7 @@ def __init__(
         assert rows_per_req == 1 or dummy_req.extend_len == rows_per_req, (rows_per_req, dummy_req.extend_len)
         self.forward = forward
         self.attn_backend = attn_backend
+        self.linear_backend = linear_backend
         self.max_graph_bs = max(cuda_graph_bs) if cuda_graph_bs else 0
         self.graph_bs_list = sorted(cuda_graph_bs)
         self.dummy_req = dummy_req
@@ -204,6 +193,8 @@ def _capture_graphs(self, max_seq_len: int):
             return logger.info_rank0("CUDA graph is disabled.")
 
         self.attn_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list, rows_per_req=self.rows_per_req)
+        if self.linear_backend is not None:
+            self.linear_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list, rows_per_req=self.rows_per_req)
 
         torch.cuda.synchronize(self.device)
         torch.cuda.empty_cache()
@@ -234,13 +225,8 @@ def _capture_graphs(self, max_seq_len: int):
                 self.batch_hook(batch)
             self.attn_backend.prepare_for_capture(batch)
             self.buffer.set_batch(batch)
-            # capture on the dummy linear-state slot so GatedDeltaNet gather/scatter
-            # touches scratch (real slot indices are written by copy_from on replay). Hybrid-
-            # radix decouples the GDN slot from table_idx -> use the GDN padding slot.
-            dummy_slot = (self.dummy_req.linear_slot_idx
-                          if self.dummy_req.linear_slot_idx is not None
-                          else self.dummy_req.table_idx)
-            self.buffer.table_idx[:bs].fill_(dummy_slot)
+            if self.linear_backend is not None:
+                self.linear_backend.prepare_for_capture(batch)
             with get_global_ctx().forward_batch(batch, **self.ctx_override):
                 self._store_output(bs, self.forward())
                 # Keep the offload cache warmed for capture. Resetting here forces
@@ -271,6 +257,8 @@ def replay(self, batch: Batch) -> Any:
         padded = batch.padded_size
         g = self.graph_map[padded]
         self.attn_backend.prepare_for_replay(batch)
+        if self.linear_backend is not None:
+            self.linear_backend.prepare_for_replay(batch)
         g.replay()
         outs = tuple(
             buf[: rows // padded * batch.size] for buf, rows in zip(self._out, self._out_rows[padded], strict=True)

python/freetoken/models/qwen4_exp/ple.py

@@ -319,10 +319,10 @@ def build_ple_metadata(
 
     The n-gram context is per-request device state (``ple_ngram_ctx`` [num_slots, ngram_size-1],
     rolled forward once per forward by ``commit_ngram_context``), so it never lags the sampled
-    token under overlap scheduling and follows the slot on COW/snapshot. A decode batch reads it
-    straight off the persistent ``linear_table_idx`` buffer, so the build is capture-safe and
-    sync-free. Reuses ``batch.fla_metadata`` (slots / indptr / fresh mask) when the scheduler
-    built it.
+    token under overlap scheduling and follows the slot on COW/snapshot. A decode batch takes its
+    slots straight off ``batch.fla_metadata.cache_indices`` (persistent under a graph), so the
+    build is capture-safe and sync-free. Reuses ``batch.fla_metadata`` (slots / indptr / fresh
+    mask) when the scheduler built it.
     """
     reqs = batch.padded_reqs
     ctx_len = args.ngram_size - 1
@@ -333,14 +333,13 @@ def build_ple_metadata(
         f"ple_ngram_ctx holds {context_pool.shape[-1]} ids, config wants {ctx_len}"
     )
     fla = getattr(batch, "fla_metadata", None)
-    slots_dev = getattr(batch, "linear_table_idx", None)
     # image rows hash as the placeholder token, never as the content pad ids
     input_ids = batch.input_ids
     if args.image_token_id is not None:
         input_ids = restore_placeholder(input_ids, args.image_token_id)
 
-    if batch.is_decode and slots_dev is not None:
-        slots = slots_dev.long()
+    if batch.is_decode and fla is not None:
+        slots = fla.cache_indices.long()
         bs = slots.numel()
         return PLEMetadata(
             input_ids=input_ids,

python/freetoken/scheduler/scheduler.py

@@ -6,7 +6,6 @@
 from typing import TYPE_CHECKING, List, NamedTuple, NoReturn, Set, Tuple, TypeAlias
 
 import torch
-from freetoken.attention.linear import build_fla_metadata
 from freetoken.core import Batch, Req
 from freetoken.env import ENV
 from freetoken.gpu_select import gpu_identity
@@ -864,28 +863,8 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput:
         input_mapping = _make_input_tuple(batch, self.device)
         write_mapping = _make_write_tuple(batch, self.device)
         batch.out_loc = self.engine.page_table[input_mapping]
-        if self.engine.linear_state_pool is not None:
-            if not batch.is_prefill:
-                # GPU GDN-state slot (one per padded request) for the decode gather/scatter;
-                # lands in the CUDA-graph input buffer via copy_from. Gate on the cache mode,
-                # NOT on whether any padded req has a linear_slot_idx -- the persistent dummy
-                # req always carries one (= padding_slot), so that test is True even for naive
-                # and would collapse all real naive reqs onto the padding slot. Hybrid: build
-                # per padded req from Req.linear_slot_idx (dummy -> padding_slot). Naive: keep
-                # the old keying = input_mapping's table_idx column (already staged, no H2D).
-                if self.cache_manager.is_hybrid:
-                    pool = self.engine.linear_state_pool
-                    slots = [r.linear_slot_idx if r.linear_slot_idx is not None
-                             else pool.padding_slot for r in batch.padded_reqs]
-                    batch.linear_table_idx = torch.tensor(
-                        slots, dtype=torch.int32, device="cpu", pin_memory=True
-                    ).to(self.device, non_blocking=True)
-                else:
-                    batch.linear_table_idx = input_mapping[0][:: batch.rows_per_req].to(torch.int32)
-            # Per-forward GDN metadata (cu_seqlens / cache_indices / continuation flags),
-            # built once here instead of rebuilt in each of the 30 GDN layers. For decode
-            # under CUDA graph the persistent cu_seqlens buffer is supplied by set_batch.
-            batch.fla_metadata = build_fla_metadata(batch, self.device)
+        if self.engine.linear_backend is not None:
+            self.engine.linear_backend.prepare_metadata(batch)
         if not batch.is_prefill:
             # This batch's padded per-row page-table rows. Backends that snapshot the table for
             # a captured replay (DSV4) read them in prepare_metadata / prepare_for_replay.

tests/models/qwen4_exp/test_gdn.py

@@ -105,9 +105,6 @@ def _prefill(op, ctx: Context, lengths: list[int], seed: int):
 def _decode(op, ctx: Context, reqs, hidden: torch.Tensor) -> torch.Tensor:
     batch = Batch(reqs=reqs, phase="decode")
     batch.padded_reqs = reqs
-    batch.linear_table_idx = torch.tensor(
-        [r.table_idx for r in reqs], dtype=torch.int32, device=DEV
-    )
     with ctx.forward_batch(batch):
         return op.forward(hidden)