@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 239 lines 7.48 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import math
5 from pathlib import Path
6
7 import pytest
8
9 from tests.conftest import VllmRunner
10 from tests.utils import create_new_process_for_each_test
11 from vllm import SamplingParams, TokensPrompt
12 from vllm.config import CUDAGraphMode
13 from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config
14
15 MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
16
17
18 def _make_tiny_overrides() -> dict:
19 linear_attn_config = {
20 "full_attn_layers": [4],
21 "head_dim": 128,
22 "kda_layers": [1, 2, 3],
23 "num_heads": 8,
24 "short_conv_kernel_size": 4,
25 "use_full_rank_gate": True,
26 }
27
28 return {
29 "bos_token_id": 1,
30 "eos_token_id": 2,
31 "pad_token_id": 0,
32 "vocab_size": 256,
33 "hidden_size": 256,
34 "head_dim": 32,
35 "intermediate_size": 512,
36 "num_hidden_layers": 4,
37 "num_attention_heads": 8,
38 "num_key_value_heads": 8,
39 "num_experts": None,
40 "num_experts_per_token": None,
41 "num_shared_experts": 0,
42 "q_lora_rank": 128,
43 "kv_lora_rank": 512,
44 "qk_nope_head_dim": 128,
45 "qk_rope_head_dim": 64,
46 "v_head_dim": 128,
47 "mla_use_nope": True,
48 "mla_use_output_gate": True,
49 "linear_attn_config": linear_attn_config,
50 "max_position_embeddings": 512,
51 "model_max_length": 512,
52 }
53
54
55 def _run_tiny_model(
56 vllm_runner: type[VllmRunner], dcp_size: int
57 ) -> list[tuple[list[int], list[float]]]:
58 with vllm_runner(
59 model_name=MODEL,
60 skip_tokenizer_init=True,
61 load_format="dummy",
62 hf_overrides=_make_tiny_overrides(),
63 tensor_parallel_size=2,
64 decode_context_parallel_size=dcp_size,
65 cp_kv_cache_interleave_size=1,
66 distributed_executor_backend="mp",
67 dtype="bfloat16",
68 seed=0,
69 enforce_eager=True,
70 max_model_len=512,
71 max_num_seqs=8,
72 max_num_batched_tokens=1024,
73 gpu_memory_utilization=0.25,
74 enable_chunked_prefill=True,
75 enable_prefix_caching=False,
76 ) as runner:
77 lengths = [1, 32, 33, 129]
78 prompts = [
79 TokensPrompt(
80 prompt_token_ids=[
81 3 + ((request_idx + token_idx) % 251) for token_idx in range(length)
82 ]
83 )
84 for request_idx, length in enumerate(lengths)
85 ]
86 outputs = runner.llm.generate(
87 prompts,
88 SamplingParams(
89 temperature=0,
90 max_tokens=8,
91 seed=0,
92 ignore_eos=True,
93 logprobs=20,
94 ),
95 use_tqdm=False,
96 )
97
98 results = []
99 for request_output in outputs:
100 completion = request_output.outputs[0]
101 token_ids = list(completion.token_ids)
102 assert completion.logprobs is not None
103 selected_logprobs = [
104 step_logprobs[token_id].logprob
105 for token_id, step_logprobs in zip(token_ids, completion.logprobs)
106 ]
107 results.append((token_ids, selected_logprobs))
108 return results
109
110
111 @create_new_process_for_each_test()
112 @pytest.mark.distributed(num_gpus=2)
113 def test_kimi_linear_dcp_tiny(
114 vllm_runner: type[VllmRunner],
115 num_gpus_available: int,
116 ) -> None:
117 if num_gpus_available < 2:
118 pytest.skip("Need at least 2 GPUs")
119
120 baseline = _run_tiny_model(vllm_runner, dcp_size=1)
121 dcp = _run_tiny_model(vllm_runner, dcp_size=2)
122
123 assert [tokens for tokens, _ in dcp] == [tokens for tokens, _ in baseline]
124 logprob_drifts = []
125 for (_, baseline_logprobs), (_, dcp_logprobs) in zip(baseline, dcp):
126 for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs):
127 assert math.isfinite(baseline_logprob)
128 assert math.isfinite(dcp_logprob)
129 logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
130
131 assert max(logprob_drifts) <= 1e-2
132
133
134 def _make_tiny_k3_config(model_dir: Path) -> str:
135 config = KimiK3Config(
136 text_config=_make_tiny_overrides(),
137 vision_config={
138 "vt_num_attention_heads": 2,
139 "vt_num_hidden_layers": 1,
140 "vt_hidden_size": 32,
141 "vt_intermediate_size": 64,
142 "qkv_hidden_size": 48,
143 },
144 architectures=["KimiK3ForConditionalGeneration"],
145 )
146 config.save_pretrained(model_dir)
147 return str(model_dir)
148
149
150 def _run_k3_partial_prefix_reuse(
151 vllm_runner: type[VllmRunner], model_name: str, dcp_size: int
152 ) -> tuple[list[int], list[float], int]:
153 block_size = 256
154 common_prefix = [3 + (token_idx % 251) for token_idx in range(block_size)]
155 prime_prompt = TokensPrompt(prompt_token_ids=common_prefix + [17])
156 replay_prompt = TokensPrompt(prompt_token_ids=common_prefix + [18, 19])
157 sampling_params = SamplingParams(
158 temperature=0,
159 max_tokens=8,
160 seed=0,
161 ignore_eos=True,
162 logprobs=20,
163 )
164
165 with vllm_runner(
166 model_name=model_name,
167 skip_tokenizer_init=True,
168 load_format="dummy",
169 language_model_only=True,
170 tensor_parallel_size=2,
171 decode_context_parallel_size=dcp_size,
172 cp_kv_cache_interleave_size=1,
173 distributed_executor_backend="mp",
174 dtype="bfloat16",
175 seed=0,
176 max_model_len=512,
177 max_num_seqs=4,
178 max_num_batched_tokens=256,
179 gpu_memory_utilization=0.85,
180 block_size=block_size,
181 enable_chunked_prefill=True,
182 enable_prefix_caching=True,
183 mamba_cache_mode="align",
184 ) as runner:
185 assert (
186 runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode
187 == CUDAGraphMode.FULL_AND_PIECEWISE
188 )
189 runner.llm.generate(
190 [prime_prompt],
191 sampling_params,
192 use_tqdm=False,
193 )
194 replay_output = runner.llm.generate(
195 [replay_prompt],
196 sampling_params,
197 use_tqdm=False,
198 )[0]
199
200 completion = replay_output.outputs[0]
201 token_ids = list(completion.token_ids)
202 assert len(token_ids) == sampling_params.max_tokens
203 assert completion.logprobs is not None
204 assert len(completion.logprobs) == sampling_params.max_tokens
205 selected_logprobs = [
206 step_logprobs[token_id].logprob
207 for token_id, step_logprobs in zip(token_ids, completion.logprobs)
208 ]
209 return token_ids, selected_logprobs, replay_output.num_cached_tokens
210
211
212 @create_new_process_for_each_test()
213 @pytest.mark.distributed(num_gpus=2)
214 def test_kimi_k3_dcp_partial_prefix_reuse(
215 vllm_runner: type[VllmRunner],
216 num_gpus_available: int,
217 tmp_path: Path,
218 ) -> None:
219 if num_gpus_available < 2:
220 pytest.skip("Need at least 2 GPUs")
221
222 model_name = _make_tiny_k3_config(tmp_path / "tiny-kimi-k3")
223 baseline_tokens, baseline_logprobs, baseline_cached = _run_k3_partial_prefix_reuse(
224 vllm_runner, model_name, dcp_size=1
225 )
226 dcp_tokens, dcp_logprobs, dcp_cached = _run_k3_partial_prefix_reuse(
227 vllm_runner, model_name, dcp_size=2
228 )
229
230 assert baseline_cached == 256
231 assert dcp_cached == 256
232 assert dcp_tokens == baseline_tokens
233 assert len(dcp_logprobs) == len(baseline_logprobs) == 8
234 logprob_drifts = []
235 for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs):
236 assert math.isfinite(baseline_logprob)
237 assert math.isfinite(dcp_logprob)
238 logprob_drifts.append(abs(baseline_logprob - dcp_logprob))
239 assert max(logprob_drifts) <= 1e-2