@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 83990f5bcc

[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>

Sy03 committed Sep 9, 2026 at 23:50 UTC 83990f5bcc0b5eb08b2f1fd1b109fa8a31b122f3
12 files changed +460 -77
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+4 -1
@@ -144,8 +144,11 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
144 assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (3, 2, 1, 4)
145 out[:num_reqs].copy_(torch.tensor([1, 4, 8], dtype=torch.int32))
146 out[num_reqs:].zero_()
147 + return out
148
148 - with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
149 + with patch.object(
150 + base_speculator, "maybe_prepare_dcp_local_seq_lens", fake_prepare
151 + ):
152 captured = _run_build(
153 fake,
154 num_reqs=3,
tests/v1/worker/test_gpu_input_batch_v2.py
+110
@@ -2,10 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the V2 model runner's InputBatch (vllm.v1.worker.gpu.input_batch)."""
4
5 +import numpy as np
6 import pytest
7 import torch
8
9 from vllm.platforms import current_platform
10 +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
11 +from vllm.v1.worker.gpu import cp_utils
12 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
13
14 DEVICE = current_platform.device_type
@@ -51,3 +54,110 @@ def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
54 assert torch.equal(
55 batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np)
56 )
57 +
58 +
59 +def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
60 + """The batch must view the caller-owned buffer, sliced to padded length.
61 +
62 + Runtime (and capture) paths all funnel through this helper; attention
63 + metadata indexes padded rows, so the view must reach
64 + num_reqs_after_padding, and it must alias the persistent buffer so CUDA
65 + graph replay sees the recomputed values.
66 + """
67 + buffers = InputBuffers(max_num_reqs=4, max_num_tokens=4, device=torch.device("cpu"))
68 + batch = InputBatch.make_dummy(2, 4, buffers)
69 + batch.num_reqs_after_padding = 4
70 +
71 + def fake_kernel(
72 + output,
73 + seq_lens,
74 + dcp_size,
75 + dcp_rank,
76 + cp_interleave,
77 + num_reqs,
78 + max_num_reqs,
79 + block_size,
80 + ):
81 + assert output is buffers.dcp_local_seq_lens
82 + assert seq_lens is batch.seq_lens
83 + assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (2, 4, 1, 16)
84 + assert (max_num_reqs, block_size) == (4, 128)
85 + output[:] = torch.tensor([1, 2, 0, 0], dtype=output.dtype)
86 +
87 + class FakeKernel:
88 + def __getitem__(self, grid):
89 + assert grid == (1,)
90 + return fake_kernel
91 +
92 + monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel())
93 + batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
94 + buffers.dcp_local_seq_lens,
95 + batch.seq_lens,
96 + batch.num_reqs,
97 + 4,
98 + 1,
99 + 16,
100 + num_reqs_padded=batch.num_reqs_after_padding,
101 + )
102 +
103 + assert batch.dcp_local_seq_lens is not None
104 + assert batch.dcp_local_seq_lens.data_ptr() == buffers.dcp_local_seq_lens.data_ptr()
105 + assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0]
106 +
107 +
108 +def test_maybe_prepare_dcp_local_seq_lens_clears_stale_metadata(monkeypatch):
109 + """dcp_size == 1 resets the field to None instead of leaving a leftover.
110 +
111 + DCP is toggled per deployment, and batches are recycled; a value from an
112 + earlier DCP batch would otherwise be consumed as if it were current.
113 + """
114 + buffers = InputBuffers(max_num_reqs=2, max_num_tokens=2, device=torch.device("cpu"))
115 + batch = InputBatch.make_dummy(1, 1, buffers)
116 + batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:1]
117 +
118 + def fail_if_called(*args, **kwargs):
119 + raise AssertionError("kernel must not run with dcp_size == 1")
120 +
121 + monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", fail_if_called)
122 + batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
123 + buffers.dcp_local_seq_lens,
124 + batch.seq_lens,
125 + batch.num_reqs,
126 + 1,
127 + 0,
128 + 1,
129 + num_reqs_padded=batch.num_reqs_after_padding,
130 + )
131 +
132 + assert batch.dcp_local_seq_lens is None
133 +
134 +
135 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
136 +@pytest.mark.parametrize("dcp_size", [2, 4])
137 +@pytest.mark.parametrize("cp_interleave", [1, 16])
138 +def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
139 + dcp_size: int, cp_interleave: int
140 +):
141 + """Every rank's local lengths must equal the reference torch formula."""
142 + device = torch.device("cuda:0")
143 + seq_lens_np = np.array([7, 16, 33, 64, 512, 1023], dtype=np.int32)
144 + buffers = InputBuffers(max_num_reqs=8, max_num_tokens=32, device=device)
145 + batch = InputBatch.make_dummy(6, 12, buffers)
146 + buffers.seq_lens[: len(seq_lens_np)] = torch.from_numpy(seq_lens_np).to(device)
147 +
148 + for dcp_rank in range(dcp_size):
149 + buffers.dcp_local_seq_lens.fill_(-1)
150 + batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
151 + buffers.dcp_local_seq_lens,
152 + batch.seq_lens,
153 + batch.num_reqs,
154 + dcp_size,
155 + dcp_rank,
156 + cp_interleave,
157 + num_reqs_padded=batch.num_reqs_after_padding,
158 + )
159 + expected = get_dcp_local_seq_lens(
160 + torch.from_numpy(seq_lens_np), dcp_size, dcp_rank, cp_interleave
161 + )
162 + assert batch.dcp_local_seq_lens is not None
163 + assert torch.equal(batch.dcp_local_seq_lens.cpu(), expected.to(torch.int32))
tests/v1/worker/test_gpu_pcp_manager.py
+105
@@ -1,11 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from dataclasses import replace
4 +
5 import numpy as np
6 import pytest
7 import torch
8
9 +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
10 +from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
11 from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
12 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
13 from vllm.v1.worker.gpu.pcp_manager import PCPManager
14 +from vllm.v1.worker.gpu.states import RequestState
15
16
17 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):
113 query_start_loc_np=np.arange(4, dtype=np.int32),
114 padded_num_tokens=2,
115 )
116 +
117 +
118 +def _make_global_decode_batch(
119 + num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device
120 +) -> InputBatch:
121 + """A replicate-decode global batch as `prepare_inputs` would build it."""
122 + num_reqs = len(num_computed_tokens)
123 + num_tokens = num_reqs
124 + seq_lens_np = np.asarray(num_computed_tokens, dtype=np.int32) + 1
125 +
126 + base = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
127 + buffers.seq_lens[:num_reqs] = torch.from_numpy(seq_lens_np).to(device)
128 + query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32)
129 + buffers.query_start_loc[: num_reqs + 1] = torch.from_numpy(query_start_loc_np).to(
130 + device
131 + )
132 +
133 + return replace(
134 + base,
135 + req_ids=[f"req_{i}" for i in range(num_reqs)],
136 + num_reqs=num_reqs,
137 + num_reqs_after_padding=num_reqs,
138 + idx_mapping=torch.arange(num_reqs, dtype=torch.int32, device=device),
139 + idx_mapping_np=np.arange(num_reqs, dtype=np.int32),
140 + num_scheduled_tokens=np.ones(num_reqs, dtype=np.int32),
141 + num_tokens=num_tokens,
142 + num_tokens_after_padding=num_tokens,
143 + num_draft_tokens=0,
144 + num_draft_tokens_per_req=np.zeros(num_reqs, dtype=np.int32),
145 + query_start_loc=buffers.query_start_loc[: num_reqs + 1],
146 + query_start_loc_np=query_start_loc_np,
147 + seq_lens=buffers.seq_lens[:num_reqs],
148 + seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_np),
149 + dcp_local_seq_lens=None,
150 + num_computed_tokens_np=np.asarray(num_computed_tokens, dtype=np.int32),
151 + prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
152 + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
153 + is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
154 + input_ids=buffers.input_ids[:num_tokens],
155 + positions=buffers.positions[:num_tokens],
156 + is_padding=buffers.is_padding[:num_tokens],
157 + prompt_lens=None,
158 + )
159 +
160 +
161 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
162 +def test_partition_defers_dcp_metadata_to_post_partition_batch():
163 + """DCP-local lengths must derive from the partitioned batch, not the
164 + global one: the partition replaces seq_lens, so a pre-partition value is
165 + stale. partition_batch therefore returns None, and the runtime populates
166 + the field afterwards from the PCP-owned buffers.
167 + """
168 + device = torch.device("cuda:0")
169 + req_states = RequestState(
170 + max_num_reqs=4,
171 + max_model_len=64,
172 + max_num_batched_tokens=8,
173 + num_speculative_steps=1,
174 + vocab_size=8,
175 + device=device,
176 + )
177 + manager = PCPManager(
178 + pcp_world_size=2,
179 + pcp_rank=0,
180 + device=device,
181 + req_states=req_states,
182 + max_num_reqs=4,
183 + max_num_tokens=8,
184 + dcp_world_size=2,
185 + dcp_rank=0,
186 + )
187 +
188 + global_buffers = InputBuffers(4, 8, device)
189 + global_batch = _make_global_decode_batch([16, 24], global_buffers, device)
190 + # A leftover from an earlier DCP batch must not survive the partition.
191 + global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2]
192 + global_batch.dcp_local_seq_lens.fill_(-1)
193 +
194 + local_batch = manager.partition_batch(global_batch, padded_num_tokens=2)
195 +
196 + assert local_batch.dcp_local_seq_lens is None
197 + assert local_batch.seq_lens.tolist() == [17, 25]
198 +
199 + # What execute_model does next: derive DCP metadata from the final batch
200 + # on the PCP-owned buffers.
201 + local_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
202 + manager.input_buffers.dcp_local_seq_lens,
203 + local_batch.seq_lens,
204 + local_batch.num_reqs,
205 + dcp_size=2,
206 + dcp_rank=0,
207 + cp_interleave=1,
208 + num_reqs_padded=local_batch.num_reqs_after_padding,
209 + )
210 + expected = get_dcp_local_seq_lens(
211 + torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1
212 + )
213 + assert local_batch.dcp_local_seq_lens is not None
214 + assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)
tests/v1/worker/test_gpu_ubatch_slicing.py
+137
@@ -22,7 +22,9 @@ import torch
22 from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
23 from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig
24 from vllm.forward_context import create_forward_context
25 +from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
26 from vllm.v1.kv_cache_interface import KVCacheConfig
27 +from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
28 from vllm.v1.worker.gpu import dp_utils
29 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
30 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
@@ -581,3 +583,138 @@ def test_ubatch_runner_names_the_microbatch_that_failed():
583 assert isinstance(result["error"], RuntimeError)
584 assert "Microbatch 0" in str(result["error"])
585 assert isinstance(result["error"].__cause__, ValueError)
586 +
587 +
588 +# DCP x DBO: each microbatch recomputes DCP-local seq_lens from its own
589 +# truncated seq_lens into its own persistent buffer (the parent's values for
590 +# a straddling request are wrong for the leading microbatch).
591 +
592 +DCP_SIZE = 2
593 +CP_INTERLEAVE = 1
594 +
595 +
596 +def _make_cuda_input_batch(
597 + query_lens: list[int], seq_lens: list[int]
598 +) -> tuple[InputBatch, InputBuffers]:
599 + buffers = InputBuffers(
600 + max_num_reqs=MAX_NUM_REQS,
601 + max_num_tokens=MAX_NUM_TOKENS,
602 + device=torch.device("cuda:0"),
603 + )
604 + return _make_input_batch(query_lens, seq_lens, buffers), buffers
605 +
606 +
607 +def _make_dcp_ubatch_buffers(
608 + num_ubatches: int = 2,
609 +) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
610 + """Per-microbatch (query_start_loc, seq_lens, dcp_local_seq_lens)."""
611 + return [
612 + (
613 + torch.zeros(MAX_NUM_REQS + 1, dtype=torch.int32, device="cuda:0"),
614 + torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device="cuda:0"),
615 + torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device="cuda:0"),
616 + )
617 + for _ in range(num_ubatches)
618 + ]
619 +
620 +
621 +def _slice_with_dcp(
622 + input_batch: InputBatch, dcp_rank: int, num_ubatches: int = 2
623 +) -> tuple[list[InputBatch], list[torch.Tensor]]:
624 + dcp_buffers = [
625 + torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device="cuda:0")
626 + for _ in range(num_ubatches)
627 + ]
628 + ubatch_buffers = _make_dcp_ubatch_buffers(num_ubatches)
629 + ubatches = [
630 + _slice_input_batch(
631 + input_batch,
632 + ubatch_slice,
633 + ubatch_buffers[i][0],
634 + ubatch_buffers[i][1],
635 + dcp_buffers[i],
636 + dcp_size=DCP_SIZE,
637 + dcp_rank=dcp_rank,
638 + cp_interleave=CP_INTERLEAVE,
639 + )
640 + for i, ubatch_slice in enumerate(
641 + create_ubatch_slices(input_batch, num_ubatches)
642 + )
643 + ]
644 + return ubatches, dcp_buffers
645 +
646 +
647 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
648 +@pytest.mark.parametrize("dcp_rank", [0, 1])
649 +def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int):
650 + """A boundary-split request keeps a truncated seq_len on the leading
651 + microbatch, whose DCP-local length may differ (in parity) from the full
652 + batch's. Microbatches must therefore derive DCP lengths from their own
653 + seq_lens, not reuse the parent's row.
654 + """
655 + input_batch, buffers = _make_cuda_input_batch(
656 + [1, 1, 10, 1, 1], [64, 96, 512, 32, 48]
657 + )
658 + # execute_model has populated the merged batch's DCP metadata already.
659 + input_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
660 + buffers.dcp_local_seq_lens,
661 + input_batch.seq_lens,
662 + input_batch.num_reqs,
663 + DCP_SIZE,
664 + dcp_rank,
665 + CP_INTERLEAVE,
666 + num_reqs_padded=input_batch.num_reqs_after_padding,
667 + )
668 + parent_lens = input_batch.dcp_local_seq_lens.clone()
669 +
670 + ubatches, _ = _slice_with_dcp(input_batch, dcp_rank)
671 +
672 + for ubatch in ubatches:
673 + assert ubatch.dcp_local_seq_lens is not None
674 + expected = get_dcp_local_seq_lens(
675 + ubatch.seq_lens.cpu(), DCP_SIZE, dcp_rank, CP_INTERLEAVE
676 + )
677 + assert torch.equal(
678 + ubatch.dcp_local_seq_lens[: ubatch.num_reqs].cpu(),
679 + expected.to(torch.int32),
680 + )
681 +
682 + # The straddling request (512 tokens, 5 of its 10 query tokens truncated)
683 + # lands on this rank with a different local length than the full batch
684 + # reports (507 vs 512 -> 254 vs 256 for dcp_size=2, interleave=1); copying
685 + # the parent's row would ship the stale value.
686 + assert ubatches[0].num_reqs == 3
687 + assert ubatches[0].seq_lens[2].item() == 507
688 + assert ubatches[0].dcp_local_seq_lens[2].item() != parent_lens[2].item()
689 +
690 +
691 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
692 +def test_each_microbatch_owns_its_dcp_buffer():
693 + """Microbatch DCP lengths live in per-microbatch persistent buffers so
694 + CPU-visible shapes stay valid under CUDA graph capture and replay."""
695 + input_batch, _ = _make_cuda_input_batch([4, 4, 4, 4], [8, 9, 10, 11])
696 +
697 + ubatches, dcp_buffers = _slice_with_dcp(input_batch, dcp_rank=0)
698 +
699 + for ubatch, dcp_buffer in zip(ubatches, dcp_buffers):
700 + assert ubatch.dcp_local_seq_lens is not None
701 + assert ubatch.dcp_local_seq_lens.data_ptr() == dcp_buffer.data_ptr()
702 + assert dcp_buffers[0].data_ptr() != dcp_buffers[1].data_ptr()
703 +
704 +
705 +def test_slicing_drops_stale_dcp_metadata_when_dcp_is_off():
706 + """dcp_size == 1 yields None even when the parent carries a value, so a
707 + microbatch can never consume DCP metadata this deployment did not
708 + enable.
709 + """
710 + buffers = _make_buffers()
711 + input_batch = _make_input_batch([4, 4, 4, 4], [8, 9, 10, 11], buffers)
712 + input_batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:4]
713 +
714 + ubatch_buffers = _make_ubatch_buffers()
715 + ubatches = [
716 + _slice_input_batch(input_batch, ubatch_slice, *ubatch_buffers[i])
717 + for i, ubatch_slice in enumerate(create_ubatch_slices(input_batch, 2))
718 + ]
719 +
720 + assert all(ubatch.dcp_local_seq_lens is None for ubatch in ubatches)
vllm/v1/attention/backends/flashinfer.py
+20 -5
@@ -628,6 +628,9 @@ class FlashInferTrtllmAPIDecode:
628 mask: torch.Tensor | None = None
629 """Packed XQA draft mask."""
630
631 + dcp_query_start_loc: torch.Tensor | None = None
632 + """Decode query boundaries for masking empty DCP shards during reduction."""
633 +
634
635 @dataclass
636 class FlashInferMetadata:
@@ -741,7 +744,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
744 self.dcp_world_size = get_dcp_group().world_size
745 self.dcp_rank = get_dcp_group().rank_in_group
746 self.dcp_kv_cache_interleave_size = (
744 - vllm_config.parallel_config.dcp_kv_cache_interleave_size
747 + vllm_config.parallel_config.cp_kv_cache_interleave_size
748 )
749 except AssertionError:
750 # DCP might not be initialized in testing
@@ -1407,17 +1410,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1410 if needs_seq_lens_cpu:
1411 with gpu_sync_allowed():
1412 seq_lens_cpu = common_attn_metadata.seq_lens_cpu
1410 - seq_lens_np = seq_lens_cpu.numpy()
1411 - num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size
1413 else:
1414 seq_lens_cpu = None
1414 - seq_lens_np = None
1415 - num_blocks_np = None
1415
1416 # Adjust seq_lens_cpu for DCP
1417 if self.use_dcp:
1418 assert seq_lens_cpu is not None
1419 if num_prefills > 0:
1420 + # Other attention groups may reuse the same common metadata.
1421 + seq_lens_cpu = seq_lens_cpu.clone()
1422 qo_indptr_prefill_cpu = (
1423 qo_indptr_cpu[num_decodes:] - qo_indptr_cpu[num_decodes]
1424 )
@@ -1435,6 +1436,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1436 self.dcp_kv_cache_interleave_size,
1437 )
1438
1439 + # Native paged attention consumes rank-local lengths (context only
1440 + # for DCP prefills), so derive its page counts after the conversion.
1441 + if seq_lens_cpu is not None:
1442 + seq_lens_np = seq_lens_cpu.numpy()
1443 + num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size
1444 + else:
1445 + seq_lens_np = None
1446 + num_blocks_np = None
1447 +
1448 # Adjust num_block_np for cascade attention
1449 if use_cascade:
1450 assert num_blocks_np is not None
@@ -1700,6 +1710,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1710 q_len_per_req=q_len_per_req,
1711 q_cu_seq_lens=q_cu_seq_lens,
1712 mask=decode_mask,
1713 + dcp_query_start_loc=(
1714 + qo_indptr[: num_decodes + 1] if self.use_dcp else None
1715 + ),
1716 )
1717 else:
1718 assert seq_lens_cpu is not None
@@ -2566,6 +2579,8 @@ class FlashInferImpl(AttentionImpl):
2579 out,
2580 lse,
2581 get_dcp_group(),
2582 + seq_lens=seq_lens_decode,
2583 + query_start_loc=attn_metadata.decode.dcp_query_start_loc,
2584 )
2585 elif needs_fp8_out:
2586 output[:num_decode_tokens].copy_(out)
vllm/v1/worker/gpu/cp_utils.py
+7 -4
@@ -5,17 +5,19 @@ import torch
5 from vllm.triton_utils import tl, triton
6
7
8 -def prepare_dcp_local_seq_lens(
8 +def maybe_prepare_dcp_local_seq_lens(
9 dcp_local_seq_lens: torch.Tensor,
10 seq_lens: torch.Tensor,
11 num_reqs: int,
12 dcp_size: int,
13 dcp_rank: int,
14 cp_interleave: int,
15 -) -> None:
16 - """Populate the persistent DCP local seq_lens buffer (CUDA graph safe)."""
15 + *,
16 + num_reqs_padded: int | None = None,
17 +) -> torch.Tensor | None:
18 + """Populate caller-owned storage and return its padded view, or None without DCP."""
19 if dcp_size == 1:
18 - return
20 + return None
21
22 max_num_reqs = dcp_local_seq_lens.shape[0]
23 BLOCK_SIZE = 128
@@ -30,6 +32,7 @@ def prepare_dcp_local_seq_lens(
32 max_num_reqs,
33 BLOCK_SIZE,
34 )
35 + return dcp_local_seq_lens[:num_reqs_padded]
36
37
38 @triton.jit
vllm/v1/worker/gpu/cudagraph_utils.py
+10 -12
@@ -38,7 +38,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
38 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
39 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
40 from vllm.v1.worker.gpu.block_table import BlockTables
41 -from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
41 +from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
42 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
43 from vllm.v1.worker.gpu.model_states.interface import ModelState
44 from vllm.v1.worker.utils import AttentionGroup, clear_layer_kv_caches
@@ -699,17 +699,15 @@ def prepare_inputs_to_capture(
699 slot_mappings, kv_cache_config
700 )
701
702 - # HACK(woosuk): Special handling for DCP.
703 - if block_tables.cp_size > 1:
704 - prepare_dcp_local_seq_lens(
705 - input_buffers.dcp_local_seq_lens,
706 - input_batch.seq_lens,
707 - num_reqs,
708 - block_tables.cp_size,
709 - block_tables.cp_rank,
710 - block_tables.cp_interleave,
711 - )
712 - input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
702 + input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
703 + input_buffers.dcp_local_seq_lens,
704 + input_batch.seq_lens,
705 + input_batch.num_reqs,
706 + block_tables.cp_size,
707 + block_tables.cp_rank,
708 + block_tables.cp_interleave,
709 + num_reqs_padded=input_batch.num_reqs_after_padding,
710 + )
711
712 # NOTE(woosuk): Attention metadata is required not just by standard attention
713 # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
vllm/v1/worker/gpu/model_runner.py
+18 -16
@@ -99,7 +99,7 @@ from vllm.v1.worker.gpu.buffer_utils import (
99 async_copy_to_gpu,
100 set_default_max_concurrency,
101 )
102 -from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
102 +from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
103 from vllm.v1.worker.gpu.cudagraph_utils import (
104 BatchExecutionDescriptor,
105 ModelCudaGraphManager,
@@ -1324,19 +1324,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1324 )
1325 seq_lens = self.input_buffers.seq_lens[:num_reqs_padded]
1326
1327 - dcp_local_seq_lens = None
1328 - if self.use_dcp:
1329 - # Prepare dcp local seq_lens.
1330 - prepare_dcp_local_seq_lens(
1331 - self.input_buffers.dcp_local_seq_lens,
1332 - self.input_buffers.seq_lens,
1333 - num_reqs,
1334 - self.dcp_size,
1335 - self.dcp_rank,
1336 - self.cp_interleave,
1337 - )
1338 - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded]
1339 -
1327 # Some input token ids are directly read from the last sampled tokens
1328 # and draft tokens. Also, get the logits indices to sample tokens from.
1329 logits_indices = combine_sampled_and_draft_tokens(
@@ -1384,7 +1371,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1371 query_start_loc_np=query_start_loc_np,
1372 seq_lens=seq_lens,
1373 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
1387 - dcp_local_seq_lens=dcp_local_seq_lens,
1374 + dcp_local_seq_lens=None,
1375 num_computed_tokens_np=num_computed_tokens_np,
1376 prefill_len_np=batch_req_state.prefill_len_np,
1377 num_computed_prefill_tokens_np=batch_req_state.num_computed_prefill_tokens_np,
@@ -1404,11 +1391,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1391 else None
1392 ),
1393 )
1407 - return pcp.maybe_partition_pcp_batch(
1394 + input_batch = pcp.maybe_partition_pcp_batch(
1395 self.pcp_manager,
1396 input_batch,
1397 padded_num_tokens=batch_desc.num_tokens,
1398 )
1399 + return input_batch
1400
1401 def prepare_attn(
1402 self, input_batch: InputBatch
@@ -1705,6 +1693,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1693
1694 attn_metadata = None
1695 slot_mappings_by_layer = None
1696 + if not (dummy_run and skip_attn_for_dummy_run):
1697 + input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
1698 + (
1699 + self.pcp_manager.input_buffers
1700 + if self.pcp_manager is not None
1701 + else self.input_buffers
1702 + ).dcp_local_seq_lens,
1703 + input_batch.seq_lens,
1704 + input_batch.num_reqs,
1705 + self.dcp_size,
1706 + self.dcp_rank,
1707 + self.cp_interleave,
1708 + num_reqs_padded=input_batch.num_reqs_after_padding,
1709 + )
1710 ubatch_state: UBatchState | None = None
1711 if batch_desc.num_ubatches > 1:
1712 assert self.ubatch_runner is not None
vllm/v1/worker/gpu/pcp_manager.py
+1 -14
@@ -12,7 +12,6 @@ from vllm.logger import init_logger
12 from vllm.v1.attention.backends.utils import PAD_SLOT_ID
13 from vllm.v1.worker.gpu.block_table import BlockTables
14 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
15 -from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
15 from vllm.v1.worker.gpu.input_batch import (
16 InputBatch,
17 InputBuffers,
@@ -537,18 +536,6 @@ class PCPManager:
536 seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)
537 seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens
538
540 - dcp_local_seq_lens = None
541 - if self.dcp_world_size > 1:
542 - prepare_dcp_local_seq_lens(
543 - input_buffers.dcp_local_seq_lens,
544 - seq_lens,
545 - num_local_reqs,
546 - self.dcp_world_size,
547 - self.dcp_rank,
548 - self.cp_interleave,
549 - )
550 - dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_local_reqs]
551 -
539 return replace(
540 input_batch,
541 req_ids=local_req_ids,
@@ -569,7 +556,7 @@ class PCPManager:
556 query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1],
557 seq_lens=seq_lens,
558 seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),
572 - dcp_local_seq_lens=dcp_local_seq_lens,
559 + dcp_local_seq_lens=None,
560 num_computed_tokens_np=local_start_pos_np,
561 prefill_len_np=local_prefill_len_np,
562 num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
+11 -13
@@ -11,7 +11,7 @@ from vllm.v1.worker.gpu.attn_utils import (
11 build_slot_mappings_by_layer,
12 )
13 from vllm.v1.worker.gpu.block_table import BlockTables
14 -from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
14 +from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
15 from vllm.v1.worker.gpu.cudagraph_utils import (
16 AttentionState,
17 BatchExecutionDescriptor,
@@ -42,17 +42,15 @@ def _prepare_dflash_inputs_to_capture(
42 attn_metadata = None
43 if not skip_attn:
44 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
45 - dcp_local_seq_lens = None
46 - if block_tables.cp_size > 1:
47 - prepare_dcp_local_seq_lens(
48 - input_buffers.dcp_local_seq_lens,
49 - input_buffers.seq_lens,
50 - num_reqs,
51 - block_tables.cp_size,
52 - block_tables.cp_rank,
53 - block_tables.cp_interleave,
54 - )
55 - dcp_local_seq_lens = input_buffers.dcp_local_seq_lens
45 + input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
46 + input_buffers.dcp_local_seq_lens,
47 + input_batch.seq_lens,
48 + input_batch.num_reqs,
49 + block_tables.cp_size,
50 + block_tables.cp_rank,
51 + block_tables.cp_interleave,
52 + num_reqs_padded=input_batch.num_reqs_after_padding,
53 + )
54 attn_metadata = build_attn_metadata(
55 attn_groups=attn_groups,
56 num_reqs=num_reqs,
@@ -61,7 +59,7 @@ def _prepare_dflash_inputs_to_capture(
59 query_start_loc_cpu=query_start_loc_cpu,
60 max_query_len=num_tokens // num_reqs,
61 seq_lens=input_batch.seq_lens,
64 - dcp_local_seq_lens=dcp_local_seq_lens,
62 + dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
63 max_seq_len=max_model_len,
64 block_tables=input_block_tables,
65 slot_mappings=slot_mappings,
vllm/v1/worker/gpu/spec_decode/speculator.py
+2 -3
@@ -22,7 +22,7 @@ from vllm.v1.worker.gpu.attn_utils import (
22 init_attn_backend,
23 )
24 from vllm.v1.worker.gpu.block_table import BlockTables
25 -from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
25 +from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
26 from vllm.v1.worker.gpu.dp_utils import DPSyncState
27 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
28 from vllm.v1.worker.gpu.model_states.interface import ModelState
@@ -297,7 +297,7 @@ class DraftModelSpeculator(BaseSpeculator):
297 if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
298 # Draft steps advance and rewind their own global sequence lengths,
299 # so the target model's DCP-local lengths may already be stale.
300 - prepare_dcp_local_seq_lens(
300 + dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
301 self.input_buffers.dcp_local_seq_lens,
302 self.input_buffers.seq_lens,
303 num_reqs,
@@ -305,7 +305,6 @@ class DraftModelSpeculator(BaseSpeculator):
305 self.block_tables.cp_rank,
306 self.block_tables.cp_interleave,
307 )
308 - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens
308 attn_metadata = build_attn_metadata(
309 attn_groups=self.attn_groups,
310 num_reqs=num_reqs_padded,
vllm/v1/worker/gpu/ubatch_utils.py
+35 -9
@@ -10,6 +10,7 @@ import numpy as np
10 import torch
11
12 from vllm.config import CUDAGraphMode, VllmConfig
13 +from vllm.distributed import get_dcp_group
14 from vllm.forward_context import (
15 DPMetadata,
16 ForwardContext,
@@ -21,6 +22,7 @@ from vllm.sequence import IntermediateTensors
22 from vllm.utils.torch_utils import current_stream
23 from vllm.v1.kv_cache_interface import KVCacheConfig
24 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
25 +from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
26 from vllm.v1.worker.gpu.input_batch import InputBatch
27 from vllm.v1.worker.gpu.model_states.interface import ModelState
28 from vllm.v1.worker.ubatch_utils import (
@@ -76,16 +78,20 @@ def _slice_input_batch(
78 ubatch_slice: UBatchSlice,
79 query_start_loc_buf: torch.Tensor,
80 seq_lens_buf: torch.Tensor,
81 + dcp_local_seq_lens_buf: torch.Tensor | None = None,
82 + *,
83 + dcp_size: int = 1,
84 + dcp_rank: int = 0,
85 + cp_interleave: int = 1,
86 ) -> InputBatch:
87 """Build the sub-`InputBatch` a single microbatch runs on.
88
89 A request straddling the boundary appears in both microbatches, its query
90 truncated to the tokens each one owns.
91
85 - Everything except `query_start_loc` and `seq_lens` is a view of the full
86 - batch's buffers, so slicing costs no allocation. Those two are rebased onto
87 - the microbatch's token range and written to the caller's buffers, which
88 - keeps them valid under CUDA graph capture and replay.
92 + Query offsets and sequence lengths (including DCP-local lengths) are
93 + written to the caller's persistent buffers after truncation. Other GPU
94 + inputs are views of the full batch's buffers.
95
96 The sub-batch describes the forward pass only. Sampling runs once over the
97 merged batch, so `logits_indices`, `cu_num_logits` and the draft-token
@@ -139,11 +145,18 @@ def _slice_input_batch(
145 # max_query_len from this array see the microbatch's own lengths.
146 num_scheduled_tokens = np.diff(query_start_loc_np)[:num_reqs]
147
142 - dcp_local_seq_lens = input_batch.dcp_local_seq_lens
143 - if dcp_local_seq_lens is not None:
144 - # NOTE: a request split across microbatches keeps its full local
145 - # seq_len here. DCP is not yet validated with DBO.
146 - dcp_local_seq_lens = dcp_local_seq_lens[req_start:req_stop]
148 + dcp_local_seq_lens = None
149 + if dcp_size > 1:
150 + assert dcp_local_seq_lens_buf is not None
151 + dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
152 + dcp_local_seq_lens_buf,
153 + seq_lens,
154 + num_reqs,
155 + dcp_size,
156 + dcp_rank,
157 + cp_interleave,
158 + num_reqs_padded=num_reqs_padded,
159 + )
160
161 return replace(
162 input_batch,
@@ -278,6 +291,9 @@ class UBatchRunner:
291 self.vllm_config = vllm_config
292 self.parallel_config = vllm_config.parallel_config
293 self.num_ubatches = self.parallel_config.num_ubatches
294 + self.dcp_size = self.parallel_config.decode_context_parallel_size
295 + self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_size > 1 else 0
296 + self.cp_interleave = self.parallel_config.cp_kv_cache_interleave_size
297 self.device = device
298 self.model_state = model_state
299 self.attn_groups = attn_groups
@@ -293,6 +309,12 @@ class UBatchRunner:
309 torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
310 for _ in range(self.num_ubatches)
311 ]
312 + self.ubatch_dcp_local_seq_lens = [
313 + torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
314 + if self.dcp_size > 1
315 + else None
316 + for _ in range(self.num_ubatches)
317 + ]
318 self.comm_stream = torch.cuda.Stream(device=device)
319 # The microbatch threads plus the thread that starts them.
320 self.ready_barrier = threading.Barrier(self.num_ubatches + 1)
@@ -320,6 +342,10 @@ class UBatchRunner:
342 ubatch_slice,
343 self.ubatch_query_start_loc[i],
344 self.ubatch_seq_lens[i],
345 + self.ubatch_dcp_local_seq_lens[i],
346 + dcp_size=self.dcp_size,
347 + dcp_rank=self.dcp_rank,
348 + cp_interleave=self.cp_interleave,
349 )
350 ubatch_slot_mappings = slot_mappings[:, ubatch_slice.token_slice]
351 ubatch_block_tables = tuple(