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 10, 11:30 UTC
[Misc] Tidy up some spec decode logic in GPUModelRunner (#31591)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jan 8, 2026 at 09:10 UTC
a3d909ad2b2936333af88f308ae2147e1dd7a6d8
3 files changed
+51
-50
vllm/v1/worker/gpu/model_runner.py
+4
@@ -139,6 +139,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
139
# CUDA graphs.
140
self.cudagraph_manager = CudaGraphManager(self.vllm_config, self.device)
141
142
+ def update_max_model_len(self, max_model_len: int) -> None:
143
+ self.max_model_len = max_model_len
144
+ self.req_states.max_model_len = max_model_len
145
+
146
def get_supported_tasks(self) -> tuple[str]:
147
return ("generate",)
148
vllm/v1/worker/gpu_model_runner.py
+46
-49
@@ -452,6 +452,11 @@ class GPUModelRunner(
452
self.num_spec_tokens = 0
453
if self.speculative_config:
454
self.num_spec_tokens = self.speculative_config.num_speculative_tokens
455
+ draft_config = self.speculative_config.draft_model_config
456
+ if draft_config is not None and draft_config.max_model_len is not None:
457
+ self.effective_drafter_max_model_len = draft_config.max_model_len
458
+ else:
459
+ self.effective_drafter_max_model_len = self.max_model_len
460
461
# Request states.
462
self.requests: dict[str, CachedRequestState] = {}
@@ -674,6 +679,13 @@ class GPUModelRunner(
679
self.kv_connector_output: KVConnectorOutput | None = None
680
self.layerwise_nvtx_hooks_registered = False
681
682
+ def update_max_model_len(self, max_model_len: int) -> None:
683
+ self.max_model_len = max_model_len
684
+ if self.speculative_config:
685
+ draft_config = self.speculative_config.draft_model_config
686
+ if draft_config is None or draft_config.max_model_len is None:
687
+ self.effective_drafter_max_model_len = self.max_model_len
688
+
689
def reset_mm_cache(self) -> None:
690
if self.mm_budget:
691
self.mm_budget.reset_cache()
@@ -3399,54 +3411,41 @@ class GPUModelRunner(
3411
self._copy_draft_token_ids_to_cpu(scheduler_output)
3412
3413
spec_config = self.speculative_config
3402
- use_padded_batch_for_eagle = (
3403
- spec_config is not None
3404
- and spec_config.use_eagle()
3405
- and not spec_config.disable_padded_drafter_batch
3406
- )
3407
- effective_drafter_max_model_len = self.max_model_len
3408
- if effective_drafter_max_model_len is None:
3409
- effective_drafter_max_model_len = self.model_config.max_model_len
3410
- if (
3411
- spec_config is not None
3412
- and spec_config.draft_model_config is not None
3413
- and spec_config.draft_model_config.max_model_len is not None
3414
- ):
3415
- effective_drafter_max_model_len = (
3416
- spec_config.draft_model_config.max_model_len
3414
+ propose_drafts_after_bookkeeping = False
3415
+ if spec_config is not None:
3416
+ input_fits_in_drafter = spec_decode_common_attn_metadata is not None and (
3417
+ spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
3418
+ <= self.effective_drafter_max_model_len
3419
)
3418
- input_fits_in_drafter = spec_decode_common_attn_metadata and (
3419
- spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens
3420
- <= effective_drafter_max_model_len
3421
- )
3422
- if use_padded_batch_for_eagle:
3423
- assert self.speculative_config is not None
3424
- assert isinstance(self.drafter, EagleProposer)
3425
- sampled_token_ids = sampler_output.sampled_token_ids
3426
- if input_fits_in_drafter:
3420
+ if spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch:
3421
# EAGLE speculative decoding can use the GPU sampled tokens
3422
# as inputs, and does not need to wait for bookkeeping to finish.
3429
- propose_draft_token_ids(sampled_token_ids)
3430
- elif self.valid_sampled_token_count_event is not None:
3431
- assert spec_decode_common_attn_metadata is not None
3432
- next_token_ids, valid_sampled_tokens_count = (
3433
- self.drafter.prepare_next_token_ids_padded(
3434
- spec_decode_common_attn_metadata,
3435
- sampled_token_ids,
3436
- self.requests,
3437
- self.input_batch,
3438
- self.discard_request_mask.gpu,
3423
+ assert isinstance(self.drafter, EagleProposer)
3424
+ sampled_token_ids = sampler_output.sampled_token_ids
3425
+ if input_fits_in_drafter:
3426
+ propose_draft_token_ids(sampled_token_ids)
3427
+ elif self.valid_sampled_token_count_event is not None:
3428
+ assert spec_decode_common_attn_metadata is not None
3429
+ next_token_ids, valid_sampled_tokens_count = (
3430
+ self.drafter.prepare_next_token_ids_padded(
3431
+ spec_decode_common_attn_metadata,
3432
+ sampled_token_ids,
3433
+ self.requests,
3434
+ self.input_batch,
3435
+ self.discard_request_mask.gpu,
3436
+ )
3437
)
3440
- )
3441
- self._copy_valid_sampled_token_count(
3442
- next_token_ids, valid_sampled_tokens_count
3443
- )
3444
- # Since we couldn't run the drafter,
3445
- # just use zeros for the draft tokens.
3446
- self._draft_token_ids = torch.zeros(
3447
- 1, device=self.device, dtype=torch.int32
3448
- ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
3449
- self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
3438
+ self._copy_valid_sampled_token_count(
3439
+ next_token_ids, valid_sampled_tokens_count
3440
+ )
3441
+ # Since we couldn't run the drafter,
3442
+ # just use zeros for the draft tokens.
3443
+ self._draft_token_ids = torch.zeros(
3444
+ 1, device=self.device, dtype=torch.int32
3445
+ ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)
3446
+ self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)
3447
+ else:
3448
+ propose_drafts_after_bookkeeping = input_fits_in_drafter
3449
3450
with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
3451
(
@@ -3466,17 +3465,14 @@ class GPUModelRunner(
3465
spec_decode_metadata,
3466
)
3467
3469
- if (
3470
- self.speculative_config
3471
- and not use_padded_batch_for_eagle
3472
- and input_fits_in_drafter
3473
- ):
3468
+ if propose_drafts_after_bookkeeping:
3469
# ngram and other speculative decoding methods use the sampled
3470
# tokens on the CPU, so they are run after bookkeeping.
3471
propose_draft_token_ids(valid_sampled_token_ids)
3472
3473
with record_function_or_nullcontext("gpu_model_runner: eplb"):
3474
self.eplb_step()
3475
+
3476
with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
3477
output = ModelRunnerOutput(
3478
req_ids=req_ids_output_copy,
@@ -3494,6 +3490,7 @@ class GPUModelRunner(
3490
3491
if not self.use_async_scheduling:
3492
return output
3493
+
3494
with record_function_or_nullcontext(
3495
"gpu_model_runner: AsyncGPUModelRunnerOutput"
3496
):
vllm/v1/worker/gpu_worker.py
+1
-1
@@ -390,7 +390,7 @@ class Worker(WorkerBase):
390
"""
391
self.model_config.max_model_len = max_model_len
392
if self.model_runner is not None:
393
- self.model_runner.max_model_len = max_model_len
393
+ self.model_runner.update_max_model_len(max_model_len)
394
logger.debug("Updated max_model_len to %d", max_model_len)
395
396
def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: