diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index f33c582554..3c65611966 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -110,12 +110,17 @@ steps: - csrc/ - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml + - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml - tests/evals/gsm8k/configs/models-pcp.txt - vllm/model_executor/layers/quantization - vllm/config/parallel.py - vllm/distributed/parallel_state.py - vllm/model_executor/layers/attention/mla_attention.py - vllm/model_executor/layers/attention/pcp.py + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - vllm/model_executor/layers/sparse_attn_indexer.py + - vllm/v1/attention/backends/mla/flashmla_sparse.py + - vllm/v1/attention/backends/mla/indexer.py - vllm/v1/worker/gpu/model_runner.py - vllm/v1/worker/gpu/pcp_manager.py autorun_on_main: true diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml new file mode 100644 index 0000000000..f57467ac74 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml @@ -0,0 +1,19 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +max_concurrency: 100 +server_args: >- + --compilation-config '{"cudagraph_mode":"PIECEWISE"}' + --max-model-len 4096 + --max-num-batched-tokens 32768 + --safetensors-load-strategy prefetch + --moe-backend flashinfer_cutlass + --prefill-context-parallel-size 4 + --decode-context-parallel-size 4 + --enable-expert-parallel + --kv-cache-dtype fp8 +env: + PYTORCH_CUDA_ALLOC_CONF: "expandable_segments:True" + VLLM_LOGGING_LEVEL: "DEBUG" + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/models-pcp.txt b/tests/evals/gsm8k/configs/models-pcp.txt index df645edad5..e6d5bc9d39 100644 --- a/tests/evals/gsm8k/configs/models-pcp.txt +++ b/tests/evals/gsm8k/configs/models-pcp.txt @@ -1,2 +1,3 @@ GLM-5.2-NVFP4-TP2-PCP2-EP.yaml GLM-5.2-NVFP4-TP1-PCP4-EP.yaml +GLM-5.2-NVFP4-TP1-PCP4-DCP4-EP.yaml diff --git a/tests/v1/attention/test_indexer_dcp_localize.py b/tests/v1/attention/test_indexer_dcp_localize.py index 7f7c20d53e..ffcfcbfac1 100644 --- a/tests/v1/attention/test_indexer_dcp_localize.py +++ b/tests/v1/attention/test_indexer_dcp_localize.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np import pytest import torch @@ -962,3 +963,37 @@ def test_sparse_decode_dcp_short_context_matches_non_dcp(): dcp_out, dcp_lse = _dcp_lse_merge(local_outs, local_lses) torch.testing.assert_close(dcp_out, ref_out, atol=1e-5, rtol=1e-5) torch.testing.assert_close(dcp_lse, ref_lse, atol=1e-5, rtol=1e-5) + + +@pytest.mark.parametrize("dcp_world_size", [2, 4, 8]) +@pytest.mark.parametrize("req_lens", [[7], [8, 8], [1, 9], [256, 1, 2730]]) +def test_pcp_plan_deinterleave_restores_global_order(dcp_world_size, req_lens): + """The index gather must undo DCP sharding per request""" + from vllm.v1.attention.backends.mla.indexer import build_pcp_global_chunk_plan + + scheduled = np.array(req_lens, dtype=np.int64) + shard_rows = -(-scheduled // dcp_world_size) + rows = np.arange(len(scheduled)) + plan = build_pcp_global_chunk_plan( + rows, shard_rows, dcp_world_size, torch.device("cpu") + ) + + starts = np.concatenate([[0], np.cumsum(scheduled)]) + + # Build each rank's padded shard exactly as the cache gather would. + padded_cu = plan.padded_local_cu.tolist() + shards = torch.zeros(dcp_world_size, plan.padded_local_total, 1) + for r in range(dcp_world_size): + for i, g in enumerate(req_lens): + for t in range(r, g, dcp_world_size): + shards[r, padded_cu[i] + t // dcp_world_size, 0] = starts[i] + t + + gathered = shards.reshape(dcp_world_size * plan.padded_local_total, 1) + restored = gathered[plan.deinterleave_idx] + # The layout is padded per request; only the real positions are read. + row_start = plan.row_start_cu.tolist() + for i, g in enumerate(req_lens): + torch.testing.assert_close( + restored[row_start[i] : row_start[i] + g, 0], + torch.arange(starts[i], starts[i] + g, dtype=torch.float32), + ) diff --git a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py index aabe88eeee..49462ea8fb 100644 --- a/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py +++ b/tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py @@ -212,6 +212,7 @@ def test_index_conversion_warmup_uses_physical_block_stride(): ), parallel_config=SimpleNamespace( decode_context_parallel_size=1, + prefill_context_parallel_size=1, cp_kv_cache_interleave_size=1, ), ) diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index e7239320c2..0766584f54 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1080,6 +1080,7 @@ def test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens(): _convert_logical_to_physical_topk=_convert_topk, index_group=None, index_group_index=0, + pcp_dcp_kv_gather=False, ) out, _ = FlashMLASparseImpl._forward_bf16_kv( @@ -3257,6 +3258,7 @@ def test_flashmla_fp8_metadata_excludes_zero_token_decode_padding(monkeypatch): device=torch.device(DEVICE_TYPE), dummy_block_table=torch.zeros(7, 1, device=DEVICE_TYPE), max_model_len_tensor=torch.zeros(7, device=DEVICE_TYPE), + pcp_dcp_kv_gather=False, ) query_start_loc_cpu = torch.tensor([0, 110, 220, 330, 440, 550, 660, 660]) common_metadata = SimpleNamespace( @@ -3363,6 +3365,7 @@ def test_flashmla_fp8_paths_accept_decode_subset(monkeypatch, use_mixed_batch: b index_group=None, index_group_index=0, dcp_world_size=1, + pcp_dcp_kv_gather=False, need_to_return_lse_for_decode=False, _fp8_flash_mla_kernel=run_kernel, _convert_logical_to_physical_topk=( diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index 2d8941f6ff..dc935929d4 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -114,6 +114,155 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): ) +def _rank_rows( + pcp_rank: int, + pcp_world_size: int, + num_scheduled_tokens: np.ndarray, + num_computed_tokens: np.ndarray, + is_prefilling: np.ndarray, +) -> tuple[np.ndarray, np.ndarray]: + """One PCP+DCP rank's rows as (global request, extent).""" + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=pcp_rank, + device=torch.device("cpu"), + dcp_world_size=2, + ) + query_start_loc_np = np.concatenate([[0], np.cumsum(num_scheduled_tokens)]).astype( + np.int32 + ) + segments = manager._get_rank_segments( + pcp_rank, num_scheduled_tokens, is_prefilling, query_start_loc_np + ) + rows = np.array([segment.global_batch_req_idx for segment in segments]) + extents = (num_computed_tokens + num_scheduled_tokens)[rows] + return rows, extents + + +@pytest.mark.parametrize( + ("num_scheduled_tokens", "num_computed_tokens"), + [ + ([32, 32], [0, 0]), + ([32, 32, 32, 32], [0, 0, 0, 0]), + ([32, 32], [100, 0]), # one continued, one fresh + ([64, 32, 48], [0, 512, 0]), # ragged lengths and mixed contexts + # A replicated prefill BEHIND a split one. Sorting replicated rows first + # put request 1 ahead of request 0. + ([64, 3], [0, 0]), + ([64, 3, 64], [0, 0, 0]), + ], +) +@pytest.mark.parametrize("pcp_world_size", [2, 4]) +def test_published_row_order_is_identical_on_every_pcp_rank( + num_scheduled_tokens, num_computed_tokens, pcp_world_size +): + """Every rank must map its rows to the same global requests, in the same order.""" + num_scheduled_tokens = np.array(num_scheduled_tokens, dtype=np.int32) + num_computed_tokens = np.array(num_computed_tokens, dtype=np.int32) + is_prefilling = np.ones(len(num_scheduled_tokens), dtype=np.bool_) + + orders = [ + _rank_rows( + rank, + pcp_world_size, + num_scheduled_tokens, + num_computed_tokens, + is_prefilling, + )[0] + for rank in range(pcp_world_size) + ] + for rank, order in enumerate(orders[1:], start=1): + assert np.array_equal(orders[0], order), ( + f"rank {rank} rows map to {order.tolist()}, rank 0 to {orders[0].tolist()}" + ) + # Grouped by request and ascending: what the indexer plan indexes by. + assert np.all(np.diff(orders[0]) >= 0) + + +def test_published_row_order_puts_every_decode_before_every_prefill(): + """split_decodes_and_prefills takes the FIRST prefilling row as the boundary.""" + # req 0 is a continued prefill, req 1 is a decode. + req_idx, _ = _rank_rows( + pcp_rank=0, + pcp_world_size=2, + num_scheduled_tokens=np.array([32, 1], dtype=np.int32), + num_computed_tokens=np.array([100, 20], dtype=np.int32), + is_prefilling=np.array([True, False], dtype=np.bool_), + ) + # The decode (request 1) must be row 0. + assert req_idx[0] == 1 + + +def test_split_prefill_rows_repeat_the_request_and_its_full_extent(): + """Equal adjacent request indices are what let a backend share a KV region.""" + req_idx, extents = _rank_rows( + pcp_rank=1, + pcp_world_size=2, + num_scheduled_tokens=np.array([64, 3], dtype=np.int32), + num_computed_tokens=np.array([100, 0], dtype=np.int32), + is_prefilling=np.ones(2, dtype=np.bool_), + ) + assert req_idx.tolist() == [0, 0, 1] + assert extents.tolist() == [164, 164, 3] + + +@pytest.mark.parametrize("query_len", [1, 2, 3, 5, 6, 9]) +def test_dcp_replicates_prefills_too_short_to_split(query_len): + """A prefill that cannot fill 2*pcp chunks is replicated, not split.""" + pcp_world_size = 2 + num_scheduled_tokens = np.array([query_len], dtype=np.int32) + is_prefilling = np.ones(1, dtype=np.bool_) + + rows_per_rank = [] + for rank in range(pcp_world_size): + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=rank, + device=torch.device("cpu"), + dcp_world_size=2, + ) + assert manager.replicated_requests(num_scheduled_tokens, is_prefilling)[0] + rows = list( + manager._iter_rank_chunks(rank, num_scheduled_tokens, is_prefilling) + ) + assert rows == [(0, 0, query_len)] + rows_per_rank.append(rows) + assert rows_per_rank[0] == rows_per_rank[1] + + +@pytest.mark.parametrize("pcp_world_size", [2, 4, 8]) +@pytest.mark.parametrize( + "query_len", [16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 1000, 4097] +) +def test_pcp_first_chunk_row_is_never_short(pcp_world_size, query_len): + num_scheduled_tokens = np.array([query_len], dtype=np.int32) + is_prefilling = np.ones(1, dtype=np.bool_) + + charged_per_rank = [] + for rank in range(pcp_world_size): + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=rank, + device=torch.device("cpu"), + dcp_world_size=pcp_world_size, + ) + chunk_lens = [ + chunk_len + for _, _, chunk_len in manager._iter_rank_chunks( + rank, num_scheduled_tokens, is_prefilling + ) + ] + assert chunk_lens, f"rank {rank} got no rows for {query_len=}" + assert chunk_lens[0] == max(chunk_lens), ( + f"rank {rank} emitted a short first chunk for {query_len=}: {chunk_lens}" + ) + charged_per_rank.append(len(chunk_lens) * chunk_lens[0]) + + assert len(set(charged_per_rank)) == 1, ( + f"ranks would chunk differently for {query_len=}: {charged_per_rank}" + ) + + def _make_global_decode_batch( num_computed_tokens: list[int], buffers: InputBuffers, device: torch.device ) -> InputBatch: diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 21596eb043..1c2e5b3962 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -420,6 +420,9 @@ class MLAAttention(nn.Module, AttentionLayerBase): """ supports_dense_mha_prefill: ClassVar[bool] = True + # Under PCP+DCP only the decode rows carry an LSE; the base forward + # merges a full-batch LSE, so subclasses opt in with their own forward. + supports_pcp_dcp: ClassVar[bool] = False def __init__( self, @@ -660,6 +663,10 @@ class MLAAttention(nn.Module, AttentionLayerBase): _COMPUTE_PREFILL_METADATA_KERNEL.register_warmup() + if self.use_pcp and self.impl.dcp_world_size > 1 and not self.supports_pcp_dcp: + raise NotImplementedError( + f"{type(self).__name__} does not support PCP+DCP." + ) self.dcp_manager: MLADCPManager | None = None if self.impl.dcp_world_size > 1: query_dtype = ( diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 5baddcd0a5..79ca612032 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -190,7 +190,8 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): pin_memory=PIN_MEMORY, ) parallel_config = vllm_config.parallel_config - self.use_pcp = parallel_config.prefill_context_parallel_size > 1 + self.pcp_world_size = parallel_config.prefill_context_parallel_size + self.use_pcp = self.pcp_world_size > 1 try: self.dcp_world_size = get_dcp_group().world_size except AssertionError: @@ -409,8 +410,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): q_data_type=self.model_config.dtype, output_dtype=self.model_config.dtype, prefill_backend=self._prefill_backend, - use_dense_mha=_use_dense_mha_prefill( - self.vllm_config, prefill_max_seq_len, self.topk_tokens + use_dense_mha=( + _use_dense_mha_prefill( + self.vllm_config, prefill_max_seq_len, self.topk_tokens + ) + # Dense MHA cannot be used under PCP+DCP. + and not (self.use_pcp and self.dcp_world_size > 1) ), topk_mask_workspace=self.topk_mask_workspace, host_staging_plan=staging_plan, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 817f8ad361..d663c54de3 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -3,6 +3,7 @@ """Custom Sparse Attention Indexer layers.""" import torch +import torch.distributed as dist import vllm.envs as envs from vllm import _custom_ops as ops @@ -130,6 +131,23 @@ def _merge_dcp_topk_global( ) +def dcp_gather_kv_rows( + local_padded: torch.Tensor, + deinterleave_idx: torch.Tensor, + gathered_buf: torch.Tensor, + out_buf: torch.Tensor, +) -> torch.Tensor: + """All-gather this rank's KV shard and return it in global token order.""" + + gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]] + dist.all_gather_into_tensor( + gathered, local_padded.contiguous(), group=get_dcp_group().device_group + ) + return torch.index_select( + gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]] + ) + + @triton.jit def _fused_indexer_q_rope_quant_kernel( positions, @@ -346,11 +364,19 @@ def sparse_attn_indexer( values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, fp8_dtype, use_fp4_cache ) - current_workspace_manager().get_simultaneous( + profile_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ values_spec, scales_spec, ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), - ) + ] + if use_pcp and dcp_world_size > 1: + # The PCP+DCP path takes an all-gather destination and a + # de-interleaved result. + gather_spec = _gather_workspace_shapes( + total_seq_lens, head_dim, fp8_dtype, use_fp4_cache + ) + profile_specs.extend(gather_spec * 2) + current_workspace_manager().get_simultaneous(*profile_specs) # Dummy allocation to simulate for peak logits tensor memory during inference. # FP8 elements so elements == bytes @@ -462,9 +488,27 @@ def sparse_attn_indexer( values_spec, scales_spec = _gather_workspace_shapes( total_seq_lens, head_dim, fp8_dtype, use_fp4_cache ) - k_quant_full, k_scale_full = workspace_manager.get_simultaneous( + # PCP + DCP needs two more pairs: the rank-major all-gather destination + # and the de-interleaved result. + pcp_chunks = [ + (c, c.pcp_deinterleave_idx) + for c in prefill_metadata.chunks + if c.pcp_deinterleave_idx is not None + ] + gather_specs: list[tuple[tuple[int, int], torch.dtype]] = [] + if pcp_chunks: + gathered_rows = max( + dcp_world_size * c.local_total_seq_lens for c, _ in pcp_chunks + ) + deinterleaved_rows = max(idx.shape[0] for _, idx in pcp_chunks) + for rows in (gathered_rows, deinterleaved_rows): + gather_specs.extend( + _gather_workspace_shapes(rows, head_dim, fp8_dtype, use_fp4_cache) + ) + k_quant_full, k_scale_full, *gather_bufs = workspace_manager.get_simultaneous( values_spec, scales_spec, + *gather_specs, ) for chunk in prefill_metadata.chunks: cu_seqlen_ks = chunk.cu_seqlen_ks @@ -481,6 +525,23 @@ def sparse_attn_indexer( chunk.local_cu_seq_lens, ) + # PCP + DCP KV all-gather. + deinterleave_idx = chunk.pcp_deinterleave_idx + if deinterleave_idx is not None: + # local_total_seq_lens is the PCP-padded extent here, so every + # rank contributes an identically shaped shard. + gathered_values, gathered_scales, out_values, out_scales = gather_bufs + shard = k_quant[: chunk.local_total_seq_lens] + k_quant = dcp_gather_kv_rows( + shard, deinterleave_idx, gathered_values, out_values + ) + k_scale = dcp_gather_kv_rows( + k_scale[: chunk.local_total_seq_lens], + deinterleave_idx, + gathered_scales, + out_scales, + ) + q_slice = q_quant[chunk.token_start : chunk.token_end] q_scale_slice = ( q_scale[chunk.token_start : chunk.token_end] @@ -491,7 +552,7 @@ def sparse_attn_indexer( chunk.token_start : chunk.token_end, :topk_tokens ] - if chunk.local_total_seq_lens == 0: + if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) topk_indices.fill_(-1) else: @@ -561,15 +622,18 @@ def sparse_attn_indexer( topk_tokens, ) - _merge_dcp_topk_global( - logits, - topk_indices, - topk_tokens, - dcp_rank, - dcp_world_size, - cp_kv_cache_interleave_size, - row_starts=chunk.cu_seqlen_ks, - ) + if deinterleave_idx is None: + # Under the PCP path the top-k already ran over the whole + # context. + _merge_dcp_topk_global( + logits, + topk_indices, + topk_tokens, + dcp_rank, + dcp_world_size, + cp_kv_cache_interleave_size, + row_starts=chunk.cu_seqlen_ks, + ) if has_decode: decode_metadata = attn_metadata_narrowed.decode @@ -910,6 +974,10 @@ class SparseAttnIndexer(CustomOp): if current_platform.is_cuda() or current_platform.is_xpu(): return self.forward_cuda(hidden_states, q_quant, k, weights) elif current_platform.is_rocm(): + if self.use_pcp and self.dcp_world_size > 1: + raise NotImplementedError( + "The ROCm sparse-indexer path does not support PCP+DCP." + ) return self.forward_hip(hidden_states, q_quant, k, weights) elif current_platform.is_cpu(): return self.forward_cpu(hidden_states, q_quant, k, weights) diff --git a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py index ba7a5adff1..2a314042de 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer_kpool.py +++ b/vllm/model_executor/layers/sparse_attn_indexer_kpool.py @@ -919,6 +919,16 @@ class SparseAttnIndexerKpool(CustomOp): raise RuntimeError( "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." ) + _cfg = get_current_vllm_config_or_none() + _parallel = _cfg.parallel_config if _cfg is not None else None + if ( + _parallel is not None + and _parallel.prefill_context_parallel_size > 1 + and _parallel.decode_context_parallel_size > 1 + ): + raise NotImplementedError( + "SparseAttnIndexerKpool does not support PCP+DCP." + ) def forward_native( self, diff --git a/vllm/models/deepseek_v32/attention.py b/vllm/models/deepseek_v32/attention.py index b0fe554e7e..00dd2e18a8 100644 --- a/vllm/models/deepseek_v32/attention.py +++ b/vllm/models/deepseek_v32/attention.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import TYPE_CHECKING, cast +from typing import TYPE_CHECKING, ClassVar, cast import torch import torch.nn as nn @@ -121,6 +121,7 @@ class DeepseekV32Indexer(nn.Module): class DeepseekV32Attention(MLAAttention): indexer: "DeepseekV32Indexer | None" indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer + supports_pcp_dcp: ClassVar[bool] = True def __init__( self, @@ -569,8 +570,9 @@ class DeepseekV32Attention(MLAAttention): seq_lens: torch.Tensor | None query_start_loc: torch.Tensor | None if self.use_pcp: - if attn_metadata.decode is not None: - seq_lens = attn_metadata.decode.seq_lens + decode_metadata = getattr(attn_metadata, "decode", None) + if decode_metadata is not None: + seq_lens = decode_metadata.seq_lens else: all_seq_lens = cast( torch.Tensor, @@ -585,12 +587,23 @@ class DeepseekV32Attention(MLAAttention): # PCP-only empty-shard metadata is needed. seq_lens = None query_start_loc = None - attn_out = self.dcp_manager.combine( - attn_out, - lse, - seq_lens=seq_lens, # type: ignore[arg-type] - query_start_loc=query_start_loc, # type: ignore[arg-type] - ) + # Under PCP+DCP the prefill rows attended over the gathered KV, so + # only the decode rows carry an LSE and take part in the merge. + num_merge_rows = lse.shape[0] + if num_merge_rows == attn_out.shape[0]: + attn_out = self.dcp_manager.combine( + attn_out, + lse, + seq_lens=seq_lens, # type: ignore[arg-type] + query_start_loc=query_start_loc, # type: ignore[arg-type] + ) + elif num_merge_rows > 0: + attn_out[:num_merge_rows] = self.dcp_manager.combine( + attn_out[:num_merge_rows], + lse, + seq_lens=seq_lens, # type: ignore[arg-type] + query_start_loc=query_start_loc, # type: ignore[arg-type] + ) if self.use_pcp: attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index addd68ec6a..c427d3a1de 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -415,9 +415,14 @@ class CommonAttentionMetadata: encoder_seq_lens_cpu: np.ndarray | None = None dcp_local_seq_lens: torch.Tensor | None = None - dcp_local_seq_lens_cpu: torch.Tensor | None = None """Sequence lengths of the local rank in decode context parallelism world""" + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + """(batch_size,) CPU upper bound on dcp_local_seq_lens. Under PCP+DCP it + holds, on every row, the largest shard any DCP rank has of the row's + whole request, identical on every PCP rank, so the sparse backends pad + their KV gather to it.""" + positions: torch.Tensor | None = None """(num_actual_tokens,) token positions. Optional; set when the caller has positions available so that builders can pre-compute position-dependent @@ -434,6 +439,11 @@ class CommonAttentionMetadata: decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.""" + req_idx: np.ndarray | None = None + """(batch_size,) index of each row's request in the runner's request + table. Rows of one request are adjacent, so equal neighbours are PCP + chunks sharing one KV context.""" + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None """PrefixLM bidirectional ranges for multimodal tokens. Maps request index to list of (start, end) token position ranges @@ -525,9 +535,12 @@ class CommonAttentionMetadata: encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu), dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens), - dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu), + dcp_local_seq_lens_cpu_upper_bound=maybe_slice_reqs( + self.dcp_local_seq_lens_cpu_upper_bound + ), seq_lens_cpu_upper_bound=maybe_slice_reqs(self.seq_lens_cpu_upper_bound), is_prefilling=maybe_slice_reqs(self.is_prefilling), + req_idx=maybe_slice_reqs(self.req_idx), rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens), replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu), ) diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 4ecf559683..433cded567 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py @@ -3,11 +3,14 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar +import numpy as np import torch +import torch.distributed as dist from vllm import _custom_ops as ops from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType +from vllm.distributed.parallel_state import get_dcp_group from vllm.logger import init_logger from vllm.model_executor.layers.attention.sparse_mla_attention import ( SparseMLACommonImpl, @@ -28,10 +31,12 @@ from vllm.v1.attention.backend import ( from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, + request_row_bounds, triton_convert_req_index_to_global_index, triton_filter_and_convert_dcp_index, ) from vllm.v1.attention.backends.utils import ( + get_dcp_local_seq_lens, reshape_attn_output_for_spec_decode, reshape_query_for_spec_decode, split_prefill_chunks, @@ -43,6 +48,7 @@ from vllm.v1.attention.ops.flashmla import ( get_mla_metadata, ) from vllm.v1.kv_cache_interface import AttentionSpec +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu from vllm.v1.worker.workspace import current_workspace_manager if TYPE_CHECKING: @@ -171,6 +177,28 @@ class FlashMLASparseBackend(AttentionBackend): return None +def gathered_prefill_shards( + row_req_idx: np.ndarray, row_shard_rows: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Under PCP+DCP, group this rank's prefill rows by request and take each + request's slice of this rank's KV shard. + + PCP gives a rank two adjacent chunk rows of a split prefill; they share + one context, so the workspace holds it once. ``row_shard_rows`` is + ``CommonAttentionMetadata.dcp_local_seq_lens_cpu_upper_bound``: the + largest DCP shard of the whole request on every row, identical on every + PCP rank, so the all-gather is shaped the same everywhere. + + Returns the request row bounds and, per request, its shard rows. + """ + row_bounds = request_row_bounds(row_req_idx) + shard_rows = row_shard_rows[row_bounds[:-1]].astype(np.int64) + assert np.all(shard_rows > 0), ( + f"PCP+DCP prefill got an empty context: {shard_rows.tolist()}" + ) + return row_bounds, shard_rows + + @dataclass class FlashMLASparseMetadata(SparseMLACommonMetadata): @dataclass @@ -331,8 +359,13 @@ class FlashMLASparseMetadataBuilder( device=device, ) + # PCP+DCP attends prefill rows over the DCP-gathered KV, which only the + # separate prefill/decode path can do. + self.pcp_dcp_kv_gather = self.use_pcp and self.dcp_world_size > 1 self.fp8_use_mixed_batch = ( - self.num_heads < MIN_HEADS_FOR_BF16_PREFILL and not self.use_hisparse + self.num_heads < MIN_HEADS_FOR_BF16_PREFILL + and not self.use_hisparse + and not self.pcp_dcp_kv_gather ) if parallel_config.decode_context_parallel_size > 1: @@ -342,7 +375,12 @@ class FlashMLASparseMetadataBuilder( "default 'ag_rs' DCP comm backend; got " f"'{parallel_config.dcp_comm_backend}'" ) - if not self.fp8_use_mixed_batch: + if self.pcp_dcp_kv_gather and cache_config.cache_dtype != "fp8_ds_mla": + raise NotImplementedError( + "PCP+DCP sparse prefill gathers the KV through the fp8_ds_mla " + f"upconvert; got a {cache_config.cache_dtype} cache" + ) + if not self.fp8_use_mixed_batch and not self.pcp_dcp_kv_gather: raise NotImplementedError( "DCP for FlashMLA sparse is only supported on the " "mixed-batch fp8 path (num_heads < " @@ -353,9 +391,16 @@ class FlashMLASparseMetadataBuilder( # Head padding (and the tile-scheduler metadata sized from it) is # computed from the local head count, but the kernel runs on the # DCP-gathered heads. - gathered_num_heads = ( - self.num_heads * parallel_config.decode_context_parallel_size - ) + if self.use_pcp: + gathered_num_heads = ( + self.num_heads * parallel_config.tensor_parallel_size + if self.dcp_world_size > self.pcp_world_size + else self.num_heads + ) + else: + gathered_num_heads = ( + self.num_heads * parallel_config.decode_context_parallel_size + ) gathered_padded_heads = FlashMLASparseImpl._compute_fp8_decode_padded_heads( gathered_num_heads ) @@ -433,6 +478,40 @@ class FlashMLASparseMetadataBuilder( query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu prefill_seq_lens_cpu = seq_lens_cpu[num_decodes:] + # Chunk prefill requests to fit within workspace size + max_prefill_buffer_size = get_prefill_workspace_size( + self.vllm_config.model_config.max_model_len + ) + # Workspace rows per entry, and the prefill rows each entry covers. + workspace_rows = prefill_seq_lens_cpu + row_bounds = np.arange(num_prefills + 1) + prefill_block_table = common_attn_metadata.block_table_tensor[num_decodes:] + prefill_seq_lens = common_attn_metadata.seq_lens[num_decodes:] + if self.pcp_dcp_kv_gather: + # One entry per request, holding this rank's KV shard: the + # all-gather of every rank's shard is the whole context. + # A dummy batch bypasses the PCP manager: one row per request, + # so the row's own extent is the request's. + row_req_idx = common_attn_metadata.req_idx + if row_req_idx is None: + row_req_idx = np.arange(common_attn_metadata.num_reqs) + shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound + if shard_rows_cpu is None: + shard_rows_cpu = get_dcp_local_seq_lens( + seq_lens_cpu, + self.dcp_world_size, + 0, + self.cp_kv_cache_interleave_size, + ) + row_bounds, rows_per_rank = gathered_prefill_shards( + row_req_idx[num_decodes:], shard_rows_cpu[num_decodes:].numpy() + ) + workspace_rows = torch.from_numpy(rows_per_rank.astype(np.int32)) + max_prefill_buffer_size //= self.dcp_world_size + entry_rows = async_copy_to_gpu(row_bounds[:-1], device=self.device) + prefill_block_table = prefill_block_table.index_select(0, entry_rows) + prefill_seq_lens = prefill_seq_lens.index_select(0, entry_rows) + num_entries = len(workspace_rows) # Build prefill_request_id: -1 for decode, request index for # prefill. This enables a single @@ -441,33 +520,28 @@ class FlashMLASparseMetadataBuilder( (num_tokens,), -1, dtype=torch.int32, device=self.device ) # Map prefill tokens to their request IDs (0, 1, 2, ...) - for req_idx in range(num_prefills): + for req_idx in range(num_entries): # Get query token range for this prefill request - global_req_idx = num_decodes + req_idx - req_query_start = query_start_loc_cpu[global_req_idx] - req_query_end = query_start_loc_cpu[global_req_idx + 1] + req_query_start = query_start_loc_cpu[ + num_decodes + int(row_bounds[req_idx]) + ] + req_query_end = query_start_loc_cpu[ + num_decodes + int(row_bounds[req_idx + 1]) + ] prefill_request_id[req_query_start:req_query_end] = req_idx # will be adjusted by chunk loop prefill_workspace_starts_cpu = torch.zeros( - num_prefills, dtype=torch.int32, pin_memory=True - ) - prefill_workspace_starts_cpu[1:] = torch.cumsum( - prefill_seq_lens_cpu[:-1], dim=0 + num_entries, dtype=torch.int32, pin_memory=True ) + prefill_workspace_starts_cpu[1:] = torch.cumsum(workspace_rows[:-1], dim=0) # populated by non-blocking copy after prefill_workspace_starts_cpu is # updated by each chunk prefill_workspace_starts = torch.empty( - num_prefills, dtype=torch.int32, device=self.device + num_entries, dtype=torch.int32, device=self.device ) - # Chunk prefill requests to fit within workspace size - max_prefill_buffer_size = get_prefill_workspace_size( - self.vllm_config.model_config.max_model_len - ) - chunk_bounds = split_prefill_chunks( - prefill_seq_lens_cpu, max_prefill_buffer_size - ) + chunk_bounds = split_prefill_chunks(workspace_rows, max_prefill_buffer_size) prefill_chunks = [] for chunk_start, chunk_end in chunk_bounds: @@ -480,19 +554,19 @@ class FlashMLASparseMetadataBuilder( offset = prefill_workspace_starts_cpu[chunk_start].item() prefill_workspace_starts_cpu[chunk_start:chunk_end] -= offset - chunk_tot_seqlen = prefill_seq_lens_cpu[chunk_start:chunk_end].sum() - token_start = query_start_loc_cpu[num_decodes + chunk_start].item() - token_end = query_start_loc_cpu[num_decodes + chunk_end].item() + chunk_tot_seqlen = workspace_rows[chunk_start:chunk_end].sum() + token_start = query_start_loc_cpu[ + num_decodes + int(row_bounds[chunk_start]) + ].item() + token_end = query_start_loc_cpu[ + num_decodes + int(row_bounds[chunk_end]) + ].item() tokens_slice = slice(token_start, token_end) # Create chunk view of gpu tensor chunk_workspace_starts = prefill_workspace_starts[chunk_start:chunk_end] - chunk_block_table = common_attn_metadata.block_table_tensor[ - num_decodes + chunk_start : num_decodes + chunk_end - ] - chunk_seq_lens = common_attn_metadata.seq_lens[ - num_decodes + chunk_start : num_decodes + chunk_end - ] + chunk_block_table = prefill_block_table[chunk_start:chunk_end] + chunk_seq_lens = prefill_seq_lens[chunk_start:chunk_end] prefill_chunks.append( FP8Meta.Prefill.Chunk( @@ -516,14 +590,19 @@ class FlashMLASparseMetadataBuilder( ) if num_decodes > 0: - # Use padded head count since that's what the kernel will see - scheduler_metadata, _ = get_mla_metadata() + if self.pcp_dcp_kv_gather: + # Decode rows take the mixed-batch kernel call, which returns + # the LSE the DCP merge needs. + kernel_meta = self._build_fp8_mixed_decode_prefill() + else: + # Use padded head count since that's what the kernel will see + scheduler_metadata, _ = get_mla_metadata() - kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata( - scheduler_metadata=scheduler_metadata, - dummy_block_table=self.dummy_block_table[:active_num_decodes], - cache_lens=self.max_model_len_tensor[:active_num_decodes], - ) + kernel_meta = FlashMLASparseMetadata.FP8KernelMetadata( + scheduler_metadata=scheduler_metadata, + dummy_block_table=self.dummy_block_table[:active_num_decodes], + cache_lens=self.max_model_len_tensor[:active_num_decodes], + ) fp8_metadata.decode = FP8Meta.Decode( seq_lens=common_attn_metadata.seq_lens[:active_num_decodes], kernel_metadata=kernel_meta, @@ -626,19 +705,35 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): "the bf16 sparse path is not supported under DCP." ) + self.pcp_dcp_kv_gather = False + self.gathered_kv_workspace: torch.Tensor | None = None if kv_cache_dtype in QUANTIZED_DS_MLA_CACHE_FORMATS: # Reserve workspace during initialization assert vllm_config is not None and vllm_config.model_config is not None prefill_workspace_size = get_prefill_workspace_size( vllm_config.model_config.max_model_len ) - self.prefill_workspace_shape = (prefill_workspace_size, head_size) - self.q_concat_buffer, self.prefill_bf16_workspace = ( - current_workspace_manager().get_simultaneous( - (q_concat_shape, torch.bfloat16), - (self.prefill_workspace_shape, torch.bfloat16), - ) + parallel_config = vllm_config.parallel_config + self.pcp_dcp_kv_gather = ( + parallel_config.prefill_context_parallel_size > 1 + and parallel_config.decode_context_parallel_size > 1 ) + shard_rows = prefill_workspace_size + if self.pcp_dcp_kv_gather: + # PCP+DCP upconverts this rank's KV shard, then all-gathers the + # shards into a workspace of the full prefill size. + shard_rows //= parallel_config.decode_context_parallel_size + self.prefill_workspace_shape = (shard_rows, head_size) + specs = [ + (q_concat_shape, torch.bfloat16), + (self.prefill_workspace_shape, torch.bfloat16), + ] + if self.pcp_dcp_kv_gather: + specs.append(((prefill_workspace_size, head_size), torch.bfloat16)) + buffers = current_workspace_manager().get_simultaneous(*specs) + self.q_concat_buffer, self.prefill_bf16_workspace = buffers[:2] + if self.pcp_dcp_kv_gather: + self.gathered_kv_workspace = buffers[2] else: (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous( (q_concat_shape, torch.bfloat16), @@ -729,13 +824,49 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): return attn_out, lse return torch.cat([decode_out, attn_out], dim=0), None + def _gather_prefill_chunk( + self, + chunk: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill.Chunk", + shard: torch.Tensor, + topk_indices: torch.Tensor, + attn_metadata: FlashMLASparseMetadata, + prefill_meta: "FlashMLASparseMetadata.FP8SeparatePrefillDecode.Prefill", + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """All-gather this rank's upconverted KV shard so the chunk's rows attend + the whole context, and map their top-k onto the rank-major result.""" + shard_rows = int(chunk.chunk_tot_seqlen) + assert self.gathered_kv_workspace is not None + gathered_kv = self.gathered_kv_workspace[: self.dcp_world_size * shard_rows] + dist.all_gather_into_tensor( + gathered_kv, shard, group=get_dcp_group().device_group + ) + + tokens = chunk.tokens_slice + topk_indices, topk_length = triton_convert_req_index_to_global_index( + attn_metadata.req_id_per_token[tokens], + attn_metadata.block_table, + topk_indices, + BLOCK_SIZE=attn_metadata.block_size, + NUM_TOPK_TOKENS=topk_indices.shape[1], + HAS_PREFILL_WORKSPACE=True, + prefill_workspace_request_ids=prefill_meta.request_ids[tokens], + prefill_workspace_starts=prefill_meta.workspace_starts, + prefill_workspace_rank_stride=shard_rows, + dcp_size=self.dcp_world_size, + dcp_rank=self.dcp_rank, + cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size, + return_valid_counts=True, + ) + return gathered_kv, topk_indices, topk_length + def _forward_fp8_kv_separate_prefill_decode( self, q: torch.Tensor, kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """The lse covers the decode rows and is only returned under PCP+DCP.""" fp8_metadata = attn_metadata.fp8_extra_metadata assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8SeparatePrefillDecode) num_decodes = fp8_metadata.num_decodes @@ -783,8 +914,13 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): # For BF16 cache: always use global cache slots (no workspace) # prefill_workspace_starts has been adjusted in-place per chunk so # prefill indices automatically come out chunk-local - topk_length = None - if num_prefill_tokens == 0 and not uses_host_cache: + # Under PCP+DCP the decode rows are converted by the mixed-batch call + # and the prefill rows per chunk, against the gathered workspace. + lse: torch.Tensor | None = None + topk_length: torch.Tensor | None = None + if self.pcp_dcp_kv_gather: + pass + elif num_prefill_tokens == 0 and not uses_host_cache: topk_indices, topk_length = self._convert_logical_to_physical_topk( topk_indices, attn_metadata, @@ -812,6 +948,17 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): topk_indices: torch.Tensor, ) -> torch.Tensor: assert fp8_metadata.decode is not None + if self.pcp_dcp_kv_gather: + # Attend this rank's KV shard; the LSE feeds the DCP merge. + nonlocal lse + out, lse = self._forward_fp8_kv_mixed_batch( + q, + kv_c_and_k_pe_cache, + topk_indices, + attn_metadata, + kernel_metadata=fp8_metadata.decode.kernel_metadata, + ) + return out if uses_host_cache: return self._host_backed_fp8_decode( q, @@ -896,9 +1043,20 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): ) chunk_q = q[chunk.tokens_slice] - chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] - assert topk_length is not None - chunk_topk_length = topk_length[chunk.tokens_slice] + if self.pcp_dcp_kv_gather: + chunk_workspace, chunk_topk_indices_workspace, chunk_topk_length = ( + self._gather_prefill_chunk( + chunk, + chunk_workspace, + topk_indices[chunk.tokens_slice], + attn_metadata, + fp8_metadata.prefill, + ) + ) + else: + chunk_topk_indices_workspace = topk_indices[chunk.tokens_slice] + assert topk_length is not None + chunk_topk_length = topk_length[chunk.tokens_slice] attn_out[chunk.tokens_slice], _ = self._bf16_flash_mla_kernel( chunk_q, @@ -907,7 +1065,10 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): chunk_topk_length, ) - return attn_out + if self.pcp_dcp_kv_gather and lse is None: + # No decode rows: the DCP merge still expects an LSE, of no rows. + lse = q.new_empty((0, q.shape[1]), dtype=torch.float32) + return attn_out, lse def _forward_fp8_kv_mixed_batch( self, @@ -915,6 +1076,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): kv_c_and_k_pe_cache: torch.Tensor, topk_indices: torch.Tensor, attn_metadata: FlashMLASparseMetadata, + kernel_metadata: "FlashMLASparseMetadata.FP8KernelMetadata | None" = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: """Mixed batch FP8 forward path that treats all tokens as one batch. @@ -924,12 +1086,8 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): The lse is only returned when DCP needs it, otherwise None. """ - assert attn_metadata.fp8_extra_metadata is not None - assert isinstance( - attn_metadata.fp8_extra_metadata, - FlashMLASparseMetadata.FP8KernelMetadata, - ) - fp8_metadata = attn_metadata.fp8_extra_metadata + fp8_metadata = kernel_metadata or attn_metadata.fp8_extra_metadata + assert isinstance(fp8_metadata, FlashMLASparseMetadata.FP8KernelMetadata) block_table = attn_metadata.block_table # req_id_per_token covers the whole batch; slice it to the MQA tokens @@ -1160,7 +1318,7 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]): q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) else: - attn_out = self._forward_fp8_kv_separate_prefill_decode( + attn_out, lse = self._forward_fp8_kv_separate_prefill_decode( q, kv_c_and_k_pe_cache, topk_indices, attn_metadata ) diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 7bc0b442c3..164a3f8916 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -3,6 +3,7 @@ from dataclasses import dataclass from typing import Any +import numpy as np import torch import vllm.envs as envs @@ -25,6 +26,7 @@ from vllm.utils.deep_gemm import ( native_next_n_supported, ) from vllm.utils.platform_utils import num_compute_units +from vllm.utils.torch_utils import PIN_MEMORY from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -33,6 +35,7 @@ from vllm.v1.attention.backend import ( MultipleOf, ) from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping +from vllm.v1.attention.backends.mla.sparse_utils import request_row_bounds from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, split_decodes_and_prefills, @@ -43,6 +46,7 @@ from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, ) +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu logger = init_logger(__name__) @@ -259,6 +263,87 @@ class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend): return [64 if current_platform.is_device_capability_family(90) else 128] +@dataclass(frozen=True) +class PCPGlobalChunkPlan: + """PCP packing for one indexer prefill chunk under PCP + DCP.""" + + # [num_reqs+1] cumsum of the padded per-request context (shard rows x W): + # the global layout. + row_start_cu: torch.Tensor + # [num_reqs+1] like row_start_cu, but only a region's FIRST row carries its + # extent. + global_cu: torch.Tensor + # [num_reqs+1] cumsum of shard rows: this rank's padded layout. + padded_local_cu: torch.Tensor + padded_local_total: int + total: int + # [total] for each global row, its position in the rank-major gathered buffer. + deinterleave_idx: torch.Tensor + + +def build_pcp_global_chunk_plan( + row_req_idx: np.ndarray, + row_shard_rows: np.ndarray, + dcp_world_size: int, + device: torch.device, +) -> PCPGlobalChunkPlan: + """Plan the PCP packing for one chunk from its requests' KV shard rows. + + The global layout is padded to ``shard rows x W`` per request. Positions + past a request's true extent lie beyond every token's causal bound, so + the kernel never reads them. + """ + shard_rows = np.ascontiguousarray(row_shard_rows, dtype=np.int64) + assert row_req_idx.shape == shard_rows.shape + num_rows = len(shard_rows) + + row_bounds = request_row_bounds(row_req_idx) + region_first_row = row_bounds[:-1] + region_of_row = np.repeat(np.arange(len(region_first_row)), np.diff(row_bounds)) + + region_padded = shard_rows[region_first_row] + assert np.all(region_padded > 0), ( + f"PCP+DCP prefill got an empty context: {region_padded.tolist()}" + ) + + # Pinned host buffers, filled in place so the copies below are async. + cu = torch.zeros((3, num_rows + 1), dtype=torch.int32, pin_memory=PIN_MEMORY) + row_start_cu, global_cu, padded_cu = cu + # Only a region's first row carries its extent. + first_row = torch.from_numpy(region_first_row) + row_padded = torch.zeros(num_rows, dtype=torch.int32) + row_padded[first_row] = torch.from_numpy(region_padded).int() + torch.cumsum(row_padded * dcp_world_size, 0, out=global_cu[1:]) + torch.cumsum(row_padded, 0, out=padded_cu[1:]) + row_start_cu[:num_rows] = global_cu[first_row][torch.from_numpy(region_of_row)] + row_start_cu[num_rows] = global_cu[num_rows] + total = int(global_cu[num_rows]) + padded_total = int(padded_cu[num_rows]) + + idx = torch.empty(total, dtype=torch.int64, pin_memory=PIN_MEMORY) + idx_np = idx.numpy() + region_start = global_cu[first_row].numpy() + region_padded_start = padded_cu[first_row].numpy() + for i in range(len(region_first_row)): + g = int(region_padded[i]) * dcp_world_size + t = np.arange(g, dtype=np.int64) + idx_np[region_start[i] : region_start[i] + g] = ( + (t % dcp_world_size) * padded_total + + region_padded_start[i] + + t // dcp_world_size + ) + + cu = async_copy_to_gpu(cu, device=device) + return PCPGlobalChunkPlan( + row_start_cu=cu[0], + global_cu=cu[1], + padded_local_cu=cu[2], + padded_local_total=padded_total, + total=total, + deinterleave_idx=async_copy_to_gpu(idx, device=device), + ) + + @dataclass class DeepseekV32IndexerPrefillChunkMetadata: block_table: torch.Tensor @@ -277,6 +362,8 @@ class DeepseekV32IndexerPrefillChunkMetadata: local_total_seq_lens: int = 0 max_local_total_seq_lens: int = 0 + pcp_deinterleave_idx: torch.Tensor | None = None + class BuildPrefillChunkMetadataKernel( VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"] @@ -960,6 +1047,48 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0 return self.global_decode_seq_lens_buffer[:num_decode_tokens] + def _split_pcp_dcp_prefill_chunks( + self, + row_req_idx: np.ndarray, + row_shard_rows: np.ndarray, + row_query_lens_cpu: torch.Tensor, + max_logits_bytes: int, + request_offset: int, + ) -> list[tuple[slice, slice]]: + """Chunk by request rather than by row, so a split prefill's two rows + charge their shared context once, then widen each chunk back to rows. + + ``row_shard_rows`` holds the whole request's largest DCP shard on + every row, so the plan is identical on every PCP rank. + """ + row_bounds = request_row_bounds(row_req_idx) + first_rows = row_bounds[:-1] + # Each request's context, padded to whole DCP shards. + seq_lens = row_shard_rows[first_rows] * self.dcp_world_size + # Every rank holds a full-size chunk of a split request (shorter ones + # are replicated), so rows x longest row is the same on every rank, + # whichever row holds the short tail. + row_query_lens = row_query_lens_cpu.numpy() + query_lens = np.diff(row_bounds) * np.maximum.reduceat( + row_query_lens, first_rows + ) + chunk_specs = self._split_indexer_prefill_chunks( + torch.from_numpy(seq_lens.astype(np.int32)), + torch.from_numpy(query_lens.astype(np.int32)), + self.max_prefill_buffer_size, + max_logits_bytes, + ) + return [ + ( + slice( + request_offset + int(row_bounds[request_slice.start]), + request_offset + int(row_bounds[request_slice.stop]), + ), + query_slice, + ) + for request_slice, query_slice in chunk_specs + ] + @staticmethod def _split_indexer_prefill_chunks( compressed_seq_lens_cpu: torch.Tensor, @@ -1094,16 +1223,52 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): # slice below). assert common_attn_metadata.seq_lens_cpu_upper_bound is not None seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound - chunk_specs = self._split_indexer_prefill_chunks( - compressed_seq_lens_cpu[num_decodes:], - prefill_query_lens_cpu, - self.max_prefill_buffer_size, - max_logits_bytes, - request_offset=num_decodes, - ) + req_idx = None + shard_rows = None + if self.use_pcp and self.dcp_world_size > 1: + # The gathered KV must be packed identically on every PCP rank: + # chunk by request from its DCP shard rows, which every rank + # holds. A dummy batch bypasses the PCP manager and has one + # row per request, so its own extent is the request's. + req_idx = common_attn_metadata.req_idx + if req_idx is None: + req_idx = np.arange(num_reqs) + shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound + if shard_rows_cpu is None: + shard_rows_cpu = get_dcp_local_seq_lens( + seq_lens_cpu, + self.dcp_world_size, + 0, + self.cp_kv_cache_interleave_size, + ) + shard_rows = shard_rows_cpu.numpy() + chunk_specs = self._split_pcp_dcp_prefill_chunks( + req_idx[num_decodes:], + shard_rows[num_decodes:], + prefill_query_lens_cpu, + max_logits_bytes, + request_offset=num_decodes, + ) + else: + chunk_specs = self._split_indexer_prefill_chunks( + compressed_seq_lens_cpu[num_decodes:], + prefill_query_lens_cpu, + self.max_prefill_buffer_size, + max_logits_bytes, + request_offset=num_decodes, + ) chunks = [] for req_slice, query_slice in chunk_specs: + pcp_plan = None + if req_idx is not None: + assert shard_rows is not None + pcp_plan = build_pcp_global_chunk_plan( + req_idx[req_slice], + shard_rows[req_slice], + self.dcp_world_size, + self.device, + ) metadata = build_prefill_chunk_metadata( req_slice.start, req_slice.stop, @@ -1119,6 +1284,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): dcp_rank=self.dcp_rank, dcp_world_size=self.dcp_world_size, cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, + pcp_plan=pcp_plan, ) # Skip when total_seq_lens is 0 (i.e., no compressed token). if metadata is not None: @@ -1353,8 +1519,12 @@ def build_prefill_chunk_metadata( dcp_rank: int = 0, dcp_world_size: int = 1, cp_kv_cache_interleave_size: int = 1, + pcp_plan: PCPGlobalChunkPlan | None = None, ) -> DeepseekV32IndexerPrefillChunkMetadata | None: - total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item() + if pcp_plan is not None: + total_seq_lens = pcp_plan.total + else: + total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item() if total_seq_lens == 0: return None @@ -1362,15 +1532,22 @@ def build_prefill_chunk_metadata( device = block_table.device token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device) - cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) - # Assigning to slice avoids cpu sync. - cu_seq_lens[:1] = 0 - torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:]) + if pcp_plan is not None: + cu_seq_lens = pcp_plan.global_cu + else: + cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) + # Assigning to slice avoids cpu sync. + cu_seq_lens[:1] = 0 + torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:]) local_cu_seq_lens = cu_seq_lens local_total_seq_lens = total_seq_lens max_local_total_seq_lens = total_seq_lens - if dcp_world_size > 1: + if pcp_plan is not None: + local_cu_seq_lens = pcp_plan.padded_local_cu + local_total_seq_lens = pcp_plan.padded_local_total + max_local_total_seq_lens = pcp_plan.padded_local_total + elif dcp_world_size > 1: # Per-rank local KV length under interleave-aware DCP sharding, shape # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the # sharding matches the rest of the DCP pipeline (decode/prefill). @@ -1394,35 +1571,47 @@ def build_prefill_chunk_metadata( (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item() ) if query_slice is not None: - qs_start = query_slice.start - qs_stop = query_slice.stop + qs_start = min(query_slice.start, total_query_len) + qs_stop = min(query_slice.stop, total_query_len) else: qs_start = 0 qs_stop = total_query_len output_query_len = qs_stop - qs_start + assert 0 <= output_query_len <= total_query_len - qs_start, ( + f"query slice [{qs_start}, {qs_stop}) does not fit this rank's " + f"{total_query_len} tokens; the metadata tail would stay uninitialized" + ) cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device) cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device) - # Under DCP the kernel writes this rank's local row bounds into - # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. + if pcp_plan is not None: + row_start_cu = pcp_plan.row_start_cu + kernel_dcp_rank, kernel_dcp_world = 0, 1 + else: + # Under DCP the kernel writes this rank's local row bounds into + # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. + row_start_cu = local_cu_seq_lens + kernel_dcp_rank, kernel_dcp_world = dcp_rank, dcp_world_size _BUILD_PREFILL_CHUNK_METADATA_KERNEL( query_start_loc, uncompressed_seq_lens[start_idx:end_idx], cu_seq_lens, - local_cu_seq_lens, + row_start_cu, token_to_seq, cu_seq_len_ks, cu_seq_len_ke, qs_start, qs_stop, - dcp_rank, - dcp_world_size, + kernel_dcp_rank, + kernel_dcp_world, cp_kv_cache_interleave_size, num_reqs=num_reqs, COMPRESS_RATIO=compress_ratio, ) + pcp_deinterleave_idx = pcp_plan.deinterleave_idx if pcp_plan is not None else None + token_start = query_start_loc_cpu[start_idx].item() if query_slice is not None: token_end = token_start + qs_stop @@ -1432,6 +1621,7 @@ def build_prefill_chunk_metadata( token_end = query_start_loc_cpu[end_idx].item() return DeepseekV32IndexerPrefillChunkMetadata( + pcp_deinterleave_idx=pcp_deinterleave_idx, cu_seqlen_ks=cu_seq_len_ks, cu_seqlen_ke=cu_seq_len_ke, cu_seq_lens=cu_seq_lens, diff --git a/vllm/v1/attention/backends/mla/sparse_utils.py b/vllm/v1/attention/backends/mla/sparse_utils.py index d9b45892a6..242c8f3b78 100644 --- a/vllm/v1/attention/backends/mla/sparse_utils.py +++ b/vllm/v1/attention/backends/mla/sparse_utils.py @@ -5,6 +5,7 @@ from dataclasses import dataclass from typing import Any +import numpy as np import torch from vllm.config import VllmConfig @@ -24,6 +25,21 @@ from vllm.utils.math_utils import cdiv from vllm.v1.worker.block_table import get_block_table_width +def request_row_bounds(req_idx: np.ndarray) -> np.ndarray: + """Bounds of the runs of adjacent rows that belong to one request: run + ``r`` is rows ``[bounds[r], bounds[r + 1])``. + + Under PCP a rank holds two adjacent chunk rows of a split prefill; the + sparse backends give such a run one KV region. + """ + assert req_idx.size > 0 + bounds = np.flatnonzero(np.r_[True, req_idx[1:] != req_idx[:-1], True]) + assert bounds.size - 1 == np.unique(req_idx).size, ( + "rows of one request must be adjacent" + ) + return bounds + + def flat_kv_row_view( kv_cache: torch.Tensor, # [num_blocks, block_size, head_dim] block_size: int, @@ -69,7 +85,7 @@ class ConvertReqIndexToGlobalIndexKernel( block_table_stride: int @staticmethod - @triton.jit(do_not_specialize=["max_num_blocks_per_req"]) + @triton.jit(do_not_specialize=["max_num_blocks_per_req", "workspace_rank_stride"]) def kernel( req_id_ptr, # int32 [num_tokens] block_table_ptr, # int32 [num_requests, max_num_blocks_per_req] @@ -78,6 +94,7 @@ class ConvertReqIndexToGlobalIndexKernel( valid_count_ptr, # int32 [num_tokens] - output valid count per row prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr + workspace_rank_stride, # rows per DCP rank of a gathered prefill workspace # shapes (compile-time where possible) max_num_blocks_per_req, BLOCK_SIZE: tl.constexpr, @@ -146,7 +163,7 @@ class ConvertReqIndexToGlobalIndexKernel( # Guard block_table access valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0) bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1 - is_invalid_tok |= ~valid_block | is_remote + is_invalid_tok |= ~valid_block | (is_remote & ~is_prefill) base = tl.load(bt_ptr, mask=valid_block & ~is_prefill & ~is_remote, other=0) out_val = base * BLOCK_STRIDE_ROWS + inblock_off @@ -155,7 +172,11 @@ class ConvertReqIndexToGlobalIndexKernel( workspace_start = tl.load( workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0 ) - prefill_out = workspace_start + tok + # Under DCP the prefill workspace is the all-gather of every rank's + # shard, rank-major. With DCP_SIZE == 1 this is workspace_start + tok. + prefill_out = ( + owning_rank * workspace_rank_stride + workspace_start + local_idx + ) out_val = tl.where(is_prefill, prefill_out, out_val) out_val = tl.where(is_invalid_tok, -1, out_val) @@ -238,6 +259,22 @@ class ConvertReqIndexToGlobalIndexKernel( block_size * dcp_size, ) max_num_blocks_per_req = get_block_table_width(max_num_blocks, block_size) + pcp_dcp_combos = [] + if ( + vllm_config.parallel_config.prefill_context_parallel_size > 1 + and dcp_size > 1 + ): + # PCP + DCP attends prefill rows over a DCP-gathered workspace. + pcp_dcp_combos.append( + dict( + HAS_PREFILL_WORKSPACE=True, + COUNT_VALID=True, + COMPACT_TO_FRONT=True, + DCP_SIZE=dcp_size, + DCP_RANK=dcp_rank, + DCP_INTERLEAVE=dcp_interleave, + ) + ) return self._trace_dispatch(self.dispatch)( zip_inputs( dict( @@ -280,6 +317,7 @@ class ConvertReqIndexToGlobalIndexKernel( DCP_RANK=dcp_rank, DCP_INTERLEAVE=dcp_interleave, ), + *pcp_dcp_combos, ), BLOCK_SIZE=block_size, BLOCK_STRIDE_ROWS=block_stride_rows, @@ -316,6 +354,7 @@ class ConvertReqIndexToGlobalIndexKernel( int32_ptr if compile_key.has_prefill_workspace else None ), max_num_blocks_per_req=1, + workspace_rank_stride=1, BLOCK_SIZE=compile_key.block_size, BLOCK_STRIDE_ROWS=compile_key.block_stride_rows, BLOCK_N=compile_key.block_n, @@ -339,6 +378,7 @@ class ConvertReqIndexToGlobalIndexKernel( prefill_workspace_request_ids: torch.Tensor | None, prefill_workspace_starts: torch.Tensor | None, *, + workspace_rank_stride: int, max_num_blocks_per_req: int, BLOCK_SIZE: int, BLOCK_STRIDE_ROWS: int, @@ -406,6 +446,10 @@ def triton_convert_req_index_to_global_index( HAS_PREFILL_WORKSPACE: bool = False, prefill_workspace_request_ids: torch.Tensor | None = None, prefill_workspace_starts: torch.Tensor | None = None, + prefill_workspace_rank_stride: int | None = None, + dcp_size: int = 1, + dcp_rank: int = 0, + cp_kv_cache_interleave_size: int = 1, return_valid_counts: bool = False, out: torch.Tensor | None = None, valid_counts_out: torch.Tensor | None = None, @@ -511,6 +555,7 @@ def triton_convert_req_index_to_global_index( prefill_workspace_request_ids, prefill_workspace_starts, # shapes / constexprs + workspace_rank_stride=prefill_workspace_rank_stride or 0, max_num_blocks_per_req=max_num_blocks_per_req, BLOCK_SIZE=BLOCK_SIZE, BLOCK_STRIDE_ROWS=( @@ -521,10 +566,9 @@ def triton_convert_req_index_to_global_index( HAS_PREFILL_WORKSPACE=HAS_PREFILL_WORKSPACE, COUNT_VALID=return_valid_counts, COMPACT_TO_FRONT=return_valid_counts, - # DCP disabled (no-op de-interleave) - DCP_SIZE=1, - DCP_RANK=0, - DCP_INTERLEAVE=1, + DCP_SIZE=dcp_size, + DCP_RANK=dcp_rank, + DCP_INTERLEAVE=cp_kv_cache_interleave_size, ) if return_valid_counts: @@ -618,6 +662,7 @@ def triton_filter_and_convert_dcp_index( # No prefill workspace on the DCP decode path. None, None, + workspace_rank_stride=0, max_num_blocks_per_req=max_num_blocks_per_req, BLOCK_SIZE=BLOCK_SIZE, BLOCK_STRIDE_ROWS=( diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 7e7a212b5f..a3f6541b08 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -400,6 +400,7 @@ def build_attn_metadata( kv_cache_config: KVCacheConfig, seq_lens_cpu_upper_bound: torch.Tensor | None = None, dcp_local_seq_lens: torch.Tensor | None = None, + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None, positions: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None, @@ -409,10 +410,15 @@ def build_attn_metadata( rswa_prefix_lens: torch.Tensor | None = None, ubatch_idx: int = 0, fast_prefill: FastPrefillBatchMetadata | None = None, + req_idx: np.ndarray | None = None, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: dcp_local_seq_lens = dcp_local_seq_lens[:num_reqs] + if dcp_local_seq_lens_cpu_upper_bound is not None: + dcp_local_seq_lens_cpu_upper_bound = dcp_local_seq_lens_cpu_upper_bound[ + :num_reqs + ] if seq_lens_cpu_upper_bound is not None: seq_lens_cpu_upper_bound = seq_lens_cpu_upper_bound[:num_reqs] @@ -457,10 +463,12 @@ def build_attn_metadata( slot_mapping=slot_mapping, causal=group_causal, dcp_local_seq_lens=dcp_local_seq_lens, + dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound, positions=positions, is_prefilling=group_is_prefilling, mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, + req_idx=req_idx, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 30b9b8089c..5cbd642d25 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -114,6 +114,9 @@ class InputBatch: # (cudagraph capture) batches, which run the KV-sharing layers in full. fast_prefill: "FastPrefillBatchMetadata | None" = None + # [num_reqs] set only under PCP+DCP (see CommonAttentionMetadata). + dcp_local_seq_lens_cpu_upper_bound: torch.Tensor | None = None + @classmethod def make_dummy( cls, diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 6baceac9bc..27a9591e4b 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -220,11 +220,13 @@ class DefaultModelState(ModelState): kv_cache_config=kv_cache_config, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, + dcp_local_seq_lens_cpu_upper_bound=input_batch.dcp_local_seq_lens_cpu_upper_bound, positions=input_batch.positions, is_prefilling=torch.from_numpy(input_batch.is_prefilling_np), mm_req_doc_ranges=req_doc_ranges, for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, + req_idx=input_batch.idx_mapping_np, ubatch_idx=ubatch_idx, fast_prefill=input_batch.fast_prefill, ) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 6b97254621..3a51532f57 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -6,10 +6,10 @@ from dataclasses import dataclass, replace import numpy as np import torch -from vllm.config import VllmConfig +from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group from vllm.logger import init_logger -from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens 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.input_batch import ( @@ -154,30 +154,48 @@ class PCPManager: "MRV2 PCP only supports DSpark or single-module MTP " "speculative decoding." ) - if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): + cudagraph_mode = vllm_config.compilation_config.cudagraph_mode + is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk") + if parallel_config.decode_context_parallel_size > 1 and not is_sparse_mla: + # Dense MLA prefill sizes its DCP KV gather from each rank's own + # chunk rows, so the ranks' collectives diverge (#53573). + raise NotImplementedError("MRV2 PCP + DCP supports sparse MLA models only.") + if ( + is_sparse_mla + and parallel_config.decode_context_parallel_size == 1 + and cudagraph_mode != CUDAGraphMode.NONE + ): + raise NotImplementedError( + "MRV2 sparse MLA PCP does not support CUDA graphs yet. " + "Set -cc.cudagraph_mode=NONE." + ) + if cudagraph_mode.has_full_cudagraphs(): raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + if ( + parallel_config.decode_context_parallel_size > 1 + and parallel_config.dcp_comm_backend != "ag_rs" + ): + raise NotImplementedError( + "MRV2 PCP + DCP requires dcp_comm_backend='ag_rs'; got " + f"'{parallel_config.dcp_comm_backend}'." + ) @staticmethod def _reorder_segments( segments: list[RankSegment], - num_computed_tokens: np.ndarray, is_prefilling: np.ndarray, - query_start_loc_np: np.ndarray, ) -> list[RankSegment]: - """Move pure prefills last to match the batch ordering expected by - attention backends like MLA and sparse MLA. - """ + """Order this rank's rows decodes-first, then prefills, canonically.""" - def is_pure_prefill(segment: RankSegment) -> bool: + def sort_key(segment: RankSegment) -> tuple[bool, int, int]: req_idx = segment.global_batch_req_idx - start_pos = ( - num_computed_tokens[req_idx] - + segment.global_batch_slice.start - - query_start_loc_np[req_idx] + return ( + bool(is_prefilling[req_idx]), + req_idx, + segment.global_batch_slice.start, ) - return is_prefilling[req_idx] and start_pos == 0 - segments.sort(key=is_pure_prefill) + segments.sort(key=sort_key) rank_offset = 0 for index, segment in enumerate(segments): segments[index] = replace( @@ -189,6 +207,21 @@ class PCPManager: rank_offset += segment.num_tokens return segments + def replicated_requests( + self, + num_scheduled_tokens: np.ndarray, + is_prefilling: np.ndarray, + ) -> np.ndarray: + """Per global request, whether every PCP rank gets the whole query.""" + num_chunks = 2 * self.pcp_world_size + query_lens = np.asarray(num_scheduled_tokens, dtype=np.int64) + replicated = ~np.asarray(is_prefilling, dtype=np.bool_) + if self.dcp_world_size > 1: + chunk_sizes = (query_lens + num_chunks - 1) // num_chunks + drops_a_chunk = (num_chunks - 1) * chunk_sizes >= query_lens + replicated |= drops_a_chunk + return replicated + def _iter_rank_chunks( self, rank: int, @@ -204,17 +237,20 @@ class PCPManager: rank 1: 1 6 rank 2: 2 5 rank 3: 3 4 + + Decodes, and prefills too short to fill all eight, are replicated instead. """ num_chunks = 2 * self.pcp_world_size + replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling) for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens): query_len = int(num_tokens) if query_len == 0: continue chunk_indices: tuple[int, ...] - if bool(is_prefilling[global_batch_req_idx]): + if not replicated[global_batch_req_idx]: chunk_size = (query_len + num_chunks - 1) // num_chunks chunk_indices = (rank, num_chunks - 1 - rank) - else: # decodes are replicated + else: # decodes, and short prefills under DCP, are replicated chunk_size = query_len chunk_indices = (0,) @@ -229,7 +265,6 @@ class PCPManager: self, rank: int, num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, is_prefilling: np.ndarray, query_start_loc_np: np.ndarray, ) -> list[RankSegment]: @@ -248,12 +283,7 @@ class PCPManager: ) ) rank_offset += chunk_len - return self._reorder_segments( - rank_segments, - num_computed_tokens, - is_prefilling, - query_start_loc_np, - ) + return self._reorder_segments(rank_segments, is_prefilling) def _build_batch_layout( self, @@ -263,13 +293,13 @@ class PCPManager: query_start_loc_np: np.ndarray, padded_num_tokens: int | None = None, ) -> tuple[list[list[RankSegment]], list[int]]: + replicated = self.replicated_requests(num_scheduled_tokens, is_prefilling) segments_by_rank = [] per_rank_num_tokens = [] for rank in range(self.pcp_world_size): segments = self._get_rank_segments( rank, num_scheduled_tokens, - num_computed_tokens, is_prefilling, query_start_loc_np, ) @@ -309,7 +339,7 @@ class PCPManager: dtype=np.int64, ) # Cache insertion pairs one slot entry with each rank's local decode. - if not bool(is_prefilling[segment.global_batch_req_idx]) and rank != 0: + if replicated[segment.global_batch_req_idx] and rank != 0: continue gathered_kv_write_mask[padded_gathered_slice] = True hidden_restore_idx[segment.global_batch_slice] = np.arange( @@ -527,6 +557,19 @@ 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_cpu_upper_bound = None + if self.dcp_world_size > 1: + # The largest DCP shard of each row's whole request, identical on + # every PCP rank: the sparse backends pad their KV gather to it. + request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[ + local_to_global_batch_req_idx_np + ] + dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( + torch.from_numpy(request_seq_lens.astype(np.int32)), + self.dcp_world_size, + 0, + self.cp_interleave, + ) self._local_batch = replace( input_batch, @@ -549,6 +592,7 @@ class PCPManager: seq_lens=seq_lens, seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np), dcp_local_seq_lens=None, + dcp_local_seq_lens_cpu_upper_bound=dcp_local_seq_lens_cpu_upper_bound, 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_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 4ff50f82b5..d6a27b1f8b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2486,7 +2486,7 @@ class GPUModelRunner( self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded) cm_base.dcp_local_seq_lens = self.dcp_local_seq_lens.gpu[:num_reqs_padded] - cm_base.dcp_local_seq_lens_cpu = self.dcp_local_seq_lens.cpu[ + cm_base.dcp_local_seq_lens_cpu_upper_bound = self.dcp_local_seq_lens.cpu[ :num_reqs_padded ]