refactor: add engine and scheduler seams for multi-row decode steps

python/freetoken/attention/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Protocol
+from typing import TYPE_CHECKING, Literal, Protocol
 
 from freetoken.utils import Registry, init_logger
 
@@ -33,6 +33,8 @@ class BackendInfo:
     # Whether forward() honors a per-call AttentionSpec (window/sm_scale/sinks).
     # Non-consumers raise on a non-None spec instead of silently dropping it.
     consumes_attn_spec: bool = False
+    # How the backend captures a graph of `rows_per_req` query rows per request: "device" keeps all metadata on device, "host_plan" plans on the host before each replay, "none" cannot.
+    uniform_extend_graph: Literal["none", "host_plan", "device"] = "none"
 
 
 SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend")
@@ -58,6 +60,7 @@ def create_trtllm_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL}),
         requires_flashinfer=True,
+        uniform_extend_graph="host_plan",
     ),
 )
 def create_fi_backend(config: ModelConfig):
@@ -71,6 +74,7 @@ def create_fi_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL}),
         requires_sgl_kernel=True,
+        uniform_extend_graph="device",
     ),
 )
 def create_fa_backend(config: ModelConfig):
@@ -84,6 +88,7 @@ def create_fa_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL, AttnType.SWA}),
         consumes_attn_spec=True,
+        uniform_extend_graph="device",
     ),
 )
 def create_triton_backend(config: ModelConfig):

python/freetoken/attention/base.py

@@ -70,7 +70,8 @@ def forward(
     def prepare_metadata(self, batch: Batch) -> None: ...
 
     @abstractmethod
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None: ...
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        """Static buffers for graphs of ``rows_per_req`` query rows per request: 1 for decode, more for a uniform extend."""
 
     @abstractmethod
     def prepare_for_capture(self, batch: Batch) -> None: ...
@@ -116,8 +117,8 @@ def prepare_metadata(self, batch: Batch) -> None:
         backend = self.prefill_backend if batch.is_prefill else self.decode_backend
         return backend.prepare_metadata(batch)
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
-        self.decode_backend.init_capture_graph(max_seq_len, bs_list)
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        self.decode_backend.init_capture_graph(max_seq_len, bs_list, rows_per_req)
 
     def prepare_for_capture(self, batch: Batch) -> None:
         self.decode_backend.prepare_for_capture(batch)

python/freetoken/attention/dsa.py

@@ -341,7 +341,9 @@ def _prefill(self, md, layer_id, q_nope, q_pe, batch, inputs) -> torch.Tensor:
         return o
 
     # ----- CUDA graph (decode) ----------------------------------------------------------
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        if rows_per_req != 1:
+            raise NotImplementedError("dsa captures decode graphs only")
         self.max_seq_len = max_seq_len
         self.capture_bs = sorted(bs_list)
         max_bs = max(bs_list)

python/freetoken/attention/dsa_indexer_kpool.py

@@ -56,8 +56,8 @@ def __init__(self, config: "ModelConfig") -> None:
     # ----- CUDA-graph decode staging ------------------------------------------------------
     # The tail rings are keyed by Req.table_idx; a captured decode step must read it
     # from a static buffer restaged per replay (rows/kvlen precedent in the parent).
-    def init_capture_graph(self, max_seq_len: int, bs_list) -> None:
-        super().init_capture_graph(max_seq_len, bs_list)
+    def init_capture_graph(self, max_seq_len: int, bs_list, rows_per_req: int = 1) -> None:
+        super().init_capture_graph(max_seq_len, bs_list, rows_per_req)
         self._slots_buf = torch.zeros(max(bs_list), dtype=torch.int64, device=self.device)
 
     def _stage_decode(self, batch: "Batch", bs: int, table_idx: torch.Tensor) -> None:

python/freetoken/attention/dsv41_sparse.py

@@ -212,7 +212,9 @@ def _prefill_metadata(self, batch: Batch) -> DSV41AttnMetadata:
         decoder_rows = torch.cat(rows).to(self.device, non_blocking=True) if rows else torch.empty(0, dtype=torch.int64, device=self.device)
         return DSV41AttnMetadata(last_indices=last_indices, segments=segments, decoder_segments=dec, decoder_rows=decoder_rows, decoder_pad=pad)
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        if rows_per_req != 1:
+            raise NotImplementedError("dsv41_sparse captures decode graphs only")
         assert self.capture is None, "Capture already initialized."
         self.max_graph_bs = max(bs_list)
         self.capture = DSV41CaptureData.create(self.max_graph_bs, max_seq_len, self.device)

python/freetoken/attention/dsv4_sparse.py

@@ -190,7 +190,9 @@ def prepare_metadata(self, batch: Batch) -> None:
             last_indices=last, table_rows=self._table_rows(batch), window_ar=self._window_ar
         )
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        if rows_per_req != 1:
+            raise NotImplementedError("dsv4_sparse captures decode graphs only")
         assert self.capture is None, "Capture already initialized."
         self.max_graph_bs = max(bs_list)
         self.capture = DSV4CaptureData.create(self.max_graph_bs, max_seq_len, self.device)

python/freetoken/attention/fa.py

@@ -46,6 +46,7 @@ def __init__(self, config: ModelConfig):
         self.kvcache = ctx.kv_cache
         self.page_size = ctx.page_size
         self.capture: FACaptureData | None = None
+        self.capture_rows_per_req = 1
         self.max_graph_bs = 0
         self.capture_bs: List[int] = []
         self.scale = config.head_dim**-0.5
@@ -120,23 +121,26 @@ def prepare_metadata(self, batch: Batch) -> None:
             page_table=new_page_table,
         )
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    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."
         max_bs = max(bs_list)
-        capture = FACaptureData.create(max_bs, max_seq_len // self.page_size, self.kvcache.device)
+        capture = FACaptureData.create(max_bs, max_seq_len // self.page_size, self.kvcache.device, rows_per_req=rows_per_req)
         self.max_graph_bs = max_bs
         self.capture = capture
         self.capture_bs = sorted(bs_list)
+        self.capture_rows_per_req = rows_per_req
 
     def prepare_for_capture(self, batch: Batch) -> None:
         assert (bs := batch.size) in self.capture_bs and self.capture
         capture = self.capture
+        rows_per_req = self.capture_rows_per_req
+        capture.seq_lens[:bs].fill_(rows_per_req)
         metadata = FAMetadata(
             cu_seqlens_k=capture.cu_seqlens_k[: bs + 1],
             cu_seqlens_q=capture.cu_seqlens_q[: bs + 1],
             cache_seqlens=capture.seq_lens[:bs],
             max_seqlen_k=capture.page_table.size(1) * self.page_size,
-            max_seqlen_q=1,  # decode only
+            max_seqlen_q=rows_per_req,
             page_table=capture.page_table[:bs, :],
         )
         batch.attn_metadata = metadata
@@ -145,7 +149,7 @@ def prepare_for_replay(self, batch: Batch) -> None:
         metadata, bs = batch.attn_metadata, batch.padded_size
         assert isinstance(metadata, FAMetadata)
         assert self.capture is not None and bs in self.capture_bs
-        # cu_seqlens_q is always [0, 1, 2, ..., bs] for decode (i.e. no-op)
+        # cu_seqlens_q is the constant arange * rows_per_req under a graph (i.e. no-op)
         table_len = metadata.page_table.size(1)
         self.capture.cu_seqlens_k[: bs + 1].copy_(metadata.cu_seqlens_k)
         self.capture.seq_lens[:bs].copy_(metadata.cache_seqlens)

python/freetoken/attention/fi.py

@@ -139,8 +139,11 @@ def __init__(self, config: ModelConfig) -> None:
         # for cuda graph
         self.capture_bs: List[int] = []
         self.max_graph_bs = 0
-        self.graph_wrappers: Dict[int, CUDAGraphBatchDecodeWithPagedKVCacheWrapper] = {}
+        self.graph_wrappers: Dict[
+            int, CUDAGraphBatchDecodeWithPagedKVCacheWrapper | BatchPrefillWithPagedKVCacheWrapper
+        ] = {}
         self.capture: FICaptureData | None = None
+        self.capture_rows_per_req = 1
         self.last_event = torch.cuda.Event()
         self.last_event.record()
 
@@ -266,14 +269,15 @@ def reset_capture(self) -> None:
         super().reset_capture()
         self.graph_wrappers = {}
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    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."
         max_bs = max(bs_list)
-        capture = FICaptureData.create(max_bs, max_seq_len, self.kvcache.device)
+        capture = FICaptureData.create(max_bs, max_seq_len, self.kvcache.device, rows_per_req=rows_per_req)
         capture.page_table = capture.page_table.view(-1)  # use 1D as ragged indices
         self.max_graph_bs = max_bs
         self.capture = capture
         self.capture_bs = sorted(bs_list)
+        self.capture_rows_per_req = rows_per_req
 
     @cached_property
     def use_tensor_cores(self) -> bool:
@@ -283,22 +287,41 @@ def use_tensor_cores(self) -> bool:
         GQA = self.config.num_qo_heads // self.config.num_kv_heads
         return GQA >= 4
 
-    def prepare_for_capture(self, batch: Batch) -> None:
-        from flashinfer import CUDAGraphBatchDecodeWithPagedKVCacheWrapper
+    def _make_graph_wrapper(
+        self, bs: int
+    ) -> CUDAGraphBatchDecodeWithPagedKVCacheWrapper | BatchPrefillWithPagedKVCacheWrapper:
+        from flashinfer import BatchPrefillWithPagedKVCacheWrapper, CUDAGraphBatchDecodeWithPagedKVCacheWrapper
+
+        capture = self.capture
+        assert capture is not None
+        if self.capture_rows_per_req > 1:
+            wrapper = BatchPrefillWithPagedKVCacheWrapper(
+                self.float_workspace_buffer,
+                kv_layout="NHD",
+                backend="fa2",
+                use_cuda_graph=True,
+                qo_indptr_buf=capture.cu_seqlens_q[: bs + 1],
+                paged_kv_indptr_buf=capture.cu_seqlens_k[: bs + 1],
+                paged_kv_indices_buf=capture.indices,
+                paged_kv_last_page_len_buf=capture.one_tensor[:bs],
+            )
+        else:
+            wrapper = CUDAGraphBatchDecodeWithPagedKVCacheWrapper(
+                self.float_workspace_buffer,
+                kv_layout="NHD",
+                use_tensor_cores=self.use_tensor_cores,
+                indptr_buffer=capture.cu_seqlens_k[: bs + 1],
+                indices_buffer=capture.indices,
+                last_page_len_buffer=capture.one_tensor[:bs],
+            )
+            wrapper._backend = "fa2"
+        wrapper._int_workspace_buffer = self.int_workspace_buffer
+        return wrapper
 
+    def prepare_for_capture(self, batch: Batch) -> None:
         bs = batch.size
         assert bs in self.capture_bs and bs not in self.graph_wrappers and self.capture
-        capture = self.capture
-        self.graph_wrappers[bs] = CUDAGraphBatchDecodeWithPagedKVCacheWrapper(
-            self.float_workspace_buffer,
-            kv_layout="NHD",
-            use_tensor_cores=self.use_tensor_cores,
-            indptr_buffer=capture.cu_seqlens_k[: bs + 1],
-            indices_buffer=capture.indices,
-            last_page_len_buffer=capture.one_tensor[:bs],
-        )
-        self.graph_wrappers[bs]._backend = "fa2"
-        self.graph_wrappers[bs]._int_workspace_buffer = self.int_workspace_buffer
+        self.graph_wrappers[bs] = self._make_graph_wrapper(bs)
         self.prepare_metadata(batch)
         metadata = batch.attn_metadata
         assert isinstance(metadata, FIMetadata)

python/freetoken/attention/m3_sparse.py

@@ -368,7 +368,9 @@ def _prefill(self, md, layer_id: int, slot: int, q, index_q, batch: Batch) -> to
         return out
 
     # ----- CUDA graph (decode) ------------------------------------------------------------
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        if rows_per_req != 1:
+            raise NotImplementedError("m3_sparse captures decode graphs only")
         self.inner.init_capture_graph(max_seq_len, bs_list)
         self.capture_bs = sorted(bs_list)
         max_bs = max(bs_list)

python/freetoken/attention/qsa_sparse.py

@@ -489,7 +489,9 @@ def _scratch(self, name: str, rows: int, *shape: int, dtype: torch.dtype) -> tor
         return torch.empty((rows, *shape), dtype=dtype, device=self.device)
 
     # ----- CUDA graph (decode) --------------------------------------------------------------
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        if rows_per_req != 1:
+            raise NotImplementedError("qsa_sparse captures decode graphs only")
         self.capture_bs = sorted(bs_list)
         max_bs = max(bs_list)
         width = get_global_ctx().page_table.shape[1]

python/freetoken/attention/triton.py

@@ -32,15 +32,16 @@ def create(
         num_q_heads: int,
         max_head_dim: int,
         scratch_rows: int,
+        rows_per_req: int = 1,
         **kwargs,
     ):
         return cls(
             seq_lens=torch.ones((max_bs,), dtype=torch.int32, device=device),
-            positions=torch.zeros((max_bs,), dtype=torch.int32, device=device),
+            positions=torch.zeros((max_bs * rows_per_req,), dtype=torch.int32, device=device),
             cu_seqlens_k=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
-            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
+            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device) * rows_per_req,
             page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device),
-            q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device),
+            q_to_req=torch.arange(max_bs, dtype=torch.int32, device=device).repeat_interleave(rows_per_req),
             attn_logits=torch.empty((scratch_rows * num_q_heads * max_head_dim,), dtype=torch.float32, device=device),
             attn_lse=torch.empty((scratch_rows * num_q_heads,), dtype=torch.float32, device=device),
             **kwargs,
@@ -71,6 +72,7 @@ def __init__(self, config: ModelConfig):
         self.kvcache = get_global_ctx().kv_cache
         self.device = self.kvcache.device
         self.capture: TritonCaptureData | None = None
+        self.capture_rows_per_req = 1
         self.capture_bs: List[int] = []
         self.max_graph_bs = 0
         self.prefill_tile_min_q = 128
@@ -253,7 +255,7 @@ def prepare_metadata(self, batch: Batch) -> None:
             swa_indices=swa_indices,
         )
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    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."
         max_bs = max(bs_list)
         self.capture = TritonCaptureData.create(
@@ -263,7 +265,9 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
             num_q_heads=max(1, self.num_q_heads),
             max_head_dim=max(1, self.max_head_dim),
             scratch_rows=max(bs * self._split_cap(bs) for bs in bs_list),
+            rows_per_req=rows_per_req,
         )
+        self.capture_rows_per_req = rows_per_req
         if self._swa_capture_enabled():
             self.capture.swa_page_table = torch.zeros(
                 (max_bs, max_seq_len),
@@ -290,6 +294,10 @@ def _point_to_capture(self, metadata: TritonMetadata, bs: int) -> None:
         indices = self.capture.page_table.view(-1)
         self.capture.cu_seqlens_q[: bs + 1].copy_(metadata.cu_seqlens_q_gpu)
         self.capture.cu_seqlens_k[: bs + 1].copy_(metadata.indptr)
+        if not metadata.is_decode:
+            # the extend kernel reads prefix_lens; seq_lens is otherwise unused by this backend, so it holds them
+            self.capture.seq_lens[:bs].copy_(metadata.prefix_lens)
+            metadata.prefix_lens = self.capture.seq_lens[:bs]
         total = metadata.indices.numel()
         indices[:total].copy_(metadata.indices)
         if metadata.swa_indices is not None:
@@ -319,6 +327,24 @@ def prepare_for_capture(self, batch: Batch) -> None:
         bs = batch.size
         assert bs in self.capture_bs and self.capture is not None
         capture = self.capture
+        rows_per_req = self.capture_rows_per_req
+        if rows_per_req > 1:
+            n = bs * rows_per_req
+            # placeholders for a fresh request (kv rows == query rows, no prefix); replay overwrites both
+            capture.cu_seqlens_k[: bs + 1].copy_(capture.cu_seqlens_q[: bs + 1])
+            capture.seq_lens[:bs].fill_(0)
+            batch.attn_metadata = TritonMetadata(
+                cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1],
+                indptr=capture.cu_seqlens_k[: bs + 1],
+                indices=capture.page_table.view(-1),
+                q_to_req=capture.q_to_req[:n],
+                q_positions=capture.positions[:n],
+                is_decode=False,
+                prefix_lens=capture.seq_lens[:bs],
+                max_q_len=rows_per_req,
+                swa_indices=self._capture_swa_indices(),
+            )
+            return
         attn_logits, attn_lse = self._capture_scratch(bs)
         batch.attn_metadata = TritonMetadata(
             cu_seqlens_q_gpu=capture.cu_seqlens_q[: bs + 1],

python/freetoken/attention/trtllm.py

@@ -138,7 +138,9 @@ def prepare_metadata(self, batch: Batch) -> None:
             page_table=new_page_table,
         )
 
-    def init_capture_graph(self, max_seq_len: int, bs_list: List[int]) -> None:
+    def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:
+        if rows_per_req != 1:
+            raise NotImplementedError("trtllm captures decode graphs only")
         assert self.capture is None, "Capture already initialized."
         max_bs = max(bs_list)
         capture = TRTLLMCaptureData.create(

python/freetoken/attention/utils.py

@@ -12,12 +12,13 @@ class BaseCaptureData:
     page_table: torch.Tensor
 
     @classmethod
-    def create(cls, max_bs: int, max_seq_len: int, device: torch.device, **kwargs):
+    def create(cls, max_bs: int, max_seq_len: int, device: torch.device, rows_per_req: int = 1, **kwargs):
         return cls(
             seq_lens=torch.ones((max_bs,), dtype=torch.int32, device=device),
-            positions=torch.zeros((max_bs,), dtype=torch.int32, device=device),
+            positions=torch.zeros((max_bs * rows_per_req,), dtype=torch.int32, device=device),
             cu_seqlens_k=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
-            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device),
+            # constant under a graph: every request has exactly `rows_per_req` query rows
+            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device) * rows_per_req,
             page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device),
             **kwargs,
         )