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

[Bugfix][PCP] Align MoE routing padding with gathered tokens (#60946)

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: wangyicong <wangyicong@bytedance.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
This commit is contained in:
Yicong Wang
2026-10-14 21:04:00 -04:00
committed by GitHub
co-authored by Lucas Wilkinson Claude Opus 5.5 OpenAI Codex
parent 8530aed2b7
commit bfdbfa1632
7 changed files with 180 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/cudagraph/test_cudagraph_manager.py
+1
View File
@@ -132,6 +132,7 @@ def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
pcp_block_tables = SimpleNamespace(
num_kv_cache_groups=1,
input_block_tables=(torch.zeros(num_reqs * 2, 1, dtype=torch.int32),),
cp_interleave=1,
)
pcp_manager = PCPManager(
pcp_world_size=pcp_world_size,
tests/v1/worker/test_gpu_pcp_manager.py
+63 -5
View File
@@ -7,7 +7,9 @@ import numpy as np
import pytest
import torch
from vllm.config import CUDAGraphMode, ParallelConfig
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
from vllm.forward_context import set_forward_context
from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk
from vllm.v1.attention.backends.utils import PAD_SLOT_ID, get_dcp_local_seq_lens
from vllm.v1.attention.ops import pcp as attention_pcp
from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
@@ -48,6 +50,7 @@ def _make_capture_manager(block_table: torch.Tensor):
input_block_tables=(block_table,),
num_kv_cache_groups=1,
block_sizes=(2,),
cp_interleave=1,
)
manager = PCPManager(
pcp_world_size=2,
@@ -240,9 +243,10 @@ def test_dummy_draft_does_not_reuse_previous_graph_batch():
manager.prepare_draft_prefill(dummy_batch, dummy_batch.input_ids)
assert (
manager.get_draft_input_buffers(manager.input_buffers) is manager.input_buffers
draft_inputs, _ = manager.get_draft_input_buffers(
manager.input_buffers, dummy_batch.num_tokens_after_padding
)
assert draft_inputs is manager.input_buffers
def test_capture_uses_pcp_persistent_inputs():
@@ -256,7 +260,7 @@ def test_capture_uses_pcp_persistent_inputs():
)
dummy_batch.input_ids.copy_(torch.arange(4))
dummy_batch.positions.fill_(7)
input_batch = manager.prepare_inputs_to_capture(dummy_batch)
input_batch = manager.prepare_inputs_for_capture_or_profile(dummy_batch)
assert input_batch is not dummy_batch
assert input_batch.req_ids == dummy_batch.req_ids
@@ -266,6 +270,17 @@ def test_capture_uses_pcp_persistent_inputs():
assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr()
def test_block_table_rebind_updates_interleave():
manager, block_tables = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
final_block_tables = SimpleNamespace(**vars(block_tables))
final_block_tables.cp_interleave = 64
manager.set_block_tables(final_block_tables)
assert manager._block_tables is final_block_tables
assert manager.cp_interleave == 64
def test_dummy_context_updates_pcp_local_block_tables():
global_block_table = torch.full((4, 4), -1, dtype=torch.int32)
manager, block_tables = _make_capture_manager(global_block_table)
@@ -275,7 +290,7 @@ def test_dummy_context_updates_pcp_local_block_tables():
input_buffers=manager.input_buffers,
max_query_len=1,
)
input_batch = manager.prepare_inputs_to_capture(input_batch)
input_batch = manager.prepare_inputs_for_capture_or_profile(input_batch)
local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs)
set_dummy_context(
@@ -736,6 +751,49 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
def test_moe_padding_matches_pcp_all_gather(monkeypatch):
"""The PCP batch carries the all-gather of every rank's local padding mask,
the rows the MoE routes, even when ranks own different numbers of tokens."""
monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1")
device = torch.device("cuda:0")
local_masks, batch_masks = [], []
for pcp_rank in range(2):
manager = PCPManager(
pcp_world_size=2,
pcp_rank=pcp_rank,
device=device,
shard_decode_requests=True,
max_num_reqs=4,
max_num_tokens=8,
moe_gathers_tokens=True,
)
global_batch = _make_global_decode_batch(
[16, 24, 32], InputBuffers(4, 8, device), device
)
local_batch = manager.partition_batch(
global_batch,
BatchExecutionDescriptor(
cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=4, num_reqs=None
),
)
num_tokens = local_batch.num_tokens_after_padding
local_masks.append(manager.input_buffers.is_padding[:num_tokens].cpu())
batch_masks.append(local_batch.is_padding.clone())
# Sharded decode puts two requests on rank 0 and one on rank 1.
assert [mask.logical_not().sum().item() for mask in local_masks] == [2, 1]
for batch_mask in batch_masks:
assert torch.equal(batch_mask.cpu(), torch.cat(local_masks))
num_tokens = batch_mask.numel()
hidden_states = torch.randn(num_tokens, 4, device=device)
router_logits = torch.randn(num_tokens, 8, device=device)
with set_forward_context(None, VllmConfig(), is_padding=batch_mask):
_, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False)
assert (topk_ids[~batch_mask] >= 0).all()
assert (topk_ids[batch_mask] == -1).all()
@pytest.mark.parametrize(
("pcp_world_size", "dcp_world_size", "expected"),
[(1, 1, False), (2, 1, True), (2, 2, False)],
tests/v1/worker/test_gpu_target_dependent_ar_speculator.py
+47
View File
@@ -22,6 +22,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
from vllm.v1.attention.backends.utils import split_decodes_and_prefills
from vllm.v1.kv_cache_interface import MLAAttentionSpec
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
from vllm.v1.worker.gpu.pcp_manager import PCPManager
from vllm.v1.worker.gpu.spec_decode import speculator as base_spec_module
from vllm.v1.worker.gpu.spec_decode.multi_module_mtp.speculator import (
MultiModuleMTPSpeculator,
@@ -397,6 +398,52 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):
assert actual_feedback_hidden is hidden
@pytest.mark.parametrize("partitioned", [False, True])
def test_run_model_uses_complete_pcp_padding(monkeypatch, partitioned):
hidden = torch.full((4, 3), 1.0)
speculator = _make_speculator(monkeypatch, hidden)
gathered_is_padding = torch.tensor(
[False, False, False, True, False, False, True, True]
)
manager = PCPManager(
pcp_world_size=2,
pcp_rank=0,
device=torch.device("cpu"),
shard_decode_requests=True,
max_num_tokens=4,
moe_gathers_tokens=True,
)
if partitioned:
manager.draft_prefill_batch = SimpleNamespace(
input_ids=torch.arange(4),
positions=torch.arange(4),
is_padding=gathered_is_padding,
)
else:
local_is_padding = torch.tensor([False, False, False, True])
speculator.input_buffers.is_padding = local_is_padding
gathered_is_padding = torch.cat([local_is_padding, local_is_padding])
seen = {}
monkeypatch.setattr(
spec_module,
"set_forward_context",
lambda *args, **kwargs: (seen.update(kwargs), nullcontext())[1],
)
speculator.pcp_manager = manager
speculator._run_model(
4,
attn_metadata=None,
slot_mappings=None,
num_tokens_across_dp=None,
cudagraph_runtime_mode=CUDAGraphMode.NONE,
)
torch.testing.assert_close(seen["is_padding"], gathered_is_padding)
if partitioned:
assert seen["is_padding"] is gathered_is_padding
@pytest.mark.parametrize(
(
"method_name",
vllm/v1/worker/gpu/cudagraph_utils.py
+5 -2
View File
@@ -743,6 +743,9 @@ class ModelCudaGraphManager(CudaGraphManager):
# Capture with dummy rows marked as padding.
input_buffers.is_padding.fill_(True)
is_padding = input_buffers.is_padding[:num_tokens]
if pcp_manager is not None:
is_padding = pcp_manager.moe_is_padding(is_padding)
def forward_fn(cg_mode: CUDAGraphMode) -> None:
batch_descriptor = None
@@ -760,7 +763,7 @@ class ModelCudaGraphManager(CudaGraphManager):
num_tokens_across_dp=num_tokens_across_dp,
slot_mapping=slot_mappings,
batch_descriptor=batch_descriptor,
is_padding=input_buffers.is_padding[:num_tokens],
is_padding=is_padding,
):
if cg_mode == CUDAGraphMode.PIECEWISE:
# PIECEWISE graph (compiled PW or breakable, chosen inside
@@ -837,7 +840,7 @@ def prepare_inputs_to_capture(
num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
)
if pcp_manager is not None:
input_batch = pcp_manager.prepare_inputs_to_capture(input_batch)
input_batch = pcp_manager.prepare_inputs_for_capture_or_profile(input_batch)
block_table_provider = pcp_manager or block_tables
input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs)
vllm/v1/worker/gpu/model_runner.py
+12 -9
View File
@@ -305,7 +305,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# Draft tokens propagation - for spec-dec + struct outputs.
self.draft_tokens_handler = DraftTokensHandler(self.device)
self.pcp_manager: pcp.PCPManager | None = None
# Built before memory profiling; block tables are attached at KV init.
self.pcp_manager = pcp.maybe_build_pcp_manager(
self.vllm_config,
self.device,
self.supports_mm_inputs,
cls=self.pcp_manager_cls,
)
# Pooling models.
self.is_pooling_model = self.model_config.runner_type == "pooling"
@@ -700,13 +706,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
cp_rank=self.dcp_rank,
cp_interleave=self.cp_interleave,
)
self.pcp_manager = pcp.maybe_build_pcp_manager(
self.vllm_config,
self.device,
self.supports_mm_inputs,
self.block_tables,
cls=self.pcp_manager_cls,
)
if self.pcp_manager is not None:
self.pcp_manager.set_block_tables(self.block_tables)
self.ubatch_runner = maybe_build_ubatch_runner(
self.vllm_config,
self.device,
@@ -1848,7 +1849,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# All-zero input_ids route every token to the same experts.
input_batch.input_ids.random_(0, self.vocab_size)
if self.pcp_manager is not None:
input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
input_batch = self.pcp_manager.prepare_inputs_for_capture_or_profile(
input_batch
)
if skip_attn_for_dummy_run:
assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
"Attention metadata must be prepared for dummy runs when using "
vllm/v1/worker/gpu/pcp_manager.py
+49 -8
View File
@@ -56,6 +56,7 @@ class PCPManager:
dcp_world_size: int = 1,
dcp_rank: int = 0,
cp_interleave: int = 1,
moe_gathers_tokens: bool = False,
) -> None:
self.pcp_world_size = pcp_world_size
self.pcp_rank = pcp_rank
@@ -69,11 +70,16 @@ class PCPManager:
self._local_batch: InputBatch | None = None
self._local_gather_idx: torch.Tensor | None = None
self.draft_prefill_batch: InputBatch | None = None
self._block_tables = block_tables
self._hidden_restore_idx: torch.Tensor | None = None
self._padded_gather_idx: torch.Tensor | None = None
self._gathered_kv_write_mask: torch.Tensor | None = None
self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device)
# Padding mask over the PCP-gathered rows the MoE routes (rank-major).
self._moe_is_padding = (
torch.ones(max_num_tokens * pcp_world_size, dtype=torch.bool, device=device)
if moe_gathers_tokens and max_num_tokens is not None
else None
)
max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None
self._input_buffers = (
@@ -81,6 +87,18 @@ class PCPManager:
if max_num_local_reqs is not None and max_num_tokens is not None
else None
)
self._max_num_local_reqs = max_num_local_reqs
self._max_num_tokens = max_num_tokens
self.set_block_tables(block_tables)
def set_block_tables(self, block_tables: BlockTables | None) -> None:
"""Attach the KV block tables, which only exist after KV cache init."""
self._block_tables = block_tables
if block_tables is not None:
self.cp_interleave = block_tables.cp_interleave
max_num_local_reqs = self._max_num_local_reqs
max_num_tokens = self._max_num_tokens
device, pcp_world_size = self.device, self.pcp_world_size
self._local_block_tables: tuple[torch.Tensor, ...] | None
self._local_block_table_ptrs: torch.Tensor | None
if block_tables is not None and max_num_local_reqs is not None:
@@ -593,6 +611,12 @@ class PCPManager:
input_buffers.positions[:num_local_tokens_padded].masked_fill_(
is_padding, 0
)
if self._moe_is_padding is not None:
rows = np.arange(num_local_tokens_padded)
gathered = rows >= np.array(per_rank_num_tokens)[:, None]
is_padding = async_tensor_h2d(
gathered.ravel(), out=self._moe_is_padding[: gathered.size]
)
total_num_logits = num_local_reqs if num_local_tokens > 0 else 0
if total_num_logits > 0:
@@ -683,7 +707,9 @@ class PCPManager:
)
return self._local_batch
def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch:
def prepare_inputs_for_capture_or_profile(
self, input_batch: InputBatch
) -> InputBatch:
"""Stage a capture or dummy batch in persistent PCP input buffers."""
input_buffers = self.input_buffers
num_reqs = input_batch.num_reqs_after_padding
@@ -692,8 +718,8 @@ class PCPManager:
input_batch,
input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids),
positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions),
is_padding=input_buffers.is_padding[:num_tokens].copy_(
input_batch.is_padding
is_padding=self.moe_is_padding(
input_buffers.is_padding[:num_tokens].copy_(input_batch.is_padding)
),
query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_(
input_batch.query_start_loc
@@ -702,6 +728,15 @@ class PCPManager:
)
return input_batch
def moe_is_padding(self, is_padding: torch.Tensor) -> torch.Tensor:
"""Expand a capture/dummy padding mask, identical on every PCP rank, to
the rows the MoE routes."""
if self._moe_is_padding is None:
return is_padding
out = self._moe_is_padding[: self.pcp_world_size * is_padding.shape[0]]
out.view(self.pcp_world_size, -1).copy_(is_padding)
return out
def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
assert self._local_block_tables is not None
return tuple(
@@ -777,9 +812,13 @@ class PCPManager:
return gathered[self._hidden_restore_idx]
def get_draft_input_buffers(
self, input_buffers: InputBuffers
) -> InputBatch | InputBuffers:
return self.draft_prefill_batch or input_buffers
self, input_buffers: InputBuffers, num_tokens: int
) -> tuple[InputBatch | InputBuffers, torch.Tensor]:
"""Return draft inputs and their MoE routing padding mask."""
batch = self.draft_prefill_batch
if batch is not None:
return batch, batch.is_padding
return input_buffers, self.moe_is_padding(input_buffers.is_padding[:num_tokens])
def prepare_draft_prefill(
self, input_batch: InputBatch, input_ids: torch.Tensor
@@ -831,7 +870,7 @@ def maybe_build_pcp_manager(
vllm_config: VllmConfig,
device: torch.device,
supports_mm_inputs: bool,
block_tables: BlockTables,
block_tables: BlockTables | None = None,
cls: type[PCPManager] = PCPManager,
) -> PCPManager | None:
parallel_config = vllm_config.parallel_config
@@ -856,4 +895,6 @@ def maybe_build_pcp_manager(
dcp_world_size=dcp_size,
dcp_rank=dcp_rank,
cp_interleave=parallel_config.cp_kv_cache_interleave_size,
# Without EP all2all kernels, MoERunner all-gathers PCP tokens to route.
moe_gathers_tokens=not parallel_config.enable_expert_parallel,
)
vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py
+3 -2
View File
@@ -376,8 +376,9 @@ class TargetDependentARSpeculator(DraftModelSpeculator):
input_buffers: InputBatch | InputBuffers = self.input_buffers
is_padding = None
if self.pcp_manager is not None:
input_buffers = self.pcp_manager.get_draft_input_buffers(self.input_buffers)
is_padding = input_buffers.is_padding[:num_tokens]
input_buffers, is_padding = self.pcp_manager.get_draft_input_buffers(
self.input_buffers, num_tokens
)
batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
with set_forward_context(
attn_metadata,