diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 6f48fb817..883e38d28 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -53,6 +53,7 @@ from sglang.srt.layers.dp_attention import ( ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.utils import ceil_align @@ -676,11 +677,22 @@ class DeepseekV4AttnBackend( max_seq_len = int(seq_lens_cpu.max().item()) if forward_batch.forward_mode.is_decode_or_idle(): + # DSv4 bakes this step's KV write target (c4/c128) into metadata, + # so slice the shared multi-step out_cache_loc now rather than at + # forward time. + out_cache_loc = forward_batch.out_cache_loc + if self.topk > 0 and self.speculative_num_steps > 1: + out_cache_loc = per_step_draft_out_cache_loc( + out_cache_loc, + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, + )[self.speculative_step_id] metadata = self.init_forward_metadata_decode( max_seq_len=max_seq_len, req_pool_indices=req_pool_indices, seq_lens=seq_lens, - out_cache_loc=forward_batch.out_cache_loc, + out_cache_loc=out_cache_loc, ) elif forward_batch.forward_mode.is_target_verify(): metadata = self.init_forward_metadata_target_verify( diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 14f8fe340..a350bb2c4 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -23,6 +23,30 @@ if _is_cuda or _is_hip or _is_musa: ) +def per_step_draft_out_cache_loc( + out_cache_loc: torch.Tensor, + batch_size: int, + topk: int, + num_steps: int, +) -> torch.Tensor: + """Per-step slice of the multi-step EAGLE draft out_cache_loc buffer. + + Single source of truth for the layout shared by EagleWorkerV2.draft_forward + (per-step write target) and DeepseekV4AttnBackend (per-step compression + write target baked into metadata). + """ + expected = batch_size * topk * num_steps + assert out_cache_loc.shape[0] == expected, ( + f"out_cache_loc.shape[0]={out_cache_loc.shape[0]} != " + f"batch_size * topk * num_steps = {batch_size}*{topk}*{num_steps}={expected}" + ) + return ( + out_cache_loc.view(batch_size, topk, num_steps) + .permute(2, 0, 1) + .reshape(num_steps, -1) + ) + + def apply_eagle_prefill_input_rotation( batch: ScheduleBatch, next_token_ids: torch.Tensor ) -> None: diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index a3d14af91..18a947e1a 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -53,7 +53,11 @@ from sglang.srt.speculative.eagle_info_v2 import ( fill_accepted_out_cache_loc, fill_bonus_tokens, ) -from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient +from sglang.srt.speculative.eagle_utils import ( + TreeMaskMode, + build_tree_kernel_efficient, + per_step_draft_out_cache_loc, +) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( draft_tp_context, @@ -438,11 +442,11 @@ class EagleDraftWorker(BaseDraftWorker): if self.hot_token_id is not None: topk_index = self.hot_token_id[topk_index] - out_cache_loc = out_cache_loc.reshape( - forward_batch.batch_size, self.topk, self.speculative_num_steps - ) - out_cache_loc = out_cache_loc.permute((2, 0, 1)).reshape( - self.speculative_num_steps, -1 + out_cache_loc = per_step_draft_out_cache_loc( + out_cache_loc, + forward_batch.batch_size, + self.topk, + self.speculative_num_steps, ) # Return values