iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[MRV2] Revert explicit Triton JIT warmup migration (#56654)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Woosuk Kwon
2026-09-13 04:07:56 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent c711f740b7
commit 2671fedfc7
6 changed files with 96 additions and 227 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/spec_decode/test_dflash_prepare_inputs.py
+1 -1
View File
@@ -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
tests/v1/worker/test_gpu_autoregressive_speculator.py
-1
View File
@@ -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
vllm/v1/worker/gpu/model_runner.py
+43 -52
View File
@@ -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:
vllm/v1/worker/gpu/sample/states.py
+1 -5
View File
@@ -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)
vllm/v1/worker/gpu/spec_decode/dflash/speculator.py
+39 -92
View File
@@ -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,
)
vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py
+12 -76
View File
@@ -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,
)