[GLM 5.3 Perf] Sorted-A TMA for FP8 MoE down projection on H200, 6.0%~ 7.6% E2E Throughput Improvement

tests/kernels/moe/test_block_fp8.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from unittest.mock import patch
+
 import pytest
 import torch
 
@@ -176,6 +178,55 @@ def setup_cuda():
     torch.set_default_device("cuda")
 
 
+@pytest.mark.skipif(
+    "H200" not in torch.cuda.get_device_name(),
+    reason="The sorted-A TMA path is tuned for H200",
+)
+def test_sorted_a_tma_moe_matches_route_major(workspace_init):
+    """The H200 sorted path preserves routed FP8 block outputs exactly."""
+    torch.manual_seed(123)
+    m, experts, hidden, intermediate, topk = 1024, 256, 4096, 256, 8
+    x = (torch.randn((m, hidden)) / 10).to(torch.bfloat16)
+    w1, w2, quant_config = make_test_quant_config(
+        experts,
+        intermediate,
+        hidden,
+        torch.bfloat16,
+        torch.float8_e4m3fn,
+        block_shape=[128, 128],
+    )
+    moe_config = make_dummy_moe_config(
+        num_experts=experts,
+        experts_per_token=topk,
+        hidden_dim=hidden,
+        intermediate_size=intermediate,
+        max_num_tokens=m,
+    )
+    kernel = modular_triton_fused_moe(moe_config, quant_config)
+    triton_experts = kernel.fused_experts
+    assert isinstance(triton_experts, triton_moe.TritonExperts)
+    topk_ids = torch.randint(experts, (m, topk))
+    topk_weights = torch.softmax(torch.randn((m, topk)), dim=-1)
+    kwargs = dict(
+        activation=MoEActivation.SILU,
+        global_num_experts=experts,
+        expert_map=None,
+        apply_router_weight_on_input=False,
+    )
+    with set_current_vllm_config(vllm_config):
+        triton_experts._enable_sorted_a_tma = False
+        expected = kernel.apply(x, w1, w2, topk_weights, topk_ids, **kwargs)
+        triton_experts._enable_sorted_a_tma = True
+        with patch.object(
+            triton_moe,
+            "invoke_fused_moe_triton_kernel",
+            wraps=triton_moe.invoke_fused_moe_triton_kernel,
+        ) as calls:
+            actual = kernel.apply(x, w1, w2, topk_weights, topk_ids, **kwargs)
+    assert any(call.kwargs.get("A_is_sorted") for call in calls.call_args_list)
+    assert torch.equal(expected, actual)
+
+
 @pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS)
 @pytest.mark.parametrize("E", E)
 @pytest.mark.parametrize("topk", TOP_KS)

vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128]_down.json

@@ -0,0 +1,22 @@
+{
+    "512": {"USE_TMA": false},
+    "1024": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "3072": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "4096": {"USE_TMA": false}
+}

vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128]_down.json

@@ -0,0 +1,23 @@
+{
+    "512": {"USE_TMA": false},
+    "1023": {"USE_TMA": false},
+    "1024": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "3072": {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 128,
+        "BLOCK_SIZE_K": 128,
+        "GROUP_SIZE_M": 64,
+        "num_warps": 4,
+        "num_stages": 3,
+        "USE_TMA": true
+    },
+    "4096": {"USE_TMA": false}
+}

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

@@ -21,6 +21,7 @@
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     _prepare_expert_assignment,
     _weights_host_resident,
+    get_moe_configs,
     invoke_fused_moe_triton_kernel,
     invoke_fused_moe_wna16_triton_kernel,
     try_get_optimal_moe_config,
@@ -60,10 +61,45 @@
     kInt8StaticTensorSym,
 )
 from vllm.platforms import current_platform
-from vllm.triton_utils import tl
+from vllm.triton_utils import tl, triton
 from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
 
 
+@triton.jit
+def _silu_mul_quant_sorted_kernel(
+    input_ptr,
+    output_ptr,
+    scale_ptr,
+    num_tokens_post_padded_ptr,
+    intermediate: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    GROUP_K: tl.constexpr,
+):
+    row_start = tl.program_id(0) * BLOCK_M
+    if row_start >= tl.load(num_tokens_post_padded_ptr):
+        return
+    rows = row_start + tl.arange(0, BLOCK_M)
+    cols = tl.program_id(1) * GROUP_K + tl.arange(0, GROUP_K)
+    gate = tl.load(input_ptr + rows[:, None] * 2 * intermediate + cols[None, :])
+    up = tl.load(
+        input_ptr + rows[:, None] * 2 * intermediate + intermediate + cols[None, :]
+    )
+    gate = gate.to(tl.float32)
+    up = up.to(tl.float32)
+    # match silu_and_mul: round SiLU before multiplying the other half.
+    silu = tl.div_rn(gate, 1.0 + tl.extra.cuda.libdevice.exp(-gate))
+    silu = silu.to(input_ptr.dtype.element_ty).to(tl.float32)
+    result = (silu * up).to(input_ptr.dtype.element_ty).to(tl.float32)
+    scale_raw = tl.maximum(tl.max(tl.abs(result), 1), 1e-10) * (1.0 / 448.0)
+    scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw)))
+    quantized = tl.minimum(tl.maximum(tl.div_rn(result, scale[:, None]), -448.0), 448.0)
+    tl.store(
+        output_ptr + rows[:, None] * intermediate + cols[None, :],
+        quantized.to(tl.float8e4nv),
+    )
+    tl.store(scale_ptr + rows * (intermediate // GROUP_K) + tl.program_id(1), scale)
+
+
 class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
     """Triton-based fused MoE expert implementation."""
 
@@ -100,6 +136,7 @@ def __init__(
         # Whether quantized MOE runs natively, or through
         # higher-precision + activation QDQ.
         self.quantization_emulation = False
+        self._enable_sorted_a_tma = True
         super().__init__(moe_config, quant_config)
 
     @staticmethod
@@ -394,6 +431,56 @@ def apply(
             )
         )
 
+        from vllm.model_executor.layers.fused_moe import get_config
+
+        down_configs = (
+            get_moe_configs(
+                E,
+                cache2_dim,
+                self.quant_config.config_name(hidden_states.dtype),
+                128,
+                128,
+                down_moe=True,
+            )
+            if self.quant_config.use_fp8_w8a8
+            and self.block_shape == [128, 128]
+            and not get_config()
+            else None
+        )
+        tma_down_config = (
+            dict(down_configs[min(down_configs, key=lambda m: abs(m - num_tokens))])
+            if down_configs
+            else None
+        )
+
+        use_sorted_a = bool(
+            tma_down_config
+            and tma_down_config.pop("USE_TMA", False)
+            and self._enable_sorted_a_tma
+            and activation == MoEActivation.SILU
+            and self.activation_config.clamp_limit is None
+            and self._lora_context is None
+            and not hooked
+            and not self.quantization_emulation
+            and is_deep_gemm_e8m0_used()
+            and a2_scale is None
+            and sorted_token_ids is not None
+            and expert_map is None
+            and workspace2.dtype == torch.bfloat16
+            and current_platform.is_cuda()
+            and current_platform.is_device_capability(90)
+            and config["BLOCK_SIZE_M"] % 16 == 0
+            and w2.is_contiguous()
+            and w2.size(1) >= 4096
+            and w2.size(1) % tma_down_config["BLOCK_SIZE_N"] == 0
+            and cache2_dim % 128 == 0
+        )
+        if use_sorted_a:
+            assert sorted_token_ids is not None
+            num_sorted_rows = sorted_token_ids.numel()
+            # GEMM fills padding inside active blocks; quantization skips the tail.
+            intermediate_cache1 = workspace2.new_empty((num_sorted_rows, 1, N))
+
         # LoRA w13: applied to intermediate_cache1 before activation. When
         # the LoRA layer requested a dual-stream schedule, we run base w13
         # GEMM on the default stream and the LoRA fast-path on aux_stream;
@@ -455,6 +542,7 @@ def _base_w13_fn():
                 per_channel_quant=self.per_act_token_quant,
                 block_shape=self.block_shape,
                 B_bias=self.w1_bias,
+                C_is_sorted=use_sorted_a,
             )
 
         if lora_context is not None and lora_context.aux_stream is not None:
@@ -522,7 +610,27 @@ def _lora_w13_fn():
         # group_size=128, no LoRA requiring the BF16 intermediate).
         # The fused kernel has no clamp parameter, so a configured
         # SwiGLU clamp limit falls through to the unfused path.
-        if (
+        if use_sorted_a:
+            qintermediate_cache2 = intermediate_cache1.new_empty(
+                (num_sorted_rows, cache2_dim),
+                dtype=current_platform.fp8_dtype(),
+            )
+            a2q_scale = qintermediate_cache2.new_empty(
+                (num_sorted_rows, cache2_dim // 128),
+                dtype=torch.float32,
+            )
+            _silu_mul_quant_sorted_kernel[
+                (triton.cdiv(num_sorted_rows, 16), cache2_dim // 128)
+            ](
+                intermediate_cache1,
+                qintermediate_cache2,
+                a2q_scale,
+                num_tokens_post_padded,
+                cache2_dim,
+                BLOCK_M=16,
+                GROUP_K=128,
+            )
+        elif (
             activation == MoEActivation.SILU
             and self.activation_config.clamp_limit is None
             and self.quant_config.use_fp8_w8a8
@@ -555,6 +663,14 @@ def _lora_w13_fn():
         # the w13 pair: base GEMM on default stream, LoRA delta on aux,
         # join via .add_() into intermediate_cache3.
         def _base_w2_fn():
+            if use_sorted_a:
+                assert tma_down_config is not None
+                down_config = {
+                    **tma_down_config,
+                    "BLOCK_SIZE_M": config["BLOCK_SIZE_M"],
+                }
+            else:
+                down_config = config
             invoke_fused_moe_triton_kernel(
                 qintermediate_cache2,
                 w2,
@@ -568,7 +684,7 @@ def _base_w2_fn():
                 not hooked
                 and self._apply_router_weight_in_gemm2(apply_router_weight_on_input),
                 1,
-                config,
+                down_config,
                 compute_type=compute_type,
                 use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
                 use_int8_w8a8=self.quant_config.use_int8_w8a8,
@@ -577,6 +693,7 @@ def _base_w2_fn():
                 per_channel_quant=self.per_act_token_quant,
                 block_shape=self.block_shape,
                 B_bias=self.w2_bias,
+                A_is_sorted=use_sorted_a,
             )
 
         if lora_context is not None and lora_context.aux_stream is not None: