@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1af6f78ae5

Fork 0

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 5, 12:15 UTC

[Perf] Change Trtllm fp8 MoE to use Shuffled Weights and BlockMajorK Layout (#38993)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

Wei Zhao committed Apr 5, 2026 at 10:54 UTC 1af6f78ae5a1bd3d70f32d47fe4901bda5c97fdd

3 files changed +72 -13

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

+28 -5

@@ -112,6 +112,24 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

112 ]

113 return (weight_key, activation_key) in SUPPORTED_W_A

114

115 + def moe_problem_size(

116 + self,

117 + a1: torch.Tensor,

118 + w1: torch.Tensor,

119 + w2: torch.Tensor,

120 + topk_ids: torch.Tensor,

121 + ) -> tuple[int, int, int, int, int]:

122 + """Override to handle 4D BlockMajorK weights (E, K/bk, Mn, bk)."""

123 + if w1.dim() == 4:

124 + # BlockMajorK: (E, K/bk, Mn, bk)

125 + E = w1.shape[0]

126 + N = w1.shape[2]

127 + K = a1.size(-1)

128 + M = a1.size(0) if a1.dim() == 2 else a1.size(1)

129 + topk = topk_ids.size(1)

130 + return E, M, N, K, topk

131 + return super().moe_problem_size(a1, w1, w2, topk_ids)

132 +

133 def workspace_shapes(

134 self,

135 M: int,

@@ -152,7 +170,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

170 apply_router_weight_on_input: bool,

171 ):

172 import flashinfer

155 - from flashinfer.fused_moe import Fp8QuantizationType

173 + from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout

174

175 # Pack topk ids and weights into format expected by the kernel.

176 packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

@@ -170,10 +188,12 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

188 if is_mxfp8:

189 fp8_quant_type = Fp8QuantizationType.MxFp8

190 use_shuffled_weight = True

191 + weight_layout = WeightLayout.MajorK

192 hidden_states_scale = a1q_scale

193 else:

194 fp8_quant_type = Fp8QuantizationType.DeepSeekFp8

176 - use_shuffled_weight = False

195 + use_shuffled_weight = True

196 + weight_layout = WeightLayout.BlockMajorK

197 hidden_states_scale = a1q_scale.t().contiguous()

198

199 # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the

@@ -199,7 +219,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

219 routed_scaling_factor=None,

220 routing_method_type=1,

221 use_shuffled_weight=use_shuffled_weight,

202 - weight_layout=0,

222 + weight_layout=weight_layout,

223 fp8_quantization_type=fp8_quant_type,

224 # output=output,

225 )

@@ -322,7 +342,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit

342 topk_group: int | None = None,

343 ) -> torch.Tensor:

344 import flashinfer

325 - from flashinfer.fused_moe import Fp8QuantizationType

345 + from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout

346

347 assert not apply_router_weight_on_input

348 assert activation == MoEActivation.SILU

@@ -342,10 +362,12 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit

362 if is_mxfp8:

363 fp8_quant_type = Fp8QuantizationType.MxFp8

364 use_shuffled_weight = True

365 + weight_layout = WeightLayout.MajorK

366 hidden_states_scale = a1q_scale

367 else:

368 fp8_quant_type = Fp8QuantizationType.DeepSeekFp8

348 - use_shuffled_weight = False

369 + use_shuffled_weight = True

370 + weight_layout = WeightLayout.BlockMajorK

371 hidden_states_scale = a1q_scale.t().contiguous()

372

373 return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(

@@ -367,6 +389,7 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit

389 routed_scaling_factor=routed_scaling_factor,

390 routing_method_type=self.routing_method_type,

391 use_shuffled_weight=use_shuffled_weight,

392 + weight_layout=weight_layout,

393 fp8_quantization_type=fp8_quant_type,

394 )

395

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

+38

@@ -305,6 +305,39 @@ def align_fp8_moe_weights_for_fi(

305 return padded_w13, padded_w2, padded_intermediate

306

307

308 +def _shuffle_deepseek_fp8_moe_weights(

309 + w13: torch.Tensor,

310 + w2: torch.Tensor,

311 +) -> tuple[torch.Tensor, torch.Tensor]:

312 + """Preprocess DeepSeek FP8 block-scale weights for the FlashInfer TRT-LLM

313 + kernel using the shuffle + BlockMajorK layout variant.

314 +

315 + Returns 4D weight tensors in BlockMajorK layout

316 + (E, K/block_k, Mn, block_k)

317 + """

318 + from flashinfer import shuffle_matrix_a

319 + from flashinfer.fused_moe import convert_to_block_layout

320 +

321 + epilogue_tile_m = 64

322 + block_k = 128

323 + num_experts = w13.shape[0]

324 +

325 + w13_shuffled: list[torch.Tensor] = []

326 + w2_shuffled: list[torch.Tensor] = []

327 + for i in range(num_experts):

328 + t13 = shuffle_matrix_a(w13[i].view(torch.uint8), epilogue_tile_m)

329 + t13 = convert_to_block_layout(t13, block_k)

330 + w13_shuffled.append(t13)

331 +

332 + t2 = shuffle_matrix_a(w2[i].view(torch.uint8), epilogue_tile_m)

333 + t2 = convert_to_block_layout(t2, block_k)

334 + w2_shuffled.append(t2)

335 +

336 + w13_out = torch.stack(w13_shuffled).view(torch.float8_e4m3fn)

337 + w2_out = torch.stack(w2_shuffled).view(torch.float8_e4m3fn)

338 + return w13_out, w2_out

339 +

340 +

341 def _shuffle_mxfp8_moe_weights(

342 w13: torch.Tensor,

343 w2: torch.Tensor,

@@ -405,6 +438,7 @@ def prepare_fp8_moe_layer_for_fi(

438 hasattr(layer, "weight_block_size") and layer.weight_block_size is not None

439 )

440 is_mxfp8 = block_quant and w13_scale.dtype == torch.uint8

441 + is_deepseek_fp8 = block_quant and not is_mxfp8

442 is_gated = layer.activation.is_gated

443

444 # MXFP8 TRT-LLM requires W31 swap + reorder + shuffle.

@@ -447,6 +481,10 @@ def prepare_fp8_moe_layer_for_fi(

481 if block_quant:

482 w13_scale = swap_w13_to_w31(w13_scale)

483

484 + # DeepSeekFp8 TRT-LLM: shuffle weights into BlockMajorK layout.

485 + if is_deepseek_fp8 and is_trtllm:

486 + w13, w2 = _shuffle_deepseek_fp8_moe_weights(w13, w2)

487 +

488 # FI TRT-LLM FP8 per-tensor MoE kernel requires weight shuffle

489 # and registration of alpha scales.

490 if is_trtllm and not block_quant:

vllm/model_executor/warmup/deep_gemm_warmup.py

+6 -8

@@ -13,7 +13,7 @@ import vllm.envs as envs

13 from vllm.distributed.parallel_state import get_dp_group, is_global_first_rank

14 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts

15 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import compute_aligned_M

16 -from vllm.model_executor.layers.fused_moe.layer import FusedMoE, FusedMoEModularMethod

16 +from vllm.model_executor.layers.fused_moe.layer import FusedMoE

17 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (

18 TritonOrDeepGemmExperts,

19 )

@@ -168,14 +168,12 @@ def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:

168 ):

169 return False

170

171 - if not isinstance(module.quant_method, FusedMoEModularMethod):

172 - # modular kernels could invoke deep_gemm_moe_fp8

173 - return True

171 + moe_kernel = getattr(module.quant_method, "moe_kernel", None)

172 + if moe_kernel is None:

173 + return False

174

175 - # Further check if the ModularKernel implementation uses the DeepGemmExperts

176 - return isinstance(

177 - module.quant_method.moe_kernel, (DeepGemmExperts, TritonOrDeepGemmExperts)

178 - )

175 + fused_experts = moe_kernel.impl.fused_experts

176 + return isinstance(fused_experts, (DeepGemmExperts, TritonOrDeepGemmExperts))

177

178

179 FP8_GEMM_NT_WARMUP_CACHE: set[torch.Size] = set()