diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index c4a33f091..fc5dac537 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 68505844f..6ba5e529e 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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 + ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index d6e5fc2f0..031ab4fe8 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index 1a902303e..a51c5b51d 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/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: diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index c801bbe8c..f43e2cfc7 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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 ( diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index af71fc36f..1649c4b1d 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index c1abd35b3..68adadee6 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index bcb63fddd..d65b902da 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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): diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 106600357..d947fcb36 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/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, diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index 2bf6fc5d9..de09706bc 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/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( diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py index bfbabf1d6..6d08a428c 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_runner.py +++ b/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): diff --git a/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py b/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py index 4c853bcf8..644a77744 100644 --- a/test/registered/unit/spec/test_eagle_draft_cuda_graph_runner.py +++ b/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