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 8169cefc3c..22f0e3df2a 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 @@ -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, diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index c953afd27a..8683f1ec4e 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/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) diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index c5fe0495ee..0df3c618fe 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/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,