iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[PCP][DCP] Enable PCP+DCP on sparse-MLA models (#56157)

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