From 2671fedfc7ae604761990603fc736c0c4f21de57 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Sat, 12 Sep 2026 21:07:56 -0700 Subject: [PATCH] [MRV2] Revert explicit Triton JIT warmup migration (#56654) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- .../spec_decode/test_dflash_prepare_inputs.py | 2 +- .../test_gpu_autoregressive_speculator.py | 1 - vllm/v1/worker/gpu/model_runner.py | 95 ++++++------- vllm/v1/worker/gpu/sample/states.py | 6 +- .../gpu/spec_decode/dflash/speculator.py | 131 ++++++------------ .../multi_module_mtp/speculator.py | 88 ++---------- 6 files changed, 96 insertions(+), 227 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_prepare_inputs.py b/tests/v1/spec_decode/test_dflash_prepare_inputs.py index ce47cb8ae6..16d6d8e516 100644 --- a/tests/v1/spec_decode/test_dflash_prepare_inputs.py +++ b/tests/v1/spec_decode/test_dflash_prepare_inputs.py @@ -43,7 +43,7 @@ def _run_prepare( num_scheduled_tokens=np.array([4], dtype=np.int32), positions=torch.tensor(target_positions, dtype=torch.int64, device=device), query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device=device), - idx_mapping=torch.tensor([2], dtype=torch.int64, device=device), + idx_mapping=torch.tensor([2], dtype=torch.int32, device=device), ) query_slot_mapping = torch.full( (max_num_tokens,), -2, dtype=torch.int64, device=device diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index a38265a6f7..920c9a25f0 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/tests/v1/worker/test_gpu_autoregressive_speculator.py @@ -157,7 +157,6 @@ def test_mm_support_configured_after_model_load(monkeypatch): speculator.device = device speculator.max_num_tokens = 4 speculator.max_num_reqs = 2 - speculator.num_speculative_steps = 1 speculator.hidden_size = 3 speculator.dtype = torch.float32 speculator.draft_model_config = draft_model_config diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 14e0ac4629..0ce0288f31 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -188,7 +188,6 @@ logger = init_logger(__name__) class GPUModelRunner(LoRAModelRunnerMixin): - @JitWarmupRegistry.capture def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config self.model_config = vllm_config.model_config @@ -203,6 +202,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.speculative_config is not None and self.speculative_config.use_dspark() ) self.observability_config = vllm_config.observability_config + self.jit_warmup_registry = JitWarmupRegistry(vllm_config) self.device = device self.dtype = self.model_config.dtype @@ -399,7 +399,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): logger.info_once("Loading model from scratch...") # Capture warmup providers selected while constructing the model. - with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] + with self.jit_warmup_registry.activate(): self.model = model_loader.load_model( vllm_config=self.vllm_config, model_config=self.vllm_config.model_config, @@ -461,50 +461,47 @@ class GPUModelRunner(LoRAModelRunnerMixin): # Initialize samplers. Model states may override via custom_sampler(). if self.is_last_pp_rank and not self.is_pooling_model: - with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] - sampler_kwargs: dict[str, Any] = { - "max_num_reqs": self.max_num_reqs, - "vocab_size": self.vocab_size, - "device": self.device, - "req_states": self.req_states, - "logprobs_mode": self.model_config.logprobs_mode, - "num_speculative_tokens": self.decode_query_len, - "use_fp64_gumbel": self.model_config.use_fp64_gumbel, - "enable_trace_replay": self.model_config.enable_trace_replay, - "reasoning_config": self.vllm_config.reasoning_config, - "return_sampling_mask": self.model_config.return_sampling_mask, - } - if self.vllm_config.watermark_config is None: - self.sampler = Sampler(**sampler_kwargs) - else: - wm_config = self.vllm_config.watermark_config - watermarker = create_watermarker(wm_config) - if self.speculative_config is not None: - watermarker = create_speculative_target_watermarker(watermarker) - self.sampler = GPUWatermarkSampler( - watermarker, - deduplicate_contexts=wm_config.deduplicate_contexts, - deduplicate_contexts_max_history=( - wm_config.deduplicate_contexts_max_history - ), - **sampler_kwargs, - ) - custom = self.model_state.custom_sampler(self.sampler) + sampler_kwargs: dict[str, Any] = { + "max_num_reqs": self.max_num_reqs, + "vocab_size": self.vocab_size, + "device": self.device, + "req_states": self.req_states, + "logprobs_mode": self.model_config.logprobs_mode, + "num_speculative_tokens": self.decode_query_len, + "use_fp64_gumbel": self.model_config.use_fp64_gumbel, + "enable_trace_replay": self.model_config.enable_trace_replay, + "reasoning_config": self.vllm_config.reasoning_config, + "return_sampling_mask": self.model_config.return_sampling_mask, + } + if self.vllm_config.watermark_config is None: + self.sampler = Sampler(**sampler_kwargs) + else: + wm_config = self.vllm_config.watermark_config + watermarker = create_watermarker(wm_config) + if self.speculative_config is not None: + watermarker = create_speculative_target_watermarker(watermarker) + self.sampler = GPUWatermarkSampler( + watermarker, + deduplicate_contexts=wm_config.deduplicate_contexts, + deduplicate_contexts_max_history=( + wm_config.deduplicate_contexts_max_history + ), + **sampler_kwargs, + ) + custom = self.model_state.custom_sampler(self.sampler) - if custom: - self.vllm_config._check_watermarking_unsupported( - custom_sampler=True - ) - self.sampler, self.rejection_sampler = custom - elif self.speculative_config is not None: - self.rejection_sampler = RejectionSampler( - self.sampler, - self.speculative_config, - self.device, - watermark_key=speculative_target_watermark_key( - self.vllm_config.watermark_config - ), - ) + if custom: + self.vllm_config._check_watermarking_unsupported(custom_sampler=True) + self.sampler, self.rejection_sampler = custom + elif self.speculative_config is not None: + self.rejection_sampler = RejectionSampler( + self.sampler, + self.speculative_config, + self.device, + watermark_key=speculative_target_watermark_key( + self.vllm_config.watermark_config + ), + ) self.prompt_logprobs_worker = PromptLogprobsWorker( self.max_num_reqs, logprobs_mode=self.model_config.logprobs_mode, @@ -741,7 +738,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.speculator.init_cudagraph_manager(cudagraph_mode) # Capture warmup providers that depend on allocated KV-cache strides. - with self.jit_warmup_registry.activate(): # type: ignore[attr-defined] + with self.jit_warmup_registry.activate(): kv_caches_dict = init_kv_cache( self.compilation_config.static_forward_context, self.kv_cache_config, @@ -927,12 +924,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): # during actual execution. assert self.sampler is not None self.sampler(logits, dummy_input_batch) - if self.rejection_sampler is not None: - # Warm head-dtype logits; scheduler warmup covers processed FP32 logits. - assert self.speculator is not None - self.rejection_sampler( - logits, dummy_input_batch, self.speculator.draft_logits - ) @torch.inference_mode() def _dummy_pooler_run(self, hidden_states: torch.Tensor) -> None: diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index 08cae82128..fe4dee6a6b 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -4,10 +4,7 @@ import numpy as np import torch from vllm.sampling_params import SamplingParams -from vllm.v1.sample.ops.topk_topp_sampler import ( - apply_top_k_top_p, - register_top_k_top_p_warmups, -) +from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p from vllm.v1.worker.gpu.buffer_utils import UvaBackedTensor from vllm.v1.worker.gpu.sample.gumbel import apply_temperature from vllm.v1.worker.gpu.sample.min_p import apply_min_p @@ -21,7 +18,6 @@ class SamplingStates: def __init__(self, max_num_reqs: int, vocab_size: int): self.max_num_reqs = max_num_reqs self.vocab_size = vocab_size - register_top_k_top_p_warmups() self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 10f4346b54..a5710829e0 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -1,10 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import copy -from types import SimpleNamespace from typing import Any -import numpy as np import torch import torch.nn as nn @@ -12,15 +10,6 @@ from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup import ( - WarmupChoices, - WarmupIntRange, -) -from vllm.model_executor.warmup.jit_warmup_triton_helper import ( - DispatchSpec, - TritonWarmupTensor, - triton_kernel_dispatcher_with_warmup, -) from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -119,7 +108,6 @@ class DFlashSpeculator(DraftModelSpeculator): self.query_cudagraph_manager: DFlashCudaGraphManager | None = None self.draft_kv_cache_group_id: int = -1 - prepare_dflash_inputs.register_warmup(speculator=self) @property def attn_vllm_config(self) -> VllmConfig: @@ -694,64 +682,6 @@ def _prepare_dflash_inputs_kernel( tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) -def _prepare_dflash_warmup_inputs(*, speculator: DFlashSpeculator) -> dict[str, Any]: - int32 = TritonWarmupTensor(torch.int32) - int64 = TritonWarmupTensor(torch.int64) - float32 = TritonWarmupTensor(torch.float32) - input_buffers = SimpleNamespace( - input_ids=int32, positions=int64, query_start_loc=int32, seq_lens=int32 - ) - gid: Any = WarmupChoices(*speculator.draft_kv_cache_group_ids) - max_target_query_len: Any = WarmupIntRange(1, 257, advance=lambda value: value * 2) - block_table_stride = int(speculator.block_tables.input_block_tables[gid].stride(0)) - block_table = TritonWarmupTensor( - torch.int32, - shape=(1, block_table_stride), - strides=(block_table_stride, 1), - ) - input_batch = SimpleNamespace( - num_reqs=1, - num_scheduled_tokens=np.array([max_target_query_len]), - positions=int64, - query_start_loc=int32, - idx_mapping=int64, - ) - return dict( - input_buffers=input_buffers, - query_slot_mapping=int64, - context_positions=int64, - context_slot_mapping=int64, - sample_indices=int64, - sample_pos=int64, - sample_idx_mapping=int32, - temperature=float32, - seeds=int64, - input_batch=input_batch, - num_sampled=int32, - num_rejected=int32, - last_sampled=int64, - next_prefill_tokens=int32, - input_temperature=float32, - input_seeds=int64, - block_table=block_table, - block_size=speculator.block_tables.kernel_block_sizes[gid], - cp_rank=speculator.block_tables.cp_rank, - cp_size=speculator.block_tables.cp_size, - cp_interleave=speculator.block_tables.cp_interleave, - parallel_drafting_token_id=speculator.parallel_drafting_token_id, - num_query_per_req=speculator.num_query_per_req, - num_speculative_steps=speculator.num_speculative_steps, - max_num_reqs=speculator.max_num_reqs, - max_num_tokens=speculator.max_num_tokens, - max_model_len=speculator.max_model_len, - sample_from_anchor=speculator.sample_from_anchor, - ) - - -@triton_kernel_dispatcher_with_warmup( - kernel=_prepare_dflash_inputs_kernel, - warmup_inputs=_prepare_dflash_warmup_inputs, -) def prepare_dflash_inputs( input_buffers: InputBuffers, query_slot_mapping: torch.Tensor, @@ -788,33 +718,50 @@ def prepare_dflash_inputs( max_num_tokens: int, max_model_len: int, sample_from_anchor: bool = False, -) -> DispatchSpec: +) -> None: num_reqs = input_batch.num_reqs assert num_reqs > 0 # Cover the longest possible per-request span (ctx + query). Use the max # per-request query length, not the total token count across the batch. max_target_query_len = int(input_batch.num_scheduled_tokens.max()) max_tokens_per_req = max_target_query_len + num_query_per_req - block = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) - num_blocks = triton.cdiv(max_tokens_per_req, block) - return (num_reqs, num_blocks), dict( - out_input_ids_ptr=input_buffers.input_ids, - out_query_positions_ptr=input_buffers.positions, - out_query_start_loc_ptr=input_buffers.query_start_loc, - out_seq_lens_ptr=input_buffers.seq_lens, - out_query_slot_mapping_ptr=query_slot_mapping, - out_context_positions_ptr=context_positions, - out_context_slot_mapping_ptr=context_slot_mapping, - out_sample_indices_ptr=sample_indices, - out_sample_pos_ptr=sample_pos, - out_sample_idx_mapping_ptr=sample_idx_mapping, - out_temperature_ptr=temperature, - out_seeds_ptr=seeds, - target_positions_ptr=input_batch.positions, - target_query_start_loc_ptr=input_batch.query_start_loc, - idx_mapping_ptr=input_batch.idx_mapping, - temperature_ptr=input_temperature, - seeds_ptr=input_seeds, + BLOCK_SIZE = min(256, triton.next_power_of_2(max(1, max_tokens_per_req))) + num_blocks = triton.cdiv(max_tokens_per_req, BLOCK_SIZE) + _prepare_dflash_inputs_kernel[(num_reqs, num_blocks)]( + input_buffers.input_ids, + input_buffers.positions, + input_buffers.query_start_loc, + input_buffers.seq_lens, + query_slot_mapping, + context_positions, + context_slot_mapping, + sample_indices, + sample_pos, + sample_idx_mapping, + temperature, + seeds, + input_batch.positions, + input_batch.query_start_loc, + input_batch.idx_mapping, + last_sampled, + next_prefill_tokens, + num_sampled, + num_rejected, + input_temperature, + input_seeds, + block_table, + block_table.stride(0), + parallel_drafting_token_id, + block_size, + num_query_per_req, + num_speculative_steps, + max_num_reqs, + max_num_tokens, + max_model_len, + cp_rank, + SAMPLE_FROM_ANCHOR=sample_from_anchor, PAD_SLOT_ID=PAD_SLOT_ID, - BLOCK_SIZE=block, + CP_SIZE=cp_size, + CP_INTERLEAVE=cp_interleave, + BLOCK_SIZE=BLOCK_SIZE, ) diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index 1c37ccb6c2..1c6ae315cf 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py @@ -9,11 +9,6 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger -from vllm.model_executor.warmup.jit_warmup_triton_helper import ( - DispatchSpec, - TritonWarmupTensor, - triton_kernel_dispatcher_with_warmup, -) from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -74,9 +69,6 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator): self.cudagraph_manager: SpeculatorCudaGraphManager | None = None - _shift_input_ids.register_warmup() - _shift_input_embeds.register_warmup(speculator=self) - def load_draft_model( self, target_model: nn.Module, @@ -1016,32 +1008,6 @@ def _shift_input_ids_kernel( tl.store(input_ids_ptr + last_token_index, draft_token) -def _shift_input_ids_warmup_inputs() -> dict[str, Any]: - return dict( - input_ids=TritonWarmupTensor(torch.int32), - idx_mapping=TritonWarmupTensor(torch.int32), - query_start_loc=TritonWarmupTensor(torch.int32), - last_token_indices=TritonWarmupTensor(torch.int64), - draft_tokens=TritonWarmupTensor(torch.int64), - num_reqs=1, - ) - - -@triton_kernel_dispatcher_with_warmup( - kernel=_shift_input_ids_kernel, - warmup_inputs=_shift_input_ids_warmup_inputs, -) -def _shift_input_ids( - input_ids: torch.Tensor, - idx_mapping: torch.Tensor, - query_start_loc: torch.Tensor, - last_token_indices: torch.Tensor, - draft_tokens: torch.Tensor, - num_reqs: int, -) -> DispatchSpec: - return (num_reqs,), dict(BLOCK_SIZE=1024) - - @triton.jit def _shift_input_embeds_kernel( input_embeds_ptr, @@ -1099,44 +1065,6 @@ def _shift_input_embeds_kernel( ) -def _shift_input_embeds_warmup_inputs( - *, speculator: MultiModuleMTPSpeculator -) -> dict[str, Any]: - hidden = TritonWarmupTensor( - speculator.dtype, - shape=(1, speculator.hidden_size), - strides=(speculator.hidden_size, 1), - ) - return dict( - input_embeds=hidden, - draft_embeds=hidden, - idx_mapping=TritonWarmupTensor(torch.int32), - query_start_loc=TritonWarmupTensor(torch.int32), - last_token_indices=TritonWarmupTensor(torch.int64), - num_reqs=1, - ) - - -@triton_kernel_dispatcher_with_warmup( - kernel=_shift_input_embeds_kernel, - warmup_inputs=_shift_input_embeds_warmup_inputs, -) -def _shift_input_embeds( - input_embeds: torch.Tensor, - draft_embeds: torch.Tensor, - idx_mapping: torch.Tensor, - query_start_loc: torch.Tensor, - last_token_indices: torch.Tensor, - num_reqs: int, -) -> DispatchSpec: - hidden_size = input_embeds.shape[-1] - return (num_reqs, triton.cdiv(hidden_size, 256)), dict( - hidden_size=hidden_size, - BLOCK_SIZE_Q=16, - BLOCK_SIZE_H=256, - ) - - def update_draft_inputs( draft_tokens: torch.Tensor, draft_embeds: torch.Tensor | None, @@ -1146,21 +1074,29 @@ def update_draft_inputs( idx_mapping: torch.Tensor, num_reqs: int, ) -> None: - _shift_input_ids( + _shift_input_ids_kernel[(num_reqs,)]( input_buffers.input_ids, idx_mapping, input_buffers.query_start_loc, last_token_indices, draft_tokens, - num_reqs=num_reqs, + BLOCK_SIZE=1024, ) if input_embeds is not None: assert draft_embeds is not None - _shift_input_embeds( + hidden_size = input_embeds.shape[-1] + hidden_block_size = 256 + _shift_input_embeds_kernel[ + (num_reqs, triton.cdiv(hidden_size, hidden_block_size)) + ]( input_embeds, + input_embeds.stride(0), draft_embeds, + draft_embeds.stride(0), idx_mapping, input_buffers.query_start_loc, last_token_indices, - num_reqs=num_reqs, + hidden_size, + BLOCK_SIZE_Q=16, + BLOCK_SIZE_H=hidden_block_size, )