[Bugfix][Model] Fix K3 DSpark config for 96-head drafts (#54005)
Signed-off-by: zixi-qi <zixi@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
qizixi committed
Aug 27, 2026 at 15:58 UTC
e5377e6727454f50d15d8b6225f84b76a8dee5f6
2 files changed
+19
-8
tests/transformers_utils/test_dspark_mla_config.py
+16
-5
@@ -18,8 +18,8 @@ def _write_dspark_config(path, **overrides):
18
"hidden_size": 7168,
19
"intermediate_size": 14336,
20
"num_hidden_layers": 5,
21
- "num_attention_heads": 64,
22
- "num_key_value_heads": 64,
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,
@@ -97,9 +97,20 @@ def test_dspark_mla_rejects_unsupported_checkpoint_options(tmp_path, overrides):
97
get_config(draft_path, trust_remote_code=False)
98
99
100
-def test_dspark_mla_uses_latent_kv_geometry(tmp_path):
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"
102
- _write_dspark_config(draft_path)
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",
@@ -116,7 +127,7 @@ def test_dspark_mla_uses_latent_kv_geometry(tmp_path):
127
tensor_parallel_size=8, distributed_executor_backend="external_launcher"
128
)
129
assert model_config.get_num_kv_heads(parallel_config) == 1
119
- assert model_config.get_num_attention_heads(parallel_config) == 8
130
+ assert model_config.get_num_attention_heads(parallel_config) == expected_local_heads
131
assert model_config.get_num_experts() == 0
132
133
vllm/transformers_utils/configs/k3_dspark.py
+3
-3
@@ -1,10 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from transformers import DeepseekV2Config
4
+from transformers import DeepseekV3Config
5
6
7
-class K3DSparkConfig(DeepseekV2Config):
7
+class K3DSparkConfig(DeepseekV3Config):
8
"""Configuration for a dense MLA DSpark draft model."""
9
10
model_type = "k3_dspark"
@@ -18,7 +18,7 @@ class K3DSparkConfig(DeepseekV2Config):
18
rope_theta: float = 50000.0,
19
**kwargs,
20
) -> None:
21
- # DeepseekV2Config defaults to a MoE topology. Zero these fields so
21
+ # DeepseekV3Config defaults to a MoE topology. Zero these fields so
22
# generic vLLM config logic also recognizes this draft as dense.
23
kwargs.setdefault("n_routed_experts", 0)
24
kwargs.setdefault("n_shared_experts", 0)