admin/1Cat-vLLM
1
0 Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][Core] MRV2 warmup: cover every request count, mixed steps and a full chunk (#1402)

Browse Source 
A fresh engine paid about 1 s of TTFT on the first batch at each new
concurrency: several Triton kernels specialize on the request count, a
decode step that also admits a prefill takes its own route, and a full
4096-token chunk reaches long-row and divisible-by-16 variants the short
warmup prompts never did. Warmup ran only the max_num_seqs batch.

Warmup now also runs, for every smaller request count, one cycle with a
pure decode step and one whose decode step admits a new prefill, plus a
single-request prefill of a full chunk. Each cycle reuses the existing
prefill -> one step -> finish pattern.

Qwen3.8-27B NVFP4 + DFlash2, 4xV100, first batch after start (T=0):
c2 2.0 -> 0.8 s, c8 1.9-3.2 -> 1.33 s (steady state 0.47 / 1.34 s);
kernels compiled while serving 7 -> 1. Startup grows 115 -> 129 s.

Co-authored-by: mac <mac@MacBook-xinyi.local>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
yangzhuxinyzxmacClaude Opus 5.5
 authored and GitHub committed 2026-11-06 23:07:00 +08:00
1 parent 4a6e6124d9
commit ad339d82ed
2 files changed
+181 -105
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/v1/worker/test_gpu_warmup_blocks.py
+18 -8
View File
@@ -81,16 +81,26 @@ def test_kernel_prefill_warmup_runs_default_batch_and_extra_profile(
warmup_kernels(runner, executions.append, samples.append)
assert [output.total_num_scheduled_tokens for output in executions] == [
24,
20,
0,
33,
5,
0,
tokens = [output.total_num_scheduled_tokens for output in executions]
# Full batch, then the 33-token kernel profile and a full prefill chunk
# (capped so the decode step fits max_model_len), one request each.
assert tokens[:9] == [24, 20, 0, 33, 5, 0, 59, 5, 0]
# Every smaller request count: a pure decode step, then a decode step
# that also admits a 6-token prefill.
assert tokens[9:] == [
*(18, 15, 0, 18, 21, 0),
*(12, 10, 0, 12, 16, 0),
*(6, 5, 0, 6, 11, 0),
]
mixed = [output for output in executions if output.scheduled_new_reqs][-1]
assert len(mixed.scheduled_cached_reqs.req_ids) == 1
assert [r.req_id for r in mixed.scheduled_new_reqs] == ["_warmup_c1m_new_"]
# Every request a cycle starts is finished by its cleanup step.
started = {r.req_id for output in executions for r in output.scheduled_new_reqs}
finished = set().union(*(output.finished_req_ids for output in executions))
assert started == finished
assert connector_states == [True, False]
assert len(samples) == 4
assert len(samples) == 6 + 6 * 2
def _speculative_config(method: str) -> SpeculativeConfig:
vllm/v1/worker/gpu/warmup.py
+163 -97
View File
@@ -89,6 +89,119 @@ def _warmup_block_counter(
return block_count
def _warmup_cycle(
model_runner: GPUModelRunner,
worker_execute_model: Callable[[SchedulerOutput], Any],
worker_sample_tokens: Callable[[GrammarOutput | None], Any],
tag: str,
num_reqs: int,
prompt_len: int,
sampling_params: SamplingParams | None,
pooling_params: PoolingParams | None,
block_count: Callable[[int, KVCacheSpec], int],
kv_cache_specs: list[KVCacheSpec],
*,
with_grammar: bool,
admit_prefill: bool = False,
) -> None:
"""Prefill ``num_reqs`` requests, run one decode step, then finish them.
With ``admit_prefill`` the decode step also prefills one new request, the
mixed step a request arriving beside running ones produces.
"""
num_spec_steps = model_runner.num_speculative_steps
num_kv_cache_groups = len(kv_cache_specs)
prompt_token_ids = list(range(prompt_len))
decode_len = prompt_len + 1 + num_spec_steps
prefill_block_counts = [block_count(prompt_len, spec) for spec in kv_cache_specs]
decode_block_counts = [block_count(decode_len, spec) for spec in kv_cache_specs]
decode_block_deltas = [
d - p for d, p in zip(decode_block_counts, prefill_block_counts)
]
req_ids = [f"_warmup_{tag}_{i}_" for i in range(num_reqs)]
next_block_id = 1
def _alloc_blocks(num_blocks: int) -> list[int]:
nonlocal next_block_id
return list(range(next_block_id, next_block_id := next_block_id + num_blocks))
def _new_requests(ids: list[str]) -> list[NewRequestData]:
return [
NewRequestData.from_request(
Request(req_id, prompt_token_ids, sampling_params, pooling_params),
block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
prefill_token_ids=prompt_token_ids,
)
for req_id in ids
]
prefill_output = SchedulerOutput.make_empty()
prefill_output.scheduled_new_reqs = _new_requests(req_ids)
prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids}
prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs
prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
worker_execute_model(prefill_output)
finished = set(req_ids)
if not model_runner.is_pooling_model:
grammar_output = None
if with_grammar and model_runner.is_last_pp_rank:
# Exercise the structured-output bitmask once; extra
# operator profiles only need the model path.
vocab_size = model_runner.model_config.get_vocab_size()
bitmask_width = (vocab_size + 31) // 32
grammar_bitmask = np.full(
(len(req_ids), bitmask_width),
fill_value=-1,
dtype=np.int32,
)
grammar_output = GrammarOutput(
structured_output_request_ids=req_ids,
grammar_bitmask=grammar_bitmask,
)
worker_sample_tokens(grammar_output)
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)
]
decode_output = SchedulerOutput.make_empty()
decode_output.scheduled_cached_reqs = cached_req_data
decode_output.num_scheduled_tokens = {
req_id: 1 + num_spec_steps 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
}
if admit_prefill:
arriving = f"_warmup_{tag}_new_"
decode_output.scheduled_new_reqs = _new_requests([arriving])
decode_output.num_scheduled_tokens[arriving] = prompt_len
finished.add(arriving)
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)
cleanup_output = SchedulerOutput.make_empty()
cleanup_output.finished_req_ids = finished
worker_execute_model(cleanup_output)
@torch.inference_mode()
def warmup_kernels(
model_runner: GPUModelRunner,
@@ -111,10 +224,17 @@ def warmup_kernels(
prompt_lengths = _kernel_prefill_warmup_token_counts(
model_runner, default_prompt_len
)
if not model_runner.is_pooling_model:
# A full prefill chunk: long-row kernel variants and Triton's
# divisible-by-16 specializations that the short prompts never reach.
full_chunk = min(
model_runner.scheduler_config.max_num_batched_tokens,
model_runner.max_model_len - 1 - num_spec_steps,
)
if full_chunk > default_prompt_len:
prompt_lengths = tuple(sorted({*prompt_lengths, full_chunk}))
kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups
num_kv_cache_groups = len(kv_cache_groups)
block_count = _warmup_block_counter(model_runner)
kv_cache_specs = [group.kv_cache_spec for group in kv_cache_groups]
@@ -131,20 +251,13 @@ def warmup_kernels(
# work by max_num_seqs.
model_runner.kv_connector.set_disabled(True)
try:
num_reqs = 0
for profile_idx, prompt_len in enumerate(prompt_lengths):
prompt_token_ids = list(range(prompt_len))
decode_len = prompt_len + 1 + num_spec_steps
prefill_block_counts = [
block_count(prompt_len, spec) for spec in kv_cache_specs
]
decode_block_counts = [
block_count(decode_len, spec) for spec 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(
max_blocks_per_req = sum(
block_count(prompt_len + 1 + num_spec_steps, spec)
for spec in kv_cache_specs
)
profile_reqs = min(
model_runner.scheduler_config.max_num_seqs,
model_runner.scheduler_config.max_num_batched_tokens
// max(prompt_len, 1 + num_spec_steps),
@@ -155,92 +268,45 @@ def warmup_kernels(
),
)
if profile_idx:
num_reqs = min(num_reqs, 1)
if num_reqs <= 0:
profile_reqs = min(profile_reqs, 1)
else:
num_reqs = profile_reqs
if profile_reqs <= 0:
continue
req_ids = [f"_warmup_{profile_idx}_{i}_" for i in range(num_reqs)]
next_block_id = 1
def _alloc_blocks(num_blocks: int) -> list[int]:
nonlocal next_block_id
return list(
range(next_block_id, next_block_id := next_block_id + num_blocks)
)
new_reqs = [
NewRequestData.from_request(
Request(
req_ids[i],
prompt_token_ids,
_warmup_cycle(
model_runner,
worker_execute_model,
worker_sample_tokens,
f"{profile_idx}",
profile_reqs,
prompt_len,
sampling_params,
pooling_params,
block_count,
kv_cache_specs,
with_grammar=profile_idx == 0,
)
if not model_runner.is_pooling_model:
# Several kernels specialize on the request count, and a step that
# admits a prefill beside decoding requests takes its own route.
# Without this, the first batch at each concurrency compiles them
# while serving it (about 1 s of TTFT per new shape).
for count in range(num_reqs - 1, 0, -1):
for mixed in (False, True):
_warmup_cycle(
model_runner,
worker_execute_model,
worker_sample_tokens,
f"c{count}{'m' if mixed else 'd'}",
count,
default_prompt_len,
sampling_params,
pooling_params,
),
block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts),
prefill_token_ids=prompt_token_ids,
)
for i in range(num_reqs)
]
prefill_output = SchedulerOutput.make_empty()
prefill_output.scheduled_new_reqs = new_reqs
prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids}
prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs
prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups
worker_execute_model(prefill_output)
if not model_runner.is_pooling_model:
grammar_output = None
if profile_idx == 0 and model_runner.is_last_pp_rank:
# Exercise the structured-output bitmask once; extra
# operator profiles only need the model path.
vocab_size = model_runner.model_config.get_vocab_size()
bitmask_width = (vocab_size + 31) // 32
grammar_bitmask = np.full(
(len(req_ids), bitmask_width),
fill_value=-1,
dtype=np.int32,
block_count,
kv_cache_specs,
with_grammar=False,
admit_prefill=mixed,
)
grammar_output = GrammarOutput(
structured_output_request_ids=req_ids,
grammar_bitmask=grammar_bitmask,
)
worker_sample_tokens(grammar_output)
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)
]
decode_output = SchedulerOutput.make_empty()
decode_output.scheduled_cached_reqs = cached_req_data
decode_output.num_scheduled_tokens = {
req_id: 1 + num_spec_steps 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)
cleanup_output = SchedulerOutput.make_empty()
cleanup_output.finished_req_ids = set(req_ids)
worker_execute_model(cleanup_output)
finally:
model_runner.kv_connector.set_disabled(False)
torch.accelerator.synchronize()