diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py index 0ccb80ae05..37b741562a 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep/deepgemm.py +++ b/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,6 +244,17 @@ 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: @@ -247,11 +262,6 @@ def chunked_expanded_moe_forward( 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, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py index ef47544733..1434dbac0b 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul.py +++ b/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, ) diff --git a/unit_tests/common/fused_moe/test_prefill_swiglu_fp8.py b/unit_tests/common/fused_moe/test_prefill_swiglu_fp8.py new file mode 100644 index 0000000000..986213baa6 --- /dev/null +++ b/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)