From b6ff8a2f509cc7ac9c58176f5115a836aa1e08bd Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Sun, 19 Jul 2026 04:53:15 +0200 Subject: [PATCH] [Core] Add MRV2 virtual-batch PCP for MLA (#46570) Signed-off-by: Lucas Wilkinson Signed-off-by: Matthew Bonanni Co-authored-by: Codex Co-authored-by: OpenAI Codex Co-authored-by: Matthew Bonanni --- .buildkite/test_areas/lm_eval.yaml | 22 + .../configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml | 16 + .../configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml | 17 + tests/evals/gsm8k/configs/models-pcp.txt | 2 + tests/evals/gsm8k/gsm8k_eval.py | 16 +- tests/evals/gsm8k/test_gsm8k_correctness.py | 1 + .../v1/attention/test_attention_splitting.py | 16 + .../v1/attention/test_sparse_mla_backends.py | 7 + tests/v1/kv_offload/test_factory.py | 8 +- tests/v1/simple_kv_offload/test_scheduler.py | 40 +- vllm/config/parallel.py | 62 +- vllm/config/vllm.py | 5 +- .../device_communicators/all2all.py | 44 +- .../base_device_communicator.py | 11 +- .../kv_connector/v1/offloading/config.py | 9 +- vllm/distributed/parallel_state.py | 25 +- .../layers/attention/mla_attention.py | 72 +- vllm/model_executor/layers/attention/pcp.py | 92 +++ .../layers/attention/sparse_mla_attention.py | 4 +- .../model_executor/layers/fused_moe/config.py | 4 +- .../layers/fused_moe/runner/moe_runner.py | 28 +- .../layers/sparse_attn_indexer.py | 29 +- vllm/v1/attention/backend.py | 16 +- vllm/v1/attention/backends/mla/indexer.py | 23 +- vllm/v1/attention/backends/utils.py | 4 +- vllm/v1/attention/selector.py | 5 +- vllm/v1/core/kv_cache_coordinator.py | 2 - vllm/v1/core/kv_cache_utils.py | 11 +- vllm/v1/core/sched/scheduler.py | 2 +- vllm/v1/core/single_type_kv_cache_manager.py | 12 +- vllm/v1/kv_cache_interface.py | 16 +- vllm/v1/simple_kv_offload/manager.py | 7 +- vllm/v1/worker/block_table.py | 6 +- vllm/v1/worker/cp_utils.py | 25 +- vllm/v1/worker/gpu/attn_utils.py | 7 + vllm/v1/worker/gpu/block_table.py | 22 +- vllm/v1/worker/gpu/model_runner.py | 24 +- vllm/v1/worker/gpu/model_states/default.py | 8 +- vllm/v1/worker/gpu/pcp_manager.py | 680 ++++++++++++++++++ 39 files changed, 1204 insertions(+), 196 deletions(-) create mode 100644 tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml create mode 100644 tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml create mode 100644 tests/evals/gsm8k/configs/models-pcp.txt create mode 100644 vllm/model_executor/layers/attention/pcp.py create mode 100644 vllm/v1/worker/gpu/pcp_manager.py diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 9c08c96e4c..95a8bafdd5 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -78,6 +78,28 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt +- label: LM Eval PCP (4xB200) + key: lm-eval-pcp-4xb200 + timeout_in_minutes: 360 + device: b200-k8s + num_devices: 4 + optional: true + source_file_dependencies: + - 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/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/v1/worker/gpu/model_runner.py + - vllm/v1/worker/gpu/pcp_manager.py + autorun_on_main: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt + - label: LM Eval Large Models EP (2xB200) key: lm-eval-large-models-ep-2xb200 timeout_in_minutes: 60 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml new file mode 100644 index 0000000000..13d71af20e --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml @@ -0,0 +1,16 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +max_concurrency: 100 +server_args: >- + --enforce-eager + --max-model-len 4096 + --safetensors-load-strategy prefetch + --moe-backend flashinfer_cutlass + --prefill-context-parallel-size 4 + --enable-expert-parallel + --kv-cache-dtype fp8 +env: + VLLM_LOGGING_LEVEL: "DEBUG" + VLLM_USE_V2_MODEL_RUNNER: "1" diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml new file mode 100644 index 0000000000..b21d0b12f0 --- /dev/null +++ b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml @@ -0,0 +1,17 @@ +model_name: "nvidia/GLM-5.2-NVFP4" +accuracy_threshold: 0.90 +num_questions: 1319 +num_fewshot: 5 +max_concurrency: 100 +server_args: >- + --enforce-eager + --max-model-len 4096 + --safetensors-load-strategy prefetch + --moe-backend flashinfer_cutlass + --tensor-parallel-size 2 + --prefill-context-parallel-size 2 + --enable-expert-parallel + --kv-cache-dtype fp8 +env: + 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 new file mode 100644 index 0000000000..df645edad5 --- /dev/null +++ b/tests/evals/gsm8k/configs/models-pcp.txt @@ -0,0 +1,2 @@ +GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +GLM-5.2-NVFP4-TP1-PCP4-EP.yaml diff --git a/tests/evals/gsm8k/gsm8k_eval.py b/tests/evals/gsm8k/gsm8k_eval.py index 9c47826850..a8ee6833ca 100644 --- a/tests/evals/gsm8k/gsm8k_eval.py +++ b/tests/evals/gsm8k/gsm8k_eval.py @@ -217,6 +217,7 @@ def evaluate_gsm8k( seed: int | None = 42, request_timeout_seconds: float = 600, gen_prefix: str = "", + max_concurrency: int | None = None, ) -> dict[str, float | int]: """ Evaluate GSM8K accuracy using vLLM serve endpoint. @@ -261,7 +262,14 @@ def evaluate_gsm8k( return answer, tokens timeout = aiohttp.ClientTimeout(total=request_timeout_seconds) - async with aiohttp.ClientSession(timeout=timeout) as session: + connector = ( + aiohttp.TCPConnector(limit=max_concurrency) + if max_concurrency is not None + else None + ) + async with aiohttp.ClientSession( + timeout=timeout, connector=connector + ) as session: tasks = [get_answer(session, i) for i in range(num_questions)] await tqdm.gather(*tasks, desc="Evaluating") @@ -343,6 +351,11 @@ def main() -> None: parser.add_argument( "--seed", type=int, default=42, help="Random seed for reproducibility" ) + parser.add_argument( + "--max-concurrency", + type=int, + help="Maximum number of concurrent requests", + ) parser.add_argument("--save-results", type=str, help="Save results to JSON file") args = parser.parse_args() @@ -355,6 +368,7 @@ def main() -> None: port=args.port, temperature=args.temperature, seed=args.seed, + max_concurrency=args.max_concurrency, ) # Print results to terminal diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index 4b2b6ad381..e1ae6ec542 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -59,6 +59,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict: seed=eval_config.get("seed", 42), request_timeout_seconds=request_timeout_seconds, gen_prefix=eval_config.get("gen_prefix", ""), + max_concurrency=eval_config.get("max_concurrency"), ) return results diff --git a/tests/v1/attention/test_attention_splitting.py b/tests/v1/attention/test_attention_splitting.py index d6fc59576b..acbc22a027 100644 --- a/tests/v1/attention/test_attention_splitting.py +++ b/tests/v1/attention/test_attention_splitting.py @@ -160,6 +160,8 @@ def apply_split_decodes_and_prefills( decode_threshold: int, require_uniform: bool, padded_num_tokens: int | None = None, + is_prefilling: list[bool] | None = None, + treat_short_extends_as_decodes: bool = True, ): """Helper function to apply split_decodes_and_prefills and return the results.""" @@ -173,11 +175,14 @@ def apply_split_decodes_and_prefills( if padded_num_tokens is not None: common_metadata.num_actual_tokens = padded_num_tokens + if is_prefilling is not None: + common_metadata.is_prefilling = torch.tensor(is_prefilling) return split_decodes_and_prefills( common_metadata, decode_threshold=decode_threshold, require_uniform=require_uniform, + treat_short_extends_as_decodes=treat_short_extends_as_decodes, ) @@ -236,6 +241,17 @@ def test_split_decodes_and_prefills_uniform_all_ones(): assert num_prefill_tokens == 0 +def test_split_decodes_and_prefills_uniform_short_extend(): + result = apply_split_decodes_and_prefills( + [1, 1], + decode_threshold=1, + require_uniform=True, + is_prefilling=[False, True], + treat_short_extends_as_decodes=False, + ) + assert result == (1, 1, 1, 1) + + def test_split_decodes_and_prefills_uniform_all_short_decodes(): query_lens = [2, 2, 1, 3, 2, 1, 2] num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 6ec9fc25ed..3fe5df918a 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1114,6 +1114,13 @@ def test_sparse_backend_prefill_correctness( @pytest.mark.parametrize( "seq_lens,query_lens,workspace_size,max_logits_bytes,expected", [ + ( + torch.tensor([0]), + torch.tensor([0]), + 100, + 1000, + [], + ), # Logits constraint triggers split (M*N exceeds budget) # req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000 # req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250 diff --git a/tests/v1/kv_offload/test_factory.py b/tests/v1/kv_offload/test_factory.py index 570924cfcc..d08dbee576 100644 --- a/tests/v1/kv_offload/test_factory.py +++ b/tests/v1/kv_offload/test_factory.py @@ -385,7 +385,7 @@ def test_tiering_spec_aligns_row_size(): assert spec.num_blocks == cpu_bytes_to_use // alignment -def test_offloading_spec_resolves_prefill_context_parallel_block_sizes(): +def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel(): config = _make_layout_vllm_config( cpu_bytes_to_use=65536, extra_config={"block_size": 64}, @@ -394,9 +394,9 @@ def test_offloading_spec_resolves_prefill_context_parallel_block_sizes(): spec = _create_spec(config, _make_kv_cache_config()) - assert spec.tokens_per_block == (32,) - assert spec.tokens_per_hash == 32 - assert spec.blocks_per_chunk == 2 + assert spec.tokens_per_block == (16,) + assert spec.tokens_per_hash == 16 + assert spec.blocks_per_chunk == 4 def test_offloading_config_preserves_data_parallel_index(): diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 09586f5e6b..4f7761cf35 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -1551,9 +1551,8 @@ def _make_cp_scheduler( num_gpu_blocks: int = 16, lazy: bool = False, ) -> SchedulerFixture: - """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size.""" - cp_world_size = dcp_world_size * pcp_world_size - virtual_block_size = BLOCK_SIZE * cp_world_size + """Build a SimpleCPUOffloadScheduler with DCP-scaled block size.""" + virtual_block_size = BLOCK_SIZE * dcp_world_size kv_cache_config = _make_kv_cache_config(num_gpu_blocks) vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size) @@ -1645,17 +1644,14 @@ def _allocate_cp_gpu_blocks( ], ) def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None: - """Verify that the scheduler's block_size and cp_world_size are correctly - scaled when context parallelism is enabled.""" + """Verify block size scaling follows DCP ownership.""" fix = _make_cp_scheduler( dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size ) sched = fix.scheduler - expected_cp = dcp_world_size * pcp_world_size - assert sched.cp_world_size == expected_cp - assert sched.block_size == BLOCK_SIZE * expected_cp - assert sched.fa_block_size == BLOCK_SIZE * expected_cp + assert sched.cp_world_size == dcp_world_size + assert sched.block_size == BLOCK_SIZE * dcp_world_size # --------------------------------------------------------------------------- @@ -1682,8 +1678,7 @@ def test_cp_eager_store_and_load_roundtrip( lazy=False, ) sched = fix.scheduler - cp = dcp_world_size * pcp_world_size - vbs = BLOCK_SIZE * cp + vbs = BLOCK_SIZE * dcp_world_size num_blocks = 2 req = _make_cp_request(num_blocks, vbs) @@ -1749,18 +1744,16 @@ def test_cp_eager_store_and_load_roundtrip( def test_cp_effective_block_size_store_and_load( dcp_world_size: int, pcp_world_size: int ) -> None: - """Verify ready_blocks_g (store) and n_take_g (load) use the effective - (physical * cp) block size, not the per-rank physical size.""" + """Verify store/load physical block counts scale with DCP, not PCP.""" fix = _make_cp_scheduler( dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size ) sched = fix.scheduler gpu_pool = fix.gpu_block_pool - cp = dcp_world_size * pcp_world_size - vbs = BLOCK_SIZE * cp + vbs = BLOCK_SIZE * dcp_world_size + expected_blocks = 1 - # Store: allocate 2 blocks, confirm only 1. Without the fix, - # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks. + # Store one DCP-scaled logical block worth of tokens. req = _make_cp_request(num_blocks=2, virtual_block_size=vbs) gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs) kv = KVCacheBlocks(blocks=(gpu_blocks,)) @@ -1772,12 +1765,11 @@ def test_cp_effective_block_size_store_and_load( new_reqs={req.request_id: kv.get_block_ids()}, ) ) - assert len(m1.store_gpu_blocks) == 1 - assert len(m1.store_cpu_blocks) == 1 + assert len(m1.store_gpu_blocks) == expected_blocks + assert len(m1.store_cpu_blocks) == expected_blocks simulate_store_completion(sched, m1.store_event) - # Load: store 2 blocks from a second request, accept only 1 as external. - # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both. + # Load one DCP-scaled logical block from a two-block CPU hit. req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs) kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),)) req2.num_computed_tokens = 2 * vbs @@ -1801,7 +1793,7 @@ def test_cp_effective_block_size_store_and_load( hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0) assert hit == 2 * vbs - kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),)) + kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(expected_blocks),)) sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs) m3 = sched.build_connector_meta( make_scheduler_output( @@ -1810,8 +1802,8 @@ def test_cp_effective_block_size_store_and_load( ) ) assert m3.load_event >= 0 - assert len(m3.load_gpu_blocks) == 1 - assert len(m3.load_cpu_blocks) == 1 + assert len(m3.load_gpu_blocks) == expected_blocks + assert len(m3.load_cpu_blocks) == expected_blocks assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]] diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 7c270b0c0e..53688c05d9 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -122,10 +122,11 @@ class ParallelConfig: tensor_parallel_size: int = Field(default=1, ge=1) """Number of tensor parallel groups.""" prefill_context_parallel_size: int = Field(default=1, ge=1) - """Number of prefill context parallel groups.""" + """Number of ranks that split prefill sequence computation. PCP expands + the process world size but does not increase the KV-cache shard count.""" data_parallel_size: int = Field(default=1, ge=1) """Number of data parallel groups. MoE layers will be sharded according to - the product of the tensor parallel size and data parallel size.""" + the product of the tensor, prefill-context, and data parallel sizes.""" data_parallel_size_local: int = Field(default=1, ge=0) """Number of local data parallel groups. A value of 0 is a sentinel used by the engine-args layer to signal that data parallelism was specified @@ -337,9 +338,9 @@ class ParallelConfig: connect as clients to exchange self-picked group ports at runtime.""" decode_context_parallel_size: int = Field(default=1, ge=1) - """Number of decode context parallel groups, because the world size does - not change by dcp, it simply reuse the GPUs of TP group, and tp_size - needs to be divisible by dcp_size.""" + """Number of ranks that shard the decode KV cache. DCP does not expand + the process world size. Without PCP, DCP reuses TP ranks. With PCP, DCP + either spans the PCP axis or the full TP x PCP block.""" dcp_kv_cache_interleave_size: int = 1 """ @@ -357,13 +358,11 @@ class ParallelConfig: """ cp_kv_cache_interleave_size: int = 1 - """Interleave size of kv_cache storage while using DCP or PCP. - For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`, - and `total_cp_world_size = pcp_world_size * dcp_world_size`. - store interleave_size tokens on total_cp_rank i, - then store next interleave_size tokens on total_cp_rank i+1. + """Interleave size of kv_cache storage while using DCP. + Store interleave_size tokens on dcp_rank i, then store next + interleave_size tokens on dcp_rank i+1. Interleave_size=1: token-level alignment, where token `i` is stored on - total_cp_rank `i % total_cp_world_size`. + dcp_rank `i % dcp_world_size`. Interleave_size=block_size: block-level alignment, where tokens are first populated to the preceding ranks. Tokens are then stored in (rank i+1, block j) only after (rank i, block j) is fully occupied. @@ -480,11 +479,19 @@ class ParallelConfig: ) if not self.enable_expert_parallel: raise ValueError("enable_expert_parallel must be True to use EPLB.") - if self.tensor_parallel_size * self.data_parallel_size <= 1: + # The EP group spans the TP x PCP x DP ranks. EPLB therefore needs + # TP, PCP, or DP > 1. + if ( + self.tensor_parallel_size + * self.prefill_context_parallel_size + * self.data_parallel_size + <= 1 + ): raise ValueError( - "EPLB requires tensor_parallel_size or data_parallel_size " - f"to be greater than 1, but got " - f"TP={self.tensor_parallel_size},DP={self.data_parallel_size}." + "EPLB requires tensor, prefill-context, or data parallelism, " + f"but got TP={self.tensor_parallel_size}, " + f"PCP={self.prefill_context_parallel_size}, " + f"DP={self.data_parallel_size}." ) else: if self.eplb_config.num_redundant_experts != 0: @@ -495,15 +502,21 @@ class ParallelConfig: "num_redundant_experts." ) - # Note(hc): In the current implementation of decode context - # parallel(DCP), tp_size needs to be divisible by dcp_size, - # because the world size does not change by dcp, it simply - # reuses the GPUs of TP group, and split one TP group into - # tp_size//dcp_size DCP groups. - if self.tensor_parallel_size % self.decode_context_parallel_size != 0: + tp = self.tensor_parallel_size + pcp = self.prefill_context_parallel_size + dcp = self.decode_context_parallel_size + if pcp > 1 and self.data_parallel_size > 1: + raise ValueError("PCP does not support data parallelism yet.") + if pcp == 1: + # DCP reuses the TP ranks when PCP is disabled. + if tp % dcp != 0: + raise ValueError(f"tp_size={tp} must be divisible by dcp_size={dcp}.") + elif dcp not in (1, pcp, tp * pcp): raise ValueError( - f"tp_size={self.tensor_parallel_size} must be divisible by" - f"dcp_size={self.decode_context_parallel_size}." + "When PCP is enabled, DCP must be disabled, span the PCP " + "axis, or span the full TP x PCP axis. " + f"Got TP={tp}, PCP={pcp}, DCP={dcp}; valid DCP sizes are " + f"{sorted({1, pcp, tp * pcp})}." ) if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1: @@ -515,8 +528,7 @@ class ParallelConfig: @property def world_size_across_dp(self) -> int: - """world_size_across_dp is TPxPPxDP, it is the size of the world - including data parallelism.""" + """Process world size across TP, PCP, PP, and DP.""" return self.world_size * self.data_parallel_size @property diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index cca9adbcf5..4b4a97f41b 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2130,9 +2130,10 @@ class VllmConfig: model_config = self.model_config speculative_config = self.speculative_config - if self.parallel_config.prefill_context_parallel_size > 1: + if self.parallel_config.prefill_context_parallel_size > 1 and not ( + model_config is not None and model_config.use_mla + ): unsupported.append("prefill context parallelism") - if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE: unsupported.append("stock torch.compile") diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 7f540bc4b1..33841de306 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -8,7 +8,7 @@ import torch import torch.distributed as dist import vllm.envs as envs -from vllm.distributed import get_dp_group, get_ep_group +from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group from vllm.distributed.utils import StatelessProcessGroup from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -49,6 +49,23 @@ class AgRsAll2AllManager(All2AllManagerBase): def __init__(self, cpu_group, tcp_store_group=None): super().__init__(cpu_group, tcp_store_group) + def _get_comm_group(self, is_sequence_parallel: bool) -> Any: + if is_sequence_parallel: + return get_ep_group() + if self.dp_world_size > 1: + return get_dp_group() + return get_pcp_group() + + def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]: + if self.dp_world_size == 1: + return [num_local_tokens] * comm_group.world_size + + dp_metadata = get_forward_context().dp_metadata + assert dp_metadata is not None + sizes = dp_metadata.get_chunk_sizes_across_dp_rank() + assert sizes is not None + return sizes + def dispatch_router_logits( self, hidden_states: torch.Tensor, @@ -62,11 +79,8 @@ class AgRsAll2AllManager(All2AllManagerBase): """ Gather hidden_states and router_logits from all dp ranks. """ - dp_metadata = get_forward_context().dp_metadata - assert dp_metadata is not None - sizes = dp_metadata.get_chunk_sizes_across_dp_rank() - assert sizes is not None - dist_group = get_ep_group() if is_sequence_parallel else get_dp_group() + dist_group = self._get_comm_group(is_sequence_parallel) + sizes = self._get_sizes(hidden_states.shape[0], dist_group) assert sizes[dist_group.rank_in_group] == hidden_states.shape[0] tensors_to_gather = [hidden_states, router_logits] @@ -97,11 +111,8 @@ class AgRsAll2AllManager(All2AllManagerBase): """ Gather hidden_states and router_logits from all dp ranks. """ - dp_metadata = get_forward_context().dp_metadata - assert dp_metadata is not None - sizes = dp_metadata.get_chunk_sizes_across_dp_rank() - assert sizes is not None - dist_group = get_ep_group() if is_sequence_parallel else get_dp_group() + dist_group = self._get_comm_group(is_sequence_parallel) + sizes = self._get_sizes(hidden_states.shape[0], dist_group) assert sizes[dist_group.rank_in_group] == hidden_states.shape[0] tensors_to_gather = [hidden_states, topk_weights, topk_ids] @@ -129,12 +140,11 @@ class AgRsAll2AllManager(All2AllManagerBase): """ Reduce-scatter hidden_states across all dp ranks. """ - dp_metadata = get_forward_context().dp_metadata - assert dp_metadata is not None - sizes = dp_metadata.get_chunk_sizes_across_dp_rank() - assert sizes is not None - - dist_group = get_ep_group() if is_sequence_parallel else get_dp_group() + dist_group = self._get_comm_group(is_sequence_parallel) + sizes = self._get_sizes( + hidden_states.shape[0] // dist_group.world_size, + dist_group, + ) hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes) return hidden_states diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index 9a443b7fc1..70f1fb5d62 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -176,11 +176,16 @@ class DeviceCommunicatorBase: config = get_current_vllm_config_or_none() if config is not None: # initialize the all2all manager for DP or sequence-parallel EP. + parallel_config = config.parallel_config use_ep = ( - config.parallel_config.data_parallel_size > 1 - or config.parallel_config.use_sequence_parallel_moe + parallel_config.data_parallel_size > 1 + or parallel_config.use_sequence_parallel_moe + or ( + parallel_config.enable_expert_parallel + and parallel_config.prefill_context_parallel_size > 1 + ) ) - all2all_backend = config.parallel_config.all2all_backend + all2all_backend = parallel_config.all2all_backend self.is_ep_communicator = unique_name.split(":")[0] == "ep" self.use_all2all = self.is_ep_communicator and use_ep diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index 5cf4dffd1a..b807b480f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -36,13 +36,12 @@ def build_offloading_config( engine_id = kv_transfer_config.engine_id parallel_config = vllm_config.parallel_config - context_parallel_factor = ( - parallel_config.decode_context_parallel_size - * parallel_config.prefill_context_parallel_size - ) groups = tuple( OffloadingGroupConfig( - tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor), + tokens_per_block=( + group.kv_cache_spec.block_size + * parallel_config.decode_context_parallel_size + ), layer_names=tuple(group.layer_names), ) for group in kv_cache_config.kv_cache_groups diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index e46ca1691c..f0647323e6 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -1784,7 +1784,7 @@ def initialize_model_parallel( get_world_group().device_group ) - # the layout order is: ExternalDP x DP x PP x TP + # the layout order is: ExternalDP x DP x PP x PCP x TP # ExternalDP is the data parallel group that is not part of the model, # every dp rank can generate independently (in verl integration). # DP is the data parallel group that is part of the model, @@ -1821,17 +1821,13 @@ def initialize_model_parallel( # Build the DCP model-parallel groups. global _DCP assert _DCP is None, "decode context model parallel group is already initialized" - # Note(hc): In the current implementation of decode context parallel, - # dcp_size must not exceed tp_size, because the world size does not - # change by DCP, it simply reuses the GPUs of TP group, and split one - # TP group into tp_size//dcp_size DCP groups. - group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0) + dcp_size = decode_context_model_parallel_size or 1 + dcp_ranks = local_all_ranks if enable_elastic_ep else all_ranks + if dcp_size > 1: + # DCP spans PCP first, then TP for full TP x PCP groups. + dcp_ranks = dcp_ranks.transpose(-1, -2) + group_ranks = dcp_ranks.reshape(-1, dcp_size).unbind(0) group_ranks = [x.tolist() for x in group_ranks] - if enable_elastic_ep: - group_ranks = local_all_ranks.reshape( - -1, decode_context_model_parallel_size - ).unbind(0) - group_ranks = [x.tolist() for x in group_ranks] _DCP = init_model_parallel_group( group_ranks, get_world_group().local_rank, @@ -2005,6 +2001,13 @@ def ensure_model_parallel_initialized( f"{pcp_world_size=} vs. " f"{prefill_context_model_parallel_size=}" ) + dcp_world_size = get_dcp_group().world_size + dcp_model_parallel_size = decode_context_model_parallel_size or 1 + assert dcp_world_size == dcp_model_parallel_size, ( + "decode context parallel group already initialized, but of unexpected size: " + f"{dcp_world_size=} vs. " + f"{dcp_model_parallel_size=}" + ) def prepare_communication_buffer_for_model(model: torch.nn.Module): diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 55ad7aedef..a450ada185 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -211,6 +211,7 @@ from vllm.config import ( from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import ( get_dcp_group, + get_tp_group, is_global_first_rank, ) from vllm.forward_context import ForwardContext, get_forward_context @@ -225,6 +226,10 @@ from vllm.model_executor.layers.attention.attention import ( from vllm.model_executor.layers.attention.kv_transfer_utils import ( maybe_transfer_kv_layer, ) +from vllm.model_executor.layers.attention.pcp import ( + finalize_mla_pcp_decode, + maybe_gather_mla_latent_cache_inputs, +) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -269,7 +274,7 @@ from vllm.v1.attention.backends.utils import ( get_dcp_local_seq_lens, split_decodes_and_prefills, ) -from vllm.v1.attention.ops.common import cp_lse_ag_out_rs +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.attention.selector import get_attn_backend @@ -499,6 +504,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): self.use_direct_call = not current_platform.opaque_attention_op() vllm_config = get_current_vllm_config() + parallel_config = vllm_config.parallel_config + self.use_pcp = parallel_config.prefill_context_parallel_size > 1 compilation_config = vllm_config.compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") @@ -611,11 +618,23 @@ class MLAAttention(nn.Module, AttentionLayerBase): assert isinstance(slot_mapping, dict), ( f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. " ) + layer_slot_mapping = slot_mapping.get(self.layer_name) + kv_for_cache, kpe_for_cache, layer_slot_mapping = ( + maybe_gather_mla_latent_cache_inputs( + kv_c_normed, + k_pe, + layer_slot_mapping, + attn_metadata.num_decode_tokens + if attn_metadata is not None + else None, + self.use_pcp, + ) + ) self.impl.do_kv_cache_update( # type: ignore[attr-defined] - kv_c_normed, - k_pe, + kv_for_cache, + kpe_for_cache, self_kv_cache, - slot_mapping.get(self.layer_name), + layer_slot_mapping, self.kv_cache_dtype, self._k_scale, ) @@ -710,6 +729,10 @@ class MLAAttention(nn.Module, AttentionLayerBase): fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) num_actual_toks = attn_metadata.num_actual_tokens + if self.use_pcp and self.impl.dcp_world_size > 1 and quant_key is not None: + raise NotImplementedError( + "MRV2 MLA PCP+DCP does not support fused output quantization yet." + ) # Inputs and outputs may be padded for CUDA graphs output_padded = output @@ -833,11 +856,15 @@ class MLAAttention(nn.Module, AttentionLayerBase): else: mqa_q = (mqa_ql_nope, mqa_q_pe) if self.impl.dcp_world_size > 1: - if isinstance(mqa_q, tuple): - # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) - mqa_q = torch.cat(mqa_q, dim=-1) - # mqa_q do allgather in head dim. - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) + if self.use_pcp: + if self.impl.dcp_world_size > self.impl.pcp_world_size: + if isinstance(mqa_q, tuple): + mqa_q = torch.cat(mqa_q, dim=-1) + mqa_q = get_tp_group().all_gather(mqa_q, dim=1) + else: + if isinstance(mqa_q, tuple): + mqa_q = torch.cat(mqa_q, dim=-1) + mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) # call decode attn if not self.impl.is_sparse: @@ -846,6 +873,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): # correct dcp attn_out with lse. if self.impl.dcp_world_size > 1: + assert lse is not None if self.dcp_a2a: attn_out = dcp_a2a_lse_reduce( attn_out, @@ -853,6 +881,13 @@ class MLAAttention(nn.Module, AttentionLayerBase): get_dcp_group(), is_lse_base_on_e=self.impl.lse_base_on_e, ) + elif self.use_pcp: + attn_out = cp_lse_ag_out_ar( + attn_out, + lse, + get_dcp_group(), + is_lse_base_on_e=self.impl.lse_base_on_e, + ) else: attn_out = cp_lse_ag_out_rs( attn_out, @@ -860,6 +895,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): get_dcp_group(), is_lse_base_on_e=self.impl.lse_base_on_e, ) + if self.use_pcp: + attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) # v_up projection self._v_up_proj(attn_out, out=mqa_output_slice) @@ -903,6 +940,8 @@ class MLAAttention(nn.Module, AttentionLayerBase): raise ValueError(f"Unsupported quant_key: {quant_key}") return quant_output + if self.use_pcp and output_padded.shape[0] > num_actual_toks: + output_padded[num_actual_toks:].zero_() return output_padded def process_weights_after_loading(self, act_dtype: torch.dtype): @@ -1082,8 +1121,17 @@ def unified_mla_kv_cache_update( the data dependency between them to ensure torch.compile preserves ordering. """ layer_name = _resolve_layer_name(layer_name) - _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name) + attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context( + layer_name + ) if layer_slot_mapping is not None: + kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs( + kv_c_normed, + k_pe, + layer_slot_mapping, + attn_metadata.num_decode_tokens if attn_metadata is not None else None, + attn_layer.use_pcp, + ) attn_layer.impl.do_kv_cache_update( # type: ignore[attr-defined] kv_c_normed, k_pe, @@ -1723,6 +1771,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): self.compilation_config = vllm_config.compilation_config self.vllm_config = vllm_config self.device = device + self.use_pcp = parallel_config.prefill_context_parallel_size > 1 self.num_heads = self.model_config.get_num_attention_heads(parallel_config) self.mla_dims = get_mla_dims(self.model_config) @@ -1735,11 +1784,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): try: self.dcp_world_size = get_dcp_group().world_size - self.dcp_rank = get_dcp_group().rank_in_group except AssertionError: # DCP might not be initialized in testing self.dcp_world_size = 1 - self.dcp_rank = 0 self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size @@ -1857,6 +1904,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): common_attn_metadata, decode_threshold=self.reorder_batch_threshold, require_uniform=(self.query_len_support != QueryLenSupport.VARLEN), + treat_short_extends_as_decodes=not self.use_pcp, ) ) diff --git a/vllm/model_executor/layers/attention/pcp.py b/vllm/model_executor/layers/attention/pcp.py new file mode 100644 index 0000000000..75ab1c9e8e --- /dev/null +++ b/vllm/model_executor/layers/attention/pcp.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + +from vllm.distributed.parallel_state import ( + get_pcp_group, + get_tp_group, +) + + +def _gather_prefill_cache_inputs( + tensors: tuple[torch.Tensor, ...], + slot_mapping: torch.Tensor, + num_decode_tokens: int, +) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + """Keep replicated decode writes local and gather partitioned prefills.""" + local_num_tokens = tensors[0].shape[0] + assert all(tensor.shape[0] == local_num_tokens for tensor in tensors) + assert 0 <= num_decode_tokens <= local_num_tokens + + if num_decode_tokens == local_num_tokens: + return tensors, slot_mapping[:num_decode_tokens] + + pcp_group = get_pcp_group() + gathered_prefills = tuple( + pcp_group.all_gather(tensor[num_decode_tokens:].contiguous(), dim=0) + for tensor in tensors + ) + pcp_size = pcp_group.world_size + gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens] + if num_decode_tokens == 0: + return gathered_prefills, gathered_slot_mapping + + cache_inputs = tuple( + torch.cat((tensor[:num_decode_tokens], gathered_prefill), dim=0) + for tensor, gathered_prefill in zip(tensors, gathered_prefills) + ) + rank_slot_mappings = gathered_slot_mapping.view(pcp_size, local_num_tokens) + cache_slot_mapping = torch.cat( + ( + rank_slot_mappings[0, :num_decode_tokens], + rank_slot_mappings[:, num_decode_tokens:].flatten(), + ) + ) + return cache_inputs, cache_slot_mapping + + +def maybe_gather_mla_latent_cache_inputs( + kv_c_normed: torch.Tensor, + k_pe: torch.Tensor, + slot_mapping: torch.Tensor | None, + num_decode_tokens: int | None, + use_pcp: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + if not use_pcp or num_decode_tokens is None: + return kv_c_normed, k_pe, slot_mapping + assert slot_mapping is not None + num_tokens = kv_c_normed.shape[0] + k_pe_flat = k_pe.reshape(num_tokens, -1) + (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs( + (kv_c_normed, k_pe_flat), + slot_mapping, + num_decode_tokens, + ) + cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:]) + return cache_kv_c, cache_k_pe, cache_slot_mapping + + +def maybe_gather_indexer_k( + k: torch.Tensor, + slot_mapping: torch.Tensor, + num_decode_tokens: int, + use_pcp: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + if not use_pcp: + return k, slot_mapping + (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs( + (k,), slot_mapping, num_decode_tokens + ) + return cache_k, cache_slot_mapping + + +def finalize_mla_pcp_decode( + output: torch.Tensor, + num_heads: int, +) -> torch.Tensor: + if output.shape[1] < num_heads: + output = get_pcp_group().all_gather(output, dim=1) + elif output.shape[1] > num_heads: + head_start = get_tp_group().rank_in_group * num_heads + output = output[:, head_start : head_start + num_heads] + return output diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 35bbd4d3d4..19cad7986b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -59,13 +59,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): device=device, ) parallel_config = vllm_config.parallel_config + self.use_pcp = parallel_config.prefill_context_parallel_size > 1 try: self.dcp_world_size = get_dcp_group().world_size - self.dcp_rank = get_dcp_group().rank_in_group except AssertionError: # DCP might not be initialized in testing self.dcp_world_size = 1 - self.dcp_rank = 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size self.dcp_local_block_size = self.cp_kv_cache_interleave_size self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size @@ -174,6 +173,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]): num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold or 1, + treat_short_extends_as_decodes=not self.use_pcp, require_uniform=self.require_uniform_decodes, ) ( diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 39b14fc249..02e6c51116 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1050,7 +1050,9 @@ class FusedMoEParallelConfig: @property def use_all2all_kernels(self): - return self.use_ep and (self.dp_size > 1 or self.is_sequence_parallel) + return self.use_ep and ( + self.dp_size > 1 or self.pcp_size > 1 or self.is_sequence_parallel + ) @property def use_deepep_ht_kernels(self): diff --git a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py index cf00fa452a..7942957f32 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -752,18 +752,12 @@ class MoERunner(MoERunnerInterface): assert len(result) == 2 hidden_states, router_logits = result - # NOTE: Similar with DP, PCP also needs dispatch and combine. For - # simplicity, AgRsAll2All was added separately for PCP here. Maybe - # we should modify All2AllManager abstraction to better support PCP. - if self.moe_config.pcp_size > 1: - hidden_states = get_pcp_group().all_gather( - hidden_states, - dim=0, - ) - router_logits = get_pcp_group().all_gather( - router_logits, - dim=0, - ) + if ( + self.moe_config.pcp_size > 1 + and not self.moe_config.moe_parallel_config.use_all2all_kernels + ): + hidden_states = get_pcp_group().all_gather(hidden_states, dim=0) + router_logits = get_pcp_group().all_gather(router_logits, dim=0) return hidden_states, router_logits @@ -777,11 +771,11 @@ class MoERunner(MoERunnerInterface): hidden_states, self.moe_config.is_sequence_parallel ) - if self.moe_config.pcp_size > 1: - hidden_states = get_pcp_group().reduce_scatter( - hidden_states, - dim=0, - ) + if ( + self.moe_config.pcp_size > 1 + and not self.moe_config.moe_parallel_config.use_all2all_kernels + ): + hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0) if self.shared_experts is not None: assert shared_output is not None diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index ceb52e5d32..5b8e2bf008 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -9,10 +9,11 @@ from vllm import _custom_ops as ops from vllm._aiter_ops import rocm_aiter_ops from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.config import get_current_vllm_config -from vllm.distributed import get_dcp_group +from vllm.distributed import get_dcp_group, get_pcp_group from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp +from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) @@ -308,6 +309,7 @@ def sparse_attn_indexer( total_seq_lens: int, topk_indices_buffer: torch.Tensor, skip_k_cache_insert: bool, + use_pcp: bool, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -354,6 +356,7 @@ def sparse_attn_indexer( total_seq_lens, topk_indices_buffer, skip_k_cache_insert, + use_pcp, use_fp4_cache, ) attn_metadata_narrowed = attn_metadata[k_cache_prefix] @@ -373,18 +376,28 @@ def sparse_attn_indexer( # During speculative decoding, k may be padded to the CUDA graph batch # size while slot_mapping only covers actual tokens. Truncate k to avoid # out-of-bounds reads in the kernel. + # Keep PCP padding so every rank contributes the same all-gather shape. num_tokens = slot_mapping.shape[0] + if use_pcp: + num_tokens //= get_pcp_group().world_size if k is not None: k = k[:num_tokens] if not skip_k_cache_insert: + assert k is not None + k, slot_mapping_for_cache = maybe_gather_indexer_k( + k, + slot_mapping, + num_decode_tokens, + use_pcp, + ) # scale_fmt can be None, but the function expects str assert scale_fmt is not None assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" ops.indexer_k_quant_and_cache( k, kv_cache, - slot_mapping, + slot_mapping_for_cache, quant_block_size, scale_fmt, ) @@ -498,7 +511,14 @@ def sparse_attn_indexer( assert decode_metadata is not None kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) decode_lens = decode_metadata.decode_lens - if decode_metadata.requires_padding: + if num_decode_tokens == 0: + padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:]) + padded_q_scale = ( + q_scale[:1].reshape(1, 1, *q_scale.shape[1:]) + if q_scale is not None + else None + ) + elif decode_metadata.requires_padding: # pad in edge case where we have short chunked prefill length < # decode_threshold since we unstrictly split # prefill and decode by decode_threshold @@ -663,6 +683,7 @@ def sparse_attn_indexer_fake( total_seq_lens: int, topk_indices_buffer: torch.Tensor | None, skip_k_cache_insert: bool, + use_pcp: bool, use_fp4_cache: bool = False, dcp_rank: int = 0, dcp_world_size: int = 1, @@ -725,6 +746,7 @@ class SparseAttnIndexer(CustomOp): self.dcp_world_size = parallel_config.decode_context_parallel_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size + self.use_pcp = parallel_config.prefill_context_parallel_size > 1 if current_platform.is_cuda() and not has_deep_gemm(): raise RuntimeError( "Sparse Attention Indexer CUDA op requires DeepGEMM support in " @@ -777,6 +799,7 @@ class SparseAttnIndexer(CustomOp): self.max_total_seq_len, self.topk_indices_buffer, self.skip_k_cache_insert, + self.use_pcp, self.use_fp4_cache, self.dcp_rank, self.dcp_world_size, diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index a82d49dadb..6f23cff0b1 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -281,6 +281,13 @@ class AttentionBackend(ABC): def supports_kv_connector(cls) -> bool: return True + @classmethod + def supports_pcp(cls) -> bool: + try: + return cls.get_impl_cls().supports_pcp + except NotImplementedError: + return False + @classmethod def supports_attn_type(cls, attn_type: str) -> bool: """Check if backend supports a given attention type. @@ -327,6 +334,7 @@ class AttentionBackend(ABC): use_non_causal: bool = False, use_batch_invariant: bool = False, use_kv_connector: bool = False, + use_pcp: bool = False, ) -> list[str]: invalid_reasons = [] if not cls.supports_head_size(head_size): @@ -367,6 +375,8 @@ class AttentionBackend(ABC): invalid_reasons.append("batch invariance not supported") if use_kv_connector and not cls.supports_kv_connector(): invalid_reasons.append("KV connector not supported") + if use_pcp and not cls.supports_pcp(): + invalid_reasons.append("PCP not supported") combination_reason = cls.supports_combination( head_size, dtype, @@ -858,8 +868,8 @@ class AttentionImplBase(ABC, Generic[T]): except AssertionError: self.pcp_world_size = 1 self.pcp_rank = 0 - self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size - self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank + self.total_cp_world_size = self.dcp_world_size + self.total_cp_rank = self.dcp_rank self.need_to_return_lse_for_decode = ( self.dcp_world_size > 1 and self.can_return_lse_for_decode @@ -987,6 +997,8 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]): class MLAAttentionImpl(AttentionImplBase[T], Generic[T]): """MLA attention implementation with forward_mqa and forward_mha methods.""" + supports_pcp: bool = True + @abstractmethod def __init__( self, diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 2f578b4677..bf76cbeda8 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -6,7 +6,7 @@ import torch import vllm.envs as envs from vllm.config import VllmConfig -from vllm.distributed import get_dcp_group +from vllm.distributed import get_dcp_group, get_pcp_group from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -29,7 +29,7 @@ from vllm.v1.attention.backends.utils import ( split_decodes_and_prefills, ) from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec -from vllm.v1.worker.cp_utils import get_total_cp_world_size +from vllm.v1.worker.cp_utils import get_kv_cache_shard_count logger = init_logger(__name__) @@ -109,7 +109,7 @@ def split_indexer_prefill_chunks( end += 1 req_slice = slice(start + request_offset, end + request_offset) - max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m + max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) for q_off in range(0, chunk_m, max_q): sub_m = min(max_q, chunk_m - q_off) chunks.append((req_slice, slice(q_off, q_off + sub_m))) @@ -118,6 +118,10 @@ def split_indexer_prefill_chunks( class DeepseekV32IndexerBackend(AttentionBackend): + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_name() -> str: return "DEEPSEEK_V32_INDEXER" @@ -256,6 +260,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): parallel_config = self.vllm_config.parallel_config self.dcp_world_size = parallel_config.decode_context_parallel_size self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 + self.pcp_world_size = parallel_config.prefill_context_parallel_size + self.use_pcp = self.pcp_world_size > 1 self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size # The DCP sparse-indexer code is parameterized by interleave size, but # interleave > 1 is not yet validated end-to-end (gsm8k parity fails), @@ -339,7 +345,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): ) max_num_blocks_per_req = cdiv( self.vllm_config.model_config.max_model_len, - self.kv_cache_spec.block_size * get_total_cp_world_size(), + self.kv_cache_spec.block_size * get_kv_cache_shard_count(), ) self.expanded_block_table_buffer = torch.zeros( ( @@ -570,6 +576,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): common_attn_metadata, decode_threshold=self.decode_threshold, require_uniform=not self.use_flattening, + treat_short_extends_as_decodes=not self.use_pcp, ) ) @@ -579,6 +586,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): compressed_slot_mapping = slot_mapping compressed_seq_lens = seq_lens if self.compress_ratio > 1: + padded_num_tokens = num_tokens + if self.pcp_world_size > 1: + padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size compressed_slot_mapping = get_compressed_slot_mapping( num_tokens, query_start_loc, @@ -588,6 +598,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): self.compress_ratio, out=self.compressed_slot_mapping_buffer, ) + if self.pcp_world_size > 1: + compressed_slot_mapping = get_pcp_group().all_gather( + self.compressed_slot_mapping_buffer[:padded_num_tokens], + dim=0, + ) compressed_seq_lens = seq_lens // self.compress_ratio prefill_metadata = None diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index c8c9a7334a..c6e38c7fbf 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -612,7 +612,9 @@ def split_decodes_and_prefills( # check if we are in a padded uniform batch; this is used for full-CGs, some # requests may have a query length of 0 but since they are padding its fine # to treat them as decodes (ensures num_decodes matches the captured size) - if torch.all((query_lens == query_lens[0]) | (query_lens == 0)): + if treat_short_extends_as_decodes and torch.all( + (query_lens == query_lens[0]) | (query_lens == 0) + ): return num_reqs, 0, num_tokens, 0 # all decodes is_prefill = query_lens != query_lens[0] else: diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 387eac34c9..734136334a 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -36,6 +36,7 @@ class AttentionSelectorConfig(NamedTuple): use_non_causal: bool = False use_batch_invariant: bool = False use_kv_connector: bool = False + use_pcp: bool = False def __repr__(self): return ( @@ -52,7 +53,8 @@ class AttentionSelectorConfig(NamedTuple): f"has_sliding_window={self.has_sliding_window}, " f"use_non_causal={self.use_non_causal}, " f"use_batch_invariant={self.use_batch_invariant}, " - f"use_kv_connector={self.use_kv_connector})" + f"use_kv_connector={self.use_kv_connector}, " + f"use_pcp={self.use_pcp})" ) @@ -150,6 +152,7 @@ def get_attn_backend( use_non_causal=vllm_config.attention_config.use_non_causal, use_batch_invariant=envs.VLLM_BATCH_INVARIANT, use_kv_connector=use_kv_connector, + use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1, ) # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 7d482192d5..df8769c3f3 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -472,8 +472,6 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator): self.pcp_world_size = pcp_world_size if dcp_world_size > 1: self.block_size *= dcp_world_size - if pcp_world_size > 1: - self.block_size *= pcp_world_size # For models using only Mamba, block_size is set to max_model_len when # prefix caching is disabled, and hash_block_size validation is skipped. assert not enable_caching or (hash_block_size == self.block_size), ( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index dc5ecf3f4d..d4970d91b2 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -631,8 +631,8 @@ def resolve_kv_cache_block_sizes( - ``scheduler_block_size`` is the token-alignment invariant used by the scheduler (e.g. for ``num_computed_tokens`` rounding). Single group: - ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every - group's effective block size. Attention groups are scaled by DCP/PCP; + ``cache_config.block_size * dcp``. Multiple groups: LCM of every + group's effective block size. Attention groups are scaled by DCP; Mamba groups keep their full per-rank state and are not scaled. - ``hash_block_size`` is the granularity at which ``Request.block_hashes`` is computed. Single group: equals scheduler block size. Multiple groups: @@ -644,15 +644,14 @@ def resolve_kv_cache_block_sizes( """ cache_config = vllm_config.cache_config dcp = vllm_config.parallel_config.decode_context_parallel_size - pcp = vllm_config.parallel_config.prefill_context_parallel_size groups = kv_cache_config.kv_cache_groups - if len(groups) <= 1: # Single group: block_size * dcp * pcp - bs = cache_config.block_size * dcp * pcp + if len(groups) <= 1: + bs = cache_config.block_size * dcp return bs, bs group_block_sizes = [ - g.kv_cache_spec.block_size * dcp * pcp + g.kv_cache_spec.block_size * dcp if isinstance(g.kv_cache_spec, AttentionSpec) else g.kv_cache_spec.block_size for g in groups diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index de076e230e..46a999ef7c 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -271,7 +271,7 @@ class Scheduler(SchedulerInterface): log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, - pcp_world_size=self.pcp_world_size, + pcp_world_size=1, scheduler_block_size=self.block_size, hash_block_size=hash_block_size, metrics_collector=self.kv_metrics_collector, diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2440483909..f8578c68a3 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -75,8 +75,8 @@ class SingleTypeKVCacheManager(ABC): self.block_size = kv_cache_spec.block_size self.dcp_world_size = dcp_world_size self.pcp_world_size = pcp_world_size - if dcp_world_size * pcp_world_size > 1: - self.block_size *= dcp_world_size * pcp_world_size + if dcp_world_size > 1: + self.block_size *= dcp_world_size self.kv_cache_spec = kv_cache_spec self.block_pool = block_pool self.enable_caching = enable_caching @@ -674,10 +674,10 @@ class FullAttentionManager(SingleTypeKVCacheManager): "and chunked local attention groups" ) block_size = kv_cache_spec.block_size - if dcp_world_size * pcp_world_size > 1: - # DCP/PCP shard each block's KV across ranks; hashes must be - # viewed at the sharded (scaled) block size. - block_size *= dcp_world_size * pcp_world_size + if dcp_world_size > 1: + # DCP shards each block's KV across ranks; hashes must be viewed at + # the sharded block size. + block_size *= dcp_world_size block_hashes = resolve_block_hashes( block_hashes, block_pool.hash_block_size, diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d7f195e35b..7a44d4e1d2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -218,14 +218,9 @@ class AttentionSpec(KVCacheSpec): ) def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: - # Attention KV is token-interleaved across DCP/PCP ranks, so each rank - # only stores max_len // (dcp * pcp) tokens per request. parallel_config = vllm_config.parallel_config - total_cp_size = ( - parallel_config.decode_context_parallel_size - * parallel_config.prefill_context_parallel_size - ) - return cdiv(max_len, self.block_size * total_cp_size) + kv_shard_count = parallel_config.decode_context_parallel_size + return cdiv(max_len, self.block_size * kv_shard_count) @dataclass(frozen=True, kw_only=True) @@ -263,11 +258,8 @@ class FullAttentionSpec(AttentionSpec): def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: max_model_len = vllm_config.model_config.max_model_len dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size - # Note(hc): each dcp rank only need save - # (max_model_len//dcp_world_size) tokens locally. - if dcp_world_size * pcp_world_size > 1: - max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size) + if dcp_world_size > 1: + max_model_len = cdiv(max_model_len, dcp_world_size) return cdiv(max_model_len, self.block_size) * self.page_size_bytes @classmethod diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index 1d2b8a6b7e..515a2f19b0 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -83,8 +83,7 @@ class SimpleCPUOffloadScheduler: and vllm_config.kv_events_config.enable_kv_cache_events ) dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size - pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size - self.cp_world_size = dcp_world_size * pcp_world_size + self.cp_world_size = dcp_world_size self.block_size = scheduler_block_size self.hash_block_size = hash_block_size assert self.block_size % self.hash_block_size == 0 @@ -127,7 +126,7 @@ class SimpleCPUOffloadScheduler: enable_caching=True, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=dcp_world_size, - pcp_world_size=pcp_world_size, + pcp_world_size=1, scheduler_block_size=self.block_size, hash_block_size=self.hash_block_size, ) @@ -354,7 +353,7 @@ class SimpleCPUOffloadScheduler: ) assert num_external_tokens % g_block_size == 0, ( f"num_external_tokens={num_external_tokens} not aligned to " - f"group {g} effective block_size={g_block_size}" + f"group {g} block_size={g_block_size}" ) n_take_g = num_external_tokens // g_block_size cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g]) diff --git a/vllm/v1/worker/block_table.py b/vllm/v1/worker/block_table.py index d40887879f..332eda4cbd 100644 --- a/vllm/v1/worker/block_table.py +++ b/vllm/v1/worker/block_table.py @@ -163,8 +163,6 @@ class BlockTable: return assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT - total_cp_world_size = self.pcp_world_size * self.dcp_world_size - total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank _compute_slot_mapping_kernel[(num_reqs + 1,)]( num_tokens, self.max_num_batched_tokens, @@ -176,8 +174,8 @@ class BlockTable: self.slot_mapping.gpu, KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size, BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block, - TOTAL_CP_WORLD_SIZE=total_cp_world_size, - TOTAL_CP_RANK=total_cp_rank, + TOTAL_CP_WORLD_SIZE=self.dcp_world_size, + TOTAL_CP_RANK=self.dcp_rank, CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size, PAD_ID=PAD_SLOT_ID, BLOCK_SIZE=1024, diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 11edd86a5d..92d8383c1f 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, cast import torch from vllm.config import VllmConfig, get_layers_from_vllm_config -from vllm.distributed import get_dcp_group, get_pcp_group +from vllm.distributed import get_dcp_group from vllm.logger import init_logger from vllm.v1.attention.backend import CommonAttentionMetadata from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends @@ -27,6 +27,13 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: layer_type = cast(type[Any], AttentionLayerBase) layers = get_layers_from_vllm_config(vllm_config, layer_type) for layer in layers.values(): + get_attn_backend = getattr(layer, "get_attn_backend", None) + if pcp_size > 1 and get_attn_backend is not None: + backend = get_attn_backend() + assert backend.supports_pcp(), ( + "PCP requires attention backend support, " + f"but {backend.get_name()} does not support PCP." + ) layer_impl = getattr(layer, "impl", None) if layer_impl is None: continue @@ -44,26 +51,14 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: "--attention-backend or disable DCP." ) - if pcp_size > 1: - assert layer_impl.supports_pcp, ( - "PCP requires attention impls' support, " - f"but the impl {layer_impl.__class__.__name__} " - "does not support PCP." - ) - -def get_total_cp_world_size(): - try: - pcp_world_size = get_pcp_group().world_size - except AssertionError: - # PCP might not be initialized in testing - pcp_world_size = 1 +def get_kv_cache_shard_count() -> int: try: dcp_world_size = get_dcp_group().world_size except AssertionError: # DCP might not be initialized in testing dcp_world_size = 1 - return dcp_world_size * pcp_world_size + return dcp_world_size def get_dcp_dummy_context_len( diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 5c07860b3b..a97a4e3959 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -578,6 +578,7 @@ def build_attn_metadata( seq_lens_cpu_upper_bound: torch.Tensor | None = None, dcp_local_seq_lens: 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, model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None, for_cudagraph_capture: bool = False, @@ -605,6 +606,11 @@ def build_attn_metadata( if model_specific_attn_metadata is not None else {} ) + # Model-specific metadata (e.g. Mamba hybrid) may supply its own + # padding-aware is_prefilling, which takes precedence over the default. + group_is_prefilling = common_attn_metadata_extra_kwargs.pop( + "is_prefilling", is_prefilling + ) common_attn_metadata = CommonAttentionMetadata( query_start_loc=query_start_loc_gpu, query_start_loc_cpu=query_start_loc_cpu, @@ -619,6 +625,7 @@ def build_attn_metadata( causal=group_causal, dcp_local_seq_lens=dcp_local_seq_lens, positions=positions, + is_prefilling=group_is_prefilling, mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, **common_attn_metadata_extra_kwargs, diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 8d41ba5a36..22c4afc11b 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -135,20 +135,28 @@ class BlockTables: self, idx_mapping: torch.Tensor, num_reqs_padded: int, + out: tuple[torch.Tensor, ...] | None = None, + out_ptrs: torch.Tensor | None = None, ) -> tuple[torch.Tensor, ...]: + if out is None: + out = tuple(self.input_block_tables) + out_ptrs = self.input_block_table_ptrs + else: + assert out_ptrs is not None + assert len(out) == self.num_kv_cache_groups num_reqs = idx_mapping.shape[0] # Launch kernel with num_reqs_padded to fuse zeroing of padded rows. _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)]( idx_mapping, self.block_table_ptrs, - self.input_block_table_ptrs, + out_ptrs, self.block_table_strides, self.num_blocks.gpu, self.num_blocks.gpu.stride(0), num_reqs, BLOCK_SIZE=1024, # type: ignore ) - return tuple(bt[:num_reqs_padded] for bt in self.input_block_tables) + return tuple(bt[:num_reqs_padded] for bt in out) def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: # NOTE(woosuk): The output may be used for CUDA graph capture. @@ -163,26 +171,28 @@ class BlockTables: query_start_loc: torch.Tensor, positions: torch.Tensor, num_tokens_padded: int, + out: torch.Tensor | None = None, ) -> torch.Tensor: num_reqs = idx_mapping.shape[0] num_groups = self.num_kv_cache_groups + slot_mappings = self.slot_mappings if out is None else out _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)]( - self.max_num_batched_tokens, + slot_mappings.shape[1], idx_mapping, query_start_loc, positions, self.block_table_ptrs, self.block_table_strides, self.block_sizes_tensor, - self.slot_mappings, - self.slot_mappings.stride(0), + slot_mappings, + slot_mappings.stride(0), self.cp_rank, CP_SIZE=self.cp_size, CP_INTERLEAVE=self.cp_interleave, PAD_ID=PAD_SLOT_ID, TRITON_BLOCK_SIZE=1024, # type: ignore ) - return self.slot_mappings[:, :num_tokens_padded] + return slot_mappings[:, :num_tokens_padded] def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor: # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries. diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 518d12a6b2..293b06ff97 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -52,6 +52,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput from vllm.v1.worker.cp_utils import check_attention_cp_compatibility +from vllm.v1.worker.gpu import pcp_manager as pcp from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput from vllm.v1.worker.gpu.attn_utils import ( build_slot_mappings_by_layer, @@ -205,6 +206,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Draft tokens propagation - for spec-dec + struct outputs. self.draft_tokens_handler = DraftTokensHandler(self.device) + self.pcp_manager: pcp.PCPManager | None = None + # Pooling models. self.is_pooling_model = self.model_config.runner_type == "pooling" self.pooling_runner: PoolingRunner | None = None @@ -456,6 +459,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): cp_rank=self.dcp_rank, cp_interleave=self.cp_interleave, ) + self.pcp_manager = pcp.maybe_build_pcp_manager( + self.vllm_config, + self.device, + self.supports_mm_inputs, + self.req_states, + self.block_tables, + ) initialize_mamba_ssu_backend( self.vllm_config.mamba_config, self.kv_cache_config ) @@ -1008,7 +1018,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): # prompt_lens is only used in R-SWA case. prompt_lens = self.req_states.prompt_len.gpu[idx_mapping] - return InputBatch( + input_batch = InputBatch( req_ids=req_ids, num_reqs=num_reqs, num_reqs_after_padding=num_reqs_padded, @@ -1040,10 +1050,14 @@ class GPUModelRunner(LoRAModelRunnerMixin): has_structured_output_reqs=scheduler_output.has_structured_output_requests, prompt_lens=prompt_lens, ) + return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch) def prepare_attn( self, input_batch: InputBatch ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + if self.pcp_manager is not None: + return self.pcp_manager.prepare_attn(input_batch) + # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks]. block_tables = self.block_tables.gather_block_tables( input_batch.idx_mapping, @@ -1063,8 +1077,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): self, input_batch: InputBatch ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) - slot_mappings = self.block_tables.get_dummy_slot_mappings( - input_batch.num_tokens + slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings( + self.pcp_manager, self.block_tables, input_batch.num_tokens ) return block_tables, slot_mappings @@ -1412,6 +1426,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) # Last rank: sample tokens + hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling( + self.pcp_manager, hidden_states, input_batch + ) + sampler_output, num_sampled, num_rejected = self.sample( hidden_states, input_batch, grammar_output ) diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 18eb40640a..053f37aaf1 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -151,7 +151,10 @@ class DefaultModelState(ModelState): # For piecewise cudagraphs and eager, use unpadded sizes. num_reqs = input_batch.num_reqs num_tokens = input_batch.num_tokens - query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) + query_start_loc_cpu = torch.from_numpy( + input_batch.query_start_loc_np[: num_reqs + 1] + ) + query_start_loc_gpu = input_batch.query_start_loc[: num_reqs + 1] max_query_len = input_batch.num_scheduled_tokens.max().item() seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound if for_capture: @@ -174,7 +177,7 @@ class DefaultModelState(ModelState): attn_groups=attn_groups, num_reqs=num_reqs, num_tokens=num_tokens, - query_start_loc_gpu=input_batch.query_start_loc, + query_start_loc_gpu=query_start_loc_gpu, query_start_loc_cpu=query_start_loc_cpu, max_query_len=max_query_len, seq_lens=input_batch.seq_lens, @@ -185,6 +188,7 @@ class DefaultModelState(ModelState): seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=input_batch.dcp_local_seq_lens, 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, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py new file mode 100644 index 0000000000..f50e487494 --- /dev/null +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -0,0 +1,680 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass, replace + +import numpy as np +import torch + +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.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, + combine_sampled_and_draft_tokens, + prepare_pos_seq_lens, +) +from vllm.v1.worker.gpu.states import RequestState + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class RankSegment: + global_batch_req_idx: int + global_batch_slice: slice + rank_local_batch_slice: slice + + @property + def num_tokens(self) -> int: + return self.global_batch_slice.stop - self.global_batch_slice.start + + +class PCPManager: + """MRV2 PC batch manager. + + The model runner keeps the global scheduled batch. This manager rewrites only + the per-step InputBatch into rank-local DualChunkSwap rows and keeps the + global-batch view private to restore to the global batch shape before + sampling/postprocess. + """ + + def __init__( + self, + pcp_world_size: int, + pcp_rank: int, + device: torch.device, + req_states: RequestState | None = None, + max_num_reqs: int | None = None, + max_num_tokens: int | None = None, + block_tables: BlockTables | None = None, + dcp_world_size: int = 1, + dcp_rank: int = 0, + cp_interleave: int = 1, + ) -> None: + self.pcp_world_size = pcp_world_size + self.pcp_rank = pcp_rank + self.device = device + self.dcp_world_size = dcp_world_size + self.dcp_rank = dcp_rank + self.cp_interleave = cp_interleave + + self._global_batch: InputBatch | None = None + self._req_states = req_states + self._block_tables = block_tables + self._hidden_restore_idx: torch.Tensor | None = None + self._padded_gather_idx: torch.Tensor | None = None + self._gathered_kv_write_mask: torch.Tensor | None = None + self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device) + + max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None + self._input_buffers = ( + InputBuffers(max_num_local_reqs, max_num_tokens, device) + if max_num_local_reqs is not None and max_num_tokens is not None + else None + ) + self._local_req_idx = ( + torch.arange(max_num_local_reqs, dtype=torch.int32, device=device) + if max_num_local_reqs is not None + else None + ) + self._local_block_tables: tuple[torch.Tensor, ...] | None + self._local_block_table_ptrs: torch.Tensor | None + if block_tables is not None and max_num_local_reqs is not None: + self._local_block_tables = tuple( + table.new_zeros((max_num_local_reqs, table.shape[1])) + for table in block_tables.input_block_tables + ) + self._local_block_table_ptrs = torch.tensor( + [table.data_ptr() for table in self._local_block_tables], + dtype=torch.uint64, + device=device, + ) + else: + self._local_block_tables = None + self._local_block_table_ptrs = None + num_kv_cache_groups = ( + block_tables.num_kv_cache_groups if block_tables is not None else 0 + ) + self._global_batch_slot_mappings = ( + torch.empty( + num_kv_cache_groups, + max_num_tokens, + dtype=torch.int64, + device=device, + ) + if max_num_tokens is not None and num_kv_cache_groups > 0 + else None + ) + self._gathered_kv_slot_mappings = ( + torch.empty( + num_kv_cache_groups, + max_num_tokens * pcp_world_size, + dtype=torch.int64, + device=device, + ) + if max_num_tokens is not None and num_kv_cache_groups > 0 + else None + ) + + @staticmethod + def validate_config( + vllm_config: VllmConfig, + supports_mm_inputs: bool, + ) -> None: + parallel_config = vllm_config.parallel_config + model_config = vllm_config.model_config + pcp_size = parallel_config.prefill_context_parallel_size + if pcp_size <= 1: + return + + if not model_config.use_mla: + raise NotImplementedError("MRV2 PCP currently supports MLA models only.") + if parallel_config.pipeline_parallel_size > 1: + raise NotImplementedError("MRV2 PCP does not support PP yet.") + if model_config.is_encoder_decoder: + raise NotImplementedError( + "MRV2 PCP does not support encoder-decoder models yet." + ) + if supports_mm_inputs: + raise NotImplementedError("MRV2 PCP does not support MM inputs yet.") + if vllm_config.lora_config is not None: + raise NotImplementedError("MRV2 PCP does not support LoRA yet.") + if vllm_config.speculative_config is not None: + raise NotImplementedError( + "MRV2 PCP does not support speculative decoding yet." + ) + is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk") + if ( + is_sparse_mla + and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE + ): + raise NotImplementedError( + "MRV2 sparse MLA PCP does not support CUDA graphs yet. " + "Set -cc.cudagraph_mode=NONE." + ) + if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): + raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + + @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. + """ + + def is_pure_prefill(segment: RankSegment) -> bool: + 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 is_prefilling[req_idx] and start_pos == 0 + + segments.sort(key=is_pure_prefill) + rank_offset = 0 + for index, segment in enumerate(segments): + segments[index] = replace( + segment, + rank_local_batch_slice=slice( + rank_offset, rank_offset + segment.num_tokens + ), + ) + rank_offset += segment.num_tokens + return segments + + def _get_rank_segments( + 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]: + """Build one rank's attention-compatible DualChunkSwap rows. + + PCP=4 partitions each prefill into eight chunks: + + full: | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | + rank 0: 0 7 + rank 1: 1 6 + rank 2: 2 5 + rank 3: 3 4 + """ + rank_segments = [] + rank_offset = 0 + num_chunks = 2 * self.pcp_world_size + for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens): + query_len = int(num_tokens) + if query_len == 0: + continue + global_batch_start = int(query_start_loc_np[global_batch_req_idx]) + chunk_indices: tuple[int, ...] + if bool(is_prefilling[global_batch_req_idx]): + chunk_size = (query_len + num_chunks - 1) // num_chunks + chunk_indices = (rank, num_chunks - 1 - rank) + else: # decodes are replicated + chunk_size = query_len + chunk_indices = (0,) + + for chunk_idx in chunk_indices: + chunk_offset = chunk_idx * chunk_size + chunk_len = min(chunk_size, query_len - chunk_offset) + if chunk_len <= 0: + continue + chunk_start = global_batch_start + chunk_offset + rank_segments.append( + RankSegment( + global_batch_req_idx=global_batch_req_idx, + global_batch_slice=slice(chunk_start, chunk_start + chunk_len), + rank_local_batch_slice=slice( + rank_offset, rank_offset + chunk_len + ), + ) + ) + rank_offset += chunk_len + return self._reorder_segments( + rank_segments, + num_computed_tokens, + is_prefilling, + query_start_loc_np, + ) + + def _build_batch_layout( + self, + num_scheduled_tokens: np.ndarray, + num_computed_tokens: np.ndarray, + is_prefilling: np.ndarray, + query_start_loc_np: np.ndarray, + ) -> tuple[list[list[RankSegment]], list[int]]: + 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, + ) + num_rank_tokens = sum(segment.num_tokens for segment in segments) + segments_by_rank.append(segments) + per_rank_num_tokens.append(num_rank_tokens) + + # PCP=2 example: + # global batch: [A B C D E F G] + # rank 0 / rank 1: [A B G] / [C D E F] + # padded gathered: [A B G _ | C D E F] + # hidden_restore_idx: [0, 1, 4, 5, 6, 7, 2] + # padded_gather_idx: [0, 1, 6, 0, 2, 3, 4, 5] + # Therefore global = gathered[hidden_restore_idx] and + # padded_gathered = global[padded_gather_idx]. + hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64) + padded_num_tokens = max(per_rank_num_tokens) + num_expanded_tokens = padded_num_tokens * self.pcp_world_size + padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64) + gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_) + for rank, segments in enumerate(segments_by_rank): + expanded_rank_offset = rank * padded_num_tokens + for segment in segments: + padded_gathered_slice = slice( + expanded_rank_offset + segment.rank_local_batch_slice.start, + expanded_rank_offset + segment.rank_local_batch_slice.stop, + ) + padded_gather_idx[padded_gathered_slice] = np.arange( + segment.global_batch_slice.start, + segment.global_batch_slice.stop, + 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: + continue + gathered_kv_write_mask[padded_gathered_slice] = True + hidden_restore_idx[segment.global_batch_slice] = np.arange( + padded_gathered_slice.start, + padded_gathered_slice.stop, + dtype=np.int64, + ) + + self._hidden_restore_idx = async_copy_to_gpu( + hidden_restore_idx, device=self.device + ) + self._padded_gather_idx = async_copy_to_gpu( + padded_gather_idx, device=self.device + ) + self._gathered_kv_write_mask = async_copy_to_gpu( + gathered_kv_write_mask, device=self.device + ) + return segments_by_rank, per_rank_num_tokens + + def partition_batch(self, input_batch: InputBatch) -> InputBatch: + assert self._req_states is not None + assert self._input_buffers is not None + req_states = self._req_states + input_buffers = self._input_buffers + if input_batch.num_draft_tokens > 0: + raise NotImplementedError("MRV2 PCP does not support spec decode yet.") + + global_batch = input_batch + self._global_batch = global_batch + + num_scheduled_tokens = global_batch.num_scheduled_tokens + num_computed_tokens = global_batch.num_computed_tokens_np + is_prefilling = global_batch.is_prefilling_np + + segments_by_rank, per_rank_num_tokens = self._build_batch_layout( + num_scheduled_tokens, + num_computed_tokens, + is_prefilling, + global_batch.query_start_loc_np, + ) + + local_segments = segments_by_rank[self.pcp_rank] + if not local_segments: + local_segments = [ + RankSegment( + global_batch_req_idx=0, + global_batch_slice=slice(0, 0), + rank_local_batch_slice=slice(0, 0), + ) + ] + + num_local_reqs = len(local_segments) + if num_local_reqs > input_buffers.max_num_reqs: + raise RuntimeError( + "PCP local request count exceeds the MRV2 input buffer size: " + f"{num_local_reqs} > {input_buffers.max_num_reqs}." + ) + + local_to_global_batch_req_idx_np = np.fromiter( + (segment.global_batch_req_idx for segment in local_segments), + dtype=np.int32, + count=num_local_reqs, + ) + local_start_pos_np = np.fromiter( + ( + num_computed_tokens[segment.global_batch_req_idx] + + segment.global_batch_slice.start + - global_batch.query_start_loc_np[segment.global_batch_req_idx] + for segment in local_segments + ), + dtype=np.int32, + count=num_local_reqs, + ) + local_num_scheduled_tokens = np.fromiter( + (segment.num_tokens for segment in local_segments), + dtype=np.int32, + count=num_local_reqs, + ) + local_to_global_req_idx_np = global_batch.idx_mapping_np[ + local_to_global_batch_req_idx_np + ] + local_req_ids = [ + global_batch.req_ids[global_batch_req_idx] + for global_batch_req_idx in local_to_global_batch_req_idx_np + ] + + num_local_tokens = int(local_num_scheduled_tokens.sum()) + num_local_tokens_padded = max(per_rank_num_tokens) + fresh_prefills = int( + np.count_nonzero(is_prefilling & (num_computed_tokens == 0)) + ) + continued_prefills = int( + np.count_nonzero(is_prefilling & (num_computed_tokens > 0)) + ) + logger.debug( + "PCP batch: rank=%d global_batch_reqs=%d fresh_prefills=%d " + "continued_prefills=%d decodes=%d local_reqs=%d " + "local_tokens=%d per_rank_tokens=%s", + self.pcp_rank, + global_batch.num_reqs, + fresh_prefills, + continued_prefills, + global_batch.num_reqs - fresh_prefills - continued_prefills, + num_local_reqs, + num_local_tokens, + per_rank_num_tokens, + ) + if num_local_tokens_padded > input_buffers.max_num_tokens: + raise RuntimeError( + "PCP local token count exceeds the MRV2 input buffer size: " + f"{num_local_tokens_padded} > {input_buffers.max_num_tokens}." + ) + rank_token_start = self.pcp_rank * num_local_tokens_padded + assert self._padded_gather_idx is not None + local_gather_idx = self._padded_gather_idx[ + rank_token_start : rank_token_start + num_local_tokens_padded + ] + torch.index_select( + global_batch.input_ids, + 0, + local_gather_idx, + out=input_buffers.input_ids[:num_local_tokens_padded], + ) + + local_query_start_loc_np = np.empty( + input_buffers.max_num_reqs + 1, dtype=np.int32 + ) + local_query_start_loc_np[0] = 0 + local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1] + np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) + local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens + async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) + local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1] + + local_to_global_req_idx = async_copy_to_gpu( + local_to_global_req_idx_np, device=self.device + ) + local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device) + + assert self._local_req_idx is not None + prepare_pos_seq_lens( + self._local_req_idx[:num_local_reqs], + local_query_start_loc, + local_start_pos, + input_buffers.positions, + input_buffers.seq_lens[:num_local_reqs], + ) + seq_lens = input_buffers.seq_lens[:num_local_reqs] + is_padding = input_buffers.is_padding[:num_local_tokens_padded] + is_padding[:num_local_tokens].fill_(False) + is_padding[num_local_tokens:].fill_(True) + if num_local_tokens_padded > num_local_tokens: + input_buffers.input_ids[:num_local_tokens_padded].masked_fill_( + is_padding, 0 + ) + input_buffers.positions[:num_local_tokens_padded].masked_fill_( + is_padding, 0 + ) + + total_num_logits = num_local_reqs if num_local_tokens > 0 else 0 + if total_num_logits > 0: + cu_num_logits_np = np.arange(num_local_reqs + 1, dtype=np.int32) + cu_num_logits = torch.arange( + num_local_reqs + 1, device=self.device, dtype=torch.int32 + ) + else: + cu_num_logits_np = np.zeros(num_local_reqs + 1, dtype=np.int32) + cu_num_logits = torch.zeros( + num_local_reqs + 1, device=self.device, dtype=torch.int32 + ) + logits_indices = combine_sampled_and_draft_tokens( + input_buffers.input_ids, + local_to_global_req_idx, + req_states.last_sampled_tokens, + local_query_start_loc, + seq_lens, + req_states.prefill_len.gpu, + req_states.draft_tokens, + cu_num_logits, + total_num_logits, + 1, + ) + + local_prefill_len_np = global_batch.prefill_len_np[ + local_to_global_batch_req_idx_np + ] + local_num_computed_prefill_tokens_np = np.minimum( + local_start_pos_np, local_prefill_len_np + ) + local_is_prefilling_np = ( + local_num_computed_prefill_tokens_np < local_prefill_len_np + ) + 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, + num_reqs=num_local_reqs, + num_reqs_after_padding=num_local_reqs, + idx_mapping=local_to_global_req_idx, + idx_mapping_np=local_to_global_req_idx_np, + expanded_idx_mapping=local_to_global_req_idx, + expanded_local_pos=torch.zeros( + num_local_reqs, dtype=torch.int32, device=self.device + ), + num_scheduled_tokens=local_num_scheduled_tokens, + num_tokens=num_local_tokens, + num_tokens_after_padding=num_local_tokens_padded, + num_draft_tokens=0, + num_draft_tokens_per_req=None, + query_start_loc=local_query_start_loc, + 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, + 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, + is_prefilling_np=local_is_prefilling_np, + max_seq_len_np=global_batch.max_seq_len_np[local_to_global_batch_req_idx_np] + if global_batch.max_seq_len_np is not None + else None, + input_ids=input_buffers.input_ids[:num_local_tokens_padded], + positions=input_buffers.positions[:num_local_tokens_padded], + is_padding=is_padding, + logits_indices=logits_indices, + cu_num_logits=cu_num_logits, + cu_num_logits_np=cu_num_logits_np, + prompt_lens=None, + ) + + def prepare_attn( + self, input_batch: InputBatch + ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: + assert self._block_tables is not None + assert self._local_block_tables is not None + assert self._local_block_table_ptrs is not None + block_tables = self._block_tables.gather_block_tables( + input_batch.idx_mapping, + input_batch.num_reqs_after_padding, + out=self._local_block_tables, + out_ptrs=self._local_block_table_ptrs, + ) + slot_mappings = self.prepare_slot_mappings() + return block_tables, slot_mappings + + def prepare_slot_mappings(self) -> torch.Tensor: + assert self._block_tables is not None + assert self._global_batch_slot_mappings is not None + assert self._global_batch is not None + global_batch = self._global_batch + global_batch_slot_mappings = self._block_tables.compute_slot_mappings( + global_batch.idx_mapping, + global_batch.query_start_loc, + global_batch.positions, + global_batch.num_tokens, + out=self._global_batch_slot_mappings, + ) + return self._convert_to_gathered_slot_mappings(global_batch_slot_mappings) + + def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor: + assert self._gathered_kv_slot_mappings is not None + self._gathered_kv_slot_mappings.fill_(PAD_SLOT_ID) + return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size] + + def _convert_to_gathered_slot_mappings( + self, + global_batch_slot_mappings: torch.Tensor, + ) -> torch.Tensor: + assert self._padded_gather_idx is not None + assert self._gathered_kv_write_mask is not None + padded_gather_idx = self._padded_gather_idx + num_expanded_tokens = padded_gather_idx.shape[0] + if self._gathered_kv_slot_mappings is None: + self._gathered_kv_slot_mappings = global_batch_slot_mappings.new_empty( + global_batch_slot_mappings.shape[0], num_expanded_tokens + ) + gathered_kv_slot_mappings = self._gathered_kv_slot_mappings[ + :, :num_expanded_tokens + ] + torch.index_select( + global_batch_slot_mappings, + 1, + padded_gather_idx, + out=gathered_kv_slot_mappings, + ) + torch.where( + self._gathered_kv_write_mask.unsqueeze(0), + gathered_kv_slot_mappings, + self._pad_slot_id, + out=gathered_kv_slot_mappings, + ) + return gathered_kv_slot_mappings + + def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self._hidden_restore_idx is None: + return hidden_states + gathered = get_pcp_group().all_gather(hidden_states, dim=0) + return gathered[self._hidden_restore_idx] + + def restore_for_sampling( + self, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, InputBatch]: + assert self._global_batch is not None + return self.restore_hidden_states(hidden_states), self._global_batch + + +def maybe_partition_pcp_batch( + manager: PCPManager | None, + input_batch: InputBatch, +) -> InputBatch: + if manager is None: + return input_batch + return manager.partition_batch(input_batch) + + +def maybe_get_pcp_dummy_slot_mappings( + manager: PCPManager | None, + block_tables: BlockTables, + num_tokens: int, +) -> torch.Tensor: + if manager is None: + return block_tables.get_dummy_slot_mappings(num_tokens) + return manager.get_dummy_slot_mappings(num_tokens) + + +def maybe_restore_pcp_for_sampling( + manager: PCPManager | None, + hidden_states: torch.Tensor | None, + input_batch: InputBatch, +) -> tuple[torch.Tensor, InputBatch]: + assert hidden_states is not None + if manager is None: + return hidden_states, input_batch + return manager.restore_for_sampling(hidden_states) + + +def maybe_build_pcp_manager( + vllm_config: VllmConfig, + device: torch.device, + supports_mm_inputs: bool, + req_states: RequestState, + block_tables: BlockTables, +) -> PCPManager | None: + parallel_config = vllm_config.parallel_config + pcp_size = parallel_config.prefill_context_parallel_size + if pcp_size <= 1: + return None + + PCPManager.validate_config(vllm_config, supports_mm_inputs) + + pcp_rank = get_pcp_group().rank_in_group + dcp_size = parallel_config.decode_context_parallel_size + dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0 + + return PCPManager( + pcp_world_size=pcp_size, + pcp_rank=pcp_rank, + device=device, + req_states=req_states, + max_num_reqs=vllm_config.scheduler_config.max_num_seqs, + max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens, + block_tables=block_tables, + dcp_world_size=dcp_size, + dcp_rank=dcp_rank, + cp_interleave=parallel_config.cp_kv_cache_interleave_size, + )