diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index b0c3e26db99a..2c1aeb92dc56 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -8,6 +9,7 @@ from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import ForwardContext, override_forward_context from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, @@ -24,6 +26,7 @@ from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) +from vllm.model_executor.layers.fused_moe.runner import moe_runner from vllm.model_executor.models.llama4 import Llama4MoE from vllm.platforms import current_platform @@ -46,6 +49,61 @@ def _is_aiter_capable() -> bool: NUM_EXPERTS = [8, 16, 64] +@pytest.mark.parametrize( + "pcp_size,use_all2all,uses_input_ids", + [(1, False, False), (2, False, False), (2, False, True), (2, True, True)], +) +def test_pcp_dispatch_gathers_routing_inputs_once_when_used( + pcp_size, use_all2all, uses_input_ids, monkeypatch +): + """Hash routing must see the token IDs belonging to each gathered hidden row.""" + runner = SimpleNamespace( + do_naive_dispatch_combine=False, + router=SimpleNamespace(bias_vl=None), + moe_config=SimpleNamespace( + pcp_size=pcp_size, + has_hash_routing=uses_input_ids, + moe_parallel_config=SimpleNamespace(use_all2all_kernels=use_all2all), + ), + ) + input_id_gathers = 0 + + def all_gather(tensor, dim): + nonlocal input_id_gathers + if tensor.dtype == torch.int64: + input_id_gathers += 1 + remote = 1 - tensor if tensor.dtype == torch.uint8 else tensor + 100 + return torch.cat((tensor, remote), dim=dim) + + monkeypatch.setattr( + moe_runner, "get_pcp_group", lambda: SimpleNamespace(all_gather=all_gather) + ) + token_ids = torch.tensor([3, 7]) + hidden = token_ids.float().unsqueeze(1) + padding = torch.tensor([False, True]) + context = ForwardContext({}, {}, {}, is_padding=padding) + with override_forward_context(context): + with moe_runner.MoERunner._sequence_parallel_context(runner): + states, logits, ids = moe_runner.MoERunner._maybe_dispatch( + runner, hidden, hidden + 4, token_ids + ) + dispatched = pcp_size > 1 and not use_all2all + expected = [3, 7, 103, 107] if dispatched else [3, 7] + if dispatched and not uses_input_ids: + assert ids is None + else: + assert ids is not None + assert ids.tolist() == expected + torch.testing.assert_close(states[:, 0], torch.tensor(expected).float()) + torch.testing.assert_close(logits[:, 0], torch.tensor(expected).float() + 4) + assert context.is_padding.tolist() == ( + [False, True, True, False] if len(expected) == 4 else [False, True] + ) + moe_runner.MoERunner._maybe_dispatch(runner, hidden, hidden + 4, token_ids) + assert input_id_gathers == int(dispatched and uses_input_ids) + assert context.is_padding is padding + + def test_degenerate_grouped_config_uses_standard_topk() -> None: router = create_fused_moe_router( top_k=4, diff --git a/tests/kernels/test_compressor_kv_cache.py b/tests/kernels/test_compressor_kv_cache.py index cc4f4c72a931..1f64a5a68007 100644 --- a/tests/kernels/test_compressor_kv_cache.py +++ b/tests/kernels/test_compressor_kv_cache.py @@ -14,6 +14,7 @@ import math from types import SimpleNamespace +import numpy as np import pytest import torch @@ -157,6 +158,104 @@ def insert(): torch.testing.assert_close(kv, before, rtol=0, atol=0) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA only") +@pytest.mark.parametrize("kv_mxfp8", [False, True]) +def test_v41_pcp_swa_cache_matches_non_pcp_writer(kv_mxfp8, monkeypatch): + """PCP SWA writes must preserve the established CUDA cache bytes.""" + from vllm.models.deepseek_v41 import attention as attention_module + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + torch.manual_seed(43) + block_size = 64 + num_local_tokens = 5 + num_tokens = 2 * num_local_tokens + row_size = 528 if kv_mxfp8 else 584 + num_blocks = 2 + page_stride = math.ceil(block_size * row_size / 576) * 576 + 576 + backing = torch.full( + (num_blocks, page_stride), 165, dtype=torch.uint8, device="cuda" + ) + cache = backing.as_strided( + (num_blocks, block_size, row_size), (page_stride, row_size, 1) + ) + reference_backing = backing.clone() + reference = reference_backing.as_strided(cache.shape, cache.stride()) + + local_kv = torch.randn(num_local_tokens, 512, dtype=torch.bfloat16, device="cuda") + remote_kv = torch.randn_like(local_kv) + gathered_kv = torch.cat((local_kv, remote_kv)) + local_positions = torch.arange(num_local_tokens, device="cuda") + 3 + cache_positions = torch.cat((local_positions, local_positions + 11)) + cache_slots = torch.randperm( + num_blocks * block_size, dtype=torch.int64, device="cuda" + )[:num_tokens] + cache_slots[2] = -1 + angles = torch.randn(32, 32, device="cuda") + cos_sin = torch.cat((angles.cos(), angles.sin()), dim=-1) + num_heads = 8 + q = torch.randn( + num_local_tokens, + num_heads, + 512, + dtype=torch.bfloat16, + device="cuda", + ) + + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + torch.zeros(num_tokens, num_heads, 512, dtype=torch.bfloat16, device="cuda"), + gathered_kv, + reference.view(num_blocks, -1), + cache_slots, + cache_positions, + cos_sin, + num_heads, + 1e-20, + block_size, + False, + kv_mxfp8, + ) + + monkeypatch.setattr( + attention_module, + "get_pcp_group", + lambda: SimpleNamespace( + all_gather=lambda tensor, dim: torch.cat((tensor, remote_kv), dim=dim) + ), + ) + metadata = SimpleNamespace( + slot_mapping=torch.full( + (num_local_tokens,), -1, dtype=torch.int64, device="cuda" + ), + cache_slot_mapping=cache_slots, + cache_positions=cache_positions, + block_size=block_size, + ) + attention = SimpleNamespace( + use_pcp=True, + swa_cache_layer=SimpleNamespace( + prefix="swa", kv_cache=cache, block_size=block_size + ), + rotary_emb=SimpleNamespace(cos_sin_cache=cos_sin), + _flashinfer_fp8_kv_scale=None, + kv_mxfp8=kv_mxfp8, + accepts_unnormed_unroped_query=False, + n_local_heads=num_heads, + padded_heads=num_heads, + eps=1e-20, + head_dim=512, + ) + DeepseekV4Attention._fused_qnorm_rope_kv_insert( + attention, q, local_kv, local_positions, {"swa": metadata} + ) + + torch.testing.assert_close( + backing.view(torch.uint8), + reference_backing.view(torch.uint8), + rtol=0, + atol=0, + ) + + @pytest.mark.skipif( not current_platform.is_cuda_alike(), reason="graph capture coverage" ) @@ -796,6 +895,7 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): capacity = 8 vllm_config = MagicMock() vllm_config.scheduler_config.max_num_batched_tokens = 16 + vllm_config.parallel_config.prefill_context_parallel_size = 1 spec = CircularBufferSpec( block_size=capacity, num_kv_heads=1, @@ -832,6 +932,67 @@ def test_v41_compressor_metadata_maps_tokens_to_their_ring(): assert metadata.token_to_req_indices.tolist() == [0, 0, 0, 1, 1, 2, 2] +def test_v41_pcp_compressor_latent_is_gathered_for_cache_writes(monkeypatch): + """PCP compresses locally and publishes the gathered latent to both caches.""" + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v41 import attention as attention_module + from vllm.models.deepseek_v41.attention import DeepseekV4Attention + + local_latent = torch.arange(8, dtype=torch.float32).view(2, 4) + cache_positions = torch.tensor([0, 1, 10, 11]) + observed = {} + + class Compressor: + def __call__(self, kv_score, positions): + return local_latent + + def insert_cache(self, latent, positions, rotary_emb): + observed["insert"] = (latent, positions) + + def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale, **kwargs): + observed["indexer"] = (latent, kwargs) + return None, None, None + + monkeypatch.setattr( + attention_module, + "get_pcp_group", + lambda: SimpleNamespace( + all_gather=lambda tensor, dim: torch.cat((tensor, tensor + 10), dim=dim) + ), + ) + metadata = {"swa": SimpleNamespace(cache_positions=cache_positions)} + attention = SimpleNamespace( + use_pcp=True, + compressor=Compressor(), + indexer=prepare_indexer, + aux_stream_list=None, + ln_events=[None, None], + swa_cache_layer=SimpleNamespace(prefix="swa"), + rotary_emb=SimpleNamespace(), + indexer_rotary_emb=SimpleNamespace(), + n_local_heads=1, + head_dim=4, + _wq_b_proj=lambda qr, scale: qr, + _fused_qnorm_rope_kv_insert=lambda q, kv, pos, meta: q, + _sparse_indexer_and_attn=lambda *args: None, + ) + rows = torch.zeros(2, 4) + positions = torch.arange(2) + with override_forward_context(ForwardContext({}, metadata, {})): + DeepseekV4Attention._prepare_and_attn( + attention, rows, rows, rows, None, rows, rows, positions, rows + ) + + expected_global = torch.cat((local_latent, local_latent + 10)) + insert_latent, insert_positions = observed["insert"] + torch.testing.assert_close(insert_latent, expected_global) + torch.testing.assert_close(insert_positions, cache_positions) + indexer_latent, indexer_kwargs = observed["indexer"] + torch.testing.assert_close(indexer_latent, local_latent) + torch.testing.assert_close(indexer_kwargs["cache_latent"], expected_global) + torch.testing.assert_close(indexer_kwargs["cache_positions"], cache_positions) + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA stream coverage") @pytest.mark.parametrize( "use_aux,use_graph", [(False, False), (True, False), (True, True)] @@ -871,6 +1032,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph) context = ForwardContext({}, metadata, {}) compressor = DeepseekCompressor.__new__(DeepseekCompressor) torch.nn.Module.__init__(compressor) + compressor.pcp_size = 1 compressor.head_dim, compressor.rope_head_dim, compressor.compress_ratio = ( 512, 64, @@ -896,7 +1058,7 @@ def test_v41_attention_joins_cache_writes_before_consumption(use_aux, use_graph) use_fp4_kv=False, ) - def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale): + def prepare_indexer(qr, latent, weights, positions, rotary, qr_scale, **kwargs): DeepseekV4Indexer._produce_k(indexer, latent, positions, rotary) return None, None, None @@ -906,6 +1068,7 @@ def observe(*args): auxiliary = [torch.cuda.Stream()] if use_aux else None attention = SimpleNamespace( + use_pcp=False, compressor=compressor, indexer=prepare_indexer, aux_stream_list=auxiliary, @@ -963,6 +1126,168 @@ def reset(): previous = expected +def _build_pcp_ring_metadata(ranks, all_gather, spec, config, device, monkeypatch): + from vllm.models.deepseek_v41 import compressor as compressor_module + from vllm.models.deepseek_v41.compressor import CompressorMetadataBuilder + + monkeypatch.setattr( + compressor_module, + "get_pcp_group", + lambda: SimpleNamespace(all_gather=all_gather), + ) + for local in ranks: + builder = CompressorMetadataBuilder(spec, ["ring"], config, device) + local.metadata = builder.build(0, local.common) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Triton kernel") +@pytest.mark.parametrize("pcp_size", [2, 4]) +def test_v41_pcp_compressor_preserves_pairs_and_decode_state(pcp_size, monkeypatch): + """Pairs split across PCP chunks and decode owner changes keep exact latents.""" + from vllm.forward_context import ForwardContext, override_forward_context + from vllm.models.deepseek_v41 import compressor as compressor_module + from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( + fused_save_compress_norm, + ) + from vllm.models.deepseek_v41.compressor import DeepseekCompressor + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.kv_cache_interface import CircularBufferSpec + from vllm.v1.worker.gpu.pcp_manager import PCPManager + + torch.manual_seed(42) + device = torch.device("cuda") + norm = torch.randn(512, dtype=torch.bfloat16, device=device) + # Block 0 is reserved as the null block; requests use physical blocks 1/2. + reference_state = torch.zeros(3, 8, 1024, device=device) + rank_states = [reference_state.clone() for _ in range(pcp_size)] + config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=4), + parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size), + speculative_config=None, + ) + spec = CircularBufferSpec( + block_size=8, num_kv_heads=1, head_size=1024, head_size_v=0, dtype=torch.float32 + ) + starts = np.zeros(2, dtype=np.int32) + for lengths, prefilling in [ + ([7, 3], [True, True]), + ([9, 1], [True, False]), + ([1, 1], [False, False]), + ([1, 1], [False, False]), + ]: + lengths = np.array(lengths, dtype=np.int32) + query_start = np.concatenate(([0], np.cumsum(lengths))).astype(np.int32) + qsl = torch.from_numpy(query_start).to(device) + positions = torch.cat( + [torch.arange(s, s + n, device=device) for s, n in zip(starts, lengths)] + ) + req_ids = torch.repeat_interleave( + torch.arange(2, device=device, dtype=torch.int32), + torch.from_numpy(lengths).to(device), + ) + slots = (req_ids.long() + 1) * 8 + positions % 8 + raw = torch.randn(sum(lengths), 1024, device=device) + expected = torch.empty(sum(lengths), 512, dtype=torch.bfloat16, device=device) + fused_save_compress_norm( + raw, + positions, + reference_state, + slots, + qsl, + req_ids, + norm, + 1e-20, + 2, + expected, + ) + + ranks = [] + for rank in range(pcp_size): + manager = PCPManager(pcp_size, rank, device, shard_decode_requests=True) + segments, counts = manager._build_batch_layout( + lengths, starts, np.array(prefilling), query_start + ) + width = max(counts) + local_indices = manager._padded_gather_idx[ + rank * width : (rank + 1) * width + ] + local_lengths = [s.num_tokens for s in segments[rank]] or [0] + local_qsl = torch.tensor([0, *np.cumsum(local_lengths)], dtype=torch.int32) + local_reqs = [s.global_batch_req_idx for s in segments[rank]] or [0] + common = CommonAttentionMetadata( + query_start_loc=local_qsl.to(device), + query_start_loc_cpu=local_qsl, + seq_lens=torch.zeros(len(local_reqs), dtype=torch.int32, device=device), + num_reqs=len(local_reqs), + num_actual_tokens=counts[rank], + max_query_len=max(local_lengths), + max_seq_len=int(max(starts + lengths)), + block_table_tensor=torch.tensor( + np.array(local_reqs) + 1, dtype=torch.int32, device=device + ).unsqueeze(1), + slot_mapping=torch.full((width * pcp_size,), -1, device=device), + positions=positions[local_indices], + ) + ranks.append( + SimpleNamespace( + indices=local_indices[: counts[rank]], + raw=raw[local_indices], + positions=common.positions, + common=common, + ) + ) + + # Each rank's builder contributes its last-row candidates to a gather. + candidates: list[torch.Tensor] = [] + _build_pcp_ring_metadata( + ranks, + lambda t, dim, out=candidates: out.append(t.clone()) or t.repeat(pcp_size), + spec, + config, + device, + monkeypatch, + ) + _build_pcp_ring_metadata( + ranks, + lambda t, dim, out=candidates: torch.cat(out), + spec, + config, + device, + monkeypatch, + ) + gathered_rows = torch.cat([r.raw[r.metadata.gather_tokens] for r in ranks]) + monkeypatch.setattr( + compressor_module, + "get_pcp_group", + lambda rows=gathered_rows: SimpleNamespace(all_gather=lambda t, dim: rows), + ) + for rank, local in enumerate(ranks): + compressor = SimpleNamespace( + pcp_size=pcp_size, + head_dim=512, + compress_ratio=2, + rms_norm_eps=1e-20, + norm=SimpleNamespace(weight=norm), + state_cache=SimpleNamespace(prefix="ring", kv_cache=rank_states[rank]), + ) + context = ForwardContext({}, {"ring": local.metadata}, {}) + with override_forward_context(context): + actual = DeepseekCompressor.forward( + compressor, local.raw, local.positions + ) + assert isinstance(actual, torch.Tensor) + valid = (positions[local.indices] + 1) % 2 == 0 + torch.testing.assert_close( + actual[: local.indices.numel()][valid], + expected[local.indices][valid], + rtol=0, + atol=0, + ) + for state in rank_states[1:]: + torch.testing.assert_close(state, rank_states[0], rtol=0, atol=0) + starts += lengths + + def _on_gfx950() -> bool: if not current_platform.is_rocm(): return False diff --git a/tests/kernels/test_engram.py b/tests/kernels/test_engram.py index 15bdc9162cf6..1c9550d7c0e2 100644 --- a/tests/kernels/test_engram.py +++ b/tests/kernels/test_engram.py @@ -499,11 +499,12 @@ def fresh_cache(): @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA required") -def test_v2_model_state_gathers_lookback_window(): +@pytest.mark.parametrize("pcp", [False, True]) +def test_v2_model_state_gathers_lookback_window(pcp): """The window is gathered on device from the runner's token history.""" from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState - depth, max_num_reqs, max_model_len = 3, 4, 16 + depth, max_num_reqs, max_model_len = 3, 6, 16 state = DeepseekV41ModelState.__new__(DeepseekV41ModelState) state.rope_state = None state.supports_mm_inputs = False @@ -524,15 +525,23 @@ def test_v2_model_state_gathers_lookback_window(): # Batch rows -> request state rows: a request two decode steps in, a fresh # request, and one at the end of a 5-token prompt. input_batch = SimpleNamespace( - idx_mapping=torch.tensor([2, 0, 3], dtype=torch.int64, device="cuda") + idx_mapping=torch.tensor( + [2, 0, 3, 2] if pcp else [2, 0, 3], dtype=torch.int64, device="cuda" + ), + query_start_loc=torch.arange(5 if pcp else 4, dtype=torch.int32, device="cuda"), + positions=torch.tensor([8, 0, 6, 5] if pcp else [8, 0, 6], device="cuda"), ) window = state.prepare_inputs(input_batch, req_states)["lookback_token_ids"] - assert window.cpu().tolist() == [ + expected = [ [22, 21, 16], [-1, -1, -1], [0, 0, 0], - [-1, -1, -1], ] + if pcp: + # The second PCP chunk of request 2 starts at 5, not at its global 8. + expected.append([15, 14, 13]) + expected.extend([[-1] * depth] * (max_num_reqs - len(expected))) + assert window.cpu().tolist() == expected # Graph capture runs on the dummy inputs and replays read the same buffer. dummy = state.prepare_dummy_inputs(num_reqs=3, num_tokens=3)["lookback_token_ids"] assert dummy is window and torch.all(dummy == -1) 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 9cb7cfbb0f3b..b251b4e1c650 100644 --- a/vllm/model_executor/layers/fused_moe/runner/moe_runner.py +++ b/vllm/model_executor/layers/fused_moe/runner/moe_runner.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable, Iterable -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from typing import TYPE_CHECKING, cast import torch @@ -647,6 +647,7 @@ def _apply_quant_method( fused_out, ) + @contextmanager def _sequence_parallel_context(self): """Return a context manager for sequence-parallel token redistribution. @@ -656,7 +657,21 @@ def _sequence_parallel_context(self): returns a no-op context. """ ctx = get_forward_context() - return ( + local_is_padding = ctx.is_padding + if ( + local_is_padding is not None + and self.moe_config.pcp_size > 1 + and not self.moe_config.moe_parallel_config.use_all2all_kernels + ): + padding_key = "pcp_moe_is_padding" + if padding_key not in ctx.additional_kwargs: + ctx.additional_kwargs[padding_key] = ( + get_pcp_group() + .all_gather(local_is_padding.to(torch.uint8), dim=0) + .bool() + ) + ctx.is_padding = ctx.additional_kwargs[padding_key] + sp_context = ( ctx.dp_metadata.sp_local_sizes( self.moe_config.sp_size, pcp_size=self.moe_config.pcp_size, @@ -665,6 +680,11 @@ def _sequence_parallel_context(self): if ctx.dp_metadata else nullcontext() ) + try: + with sp_context: + yield + finally: + ctx.is_padding = local_is_padding def _maybe_add_zero_expert_output( self, @@ -803,7 +823,8 @@ def _maybe_dispatch( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: + input_ids: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: # For naive dispatch/combine Dp/Ep, dispatch the hidden states and # router logits to all experts. # NOTE: this will be removed once all kernels are migrated into the @@ -823,8 +844,22 @@ def _maybe_dispatch( ): 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 + if input_ids is not None: + uses_input_ids = ( + self.moe_config.has_hash_routing or self.router.bias_vl is not None + ) + if uses_input_ids: + ctx = get_forward_context() + input_ids_key = "pcp_moe_input_ids" + if input_ids_key not in ctx.additional_kwargs: + ctx.additional_kwargs[input_ids_key] = ( + get_pcp_group().all_gather(input_ids, dim=0) + ) + input_ids = ctx.additional_kwargs[input_ids_key] + else: + input_ids = None + + return hidden_states, router_logits, input_ids def _maybe_combine( self, @@ -913,9 +948,10 @@ def _forward_impl( # TODO(bnell): parts of the dispatch/combine steps will go away once # #32567 lands and the remaining kernels are made MKs. The PCP # code will probably remain - hidden_states, router_logits = self._maybe_dispatch( + hidden_states, router_logits, input_ids = self._maybe_dispatch( hidden_states, router_logits, + input_ids, ) shared_output, hidden_states = self._apply_quant_method( diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index cff72f12dc22..f8a9b7329a01 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -138,6 +138,10 @@ def __init__( def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: """Pick the ROCm sparse-MLA implementation for DeepSeek V4.1.""" + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + raise NotImplementedError( + "DeepSeek-V4.1 PCP currently supports NVIDIA CUDA only." + ) backend = vllm_config.attention_config.backend if backend is not None and ( backend is not AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4 diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 3aa54f9b0e97..6ce876974601 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -456,12 +456,16 @@ def _produce_k( latent: torch.Tensor | None, positions: torch.Tensor, rotary_emb: torch.nn.Module, + slot_mapping_indices: torch.Tensor | None = None, ) -> None: attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict) or latent is None: return assert self.owns_k indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) + slot_mapping = indexer_metadata.slot_mapping + if slot_mapping_indices is not None: + slot_mapping = slot_mapping[slot_mapping_indices] k_pre, _ = self.wk(latent) rocm_mxfp4_indexer_k_store( k_pre, @@ -470,7 +474,7 @@ def _produce_k( self.k_norm.weight, self.k_norm.variance_epsilon, self.k_cache.kv_cache, - indexer_metadata.slot_mapping, + slot_mapping, self.compress_ratio, self.use_fp4_kv, num_heads=self.n_head, diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 73d9c0c7233d..0be80decd8dc 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -48,7 +48,7 @@ get_current_vllm_config, ) from vllm.config.cache import CacheDType -from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.distributed import get_pcp_group, get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase @@ -288,6 +288,7 @@ def __init__( self.layer_id = layer_id self.prefix = prefix # Alias for compatibility with compressor + self.use_pcp = vllm_config.parallel_config.prefill_context_parallel_size > 1 self.hidden_size = config.hidden_size self.n_heads = config.num_attention_heads assert self.n_heads % tp_size == 0 @@ -835,7 +836,10 @@ def project_query_and_cache_kv() -> torch.Tensor: index_q_scale: torch.Tensor | None = None index_weights_out: torch.Tensor | None = None latent: torch.Tensor | None = None - aux_stream = aux_streams[0] if aux_streams is not None else None + # Keep PCP collectives on one stream in the same order on every rank. + aux_stream = ( + aux_streams[0] if aux_streams is not None and not self.use_pcp else None + ) if compressor is not None: # Q projection / KV insertion on the default stream overlaps the @@ -850,6 +854,17 @@ def project_query_and_cache_kv() -> torch.Tensor: else: q = project_query_and_cache_kv() + cache_latent = latent + cache_positions = positions + if self.use_pcp and latent is not None: + assert isinstance(attn_metadata, dict) + swa_metadata = cast( + "DeepseekSparseSWAMetadata", attn_metadata[self.swa_cache_layer.prefix] + ) + assert swa_metadata.cache_positions is not None + cache_latent = get_pcp_group().all_gather(latent, dim=0) + cache_positions = swa_metadata.cache_positions + def prepare_indexer(): if indexer is None: return None, None, None @@ -860,12 +875,16 @@ def prepare_indexer(): positions, self.indexer_rotary_emb, qr_scale, + cache_latent=cache_latent, + cache_positions=cache_positions, ) if compressor is not None: indexer_result, _ = maybe_execute_in_parallel( prepare_indexer, - lambda: compressor.insert_cache(latent, positions, self.rotary_emb), + lambda: compressor.insert_cache( + cache_latent, cache_positions, self.rotary_emb + ), self.ln_events[0], self.ln_events[1], aux_stream, @@ -1024,6 +1043,25 @@ def _fused_qnorm_rope_kv_insert( assert positions.dtype == torch.int64 cos_sin_cache = self.rotary_emb.cos_sin_cache cache_dtype = swa_kv_cache.dtype + slot_mapping = swa_metadata.slot_mapping + if self.use_pcp: + cache_kv = get_pcp_group().all_gather(kv.contiguous(), dim=0) + assert swa_metadata.cache_slot_mapping is not None + assert swa_metadata.cache_positions is not None + torch.ops._C.fused_deepseek_v4_kv_rope_insert( + cache_kv, + swa_kv_cache, + swa_metadata.cache_slot_mapping, + swa_metadata.cache_positions, + cos_sin_cache, + swa_metadata.block_size, + self._flashinfer_fp8_kv_scale + if cache_dtype == torch.float8_e4m3fn + else None, + self.kv_mxfp8, + ) + # The fused op now only prepares the rank-local queries. + slot_mapping = slot_mapping[:0] # kv is unchanged; attention reads kv solely via swa_kv_cache. if cache_dtype == torch.uint8: @@ -1045,7 +1083,7 @@ def _fused_qnorm_rope_kv_insert( q, kv, swa_kv_cache_2d, - swa_metadata.slot_mapping, + slot_mapping, positions, cos_sin_cache, pad_to, @@ -1073,7 +1111,7 @@ def _fused_qnorm_rope_kv_insert( q, kv, swa_kv_cache_3d, - swa_metadata.slot_mapping, + slot_mapping, positions, cos_sin_cache, self.eps, @@ -1089,7 +1127,7 @@ def _fused_qnorm_rope_kv_insert( kv, q_fp8, swa_kv_cache_3d, - swa_metadata.slot_mapping, + slot_mapping, positions, cos_sin_cache, self._flashinfer_fp8_kv_scale, @@ -1447,7 +1485,14 @@ def forward( positions: torch.Tensor, rotary_emb: nn.Module, qr_scale: torch.Tensor | None = None, + *, + cache_latent: torch.Tensor | None = None, + cache_positions: torch.Tensor | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: + if cache_latent is None: + cache_latent = latent + if cache_positions is None: + cache_positions = positions attn_metadata = get_forward_context().attn_metadata if isinstance(attn_metadata, dict): indexer_metadata = cast(Any, attn_metadata[self.k_cache.prefix]) @@ -1458,7 +1503,7 @@ def forward( # candidates num smaller than topk, every candidate is selected # but we still need to build k cache if self.owns_k: - self._produce_k(latent, positions, rotary_emb) + self._produce_k(cache_latent, cache_positions, rotary_emb) assert self.topk_indices_buffer is not None num_tokens = ( indexer_metadata.num_decode_tokens @@ -1478,7 +1523,7 @@ def forward( if self.owns_k: # K write must land before indexer_op reads the cache # (skip_k_cache_insert=True). - self._produce_k(latent, positions, rotary_emb) + self._produce_k(cache_latent, cache_positions, rotary_emb) return self.forward_q(qr, qr_scale, indexer_weights, positions, rotary_emb) diff --git a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py index 2ba9dbea3de7..4485430fb451 100644 --- a/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py +++ b/vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py @@ -30,6 +30,8 @@ def fused_save_compress_norm( rms_norm_eps: float, compress_ratio: int, latent_out: torch.Tensor, + prev_rows: torch.Tensor | None = None, + prev_row_indices: torch.Tensor | None = None, ) -> None: """Pool each closed group into a normalized BF16 latent; save FP32 states. @@ -58,6 +60,11 @@ def fused_save_compress_norm( rms_norm_eps: RMSNorm epsilon. compress_ratio: Group size, either 1 or 2. latent_out: BF16 [tokens, 512], written only at valid group boundaries. + prev_rows: FP32 [rows, 1024] rows gathered from all PCP ranks (CR2). A + chunk reads its predecessor here where ``prev_row_indices`` is set, + and the ring is left to ``save_ring_rows``. + prev_row_indices: [num_reqs] row in ``prev_rows`` holding the token + before each chunk's first token, or -1 to read the ring. """ assert compress_ratio in (1, 2) @@ -82,6 +89,11 @@ def fused_save_compress_norm( state_cache.shape[1], ) num_reqs = query_start_loc.numel() - 1 + if prev_rows is not None: + assert prev_row_indices is not None + assert prev_row_indices.numel() >= num_reqs + assert prev_rows.dtype == torch.float32 and prev_rows.shape[1] == 1024 + assert prev_rows.stride(1) == 1 and prev_rows.stride(0) % 16 == 0 else: state_cache = query_start_loc = token_to_req_indices = None state_stride = state_row_stride = state_block = 1 @@ -90,6 +102,8 @@ def fused_save_compress_norm( assert num_tokens <= min(kv_score.shape[0], positions.numel()) if num_tokens == 0: return + has_prev_rows = prev_rows is not None + prev_row_stride = prev_rows.stride(0) if prev_rows is not None else 1 grid = num_reqs + triton.cdiv(num_tokens, compress_ratio) _fused_save_compress_norm_kernel[(grid,)]( kv_score, @@ -100,15 +114,19 @@ def fused_save_compress_norm( token_to_req_indices, rms_norm_weight, latent_out, + prev_rows if has_prev_rows else kv_score, + prev_row_indices if has_prev_rows else slot_mapping, num_tokens, num_reqs, RAW_STRIDE=kv_score.stride(0), + PREV_ROW_STRIDE=prev_row_stride, STATE_STRIDE=state_stride, STATE_ROW_STRIDE=state_row_stride, STATE_BLOCK=state_block, COMPRESS_RATIO=compress_ratio, EPS=rms_norm_eps, JOIN_ROW_PTRS=_JOIN_ROW_PTRS, + HAS_PREV_ROWS=has_prev_rows, num_warps=4, **({"launch_pdl": False} if current_platform.is_cuda() else {}), ) @@ -124,15 +142,19 @@ def _fused_save_compress_norm_kernel( req_ids, norm_weight, latent, + prev_rows, + prev_row_indices, num_tokens, num_reqs, RAW_STRIDE: tl.constexpr, + PREV_ROW_STRIDE: tl.constexpr, STATE_STRIDE: tl.constexpr, STATE_ROW_STRIDE: tl.constexpr, STATE_BLOCK: tl.constexpr, COMPRESS_RATIO: tl.constexpr, EPS: tl.constexpr, JOIN_ROW_PTRS: tl.constexpr, + HAS_PREV_ROWS: tl.constexpr, ): pid = tl.program_id(0) d = tl.arange(0, 512) @@ -152,6 +174,10 @@ def _fused_save_compress_norm_kernel( if (position + 1) % 2 == 0: ring = state + (state_slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE prev = ring + ((position - 1) % STATE_BLOCK) * STATE_ROW_STRIDE + if HAS_PREV_ROWS: + prev_idx = tl.load(prev_row_indices + pid) + if prev_idx >= 0: + prev = prev_rows + prev_idx.to(tl.int64) * PREV_ROW_STRIDE current = raw + start.to(tl.int64) * RAW_STRIDE if JOIN_ROW_PTRS: # Joining pointers hides row alignment from Triton; restate it. @@ -177,19 +203,20 @@ def _fused_save_compress_norm_kernel( weight_prev + weight_current ) _store_latent(pooled, start, norm_weight, latent, EPS) - num_rows = tl.minimum(end - start, STATE_BLOCK) - for k in tl.range(0, num_rows): - token = end - num_rows + k - slot = tl.load(state_slots + token) - if slot >= 0: - row = ( - state - + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE - + (slot % STATE_BLOCK) * STATE_ROW_STRIDE - ) - src = raw + token.to(tl.int64) * RAW_STRIDE - tl.store(row + d, tl.load(src + d)) - tl.store(row + 512 + d, tl.load(src + 512 + d)) + if not HAS_PREV_ROWS: + num_rows = tl.minimum(end - start, STATE_BLOCK) + for k in tl.range(0, num_rows): + token = end - num_rows + k + slot = tl.load(state_slots + token) + if slot >= 0: + row = ( + state + + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE + + (slot % STATE_BLOCK) * STATE_ROW_STRIDE + ) + src = raw + token.to(tl.int64) * RAW_STRIDE + tl.store(row + d, tl.load(src + d)) + tl.store(row + 512 + d, tl.load(src + 512 + d)) return # Group program: one token (ratio 1) or the token of the pair that ends @@ -212,6 +239,50 @@ def _fused_save_compress_norm_kernel( _store_latent(pooled, t, norm_weight, latent, EPS) +def save_ring_rows( + rows: torch.Tensor, slot_mapping: torch.Tensor, state_cache: torch.Tensor +) -> None: + """Store FP32 [kv, score] rows to their ring slots, skipping slot -1.""" + assert rows.dtype == torch.float32 and rows.shape[1] == 1024 + assert rows.stride(1) == 1 and state_cache.stride(2) == 1 + assert slot_mapping.numel() == rows.shape[0] + if rows.shape[0] == 0: + return + _save_ring_rows_kernel[(rows.shape[0],)]( + rows, + slot_mapping, + state_cache, + ROW_STRIDE=rows.stride(0), + STATE_STRIDE=state_cache.stride(0), + STATE_ROW_STRIDE=state_cache.stride(1), + STATE_BLOCK=state_cache.shape[1], + num_warps=4, + ) + + +@triton.jit +def _save_ring_rows_kernel( + rows, + slot_mapping, + state, + ROW_STRIDE: tl.constexpr, + STATE_STRIDE: tl.constexpr, + STATE_ROW_STRIDE: tl.constexpr, + STATE_BLOCK: tl.constexpr, +): + pid = tl.program_id(0) + slot = tl.load(slot_mapping + pid) + if slot < 0: + return + d = tl.arange(0, 1024) + dst = ( + state + + (slot // STATE_BLOCK).to(tl.int64) * STATE_STRIDE + + (slot % STATE_BLOCK) * STATE_ROW_STRIDE + ) + tl.store(dst + d, tl.load(rows + pid.to(tl.int64) * ROW_STRIDE + d)) + + @triton.jit def _store_latent(pooled, t, norm_weight, latent, EPS: tl.constexpr): d = tl.arange(0, 512) diff --git a/vllm/models/deepseek_v41/compressor.py b/vllm/models/deepseek_v41/compressor.py index 6f56ad8d5b2d..324123fb54fe 100644 --- a/vllm/models/deepseek_v41/compressor.py +++ b/vllm/models/deepseek_v41/compressor.py @@ -8,6 +8,7 @@ from torch import nn from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import get_pcp_group from vllm.forward_context import get_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -15,6 +16,7 @@ from vllm.models.deepseek_v41.common.ops.fused_compress_quant_cache import ( fused_save_compress_norm, rope_quant_insert, + save_ring_rows, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -36,6 +38,10 @@ def __init__(self): def get_name() -> str: return "CompressorBackend" + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: return [MultipleOf(1)] @@ -55,6 +61,12 @@ class CompressorMetadata: slot_mapping: torch.Tensor query_start_loc: torch.Tensor # [num_reqs + 1] token_to_req_indices: torch.Tensor # [num_tokens] + # PCP only: every rank gathers the last row of each rank's local requests. + gather_tokens: torch.Tensor | None = None # [rows] this rank's rows + # [pcp * rows] ring slot of each request's newest row if it is open, else -1. + ring_write_slots: torch.Tensor | None = None + # [num_reqs] gathered row preceding each local chunk, or -1 to use the ring. + prev_row_indices: torch.Tensor | None = None @triton.jit(do_not_specialize=["block_table_stride", "num_actual_tokens", "num_tokens"]) @@ -91,6 +103,7 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) assert isinstance(self.kv_cache_spec, CircularBufferSpec) self.capacity = self.kv_cache_spec.block_size + self.pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size max_num_batched_tokens = ( self.vllm_config.scheduler_config.max_num_batched_tokens ) @@ -100,6 +113,8 @@ def __init__(self, *args, **kwargs): self.slot_mapping_buffer = torch.empty( max_num_batched_tokens, dtype=torch.int64, device=self.device ) + # The ring keeps only the open row under PCP, too few for spec decoding. + assert self.pcp_size == 1 or self.vllm_config.speculative_config is None def build( self, @@ -107,7 +122,7 @@ def build( common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CompressorMetadata: - num_tokens = common_attn_metadata.slot_mapping.numel() + num_tokens = common_attn_metadata.slot_mapping.numel() // self.pcp_size positions = common_attn_metadata.positions assert positions is not None token_to_req_indices = common_attn_metadata.token_to_req_indices( @@ -126,11 +141,66 @@ def build( CAPACITY=self.capacity, BLOCK=256, ) - return CompressorMetadata( + metadata = CompressorMetadata( slot_mapping=slot_mapping, query_start_loc=common_attn_metadata.query_start_loc, token_to_req_indices=token_to_req_indices, ) + if self.pcp_size > 1: + self._link_pcp_chunks(metadata, common_attn_metadata.num_reqs, positions) + return metadata + + def _link_pcp_chunks( + self, metadata: CompressorMetadata, num_reqs: int, positions: torch.Tensor + ) -> None: + """Link ratio-2 groups split across PCP chunks and pick the ring writes. + + Every rank gathers ``(ring block, position)`` of each local request's + last token; the ring block names a request on every rank. Sorted, they + give each local chunk's predecessor row and each request's newest row. + """ + num_tokens = metadata.slot_mapping.numel() + # Equal on every rank, so the gathers line up. + num_rows = min(num_tokens, 2 * self.vllm_config.scheduler_config.max_num_seqs) + num_reqs = min(num_reqs, num_rows) + query_start_loc = metadata.query_start_loc.long() + starts = query_start_loc[:num_reqs].clamp(max=num_tokens - 1) + last = (query_start_loc[1 : num_reqs + 1] - 1).clamp(min=0) + nonempty = query_start_loc[1 : num_reqs + 1] > query_start_loc[:num_reqs] + + def key(slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor: + return torch.where(slots >= 0, (slots // self.capacity) << 32 | pos, -1) + + local_keys = torch.full((num_rows,), -1, dtype=torch.int64, device=self.device) + local_keys[:num_reqs] = key( + torch.where(nonempty, metadata.slot_mapping[last], -1), positions[last] + ) + sorted_keys, order = get_pcp_group().all_gather(local_keys, dim=0).sort() + + blocks, sorted_positions = sorted_keys >> 32, sorted_keys & 0xFFFFFFFF + newest = torch.ones_like(sorted_keys, dtype=torch.bool) + newest[:-1] = blocks[:-1] != blocks[1:] + is_open = (sorted_keys >= 0) & newest & (sorted_positions % 2 == 0) + ring_write_slots = torch.empty_like(sorted_keys) + ring_write_slots[order] = torch.where( + is_open, blocks * self.capacity + sorted_positions % self.capacity, -1 + ) + + first_positions = positions[starts] + wanted = key(metadata.slot_mapping[starts], first_positions - 1) + found = torch.searchsorted(sorted_keys, wanted).clamp_(max=order.numel() - 1) + has_prev = ( + nonempty + & (wanted >= 0) + & (first_positions % 2 == 1) + & (sorted_keys[found] == wanted) + ) + + gather_tokens = torch.zeros_like(local_keys) + gather_tokens[:num_reqs] = last + metadata.gather_tokens = gather_tokens + metadata.ring_write_slots = ring_write_slots + metadata.prev_row_indices = torch.where(has_prev, order[found], -1) class CompressorStateCache(torch.nn.Module, AttentionLayerBase): @@ -205,6 +275,8 @@ def __init__( self.rotate = rotate self.prefix = prefix self.k_cache_prefix = k_cache_prefix + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 # Ratio 1 pools single tokens, so the checkpoint carries no gate. self.has_gate = compress_ratio > 1 @@ -256,15 +328,20 @@ def forward( ) -> torch.Tensor | None: """Save states and return the BF16 latent for cache insertion and indexing. - Only valid group-boundary rows are written. + Only valid group-boundary rows are written. Under PCP each rank pools + its own rows; the last row of every rank's chunks is gathered, both to + close groups split across chunks and to keep every rank's ring equal. """ attn_metadata = get_forward_context().attn_metadata if not isinstance(attn_metadata, dict): return None + prev_rows = prev_row_indices = ring_write_slots = None if self.state_cache is None: state_cache = query_start_loc = token_to_req_indices = None slot_mapping = cast(Any, attn_metadata[self.k_cache_prefix]).slot_mapping + if self.pcp_size > 1: + slot_mapping = slot_mapping.chunk(self.pcp_size)[self.pcp_rank] else: state_metadata = cast( CompressorMetadata, attn_metadata[self.state_cache.prefix] @@ -273,6 +350,13 @@ def forward( slot_mapping = state_metadata.slot_mapping query_start_loc = state_metadata.query_start_loc token_to_req_indices = state_metadata.token_to_req_indices + if self.pcp_size > 1: + assert state_metadata.gather_tokens is not None + prev_rows = get_pcp_group().all_gather( + kv_score[state_metadata.gather_tokens], dim=0 + ) + prev_row_indices = state_metadata.prev_row_indices + ring_write_slots = state_metadata.ring_write_slots latent = torch.empty( kv_score.shape[0], @@ -291,7 +375,12 @@ def forward( self.rms_norm_eps, self.compress_ratio, latent, + prev_rows=prev_rows, + prev_row_indices=prev_row_indices, ) + if prev_rows is not None: + assert state_cache is not None and ring_write_slots is not None + save_ring_rows(prev_rows, ring_write_slots, state_cache) return latent def insert_cache( diff --git a/vllm/models/deepseek_v41/nvidia/flashmla.py b/vllm/models/deepseek_v41/nvidia/flashmla.py index ce3047ceca2b..d995058216d0 100644 --- a/vllm/models/deepseek_v41/nvidia/flashmla.py +++ b/vllm/models/deepseek_v41/nvidia/flashmla.py @@ -42,6 +42,10 @@ class DeepseekSparseSWAFlashMLAMetadataBuilder(DeepseekV41SparseSWAMetadataBuild class DeepseekSparseSWAFlashMLABackend(DeepseekSparseSWABackend): + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_builder_cls() -> type[DeepseekSparseSWAFlashMLAMetadataBuilder]: return DeepseekSparseSWAFlashMLAMetadataBuilder @@ -132,6 +136,14 @@ def forward_mqa( ) assert swa_metadata is not None + if self.use_pcp: + num_tokens = ( + swa_metadata.num_decode_tokens + swa_metadata.num_prefill_tokens + ) + output[num_tokens:].zero_() + if num_tokens == 0: + return + swa_only = self.compress_ratio == 0 # SWA-only layers (compress_ratio == 0) don't have their own KV cache # allocation; consumers read the kv source's cache, which is bound on diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 94fa0b4b8b67..5e3baed48c1b 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -203,6 +203,17 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: FlashMLA path. """ backend = vllm_config.attention_config.backend + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + if backend not in ( + None, + AttentionBackendEnum.FLASHMLA_SPARSE, + AttentionBackendEnum.FLASHMLA_SPARSE_DSV4, + AttentionBackendEnum.FLASHMLA_SPARSE_DSV41, + ): + raise NotImplementedError( + "DeepSeek-V4.1 PCP requires FlashMLA sparse attention." + ) + return DeepseekV4FlashMLAAttention device_capability = current_platform.get_device_capability() if backend in ( AttentionBackendEnum.FLASHINFER_MLA_SPARSE, @@ -633,6 +644,21 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.config = config self.quant_config = quant_config self.parallel_config = vllm_config.parallel_config + if self.parallel_config.prefill_context_parallel_size > 1 and ( + self.parallel_config.pipeline_parallel_size != 1 + or self.parallel_config.data_parallel_size != 1 + or self.parallel_config.decode_context_parallel_size != 1 + or self.parallel_config.use_ubatching + or vllm_config.speculative_config is not None + or vllm_config.kv_transfer_config is not None + or vllm_config.cache_config.enable_prefix_caching + or vllm_config.kernel_config.moe_backend in MEGA_MOE_BACKENDS + ): + raise NotImplementedError( + "DeepSeek-V4.1 PCP currently requires PP=DP=DCP=1, with no " + "microbatching, speculative decoding, KV transfer, prefix " + "caching, or MegaMoE." + ) self.use_native_mega_moe = ( vllm_config.kernel_config.moe_backend in NATIVE_MEGA_MOE_BACKENDS ) diff --git a/vllm/models/deepseek_v41/nvidia/model_state.py b/vllm/models/deepseek_v41/nvidia/model_state.py index c09e4582ea34..40905637361e 100644 --- a/vllm/models/deepseek_v41/nvidia/model_state.py +++ b/vllm/models/deepseek_v41/nvidia/model_state.py @@ -41,7 +41,8 @@ def _gather_lookback_kernel( lookback_ptr, idx_mapping_ptr, - num_computed_tokens_ptr, + query_start_loc_ptr, + positions_ptr, all_token_ids_ptr, all_token_ids_stride, num_reqs, @@ -52,13 +53,16 @@ def _gather_lookback_kernel( batch_idx = tl.program_id(0) in_batch = batch_idx < num_reqs req_state_idx = tl.load(idx_mapping_ptr + batch_idx, mask=in_batch, other=0) - num_computed = tl.load(num_computed_tokens_ptr + req_state_idx) + begin = tl.load(query_start_loc_ptr + batch_idx, mask=in_batch, other=0) + end = tl.load(query_start_loc_ptr + batch_idx + 1, mask=in_batch, other=0) + in_batch = in_batch & (end > begin) + num_computed = tl.load(positions_ptr + begin, mask=in_batch, other=0) offs = tl.arange(0, BLOCK_DEPTH) pos = num_computed - 1 - offs valid = in_batch & (offs < DEPTH) & (pos >= 0) ids = tl.load( - all_token_ids_ptr + req_state_idx * all_token_ids_stride + pos, + all_token_ids_ptr + req_state_idx.to(tl.int64) * all_token_ids_stride + pos, mask=valid, other=-1, ) @@ -200,8 +204,11 @@ def __init__( self.lookback_token_ids: torch.Tensor | None = None if depth > 0: # Persistent so a captured graph can read it on replay. + max_num_reqs = self.max_num_reqs + if vllm_config.parallel_config.prefill_context_parallel_size > 1: + max_num_reqs *= 2 self.lookback_token_ids = torch.full( - (self.max_num_reqs, depth), -1, dtype=torch.int32, device=device + (max_num_reqs, depth), -1, dtype=torch.int32, device=device ) # Per request state index; batches gather from it (see prepare_attn). @@ -286,7 +293,8 @@ def prepare_inputs( _gather_lookback_kernel[(window.shape[0],)]( window, input_batch.idx_mapping, - req_states.num_computed_tokens.gpu, + input_batch.query_start_loc, + input_batch.positions, all_token_ids, all_token_ids.stride(0), input_batch.idx_mapping.shape[0], diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 1130eec0903b..3a3c26f438f0 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/vllm/models/deepseek_v41/sparse_mla.py @@ -9,6 +9,7 @@ from vllm.config import VllmConfig from vllm.config.cache import CacheDType +from vllm.distributed import get_pcp_group from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backend import ( @@ -24,6 +25,7 @@ _LAYER_TYPE_C1A, _LAYER_TYPE_C2A, _LAYER_TYPE_SWAONLY, + DeepseekSparseSWAMetadata, DeepseekSparseSWAMetadataBuilder, ) from vllm.v1.kv_cache_interface import AttentionSpec @@ -60,6 +62,12 @@ class DeepseekV41SparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) + self.pcp_size = self.vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 + if self.pcp_size > 1: + self.no_replay_start = self.no_replay_start.new_zeros( + 2 * self.vllm_config.scheduler_config.max_num_seqs + 1 + ) compress_ratios = getattr( self.vllm_config.model_config.hf_config, "compress_ratios", None ) or [0] @@ -67,6 +75,24 @@ def __init__(self, *args, **kwargs): deepseek_v41_layer_type(int(ratio)) for ratio in compress_ratios } + def build( + self, + common_prefix_len: int, + common_attn_metadata: CommonAttentionMetadata, + fast_build: bool = False, + replay_start: torch.Tensor | None = None, + ) -> DeepseekSparseSWAMetadata: + cm = common_attn_metadata + if self.pcp_size > 1: + local_slots = cm.slot_mapping.chunk(self.pcp_size)[self.pcp_rank] + cm = cm.replace(slot_mapping=local_slots) + metadata = super().build(common_prefix_len, cm, fast_build, replay_start) + if self.pcp_size > 1: + metadata.cache_slot_mapping = common_attn_metadata.slot_mapping + assert cm.positions is not None + metadata.cache_positions = get_pcp_group().all_gather(cm.positions, dim=0) + return metadata + class DeepseekV4SparseMLABackend(AttentionBackend): """DeepSeek-V4.1 sparse-MLA backend base. @@ -155,6 +181,8 @@ def __init__( ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) self.model_config = vllm_config.model_config + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 # Classify single-token queries (plus num_speculative_tokens via # supports_spec_as_decode=True) as decodes; longer queries go to prefill. self._init_reorder_batch_threshold(1, supports_spec_as_decode=True) @@ -193,9 +221,12 @@ def build( slot_mapping = cm.slot_mapping if self.compress_ratio > 1: + local_slots = slot_mapping + if self.pcp_size > 1: + local_slots = slot_mapping.chunk(self.pcp_size)[self.pcp_rank] slot_mapping = get_compressed_slot_mapping( cm.num_actual_tokens, - cm.slot_mapping, + local_slots, cm.query_start_loc, cm.seq_lens, cm.block_table_tensor.clamp_(min=0), @@ -203,6 +234,10 @@ def build( self.compress_ratio, out=self.compressed_slot_mapping_buffer, ) + if self.pcp_size > 1: + slot_mapping = get_pcp_group().all_gather( + self.compressed_slot_mapping_buffer[: local_slots.numel()], dim=0 + ) return DeepseekV4FlashMLAMetadata( num_reqs=cm.num_reqs, @@ -223,6 +258,10 @@ class DeepseekV4FlashMLAMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder): class DeepseekV4FlashMLABackend(DeepseekV4SparseMLABackend): + @classmethod + def supports_pcp(cls) -> bool: + return True + @staticmethod def get_name() -> str: return "FLASHMLA_SPARSE_DSV41" @@ -250,6 +289,10 @@ class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend): "nvfp4_ds_mla", # V4.1 fp8 SWA cache + NVFP4 compressed cache ] + @classmethod + def supports_pcp(cls) -> bool: + return False + @staticmethod def get_name() -> str: return "FLASHMLA_MEGA_ATTN_DSV41" diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index c5cbb4ef2822..840b8a6d6ff3 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -209,6 +209,9 @@ class DeepseekSparseSWAMetadata: # window attention may read (SWA bounded replay); # zeros when nothing replays. replay_start: torch.Tensor | None = None + # PCP cache writes use the gathered slots; queries use rank-local slots. + cache_slot_mapping: torch.Tensor | None = None + cache_positions: torch.Tensor | None = None # Number of decode/prefill requests/tokens (batch is reordered: decodes first) num_decodes: int = 0