iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[MoE] Shape-based padding for FlashInfer TRT-LLM MXFP8 experts (gated and non-gated) (#61293)

Browse Source 
Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
This commit is contained in:
Jhao-Ting Chen
2026-10-15 13:44:11 -04:00
committed by GitHub
co-authored by Claude Opus 5.5 Robert Shaw
parent 7cd7982dfd
commit 4f49b55c1f
2 changed files with 216 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_flashinfer.py
+125 -3
View File
@@ -52,6 +52,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
)
from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
dequant_mxfp8_to_bf16,
mxfp8_e4m3_quantize,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -65,7 +66,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
)
from vllm.model_executor.models.llama4 import Llama4MoE
from vllm.platforms import current_platform
from vllm.utils.math_utils import next_power_of_2
from vllm.utils.math_utils import next_power_of_2, round_up
from vllm.utils.torch_utils import set_random_seed
try:
@@ -788,7 +789,7 @@ def test_trtllm_fp8_swiglu_clamp_support(
def _make_mxfp8_moe_weights(
e: int, n: int, k: int
e: int, n: int, k: int, is_gated: bool = True
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
def quantize(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
q, s = zip(
@@ -796,8 +797,9 @@ def _make_mxfp8_moe_weights(
)
return torch.stack(q), torch.stack(s)
w1_rows = 2 * n if is_gated else n
w1, w1_scale = quantize(
torch.randn(e, 2 * n, k, device="cuda", dtype=torch.bfloat16) / 10
torch.randn(e, w1_rows, k, device="cuda", dtype=torch.bfloat16) / 10
)
w2, w2_scale = quantize(
torch.randn(e, k, n, device="cuda", dtype=torch.bfloat16) / 10
@@ -986,3 +988,123 @@ def test_trtllm_mxfp8_minimax2_routing_applies_routed_scale(m: int, workspace_in
)
torch.testing.assert_close(fused, reference, atol=2e-2, rtol=2e-2)
@pytest.mark.parametrize(
("activation", "n", "k"),
[
(MoEActivation.SILU, 1024, 1024),
(MoEActivation.SILU, 928, 1024),
(MoEActivation.SILU, 1024, 2688),
(MoEActivation.RELU2_NO_MUL, 1024, 1024),
# Nemotron-3.5 Lightning's non-gated MoE: intermediate 1856, hidden 2688.
(MoEActivation.RELU2_NO_MUL, 1856, 2688),
],
)
@pytest.mark.parametrize("m", [1, 32])
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="Requires TRTLLM-Gen FP8 MoE (SM100)",
)
def test_trtllm_mxfp8_moe_pads_intermediate(
m: int, activation: MoEActivation, n: int, k: int, workspace_init
):
"""TRTLLM-Gen MXFP8 experts need the intermediate size to be a multiple
of 128 and the hidden size a multiple of 256. Misaligned experts, gated or
not, are padded at load and, with the activations padded and the output
cropped as the MoE runner does, must still match the unquantized reference
on the dequantized weights."""
e, topk = 8, 2
is_gated = activation.is_gated
block_shape = [1, 32]
set_random_seed(7)
with set_current_vllm_config(vllm_config):
w1, w2, w1_scale, w2_scale = _make_mxfp8_moe_weights(e, n, k, is_gated)
w1_ref = dequant_mxfp8_to_bf16(w1, w1_scale)
w2_ref = dequant_mxfp8_to_bf16(w2, w2_scale)
layer = SimpleNamespace(
weight_block_size=block_shape,
moe_config=SimpleNamespace(
is_act_and_mul=is_gated, intermediate_size_per_partition=n
),
activation=activation,
)
w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format(
Fp8MoeBackend.FLASHINFER_TRTLLM,
layer,
w1,
w2,
w1_scale,
w2_scale,
w13_input_scale=None,
w2_input_scale=None,
)
padded_n = layer.moe_config.intermediate_size_per_partition
padded_k = layer.moe_config.hidden_dim
assert padded_n == round_up(n, 128)
assert padded_k == round_up(k, 256)
quant_config = make_fp8_moe_quant_config(
Fp8MoeBackend.FLASHINFER_TRTLLM,
w1_scale,
w2_scale,
a1_scale=None,
a2_scale=None,
block_shape=block_shape,
)
moe_config = FusedMoEConfig(
num_experts=e,
experts_per_token=topk,
hidden_dim=padded_k,
intermediate_size=padded_n,
num_local_experts=e,
num_logical_experts=e,
activation=activation,
device="cuda",
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
in_dtype=torch.bfloat16,
routing_method=RoutingMethodType.Renormalize,
max_num_tokens=next_power_of_2(m),
)
# The monolithic experts pass the activation type to the kernel; the
# modular ones always run it gated.
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
use_monolithic=True,
),
TrtLlmFp8ExpertsMonolithic(
moe_config=moe_config, quant_config=quant_config
),
)
a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
logits = torch.randn((m, e), device="cuda", dtype=torch.float32)
output = kernel.apply_monolithic(
hidden_states=torch.nn.functional.pad(a, (0, padded_k - k)),
w1=w1,
w2=w2,
router_logits=logits,
activation=activation,
global_num_experts=e,
expert_map=None,
apply_router_weight_on_input=False,
routed_scaling_factor=1.0,
routing_sink=None,
)[:, :k]
topk_weights, topk_ids, _ = fused_topk(a, logits, topk, renormalize=True)
reference = fused_experts(
a,
w1_ref,
w2_ref,
topk_weights=topk_weights,
topk_ids=topk_ids,
activation=activation,
global_num_experts=e,
expert_map=None,
)
check_accuracy(ref_output=reference, actual_output=output)
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+91 -2
View File
@@ -379,6 +379,79 @@ def align_moe_weights_for_fi(
return padded_w13, padded_w2, padded_intermediate
def align_mxfp8_moe_weights_for_fi_trtllm(
w13: torch.Tensor,
w2: torch.Tensor,
w13_scale: torch.Tensor,
w2_scale: torch.Tensor,
is_act_and_mul: bool,
min_alignment: int,
hidden_alignment: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int, int]:
"""align_moe_weights_for_fi for MXFP8 weights and their UE8M0 block scales
for the FlashInfer TRT-LLM MoE kernel, which also pads the hidden size to a
multiple of hidden_alignment.
Padded weights are zero and padded scales are 127 (UE8M0 for 1.0), so the
padded rows and columns contribute nothing. Scales may be 3D
[E, rows, cols] or flattened [E, rows * cols]; their layout is kept.
Returns the padded intermediate and hidden sizes.
"""
num_experts, hidden_size, intermediate = w2.shape
padded_hidden = round_up(hidden_size, hidden_alignment)
w13, w2, padded_intermediate = align_moe_weights_for_fi(
w13, w2, is_act_and_mul, min_alignment
)
if padded_intermediate == intermediate and padded_hidden == hidden_size:
return w13, w2, w13_scale, w2_scale, intermediate, hidden_size
if padded_hidden != hidden_size:
logger.info_once(
"Padding hidden size from %d to %d for MXFP8 TRTLLM MoE kernel.",
hidden_size,
padded_hidden,
)
padded_w13 = w13.new_zeros((num_experts, w13.shape[1], padded_hidden))
padded_w13[:, :, :hidden_size] = w13
padded_w2 = w2.new_zeros((num_experts, padded_hidden, w2.shape[2]))
padded_w2[:, :hidden_size] = w2
w13, w2 = padded_w13, padded_w2
up_mult = 2 if is_act_and_mul else 1
w13_scale_3d = w13_scale.reshape(num_experts, up_mult * intermediate, -1)
w2_scale_3d = w2_scale.reshape(num_experts, hidden_size, -1)
block = hidden_size // w13_scale_3d.shape[2]
assert padded_hidden % block == 0 and padded_intermediate % block == 0
# Pad each of the gate and up halves (or the single non-gated block) on its
# own, like the weights.
padded_w13_scale = w13_scale_3d.new_full(
(num_experts, up_mult * padded_intermediate, padded_hidden // block), 127
)
for i in range(up_mult):
dst = i * padded_intermediate
padded_w13_scale[:, dst : dst + intermediate, : w13_scale_3d.shape[2]] = (
w13_scale_3d[:, i * intermediate : (i + 1) * intermediate]
)
padded_w2_scale = w2_scale_3d.new_full(
(num_experts, padded_hidden, padded_intermediate // block), 127
)
padded_w2_scale[:, :hidden_size, : w2_scale_3d.shape[2]] = w2_scale_3d
if w13_scale.ndim == 2:
padded_w13_scale = padded_w13_scale.reshape(num_experts, -1)
if w2_scale.ndim == 2:
padded_w2_scale = padded_w2_scale.reshape(num_experts, -1)
return (
w13,
w2,
padded_w13_scale,
padded_w2_scale,
padded_intermediate,
padded_hidden,
)
def _shuffle_deepseek_fp8_moe_weights(
w13: torch.Tensor,
w2: torch.Tensor,
@@ -486,8 +559,9 @@ def prepare_fp8_moe_layer_for_fi(
Note that for trtllm we update the model state dict
with the scale format needed for these kernels.
Note that for per-tensor, we update the layer's
intermediate size if the weights needed padding.
Note that for per-tensor and for MXFP8 TRT-LLM, we update the layer's
intermediate size (and for MXFP8 TRT-LLM its hidden size) if the weights
needed padding.
"""
assert hasattr(layer.moe_config, "is_act_and_mul")
block_quant = (
@@ -499,6 +573,21 @@ def prepare_fp8_moe_layer_for_fi(
# MXFP8 TRT-LLM requires W31 swap + reorder + shuffle.
if is_mxfp8 and is_trtllm:
# The kernel needs intermediate % 128 == 0 and hidden % 256 == 0. The
# MoE runner pads the activations to the padded hidden size.
w13, w2, w13_scale, w2_scale, new_intermediate, new_hidden = (
align_mxfp8_moe_weights_for_fi_trtllm(
w13,
w2,
w13_scale,
w2_scale,
layer.moe_config.is_act_and_mul,
min_alignment=128,
hidden_alignment=256,
)
)
layer.moe_config.intermediate_size_per_partition = new_intermediate
layer.moe_config.hidden_dim = new_hidden
# FlashInfer TRT-LLM SwiGLU expects [up; gate] but vLLM stores
# [gate; up]. Swap both weights and scales before interleaving.
if layer.moe_config.is_act_and_mul: