[Spec] Anchor GLM-5.2 MTP IndexShare topk on the draft-extend step (#29787)

Co-authored-by: kpham-sgl <264503018+kpham-sgl@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -737,7 +737,7 @@ class TboForwardBatchPreparer:
             "split_index",  # for split prefill
             "orig_seq_lens",  # only used by qwen-1m, thus not care
             "return_pooled_hidden_states",
-            "reuse_mtp_topk_indices",  # forward-level flag, inherited by both child batches
+            "reuse_dsa_topk_indices",  # forward-level flag, inherited by both child batches
         ]:
             output_dict[key] = getattr(batch, key)
 

python/sglang/srt/managers/overlap_utils.py

@@ -111,6 +111,7 @@ class RelayPayload:
     topk_index: Optional[torch.Tensor] = None
     hidden_states: Optional[torch.Tensor] = None
     draft_probs: Optional[torch.Tensor] = None
+    dsa_topk_indices: Optional[torch.Tensor] = None
 
     @classmethod
     def from_draft_input(cls, draft_input: EagleDraftInput) -> RelayPayload:
@@ -120,6 +121,7 @@ class RelayPayload:
             topk_index=draft_input.topk_index,
             hidden_states=draft_input.hidden_states,
             draft_probs=getattr(draft_input, "draft_probs", None),
+            dsa_topk_indices=getattr(draft_input, "dsa_topk_indices", None),
         )
 
 
@@ -216,6 +218,15 @@ class FutureMap:
                 device=self.device,
             )
 
+        self.dsa_topk_indices_buf = None
+        if payload.dsa_topk_indices is not None:
+            seed0 = payload.dsa_topk_indices[0]
+            self.dsa_topk_indices_buf = torch.empty(
+                (self.req_pool_size, *seed0.shape),
+                dtype=payload.dsa_topk_indices.dtype,
+                device=self.device,
+            )
+
     def _resolve_spec_extras(self, batch: ScheduleBatch) -> None:
         if self.spec_algo.is_ngram():
             # FIXME: remove once precomputed draft is supported.
@@ -255,6 +266,8 @@ class FutureMap:
             draft_input.bonus_tokens = self.output_tokens_buf[indices]
         if self.need_hidden_states and not self.need_topk:
             draft_input.hidden_states = self.hidden_states_buf[indices]
+        if self.dsa_topk_indices_buf is not None:
+            draft_input.dsa_topk_indices = self.dsa_topk_indices_buf[indices]
         if _DEBUG_ASSERT:
             _assert_nonneg_and_invalidate(
                 draft_input.bonus_tokens, self.output_tokens_buf, indices
@@ -338,3 +351,10 @@ class FutureMap:
             )
         if self.draft_probs_buf is not None and payload.draft_probs is not None:
             self.draft_probs_buf[indices] = payload.draft_probs
+        if (
+            self.dsa_topk_indices_buf is not None
+            and payload.dsa_topk_indices is not None
+        ):
+            self.dsa_topk_indices_buf[indices] = payload.dsa_topk_indices.to(
+                self.dsa_topk_indices_buf.dtype
+            )

python/sglang/srt/model_executor/forward_batch_info.py

@@ -431,8 +431,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     return_hidden_states_before_norm: bool = False
 
     # Gate for reusing the first MTP draft step's indexer topk across steps;
-    # the carried topk lives on spec_info (see EagleDraftInput.mtp_topk_indices).
-    reuse_mtp_topk_indices: Optional[bool] = False
+    # the carried topk lives on spec_info (see EagleDraftInput.dsa_topk_indices).
+    reuse_dsa_topk_indices: Optional[bool] = False
 
     # === Forward-derived (built in init_new on the forward stream; FB-owned) ===
     # Position information

python/sglang/srt/models/deepseek_nextn.py

@@ -237,13 +237,22 @@ class DeepseekModelNextN(nn.Module):
                     residual,
                     zero_allocator,
                     prev_topk_indices=(
-                        forward_batch.spec_info.mtp_topk_indices
-                        if forward_batch.reuse_mtp_topk_indices
+                        forward_batch.spec_info.dsa_topk_indices
+                        if forward_batch.reuse_dsa_topk_indices
                         else None
                     ),
                 )
-                if forward_batch.reuse_mtp_topk_indices:
-                    forward_batch.spec_info.mtp_topk_indices = topk_indices
+                if forward_batch.reuse_dsa_topk_indices:
+                    forward_batch.spec_info.dsa_topk_indices = topk_indices
+
+                # MTP IndexShare: on draft-extend, publish the last-token DSA
+                # indexer top-k to seed (avoid recomputing in) the draft-decode loop.
+                if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
+                    seed_buf = forward_batch.spec_info.dsa_seed_topk_capture
+                    if seed_buf is not None and topk_indices is not None:
+                        sel = forward_batch.spec_info.dsa_seed_topk_select
+                        src = topk_indices if sel is None else topk_indices[sel]
+                        seed_buf[: src.shape[0]].copy_(src)
 
             if not forward_batch.forward_mode.is_idle():
                 if residual is not None:

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -69,6 +69,7 @@ class EagleDraftInputBuffers(ForwardInputBuffers):
     hidden_states: Optional[torch.Tensor]
     global_num_tokens_gpu: Optional[torch.Tensor]
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
+    dsa_seed_topk: Optional[torch.Tensor] = None
 
 
 class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
@@ -228,6 +229,16 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
         )
 
+        dsa_seed_topk = (
+            torch.zeros(
+                (self.max_bs, self.eagle_worker.dsa_index_topk),
+                dtype=torch.int32,
+                device=model_runner.device,
+            )
+            if self.eagle_worker.seed_dsa_topk_from_draft_extend
+            else None
+        )
+
         self.buffers = EagleDraftInputBuffers(
             input_ids=input_ids,
             req_pool_indices=req_pool_indices,
@@ -245,6 +256,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             hidden_states=hidden_states,
             global_num_tokens_gpu=global_num_tokens_gpu,
             global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+            dsa_seed_topk=dsa_seed_topk,
         )
         self.buffers.share_buffers()
 
@@ -372,6 +384,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             hidden_states=hidden_states,
             capture_hidden_mode=capture_mode,
         )
+        if self.buffers.dsa_seed_topk is not None:
+            spec_info.dsa_topk_indices = self.buffers.dsa_seed_topk[:num_seqs]
 
         sampling_info = SamplingBatchInfo(
             temperatures=self.temperatures[:num_seqs],
@@ -504,6 +518,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
                 buffers.draft_probs.zero_()
             if buffers.hidden_states is not None:
                 buffers.hidden_states.zero_()
+            if buffers.dsa_seed_topk is not None:
+                buffers.dsa_seed_topk.zero_()
             buffers.req_pool_indices.zero_()
 
         num_tokens = bs * self.num_tokens_per_bs
@@ -563,6 +579,12 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
             and forward_batch.spec_info.hidden_states is not None
         ):
             buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
+        if buffers.dsa_seed_topk is not None:
+            seed = forward_batch.spec_info.dsa_topk_indices
+            if seed is not None:
+                buffers.dsa_seed_topk[:raw_bs].copy_(seed)
+            else:
+                buffers.dsa_seed_topk[:raw_bs].zero_()
         # Only rejection sampling reads temperatures (renorm_draft_probs); skip
         # the copy otherwise to keep the non-RS path free of extra work.
         if (

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -69,6 +69,7 @@ class EagleDraftExtendInputBuffers(ForwardInputBuffers):
     next_token_logits_buffer: torch.Tensor
     global_num_tokens_gpu: Optional[torch.Tensor]
     global_num_tokens_for_logprob_gpu: Optional[torch.Tensor]
+    dsa_seed_topk_capture: Optional[torch.Tensor] = None
 
 
 class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
@@ -234,6 +235,17 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64, device="cpu"
         )
 
+        dsa_seed_topk_capture = (
+            torch.full(
+                (self.max_num_token, self.eagle_worker.dsa_index_topk),
+                -1,
+                dtype=torch.int32,
+                device=model_runner.device,
+            )
+            if self.eagle_worker.seed_dsa_topk_from_draft_extend
+            else None
+        )
+
         self.buffers = EagleDraftExtendInputBuffers(
             input_ids=input_ids,
             req_pool_indices=req_pool_indices,
@@ -249,6 +261,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             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,
+            dsa_seed_topk_capture=dsa_seed_topk_capture,
         )
         self.buffers.share_buffers()
 
@@ -385,6 +398,11 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
             padded_static_len=self.padded_static_len,
         )
 
+        if self.buffers.dsa_seed_topk_capture is not None:
+            spec_info.dsa_seed_topk_capture = self.buffers.dsa_seed_topk_capture[
+                :num_tokens
+            ]
+
         def run_once():
             self.draft_extend_attn_backend.init_forward_metadata_in_graph(forward_batch)
 

python/sglang/srt/speculative/eagle_info.py

@@ -161,7 +161,7 @@ class EagleDraftInput(SpecInput):
 
     # Survives across draft steps: spec_info is shared by reference across the
     # per-step forwards (each runs on a copied ForwardBatch, dropping writebacks).
-    mtp_topk_indices: Optional[torch.Tensor] = None
+    dsa_topk_indices: Optional[torch.Tensor] = None
 
     # Per-req bonus token (the "+1" target prediction at end of each accept
     # chain); the worker copies it here post-extend for next iter's draft.
@@ -233,6 +233,8 @@ class EagleDraftInput(SpecInput):
             if self.hidden_states is not None:
                 self.hidden_states = self.hidden_states[: len(new_indices)]
             self.bonus_tokens = self.bonus_tokens[: len(new_indices)]
+            if self.dsa_topk_indices is not None:
+                self.dsa_topk_indices = self.dsa_topk_indices[: len(new_indices)]
         else:
             # in some cases(e.g draft_extend), we have not filtered the batch by `unfinished_index`
             self.topk_p = self.topk_p[new_indices]
@@ -242,6 +244,8 @@ class EagleDraftInput(SpecInput):
             if self.hidden_states is not None:
                 self.hidden_states = self.hidden_states[new_indices]
             self.bonus_tokens = self.bonus_tokens[new_indices]
+            if self.dsa_topk_indices is not None:
+                self.dsa_topk_indices = self.dsa_topk_indices[new_indices]
 
     def merge_batch(self, spec_info: "EagleDraftInput"):
         if self.future_indices is not None:
@@ -260,6 +264,7 @@ class EagleDraftInput(SpecInput):
             self.topk_p = spec_info.topk_p
             self.topk_index = spec_info.topk_index
             self.draft_probs = spec_info.draft_probs
+            self.dsa_topk_indices = spec_info.dsa_topk_indices
             return
         if len(spec_info.topk_index) == 0:
             return
@@ -272,6 +277,12 @@ class EagleDraftInput(SpecInput):
         )
         self.topk_p = torch.cat([self.topk_p, spec_info.topk_p])
         self.topk_index = torch.cat([self.topk_index, spec_info.topk_index])
+        if self.dsa_topk_indices is not None and spec_info.dsa_topk_indices is not None:
+            self.dsa_topk_indices = torch.cat(
+                [self.dsa_topk_indices, spec_info.dsa_topk_indices]
+            )
+        else:
+            self.dsa_topk_indices = None
         if self.draft_probs is not None and spec_info.draft_probs is not None:
             self.draft_probs = torch.cat([self.draft_probs, spec_info.draft_probs])
 
@@ -317,6 +328,9 @@ class EagleDraftExtendInput(SpecInput):
     num_tokens_per_req: int = -1
     num_tokens_for_logprob_per_req: int = 1
 
+    dsa_seed_topk_capture: Optional[torch.Tensor] = None
+    dsa_seed_topk_select: Optional[torch.Tensor] = None
+
     # None for draft-extend's idle batch; attention backends fall back to
     # rebuilding plain metadata from seq_lens when this is None.
     kv_indptr: torch.Tensor = None

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -14,6 +14,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner i
 )
 from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
 from sglang.srt.kv_canary.runner.canary_manager import context_tuple
+from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
 from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
 from sglang.srt.layers.attention.tokenspeed_mla_backend import TokenspeedMLABackend
 from sglang.srt.layers.attention.triton_backend import TritonAttnBackend
@@ -191,16 +192,9 @@ class EagleDraftWorker(EagleDraftWorkerBase):
 
         # Alias for better readability
         self.draft_runner = self.draft_worker.model_runner
-        # Reuse the first draft step's NSA/DSA indexer topk across the rest;
-        # topk == 1 only (select_top_k_tokens reorders rows, desyncing indices).
-        self.index_share_for_mtp_iteration = (
-            getattr(
-                self.draft_runner.model_config.hf_config,
-                "index_share_for_mtp_iteration",
-                False,
-            )
-            and self.topk == 1
-        )
+        self._init_dsa_index_share_state()
+        # Eager draft-extend seed buffer (graph paths use their own static ones).
+        self.dsa_extend_topk_buf: Optional[torch.Tensor] = None
         self.draft_tp_context = (
             draft_tp_context if server_args.enable_dp_attention else empty_context
         )
@@ -264,6 +258,23 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         if (c := self.draft_runner.canary_manager) is not None:
             c.mark_init_finished()
 
+    def _init_dsa_index_share_state(self) -> None:
+        # Populate DSA index-share fields from the draft runner's hf_config.
+        # Reused by the attention unit-test harnesses, which skip __init__.
+        hf_config = self.draft_runner.model_config.hf_config
+        # Reuse the first draft step's DSA indexer topk across the rest;
+        # topk == 1 only (select_top_k_tokens reorders rows, desyncing indices).
+        self.index_share_for_mtp_iteration = (
+            getattr(hf_config, "index_share_for_mtp_iteration", False)
+            and self.topk == 1
+        )
+        # GLM-5.2 MTP IndexShare: seed reused indexer top-k from draft-extend
+        # (last verified token), not draft-decode step 0.
+        self.dsa_index_topk = getattr(hf_config, "index_topk", None)
+        self.seed_dsa_topk_from_draft_extend = (
+            self.index_share_for_mtp_iteration and self.dsa_index_topk is not None
+        )
+
     def _rebuild_topk1_chain_buffers(self) -> None:
         # For topk=1 the draft tree degenerates to a chain, so parent_list and
         # top_scores_index are runtime-invariant. Must be rebuilt after any
@@ -629,8 +640,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         # Forward multiple steps
         scores = None
         if self.index_share_for_mtp_iteration:
-            forward_batch.reuse_mtp_topk_indices = True
-            spec_info.mtp_topk_indices = None
+            forward_batch.reuse_dsa_topk_indices = True
+            # Keep the draft-extend seed so step 0 reuses it; else recompute it.
+            if not (
+                self.seed_dsa_topk_from_draft_extend
+                and spec_info.dsa_topk_indices is not None
+            ):
+                spec_info.dsa_topk_indices = None
         for i in range(self.speculative_num_steps):
             input_ids, hidden_states, scores, tree_info = select_top_k_tokens(
                 i, topk_p, topk_index, hidden_states, scores, self.topk
@@ -706,8 +722,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             forward_batch.positions.add_(1)
 
         if self.index_share_for_mtp_iteration:
-            spec_info.mtp_topk_indices = None
-            forward_batch.reuse_mtp_topk_indices = False
+            spec_info.dsa_topk_indices = None
+            forward_batch.reuse_dsa_topk_indices = False
 
         # Organize the results
         if (
@@ -794,6 +810,23 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         if mm_input_embeds is not None:
             forward_batch.mm_input_embeds = mm_input_embeds
 
+        # Seed the first draft-decode loop from each request's last prefill
+        # position. Gather last-per-req before the copy (prefill can be long).
+        # Skipped under context-parallel prefill (token layout wouldn't match).
+        seed_from_extend = (
+            self.seed_dsa_topk_from_draft_extend
+            and not forward_batch.forward_mode.is_idle()
+            and not dsa_use_prefill_cp(forward_batch)
+        )
+        if seed_from_extend:
+            bs = forward_batch.batch_size
+            forward_batch.spec_info.dsa_seed_topk_capture = (
+                self._get_dsa_extend_topk_buf(bs)
+            )
+            forward_batch.spec_info.dsa_seed_topk_select = (
+                torch.cumsum(forward_batch.extend_seq_lens, dim=0) - 1
+            ).long()
+
         canary_ctx = (
             context_tuple(
                 c.with_ops_outside_graph(
@@ -810,6 +843,10 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         maybe_detect_nan(logits_output.next_token_logits, "draft_extend_for_prefill")
         maybe_detect_inf(logits_output.next_token_logits, "draft_extend_for_prefill")
 
+        prefill_dsa_topk = None
+        if seed_from_extend:
+            prefill_dsa_topk = self.dsa_extend_topk_buf[:bs].clone()
+
         # Assemble the next-iter draft spec_info from the extend output.
         use_rejection_sampling = self.server_args.speculative_use_rejection_sampling
         probs = renorm_draft_probs(
@@ -829,8 +866,22 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             bonus_tokens=next_token_ids,
             num_tokens_per_req=1,
             num_tokens_for_logprob_per_req=1,
+            dsa_topk_indices=prefill_dsa_topk,
         )
 
+    def _get_dsa_extend_topk_buf(self, num_tokens: int) -> torch.Tensor:
+        """Lazily-grown int32 [num_tokens, index_topk] eager draft-extend seed buffer."""
+        buf = self.dsa_extend_topk_buf
+        if buf is None or buf.shape[0] < num_tokens:
+            buf = torch.full(
+                (num_tokens, self.dsa_index_topk),
+                -1,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            self.dsa_extend_topk_buf = buf
+        return buf[:num_tokens]
+
     def _draft_extend_for_decode(
         self, batch: ScheduleBatch, batch_result: GenerationBatchResult
     ):
@@ -882,6 +933,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             and self.cuda_graph_runner_for_draft_extend.can_run_graph(forward_batch)
         )
 
+        # Eager path publishes the indexer top-k into a worker buffer (the graph
+        # path uses the runner's static buffer). Gathered at select_index below.
+        if self.seed_dsa_topk_from_draft_extend and not can_cuda_graph:
+            forward_batch.spec_info.dsa_seed_topk_capture = (
+                self._get_dsa_extend_topk_buf(forward_batch.input_ids.shape[0])
+            )
+
         canary_ctx = (
             context_tuple(
                 c.with_ops_outside_graph(
@@ -912,6 +970,19 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             f"draft_extend_for_decode (cuda_graph={can_cuda_graph})",
         )
 
+        # Gather the per-request last-position indexer top-k as the next loop's
+        # seed (select_index already picks the last accepted position per req).
+        dsa_seed_topk_indices = None
+        if self.seed_dsa_topk_from_draft_extend:
+            if can_cuda_graph:
+                dsa_extend_topk_capture = (
+                    self.cuda_graph_runner_for_draft_extend.buffers.dsa_seed_topk_capture
+                )
+            else:
+                dsa_extend_topk_capture = forward_batch.spec_info.dsa_seed_topk_capture
+            # Fancy indexing returns a fresh tensor (detached from the buffer).
+            dsa_seed_topk_indices = dsa_extend_topk_capture[select_index]
+
         # Reorganize the spec info for the next batch
         draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[
             select_index
@@ -961,6 +1032,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
         )
         if self.server_args.speculative_use_rejection_sampling:
             next_draft_input.draft_probs = ret_draft_probs
+        if self.seed_dsa_topk_from_draft_extend:
+            next_draft_input.dsa_topk_indices = dsa_seed_topk_indices
 
 
 class EAGLEWorkerV2(BaseSpecWorker):

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -113,6 +113,9 @@ class StandaloneDraftWorker(EagleDraftWorker):
             )
             and self.topk == 1
         )
+        self.dsa_index_topk = None
+        self.seed_dsa_topk_from_draft_extend = False
+        self.dsa_extend_topk_buf = None
 
     def alloc_memory_pool(
         self,

python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py

@@ -20,6 +20,7 @@ from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
     EAGLEDraftExtendCudaGraphRunner,
 )
 from sglang.srt.speculative.eagle_info import EagleDraftExtendInput
+from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import fast_topk
 
@@ -496,6 +497,7 @@ class _EagleDraftExtendV2WorkerHarness:
         self.eagle_use_aux_hidden_state = False
         self.hot_token_id = None
         self.draft_runner.model = model_forward
+        EagleDraftWorker._init_dsa_index_share_state(self)
 
 
 def _build_eagle_draft_extend_fixture(

python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py

@@ -188,15 +188,7 @@ class _EagleDraftWorkerHarness:
         self._topk1_parents_prealloc = None
         self._topk1_score_indices_prealloc = None
         EagleDraftWorker._rebuild_topk1_chain_buffers(self)
-        # draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
-        self.index_share_for_mtp_iteration = (
-            getattr(
-                self.model_config.hf_config,
-                "index_share_for_mtp_iteration",
-                False,
-            )
-            and self.topk == 1
-        )
+        EagleDraftWorker._init_dsa_index_share_state(self)
 
     @property
     def draft_model_runner(self):

test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py

@@ -74,6 +74,7 @@ class TestEagleDraftCudaGraphRunner(CustomTestCase):
             hidden_states=torch.empty(CAPTURE_BS, 2, dtype=torch.float32),
             req_pool_indices=torch.empty(CAPTURE_BS, dtype=torch.int32),
             seq_lens_cpu=torch.empty(CAPTURE_BS, dtype=torch.int32),
+            dsa_seed_topk=None,
         )
         runner.capture_bs = [1, CAPTURE_BS]
         runner.num_tokens_per_bs = 1