Refactor / simplify MultiLayerEagleDraftExtendCudaGraphRunner to use rotation (#28492)

python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py

@@ -15,32 +15,23 @@
 from __future__ import annotations
 
 import logging
-import time
+from typing import TYPE_CHECKING
-from typing import TYPE_CHECKING, List, Optional
-
-import torch
 
 from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
     MultiLayerEagleDraftExtendCudaGraphRunner,
     MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
 )
-from sglang.srt.utils import get_available_gpu_memory
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
-    from sglang.srt.speculative.multi_layer_eagle_worker import (
+    pass
-        MultiLayerEagleDraftWorker,
-    )
 
 
 class MultiLayerEagleDraftExtendNpuGraphRunner(
     MultiLayerEagleDraftExtendCudaGraphRunner
 ):
-    def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int):
-        super().__init__(eagle_worker, step)
-
     def _replay_graph(self, shape_key, forward_batch):
         seq_lens = self.buffers.seq_lens_cpu[: self.raw_bs].tolist() + [0] * (
             self.bs - self.raw_bs
@@ -56,86 +47,8 @@ class MultiLayerEagleDraftExtendNpuGraphRunner(
 class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
     MultiLayerEagleMultiStepDraftExtendCudaGraphRunner
 ):
-    def __init__(self, eagle_worker: MultiLayerEagleDraftWorker):
+    def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendNpuGraphRunner:
-        super().__init__(eagle_worker)
+        return MultiLayerEagleDraftExtendNpuGraphRunner(self.eagle_worker, step)
 
-    def _init_and_capture(self):
+    def _cuda_graph_disabled(self) -> bool:
-        if cuda_graph_fully_disabled():
+        return cuda_graph_fully_disabled()
-            self.runners = [None] * self.speculative_num_steps
-            return
-
-        self.runners: List[Optional[MultiLayerEagleDraftExtendNpuGraphRunner]] = []
-        buffer_len_list: List[int] = []
-
-        for step in range(self.speculative_num_steps):
-            if self.draft_extend_attn_backend_list[step]:
-                runner = MultiLayerEagleDraftExtendNpuGraphRunner(
-                    self.eagle_worker, step
-                )
-                self.runners.append(runner)
-
-                self.seq_len_fill_value = runner.seq_len_fill_value
-                self.max_bs = runner.max_bs
-                buffer_len_list.append(runner.max_num_token)
-                self.offsets.append(self.offsets[-1] + runner.max_num_token)
-            else:
-                self.runners.append(None)
-
-        self.cuda_graph_buffers["seq_lens_cpu"] = torch.full(
-            (self.max_bs,),
-            self.seq_len_fill_value,
-            dtype=torch.int32,
-        )
-
-        with torch.device(self.device):
-            self.cuda_graph_buffers["input_ids"] = torch.zeros(
-                (self.offsets[-1],), dtype=torch.int64
-            )
-            self.cuda_graph_buffers["out_cache_loc"] = torch.ones(
-                (self.offsets[-1],), dtype=torch.int64
-            )
-            self.cuda_graph_buffers["positions"] = torch.zeros(
-                (self.offsets[-1],), dtype=torch.int64
-            )
-
-            self.cuda_graph_buffers["seq_lens"] = torch.full(
-                (self.max_bs,),
-                self.seq_len_fill_value,
-                dtype=torch.int32,
-            )
-            self.cuda_graph_buffers["req_pool_indices"] = torch.zeros(
-                (self.max_bs,), dtype=torch.int64
-            )
-            self.cuda_graph_buffers["num_correct_drafts"] = torch.full(
-                (self.max_bs,), 1, dtype=torch.int32
-            )
-            self.cuda_graph_buffers["num_accept_tokens"] = torch.full(
-                (self.max_bs,), 1, dtype=torch.int32
-            )
-
-        for step in range(self.speculative_num_steps - 1, -1, -1):
-            if self.runners[step] is not None:
-                tic = time.perf_counter()
-                before_mem = get_available_gpu_memory(self.device, self.gpu_id)
-                logger.info(
-                    f"Capture draft extend CUDA graph begin. step={step}, "
-                    f"avail mem={before_mem:.2f} GB"
-                )
-
-                self.runners[step].init_buffers_and_capture(
-                    self.cuda_graph_buffers,
-                    self.offsets[step],
-                    (
-                        self.runners[step + 1]
-                        if step + 1 < self.speculative_num_steps
-                        else None
-                    ),
-                )
-
-                after_mem = get_available_gpu_memory(self.device, self.gpu_id)
-                logger.info(
-                    "Capture draft extend CUDA graph end. "
-                    f"step={step}, elapsed={time.perf_counter() - tic:.2f} s, "
-                    f"mem usage={(before_mem - after_mem):.2f} GB, "
-                    f"avail mem={after_mem:.2f} GB."
-                )

python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py

@@ -14,9 +14,11 @@
 
 from __future__ import annotations
 
+import contextlib
 import logging
 import time
 from dataclasses import dataclass
+from types import SimpleNamespace
 from typing import TYPE_CHECKING, Callable, List, Optional
 
 import torch
@@ -41,7 +43,6 @@ from sglang.srt.model_executor.forward_batch_info import (
 from sglang.srt.model_executor.forward_context import (
     ForwardContext,
     forward_context,
-    get_req_to_token_pool,
 )
 from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
 from sglang.srt.model_executor.runner import (
@@ -56,7 +57,6 @@ from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
 )
 from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
-from sglang.srt.speculative.multi_layer_eagle_utils import assign_new_state_triton
 from sglang.srt.speculative.spec_utils import fast_topk
 from sglang.srt.utils import (
     get_available_gpu_memory,
@@ -77,19 +77,22 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers):
-    # Sliced from shared parent buffers
+    """A single persistent buffer set shared by every MTP draft step."""
+
     input_ids: torch.Tensor
     out_cache_loc: torch.Tensor
     positions: torch.Tensor
-    # Shared from parent
     seq_lens: torch.Tensor
     seq_lens_cpu: torch.Tensor
     req_pool_indices: torch.Tensor
     num_correct_drafts: torch.Tensor
     num_accept_tokens: torch.Tensor
-    # Per-step buffers
     extend_seq_lens: torch.Tensor
     extend_start_loc: torch.Tensor
+    # Flat index (into the token dimension) of each request's last accepted
+    # token. Used both by the in-graph top-k gather and by the worker's
+    # per-step input_ids rotation.
+    select_index: torch.Tensor
     mrope_positions: torch.Tensor
     hidden_states: torch.Tensor
     next_token_logits_buffer: torch.Tensor
@@ -100,12 +103,11 @@ class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers):
 class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
     """Per-step multi-layer EAGLE draft-extend runner.
 
-    Subclasses DecodeCudaGraphRunner. Shares buffers across steps
+    Subclasses DecodeCudaGraphRunner. All steps share a single buffer set
-    via the composite MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
+    owned by the composite MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
-    so initialization is split: __init__ does basic field setup,
+    so initialization is split: __init__ does basic field setup, and
-    init_buffers_and_capture (called by the composite once shared
+    init_buffers_and_capture (called by the composite once the shared buffers
-    buffers exist) finishes by allocating per-step buffers and running
+    exist) attaches them and runs capture.
-    capture.
     """
 
     def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int):
@@ -151,10 +153,12 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         self.capture_bs, _ = get_batch_sizes_to_capture(model_runner)
         self.padded_static_len = -1
 
-        # For Attention Backend
+        # Fixed window: every step extends each request by the same number of
-        self.num_tokens_per_bs = self.speculative_num_steps + 1 + step
+        # tokens, which lets all steps share one buffer set.
+        self.num_tokens_per_bs = self.speculative_num_draft_tokens
         self.max_bs = max(self.capture_bs)
         self.max_num_token = self.max_bs * self.num_tokens_per_bs
+        self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs
 
         self.eagle_worker.draft_extend_attn_backend_list[
             self.step
@@ -163,116 +167,13 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             self.step
         ].get_cuda_graph_seq_len_fill_value()
 
-    def init_buffers_and_capture(
+    def init_buffers_and_capture(self, buffers: MultiLayerEagleDraftExtendInputBuffers):
-        self,
+        """Attach the shared buffer set and capture this step's graphs."""
-        cuda_graph_buffers,
+        self.buffers = buffers
-        offset,
-        next_cuda_graph_runner,
-    ):
-        self.next_cuda_graph_runner = next_cuda_graph_runner
-        seq_lens_cpu = cuda_graph_buffers["seq_lens_cpu"]
-        self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs
 
         if self.enable_torch_compile:
             set_torch_compile_config()
 
-        with torch.device(self.model_runner.device):
-            input_ids = cuda_graph_buffers["input_ids"][
-                offset : offset + self.max_num_token
-            ]
-            out_cache_loc = cuda_graph_buffers["out_cache_loc"][
-                offset : offset + self.max_num_token
-            ]
-            positions = cuda_graph_buffers["positions"][
-                offset : offset + self.max_num_token
-            ]
-
-            seq_lens = cuda_graph_buffers["seq_lens"]
-            req_pool_indices = cuda_graph_buffers["req_pool_indices"]
-            num_correct_drafts = cuda_graph_buffers["num_correct_drafts"]
-            num_accept_tokens = cuda_graph_buffers["num_accept_tokens"]
-
-            extend_seq_lens = torch.full(
-                (self.max_bs,),
-                self.num_tokens_per_bs,
-                dtype=torch.int32,
-            )
-            extend_start_loc = torch.arange(
-                0,
-                self.max_bs * self.num_tokens_per_bs,
-                step=self.num_tokens_per_bs,
-                dtype=torch.int32,
-            )
-
-            mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64)
-
-            hidden_states = torch.zeros(
-                (
-                    self.max_num_token,
-                    EagleDraftExtendInput.hidden_size_for(self.eagle_worker),
-                ),
-                dtype=EagleDraftExtendInput.dtype_for(self.eagle_worker),
-            )
-
-            if self.require_gathered_buffer:
-                if self.require_mlp_tp_gather:
-                    global_num_tokens_gpu = torch.zeros(
-                        (self.dp_size,), dtype=torch.int32
-                    )
-                    global_num_tokens_for_logprob_gpu = torch.zeros(
-                        (self.dp_size,), dtype=torch.int32
-                    )
-                else:
-                    assert self.require_attn_tp_gather
-                    global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
-                    global_num_tokens_for_logprob_gpu = torch.zeros(
-                        (1,), dtype=torch.int32
-                    )
-            else:
-                global_num_tokens_gpu = None
-                global_num_tokens_for_logprob_gpu = None
-
-            if hasattr(
-                self.model_runner.model_config.hf_config, "draft_vocab_size"
-            ):  # llama_eagle
-                vocab_size = self.model_runner.model_config.hf_config.draft_vocab_size
-            elif hasattr(
-                self.model_runner.model_config.hf_config, "hot_vocab_size"
-            ):  # llama_eagle3
-                vocab_size = self.model_runner.model_config.hf_config.hot_vocab_size
-            else:
-                vocab_size = self.model_runner.model_config.vocab_size
-
-            next_token_logits_buffer = torch.zeros(
-                (
-                    (
-                        self.max_bs * self.num_tokens_per_bs
-                        if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2
-                        else self.max_bs
-                    ),
-                    vocab_size,
-                ),
-                dtype=torch.float,
-            )
-
-        self.buffers = MultiLayerEagleDraftExtendInputBuffers(
-            input_ids=input_ids,
-            out_cache_loc=out_cache_loc,
-            positions=positions,
-            seq_lens=seq_lens,
-            seq_lens_cpu=seq_lens_cpu,
-            req_pool_indices=req_pool_indices,
-            num_correct_drafts=num_correct_drafts,
-            num_accept_tokens=num_accept_tokens,
-            extend_seq_lens=extend_seq_lens,
-            extend_start_loc=extend_start_loc,
-            mrope_positions=mrope_positions,
-            hidden_states=hidden_states,
-            next_token_logits_buffer=next_token_logits_buffer,
-            global_num_tokens_gpu=global_num_tokens_gpu,
-            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
-        )
-
         self.backend = resolve_decode_backend(self)
 
         try:
@@ -403,6 +304,19 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         )
         return forward_batch
 
+    def _postprocess_forward_batch(self, forward_batch: ForwardBatch, bs: int):
+        """Hook for subclasses to mutate the captured forward batch."""
+        return forward_batch
+
+    def _compute_topk(self, ret, bs: int):
+        """Compute top-k on the last accepted token's logits and attach it to
+        ``ret``. The gather index lives in a persistent buffer, so the captured
+        graph reads the right rows on each replay. Overridable so distributed
+        (vocab-sharded) builds can plug in an all-reduce-aware sampler."""
+        buffers = self.buffers
+        probs = torch.softmax(ret.next_token_logits[buffers.select_index[:bs]], dim=-1)
+        ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1)
+
     def capture_one_shape(
         self,
         size: int,
@@ -415,6 +329,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
 
         num_tokens = bs * self.num_tokens_per_bs
         forward_batch = self.get_forward_batch(bs)
+        forward_batch = self._postprocess_forward_batch(forward_batch, bs)
         attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step]
 
         def run_once():
@@ -443,47 +358,8 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             ):
                 buffers.hidden_states[:num_tokens].copy_(ret.hidden_states[:num_tokens])
 
-            # num_correct_drafts is drafts-only; the last accepted draft sits at index
+            self._compute_topk(ret, bs)
-            # `num_correct_drafts` within the (current_token + drafts) slot range.
-            select_index = (
-                torch.arange(bs, device=self.model_runner.device)
-                * (self.speculative_num_draft_tokens + self.step)
-                + buffers.num_correct_drafts[:bs]
-                + self.step
-            )
 
-            probs = torch.softmax(ret.next_token_logits[select_index], dim=-1)
-            ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1)
-
-            if self.next_cuda_graph_runner is not None:
-                next_buffers = self.next_cuda_graph_runner.buffers
-                # rejected drafts = proposed drafts - accepted drafts.
-                # speculative_num_draft_tokens includes the current-token slot, so -1.
-                padding_lens = (
-                    self.speculative_num_draft_tokens - 1
-                ) - buffers.num_correct_drafts[:bs]
-                assign_new_state_triton(
-                    ret.topk_index,
-                    buffers.input_ids,
-                    buffers.positions,
-                    buffers.hidden_states,
-                    buffers.out_cache_loc,
-                    buffers.extend_seq_lens,
-                    buffers.extend_start_loc,
-                    next_buffers.input_ids,
-                    next_buffers.positions,
-                    next_buffers.hidden_states,
-                    next_buffers.out_cache_loc,
-                    next_buffers.extend_seq_lens,
-                    next_buffers.extend_start_loc,
-                    next_buffers.seq_lens,
-                    padding_lens,
-                    forward_batch.batch_size,
-                    self.step,
-                    forward_batch.req_pool_indices,
-                    get_req_to_token_pool().req_to_token,
-                    self.eagle_worker.req_to_hidden_states_pool,
-                )
             forward_batch.out_cache_loc = output_cache_loc_backup
             forward_batch.spec_info.hidden_states = hidden_states_backup
             return ret
@@ -501,123 +377,53 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
                 ),
             )
 
-    def init_replay_state(
+    def replay(self, bs: int, seq_lens_sum: int, spec_info: EagleDraftExtendInput):
-        self, forward_batch: ForwardBatch, bs: int, raw_bs: int, num_tokens: int
+        """Init this step's attention metadata for the prepared bucket and
-    ):
+        replay its graph. Buffers must already be populated by the composite
-        buffers = self.buffers
+        runner's ``prepare`` (step 0) or by the previous step's in-graph chain
-        buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids)
+        write + worker-side rotation (steps > 0)."""
-        buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
-        if forward_batch.extend_seq_lens is not None:
-            buffers.extend_seq_lens[:raw_bs].copy_(forward_batch.extend_seq_lens)
-            buffers.extend_start_loc[:raw_bs].copy_(forward_batch.extend_start_loc)
-        buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
-        buffers.positions[:num_tokens].copy_(forward_batch.positions)
-        if (
-            forward_batch.spec_info.hidden_states.shape[1]
-            == buffers.hidden_states.shape[1]
-        ):
-            buffers.hidden_states[:num_tokens].copy_(
-                forward_batch.spec_info.hidden_states
-            )
-        if forward_batch.spec_info.num_correct_drafts is not None:
-            buffers.num_correct_drafts[:raw_bs].copy_(
-                forward_batch.spec_info.num_correct_drafts
-            )
-            buffers.num_accept_tokens[:raw_bs].copy_(
-                forward_batch.spec_info.num_accept_tokens
-            )
-        buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
-
-        if forward_batch.seq_lens_cpu is not None:
-            if bs != raw_bs:
-                buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
-            buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
-
-        if forward_batch.extend_seq_lens_cpu is not None:
-            self.extend_seq_lens_cpu[:raw_bs] = forward_batch.extend_seq_lens_cpu
-
-    def execute(self, forward_batch: ForwardBatch, init_state: bool = True):
-        assert forward_batch.out_cache_loc is not None
         self.deepep_adapter.replay()
         buffers = self.buffers
-
+        num_tokens = bs * self.num_tokens_per_bs
-        raw_bs = forward_batch.batch_size
-        num_tokens = raw_bs * self.num_tokens_per_bs
-        if self.require_mlp_tp_gather:
-            max_batch_size = max(forward_batch.original_global_num_tokens_cpu)
-            bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs)
-        else:
-            bs = self._pad_to_bucket(raw_bs, self.capture_bs)
-
-        if init_state:
-            self.init_replay_state(forward_batch, bs, raw_bs, num_tokens)
 
         if self.require_gathered_buffer:
-            buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
+            buffers.global_num_tokens_gpu.fill_(num_tokens)
-            buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(num_tokens)
-
-        forward_batch.spec_info.hidden_states = buffers.hidden_states[:num_tokens]
-        forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[:bs]
-        forward_batch.spec_info.num_accept_tokens = buffers.num_accept_tokens[:bs]
-        forward_batch.spec_info.num_tokens_per_req = self.num_tokens_per_bs
-        forward_batch.spec_info.num_tokens_for_logprob_per_req = 1
-        forward_batch.spec_info.positions = buffers.positions[:num_tokens]
-        forward_batch.spec_info.extend_seq_lens_tensor = buffers.extend_seq_lens[:bs]
-
-        from types import SimpleNamespace
 
         fb_view = SimpleNamespace(
             batch_size=bs,
             forward_mode=self.forward_mode,
-            input_ids=getattr(forward_batch, "input_ids", None),
+            input_ids=buffers.input_ids[:num_tokens],
             req_pool_indices=buffers.req_pool_indices,
             seq_lens=buffers.seq_lens,
-            seq_lens_sum=forward_batch.seq_lens_sum
+            seq_lens_sum=seq_lens_sum,
-            + (bs - raw_bs) * self.seq_len_fill_value,
             seq_lens_cpu=buffers.seq_lens_cpu,
             encoder_lens=None,
-            # per-step write target (advanced in-graph by assign_new_state);
+            # per-step write target; out_cache_loc is frozen at prepare() time.
-            # forward_batch.out_cache_loc is frozen at step 0.
             out_cache_loc=buffers.out_cache_loc[:num_tokens],
-            spec_info=forward_batch.spec_info,
+            spec_info=spec_info,
         )
         self.eagle_worker.draft_extend_attn_backend_list[
             self.step
         ].init_forward_metadata_out_graph(fb_view)
 
-        self.raw_bs = raw_bs
         self.bs = bs
         shape_key = self._make_graph_key(bs)
-        out = self._replay_graph(shape_key, forward_batch)
+        return self._replay_graph(shape_key, fb_view)
-
-        if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2:
-            unpadding_bs = num_tokens
-        elif bs != raw_bs:
-            forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[
-                :raw_bs
-            ]
-            forward_batch.spec_info.num_accept_tokens = buffers.num_accept_tokens[
-                :raw_bs
-            ]
-            unpadding_bs = raw_bs
-        else:
-            unpadding_bs = None
-
-        if unpadding_bs is not None:
-            out_copy = out
-            out = LogitsProcessorOutput(
-                next_token_logits=out.next_token_logits[:unpadding_bs],
-                hidden_states=out.hidden_states[:unpadding_bs],
-            )
-            out.topk_p = out_copy.topk_p[:raw_bs]
-            out.topk_index = out_copy.topk_index[:raw_bs]
-        return out
 
 
 class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
-    """Composite orchestrator that owns speculative_num_steps per-step
+    """Owns one shared buffer set and the per-step runners.
-    runners with shared input buffers. Not itself a
+
-    DecodeCudaGraphRunner — it only routes work to the per-step
+    Usage from the worker::
+
+        runner.prepare(forward_batch)
+        for step in range(num_steps):
+            _, topk_p, topk_index = runner.replay(step)
+            if step < num_steps - 1:
+                rotate_input_ids_triton(...)  # advance the draft chain
+
+    Not itself a DecodeCudaGraphRunner -- it only routes work to the per-step
     runners.
     """
 
@@ -630,72 +436,57 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
             eagle_worker.draft_extend_attn_backend_list
         )
 
-        self.runners = []
+        self.runners: List[Optional[MultiLayerEagleDraftExtendCudaGraphRunner]] = []
-        self.cuda_graph_buffers = {}
         self.seq_len_fill_value = 1
         self.max_bs = 1
-        self.offsets = [0]
+        self.num_tokens_per_bs = 1
 
         self._init_and_capture()
 
+    def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendCudaGraphRunner:
+        return MultiLayerEagleDraftExtendCudaGraphRunner(self.eagle_worker, step)
+
+    def _capture_context(self, step: int):
+        """Context manager active while capturing ``step``'s graphs. Subclasses
+        can use it e.g. to temporarily expose a sharded local vocab size."""
+        return contextlib.nullcontext()
+
+    def _on_runners_created(self):
+        """Hook called after all per-step runners exist but before buffers are
+        allocated/captured (e.g. to allocate shared sconv buffers)."""
+
+    def _cuda_graph_disabled(self) -> bool:
+        return check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED)
+
     def _init_and_capture(self):
-        if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED):
+        if self._cuda_graph_disabled():
             self.runners = [None] * self.speculative_num_steps
             return
 
-        self.runners: List[Optional[MultiLayerEagleDraftExtendCudaGraphRunner]] = []
+        self.runners = []
-        buffer_len_list: List[int] = []
 
-        # 1. Construct per-step runners (cheap setup only).
+        # 1. Construct per-step runners (each initializes its own attn cuda
+        #    graph state). They share the same fixed window size.
         for step in range(self.speculative_num_steps):
             if self.draft_extend_attn_backend_list[step]:
-                runner = MultiLayerEagleDraftExtendCudaGraphRunner(
+                runner = self._create_runner(step)
-                    self.eagle_worker, step
-                )
                 self.runners.append(runner)
-
                 self.seq_len_fill_value = runner.seq_len_fill_value
                 self.max_bs = runner.max_bs
-                buffer_len_list.append(runner.max_num_token)
+                self.num_tokens_per_bs = runner.num_tokens_per_bs
-                self.offsets.append(self.offsets[-1] + runner.max_num_token)
+                self.capture_bs = runner.capture_bs
+                self.require_gathered_buffer = runner.require_gathered_buffer
+                self.require_mlp_tp_gather = runner.require_mlp_tp_gather
+                self.require_mlp_sync = runner.require_mlp_sync
+                self.disable_padding = runner.disable_padding
             else:
                 self.runners.append(None)
 
-        # 2. Allocate shared buffers.
+        self._on_runners_created()
-        self.cuda_graph_buffers["seq_lens_cpu"] = torch.full(
-            (self.max_bs,),
-            self.seq_len_fill_value,
-            dtype=torch.int32,
-        )
 
-        with torch.device(self.device):
+        # 2. Allocate the single shared buffer set and capture each step in
-            self.cuda_graph_buffers["input_ids"] = torch.zeros(
+        #    reverse order.
-                (self.offsets[-1],), dtype=torch.int64
+        self.buffers = self._allocate_buffers()
-            )
-            self.cuda_graph_buffers["out_cache_loc"] = torch.ones(
-                (self.offsets[-1],), dtype=torch.int64
-            )
-            self.cuda_graph_buffers["positions"] = torch.zeros(
-                (self.offsets[-1],), dtype=torch.int64
-            )
-
-            self.cuda_graph_buffers["seq_lens"] = torch.full(
-                (self.max_bs,),
-                self.seq_len_fill_value,
-                dtype=torch.int32,
-            )
-            self.cuda_graph_buffers["req_pool_indices"] = torch.zeros(
-                (self.max_bs,), dtype=torch.int64
-            )
-            self.cuda_graph_buffers["num_correct_drafts"] = torch.full(
-                (self.max_bs,), 1, dtype=torch.int32
-            )
-            self.cuda_graph_buffers["num_accept_tokens"] = torch.full(
-                (self.max_bs,), 1, dtype=torch.int32
-            )
-
-        # 3. Per-step capture, in reverse order so that next_cuda_graph_runner
-        # is already initialized when this step references it.
         for step in range(self.speculative_num_steps - 1, -1, -1):
             if self.runners[step] is not None:
                 tic = time.perf_counter()
@@ -705,15 +496,8 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
                     f"avail mem={before_mem:.2f} GB"
                 )
 
-                self.runners[step].init_buffers_and_capture(
+                with self._capture_context(step):
-                    self.cuda_graph_buffers,
+                    self.runners[step].init_buffers_and_capture(self.buffers)
-                    self.offsets[step],
-                    (
-                        self.runners[step + 1]
-                        if step + 1 < self.speculative_num_steps
-                        else None
-                    ),
-                )
 
                 after_mem = get_available_gpu_memory(self.device, self.gpu_id)
                 logger.info(
@@ -723,18 +507,190 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner:
                     f"avail mem={after_mem:.2f} GB."
                 )
 
-    def reset_buffers(self, forward_batch, batch_result):
+    def _vocab_size(self) -> int:
-        self.cuda_graph_buffers["input_ids"].zero_()
+        model_runner = self.eagle_worker.mtp_model_runner(0)
-        self.cuda_graph_buffers["seq_lens"].fill_(self.seq_len_fill_value)
+        if hasattr(model_runner.model_config.hf_config, "draft_vocab_size"):
-        self.cuda_graph_buffers["out_cache_loc"].zero_()
+            return model_runner.model_config.hf_config.draft_vocab_size
-        self.cuda_graph_buffers["positions"].zero_()
+        if hasattr(model_runner.model_config.hf_config, "hot_vocab_size"):
-        # `batch_result.accept_lens` is drafts + bonus.
+            return model_runner.model_config.hf_config.hot_vocab_size
-        bs = forward_batch.batch_size
+        return model_runner.model_config.vocab_size
-        self.cuda_graph_buffers["num_correct_drafts"][:bs].copy_(
+
-            batch_result.accept_lens - 1
+    def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers:
+        runner = next(r for r in self.runners if r is not None)
+        max_bs = self.max_bs
+        num_tokens_per_bs = self.num_tokens_per_bs
+        max_num_token = max_bs * num_tokens_per_bs
+        hidden_size = EagleDraftExtendInput.hidden_size_for(self.eagle_worker)
+        dtype = EagleDraftExtendInput.dtype_for(self.eagle_worker)
+        vocab_size = self._vocab_size()
+
+        seq_lens_cpu = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32)
+
+        with torch.device(self.device):
+            input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
+            out_cache_loc = torch.ones((max_num_token,), dtype=torch.int64)
+            positions = torch.zeros((max_num_token,), dtype=torch.int64)
+            mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
+            hidden_states = torch.zeros((max_num_token, hidden_size), dtype=dtype)
+
+            seq_lens = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32)
+            req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
+            num_correct_drafts = torch.full((max_bs,), 1, dtype=torch.int32)
+            num_accept_tokens = torch.full((max_bs,), 1, dtype=torch.int32)
+
+            # Fixed window: every request extends by exactly num_tokens_per_bs
+            # tokens, and start locs are a constant arange.
+            extend_seq_lens = torch.full(
+                (max_bs,), num_tokens_per_bs, dtype=torch.int32
             )
-        self.cuda_graph_buffers["num_accept_tokens"][:bs].copy_(
+            extend_start_loc = torch.arange(
-            batch_result.accept_lens
+                0, max_num_token, step=num_tokens_per_bs, dtype=torch.int32
+            )
+            select_index = torch.zeros((max_bs,), dtype=torch.int64)
+
+            next_token_logits_buffer = torch.zeros(
+                (max_num_token, vocab_size), dtype=torch.float
+            )
+
+            if self.require_gathered_buffer:
+                if self.require_mlp_tp_gather:
+                    dp_size = runner.dp_size
+                    global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (dp_size,), dtype=torch.int32
+                    )
+                else:
+                    global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+                    global_num_tokens_for_logprob_gpu = torch.zeros(
+                        (1,), dtype=torch.int32
+                    )
+            else:
+                global_num_tokens_gpu = None
+                global_num_tokens_for_logprob_gpu = None
+
+        return MultiLayerEagleDraftExtendInputBuffers(
+            input_ids=input_ids,
+            out_cache_loc=out_cache_loc,
+            positions=positions,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            req_pool_indices=req_pool_indices,
+            num_correct_drafts=num_correct_drafts,
+            num_accept_tokens=num_accept_tokens,
+            extend_seq_lens=extend_seq_lens,
+            extend_start_loc=extend_start_loc,
+            select_index=select_index,
+            mrope_positions=mrope_positions,
+            hidden_states=hidden_states,
+            next_token_logits_buffer=next_token_logits_buffer,
+            global_num_tokens_gpu=global_num_tokens_gpu,
+            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+        )
+
+    def _prepare_extra(self, forward_batch: ForwardBatch) -> None:
+        """Hook for subclasses to populate extra per-call buffers (e.g. sconv)."""
+
+    def prepare(self, forward_batch: ForwardBatch):
+        """Populate the shared buffers once from ``forward_batch`` and bucketize
+        the batch size. Subsequent ``replay(step)`` calls reuse this state."""
+        buffers = self.buffers
+        raw_bs = forward_batch.batch_size
+        num_tokens = raw_bs * self.num_tokens_per_bs
+
+        # Bucketize to a captured batch size (padding the tail).
+        if self.require_mlp_tp_gather:
+            max_batch_size = max(forward_batch.original_global_num_tokens_cpu)
+            bs = self.get_runner(0)._pad_to_bucket(int(max_batch_size), self.capture_bs)
+        else:
+            bs = self.get_runner(0)._pad_to_bucket(raw_bs, self.capture_bs)
+
+        # Reset padded slots, then copy the real values in.
+        buffers.input_ids.zero_()
+        buffers.out_cache_loc.zero_()
+        buffers.positions.zero_()
+        buffers.seq_lens.fill_(self.seq_len_fill_value)
+
+        buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids)
+        buffers.positions[:num_tokens].copy_(forward_batch.positions)
+        buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
+        buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
+        buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
+
+        if (
+            forward_batch.spec_info.hidden_states.shape[1]
+            == buffers.hidden_states.shape[1]
+        ):
+            buffers.hidden_states[:num_tokens].copy_(
+                forward_batch.spec_info.hidden_states
+            )
+
+        buffers.num_correct_drafts[:raw_bs].copy_(
+            forward_batch.spec_info.num_correct_drafts
+        )
+        buffers.num_accept_tokens[:raw_bs].copy_(
+            forward_batch.spec_info.num_accept_tokens
+        )
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                buffers.seq_lens_cpu.fill_(self.seq_len_fill_value)
+            buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+
+        # select_index[i] = i * window + num_correct_drafts[i]: the flat index
+        # of request i's last accepted token. Used by the in-graph top-k gather
+        # and by the worker's rotation.
+        arange = torch.arange(bs, device=self.device, dtype=torch.int64)
+        buffers.select_index[:bs].copy_(
+            arange * self.num_tokens_per_bs + buffers.num_correct_drafts[:bs]
+        )
+
+        if self.require_gathered_buffer:
+            buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs)
+            buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs)
+
+        # Reusable spec_info for per-step attention metadata.
+        padded_num_tokens = bs * self.num_tokens_per_bs
+        spec_info = EagleDraftExtendInput(
+            hidden_states=buffers.hidden_states[:padded_num_tokens],
+            num_correct_drafts=buffers.num_correct_drafts[:bs],
+            num_accept_tokens=buffers.num_accept_tokens[:bs],
+        )
+        spec_info.num_tokens_per_req = self.num_tokens_per_bs
+        spec_info.num_tokens_for_logprob_per_req = 1
+        spec_info.positions = buffers.positions[:padded_num_tokens]
+        spec_info.extend_seq_lens_tensor = buffers.extend_seq_lens[:bs]
+        self._replay_spec_info = spec_info
+
+        self.raw_bs = raw_bs
+        self.bs = bs
+        self.raw_num_tokens = num_tokens
+        self.seq_lens_sum = (
+            forward_batch.seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value
+        )
+
+        self._prepare_extra(forward_batch)
+
+    def replay(self, step: int):
+        """Replay ``step``'s graph at the prepared bucket. Returns
+        ``(LogitsProcessorOutput, topk_p, topk_index)`` sliced to the real
+        batch size."""
+        runner = self.runners[step]
+        runner.raw_bs = self.raw_bs
+        out = runner.replay(self.bs, self.seq_lens_sum, self._replay_spec_info)
+        raw_bs = self.raw_bs
+        raw_num_tokens = self.raw_num_tokens
+        logits_output = LogitsProcessorOutput(
+            next_token_logits=out.next_token_logits[:raw_num_tokens],
+            hidden_states=(
+                out.hidden_states[:raw_num_tokens]
+                if out.hidden_states is not None
+                else None
+            ),
+        )
+        return (
+            logits_output,
+            out.topk_p[:raw_bs],
+            out.topk_index[:raw_bs],
         )
 
     def get_runner(self, step):

python/sglang/srt/speculative/multi_layer_eagle_utils.py

@@ -13,21 +13,11 @@
 # ==============================================================================
 
 from sglang.srt.speculative.triton_ops.multi_layer_eagle import (
-    assign_hidden_states_pool_kernel,
-    assign_hidden_states_pool_torch,
-    assign_hidden_states_pool_triton,
-    assign_new_state_kernel,
-    assign_new_state_triton,
     rotate_input_ids_kernel,
     rotate_input_ids_triton,
 )
 
 __all__ = [
-    "assign_hidden_states_pool_kernel",
-    "assign_hidden_states_pool_torch",
-    "assign_hidden_states_pool_triton",
-    "assign_new_state_kernel",
-    "assign_new_state_triton",
     "rotate_input_ids_kernel",
     "rotate_input_ids_triton",
 ]

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -58,10 +58,7 @@ from sglang.srt.speculative.eagle_utils import (
 from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
     MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
 )
-from sglang.srt.speculative.multi_layer_eagle_utils import (
+from sglang.srt.speculative.multi_layer_eagle_utils import rotate_input_ids_triton
-    assign_hidden_states_pool_triton,
-    rotate_input_ids_triton,
-)
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import (
     draft_tp_context,
@@ -188,18 +185,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         )
         self.init_lm_head()
 
-        # KV cache reversion buffer; sized to mirror req_to_token (indexed by
-        # req_pool_idx).
-        self.req_to_hidden_states_pool = torch.empty(
-            (
-                self.req_to_token_pool.req_to_token.shape[0],
-                self.speculative_num_steps - 1,
-                self.model_config.hidden_size,
-            ),
-            dtype=self.model_config.dtype,
-            device=self.device,
-        )
-
     def init_attention_backends(self):
         with self.draft_tp_context(
             self.draft_runner_list[0].tp_group
@@ -255,12 +240,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                 MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(self)
             )
 
-    def reset_cuda_graph_buffers(self, forward_batch, batch_result):
-        if self.cuda_graph_runner_for_draft_extend:
-            self.cuda_graph_runner_for_draft_extend.reset_buffers(
-                forward_batch, batch_result
-            )
-
     def draft(self, batch: ScheduleBatch):
         draft_input: EagleDraftInput = batch.spec_info
         forward_batch, can_cuda_graph = self.prepare_for_draft(
@@ -492,6 +471,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                     forward_batch.extend_seq_lens,
                     topk_index,
                 )
+
         next_draft_input = EagleDraftInput(
             topk_p=torch.cat(topk_p_list, dim=1),
             topk_index=torch.cat(topk_index_list, dim=1),
@@ -503,17 +483,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
             num_tokens_for_logprob_per_req=1,
         )
 
-        # Update req_to_hidden_states_pool for KV Cache reversion
-        if forward_batch.extend_seq_lens is not None:
-            assign_hidden_states_pool_triton(
-                target_hidden_states,
-                forward_batch.req_pool_indices,
-                self.req_to_hidden_states_pool,
-                self.speculative_num_steps - 1,
-                forward_batch.batch_size,
-                forward_batch.extend_seq_lens,
-                forward_batch.extend_start_loc,
-            )
         return next_draft_input
 
     def _draft_extend_for_decode(
@@ -543,6 +512,11 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
             torch.get_device_module(self.device).current_stream().wait_stream(
                 self.plan_stream
             )
+        # `batch_result.accept_lens` includes the bonus token, so drafts-only
+        # is accept_lens - 1. Stash on spec_info for the cuda-graph prepare().
+        forward_batch.spec_info.num_correct_drafts = batch_result.accept_lens - 1
+        forward_batch.spec_info.num_accept_tokens = batch_result.accept_lens
+
         # Run draft extend batch in the main compute stream
         can_cuda_graph = (
             self.cuda_graph_runner_for_draft_extend
@@ -553,7 +527,24 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
         next_token_ids_backup = batch_result.next_token_ids.clone()
 
         if can_cuda_graph:
-            self.reset_cuda_graph_buffers(forward_batch, batch_result)
+            cgr = self.cuda_graph_runner_for_draft_extend
+            # Populate the single shared buffer set once; each step replays
+            # against it and the chain is advanced in place between steps.
+            cgr.prepare(forward_batch)
+            for step in range(self.speculative_num_steps):
+                _, ret_topk_p, ret_topk_index = cgr.replay(step)
+                ret_topk_p_list.append(ret_topk_p.clone())
+                ret_topk_index_list.append(ret_topk_index.clone())
+                # Advance the draft chain by rotating the shared input_ids window
+                # in place; step N+1's graph then reads the rotated values.
+                if step < self.speculative_num_steps - 1:
+                    rotate_input_ids_triton(
+                        cgr.buffers.input_ids[: cgr.raw_num_tokens],
+                        cgr.buffers.extend_start_loc[: cgr.raw_bs],
+                        cgr.buffers.extend_seq_lens[: cgr.raw_bs],
+                        ret_topk_index,
+                        cgr.buffers.select_index[: cgr.raw_bs],
+                    )
         else:
             logger.warning_once(
                 "can't use cuda graph for draft extend! may have correctness issue!"
@@ -572,24 +563,10 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
             # its own metadata in forward_extend (post-pad), otherwise
             # per-runner attn_backend.forward_metadata is never initialized for
             # draft_runner_list[1+].
-            if not _is_npu or can_cuda_graph:
+            if not _is_npu:
                 forward_batch.mark_forward_metadata_ready()
 
             for step in range(self.speculative_num_steps):
-            # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}")
-            if can_cuda_graph:
-                draft_logits_output = (
-                    self.cuda_graph_runner_for_draft_extend.get_runner(step).execute(
-                        forward_batch, init_state=(step == 0)
-                    )
-                )
-                ret_topk_p, ret_topk_index = (
-                    draft_logits_output.topk_p,
-                    draft_logits_output.topk_index,
-                )
-            else:
-                # Skip relies on the unconditional mark above (pre-existing
-                # no-pre-plan behavior preserved verbatim).
                 draft_logits_output = self.draft_runner_list[step].forward(
                     forward_batch
                 )
@@ -618,39 +595,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
                 ret_topk_p_list.append(ret_topk_p)
                 ret_topk_index_list.append(ret_topk_index)
 
-        # Update req_to_hidden_states_pool for KV Cache reversion
-        if (
-            forward_batch.extend_seq_lens is not None
-            and self.cuda_graph_runner_for_draft_extend is not None
-        ):
-            if can_cuda_graph:
-                last_runner = self.cuda_graph_runner_for_draft_extend.get_last_runner()
-                hidden_states = last_runner.buffers.hidden_states
-                req_pool_indices = last_runner.buffers.req_pool_indices
-                extend_seq_lens = last_runner.buffers.extend_seq_lens
-                extend_start_loc = last_runner.buffers.extend_start_loc
-            else:
-                hidden_states = draft_logits_output.logits_output.hidden_states
-                req_pool_indices = forward_batch.req_pool_indices
-                extend_seq_lens = forward_batch.extend_seq_lens
-                extend_start_loc = forward_batch.extend_start_loc
-            assign_hidden_states_pool_triton(
-                hidden_states,
-                req_pool_indices,
-                self.req_to_hidden_states_pool,
-                self.speculative_num_steps - 1,
-                forward_batch.batch_size,
-                extend_seq_lens,
-                extend_start_loc,
-            )
-
-        # Reorganize the spec info for the next batch
-        # draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[
-        #     select_index
-        # ]
-        # draft_logits_output.hidden_states = draft_logits_output.hidden_states[
-        #     select_index
-        # ]
         batch_result.next_token_ids = next_token_ids_backup
         # Construct the return values
         next_draft_input = batch_result.next_draft_input

python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py

@@ -12,7 +12,6 @@
 # limitations under the License.
 # ==============================================================================
 
-import torch
 import triton
 import triton.language as tl
 
@@ -70,281 +69,3 @@ def rotate_input_ids_triton(
         BLOCK_SIZE=BLOCK_SIZE,
     )
     return input_ids
-
-
-@triton.jit
-def assign_new_state_kernel(
-    # Source pointers
-    old_input_ids_ptr,
-    old_positions_ptr,
-    old_hidden_states_ptr,
-    old_out_cache_loc_ptr,
-    old_extend_seq_lens_ptr,
-    old_extend_start_loc_ptr,
-    # Destination pointers
-    input_ids_ptr,
-    positions_ptr,
-    hidden_states_ptr,
-    out_cache_loc_ptr,
-    extend_seq_lens_ptr,
-    extend_start_loc_ptr,
-    # Auxiliary data pointers
-    next_token_ids_ptr,
-    seq_lens_ptr,
-    padding_lens_ptr,
-    req_pool_indices_ptr,
-    req_to_token_ptr,
-    req_to_hidden_states_pool_ptr,
-    # Scalars and Strides
-    step,
-    stride_hidden_seq,
-    stride_hidden_dim,  # hidden_states strides
-    stride_pool_req,
-    stride_pool_step,
-    stride_pool_dim,  # pool strides
-    stride_req_token_0,
-    stride_req_token_1,  # req_to_token strides
-    # Meta-parameters
-    HIDDEN_DIM: tl.constexpr,
-    BLOCK_SEQ: tl.constexpr,
-    BLOCK_HID: tl.constexpr,
-):
-    pid = tl.program_id(0)
-
-    seq_len: tl.tensor = tl.load(seq_lens_ptr + pid)
-    old_extend_len = tl.load(old_extend_seq_lens_ptr + pid)
-    old_start = tl.load(old_extend_start_loc_ptr + pid)
-    new_extend_len = old_extend_len + 1
-    new_start = old_start + pid
-
-    tl.store(extend_seq_lens_ptr + pid, new_extend_len)
-    tl.store(extend_start_loc_ptr + pid, new_start)
-
-    offs_seq = tl.arange(0, BLOCK_SEQ)
-    mask_seq = offs_seq < old_extend_len
-
-    old_ids = tl.load(old_input_ids_ptr + old_start + offs_seq, mask=mask_seq)
-    tl.store(input_ids_ptr + new_start + offs_seq, old_ids, mask=mask_seq)
-    padding_len = tl.load(padding_lens_ptr + pid)
-    tl.store(
-        input_ids_ptr + new_start + old_extend_len - padding_len,
-        tl.load(next_token_ids_ptr + pid),
-    )
-
-    old_pos = tl.load(old_positions_ptr + old_start + offs_seq, mask=mask_seq)
-    tl.store(positions_ptr + new_start + 1 + offs_seq, old_pos, mask=mask_seq)
-    tl.store(
-        positions_ptr + new_start, max(tl.load(old_positions_ptr + old_start) - 1, 0)
-    )
-
-    old_cache = tl.load(old_out_cache_loc_ptr + old_start + offs_seq, mask=mask_seq)
-    tl.store(out_cache_loc_ptr + new_start + 1 + offs_seq, old_cache, mask=mask_seq)
-
-    req_idx = tl.load(req_pool_indices_ptr + pid)
-    token_idx_col = seq_len - old_extend_len - 1
-    if token_idx_col >= 0:
-        req_token_ptr_loc = (
-            req_to_token_ptr
-            + (req_idx * stride_req_token_0)
-            + (token_idx_col * stride_req_token_1)
-        )
-        last_cache_loc = tl.load(req_token_ptr_loc)
-        tl.store(out_cache_loc_ptr + new_start, last_cache_loc)
-
-    pool_vec_offset_base = ((req_idx + 1) * stride_pool_req) + (
-        -(step + 1) * stride_pool_step
-    )
-
-    for off_h in range(0, HIDDEN_DIM, BLOCK_HID):
-        offs_h = off_h + tl.arange(0, BLOCK_HID)
-        mask_h = offs_h < HIDDEN_DIM
-
-        for i in range(BLOCK_SEQ):
-            if i < old_extend_len:
-                old_h_ptr = (
-                    old_hidden_states_ptr
-                    + (old_start + i) * stride_hidden_seq
-                    + (offs_h * stride_hidden_dim)
-                )
-                new_h_ptr = (
-                    hidden_states_ptr
-                    + (new_start + 1 + i) * stride_hidden_seq
-                    + (offs_h * stride_hidden_dim)
-                )
-
-                chunk_old = tl.load(old_h_ptr, mask=mask_h)
-                tl.store(new_h_ptr, chunk_old, mask=mask_h)
-
-        pool_ptrs = (
-            req_to_hidden_states_pool_ptr
-            + pool_vec_offset_base
-            + (offs_h * stride_pool_dim)
-        )
-        pool_val = tl.load(pool_ptrs, mask=mask_h)
-
-        new_h_start_ptrs = (
-            hidden_states_ptr
-            + (new_start * stride_hidden_seq)
-            + (offs_h * stride_hidden_dim)
-        )
-        tl.store(new_h_start_ptrs, pool_val, mask=mask_h)
-
-
-def assign_new_state_triton(
-    next_token_ids: torch.Tensor,
-    old_input_ids: torch.Tensor,
-    old_positions: torch.Tensor,
-    old_hidden_states: torch.Tensor,
-    old_out_cache_loc: torch.Tensor,
-    old_extend_seq_lens: torch.Tensor,
-    old_extend_start_loc: torch.Tensor,
-    input_ids: torch.Tensor,
-    positions: torch.Tensor,
-    hidden_states: torch.Tensor,
-    out_cache_loc: torch.Tensor,
-    extend_seq_lens: torch.Tensor,
-    extend_start_loc: torch.Tensor,
-    seq_lens: torch.Tensor,
-    padding_lens: torch.Tensor,
-    num_seqs: int,
-    step: int,
-    req_pool_indices: torch.Tensor,
-    req_to_token: torch.Tensor,
-    req_to_hidden_states_pool: torch.Tensor,
-):
-    """
-    Wrapper function to calculate offsets and launch the Triton kernel.
-    """
-    hidden_dim = hidden_states.shape[1]
-
-    BLOCK_SEQ = 8
-    BLOCK_HID = 64
-
-    grid = (num_seqs,)
-
-    assign_new_state_kernel[grid](
-        # Pointers
-        old_input_ids,
-        old_positions,
-        old_hidden_states,
-        old_out_cache_loc,
-        old_extend_seq_lens,
-        old_extend_start_loc,
-        input_ids,
-        positions,
-        hidden_states,
-        out_cache_loc,
-        extend_seq_lens,
-        extend_start_loc,
-        next_token_ids,
-        seq_lens,
-        padding_lens,
-        req_pool_indices,
-        req_to_token,
-        req_to_hidden_states_pool,
-        # Constants/Strides
-        step,
-        old_hidden_states.stride(0),
-        old_hidden_states.stride(1),
-        req_to_hidden_states_pool.stride(0),
-        req_to_hidden_states_pool.stride(1),
-        req_to_hidden_states_pool.stride(2),
-        req_to_token.stride(0),
-        req_to_token.stride(1),
-        # Meta
-        HIDDEN_DIM=hidden_dim,
-        BLOCK_SEQ=BLOCK_SEQ,
-        BLOCK_HID=BLOCK_HID,
-    )
-
-
-@triton.jit
-def assign_hidden_states_pool_kernel(
-    hidden_states_ptr,
-    req_pool_indices_ptr,
-    req_to_hidden_states_pool_ptr,
-    extend_seq_lens_ptr,
-    extend_start_loc_ptr,
-    stride_hidden_seq,
-    stride_hidden_dim,
-    stride_pool_req,
-    stride_pool_step,
-    stride_pool_dim,
-    HIDDEN_DIM: tl.constexpr,
-    pool_size: tl.constexpr,
-    BLOCK_HID: tl.constexpr,
-):
-    pid = tl.program_id(0)
-
-    extend_len = tl.load(extend_seq_lens_ptr + pid)
-    start_loc = tl.load(extend_start_loc_ptr + pid)
-    end_loc = start_loc + extend_len
-
-    req_idx = tl.load(req_pool_indices_ptr + pid)
-    pool_vec_offset_base = req_idx * stride_pool_req
-
-    for i in range(pool_size):
-        for off_h in range(0, HIDDEN_DIM, BLOCK_HID):
-            offs_h = off_h + tl.arange(0, BLOCK_HID)
-            mask_h = offs_h < HIDDEN_DIM
-
-            hid_ptr = (
-                hidden_states_ptr
-                + (end_loc - pool_size + i) * stride_hidden_seq
-                + offs_h * stride_hidden_dim
-            )
-            hid_val = tl.load(hid_ptr, mask=mask_h)
-
-            pool_ptr = (
-                req_to_hidden_states_pool_ptr
-                + pool_vec_offset_base
-                + i * stride_pool_step
-                + offs_h * stride_pool_dim
-            )
-            tl.store(pool_ptr, hid_val, mask=mask_h)
-
-
-def assign_hidden_states_pool_triton(
-    hidden_states: torch.Tensor,
-    req_pool_indices: torch.Tensor,
-    req_to_hidden_states_pool: torch.Tensor,
-    pool_size: int,
-    num_seqs: int,
-    extend_seq_lens: torch.Tensor,
-    extend_start_loc: torch.Tensor,
-):
-    grid = (num_seqs,)
-    assign_hidden_states_pool_kernel[grid](
-        hidden_states,
-        req_pool_indices,
-        req_to_hidden_states_pool,
-        extend_seq_lens,
-        extend_start_loc,
-        hidden_states.stride(0),
-        hidden_states.stride(1),
-        req_to_hidden_states_pool.stride(0),
-        req_to_hidden_states_pool.stride(1),
-        req_to_hidden_states_pool.stride(2),
-        HIDDEN_DIM=hidden_states.shape[1],
-        pool_size=pool_size,
-        BLOCK_HID=64,
-    )
-
-
-def assign_hidden_states_pool_torch(
-    hidden_states: torch.Tensor,
-    req_pool_indices: torch.Tensor,
-    req_to_hidden_states_pool: torch.Tensor,
-    pool_size: int,
-    num_seqs: int,
-    extend_seq_lens: torch.Tensor,
-    extend_start_loc: torch.Tensor,
-):
-    for req in range(num_seqs):
-        pool_idx = req_pool_indices[req]
-        extend_len = extend_seq_lens[req]
-        start_loc = extend_start_loc[req]
-        end_loc = start_loc + extend_len
-        req_to_hidden_states_pool[pool_idx, :pool_size, :].copy_(
-            hidden_states[end_loc - pool_size : end_loc, :]
-        )