A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI] Fix Pre-commit Issue (#25497)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 23, 2025 at 16:09 UTC
8b8a8afc895ab200cf9aa27ea47c5fc6dbef443b
1 file changed
+7
-2
vllm/v1/worker/gpu_model_runner.py
+7
-2
@@ -2367,7 +2367,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2367
sampling_metadata: SamplingMetadata,
2368
hidden_states: torch.Tensor,
2369
sample_hidden_states: torch.Tensor,
2370
- aux_hidden_states: Optional[torch.Tensor],
2370
+ aux_hidden_states: Optional[list[torch.Tensor]],
2371
spec_decode_metadata: Optional[SpecDecodeMetadata],
2372
common_attn_metadata: CommonAttentionMetadata,
2373
) -> Union[list[list[int]], torch.Tensor]:
@@ -2387,6 +2387,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2387
else:
2388
indices = []
2389
offset = 0
2390
+ assert spec_decode_metadata is not None
2391
for num_draft, tokens in zip(
2392
spec_decode_metadata.num_draft_tokens,
2393
sampled_token_ids):
@@ -2437,6 +2438,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2438
# TODO(woosuk): Support M-RoPE.
2439
target_positions = self.positions.gpu[:num_scheduled_tokens]
2440
if self.use_aux_hidden_state_outputs:
2441
+ assert aux_hidden_states is not None
2442
target_hidden_states = torch.cat(
2443
[h[:num_scheduled_tokens] for h in aux_hidden_states],
2444
dim=-1)
@@ -2462,6 +2464,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2464
# TODO(woosuk): Support M-RoPE.
2465
target_positions = self.positions.gpu[token_indices]
2466
if self.use_aux_hidden_state_outputs:
2467
+ assert aux_hidden_states is not None
2468
target_hidden_states = torch.cat(
2469
[h[token_indices] for h in aux_hidden_states], dim=-1)
2470
else:
@@ -2897,7 +2900,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2900
assert not create_mixed_batch
2901
num_reqs = cdiv(num_tokens, max_query_len)
2902
assert num_reqs <= max_num_reqs, \
2900
- "Do not capture num_reqs > max_num_reqs for uniform batch"
2903
+ f"Do not capture num_reqs {num_reqs} > max_num_reqs " \
2904
+ f"{max_num_reqs} for uniform batch. Num tokens: " \
2905
+ f"{num_tokens}, max_query_len: {max_query_len}"
2906
num_scheduled_tokens_list = [max_query_len] * num_reqs
2907
if num_tokens % max_query_len != 0:
2908
num_scheduled_tokens_list[-1] = num_tokens % max_query_len