[Bug Fix] Fix DSv4 EAGLE draft CUDA graph replay metadata

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -805,7 +805,6 @@ def init_forward_metadata_replay_cuda_graph(
             device = seq_lens.device
             seq_lens = torch.ones(bs, dtype=seq_lens.dtype, device=device)
             seq_lens_cpu = torch.ones(bs, dtype=torch.int64)
-            seq_lens_sum = bs
             req_pool_indices = torch.zeros(
                 bs, dtype=req_pool_indices.dtype, device=device
             )
@@ -823,6 +822,12 @@ def init_forward_metadata_replay_cuda_graph(
         if bucket == _GraphBucket.DECODE_OR_IDLE:
             assert out_cache_loc is not None
             assert len(out_cache_loc.shape) == 1, f"{out_cache_loc.shape=}"
+            if len(out_cache_loc) > bs:
+                raise ValueError(
+                    "DSv4 decode replay metadata expects one out_cache_loc per "
+                    f"graph batch entry; got {len(out_cache_loc)} locations for "
+                    f"{bs=}. Multi-step EAGLE replay must pass a per-step slice."
+                )
             out_cache_loc_padded = torch.nn.functional.pad(
                 out_cache_loc,
                 pad=(0, bs - len(out_cache_loc)),
@@ -1227,26 +1232,63 @@ def init_forward_metadata_replay_cuda_graph(
         if self.speculative_num_steps == 1:
             return
 
-        self.attn_backends[0]._replay_forward_batch = forward_batch
-        self.attn_backends[0].init_forward_metadata_replay_cuda_graph(
-            bs=bs,
-            req_pool_indices=forward_batch.req_pool_indices,
-            seq_lens=forward_batch.seq_lens,
-            seq_lens_sum=forward_batch.seq_lens_sum,
-            encoder_lens=None,
-            forward_mode=ForwardMode.DECODE,
-            spec_info=forward_batch.spec_info,
-            seq_lens_cpu=forward_batch.seq_lens_cpu,
+        # EAGLE stores draft cache locations request-major, while each captured
+        # decode step consumes one step-major slice.
+        step_out_cache_locs = _split_eagle_draft_out_cache_loc_by_step(
+            forward_batch.out_cache_loc,
+            batch_size=bs,
+            topk=self.topk,
+            speculative_num_steps=self.speculative_num_steps,
         )
-        self.attn_backends[0]._replay_forward_batch = None
-        temp_metadata = self.attn_backends[0].forward_metadata
-
-        for i in range(1, self.speculative_num_steps - 1):
-            self.attn_backends[i].replay_cuda_graph_metadata_from(
-                bs=bs,
-                temp_metadata=temp_metadata,
-                bucket=_GraphBucket.DECODE_OR_IDLE,
-            )
+        original_out_cache_loc = forward_batch.out_cache_loc
+        try:
+            for i in range(self.speculative_num_steps - 1):
+                backend = self.attn_backends[i]
+                forward_batch.out_cache_loc = step_out_cache_locs[i]
+                backend._replay_forward_batch = forward_batch
+                backend.init_forward_metadata_replay_cuda_graph(
+                    bs=bs,
+                    req_pool_indices=forward_batch.req_pool_indices,
+                    seq_lens=forward_batch.seq_lens,
+                    seq_lens_sum=forward_batch.seq_lens_sum,
+                    encoder_lens=None,
+                    forward_mode=ForwardMode.DECODE,
+                    spec_info=forward_batch.spec_info,
+                    seq_lens_cpu=forward_batch.seq_lens_cpu,
+                )
+                backend._replay_forward_batch = None
+        finally:
+            forward_batch.out_cache_loc = original_out_cache_loc
+            for backend in self.attn_backends:
+                backend._replay_forward_batch = None
+
+
+def _split_eagle_draft_out_cache_loc_by_step(
+    out_cache_loc: torch.Tensor,
+    *,
+    batch_size: int,
+    topk: int,
+    speculative_num_steps: int,
+) -> torch.Tensor:
+    """Match the request-major EAGLE draft layout used by draft_forward.
+
+    Keep this in sync with the equivalent reshape/permute in eagle_worker.py
+    and eagle_worker_v2.py.
+    """
+    expected_num_locs = batch_size * topk * speculative_num_steps
+    if out_cache_loc.numel() != expected_num_locs:
+        raise ValueError(
+            "EAGLE draft out_cache_loc must be padded to the CUDA graph batch "
+            f"before DSv4 replay metadata is initialized: got {out_cache_loc.numel()} "
+            f"locations, expected {expected_num_locs} for {batch_size=} {topk=} "
+            f"{speculative_num_steps=}."
+        )
+
+    return (
+        out_cache_loc.reshape(batch_size, topk, speculative_num_steps)
+        .permute(2, 0, 1)
+        .reshape(speculative_num_steps, batch_size * topk)
+    )
 
 
 def _pad_tensor_to_size(tensor: torch.Tensor, size: int, *, value: int = 0):

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -398,6 +398,12 @@ def replay(self, forward_batch: ForwardBatch):
 
         raw_bs = forward_batch.batch_size
         raw_num_token = raw_bs * self.num_tokens_per_bs
+        original_batch_size = forward_batch.batch_size
+        original_positions = forward_batch.positions
+        original_seq_lens = forward_batch.seq_lens
+        original_req_pool_indices = forward_batch.req_pool_indices
+        original_out_cache_loc = forward_batch.out_cache_loc
+        original_seq_lens_cpu = forward_batch.seq_lens_cpu
 
         # Pad
         if self.require_mlp_tp_gather:
@@ -415,6 +421,8 @@ def replay(self, forward_batch: ForwardBatch):
         bs = self.capture_bs[index]
         if bs != raw_bs:
             buffers.seq_lens.fill_(self.seq_len_fill_value)
+            # Padded graph lanes must write to reserved cache slot 0.
+            # DSv4 replay metadata relies on these zero-filled padding slots.
             buffers.out_cache_loc.zero_()
             buffers.positions.zero_()
             buffers.topk_p.zero_()
@@ -461,6 +469,11 @@ def replay(self, forward_batch: ForwardBatch):
             forward_batch.batch_size = bs
             forward_batch.seq_lens = buffers.seq_lens[:bs]
             forward_batch.req_pool_indices = buffers.req_pool_indices[:bs]
+            # Replay metadata must see the same padded cache-location layout as
+            # the captured graph, not the raw request batch.
+            forward_batch.out_cache_loc = buffers.out_cache_loc[
+                : num_tokens * self.speculative_num_steps
+            ]
             forward_batch.positions = buffers.positions[:num_tokens]
 
         if forward_batch.seq_lens_cpu is not None:
@@ -469,24 +482,28 @@ def replay(self, forward_batch: ForwardBatch):
             buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
             forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:bs]
 
-        self.draft_attn_backend.init_forward_metadata_replay_cuda_graph(
-            forward_batch, bs
-        )
-        self.raw_bs = raw_bs
-        self.bs = bs
-        # TODO: The forward_batch.seq_len_sum might need to be updated to reflect the padding in the cuda graph
+        try:
+            self.draft_attn_backend.init_forward_metadata_replay_cuda_graph(
+                forward_batch, bs
+            )
+            self.raw_bs = raw_bs
+            self.bs = bs
+            # TODO: The forward_batch.seq_len_sum might need to be updated to reflect the padding in the cuda graph
+
+            # Replay
+            self._replay(forward_batch)
+        finally:
+            if bs != raw_bs:
+                forward_batch.batch_size = original_batch_size
+                forward_batch.positions = original_positions
+                forward_batch.seq_lens = original_seq_lens
+                forward_batch.req_pool_indices = original_req_pool_indices
+                forward_batch.out_cache_loc = original_out_cache_loc
+                forward_batch.seq_lens_cpu = original_seq_lens_cpu
 
-        # Replay
-        self._replay(forward_batch)
         out = self.output_buffers[bs]
 
         if bs != raw_bs:
             out = self._postprocess_output_to_raw_bs(out, raw_bs)
-            forward_batch.batch_size = raw_bs
-            forward_batch.positions = buffers.positions[:raw_num_token]
-            forward_batch.seq_lens = buffers.seq_lens[:raw_bs]
-            forward_batch.req_pool_indices = buffers.req_pool_indices[:raw_bs]
-            if forward_batch.seq_lens_cpu is not None:
-                forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:raw_bs]
 
         return out