@tduncan / Nvidia-Optimization-Stack-VLLM / commits / afec265bca

[MRV2] Miscellaneous code cleanup (#57980)

Nick Hill committed Sep 23, 2026 at 01:51 UTC afec265bca16d02350a2da1855cbb90dd4c7123b
24 files changed +163 -326
tests/v1/spec_decode/test_eagle_draft_attn_metadata.py
+1 -3
@@ -166,9 +166,7 @@ def test_build_draft_attn_metadata_recomputes_dcp_local_seq_lens():
166 out[num_reqs:].zero_()
167 return out
168
169 - with patch.object(
170 - base_speculator, "maybe_prepare_dcp_local_seq_lens", fake_prepare
171 - ):
169 + with patch.object(base_speculator, "prepare_dcp_local_seq_lens", fake_prepare):
170 captured = _run_build(
171 fake,
172 num_reqs=3,
tests/v1/worker/test_gpu_input_batch_v2.py
+4 -31
@@ -81,7 +81,7 @@ def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool):
81 assert bool((topk_ids >= 0).all()) is not is_padding
82
83
84 -def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
84 +def test_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
85 """The batch must view the caller-owned buffer, sliced to padded length.
86
87 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):
115 return fake_kernel
116
117 monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel())
118 - batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
118 + batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
119 buffers.dcp_local_seq_lens,
120 batch.seq_lens,
121 batch.num_reqs,
@@ -130,37 +130,10 @@ def test_maybe_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch):
130 assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0]
131
132
133 -def test_maybe_prepare_dcp_local_seq_lens_clears_stale_metadata(monkeypatch):
134 - """dcp_size == 1 resets the field to None instead of leaving a leftover.
135 -
136 - DCP is toggled per deployment, and batches are recycled; a value from an
137 - earlier DCP batch would otherwise be consumed as if it were current.
138 - """
139 - buffers = InputBuffers(max_num_reqs=2, max_num_tokens=2, device=torch.device("cpu"))
140 - batch = InputBatch.make_dummy(1, 1, buffers)
141 - batch.dcp_local_seq_lens = buffers.dcp_local_seq_lens[:1]
142 -
143 - def fail_if_called(*args, **kwargs):
144 - raise AssertionError("kernel must not run with dcp_size == 1")
145 -
146 - monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", fail_if_called)
147 - batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
148 - buffers.dcp_local_seq_lens,
149 - batch.seq_lens,
150 - batch.num_reqs,
151 - 1,
152 - 0,
153 - 1,
154 - num_reqs_padded=batch.num_reqs_after_padding,
155 - )
156 -
157 - assert batch.dcp_local_seq_lens is None
158 -
159 -
133 @pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA")
134 @pytest.mark.parametrize("dcp_size", [2, 4])
135 @pytest.mark.parametrize("cp_interleave", [1, 16])
163 -def test_maybe_prepare_dcp_local_seq_lens_matches_reference(
136 +def test_prepare_dcp_local_seq_lens_matches_reference(
137 dcp_size: int, cp_interleave: int
138 ):
139 """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(
145
146 for dcp_rank in range(dcp_size):
147 buffers.dcp_local_seq_lens.fill_(-1)
175 - batch.dcp_local_seq_lens = cp_utils.maybe_prepare_dcp_local_seq_lens(
148 + batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens(
149 buffers.dcp_local_seq_lens,
150 batch.seq_lens,
151 batch.num_reqs,
tests/v1/worker/test_gpu_pcp_manager.py
+31 -23
@@ -2,7 +2,6 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import replace
4 from types import SimpleNamespace
5 -from unittest.mock import MagicMock
5
6 import numpy as np
7 import pytest
@@ -170,39 +169,45 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
169 )
170
171
172 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU kernels")
173 @pytest.mark.parametrize(
174 ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"),
175 [
176 - (CUDAGraphMode.NONE, 4, None, None),
177 - (CUDAGraphMode.PIECEWISE, None, 8, None),
176 + (CUDAGraphMode.NONE, 4, 2, 2),
177 + (CUDAGraphMode.PIECEWISE, None, 8, 2),
178 (CUDAGraphMode.FULL, 4, 8, 4),
179 ],
180 )
181 def test_partition_padding_is_derived_from_batch_descriptor(
182 cg_mode, num_reqs, expected_tokens, expected_reqs
183 ):
184 - manager = MagicMock()
185 - input_batch = MagicMock()
186 - manager.partition_batch.return_value = input_batch
187 - batch_desc = BatchExecutionDescriptor(
188 - cg_mode=cg_mode,
189 - num_tokens=8,
190 - num_reqs=num_reqs,
184 + """partition_batch derives graph padding from the batch descriptor: no
185 + padding without graphs, token padding only for PIECEWISE, and request
186 + padding on top for FULL."""
187 + device = torch.device("cuda:0")
188 + manager = PCPManager(
189 + pcp_world_size=2,
190 + pcp_rank=0,
191 + device=device,
192 + max_num_reqs=4,
193 + max_num_tokens=8,
194 )
192 -
193 - result = pcp_manager_module.maybe_partition_pcp_batch(
194 - manager,
195 - input_batch,
196 - batch_desc,
195 + global_batch = _make_global_decode_batch(
196 + [16, 24], InputBuffers(4, 8, device), device
197 )
198
199 - assert result is input_batch
200 - manager.partition_batch.assert_called_once_with(
201 - input_batch,
202 - padded_num_tokens=expected_tokens,
203 - padded_num_reqs=expected_reqs,
199 + local_batch = manager.partition_batch(
200 + global_batch,
201 + BatchExecutionDescriptor(
202 + cg_mode=cg_mode,
203 + num_tokens=8,
204 + num_reqs=num_reqs,
205 + ),
206 )
207
208 + assert local_batch.num_tokens_after_padding == expected_tokens
209 + assert local_batch.num_reqs_after_padding == expected_reqs
210 +
211
212 def test_capture_uses_pcp_persistent_inputs():
213 manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
@@ -472,8 +477,11 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
477
478 local_batch = manager.partition_batch(
479 global_batch,
475 - padded_num_tokens=4,
476 - padded_num_reqs=4,
480 + BatchExecutionDescriptor(
481 + cg_mode=CUDAGraphMode.FULL,
482 + num_tokens=4,
483 + num_reqs=4,
484 + ),
485 )
486
487 assert local_batch.dcp_local_seq_lens is None
@@ -489,7 +497,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
497
498 # What execute_model does next: derive DCP metadata from the final batch
499 # on the PCP-owned buffers.
492 - local_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
500 + local_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
501 manager.input_buffers.dcp_local_seq_lens,
502 local_batch.seq_lens,
503 local_batch.num_reqs,
tests/v1/worker/test_gpu_ubatch_slicing.py
+1 -1
@@ -756,7 +756,7 @@ def test_microbatches_recompute_dcp_lens_from_truncated_seq_lens(dcp_rank: int):
756 [1, 1, 10, 1, 1], [64, 96, 512, 32, 48]
757 )
758 # execute_model has populated the merged batch's DCP metadata already.
759 - input_batch.dcp_local_seq_lens = gpu_cp_utils.maybe_prepare_dcp_local_seq_lens(
759 + input_batch.dcp_local_seq_lens = gpu_cp_utils.prepare_dcp_local_seq_lens(
760 buffers.dcp_local_seq_lens,
761 input_batch.seq_lens,
762 input_batch.num_reqs,
vllm/v1/worker/gpu/block_table.py
+1 -2
@@ -59,8 +59,7 @@ class BlockTables:
59 self.block_tables.append(block_table)
60
61 self.num_blocks = UvaBackedTensor(
62 - (self.num_kv_cache_groups, self.max_num_reqs),
63 - dtype=torch.int32,
62 + (self.num_kv_cache_groups, self.max_num_reqs), dtype=torch.int32
63 )
64 self.fused_writer: FusedStagedWriter | None = None
65 if self.num_kv_cache_groups > 1:
vllm/v1/worker/gpu/cp_utils.py
+2 -3
@@ -5,7 +5,7 @@ import torch
5 from vllm.triton_utils import tl, triton
6
7
8 -def maybe_prepare_dcp_local_seq_lens(
8 +def prepare_dcp_local_seq_lens(
9 dcp_local_seq_lens: torch.Tensor,
10 seq_lens: torch.Tensor,
11 num_reqs: int,
@@ -16,8 +16,7 @@ def maybe_prepare_dcp_local_seq_lens(
16 num_reqs_padded: int | None = None,
17 ) -> torch.Tensor | None:
18 """Populate caller-owned storage and return its padded view, or None without DCP."""
19 - if dcp_size == 1:
20 - return None
19 + assert dcp_size > 1
20
21 max_num_reqs = dcp_local_seq_lens.shape[0]
22 BLOCK_SIZE = 128
vllm/v1/worker/gpu/cudagraph_utils.py
+12 -12
@@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
40 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
41 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
42 from vllm.v1.worker.gpu.block_table import BlockTables
43 -from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
43 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
44 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
45 from vllm.v1.worker.gpu.model_states.interface import ModelState
46 from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches
@@ -643,8 +643,7 @@ class ModelCudaGraphManager(CudaGraphManager):
643 self.intermediate_tensors[k][:num_tokens] = v
644
645 def create_forward_fn(
646 - desc: BatchExecutionDescriptor,
647 - warmup: bool,
646 + desc: BatchExecutionDescriptor, warmup: bool
647 ) -> Callable[[CUDAGraphMode], None]:
648 num_tokens = desc.num_tokens
649 num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
@@ -789,15 +788,16 @@ def prepare_inputs_to_capture(
788 slot_mappings, kv_cache_config
789 )
790
792 - input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
793 - input_buffers.dcp_local_seq_lens,
794 - input_batch.seq_lens,
795 - input_batch.num_reqs,
796 - block_tables.cp_size,
797 - block_tables.cp_rank,
798 - block_tables.cp_interleave,
799 - num_reqs_padded=input_batch.num_reqs_after_padding,
800 - )
791 + if block_tables.cp_size > 1:
792 + input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
793 + input_buffers.dcp_local_seq_lens,
794 + input_batch.seq_lens,
795 + input_batch.num_reqs,
796 + block_tables.cp_size,
797 + block_tables.cp_rank,
798 + block_tables.cp_interleave,
799 + num_reqs_padded=input_batch.num_reqs_after_padding,
800 + )
801
802 # NOTE(woosuk): Attention metadata is required not just by standard attention
803 # kernels, but also by specialized attention-like operations (e.g., Inkling's sconv,
vllm/v1/worker/gpu/ec_connector.py
+3 -10
@@ -29,10 +29,7 @@ class ECConnector:
29 ) -> Generator[ECConnectorOutput | None, None, None]:
30 yield None
31
32 - def no_forward(
33 - self,
34 - scheduler_output: "SchedulerOutput",
35 - ) -> ModelRunnerOutput:
32 + def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
33 return EMPTY_MODEL_RUNNER_OUTPUT
34
35
@@ -84,10 +81,7 @@ class ActiveECConnector(ECConnector):
81 output.ec_connector_worker_meta = ec_connector.build_connector_worker_meta()
82 ec_connector.clear_connector_metadata()
83
87 - def no_forward(
88 - self,
89 - scheduler_output: "SchedulerOutput",
90 - ) -> ModelRunnerOutput:
84 + def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
85 # EC send/recv even if no work to do.
86 with self.maybe_get_output(scheduler_output) as ec_connector_output:
87 pass
@@ -99,8 +93,7 @@ NO_OP_EC_CONNECTOR = ECConnector()
93
94
95 def get_ec_connector(
102 - vllm_config: VllmConfig,
103 - encoder_cache: "EncoderCache | None",
96 + vllm_config: VllmConfig, encoder_cache: "EncoderCache | None"
97 ) -> ECConnector:
98 if (
99 not has_ec_transfer()
vllm/v1/worker/gpu/eplb_utils.py
+5 -18
@@ -77,19 +77,13 @@ class EPLBController:
77 assert speculative_config is not None
78 assert speculative_config.draft_model_config is not None
79 assert self.state is not None
80 - self.state.add_model(
81 - draft_moe_model,
82 - speculative_config.draft_model_config,
83 - )
80 + self.state.add_model(draft_moe_model, speculative_config.draft_model_config)
81 speculator.set_eplb_state(self.state)
82 self._has_registered_models = True
83 return True
84
85 def maybe_register_model(
89 - self,
90 - model: nn.Module,
91 - model_config: Any,
92 - load_dummy_weights: bool,
86 + self, model: nn.Module, model_config: Any, load_dummy_weights: bool
87 ) -> bool:
88 if not self.parallel_config.enable_eplb or load_dummy_weights:
89 return False
@@ -110,11 +104,7 @@ class EPLBController:
104 if eplb_models_added and self.state is not None and self.state.is_async:
105 self.state.start_async_loop()
106
113 - def step(
114 - self,
115 - is_dummy: bool = False,
116 - is_profile: bool = False,
117 - ) -> None:
107 + def step(self, is_dummy: bool = False, is_profile: bool = False) -> None:
108 if (
109 not self.parallel_config.enable_eplb
110 or self.suppressed
@@ -123,11 +113,8 @@ class EPLBController:
113 ):
114 return
115
126 - self.state.step(
127 - is_dummy,
128 - is_profile,
129 - log_stats=self.parallel_config.eplb_config.log_balancedness,
130 - )
116 + log_stats = self.parallel_config.eplb_config.log_balancedness
117 + self.state.step(is_dummy, is_profile, log_stats=log_stats)
118
119 def prepare_forward(
120 self,
vllm/v1/worker/gpu/input_batch.py
+3 -14
@@ -17,12 +17,7 @@ if TYPE_CHECKING:
17
18
19 class InputBuffers:
20 - def __init__(
21 - self,
22 - max_num_reqs: int,
23 - max_num_tokens: int,
24 - device: torch.device,
25 - ):
20 + def __init__(self, max_num_reqs: int, max_num_tokens: int, device: torch.device):
21 self.max_num_reqs = max_num_reqs
22 self.max_num_tokens = max_num_tokens
23 self.device = device
@@ -474,11 +469,7 @@ def combine_sampled_and_draft_tokens(
469 num_reqs = idx_mapping.shape[0]
470 num_speculative_steps = draft_tokens.shape[-1]
471
477 - logits_indices = torch.empty(
478 - num_logits,
479 - dtype=torch.int64,
480 - device=input_ids.device,
481 - )
472 + logits_indices = torch.empty(num_logits, dtype=torch.int64, device=input_ids.device)
473 _combine_sampled_and_draft_tokens_kernel[(num_reqs,)](
474 input_ids,
475 idx_mapping,
@@ -711,9 +702,7 @@ def expand_idx_mapping(
702 ) -> tuple[torch.Tensor, torch.Tensor]:
703 num_reqs = idx_mapping.shape[0]
704 expanded_idx_mapping = idx_mapping.new_empty(total_num_logits)
714 - expanded_local_pos = torch.empty(
715 - total_num_logits, dtype=torch.int32, device=idx_mapping.device
716 - )
705 + expanded_local_pos = idx_mapping.new_empty(total_num_logits)
706 _expand_idx_mapping_kernel[(num_reqs,)](
707 idx_mapping,
708 expanded_idx_mapping,
vllm/v1/worker/gpu/model_runner.py
+44 -71
@@ -36,10 +36,7 @@ from vllm.distributed.aux_output_connector.worker import (
36 AuxOutputWorkerConnector,
37 get_aux_output_connector,
38 )
39 -from vllm.distributed.parallel_state import (
40 - get_dcp_group,
41 - get_pp_group,
42 -)
39 +from vllm.distributed.parallel_state import get_dcp_group, get_pp_group
40 from vllm.forward_context import BatchDescriptor, set_forward_context
41 from vllm.logger import init_logger
42 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 (
94 )
95 from vllm.v1.worker.gpu.block_table import BlockTables
96 from vllm.v1.worker.gpu.buffer_utils import set_default_max_concurrency
100 -from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
97 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
98 from vllm.v1.worker.gpu.cudagraph_utils import (
99 BatchExecutionDescriptor,
100 ModelCudaGraphManager,
@@ -494,8 +491,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
491 ),
492 )
493 self.prompt_logprobs_worker = PromptLogprobsWorker(
497 - self.max_num_reqs,
498 - logprobs_mode=self.model_config.logprobs_mode,
494 + self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode
495 )
496 self.structured_outputs_worker = StructuredOutputsWorker(
497 max_num_logits=self.max_num_reqs * self.decode_query_len,
@@ -508,9 +504,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
504 if self.is_pooling_model and self.is_last_pp_rank:
505 self.pooling_runner = PoolingRunner(self.model, self.vllm_config)
506 eplb_models_added |= self.eplb.maybe_register_model(
511 - self.model,
512 - self.model_config,
513 - load_dummy_weights,
507 + self.model, self.model_config, load_dummy_weights
508 )
509 self.eplb.maybe_start_async_loop(eplb_models_added)
510
@@ -869,14 +863,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
863 self.step_timing.drafter_start()
864 mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
865 if self.speculator.supports_mm_inputs:
872 - mm_inputs = (
873 - [],
874 - torch.zeros(
875 - input_batch.num_tokens,
876 - dtype=torch.bool,
877 - device="cpu",
878 - ),
866 + all_false = torch.zeros(
867 + input_batch.num_tokens, dtype=torch.bool, device="cpu"
868 )
869 + mm_inputs = [], all_false
870
871 # Let the target override the hidden state fed to the drafter
872 # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The
@@ -945,9 +935,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
935 mm_config = self.model_config.multimodal_config
936 if mm_config is not None and not mm_config.skip_mm_profiling:
937 mm_budget = MultiModalBudget(
948 - self.vllm_config,
949 - self.mm_registry,
950 - enable_cache=False,
938 + self.vllm_config, self.mm_registry, enable_cache=False
939 )
940 dummy_mm_inputs = mm_budget.get_dummy_encoder_profile_inputs()
941 self.model_state.encoder_runner.profile_encoder_cache(
@@ -1294,7 +1282,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1282 req_ids = batch_req_state.req_ids
1283 num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens
1284 idx_mapping_np = batch_req_state.idx_mapping_np
1297 - idx_mapping = async_tensor_h2d(idx_mapping_np, device=self.device)
1285 + idx_mapping = async_tensor_h2d(
1286 + idx_mapping_np, device=self.device, dtype=torch.int32
1287 + )
1288 num_reqs = len(req_ids)
1289
1290 # Get the number of draft tokens for each request.
@@ -1471,11 +1461,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1461 else None
1462 ),
1463 )
1474 - input_batch = pcp.maybe_partition_pcp_batch(
1475 - self.pcp_manager,
1476 - input_batch,
1477 - batch_desc,
1478 - )
1464 + if self.pcp_manager is not None:
1465 + input_batch = self.pcp_manager.partition_batch(input_batch, batch_desc)
1466 return input_batch
1467
1468 def prepare_attn(
@@ -1486,8 +1473,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1473
1474 # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
1475 block_tables = self.block_tables.gather_block_tables(
1489 - input_batch.idx_mapping,
1490 - num_reqs_padded=input_batch.num_reqs_after_padding,
1476 + input_batch.idx_mapping, num_reqs_padded=input_batch.num_reqs_after_padding
1477 )
1478 # Slot mappings: [num_kv_cache_groups, num_tokens_padded].
1479 # Kernel pads beyond num_tokens with PAD_SLOT_ID.
@@ -1506,15 +1492,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1492 block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
1493 if valid_state_slots:
1494 state_slots = torch.arange(
1509 - 1,
1510 - input_batch.num_reqs + 1,
1511 - dtype=torch.int32,
1512 - device=self.device,
1495 + 1, input_batch.num_reqs + 1, dtype=torch.int32, device=self.device
1496 )
1497 for block_table in block_tables:
1498 block_table[:, 0].copy_(state_slots)
1516 - slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
1517 - self.pcp_manager, self.block_tables, input_batch.num_tokens
1499 + slot_mappings = block_table_provider.get_dummy_slot_mappings(
1500 + input_batch.num_tokens
1501 )
1502 return block_tables, slot_mappings
1503
@@ -1607,11 +1590,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1590 query_start_loc: torch.Tensor | None = None,
1591 ) -> None:
1592 # Update the number of computed tokens.
1593 + output_bin_counts = None
1594 if self.is_last_pp_rank:
1595 assert self.sampler is not None
1596 output_bin_counts = self.sampler.penalties_state.output_bin_counts
1613 - else:
1614 - output_bin_counts = None
1597 post_update(
1598 idx_mapping,
1599 self.req_states.num_computed_tokens.gpu,
@@ -1769,32 +1751,27 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1751 "Attention metadata must be prepared for dummy runs when using "
1752 "FULL cudagraph mode."
1753 )
1772 - block_tables = None
1773 - slot_mappings = None
1754 + block_tables, slot_mappings = None, None
1755 else:
1756 block_tables, slot_mappings = self.prepare_dummy_attn(
1757 input_batch, valid_dummy_state_slots
1758 )
1778 - if not skip_attn_for_dummy_run and context_len:
1779 - assert block_tables is not None
1780 - set_dummy_context(
1781 - input_batch,
1782 - self.block_tables,
1783 - context_len,
1784 - self.kv_cache_config.num_blocks,
1785 - self.max_model_len,
1786 - input_block_tables=block_tables,
1787 - )
1759 + if context_len:
1760 + set_dummy_context(
1761 + input_batch,
1762 + self.block_tables,
1763 + context_len,
1764 + self.kv_cache_config.num_blocks,
1765 + self.max_model_len,
1766 + input_block_tables=block_tables,
1767 + )
1768
1789 - attn_metadata = None
1790 - slot_mappings_by_layer = None
1791 - if not (dummy_run and skip_attn_for_dummy_run):
1792 - input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
1793 - (
1794 - self.pcp_manager.input_buffers
1795 - if self.pcp_manager is not None
1796 - else self.input_buffers
1797 - ).dcp_local_seq_lens,
1769 + if self.dcp_size > 1 and not (dummy_run and skip_attn_for_dummy_run):
1770 + input_buffers = self.input_buffers
1771 + if self.pcp_manager is not None:
1772 + input_buffers = self.pcp_manager.input_buffers
1773 + input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
1774 + input_buffers.dcp_local_seq_lens,
1775 input_batch.seq_lens,
1776 input_batch.num_reqs,
1777 self.dcp_size,
@@ -1802,6 +1779,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1779 self.cp_interleave,
1780 num_reqs_padded=input_batch.num_reqs_after_padding,
1781 )
1782 +
1783 + attn_metadata = None
1784 + slot_mappings_by_layer = None
1785 ubatch_state: UBatchState | None = None
1786 if batch_desc.num_ubatches > 1:
1787 assert self.ubatch_runner is not None
@@ -2046,16 +2026,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2026 return ModelRunnerOutput.with_ec_conn_output(output, ec_connector_output)
2027
2028 # Last rank: sample tokens
2029 + assert hidden_states is not None
2030 draft_hidden_states = hidden_states
2050 - assert draft_hidden_states is not None
2051 - hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
2052 - self.pcp_manager, hidden_states, input_batch
2053 - )
2054 - if self.pcp_manager is not None and aux_hidden_states is not None:
2055 - aux_hidden_states = [
2056 - self.pcp_manager.restore_hidden_states(states)
2057 - for states in aux_hidden_states
2058 - ]
2031 + if self.pcp_manager is not None:
2032 + hidden_states, aux_hidden_states, input_batch = (
2033 + self.pcp_manager.restore_for_sampling(hidden_states, aux_hidden_states)
2034 + )
2035
2036 sampler_output, num_sampled, num_rejected = self.sample(
2037 hidden_states, input_batch, grammar_output
@@ -2305,13 +2281,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2281 self.eplb.suppressed = suppressed
2282
2283 def setup_eplb_from_mapping(
2308 - self,
2309 - expanded_physical_to_logical: torch.Tensor,
2284 + self, expanded_physical_to_logical: torch.Tensor
2285 ) -> None:
2286 self.eplb.setup_from_mapping(
2312 - self.model,
2313 - self.model_config,
2314 - expanded_physical_to_logical,
2287 + self.model, self.model_config, expanded_physical_to_logical
2288 )
2289
2290 ########### EPLB methods end ###########
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+1 -3
@@ -27,9 +27,7 @@ class EncoderDecoderAttnMetadata(ModelSpecificAttnMetadata):
27 encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
28
29 def get_extra_common_attn_kwargs(
30 - self,
31 - kv_cache_group_id: int,
32 - num_reqs: int,
30 + self, kv_cache_group_id: int, num_reqs: int
31 ) -> dict[str, Any]:
32 encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id)
33 if encoder_seq_lens is None:
vllm/v1/worker/gpu/model_states/encoder_only.py
+1 -4
@@ -188,10 +188,7 @@ class EncoderOnlyModelState(DefaultModelState):
188 return attn_metadata
189
190 def _build_encoder_attn_metadata(
191 - self,
192 - input_batch: InputBatch,
193 - cudagraph_mode: CUDAGraphMode,
194 - for_capture: bool,
191 + self, input_batch: InputBatch, cudagraph_mode: CUDAGraphMode, for_capture: bool
192 ) -> dict[str, Any]:
193 if cudagraph_mode == CUDAGraphMode.FULL:
194 num_reqs = input_batch.num_reqs_after_padding
vllm/v1/worker/gpu/model_states/interface.py
+2 -6
@@ -28,16 +28,12 @@ class ModelSpecificAttnMetadata:
28 """Base class for model-specific attention metadata."""
29
30 def get_extra_common_attn_kwargs(
31 - self,
32 - kv_cache_group_id: int,
33 - num_reqs: int,
31 + self, kv_cache_group_id: int, num_reqs: int
32 ) -> dict[str, Any]:
33 return {}
34
35 def get_extra_attn_kwargs(
38 - self,
39 - attn_metadata_builder: Any,
40 - num_reqs: int,
36 + self, attn_metadata_builder: Any, num_reqs: int
37 ) -> dict[str, Any]:
38 return {}
39
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+3 -9
@@ -43,16 +43,12 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
43 num_decode_draft_tokens_cpu: torch.Tensor | None = None
44
45 def get_extra_common_attn_kwargs(
46 - self,
47 - kv_cache_group_id: int,
48 - num_reqs: int,
46 + self, kv_cache_group_id: int, num_reqs: int
47 ) -> dict[str, Any]:
48 return {"is_prefilling": self.is_prefilling[:num_reqs]}
49
50 def get_extra_attn_kwargs(
53 - self,
54 - attn_metadata_builder: Any,
55 - num_reqs: int,
51 + self, attn_metadata_builder: Any, num_reqs: int
52 ) -> dict[str, Any]:
53 if not isinstance(
54 attn_metadata_builder,
@@ -338,9 +334,7 @@ class MambaHybridModelState(DefaultModelState):
334 )
335 if self.recoverssm is not None:
336 self.recoverssm.record_step(
341 - attn_metadata,
342 - attn_groups,
343 - for_capture=for_capture,
337 + attn_metadata, attn_groups, for_capture=for_capture
338 )
339 return attn_metadata
340
vllm/v1/worker/gpu/pcp_manager.py
+21 -64
@@ -192,8 +192,7 @@ class PCPManager:
192
193 @staticmethod
194 def _reorder_segments(
195 - segments: list[RankSegment],
196 - is_prefilling: np.ndarray,
195 + segments: list[RankSegment], is_prefilling: np.ndarray
196 ) -> list[RankSegment]:
197 """Order this rank's rows decodes-first, then prefills, canonically."""
198
@@ -218,9 +217,7 @@ class PCPManager:
217 return segments
218
219 def replicated_requests(
221 - self,
222 - num_scheduled_tokens: np.ndarray,
223 - is_prefilling: np.ndarray,
220 + self, num_scheduled_tokens: np.ndarray, is_prefilling: np.ndarray
221 ) -> np.ndarray:
222 """Per global request, whether every PCP rank gets the whole query."""
223 num_chunks = 2 * self.pcp_world_size
@@ -407,10 +404,7 @@ class PCPManager:
404 return self._input_buffers
405
406 def partition_batch(
410 - self,
411 - input_batch: InputBatch,
412 - padded_num_tokens: int | None = None,
413 - padded_num_reqs: int | None = None,
407 + self, input_batch: InputBatch, batch_desc: "BatchExecutionDescriptor"
408 ) -> InputBatch:
409 assert self._input_buffers is not None
410 input_buffers = self._input_buffers
@@ -422,6 +416,13 @@ class PCPManager:
416 num_computed_tokens = global_batch.num_computed_tokens_np
417 is_prefilling = global_batch.is_prefilling_np
418
419 + padded_num_tokens = None
420 + padded_num_reqs = None
421 + if batch_desc.cg_mode != CUDAGraphMode.NONE:
422 + padded_num_tokens = batch_desc.num_tokens
423 + if batch_desc.cg_mode == CUDAGraphMode.FULL:
424 + padded_num_reqs = batch_desc.num_reqs
425 +
426 segments_by_rank, per_rank_num_tokens = self._build_batch_layout(
427 num_scheduled_tokens,
428 num_computed_tokens,
@@ -442,9 +443,7 @@ class PCPManager:
443
444 num_local_reqs = len(local_segments)
445 num_reqs_after_padding = self._resolve_num_reqs_after_padding(
445 - global_batch,
446 - padded_num_reqs,
447 - num_local_reqs,
446 + global_batch, padded_num_reqs, num_local_reqs
447 )
448 if num_reqs_after_padding > input_buffers.max_num_reqs:
449 raise RuntimeError(
@@ -714,8 +713,7 @@ class PCPManager:
713 return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size]
714
715 def _convert_to_gathered_slot_mappings(
717 - self,
718 - global_batch_slot_mappings: torch.Tensor,
716 + self, global_batch_slot_mappings: torch.Tensor
717 ) -> torch.Tensor:
718 assert self._padded_gather_idx is not None
719 assert self._gathered_kv_write_mask is not None
@@ -754,9 +752,7 @@ class PCPManager:
752 return self.draft_prefill_batch or input_buffers
753
754 def prepare_draft_prefill(
757 - self,
758 - input_batch: InputBatch,
759 - input_ids: torch.Tensor,
755 + self, input_batch: InputBatch, input_ids: torch.Tensor
756 ) -> None:
757 if input_batch is not self._global_batch or self._local_batch is None:
758 return
@@ -789,54 +785,15 @@ class PCPManager:
785 return last_hidden_states, hidden_states
786
787 def restore_for_sampling(
792 - self,
793 - hidden_states: torch.Tensor,
794 - ) -> tuple[torch.Tensor, InputBatch]:
788 + self, hidden_states: torch.Tensor, aux_hidden_states: list[torch.Tensor] | None
789 + ) -> tuple[torch.Tensor, list[torch.Tensor] | None, InputBatch]:
790 assert self._global_batch is not None
796 - return self.restore_hidden_states(hidden_states), self._global_batch
797 -
798 -
799 -def maybe_partition_pcp_batch(
800 - manager: PCPManager | None,
801 - input_batch: InputBatch,
802 - batch_desc: "BatchExecutionDescriptor",
803 -) -> InputBatch:
804 - if manager is None:
805 - return input_batch
806 -
807 - padded_num_tokens = None
808 - padded_num_reqs = None
809 - if batch_desc.cg_mode != CUDAGraphMode.NONE:
810 - padded_num_tokens = batch_desc.num_tokens
811 - if batch_desc.cg_mode == CUDAGraphMode.FULL:
812 - padded_num_reqs = batch_desc.num_reqs
813 -
814 - return manager.partition_batch(
815 - input_batch,
816 - padded_num_tokens=padded_num_tokens,
817 - padded_num_reqs=padded_num_reqs,
818 - )
819 -
820 -
821 -def maybe_get_pcp_dummy_slot_mappings(
822 - manager: PCPManager | None,
823 - block_tables: BlockTables,
824 - num_tokens: int,
825 -) -> torch.Tensor:
826 - if manager is None:
827 - return block_tables.get_dummy_slot_mappings(num_tokens)
828 - return manager.get_dummy_slot_mappings(num_tokens)
829 -
830 -
831 -def maybe_restore_pcp_for_sampling(
832 - manager: PCPManager | None,
833 - hidden_states: torch.Tensor | None,
834 - input_batch: InputBatch,
835 -) -> tuple[torch.Tensor, InputBatch]:
836 - assert hidden_states is not None
837 - if manager is None:
838 - return hidden_states, input_batch
839 - return manager.restore_for_sampling(hidden_states)
791 + hidden_states = self.restore_hidden_states(hidden_states)
792 + if aux_hidden_states is not None:
793 + aux_hidden_states = [
794 + self.restore_hidden_states(states) for states in aux_hidden_states
795 + ]
796 + return hidden_states, aux_hidden_states, self._global_batch
797
798
799 def maybe_build_pcp_manager(
vllm/v1/worker/gpu/spec_decode/adaptive_verification.py
+4 -7
@@ -255,9 +255,7 @@ class AdaptiveVerificationManager:
255 logger.debug("DSpark cost tables: %s", self.cost_tables)
256
257 def record_confidences(
258 - self,
259 - confidence_probs: torch.Tensor,
260 - input_batch: "InputBatch",
258 + self, confidence_probs: torch.Tensor, input_batch: "InputBatch"
259 ) -> None:
260 """Publish this step's raw confidences for the ranking kernel and start
261 copying them to the CPU, where a later step's budget reads them."""
@@ -305,9 +303,9 @@ class AdaptiveVerificationManager:
303 )
304 num_non_draft_tokens = scheduled_tokens - scheduled_drafts
305 slots = np.fromiter(
308 - (self.req_states.req_id_to_index[req_id] for req_id in req_ids),
306 + map(self.req_states.req_id_to_index.__getitem__, req_ids),
307 dtype=np.int32,
310 - count=len(req_ids),
308 + count=num_reqs,
309 )
310 stale_confidences = self._stale_confidences[self._stale_idx].np[slots]
311 survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1)
@@ -469,8 +467,7 @@ def maybe_create_adaptive_verification_manager(
467 # The selector rejects unsupported backends, but models that
468 # hard-wire theirs (e.g. DeepSeek-V4) never go through it.
469 backend = get_query_lens_mismatch_unsupported_backend(
472 - attn_groups,
473 - checked_layer_names=target_layer_names,
470 + attn_groups, checked_layer_names=target_layer_names
471 )
472 if backend is not None:
473 raise ValueError(
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py
+3 -10
@@ -53,10 +53,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
53 return
54
55 self.inputs_embeds = torch.zeros(
56 - self.max_num_tokens,
57 - self.hidden_size,
58 - dtype=self.dtype,
59 - device=self.device,
56 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=self.device
57 )
58
59 # Lifecycle hooks for model-specific optimizations. Subclasses override
@@ -140,10 +137,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
137
138 # Initialize cudagraph manager for draft decodes (draft positions > 0).
139 self.decode_cudagraph_manager = SpeculatorCudaGraphManager(
143 - self.vllm_config,
144 - self.device,
145 - cudagraph_mode,
146 - decode_query_len=1,
140 + self.vllm_config, self.device, cudagraph_mode, decode_query_len=1
141 )
142
143 def capture(self) -> None:
@@ -272,8 +266,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
266
267 if self.pcp_manager is not None:
268 self.pcp_manager.prepare_draft_prefill(
275 - input_batch,
276 - self.input_buffers.input_ids[:num_tokens_padded],
269 + input_batch, self.input_buffers.input_ids[:num_tokens_padded]
270 )
271 prefill = self.pcp_manager.draft_prefill_batch
272 if prefill is not None:
vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py
+11 -10
@@ -11,7 +11,7 @@ from vllm.v1.worker.gpu.attn_utils import (
11 build_slot_mappings_by_layer,
12 )
13 from vllm.v1.worker.gpu.block_table import BlockTables
14 -from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
14 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
15 from vllm.v1.worker.gpu.cudagraph_utils import (
16 AttentionState,
17 BatchExecutionDescriptor,
@@ -42,15 +42,16 @@ def _prepare_dflash_inputs_to_capture(
42 attn_metadata = None
43 if not skip_attn:
44 query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
45 - input_batch.dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
46 - input_buffers.dcp_local_seq_lens,
47 - input_batch.seq_lens,
48 - input_batch.num_reqs,
49 - block_tables.cp_size,
50 - block_tables.cp_rank,
51 - block_tables.cp_interleave,
52 - num_reqs_padded=input_batch.num_reqs_after_padding,
53 - )
45 + if block_tables.cp_size > 1:
46 + input_batch.dcp_local_seq_lens = prepare_dcp_local_seq_lens(
47 + input_buffers.dcp_local_seq_lens,
48 + input_batch.seq_lens,
49 + input_batch.num_reqs,
50 + block_tables.cp_size,
51 + block_tables.cp_rank,
52 + block_tables.cp_interleave,
53 + num_reqs_padded=input_batch.num_reqs_after_padding,
54 + )
55 attn_metadata = build_attn_metadata(
56 attn_groups=attn_groups,
57 num_reqs=num_reqs,
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+1 -4
@@ -149,10 +149,7 @@ class DSparkSpeculator(DFlashSpeculator):
149 )
150 if self.draft_watermarker is not None:
151 sampled = self.draft_watermarker.sample(
152 - logits,
153 - sampled,
154 - idx_map,
155 - self.temperature,
152 + logits, sampled, idx_map, self.temperature
153 )
154 return sampled
155
vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+1 -2
@@ -134,8 +134,7 @@ class RejectionSampler:
134 if expanded_logits:
135 if self.enable_adaptive_verification:
136 # Adaptive verification keeps the true per-request boundaries
137 - # on device only; cu_num_logits_np holds the pre-compacted
138 - # layout.
137 + # on device only; cu_num_logits_np holds the pre-compacted layout.
138 cu_num_generated_tokens = cu_num_logits.clone()
139 else:
140 cu_num_generated_tokens = cu_num_logits_np.tolist()
vllm/v1/worker/gpu/spec_decode/speculator.py
+5 -13
@@ -26,7 +26,7 @@ from vllm.v1.worker.gpu.attn_utils import (
26 init_attn_backend,
27 )
28 from vllm.v1.worker.gpu.block_table import BlockTables
29 -from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
29 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
30 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
31 from vllm.v1.worker.gpu.dp_utils import DPSyncState
32 from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
@@ -202,9 +202,7 @@ class DraftModelSpeculator(BaseSpeculator):
202
203 @abstractmethod
204 def load_draft_model(
205 - self,
206 - target_model: nn.Module,
207 - target_attn_layer_names: set[str],
205 + self, target_model: nn.Module, target_attn_layer_names: set[str]
206 ) -> nn.Module:
207 pass
208
@@ -337,7 +335,7 @@ class DraftModelSpeculator(BaseSpeculator):
335 if dcp_local_seq_lens is None and self.block_tables.cp_size > 1:
336 # Draft steps advance and rewind their own global sequence lengths,
337 # so the target model's DCP-local lengths may already be stale.
340 - dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
338 + dcp_local_seq_lens = prepare_dcp_local_seq_lens(
339 self.input_buffers.dcp_local_seq_lens,
340 self.input_buffers.seq_lens,
341 num_reqs,
@@ -422,10 +420,7 @@ class DraftModelSpeculator(BaseSpeculator):
420 )
421 if self.draft_watermarker is not None:
422 sampled = self.draft_watermarker.sample(
425 - logits,
426 - sampled,
427 - idx_mapping,
428 - temperature,
423 + logits, sampled, idx_mapping, temperature
424 )
425 elif self.use_local_argmax_reduction:
426 return self.model.get_top_tokens(hidden_states)
@@ -436,10 +431,7 @@ class DraftModelSpeculator(BaseSpeculator):
431 return sampled
432
433 def _maybe_predict_acceptance(
439 - self,
440 - logits: torch.Tensor,
441 - idx_mapping: torch.Tensor,
442 - draft_step: torch.Tensor,
434 + self, logits: torch.Tensor, idx_mapping: torch.Tensor, draft_step: torch.Tensor
435 ) -> None:
436 if self.acceptance_estimator is not None:
437 self.acceptance_estimator.predict(
vllm/v1/worker/gpu/states.py
+1 -4
@@ -78,10 +78,7 @@ class RequestState:
78 )
79
80 self.next_prefill_tokens = torch.zeros(
81 - num_prefill_lookahead,
82 - self.max_num_reqs,
83 - dtype=torch.int32,
84 - device=device,
81 + num_prefill_lookahead, self.max_num_reqs, dtype=torch.int32, device=device
82 )
83
84 @property
vllm/v1/worker/gpu/ubatch_utils.py
+2 -2
@@ -24,7 +24,7 @@ from vllm.sequence import IntermediateTensors
24 from vllm.utils.torch_utils import current_stream
25 from vllm.v1.kv_cache_interface import KVCacheConfig
26 from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
27 -from vllm.v1.worker.gpu.cp_utils import maybe_prepare_dcp_local_seq_lens
27 +from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
28 from vllm.v1.worker.gpu.input_batch import InputBatch
29 from vllm.v1.worker.gpu.model_states.interface import ModelState
30 from vllm.v1.worker.ubatch_utils import (
@@ -150,7 +150,7 @@ def _slice_input_batch(
150 dcp_local_seq_lens = None
151 if dcp_size > 1:
152 assert dcp_local_seq_lens_buf is not None
153 - dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
153 + dcp_local_seq_lens = prepare_dcp_local_seq_lens(
154 dcp_local_seq_lens_buf,
155 seq_lens,
156 num_reqs,