perf(quant): 融合 SwiGLU 激活与 FP8 输入量化 (#1679)

Author: WANDY666

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

@@ -6,6 +6,93 @@
 from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu
 
 
+@triton.jit
+def _silu_and_mul_quant_fp8(
+    input_ptr,
+    quant_ptr,
+    scale_ptr,
+    input_stride: tl.constexpr,
+    size_n: tl.constexpr,
+    groups,
+    scale_stride_n,
+    limit,
+    fp8_max: tl.constexpr,
+    GROUP_SIZE: tl.constexpr,
+    GROUPS_PER_CTA: tl.constexpr,
+    USE_LIMIT: tl.constexpr,
+    USE_GELU: tl.constexpr,
+):
+    ids = tl.program_id(0) * GROUPS_PER_CTA + tl.arange(0, GROUPS_PER_CTA)
+    rows = ids // (size_n // GROUP_SIZE)
+    group_cols = ids % (size_n // GROUP_SIZE)
+    cols = tl.arange(0, GROUP_SIZE)[None, :]
+    offsets = rows[:, None] * input_stride + group_cols[:, None] * GROUP_SIZE + cols
+    mask = ids[:, None] < groups
+    gate = tl.load(input_ptr + offsets, mask, 0).to(tl.float32)
+    up = tl.load(input_ptr + offsets + size_n, mask, 0).to(tl.float32)
+    if USE_LIMIT:
+        gate = tl.minimum(gate, limit)
+        up = tl.minimum(tl.maximum(up, -limit), limit)
+    if 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))
+    else:
+        gate = gate / (1 + tl.exp(-gate))
+    gate = gate.to(input_ptr.dtype.element_ty)
+    # Preserve both rounding boundaries of activation -> BF16 buffer -> FP8.
+    value = (up * gate.to(tl.float32)).to(input_ptr.dtype.element_ty).to(tl.float32)
+    scale = tl.maximum(tl.max(tl.abs(value), axis=1), 1e-10) / fp8_max
+    scale = tl.exp2(tl.ceil(tl.log2(scale)))
+    quant = tl.clamp(value / scale[:, None], -fp8_max, fp8_max).to(quant_ptr.dtype.element_ty)
+    tl.store(quant_ptr + ids[:, None] * GROUP_SIZE + cols, quant, mask)
+    tl.store(scale_ptr + rows + group_cols * scale_stride_n, scale, ids < groups)
+
+
+@autotune(
+    kernel_name="silu_and_mul_quant_fp8:v1",
+    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: {
+        "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,
+    },
+    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):
+    """Fuse blocked activation with FP8 quantization and TMA-aligned UE8M0 scales."""
+    assert input.ndim == 2 and input.stride(1) == 1
+    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)
+    aligned_m = triton.cdiv(m, 4) * 4
+    scale = alloc_func((n // group_size, aligned_m), device=input.device, dtype=torch.float32).T[:m, :]
+    if run_config is None:
+        run_config = {"GROUPS_PER_CTA": 4, "num_warps": 4}
+    groups = m * n // group_size
+    _silu_and_mul_quant_fp8[(triton.cdiv(groups, run_config["GROUPS_PER_CTA"]),)](
+        input,
+        quant,
+        scale,
+        input.stride(0),
+        n,
+        groups,
+        scale.stride(1),
+        limit,
+        torch.finfo(dtype).max,
+        group_size,
+        USE_LIMIT=limit is not None,
+        USE_GELU=ffn_use_tanh_approximate_gelu() and limit is None,
+        **run_config,
+    )
+    return quant, scale
+
+
 @triton.jit
 def _silu_and_mul_kernel_fast(
     input_ptr,

lightllm/common/quantization/deepgemm.py

@@ -97,7 +97,7 @@ def apply(
     ) -> torch.Tensor:
         qweight = weight_pack.weight
         weight_scale = weight_pack.weight_scale
-        input_scale = None
+        input_scale = getattr(input_tensor, "scale", None)
         alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty
         m, k = input_tensor.shape
         n = qweight.shape[0]
@@ -112,6 +112,9 @@ def apply(
                 use_ue8m0_scales=self.use_ue8m0_scales,
                 use_packed_ue8m0=self.use_packed_ue8m0,
             )
+        else:
+            assert out is not None, "prequantized input requires an explicit floating-point output"
+            qinput_tensor = input_tensor
 
         if out is None:
             out = alloc_func((m, n), dtype=input_tensor.dtype, device=input_tensor.device)

lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py

@@ -4,7 +4,11 @@
 from lightllm.common.basemodel.attention.base_att import AttControl
 from lightllm.common.basemodel.triton_kernel.mhc import hc_post, hc_post_pre_norm, hc_pre_norm
 from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import use_sm100_mega_moe
-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.quantization.deepgemm import DeepGEMMFP8w8a8B128QuantizationMethod
 from lightllm.models.deepseek_v4.layer_weights.transformer_layer_weight import DeepseekV4TransformerLayerWeight
 from lightllm.utils.envs_utils import get_env_start_args
 from lightllm.utils.dist_utils import get_global_world_size
@@ -534,6 +538,23 @@ def _token_attention_kernel(
     def _ffn_tp(self, input, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
         input = input.view(-1, self.embed_dim_)
         gate_up = layer_weight.gate_up_proj.mm(input)
+        if (
+            isinstance(layer_weight.down_proj.quant_method, DeepGEMMFP8w8a8B128QuantizationMethod)
+            and layer_weight.down_proj.quant_method.use_ue8m0_scales
+            and not layer_weight.down_proj.quant_method.use_packed_ue8m0
+        ):
+            shared, scale = silu_and_mul_quant_fp8(
+                gate_up,
+                layer_weight.down_proj.quant_method.block_size,
+                layer_weight.down_proj.mm_param.weight.dtype,
+                limit=self.swiglu_limit,
+                alloc_func=self.alloc_tensor,
+            )
+            shared.scale = scale
+            del gate_up
+            # FP8 输入不再携带原始激活 dtype，由调用方保留浮点输出类型。
+            out = self.alloc_tensor((input.size(0), layer_weight.down_proj.mm_param.weight.shape[0]), input.dtype)
+            return layer_weight.down_proj.mm(shared, out=out)
         shared = self.alloc_tensor((input.size(0), gate_up.size(1) // 2), input.dtype)
         silu_and_mul_fwd(
             gate_up,