@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 153 lines 4.89 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import json
5
6 import pytest
7
8 from vllm.config import ModelConfig, ParallelConfig, SpeculativeConfig
9 from vllm.transformers_utils.config import get_config
10 from vllm.transformers_utils.configs.k3_dspark import K3DSparkConfig
11
12
13 def _write_dspark_config(path, **overrides):
14 path.mkdir()
15 config = {
16 "architectures": ["K3DSparkModel"],
17 "model_type": "k3_dspark",
18 "hidden_size": 7168,
19 "intermediate_size": 14336,
20 "num_hidden_layers": 5,
21 "num_attention_heads": 96,
22 "num_key_value_heads": 96,
23 "q_lora_rank": 1536,
24 "kv_lora_rank": 512,
25 "qk_nope_head_dim": 128,
26 "qk_rope_head_dim": 64,
27 "v_head_dim": 128,
28 "vocab_size": 163840,
29 "rms_norm_eps": 1e-5,
30 "max_position_embeddings": 32768,
31 "rope_theta": 50000.0,
32 "num_target_layers": 5,
33 "target_hidden_size": 7168,
34 "target_num_hidden_layers": 93,
35 "target_layer_ids": [2, 23, 47, 71, 89],
36 "markov_rank": 256,
37 "draft_vocab_size": 163840,
38 "torch_dtype": "bfloat16",
39 }
40 config.update(overrides)
41 (path / "config.json").write_text(json.dumps(config))
42
43
44 def _write_target_config(path):
45 path.mkdir()
46 config = {
47 "architectures": ["LlamaForCausalLM"],
48 "model_type": "llama",
49 "hidden_size": 7168,
50 "intermediate_size": 14336,
51 "num_hidden_layers": 93,
52 "num_attention_heads": 56,
53 "num_key_value_heads": 8,
54 "vocab_size": 163840,
55 "max_position_embeddings": 32768,
56 "torch_dtype": "bfloat16",
57 }
58 (path / "config.json").write_text(json.dumps(config))
59
60
61 def test_dspark_mla_config_loads_from_local_json(tmp_path):
62 draft_path = tmp_path / "draft"
63 _write_dspark_config(draft_path)
64
65 config = get_config(draft_path, trust_remote_code=False)
66
67 assert isinstance(config, K3DSparkConfig)
68 assert config.model_type == "k3_dspark"
69 assert config.architectures == ["K3DSparkModel"]
70 assert config.hidden_act == "silu"
71 assert config.rope_parameters == {
72 "rope_type": "default",
73 "rope_theta": 50000.0,
74 }
75 assert config.n_routed_experts == 0
76 assert config.draft_vocab_size == config.vocab_size
77
78
79 @pytest.mark.parametrize(
80 "overrides",
81 [
82 {"mla_use_nope": True},
83 {"mla_use_output_gate": True},
84 {"mla_use_qk_norm": True},
85 {"dspark_bonus_anchor": True},
86 {"q_lora_rank": None},
87 {"draft_vocab_size": 8192},
88 {"target_layer_ids": []},
89 {"num_target_layers": 4},
90 ],
91 )
92 def test_dspark_mla_rejects_unsupported_checkpoint_options(tmp_path, overrides):
93 draft_path = tmp_path / "draft"
94 _write_dspark_config(draft_path, **overrides)
95
96 with pytest.raises(ValueError, match="MLA DSpark"):
97 get_config(draft_path, trust_remote_code=False)
98
99
100 @pytest.mark.parametrize(
101 ("num_attention_heads", "expected_local_heads"),
102 [(64, 8), (96, 12)],
103 ids=["released-64-head", "block5-96-head"],
104 )
105 def test_dspark_mla_uses_latent_kv_geometry(
106 tmp_path, num_attention_heads, expected_local_heads
107 ):
108 draft_path = tmp_path / "draft"
109 _write_dspark_config(
110 draft_path,
111 num_attention_heads=num_attention_heads,
112 num_key_value_heads=num_attention_heads,
113 )
114 model_config = ModelConfig(
115 model=str(draft_path),
116 tokenizer_mode="skip",
117 runner="draft",
118 max_model_len=32768,
119 )
120
121 assert model_config.is_deepseek_mla
122 assert model_config.use_mla
123 assert model_config.get_head_size() == 576
124 # external_launcher skips ParallelConfig's local-GPU-count check so the
125 # config logic can be exercised at TP8 on a single-GPU test node.
126 parallel_config = ParallelConfig(
127 tensor_parallel_size=8, distributed_executor_backend="external_launcher"
128 )
129 assert model_config.get_num_kv_heads(parallel_config) == 1
130 assert model_config.get_num_attention_heads(parallel_config) == expected_local_heads
131 assert model_config.get_num_experts() == 0
132
133
134 def test_dspark_mla_speculative_config_preserves_architecture(tmp_path):
135 target_path = tmp_path / "target"
136 draft_path = tmp_path / "draft"
137 _write_target_config(target_path)
138 _write_dspark_config(draft_path)
139 target_config = ModelConfig(
140 model=str(target_path), tokenizer_mode="skip", max_model_len=32768
141 )
142 speculative_config = SpeculativeConfig(
143 model=str(draft_path),
144 method="dspark",
145 num_speculative_tokens=8,
146 target_model_config=target_config,
147 target_parallel_config=ParallelConfig(),
148 )
149
150 assert speculative_config.parallel_drafting
151 assert speculative_config.draft_model_config.architectures == ["K3DSparkModel"]
152 assert speculative_config.draft_model_config.hf_config.model_type == "k3_dspark"
153 assert speculative_config.draft_model_config.use_mla