[MRV2] Minor model_runner.py code cleanup (#58610)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Sep 25, 2026 at 07:59 UTC
234df712ed6da81d8c8213a16cd913a747380214
2 files changed
+3
-27
tests/v1/worker/test_gpu_model_runner_v2.py
-14
@@ -21,20 +21,6 @@ from vllm.v1.worker.gpu.block_table import BlockTables
21
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
22
23
24
-def test_prepare_padding_mask_marks_sequence_parallel_padding():
25
- runner = GPUModelRunner.__new__(GPUModelRunner)
26
- runner.input_buffers = SimpleNamespace(is_padding=torch.empty(8, dtype=torch.bool))
27
-
28
- mask = runner._prepare_padding_mask(1, 8)
29
-
30
- assert mask.tolist() == [False, True, True, True, True, True, True, True]
31
- assert mask.data_ptr() == runner.input_buffers.is_padding.data_ptr()
32
-
33
- mask = runner._prepare_padding_mask(0, 8)
34
-
35
- assert mask.all()
36
-
37
-
24
def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
25
runner = GPUModelRunner.__new__(GPUModelRunner)
26
runner.max_model_len = 262144
vllm/v1/worker/gpu/model_runner.py
+3
-13
@@ -797,8 +797,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
797
assert num_tokens % self.decode_query_len == 0
798
elif self.speculator is not None:
799
num_reqs = min(
800
- num_reqs,
801
- self.max_num_tokens // self.speculator.num_query_per_req,
800
+ num_reqs, self.max_num_tokens // self.speculator.num_query_per_req
801
)
802
# Distribute the remainder evenly so no dummy request exceeds
803
# ceil(num_tokens / num_reqs) <= max_model_len tokens.
@@ -1257,14 +1256,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1256
num_reqs, num_toks, max_query_len, batch_state.has_prefill
1257
)
1258
1260
- def _prepare_padding_mask(
1261
- self, num_tokens: int, num_tokens_after_padding: int
1262
- ) -> torch.Tensor:
1263
- is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
1264
- is_padding[:num_tokens].fill_(False)
1265
- is_padding[num_tokens:].fill_(True)
1266
- return is_padding
1267
-
1259
def prepare_inputs(
1260
self,
1261
scheduler_output: SchedulerOutput,
@@ -1277,9 +1268,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1268
assert num_tokens > 0
1269
is_padding = self.input_buffers.is_padding[:num_tokens_after_padding]
1270
if envs.VLLM_MOE_SKIP_PADDING:
1280
- is_padding = self._prepare_padding_mask(
1281
- num_tokens, num_tokens_after_padding
1282
- )
1271
+ is_padding[:num_tokens].fill_(False)
1272
+ is_padding[num_tokens:].fill_(True)
1273
1274
req_ids = batch_req_state.req_ids
1275
num_scheduled_tokens_np = batch_req_state.num_scheduled_tokens