| 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 | |
| 9 | from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( |
| 10 | bind_routed_experts_capturer, |
| 11 | ) |
| 12 | from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM |
| 13 | from vllm.models.deepseek_v4.nvidia.model import ( |
| 14 | DeepseekV4ForCausalLM, |
| 15 | DeepseekV4MegaMoEExperts, |
| 16 | DeepseekV4MoE, |
| 17 | make_deepseek_v4_expert_params_mapping, |
| 18 | prepare_mega_gate_routing_metadata, |
| 19 | ) |
| 20 | from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP |
| 21 | from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs |
| 22 | from vllm.models.deepseek_v41.common.mm_preprocess import IMAGE_SENTINEL_BASE_ID |
| 23 | from vllm.models.deepseek_v41.nvidia.model import DeepseekV4MoE as DeepseekV41MoE |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.transformers_utils.configs.deepseek_v41 import DeepseekV41Config |
| 26 | from vllm.utils.torch_utils import set_default_torch_dtype |
| 27 | |
| 28 | pytestmark = pytest.mark.skipif( |
| 29 | not current_platform.is_cuda(), |
| 30 | reason="DeepSeek V4 MegaMoE requires CUDA", |
| 31 | ) |
| 32 | |
| 33 | |
| 34 | @pytest.fixture |
| 35 | def v41_moe_config(dist_init): |
| 36 | return SimpleNamespace( |
| 37 | model_config=SimpleNamespace( |
| 38 | dtype=torch.bfloat16, |
| 39 | hf_config=DeepseekV41Config( |
| 40 | text_config=dict( |
| 41 | hidden_size=128, |
| 42 | num_hidden_layers=2, |
| 43 | n_routed_experts=8, |
| 44 | num_experts_per_tok=2, |
| 45 | dspark_n_routed_experts=4, |
| 46 | dspark_num_experts_per_tok=3, |
| 47 | n_shared_experts=1, |
| 48 | moe_intermediate_size=128, |
| 49 | hidden_act="silu", |
| 50 | swiglu_limit=10.0, |
| 51 | norm_topk_prob=True, |
| 52 | topk_method="noaux_tc", |
| 53 | routed_scaling_factor=1.5, |
| 54 | ), |
| 55 | ), |
| 56 | ), |
| 57 | quant_config=None, |
| 58 | kernel_config=SimpleNamespace(moe_backend="deep_gemm_mega_moe"), |
| 59 | parallel_config=SimpleNamespace( |
| 60 | enable_expert_parallel=True, |
| 61 | enable_eplb=False, |
| 62 | eplb_config=SimpleNamespace(num_redundant_experts=0), |
| 63 | ), |
| 64 | scheduler_config=SimpleNamespace(max_num_batched_tokens=4), |
| 65 | compilation_config=SimpleNamespace(static_forward_context={}), |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | @pytest.mark.parametrize("use_cudagraph", [False, True]) |
| 70 | @pytest.mark.parametrize("above_threshold", [False, True]) |
| 71 | @pytest.mark.parametrize("vision", [False, True]) |
| 72 | @pytest.mark.parametrize("layer_id,num_experts,top_k", [(0, 384, 6), (2, 128, 3)]) |
| 73 | def test_deepseek_v41_moe_routes_without_hash_table( |
| 74 | v41_moe_config, |
| 75 | monkeypatch, |
| 76 | vision, |
| 77 | layer_id, |
| 78 | num_experts, |
| 79 | top_k, |
| 80 | above_threshold, |
| 81 | use_cudagraph, |
| 82 | ): |
| 83 | """Main and draft layers select experts by score and preserve image routing.""" |
| 84 | if not current_platform.is_device_capability_family(100): |
| 85 | pytest.skip("DeepGEMM Mega Gate requires SM100") |
| 86 | |
| 87 | config = v41_moe_config.model_config.hf_config |
| 88 | config.hidden_size = 256 |
| 89 | config.n_routed_experts = 384 |
| 90 | config.num_experts_per_tok = 6 |
| 91 | config.dspark_n_routed_experts = 128 |
| 92 | num_tokens = (1 if num_experts == 128 else 16) + int(above_threshold) |
| 93 | config.vision_n_layers = int(vision) |
| 94 | with ( |
| 95 | set_default_torch_dtype(v41_moe_config.model_config.dtype), |
| 96 | torch.device("cuda"), |
| 97 | ): |
| 98 | moe = DeepseekV41MoE(v41_moe_config, prefix=f"model.layers.{layer_id}.ffn") |
| 99 | hidden_states = torch.randn( |
| 100 | num_tokens, config.hidden_size, dtype=torch.bfloat16 |
| 101 | ) |
| 102 | input_ids = ( |
| 103 | torch.tensor( |
| 104 | [42, IMAGE_SENTINEL_BASE_ID, IMAGE_SENTINEL_BASE_ID, 129257] |
| 105 | ).repeat((num_tokens + 3) // 4)[:num_tokens] |
| 106 | if vision |
| 107 | else None |
| 108 | ) |
| 109 | assert moe.gate.tid2eid is None |
| 110 | assert moe.gate.weight.shape == (num_experts, config.hidden_size) |
| 111 | assert isinstance(moe.experts, DeepseekV4MegaMoEExperts) |
| 112 | assert moe.experts.w13_weight.shape == (num_experts, 256, 128) |
| 113 | assert moe.experts.top_k == top_k |
| 114 | assert (config.n_routed_experts, config.num_experts_per_tok) == (384, 6) |
| 115 | |
| 116 | with torch.no_grad(): |
| 117 | moe.gate.weight.normal_(std=0.01) |
| 118 | moe.gate.e_score_correction_bias.copy_(torch.arange(num_experts, device="cuda")) |
| 119 | if vision: |
| 120 | moe.gate.bias_vl.copy_(-moe.gate.e_score_correction_bias) |
| 121 | |
| 122 | scores = torch.nn.functional.softplus( |
| 123 | torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32) |
| 124 | ).sqrt() |
| 125 | bias = moe.gate.e_score_correction_bias |
| 126 | if vision: |
| 127 | image_mask = input_ids == IMAGE_SENTINEL_BASE_ID |
| 128 | bias = torch.where(image_mask[:, None], moe.gate.bias_vl, bias) |
| 129 | expected_ids = (scores + bias).topk(top_k, dim=-1).indices |
| 130 | expected_weights = scores.gather(1, expected_ids) |
| 131 | expected_weights *= config.routed_scaling_factor / expected_weights.sum( |
| 132 | dim=-1, keepdim=True |
| 133 | ) |
| 134 | |
| 135 | routed = {} |
| 136 | |
| 137 | def check_routing(x, weights, ids, *, activation_clamp): |
| 138 | routed["ids"], routed["weights"] = ids, weights |
| 139 | assert activation_clamp == config.swiglu_limit |
| 140 | return x.clone() |
| 141 | |
| 142 | monkeypatch.setattr(moe.experts, "forward", check_routing) |
| 143 | monkeypatch.setattr(moe.shared_experts, "forward", lambda x: 2 * x) |
| 144 | routing_input_ids = ( |
| 145 | input_ids |
| 146 | if input_ids is not None |
| 147 | else torch.zeros(num_tokens, dtype=torch.int64, device="cuda") |
| 148 | ) |
| 149 | metadata = prepare_mega_gate_routing_metadata( |
| 150 | routing_input_ids, |
| 151 | has_hash_routing=False, |
| 152 | image_sentinel_base_id=IMAGE_SENTINEL_BASE_ID if vision else None, |
| 153 | ) |
| 154 | if use_cudagraph: |
| 155 | stream = torch.cuda.Stream() |
| 156 | stream.wait_stream(torch.cuda.current_stream()) |
| 157 | with torch.cuda.stream(stream): |
| 158 | for _ in range(3): |
| 159 | moe(hidden_states, input_ids, metadata) |
| 160 | stream.synchronize() |
| 161 | graph = torch.cuda.CUDAGraph() |
| 162 | with torch.cuda.graph(graph, stream=stream): |
| 163 | output = moe(hidden_states, input_ids, metadata) |
| 164 | graph.replay() |
| 165 | torch.accelerator.synchronize() |
| 166 | else: |
| 167 | output = moe(hidden_states, input_ids, metadata) |
| 168 | |
| 169 | ids, weights = routed["ids"], routed["weights"] |
| 170 | torch.testing.assert_close(ids, expected_ids) |
| 171 | torch.testing.assert_close( |
| 172 | weights, |
| 173 | expected_weights, |
| 174 | rtol=1e-3, |
| 175 | atol=1e-4, |
| 176 | ) |
| 177 | torch.testing.assert_close(output, 3 * hidden_states) |
| 178 | |
| 179 | |
| 180 | @pytest.mark.parametrize( |
| 181 | "draft_experts,draft_top_k,expected", [(0, 0, (8, 2)), (4, 3, (4, 3))] |
| 182 | ) |
| 183 | def test_deepseek_v41_fused_moe_uses_draft_counts_or_main_defaults( |
| 184 | v41_moe_config, monkeypatch, draft_experts, draft_top_k, expected |
| 185 | ): |
| 186 | config = v41_moe_config.model_config.hf_config |
| 187 | config.dspark_n_routed_experts = draft_experts |
| 188 | config.dspark_num_experts_per_tok = draft_top_k |
| 189 | v41_moe_config.kernel_config.moe_backend = "auto" |
| 190 | captured = {} |
| 191 | |
| 192 | def make_experts(**kwargs): |
| 193 | captured.update(kwargs) |
| 194 | return torch.nn.Identity() |
| 195 | |
| 196 | monkeypatch.setattr( |
| 197 | "vllm.models.deepseek_v4.nvidia.model.FusedMoEFactory", make_experts |
| 198 | ) |
| 199 | moe = DeepseekV41MoE(v41_moe_config, prefix="model.layers.2.ffn") |
| 200 | assert moe.gate.weight.shape == (expected[0], 128) |
| 201 | assert (captured["num_experts"], captured["top_k"]) == expected |
| 202 | assert captured["hash_indices_table"] is None |
| 203 | |
| 204 | |
| 205 | def test_deepseek_v4_moe_preserves_configured_hash_layers(v41_moe_config): |
| 206 | config = v41_moe_config.model_config.hf_config |
| 207 | config.num_hash_layers = 1 |
| 208 | config.vocab_size = 32 |
| 209 | moe = DeepseekV4MoE( |
| 210 | v41_moe_config, |
| 211 | prefix="model.layers.0.ffn", |
| 212 | num_hash_layers=config.num_hash_layers, |
| 213 | ) |
| 214 | assert moe.gate.tid2eid.shape == (32, 2) |
| 215 | assert moe.gate.e_score_correction_bias is None |
| 216 | with pytest.raises(ValueError, match="hash MoE routing requires input_ids"): |
| 217 | moe(torch.zeros(1, 128)) |
| 218 | |
| 219 | |
| 220 | @pytest.mark.parametrize("num_tokens", [16, 17]) |
| 221 | def test_deepseek_v4_mega_gate_hash_routing_correctness( |
| 222 | v41_moe_config, monkeypatch, num_tokens |
| 223 | ): |
| 224 | if not current_platform.is_device_capability_family(100): |
| 225 | pytest.skip("DeepGEMM Mega Gate requires SM100") |
| 226 | |
| 227 | config = v41_moe_config.model_config.hf_config |
| 228 | config.hidden_size = 256 |
| 229 | config.num_hash_layers = 1 |
| 230 | config.vocab_size = 32 |
| 231 | with ( |
| 232 | set_default_torch_dtype(v41_moe_config.model_config.dtype), |
| 233 | torch.device("cuda"), |
| 234 | ): |
| 235 | moe = DeepseekV4MoE( |
| 236 | v41_moe_config, |
| 237 | prefix="model.layers.0.ffn", |
| 238 | num_hash_layers=config.num_hash_layers, |
| 239 | ) |
| 240 | hidden_states = torch.randn( |
| 241 | num_tokens, config.hidden_size, dtype=torch.bfloat16 |
| 242 | ) |
| 243 | input_ids = torch.arange(num_tokens) |
| 244 | |
| 245 | token_ids = torch.arange(config.vocab_size, device="cuda") |
| 246 | fixed_ids = torch.stack( |
| 247 | ( |
| 248 | token_ids % config.n_routed_experts, |
| 249 | (token_ids + 3) % config.n_routed_experts, |
| 250 | ), |
| 251 | dim=1, |
| 252 | ) |
| 253 | with torch.no_grad(): |
| 254 | moe.gate.weight.normal_(std=0.01) |
| 255 | moe.gate.tid2eid.copy_(fixed_ids) |
| 256 | |
| 257 | expected_ids = fixed_ids[input_ids] |
| 258 | scores = torch.nn.functional.softplus( |
| 259 | torch.mm(hidden_states, moe.gate.weight.t(), out_dtype=torch.float32) |
| 260 | ).sqrt() |
| 261 | expected_weights = scores.gather(1, expected_ids) |
| 262 | expected_weights *= config.routed_scaling_factor / expected_weights.sum( |
| 263 | dim=-1, keepdim=True |
| 264 | ) |
| 265 | |
| 266 | def check_routing(x, weights, ids, *, activation_clamp): |
| 267 | torch.testing.assert_close(ids, expected_ids) |
| 268 | torch.testing.assert_close(weights, expected_weights, rtol=1e-3, atol=1e-4) |
| 269 | return x.clone() |
| 270 | |
| 271 | monkeypatch.setattr(moe.experts, "forward", check_routing) |
| 272 | monkeypatch.setattr(moe.shared_experts, "forward", torch.zeros_like) |
| 273 | metadata = prepare_mega_gate_routing_metadata( |
| 274 | input_ids, |
| 275 | has_hash_routing=True, |
| 276 | image_sentinel_base_id=None, |
| 277 | ) |
| 278 | torch.testing.assert_close(moe(hidden_states, input_ids, metadata), hidden_states) |
| 279 | |
| 280 | |
| 281 | def test_deepseek_v4_mega_moe_expert_mapping(): |
| 282 | mapping = make_deepseek_v4_expert_params_mapping(2) |
| 283 | |
| 284 | assert mapping == [ |
| 285 | ("experts.w13_", "experts.0.w1.", 0, "w1"), |
| 286 | ("experts.w2_", "experts.0.w2.", 0, "w2"), |
| 287 | ("experts.w13_", "experts.0.w3.", 0, "w3"), |
| 288 | ("experts.w13_", "experts.1.w1.", 1, "w1"), |
| 289 | ("experts.w2_", "experts.1.w2.", 1, "w2"), |
| 290 | ("experts.w13_", "experts.1.w3.", 1, "w3"), |
| 291 | ] |
| 292 | |
| 293 | |
| 294 | def test_deepseek_v4_mega_moe_ue8m0_uint8_to_float(): |
| 295 | raw = torch.tensor([0, 126, 127, 128], dtype=torch.uint8) |
| 296 | |
| 297 | decoded = DeepseekV4MegaMoEExperts._ue8m0_uint8_to_float(raw) |
| 298 | |
| 299 | assert torch.equal(decoded.view(torch.int32), raw.to(torch.int32) << 23) |
| 300 | assert decoded[0].item() == 0.0 |
| 301 | assert decoded[1].item() == 0.5 |
| 302 | assert decoded[2].item() == 1.0 |
| 303 | assert decoded[3].item() == 2.0 |
| 304 | |
| 305 | |
| 306 | @pytest.mark.parametrize("use_kimi", [False, True]) |
| 307 | def test_deep_gemm_mega_moe_capture_precedes_eplb(monkeypatch, use_kimi): |
| 308 | experts_cls = DeepseekV4MegaMoEExperts |
| 309 | if use_kimi: |
| 310 | from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts |
| 311 | |
| 312 | experts_cls = KimiK3MegaMoEExperts |
| 313 | |
| 314 | experts = experts_cls.__new__(experts_cls) |
| 315 | torch.nn.Module.__init__(experts) |
| 316 | if use_kimi: |
| 317 | experts.synchronize_first_launch = lambda: None |
| 318 | experts.prefix = "model.layers.3.ffn.experts" |
| 319 | experts.max_num_tokens = 4 |
| 320 | experts.capture_fn = None |
| 321 | experts.get_symm_buffer = lambda: object() |
| 322 | experts.eplb_state = SimpleNamespace( |
| 323 | logical_to_physical_map=torch.empty(1), |
| 324 | expert_load_view=torch.empty(1), |
| 325 | logical_replica_count=torch.empty(1), |
| 326 | should_record_tensor=torch.empty(1), |
| 327 | num_unpadded_tokens_tensors=None, |
| 328 | ) |
| 329 | |
| 330 | topk_ids = torch.tensor([[1, 2], [3, 4]]) |
| 331 | captured: list[tuple[int, torch.Tensor]] = [] |
| 332 | bind_routed_experts_capturer( |
| 333 | SimpleNamespace(modules=lambda: [experts]), |
| 334 | SimpleNamespace(capture=lambda layer_id, ids: captured.append((layer_id, ids))), |
| 335 | ) |
| 336 | |
| 337 | class MappingReached(Exception): |
| 338 | pass |
| 339 | |
| 340 | def map_ids(**kwargs): |
| 341 | assert captured == [(3, topk_ids)] |
| 342 | raise MappingReached |
| 343 | |
| 344 | monkeypatch.setattr( |
| 345 | f"{experts_cls.__module__}.eplb_map_to_physical_and_record", |
| 346 | map_ids, |
| 347 | ) |
| 348 | monkeypatch.setattr( |
| 349 | "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace() |
| 350 | ) |
| 351 | |
| 352 | with pytest.raises(MappingReached): |
| 353 | experts( |
| 354 | torch.empty(2, 8), |
| 355 | torch.empty(2, 2), |
| 356 | topk_ids, |
| 357 | activation_clamp=None, |
| 358 | ) |
| 359 | |
| 360 | |
| 361 | def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership(): |
| 362 | vllm_config = SimpleNamespace( |
| 363 | scheduler_config=SimpleNamespace(max_num_batched_tokens=4), |
| 364 | compilation_config=SimpleNamespace(static_forward_context={}), |
| 365 | ) |
| 366 | experts = DeepseekV4MegaMoEExperts( |
| 367 | vllm_config, |
| 368 | num_experts=4, |
| 369 | num_local_experts=2, |
| 370 | experts_start_idx=2, |
| 371 | top_k=2, |
| 372 | hidden_size=128, |
| 373 | intermediate_size=128, |
| 374 | ) |
| 375 | |
| 376 | nonlocal_weight = torch.ones(128, 64, dtype=torch.uint8) |
| 377 | assert ( |
| 378 | experts.weight_loader( |
| 379 | experts.w13_weight, |
| 380 | nonlocal_weight, |
| 381 | "experts.w13_weight", |
| 382 | shard_id="w1", |
| 383 | expert_id=1, |
| 384 | return_success=True, |
| 385 | ) |
| 386 | is False |
| 387 | ) |
| 388 | |
| 389 | w1 = torch.full((128, 64), 3, dtype=torch.uint8) |
| 390 | w3 = torch.full((128, 64), 7, dtype=torch.uint8) |
| 391 | w2 = torch.full((128, 64), 11, dtype=torch.uint8) |
| 392 | |
| 393 | assert experts.weight_loader( |
| 394 | experts.w13_weight, |
| 395 | w1, |
| 396 | "experts.w13_weight", |
| 397 | shard_id="w1", |
| 398 | expert_id=2, |
| 399 | return_success=True, |
| 400 | ) |
| 401 | assert experts.weight_loader( |
| 402 | experts.w13_weight, |
| 403 | w3, |
| 404 | "experts.w13_weight", |
| 405 | shard_id="w3", |
| 406 | expert_id=2, |
| 407 | return_success=True, |
| 408 | ) |
| 409 | assert experts.weight_loader( |
| 410 | experts.w2_weight, |
| 411 | w2, |
| 412 | "experts.w2_weight", |
| 413 | shard_id="w2", |
| 414 | expert_id=2, |
| 415 | return_success=True, |
| 416 | ) |
| 417 | |
| 418 | assert torch.equal(experts.w13_weight[0, :128], w1) |
| 419 | assert torch.equal(experts.w13_weight[0, 128:], w3) |
| 420 | assert torch.equal(experts.w2_weight[0], w2) |
| 421 | assert torch.count_nonzero(experts.w13_weight[1]) == 0 |
| 422 | |
| 423 | |
| 424 | @pytest.mark.parametrize("intermediate_size", [512, 2304]) |
| 425 | def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( |
| 426 | monkeypatch, intermediate_size |
| 427 | ): |
| 428 | """Keep native widths so V4.1 routed and shared experts can fuse.""" |
| 429 | experts = DeepseekV4MegaMoEExperts( |
| 430 | SimpleNamespace( |
| 431 | scheduler_config=SimpleNamespace(max_num_batched_tokens=4), |
| 432 | compilation_config=SimpleNamespace(static_forward_context={}), |
| 433 | ), |
| 434 | num_experts=2, |
| 435 | num_local_experts=2, |
| 436 | experts_start_idx=0, |
| 437 | top_k=1, |
| 438 | hidden_size=128, |
| 439 | intermediate_size=intermediate_size, |
| 440 | ) |
| 441 | originals = [] |
| 442 | for param in ( |
| 443 | experts.w13_weight, |
| 444 | experts.w13_weight_scale, |
| 445 | experts.w2_weight, |
| 446 | experts.w2_weight_scale, |
| 447 | ): |
| 448 | param.data.random_(1, 128) |
| 449 | originals.append(param.data.clone()) |
| 450 | monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None) |
| 451 | monkeypatch.setattr( |
| 452 | "vllm.utils.deep_gemm._import_deep_gemm", |
| 453 | lambda: SimpleNamespace( |
| 454 | transform_sf_into_required_layout=lambda sf, *args: sf, |
| 455 | transform_weights_for_mega_moe=lambda l1, l2: (l1, l2), |
| 456 | ), |
| 457 | ) |
| 458 | |
| 459 | experts.finalize_weights() |
| 460 | assert experts.intermediate_size == intermediate_size |
| 461 | transformed = (*experts._transformed_l1_weights, *experts._transformed_l2_weights) |
| 462 | for actual, original in zip(transformed, originals): |
| 463 | expected = ( |
| 464 | original |
| 465 | if actual.dtype == torch.int8 |
| 466 | else experts._ue8m0_uint8_to_float(original) |
| 467 | ) |
| 468 | assert torch.equal(actual.view(expected.dtype), expected) |
| 469 | |
| 470 | transformed_l1 = experts._transformed_l1_weights |
| 471 | experts.finalize_weights() |
| 472 | assert experts._transformed_l1_weights is transformed_l1 |
| 473 | |
| 474 | |
| 475 | @pytest.mark.parametrize( |
| 476 | "hidden_size,intermediate_size,block_size,mxfp8", |
| 477 | [ |
| 478 | pytest.param(128, 512, 128, False, id="aligned-block128"), |
| 479 | pytest.param(128, 512, 128, True, id="aligned-block128-mxfp8"), |
| 480 | pytest.param(5120, 2304, 32, False, id="deepseek-v41-flash"), |
| 481 | pytest.param(128, 2304, 128, False, id="native-block128"), |
| 482 | pytest.param(5120, 2304, 32, True, id="deepseek-v41-flash-mxfp8"), |
| 483 | ], |
| 484 | ) |
| 485 | def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( |
| 486 | monkeypatch, hidden_size, intermediate_size, block_size, mxfp8 |
| 487 | ): |
| 488 | """Shared fusion preserves checkpoint weights and scales at native widths.""" |
| 489 | |
| 490 | class FakeDeepGemm: |
| 491 | transformed_dims: list[tuple[int, int]] = [] |
| 492 | scale_inputs: list[tuple[int, ...]] = [] |
| 493 | scale_values: list[torch.Tensor] = [] |
| 494 | |
| 495 | @staticmethod |
| 496 | def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs): |
| 497 | return None |
| 498 | |
| 499 | @staticmethod |
| 500 | def get_block_m_for_mega_moe(*args, **kwargs): |
| 501 | return 128 |
| 502 | |
| 503 | @staticmethod |
| 504 | def fp8_fp4_mega_moe( |
| 505 | y, |
| 506 | l1_weights, |
| 507 | l2_weights, |
| 508 | sym_buffer, |
| 509 | shared_l1_weights=None, |
| 510 | shared_l2_weights=None, |
| 511 | **kwargs, |
| 512 | ): |
| 513 | return None |
| 514 | |
| 515 | @classmethod |
| 516 | def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs): |
| 517 | cls.scale_inputs.append(tuple(sf.shape)) |
| 518 | cls.scale_values.append(sf) |
| 519 | return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32) |
| 520 | |
| 521 | @classmethod |
| 522 | def transform_weights_for_mega_moe(cls, l1_weights, l2_weights): |
| 523 | cls.transformed_dims.append((l1_weights[0].dim(), l2_weights[0].dim())) |
| 524 | if l1_weights[0].dim() == 2: |
| 525 | return (l1_weights[0].clone(), l1_weights[1]), l2_weights |
| 526 | return l1_weights, l2_weights |
| 527 | |
| 528 | vllm_config = SimpleNamespace( |
| 529 | scheduler_config=SimpleNamespace(max_num_batched_tokens=4), |
| 530 | compilation_config=SimpleNamespace(static_forward_context={}), |
| 531 | ) |
| 532 | experts = DeepseekV4MegaMoEExperts( |
| 533 | vllm_config, |
| 534 | num_experts=2, |
| 535 | num_local_experts=1, |
| 536 | experts_start_idx=0, |
| 537 | top_k=1, |
| 538 | hidden_size=hidden_size, |
| 539 | intermediate_size=intermediate_size, |
| 540 | num_shared_experts=1, |
| 541 | ) |
| 542 | experts._check_runtime_supported = lambda: None |
| 543 | |
| 544 | def fp8_parameter(*shape): |
| 545 | return torch.nn.Parameter( |
| 546 | torch.empty(*shape, dtype=torch.uint8) |
| 547 | .random_(1, 119) |
| 548 | .view(torch.float8_e4m3fn), |
| 549 | requires_grad=False, |
| 550 | ) |
| 551 | |
| 552 | def scale_parameter(*shape, dtype=torch.int32): |
| 553 | return torch.nn.Parameter(torch.ones(*shape, dtype=dtype), requires_grad=False) |
| 554 | |
| 555 | shared_experts = SimpleNamespace( |
| 556 | gate_up_proj=SimpleNamespace( |
| 557 | weight=fp8_parameter(2 * intermediate_size, hidden_size), |
| 558 | weight_block_size=(block_size, block_size), |
| 559 | weight_scale_inv=scale_parameter( |
| 560 | 2 * intermediate_size // block_size, |
| 561 | hidden_size // block_size, |
| 562 | dtype=torch.float8_e8m0fnu, |
| 563 | ), |
| 564 | ), |
| 565 | down_proj=SimpleNamespace( |
| 566 | weight=fp8_parameter(hidden_size, intermediate_size), |
| 567 | weight_block_size=(block_size, block_size), |
| 568 | weight_scale_inv=scale_parameter( |
| 569 | hidden_size // block_size, |
| 570 | intermediate_size // block_size, |
| 571 | dtype=torch.float8_e8m0fnu, |
| 572 | ), |
| 573 | ), |
| 574 | ) |
| 575 | if mxfp8: |
| 576 | for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): |
| 577 | linear.weight_block_size = (1, 32) |
| 578 | linear.weight_scale = torch.nn.Parameter( |
| 579 | linear.weight_scale_inv.view(torch.uint8) |
| 580 | .repeat_interleave(block_size, dim=0) |
| 581 | .repeat_interleave(block_size // 32, dim=1), |
| 582 | requires_grad=False, |
| 583 | ) |
| 584 | del linear.weight_scale_inv |
| 585 | originals = [] |
| 586 | for linear in (shared_experts.gate_up_proj, shared_experts.down_proj): |
| 587 | scale = linear.weight_scale if mxfp8 else linear.weight_scale_inv |
| 588 | scale.data.view(torch.uint8).random_(124, 130) |
| 589 | block_m, block_k = linear.weight_block_size |
| 590 | scale_1x32 = ( |
| 591 | experts._ue8m0_uint8_to_float(scale.view(torch.uint8)) |
| 592 | .repeat_interleave(block_m, dim=0) |
| 593 | .repeat_interleave(block_k // 32, dim=1) |
| 594 | ) |
| 595 | originals.append((linear.weight.view(torch.uint8).clone(), scale_1x32)) |
| 596 | monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm) |
| 597 | |
| 598 | original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr() |
| 599 | experts.finalize_weights(shared_experts) |
| 600 | |
| 601 | assert experts.has_fused_shared_experts |
| 602 | assert experts.intermediate_size == intermediate_size |
| 603 | assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)] |
| 604 | assert FakeDeepGemm.scale_inputs[-2:] == [ |
| 605 | (1, 2 * intermediate_size, hidden_size // 32), |
| 606 | (1, hidden_size, intermediate_size // 32), |
| 607 | ] |
| 608 | for transformed, actual_scale, (weight, scale) in zip( |
| 609 | ( |
| 610 | experts._transformed_shared_l1_weights, |
| 611 | experts._transformed_shared_l2_weights, |
| 612 | ), |
| 613 | FakeDeepGemm.scale_values[-2:], |
| 614 | originals, |
| 615 | ): |
| 616 | assert torch.equal(transformed[0].view(torch.uint8), weight) |
| 617 | assert torch.equal(actual_scale[0], scale) |
| 618 | assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr |
| 619 | assert ( |
| 620 | experts._transformed_shared_l1_weights[0].data_ptr() |
| 621 | == shared_experts.gate_up_proj.weight.data_ptr() |
| 622 | ) |
| 623 | assert ( |
| 624 | experts._transformed_shared_l2_weights[0].data_ptr() |
| 625 | == shared_experts.down_proj.weight.data_ptr() |
| 626 | ) |
| 627 | |
| 628 | transformed_l1 = experts._transformed_shared_l1_weights |
| 629 | experts.finalize_weights(shared_experts) |
| 630 | assert experts._transformed_shared_l1_weights is transformed_l1 |
| 631 | |
| 632 | |
| 633 | @pytest.mark.parametrize("fused", [False, True]) |
| 634 | def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( |
| 635 | monkeypatch, fused |
| 636 | ): |
| 637 | class FakeGate(torch.nn.Module): |
| 638 | weight = torch.empty(2, 128) |
| 639 | tid2eid = None |
| 640 | e_score_correction_bias = None |
| 641 | |
| 642 | def forward(self, hidden_states): |
| 643 | return torch.empty(hidden_states.shape[0], 2), None |
| 644 | |
| 645 | class FakeExperts(torch.nn.Module): |
| 646 | has_fused_shared_experts = fused |
| 647 | |
| 648 | def forward(self, hidden_states, *args, **kwargs): |
| 649 | return torch.ones_like(hidden_states) |
| 650 | |
| 651 | class FakeSharedExperts(torch.nn.Module): |
| 652 | calls = 0 |
| 653 | |
| 654 | def forward(self, hidden_states): |
| 655 | self.calls += 1 |
| 656 | return torch.full_like(hidden_states, 2) |
| 657 | |
| 658 | moe = DeepseekV4MoE.__new__(DeepseekV4MoE) |
| 659 | torch.nn.Module.__init__(moe) |
| 660 | moe.use_mega_moe = True |
| 661 | moe.gate = FakeGate() |
| 662 | moe.experts = FakeExperts() |
| 663 | moe.shared_experts = FakeSharedExperts() |
| 664 | moe.scoring_func = "sqrtsoftplus" |
| 665 | moe.n_activated_experts = 1 |
| 666 | moe.renormalize = True |
| 667 | moe.hash_indices_dtype = torch.int64 |
| 668 | moe.routed_scaling_factor = 1.0 |
| 669 | moe.swiglu_limit = 10.0 |
| 670 | monkeypatch.setattr( |
| 671 | "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias", |
| 672 | lambda **kwargs: ( |
| 673 | torch.ones(kwargs["hidden_states"].shape[0], 1), |
| 674 | torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64), |
| 675 | ), |
| 676 | ) |
| 677 | |
| 678 | output = moe(torch.zeros(2, 128)) |
| 679 | |
| 680 | expected = 1 if fused else 3 |
| 681 | assert torch.all(output == expected) |
| 682 | assert moe.shared_experts.calls == (0 if fused else 1) |
| 683 | |
| 684 | |
| 685 | @pytest.mark.skipif( |
| 686 | not torch.cuda.is_available(), |
| 687 | reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", |
| 688 | ) |
| 689 | @pytest.mark.parametrize("num_tokens", [7, 63, 64, 65, 127, 128, 135, 434, 16384]) |
| 690 | @pytest.mark.parametrize("hidden_size", [256, 5120]) |
| 691 | def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact( |
| 692 | num_tokens, hidden_size |
| 693 | ): |
| 694 | from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 |
| 695 | |
| 696 | device = torch.device("cuda") |
| 697 | top_k = 8 |
| 698 | |
| 699 | generator = torch.Generator(device=device) |
| 700 | generator.manual_seed(0) |
| 701 | hidden_states = ( |
| 702 | torch.randn( |
| 703 | num_tokens, |
| 704 | hidden_size, |
| 705 | device=device, |
| 706 | dtype=torch.float32, |
| 707 | generator=generator, |
| 708 | ) |
| 709 | * 17.0 |
| 710 | ).to(torch.bfloat16) |
| 711 | hidden_states[0, :32] = 0 |
| 712 | hidden_states[1, 32:64] = 1.0e-6 |
| 713 | hidden_states[2, 64:96] = -1.0e-6 |
| 714 | |
| 715 | topk_ids = torch.randint( |
| 716 | 0, |
| 717 | 256, |
| 718 | (num_tokens, top_k), |
| 719 | device=device, |
| 720 | dtype=torch.int32, |
| 721 | generator=generator, |
| 722 | ) |
| 723 | topk_weights = torch.randn( |
| 724 | num_tokens, |
| 725 | top_k, |
| 726 | device=device, |
| 727 | dtype=torch.float32, |
| 728 | generator=generator, |
| 729 | ) |
| 730 | |
| 731 | ref_x, ref_x_sf = per_token_cast_to_fp8( |
| 732 | hidden_states, |
| 733 | use_ue8m0=True, |
| 734 | gran_k=32, |
| 735 | use_packed_ue8m0=True, |
| 736 | ) |
| 737 | ref_topk_idx = topk_ids.to(torch.int64) |
| 738 | ref_topk_weights = topk_weights.clone() |
| 739 | |
| 740 | fused_x = torch.empty_like(ref_x) |
| 741 | fused_x_sf = torch.empty_like(ref_x_sf) |
| 742 | fused_topk_idx = torch.empty_like(ref_topk_idx) |
| 743 | fused_topk_weights = torch.empty_like(ref_topk_weights) |
| 744 | |
| 745 | prepare_megamoe_inputs( |
| 746 | hidden_states, |
| 747 | topk_weights, |
| 748 | topk_ids, |
| 749 | fused_x, |
| 750 | fused_x_sf, |
| 751 | fused_topk_idx, |
| 752 | fused_topk_weights, |
| 753 | ) |
| 754 | torch.accelerator.synchronize() |
| 755 | |
| 756 | assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) |
| 757 | assert torch.equal(fused_x_sf, ref_x_sf) |
| 758 | assert torch.equal(fused_topk_idx, ref_topk_idx) |
| 759 | assert torch.equal( |
| 760 | fused_topk_weights.view(torch.uint8), |
| 761 | ref_topk_weights.view(torch.uint8), |
| 762 | ) |
| 763 | |
| 764 | |
| 765 | @pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192]) |
| 766 | def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m): |
| 767 | from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 |
| 768 | |
| 769 | device = torch.device("cuda") |
| 770 | num_tokens = shared_block_m + 7 |
| 771 | hidden_size = 256 |
| 772 | top_k = 8 |
| 773 | generator = torch.Generator(device=device) |
| 774 | generator.manual_seed(shared_block_m) |
| 775 | hidden_states = torch.randn( |
| 776 | num_tokens, |
| 777 | hidden_size, |
| 778 | device=device, |
| 779 | dtype=torch.bfloat16, |
| 780 | generator=generator, |
| 781 | ) |
| 782 | topk_ids = torch.randint( |
| 783 | 0, |
| 784 | 256, |
| 785 | (num_tokens, top_k), |
| 786 | device=device, |
| 787 | dtype=torch.int32, |
| 788 | generator=generator, |
| 789 | ) |
| 790 | topk_weights = torch.randn( |
| 791 | num_tokens, |
| 792 | top_k, |
| 793 | device=device, |
| 794 | dtype=torch.float32, |
| 795 | generator=generator, |
| 796 | ) |
| 797 | |
| 798 | ref_x, ref_x_sf = per_token_cast_to_fp8( |
| 799 | hidden_states, |
| 800 | use_ue8m0=True, |
| 801 | gran_k=32, |
| 802 | use_packed_ue8m0=True, |
| 803 | ) |
| 804 | aligned_block_m = ((shared_block_m + 127) // 128) * 128 |
| 805 | num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * ( |
| 806 | aligned_block_m |
| 807 | ) |
| 808 | ref_shared_x_sf = torch.zeros( |
| 809 | num_shared_rows, |
| 810 | hidden_size // 128, |
| 811 | dtype=torch.int32, |
| 812 | device=device, |
| 813 | ) |
| 814 | for token_id in range(num_tokens): |
| 815 | m_in_block = token_id % shared_block_m |
| 816 | transposed_m = ( |
| 817 | (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32 |
| 818 | ) |
| 819 | shared_row = token_id // shared_block_m * aligned_block_m + transposed_m |
| 820 | ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id]) |
| 821 | |
| 822 | fused_x = torch.empty_like(ref_x) |
| 823 | fused_x_sf = torch.empty_like(ref_x_sf) |
| 824 | fused_shared_storage = torch.full( |
| 825 | (hidden_size // 128, num_shared_rows), |
| 826 | -1, |
| 827 | dtype=torch.int32, |
| 828 | device=device, |
| 829 | ) |
| 830 | fused_shared_x_sf = fused_shared_storage.t() |
| 831 | fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64) |
| 832 | fused_topk_weights = torch.empty_like(topk_weights) |
| 833 | |
| 834 | prepare_megamoe_inputs( |
| 835 | hidden_states, |
| 836 | topk_weights, |
| 837 | topk_ids, |
| 838 | fused_x, |
| 839 | fused_x_sf, |
| 840 | fused_topk_idx, |
| 841 | fused_topk_weights, |
| 842 | shared_x_sf=fused_shared_x_sf, |
| 843 | shared_block_m=shared_block_m, |
| 844 | ) |
| 845 | torch.accelerator.synchronize() |
| 846 | |
| 847 | populated = ref_shared_x_sf != 0 |
| 848 | assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) |
| 849 | assert torch.equal(fused_x_sf, ref_x_sf) |
| 850 | assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated]) |
| 851 | assert torch.all(fused_shared_x_sf[~populated] == -1) |
| 852 | |
| 853 | |
| 854 | def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast(): |
| 855 | """The loader invokes the model-level PWAL hook for every load format, |
| 856 | so it must finalize megamoe + mhc broadcast weights to cover dummy |
| 857 | load, which skips load_weights().""" |
| 858 | calls = [] |
| 859 | stub = SimpleNamespace( |
| 860 | model=SimpleNamespace( |
| 861 | finalize_mega_moe_weights=lambda: calls.append("mega_moe"), |
| 862 | finalize_mhc_broadcast_weights=lambda: calls.append("mhc"), |
| 863 | ) |
| 864 | ) |
| 865 | |
| 866 | DeepseekV4ForCausalLM.process_weights_after_loading(stub) |
| 867 | |
| 868 | assert calls == ["mega_moe", "mhc"] |
| 869 | |
| 870 | |
| 871 | def test_deepseek_v4_drafter_pwal_hooks_finalize_mega_moe(): |
| 872 | """MTP/DSpark drafters load as their own top-level models, so each needs |
| 873 | its own PWAL hook now that the megamoe forward no longer finalizes |
| 874 | weights lazily on first use.""" |
| 875 | calls = [] |
| 876 | mtp = SimpleNamespace(finalize_mega_moe_weights=lambda: calls.append("mtp")) |
| 877 | DeepSeekV4MTP.process_weights_after_loading(mtp) |
| 878 | |
| 879 | dspark = SimpleNamespace(_finalize_moe=lambda: calls.append("dspark")) |
| 880 | DSparkDeepseekV4ForCausalLM.process_weights_after_loading(dspark) |
| 881 | |
| 882 | assert calls == ["mtp", "dspark"] |
| 883 | |
| 884 | |
| 885 | @pytest.mark.skipif( |
| 886 | not torch.cuda.is_available(), |
| 887 | reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.", |
| 888 | ) |
| 889 | @pytest.mark.parametrize("nonfinite_padding", [False, True]) |
| 890 | @pytest.mark.parametrize("num_tokens", [7, 65, 135]) |
| 891 | def test_deepseek_v4_mega_moe_fused_input_staging_masks_padding( |
| 892 | nonfinite_padding, num_tokens |
| 893 | ): |
| 894 | from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8 |
| 895 | |
| 896 | device = torch.device("cuda") |
| 897 | hidden_size = 256 |
| 898 | top_k = 8 |
| 899 | |
| 900 | generator = torch.Generator(device=device) |
| 901 | generator.manual_seed(1) |
| 902 | hidden_states = torch.randn( |
| 903 | num_tokens, |
| 904 | hidden_size, |
| 905 | device=device, |
| 906 | dtype=torch.bfloat16, |
| 907 | generator=generator, |
| 908 | ) |
| 909 | topk_ids = torch.randint( |
| 910 | 0, |
| 911 | 256, |
| 912 | (num_tokens, top_k), |
| 913 | device=device, |
| 914 | dtype=torch.int32, |
| 915 | generator=generator, |
| 916 | ) |
| 917 | topk_weights = torch.randn( |
| 918 | num_tokens, |
| 919 | top_k, |
| 920 | device=device, |
| 921 | dtype=torch.float32, |
| 922 | generator=generator, |
| 923 | ) |
| 924 | is_padding = torch.zeros(num_tokens, device=device, dtype=torch.bool) |
| 925 | is_padding[1::3] = True |
| 926 | is_padding[-1] = True |
| 927 | if nonfinite_padding: |
| 928 | topk_weights[is_padding] = float("nan") |
| 929 | |
| 930 | ref_x, ref_x_sf = per_token_cast_to_fp8( |
| 931 | hidden_states, |
| 932 | use_ue8m0=True, |
| 933 | gran_k=32, |
| 934 | use_packed_ue8m0=True, |
| 935 | ) |
| 936 | ref_topk_idx = topk_ids.to(torch.int64) |
| 937 | ref_topk_idx[is_padding] = -1 |
| 938 | ref_topk_weights = topk_weights.clone() |
| 939 | ref_topk_weights[is_padding] = 0.0 |
| 940 | |
| 941 | fused_x = torch.empty_like(ref_x) |
| 942 | fused_x_sf = torch.empty_like(ref_x_sf) |
| 943 | fused_topk_idx = torch.empty_like(ref_topk_idx) |
| 944 | fused_topk_weights = torch.empty_like(ref_topk_weights) |
| 945 | |
| 946 | prepare_megamoe_inputs( |
| 947 | hidden_states, |
| 948 | topk_weights, |
| 949 | topk_ids, |
| 950 | fused_x, |
| 951 | fused_x_sf, |
| 952 | fused_topk_idx, |
| 953 | fused_topk_weights, |
| 954 | is_padding=is_padding, |
| 955 | ) |
| 956 | torch.accelerator.synchronize() |
| 957 | |
| 958 | assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8)) |
| 959 | assert torch.equal(fused_x_sf, ref_x_sf) |
| 960 | assert torch.equal(fused_topk_idx, ref_topk_idx) |
| 961 | assert torch.equal( |
| 962 | fused_topk_weights.view(torch.uint8), |
| 963 | ref_topk_weights.view(torch.uint8), |
| 964 | ) |