diff --git a/python/freetoken/attention/linear.py b/python/freetoken/attention/linear.py index f3717e58e7..1c4c34a6af 100644 --- a/python/freetoken/attention/linear.py +++ b/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"] diff --git a/python/freetoken/core.py b/python/freetoken/core.py index f37ef84f89..a7e41dd9ae 100644 --- a/python/freetoken/core.py +++ b/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. diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index eca7bddd69..121e1070e8 100644 --- a/python/freetoken/engine/engine.py +++ b/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: diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index 42306918be..1282a15712 100644 --- a/python/freetoken/engine/graph.py +++ b/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,14 +41,10 @@ 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] @@ -58,11 +52,6 @@ def set_batch(self, batch: Batch) -> None: 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) diff --git a/python/freetoken/models/qwen4_exp/ple.py b/python/freetoken/models/qwen4_exp/ple.py index c944c4946f..5a3b321f08 100644 --- a/python/freetoken/models/qwen4_exp/ple.py +++ b/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, diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 21ef19d9f8..452f8f225e 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/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. diff --git a/tests/models/qwen4_exp/test_gdn.py b/tests/models/qwen4_exp/test_gdn.py index 81dd0e7673..a854626a80 100644 --- a/tests/models/qwen4_exp/test_gdn.py +++ b/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) diff --git a/tests/models/qwen4_exp/test_ple.py b/tests/models/qwen4_exp/test_ple.py index 0d2461427e..0180e4b415 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 +from freetoken.attention.linear import FLAMetadata from freetoken.models.config import ModelConfig from freetoken.models.qwen4_exp.config import parse_config from freetoken.models.qwen4_exp.ple import ( @@ -561,8 +562,9 @@ def _fake_batch(reqs, *, decode, input_ids, positions=None, table_idx=None, devi is_prefill=not decode, input_ids=torch.tensor(input_ids, dtype=torch.int64, device=device), positions=None if positions is None else torch.tensor(positions, dtype=torch.int32, device=device), - linear_table_idx=( - None if table_idx is None else torch.tensor(table_idx, dtype=torch.int32, device=device) + fla_metadata=None if table_idx is None else FLAMetadata( + cu_seqlens=torch.arange(len(table_idx) + 1, dtype=torch.int32, device=device), + cache_indices=torch.tensor(table_idx, dtype=torch.int32, device=device), ), ) @@ -648,7 +650,8 @@ def test_decode_graph_replay_matches_eager(): batch = SimpleNamespace( padded_reqs=[None] * bs, is_decode=True, is_prefill=False, input_ids=torch.randint(0, VOCAB, (bs,), device="cuda", dtype=torch.int32), - positions=positions, linear_table_idx=slots, + positions=positions, + fla_metadata=FLAMetadata(cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device="cuda"), cache_indices=slots), ) R = torch.randn(bs, args.ple_state_width, device="cuda", dtype=torch.bfloat16) states0 = torch.randn(bs + 1, args.ple_state_width, args.ple_conv_state_len,