[Perf] Use TMA for MOE up/down proj, 3.1~3.4% E2E Throughput Improvement for GLM5.3 (#63820)

Signed-off-by: yewentao256 <zhyanwentao@126.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/kernels/moe/test_block_fp8.py

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import (
     fused_experts,
     fused_topk,
+    override_config,
 )
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.all2all_utils import (
@@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
     TritonOrDeepGemmExperts,
 )
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    try_get_optimal_moe_config,
+)
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8,
 )
@@ -260,6 +264,57 @@ def test_w8a8_block_fp8_fused_moe(
     torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol)
 
 
+@pytest.mark.parametrize("M", [24, 1024])
+@pytest.mark.parametrize("modular", [False, True])
+@torch.inference_mode()
+def test_h200_block_fp8_up_and_down_tma_matches_pointer(M, modular, workspace_init):
+    if "H200" not in torch.cuda.get_device_name():
+        pytest.skip("This block FP8 TMA test targets H200")
+
+    E, N, K, topk = 256, 256, 128, 8
+    torch.manual_seed(123)
+    a = torch.randn((M, K), dtype=torch.bfloat16) / 10
+    score = torch.randn((M, E), dtype=torch.bfloat16)
+    w1, w2, quant_config = make_test_quant_config(
+        E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=[128, 128]
+    )
+    topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
+
+    config_args = (w1.shape, w2.shape, topk, "fp8_w8a8", M)
+    config = try_get_optimal_moe_config(*config_args, block_shape=[128, 128])
+    assert w2.is_contiguous()
+
+    if modular:
+        moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config)
+
+        def run():
+            return moe.apply(
+                a,
+                w1,
+                w2,
+                topk_weights,
+                topk_ids,
+                activation=MoEActivation.SILU,
+                apply_router_weight_on_input=False,
+                expert_map=None,
+                global_num_experts=E,
+            )
+
+    else:
+
+        def run():
+            return fused_experts(
+                a, w1, w2, topk_weights, topk_ids, quant_config=quant_config
+            )
+
+    with set_current_vllm_config(vllm_config):
+        with override_config({**config, "USE_TMA": False}):
+            pointer_output = run()
+        tma_output = run()
+
+    torch.testing.assert_close(tma_output, pointer_output, atol=0, rtol=0)
+
+
 @pytest.mark.parametrize("clamp_limit", [None, 10.0])
 @torch.inference_mode()
 def test_w8a8_block_fp8_clamp_skips_fused_silu_quant(

tests/kernels/moe/test_moe.py

@@ -75,7 +75,7 @@ def test_triton_moe_launcher_passes_scalar_scale_as_pointer(monkeypatch) -> None
     class FakeKernel:
         def __getitem__(self, grid):
             def launch(*args, **kwargs) -> None:
-                captured["a_scale"] = args[4]
+                captured["a_scale"] = args[5]
 
             return launch
 

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -299,6 +299,7 @@ def fused_moe_kernel(
     # Pointers to matrices
     a_ptr,
     b_ptr,
+    b_desc,
     c_ptr,
     b_bias_ptr,
     a_scale_ptr,
@@ -351,6 +352,7 @@ def fused_moe_kernel(
     SWAP_AB: tl.constexpr,
     # Tensor-descriptor path for the A gather and B load in the K-loop.
     USE_TD: tl.constexpr = False,
+    USE_TMA_B: tl.constexpr = False,
 ):
     """Implements the fused computation for a Mixture of Experts (MOE) using
     token and expert matrices.
@@ -441,6 +443,7 @@ def fused_moe_kernel(
     offs_k = tl.arange(0, BLOCK_SIZE_K)
     # TD gather and the SWAP_AB accumulator layout are mutually exclusive.
     tl.static_assert(not (USE_TD and SWAP_AB))
+    tl.static_assert(not (USE_TD and USE_TMA_B))
     if USE_TD:
         # ``tt.descriptor_gather`` requires block_shape[0] == 1 and i32 idx.
         m_td = num_valid_tokens // top_k
@@ -523,16 +526,26 @@ def fused_moe_kernel(
             b = b_desc.load([pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]).T
         elif SWAP_AB:
             a_mask = (offs_k[:, None] < K - k * BLOCK_SIZE_K) & token_mask[None, :]
-            b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K
             a = tl.load(a_ptrs, mask=a_mask, other=0.0)
-            b = tl.load(b_ptrs, mask=b_mask, other=0.0)
+            if not USE_TMA_B:
+                b_mask = offs_k[None, :] < K - k * BLOCK_SIZE_K
+                b = tl.load(b_ptrs, mask=b_mask, other=0.0)
         else:
             a = tl.load(
                 a_ptrs,
                 mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
                 other=0.0,
             )
-            b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+            if not USE_TMA_B:
+                b = tl.load(
+                    b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0
+                )
+        if USE_TMA_B:
+            b = b_desc.load(
+                [off_experts.to(tl.int32), pid_n * BLOCK_SIZE_N, k * BLOCK_SIZE_K]
+            ).reshape((BLOCK_SIZE_N, BLOCK_SIZE_K))
+            if not SWAP_AB:
+                b = b.T
         # We accumulate along the K dimension.
         if use_int8_w8a16:
             accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
@@ -842,6 +855,26 @@ def invoke_fused_moe_triton_kernel(
     BLOCK_SIZE_K = config.pop("BLOCK_SIZE_K")
     if block_shape is not None:
         BLOCK_SIZE_K = min(BLOCK_SIZE_K, min(block_shape[0], block_shape[1]))
+    use_tma_b = (
+        bool(config.pop("USE_TMA", True))
+        and use_fp8_w8a8
+        and block_shape is not None
+        and all(block_shape)
+        and current_platform.is_cuda()
+        and current_platform.is_device_capability(90)
+        and B.device.type == "cuda"
+        and B.is_contiguous()
+        and B.size(1) % config["BLOCK_SIZE_N"] == 0
+        and B.size(2) % BLOCK_SIZE_K == 0
+    )
+    b_desc = None
+    if use_tma_b:
+        from triton.tools.tensor_descriptor import TensorDescriptor
+
+        set_triton_allocator(B.device)
+        b_desc = TensorDescriptor.from_tensor(
+            B, [1, config["BLOCK_SIZE_N"], BLOCK_SIZE_K]
+        )
     if use_td and A.size(1) % BLOCK_SIZE_K != 0:
         # TD gather/load feeding tl.dot with a non-block-aligned K
         # miscompiles (~74% of output elements wrong) on real HW;
@@ -864,6 +897,7 @@ def invoke_fused_moe_triton_kernel(
     fused_moe_kernel[grid](
         A,
         B,
+        b_desc,
         C,
         B_bias,
         A_scale,
@@ -906,6 +940,7 @@ def invoke_fused_moe_triton_kernel(
         BLOCK_SIZE_K=BLOCK_SIZE_K,
         SWAP_AB=SWAP_AB,
         USE_TD=use_td,
+        USE_TMA_B=use_tma_b,
         **config,
     )