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

[Perf][MoE] Skip the TRTLLM fp8-block activation-scale transpose copy (#61773)

Browse Source 
Signed-off-by: Thien Tran <gau.nernst@yahoo.com.sg>
This commit is contained in:
Thien Tran
2026-10-15 12:09:07 -07:00
committed by GitHub
parent 4d58fa254a
commit 82b9ccaa79
5 changed files with 77 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_flashinfer.py
+54 -12
View File
@@ -808,16 +808,20 @@ def _make_mxfp8_moe_weights(
@pytest.mark.parametrize("block_shape", [[128, 128], [1, 32]], ids=["block", "mxfp8"])
@pytest.mark.parametrize("transposed", [True, False], ids=["transposed", "row-major"])
@pytest.mark.parametrize("m", [1, 16])
@pytest.mark.skipif(
not current_platform.is_device_capability_family(100),
reason="Requires TRTLLM-Gen FP8 MoE (SM100)",
)
def test_trtllm_fp8_block_moe_deferred_finalize(
m: int, block_shape: list[int], workspace_init
m: int, block_shape: list[int], transposed: bool, workspace_init
):
"""TRTLLM-Gen block-FP8 and MXFP8 modular experts can leave the top-k
finalize to the caller."""
finalize to the caller. For fp8-block, ``transposed`` toggles the oracle's
is_scale_transposed column-major activation-scale layout."""
if block_shape != [128, 128] and transposed:
pytest.skip("column-major scales apply to fp8 [128, 128] block only")
e, topk, n, k = 32, 4, 1024, 1024
set_random_seed(7)
with set_current_vllm_config(vllm_config):
@@ -851,6 +855,9 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
a2_scale=None,
block_shape=block_shape,
)
is_block = block_shape == [128, 128]
if is_block and not transposed:
quant_config = replace(quant_config, is_scale_transposed=False)
# One rank of a TP group, which deferral needs.
moe_config = FusedMoEConfig(
num_experts=e,
@@ -868,19 +875,27 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
routing_method=RoutingMethodType.TopK,
max_num_tokens=next_power_of_2(m),
)
kernel = mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config, quant_config=quant_config, allow_new_interface=True
),
TrtLlmFp8ExpertsModular(moe_config=moe_config, quant_config=quant_config),
)
def make_kernel(quant_config: FusedMoEQuantConfig) -> mk.FusedMoEKernel:
return mk.FusedMoEKernel(
maybe_make_prepare_finalize(
moe=moe_config,
quant_config=quant_config,
allow_new_interface=True,
),
TrtLlmFp8ExpertsModular(
moe_config=moe_config, quant_config=quant_config
),
)
kernel = make_kernel(quant_config)
a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
check_deferred_moe_finalize(
moe_config,
lambda: kernel.apply(
def apply(k_: mk.FusedMoEKernel) -> torch.Tensor:
return k_.apply(
hidden_states=a,
w1=w1,
w2=w2,
@@ -890,7 +905,34 @@ def test_trtllm_fp8_block_moe_deferred_finalize(
global_num_experts=e,
expert_map=None,
apply_router_weight_on_input=False,
),
)
if is_block and transposed:
# The oracle gates column-major scales on TRTLLM fp8-block.
assert quant_config.is_scale_transposed
assert not make_fp8_moe_quant_config(
Fp8MoeBackend.DEEPGEMM,
w1_scale,
w2_scale,
a1_scale=None,
a2_scale=None,
block_shape=block_shape,
).is_scale_transposed
# no_dp_ep emits column-major [M, G] scales (physical [G, M]), so
# the experts' a1q_scale.t().contiguous() is a no-op.
_, a1q_scale, _, _, _ = kernel.prepare_finalize.prepare(
a, topk_weights, topk_ids, e, None, False, quant_config
)
assert a1q_scale.t().is_contiguous()
# Bitwise-equal to running with row-major scales.
row_major = replace(quant_config, is_scale_transposed=False)
torch.testing.assert_close(
apply(kernel), apply(make_kernel(row_major)), atol=0, rtol=0
)
check_deferred_moe_finalize(
moe_config,
lambda: apply(kernel),
router_weights=topk_weights,
)
vllm/model_executor/layers/fused_moe/config.py
+9
View File
@@ -249,6 +249,8 @@ class FusedMoEQuantConfig:
_w1: FusedMoEQuantDesc
_w2: FusedMoEQuantDesc
is_scale_swizzled: bool = True
# Emit block-quant activation scales column-major (physical [G, M] layout).
is_scale_transposed: bool = False
# MXFP4-specific TRTLLM parameters for SwiGLU activation clamping.
# These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit
@@ -503,6 +505,7 @@ class FusedMoEQuantConfig:
w2_zp: torch.Tensor | None = None,
weight_dtype: torch.dtype | str | None = None,
is_scale_swizzled: bool = True,
is_scale_transposed: bool = False,
gemm1_alpha: float | None = None,
gemm1_beta: float | None = None,
gemm1_clamp_limit: float | None = None,
@@ -537,6 +540,9 @@ class FusedMoEQuantConfig:
- is_scale_swizzled: Whether the activation scale-factor layout is
swizzled. Pass through to the underlying quantization kernel for
dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True.
- is_scale_transposed: Whether block-quantized activation scales are
emitted column-major (physical [num_groups, num_tokens] layout) by
the quantize step. Defaults to False.
- gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter.
- gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter.
- gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit.
@@ -573,6 +579,7 @@ class FusedMoEQuantConfig:
weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
),
is_scale_swizzled=is_scale_swizzled,
is_scale_transposed=is_scale_transposed,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=gemm1_clamp_limit,
@@ -597,6 +604,7 @@ def fp8_w8a8_moe_quant_config(
a2_gscale: torch.Tensor | None = None,
g1_alphas: torch.Tensor | None = None,
g2_alphas: torch.Tensor | None = None,
is_scale_transposed: bool = False,
gemm1_alpha: float | None = None,
gemm1_beta: float | None = None,
gemm1_clamp_limit: float | None = None,
@@ -617,6 +625,7 @@ def fp8_w8a8_moe_quant_config(
per_act_token_quant=per_act_token_quant,
per_out_ch_quant=per_out_ch_quant,
block_shape=block_shape,
is_scale_transposed=is_scale_transposed,
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=gemm1_clamp_limit,
vllm/model_executor/layers/fused_moe/oracle/fp8.py
+5
View File
@@ -841,6 +841,11 @@ def make_fp8_moe_quant_config(
block_shape=block_shape,
per_act_token_quant=per_act_token_quant,
per_out_ch_quant=per_out_ch_quant,
# TRTLLM fp8-block consumes activation scales as [K/128, M]; emit them
# transposed from the quant kernel instead of copying afterwards.
is_scale_transposed=(
fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM and block_shape == [128, 128]
),
gemm1_alpha=gemm1_alpha,
gemm1_beta=gemm1_beta,
gemm1_clamp_limit=swiglu_limit,
vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py
+1
View File
@@ -31,6 +31,7 @@ def _quantize_input(
per_act_token_quant=quant_config.per_act_token_quant,
block_shape=quant_config.block_shape,
is_scale_swizzled=quant_config.is_scale_swizzled,
is_scale_transposed=quant_config.is_scale_transposed,
mx_alignment=quant_config.mx_alignment,
)
vllm/model_executor/layers/fused_moe/utils.py
+8 -2
View File
@@ -325,6 +325,7 @@ def _fp8_quantize(
A_scale: torch.Tensor | None,
per_act_token: bool,
block_shape: list[int] | None = None,
is_scale_transposed: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Perform fp8 quantization on the inputs. If a block_shape
is provided, the output will be blocked.
@@ -339,7 +340,9 @@ def _fp8_quantize(
assert not per_act_token
assert len(block_shape) == 2
_, block_k = block_shape[0], block_shape[1]
A, A_scale = per_token_group_quant_fp8(A, block_k)
A, A_scale = per_token_group_quant_fp8(
A, block_k, column_major_scales=is_scale_transposed
)
assert cdiv(A.size(-1), block_k) == A_scale.size(-1)
return A, A_scale
@@ -462,6 +465,7 @@ def moe_kernel_quantize_input(
per_act_token_quant: bool,
block_shape: list[int] | None = None,
is_scale_swizzled: bool = True,
is_scale_transposed: bool = False,
quantization_emulation: bool = False,
mx_alignment: int = 0,
) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -469,7 +473,9 @@ def moe_kernel_quantize_input(
if quantization_emulation:
return _fp8_quantize_dequantize(A, A_scale)
else:
return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
return _fp8_quantize(
A, A_scale, per_act_token_quant, block_shape, is_scale_transposed
)
elif quant_dtype == torch.int8:
if quantization_emulation:
raise NotImplementedError(