@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 141 lines 5.19 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Catch runtime (post-warmup) JIT compilations for JIT-heavy backends running
4 e2e tests on popular models, for which we only load a few blocks for performance.
5
6 NOTE(NickLucche) With cuda graphs on, kernels fully covered by graphs captured during
7 warmup do not re-trigger the Python JIT hooks. The targeted paths (prefill
8 MoE/MLA/SSM and the sampler) run mixed, so they are unaffected.
9 """
10
11 from dataclasses import dataclass
12
13 import pytest
14
15 from vllm import LLM, SamplingParams
16 from vllm.inputs import TokensPrompt
17
18 from ..models.utils import dummy_hf_overrides
19 from ..utils import create_new_process_for_each_test
20
21 # Warmup coverage is still incomplete for these backends, so the monitor fires
22 # during inference. Tracked in https://github.com/vllm-project/vllm/issues/49349;
23 # drop this once the warmup contract migrations land.
24 pytestmark = pytest.mark.skip(reason="Kernel warmup coverage is still incomplete")
25
26
27 @dataclass(frozen=True)
28 class JitModel:
29 model: str
30 draft: str | None = None
31 trust_remote_code: bool = False
32
33
34 JIT_MONITOR_MODELS = [
35 JitModel("Qwen/Qwen3-0.6B"),
36 JitModel("deepseek-ai/DeepSeek-V2-Lite-Chat", trust_remote_code=True),
37 JitModel("deepseek-ai/DeepSeek-V3", trust_remote_code=True),
38 JitModel("ibm-granite/granite-4.0-tiny-preview"),
39 JitModel(
40 "luccafong/deepseek_mtp_main_random",
41 draft="luccafong/deepseek_mtp_draft_random",
42 trust_remote_code=True,
43 ),
44 JitModel(
45 "eagle618/deepseek-v3-random",
46 draft="eagle618/eagle-deepseek-v3-random",
47 trust_remote_code=True,
48 ),
49 ]
50
51
52 def _run_shape_battery(llm: LLM) -> None:
53 """Exercise diverse compile keys so missing warmup keys surface.
54
55 Token-id prompts keep shapes exact and avoid depending on a tokenizer.
56 Outputs are meaningless under dummy weights; we assert only that no JIT
57 fired.
58 """
59 short = TokensPrompt(prompt_token_ids=[1, 2, 3, 4])
60 medium = TokensPrompt(prompt_token_ids=list(range(1, 33)))
61 long = TokensPrompt(prompt_token_ids=list(range(1, 129)))
62
63 # Greedy single-sequence multi-step decode: prefill + autoregressive decode
64 # + greedy sampler.
65 llm.generate(medium, SamplingParams(temperature=0.0, max_tokens=16))
66
67 # Batched prefill with mixed lengths: varlen prefill + padded decode.
68 llm.generate([short, medium, long], SamplingParams(temperature=0.0, max_tokens=8))
69
70 # Triton sampler kernels: top_k / top_p / min_p each specialize.
71 for sampling_params in (
72 SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0),
73 SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0),
74 SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0),
75 SamplingParams(
76 temperature=0.8, top_k=20, top_p=0.9, min_p=0.1, max_tokens=8, seed=0
77 ),
78 ):
79 llm.generate(medium, sampling_params)
80
81 # Heterogeneous SamplingParams in one step, where missing sampler warmup
82 # keys most often hide.
83 llm.generate(
84 [medium] * 4,
85 [
86 SamplingParams(temperature=0.0, max_tokens=8),
87 SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0),
88 SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0),
89 SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0),
90 ],
91 )
92
93
94 @create_new_process_for_each_test("spawn")
95 def can_run_without_jit(spec: JitModel):
96 """Boot ``spec`` with the monitor armed and run the shape battery.
97
98 A subprocess per model is required: the monitor's hooks are process-global
99 and, once armed in ``error`` mode, stay armed. It must be spawned rather
100 than forked, since forking a pytest process that already initialized CUDA
101 poisons the child.
102 """
103 llm = LLM(
104 spec.model,
105 trust_remote_code=spec.trust_remote_code,
106 max_model_len=2048,
107 max_num_seqs=8,
108 gpu_memory_utilization=0.80,
109 load_format="dummy",
110 hf_overrides=dummy_hf_overrides,
111 # cuda graphs cover captured decode shapes, run eager.
112 enforce_eager=False,
113 jit_monitor_mode="error",
114 speculative_config={
115 "model": spec.draft,
116 "num_speculative_tokens": 2,
117 }
118 if spec.draft
119 else None,
120 )
121
122 try:
123 _run_shape_battery(llm)
124 except Exception as e:
125 # The monitor's message contains "during inference"; distinguish a real
126 # JIT miss from an unrelated crash.
127 if "during inference" in str(e):
128 pytest.fail(
129 f"{spec.model}: post-warmup JIT compilation detected - a warmup "
130 f"key is missing for a shape in the battery.\n{e}"
131 )
132 raise
133
134
135 @pytest.mark.parametrize("spec", JIT_MONITOR_MODELS, ids=lambda s: s.model)
136 def test_no_runtime_jit(spec: JitModel, monkeypatch: pytest.MonkeyPatch):
137 """Assert JIT-heavy backends do not JIT-compile during inference."""
138 # Set here rather than in the child so the spawned process inherits it:
139 # the engine core must not be forked once the test process has CUDA up.
140 monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
141 can_run_without_jit(spec)