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

[Core] Warm up runner-owned Triton kernels before the first request (#49903)

Browse Source
This commit is contained in:
Nick Hill
2026-07-28 07:58:02 -07:00
committed by GitHub
parent 1e81853afc
commit 0d0504b54c
7 changed files with 183 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_kv_block_zeroer.py
+65 -1
View File
@@ -4,7 +4,7 @@
import pytest
import torch
from vllm.v1.worker.utils import KVBlockZeroer
from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@@ -86,3 +86,67 @@ def test_non_uniform_page_sizes():
assert torch.all(storage[1] == 0)
assert torch.all(storage[2] == 0)
assert torch.all(storage[3] == 1)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_warmup_compiles_every_n_blocks_specialization():
"""After warmup, no launch should trigger a first-request JIT compile.
``n_blocks`` is ``do_not_specialize``, so a single warmup launch must
cover every block count.
"""
device = torch.device("cuda")
num_blocks = 64
page_size_el = 4
storage = torch.ones((num_blocks, page_size_el), dtype=torch.int32, device=device)
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
zeroer.device = device
zeroer._meta = (
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
1, # max_chunks
page_size_el, # blk_size
1, # n_segs
)
def compiled_variants() -> set:
return {
key
for caches in _zero_kv_blocks_kernel.device_caches.values()
for key in caches[0]
}
zeroer.warmup(num_blocks)
torch.accelerator.synchronize()
warmed = compiled_variants()
assert warmed
for n_blocks in (1, 2, 3, 16, 32):
zeroer.zero_block_ids(list(range(n_blocks)))
torch.accelerator.synchronize()
assert compiled_variants() == warmed
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_warmup_respects_available_block_count():
"""An empty KV cache must not be warmed with out-of-range block IDs."""
device = torch.device("cuda")
page_size_el = 4
storage = torch.ones((1, page_size_el), dtype=torch.int32, device=device)
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
zeroer.device = device
zeroer._meta = (
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
1,
page_size_el,
1,
)
zeroer.warmup(0)
torch.accelerator.synchronize()
assert torch.all(storage == 1)
vllm/model_executor/warmup/kernel_warmup.py
+10 -6
View File
@@ -98,12 +98,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
minimax_m3_msa_warmup,
)
# Pooling models do not use the generation slot-mapping path.
if not worker.use_v2_model_runner and not worker.model_runner.is_pooling_model:
warm_v1_block_table_kernels(
getattr(worker.model_runner, "device", torch.device("cuda")),
worker.scheduler_config.max_num_batched_tokens,
)
if not worker.use_v2_model_runner:
# Pooling models do not use the generation slot-mapping path.
if not worker.model_runner.is_pooling_model:
warm_v1_block_table_kernels(worker.model_runner)
# The KV-block zeroing kernel is driven by the scheduler's
# `new_block_ids_to_zero`, so no dummy run ever reaches it.
zeroer = getattr(worker.model_runner, "_kv_block_zeroer", None)
if zeroer is not None:
zeroer.warmup(worker.model_runner.kv_cache_config.num_blocks)
qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
# DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
vllm/model_executor/warmup/qwen_triton_warmup.py
-114
View File
@@ -24,24 +24,10 @@ _QWEN_MODEL_TYPES = frozenset(
}
)
_ZERO_KV_N_BLOCKS = (1, 2)
_SLOT_MAPPING_KV_BLOCK_SIZE = 16
_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE = 1
_SLOT_MAPPING_BLOCK_TABLE_STRIDES = (1, 3)
# Covers L=1 constexpr, non-divisible runtime L, and divisible runtime L.
_FLA_POST_CONV_WARMUP_LENGTHS = (1, 2, 16)
@dataclass(frozen=True)
class _ZeroKvWarmupConfig:
seg_page_sizes: torch.Tensor
max_chunks: int
block_size: int
n_segs: int
@dataclass(frozen=True)
class _QwenGDNWarmupConfig:
h: int
@@ -148,96 +134,6 @@ def _qwen_gdn_warmup_config(
return None
def _get_kv_block_zeroer(runner: object) -> object | None:
zeroer = getattr(runner, "kv_block_zeroer", None)
if zeroer is None:
zeroer = getattr(runner, "_kv_block_zeroer", None)
return zeroer
def _zero_kv_warmup_config(runner: object) -> _ZeroKvWarmupConfig | None:
zeroer = _get_kv_block_zeroer(runner)
meta = getattr(zeroer, "_meta", None)
if meta is None:
return None
_, seg_page_sizes, max_chunks, block_size, n_segs = meta
return _ZeroKvWarmupConfig(
seg_page_sizes=seg_page_sizes,
max_chunks=int(max_chunks),
block_size=int(block_size),
n_segs=int(n_segs),
)
def _warm_zero_kv_blocks_with_runner_zeroer(runner: object) -> bool:
zeroer = _get_kv_block_zeroer(runner)
zero_block_ids = getattr(zeroer, "zero_block_ids", None)
if not callable(zero_block_ids):
return False
for n_blocks in _ZERO_KV_N_BLOCKS:
zero_block_ids(list(range(n_blocks)))
return True
def _warm_zero_kv_blocks_kernel(
device: torch.device, config: _ZeroKvWarmupConfig
) -> None:
from vllm.v1.worker.utils import _zero_kv_blocks_kernel
max_n_blocks = max(_ZERO_KV_N_BLOCKS)
max_page_size = int(config.seg_page_sizes.max().item())
scratch = torch.empty(
max_n_blocks * max_page_size,
dtype=torch.int32,
device=device,
)
seg_addrs = torch.tensor(
[scratch.data_ptr()] * config.n_segs,
dtype=torch.uint64,
device=device,
)
for n_blocks in _ZERO_KV_N_BLOCKS:
block_ids = torch.arange(n_blocks, dtype=torch.int64, device=device)
grid = (n_blocks * config.n_segs * config.max_chunks,)
_zero_kv_blocks_kernel[grid](
seg_addrs,
config.seg_page_sizes,
block_ids,
n_blocks,
N_SEGS=config.n_segs,
MAX_CHUNKS=config.max_chunks,
BLOCK_SIZE=config.block_size,
)
def _warm_compute_slot_mapping_kernel(device: torch.device) -> None:
from vllm.v1.worker.block_table import BlockTable
# num_tokens/max_num_tokens are do_not_specialize; keep the launch tiny.
num_tokens = 1
query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
for block_table_stride in _SLOT_MAPPING_BLOCK_TABLE_STRIDES:
# Use BlockTable so the JIT key matches the production slot-mapping call.
block_table = BlockTable(
block_size=_SLOT_MAPPING_KV_BLOCK_SIZE,
max_num_reqs=1,
max_num_blocks_per_req=block_table_stride,
max_num_batched_tokens=num_tokens,
pin_memory=False,
device=device,
kernel_block_size=_SLOT_MAPPING_KV_BLOCK_SIZE,
cp_kv_cache_interleave_size=_SLOT_MAPPING_CP_KV_CACHE_INTERLEAVE_SIZE,
)
block_table.add_row(list(range(block_table_stride)), 0)
block_table.commit_block_table(num_reqs=1)
block_table.compute_slot_mapping(1, query_start_loc, positions)
def _warm_causal_conv1d_fwd_kernel(
device: torch.device, config: _QwenGDNWarmupConfig
) -> None:
@@ -373,16 +269,6 @@ def qwen_triton_warmup(
device = getattr(runner, "device", torch.device("cuda"))
logger.info("Warming up Qwen Triton kernels for model_type=%s.", model_type)
zero_config = _zero_kv_warmup_config(runner)
warmed_zeroer = _warm_zero_kv_blocks_with_runner_zeroer(runner)
if zero_config is not None:
_warm_zero_kv_blocks_kernel(device, zero_config)
elif not warmed_zeroer:
logger.info("Skipping Qwen zero-kv warmup: no KVBlockZeroer metadata.")
_warm_compute_slot_mapping_kernel(device)
_synchronize_device(device)
compilation_config = getattr(runner, "compilation_config", None)
static_forward_context = getattr(compilation_config, "static_forward_context", None)
gdn_config = _qwen_gdn_warmup_config(static_forward_context)
vllm/model_executor/warmup/v1_block_table_warmup.py
+14 -28
View File
@@ -2,42 +2,28 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Warm up v1 block-table Triton kernels."""
from typing import TYPE_CHECKING
import torch
if TYPE_CHECKING:
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
_SLOT_MAPPING_WARMUP_TOKENS = 8
_SLOT_MAPPING_WARMUP_BLOCK_SIZES = (3, 16)
_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE = 1
def warm_v1_block_table_kernels(
device: torch.device,
max_tokens: int,
) -> None:
from vllm.v1.worker.block_table import BlockTable
def warm_v1_block_table_kernels(runner: "GPUModelRunner") -> None:
"""JIT-compile ``_compute_slot_mapping_kernel`` for the real block tables."""
num_tokens = max(0, min(_SLOT_MAPPING_WARMUP_TOKENS, max_tokens))
device = runner.device
block_table = runner.input_batch.block_table
num_tokens = min(
_SLOT_MAPPING_WARMUP_TOKENS,
runner.scheduler_config.max_num_batched_tokens,
)
if num_tokens <= 0:
return
query_start_loc = torch.tensor([0, num_tokens], dtype=torch.int32, device=device)
positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
for block_size in _SLOT_MAPPING_WARMUP_BLOCK_SIZES:
max_num_blocks_per_req = max(
1, (max(num_tokens, max_tokens) + block_size - 1) // block_size
)
max_num_blocks_per_req = ((max_num_blocks_per_req + 15) // 16) * 16
block_table = BlockTable(
block_size=block_size,
max_num_reqs=1,
max_num_blocks_per_req=max_num_blocks_per_req,
max_num_batched_tokens=max(num_tokens, max_tokens),
pin_memory=False,
device=device,
kernel_block_size=block_size,
cp_kv_cache_interleave_size=(
_SLOT_MAPPING_WARMUP_CP_KV_CACHE_INTERLEAVE_SIZE
),
)
block_table.add_row(list(range(max_num_blocks_per_req)), 0)
block_table.commit_block_table(1)
block_table.compute_slot_mapping(1, query_start_loc, positions)
block_table.compute_slot_mapping(1, query_start_loc, positions)
vllm/v1/worker/gpu/warmup.py
+85 -36
View File
@@ -157,13 +157,10 @@ def warmup_kernels(
worker_execute_model: Callable[[SchedulerOutput], Any],
worker_sample_tokens: Callable[[GrammarOutput | None], Any],
) -> None:
"""Run two execute_model + sample_tokens iterations to JIT compile
triton kernels. We must call the provided worker's execute_model for
pipeline parallel coordination.
"""Run scheduler-realistic prefill and decode steps to JIT compile kernels.
The first iteration simulates a prefill with requests of
decode_query_len + 1 prompt tokens each. The second iteration simulates
a decode step with all requests generating decode_query_len tokens.
We must call the provided worker's execute_model for pipeline parallel
coordination.
"""
num_spec_steps = model_runner.num_speculative_steps
decode_query_len = model_runner.decode_query_len
@@ -172,8 +169,13 @@ def warmup_kernels(
# a uniform decode batch.
prompt_len = decode_query_len + 1
prompt_token_ids = list(range(prompt_len))
# After prefill, decode generates decode_query_len tokens.
decode_len = prompt_len + decode_query_len
# Upper bound on the decode steps built in `decode_steps` below.
num_decode_steps = 1
if not model_runner.is_pooling_model:
num_decode_steps = 5 if num_spec_steps > 0 else 3
# Size the block allocation for the worst case: every request advancing
# decode_query_len tokens on every decode step.
decode_len = prompt_len + num_decode_steps * decode_query_len
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
num_kv_cache_groups = len(kv_cache_groups)
@@ -208,9 +210,6 @@ def warmup_kernels(
kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs]
decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs]
decode_block_deltas = [
d - p for d, p in zip(decode_block_counts, prefill_block_counts)
]
max_blocks_per_req = sum(decode_block_counts)
num_reqs = min(
@@ -243,6 +242,11 @@ def warmup_kernels(
nonlocal next_block_id
return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
# The KV-block zeroing kernel is driven by the scheduler's
# new_block_ids_to_zero, so none of the steps below reach it.
if model_runner.kv_block_zeroer is not None:
model_runner.kv_block_zeroer.warmup(model_runner.kv_cache_config.num_blocks)
# Step 1: Prefill all requests with 1 + decode_query_len prompt tokens each.
new_reqs = [
NewRequestData.from_request(
@@ -287,33 +291,78 @@ def warmup_kernels(
worker_sample_tokens(grammar_output)
# Step 2: Decode all requests with decode_query_len tokens each.
cached_req_data = CachedRequestData.make_empty()
cached_req_data.req_ids = list(req_ids)
cached_req_data.num_computed_tokens = [prompt_len] * num_reqs
cached_req_data.num_output_tokens = [1] * num_reqs
new_block = any(decode_block_deltas)
cached_req_data.new_block_ids = [
tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None
for _ in range(num_reqs)
# Per-request state carried across the decode steps.
req_computed = [prompt_len] * num_reqs
req_blocks = [list(prefill_block_counts) for _ in range(num_reqs)]
def _run_decode_step(indices: list[int], spec_flags: list[bool]) -> None:
"""Decode `indices`, spec-decoding the ones flagged in `spec_flags`."""
cached_req_data = CachedRequestData.make_empty()
cached_req_data.req_ids = [req_ids[i] for i in indices]
cached_req_data.num_computed_tokens = [req_computed[i] for i in indices]
cached_req_data.num_output_tokens = [1] * len(indices)
cached_req_data.new_block_ids = []
step_num_scheduled_tokens: dict[str, int] = {}
step_spec_tokens: dict[str, list[int]] = {}
for i, use_spec in zip(indices, spec_flags):
num_tokens = decode_query_len if use_spec else 1
after = req_computed[i] + num_tokens
deltas = [
_warmup_block_count(after, spec) - held
for spec, held in zip(kv_cache_specs, req_blocks[i])
]
cached_req_data.new_block_ids.append(
tuple(_alloc_blocks(n) for n in deltas) if any(deltas) else None
)
req_blocks[i] = [
held + delta for held, delta in zip(req_blocks[i], deltas)
]
step_num_scheduled_tokens[req_ids[i]] = num_tokens
if use_spec:
step_spec_tokens[req_ids[i]] = [0] * num_spec_steps
decode_output = SchedulerOutput.make_empty()
decode_output.scheduled_cached_reqs = cached_req_data
decode_output.num_scheduled_tokens = step_num_scheduled_tokens
decode_output.scheduled_spec_decode_tokens = step_spec_tokens
decode_output.total_num_scheduled_tokens = sum(
step_num_scheduled_tokens.values()
)
decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
worker_execute_model(decode_output)
worker_sample_tokens(None)
for i, use_spec in zip(indices, spec_flags):
req_computed[i] += decode_query_len if use_spec else 1
all_indices = list(range(num_reqs))
use_spec_decode = num_spec_steps > 0
# Decode steps to warm, as (request indices, per-request spec flag).
# Under spec decoding the scheduler drops requests the drafter proposed
# nothing for, so warm each batch shape with and without draft tokens.
decode_steps: list[tuple[list[int], list[bool]]] = [
(all_indices, [use_spec_decode] * num_reqs),
]
if num_reqs >= 2:
# Mixed spec / non-spec: GDN and KDA reclassify the non-spec decode
# as a prefill and split the batch into spec/non-spec token indices.
decode_steps.append(([0, 1], [use_spec_decode, False]))
if use_spec_decode:
# Exercise the model paths that split a batch by whether each
# request received draft tokens.
decode_steps.append(([0, 1], [False, False]))
if num_reqs > 1:
decode_steps.append(([0], [use_spec_decode]))
if use_spec_decode:
decode_steps.append(([0], [False]))
elif use_spec_decode:
decode_steps.append(([0], [False]))
decode_output = SchedulerOutput.make_empty()
decode_output.scheduled_cached_reqs = cached_req_data
decode_output.num_scheduled_tokens = {
req_id: decode_query_len for req_id in req_ids
}
if num_spec_steps > 0:
decode_output.scheduled_spec_decode_tokens = {
req_id: [0] * num_spec_steps for req_id in req_ids
}
decode_output.total_num_scheduled_tokens = sum(
decode_output.num_scheduled_tokens.values()
)
decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
worker_execute_model(decode_output)
worker_sample_tokens(None)
for step_indices, step_spec_flags in decode_steps:
_run_decode_step(step_indices, step_spec_flags)
# Clean up - process finish_req_ids.
cleanup_output = SchedulerOutput.make_empty()
vllm/v1/worker/mamba_utils.py
+3 -3
View File
@@ -150,7 +150,7 @@ def _copy_mamba_state_block(
tl.store(tail_dst + tail_off, tail_data, mask=tail_mask)
@triton.jit
@triton.jit(do_not_specialize=["num_reqs"])
def postprocess_mamba_fused_kernel(
# Decision inputs (per-request)
num_accepted_tokens_ptr,
@@ -280,7 +280,7 @@ def postprocess_mamba_fused_kernel(
)
@triton.jit
@triton.jit(do_not_specialize=["num_reqs"])
def preprocess_mamba_align_fused_kernel(
idx_mapping_ptr,
state_idx_ptr,
@@ -326,7 +326,7 @@ def preprocess_mamba_align_fused_kernel(
tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset)
@triton.jit
@triton.jit(do_not_specialize=["num_reqs"])
def precopy_mamba_align_fused_kernel(
# Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by
# the V2 fused align preprocess kernel for the current step:
vllm/v1/worker/utils.py
+6 -1
View File
@@ -40,7 +40,7 @@ from vllm.v1.kv_cache_interface import (
logger = init_logger(__name__)
@triton.jit
@triton.jit(do_not_specialize=["n_blocks"])
def _zero_kv_blocks_kernel(
seg_addrs_ptr,
seg_page_sizes_ptr,
@@ -206,6 +206,11 @@ class KVBlockZeroer:
BLOCK_SIZE=blk_size,
)
def warmup(self, num_kv_blocks: int) -> None:
"""JIT-compile the zeroing kernel before the first real request."""
if num_kv_blocks > 0:
self.zero_block_ids([0])
@dataclass
class AttentionGroup: