A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 01:30 UTC
Add support for ModelOpt MXFP8 MoE models (#35986)
Signed-off-by: Daniel Serebrenik <daserebrenik@nvidia.com>
danisereb committed
Mar 8, 2026 at 22:00 UTC
0a6a3a12906bd581fb2983c81b4d51dc60e0bb4a
4 files changed
+597
-18
tests/kernels/moe/test_ocp_mx_moe.py
+185
-2
@@ -20,6 +20,8 @@ TRTLLM_GEN_MXFP4_AVAILABLE = (
20
current_platform.is_cuda() and current_platform.is_device_capability_family(100)
21
)
22
23
+TRTLLM_GEN_MXFP8_AVAILABLE = TRTLLM_GEN_MXFP4_AVAILABLE
24
+
25
HOPPER_MXFP4_BF16_AVAILABLE = (
26
current_platform.is_cuda()
27
and current_platform.is_device_capability(90)
@@ -34,9 +36,15 @@ if TRTLLM_GEN_MXFP4_AVAILABLE:
36
shuffle_matrix_a,
37
shuffle_matrix_sf_a,
38
trtllm_fp4_block_scale_moe,
39
+ trtllm_fp8_block_scale_moe,
40
)
41
from flashinfer.fp4_quantization import nvfp4_block_scale_interleave
39
- from flashinfer.fused_moe.core import get_w2_permute_indices_with_cache
42
+
43
+if TRTLLM_GEN_MXFP8_AVAILABLE:
44
+ from flashinfer.fused_moe.core import (
45
+ Fp8QuantizationType,
46
+ get_w2_permute_indices_with_cache,
47
+ )
48
49
50
@dataclass
@@ -160,6 +168,7 @@ def reference_moe(
168
beta,
169
limit,
170
act_type,
171
+ is_gated,
172
):
173
# renormalize routing
174
experts = torch.topk(roouting_logits, k=topk, dim=-1, sorted=True)
@@ -170,7 +179,12 @@ def reference_moe(
179
mlp1_weight = w13[expert_indices, ...]
180
mlp1_bias = bias13[expert_indices, ...]
181
t = torch.einsum("beck,bk->bec", mlp1_weight, t) + mlp1_bias
173
- t = swiglu(t, alpha=alpha, beta=beta, limit=limit)
182
+ if is_gated:
183
+ t = swiglu(t, alpha=alpha, beta=beta, limit=limit)
184
+ else:
185
+ # RELU2_NO_MUL: relu(x)^2
186
+ t = torch.relu(t)
187
+ t = t * t
188
189
if act_type == "mxfp8":
190
t_quantized, t_scale = mxfp8_quantize(
@@ -569,6 +583,7 @@ def test_trtllm_gen_mxfp4_fused_moe(
583
beta,
584
limit,
585
act_type,
586
+ is_gated=True,
587
)
588
ref_result[start_idx:end_idx].copy_(chunk_result)
589
@@ -705,6 +720,7 @@ def test_flashinfer_cutlass_mxfp4_fused_moe(
720
beta,
721
limit,
722
"bf16",
723
+ is_gated=True,
724
)
725
726
from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
@@ -890,6 +906,7 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe(
906
beta,
907
limit,
908
"mxfp8",
909
+ is_gated=True,
910
)
911
912
# Prepare inputs for FlashInfer CUTLASS fused MoE
@@ -965,3 +982,169 @@ def test_flashinfer_cutlass_mxfp4_mxfp8_fused_moe(
982
983
# Allow some mismatch due to MXFP4 quantization
984
check_accuracy(ref, out, atol=0, rtol=0.3, percent=0.8)
985
+
986
+
987
+@pytest.mark.parametrize("topk", [1, 4])
988
+@pytest.mark.parametrize("num_experts", [32])
989
+@pytest.mark.parametrize("num_tokens", [1, 128])
990
+@pytest.mark.parametrize("intermediate_size,hidden_size", [(3072, 3072)])
991
+@pytest.mark.parametrize("is_gated", [True], ids=["gated"])
992
+@pytest.mark.skipif(
993
+ not TRTLLM_GEN_MXFP8_AVAILABLE,
994
+ reason="nvidia gpu and compute capability sm100 is required for this test",
995
+)
996
+def test_trtllm_gen_mxfp8_block_scale_moe(
997
+ topk: int,
998
+ num_experts: int,
999
+ num_tokens: int,
1000
+ intermediate_size: int,
1001
+ hidden_size: int,
1002
+ is_gated: bool,
1003
+):
1004
+ torch.manual_seed(42)
1005
+ device = "cuda:0"
1006
+
1007
+ inter_size = intermediate_size * (2 if is_gated else 1)
1008
+
1009
+ hidden_states = (
1010
+ torch.randn(num_tokens, hidden_size, device=device, dtype=torch.bfloat16) / 20
1011
+ )
1012
+ w13 = (
1013
+ torch.randn(
1014
+ num_experts,
1015
+ inter_size,
1016
+ hidden_size,
1017
+ device=device,
1018
+ dtype=torch.bfloat16,
1019
+ )
1020
+ / 20
1021
+ )
1022
+ w2 = (
1023
+ torch.randn(
1024
+ num_experts,
1025
+ hidden_size,
1026
+ intermediate_size,
1027
+ device=device,
1028
+ dtype=torch.bfloat16,
1029
+ )
1030
+ / 20
1031
+ )
1032
+ router_logits = torch.rand(
1033
+ num_tokens, num_experts, dtype=torch.float32, device=device
1034
+ )
1035
+ router_logits_kernel = router_logits.to(torch.bfloat16)
1036
+
1037
+ # Quantize weights to MXFP8 and normalize scales to [E, M, K//32].
1038
+ w13_q, w13_scale = mxfp8_quantize(w13, is_sf_swizzled_layout=False)
1039
+ w2_q, w2_scale = mxfp8_quantize(w2, is_sf_swizzled_layout=False)
1040
+ if w13_scale.ndim == 1:
1041
+ w13_scale = w13_scale.view(
1042
+ num_experts,
1043
+ inter_size,
1044
+ hidden_size // 32,
1045
+ )
1046
+ if w2_scale.ndim == 1:
1047
+ w2_scale = w2_scale.view(num_experts, hidden_size, intermediate_size // 32)
1048
+
1049
+ # Quantize activations to MXFP8.
1050
+ hidden_states_q, hidden_states_scale = mxfp8_quantize(
1051
+ hidden_states, is_sf_swizzled_layout=False
1052
+ )
1053
+ if hidden_states_scale.ndim == 1:
1054
+ hidden_states_scale = hidden_states_scale.view(num_tokens, hidden_size // 32)
1055
+
1056
+ # Reference output using dequantized tensors + MXFP8 intermediate quantization.
1057
+ w13_ref = mxfp8_dequantize(w13_q, w13_scale).to(torch.float32)
1058
+ w2_ref = mxfp8_dequantize(w2_q, w2_scale).to(torch.float32)
1059
+ hidden_states_ref = mxfp8_dequantize(hidden_states_q, hidden_states_scale).to(
1060
+ torch.float32
1061
+ )
1062
+ bias13 = torch.zeros(
1063
+ num_experts,
1064
+ intermediate_size * (2 if is_gated else 1),
1065
+ device=device,
1066
+ )
1067
+ bias2 = torch.zeros(num_experts, hidden_size, device=device)
1068
+ ref = reference_moe(
1069
+ router_logits_kernel.to(torch.float32),
1070
+ topk,
1071
+ num_experts,
1072
+ hidden_states_ref,
1073
+ w13_ref,
1074
+ bias13,
1075
+ w2_ref,
1076
+ bias2,
1077
+ alpha=1.0,
1078
+ beta=0.0,
1079
+ limit=None,
1080
+ act_type="mxfp8",
1081
+ is_gated=is_gated,
1082
+ )
1083
+
1084
+ # Shuffle weights/scales with the same indexed layout used by TRTLLM kernels.
1085
+ epilogue_tile_m = 128
1086
+ gemm1_weights_shuffled = []
1087
+ gemm1_scales_shuffled = []
1088
+ gemm2_weights_shuffled = []
1089
+ gemm2_scales_shuffled = []
1090
+ for i in range(num_experts):
1091
+ w13_rows = intermediate_size * (2 if is_gated else 1)
1092
+ w13_interleaved = w13_q[i].clone().reshape(w13_rows, -1)
1093
+ w13_scale_interleaved = w13_scale[i].clone().reshape(w13_rows, -1)
1094
+ if is_gated:
1095
+ w13_interleaved = reorder_rows_for_gated_act_gemm(w13_interleaved)
1096
+ w13_scale_interleaved = reorder_rows_for_gated_act_gemm(
1097
+ w13_scale_interleaved
1098
+ )
1099
+ gemm1_weights_shuffled.append(
1100
+ shuffle_matrix_a(w13_interleaved.view(torch.uint8), epilogue_tile_m)
1101
+ .contiguous()
1102
+ .view(w13_q.dtype)
1103
+ )
1104
+ gemm2_weights_shuffled.append(
1105
+ shuffle_matrix_a(w2_q[i].view(torch.uint8), epilogue_tile_m)
1106
+ .contiguous()
1107
+ .view(w2_q.dtype)
1108
+ )
1109
+
1110
+ gemm1_scales_shuffled.append(
1111
+ shuffle_matrix_sf_a(
1112
+ w13_scale_interleaved.view(torch.uint8).reshape(w13_rows, -1),
1113
+ epilogue_tile_m,
1114
+ )
1115
+ .contiguous()
1116
+ .view(w13_scale.dtype)
1117
+ )
1118
+ gemm2_scales_shuffled.append(
1119
+ shuffle_matrix_sf_a(
1120
+ w2_scale[i].view(torch.uint8).reshape(hidden_size, -1), epilogue_tile_m
1121
+ )
1122
+ .contiguous()
1123
+ .view(w2_scale.dtype)
1124
+ )
1125
+
1126
+ out = trtllm_fp8_block_scale_moe(
1127
+ routing_logits=router_logits_kernel,
1128
+ routing_bias=None,
1129
+ hidden_states=hidden_states_q,
1130
+ hidden_states_scale=hidden_states_scale,
1131
+ gemm1_weights=torch.stack(gemm1_weights_shuffled),
1132
+ gemm1_weights_scale=torch.stack(gemm1_scales_shuffled),
1133
+ gemm2_weights=torch.stack(gemm2_weights_shuffled),
1134
+ gemm2_weights_scale=torch.stack(gemm2_scales_shuffled),
1135
+ num_experts=num_experts,
1136
+ top_k=topk,
1137
+ n_group=None,
1138
+ topk_group=None,
1139
+ intermediate_size=intermediate_size,
1140
+ local_expert_offset=0,
1141
+ local_num_experts=num_experts,
1142
+ routed_scaling_factor=None,
1143
+ routing_method_type=1, # renormalize routing
1144
+ use_shuffled_weight=True,
1145
+ weight_layout=0, # MajorK
1146
+ fp8_quantization_type=Fp8QuantizationType.MxFp8,
1147
+ )
1148
+
1149
+ # Block-scale MXFP8 kernels are approximate; require majority close.
1150
+ check_accuracy(ref, out, atol=0.1, rtol=0.85, percent=0.8)
vllm/model_executor/layers/fused_moe/layer.py
+9
@@ -1204,17 +1204,26 @@ class FusedMoE(CustomOp):
1204
# Determine per-tensor weight scale patterns based on variant
1205
# Use the dedicated method instead of brittle string matching
1206
uses_weight_scale_2 = self.quant_method.uses_weight_scale_2_pattern()
1207
+ quant_method = getattr(param, "quant_method", None)
1208
1209
# Call _load_per_tensor_weight_scale() to load per-tensor (scalar)
1210
# weights scales.
1211
# Input scales are always per-tensor.
1212
# Weight scales: FP4 uses "weight_scale_2" and FP8 uses
1213
# "weight_scale" for per-tensor scales.
1214
+ # NOTE: ModelOpt MXFP8 MoE uses block scales in weight_scale
1215
+ # tensors (quant_method=BLOCK), so those must not be treated
1216
+ # as per-tensor scalars here.
1217
+ is_block_weight_scale = (
1218
+ "weight_scale" in weight_name
1219
+ and quant_method == FusedMoeWeightScaleSupported.BLOCK.value
1220
+ )
1221
is_per_tensor = (
1222
"weight_scale_2" in weight_name
1223
if uses_weight_scale_2
1224
else "weight_scale" in weight_name
1225
) or "input_scale" in weight_name
1226
+ is_per_tensor = is_per_tensor and not is_block_weight_scale
1227
if is_per_tensor:
1228
self._load_per_tensor_weight_scale(
1229
shard_id=shard_id,
vllm/model_executor/layers/fused_moe/oracle/mxfp8.py
new
+44
@@ -0,0 +1,44 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from enum import Enum
4
+
5
+from vllm.logger import init_logger
6
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
7
+
8
+logger = init_logger(__name__)
9
+
10
+
11
+class MxFp8MoeBackend(Enum):
12
+ FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
13
+
14
+
15
+def select_mxfp8_moe_backend(
16
+ config: FusedMoEConfig,
17
+) -> MxFp8MoeBackend:
18
+ if config.is_lora_enabled:
19
+ raise NotImplementedError("LoRA is not supported for MXFP8 MoE.")
20
+
21
+ AVAILABLE_BACKENDS = [
22
+ MxFp8MoeBackend.FLASHINFER_TRTLLM,
23
+ ]
24
+
25
+ runner_backend = config.moe_backend
26
+ if runner_backend != "auto":
27
+ mapping = {
28
+ "flashinfer_trtllm": MxFp8MoeBackend.FLASHINFER_TRTLLM,
29
+ }
30
+ if backend := mapping.get(runner_backend):
31
+ logger.info_once(
32
+ "Using '%s' MxFp8 MoE backend (user-requested).",
33
+ backend.value,
34
+ )
35
+ return backend
36
+ raise ValueError(
37
+ f"moe_backend='{runner_backend}' is not supported for MXFP8 MoE. "
38
+ f"Expected one of {list(mapping.keys())}."
39
+ )
40
+
41
+ # Auto-select: only one backend available for now.
42
+ backend = AVAILABLE_BACKENDS[0]
43
+ logger.info_once("Using '%s' MxFp8 MoE backend.", backend.value)
44
+ return backend
vllm/model_executor/layers/quantization/modelopt.py
+359
-16
@@ -9,17 +9,19 @@ from torch.nn.parameter import Parameter
9
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.logger import init_logger
12
-from vllm.model_executor.kernels.linear import (
13
- init_fp8_linear_kernel,
14
-)
12
+from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
13
from vllm.model_executor.layers.attention import Attention, MLAAttention
14
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15
from vllm.model_executor.layers.fused_moe.config import (
16
FusedMoEConfig,
17
FusedMoEQuantConfig,
18
+ RoutingMethodType,
19
+)
20
+from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
21
+ FusedMoEMethodBase,
22
)
23
from vllm.model_executor.layers.fused_moe.layer import (
24
FusedMoE,
22
- FusedMoEMethodBase,
25
FusedMoeWeightScaleSupported,
26
)
27
from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
@@ -28,6 +30,10 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
30
make_fp8_moe_quant_config,
31
select_fp8_moe_backend,
32
)
33
+from vllm.model_executor.layers.fused_moe.oracle.mxfp8 import (
34
+ MxFp8MoeBackend,
35
+ select_mxfp8_moe_backend,
36
+)
37
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
38
convert_to_nvfp4_moe_kernel_format,
39
is_global_sf_supported_for_nvfp4_backend,
@@ -46,6 +52,9 @@ from vllm.model_executor.layers.quantization.base_config import (
52
QuantizeMethodBase,
53
)
54
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
55
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
56
+ swap_w13_to_w31,
57
+)
58
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
59
W8A8BlockFp8LinearOp,
60
process_fp8_input_tensor_strategy_moe,
@@ -60,6 +69,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
69
MXFP8_VALUE_DTYPE,
70
Mxfp8LinearBackend,
71
Mxfp8LinearOp,
72
+ mxfp8_e4m3_quantize,
73
swizzle_mxfp8_scale,
74
)
75
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
@@ -86,7 +96,8 @@ from vllm.model_executor.parameter import (
96
ModelWeightParameter,
97
PerTensorScaleParameter,
98
)
89
-from vllm.model_executor.utils import replace_parameter
99
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
100
+from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
101
102
if TYPE_CHECKING:
103
from vllm.model_executor.models.utils import WeightsMapper
@@ -1487,17 +1498,6 @@ class ModelOptMxFp8Config(ModelOptQuantConfigBase):
1498
# MXFP8 hardware acceleration requires Blackwell (SM100) or newer
1499
return 100
1500
1490
- def get_quant_method(
1491
- self, layer: torch.nn.Module, prefix: str
1492
- ) -> "QuantizeMethodBase | None":
1493
- # MXFP8 does not yet support MoE models
1494
- if isinstance(layer, FusedMoE):
1495
- raise NotImplementedError(
1496
- "MXFP8 quantization does not yet support MoE models. "
1497
- "Please use FP8 or NVFP4 quantization for MoE models."
1498
- )
1499
- return super().get_quant_method(layer, prefix)
1500
-
1501
@classmethod
1502
def override_quantization_method(
1503
cls, hf_quant_cfg, user_quant
@@ -1699,8 +1699,351 @@ class ModelOptMxFp8LinearMethod(LinearMethodBase):
1699
)
1700
1701
1702
+class ModelOptMxFp8FusedMoE(FusedMoEMethodBase):
1703
+ """FlashInfer TRTLLM MXFP8 block-scale MoE for ModelOpt checkpoints."""
1704
+
1705
+ def __init__(
1706
+ self,
1707
+ quant_config: ModelOptMxFp8Config,
1708
+ moe_config: FusedMoEConfig,
1709
+ ) -> None:
1710
+ super().__init__(moe_config)
1711
+ self.quant_config = quant_config
1712
+ assert self.quant_config.is_checkpoint_mxfp8_serialized
1713
+
1714
+ # Select MXFP8 MoE backend
1715
+ self.mxfp8_backend = select_mxfp8_moe_backend(self.moe)
1716
+
1717
+ def create_weights(
1718
+ self,
1719
+ layer: torch.nn.Module,
1720
+ num_experts: int,
1721
+ hidden_size: int,
1722
+ intermediate_size_per_partition: int,
1723
+ params_dtype: torch.dtype,
1724
+ **extra_weight_attrs,
1725
+ ):
1726
+ layer.intermediate_size_per_partition = intermediate_size_per_partition
1727
+ layer.hidden_size = hidden_size
1728
+ layer.orig_dtype = params_dtype
1729
+
1730
+ if hidden_size % MXFP8_BLOCK_SIZE != 0:
1731
+ raise ValueError(
1732
+ f"MXFP8 MoE requires hidden_size divisible by {MXFP8_BLOCK_SIZE}, "
1733
+ f"got {hidden_size}."
1734
+ )
1735
+ if intermediate_size_per_partition % MXFP8_BLOCK_SIZE != 0:
1736
+ raise ValueError(
1737
+ "MXFP8 MoE requires intermediate_size_per_partition divisible by "
1738
+ f"{MXFP8_BLOCK_SIZE}, got {intermediate_size_per_partition}."
1739
+ )
1740
+
1741
+ layer.num_experts = num_experts
1742
+ weight_loader = extra_weight_attrs.get("weight_loader")
1743
+ w13_num_shards = 2 if self.moe.is_act_and_mul else 1
1744
+
1745
+ # GEMM 1 weights: [E, (2I or I), H]
1746
+ w13_weight = ModelWeightParameter(
1747
+ data=torch.empty(
1748
+ num_experts,
1749
+ w13_num_shards * intermediate_size_per_partition,
1750
+ hidden_size,
1751
+ dtype=MXFP8_VALUE_DTYPE,
1752
+ ),
1753
+ input_dim=2,
1754
+ output_dim=1,
1755
+ weight_loader=weight_loader,
1756
+ )
1757
+ layer.register_parameter("w13_weight", w13_weight)
1758
+
1759
+ # GEMM 2 weights: [E, H, I]
1760
+ w2_weight = ModelWeightParameter(
1761
+ data=torch.empty(
1762
+ num_experts,
1763
+ hidden_size,
1764
+ intermediate_size_per_partition,
1765
+ dtype=MXFP8_VALUE_DTYPE,
1766
+ ),
1767
+ input_dim=2,
1768
+ output_dim=1,
1769
+ weight_loader=weight_loader,
1770
+ )
1771
+ layer.register_parameter("w2_weight", w2_weight)
1772
+
1773
+ # Per-block (K=32) E8M0 scales.
1774
+ w13_weight_scale = ModelWeightParameter(
1775
+ data=torch.empty(
1776
+ num_experts,
1777
+ w13_num_shards * intermediate_size_per_partition,
1778
+ hidden_size // MXFP8_BLOCK_SIZE,
1779
+ dtype=MXFP8_SCALE_DTYPE,
1780
+ ),
1781
+ input_dim=2,
1782
+ output_dim=1,
1783
+ weight_loader=weight_loader,
1784
+ )
1785
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
1786
+
1787
+ w2_weight_scale = ModelWeightParameter(
1788
+ data=torch.empty(
1789
+ num_experts,
1790
+ hidden_size,
1791
+ intermediate_size_per_partition // MXFP8_BLOCK_SIZE,
1792
+ dtype=MXFP8_SCALE_DTYPE,
1793
+ ),
1794
+ input_dim=2,
1795
+ output_dim=1,
1796
+ weight_loader=weight_loader,
1797
+ )
1798
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
1799
+
1800
+ # Ensure the generic MoE weight-loader treats these as block scales.
1801
+ set_weight_attrs(
1802
+ layer.w13_weight_scale,
1803
+ {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value},
1804
+ )
1805
+ set_weight_attrs(
1806
+ layer.w2_weight_scale,
1807
+ {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value},
1808
+ )
1809
+
1810
+ @staticmethod
1811
+ def _check_weight_dtypes(layer: torch.nn.Module) -> None:
1812
+ """Validate weight and scale dtypes before processing."""
1813
+ expected = {
1814
+ "w13_weight": MXFP8_VALUE_DTYPE,
1815
+ "w2_weight": MXFP8_VALUE_DTYPE,
1816
+ "w13_weight_scale": MXFP8_SCALE_DTYPE,
1817
+ "w2_weight_scale": MXFP8_SCALE_DTYPE,
1818
+ }
1819
+ for name, expected_dtype in expected.items():
1820
+ actual = getattr(layer, name).dtype
1821
+ if actual != expected_dtype:
1822
+ raise ValueError(
1823
+ f"Expected {name} dtype {expected_dtype}, got {actual}."
1824
+ )
1825
+
1826
+ def _shuffle_weights_for_trtllm(self, layer: torch.nn.Module) -> None:
1827
+ """Shuffle weights and scales into FlashInfer TRTLLM MXFP8 layout."""
1828
+ from flashinfer import (
1829
+ reorder_rows_for_gated_act_gemm,
1830
+ shuffle_matrix_a,
1831
+ shuffle_matrix_sf_a,
1832
+ )
1833
+
1834
+ epilogue_tile_m = 128
1835
+ num_experts = layer.w13_weight.shape[0]
1836
+ is_gated = self.moe.is_act_and_mul
1837
+ intermediate_size_factor = 2 if is_gated else 1
1838
+
1839
+ w13_weight = layer.w13_weight.data
1840
+ w13_scale = layer.w13_weight_scale.data
1841
+ if is_gated:
1842
+ # FI TRTLLM gated kernels use W31 ordering. Model checkpoints store
1843
+ # gated projection as W13, so convert once before shuffling.
1844
+ w13_weight = swap_w13_to_w31(w13_weight)
1845
+ w13_scale = swap_w13_to_w31(w13_scale)
1846
+
1847
+ w13_weight_shuffled = []
1848
+ w2_weight_shuffled = []
1849
+ w13_scale_shuffled = []
1850
+ w2_scale_shuffled = []
1851
+ for i in range(num_experts):
1852
+ w13_i = w13_weight[i].reshape(
1853
+ intermediate_size_factor * layer.intermediate_size_per_partition, -1
1854
+ )
1855
+ w13_sf_i = w13_scale[i].reshape(
1856
+ intermediate_size_factor * layer.intermediate_size_per_partition, -1
1857
+ )
1858
+ if is_gated:
1859
+ # Reorder rows for gated activation layout expected by TRTLLM.
1860
+ w13_i = reorder_rows_for_gated_act_gemm(w13_i.clone())
1861
+ w13_sf_i = reorder_rows_for_gated_act_gemm(w13_sf_i.clone())
1862
+
1863
+ w13_shuffled_i = shuffle_matrix_a(w13_i.view(torch.uint8), epilogue_tile_m)
1864
+ w2_shuffled_i = shuffle_matrix_a(
1865
+ layer.w2_weight.data[i].view(torch.uint8), epilogue_tile_m
1866
+ )
1867
+ w13_weight_shuffled.append(
1868
+ w13_shuffled_i.contiguous().view(MXFP8_VALUE_DTYPE)
1869
+ )
1870
+ w2_weight_shuffled.append(
1871
+ w2_shuffled_i.contiguous().view(MXFP8_VALUE_DTYPE)
1872
+ )
1873
+ w13_sf_shuffled_i = shuffle_matrix_sf_a(
1874
+ w13_sf_i.view(torch.uint8).reshape(
1875
+ intermediate_size_factor * layer.intermediate_size_per_partition,
1876
+ -1,
1877
+ ),
1878
+ epilogue_tile_m,
1879
+ )
1880
+ w2_sf_shuffled_i = shuffle_matrix_sf_a(
1881
+ layer.w2_weight_scale.data[i]
1882
+ .view(torch.uint8)
1883
+ .reshape(layer.hidden_size, -1),
1884
+ epilogue_tile_m,
1885
+ )
1886
+ w13_scale_shuffled.append(
1887
+ w13_sf_shuffled_i.contiguous().view(MXFP8_SCALE_DTYPE)
1888
+ )
1889
+ w2_scale_shuffled.append(
1890
+ w2_sf_shuffled_i.contiguous().view(MXFP8_SCALE_DTYPE)
1891
+ )
1892
+
1893
+ replace_parameter(
1894
+ layer, "w13_weight", torch.stack(w13_weight_shuffled).contiguous()
1895
+ )
1896
+ replace_parameter(
1897
+ layer, "w2_weight", torch.stack(w2_weight_shuffled).contiguous()
1898
+ )
1899
+ replace_parameter(
1900
+ layer,
1901
+ "w13_weight_scale",
1902
+ torch.stack(w13_scale_shuffled).contiguous(),
1903
+ )
1904
+ replace_parameter(
1905
+ layer,
1906
+ "w2_weight_scale",
1907
+ torch.stack(w2_scale_shuffled).contiguous(),
1908
+ )
1909
+
1910
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1911
+ if getattr(layer, "_already_called_process_weights_after_loading", False):
1912
+ return
1913
+
1914
+ self._check_weight_dtypes(layer)
1915
+ self._shuffle_weights_for_trtllm(layer)
1916
+ layer._already_called_process_weights_after_loading = True
1917
+
1918
+ def maybe_make_prepare_finalize(
1919
+ self,
1920
+ routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
1921
+ ) -> mk.FusedMoEPrepareAndFinalizeModular | None:
1922
+ raise ValueError(
1923
+ f"{self.__class__.__name__} uses the new modular kernel initialization "
1924
+ "logic. This function should not be called."
1925
+ )
1926
+
1927
+ def select_gemm_impl(
1928
+ self,
1929
+ prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
1930
+ layer: torch.nn.Module,
1931
+ ) -> mk.FusedMoEExpertsModular:
1932
+ raise ValueError(
1933
+ f"{self.__class__.__name__} uses the new modular kernel initialization "
1934
+ "logic. This function should not be called."
1935
+ )
1936
+
1937
+ def get_fused_moe_quant_config(
1938
+ self, layer: torch.nn.Module
1939
+ ) -> FusedMoEQuantConfig | None:
1940
+ # TRTLLM MXFP8 path is monolithic and does not use modular kernel config.
1941
+ return None
1942
+
1943
+ @property
1944
+ def is_monolithic(self) -> bool:
1945
+ return self.mxfp8_backend == MxFp8MoeBackend.FLASHINFER_TRTLLM
1946
+
1947
+ def apply_monolithic(
1948
+ self,
1949
+ layer: FusedMoE,
1950
+ x: torch.Tensor,
1951
+ router_logits: torch.Tensor,
1952
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1953
+ from flashinfer.fused_moe.core import (
1954
+ ActivationType,
1955
+ Fp8QuantizationType,
1956
+ )
1957
+
1958
+ assert self.mxfp8_backend == MxFp8MoeBackend.FLASHINFER_TRTLLM
1959
+
1960
+ if layer.enable_eplb:
1961
+ raise NotImplementedError(
1962
+ "EPLB is not supported for FlashInfer TRTLLM MXFP8 MoE backend."
1963
+ )
1964
+
1965
+ supported_activations = [MoEActivation.SILU]
1966
+ if layer.activation not in supported_activations:
1967
+ raise NotImplementedError(
1968
+ "FlashInfer TRTLLM MXFP8 MoE supports only "
1969
+ f"{supported_activations}, got {layer.activation}."
1970
+ )
1971
+
1972
+ # Map vLLM MoEActivation to FlashInfer ActivationType.
1973
+ activation_map = {
1974
+ MoEActivation.SILU: ActivationType.Swiglu,
1975
+ MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
1976
+ }
1977
+ fi_activation_type: ActivationType = activation_map[layer.activation]
1978
+
1979
+ # DeepSeekV3 routing requires float32 logits; others expect bfloat16.
1980
+ if layer.routing_method_type == RoutingMethodType.DeepSeekV3:
1981
+ assert router_logits.dtype == torch.float32, (
1982
+ "DeepSeekV3 routing requires float32 router_logits, "
1983
+ f"got {router_logits.dtype}."
1984
+ )
1985
+ else:
1986
+ router_logits = router_logits.to(torch.bfloat16)
1987
+
1988
+ # Treat 0 as "unset" for compatibility with ungrouped routing configs.
1989
+ n_group = layer.num_expert_group or None
1990
+ topk_group = layer.topk_group or None
1991
+
1992
+ hidden_states_mxfp8, hidden_states_scale = mxfp8_e4m3_quantize(
1993
+ x,
1994
+ is_sf_swizzled_layout=False,
1995
+ )
1996
+
1997
+ kwargs: dict = dict(
1998
+ routing_logits=router_logits,
1999
+ routing_bias=layer.e_score_correction_bias,
2000
+ hidden_states=hidden_states_mxfp8,
2001
+ hidden_states_scale=hidden_states_scale,
2002
+ gemm1_weights=layer.w13_weight,
2003
+ gemm1_weights_scale=layer.w13_weight_scale,
2004
+ gemm2_weights=layer.w2_weight,
2005
+ gemm2_weights_scale=layer.w2_weight_scale,
2006
+ num_experts=layer.global_num_experts,
2007
+ top_k=layer.top_k,
2008
+ # Keep Optional semantics: FlashInfer expects None for non-grouped
2009
+ # routing (e.g. Qwen3 Renormalize), not 0.
2010
+ n_group=n_group,
2011
+ topk_group=topk_group,
2012
+ intermediate_size=layer.intermediate_size_per_partition,
2013
+ local_expert_offset=layer.ep_rank * layer.local_num_experts,
2014
+ local_num_experts=layer.local_num_experts,
2015
+ routed_scaling_factor=layer.routed_scaling_factor,
2016
+ routing_method_type=layer.routing_method_type,
2017
+ use_shuffled_weight=True,
2018
+ weight_layout=0,
2019
+ fp8_quantization_type=Fp8QuantizationType.MxFp8,
2020
+ )
2021
+
2022
+ if fi_activation_type != ActivationType.Swiglu:
2023
+ raise NotImplementedError(
2024
+ "FlashInfer TRTLLM MXFP8 MoE supports only Swiglu activation, "
2025
+ f"got {fi_activation_type}."
2026
+ )
2027
+
2028
+ return flashinfer_trtllm_fp8_block_scale_moe(**kwargs)
2029
+
2030
+ def apply(
2031
+ self,
2032
+ layer: FusedMoE,
2033
+ x: torch.Tensor,
2034
+ topk_weights: torch.Tensor,
2035
+ topk_ids: torch.Tensor,
2036
+ shared_experts_input: torch.Tensor | None,
2037
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
2038
+ assert not self.is_monolithic
2039
+ raise NotImplementedError(
2040
+ "Non-monolithic MXFP8 MoE path is not yet implemented."
2041
+ )
2042
+
2043
+
2044
# Register the method classes for ModelOptMxFp8Config
2045
ModelOptMxFp8Config.LinearMethodCls = ModelOptMxFp8LinearMethod
2046
+ModelOptMxFp8Config.FusedMoEMethodCls = ModelOptMxFp8FusedMoE
2047
ModelOptMxFp8Config.KVCacheMethodCls = ModelOptFp8KVCacheMethod
2048
2049