[Bugfix][MRV2] Initialize DCP metadata after batch partitioning (#55212)

Signed-off-by: Raphael Rialland <raphael.rialland@mistral.ai>
Signed-off-by: Ziyang Zhang <hafuhafu@qq.com>
Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Raphael Rialland <raphael.rialland@mistral.ai>
Co-authored-by: Ziyang Zhang <hafuhafu@qq.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/v1/spec_decode/test_eagle_draft_attn_metadata.py

@@ -144,8 +144,11 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_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_()
+        return out
 
-    with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
+    with patch.object(
+        base_speculator, "maybe_prepare_dcp_local_seq_lens", fake_prepare
+    ):
         captured = _run_build(
             fake,
             num_reqs=3,

tests/v1/worker/test_gpu_input_batch_v2.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Tests for the V2 model runner's InputBatch (vllm.v1.worker.gpu.input_batch)."""
 
+import numpy as np
 import pytest
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
+from vllm.v1.worker.gpu import cp_utils
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 
 DEVICE = current_platform.device_type
@@ -51,3 +54,110 @@ def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
     assert torch.equal(
         batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np)
     )
+
+
+def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
+    """The batch must view the caller-owned buffer, sliced to padded length.
+
+    Runtime (and capture) paths all funnel through this helper; attention
+    metadata indexes padded rows, so the view must reach
+    num_reqs_after_padding, and it must alias the persistent buffer so CUDA
+    graph replay sees the recomputed values.
+    """
+    buffers = InputBuffers(max_num_reqs=4, max_num_tokens=4, device=torch.device("cpu"))
+    batch = InputBatch.make_dummy(2, 4, buffers)
+    batch.num_reqs_after_padding = 4
+
+    def fake_kernel(
+        output,
+        seq_lens,
+        dcp_size,
+        dcp_rank,
+        cp_interleave,
+        num_reqs,
+        max_num_reqs,
+        block_size,
+    ):
+        assert output is buffers.dcp_local_seq_lens
+        assert seq_lens is batch.seq_lens
+        assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (2, 4, 1, 16)
+        assert (max_num_reqs, block_size) == (4, 128)
+        output[:] = torch.tensor([1, 2, 0, 0], dtype=output.dtype)
+
+    class FakeKernel:
+        def __getitem__(self, grid):
+            assert grid == (1,)
+            return fake_kernel
+
+    monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel())
+    batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+        buffers.dcp_local_seq_lens,
+        batch.seq_lens,
+        batch.num_reqs,
+        4,
+        1,
+        16,
+        num_reqs_padded=batch.num_reqs_after_padding,
+    )
+
+    assert batch.dcp_local_seq_lens is not None
+    assert batch.dcp_local_seq_lens.data_ptr() == buffers.dcp_local_seq_lens.data_ptr()
+    assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0]
+
+
+def test_maybe_prepare_dcp_local_seq_lens_clears_stale_metadata(monkeypatch):
+    """dcp_size == 1 resets the field to None instead of leaving a leftover.
+
+    DCP is toggled per deployment, and batches are recycled; a value from an
+    earlier DCP batch would otherwise be consumed as if it were current.
+    """
+    buffers = InputBuffers(max_num_reqs=2, max_num_tokens=2, device=torch.device("cpu"))
+    batch = InputBatch.make_dummy(1, 1, buffers)
+    batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:1]
+
+    def fail_if_called(*args, **kwargs):
+        raise AssertionError("kernel must not run with dcp_size == 1")
+
+    monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", fail_if_called)
+    batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+        buffers.dcp_local_seq_lens,
+        batch.seq_lens,
+        batch.num_reqs,
+        1,
+        0,
+        1,
+        num_reqs_padded=batch.num_reqs_after_padding,
+    )
+
+    assert batch.dcp_local_seq_lens is None
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
+@pytest.mark.parametrize("dcp_size", [2, 4])
+@pytest.mark.parametrize("cp_interleave", [1, 16])
+def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
+    dcp_size: int, cp_interleave: int
+):
+    """Every rank's local lengths must equal the reference torch formula."""
+    device = torch.device("cuda:0")
+    seq_lens_np = np.array([7, 16, 33, 64, 512, 1023], dtype=np.int32)
+    buffers = InputBuffers(max_num_reqs=8, max_num_tokens=32, device=device)
+    batch = InputBatch.make_dummy(6, 12, buffers)
+    buffers.seq_lens[: len(seq_lens_np)] = torch.from_numpy(seq_lens_np).to(device)
+
+    for dcp_rank in range(dcp_size):
+        buffers.dcp_local_seq_lens.fill_(-1)
+        batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
+            buffers.dcp_local_seq_lens,
+            batch.seq_lens,
+            batch.num_reqs,
+            dcp_size,
+            dcp_rank,
+            cp_interleave,
+            num_reqs_padded=batch.num_reqs_after_padding,
+        )
+        expected = get_dcp_local_seq_lens(
+            torch.from_numpy(seq_lens_np), dcp_size, dcp_rank, cp_interleave
+        )
+        assert batch.dcp_local_seq_lens is not None
+        assert torch.equal(batch.dcp_local_seq_lens.cpu(), expected.to(torch.int32))

tests/v1/worker/test_gpu_pcp_manager.py

@@ -1,11 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import replace
+
 import numpy as np
 import pytest
 import torch
 
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
+from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
 from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.pcp_manager import PCPManager
+from vllm.v1.worker.gpu.states import RequestState
 
 
 def _copy_to_cpu(value, out=None, device=None):
@@ -107,3 +113,102 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
             query_start_loc_np=np.arange(4, dtype=np.int32),
             padded_num_tokens=2,
         )
+
+
+def _make_global_decode_batch(
+    num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device
+) -> InputBatch:
+    """A replicate-decode global batch as `prepare_inputs` would build it."""
+    num_reqs = len(num_computed_tokens)
+    num_tokens = num_reqs
+    seq_lens_np = np.asarray(num_computed_tokens, dtype=np.int32) + 1
+
+    base = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
+    buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device)
+    query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32)
+    buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to(
+        device
+    )
+
+    return replace(
+        base,
+        req_ids=[f"req_{i}" for i in range(num_reqs)],
+        num_reqs=num_reqs,
+        num_reqs_after_padding=num_reqs,
+        idx_mapping=torch.arange(num_reqs, dtype=torch.int32, device=device),
+        idx_mapping_np=np.arange(num_reqs, dtype=np.int32),
+        num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32),
+        num_tokens=num_tokens,
+        num_tokens_after_padding=num_tokens,
+        num_draft_tokens=0,
+        num_draft_tokens_per_req=np.zeros(num_reqs, dtype=np.int32),
+        query_start_loc=buffers.query_start_loc[: num_reqs + 1],
+        query_start_loc_np=query_start_loc_np,
+        seq_lens=buffers.seq_lens[:num_reqs],
+        seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_np),
+        dcp_local_seq_lens=None,
+        num_computed_tokens_np=np.asarray(num_computed_tokens, dtype=np.int32),
+        prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
+        num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
+        is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
+        input_ids=buffers.input_ids[:num_tokens],
+        positions=buffers.positions[:num_tokens],
+        is_padding=buffers.is_padding[:num_tokens],
+        prompt_lens=None,
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
+def test_partition_defers_dcp_metadata_to_post_partition_batch():
+    """DCP-local lengths must derive from the partitioned batch, not the
+    global one: the partition replaces seq_lens, so a pre-partition value is
+    stale. partition_batch therefore returns None, and the runtime populates
+    the field afterwards from the PCP-owned buffers.
+    """
+    device = torch.device("cuda:0")
+    req_states = RequestState(
+        max_num_reqs=4,
+        max_model_len=64,
+        max_num_batched_tokens=8,
+        num_speculative_steps=1,
+        vocab_size=8,
+        device=device,
+    )
+    manager = PCPManager(
+        pcp_world_size=2,
+        pcp_rank=0,
+        device=device,
+        req_states=req_states,
+        max_num_reqs=4,
+        max_num_tokens=8,
+        dcp_world_size=2,
+        dcp_rank=0,
+    )
+
+    global_buffers = InputBuffers(4, 8, device)
+    global_batch = _make_global_decode_batch([16, 24], global_buffers, device)
+    # A leftover from an earlier DCP batch must not survive the partition.
+    global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2]
+    global_batch.dcp_local_seq_lens.fill_(-1)
+
+    local_batch = manager.partition_batch(global_batch, padded_num_tokens=2)
+
+    assert local_batch.dcp_local_seq_lens is None
+    assert local_batch.seq_lens.tolist() == [17, 25]
+
+    # What execute_model does next: derive DCP metadata from the final batch
+    # on the PCP-owned buffers.
+    local_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
+        manager.input_buffers.dcp_local_seq_lens,
+        local_batch.seq_lens,
+        local_batch.num_reqs,
+        dcp_size=2,
+        dcp_rank=0,
+        cp_interleave=1,
+        num_reqs_padded=local_batch.num_reqs_after_padding,
+    )
+    expected = get_dcp_local_seq_lens(
+        torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1
+    )
+    assert local_batch.dcp_local_seq_lens is not None
+    assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)

tests/v1/worker/test_gpu_ubatch_slicing.py

@@ -22,7 +22,9 @@ import torch
 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
 from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig
 from vllm.forward_context import create_forward_context
+from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
 from vllm.v1.worker.gpu import dp_utils
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
@@ -581,3 +583,138 @@ def test_ubatch_runner_names_the_microbatch_that_failed():
     assert isinstance(result["error"], RuntimeError)
     assert "Microbatch 0" in str(result["error"])
     assert isinstance(result["error"].__cause__, ValueError)
+
+
+# DCP x DBO: each microbatch recomputes DCP-local seq_lens from its own
+# truncated seq_lens into its own persistent buffer (the parent's values for
+# a straddling request are wrong for the leading microbatch).
+
+DCP_SIZE = 2
+CP_INTERLEAVE = 1
+
+
+def _make_cuda_input_batch(
+    query_lens: list[int], seq_lens: list[int]
+) -> tuple[InputBatch, InputBuffers]:
+    buffers = InputBuffers(
+        max_num_reqs=MAX_NUM_REQS,
+        max_num_tokens=MAX_NUM_TOKENS,
+        device=torch.device("cuda:0"),
+    )
+    return _make_input_batch(query_lens, seq_lens, buffers), buffers
+
+
+def _make_dcp_ubatch_buffers(
+    num_ubatches: int = 2,
+) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """Per-microbatch (query_start_loc, seq_lens, dcp_local_seq_lens)."""
+    return [
+        (
+            torch.zeros(MAX_NUM_REQS + 1, dtype=torch.int32, device="cuda:0"),
+            torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device="cuda:0"),
+            torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device="cuda:0"),
+        )
+        for _ in range(num_ubatches)
+    ]
+
+
+def _slice_with_dcp(
+    input_batch: InputBatch, dcp_rank: int, num_ubatches: int = 2
+) -> tuple[list[InputBatch], list[torch.Tensor]]:
+    dcp_buffers = [
+        torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device="cuda:0")
+        for _ in range(num_ubatches)
+    ]
+    ubatch_buffers = _make_dcp_ubatch_buffers(num_ubatches)
+    ubatches = [
+        _slice_input_batch(
+            input_batch,
+            ubatch_slice,
+            ubatch_buffers[i][0],
+            ubatch_buffers[i][1],
+            dcp_buffers[i],
+            dcp_size=DCP_SIZE,
+            dcp_rank=dcp_rank,
+            cp_interleave=CP_INTERLEAVE,
+        )
+        for i, ubatch_slice in enumerate(
+            create_ubatch_slices(input_batch, num_ubatches)
+        )
+    ]
+    return ubatches, dcp_buffers
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
+@pytest.mark.parametrize("dcp_rank", [0, 1])
+def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int):
+    """A boundary-split request keeps a truncated seq_len on the leading
+    microbatch, whose DCP-local length may differ (in parity) from the full
+    batch's. Microbatches must therefore derive DCP lengths from their own
+    seq_lens, not reuse the parent's row.
+    """
+    input_batch, buffers = _make_cuda_input_batch(
+        [1, 1, 10, 1, 1], [64, 96, 512, 32, 48]
+    )
+    # execute_model has populated the merged batch's DCP metadata already.
+    input_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
+        buffers.dcp_local_seq_lens,
+        input_batch.seq_lens,
+        input_batch.num_reqs,
+        DCP_SIZE,
+        dcp_rank,
+        CP_INTERLEAVE,
+        num_reqs_padded=input_batch.num_reqs_after_padding,
+    )
+    parent_lens = input_batch.dcp_local_seq_lens.clone()
+
+    ubatches, _ = _slice_with_dcp(input_batch, dcp_rank)
+
+    for ubatch in ubatches:
+        assert ubatch.dcp_local_seq_lens is not None
+        expected = get_dcp_local_seq_lens(
+            ubatch.seq_lens.cpu(), DCP_SIZE, dcp_rank, CP_INTERLEAVE
+        )
+        assert torch.equal(
+            ubatch.dcp_local_seq_lens[: ubatch.num_reqs].cpu(),
+            expected.to(torch.int32),
+        )
+
+    # The straddling request (512 tokens, 5 of its 10 query tokens truncated)
+    # lands on this rank with a different local length than the full batch
+    # reports (507 vs 512 -> 254 vs 256 for dcp_size=2, interleave=1); copying
+    # the parent's row would ship the stale value.
+    assert ubatches[0].num_reqs == 3
+    assert ubatches[0].seq_lens[2].item() == 507
+    assert ubatches[0].dcp_local_seq_lens[2].item() != parent_lens[2].item()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
+def test_each_microbatch_owns_its_dcp_buffer():
+    """Microbatch DCP lengths live in per-microbatch persistent buffers so
+    CPU-visible shapes stay valid under CUDA graph capture and replay."""
+    input_batch, _ = _make_cuda_input_batch([4, 4, 4, 4], [8, 9, 10, 11])
+
+    ubatches, dcp_buffers = _slice_with_dcp(input_batch, dcp_rank=0)
+
+    for ubatch, dcp_buffer in zip(ubatches, dcp_buffers):
+        assert ubatch.dcp_local_seq_lens is not None
+        assert ubatch.dcp_local_seq_lens.data_ptr() == dcp_buffer.data_ptr()
+    assert dcp_buffers[0].data_ptr() != dcp_buffers[1].data_ptr()
+
+
+def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off():
+    """dcp_size == 1 yields None even when the parent carries a value, so a
+    microbatch can never consume DCP metadata this deployment did not
+    enable.
+    """
+    buffers = _make_buffers()
+    input_batch = _make_input_batch([4, 4, 4, 4], [8, 9, 10, 11], buffers)
+    input_batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:4]
+
+    ubatch_buffers = _make_ubatch_buffers()
+    ubatches = [
+        _slice_input_batch(input_batch, ubatch_slice, *ubatch_buffers[i])
+        for i, ubatch_slice in enumerate(create_ubatch_slices(input_batch, 2))
+    ]
+
+    assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches)

vllm/v1/attention/backends/flashinfer.py

@@ -628,6 +628,9 @@ class FlashInferTrtllmAPIDecode:
     mask: torch.Tensor | None = None
     """Packed XQA draft mask."""
 
+    dcp_query_start_loc: torch.Tensor | None = None
+    """Decode query boundaries for masking empty DCP shards during reduction."""
+
 
 @dataclass
 class FlashInferMetadata:
@@ -741,7 +744,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             self.dcp_world_size = get_dcp_group().world_size
             self.dcp_rank = get_dcp_group().rank_in_group
             self.dcp_kv_cache_interleave_size = (
-                vllm_config.parallel_config.dcp_kv_cache_interleave_size
+                vllm_config.parallel_config.cp_kv_cache_interleave_size
             )
         except AssertionError:
             # DCP might not be initialized in testing
@@ -1407,17 +1410,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
         if needs_seq_lens_cpu:
             with gpu_sync_allowed():
                 seq_lens_cpu = common_attn_metadata.seq_lens_cpu
-            seq_lens_np = seq_lens_cpu.numpy()
-            num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size
         else:
             seq_lens_cpu = None
-            seq_lens_np = None
-            num_blocks_np = None
 
         # Adjust seq_lens_cpu for DCP
         if self.use_dcp:
             assert seq_lens_cpu is not None
             if num_prefills > 0:
+                # Other attention groups may reuse the same common metadata.
+                seq_lens_cpu = seq_lens_cpu.clone()
                 qo_indptr_prefill_cpu = (
                     qo_indptr_cpu[num_decodes:] - qo_indptr_cpu[num_decodes]
                 )
@@ -1435,6 +1436,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
                 self.dcp_kv_cache_interleave_size,
             )
 
+        # Native paged attention consumes rank-local lengths (context only
+        # for DCP prefills), so derive its page counts after the conversion.
+        if seq_lens_cpu is not None:
+            seq_lens_np = seq_lens_cpu.numpy()
+            num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size
+        else:
+            seq_lens_np = None
+            num_blocks_np = None
+
         # Adjust num_block_np for cascade attention
         if use_cascade:
             assert num_blocks_np is not None
@@ -1700,6 +1710,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
                     q_len_per_req=q_len_per_req,
                     q_cu_seq_lens=q_cu_seq_lens,
                     mask=decode_mask,
+                    dcp_query_start_loc=(
+                        qo_indptr[: num_decodes + 1] if self.use_dcp else None
+                    ),
                 )
             else:
                 assert seq_lens_cpu is not None
@@ -2566,6 +2579,8 @@ class FlashInferImpl(AttentionImpl):
                         out,
                         lse,
                         get_dcp_group(),
+                        seq_lens=seq_lens_decode,
+                        query_start_loc=attn_metadata.decode.dcp_query_start_loc,
                     )
                 elif needs_fp8_out:
                     output[:num_decode_tokens].copy_(out)

vllm/v1/worker/gpu/cp_utils.py

@@ -5,17 +5,19 @@ import torch
 from vllm.triton_utils import tl, triton
 
 
-def prepare_dcp_local_seq_lens(
+def maybe_prepare_dcp_local_seq_lens(
     dcp_local_seq_lens: torch.Tensor,
     seq_lens: torch.Tensor,
     num_reqs: int,
     dcp_size: int,
     dcp_rank: int,
     cp_interleave: int,
-) -> None:
-    """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
+    *,
+    num_reqs_padded: int | None = None,
+) -> torch.Tensor | None:
+    """Populate caller-owned storage and return its padded view, or None without DCP."""
     if dcp_size == 1:
-        return
+        return None
 
     max_num_reqs = dcp_local_seq_lens.shape[0]
     BLOCK_SIZE = 128
@@ -30,6 +32,7 @@ def prepare_dcp_local_seq_lens(
         max_num_reqs,
         BLOCK_SIZE,
     )
+    return dcp_local_seq_lens[:num_reqs_padded]
 
 
 @triton.jit

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -38,7 +38,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
 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 prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches
@@ -699,17 +699,15 @@ def prepare_inputs_to_capture(
         slot_mappings, kv_cache_config
     )
 
-    # HACK(woosuk): Special handling for DCP.
-    if block_tables.cp_size > 1:
-        prepare_dcp_local_seq_lens(
-            input_buffers.dcp_local_seq_lens,
-            input_batch.seq_lens,
-            num_reqs,
-            block_tables.cp_size,
-            block_tables.cp_rank,
-            block_tables.cp_interleave,
-        )
-        input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
+    input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+        input_buffers.dcp_local_seq_lens,
+        input_batch.seq_lens,
+        input_batch.num_reqs,
+        block_tables.cp_size,
+        block_tables.cp_rank,
+        block_tables.cp_interleave,
+        num_reqs_padded=input_batch.num_reqs_after_padding,
+    )
 
     # NOTE(woosuk): Attention metadata is required not just by standard attention
     # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,

vllm/v1/worker/gpu/model_runner.py

@@ -99,7 +99,7 @@ from vllm.v1.worker.gpu.buffer_utils import (
     async_copy_to_gpu,
     set_default_max_concurrency,
 )
-from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
     ModelCudaGraphManager,
@@ -1324,19 +1324,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         )
         seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
 
-        dcp_local_seq_lens = None
-        if self.use_dcp:
-            # Prepare dcp local seq_lens.
-            prepare_dcp_local_seq_lens(
-                self.input_buffers.dcp_local_seq_lens,
-                self.input_buffers.seq_lens,
-                num_reqs,
-                self.dcp_size,
-                self.dcp_rank,
-                self.cp_interleave,
-            )
-            dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded]
-
         # Some input token ids are directly read from the last sampled tokens
         # and draft tokens. Also, get the logits indices to sample tokens from.
         logits_indices = combine_sampled_and_draft_tokens(
@@ -1384,7 +1371,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             query_start_loc_np=query_start_loc_np,
             seq_lens=seq_lens,
             seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
-            dcp_local_seq_lens=dcp_local_seq_lens,
+            dcp_local_seq_lens=None,
             num_computed_tokens_np=num_computed_tokens_np,
             prefill_len_np=batch_req_state.prefill_len_np,
             num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
@@ -1404,11 +1391,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 else None
             ),
         )
-        return pcp.maybe_partition_pcp_batch(
+        input_batch = pcp.maybe_partition_pcp_batch(
             self.pcp_manager,
             input_batch,
             padded_num_tokens=batch_desc.num_tokens,
         )
+        return input_batch
 
     def prepare_attn(
         self, input_batch: InputBatch
@@ -1705,6 +1693,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         attn_metadata = None
         slot_mappings_by_layer = None
+        if not (dummy_run and skip_attn_for_dummy_run):
+            input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+                (
+                    self.pcp_manager.input_buffers
+                    if self.pcp_manager is not None
+                    else self.input_buffers
+                ).dcp_local_seq_lens,
+                input_batch.seq_lens,
+                input_batch.num_reqs,
+                self.dcp_size,
+                self.dcp_rank,
+                self.cp_interleave,
+                num_reqs_padded=input_batch.num_reqs_after_padding,
+            )
         ubatch_state: UBatchState | None = None
         if batch_desc.num_ubatches > 1:
             assert self.ubatch_runner is not None

vllm/v1/worker/gpu/pcp_manager.py

@@ -12,7 +12,6 @@ from vllm.logger import init_logger
 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
-from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import (
     InputBatch,
     InputBuffers,
@@ -537,18 +536,6 @@ class PCPManager:
         seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)
         seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens
 
-        dcp_local_seq_lens = None
-        if self.dcp_world_size > 1:
-            prepare_dcp_local_seq_lens(
-                input_buffers.dcp_local_seq_lens,
-                seq_lens,
-                num_local_reqs,
-                self.dcp_world_size,
-                self.dcp_rank,
-                self.cp_interleave,
-            )
-            dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_local_reqs]
-
         return replace(
             input_batch,
             req_ids=local_req_ids,
@@ -569,7 +556,7 @@ class PCPManager:
             query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1],
             seq_lens=seq_lens,
             seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),
-            dcp_local_seq_lens=dcp_local_seq_lens,
+            dcp_local_seq_lens=None,
             num_computed_tokens_np=local_start_pos_np,
             prefill_len_np=local_prefill_len_np,
             num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,

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

@@ -11,7 +11,7 @@ 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 prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.cudagraph_utils import (
     AttentionState,
     BatchExecutionDescriptor,
@@ -42,17 +42,15 @@ def _prepare_dflash_inputs_to_capture(
     attn_metadata = None
     if not skip_attn:
         query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
-        dcp_local_seq_lens = None
-        if block_tables.cp_size > 1:
-            prepare_dcp_local_seq_lens(
-                input_buffers.dcp_local_seq_lens,
-                input_buffers.seq_lens,
-                num_reqs,
-                block_tables.cp_size,
-                block_tables.cp_rank,
-                block_tables.cp_interleave,
-            )
-            dcp_local_seq_lens = input_buffers.dcp_local_seq_lens
+        input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+            input_buffers.dcp_local_seq_lens,
+            input_batch.seq_lens,
+            input_batch.num_reqs,
+            block_tables.cp_size,
+            block_tables.cp_rank,
+            block_tables.cp_interleave,
+            num_reqs_padded=input_batch.num_reqs_after_padding,
+        )
         attn_metadata = build_attn_metadata(
             attn_groups=attn_groups,
             num_reqs=num_reqs,
@@ -61,7 +59,7 @@ def _prepare_dflash_inputs_to_capture(
             query_start_loc_cpu=query_start_loc_cpu,
             max_query_len=num_tokens // num_reqs,
             seq_lens=input_batch.seq_lens,
-            dcp_local_seq_lens=dcp_local_seq_lens,
+            dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
             max_seq_len=max_model_len,
             block_tables=input_block_tables,
             slot_mappings=slot_mappings,

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

@@ -22,7 +22,7 @@ from vllm.v1.worker.gpu.attn_utils import (
     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.cp_utils import maybe_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
@@ -297,7 +297,7 @@ class DraftModelSpeculator(BaseSpeculator):
         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(
+            dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
                 self.input_buffers.dcp_local_seq_lens,
                 self.input_buffers.seq_lens,
                 num_reqs,
@@ -305,7 +305,6 @@ class DraftModelSpeculator(BaseSpeculator):
                 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,

vllm/v1/worker/gpu/ubatch_utils.py

@@ -10,6 +10,7 @@ import numpy as np
 import torch
 
 from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.distributed import get_dcp_group
 from vllm.forward_context import (
     DPMetadata,
     ForwardContext,
@@ -21,6 +22,7 @@ from vllm.sequence import IntermediateTensors
 from vllm.utils.torch_utils import current_stream
 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.cp_utils import maybe_prepare_dcp_local_seq_lens
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.ubatch_utils import (
@@ -76,16 +78,20 @@ def _slice_input_batch(
     ubatch_slice: UBatchSlice,
     query_start_loc_buf: torch.Tensor,
     seq_lens_buf: torch.Tensor,
+    dcp_local_seq_lens_buf: torch.Tensor | None = None,
+    *,
+    dcp_size: int = 1,
+    dcp_rank: int = 0,
+    cp_interleave: int = 1,
 ) -> InputBatch:
     """Build the sub-`InputBatch` a single microbatch runs on.
 
     A request straddling the boundary appears in both microbatches, its query
     truncated to the tokens each one owns.
 
-    Everything except `query_start_loc` and `seq_lens` is a view of the full
-    batch's buffers, so slicing costs no allocation. Those two are rebased onto
-    the microbatch's token range and written to the caller's buffers, which
-    keeps them valid under CUDA graph capture and replay.
+    Query offsets and sequence lengths (including DCP-local lengths) are
+    written to the caller's persistent buffers after truncation. Other GPU
+    inputs are views of the full batch's buffers.
 
     The sub-batch describes the forward pass only. Sampling runs once over the
     merged batch, so `logits_indices`, `cu_num_logits` and the draft-token
@@ -139,11 +145,18 @@ def _slice_input_batch(
     # max_query_len from this array see the microbatch's own lengths.
     num_scheduled_tokens = np.diff(query_start_loc_np)[:num_reqs]
 
-    dcp_local_seq_lens = input_batch.dcp_local_seq_lens
-    if dcp_local_seq_lens is not None:
-        # NOTE: a request split across microbatches keeps its full local
-        # seq_len here. DCP is not yet validated with DBO.
-        dcp_local_seq_lens = dcp_local_seq_lens[req_start:req_stop]
+    dcp_local_seq_lens = None
+    if dcp_size > 1:
+        assert dcp_local_seq_lens_buf is not None
+        dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
+            dcp_local_seq_lens_buf,
+            seq_lens,
+            num_reqs,
+            dcp_size,
+            dcp_rank,
+            cp_interleave,
+            num_reqs_padded=num_reqs_padded,
+        )
 
     return replace(
         input_batch,
@@ -278,6 +291,9 @@ class UBatchRunner:
         self.vllm_config = vllm_config
         self.parallel_config = vllm_config.parallel_config
         self.num_ubatches = self.parallel_config.num_ubatches
+        self.dcp_size = self.parallel_config.decode_context_parallel_size
+        self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0
+        self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
         self.device = device
         self.model_state = model_state
         self.attn_groups = attn_groups
@@ -293,6 +309,12 @@ class UBatchRunner:
             torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
             for _ in range(self.num_ubatches)
         ]
+        self.ubatch_dcp_local_seq_lens = [
+            torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
+            if self.dcp_size > 1
+            else None
+            for _ in range(self.num_ubatches)
+        ]
         self.comm_stream = torch.cuda.Stream(device=device)
         # The microbatch threads plus the thread that starts them.
         self.ready_barrier = threading.Barrier(self.num_ubatches + 1)
@@ -320,6 +342,10 @@ class UBatchRunner:
                 ubatch_slice,
                 self.ubatch_query_start_loc[i],
                 self.ubatch_seq_lens[i],
+                self.ubatch_dcp_local_seq_lens[i],
+                dcp_size=self.dcp_size,
+                dcp_rank=self.dcp_rank,
+                cp_interleave=self.cp_interleave,
             )
             ubatch_slot_mappings = slot_mappings[:, ubatch_slice.token_slice]
             ubatch_block_tables = tuple(