From b2d0dba1b8c18001d1f6bba19b079445606cdf35 Mon Sep 17 00:00:00 2001 From: baishihao Date: Thu, 12 Nov 2026 19:10:33 +0800 Subject: [PATCH] perf(glm): fuse dense and shared FFN activation quantization --- .../fused_moe/moe_silu_and_mul.py | 14 ++- .../layer_infer/transformer_layer_infer.py | 21 ++++ .../models/glm5_next/test_ffn_swiglu_quant.py | 112 ++++++++++++++++++ 3 files changed, 144 insertions(+), 3 deletions(-) create mode 100644 unit_tests/models/glm5_next/test_ffn_swiglu_quant.py 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/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py index be7a2fa781..67939f18d9 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -11,7 +11,9 @@ from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward 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.common.basemodel.triton_kernel.mhc import ( hc_contract, hc_post, @@ -89,6 +91,25 @@ def _ffn_tp(self, input, infer_state, layer_weight): input = input.view(-1, self.embed_dim_) up_gate_out = layer_weight.gate_up_proj.mm(input) + quant_method = layer_weight.down_proj.quant_method + if ( + isinstance(quant_method, DeepGEMMFP8w8a8B128QuantizationMethod) + and quant_method.use_ue8m0_scales + and not quant_method.use_packed_ue8m0 + and up_gate_out.dtype == torch.bfloat16 + ): + activated, scale = silu_and_mul_quant_fp8( + input=up_gate_out, + group_size=quant_method.block_size, + dtype=layer_weight.down_proj.mm_param.weight.dtype, + limit=self.swiglu_limit, + alpha=1.0, + alloc_func=self.alloc_tensor, + ) + activated.scale = scale + del up_gate_out + out = self.alloc_tensor(shape=(input.size(0), self.embed_dim_), dtype=input.dtype) + return layer_weight.down_proj.mm(input_tensor=activated, out=out) ffn1_out = self.alloc_tensor((input.size(0), up_gate_out.size(1) // 2), input.dtype) silu_and_mul_fwd( up_gate_out, diff --git a/unit_tests/models/glm5_next/test_ffn_swiglu_quant.py b/unit_tests/models/glm5_next/test_ffn_swiglu_quant.py new file mode 100644 index 0000000000..72ba02b43d --- /dev/null +++ b/unit_tests/models/glm5_next/test_ffn_swiglu_quant.py @@ -0,0 +1,112 @@ +from collections.abc import Callable +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightTpl +from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 +from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant +from lightllm.common.quantization.deepgemm import DeepGEMMFP8w8a8B128QuantizationMethod +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.common.quantization.quantize_method import WeightPack +from lightllm.models.glm5_next.layer_infer import transformer_layer_infer as glm + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() != (9, 0), + reason="Hopper and DeepGEMM required", +) + + +def make_weight(values: torch.Tensor, quantized: bool, ue8m0: bool) -> MMWeightTpl: + owner = MMWeightTpl.__new__(MMWeightTpl) + owner.bias = None + if quantized: + weight, scale = weight_quant(x=values, block_size=128, use_ue8m0_scales=ue8m0) + owner.mm_param = WeightPack(weight=weight, weight_scale=scale) + owner.quant_method = DeepGEMMFP8w8a8B128QuantizationMethod() + owner.quant_method.use_ue8m0_scales = ue8m0 + owner.quant_method.use_packed_ue8m0 = False + else: + owner.mm_param = WeightPack(weight=values) + owner.quant_method = NoQuantization() + return owner + + +@pytest.mark.parametrize("rows,intermediate", [(1, 256), (7, 2048), (33, 2048), (256, 2048), (129, 12288)]) +@pytest.mark.parametrize("quantized,ue8m0", [(True, True), (True, False), (False, False)]) +def test_ffn_matches_main_and_preserves_quantization( + rows: int, intermediate: int, quantized: bool, ue8m0: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + torch.manual_seed(seed=5512 + rows) + hidden = 256 + x = torch.randn(size=(rows, hidden), device="cuda", dtype=torch.bfloat16) * 8 + gate_up = make_weight( + values=torch.randn(size=(2 * intermediate, hidden), device="cuda", dtype=x.dtype) * 0.1, + quantized=quantized, + ue8m0=ue8m0, + ) + down = make_weight( + values=torch.randn(size=(hidden, intermediate), device="cuda", dtype=x.dtype) * 0.05, + quantized=quantized, + ue8m0=ue8m0, + ) + layer = glm.Glm5NextTransformerLayerInfer.__new__(glm.Glm5NextTransformerLayerInfer) + layer.embed_dim_, layer.swiglu_limit = hidden, 10.0 + weights = SimpleNamespace(gate_up_proj=gate_up, down_proj=down) + projected = gate_up.mm(input_tensor=x) + activated = torch.empty(size=(rows, intermediate), device=x.device, dtype=x.dtype) + glm.silu_and_mul_fwd(input=projected, output=activated, limit=10.0, alpha=1.0, clamp_up_add_one=False) + expected = down.mm(input_tensor=activated) + seen = [] + fused = glm.silu_and_mul_quant_fp8 + + def capture( + input: torch.Tensor, group_size: int, dtype: torch.dtype, limit: float, alpha: float, alloc_func: Callable + ) -> tuple[torch.Tensor, torch.Tensor]: + q, s = fused(input=input, group_size=group_size, dtype=dtype, limit=limit, alpha=alpha, alloc_func=alloc_func) + seen.append((q.clone(), s.clone())) + return q, s + + monkeypatch.setattr(target=glm, name="silu_and_mul_quant_fp8", value=capture) + actual = layer._ffn_tp(input=x, infer_state=None, layer_weight=weights) + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0) + assert len(seen) == int(quantized and ue8m0) + if seen: + q, s = per_token_group_quant_fp8( + x=activated, group_size=128, column_major_scales=True, scale_tma_aligned=True, use_ue8m0_scales=True + ) + torch.testing.assert_close(actual=seen[0][0].view(torch.uint8), expected=q.view(torch.uint8), atol=0, rtol=0) + torch.testing.assert_close(actual=seen[0][1], expected=s, atol=0, rtol=0) + + +def test_ffn_cuda_graph_reads_changed_inputs() -> None: + torch.manual_seed(seed=5513) + hidden, intermediate, rows = 256, 2048, 7 + x = torch.randn(size=(rows, hidden), device="cuda", dtype=torch.bfloat16) + layer = glm.Glm5NextTransformerLayerInfer.__new__(glm.Glm5NextTransformerLayerInfer) + layer.embed_dim_, layer.swiglu_limit = hidden, 10.0 + weights = SimpleNamespace( + gate_up_proj=make_weight( + values=torch.randn(size=(2 * intermediate, hidden), device="cuda", dtype=x.dtype) * 0.1, + quantized=True, + ue8m0=True, + ), + down_proj=make_weight( + values=torch.randn(size=(hidden, intermediate), device="cuda", dtype=x.dtype) * 0.05, + quantized=True, + ue8m0=True, + ), + ) + layer._ffn_tp(input=x, infer_state=None, layer_weight=weights) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(cuda_graph=graph): + actual = layer._ffn_tp(input=x, infer_state=None, layer_weight=weights) + for factor in (1.0, 12.0): + x.normal_().mul_(factor) + projected = weights.gate_up_proj.mm(input_tensor=x) + activated = torch.empty(size=(rows, intermediate), device=x.device, dtype=x.dtype) + glm.silu_and_mul_fwd(input=projected, output=activated, limit=10.0, alpha=1.0, clamp_up_add_one=False) + expected = weights.down_proj.mm(input_tensor=activated) + graph.replay() + torch.testing.assert_close(actual=actual, expected=expected, atol=0, rtol=0)