@tduncan / Nvidia-Optimization-Stack-VLLM / commits / de69e821b7

[ROCm][Perf][DeepSeek V4] Fuse native FP8 shared expert with MXFP4 routed experts (#53161)

Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Cursor Agent <agent@cursor.com>

Fangzhou Ai committed Sep 6, 2026 at 22:42 UTC de69e821b7c834bcf1ec96325ec8c410ac183680
4 files changed +839 -6
tests/model_executor/layers/test_fused_shared_expert.py
+340 -1
@@ -2,8 +2,10 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import importlib
5 +import inspect
6 +import sys
7 from copy import deepcopy
6 -from types import SimpleNamespace
8 +from types import ModuleType, SimpleNamespace
9 from typing import Any, cast
10 from unittest.mock import patch
11
@@ -414,6 +416,343 @@ def test_deepseek_v4_shared_expert_fse_uses_mtp_quantization_config_prefix(
416 assert reason is None
417
418
419 +def test_deepseek_v4_heterogeneous_fhmoe_keeps_native_intermediate_width() -> None:
420 + from vllm.models.deepseek_v4.amd.model import _prepare_native_fp8_shared_expert
421 +
422 + hidden_size = 7168
423 + intermediate_size = 384
424 + w13 = torch.empty((2 * intermediate_size, hidden_size), dtype=torch.float8_e4m3fn)
425 + w2 = torch.empty((hidden_size, intermediate_size), dtype=torch.float8_e4m3fn)
426 + w13_scale_bytes = (torch.arange(6 * 56).reshape(6, 56).remainder(20) + 0x60).to(
427 + torch.uint8
428 + )
429 + w2_scale_bytes = (torch.arange(56 * 3).reshape(56, 3).remainder(30) + 0x50).to(
430 + torch.uint8
431 + )
432 + w13_scale = w13_scale_bytes.view(torch.float8_e8m0fnu)
433 + w2_scale = w2_scale_bytes.view(torch.float8_e8m0fnu)
434 +
435 + prepared = _prepare_native_fp8_shared_expert(
436 + w13, w2, w13_scale, w2_scale, intermediate_size
437 + )
438 +
439 + assert prepared[0].shape == (1, 768, 7168)
440 + assert prepared[1].shape == (1, 7168, 384)
441 + assert prepared[2].shape == (768, 224)
442 + assert prepared[3].shape == (7168, 16)
443 + expected_w13_scale = w13_scale_bytes.repeat_interleave(
444 + 128, dim=0
445 + ).repeat_interleave(4, dim=1)
446 + expected_w2_scale = w2_scale_bytes.repeat_interleave(128, dim=0).repeat_interleave(
447 + 4, dim=1
448 + )
449 + assert torch.equal(prepared[2].view(torch.uint8), expected_w13_scale)
450 + assert torch.equal(prepared[3].view(torch.uint8)[:, :12], expected_w2_scale)
451 + assert torch.all(prepared[3].view(torch.uint8)[:, 12:] == 0x7F)
452 +
453 +
454 +@pytest.mark.parametrize(
455 + ("num_tokens", "supported_through", "expected"),
456 + [
457 + (0, 4096, False),
458 + (1, 4096, True),
459 + (1536, 4096, True),
460 + (2048, 4096, True),
461 + (2049, 4096, True),
462 + (4096, 4096, True),
463 + (4097, 4096, False),
464 + (1536, 0, False),
465 + ],
466 +)
467 +def test_deepseek_v4_heterogeneous_fhmoe_token_policy(
468 + monkeypatch: pytest.MonkeyPatch,
469 + num_tokens: int,
470 + supported_through: int,
471 + expected: bool,
472 +) -> None:
473 + from vllm.models.deepseek_v4.amd import model as deepseek_v4_model
474 +
475 + checked_tokens: list[int] = []
476 +
477 + def supports(num_tokens: int) -> bool:
478 + checked_tokens.append(num_tokens)
479 + return num_tokens <= supported_through
480 +
481 + monkeypatch.setattr(
482 + deepseek_v4_model.rocm_aiter_ops,
483 + "fused_moe_supports_heterogeneous_shared_expert",
484 + supports,
485 + )
486 +
487 + assert deepseek_v4_model._use_heterogeneous_fhmoe(num_tokens) is expected
488 + assert checked_tokens == ([] if num_tokens == 0 else [num_tokens])
489 +
490 +
491 +def _supported_fhmoe_signature(
492 + shared_w1=None,
493 + shared_w2=None,
494 + shared_w1_scale=None,
495 + shared_w2_scale=None,
496 + shared_expert_id=-1,
497 +) -> None:
498 + pass
499 +
500 +
501 +def _incomplete_fhmoe_signature(
502 + shared_w1=None,
503 + shared_w2=None,
504 + shared_w1_scale=None,
505 + shared_w2_scale=None,
506 +) -> None:
507 + pass
508 +
509 +
510 +def _install_fake_aiter_fhmoe(
511 + monkeypatch: pytest.MonkeyPatch,
512 + supports_dsv4_i384_fhmoe: object,
513 + fused_moe: object,
514 +) -> None:
515 + fake_aiter = ModuleType("aiter")
516 + fake_aiter.__path__ = []
517 + fake_fhmoe = ModuleType("aiter.fhmoe")
518 + if supports_dsv4_i384_fhmoe is not None:
519 + fake_fhmoe.__dict__["supports_dsv4_i384_fhmoe"] = supports_dsv4_i384_fhmoe
520 + fake_fused_moe = ModuleType("aiter.fused_moe")
521 + fake_fused_moe.__dict__["fused_moe"] = fused_moe
522 + fake_aiter.__dict__["fhmoe"] = fake_fhmoe
523 + fake_aiter.__dict__["fused_moe"] = fake_fused_moe
524 + monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
525 + monkeypatch.setitem(sys.modules, "aiter.fhmoe", fake_fhmoe)
526 + monkeypatch.setitem(sys.modules, "aiter.fused_moe", fake_fused_moe)
527 +
528 +
529 +def _supports_fhmoe_through_2047(max_tokens: int) -> bool:
530 + return max_tokens <= 2047
531 +
532 +
533 +def _supports_fhmoe_through_2048(max_tokens: int) -> bool:
534 + return max_tokens <= 2048
535 +
536 +
537 +def _supports_fhmoe_through_4096(max_tokens: int) -> bool:
538 + return max_tokens <= 4096
539 +
540 +
541 +def _raises_fhmoe_config_error(max_tokens: int) -> bool:
542 + raise OSError
543 +
544 +
545 +def _returns_truthy_non_bool(max_tokens: int) -> int:
546 + return 1
547 +
548 +
549 +@pytest.mark.parametrize(
550 + ("num_tokens", "capability", "fused_moe", "expected"),
551 + [
552 + (0, _supports_fhmoe_through_2048, _supported_fhmoe_signature, False),
553 + (True, _supports_fhmoe_through_2048, _supported_fhmoe_signature, False),
554 + (2048, None, _supported_fhmoe_signature, False),
555 + (2048, True, _supported_fhmoe_signature, False),
556 + (2048, _raises_fhmoe_config_error, _supported_fhmoe_signature, False),
557 + (2048, _returns_truthy_non_bool, _supported_fhmoe_signature, False),
558 + (2048, _supports_fhmoe_through_2047, _supported_fhmoe_signature, False),
559 + (2048, _supports_fhmoe_through_2048, _incomplete_fhmoe_signature, False),
560 + (2048, _supports_fhmoe_through_2048, _supported_fhmoe_signature, True),
561 + (2049, _supports_fhmoe_through_2048, _supported_fhmoe_signature, False),
562 + (4096, _supports_fhmoe_through_4096, _supported_fhmoe_signature, True),
563 + (4097, _supports_fhmoe_through_4096, _supported_fhmoe_signature, False),
564 + ],
565 +)
566 +def test_deepseek_v4_heterogeneous_fhmoe_aiter_capability(
567 + monkeypatch: pytest.MonkeyPatch,
568 + num_tokens: object,
569 + capability: object,
570 + fused_moe: object,
571 + expected: bool,
572 +) -> None:
573 + from vllm._aiter_ops import rocm_aiter_ops
574 +
575 + _install_fake_aiter_fhmoe(monkeypatch, capability, fused_moe)
576 +
577 + assert (
578 + rocm_aiter_ops._probe_dsv4_i384_fhmoe_capability(cast(int, num_tokens))
579 + is expected
580 + )
581 +
582 +
583 +def test_deepseek_v4_heterogeneous_fhmoe_aiter_capability_catches_import_error(
584 + monkeypatch: pytest.MonkeyPatch,
585 +) -> None:
586 + from vllm._aiter_ops import rocm_aiter_ops
587 +
588 + fake_aiter = ModuleType("aiter")
589 + fake_aiter.__path__ = []
590 + monkeypatch.setitem(sys.modules, "aiter", fake_aiter)
591 + monkeypatch.setitem(sys.modules, "aiter.fhmoe", None)
592 +
593 + assert not rocm_aiter_ops._probe_dsv4_i384_fhmoe_capability(2048)
594 +
595 +
596 +@pytest.mark.parametrize("error", [TypeError, ValueError])
597 +def test_deepseek_v4_heterogeneous_fhmoe_aiter_capability_catches_signature_error(
598 + monkeypatch: pytest.MonkeyPatch,
599 + error: type[Exception],
600 +) -> None:
601 + from vllm._aiter_ops import rocm_aiter_ops
602 +
603 + _install_fake_aiter_fhmoe(
604 + monkeypatch, _supports_fhmoe_through_2048, _supported_fhmoe_signature
605 + )
606 +
607 + def raise_signature_error(_: object) -> inspect.Signature:
608 + raise error
609 +
610 + monkeypatch.setattr(inspect, "signature", raise_signature_error)
611 +
612 + assert not rocm_aiter_ops._probe_dsv4_i384_fhmoe_capability(2048)
613 +
614 +
615 +@pytest.mark.parametrize(
616 + ("setting", "value", "expected"),
617 + [
618 + ("data_parallel_size", 1, True),
619 + ("data_parallel_size", 2, False),
620 + ("prefill_context_parallel_size", 2, False),
621 + ("topk_method", "greedy", False),
622 + ("fhmoe_supported", False, False),
623 + ],
624 +)
625 +def test_deepseek_v4_heterogeneous_fhmoe_compatibility_gates(
626 + monkeypatch: pytest.MonkeyPatch,
627 + setting: str,
628 + value: object,
629 + expected: bool,
630 +) -> None:
631 + from vllm.models.deepseek_v4.amd import model as deepseek_v4_model
632 +
633 + hf_config = SimpleNamespace(
634 + n_routed_experts=384,
635 + num_experts_per_tok=6,
636 + n_shared_experts=1,
637 + hidden_size=7168,
638 + moe_intermediate_size=3072,
639 + hidden_act="silu",
640 + expert_dtype="fp4",
641 + topk_method="noaux_tc",
642 + )
643 + parallel_config = SimpleNamespace(
644 + enable_expert_parallel=False,
645 + enable_eplb=False,
646 + tensor_parallel_size=8,
647 + data_parallel_size=1,
648 + prefill_context_parallel_size=1,
649 + )
650 + if setting == "topk_method":
651 + hf_config.topk_method = value
652 + elif setting != "fhmoe_supported":
653 + setattr(parallel_config, setting, value)
654 +
655 + quant_config = SimpleNamespace(
656 + get_name=lambda: "deepseek_v4_fp8",
657 + moe_quant_algo="",
658 + weight_block_size=[128, 128],
659 + is_checkpoint_fp8_serialized=True,
660 + is_scale_e8m0=True,
661 + ignored_layers=None,
662 + )
663 + vllm_config = SimpleNamespace(
664 + model_config=SimpleNamespace(hf_config=hf_config, dtype=torch.bfloat16),
665 + quant_config=quant_config,
666 + parallel_config=parallel_config,
667 + kernel_config=SimpleNamespace(moe_backend="aiter"),
668 + offload_config=None,
669 + )
670 + monkeypatch.setattr(deepseek_v4_model.current_platform, "is_rocm", lambda: True)
671 + monkeypatch.setattr(
672 + deepseek_v4_model.envs,
673 + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS",
674 + True,
675 + )
676 + monkeypatch.setattr(deepseek_v4_model, "on_gfx950", lambda: True)
677 + monkeypatch.setattr(
678 + deepseek_v4_model.rocm_aiter_ops,
679 + "is_fusion_moe_shared_experts_enabled",
680 + lambda: True,
681 + )
682 + checked_tokens: list[int] = []
683 + fhmoe_supported = setting != "fhmoe_supported" or value is True
684 +
685 + def supports(num_tokens: int) -> bool:
686 + checked_tokens.append(num_tokens)
687 + return fhmoe_supported
688 +
689 + monkeypatch.setattr(
690 + deepseek_v4_model.rocm_aiter_ops,
691 + "fused_moe_supports_heterogeneous_shared_expert",
692 + supports,
693 + )
694 +
695 + assert (
696 + deepseek_v4_model._heterogeneous_shared_expert_enabled(
697 + cast(VllmConfig, vllm_config)
698 + )
699 + is expected
700 + )
701 + assert checked_tokens == [1]
702 +
703 +
704 +@pytest.mark.parametrize(
705 + ("weights_shape", "ids_shape", "message"),
706 + [
707 + ((2, 7, 1), (2, 7, 1), "equal two-dimensional shapes"),
708 + ((2, 7), (3, 7), "equal two-dimensional shapes"),
709 + ((2, 6), (2, 6), "exactly 7 columns"),
710 + ((2, 8), (2, 8), "exactly 7 columns"),
711 + ],
712 +)
713 +def test_deepseek_v4_heterogeneous_fhmoe_rejects_invalid_routes(
714 + weights_shape: tuple[int, ...],
715 + ids_shape: tuple[int, ...],
716 + message: str,
717 +) -> None:
718 + from vllm.models.deepseek_v4.amd.model import _validate_heterogeneous_routes
719 +
720 + with pytest.raises(ValueError, match=message):
721 + _validate_heterogeneous_routes(
722 + torch.empty(weights_shape),
723 + torch.empty(ids_shape, dtype=torch.int64),
724 + experts_per_token=6,
725 + )
726 +
727 +
728 +def test_deepseek_v4_heterogeneous_fhmoe_accepts_appended_shared_route() -> None:
729 + from vllm.models.deepseek_v4.amd.model import _validate_heterogeneous_routes
730 +
731 + _validate_heterogeneous_routes(
732 + torch.empty((2, 7)),
733 + torch.empty((2, 7), dtype=torch.int64),
734 + experts_per_token=6,
735 + )
736 +
737 +
738 +def test_aiter_fused_moe_validates_shared_expert_arguments() -> None:
739 + from vllm._aiter_ops import (
740 + _validate_rocm_aiter_fused_moe_shared_expert_args,
741 + )
742 +
743 + validate = _validate_rocm_aiter_fused_moe_shared_expert_args
744 + shared_tensor = torch.empty(0)
745 +
746 + assert not validate(None, None, None, None, -1)
747 + with pytest.raises(ValueError, match="requires shared weights and scales"):
748 + validate(None, None, None, None, 0)
749 + with pytest.raises(ValueError, match="both shared weights and scales"):
750 + validate(shared_tensor, None, None, None, 0)
751 + with pytest.raises(ValueError, match="non-negative shared expert ID"):
752 + validate(shared_tensor, shared_tensor, shared_tensor, shared_tensor, -1)
753 + assert validate(shared_tensor, shared_tensor, shared_tensor, shared_tensor, 0)
754 +
755 +
756 def test_is_model_fused_shared_expert_compatible() -> None:
757 class MoE(nn.Module):
758 def __init__(self, enabled: bool) -> None:
vllm/_aiter_ops.py
+107
@@ -211,6 +211,31 @@ def if_aiter_supported(func: Callable) -> Callable:
211 return wrapper
212
213
214 +def _validate_rocm_aiter_fused_moe_shared_expert_args(
215 + shared_w1: torch.Tensor | None,
216 + shared_w2: torch.Tensor | None,
217 + shared_w1_scale: torch.Tensor | None,
218 + shared_w2_scale: torch.Tensor | None,
219 + shared_expert_id: int,
220 +) -> bool:
221 + shared_tensors = (shared_w1, shared_w2, shared_w1_scale, shared_w2_scale)
222 + has_shared_expert = any(tensor is not None for tensor in shared_tensors)
223 + if has_shared_expert:
224 + if not all(tensor is not None for tensor in shared_tensors):
225 + raise ValueError(
226 + "Heterogeneous fused MoE requires both shared weights and scales."
227 + )
228 + if shared_expert_id < 0:
229 + raise ValueError(
230 + "Heterogeneous fused MoE requires a non-negative shared expert ID."
231 + )
232 + elif shared_expert_id >= 0:
233 + raise ValueError(
234 + "A non-negative shared expert ID requires shared weights and scales."
235 + )
236 + return has_shared_expert
237 +
238 +
239 def _rocm_aiter_fused_moe_impl(
240 hidden_states: torch.Tensor,
241 w1: torch.Tensor,
@@ -236,7 +261,20 @@ def _rocm_aiter_fused_moe_impl(
261 swiglu_limit: float = 0.0,
262 beta: float | None = None,
263 linear_beta: float | None = None,
264 + shared_w1: torch.Tensor | None = None,
265 + shared_w2: torch.Tensor | None = None,
266 + shared_w1_scale: torch.Tensor | None = None,
267 + shared_w2_scale: torch.Tensor | None = None,
268 + shared_expert_id: int = -1,
269 ) -> torch.Tensor:
270 + has_shared_expert = _validate_rocm_aiter_fused_moe_shared_expert_args(
271 + shared_w1,
272 + shared_w2,
273 + shared_w1_scale,
274 + shared_w2_scale,
275 + shared_expert_id,
276 + )
277 +
278 from aiter import ActivationType, QuantType
279 from aiter.fused_moe import fused_moe
280
@@ -253,6 +291,15 @@ def _rocm_aiter_fused_moe_impl(
291 extra_kwargs["beta"] = beta
292 extra_kwargs["linear_beta"] = linear_beta
293
294 + if has_shared_expert:
295 + extra_kwargs.update(
296 + shared_w1=shared_w1,
297 + shared_w2=shared_w2,
298 + shared_w1_scale=shared_w1_scale,
299 + shared_w2_scale=shared_w2_scale,
300 + shared_expert_id=shared_expert_id,
301 + )
302 +
303 return fused_moe(
304 hidden_states,
305 w1,
@@ -304,6 +351,11 @@ def _rocm_aiter_fused_moe_fake(
351 swiglu_limit: float = 0.0,
352 beta: float | None = None,
353 linear_beta: float | None = None,
354 + shared_w1: torch.Tensor | None = None,
355 + shared_w2: torch.Tensor | None = None,
356 + shared_w1_scale: torch.Tensor | None = None,
357 + shared_w2_scale: torch.Tensor | None = None,
358 + shared_expert_id: int = -1,
359 ) -> torch.Tensor:
360 if output_dtype is not None:
361 return torch.empty_like(hidden_states, dtype=output_dtype)
@@ -1997,6 +2049,44 @@ class rocm_aiter_ops:
2049
2050 return "gate_mode" in inspect.signature(fused_moe).parameters
2051
2052 + @staticmethod
2053 + def _probe_dsv4_i384_fhmoe_capability(num_tokens: int) -> bool:
2054 + """Probe AITER's CSV-backed DSV4 native-I384 FHMoE contract."""
2055 + if type(num_tokens) is not int or num_tokens <= 0:
2056 + return False
2057 +
2058 + try:
2059 + import inspect
2060 +
2061 + import aiter.fhmoe as fhmoe
2062 + from aiter.fused_moe import fused_moe
2063 +
2064 + params = inspect.signature(fused_moe).parameters
2065 + supports_dsv4_i384_fhmoe = getattr(fhmoe, "supports_dsv4_i384_fhmoe", None)
2066 + if not callable(supports_dsv4_i384_fhmoe):
2067 + return False
2068 + supports_num_tokens = supports_dsv4_i384_fhmoe(num_tokens)
2069 + except Exception:
2070 + return False
2071 +
2072 + return supports_num_tokens is True and all(
2073 + name in params
2074 + for name in (
2075 + "shared_w1",
2076 + "shared_w2",
2077 + "shared_w1_scale",
2078 + "shared_w2_scale",
2079 + "shared_expert_id",
2080 + )
2081 + )
2082 +
2083 + @classmethod
2084 + @if_aiter_supported
2085 + @functools.cache
2086 + def fused_moe_supports_heterogeneous_shared_expert(cls, num_tokens: int) -> bool:
2087 + """Whether AITER has DSV4 native-I384 configs through the given M."""
2088 + return cls._probe_dsv4_i384_fhmoe_capability(num_tokens)
2089 +
2090 @staticmethod
2091 def register_ops_once() -> None:
2092 global _OPS_REGISTERED
@@ -2465,6 +2555,11 @@ class rocm_aiter_ops:
2555 swiglu_limit: float = 0.0,
2556 beta: float | None = None,
2557 linear_beta: float | None = None,
2558 + shared_w1: torch.Tensor | None = None,
2559 + shared_w2: torch.Tensor | None = None,
2560 + shared_w1_scale: torch.Tensor | None = None,
2561 + shared_w2_scale: torch.Tensor | None = None,
2562 + shared_expert_id: int = -1,
2563 ) -> torch.Tensor:
2564 return torch.ops.vllm.rocm_aiter_fused_moe(
2565 hidden_states,
@@ -2491,6 +2586,11 @@ class rocm_aiter_ops:
2586 swiglu_limit,
2587 beta,
2588 linear_beta,
2589 + shared_w1,
2590 + shared_w2,
2591 + shared_w1_scale,
2592 + shared_w2_scale,
2593 + shared_expert_id,
2594 )
2595
2596 @staticmethod
@@ -3074,6 +3174,13 @@ class rocm_aiter_ops:
3174
3175 return shuffle_scale_a16w4(tensor, num_experts, gate_up)
3176
3177 + @staticmethod
3178 + def shuffle_scale(tensor: torch.Tensor) -> torch.Tensor:
3179 + """Shuffle a non-gate/up E8M0 scale tensor for AITER."""
3180 + from aiter.ops.shuffle import shuffle_scale
3181 +
3182 + return shuffle_scale(tensor)
3183 +
3184 @staticmethod
3185 def shuffle_weights(
3186 *tensors: torch.Tensor, layout: tuple[int, int] = (16, 16)
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
+15
@@ -250,6 +250,11 @@ def rocm_aiter_fused_experts(
250 num_local_tokens: torch.Tensor | None = None,
251 output_dtype: torch.dtype | None = None,
252 moe_sorting_dispatch_policy: int = 0,
253 + shared_w1: torch.Tensor | None = None,
254 + shared_w2: torch.Tensor | None = None,
255 + shared_w1_scale: torch.Tensor | None = None,
256 + shared_w2_scale: torch.Tensor | None = None,
257 + shared_expert_id: int = -1,
258 ) -> torch.Tensor:
259 """ROCm AITER fused MoE expert computation."""
260 if quant_config is None:
@@ -419,8 +424,18 @@ def rocm_aiter_fused_experts(
424 bias1=quant_config.w1_bias if quant_config.use_mxfp4_w4a16 else None,
425 bias2=quant_config.w2_bias if quant_config.use_mxfp4_w4a16 else None,
426 moe_sorting_dispatch_policy=moe_sorting_dispatch_policy,
427 + swiglu_limit=(
428 + 0.0
429 + if moe_config.swiglu_limit is None
430 + else float(moe_config.swiglu_limit)
431 + ),
432 beta=moe_config.activation_situ_beta,
433 linear_beta=moe_config.activation_situ_linear_beta,
434 + shared_w1=shared_w1,
435 + shared_w2=shared_w2,
436 + shared_w1_scale=shared_w1_scale,
437 + shared_w2_scale=shared_w2_scale,
438 + shared_expert_id=shared_expert_id,
439 )
440
441
vllm/models/deepseek_v4/amd/model.py
+377 -5
@@ -1,7 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import typing
4 +import weakref
5 from collections.abc import Callable, Iterable
6 +from dataclasses import replace
7 from itertools import islice
8
9 import regex as re
@@ -20,9 +22,14 @@ from vllm.logger import init_logger
22 from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
23 from vllm.model_executor.layers.fused_moe import (
24 FusedMoEFactory,
25 + FusedMoEQuantConfig,
26 GateLinear,
27 + RoutedExperts,
28 fused_moe_make_expert_params_mapping,
29 )
30 +from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
31 + rocm_aiter_fused_experts,
32 +)
33 from vllm.model_executor.layers.fused_moe.utils import (
34 is_model_fused_shared_expert_compatible,
35 )
@@ -41,6 +48,7 @@ from vllm.model_executor.layers.mhc import (
48 MHCPreOp,
49 )
50 from vllm.model_executor.layers.quantization import QuantizationConfig
51 +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
52 from vllm.model_executor.layers.quantization.utils.config_utils import (
53 is_shared_expert_quant_fse_compatible,
54 )
@@ -65,6 +73,7 @@ from vllm.model_executor.models.utils import (
73 )
74 from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
75 from vllm.platforms import current_platform
76 +from vllm.platforms.rocm import on_gfx950
77 from vllm.sequence import IntermediateTensors
78
79 logger = init_logger(__name__)
@@ -163,6 +172,336 @@ class DeepseekV4MLP(nn.Module):
172 return x
173
174
175 +def _use_heterogeneous_fhmoe(num_tokens: int) -> bool:
176 + return num_tokens > 0 and (
177 + rocm_aiter_ops.fused_moe_supports_heterogeneous_shared_expert(num_tokens)
178 + is True
179 + )
180 +
181 +
182 +def _validate_heterogeneous_routes(
183 + topk_weights: torch.Tensor,
184 + topk_ids: torch.Tensor,
185 + experts_per_token: int,
186 +) -> None:
187 + if (
188 + topk_weights.ndim != 2
189 + or topk_ids.ndim != 2
190 + or topk_weights.shape != topk_ids.shape
191 + ):
192 + raise ValueError(
193 + "Heterogeneous fused MoE routes must have equal two-dimensional shapes."
194 + )
195 + expected_columns = experts_per_token + 1
196 + if topk_weights.shape[1] != expected_columns:
197 + raise ValueError(
198 + "Heterogeneous fused MoE routes must have exactly "
199 + f"{expected_columns} columns, got {topk_weights.shape[1]}."
200 + )
201 +
202 +
203 +def _heterogeneous_shared_expert_enabled(vllm_config: VllmConfig) -> bool:
204 + config = vllm_config.model_config.hf_config
205 + quant_config = vllm_config.quant_config
206 + parallel_config = vllm_config.parallel_config
207 + reasons: list[str] = []
208 +
209 + if not rocm_aiter_ops.is_fusion_moe_shared_experts_enabled():
210 + return False
211 + if not on_gfx950():
212 + reasons.append("the device is not gfx950")
213 + if parallel_config.enable_expert_parallel:
214 + reasons.append("expert parallelism is enabled")
215 + if parallel_config.enable_eplb:
216 + reasons.append("EPLB is enabled")
217 + if parallel_config.tensor_parallel_size != 8:
218 + reasons.append("tensor parallelism is not 8")
219 + if parallel_config.data_parallel_size != 1:
220 + reasons.append("data parallelism is enabled")
221 + if parallel_config.prefill_context_parallel_size != 1:
222 + reasons.append("prefill context parallelism is enabled")
223 + if vllm_config.kernel_config.moe_backend != "aiter":
224 + reasons.append("the MoE backend is not AITER")
225 + if vllm_config.model_config.dtype != torch.bfloat16:
226 + reasons.append("the model dtype is not BF16")
227 + if not rocm_aiter_ops.fused_moe_supports_heterogeneous_shared_expert(1):
228 + reasons.append(
229 + "AITER lacks CSV-backed DSV4 I384 heterogeneous shared-expert support"
230 + )
231 +
232 + expected_model = {
233 + "n_routed_experts": 384,
234 + "num_experts_per_tok": 6,
235 + "n_shared_experts": 1,
236 + "hidden_size": 7168,
237 + "moe_intermediate_size": 3072,
238 + "hidden_act": "silu",
239 + "expert_dtype": "fp4",
240 + "topk_method": "noaux_tc",
241 + }
242 + for name, expected in expected_model.items():
243 + if getattr(config, name, None) != expected:
244 + reasons.append(f"{name} is not {expected!r}")
245 +
246 + if quant_config is None or quant_config.get_name() != "deepseek_v4_fp8":
247 + reasons.append("the DeepSeek V4 FP8 quantization config is not active")
248 + else:
249 + if getattr(quant_config, "moe_quant_algo", "").upper() == "NVFP4":
250 + reasons.append("routed experts use NVFP4")
251 + if getattr(quant_config, "weight_block_size", None) != [128, 128]:
252 + reasons.append("shared experts are not block-128 FP8")
253 + if not getattr(quant_config, "is_checkpoint_fp8_serialized", False):
254 + reasons.append("shared experts are not serialized as FP8")
255 + if not getattr(quant_config, "is_scale_e8m0", False):
256 + reasons.append("shared-expert scales are not E8M0")
257 + if getattr(quant_config, "ignored_layers", None):
258 + reasons.append("the quantization config has ignored layers")
259 +
260 + offload_config = getattr(vllm_config, "offload_config", None)
261 + if offload_config is not None and (
262 + getattr(getattr(offload_config, "uva", None), "cpu_offload_gb", 0) > 0
263 + or getattr(
264 + getattr(offload_config, "prefetch", None),
265 + "offload_group_size",
266 + 0,
267 + )
268 + > 0
269 + ):
270 + reasons.append("weight offloading is enabled")
271 +
272 + if reasons:
273 + logger.debug_once(
274 + "DeepSeek V4 heterogeneous shared-expert fusion is unavailable: %s. "
275 + "Continuing with the existing shared-expert path.",
276 + "; ".join(reasons),
277 + )
278 + return False
279 +
280 + logger.info_once(
281 + "Using AITER DeepSeek V4 heterogeneous MoE for CSV-covered MoE input "
282 + "rows (M) per invocation (native intermediate width 384)."
283 + )
284 + return True
285 +
286 +
287 +@torch.no_grad()
288 +def _prepare_native_fp8_shared_expert(
289 + w13: torch.Tensor,
290 + w2: torch.Tensor,
291 + w13_scale: torch.Tensor,
292 + w2_scale: torch.Tensor,
293 + intermediate_size: int,
294 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
295 + """Expand native block-128 E8M0 scales to FHMoE's 1x32 layout."""
296 + if w13.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn:
297 + raise ValueError("Heterogeneous shared-expert weights must be FP8 E4M3.")
298 + hidden_size = w13.shape[1]
299 + if w13.shape != (2 * intermediate_size, hidden_size):
300 + raise ValueError(f"Unexpected shared W13 shape {tuple(w13.shape)}.")
301 + if w2.shape != (hidden_size, intermediate_size):
302 + raise ValueError(f"Unexpected shared W2 shape {tuple(w2.shape)}.")
303 +
304 + def scale_bytes(scale: torch.Tensor) -> torch.Tensor:
305 + if scale.dtype == torch.float8_e8m0fnu:
306 + return scale.contiguous().view(torch.uint8)
307 + if scale.dtype == torch.uint8:
308 + return scale.contiguous()
309 + if scale.dtype == torch.float32:
310 + bits = scale.contiguous().view(torch.int32)
311 + exponent = (bits >> 23).to(torch.uint8)
312 + if not torch.equal(bits, exponent.to(torch.int32) << 23):
313 + raise ValueError("FP32 shared-expert scales are not exact E8M0.")
314 + return exponent
315 + raise ValueError(f"Unsupported shared-expert scale dtype {scale.dtype}.")
316 +
317 + w13_scale = scale_bytes(w13_scale)
318 + w2_scale = scale_bytes(w2_scale)
319 + expected_w13_scale = (2 * intermediate_size // 128, hidden_size // 128)
320 + expected_w2_scale = (hidden_size // 128, intermediate_size // 128)
321 + if tuple(w13_scale.shape) != expected_w13_scale:
322 + raise ValueError(
323 + f"Expected shared W13 scale shape {expected_w13_scale}, got "
324 + f"{tuple(w13_scale.shape)}."
325 + )
326 + if tuple(w2_scale.shape) != expected_w2_scale:
327 + raise ValueError(
328 + f"Expected shared W2 scale shape {expected_w2_scale}, got "
329 + f"{tuple(w2_scale.shape)}."
330 + )
331 +
332 + w13_scale = w13_scale.repeat_interleave(128, dim=0).repeat_interleave(4, dim=1)
333 + w2_scale = w2_scale.repeat_interleave(128, dim=0).repeat_interleave(4, dim=1)
334 +
335 + scale_cols = ((w2_scale.shape[1] + 7) // 8) * 8
336 + if scale_cols != w2_scale.shape[1]:
337 + padded_w2_scale = torch.full(
338 + (w2_scale.shape[0], scale_cols),
339 + 0x7F,
340 + dtype=torch.uint8,
341 + device=w2_scale.device,
342 + )
343 + padded_w2_scale[:, : w2_scale.shape[1]].copy_(w2_scale)
344 + w2_scale = padded_w2_scale
345 +
346 + return (
347 + w13.unsqueeze(0),
348 + w2.unsqueeze(0),
349 + w13_scale.view(torch.float8_e8m0fnu),
350 + w2_scale.view(torch.float8_e8m0fnu),
351 + )
352 +
353 +
354 +class DeepseekV4HeterogeneousMxfp4MoEMethod(Mxfp4MoEMethod):
355 + def process_weights_after_loading(self, layer: RoutedExperts) -> None:
356 + super().process_weights_after_loading(layer)
357 + assert isinstance(layer, DeepseekV4HeterogeneousSharedRoutedExperts)
358 + layer.prepare_heterogeneous_shared_expert()
359 +
360 +
361 +class DeepseekV4HeterogeneousSharedRoutedExperts(RoutedExperts):
362 + def __init__(
363 + self,
364 + *args,
365 + shared_expert: DeepseekV4MLP,
366 + shared_expert_id: int,
367 + **kwargs,
368 + ):
369 + self._shared_expert_ref = weakref.ref(shared_expert)
370 + self.shared_expert_id = shared_expert_id
371 + self._routed_quant_config: FusedMoEQuantConfig | None = None
372 + super().__init__(*args, **kwargs)
373 + if self.expert_map_manager.num_fused_shared_experts != 1:
374 + raise ValueError("Heterogeneous fusion requires one appended expert.")
375 + if self.shared_expert_id != self.global_num_experts:
376 + raise ValueError("The shared expert ID must name the appended slot.")
377 + self.register_buffer("shared_w1", None, persistent=False)
378 + self.register_buffer("shared_w2", None, persistent=False)
379 + self.register_buffer("shared_w1_scale", None, persistent=False)
380 + self.register_buffer("shared_w2_scale", None, persistent=False)
381 +
382 + def _get_quant_method(self, prefix, quant_config, moe_config):
383 + quant_method = super()._get_quant_method(prefix, quant_config, moe_config)
384 + if not isinstance(quant_method, Mxfp4MoEMethod):
385 + raise ValueError("Heterogeneous fusion requires MXFP4 routed experts.")
386 + return DeepseekV4HeterogeneousMxfp4MoEMethod(moe_config)
387 +
388 + @torch.no_grad()
389 + def prepare_heterogeneous_shared_expert(self) -> None:
390 + shared_expert = self._shared_expert_ref()
391 + if shared_expert is None:
392 + raise RuntimeError("The native shared-expert module was released.")
393 +
394 + shared = _prepare_native_fp8_shared_expert(
395 + shared_expert.gate_up_proj.weight,
396 + shared_expert.down_proj.weight,
397 + shared_expert.gate_up_proj.weight_scale_inv,
398 + shared_expert.down_proj.weight_scale_inv,
399 + self.moe_config.intermediate_size_per_partition,
400 + )
401 + self.shared_w1 = rocm_aiter_ops.shuffle_weight_a16w4(shared[0], 16, True)
402 + self.shared_w2 = rocm_aiter_ops.shuffle_weight_a16w4(shared[1], 16, False)
403 + self.shared_w1_scale = rocm_aiter_ops.shuffle_scale_a16w4(shared[2], 1, True)
404 + self.shared_w2_scale = rocm_aiter_ops.shuffle_scale(shared[3])
405 +
406 + self._ensure_moe_quant_config_init()
407 + quant_config = self.quant_method.moe_quant_config
408 + if quant_config is None:
409 + raise RuntimeError("The routed MXFP4 quantization config is unavailable.")
410 + if quant_config.w1_scale is None or quant_config.w2_scale is None:
411 + raise RuntimeError("The routed MXFP4 scales are unavailable.")
412 +
413 + def routed_scales(scale: torch.Tensor) -> torch.Tensor:
414 + scale = scale.reshape(self.local_num_experts, -1, scale.shape[-1])[
415 + : self.shared_expert_id
416 + ]
417 + return scale.reshape(-1, scale.shape[-1]).contiguous()
418 +
419 + self._routed_quant_config = replace(
420 + quant_config,
421 + _w1=replace(
422 + quant_config._w1,
423 + scale=routed_scales(quant_config.w1_scale),
424 + ),
425 + _w2=replace(
426 + quant_config._w2,
427 + scale=routed_scales(quant_config.w2_scale),
428 + ),
429 + )
430 +
431 + def forward_modular(
432 + self,
433 + x: torch.Tensor,
434 + topk_weights: torch.Tensor,
435 + topk_ids: torch.Tensor,
436 + shared_experts=None,
437 + shared_experts_input: torch.Tensor | None = None,
438 + ) -> torch.Tensor:
439 + if shared_experts is not None or shared_experts_input is not None:
440 + raise ValueError("The heterogeneous expert owns the shared path.")
441 + _validate_heterogeneous_routes(
442 + topk_weights,
443 + topk_ids,
444 + self.moe_config.experts_per_token,
445 + )
446 + if any(
447 + tensor is None
448 + for tensor in (
449 + self.shared_w1,
450 + self.shared_w2,
451 + self.shared_w1_scale,
452 + self.shared_w2_scale,
453 + )
454 + ):
455 + raise RuntimeError("Heterogeneous shared weights were not prepared.")
456 +
457 + if _use_heterogeneous_fhmoe(x.shape[0]):
458 + self._ensure_moe_quant_config_init()
459 + quant_config = self.quant_method.moe_quant_config
460 + if quant_config is None:
461 + raise RuntimeError("The routed MXFP4 config is unavailable.")
462 + return rocm_aiter_fused_experts(
463 + hidden_states=x,
464 + w1=self.w13_weight,
465 + w2=self.w2_weight,
466 + topk_weights=topk_weights,
467 + topk_ids=topk_ids,
468 + moe_config=self.moe_config,
469 + activation=self.activation,
470 + expert_mask=self.expert_map,
471 + quant_config=quant_config,
472 + output_dtype=x.dtype,
473 + moe_sorting_dispatch_policy=(rocm_aiter_ops.get_moe_dispatch_policy()),
474 + shared_w1=self.shared_w1,
475 + shared_w2=self.shared_w2,
476 + shared_w1_scale=self.shared_w1_scale,
477 + shared_w2_scale=self.shared_w2_scale,
478 + shared_expert_id=self.shared_expert_id,
479 + )
480 +
481 + shared_expert = self._shared_expert_ref()
482 + if shared_expert is None or self._routed_quant_config is None:
483 + raise RuntimeError("The separate shared-expert fallback is unavailable.")
484 + shared_out = shared_expert(x)
485 + routed_w1 = self.w13_weight[: self.shared_expert_id]
486 + routed_w2 = self.w2_weight[: self.shared_expert_id]
487 + routed_w1.is_shuffled = True
488 + routed_w2.is_shuffled = True
489 + routed = rocm_aiter_fused_experts(
490 + hidden_states=x,
491 + w1=routed_w1,
492 + w2=routed_w2,
493 + topk_weights=topk_weights[:, :-1].contiguous(),
494 + topk_ids=topk_ids[:, :-1].contiguous(),
495 + moe_config=self.moe_config,
496 + activation=self.activation,
497 + expert_mask=self.expert_map,
498 + quant_config=self._routed_quant_config,
499 + output_dtype=x.dtype,
500 + moe_sorting_dispatch_policy=rocm_aiter_ops.get_moe_dispatch_policy(),
501 + )
502 + return routed + shared_out
503 +
504 +
505 def _fuse_shared_experts_enabled(config) -> bool:
506 return bool(
507 getattr(config, "n_shared_experts", None)
@@ -176,6 +515,7 @@ class DeepseekV4MoE(nn.Module):
515 self,
516 vllm_config: VllmConfig,
517 prefix: str = "",
518 + fuse_heterogeneous_shared_expert: bool = False,
519 ):
520 super().__init__()
521
@@ -230,7 +570,12 @@ class DeepseekV4MoE(nn.Module):
570 # TODO: Historically, only `VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1`
571 # is checked to enable FSE for DeepSeek-v4, despite AITER not being used.
572 # This should be cleaned up and use `resolve_layer_fused_shared_expert`.
233 - fse_requested = _fuse_shared_experts_enabled(config)
573 + self.fuse_heterogeneous_shared_expert = fuse_heterogeneous_shared_expert
574 + fse_requested = (
575 + _fuse_shared_experts_enabled(config)
576 + and not self.fuse_heterogeneous_shared_expert
577 + )
578 + fse_compatible = False
579 if fse_requested:
580 fse_compatible, fse_reason = is_shared_expert_quant_fse_compatible(
581 quant_config,
@@ -267,12 +612,15 @@ class DeepseekV4MoE(nn.Module):
612 self.experts_start_idx = self.tp_rank * self.n_local_experts
613 self.experts_end_idx = self.experts_start_idx + self.n_local_experts
614
615 + fuse_shared_into_routed = (
616 + self.is_fused_shared_expert_enabled or self.fuse_heterogeneous_shared_expert
617 + )
618 self.experts = FusedMoEFactory(
271 - shared_experts=self.shared_experts,
619 + shared_experts=(None if fuse_shared_into_routed else self.shared_experts),
620 n_shared_experts=(
273 - config.n_shared_experts if self.is_fused_shared_expert_enabled else None
621 + config.n_shared_experts if fuse_shared_into_routed else None
622 ),
275 - fuse_shared_experts=self.is_fused_shared_expert_enabled,
623 + fuse_shared_experts=fuse_shared_into_routed,
624 gate=self.gate,
625 num_experts=config.n_routed_experts,
626 top_k=config.num_experts_per_tok,
@@ -287,6 +635,19 @@ class DeepseekV4MoE(nn.Module):
635 hash_indices_table=self.gate.tid2eid,
636 swiglu_limit=self.swiglu_limit,
637 router_logits_dtype=torch.float32,
638 + routed_experts_cls=(
639 + DeepseekV4HeterogeneousSharedRoutedExperts
640 + if self.fuse_heterogeneous_shared_expert
641 + else None
642 + ),
643 + routed_experts_args=(
644 + {
645 + "shared_expert": self.shared_experts,
646 + "shared_expert_id": config.n_routed_experts,
647 + }
648 + if self.fuse_heterogeneous_shared_expert
649 + else None
650 + ),
651 )
652
653 def forward(
@@ -316,6 +677,7 @@ class DeepseekV4DecoderLayer(nn.Module):
677 prefix,
678 topk_indices_buffer: torch.Tensor | None = None,
679 aux_stream_list: list[torch.cuda.Stream] | None = None,
680 + fuse_heterogeneous_shared_expert: bool = False,
681 ):
682 super().__init__()
683
@@ -333,7 +695,11 @@ class DeepseekV4DecoderLayer(nn.Module):
695 topk_indices_buffer=topk_indices_buffer,
696 aux_stream_list=aux_stream_list,
697 )
336 - self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")
698 + self.ffn = DeepseekV4MoE(
699 + vllm_config,
700 + prefix=f"{prefix}.ffn",
701 + fuse_heterogeneous_shared_expert=fuse_heterogeneous_shared_expert,
702 + )
703
704 self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
705 self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)
@@ -598,6 +964,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
964 else:
965 self.embed_tokens = PPMissingLayer()
966
967 + self.fuse_heterogeneous_shared_expert = _heterogeneous_shared_expert_enabled(
968 + vllm_config
969 + )
970 self.start_layer, self.end_layer, self.layers = make_layers(
971 config.num_hidden_layers,
972 lambda prefix: DeepseekV4DecoderLayer(
@@ -605,6 +974,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
974 prefix=prefix,
975 topk_indices_buffer=self.topk_indices_buffer,
976 aux_stream_list=aux_stream_list,
977 + fuse_heterogeneous_shared_expert=(
978 + self.fuse_heterogeneous_shared_expert
979 + ),
980 ),
981 prefix=f"{prefix}.layers",
982 )