perf(moe): fuse GLM prefill SwiGLU and UE8M0 quantization

lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py

@@ -3,7 +3,10 @@
 from typing import List, Optional, Tuple
 
 from lightllm.utils.log_utils import init_logger
-from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import (
+    silu_and_mul_fwd,
+    silu_and_mul_quant_fp8,
+)
 from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul_mix_quant_ep import (
     silu_and_mul_psum_post_quant_fwd,
 )
@@ -17,6 +20,7 @@
     ep_zero_padding,
 )
 from lightllm.utils.envs_utils import get_env_use_ue8m0_scales
+from lightllm.utils.device_utils import is_hopper
 from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType
 
 logger = init_logger(__name__)
@@ -240,18 +244,24 @@ def chunked_expanded_moe_forward(
         dtype=hidden_dtype,
         device=recv_x[0].device,
     )
+    fuse_activation = (
+        not is_autotune_warmup
+        and get_env_use_ue8m0_scales()
+        and hidden_dtype == torch.bfloat16
+        and w2.dtype == torch.float8_e4m3fn
+        and (hidden_size, intermediate_size, block_size_k) == (4096, 2048, 128)
+        and alpha == 1.0
+        and limit == 10.0
+        and not clamp_up_add_one
+        and is_hopper()
+    )
     try:
         for chunk_index, chunk_start in enumerate(range(0, all_tokens, max_chunk_rows)):
             if is_autotune_warmup and chunk_index == 1:
                 Autotuner.end_autotune_warmup()
 
             chunk_end = min(chunk_start + max_chunk_rows, all_tokens)
             chunk_rows = chunk_end - chunk_start
-            silu_out = torch.empty(
-                (chunk_rows, intermediate_size),
-                dtype=hidden_dtype,
-                device=recv_x[0].device,
-            )
             gemm_out_a = torch.empty(
                 (chunk_rows, intermediate_twice),
                 dtype=hidden_dtype,
@@ -263,19 +273,29 @@ def chunked_expanded_moe_forward(
                 gemm_out_a,
                 m_indices[chunk_start:chunk_end],
             )
-            silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one)
+            if fuse_activation and 896 <= chunk_rows <= 32768:
+                qsilu_out, qsilu_out_scale = silu_and_mul_quant_fp8(
+                    input=gemm_out_a,
+                    group_size=block_size_k,
+                    dtype=w2.dtype,
+                    limit=limit,
+                    alpha=alpha,
+                    run_config={"GROUPS_PER_CTA": 4, "num_warps": 1},
+                )
+            else:
+                silu_out = torch.empty((chunk_rows, intermediate_size), dtype=hidden_dtype, device=recv_x[0].device)
+                silu_and_mul_fwd(gemm_out_a, silu_out, alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one)
+                qsilu_out, qsilu_out_scale = per_token_group_quant_fp8(
+                    silu_out,
+                    block_size_k,
+                    dtype=w2.dtype,
+                    column_major_scales=True,
+                    scale_tma_aligned=True,
+                    use_ue8m0_scales=get_env_use_ue8m0_scales(),
+                )
+                del silu_out
             del gemm_out_a
 
-            qsilu_out, qsilu_out_scale = per_token_group_quant_fp8(
-                silu_out,
-                block_size_k,
-                dtype=w2.dtype,
-                column_major_scales=True,
-                scale_tma_aligned=True,
-                use_ue8m0_scales=get_env_use_ue8m0_scales(),
-            )
-            del silu_out
-
             gemm_out_b = torch.empty(
                 (chunk_rows, hidden_size),
                 dtype=hidden_dtype,

lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py

@@ -16,10 +16,12 @@ def _silu_and_mul_quant_fp8(
     groups,
     scale_stride_n,
     limit,
+    alpha,
     fp8_max: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
     GROUPS_PER_CTA: tl.constexpr,
     USE_LIMIT: tl.constexpr,
+    USE_ALPHA: tl.constexpr,
     USE_GELU: tl.constexpr,
 ):
     ids = tl.program_id(0) * GROUPS_PER_CTA + tl.arange(0, GROUPS_PER_CTA)
@@ -33,7 +35,9 @@ def _silu_and_mul_quant_fp8(
     if USE_LIMIT:
         gate = tl.minimum(gate, limit)
         up = tl.minimum(tl.maximum(up, -limit), limit)
-    if USE_GELU:
+    if USE_ALPHA:
+        gate = 1 / (1 + tl.exp(-gate * alpha)) * gate
+    elif USE_GELU:
         cubed = gate * gate * gate
         arg = 0.7978845608028654 * (gate + 0.044715 * cubed)
         gate = 0.5 * gate * (1.0 + (2.0 / (1.0 + tl.exp(-2.0 * arg)) - 1.0))
@@ -54,19 +58,21 @@ def _silu_and_mul_quant_fp8(
     configs_gen_func=lambda: [
         {"GROUPS_PER_CTA": g, "num_warps": w} for g, w in [(1, 1), (4, 1), (4, 4), (8, 4), (16, 4)]
     ],
-    static_key_func=lambda input, group_size, dtype, limit: {
+    static_key_func=lambda input, group_size, dtype, limit, alpha: {
         "N": input.shape[1] // 2,
         "input_dtype": str(input.dtype),
         "output_dtype": str(dtype),
         "group_size": group_size,
         "clamp": limit is not None,
         "gelu": ffn_use_tanh_approximate_gelu() and limit is None,
+        **({"alpha": alpha} if alpha is not None else {}),
     },
     run_key_func=lambda input: input.shape[0],
 )
-def silu_and_mul_quant_fp8(input, group_size, dtype, limit=None, alloc_func=torch.empty, run_config=None):
+def silu_and_mul_quant_fp8(input, group_size, dtype, limit=None, alloc_func=torch.empty, run_config=None, alpha=None):
     """Fuse blocked activation with FP8 quantization and TMA-aligned UE8M0 scales."""
     assert input.ndim == 2 and input.stride(1) == 1
+    assert alpha is None or limit is not None
     assert input.shape[1] % (2 * group_size) == 0 and group_size == triton.next_power_of_2(group_size)
     m, n = input.shape[0], input.shape[1] // 2
     quant = alloc_func((m, n), device=input.device, dtype=dtype)
@@ -84,9 +90,11 @@ def silu_and_mul_quant_fp8(input, group_size, dtype, limit=None, alloc_func=torc
         groups,
         scale.stride(1),
         limit,
+        alpha,
         torch.finfo(dtype).max,
         group_size,
         USE_LIMIT=limit is not None,
+        USE_ALPHA=alpha is not None,
         USE_GELU=ffn_use_tanh_approximate_gelu() and limit is None,
         **run_config,
     )

unit_tests/common/fused_moe/test_prefill_swiglu_fp8.py

@@ -0,0 +1,49 @@
+import pytest
+import torch
+
+from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import (
+    silu_and_mul_fwd,
+    silu_and_mul_quant_fp8,
+)
+from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
+
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.mark.parametrize("rows", [1, 7, 896, 4096, 8192, 23168, 32768])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("row_gap", [0, 16])
+def test_explicit_alpha_preserves_activation_rounding_and_ue8m0(rows: int, dtype: torch.dtype, row_gap: int) -> None:
+    torch.manual_seed(358)
+    intermediate = 2048
+    x = (torch.randn(rows, 2 * intermediate + row_gap, device="cuda", dtype=dtype) * 6)[:, : 2 * intermediate]
+    # Zero and tiny groups exercise the scale floor as well as clamped values.
+    x[0, :128] = 0
+    x[0, 128:256] = 1e-6
+    activated = torch.empty(rows, intermediate, device="cuda", dtype=dtype)
+    silu_and_mul_fwd(x, activated, alpha=1.0, limit=10.0, clamp_up_add_one=False)
+    expected_q, expected_scale = per_token_group_quant_fp8(
+        activated, 128, column_major_scales=True, scale_tma_aligned=True, use_ue8m0_scales=True
+    )
+    actual_q, actual_scale = silu_and_mul_quant_fp8(
+        x, 128, torch.float8_e4m3fn, limit=10.0, alpha=1.0, run_config={"GROUPS_PER_CTA": 4, "num_warps": 1}
+    )
+    assert torch.equal(actual_q.view(torch.uint8), expected_q.view(torch.uint8))
+    assert torch.equal(actual_scale, expected_scale)
+    assert actual_scale.stride() == expected_scale.stride()
+    assert torch.equal(torch.log2(actual_scale), torch.log2(actual_scale).round())
+
+
+@pytest.mark.parametrize("rows", [1, 128, 2048])
+def test_default_activation_path_still_matches_unfused(rows: int) -> None:
+    torch.manual_seed(1679)
+    x = torch.randn(rows, 4096, device="cuda", dtype=torch.bfloat16)
+    activated = torch.empty(rows, 2048, device="cuda", dtype=x.dtype)
+    silu_and_mul_fwd(x, activated)
+    expected_q, expected_scale = per_token_group_quant_fp8(
+        activated, 128, column_major_scales=True, scale_tma_aligned=True, use_ue8m0_scales=True
+    )
+    actual_q, actual_scale = silu_and_mul_quant_fp8(x, 128, torch.float8_e4m3fn)
+    assert torch.equal(actual_q.view(torch.uint8), expected_q.view(torch.uint8))
+    assert torch.equal(actual_scale, expected_scale)