| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | import torch.nn as nn |
| 9 | |
| 10 | from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper |
| 11 | from vllm.model_executor.layers.logits_processor import LogitsProcessor |
| 12 | from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead |
| 13 | from vllm.model_executor.models.registry import ModelRegistry |
| 14 | from vllm.models.deepseek_v4.nvidia import dspark as dsv4_dspark |
| 15 | from vllm.models.kimi_k3.nvidia import dspark_mla |
| 16 | from vllm.models.kimi_k3.nvidia.dspark_mla import K3DSparkForCausalLM, K3DSparkModel |
| 17 | |
| 18 | |
| 19 | def test_dspark_mla_uses_compile_free_model_entrypoint(): |
| 20 | assert ModelRegistry._try_load_model_cls("K3DSparkModel") is K3DSparkForCausalLM |
| 21 | assert not issubclass(K3DSparkModel, TorchCompileWithNoGuardsWrapper) |
| 22 | |
| 23 | |
| 24 | @pytest.mark.parametrize( |
| 25 | ("checkpoint_name", "runtime_name", "shard_id"), |
| 26 | [ |
| 27 | ( |
| 28 | "layers.0.self_attn.q_a_proj.weight", |
| 29 | "model.layers.0.self_attn.fused_qkv_a_proj.weight", |
| 30 | 0, |
| 31 | ), |
| 32 | ( |
| 33 | "layers.0.self_attn.kv_a_proj_with_mqa.weight", |
| 34 | "model.layers.0.self_attn.fused_qkv_a_proj.weight", |
| 35 | 1, |
| 36 | ), |
| 37 | ( |
| 38 | "layers.0.mlp.gate_proj.weight", |
| 39 | "model.layers.0.mlp.gate_up_proj.weight", |
| 40 | 0, |
| 41 | ), |
| 42 | ( |
| 43 | "layers.0.mlp.up_proj.weight", |
| 44 | "model.layers.0.mlp.gate_up_proj.weight", |
| 45 | 1, |
| 46 | ), |
| 47 | ("context_proj.weight", "model.context_proj.weight", None), |
| 48 | ], |
| 49 | ) |
| 50 | def test_dspark_mla_checkpoint_weight_mapping(checkpoint_name, runtime_name, shard_id): |
| 51 | assert K3DSparkForCausalLM.hf_to_vllm_mapper._map_name_with_shard( |
| 52 | checkpoint_name |
| 53 | ) == (runtime_name, shard_id) |
| 54 | |
| 55 | |
| 56 | def test_dspark_mla_shares_frozen_target_weights_and_skips_training_head(): |
| 57 | assert not K3DSparkForCausalLM.has_own_embed_tokens |
| 58 | assert not K3DSparkForCausalLM.has_own_lm_head |
| 59 | mapper = K3DSparkForCausalLM.hf_to_vllm_mapper |
| 60 | for name in ("confidence_head.weight", "embed_tokens.weight", "lm_head.weight"): |
| 61 | assert mapper._map_name(name) is None |
| 62 | |
| 63 | |
| 64 | @pytest.mark.cpu_test |
| 65 | def test_dspark_markov_head_is_replicated( |
| 66 | monkeypatch: pytest.MonkeyPatch, |
| 67 | ): |
| 68 | from vllm.model_executor.layers import logits_processor, vocab_parallel_embedding |
| 69 | |
| 70 | monkeypatch.setattr( |
| 71 | vocab_parallel_embedding, "get_tensor_model_parallel_rank", lambda: 3 |
| 72 | ) |
| 73 | monkeypatch.setattr( |
| 74 | vocab_parallel_embedding, |
| 75 | "get_tensor_model_parallel_world_size", |
| 76 | lambda: 8, |
| 77 | ) |
| 78 | monkeypatch.setattr( |
| 79 | logits_processor, |
| 80 | "get_current_vllm_config", |
| 81 | lambda: SimpleNamespace(model_config=None), |
| 82 | ) |
| 83 | |
| 84 | head = DSparkMarkovHead(128, 128, 8, prefix="markov_head") |
| 85 | assert head.markov_w2.tp_size == 1 |
| 86 | assert head.markov_w1.weight.shape == (128, 8) |
| 87 | assert head.markov_w2.weight.shape == (128, 8) |
| 88 | |
| 89 | def fail_collective(*args, **kwargs): |
| 90 | raise AssertionError("replicated Markov head must not invoke TP collectives") |
| 91 | |
| 92 | monkeypatch.setattr( |
| 93 | vocab_parallel_embedding, |
| 94 | "tensor_model_parallel_all_reduce", |
| 95 | fail_collective, |
| 96 | ) |
| 97 | logits_processor = LogitsProcessor(128) |
| 98 | monkeypatch.setattr(logits_processor, "_gather_logits", fail_collective) |
| 99 | |
| 100 | markov_embed = head.embed(torch.tensor([1, 2])) |
| 101 | bias = head.bias(markov_embed, logits_processor) |
| 102 | assert markov_embed.shape == (2, 8) |
| 103 | assert bias.shape == (2, 128) |
| 104 | |
| 105 | |
| 106 | @pytest.mark.cpu_test |
| 107 | def test_k3_dspark_uses_replicated_markov_head(monkeypatch: pytest.MonkeyPatch): |
| 108 | markov_head_calls = [] |
| 109 | context_kv_proj_calls = [] |
| 110 | |
| 111 | class DummyModule(nn.Module): |
| 112 | def __init__(self, *args, **kwargs): |
| 113 | super().__init__() |
| 114 | |
| 115 | def make_markov_head(*args, **kwargs): |
| 116 | markov_head_calls.append((args, kwargs)) |
| 117 | return DummyModule() |
| 118 | |
| 119 | def make_context_kv_proj(*args, **kwargs): |
| 120 | context_kv_proj_calls.append((args, kwargs)) |
| 121 | return DummyModule() |
| 122 | |
| 123 | monkeypatch.setattr(dspark_mla, "get_draft_quant_config", lambda _: None) |
| 124 | monkeypatch.setattr(dspark_mla, "ReplicatedLinear", DummyModule) |
| 125 | monkeypatch.setattr(dspark_mla, "MergedColumnParallelLinear", make_context_kv_proj) |
| 126 | monkeypatch.setattr(dspark_mla, "RMSNorm", DummyModule) |
| 127 | monkeypatch.setattr(dspark_mla, "K3DSparkDecoderLayer", DummyModule) |
| 128 | monkeypatch.setattr(dspark_mla, "DSparkMarkovHead", make_markov_head) |
| 129 | |
| 130 | config = SimpleNamespace( |
| 131 | target_hidden_size=16, |
| 132 | num_target_layers=2, |
| 133 | hidden_size=8, |
| 134 | kv_lora_rank=3, |
| 135 | qk_rope_head_dim=1, |
| 136 | rms_norm_eps=1e-6, |
| 137 | num_hidden_layers=1, |
| 138 | vocab_size=128, |
| 139 | draft_vocab_size=128, |
| 140 | markov_rank=4, |
| 141 | ) |
| 142 | vllm_config = SimpleNamespace( |
| 143 | speculative_config=SimpleNamespace( |
| 144 | draft_model_config=SimpleNamespace(hf_config=config) |
| 145 | ), |
| 146 | scheduler_config=SimpleNamespace(max_num_batched_tokens=16), |
| 147 | ) |
| 148 | |
| 149 | K3DSparkModel(vllm_config=vllm_config, start_layer_id=0, prefix="model") |
| 150 | |
| 151 | assert len(markov_head_calls) == 1 |
| 152 | assert context_kv_proj_calls == [ |
| 153 | ( |
| 154 | (8, [4]), |
| 155 | { |
| 156 | "bias": False, |
| 157 | "return_bias": False, |
| 158 | "quant_config": None, |
| 159 | "prefix": "model.layers.0.self_attn.fused_qkv_a_proj", |
| 160 | "disable_tp": True, |
| 161 | }, |
| 162 | ) |
| 163 | ] |
| 164 | |
| 165 | |
| 166 | def test_context_kv_weights_are_loaded_as_merged_linear_shards(): |
| 167 | weights = [ |
| 168 | ( |
| 169 | "layers.0.self_attn.kv_a_proj_with_mqa.weight_packed", |
| 170 | torch.arange(4), |
| 171 | ), |
| 172 | ( |
| 173 | "layers.1.self_attn.kv_a_proj_with_mqa.weight_scale", |
| 174 | torch.tensor(0.5), |
| 175 | ), |
| 176 | ] |
| 177 | |
| 178 | duplicated = dspark_mla._duplicate_context_kv_weights(weights, 2) |
| 179 | mapped = list(K3DSparkForCausalLM.hf_to_vllm_mapper.apply(duplicated)) |
| 180 | |
| 181 | assert [name for name, _ in mapped] == [ |
| 182 | "model.layers.0.self_attn.fused_qkv_a_proj.weight_packed", |
| 183 | "model.context_kv_proj.weight_packed", |
| 184 | "model.layers.1.self_attn.fused_qkv_a_proj.weight_scale", |
| 185 | "model.context_kv_proj.weight_scale", |
| 186 | ] |
| 187 | assert [weight.shard_id for _, weight in mapped] == [1, 0, 1, 1] |
| 188 | assert mapped[0][1].data_ptr() == mapped[1][1].data_ptr() |
| 189 | assert mapped[2][1].data_ptr() == mapped[3][1].data_ptr() |
| 190 | |
| 191 | |
| 192 | @pytest.mark.cpu_test |
| 193 | @pytest.mark.parametrize( |
| 194 | "scale_dtype", [torch.uint8, torch.float8_e8m0fnu, torch.float32] |
| 195 | ) |
| 196 | @pytest.mark.parametrize( |
| 197 | ("checkpoint_name", "runtime_module", "shard_id"), |
| 198 | [ |
| 199 | ("mtp.0.attn.wq_a.scale", "model.layers.0.attn.fused_wqa_wkv", 0), |
| 200 | ("mtp.0.attn.wkv.scale", "model.layers.0.attn.fused_wqa_wkv", 1), |
| 201 | ("mtp.0.main_proj.scale", "model.main_proj", None), |
| 202 | ( |
| 203 | "mtp.0.ffn.shared_experts.w1.scale", |
| 204 | "model.layers.0.ffn.shared_experts.gate_up_proj", |
| 205 | 0, |
| 206 | ), |
| 207 | ( |
| 208 | "mtp.0.ffn.shared_experts.w2.scale", |
| 209 | "model.layers.0.ffn.shared_experts.down_proj", |
| 210 | None, |
| 211 | ), |
| 212 | ], |
| 213 | ) |
| 214 | def test_v41_dspark_loads_linear_scales( |
| 215 | monkeypatch, scale_dtype, checkpoint_name, runtime_module, shard_id |
| 216 | ): |
| 217 | """Checkpoint ``.scale`` maps to the quant method's scale parameter and |
| 218 | loads untouched. MXFP8 block-scale expansion lives in the KMxfp8Static |
| 219 | loader (see tests/quantization/test_modelopt.py), not in load_weights.""" |
| 220 | from vllm.models.deepseek_v41.nvidia import dspark |
| 221 | |
| 222 | mxfp8 = scale_dtype != torch.float32 |
| 223 | scale_name = "weight_scale" if mxfp8 else "weight_scale_inv" |
| 224 | runtime_name = f"{runtime_module}.{scale_name}" |
| 225 | raw = torch.tensor([[120, 127], [128, 130]], dtype=torch.uint8) |
| 226 | checkpoint_scale = raw.view(scale_dtype) if mxfp8 else raw.float() |
| 227 | param = nn.Parameter(torch.empty_like(checkpoint_scale), requires_grad=False) |
| 228 | shards = [] |
| 229 | |
| 230 | def load_scale(param, weight, *args): |
| 231 | shards.append(args) |
| 232 | assert weight.dtype == checkpoint_scale.dtype |
| 233 | param.copy_(weight) |
| 234 | |
| 235 | param.weight_loader = load_scale |
| 236 | draft = SimpleNamespace( |
| 237 | config=SimpleNamespace(num_attention_heads=4, n_routed_experts=1), |
| 238 | quant_config=SimpleNamespace( |
| 239 | weight_block_size=[32, 32] if mxfp8 else [128, 128] |
| 240 | ), |
| 241 | linear_scale_name=scale_name, |
| 242 | pad_shared_expert=False, |
| 243 | model=SimpleNamespace( |
| 244 | layers=[SimpleNamespace(ffn=SimpleNamespace(use_mega_moe=False))], |
| 245 | confidence_head=None, |
| 246 | ), |
| 247 | named_parameters=lambda: [(runtime_name, param)], |
| 248 | process_weights_after_loading=lambda: None, |
| 249 | ) |
| 250 | draft._remap_dspark_name = lambda name: ( |
| 251 | dspark.DSparkDeepseekV4ForCausalLM._remap_dspark_name(draft, name) |
| 252 | ) |
| 253 | monkeypatch.setattr(dspark, "get_tensor_model_parallel_world_size", lambda: 4) |
| 254 | monkeypatch.setattr(dspark, "get_tensor_model_parallel_rank", lambda: 0) |
| 255 | monkeypatch.setattr( |
| 256 | dspark, "fused_moe_make_expert_params_mapping", lambda *a, **kw: [] |
| 257 | ) |
| 258 | |
| 259 | loaded = dspark.DSparkDeepseekV4ForCausalLM.load_weights( |
| 260 | draft, [(checkpoint_name, checkpoint_scale)] |
| 261 | ) |
| 262 | |
| 263 | assert loaded == {runtime_name} |
| 264 | assert shards == [() if shard_id is None else (shard_id,)] |
| 265 | torch.testing.assert_close(param, checkpoint_scale) |
| 266 | |
| 267 | |
| 268 | def test_dsv4_context_wkv_weights_are_duplicated_by_draft_layer(): |
| 269 | weights = [ |
| 270 | ("mtp.0.attn.wkv.weight", torch.arange(4)), |
| 271 | ("mtp.1.attn.wq_a.weight", torch.arange(3)), |
| 272 | ("mtp.2.attn.wkv.scale", torch.tensor(0.5)), |
| 273 | ("mtp.3.attn.wkv.weight", torch.arange(2)), |
| 274 | ] |
| 275 | |
| 276 | duplicated = list(dsv4_dspark._duplicate_context_wkv_weights(weights, 3)) |
| 277 | |
| 278 | assert [name for name, _ in duplicated] == [ |
| 279 | "mtp.0.attn.wkv.weight", |
| 280 | "context_wkv_proj.weight", |
| 281 | "mtp.1.attn.wq_a.weight", |
| 282 | "mtp.2.attn.wkv.scale", |
| 283 | "context_wkv_proj.scale", |
| 284 | "mtp.3.attn.wkv.weight", |
| 285 | ] |
| 286 | assert duplicated[1][1].shard_id == 0 |
| 287 | assert duplicated[4][1].shard_id == 2 |
| 288 | assert duplicated[0][1].data_ptr() == duplicated[1][1].data_ptr() |
| 289 | assert duplicated[3][1].data_ptr() == duplicated[4][1].data_ptr() |
| 290 | |
| 291 | |
| 292 | def test_dsv4_context_kv_uses_one_stacked_wkv_projection(monkeypatch): |
| 293 | calls = [] |
| 294 | stacked_output = torch.arange(24, dtype=torch.float32).view(2, 12) |
| 295 | |
| 296 | class StackedProjection: |
| 297 | def __init__(self): |
| 298 | self.calls = 0 |
| 299 | |
| 300 | def __call__(self, main_x): |
| 301 | self.calls += 1 |
| 302 | assert main_x.shape == (2, 5) |
| 303 | return stacked_output |
| 304 | |
| 305 | projection = StackedProjection() |
| 306 | layers = [ |
| 307 | SimpleNamespace(attn=SimpleNamespace(kv_norm=lambda kv, offset=i: kv + offset)) |
| 308 | for i in range(3) |
| 309 | ] |
| 310 | model = SimpleNamespace( |
| 311 | config=SimpleNamespace(head_dim=4), |
| 312 | context_wkv_proj=projection, |
| 313 | layers=layers, |
| 314 | num_dspark_layers=3, |
| 315 | ) |
| 316 | slot_mappings = [torch.tensor([0, 1]), None, torch.tensor([4, 5])] |
| 317 | monkeypatch.setattr( |
| 318 | dsv4_dspark, |
| 319 | "_insert_context_kv", |
| 320 | lambda attn, kv, positions, slots: calls.append( |
| 321 | (attn, kv.clone(), positions, slots) |
| 322 | ), |
| 323 | ) |
| 324 | |
| 325 | dsv4_dspark.DSparkDeepseekV4Model.precompute_and_store_context_kv( |
| 326 | model, |
| 327 | torch.zeros(2, 5), |
| 328 | torch.tensor([7, 8]), |
| 329 | slot_mappings, |
| 330 | ) |
| 331 | |
| 332 | assert projection.calls == 1 |
| 333 | assert len(calls) == 2 |
| 334 | assert torch.equal(calls[0][1], stacked_output.view(2, 3, 4)[:, 0]) |
| 335 | assert torch.equal(calls[1][1], stacked_output.view(2, 3, 4)[:, 2] + 2) |
| 336 | assert calls[0][3] is slot_mappings[0] |
| 337 | assert calls[1][3] is slot_mappings[2] |