@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 185 lines 6.18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import math
5
6 import pytest
7
8 from tests.utils import multi_gpu_marks
9 from vllm import SamplingParams
10 from vllm.config import CompilationConfig
11 from vllm.config.compilation import CompilationMode, CUDAGraphMode
12 from vllm.platforms import current_platform
13 from vllm.v1.metrics.reader import Counter
14
15 pytestmark = pytest.mark.skipif(
16 not current_platform.is_rocm(), reason="ROCm-specific tests"
17 )
18
19
20 def _small_hyv4_config(config):
21 # Apply to both target and MTP draft while retaining MLA/indexer dimensions.
22 config.update(
23 {
24 "hidden_size": 512,
25 "num_hidden_layers": 2,
26 "num_attention_heads": 64,
27 "q_lora_rank": 128,
28 "intermediate_size": 512,
29 "moe_intermediate_size": 256,
30 "n_routed_experts": 4,
31 "num_experts_per_tok": 2,
32 "n_shared_experts": 1,
33 "enable_ihc": True,
34 "index_n_heads": 32,
35 "index_topk": 2048,
36 "layer_types": ["deepseek_sparse_attention"] * 2,
37 "indexer_types": ["full", "shared"],
38 "num_nextn_predict_layers": 1,
39 "vocab_size": 256,
40 "bos_token_id": 1,
41 "eos_token_id": 2,
42 "pad_token_id": 0,
43 }
44 )
45 return config
46
47
48 def _num_captured_graphs(_worker):
49 from vllm.compilation.counter import compilation_counter
50
51 return compilation_counter.num_cudagraph_captured
52
53
54 def _piecewise_compile_state(worker):
55 from vllm.compilation.counter import compilation_counter
56
57 # `cudagraph_mode` is returned by name: the enum does not survive the
58 # collective_rpc round trip, it arrives as its bare value.
59 return (
60 compilation_counter.num_models_seen,
61 compilation_counter.num_cudagraph_captured,
62 worker.vllm_config.compilation_config.cudagraph_mode.name,
63 )
64
65
66 @pytest.mark.parametrize(
67 "tp,enforce_eager,mtp,dtype,kv_cache_dtype",
68 [
69 pytest.param(1, True, False, "bfloat16", "auto", id="eager-tp1"),
70 pytest.param(1, True, False, "float16", "auto", id="eager-fp16"),
71 pytest.param(1, True, False, "bfloat16", "fp8", id="eager-fp8-kv"),
72 pytest.param(
73 2,
74 False,
75 False,
76 "bfloat16",
77 "auto",
78 id="graphs-tp2",
79 marks=multi_gpu_marks(num_gpus=2),
80 ),
81 pytest.param(1, True, True, "bfloat16", "auto", id="eager-mtp"),
82 ],
83 )
84 def test_hyv4_dummy_generation(
85 vllm_runner,
86 enable_pickle,
87 tp: int,
88 enforce_eager: bool,
89 mtp: bool,
90 dtype: str,
91 kv_cache_dtype: str,
92 ) -> None:
93 """Exercise prefill, decode, and batch consistency with dummy HY weights."""
94 with vllm_runner(
95 "tencent/Hy4-preview",
96 trust_remote_code=True,
97 skip_tokenizer_init=True,
98 load_format="dummy",
99 disable_log_stats=not mtp,
100 dtype=dtype,
101 kv_cache_dtype=kv_cache_dtype,
102 hf_overrides=_small_hyv4_config,
103 max_model_len=256,
104 max_num_batched_tokens=128,
105 max_num_seqs=8,
106 enable_chunked_prefill=True,
107 enable_prefix_caching=False,
108 kv_cache_memory_bytes=128 * 1024 * 1024,
109 enforce_eager=enforce_eager,
110 tensor_parallel_size=tp,
111 speculative_config={"method": "mtp", "num_speculative_tokens": 2}
112 if mtp
113 else None,
114 seed=17,
115 # Exercise supported decode graphs independently of default graph policy.
116 compilation_config=CompilationConfig(
117 mode=CompilationMode.NONE,
118 cudagraph_mode=CUDAGraphMode.NONE
119 if enforce_eager
120 else CUDAGraphMode.FULL_DECODE_ONLY,
121 cudagraph_capture_sizes=[] if enforce_eager else [1, 2, 4],
122 ),
123 ) as model:
124 if not enforce_eager:
125 assert model.collective_rpc(_num_captured_graphs) == [3] * tp
126 prompts = [
127 {"prompt_token_ids": [11, 12, 13, 14]},
128 {"prompt_token_ids": list(range(20, 84))},
129 ]
130 params = SamplingParams(
131 temperature=0, max_tokens=4, ignore_eos=True, logprobs=1
132 )
133 outputs = model.llm.generate(prompts, params)
134 assert len(outputs) == len(prompts)
135 for prompt, output in zip(prompts, outputs):
136 completion = output.outputs[0]
137 assert len(completion.token_ids) == 4
138 assert completion.cumulative_logprob is not None
139 assert math.isfinite(completion.cumulative_logprob)
140 single = model.llm.generate([prompt], params)[0].outputs[0]
141 assert single.token_ids == completion.token_ids
142
143 if mtp:
144 assert any(
145 isinstance(metric, Counter)
146 and metric.name == "vllm:spec_decode_num_drafts"
147 and metric.value > 0
148 for metric in model.llm.get_metrics()
149 ), "MTP drafter did not run"
150
151
152 def test_hyv4_backbone_reaches_the_compiler(vllm_runner, enable_pickle) -> None:
153 """Guard against the backbone losing `@support_torch_compile`.
154
155 An undecorated backbone is not a hard failure: vLLM downgrades
156 `cudagraph_mode` to NONE and serves the whole model eager, which on TP8
157 gfx942 cost ~4x decode latency.
158 """
159 with vllm_runner(
160 "tencent/Hy4-preview",
161 trust_remote_code=True,
162 skip_tokenizer_init=True,
163 load_format="dummy",
164 dtype="bfloat16",
165 hf_overrides=_small_hyv4_config,
166 max_model_len=256,
167 max_num_batched_tokens=128,
168 max_num_seqs=8,
169 enable_chunked_prefill=True,
170 enable_prefix_caching=False,
171 kv_cache_memory_bytes=128 * 1024 * 1024,
172 enforce_eager=False,
173 seed=17,
174 compilation_config=CompilationConfig(
175 mode=CompilationMode.VLLM_COMPILE,
176 cudagraph_mode=CUDAGraphMode.PIECEWISE,
177 cudagraph_capture_sizes=[1, 2, 4],
178 ),
179 ) as model:
180 for models_seen, captured, mode in model.collective_rpc(
181 _piecewise_compile_state
182 ):
183 assert models_seen > 0, "backbone was never handed to the compiler"
184 assert mode == CUDAGraphMode.PIECEWISE.name
185 assert captured > 0