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

[MRV2] Miscellaneous code cleanup (#57980)

Browse Source
This commit is contained in:
Nick Hill
2026-09-23 01:51:52 -07:00
committed by GitHub
parent 36bc694889
commit afec265bca
24 changed files with 163 additions and 326 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+1 -3
View File
@@ -166,9 +166,7 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
out[num_reqs:].zero_()
return out
with patch.object(
base_speculator, "maybe_prepare_dcp_local_seq_lens", fake_prepare
):
with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
captured = _run_build(
fake,
num_reqs=3,
tests/v1/worker/test_gpu_input_batch_v2.py
+4 -31
View File
@@ -81,7 +81,7 @@ def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool):
assert bool((topk_ids >= 0).all()) is not is_padding
def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
def test_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
"""The batch must view the caller-owned buffer, sliced to padded length.
Runtime (and capture) paths all funnel through this helper; attention
@@ -115,7 +115,7 @@ def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
return fake_kernel
monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel())
batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
buffers.dcp_local_seq_lens,
batch.seq_lens,
batch.num_reqs,
@@ -130,37 +130,10 @@ def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0]
def test_maybe_prepare_dcp_local_seq_lens_clears_stale_metadata(monkeypatch):
"""dcp_size == 1 resets the field to None instead of leaving a leftover.
DCP is toggled per deployment, and batches are recycled; a value from an
earlier DCP batch would otherwise be consumed as if it were current.
"""
buffers = InputBuffers(max_num_reqs=2, max_num_tokens=2, device=torch.device("cpu"))
batch = InputBatch.make_dummy(1, 1, buffers)
batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:1]
def fail_if_called(*args, **kwargs):
raise AssertionError("kernel must not run with dcp_size == 1")
monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", fail_if_called)
batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
buffers.dcp_local_seq_lens,
batch.seq_lens,
batch.num_reqs,
1,
0,
1,
num_reqs_padded=batch.num_reqs_after_padding,
)
assert batch.dcp_local_seq_lens is None
@pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
@pytest.mark.parametrize("dcp_size", [2, 4])
@pytest.mark.parametrize("cp_interleave", [1, 16])
def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
def test_prepare_dcp_local_seq_lens_matches_reference(
dcp_size: int, cp_interleave: int
):
"""Every rank's local lengths must equal the reference torch formula."""
@@ -172,7 +145,7 @@ def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
for dcp_rank in range(dcp_size):
buffers.dcp_local_seq_lens.fill_(-1)
batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
buffers.dcp_local_seq_lens,
batch.seq_lens,
batch.num_reqs,
tests/v1/worker/test_gpu_pcp_manager.py
+31 -23
View File
@@ -2,7 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import MagicMock
import numpy as np
import pytest
@@ -170,38 +169,44 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
@pytest.mark.parametrize(
("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"),
[
(CUDAGraphMode.NONE, 4, None, None),
(CUDAGraphMode.PIECEWISE, None, 8, None),
(CUDAGraphMode.NONE, 4, 2, 2),
(CUDAGraphMode.PIECEWISE, None, 8, 2),
(CUDAGraphMode.FULL, 4, 8, 4),
],
)
def test_partition_padding_is_derived_from_batch_descriptor(
cg_mode, num_reqs, expected_tokens, expected_reqs
):
manager = MagicMock()
input_batch = MagicMock()
manager.partition_batch.return_value = input_batch
batch_desc = BatchExecutionDescriptor(
cg_mode=cg_mode,
num_tokens=8,
num_reqs=num_reqs,
"""partition_batch derives graph padding from the batch descriptor: no
padding without graphs, token padding only for PIECEWISE, and request
padding on top for FULL."""
device = torch.device("cuda:0")
manager = PCPManager(
pcp_world_size=2,
pcp_rank=0,
device=device,
max_num_reqs=4,
max_num_tokens=8,
)
global_batch = _make_global_decode_batch(
[16, 24], InputBuffers(4, 8, device), device
)
result = pcp_manager_module.maybe_partition_pcp_batch(
manager,
input_batch,
batch_desc,
local_batch = manager.partition_batch(
global_batch,
BatchExecutionDescriptor(
cg_mode=cg_mode,
num_tokens=8,
num_reqs=num_reqs,
),
)
assert result is input_batch
manager.partition_batch.assert_called_once_with(
input_batch,
padded_num_tokens=expected_tokens,
padded_num_reqs=expected_reqs,
)
assert local_batch.num_tokens_after_padding == expected_tokens
assert local_batch.num_reqs_after_padding == expected_reqs
def test_capture_uses_pcp_persistent_inputs():
@@ -472,8 +477,11 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
local_batch = manager.partition_batch(
global_batch,
padded_num_tokens=4,
padded_num_reqs=4,
BatchExecutionDescriptor(
cg_mode=CUDAGraphMode.FULL,
num_tokens=4,
num_reqs=4,
),
)
assert local_batch.dcp_local_seq_lens is None
@@ -489,7 +497,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
# What execute_model does next: derive DCP metadata from the final batch
# on the PCP-owned buffers.
local_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
local_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
manager.input_buffers.dcp_local_seq_lens,
local_batch.seq_lens,
local_batch.num_reqs,
tests/v1/worker/test_gpu_ubatch_slicing.py
+1 -1
View File
@@ -756,7 +756,7 @@ def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int):
[1, 1, 10, 1, 1], [64, 96, 512, 32, 48]
)
# execute_model has populated the merged batch's DCP metadata already.
input_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
input_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
buffers.dcp_local_seq_lens,
input_batch.seq_lens,
input_batch.num_reqs,
vllm/v1/worker/gpu/block_table.py
+1 -2
View File
@@ -59,8 +59,7 @@ class BlockTables:
self.block_tables.append(block_table)
self.num_blocks = UvaBackedTensor(
(self.num_kv_cache_groups, self.max_num_reqs),
dtype=torch.int32,
(self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32
)
self.fused_writer: FusedStagedWriter | None = None
if self.num_kv_cache_groups > 1:
vllm/v1/worker/gpu/cp_utils.py
+2 -3
View File
@@ -5,7 +5,7 @@ import torch
from vllm.triton_utils import tl, triton
def maybe_prepare_dcp_local_seq_lens(
def prepare_dcp_local_seq_lens(
dcp_local_seq_lens: torch.Tensor,
seq_lens: torch.Tensor,
num_reqs: int,
@@ -16,8 +16,7 @@ def maybe_prepare_dcp_local_seq_lens(
num_reqs_padded: int | None = None,
) -> torch.Tensor | None:
"""Populate caller-owned storage and return its padded view, or None without DCP."""
if dcp_size == 1:
return None
assert dcp_size > 1
max_num_reqs = dcp_local_seq_lens.shape[0]
BLOCK_SIZE = 128
vllm/v1/worker/gpu/cudagraph_utils.py
+12 -12
View File
@@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
from vllm.v1.worker.gpu.model_states.interface import ModelState
from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches
@@ -643,8 +643,7 @@ class ModelCudaGraphManager(CudaGraphManager):
self.intermediate_tensors[k][:num_tokens] = v
def create_forward_fn(
desc: BatchExecutionDescriptor,
warmup: bool,
desc: BatchExecutionDescriptor, warmup: bool
) -> Callable[[CUDAGraphMode], None]:
num_tokens = desc.num_tokens
num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
@@ -789,15 +788,16 @@ def prepare_inputs_to_capture(
slot_mappings, kv_cache_config
)
input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
input_buffers.dcp_local_seq_lens,
input_batch.seq_lens,
input_batch.num_reqs,
block_tables.cp_size,
block_tables.cp_rank,
block_tables.cp_interleave,
num_reqs_padded=input_batch.num_reqs_after_padding,
)
if block_tables.cp_size > 1:
input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
input_buffers.dcp_local_seq_lens,
input_batch.seq_lens,
input_batch.num_reqs,
block_tables.cp_size,
block_tables.cp_rank,
block_tables.cp_interleave,
num_reqs_padded=input_batch.num_reqs_after_padding,
)
# NOTE(woosuk): Attention metadata is required not just by standard attention
# kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
vllm/v1/worker/gpu/ec_connector.py
+3 -10
View File
@@ -29,10 +29,7 @@ class ECConnector:
) -> Generator[ECConnectorOutput | None, None, None]:
yield None
def no_forward(
self,
scheduler_output: "SchedulerOutput",
) -> ModelRunnerOutput:
def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
return EMPTY_MODEL_RUNNER_OUTPUT
@@ -84,10 +81,7 @@ class ActiveECConnector(ECConnector):
output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta()
ec_connector.clear_connector_metadata()
def no_forward(
self,
scheduler_output: "SchedulerOutput",
) -> ModelRunnerOutput:
def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
# EC send/recv even if no work to do.
with self.maybe_get_output(scheduler_output) as ec_connector_output:
pass
@@ -99,8 +93,7 @@ NO_OP_EC_CONNECTOR = ECConnector()
def get_ec_connector(
vllm_config: VllmConfig,
encoder_cache: "EncoderCache | None",
vllm_config: VllmConfig, encoder_cache: "EncoderCache | None"
) -> ECConnector:
if (
not has_ec_transfer()
vllm/v1/worker/gpu/eplb_utils.py
+5 -18
View File
@@ -77,19 +77,13 @@ class EPLBController:
assert speculative_config is not None
assert speculative_config.draft_model_config is not None
assert self.state is not None
self.state.add_model(
draft_moe_model,
speculative_config.draft_model_config,
)
self.state.add_model(draft_moe_model, speculative_config.draft_model_config)
speculator.set_eplb_state(self.state)
self._has_registered_models = True
return True
def maybe_register_model(
self,
model: nn.Module,
model_config: Any,
load_dummy_weights: bool,
self, model: nn.Module, model_config: Any, load_dummy_weights: bool
) -> bool:
if not self.parallel_config.enable_eplb or load_dummy_weights:
return False
@@ -110,11 +104,7 @@ class EPLBController:
if eplb_models_added and self.state is not None and self.state.is_async:
self.state.start_async_loop()
def step(
self,
is_dummy: bool = False,
is_profile: bool = False,
) -> None:
def step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
if (
not self.parallel_config.enable_eplb
or self.suppressed
@@ -123,11 +113,8 @@ class EPLBController:
):
return
self.state.step(
is_dummy,
is_profile,
log_stats=self.parallel_config.eplb_config.log_balancedness,
)
log_stats = self.parallel_config.eplb_config.log_balancedness
self.state.step(is_dummy, is_profile, log_stats=log_stats)
def prepare_forward(
self,
vllm/v1/worker/gpu/input_batch.py
+3 -14
View File
@@ -17,12 +17,7 @@ if TYPE_CHECKING:
class InputBuffers:
def __init__(
self,
max_num_reqs: int,
max_num_tokens: int,
device: torch.device,
):
def __init__(self, max_num_reqs: int, max_num_tokens: int, device: torch.device):
self.max_num_reqs = max_num_reqs
self.max_num_tokens = max_num_tokens
self.device = device
@@ -474,11 +469,7 @@ def combine_sampled_and_draft_tokens(
num_reqs = idx_mapping.shape[0]
num_speculative_steps = draft_tokens.shape[-1]
logits_indices = torch.empty(
num_logits,
dtype=torch.int64,
device=input_ids.device,
)
logits_indices = torch.empty(num_logits, dtype=torch.int64, device=input_ids.device)
_combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
input_ids,
idx_mapping,
@@ -711,9 +702,7 @@ def expand_idx_mapping(
) -> tuple[torch.Tensor, torch.Tensor]:
num_reqs = idx_mapping.shape[0]
expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)
expanded_local_pos = torch.empty(
total_num_logits, dtype=torch.int32, device=idx_mapping.device
)
expanded_local_pos = idx_mapping.new_empty(total_num_logits)
_expand_idx_mapping_kernel[(num_reqs,)](
idx_mapping,
expanded_idx_mapping,
vllm/v1/worker/gpu/model_runner.py
+44 -71
View File
@@ -36,10 +36,7 @@ from vllm.distributed.aux_output_connector.worker import (
AuxOutputWorkerConnector,
get_aux_output_connector,
)
from vllm.distributed.parallel_state import (
get_dcp_group,
get_pp_group,
)
from vllm.distributed.parallel_state import get_dcp_group, get_pp_group
from vllm.forward_context import BatchDescriptor, set_forward_context
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.all2all_utils import get_ep_all2all_manager
@@ -97,7 +94,7 @@ from vllm.v1.worker.gpu.attn_utils import (
)
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cudagraph_utils import (
BatchExecutionDescriptor,
ModelCudaGraphManager,
@@ -494,8 +491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
),
)
self.prompt_logprobs_worker = PromptLogprobsWorker(
self.max_num_reqs,
logprobs_mode=self.model_config.logprobs_mode,
self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
)
self.structured_outputs_worker = StructuredOutputsWorker(
max_num_logits=self.max_num_reqs * self.decode_query_len,
@@ -508,9 +504,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
if self.is_pooling_model and self.is_last_pp_rank:
self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
eplb_models_added |= self.eplb.maybe_register_model(
self.model,
self.model_config,
load_dummy_weights,
self.model, self.model_config, load_dummy_weights
)
self.eplb.maybe_start_async_loop(eplb_models_added)
@@ -869,14 +863,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.step_timing.drafter_start()
mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
if self.speculator.supports_mm_inputs:
mm_inputs = (
[],
torch.zeros(
input_batch.num_tokens,
dtype=torch.bool,
device="cpu",
),
all_false = torch.zeros(
input_batch.num_tokens, dtype=torch.bool, device="cpu"
)
mm_inputs = [], all_false
# Let the target override the hidden state fed to the drafter
# (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
@@ -945,9 +935,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
mm_config = self.model_config.multimodal_config
if mm_config is not None and not mm_config.skip_mm_profiling:
mm_budget = MultiModalBudget(
self.vllm_config,
self.mm_registry,
enable_cache=False,
self.vllm_config, self.mm_registry, enable_cache=False
)
dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs()
self.model_state.encoder_runner.profile_encoder_cache(
@@ -1294,7 +1282,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
req_ids = batch_req_state.req_ids
num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
idx_mapping_np = batch_req_state.idx_mapping_np
idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)
idx_mapping = async_tensor_h2d(
idx_mapping_np, device=self.device, dtype=torch.int32
)
num_reqs = len(req_ids)
# Get the number of draft tokens for each request.
@@ -1471,11 +1461,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
else None
),
)
input_batch = pcp.maybe_partition_pcp_batch(
self.pcp_manager,
input_batch,
batch_desc,
)
if self.pcp_manager is not None:
input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc)
return input_batch
def prepare_attn(
@@ -1486,8 +1473,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
# Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
block_tables = self.block_tables.gather_block_tables(
input_batch.idx_mapping,
num_reqs_padded=input_batch.num_reqs_after_padding,
input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding
)
# Slot mappings: [num_kv_cache_groups, num_tokens_padded].
# Kernel pads beyond num_tokens with PAD_SLOT_ID.
@@ -1506,15 +1492,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
if valid_state_slots:
state_slots = torch.arange(
1,
input_batch.num_reqs + 1,
dtype=torch.int32,
device=self.device,
1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device
)
for block_table in block_tables:
block_table[:, 0].copy_(state_slots)
slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
self.pcp_manager, self.block_tables, input_batch.num_tokens
slot_mappings = block_table_provider.get_dummy_slot_mappings(
input_batch.num_tokens
)
return block_tables, slot_mappings
@@ -1607,11 +1590,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
query_start_loc: torch.Tensor | None = None,
) -> None:
# Update the number of computed tokens.
output_bin_counts = None
if self.is_last_pp_rank:
assert self.sampler is not None
output_bin_counts = self.sampler.penalties_state.output_bin_counts
else:
output_bin_counts = None
post_update(
idx_mapping,
self.req_states.num_computed_tokens.gpu,
@@ -1769,32 +1751,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
"Attention metadata must be prepared for dummy runs when using "
"FULL cudagraph mode."
)
block_tables = None
slot_mappings = None
block_tables, slot_mappings = None, None
else:
block_tables, slot_mappings = self.prepare_dummy_attn(
input_batch, valid_dummy_state_slots
)
if not skip_attn_for_dummy_run and context_len:
assert block_tables is not None
set_dummy_context(
input_batch,
self.block_tables,
context_len,
self.kv_cache_config.num_blocks,
self.max_model_len,
input_block_tables=block_tables,
)
if context_len:
set_dummy_context(
input_batch,
self.block_tables,
context_len,
self.kv_cache_config.num_blocks,
self.max_model_len,
input_block_tables=block_tables,
)
attn_metadata = None
slot_mappings_by_layer = None
if not (dummy_run and skip_attn_for_dummy_run):
input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
(
self.pcp_manager.input_buffers
if self.pcp_manager is not None
else self.input_buffers
).dcp_local_seq_lens,
if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run):
input_buffers = self.input_buffers
if self.pcp_manager is not None:
input_buffers = self.pcp_manager.input_buffers
input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
input_buffers.dcp_local_seq_lens,
input_batch.seq_lens,
input_batch.num_reqs,
self.dcp_size,
@@ -1802,6 +1779,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.cp_interleave,
num_reqs_padded=input_batch.num_reqs_after_padding,
)
attn_metadata = None
slot_mappings_by_layer = None
ubatch_state: UBatchState | None = None
if batch_desc.num_ubatches > 1:
assert self.ubatch_runner is not None
@@ -2046,16 +2026,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
# Last rank: sample tokens
assert hidden_states is not None
draft_hidden_states = hidden_states
assert draft_hidden_states is not None
hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
self.pcp_manager, hidden_states, input_batch
)
if self.pcp_manager is not None and aux_hidden_states is not None:
aux_hidden_states = [
self.pcp_manager.restore_hidden_states(states)
for states in aux_hidden_states
]
if self.pcp_manager is not None:
hidden_states, aux_hidden_states, input_batch = (
self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states)
)
sampler_output, num_sampled, num_rejected = self.sample(
hidden_states, input_batch, grammar_output
@@ -2305,13 +2281,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.eplb.suppressed = suppressed
def setup_eplb_from_mapping(
self,
expanded_physical_to_logical: torch.Tensor,
self, expanded_physical_to_logical: torch.Tensor
) -> None:
self.eplb.setup_from_mapping(
self.model,
self.model_config,
expanded_physical_to_logical,
self.model, self.model_config, expanded_physical_to_logical
)
########### EPLB methods end ###########
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+1 -3
View File
@@ -27,9 +27,7 @@ class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
def get_extra_common_attn_kwargs(
self,
kv_cache_group_id: int,
num_reqs: int,
self, kv_cache_group_id: int, num_reqs: int
) -> dict[str, Any]:
encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id)
if encoder_seq_lens is None:
vllm/v1/worker/gpu/model_states/encoder_only.py
+1 -4
View File
@@ -188,10 +188,7 @@ class EncoderOnlyModelState(DefaultModelState):
return attn_metadata
def _build_encoder_attn_metadata(
self,
input_batch: InputBatch,
cudagraph_mode: CUDAGraphMode,
for_capture: bool,
self, input_batch: InputBatch, cudagraph_mode: CUDAGraphMode, for_capture: bool
) -> dict[str, Any]:
if cudagraph_mode == CUDAGraphMode.FULL:
num_reqs = input_batch.num_reqs_after_padding
vllm/v1/worker/gpu/model_states/interface.py
+2 -6
View File
@@ -28,16 +28,12 @@ class ModelSpecificAttnMetadata:
"""Base class for model-specific attention metadata."""
def get_extra_common_attn_kwargs(
self,
kv_cache_group_id: int,
num_reqs: int,
self, kv_cache_group_id: int, num_reqs: int
) -> dict[str, Any]:
return {}
def get_extra_attn_kwargs(
self,
attn_metadata_builder: Any,
num_reqs: int,
self, attn_metadata_builder: Any, num_reqs: int
) -> dict[str, Any]:
return {}
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+3 -9
View File
@@ -43,16 +43,12 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
num_decode_draft_tokens_cpu: torch.Tensor | None = None
def get_extra_common_attn_kwargs(
self,
kv_cache_group_id: int,
num_reqs: int,
self, kv_cache_group_id: int, num_reqs: int
) -> dict[str, Any]:
return {"is_prefilling": self.is_prefilling[:num_reqs]}
def get_extra_attn_kwargs(
self,
attn_metadata_builder: Any,
num_reqs: int,
self, attn_metadata_builder: Any, num_reqs: int
) -> dict[str, Any]:
if not isinstance(
attn_metadata_builder,
@@ -338,9 +334,7 @@ class MambaHybridModelState(DefaultModelState):
)
if self.recoverssm is not None:
self.recoverssm.record_step(
attn_metadata,
attn_groups,
for_capture=for_capture,
attn_metadata, attn_groups, for_capture=for_capture
)
return attn_metadata
vllm/v1/worker/gpu/pcp_manager.py
+21 -64
View File
@@ -192,8 +192,7 @@ class PCPManager:
@staticmethod
def _reorder_segments(
segments: list[RankSegment],
is_prefilling: np.ndarray,
segments: list[RankSegment], is_prefilling: np.ndarray
) -> list[RankSegment]:
"""Order this rank's rows decodes-first, then prefills, canonically."""
@@ -218,9 +217,7 @@ class PCPManager:
return segments
def replicated_requests(
self,
num_scheduled_tokens: np.ndarray,
is_prefilling: np.ndarray,
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
@@ -407,10 +404,7 @@ class PCPManager:
return self._input_buffers
def partition_batch(
self,
input_batch: InputBatch,
padded_num_tokens: int | None = None,
padded_num_reqs: int | None = None,
self, input_batch: InputBatch, batch_desc: "BatchExecutionDescriptor"
) -> InputBatch:
assert self._input_buffers is not None
input_buffers = self._input_buffers
@@ -422,6 +416,13 @@ class PCPManager:
num_computed_tokens = global_batch.num_computed_tokens_np
is_prefilling = global_batch.is_prefilling_np
padded_num_tokens = None
padded_num_reqs = None
if batch_desc.cg_mode != CUDAGraphMode.NONE:
padded_num_tokens = batch_desc.num_tokens
if batch_desc.cg_mode == CUDAGraphMode.FULL:
padded_num_reqs = batch_desc.num_reqs
segments_by_rank, per_rank_num_tokens = self._build_batch_layout(
num_scheduled_tokens,
num_computed_tokens,
@@ -442,9 +443,7 @@ class PCPManager:
num_local_reqs = len(local_segments)
num_reqs_after_padding = self._resolve_num_reqs_after_padding(
global_batch,
padded_num_reqs,
num_local_reqs,
global_batch, padded_num_reqs, num_local_reqs
)
if num_reqs_after_padding > input_buffers.max_num_reqs:
raise RuntimeError(
@@ -714,8 +713,7 @@ class PCPManager:
return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size]
def _convert_to_gathered_slot_mappings(
self,
global_batch_slot_mappings: torch.Tensor,
self, global_batch_slot_mappings: torch.Tensor
) -> torch.Tensor:
assert self._padded_gather_idx is not None
assert self._gathered_kv_write_mask is not None
@@ -754,9 +752,7 @@ class PCPManager:
return self.draft_prefill_batch or input_buffers
def prepare_draft_prefill(
self,
input_batch: InputBatch,
input_ids: torch.Tensor,
self, input_batch: InputBatch, input_ids: torch.Tensor
) -> None:
if input_batch is not self._global_batch or self._local_batch is None:
return
@@ -789,54 +785,15 @@ class PCPManager:
return last_hidden_states, hidden_states
def restore_for_sampling(
self,
hidden_states: torch.Tensor,
) -> tuple[torch.Tensor, InputBatch]:
self, hidden_states: torch.Tensor, aux_hidden_states: list[torch.Tensor] | None
) -> tuple[torch.Tensor, list[torch.Tensor] | None, InputBatch]:
assert self._global_batch is not None
return self.restore_hidden_states(hidden_states), self._global_batch
def maybe_partition_pcp_batch(
manager: PCPManager | None,
input_batch: InputBatch,
batch_desc: "BatchExecutionDescriptor",
) -> InputBatch:
if manager is None:
return input_batch
padded_num_tokens = None
padded_num_reqs = None
if batch_desc.cg_mode != CUDAGraphMode.NONE:
padded_num_tokens = batch_desc.num_tokens
if batch_desc.cg_mode == CUDAGraphMode.FULL:
padded_num_reqs = batch_desc.num_reqs
return manager.partition_batch(
input_batch,
padded_num_tokens=padded_num_tokens,
padded_num_reqs=padded_num_reqs,
)
def maybe_get_pcp_dummy_slot_mappings(
manager: PCPManager | None,
block_tables: BlockTables,
num_tokens: int,
) -> torch.Tensor:
if manager is None:
return block_tables.get_dummy_slot_mappings(num_tokens)
return manager.get_dummy_slot_mappings(num_tokens)
def maybe_restore_pcp_for_sampling(
manager: PCPManager | None,
hidden_states: torch.Tensor | None,
input_batch: InputBatch,
) -> tuple[torch.Tensor, InputBatch]:
assert hidden_states is not None
if manager is None:
return hidden_states, input_batch
return manager.restore_for_sampling(hidden_states)
hidden_states = self.restore_hidden_states(hidden_states)
if aux_hidden_states is not None:
aux_hidden_states = [
self.restore_hidden_states(states) for states in aux_hidden_states
]
return hidden_states, aux_hidden_states, self._global_batch
def maybe_build_pcp_manager(
vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+4 -7
View File
@@ -255,9 +255,7 @@ class AdaptiveVerificationManager:
logger.debug("DSpark cost tables: %s", self.cost_tables)
def record_confidences(
self,
confidence_probs: torch.Tensor,
input_batch: "InputBatch",
self, confidence_probs: torch.Tensor, input_batch: "InputBatch"
) -> None:
"""Publish this step's raw confidences for the ranking kernel and start
copying them to the CPU, where a later step's budget reads them."""
@@ -305,9 +303,9 @@ class AdaptiveVerificationManager:
)
num_non_draft_tokens = scheduled_tokens - scheduled_drafts
slots = np.fromiter(
(self.req_states.req_id_to_index[req_id] for req_id in req_ids),
map(self.req_states.req_id_to_index.__getitem__, req_ids),
dtype=np.int32,
count=len(req_ids),
count=num_reqs,
)
stale_confidences = self._stale_confidences[self._stale_idx].np[slots]
survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1)
@@ -469,8 +467,7 @@ def maybe_create_adaptive_verification_manager(
# The selector rejects unsupported backends, but models that
# hard-wire theirs (e.g. DeepSeek-V4) never go through it.
backend = get_query_lens_mismatch_unsupported_backend(
attn_groups,
checked_layer_names=target_layer_names,
attn_groups, checked_layer_names=target_layer_names
)
if backend is not None:
raise ValueError(
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+3 -10
View File
@@ -53,10 +53,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
return
self.inputs_embeds = torch.zeros(
self.max_num_tokens,
self.hidden_size,
dtype=self.dtype,
device=self.device,
self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=self.device
)
# Lifecycle hooks for model-specific optimizations. Subclasses override
@@ -140,10 +137,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
# Initialize cudagraph manager for draft decodes (draft positions > 0).
self.decode_cudagraph_manager = SpeculatorCudaGraphManager(
self.vllm_config,
self.device,
cudagraph_mode,
decode_query_len=1,
self.vllm_config, self.device, cudagraph_mode, decode_query_len=1
)
def capture(self) -> None:
@@ -272,8 +266,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
if self.pcp_manager is not None:
self.pcp_manager.prepare_draft_prefill(
input_batch,
self.input_buffers.input_ids[:num_tokens_padded],
input_batch, self.input_buffers.input_ids[:num_tokens_padded]
)
prefill = self.pcp_manager.draft_prefill_batch
if prefill is not None:
vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
+11 -10
View File
@@ -11,7 +11,7 @@ from vllm.v1.worker.gpu.attn_utils import (
build_slot_mappings_by_layer,
)
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cudagraph_utils import (
AttentionState,
BatchExecutionDescriptor,
@@ -42,15 +42,16 @@ def _prepare_dflash_inputs_to_capture(
attn_metadata = None
if not skip_attn:
query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
input_buffers.dcp_local_seq_lens,
input_batch.seq_lens,
input_batch.num_reqs,
block_tables.cp_size,
block_tables.cp_rank,
block_tables.cp_interleave,
num_reqs_padded=input_batch.num_reqs_after_padding,
)
if block_tables.cp_size > 1:
input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
input_buffers.dcp_local_seq_lens,
input_batch.seq_lens,
input_batch.num_reqs,
block_tables.cp_size,
block_tables.cp_rank,
block_tables.cp_interleave,
num_reqs_padded=input_batch.num_reqs_after_padding,
)
attn_metadata = build_attn_metadata(
attn_groups=attn_groups,
num_reqs=num_reqs,
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+1 -4
View File
@@ -149,10 +149,7 @@ class DSparkSpeculator(DFlashSpeculator):
)
if self.draft_watermarker is not None:
sampled = self.draft_watermarker.sample(
logits,
sampled,
idx_map,
self.temperature,
logits, sampled, idx_map, self.temperature
)
return sampled
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+1 -2
View File
@@ -134,8 +134,7 @@ class RejectionSampler:
if expanded_logits:
if self.enable_adaptive_verification:
# Adaptive verification keeps the true per-request boundaries
# on device only; cu_num_logits_np holds the pre-compacted
# layout.
# on device only; cu_num_logits_np holds the pre-compacted layout.
cu_num_generated_tokens = cu_num_logits.clone()
else:
cu_num_generated_tokens = cu_num_logits_np.tolist()
vllm/v1/worker/gpu/spec_decode/speculator.py
+5 -13
View File
@@ -26,7 +26,7 @@ from vllm.v1.worker.gpu.attn_utils import (
init_attn_backend,
)
from vllm.v1.worker.gpu.block_table import BlockTables
from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
from vllm.v1.worker.gpu.dp_utils import DPSyncState
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
@@ -202,9 +202,7 @@ class DraftModelSpeculator(BaseSpeculator):
@abstractmethod
def load_draft_model(
self,
target_model: nn.Module,
target_attn_layer_names: set[str],
self, target_model: nn.Module, target_attn_layer_names: set[str]
) -> nn.Module:
pass
@@ -337,7 +335,7 @@ class DraftModelSpeculator(BaseSpeculator):
if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
# Draft steps advance and rewind their own global sequence lengths,
# so the target model's DCP-local lengths may already be stale.
dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
dcp_local_seq_lens = prepare_dcp_local_seq_lens(
self.input_buffers.dcp_local_seq_lens,
self.input_buffers.seq_lens,
num_reqs,
@@ -422,10 +420,7 @@ class DraftModelSpeculator(BaseSpeculator):
)
if self.draft_watermarker is not None:
sampled = self.draft_watermarker.sample(
logits,
sampled,
idx_mapping,
temperature,
logits, sampled, idx_mapping, temperature
)
elif self.use_local_argmax_reduction:
return self.model.get_top_tokens(hidden_states)
@@ -436,10 +431,7 @@ class DraftModelSpeculator(BaseSpeculator):
return sampled
def _maybe_predict_acceptance(
self,
logits: torch.Tensor,
idx_mapping: torch.Tensor,
draft_step: torch.Tensor,
self, logits: torch.Tensor, idx_mapping: torch.Tensor, draft_step: torch.Tensor
) -> None:
if self.acceptance_estimator is not None:
self.acceptance_estimator.predict(
vllm/v1/worker/gpu/states.py
+1 -4
View File
@@ -78,10 +78,7 @@ class RequestState:
)
self.next_prefill_tokens = torch.zeros(
num_prefill_lookahead,
self.max_num_reqs,
dtype=torch.int32,
device=device,
num_prefill_lookahead, self.max_num_reqs, dtype=torch.int32, device=device
)
@property
vllm/v1/worker/gpu/ubatch_utils.py
+2 -2
View File
@@ -24,7 +24,7 @@ from vllm.sequence import IntermediateTensors
from vllm.utils.torch_utils import current_stream
from vllm.v1.kv_cache_interface import KVCacheConfig
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
from vllm.v1.worker.gpu.input_batch import InputBatch
from vllm.v1.worker.gpu.model_states.interface import ModelState
from vllm.v1.worker.ubatch_utils import (
@@ -150,7 +150,7 @@ def _slice_input_batch(
dcp_local_seq_lens = None
if dcp_size > 1:
assert dcp_local_seq_lens_buf is not None
dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
dcp_local_seq_lens = prepare_dcp_local_seq_lens(
dcp_local_seq_lens_buf,
seq_lens,
num_reqs,