From 83990f5bcc0b5eb08b2f1fd1b109fa8a31b122f3 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 9 Sep 2026 23:50:09 -0700 Subject: [PATCH] [Bugfix][MRV2] Initialize DCP metadata after batch partitioning (#55212) Signed-off-by: Raphael Rialland Signed-off-by: Ziyang Zhang Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Raphael Rialland Co-authored-by: Ziyang Zhang Co-authored-by: Codex Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../test_eagle_draft_attn_metadata.py | 5 +- tests/v1/worker/test_gpu_input_batch_v2.py | 110 ++++++++++++++ tests/v1/worker/test_gpu_pcp_manager.py | 105 ++++++++++++++ tests/v1/worker/test_gpu_ubatch_slicing.py | 137 ++++++++++++++++++ vllm/v1/attention/backends/flashinfer.py | 25 +++- vllm/v1/worker/gpu/cp_utils.py | 11 +- vllm/v1/worker/gpu/cudagraph_utils.py | 22 ++- vllm/v1/worker/gpu/model_runner.py | 34 +++-- vllm/v1/worker/gpu/pcp_manager.py | 15 +- .../gpu/spec_decode/dflash/cudagraph.py | 24 ++- vllm/v1/worker/gpu/spec_decode/speculator.py | 5 +- vllm/v1/worker/gpu/ubatch_utils.py | 44 ++++-- 12 files changed, 460 insertions(+), 77 deletions(-) diff --git a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py b/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py index f66d466c53..b3006081e5 100644 --- a/tests/v1/spec_decode/test_eagle_draft_attn_metadata.py +++ b/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, diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py index f517abb2b1..8937f2dc18 100644 --- a/tests/v1/worker/test_gpu_input_batch_v2.py +++ b/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)) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index efb69ab120..eea668aafb 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/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) diff --git a/tests/v1/worker/test_gpu_ubatch_slicing.py b/tests/v1/worker/test_gpu_ubatch_slicing.py index 5c76471803..f54e746710 100644 --- a/tests/v1/worker/test_gpu_ubatch_slicing.py +++ b/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) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 26e1defa3b..4b79135cab 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/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) diff --git a/vllm/v1/worker/gpu/cp_utils.py b/vllm/v1/worker/gpu/cp_utils.py index 77010990b8..470271a95b 100644 --- a/vllm/v1/worker/gpu/cp_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 4191f150ec..287e64ce30 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 1ee6a73072..ec5ce63503 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 146320be13..abec427b37 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/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, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index c68b549745..725714fcdb 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/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, diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index ecf36ce517..3c510e8885 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/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, diff --git a/vllm/v1/worker/gpu/ubatch_utils.py b/vllm/v1/worker/gpu/ubatch_utils.py index 350805707d..012ce6c471 100644 --- a/vllm/v1/worker/gpu/ubatch_utils.py +++ b/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(