From b7bddc383aef6165ed3846093f8987dfb9124190 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Mon, 26 Oct 2026 06:32:08 +0000
Subject: [PATCH 1/8] refactor: add engine and scheduler seams for multi-row
 decode steps

---
 python/freetoken/attention/__init__.py        |   7 +-
 python/freetoken/attention/base.py            |   7 +-
 python/freetoken/attention/dsa.py             |   4 +-
 .../freetoken/attention/dsa_indexer_kpool.py  |   4 +-
 python/freetoken/attention/dsv41_sparse.py    |   4 +-
 python/freetoken/attention/dsv4_sparse.py     |   4 +-
 python/freetoken/attention/fa.py              |  12 +-
 python/freetoken/attention/fi.py              |  55 +++++--
 python/freetoken/attention/m3_sparse.py       |   4 +-
 python/freetoken/attention/qsa_sparse.py      |   4 +-
 python/freetoken/attention/triton.py          |  34 ++++-
 python/freetoken/attention/trtllm.py          |   4 +-
 python/freetoken/attention/utils.py           |   7 +-
 python/freetoken/core.py                      |  63 +++++++-
 python/freetoken/engine/config.py             |   5 +
 python/freetoken/engine/engine.py             |  36 +++--
 python/freetoken/engine/graph.py              | 135 +++++++++++------
 python/freetoken/layers/embedding.py          |   8 +-
 python/freetoken/llm/llm.py                   |   6 +-
 python/freetoken/message/tokenizer.py         |   3 +-
 python/freetoken/models/blocks.py             |  15 ++
 python/freetoken/scheduler/cache.py           |  13 ++
 python/freetoken/scheduler/decode.py          |   4 +-
 python/freetoken/scheduler/prefill.py         |   8 +-
 python/freetoken/scheduler/scheduler.py       |  85 +++++------
 python/freetoken/scheduler/status.py          |   8 +-
 python/freetoken/tokenizer/detokenize.py      |   6 +-
 python/freetoken/tokenizer/server.py          |   2 +-
 tests/attention/test_capture_steps.py         | 143 ++++++++++++++++++
 tests/engine/test_context_override.py         |  47 ++++++
 tests/models/qwen4_exp/test_ple_disk.py       |   5 +-
 tests/moe/test_offload.py                     |   5 +-
 .../scheduler/test_abort_inflight_prefill.py  |   1 +
 tests/scheduler/test_cost_accounting_core.py  |   2 +-
 tests/scheduler/test_rollback_tail.py         |  59 ++++++++
 tests/server/test_message_wire.py             |   2 +-
 36 files changed, 646 insertions(+), 165 deletions(-)
 create mode 100644 tests/attention/test_capture_steps.py
 create mode 100644 tests/engine/test_context_override.py
 create mode 100644 tests/scheduler/test_rollback_tail.py

diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py
index eda3b61f55..3957cff6fb 100644
--- a/python/freetoken/attention/__init__.py
+++ b/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 `steps` 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):
diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py
index 846e99447a..d0144eacad 100644
--- a/python/freetoken/attention/base.py
+++ b/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], steps: int = 1) -> None:
+        """Static buffers for graphs of ``steps`` 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], steps: int = 1) -> None:
+        self.decode_backend.init_capture_graph(max_seq_len, bs_list, steps)
 
     def prepare_for_capture(self, batch: Batch) -> None:
         self.decode_backend.prepare_for_capture(batch)
diff --git a/python/freetoken/attention/dsa.py b/python/freetoken/attention/dsa.py
index e8c7b3c2a5..c2e56422ab 100644
--- a/python/freetoken/attention/dsa.py
+++ b/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], steps: int = 1) -> None:
+        if steps != 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)
diff --git a/python/freetoken/attention/dsa_indexer_kpool.py b/python/freetoken/attention/dsa_indexer_kpool.py
index 2b2492aa46..6164496767 100644
--- a/python/freetoken/attention/dsa_indexer_kpool.py
+++ b/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, steps: int = 1) -> None:
+        super().init_capture_graph(max_seq_len, bs_list, steps)
         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:
diff --git a/python/freetoken/attention/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py
index baaa2cf019..d88f1571e8 100644
--- a/python/freetoken/attention/dsv41_sparse.py
+++ b/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], steps: int = 1) -> None:
+        if steps != 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)
diff --git a/python/freetoken/attention/dsv4_sparse.py b/python/freetoken/attention/dsv4_sparse.py
index 2c740e2a48..984184b8cc 100644
--- a/python/freetoken/attention/dsv4_sparse.py
+++ b/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], steps: int = 1) -> None:
+        if steps != 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)
diff --git a/python/freetoken/attention/fa.py b/python/freetoken/attention/fa.py
index 782ef040e2..e3d38b308c 100644
--- a/python/freetoken/attention/fa.py
+++ b/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_steps = 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], steps: 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, steps=steps)
         self.max_graph_bs = max_bs
         self.capture = capture
         self.capture_bs = sorted(bs_list)
+        self.capture_steps = steps
 
     def prepare_for_capture(self, batch: Batch) -> None:
         assert (bs := batch.size) in self.capture_bs and self.capture
         capture = self.capture
+        steps = self.capture_steps
+        capture.seq_lens[:bs].fill_(steps)
         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=steps,
             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 * steps 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)
diff --git a/python/freetoken/attention/fi.py b/python/freetoken/attention/fi.py
index c9e58538f7..73e75347f6 100644
--- a/python/freetoken/attention/fi.py
+++ b/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_steps = 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], steps: 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, steps=steps)
         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_steps = steps
 
     @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_steps > 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)
diff --git a/python/freetoken/attention/m3_sparse.py b/python/freetoken/attention/m3_sparse.py
index 82a98a5a78..30de53da84 100644
--- a/python/freetoken/attention/m3_sparse.py
+++ b/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], steps: int = 1) -> None:
+        if steps != 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)
diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index 8ce63c4242..23089800eb 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/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], steps: int = 1) -> None:
+        if steps != 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]
diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py
index 7f29c481bb..d753158102 100644
--- a/python/freetoken/attention/triton.py
+++ b/python/freetoken/attention/triton.py
@@ -32,15 +32,16 @@ def create(
         num_q_heads: int,
         max_head_dim: int,
         scratch_rows: int,
+        steps: 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 * steps,), 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) * steps,
             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(steps),
             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_steps = 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], steps: 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),
+            steps=steps,
         )
+        self.capture_steps = steps
         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
+        steps = self.capture_steps
+        if steps > 1:
+            n = bs * steps
+            # 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=steps,
+                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],
diff --git a/python/freetoken/attention/trtllm.py b/python/freetoken/attention/trtllm.py
index 2984415628..6f76379a78 100644
--- a/python/freetoken/attention/trtllm.py
+++ b/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], steps: int = 1) -> None:
+        if steps != 1:
+            raise NotImplementedError("trtllm captures decode graphs only")
         assert self.capture is None, "Capture already initialized."
         max_bs = max(bs_list)
         capture = TRTLLMCaptureData.create(
diff --git a/python/freetoken/attention/utils.py b/python/freetoken/attention/utils.py
index c45c394148..0a8c06127b 100644
--- a/python/freetoken/attention/utils.py
+++ b/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, steps: 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 * steps,), 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 `steps` query rows
+            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device) * steps,
             page_table=torch.zeros((max_bs, max_seq_len), dtype=torch.int32, device=device),
             **kwargs,
         )
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index 812f9a4349..acea1f3e8d 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -2,7 +2,7 @@
 
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, List, Literal, Tuple
+from typing import TYPE_CHECKING, Any, List, Literal, Tuple
 
 import torch
 
@@ -11,6 +11,7 @@
     from freetoken.attention.linear import FLAMetadata
     from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool
     from freetoken.kvcache.linear_state_pool import LinearStatePool
+    from freetoken.models.blocks import FeatureSink
     from freetoken.moe.offload_cache import OffloadMoeCache
 
 
@@ -95,6 +96,15 @@ def complete_one(self) -> None:
         self.cached_len = self.device_len
         self.device_len += 1
 
+    def widen_pending(self, k: int) -> None:
+        """Extend the pending token at cached_len by k provisional draft tokens."""
+        self.device_len += k
+
+    def commit_pending(self, n: int) -> None:
+        """Make n tokens final and leave the last committed one pending at cached_len."""
+        self.cached_len += n
+        self.device_len = self.cached_len + 1
+
     def append_host(self, next_token: torch.Tensor) -> None:
         n = self.input_ids.numel()
         m = n + next_token.numel()
@@ -115,10 +125,21 @@ def __repr__(self) -> str:
 
 
 
+@dataclass
+class VerifyInfo:
+    """A uniform extend of `steps` rows per request over read-only recurrent state."""
+
+    steps: int
+    gdn: Any = None
+    model_scratch: Any = None
+
+
 @dataclass
 class Batch:
     reqs: List[Req]
-    phase: Literal["prefill", "decode"]
+    phase: Literal["prefill", "decode", "verify"]
+    # query rows per request: 1 on decode, steps on verify, None on prefill (ragged)
+    rows_per_req: int | None = None
     # these fields should be set by scheduler
     input_ids: torch.Tensor = field(init=False)
     positions: torch.Tensor = field(init=False)
@@ -160,6 +181,13 @@ class Batch:
     # _prepare_batch succeeds. Continuation chunks leave this empty, so accounting is
     # exactly-once.
     prompt_admissions: List[Tuple[int, int, int]] = field(default_factory=list, init=False)
+    # per-row hidden states a draft model reads from this forward; None when nothing reads them
+    features: "FeatureSink | None" = field(default=None, init=False)
+    verify: VerifyInfo | None = field(default=None, init=False)
+
+    def __post_init__(self) -> None:
+        if self.rows_per_req is None and self.phase == "decode":
+            self.rows_per_req = 1
 
     @property
     def is_prefill(self) -> bool:
@@ -200,11 +228,38 @@ def batch(self) -> Batch:
         return self._batch
 
     @contextmanager
-    def forward_batch(self, batch: Batch):
+    def override(
+        self,
+        *,
+        kv_cache: BaseKVCachePool | None = None,
+        attn_backend: BaseAttnBackend | None = None,
+        hide_target_state: bool = False,
+    ):
+        """Swap the KV pool and attention backend for a draft forward; hide_target_state also hides the target's GDN pool and MoE offload cache."""
+        swaps = {}
+        if kv_cache is not None:
+            swaps["kv_cache"] = kv_cache
+        if attn_backend is not None:
+            swaps["attn_backend"] = attn_backend
+        if hide_target_state:
+            swaps["linear_state_pool"] = None
+            swaps["moe_offload_cache"] = None
+        saved = {name: getattr(self, name) for name in swaps}
+        try:
+            for name, value in swaps.items():
+                setattr(self, name, value)
+            yield
+        finally:
+            for name, value in saved.items():
+                setattr(self, name, value)
+
+    @contextmanager
+    def forward_batch(self, batch: Batch, **override):
         assert self._batch is None, "Nested forward_batch is not allowed"
         try:
             self._batch = batch
-            yield
+            with self.override(**override):
+                yield
         finally:
             self._batch = None
 
diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py
index 77aaa3eefe..39379c4278 100644
--- a/python/freetoken/engine/config.py
+++ b/python/freetoken/engine/config.py
@@ -176,6 +176,11 @@ def max_seq_len(self) -> int:
     def max_forward_len(self) -> int:
         return self.max_seq_len
 
+    @property
+    def decode_rows_per_req(self) -> int:
+        """Query rows each running request takes per decode step; every per-request decode allowance derives from it."""
+        return 1
+
     @property
     def distributed_addr(self) -> str:
         return "tcp://127.0.0.1:2333"
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 2de7f47a3e..1b835648d2 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -111,6 +111,12 @@ def _page_table_width(max_seq_len: int, page_size: int) -> int:
     return align_ceil(align_ceil(max_seq_len, page_size), 32)
 
 
+def _seq_limits(config, num_tokens: int) -> tuple[int, int]:
+    """(longest request, page-table width). A decode step may widen a request by decode_rows_per_req - 1 provisional rows past its length, so requests get that much less and the table still covers the widened tail."""
+    limit = min(config.max_seq_len, num_tokens)
+    return limit - (getattr(config, "decode_rows_per_req", 1) - 1), _page_table_width(limit, config.page_size)
+
+
 def _required_attn_types(model_config) -> frozenset[AttnType]:
     """Backend-driving attention types of this model, from the group-spec walk
     (single source shared with the pool factory and the KV cost model). getattr
@@ -524,8 +530,7 @@ def __init__(self, config: EngineConfig):
 
         # ======================= Page table initialization ========================
         # NOTE: 1. aligned to 128 bytes; 2. store raw locations instead of pages
-        self.max_seq_len = min(config.max_seq_len, num_tokens)
-        aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size)
+        self.max_seq_len, aligned_max_seq_len = _seq_limits(config, num_tokens)
         self.ctx.page_table = self.page_table = torch.zeros(  # + 1 for dummy request
             (config.max_running_req + 1, aligned_max_seq_len),
             dtype=torch.int32,
@@ -564,16 +569,16 @@ def __init__(self, config: EngineConfig):
         self.graph_runner = GraphRunner(
             stream=self.stream,
             device=self.device,
-            model=self.model,
+            forward=self.model.forward,
             attn_backend=self.attn_backend,
             cuda_graph_bs=config.cuda_graph_bs,
             cuda_graph_max_bs=config.cuda_graph_max_bs,
             free_memory=init_free_memory,
             max_seq_len=aligned_max_seq_len,
-            vocab_size=config.model_config.vocab_size,
             dummy_req=self.dummy_req,
             moe_offload_cache=self.moe_offload_cache,
             mrope=config.model_config.model_is_mrope,
+            out_dtype=torch.float32,
         )
         if config.attention_backend.split(",")[0] == "triton":
             # Prefill runs on the first comma part; warm its autotune cache.
@@ -1034,6 +1039,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None:
         # Decode batches never exceed max_running_req, but CUDA-graph padding can
         # round a batch up to the largest captured size; cover both.
         max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, *(config.cuda_graph_bs or []), 1)
+        max_tokens *= config.decode_rows_per_req
         by_id = sorted(layers, key=lambda layer: layer.layer_id)
         layer_formats = [layer.quant_method.cpu_spec() for layer in by_id]
         executor = CpuMoeExecutor(
@@ -1098,8 +1104,7 @@ def _resize_kv_pool(self, config, num_pages: int, num_swa_pages: int | None) ->
 
     def _refresh_seq_state(self, config) -> None:
         num_tokens = self.num_pages * config.page_size
-        self.max_seq_len = min(config.max_seq_len, num_tokens)
-        aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size)
+        self.max_seq_len, aligned_max_seq_len = _seq_limits(config, num_tokens)
         if aligned_max_seq_len != self.page_table.shape[1]:
             # max_seq_len changed (e.g. KV grew past the startup token budget); the page table
             # columns must track it or new requests would index out of bounds. The scheduler
@@ -1244,7 +1249,7 @@ def rebuild_runtime_cache(
             self.linear_state_pool.rebuild(num_mamba_slots + 1)
         # 3. Refresh max_seq_len (+ generic page table) for the new token budget.
         self._refresh_seq_state(config)
-        aligned_max_seq_len = _page_table_width(self.max_seq_len, config.page_size)
+        aligned_max_seq_len = self.page_table.shape[1]
         # 4. Re-capture CUDA graphs against the new tensors (reset_capture above re-armed
         #    the backend; _sync_get_memory empties the cache so freed memory is reclaimed).
         gc.collect()
@@ -1252,22 +1257,20 @@ def rebuild_runtime_cache(
         self.graph_runner = GraphRunner(
             stream=self.stream,
             device=self.device,
-            model=self.model,
+            forward=self.model.forward,
             attn_backend=self.attn_backend,
             cuda_graph_bs=prior_graph_bs,  # reuse the startup-resolved set (see above)
             cuda_graph_max_bs=config.cuda_graph_max_bs,
             free_memory=free_min,
             max_seq_len=aligned_max_seq_len,
-            vocab_size=config.model_config.vocab_size,
             dummy_req=self.dummy_req,
             moe_offload_cache=self.moe_offload_cache,
             mrope=config.model_config.model_is_mrope,
+            out_dtype=torch.float32,
         )
 
-    def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
-        assert torch.cuda.current_stream() == self.stream
-        if batch.mm_gather_plan:
-            self._run_mm_encoder(batch)
+    def forward_logits(self, batch: Batch) -> torch.Tensor:
+        """The batch's logits from its captured graph, or eagerly when none fits."""
         use_graph = self.graph_runner.can_use_cuda_graph(batch)
         with self.ctx.forward_batch(batch), self.model.forward_host_ctx(batch, use_graph):
             logits = self.graph_runner.replay(batch) if use_graph else self.model.forward()
@@ -1275,6 +1278,13 @@ def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
             # One pinned read: surfaces a fired flag-handshake watchdog (dead coordinator
             # -> stale expert outputs) as a loud error instead of silent corruption.
             self.cpu_moe_executor.raise_if_unhealthy()
+        return logits
+
+    def forward_batch(self, batch: Batch, args: BatchSamplingArgs) -> ForwardOutput:
+        assert torch.cuda.current_stream() == self.stream
+        if batch.mm_gather_plan:
+            self._run_mm_encoder(batch)
+        logits = self.forward_logits(batch)
 
         for req in batch.reqs:
             req.complete_one()
diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py
index a995316ddc..ff7006e4b0 100644
--- a/python/freetoken/engine/graph.py
+++ b/python/freetoken/engine/graph.py
@@ -5,7 +5,7 @@
 import gc
 import sys
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal
 
 import torch
 from freetoken.core import Batch, Req, get_global_ctx
@@ -16,7 +16,6 @@
 
 if TYPE_CHECKING:
     from freetoken.attention import BaseAttnBackend
-    from freetoken.models import BaseLLMModel
     from freetoken.moe.offload_cache import OffloadMoeCache
 
 logger = init_logger(__name__)
@@ -27,56 +26,55 @@ class GraphCaptureBuffer:
     input_ids: torch.Tensor
     out_loc: torch.Tensor
     positions: torch.Tensor
-    # [3, bs] t/h/w rope positions; allocated only for mrope models (else None).
+    # [3, bs * steps] t/h/w rope positions; allocated only for mrope models (else None).
     mrope_positions: torch.Tensor | None
-    logits: torch.Tensor
     table_idx: torch.Tensor  # per-request slot id for GatedDeltaNet state gather/scatter
-    # Decode GDN query indptr = arange(bs+1); a constant per captured bs, filled once.
+    # GDN query indptr = arange(bs+1) * steps; a constant per captured bs, filled once.
     fla_cu_seqlens: torch.Tensor
+    steps: int = 1
 
     @classmethod
-    def init(
-        cls, bs: int, vocab_size: int, device: torch.device, mrope: bool = False
-    ) -> GraphCaptureBuffer:
+    def init(cls, bs: int, device: torch.device, mrope: bool = False, steps: int = 1) -> GraphCaptureBuffer:
+        n = bs * steps
         return GraphCaptureBuffer(
-            input_ids=torch.zeros(bs, dtype=torch.int32, device=device),
-            out_loc=torch.zeros(bs, dtype=torch.int32, device=device),
-            positions=torch.zeros(bs, dtype=torch.int32, device=device),
+            input_ids=torch.zeros(n, dtype=torch.int32, device=device),
+            out_loc=torch.zeros(n, dtype=torch.int32, device=device),
+            positions=torch.zeros(n, dtype=torch.int32, device=device),
             mrope_positions=(
-                torch.zeros(3, bs, dtype=torch.int32, device=device) if mrope else None
+                torch.zeros(3, n, dtype=torch.int32, device=device) if mrope else None
             ),
-            logits=torch.empty(bs, vocab_size, dtype=torch.float32, device=device),
             table_idx=torch.zeros(bs, dtype=torch.int32, device=device),
-            fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device),
+            fla_cu_seqlens=torch.arange(bs + 1, dtype=torch.int32, device=device) * steps,
+            steps=steps,
         )
 
     def set_batch(self, batch: Batch) -> None:
         from freetoken.attention.linear import FLAMetadata
 
-        _slice = slice(batch.padded_size)
         bs = batch.padded_size
-        batch.input_ids = self.input_ids[_slice]
-        batch.out_loc = self.out_loc[_slice]
-        batch.positions = self.positions[_slice]
+        _rows = slice(bs * self.steps)
+        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[:, _slice]
-        batch.linear_table_idx = self.table_idx[_slice]
+            batch.mrope_positions = self.mrope_positions[:, _rows]
+        batch.linear_table_idx = self.table_idx[:bs]
         # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the
         # persistent 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[_slice]
+            cu_seqlens=self.fla_cu_seqlens[: bs + 1], cache_indices=self.table_idx[:bs]
         )
 
     def copy_from(self, batch: Batch) -> None:
-        _slice = slice(batch.padded_size)
-        self.input_ids[_slice] = batch.input_ids
+        _rows = slice(batch.padded_size * self.steps)
+        self.input_ids[_rows] = batch.input_ids
         if batch.out_loc is not None:
-            self.out_loc[_slice] = batch.out_loc
-        self.positions[_slice] = batch.positions
+            self.out_loc[_rows] = batch.out_loc
+        self.positions[_rows] = batch.positions
         if self.mrope_positions is not None:
-            self.mrope_positions[:, _slice] = batch.mrope_positions
+            self.mrope_positions[:, _rows] = batch.mrope_positions
         if batch.linear_table_idx is not None:
-            self.table_idx[_slice] = batch.linear_table_idx
+            self.table_idx[: batch.padded_size] = batch.linear_table_idx
 
 
 def _determine_cuda_graph_bs(
@@ -130,26 +128,41 @@ def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> Non
 
 
 class GraphRunner:
+    """CUDA graphs of ``forward`` over padded batches of ``steps`` rows per request, one graph per batch size.
+
+    ``forward`` returns one tensor or a tuple of tensors; each is copied into a buffer sized by the largest
+    captured batch, and ``replay`` returns the real requests' rows of each.
+    """
+
     def __init__(
         self,
         stream: torch.cuda.Stream,
         device: torch.device,
-        model: BaseLLMModel,
+        forward: Callable[[], Any],
         attn_backend: BaseAttnBackend,
         cuda_graph_bs: List[int] | None,
         cuda_graph_max_bs: int | None,
         free_memory: int,
         max_seq_len: int,
-        vocab_size: int,
         dummy_req: Req,
         moe_offload_cache: OffloadMoeCache | None = None,
         mrope: bool = False,
+        *,
+        steps: int = 1,
+        phase: Literal["decode", "verify"] = "decode",
+        ctx_override: Dict[str, Any] | None = None,
+        batch_hook: Callable[[Batch], None] | None = None,
+        pool: Any = None,
+        out_dtype: torch.dtype | None = None,
     ) -> None:
         cuda_graph_bs = _determine_cuda_graph_bs(
             cuda_graph_bs=cuda_graph_bs,
             cuda_graph_max_bs=cuda_graph_max_bs,
             free_memory=free_memory,
         )
+        # the padded rows replay the dummy request, so it must carry one row per step
+        assert steps == 1 or dummy_req.extend_len == steps, (steps, dummy_req.extend_len)
+        self.forward = forward
         self.attn_backend = attn_backend
         self.max_graph_bs = max(cuda_graph_bs) if cuda_graph_bs else 0
         self.graph_bs_list = sorted(cuda_graph_bs)
@@ -158,13 +171,33 @@ def __init__(
         self.mrope = mrope
         self.stream = stream
         self.device = device
-        self._capture_graphs(max_seq_len, vocab_size, model)
+        self.steps = steps
+        self.phase = phase
+        self.ctx_override = ctx_override or {}
+        self.batch_hook = batch_hook
+        self.graph_pool = pool
+        self.out_dtype = out_dtype
+        self._out: List[torch.Tensor] | None = None
+        self._out_is_tuple = False
+        self._out_rows: Dict[int, tuple[int, ...]] = {}
+        self._capture_graphs(max_seq_len)
 
     def _reset_moe_offload_cache(self) -> None:
         if self.moe_offload_cache is not None:
             self.moe_offload_cache.reset()
 
-    def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel):
+    def _store_output(self, bs: int, out: Any) -> None:
+        self._out_is_tuple = isinstance(out, tuple)
+        outs = out if self._out_is_tuple else (out,)
+        if self._out is None:
+            # the largest batch is captured first, so its outputs size the buffers
+            assert self.out_dtype is None or len(outs) == 1, "out_dtype applies to a single output"
+            self._out = [torch.empty_like(o, dtype=self.out_dtype or o.dtype) for o in outs]
+        for buf, o in zip(self._out, outs, strict=True):
+            buf[: o.shape[0]] = o
+        self._out_rows[bs] = tuple(o.shape[0] for o in outs)
+
+    def _capture_graphs(self, max_seq_len: int):
         # Mark the post-weights "warmup" phase for /health: this stretch (graph capture — or the
         # remaining readiness work when graphs are disabled) moves no bytes, so without this the
         # loader would sit at 100% (last byte bar) until the ready ack. total=0 ⇒ the desktop
@@ -175,7 +208,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel
         if self.max_graph_bs == 0:
             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)
+        self.attn_backend.init_capture_graph(max_seq_len=max_seq_len, bs_list=self.graph_bs_list, steps=self.steps)
 
         torch.cuda.synchronize(self.device)
         torch.cuda.empty_cache()
@@ -185,9 +218,7 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel
         free_memory = get_free_memory(self.device)
         logger.info_rank0(f"Free GPU memory before capturing CUDA graphs: {mem_GB(free_memory)}")
 
-        self.buffer = GraphCaptureBuffer.init(
-            self.max_graph_bs, vocab_size, self.device, mrope=self.mrope
-        )
+        self.buffer = GraphCaptureBuffer.init(self.max_graph_bs, self.device, mrope=self.mrope, steps=self.steps)
         self._reset_moe_offload_cache()
 
         pbar = tqdm(
@@ -196,14 +227,16 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel
             unit="batch",
             disable=not get_tp_info().is_primary(),  # disable for non-primary ranks
         )
-        pool = None
+        pool = self.graph_pool
         for bs in pbar:
             free_memory = get_free_memory(self.device)
             pbar.desc = f"Capturing graphs: bs = {bs:<3} | avail_mem = {mem_GB(free_memory)}"
             pbar.refresh()
             graph = torch.cuda.CUDAGraph()
-            batch = Batch(reqs=[self.dummy_req] * bs, phase="decode")
+            batch = Batch(reqs=[self.dummy_req] * bs, phase=self.phase, rows_per_req=self.steps)
             batch.padded_reqs = batch.reqs
+            if self.batch_hook is not None:
+                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
@@ -213,32 +246,41 @@ def _capture_graphs(self, max_seq_len: int, vocab_size: int, model: BaseLLMModel
                           if self.dummy_req.linear_slot_idx is not None
                           else self.dummy_req.table_idx)
             self.buffer.table_idx[:bs].fill_(dummy_slot)
-            with get_global_ctx().forward_batch(batch):
-                self.buffer.logits[:bs] = model.forward()
+            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
                 # CUDA graph capture to replay cold-cache expert copies.
                 with torch.cuda.graph(graph, pool=pool, stream=self.stream):
-                    self.buffer.logits[:bs] = model.forward()
+                    self._store_output(bs, self.forward())
                 _upload_graph(graph, self.stream)
                 self._reset_moe_offload_cache()
             if pool is None:
                 pool = graph.pool()  # reuse cuda graph handle to reduce memory
             self.graph_map[bs] = graph
+        self.graph_pool = pool
 
         self._reset_moe_offload_cache()
         free_memory = get_free_memory(self.device)
         logger.info_rank0(f"Free GPU memory after capturing CUDA graphs: {mem_GB(free_memory)}")
 
     def can_use_cuda_graph(self, batch: Batch) -> bool:
-        return batch.is_decode and batch.size <= self.max_graph_bs
+        return (
+            batch.phase == self.phase
+            and batch.rows_per_req == self.steps
+            and batch.size <= self.max_graph_bs
+        )
 
-    def replay(self, batch: Batch) -> torch.Tensor:
-        assert self.can_use_cuda_graph(batch)
+    def replay(self, batch: Batch) -> Any:
+        assert self.can_use_cuda_graph(batch) and self._out is not None
         self.buffer.copy_from(batch)
-        g = self.graph_map[batch.padded_size]
+        padded = batch.padded_size
+        g = self.graph_map[padded]
         self.attn_backend.prepare_for_replay(batch)
         g.replay()
-        return self.buffer.logits[: batch.size]
+        outs = tuple(
+            buf[: rows // padded * batch.size] for buf, rows in zip(self._out, self._out_rows[padded], strict=True)
+        )
+        return outs if self._out_is_tuple else outs[0]
 
     def pad_batch(self, batch: Batch) -> None:
         padded_size = (  # choose the first available batch size
@@ -251,10 +293,11 @@ def pad_batch(self, batch: Batch) -> None:
     # NOTE: This must be called before freeing NCCL resources to prevent program hang
     def destroy_cuda_graphs(self) -> None:
         # Drop the CUDAGraph objects (and the shared mempool they hold) AND the static
-        # GraphCaptureBuffer tensors ([max_bs, vocab] logits + input/out_loc/positions/...).
+        # GraphCaptureBuffer tensors (input/out_loc/positions/...) and output buffers ([max_bs, vocab] logits).
         # Dropping the references is the load-bearing step; without it a runtime rebuild's
         # free-before-alloc cannot reclaim this GPU memory. empty_cache() is left to the
         # caller / next capture (GraphRunner._capture_graphs already runs it).
         self.graph_map = {}
         self.buffer = None
+        self._out = None
         gc.collect()
diff --git a/python/freetoken/layers/embedding.py b/python/freetoken/layers/embedding.py
index f86670b14c..469b752da0 100644
--- a/python/freetoken/layers/embedding.py
+++ b/python/freetoken/layers/embedding.py
@@ -129,11 +129,16 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         ctx = get_global_ctx()
         batch = ctx.batch
         bs = batch.size
+        if (features := getattr(batch, "features", None)) is not None:
+            features.final_hidden(x)
         if batch.is_prefill:
             indices = batch.attn_metadata.get_last_indices(bs)
             x = x[indices].contiguous()
             del indices
+        return self.logits(x)
 
+    def logits(self, x: torch.Tensor) -> torch.Tensor:
+        """Full-vocab logits of the rows of ``x``; reads no batch state."""
         if self.tied_embedding is not None:
             logits = self.tied_embedding.quant_method.logits(self.tied_embedding, x, self.bias)
         else:
@@ -143,7 +148,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         input_shape = logits.shape
         output_tensor = self._comm.all_gather(logits)
 
-        if bs == 1:
+        # keyed on rows, not batch size: a verify batch has several rows per request
+        if input_shape[0] == 1:
             return output_tensor.view(1, -1)[:, : self.num_embeddings]
 
         output_tensor = output_tensor.view((self.tp_size,) + input_shape)
diff --git a/python/freetoken/llm/llm.py b/python/freetoken/llm/llm.py
index 20ae087405..8b6a17a15f 100644
--- a/python/freetoken/llm/llm.py
+++ b/python/freetoken/llm/llm.py
@@ -94,8 +94,10 @@ def offline_send_result(self, reply: List[BaseTokenizerMsg]) -> None:
                 continue
             assert isinstance(msg, DetokenizeMsg)
             status = self.status_map[msg.uid]
-            if not (msg.finished and msg.next_token in self.eos_token_ids):
-                status.output_ids.append(msg.next_token)
+            tokens = msg.next_tokens
+            if msg.finished and tokens and tokens[-1] in self.eos_token_ids:
+                tokens = tokens[:-1]
+            status.output_ids.extend(tokens)
 
     def generate(
         self,
diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py
index f091d0d369..56308d6ece 100644
--- a/python/freetoken/message/tokenizer.py
+++ b/python/freetoken/message/tokenizer.py
@@ -27,7 +27,8 @@ class BatchTokenizerMsg(BaseTokenizerMsg):
 @dataclass
 class DetokenizeMsg(BaseTokenizerMsg):
     uid: int
-    next_token: int
+    # tokens committed this step, in order: one per decode step, up to k+1 per speculative verify
+    next_tokens: List[int]
     finished: bool
     finish_reason: str | None = None
     # The stop string that ended generation (if any), so the detokenizer can trim it
diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py
index 4f039fd3ba..67c6a54923 100644
--- a/python/freetoken/models/blocks.py
+++ b/python/freetoken/models/blocks.py
@@ -48,6 +48,21 @@ def place_encoder_weights(self, mode: str) -> None:
         ...
 
 
+@runtime_checkable
+class FeatureSink(Protocol):
+    """Where a forward hands per-row hidden states to a draft model; set on ``batch.features``."""
+
+    wanted_layers: frozenset[int]
+
+    def final_hidden(self, x: torch.Tensor) -> None:
+        """Every row's hidden state before the LM head; the first call in a forward wins, so a model may call it before the head does."""
+        ...
+
+    def layer_output(self, layer_id: int, x: torch.Tensor, num_streams: int = 1) -> None:
+        """A layer's output rows; called only for layers in ``wanted_layers``."""
+        ...
+
+
 def embed_input_ids(embed_tokens, input_ids: torch.Tensor, batch: Batch) -> torch.Tensor:
     """Token embeddings of the batch; on a chunk with multimodal rows, batch.mm_embeds' leading columns replace the rows batch.mm_rows."""
     if batch.mm_embeds is None:
diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py
index bd675d7262..fb1cd93a43 100644
--- a/python/freetoken/scheduler/cache.py
+++ b/python/freetoken/scheduler/cache.py
@@ -623,6 +623,19 @@ def _free(self, indices: torch.Tensor) -> None:
         if len(indices) > 0:
             self.free_slots = torch.cat([self.free_slots, indices[:: self.page_size]])
 
+    def rollback_tail(self, req: Req, keep_len: int, drop_len: int) -> None:
+        """Free the slots allocated for positions [keep_len, drop_len), rounded out to whole pages as allocate_paged allocates them.
+
+        Call inside lazy_free_region, after the forward that wrote them has drained.
+        """
+        start = div_ceil(keep_len, self.page_size) * self.page_size
+        end = div_ceil(drop_len, self.page_size) * self.page_size
+        if end > start:
+            tail = self.page_table[req.table_idx, start:end]
+            if self.swa_paged:
+                self._free_swa(tail)
+            self._free(tail)
+
     def _page_to_token(self, pages: torch.Tensor) -> torch.Tensor:
         if self.page_size == 1:
             return pages
diff --git a/python/freetoken/scheduler/decode.py b/python/freetoken/scheduler/decode.py
index 8f74c12c78..be64726e20 100644
--- a/python/freetoken/scheduler/decode.py
+++ b/python/freetoken/scheduler/decode.py
@@ -9,6 +9,8 @@
 @dataclass
 class DecodeManager:
     page_size: int
+    # slots a decode step may hold past each request's committed length (provisional rows)
+    token_headroom: int = 0
     running_reqs: Set[Req] = field(default_factory=set)
 
     def filter_reqs(self, reqs: Iterable[Req]) -> None:
@@ -26,7 +28,7 @@ def abort_req(self, uid: int) -> Req | None:
 
     @property
     def inflight_tokens(self) -> int:
-        tokens_reserved = (self.page_size - 1) * len(self.running_reqs)  # 1 page reserved
+        tokens_reserved = (self.page_size - 1 + self.token_headroom) * len(self.running_reqs)  # 1 page reserved
         return sum(req.remain_len for req in self.running_reqs) + tokens_reserved
 
     def schedule_next_batch(self) -> Batch | None:
diff --git a/python/freetoken/scheduler/prefill.py b/python/freetoken/scheduler/prefill.py
index de4bddfb4a..94e708fb1d 100644
--- a/python/freetoken/scheduler/prefill.py
+++ b/python/freetoken/scheduler/prefill.py
@@ -55,6 +55,8 @@ class PrefillAdder:
     # allocated only in allocate_paged (after the pass), so swa_available_size does not decrement
     # across the admission loop -- without this, successive admits all see the full pool.
     reserved_swa: int = 0
+    # slots a decode step may hold past a request's own length; reserved with it
+    token_headroom: int = 0
 
     def __post_init__(self) -> None:
         if not self.pass_budget:
@@ -79,7 +81,7 @@ def _try_allocate_one(self, req: PendingReq):
         # TODO: better estimate policy
         extend_len = req.input_len - cached_len
         estimated_size = self._kv_reservation_size(
-            req.input_len + req.output_len, cached_len
+            req.input_len + req.output_len + self.token_headroom, cached_len
         )
 
         if estimated_size + self.reserved_size > self.cache_manager.available_size:
@@ -201,7 +203,7 @@ def _add_one_req(
         # page span here; charging raw tokens can admit several short requests against one page
         # even though allocate_paged() needs a separate page for each request.
         self.reserved_size += self._kv_reservation_size(
-            pending_req.input_len + pending_req.output_len, cached_len
+            pending_req.input_len + pending_req.output_len + self.token_headroom, cached_len
         )
         # NOTE: update the tokens ids only; new pages will be allocated in the scheduler
         _slice = slice(cached_len, cached_len + chunk_size)
@@ -277,6 +279,7 @@ class PrefillManager:
     decode_manager: DecodeManager
     encoder_cache: EncoderCache | None = None
     keep_images_whole: bool = False
+    token_headroom: int = 0
     pending_list: List[PendingReq] = field(default_factory=list)
 
     def add_one_req(self, req: UserMsg) -> None:
@@ -303,6 +306,7 @@ def schedule_next_batch(self, prefill_budget: int) -> Batch | None:
             table_manager=self.table_manager,
             encoder_cache=self.encoder_cache,
             keep_images_whole=self.keep_images_whole,
+            token_headroom=self.token_headroom,
         )
         reqs: List[Req] = []
         chunked_list: List[PendingReq] = []
diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py
index f16608cc6a..21ef19d9f8 100644
--- a/python/freetoken/scheduler/scheduler.py
+++ b/python/freetoken/scheduler/scheduler.py
@@ -93,7 +93,8 @@ def __init__(self, config: SchedulerConfig):
                 None,
             ) or getattr(self.engine.kv_cache, "sliding_window_size", None),
         )
-        self.decode_manager = DecodeManager(config.page_size)
+        token_headroom = config.decode_rows_per_req - 1
+        self.decode_manager = DecodeManager(config.page_size, token_headroom=token_headroom)
         self._bidirectional_mm = any(getattr(g, "bidirectional_mm_blocks", False) for g in config.model_config.attention_groups)
         self.prefill_manager = PrefillManager(
             self.cache_manager,
@@ -101,6 +102,7 @@ def __init__(self, config: SchedulerConfig):
             self.decode_manager,
             encoder_cache=self.engine.encoder_cache,
             keep_images_whole=self._bidirectional_mm,
+            token_headroom=token_headroom,
         )
 
         # some alias for easy access
@@ -360,43 +362,9 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
                     # are freed below/already; shipping this token would append past the
                     # client's terminal reply.
                     continue
-                next_token = next_tokens_cpu[i]
-                req.append_host(next_token.unsqueeze(0))
-                next_token = int(next_token.item())
-                # EOS / stop-string -> "stop", output budget exhausted -> "length";
-                # EOS and stop strings win over length.
-                # Overlap can advance device_len one token ahead of this host reply.
-                hit_length = req.input_ids.numel() >= req.max_device_len
-                hit_eos = (
-                    not req.sampling_params.ignore_eos and next_token in self.eos_token_ids
-                )
-                matched_stop = (
-                    self._match_stop_str(req)
-                    if not hit_eos and req.sampling_params.stop_strs
-                    else None
-                )
-                finished = hit_length or hit_eos or matched_stop is not None
-                finish_reason = (
-                    ("stop" if (hit_eos or matched_stop is not None) else "length")
-                    if finished
-                    else None
-                )
-                if (
-                    next_token == self.toolcall_anchor_id
-                    and req.toolcall_anchor_len is None
-                    and not finished
-                ):
-                    req.toolcall_anchor_len = req.input_ids.numel()
-                reply.append(
-                    DetokenizeMsg(
-                        uid=req.uid,
-                        next_token=next_token,
-                        finished=finished,
-                        finish_reason=finish_reason,
-                        matched_stop=matched_stop,
-                        stop_strs=req.sampling_params.stop_strs or None,
-                    )
-                )
+                msg, _ = self._commit_tokens(req, next_tokens_cpu[i : i + 1])
+                reply.append(msg)
+                finished = msg.finished
 
                 # NOTE: overlap scheduling may make the request freed twice, skip second free
                 if finished and req not in self.finished_reqs:
@@ -444,6 +412,37 @@ def _process_last_data(self, last_data: ForwardData | None) -> None:
         )
         self.send_result(reply)
 
+    def _commit_tokens(self, req: Req, tokens: torch.Tensor) -> Tuple[DetokenizeMsg, int]:
+        """Append sampled tokens in order until one finishes the request; returns the reply and how many tokens were committed."""
+        params = req.sampling_params
+        committed: List[int] = []
+        finished, finish_reason, matched_stop = False, None, None
+        for j in range(tokens.numel()):
+            req.append_host(tokens[j : j + 1])
+            token = int(tokens[j].item())
+            committed.append(token)
+            # EOS / stop-string -> "stop", output budget exhausted -> "length";
+            # EOS and stop strings win over length.
+            # Overlap can advance device_len one token ahead of this host reply.
+            hit_length = req.input_ids.numel() >= req.max_device_len
+            hit_eos = not params.ignore_eos and token in self.eos_token_ids
+            matched_stop = self._match_stop_str(req) if not hit_eos and params.stop_strs else None
+            if hit_length or hit_eos or matched_stop is not None:
+                finished = True
+                finish_reason = "stop" if (hit_eos or matched_stop is not None) else "length"
+                break
+            if token == self.toolcall_anchor_id and req.toolcall_anchor_len is None:
+                req.toolcall_anchor_len = req.input_ids.numel()
+        msg = DetokenizeMsg(
+            uid=req.uid,
+            next_tokens=committed,
+            finished=finished,
+            finish_reason=finish_reason,
+            matched_stop=matched_stop,
+            stop_strs=params.stop_strs or None,
+        )
+        return msg, len(committed)
+
     def _match_stop_str(self, req: Req) -> str | None:
         """First stop string present in this request's generated tail, else None. Decodes
         only a short suffix (bounded by the longest stop string's char length, so a stop of
@@ -841,7 +840,7 @@ def _log_cache_geometry(self, event: str) -> None:
     def _prepare_batch(self, batch: Batch) -> ForwardInput:
         self.engine.graph_runner.pad_batch(batch)
         self._forward_iter += 1
-        if batch.is_decode:
+        if not batch.is_prefill:
             # Free each decoding request's now-out-of-window SWA slots BEFORE the alloc below,
             # so they can back the new token -- this is what bounds the per-request swa
             # footprint during decode. (no-op unless the model is SWA / paged swa pool.)
@@ -866,7 +865,7 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput:
         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 batch.is_decode:
+            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
@@ -882,15 +881,15 @@ def _prepare_batch(self, batch: Batch) -> ForwardInput:
                         slots, dtype=torch.int32, device="cpu", pin_memory=True
                     ).to(self.device, non_blocking=True)
                 else:
-                    batch.linear_table_idx = input_mapping[0].to(torch.int32)
+                    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 batch.is_decode:
+        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.
-            batch.active_table_idx = input_mapping[0].view(-1)
+            batch.active_table_idx = input_mapping[0][:: batch.rows_per_req].view(-1)
         self.engine.attn_backend.prepare_metadata(batch)
         return ForwardInput(
             batch=batch,
diff --git a/python/freetoken/scheduler/status.py b/python/freetoken/scheduler/status.py
index 76c20cd405..5e9ca0e0da 100644
--- a/python/freetoken/scheduler/status.py
+++ b/python/freetoken/scheduler/status.py
@@ -36,7 +36,9 @@ def report_batch(
         page_size: int,
         mamba_slots: tuple[int, int] | None = None,
         swa_tokens: tuple[int, int] | None = None,
+        generated: int | None = None,
     ) -> None:
+        """``generated``: tokens a decode step committed; None means one per request."""
         if batch.is_prefill:
             self._report_prefill(
                 batch,
@@ -47,7 +49,7 @@ def report_batch(
                 mamba_slots=mamba_slots,
                 swa_tokens=swa_tokens,
             )
-        elif batch.is_decode:
+        else:
             self._report_decode(
                 batch,
                 running_reqs=running_reqs,
@@ -57,6 +59,7 @@ def report_batch(
                 page_size=page_size,
                 mamba_slots=mamba_slots,
                 swa_tokens=swa_tokens,
+                generated=generated,
             )
 
     def _report_prefill(
@@ -103,9 +106,10 @@ def _report_decode(
         page_size: int,
         mamba_slots: tuple[int, int] | None = None,
         swa_tokens: tuple[int, int] | None = None,
+        generated: int | None = None,
     ) -> None:
         self._decode_forward_count += 1
-        self._decode_generated_tokens += len(batch.reqs)
+        self._decode_generated_tokens += len(batch.reqs) if generated is None else generated
         if self._decode_forward_count % self.decode_log_interval != 0:
             return
 
diff --git a/python/freetoken/tokenizer/detokenize.py b/python/freetoken/tokenizer/detokenize.py
index d138ac1a68..045e338918 100644
--- a/python/freetoken/tokenizer/detokenize.py
+++ b/python/freetoken/tokenizer/detokenize.py
@@ -103,8 +103,10 @@ def detokenize(self, msgs: List[DetokenizeMsg]) -> List[str]:
                     sent_offset=0,
                 )
             s = self.decode_map[msg.uid]
-            if not (msg.finished and msg.next_token in self.eos_token_ids):
-                s.decoded_ids.append(msg.next_token)
+            tokens = msg.next_tokens
+            if msg.finished and tokens and tokens[-1] in self.eos_token_ids:
+                tokens = tokens[:-1]
+            s.decoded_ids.extend(tokens)
             read_ids.append(s.decoded_ids[s.surr_offset :])
             surr_ids.append(s.decoded_ids[s.surr_offset : s.read_offset])
 
diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py
index fad66d746d..9888716050 100644
--- a/python/freetoken/tokenizer/server.py
+++ b/python/freetoken/tokenizer/server.py
@@ -236,7 +236,7 @@ def tokenize_worker(
                         finished=msg.finished,
                         finish_reason=msg.finish_reason,
                         matched_stop=msg.matched_stop,
-                        completion_tokens_delta=1,
+                        completion_tokens_delta=len(msg.next_tokens),
                         kv_used_pages=msg.kv_used_pages,
                         kv_total_pages=msg.kv_total_pages,
                         mamba_used_slots=msg.mamba_used_slots,
diff --git a/tests/attention/test_capture_steps.py b/tests/attention/test_capture_steps.py
new file mode 100644
index 0000000000..d7b37f6f58
--- /dev/null
+++ b/tests/attention/test_capture_steps.py
@@ -0,0 +1,143 @@
+"""Graphs captured for `steps` query rows per request (1 = decode, more = a uniform extend such as a speculative verify) replay what the eager path computes."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import freetoken.core as core
+from freetoken.attention import attention_backend_info
+from freetoken.core import Batch, Context, Req
+from freetoken.distributed import set_tp_info, try_get_tp_info
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+HEADS, KV_HEADS, HEAD_DIM = 8, 2, 128
+BS, WIDTH = 2, 64
+CACHED = (5, 17)
+
+
+def _backend(name: str):
+    config = SimpleNamespace(num_qo_heads=HEADS, num_kv_heads=KV_HEADS, head_dim=HEAD_DIM)
+    if name == "fa":
+        pytest.importorskip("sgl_kernel")
+        from freetoken.attention.fa import FlashAttentionBackend
+
+        return FlashAttentionBackend(config)
+    if name == "fi":
+        pytest.importorskip("flashinfer")
+        from freetoken.attention.fi import FlashInferBackend
+
+        return FlashInferBackend(config)
+    from freetoken.attention.triton import TritonAttentionBackend
+
+    return TritonAttentionBackend(config)
+
+
+def _context(monkeypatch) -> Context:
+    from freetoken.kvcache.mha_pool import MHAKVCache
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    device = torch.device("cuda")
+    ctx = Context(page_size=1)
+    # rows 0..BS-1 belong to the requests, row BS is the dummy request's (padded rows write there)
+    ctx.page_table = torch.arange((BS + 1) * WIDTH, dtype=torch.int32, device=device).view(BS + 1, WIDTH)
+    ctx.kv_cache = MHAKVCache(KV_HEADS, 1, HEAD_DIM, num_pages=(BS + 1) * WIDTH, page_size=1, dtype=torch.bfloat16, device=device)
+    n = (BS + 1) * WIDTH
+    ctx.kv_cache.store_kv(
+        torch.randn(n, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device),
+        torch.randn(n, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device=device),
+        torch.arange(n, dtype=torch.int32, device=device),
+        0,
+    )
+    monkeypatch.setattr(core, "_GLOBAL_CTX", ctx)
+    return ctx
+
+
+def _batch(reqs, steps: int, ctx: Context) -> Batch:
+    batch = Batch(reqs=reqs, phase="decode" if steps == 1 else "verify", rows_per_req=steps)
+    batch.padded_reqs = reqs
+    batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len, dtype=torch.int32) for r in reqs]).cuda()
+    batch.out_loc = torch.cat([ctx.page_table[r.table_idx, r.cached_len : r.device_len] for r in reqs])
+    return batch
+
+
+def _req(table_idx: int, cached: int, steps: int) -> Req:
+    return Req(input_ids=torch.zeros(cached + steps, dtype=torch.int32), table_idx=table_idx, cached_len=cached,
+               output_len=1, uid=table_idx, sampling_params=None, cache_handle=None)
+
+
+@pytest.mark.parametrize("steps", [1, 4])
+@pytest.mark.parametrize("name", ["fa", "triton", "fi"])
+def test_graph_replay_matches_eager_attention(monkeypatch, name, steps):
+    assert attention_backend_info(name).uniform_extend_graph != "none"
+    ctx = _context(monkeypatch)
+    backend = _backend(name)
+    rows = BS * steps
+    q = torch.randn(rows, HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    k = torch.randn(rows, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+    v = torch.randn(rows, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
+
+    eager = _batch([_req(i, CACHED[i], steps) for i in range(BS)], steps, ctx)
+    backend.prepare_metadata(eager)
+    expected = backend.forward(q, k, v, 0, eager)
+
+    backend.init_capture_graph(max_seq_len=WIDTH, bs_list=[BS], steps=steps)
+    capture = _batch([_req(BS, 0, steps)] * BS, steps, ctx)
+    backend.prepare_for_capture(capture)
+    static_q, static_k, static_v = q.clone(), k.clone(), v.clone()
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        backend.forward(static_q, static_k, static_v, 0, capture)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, stream=stream):
+            out = backend.forward(static_q, static_k, static_v, 0, capture)
+    torch.cuda.current_stream().wait_stream(stream)
+
+    replay = _batch([_req(i, CACHED[i], steps) for i in range(BS)], steps, ctx)
+    backend.prepare_metadata(replay)
+    backend.prepare_for_replay(replay)
+    capture.out_loc.copy_(replay.out_loc)
+    graph.replay()
+    torch.cuda.synchronize()
+    torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2)
+
+
+def test_graph_runner_returns_each_output_for_the_real_rows(monkeypatch):
+    from freetoken.engine.graph import GraphRunner
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1))
+    steps = 3
+    attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None, prepare_for_capture=lambda batch: None,
+                           prepare_for_replay=lambda batch: None)
+
+    def forward():
+        batch = core.get_global_ctx().batch
+        ids = batch.input_ids.float()
+        return ids.unsqueeze(1) * 2, ids.view(-1, steps).sum(dim=1)
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward, attn_backend=attn,
+                             cuda_graph_bs=[2, 4], cuda_graph_max_bs=None, free_memory=1024**3, max_seq_len=16,
+                             dummy_req=_req(0, 0, steps), steps=steps, phase="verify")
+        reqs = [_req(i, 1, steps) for i in range(3)]
+        assert not runner.can_use_cuda_graph(Batch(reqs=reqs, phase="decode"))
+        batch = Batch(reqs=reqs, phase="verify", rows_per_req=steps)
+        assert runner.can_use_cuda_graph(batch)
+        runner.pad_batch(batch)
+        assert batch.padded_size == 4
+        batch.input_ids = torch.arange(4 * steps, dtype=torch.int32, device="cuda")
+        batch.out_loc = torch.zeros(4 * steps, dtype=torch.int32, device="cuda")
+        batch.positions = torch.zeros(4 * steps, dtype=torch.int32, device="cuda")
+        per_row, per_req = runner.replay(batch)
+    torch.cuda.synchronize()
+    ids = torch.arange(3 * steps, dtype=torch.float32, device="cuda")
+    assert torch.equal(per_row, ids.unsqueeze(1) * 2)
+    assert torch.equal(per_req, ids.view(3, steps).sum(dim=1))
+    runner.destroy_cuda_graphs()
diff --git a/tests/engine/test_context_override.py b/tests/engine/test_context_override.py
new file mode 100644
index 0000000000..07754f8af4
--- /dev/null
+++ b/tests/engine/test_context_override.py
@@ -0,0 +1,47 @@
+"""Context.override: a draft forward sees its own KV pool and backend and none of the target's GDN pool or MoE offload cache, and everything comes back on exit."""
+
+import pytest
+
+from freetoken.core import Batch, Context
+
+
+def _ctx() -> Context:
+    ctx = Context(page_size=1)
+    ctx.kv_cache = "target_kv"
+    ctx.attn_backend = "target_backend"
+    ctx.linear_state_pool = "gdn_pool"
+    ctx.moe_offload_cache = "expert_cache"
+    return ctx
+
+
+def _state(ctx: Context) -> tuple:
+    return ctx.kv_cache, ctx.attn_backend, ctx.linear_state_pool, ctx.moe_offload_cache
+
+
+def test_override_hides_target_state_and_restores_on_exit_and_error():
+    ctx = _ctx()
+    before = _state(ctx)
+    with ctx.override(kv_cache="draft_kv", attn_backend="draft_backend", hide_target_state=True):
+        assert _state(ctx) == ("draft_kv", "draft_backend", None, None)
+        with ctx.override(kv_cache="inner_kv"):
+            assert _state(ctx) == ("inner_kv", "draft_backend", None, None)
+        assert _state(ctx) == ("draft_kv", "draft_backend", None, None)
+    assert _state(ctx) == before
+
+    with pytest.raises(RuntimeError):
+        with ctx.override(kv_cache="draft_kv", hide_target_state=True):
+            raise RuntimeError
+    assert _state(ctx) == before
+
+
+def test_forward_batch_scopes_the_override_with_the_batch():
+    ctx = _ctx()
+    batch = Batch(reqs=[], phase="decode")
+    with ctx.forward_batch(batch, kv_cache="draft_kv", hide_target_state=True):
+        assert ctx.batch is batch
+        assert _state(ctx) == ("draft_kv", "target_backend", None, None)
+        with pytest.raises(AssertionError):
+            with ctx.forward_batch(batch):
+                pass
+    assert ctx._batch is None
+    assert _state(ctx) == ("target_kv", "target_backend", "gdn_pool", "expert_cache")
diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py
index f71c9e6a0f..85a98dfdae 100644
--- a/tests/models/qwen4_exp/test_ple_disk.py
+++ b/tests/models/qwen4_exp/test_ple_disk.py
@@ -398,10 +398,9 @@ def forward():
     stream.wait_stream(torch.cuda.current_stream())
     with torch.cuda.stream(stream):
         # Exercise the capture-time upload from #697 as well as the PLE dispatch hook.
-        runner = GraphRunner(stream=stream, device=torch.device("cuda"), model=SimpleNamespace(forward=forward),
+        runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward,
                              attn_backend=attn, cuda_graph_bs=[8], cuda_graph_max_bs=None,
-                             free_memory=1024**3, max_seq_len=16,
-                             vocab_size=args.num_ngram_heads * disk.head_dim, dummy_req=req)
+                             free_memory=1024**3, max_seq_len=16, dummy_req=req, out_dtype=torch.float32)
         batch = Batch(reqs=[req] * 8, phase="decode")
         runner.pad_batch(batch)
         batch.positions = torch.full((8,), 2, dtype=torch.int32, device="cuda")
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index a29d94e9af..d1a89f6bfd 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -475,7 +475,7 @@ def fake_cuda_graph(graph, pool=None, stream=None):
         events.append("graph_exit")
 
     class FakeAttnBackend:
-        def init_capture_graph(self, max_seq_len, bs_list):
+        def init_capture_graph(self, max_seq_len, bs_list, steps=1):
             pass
 
         def prepare_for_capture(self, batch):
@@ -511,13 +511,12 @@ def reset(self):
     GraphRunner(
         stream=None,
         device=torch.device("cpu"),
-        model=FakeModel(),
+        forward=FakeModel().forward,
         attn_backend=FakeAttnBackend(),
         cuda_graph_bs=[1],
         cuda_graph_max_bs=None,
         free_memory=1024,
         max_seq_len=1,
-        vocab_size=3,
         dummy_req=dummy_req,
         moe_offload_cache=FakeOffloadCache(),
     )
diff --git a/tests/scheduler/test_abort_inflight_prefill.py b/tests/scheduler/test_abort_inflight_prefill.py
index 85c217859f..2ec5508a1b 100644
--- a/tests/scheduler/test_abort_inflight_prefill.py
+++ b/tests/scheduler/test_abort_inflight_prefill.py
@@ -72,6 +72,7 @@ def _setup():
         _last_data=None,
     )
     stub._free_req_resources = lambda req: Scheduler._free_req_resources(stub, req)
+    stub._commit_tokens = lambda req, tokens: Scheduler._commit_tokens(stub, req, tokens)
     return pool, cm, tm, dm, pm, sent, stub
 
 
diff --git a/tests/scheduler/test_cost_accounting_core.py b/tests/scheduler/test_cost_accounting_core.py
index ed741ef8cf..67ac53d19a 100644
--- a/tests/scheduler/test_cost_accounting_core.py
+++ b/tests/scheduler/test_cost_accounting_core.py
@@ -209,7 +209,7 @@ def test_normal_loop_sends_prior_sample_before_abort_terminal():
     scheduler._schedule_next_batch = lambda: None
     sent = []
     scheduler.send_result = lambda messages: sent.append(messages)
-    late_sample = DetokenizeMsg(uid=5, next_token=99, finished=False)
+    late_sample = DetokenizeMsg(uid=5, next_tokens=[99], finished=False)
     scheduler._process_last_data = lambda _data: scheduler.send_result([late_sample])
 
     Scheduler.normal_loop(scheduler)
diff --git a/tests/scheduler/test_rollback_tail.py b/tests/scheduler/test_rollback_tail.py
new file mode 100644
index 0000000000..59022bb026
--- /dev/null
+++ b/tests/scheduler/test_rollback_tail.py
@@ -0,0 +1,59 @@
+"""rollback_tail after a widened decode step: the request keeps exactly the pages of its committed tokens, the rest (and their SWA slots) go back, and the next step allocates no page twice."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from freetoken.core import Req
+from freetoken.scheduler.cache import CacheManager
+from freetoken.utils import div_ceil
+
+K = 3
+NUM_PAGES = 64
+
+
+class _SwaPool:
+    """Tracks which full-pool slots hold an SWA slot."""
+
+    swa_paged = True
+
+    def __init__(self):
+        self.held: set[int] = set()
+
+    def alloc_swa(self, indices: torch.Tensor) -> None:
+        new = set(indices.tolist())
+        assert not new & self.held, "an SWA slot was allocated twice"
+        self.held |= new
+
+    def free_swa(self, indices: torch.Tensor) -> None:
+        self.held -= set(indices.tolist())
+
+
+def _held_slots(cm: CacheManager, req: Req, page_size: int) -> set[int]:
+    pages = div_ceil(req.cached_len, page_size)
+    return set(cm.page_table[req.table_idx, : pages * page_size].tolist())
+
+
+@pytest.mark.parametrize("page_size", [1, 4])
+@pytest.mark.parametrize("cached", [8, 9, 11])
+@pytest.mark.parametrize("n", range(1, K + 2))
+def test_rollback_keeps_committed_pages_only(page_size, cached, n):
+    swa = _SwaPool()
+    page_table = torch.zeros(2, 64, dtype=torch.int32)
+    cm = CacheManager(NUM_PAGES, page_size, page_table, "naive", swa_pool=swa)
+    req = Req(input_ids=torch.zeros(cached, dtype=torch.int32), table_idx=0, cached_len=0, output_len=32,
+              uid=0, sampling_params=SimpleNamespace(), cache_handle=None)
+    cm.allocate_paged([req])
+    req.complete_one()  # cached_len = cached, the pending token sits at position `cached`
+
+    for _ in range(2):
+        entry = req.cached_len
+        req.widen_pending(K)
+        cm.allocate_paged([req])
+        req.commit_pending(n)
+        with cm.lazy_free_region():
+            cm.rollback_tail(req, req.cached_len, entry + K + 1)
+        kept = div_ceil(req.cached_len, page_size)
+        assert len(cm.free_slots) == NUM_PAGES - kept
+        assert swa.held == _held_slots(cm, req, page_size)
diff --git a/tests/server/test_message_wire.py b/tests/server/test_message_wire.py
index fdc095921d..76d8ecc397 100644
--- a/tests/server/test_message_wire.py
+++ b/tests/server/test_message_wire.py
@@ -106,7 +106,7 @@ def test_user_reply_token_deltas_round_trip():
 
 def test_detokenize_msg_carries_kv_usage_round_trip():
     msg = DetokenizeMsg(
-        uid=3, next_token=42, finished=True,
+        uid=3, next_tokens=[42], finished=True,
         kv_used_pages=10, kv_total_pages=256, gpu_mem_bytes=1 << 30,
         mamba_used_slots=7, mamba_total_slots=64,
         swa_used_tokens=8448, swa_total_tokens=76800,

From e46f0c0a31d7f8e1d695a294579d7bf069b55f94 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 06:08:52 +0000
Subject: [PATCH 2/8] refactor: rename the graph steps argument to rows_per_req

---
 python/freetoken/attention/__init__.py        |  2 +-
 python/freetoken/attention/base.py            |  8 ++--
 python/freetoken/attention/dsa.py             |  4 +-
 .../freetoken/attention/dsa_indexer_kpool.py  |  4 +-
 python/freetoken/attention/dsv41_sparse.py    |  4 +-
 python/freetoken/attention/dsv4_sparse.py     |  4 +-
 python/freetoken/attention/fa.py              | 16 +++----
 python/freetoken/attention/fi.py              | 10 ++---
 python/freetoken/attention/m3_sparse.py       |  4 +-
 python/freetoken/attention/qsa_sparse.py      |  4 +-
 python/freetoken/attention/triton.py          | 24 +++++-----
 python/freetoken/attention/trtllm.py          |  4 +-
 python/freetoken/attention/utils.py           |  8 ++--
 python/freetoken/core.py                      |  6 +--
 python/freetoken/engine/graph.py              | 34 +++++++-------
 ..._steps.py => test_capture_rows_per_req.py} | 44 +++++++++----------
 tests/moe/test_offload.py                     |  2 +-
 17 files changed, 91 insertions(+), 91 deletions(-)
 rename tests/attention/{test_capture_steps.py => test_capture_rows_per_req.py} (74%)

diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py
index 3957cff6fb..8e52845003 100644
--- a/python/freetoken/attention/__init__.py
+++ b/python/freetoken/attention/__init__.py
@@ -33,7 +33,7 @@ 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 `steps` query rows per request: "device" keeps all metadata on device, "host_plan" plans on the host before each replay, "none" cannot.
+    # 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"
 
 
diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py
index d0144eacad..0217299292 100644
--- a/python/freetoken/attention/base.py
+++ b/python/freetoken/attention/base.py
@@ -70,8 +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], steps: int = 1) -> None:
-        """Static buffers for graphs of ``steps`` query rows per request: 1 for decode, more for a uniform extend."""
+    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: ...
@@ -117,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], steps: int = 1) -> None:
-        self.decode_backend.init_capture_graph(max_seq_len, bs_list, steps)
+    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)
diff --git a/python/freetoken/attention/dsa.py b/python/freetoken/attention/dsa.py
index c2e56422ab..b565b01f51 100644
--- a/python/freetoken/attention/dsa.py
+++ b/python/freetoken/attention/dsa.py
@@ -341,8 +341,8 @@ 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], steps: int = 1) -> None:
-        if steps != 1:
+    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)
diff --git a/python/freetoken/attention/dsa_indexer_kpool.py b/python/freetoken/attention/dsa_indexer_kpool.py
index 6164496767..2db80ea886 100644
--- a/python/freetoken/attention/dsa_indexer_kpool.py
+++ b/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, steps: int = 1) -> None:
-        super().init_capture_graph(max_seq_len, bs_list, steps)
+    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:
diff --git a/python/freetoken/attention/dsv41_sparse.py b/python/freetoken/attention/dsv41_sparse.py
index d88f1571e8..0d827ffb3e 100644
--- a/python/freetoken/attention/dsv41_sparse.py
+++ b/python/freetoken/attention/dsv41_sparse.py
@@ -212,8 +212,8 @@ 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], steps: int = 1) -> None:
-        if steps != 1:
+    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)
diff --git a/python/freetoken/attention/dsv4_sparse.py b/python/freetoken/attention/dsv4_sparse.py
index 984184b8cc..bb68411af0 100644
--- a/python/freetoken/attention/dsv4_sparse.py
+++ b/python/freetoken/attention/dsv4_sparse.py
@@ -190,8 +190,8 @@ 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], steps: int = 1) -> None:
-        if steps != 1:
+    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)
diff --git a/python/freetoken/attention/fa.py b/python/freetoken/attention/fa.py
index e3d38b308c..07212e1373 100644
--- a/python/freetoken/attention/fa.py
+++ b/python/freetoken/attention/fa.py
@@ -46,7 +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_steps = 1
+        self.capture_rows_per_req = 1
         self.max_graph_bs = 0
         self.capture_bs: List[int] = []
         self.scale = config.head_dim**-0.5
@@ -121,26 +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], steps: int = 1) -> 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, steps=steps)
+        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_steps = steps
+        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
-        steps = self.capture_steps
-        capture.seq_lens[:bs].fill_(steps)
+        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=steps,
+            max_seqlen_q=rows_per_req,
             page_table=capture.page_table[:bs, :],
         )
         batch.attn_metadata = metadata
@@ -149,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 the constant arange * steps under a graph (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)
diff --git a/python/freetoken/attention/fi.py b/python/freetoken/attention/fi.py
index 73e75347f6..175c21ffaf 100644
--- a/python/freetoken/attention/fi.py
+++ b/python/freetoken/attention/fi.py
@@ -143,7 +143,7 @@ def __init__(self, config: ModelConfig) -> None:
             int, CUDAGraphBatchDecodeWithPagedKVCacheWrapper | BatchPrefillWithPagedKVCacheWrapper
         ] = {}
         self.capture: FICaptureData | None = None
-        self.capture_steps = 1
+        self.capture_rows_per_req = 1
         self.last_event = torch.cuda.Event()
         self.last_event.record()
 
@@ -269,15 +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], steps: int = 1) -> 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, steps=steps)
+        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_steps = steps
+        self.capture_rows_per_req = rows_per_req
 
     @cached_property
     def use_tensor_cores(self) -> bool:
@@ -294,7 +294,7 @@ def _make_graph_wrapper(
 
         capture = self.capture
         assert capture is not None
-        if self.capture_steps > 1:
+        if self.capture_rows_per_req > 1:
             wrapper = BatchPrefillWithPagedKVCacheWrapper(
                 self.float_workspace_buffer,
                 kv_layout="NHD",
diff --git a/python/freetoken/attention/m3_sparse.py b/python/freetoken/attention/m3_sparse.py
index 30de53da84..dc1869cb5c 100644
--- a/python/freetoken/attention/m3_sparse.py
+++ b/python/freetoken/attention/m3_sparse.py
@@ -368,8 +368,8 @@ 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], steps: int = 1) -> None:
-        if steps != 1:
+    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)
diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py
index 23089800eb..3e37d5f888 100644
--- a/python/freetoken/attention/qsa_sparse.py
+++ b/python/freetoken/attention/qsa_sparse.py
@@ -489,8 +489,8 @@ 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], steps: int = 1) -> None:
-        if steps != 1:
+    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)
diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py
index d753158102..071ee0d454 100644
--- a/python/freetoken/attention/triton.py
+++ b/python/freetoken/attention/triton.py
@@ -32,16 +32,16 @@ def create(
         num_q_heads: int,
         max_head_dim: int,
         scratch_rows: int,
-        steps: int = 1,
+        rows_per_req: int = 1,
         **kwargs,
     ):
         return cls(
             seq_lens=torch.ones((max_bs,), dtype=torch.int32, device=device),
-            positions=torch.zeros((max_bs * steps,), 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) * steps,
+            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).repeat_interleave(steps),
+            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,
@@ -72,7 +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_steps = 1
+        self.capture_rows_per_req = 1
         self.capture_bs: List[int] = []
         self.max_graph_bs = 0
         self.prefill_tile_min_q = 128
@@ -255,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], steps: int = 1) -> 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(
@@ -265,9 +265,9 @@ def init_capture_graph(self, max_seq_len: int, bs_list: List[int], steps: int =
             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),
-            steps=steps,
+            rows_per_req=rows_per_req,
         )
-        self.capture_steps = steps
+        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),
@@ -327,9 +327,9 @@ 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
-        steps = self.capture_steps
-        if steps > 1:
-            n = bs * steps
+        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)
@@ -341,7 +341,7 @@ def prepare_for_capture(self, batch: Batch) -> None:
                 q_positions=capture.positions[:n],
                 is_decode=False,
                 prefix_lens=capture.seq_lens[:bs],
-                max_q_len=steps,
+                max_q_len=rows_per_req,
                 swa_indices=self._capture_swa_indices(),
             )
             return
diff --git a/python/freetoken/attention/trtllm.py b/python/freetoken/attention/trtllm.py
index 6f76379a78..20eb2fb538 100644
--- a/python/freetoken/attention/trtllm.py
+++ b/python/freetoken/attention/trtllm.py
@@ -138,8 +138,8 @@ 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], steps: int = 1) -> None:
-        if steps != 1:
+    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)
diff --git a/python/freetoken/attention/utils.py b/python/freetoken/attention/utils.py
index 0a8c06127b..1de53e70d3 100644
--- a/python/freetoken/attention/utils.py
+++ b/python/freetoken/attention/utils.py
@@ -12,13 +12,13 @@ class BaseCaptureData:
     page_table: torch.Tensor
 
     @classmethod
-    def create(cls, max_bs: int, max_seq_len: int, device: torch.device, steps: int = 1, **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 * steps,), 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),
-            # constant under a graph: every request has exactly `steps` query rows
-            cu_seqlens_q=torch.arange(0, max_bs + 1, dtype=torch.int32, device=device) * steps,
+            # 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,
         )
diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index acea1f3e8d..f4506caf97 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -127,9 +127,9 @@ def __repr__(self) -> str:
 
 @dataclass
 class VerifyInfo:
-    """A uniform extend of `steps` rows per request over read-only recurrent state."""
+    """A uniform extend of `rows_per_req` rows per request over read-only recurrent state."""
 
-    steps: int
+    rows_per_req: int
     gdn: Any = None
     model_scratch: Any = None
 
@@ -138,7 +138,7 @@ class VerifyInfo:
 class Batch:
     reqs: List[Req]
     phase: Literal["prefill", "decode", "verify"]
-    # query rows per request: 1 on decode, steps on verify, None on prefill (ragged)
+    # query rows per request: 1 on decode, more on verify, None on prefill (ragged)
     rows_per_req: int | None = None
     # these fields should be set by scheduler
     input_ids: torch.Tensor = field(init=False)
diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py
index ff7006e4b0..1c490fc51e 100644
--- a/python/freetoken/engine/graph.py
+++ b/python/freetoken/engine/graph.py
@@ -26,16 +26,16 @@ class GraphCaptureBuffer:
     input_ids: torch.Tensor
     out_loc: torch.Tensor
     positions: torch.Tensor
-    # [3, bs * steps] t/h/w rope positions; allocated only for mrope models (else None).
+    # [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) * steps; a constant per captured bs, filled once.
+    # GDN query indptr = arange(bs+1) * rows_per_req; a constant per captured bs, filled once.
     fla_cu_seqlens: torch.Tensor
-    steps: int = 1
+    rows_per_req: int = 1
 
     @classmethod
-    def init(cls, bs: int, device: torch.device, mrope: bool = False, steps: int = 1) -> GraphCaptureBuffer:
-        n = bs * steps
+    def init(cls, bs: int, device: torch.device, mrope: bool = False, rows_per_req: int = 1) -> GraphCaptureBuffer:
+        n = bs * rows_per_req
         return GraphCaptureBuffer(
             input_ids=torch.zeros(n, dtype=torch.int32, device=device),
             out_loc=torch.zeros(n, dtype=torch.int32, device=device),
@@ -44,15 +44,15 @@ def init(cls, bs: int, device: torch.device, mrope: bool = False, steps: int = 1
                 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) * steps,
-            steps=steps,
+            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.steps)
+        _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]
@@ -66,7 +66,7 @@ def set_batch(self, batch: Batch) -> None:
         )
 
     def copy_from(self, batch: Batch) -> None:
-        _rows = slice(batch.padded_size * self.steps)
+        _rows = slice(batch.padded_size * self.rows_per_req)
         self.input_ids[_rows] = batch.input_ids
         if batch.out_loc is not None:
             self.out_loc[_rows] = batch.out_loc
@@ -128,7 +128,7 @@ def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> Non
 
 
 class GraphRunner:
-    """CUDA graphs of ``forward`` over padded batches of ``steps`` rows per request, one graph per batch size.
+    """CUDA graphs of ``forward`` over padded batches of ``rows_per_req`` rows per request, one graph per batch size.
 
     ``forward`` returns one tensor or a tuple of tensors; each is copied into a buffer sized by the largest
     captured batch, and ``replay`` returns the real requests' rows of each.
@@ -148,7 +148,7 @@ def __init__(
         moe_offload_cache: OffloadMoeCache | None = None,
         mrope: bool = False,
         *,
-        steps: int = 1,
+        rows_per_req: int = 1,
         phase: Literal["decode", "verify"] = "decode",
         ctx_override: Dict[str, Any] | None = None,
         batch_hook: Callable[[Batch], None] | None = None,
@@ -161,7 +161,7 @@ def __init__(
             free_memory=free_memory,
         )
         # the padded rows replay the dummy request, so it must carry one row per step
-        assert steps == 1 or dummy_req.extend_len == steps, (steps, dummy_req.extend_len)
+        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.max_graph_bs = max(cuda_graph_bs) if cuda_graph_bs else 0
@@ -171,7 +171,7 @@ def __init__(
         self.mrope = mrope
         self.stream = stream
         self.device = device
-        self.steps = steps
+        self.rows_per_req = rows_per_req
         self.phase = phase
         self.ctx_override = ctx_override or {}
         self.batch_hook = batch_hook
@@ -208,7 +208,7 @@ def _capture_graphs(self, max_seq_len: int):
         if self.max_graph_bs == 0:
             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, steps=self.steps)
+        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)
 
         torch.cuda.synchronize(self.device)
         torch.cuda.empty_cache()
@@ -218,7 +218,7 @@ def _capture_graphs(self, max_seq_len: int):
         free_memory = get_free_memory(self.device)
         logger.info_rank0(f"Free GPU memory before capturing CUDA graphs: {mem_GB(free_memory)}")
 
-        self.buffer = GraphCaptureBuffer.init(self.max_graph_bs, self.device, mrope=self.mrope, steps=self.steps)
+        self.buffer = GraphCaptureBuffer.init(self.max_graph_bs, self.device, mrope=self.mrope, rows_per_req=self.rows_per_req)
         self._reset_moe_offload_cache()
 
         pbar = tqdm(
@@ -233,7 +233,7 @@ def _capture_graphs(self, max_seq_len: int):
             pbar.desc = f"Capturing graphs: bs = {bs:<3} | avail_mem = {mem_GB(free_memory)}"
             pbar.refresh()
             graph = torch.cuda.CUDAGraph()
-            batch = Batch(reqs=[self.dummy_req] * bs, phase=self.phase, rows_per_req=self.steps)
+            batch = Batch(reqs=[self.dummy_req] * bs, phase=self.phase, rows_per_req=self.rows_per_req)
             batch.padded_reqs = batch.reqs
             if self.batch_hook is not None:
                 self.batch_hook(batch)
@@ -266,7 +266,7 @@ def _capture_graphs(self, max_seq_len: int):
     def can_use_cuda_graph(self, batch: Batch) -> bool:
         return (
             batch.phase == self.phase
-            and batch.rows_per_req == self.steps
+            and batch.rows_per_req == self.rows_per_req
             and batch.size <= self.max_graph_bs
         )
 
diff --git a/tests/attention/test_capture_steps.py b/tests/attention/test_capture_rows_per_req.py
similarity index 74%
rename from tests/attention/test_capture_steps.py
rename to tests/attention/test_capture_rows_per_req.py
index d7b37f6f58..aeb84d24ab 100644
--- a/tests/attention/test_capture_steps.py
+++ b/tests/attention/test_capture_rows_per_req.py
@@ -1,4 +1,4 @@
-"""Graphs captured for `steps` query rows per request (1 = decode, more = a uniform extend such as a speculative verify) replay what the eager path computes."""
+"""Graphs captured for `rows_per_req` query rows per request (1 = decode, more = a uniform extend such as a speculative verify) replay what the eager path computes."""
 
 from types import SimpleNamespace
 
@@ -55,36 +55,36 @@ def _context(monkeypatch) -> Context:
     return ctx
 
 
-def _batch(reqs, steps: int, ctx: Context) -> Batch:
-    batch = Batch(reqs=reqs, phase="decode" if steps == 1 else "verify", rows_per_req=steps)
+def _batch(reqs, rows_per_req: int, ctx: Context) -> Batch:
+    batch = Batch(reqs=reqs, phase="decode" if rows_per_req == 1 else "verify", rows_per_req=rows_per_req)
     batch.padded_reqs = reqs
     batch.positions = torch.cat([torch.arange(r.cached_len, r.device_len, dtype=torch.int32) for r in reqs]).cuda()
     batch.out_loc = torch.cat([ctx.page_table[r.table_idx, r.cached_len : r.device_len] for r in reqs])
     return batch
 
 
-def _req(table_idx: int, cached: int, steps: int) -> Req:
-    return Req(input_ids=torch.zeros(cached + steps, dtype=torch.int32), table_idx=table_idx, cached_len=cached,
+def _req(table_idx: int, cached: int, rows_per_req: int) -> Req:
+    return Req(input_ids=torch.zeros(cached + rows_per_req, dtype=torch.int32), table_idx=table_idx, cached_len=cached,
                output_len=1, uid=table_idx, sampling_params=None, cache_handle=None)
 
 
-@pytest.mark.parametrize("steps", [1, 4])
+@pytest.mark.parametrize("rows_per_req", [1, 4])
 @pytest.mark.parametrize("name", ["fa", "triton", "fi"])
-def test_graph_replay_matches_eager_attention(monkeypatch, name, steps):
+def test_graph_replay_matches_eager_attention(monkeypatch, name, rows_per_req):
     assert attention_backend_info(name).uniform_extend_graph != "none"
     ctx = _context(monkeypatch)
     backend = _backend(name)
-    rows = BS * steps
+    rows = BS * rows_per_req
     q = torch.randn(rows, HEADS, HEAD_DIM, dtype=torch.bfloat16, device="cuda")
     k = torch.randn(rows, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
     v = torch.randn(rows, KV_HEADS * HEAD_DIM, dtype=torch.bfloat16, device="cuda")
 
-    eager = _batch([_req(i, CACHED[i], steps) for i in range(BS)], steps, ctx)
+    eager = _batch([_req(i, CACHED[i], rows_per_req) for i in range(BS)], rows_per_req, ctx)
     backend.prepare_metadata(eager)
     expected = backend.forward(q, k, v, 0, eager)
 
-    backend.init_capture_graph(max_seq_len=WIDTH, bs_list=[BS], steps=steps)
-    capture = _batch([_req(BS, 0, steps)] * BS, steps, ctx)
+    backend.init_capture_graph(max_seq_len=WIDTH, bs_list=[BS], rows_per_req=rows_per_req)
+    capture = _batch([_req(BS, 0, rows_per_req)] * BS, rows_per_req, ctx)
     backend.prepare_for_capture(capture)
     static_q, static_k, static_v = q.clone(), k.clone(), v.clone()
     stream = torch.cuda.Stream()
@@ -96,7 +96,7 @@ def test_graph_replay_matches_eager_attention(monkeypatch, name, steps):
             out = backend.forward(static_q, static_k, static_v, 0, capture)
     torch.cuda.current_stream().wait_stream(stream)
 
-    replay = _batch([_req(i, CACHED[i], steps) for i in range(BS)], steps, ctx)
+    replay = _batch([_req(i, CACHED[i], rows_per_req) for i in range(BS)], rows_per_req, ctx)
     backend.prepare_metadata(replay)
     backend.prepare_for_replay(replay)
     capture.out_loc.copy_(replay.out_loc)
@@ -111,33 +111,33 @@ def test_graph_runner_returns_each_output_for_the_real_rows(monkeypatch):
     if try_get_tp_info() is None:
         set_tp_info(rank=0, size=1)
     monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1))
-    steps = 3
+    rows_per_req = 3
     attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None, prepare_for_capture=lambda batch: None,
                            prepare_for_replay=lambda batch: None)
 
     def forward():
         batch = core.get_global_ctx().batch
         ids = batch.input_ids.float()
-        return ids.unsqueeze(1) * 2, ids.view(-1, steps).sum(dim=1)
+        return ids.unsqueeze(1) * 2, ids.view(-1, rows_per_req).sum(dim=1)
 
     stream = torch.cuda.Stream()
     stream.wait_stream(torch.cuda.current_stream())
     with torch.cuda.stream(stream):
         runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward, attn_backend=attn,
                              cuda_graph_bs=[2, 4], cuda_graph_max_bs=None, free_memory=1024**3, max_seq_len=16,
-                             dummy_req=_req(0, 0, steps), steps=steps, phase="verify")
-        reqs = [_req(i, 1, steps) for i in range(3)]
+                             dummy_req=_req(0, 0, rows_per_req), rows_per_req=rows_per_req, phase="verify")
+        reqs = [_req(i, 1, rows_per_req) for i in range(3)]
         assert not runner.can_use_cuda_graph(Batch(reqs=reqs, phase="decode"))
-        batch = Batch(reqs=reqs, phase="verify", rows_per_req=steps)
+        batch = Batch(reqs=reqs, phase="verify", rows_per_req=rows_per_req)
         assert runner.can_use_cuda_graph(batch)
         runner.pad_batch(batch)
         assert batch.padded_size == 4
-        batch.input_ids = torch.arange(4 * steps, dtype=torch.int32, device="cuda")
-        batch.out_loc = torch.zeros(4 * steps, dtype=torch.int32, device="cuda")
-        batch.positions = torch.zeros(4 * steps, dtype=torch.int32, device="cuda")
+        batch.input_ids = torch.arange(4 * rows_per_req, dtype=torch.int32, device="cuda")
+        batch.out_loc = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
+        batch.positions = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
         per_row, per_req = runner.replay(batch)
     torch.cuda.synchronize()
-    ids = torch.arange(3 * steps, dtype=torch.float32, device="cuda")
+    ids = torch.arange(3 * rows_per_req, dtype=torch.float32, device="cuda")
     assert torch.equal(per_row, ids.unsqueeze(1) * 2)
-    assert torch.equal(per_req, ids.view(3, steps).sum(dim=1))
+    assert torch.equal(per_req, ids.view(3, rows_per_req).sum(dim=1))
     runner.destroy_cuda_graphs()
diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py
index d1a89f6bfd..3baddeb523 100644
--- a/tests/moe/test_offload.py
+++ b/tests/moe/test_offload.py
@@ -475,7 +475,7 @@ def fake_cuda_graph(graph, pool=None, stream=None):
         events.append("graph_exit")
 
     class FakeAttnBackend:
-        def init_capture_graph(self, max_seq_len, bs_list, steps=1):
+        def init_capture_graph(self, max_seq_len, bs_list, rows_per_req=1):
             pass
 
         def prepare_for_capture(self, batch):

From 55628a1a559bbf2e4891752d1b10626abea75f75 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 07:55:30 +0000
Subject: [PATCH 3/8] refactor: drop VerifyInfo and require rows_per_req on
 verify batches

---
 python/freetoken/core.py                     | 14 ++---
 python/freetoken/scheduler/decode.py         |  3 +-
 tests/attention/test_capture_rows_per_req.py | 42 +--------------
 tests/engine/test_graph_runner.py            | 55 ++++++++++++++++++++
 4 files changed, 62 insertions(+), 52 deletions(-)
 create mode 100644 tests/engine/test_graph_runner.py

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index f4506caf97..d180ca5b0d 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -2,7 +2,7 @@
 
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, List, Literal, Tuple
+from typing import TYPE_CHECKING, List, Literal, Tuple
 
 import torch
 
@@ -125,15 +125,6 @@ def __repr__(self) -> str:
 
 
 
-@dataclass
-class VerifyInfo:
-    """A uniform extend of `rows_per_req` rows per request over read-only recurrent state."""
-
-    rows_per_req: int
-    gdn: Any = None
-    model_scratch: Any = None
-
-
 @dataclass
 class Batch:
     reqs: List[Req]
@@ -183,11 +174,12 @@ class Batch:
     prompt_admissions: List[Tuple[int, int, int]] = field(default_factory=list, init=False)
     # per-row hidden states a draft model reads from this forward; None when nothing reads them
     features: "FeatureSink | None" = field(default=None, init=False)
-    verify: VerifyInfo | None = field(default=None, init=False)
 
     def __post_init__(self) -> None:
         if self.rows_per_req is None and self.phase == "decode":
             self.rows_per_req = 1
+        # without it, input_mapping[0][::None] would read every row as its own request
+        assert self.phase != "verify" or self.rows_per_req is not None, "a verify batch needs rows_per_req"
 
     @property
     def is_prefill(self) -> bool:
diff --git a/python/freetoken/scheduler/decode.py b/python/freetoken/scheduler/decode.py
index be64726e20..c5c75ef2e4 100644
--- a/python/freetoken/scheduler/decode.py
+++ b/python/freetoken/scheduler/decode.py
@@ -28,7 +28,8 @@ def abort_req(self, uid: int) -> Req | None:
 
     @property
     def inflight_tokens(self) -> int:
-        tokens_reserved = (self.page_size - 1 + self.token_headroom) * len(self.running_reqs)  # 1 page reserved
+        # each request's partly filled page, plus the provisional rows a decode step may add
+        tokens_reserved = (self.page_size - 1 + self.token_headroom) * len(self.running_reqs)
         return sum(req.remain_len for req in self.running_reqs) + tokens_reserved
 
     def schedule_next_batch(self) -> Batch | None:
diff --git a/tests/attention/test_capture_rows_per_req.py b/tests/attention/test_capture_rows_per_req.py
index aeb84d24ab..a7469f1ba5 100644
--- a/tests/attention/test_capture_rows_per_req.py
+++ b/tests/attention/test_capture_rows_per_req.py
@@ -1,4 +1,4 @@
-"""Graphs captured for `rows_per_req` query rows per request (1 = decode, more = a uniform extend such as a speculative verify) replay what the eager path computes."""
+"""Attention graphs captured for `rows_per_req` query rows per request (1 = decode, more = a uniform extend such as a speculative verify) replay what the eager path computes."""
 
 from types import SimpleNamespace
 
@@ -102,42 +102,4 @@ def test_graph_replay_matches_eager_attention(monkeypatch, name, rows_per_req):
     capture.out_loc.copy_(replay.out_loc)
     graph.replay()
     torch.cuda.synchronize()
-    torch.testing.assert_close(out, expected, atol=2e-2, rtol=2e-2)
-
-
-def test_graph_runner_returns_each_output_for_the_real_rows(monkeypatch):
-    from freetoken.engine.graph import GraphRunner
-
-    if try_get_tp_info() is None:
-        set_tp_info(rank=0, size=1)
-    monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1))
-    rows_per_req = 3
-    attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None, prepare_for_capture=lambda batch: None,
-                           prepare_for_replay=lambda batch: None)
-
-    def forward():
-        batch = core.get_global_ctx().batch
-        ids = batch.input_ids.float()
-        return ids.unsqueeze(1) * 2, ids.view(-1, rows_per_req).sum(dim=1)
-
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward, attn_backend=attn,
-                             cuda_graph_bs=[2, 4], cuda_graph_max_bs=None, free_memory=1024**3, max_seq_len=16,
-                             dummy_req=_req(0, 0, rows_per_req), rows_per_req=rows_per_req, phase="verify")
-        reqs = [_req(i, 1, rows_per_req) for i in range(3)]
-        assert not runner.can_use_cuda_graph(Batch(reqs=reqs, phase="decode"))
-        batch = Batch(reqs=reqs, phase="verify", rows_per_req=rows_per_req)
-        assert runner.can_use_cuda_graph(batch)
-        runner.pad_batch(batch)
-        assert batch.padded_size == 4
-        batch.input_ids = torch.arange(4 * rows_per_req, dtype=torch.int32, device="cuda")
-        batch.out_loc = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
-        batch.positions = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
-        per_row, per_req = runner.replay(batch)
-    torch.cuda.synchronize()
-    ids = torch.arange(3 * rows_per_req, dtype=torch.float32, device="cuda")
-    assert torch.equal(per_row, ids.unsqueeze(1) * 2)
-    assert torch.equal(per_req, ids.view(3, rows_per_req).sum(dim=1))
-    runner.destroy_cuda_graphs()
+    torch.testing.assert_close(out, expected, atol=1e-2, rtol=1e-2)
diff --git a/tests/engine/test_graph_runner.py b/tests/engine/test_graph_runner.py
new file mode 100644
index 0000000000..4cf4e20bf1
--- /dev/null
+++ b/tests/engine/test_graph_runner.py
@@ -0,0 +1,55 @@
+"""GraphRunner over `rows_per_req` rows per request: it picks graphs by phase and row count, and replay returns every output of `forward` for the real rows only."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import freetoken.core as core
+from freetoken.core import Batch, Context, Req
+from freetoken.distributed import set_tp_info, try_get_tp_info
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
+
+
+def _req(table_idx: int, cached: int, rows_per_req: int) -> Req:
+    return Req(input_ids=torch.zeros(cached + rows_per_req, dtype=torch.int32), table_idx=table_idx, cached_len=cached,
+               output_len=1, uid=table_idx, sampling_params=None, cache_handle=None)
+
+
+def test_graph_runner_returns_each_output_for_the_real_rows(monkeypatch):
+    from freetoken.engine.graph import GraphRunner
+
+    if try_get_tp_info() is None:
+        set_tp_info(rank=0, size=1)
+    monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1))
+    rows_per_req = 3
+    attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None, prepare_for_capture=lambda batch: None,
+                           prepare_for_replay=lambda batch: None)
+
+    def forward():
+        batch = core.get_global_ctx().batch
+        ids = batch.input_ids.float()
+        return ids.unsqueeze(1) * 2, ids.view(-1, rows_per_req).sum(dim=1)
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward, attn_backend=attn,
+                             cuda_graph_bs=[2, 4], cuda_graph_max_bs=None, free_memory=1024**3, max_seq_len=16,
+                             dummy_req=_req(0, 0, rows_per_req), rows_per_req=rows_per_req, phase="verify")
+        reqs = [_req(i, 1, rows_per_req) for i in range(3)]
+        assert not runner.can_use_cuda_graph(Batch(reqs=reqs, phase="decode"))
+        batch = Batch(reqs=reqs, phase="verify", rows_per_req=rows_per_req)
+        assert runner.can_use_cuda_graph(batch)
+        runner.pad_batch(batch)
+        assert batch.padded_size == 4
+        batch.input_ids = torch.arange(4 * rows_per_req, dtype=torch.int32, device="cuda")
+        batch.out_loc = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
+        batch.positions = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
+        per_row, per_req = runner.replay(batch)
+    torch.cuda.synchronize()
+    ids = torch.arange(3 * rows_per_req, dtype=torch.float32, device="cuda")
+    assert torch.equal(per_row, ids.unsqueeze(1) * 2)
+    assert torch.equal(per_req, ids.view(3, rows_per_req).sum(dim=1))
+    runner.destroy_cuda_graphs()

From bd5e7d56f95e041f9b60e1d7f83056f2b90a2d6b Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 18:48:07 +0000
Subject: [PATCH 4/8] refactor: fix stale seam comments and read
 decode_rows_per_req directly

---
 python/freetoken/engine/engine.py      |  2 +-
 python/freetoken/engine/graph.py       | 10 +++-------
 python/freetoken/models/blocks.py      |  2 +-
 tests/kvcache/test_kv_cache_rebuild.py |  2 +-
 4 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 1b835648d2..fd329fceb2 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -114,7 +114,7 @@ def _page_table_width(max_seq_len: int, page_size: int) -> int:
 def _seq_limits(config, num_tokens: int) -> tuple[int, int]:
     """(longest request, page-table width). A decode step may widen a request by decode_rows_per_req - 1 provisional rows past its length, so requests get that much less and the table still covers the widened tail."""
     limit = min(config.max_seq_len, num_tokens)
-    return limit - (getattr(config, "decode_rows_per_req", 1) - 1), _page_table_width(limit, config.page_size)
+    return limit - (config.decode_rows_per_req - 1), _page_table_width(limit, config.page_size)
 
 
 def _required_attn_types(model_config) -> frozenset[AttnType]:
diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py
index 1c490fc51e..eb71115e79 100644
--- a/python/freetoken/engine/graph.py
+++ b/python/freetoken/engine/graph.py
@@ -59,7 +59,7 @@ def set_batch(self, batch: Batch) -> None:
         if self.mrope_positions is not None:
             batch.mrope_positions = self.mrope_positions[:, _rows]
         batch.linear_table_idx = self.table_idx[:bs]
-        # Decode GDN metadata reads the persistent cu_seqlens (constant arange) and the
+        # Decode GDN metadata reads the persistent cu_seqlens (constant arange * rows_per_req) and the
         # persistent 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]
@@ -128,11 +128,7 @@ def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> Non
 
 
 class GraphRunner:
-    """CUDA graphs of ``forward`` over padded batches of ``rows_per_req`` rows per request, one graph per batch size.
-
-    ``forward`` returns one tensor or a tuple of tensors; each is copied into a buffer sized by the largest
-    captured batch, and ``replay`` returns the real requests' rows of each.
-    """
+    """CUDA graphs of ``forward`` over padded batches of ``rows_per_req`` rows per request, one graph per batch size; ``replay`` returns the real requests' rows of each output."""
 
     def __init__(
         self,
@@ -160,7 +156,7 @@ def __init__(
             cuda_graph_max_bs=cuda_graph_max_bs,
             free_memory=free_memory,
         )
-        # the padded rows replay the dummy request, so it must carry one row per step
+        # the padded rows replay the dummy request, so it must carry rows_per_req rows
         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
diff --git a/python/freetoken/models/blocks.py b/python/freetoken/models/blocks.py
index 67c6a54923..a7a095a291 100644
--- a/python/freetoken/models/blocks.py
+++ b/python/freetoken/models/blocks.py
@@ -55,7 +55,7 @@ class FeatureSink(Protocol):
     wanted_layers: frozenset[int]
 
     def final_hidden(self, x: torch.Tensor) -> None:
-        """Every row's hidden state before the LM head; the first call in a forward wins, so a model may call it before the head does."""
+        """Every row's hidden state before the LM head; the LM head calls it once per forward."""
         ...
 
     def layer_output(self, layer_id: int, x: torch.Tensor, num_streams: int = 1) -> None:
diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py
index 0e3b7050e8..be1e87c3a4 100644
--- a/tests/kvcache/test_kv_cache_rebuild.py
+++ b/tests/kvcache/test_kv_cache_rebuild.py
@@ -227,7 +227,7 @@ def test_dsv4_refresh_seq_state_tracks_page_table_width():
 
     P, mrr = 128, 4
     config = SimpleNamespace(
-        max_seq_len=65536, page_size=P, max_running_req=mrr,
+        max_seq_len=65536, page_size=P, max_running_req=mrr, decode_rows_per_req=1,
         model_config=SimpleNamespace(dsv4_args=SimpleNamespace(window_size=P)),
     )
     # attach_page_table only re-points the pool's full_loc_map; the decode snapshot belongs to

From 3b8f8735c0d2110ba6b280724e9b1f32a42d15c9 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 19:10:55 +0000
Subject: [PATCH 5/8] refactor: reword two GraphRunner comments for multi-row
 graphs

---
 python/freetoken/engine/graph.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py
index eb71115e79..42306918be 100644
--- a/python/freetoken/engine/graph.py
+++ b/python/freetoken/engine/graph.py
@@ -59,8 +59,7 @@ def set_batch(self, batch: Batch) -> None:
         if self.mrope_positions is not None:
             batch.mrope_positions = self.mrope_positions[:, _rows]
         batch.linear_table_idx = self.table_idx[:bs]
-        # Decode GDN metadata reads the persistent cu_seqlens (constant arange * rows_per_req) and the
-        # persistent table_idx slot map, so the captured kernels see stable addresses.
+        # 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]
         )
@@ -289,7 +288,7 @@ def pad_batch(self, batch: Batch) -> None:
     # NOTE: This must be called before freeing NCCL resources to prevent program hang
     def destroy_cuda_graphs(self) -> None:
         # Drop the CUDAGraph objects (and the shared mempool they hold) AND the static
-        # GraphCaptureBuffer tensors (input/out_loc/positions/...) and output buffers ([max_bs, vocab] logits).
+        # GraphCaptureBuffer tensors (input/out_loc/positions/...) and the output buffers (_out).
         # Dropping the references is the load-bearing step; without it a runtime rebuild's
         # free-before-alloc cannot reclaim this GPU memory. empty_cache() is left to the
         # caller / next capture (GraphRunner._capture_graphs already runs it).

From 6d52946910af09b05d326d3c6e0135e99b51a452 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 21:35:03 +0000
Subject: [PATCH 6/8] fix(engine): scale the GPU-decode MoE pool limits by
 decode_rows_per_req

---
 python/freetoken/engine/engine.py | 9 +++++----
 tests/engine/test_cache_budget.py | 9 ++++++++-
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index fd329fceb2..0474680067 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -47,7 +47,8 @@ def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) ->
     small = [s for s in pool_sizes if s < num_experts]
     if not small:
         return
-    limit = min(small) // config.model_config.num_experts_per_tok
+    # a decode step routes every row of every request, decode_rows_per_req rows each
+    limit = min(small) // (config.model_config.num_experts_per_tok * config.decode_rows_per_req)
     if config.max_running_req > limit:
         logger.warning_rank0(
             f"the smallest MoE slot pool has {min(small)} slots; lowering max_running_req "
@@ -810,8 +811,8 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
                 f"(--moe-cache-scope per-layer gives each layer its own)"
             )
         self._moe_slot_limit = _slot_limit(layer_methods) if layer_methods is not None else (method.slot_limit() if method is not None else None)
-        # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one token per step
-        pools = self._moe_pools(1)
+        # with GPU decode, a pool below num_experts slots bounds the batch; the startup split keeps room for one request's decode step
+        pools = self._moe_pools(config.decode_rows_per_req if decode_target == "gpu" else 1)
         size_cap = sum(p.max_num_slots for p in pools)
         if config.moe_cache_auto:
             size, pages, overlap = self._resolve_auto_moe_cache_size(config, banks, method, max_slots=size_cap)
@@ -1148,7 +1149,7 @@ def rebuild_runtime_cache(
                 )
             try:
                 # a rebuild does not lower max_running_req again, so with GPU decode every pool must hold a full batch; hybrid caps its fetches instead
-                batch = config.max_running_req if self.moe_offload_cache.decode_target == "gpu" else 1
+                batch = config.max_running_req * config.decode_rows_per_req if self.moe_offload_cache.decode_target == "gpu" else 1
                 moe_pool_sizes = split_slots(moe_cache_size, self._moe_pools(batch), self._moe_expert_counts)
                 self.moe_offload_cache.validate_rebuild(moe_cache_size, moe_pool_sizes)
             except ValueError as e:
diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py
index 44130d3549..5c1f93b348 100644
--- a/tests/engine/test_cache_budget.py
+++ b/tests/engine/test_cache_budget.py
@@ -569,7 +569,14 @@ def test_moe_pool_flags_and_decode_batch_clamp():
     # GPU decode needs batch * top_k distinct experts to fit the smallest pool below num_experts
     config = SimpleNamespace(
         model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4),
-        max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4],
+        max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], decode_rows_per_req=1,
     )
     _clamp_decode_batch_to_pools(config, [12, 16, 9])
     assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (2, 2, [1, 2])
+    # two rows per request route twice the experts: 9 // (4 * 2) leaves one request
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(num_experts=16, num_experts_per_tok=4),
+        max_running_req=4, cuda_graph_max_bs=None, cuda_graph_bs=[1, 2, 4], decode_rows_per_req=2,
+    )
+    _clamp_decode_batch_to_pools(config, [12, 16, 9])
+    assert (config.max_running_req, config.cuda_graph_max_bs, config.cuda_graph_bs) == (1, 1, [1])

From 1cb80f7b58f70ee5ed292870f6926bf79e6a0ab4 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 22:45:47 +0000
Subject: [PATCH 7/8] refactor: drop the unread Context.moe_offload_cache field

---
 python/freetoken/core.py              |  5 +----
 python/freetoken/engine/engine.py     |  1 -
 python/freetoken/moe/offload_cache.py |  2 +-
 tests/engine/test_context_override.py | 15 +++++++--------
 4 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/python/freetoken/core.py b/python/freetoken/core.py
index d180ca5b0d..f37ef84f89 100644
--- a/python/freetoken/core.py
+++ b/python/freetoken/core.py
@@ -12,7 +12,6 @@
     from freetoken.kvcache import BaseCacheHandle, BaseKVCachePool
     from freetoken.kvcache.linear_state_pool import LinearStatePool
     from freetoken.models.blocks import FeatureSink
-    from freetoken.moe.offload_cache import OffloadMoeCache
 
 
 @dataclass
@@ -207,7 +206,6 @@ class Context:
     # NOTE: this table always treat page_size = 1
     page_table: torch.Tensor = field(init=False)
     attn_backend: BaseAttnBackend = field(init=False)
-    moe_offload_cache: OffloadMoeCache | None = None
     kv_cache: BaseKVCachePool = field(init=False)
     # Per-request recurrent state for GatedDeltaNet layers; set by the engine for
     # hybrid linear-attention models, otherwise None.
@@ -227,7 +225,7 @@ def override(
         attn_backend: BaseAttnBackend | None = None,
         hide_target_state: bool = False,
     ):
-        """Swap the KV pool and attention backend for a draft forward; hide_target_state also hides the target's GDN pool and MoE offload cache."""
+        """Swap the KV pool and attention backend for a draft forward; hide_target_state also hides the target's GDN pool."""
         swaps = {}
         if kv_cache is not None:
             swaps["kv_cache"] = kv_cache
@@ -235,7 +233,6 @@ def override(
             swaps["attn_backend"] = attn_backend
         if hide_target_state:
             swaps["linear_state_pool"] = None
-            swaps["moe_offload_cache"] = None
         saved = {name: getattr(self, name) for name in swaps}
         try:
             for name, value in swaps.items():
diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py
index 0474680067..eca7bddd69 100644
--- a/python/freetoken/engine/engine.py
+++ b/python/freetoken/engine/engine.py
@@ -908,7 +908,6 @@ def _init_offload_moe_cache(self, config: EngineConfig) -> OffloadMoeCache:
             self._resolve_hybrid_fetch(config, cache)
         if locked_fetch:
             attach_locked_fetch(cache, config, staging, self._profile_layer_rows)
-        self.ctx.moe_offload_cache = cache
         self.moe_offload_cache = cache
         return cache
 
diff --git a/python/freetoken/moe/offload_cache.py b/python/freetoken/moe/offload_cache.py
index 4a6a297df9..6610b128cd 100644
--- a/python/freetoken/moe/offload_cache.py
+++ b/python/freetoken/moe/offload_cache.py
@@ -690,7 +690,7 @@ def rebuild(self, cache_size: int, pool_sizes: list[int] | None = None) -> None:
         reloads banks. Tears down prefill-overlap buffers first (their views alias
         the old ``bank_caches``), frees the old GPU tensors, then reallocates. Slots
         cold-start after rebuild. Object identity is preserved so attached layers and
-        ``ctx.moe_offload_cache`` stay valid.
+        ``Engine.moe_offload_cache`` stay valid.
         """
         assert self.bank_sources, "set_bank_sources must run before rebuild"
         sizes = [cache_size] if pool_sizes is None else list(pool_sizes)
diff --git a/tests/engine/test_context_override.py b/tests/engine/test_context_override.py
index 07754f8af4..2052bc8325 100644
--- a/tests/engine/test_context_override.py
+++ b/tests/engine/test_context_override.py
@@ -1,4 +1,4 @@
-"""Context.override: a draft forward sees its own KV pool and backend and none of the target's GDN pool or MoE offload cache, and everything comes back on exit."""
+"""Context.override: a draft forward sees its own KV pool and backend and none of the target's GDN pool, and everything comes back on exit."""
 
 import pytest
 
@@ -10,22 +10,21 @@ def _ctx() -> Context:
     ctx.kv_cache = "target_kv"
     ctx.attn_backend = "target_backend"
     ctx.linear_state_pool = "gdn_pool"
-    ctx.moe_offload_cache = "expert_cache"
     return ctx
 
 
 def _state(ctx: Context) -> tuple:
-    return ctx.kv_cache, ctx.attn_backend, ctx.linear_state_pool, ctx.moe_offload_cache
+    return ctx.kv_cache, ctx.attn_backend, ctx.linear_state_pool
 
 
 def test_override_hides_target_state_and_restores_on_exit_and_error():
     ctx = _ctx()
     before = _state(ctx)
     with ctx.override(kv_cache="draft_kv", attn_backend="draft_backend", hide_target_state=True):
-        assert _state(ctx) == ("draft_kv", "draft_backend", None, None)
+        assert _state(ctx) == ("draft_kv", "draft_backend", None)
         with ctx.override(kv_cache="inner_kv"):
-            assert _state(ctx) == ("inner_kv", "draft_backend", None, None)
-        assert _state(ctx) == ("draft_kv", "draft_backend", None, None)
+            assert _state(ctx) == ("inner_kv", "draft_backend", None)
+        assert _state(ctx) == ("draft_kv", "draft_backend", None)
     assert _state(ctx) == before
 
     with pytest.raises(RuntimeError):
@@ -39,9 +38,9 @@ def test_forward_batch_scopes_the_override_with_the_batch():
     batch = Batch(reqs=[], phase="decode")
     with ctx.forward_batch(batch, kv_cache="draft_kv", hide_target_state=True):
         assert ctx.batch is batch
-        assert _state(ctx) == ("draft_kv", "target_backend", None, None)
+        assert _state(ctx) == ("draft_kv", "target_backend", None)
         with pytest.raises(AssertionError):
             with ctx.forward_batch(batch):
                 pass
     assert ctx._batch is None
-    assert _state(ctx) == ("target_kv", "target_backend", "gdn_pool", "expert_cache")
+    assert _state(ctx) == ("target_kv", "target_backend", "gdn_pool")

From 49b52c6e4ff0120a4ff77f72ede2674534671668 Mon Sep 17 00:00:00 2001
From: Xiaoze Fan <jasonfxz@sjtu.edu.cn>
Date: Tue, 27 Oct 2026 22:54:22 +0000
Subject: [PATCH 8/8] test: trim the multi-row seam tests

---
 tests/engine/test_context_override.py | 46 ----------------------
 tests/engine/test_graph_runner.py     | 55 ---------------------------
 tests/scheduler/test_rollback_tail.py | 29 ++------------
 3 files changed, 3 insertions(+), 127 deletions(-)
 delete mode 100644 tests/engine/test_context_override.py
 delete mode 100644 tests/engine/test_graph_runner.py

diff --git a/tests/engine/test_context_override.py b/tests/engine/test_context_override.py
deleted file mode 100644
index 2052bc8325..0000000000
--- a/tests/engine/test_context_override.py
+++ /dev/null
@@ -1,46 +0,0 @@
-"""Context.override: a draft forward sees its own KV pool and backend and none of the target's GDN pool, and everything comes back on exit."""
-
-import pytest
-
-from freetoken.core import Batch, Context
-
-
-def _ctx() -> Context:
-    ctx = Context(page_size=1)
-    ctx.kv_cache = "target_kv"
-    ctx.attn_backend = "target_backend"
-    ctx.linear_state_pool = "gdn_pool"
-    return ctx
-
-
-def _state(ctx: Context) -> tuple:
-    return ctx.kv_cache, ctx.attn_backend, ctx.linear_state_pool
-
-
-def test_override_hides_target_state_and_restores_on_exit_and_error():
-    ctx = _ctx()
-    before = _state(ctx)
-    with ctx.override(kv_cache="draft_kv", attn_backend="draft_backend", hide_target_state=True):
-        assert _state(ctx) == ("draft_kv", "draft_backend", None)
-        with ctx.override(kv_cache="inner_kv"):
-            assert _state(ctx) == ("inner_kv", "draft_backend", None)
-        assert _state(ctx) == ("draft_kv", "draft_backend", None)
-    assert _state(ctx) == before
-
-    with pytest.raises(RuntimeError):
-        with ctx.override(kv_cache="draft_kv", hide_target_state=True):
-            raise RuntimeError
-    assert _state(ctx) == before
-
-
-def test_forward_batch_scopes_the_override_with_the_batch():
-    ctx = _ctx()
-    batch = Batch(reqs=[], phase="decode")
-    with ctx.forward_batch(batch, kv_cache="draft_kv", hide_target_state=True):
-        assert ctx.batch is batch
-        assert _state(ctx) == ("draft_kv", "target_backend", None)
-        with pytest.raises(AssertionError):
-            with ctx.forward_batch(batch):
-                pass
-    assert ctx._batch is None
-    assert _state(ctx) == ("target_kv", "target_backend", "gdn_pool")
diff --git a/tests/engine/test_graph_runner.py b/tests/engine/test_graph_runner.py
deleted file mode 100644
index 4cf4e20bf1..0000000000
--- a/tests/engine/test_graph_runner.py
+++ /dev/null
@@ -1,55 +0,0 @@
-"""GraphRunner over `rows_per_req` rows per request: it picks graphs by phase and row count, and replay returns every output of `forward` for the real rows only."""
-
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-import freetoken.core as core
-from freetoken.core import Batch, Context, Req
-from freetoken.distributed import set_tp_info, try_get_tp_info
-
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
-
-
-def _req(table_idx: int, cached: int, rows_per_req: int) -> Req:
-    return Req(input_ids=torch.zeros(cached + rows_per_req, dtype=torch.int32), table_idx=table_idx, cached_len=cached,
-               output_len=1, uid=table_idx, sampling_params=None, cache_handle=None)
-
-
-def test_graph_runner_returns_each_output_for_the_real_rows(monkeypatch):
-    from freetoken.engine.graph import GraphRunner
-
-    if try_get_tp_info() is None:
-        set_tp_info(rank=0, size=1)
-    monkeypatch.setattr(core, "_GLOBAL_CTX", Context(page_size=1))
-    rows_per_req = 3
-    attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None, prepare_for_capture=lambda batch: None,
-                           prepare_for_replay=lambda batch: None)
-
-    def forward():
-        batch = core.get_global_ctx().batch
-        ids = batch.input_ids.float()
-        return ids.unsqueeze(1) * 2, ids.view(-1, rows_per_req).sum(dim=1)
-
-    stream = torch.cuda.Stream()
-    stream.wait_stream(torch.cuda.current_stream())
-    with torch.cuda.stream(stream):
-        runner = GraphRunner(stream=stream, device=torch.device("cuda"), forward=forward, attn_backend=attn,
-                             cuda_graph_bs=[2, 4], cuda_graph_max_bs=None, free_memory=1024**3, max_seq_len=16,
-                             dummy_req=_req(0, 0, rows_per_req), rows_per_req=rows_per_req, phase="verify")
-        reqs = [_req(i, 1, rows_per_req) for i in range(3)]
-        assert not runner.can_use_cuda_graph(Batch(reqs=reqs, phase="decode"))
-        batch = Batch(reqs=reqs, phase="verify", rows_per_req=rows_per_req)
-        assert runner.can_use_cuda_graph(batch)
-        runner.pad_batch(batch)
-        assert batch.padded_size == 4
-        batch.input_ids = torch.arange(4 * rows_per_req, dtype=torch.int32, device="cuda")
-        batch.out_loc = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
-        batch.positions = torch.zeros(4 * rows_per_req, dtype=torch.int32, device="cuda")
-        per_row, per_req = runner.replay(batch)
-    torch.cuda.synchronize()
-    ids = torch.arange(3 * rows_per_req, dtype=torch.float32, device="cuda")
-    assert torch.equal(per_row, ids.unsqueeze(1) * 2)
-    assert torch.equal(per_req, ids.view(3, rows_per_req).sum(dim=1))
-    runner.destroy_cuda_graphs()
diff --git a/tests/scheduler/test_rollback_tail.py b/tests/scheduler/test_rollback_tail.py
index 59022bb026..da74b188c2 100644
--- a/tests/scheduler/test_rollback_tail.py
+++ b/tests/scheduler/test_rollback_tail.py
@@ -1,4 +1,4 @@
-"""rollback_tail after a widened decode step: the request keeps exactly the pages of its committed tokens, the rest (and their SWA slots) go back, and the next step allocates no page twice."""
+"""rollback_tail after a widened decode step: the request keeps exactly the pages of its committed tokens, the rest go back, and the next step allocates no page twice."""
 
 from types import SimpleNamespace
 
@@ -13,35 +13,12 @@
 NUM_PAGES = 64
 
 
-class _SwaPool:
-    """Tracks which full-pool slots hold an SWA slot."""
-
-    swa_paged = True
-
-    def __init__(self):
-        self.held: set[int] = set()
-
-    def alloc_swa(self, indices: torch.Tensor) -> None:
-        new = set(indices.tolist())
-        assert not new & self.held, "an SWA slot was allocated twice"
-        self.held |= new
-
-    def free_swa(self, indices: torch.Tensor) -> None:
-        self.held -= set(indices.tolist())
-
-
-def _held_slots(cm: CacheManager, req: Req, page_size: int) -> set[int]:
-    pages = div_ceil(req.cached_len, page_size)
-    return set(cm.page_table[req.table_idx, : pages * page_size].tolist())
-
-
 @pytest.mark.parametrize("page_size", [1, 4])
 @pytest.mark.parametrize("cached", [8, 9, 11])
 @pytest.mark.parametrize("n", range(1, K + 2))
 def test_rollback_keeps_committed_pages_only(page_size, cached, n):
-    swa = _SwaPool()
     page_table = torch.zeros(2, 64, dtype=torch.int32)
-    cm = CacheManager(NUM_PAGES, page_size, page_table, "naive", swa_pool=swa)
+    cm = CacheManager(NUM_PAGES, page_size, page_table, "naive")
     req = Req(input_ids=torch.zeros(cached, dtype=torch.int32), table_idx=0, cached_len=0, output_len=32,
               uid=0, sampling_params=SimpleNamespace(), cache_handle=None)
     cm.allocate_paged([req])
@@ -56,4 +33,4 @@ def test_rollback_keeps_committed_pages_only(page_size, cached, n):
             cm.rollback_tail(req, req.cached_len, entry + K + 1)
         kept = div_ceil(req.cached_len, page_size)
         assert len(cm.free_slots) == NUM_PAGES - kept
-        assert swa.held == _held_slots(cm, req, page_size)
+        assert not set(cm.free_slots.tolist()) & set(cm.page_table[0, : kept * page_size : page_size].tolist())