[MRV2] Revert explicit Triton JIT warmup migration (#56654)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Woosuk Kwon committed
Sep 12, 2026 at 21:07 UTC
2671fedfc7ae604761990603fc736c0c4f21de57
6 files changed
+97
-228
tests/v1/spec_decode/test_dflash_prepare_inputs.py
+1
-1
@@ -43,7 +43,7 @@ def _run_prepare(
43
num_scheduled_tokens=np.array([4], dtype=np.int32),
44
positions=torch.tensor(target_positions, dtype=torch.int64, device=device),
45
query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device),
46
- idx_mapping=torch.tensor([2], dtype=torch.int64, device=device),
46
+ idx_mapping=torch.tensor([2], dtype=torch.int32, device=device),
47
)
48
query_slot_mapping = torch.full(
49
(max_num_tokens,), -2, dtype=torch.int64, device=device
tests/v1/worker/test_gpu_autoregressive_speculator.py
-1
@@ -157,7 +157,6 @@ def test_mm_support_configured_after_model_load(monkeypatch):
157
speculator.device = device
158
speculator.max_num_tokens = 4
159
speculator.max_num_reqs = 2
160
- speculator.num_speculative_steps = 1
160
speculator.hidden_size = 3
161
speculator.dtype = torch.float32
162
speculator.draft_model_config = draft_model_config
vllm/v1/worker/gpu/model_runner.py
+44
-53
@@ -188,7 +188,6 @@ logger = init_logger(__name__)
188
189
190
class GPUModelRunner(LoRAModelRunnerMixin):
191
- @JitWarmupRegistry.capture
191
def __init__(self, vllm_config: VllmConfig, device: torch.device):
192
self.vllm_config = vllm_config
193
self.model_config = vllm_config.model_config
@@ -203,6 +202,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
202
self.speculative_config is not None and self.speculative_config.use_dspark()
203
)
204
self.observability_config = vllm_config.observability_config
205
+ self.jit_warmup_registry = JitWarmupRegistry(vllm_config)
206
207
self.device = device
208
self.dtype = self.model_config.dtype
@@ -399,7 +399,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
399
logger.info_once("Loading model from scratch...")
400
401
# Capture warmup providers selected while constructing the model.
402
- with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]
402
+ with self.jit_warmup_registry.activate():
403
self.model = model_loader.load_model(
404
vllm_config=self.vllm_config,
405
model_config=self.vllm_config.model_config,
@@ -461,50 +461,47 @@ class GPUModelRunner(LoRAModelRunnerMixin):
461
462
# Initialize samplers. Model states may override via custom_sampler().
463
if self.is_last_pp_rank and not self.is_pooling_model:
464
- with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]
465
- sampler_kwargs: dict[str, Any] = {
466
- "max_num_reqs": self.max_num_reqs,
467
- "vocab_size": self.vocab_size,
468
- "device": self.device,
469
- "req_states": self.req_states,
470
- "logprobs_mode": self.model_config.logprobs_mode,
471
- "num_speculative_tokens": self.decode_query_len,
472
- "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
473
- "enable_trace_replay": self.model_config.enable_trace_replay,
474
- "reasoning_config": self.vllm_config.reasoning_config,
475
- "return_sampling_mask": self.model_config.return_sampling_mask,
476
- }
477
- if self.vllm_config.watermark_config is None:
478
- self.sampler = Sampler(**sampler_kwargs)
479
- else:
480
- wm_config = self.vllm_config.watermark_config
481
- watermarker = create_watermarker(wm_config)
482
- if self.speculative_config is not None:
483
- watermarker = create_speculative_target_watermarker(watermarker)
484
- self.sampler = GPUWatermarkSampler(
485
- watermarker,
486
- deduplicate_contexts=wm_config.deduplicate_contexts,
487
- deduplicate_contexts_max_history=(
488
- wm_config.deduplicate_contexts_max_history
489
- ),
490
- **sampler_kwargs,
491
- )
492
- custom = self.model_state.custom_sampler(self.sampler)
493
-
494
- if custom:
495
- self.vllm_config._check_watermarking_unsupported(
496
- custom_sampler=True
497
- )
498
- self.sampler, self.rejection_sampler = custom
499
- elif self.speculative_config is not None:
500
- self.rejection_sampler = RejectionSampler(
501
- self.sampler,
502
- self.speculative_config,
503
- self.device,
504
- watermark_key=speculative_target_watermark_key(
505
- self.vllm_config.watermark_config
506
- ),
507
- )
464
+ sampler_kwargs: dict[str, Any] = {
465
+ "max_num_reqs": self.max_num_reqs,
466
+ "vocab_size": self.vocab_size,
467
+ "device": self.device,
468
+ "req_states": self.req_states,
469
+ "logprobs_mode": self.model_config.logprobs_mode,
470
+ "num_speculative_tokens": self.decode_query_len,
471
+ "use_fp64_gumbel": self.model_config.use_fp64_gumbel,
472
+ "enable_trace_replay": self.model_config.enable_trace_replay,
473
+ "reasoning_config": self.vllm_config.reasoning_config,
474
+ "return_sampling_mask": self.model_config.return_sampling_mask,
475
+ }
476
+ if self.vllm_config.watermark_config is None:
477
+ self.sampler = Sampler(**sampler_kwargs)
478
+ else:
479
+ wm_config = self.vllm_config.watermark_config
480
+ watermarker = create_watermarker(wm_config)
481
+ if self.speculative_config is not None:
482
+ watermarker = create_speculative_target_watermarker(watermarker)
483
+ self.sampler = GPUWatermarkSampler(
484
+ watermarker,
485
+ deduplicate_contexts=wm_config.deduplicate_contexts,
486
+ deduplicate_contexts_max_history=(
487
+ wm_config.deduplicate_contexts_max_history
488
+ ),
489
+ **sampler_kwargs,
490
+ )
491
+ custom = self.model_state.custom_sampler(self.sampler)
492
+
493
+ if custom:
494
+ self.vllm_config._check_watermarking_unsupported(custom_sampler=True)
495
+ self.sampler, self.rejection_sampler = custom
496
+ elif self.speculative_config is not None:
497
+ self.rejection_sampler = RejectionSampler(
498
+ self.sampler,
499
+ self.speculative_config,
500
+ self.device,
501
+ watermark_key=speculative_target_watermark_key(
502
+ self.vllm_config.watermark_config
503
+ ),
504
+ )
505
self.prompt_logprobs_worker = PromptLogprobsWorker(
506
self.max_num_reqs,
507
logprobs_mode=self.model_config.logprobs_mode,
@@ -741,7 +738,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
738
self.speculator.init_cudagraph_manager(cudagraph_mode)
739
740
# Capture warmup providers that depend on allocated KV-cache strides.
744
- with self.jit_warmup_registry.activate(): # type: ignore[attr-defined]
741
+ with self.jit_warmup_registry.activate():
742
kv_caches_dict = init_kv_cache(
743
self.compilation_config.static_forward_context,
744
self.kv_cache_config,
@@ -927,12 +924,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
924
# during actual execution.
925
assert self.sampler is not None
926
self.sampler(logits, dummy_input_batch)
930
- if self.rejection_sampler is not None:
931
- # Warm head-dtype logits; scheduler warmup covers processed FP32 logits.
932
- assert self.speculator is not None
933
- self.rejection_sampler(
934
- logits, dummy_input_batch, self.speculator.draft_logits
935
- )
927
928
@torch.inference_mode()
929
def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None:
vllm/v1/worker/gpu/sample/states.py
+1
-5
@@ -4,10 +4,7 @@ import numpy as np
4
import torch
5
6
from vllm.sampling_params import SamplingParams
7
-from vllm.v1.sample.ops.topk_topp_sampler import (
8
- apply_top_k_top_p,
9
- register_top_k_top_p_warmups,
10
-)
7
+from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
8
from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor
9
from vllm.v1.worker.gpu.sample.gumbel import apply_temperature
10
from vllm.v1.worker.gpu.sample.min_p import apply_min_p
@@ -21,7 +18,6 @@ class SamplingStates:
18
def __init__(self, max_num_reqs: int, vocab_size: int):
19
self.max_num_reqs = max_num_reqs
20
self.vocab_size = vocab_size
24
- register_top_k_top_p_warmups()
21
22
self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
23
self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+39
-92
@@ -1,10 +1,8 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import copy
4
-from types import SimpleNamespace
4
from typing import Any
5
7
-import numpy as np
6
import torch
7
import torch.nn as nn
8
@@ -12,15 +10,6 @@ from vllm.config import VllmConfig, replace
10
from vllm.config.compilation import CUDAGraphMode
11
from vllm.forward_context import BatchDescriptor, set_forward_context
12
from vllm.logger import init_logger
15
-from vllm.model_executor.warmup.jit_warmup import (
16
- WarmupChoices,
17
- WarmupIntRange,
18
-)
19
-from vllm.model_executor.warmup.jit_warmup_triton_helper import (
20
- DispatchSpec,
21
- TritonWarmupTensor,
22
- triton_kernel_dispatcher_with_warmup,
23
-)
13
from vllm.triton_utils import tl, triton
14
from vllm.v1.attention.backend import AttentionCGSupport
15
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
@@ -119,7 +108,6 @@ class DFlashSpeculator(DraftModelSpeculator):
108
109
self.query_cudagraph_manager: DFlashCudaGraphManager | None = None
110
self.draft_kv_cache_group_id: int = -1
122
- prepare_dflash_inputs.register_warmup(speculator=self)
111
112
@property
113
def attn_vllm_config(self) -> VllmConfig:
@@ -694,64 +682,6 @@ def _prepare_dflash_inputs_kernel(
682
tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask)
683
684
697
-def _prepare_dflash_warmup_inputs(*, speculator: DFlashSpeculator) -> dict[str, Any]:
698
- int32 = TritonWarmupTensor(torch.int32)
699
- int64 = TritonWarmupTensor(torch.int64)
700
- float32 = TritonWarmupTensor(torch.float32)
701
- input_buffers = SimpleNamespace(
702
- input_ids=int32, positions=int64, query_start_loc=int32, seq_lens=int32
703
- )
704
- gid: Any = WarmupChoices(*speculator.draft_kv_cache_group_ids)
705
- max_target_query_len: Any = WarmupIntRange(1, 257, advance=lambda value: value * 2)
706
- block_table_stride = int(speculator.block_tables.input_block_tables[gid].stride(0))
707
- block_table = TritonWarmupTensor(
708
- torch.int32,
709
- shape=(1, block_table_stride),
710
- strides=(block_table_stride, 1),
711
- )
712
- input_batch = SimpleNamespace(
713
- num_reqs=1,
714
- num_scheduled_tokens=np.array([max_target_query_len]),
715
- positions=int64,
716
- query_start_loc=int32,
717
- idx_mapping=int64,
718
- )
719
- return dict(
720
- input_buffers=input_buffers,
721
- query_slot_mapping=int64,
722
- context_positions=int64,
723
- context_slot_mapping=int64,
724
- sample_indices=int64,
725
- sample_pos=int64,
726
- sample_idx_mapping=int32,
727
- temperature=float32,
728
- seeds=int64,
729
- input_batch=input_batch,
730
- num_sampled=int32,
731
- num_rejected=int32,
732
- last_sampled=int64,
733
- next_prefill_tokens=int32,
734
- input_temperature=float32,
735
- input_seeds=int64,
736
- block_table=block_table,
737
- block_size=speculator.block_tables.kernel_block_sizes[gid],
738
- cp_rank=speculator.block_tables.cp_rank,
739
- cp_size=speculator.block_tables.cp_size,
740
- cp_interleave=speculator.block_tables.cp_interleave,
741
- parallel_drafting_token_id=speculator.parallel_drafting_token_id,
742
- num_query_per_req=speculator.num_query_per_req,
743
- num_speculative_steps=speculator.num_speculative_steps,
744
- max_num_reqs=speculator.max_num_reqs,
745
- max_num_tokens=speculator.max_num_tokens,
746
- max_model_len=speculator.max_model_len,
747
- sample_from_anchor=speculator.sample_from_anchor,
748
- )
749
-
750
-
751
-@triton_kernel_dispatcher_with_warmup(
752
- kernel=_prepare_dflash_inputs_kernel,
753
- warmup_inputs=_prepare_dflash_warmup_inputs,
754
-)
685
def prepare_dflash_inputs(
686
input_buffers: InputBuffers,
687
query_slot_mapping: torch.Tensor,
@@ -788,33 +718,50 @@ def prepare_dflash_inputs(
718
max_num_tokens: int,
719
max_model_len: int,
720
sample_from_anchor: bool = False,
791
-) -> DispatchSpec:
721
+) -> None:
722
num_reqs = input_batch.num_reqs
723
assert num_reqs > 0
724
# Cover the longest possible per-request span (ctx + query). Use the max
725
# per-request query length, not the total token count across the batch.
726
max_target_query_len = int(input_batch.num_scheduled_tokens.max())
727
max_tokens_per_req = max_target_query_len + num_query_per_req
798
- block = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
799
- num_blocks = triton.cdiv(max_tokens_per_req, block)
800
- return (num_reqs, num_blocks), dict(
801
- out_input_ids_ptr=input_buffers.input_ids,
802
- out_query_positions_ptr=input_buffers.positions,
803
- out_query_start_loc_ptr=input_buffers.query_start_loc,
804
- out_seq_lens_ptr=input_buffers.seq_lens,
805
- out_query_slot_mapping_ptr=query_slot_mapping,
806
- out_context_positions_ptr=context_positions,
807
- out_context_slot_mapping_ptr=context_slot_mapping,
808
- out_sample_indices_ptr=sample_indices,
809
- out_sample_pos_ptr=sample_pos,
810
- out_sample_idx_mapping_ptr=sample_idx_mapping,
811
- out_temperature_ptr=temperature,
812
- out_seeds_ptr=seeds,
813
- target_positions_ptr=input_batch.positions,
814
- target_query_start_loc_ptr=input_batch.query_start_loc,
815
- idx_mapping_ptr=input_batch.idx_mapping,
816
- temperature_ptr=input_temperature,
817
- seeds_ptr=input_seeds,
728
+ BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req)))
729
+ num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE)
730
+ _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)](
731
+ input_buffers.input_ids,
732
+ input_buffers.positions,
733
+ input_buffers.query_start_loc,
734
+ input_buffers.seq_lens,
735
+ query_slot_mapping,
736
+ context_positions,
737
+ context_slot_mapping,
738
+ sample_indices,
739
+ sample_pos,
740
+ sample_idx_mapping,
741
+ temperature,
742
+ seeds,
743
+ input_batch.positions,
744
+ input_batch.query_start_loc,
745
+ input_batch.idx_mapping,
746
+ last_sampled,
747
+ next_prefill_tokens,
748
+ num_sampled,
749
+ num_rejected,
750
+ input_temperature,
751
+ input_seeds,
752
+ block_table,
753
+ block_table.stride(0),
754
+ parallel_drafting_token_id,
755
+ block_size,
756
+ num_query_per_req,
757
+ num_speculative_steps,
758
+ max_num_reqs,
759
+ max_num_tokens,
760
+ max_model_len,
761
+ cp_rank,
762
+ SAMPLE_FROM_ANCHOR=sample_from_anchor,
763
PAD_SLOT_ID=PAD_SLOT_ID,
819
- BLOCK_SIZE=block,
764
+ CP_SIZE=cp_size,
765
+ CP_INTERLEAVE=cp_interleave,
766
+ BLOCK_SIZE=BLOCK_SIZE,
767
)
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+12
-76
@@ -9,11 +9,6 @@ from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
from vllm.forward_context import BatchDescriptor, set_forward_context
11
from vllm.logger import init_logger
12
-from vllm.model_executor.warmup.jit_warmup_triton_helper import (
13
- DispatchSpec,
14
- TritonWarmupTensor,
15
- triton_kernel_dispatcher_with_warmup,
16
-)
12
from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
14
from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
@@ -74,9 +69,6 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
69
70
self.cudagraph_manager: SpeculatorCudaGraphManager | None = None
71
77
- _shift_input_ids.register_warmup()
78
- _shift_input_embeds.register_warmup(speculator=self)
79
-
72
def load_draft_model(
73
self,
74
target_model: nn.Module,
@@ -1016,32 +1008,6 @@ def _shift_input_ids_kernel(
1008
tl.store(input_ids_ptr + last_token_index, draft_token)
1009
1010
1019
-def _shift_input_ids_warmup_inputs() -> dict[str, Any]:
1020
- return dict(
1021
- input_ids=TritonWarmupTensor(torch.int32),
1022
- idx_mapping=TritonWarmupTensor(torch.int32),
1023
- query_start_loc=TritonWarmupTensor(torch.int32),
1024
- last_token_indices=TritonWarmupTensor(torch.int64),
1025
- draft_tokens=TritonWarmupTensor(torch.int64),
1026
- num_reqs=1,
1027
- )
1028
-
1029
-
1030
-@triton_kernel_dispatcher_with_warmup(
1031
- kernel=_shift_input_ids_kernel,
1032
- warmup_inputs=_shift_input_ids_warmup_inputs,
1033
-)
1034
-def _shift_input_ids(
1035
- input_ids: torch.Tensor,
1036
- idx_mapping: torch.Tensor,
1037
- query_start_loc: torch.Tensor,
1038
- last_token_indices: torch.Tensor,
1039
- draft_tokens: torch.Tensor,
1040
- num_reqs: int,
1041
-) -> DispatchSpec:
1042
- return (num_reqs,), dict(BLOCK_SIZE=1024)
1043
-
1044
-
1011
@triton.jit
1012
def _shift_input_embeds_kernel(
1013
input_embeds_ptr,
@@ -1099,44 +1065,6 @@ def _shift_input_embeds_kernel(
1065
)
1066
1067
1102
-def _shift_input_embeds_warmup_inputs(
1103
- *, speculator: MultiModuleMTPSpeculator
1104
-) -> dict[str, Any]:
1105
- hidden = TritonWarmupTensor(
1106
- speculator.dtype,
1107
- shape=(1, speculator.hidden_size),
1108
- strides=(speculator.hidden_size, 1),
1109
- )
1110
- return dict(
1111
- input_embeds=hidden,
1112
- draft_embeds=hidden,
1113
- idx_mapping=TritonWarmupTensor(torch.int32),
1114
- query_start_loc=TritonWarmupTensor(torch.int32),
1115
- last_token_indices=TritonWarmupTensor(torch.int64),
1116
- num_reqs=1,
1117
- )
1118
-
1119
-
1120
-@triton_kernel_dispatcher_with_warmup(
1121
- kernel=_shift_input_embeds_kernel,
1122
- warmup_inputs=_shift_input_embeds_warmup_inputs,
1123
-)
1124
-def _shift_input_embeds(
1125
- input_embeds: torch.Tensor,
1126
- draft_embeds: torch.Tensor,
1127
- idx_mapping: torch.Tensor,
1128
- query_start_loc: torch.Tensor,
1129
- last_token_indices: torch.Tensor,
1130
- num_reqs: int,
1131
-) -> DispatchSpec:
1132
- hidden_size = input_embeds.shape[-1]
1133
- return (num_reqs, triton.cdiv(hidden_size, 256)), dict(
1134
- hidden_size=hidden_size,
1135
- BLOCK_SIZE_Q=16,
1136
- BLOCK_SIZE_H=256,
1137
- )
1138
-
1139
-
1068
def update_draft_inputs(
1069
draft_tokens: torch.Tensor,
1070
draft_embeds: torch.Tensor | None,
@@ -1146,21 +1074,29 @@ def update_draft_inputs(
1074
idx_mapping: torch.Tensor,
1075
num_reqs: int,
1076
) -> None:
1149
- _shift_input_ids(
1077
+ _shift_input_ids_kernel[(num_reqs,)](
1078
input_buffers.input_ids,
1079
idx_mapping,
1080
input_buffers.query_start_loc,
1081
last_token_indices,
1082
draft_tokens,
1155
- num_reqs=num_reqs,
1083
+ BLOCK_SIZE=1024,
1084
)
1085
if input_embeds is not None:
1086
assert draft_embeds is not None
1159
- _shift_input_embeds(
1087
+ hidden_size = input_embeds.shape[-1]
1088
+ hidden_block_size = 256
1089
+ _shift_input_embeds_kernel[
1090
+ (num_reqs, triton.cdiv(hidden_size, hidden_block_size))
1091
+ ](
1092
input_embeds,
1093
+ input_embeds.stride(0),
1094
draft_embeds,
1095
+ draft_embeds.stride(0),
1096
idx_mapping,
1097
input_buffers.query_start_loc,
1098
last_token_indices,
1165
- num_reqs=num_reqs,
1099
+ hidden_size,
1100
+ BLOCK_SIZE_Q=16,
1101
+ BLOCK_SIZE_H=hidden_block_size,
1102
)