[Spec Decode][DCP] Populate draft-local sequence lengths

tests/v1/spec_decode/test_eagle_draft_attn_metadata.py

@@ -37,10 +37,14 @@ def _make_fake_speculator(
     fake_input_buffers = SimpleNamespace(
         query_start_loc=torch.zeros(max_num_reqs + 1, dtype=torch.int32),
         seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
+        dcp_local_seq_lens=torch.zeros(max_num_reqs, dtype=torch.int32),
     )
     fake_block_tables = SimpleNamespace(
         input_block_tables=[torch.zeros(max_num_reqs, 4, dtype=torch.int32)],
         slot_mappings=torch.zeros(1, max_num_tokens, dtype=torch.int64),
+        cp_size=1,
+        cp_rank=0,
+        cp_interleave=1,
     )
     return SimpleNamespace(
         arange=torch.arange(max_num_reqs + 1, dtype=torch.int32, device="cpu"),
@@ -126,3 +130,32 @@ def test_build_draft_attn_metadata_clamps_to_max_model_len():
     bound = captured["seq_lens_cpu_upper_bound"]
     # 1023 + 3 = 1026 -> clamped to 1024; 500 + 3 = 503 unaffected.
     assert torch.equal(bound, torch.tensor([1024, 503], dtype=torch.int32))
+
+
+def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
+    fake = _make_fake_speculator()
+    fake.block_tables.cp_size = 2
+    fake.block_tables.cp_rank = 1
+    fake.block_tables.cp_interleave = 4
+    fake.input_buffers.seq_lens[:3] = torch.tensor([5, 9, 16])
+
+    def fake_prepare(out, seq_lens, num_reqs, dcp_size, dcp_rank, cp_interleave):
+        assert seq_lens is fake.input_buffers.seq_lens
+        assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4)
+        out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32))
+        out[num_reqs:].zero_()
+
+    with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
+        captured = _run_build(
+            fake,
+            num_reqs=3,
+            num_reqs_padded=4,
+            num_tokens_padded=4,
+            base=torch.tensor([5, 9, 16]),
+            step=0,
+        )
+
+    local = captured["dcp_local_seq_lens"]
+    assert isinstance(local, torch.Tensor)
+    assert local.data_ptr() == fake.input_buffers.dcp_local_seq_lens.data_ptr()
+    assert local.tolist() == [1, 4, 8, 0]

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -18,7 +18,7 @@
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
 from vllm.v1.worker.gpu.block_table import BlockTables
-from vllm.v1.worker.gpu.cp_utils import cp_local_slot, prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import cp_local_slot
 from vllm.v1.worker.gpu.dp_utils import DPSyncState, dispatch_cg_and_sync_dp
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -291,16 +291,6 @@ def _build_draft_attn_metadata(
         if not self.draft_attn_layer_names:
             return None
         assert num_query_per_req is None  # Omitted for DFlash, read from self instead
-        if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
-            prepare_dcp_local_seq_lens(
-                self.input_buffers.dcp_local_seq_lens,
-                self.input_buffers.seq_lens,
-                num_reqs,
-                self.block_tables.cp_size,
-                self.block_tables.cp_rank,
-                self.block_tables.cp_interleave,
-            )
-            dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
         return super()._build_draft_attn_metadata(
             num_reqs,
             num_reqs_padded,

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -21,6 +21,7 @@
     init_attn_backend,
 )
 from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.dp_utils import DPSyncState
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -274,6 +275,18 @@ def _build_draft_attn_metadata(
             out=draft_seq_lens_cpu_upper_bound[:num_reqs],
         )
         draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len)
+        if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
+            # Draft steps advance and rewind their own global sequence lengths,
+            # so the target model's DCP-local lengths may already be stale.
+            prepare_dcp_local_seq_lens(
+                self.input_buffers.dcp_local_seq_lens,
+                self.input_buffers.seq_lens,
+                num_reqs,
+                self.block_tables.cp_size,
+                self.block_tables.cp_rank,
+                self.block_tables.cp_interleave,
+            )
+            dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
         attn_metadata = build_attn_metadata(
             attn_groups=self.attn_groups,
             num_reqs=num_reqs_padded,