From 347ad611b00ec847110497d50eeff8a28eec0f6a Mon Sep 17 00:00:00 2001
From: wanzihao <1060304770@qq.com>
Date: Mon, 2 Nov 2026 17:58:34 +0800
Subject: [PATCH 1/3] perf(quant): fuse SwiGLU activation and UE8M0 input
 quantization

---
 .../meta_weights/mm_weight/mm_weight.py       | 10 +++
 .../fused_moe/moe_silu_and_mul.py             | 87 +++++++++++++++++++
 lightllm/common/quantization/deepgemm.py      | 22 +++++
 .../common/quantization/quantize_method.py    | 15 ++++
 .../layer_infer/transformer_layer_infer.py    | 12 +--
 5 files changed, 135 insertions(+), 11 deletions(-)

diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
index 5021699143..81d408c1ed 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
@@ -60,6 +60,16 @@ def mm(
             input_tensor, self.mm_param, out, use_custom_tensor_mananger=use_custom_tensor_mananger, bias=self.bias
         )
 
+    def mm_swiglu(self, input_tensor, limit=None, out=None, use_custom_tensor_mananger=True):
+        return self.quant_method.apply_swiglu(
+            input_tensor,
+            self.mm_param,
+            limit=limit,
+            out=out,
+            use_custom_tensor_mananger=use_custom_tensor_mananger,
+            bias=self.bias,
+        )
+
     def gen_weight_quant_param_names(self):
         self.quanted_weight_names = [None] * len(self.weight_names)
         self.weight_zero_point_names = [None] * len(self.weight_names)
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..c0fc3ff557 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: tl.constexpr,
+    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..92461be61f 100644
--- a/lightllm/common/quantization/deepgemm.py
+++ b/lightllm/common/quantization/deepgemm.py
@@ -118,6 +118,28 @@ def apply(
         _deepgemm_fp8_nt((qinput_tensor, input_scale), (qweight, weight_scale), out)
         return out
 
+    def apply_swiglu(self, input_tensor, weight_pack, limit=None, out=None, use_custom_tensor_mananger=True, bias=None):
+        from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_quant_fp8
+
+        if self.use_packed_ue8m0 or not self.use_ue8m0_scales:
+            return super().apply_swiglu(input_tensor, weight_pack, limit, out, use_custom_tensor_mananger, bias)
+        alloc_func = self.cache_manager.empty if use_custom_tensor_mananger else torch.empty
+        quant, scale = silu_and_mul_quant_fp8(
+            input_tensor,
+            self.block_size,
+            weight_pack.weight.dtype,
+            limit=limit,
+            alloc_func=alloc_func,
+        )
+        if out is None:
+            out = alloc_func(
+                (input_tensor.shape[0], weight_pack.weight.shape[0]),
+                dtype=input_tensor.dtype,
+                device=input_tensor.device,
+            )
+        _deepgemm_fp8_nt((quant, scale), (weight_pack.weight, weight_pack.weight_scale), out)
+        return out
+
     def _create_weight(
         self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
     ) -> Tuple[WeightPack, List[WeightPack]]:
diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py
index 95d8d806f9..c6dbdef935 100644
--- a/lightllm/common/quantization/quantize_method.py
+++ b/lightllm/common/quantization/quantize_method.py
@@ -63,6 +63,21 @@ def apply(
     def method_name(self):
         pass
 
+    def apply_swiglu(self, input_tensor, weight_pack, limit=None, out=None, use_custom_tensor_mananger=True, bias=None):
+        from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
+        from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
+
+        alloc_func = g_cache_manager.empty if use_custom_tensor_mananger else torch.empty
+        activation = alloc_func(
+            (input_tensor.shape[0], input_tensor.shape[1] // 2), dtype=input_tensor.dtype, device=input_tensor.device
+        )
+        silu_and_mul_fwd(
+            input_tensor, activation, limit=limit, alpha=1.0 if limit is not None else None, clamp_up_add_one=False
+        )
+        return self.apply(
+            activation, weight_pack, out, use_custom_tensor_mananger=use_custom_tensor_mananger, bias=bias
+        )
+
     def create_weight(
         self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int
     ) -> Tuple[WeightPack, List[WeightPack]]:
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..9ea1b70fe4 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,6 @@
 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.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,16 +533,7 @@ 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)
-        shared = self.alloc_tensor((input.size(0), gate_up.size(1) // 2), input.dtype)
-        silu_and_mul_fwd(
-            gate_up,
-            shared,
-            alpha=1.0,
-            limit=self.swiglu_limit,
-            clamp_up_add_one=False,
-        )
-        del gate_up
-        return layer_weight.down_proj.mm(shared)
+        return layer_weight.down_proj.mm_swiglu(gate_up, limit=self.swiglu_limit)
 
     def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
         x = x.view(-1, self.embed_dim_)

From 0a451dc43b1e18e70f7161b926ad28cd21e252a6 Mon Sep 17 00:00:00 2001
From: wanzihao <1060304770@qq.com>
Date: Tue, 3 Nov 2026 14:22:01 +0800
Subject: [PATCH 2/3] =?UTF-8?q?refactor(quant):=20=E9=80=9A=E8=BF=87=20Ten?=
 =?UTF-8?q?sor.scale=20=E5=A4=8D=E7=94=A8=20mm=20=E6=8E=A5=E5=8F=A3?=
 =?UTF-8?q?=E5=B9=B6=E5=87=8F=E5=B0=91=20kernel=20=E7=BC=96=E8=AF=91?=
 =?UTF-8?q?=E7=89=88=E6=9C=AC?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../meta_weights/mm_weight/mm_weight.py       | 10 ------
 .../fused_moe/moe_silu_and_mul.py             |  2 +-
 lightllm/common/quantization/deepgemm.py      | 33 ++++---------------
 .../common/quantization/quantize_method.py    | 15 ---------
 .../layer_infer/transformer_layer_infer.py    | 29 +++++++++++++++-
 5 files changed, 36 insertions(+), 53 deletions(-)

diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
index 81d408c1ed..5021699143 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
+++ b/lightllm/common/basemodel/layer_weights/meta_weights/mm_weight/mm_weight.py
@@ -60,16 +60,6 @@ def mm(
             input_tensor, self.mm_param, out, use_custom_tensor_mananger=use_custom_tensor_mananger, bias=self.bias
         )
 
-    def mm_swiglu(self, input_tensor, limit=None, out=None, use_custom_tensor_mananger=True):
-        return self.quant_method.apply_swiglu(
-            input_tensor,
-            self.mm_param,
-            limit=limit,
-            out=out,
-            use_custom_tensor_mananger=use_custom_tensor_mananger,
-            bias=self.bias,
-        )
-
     def gen_weight_quant_param_names(self):
         self.quanted_weight_names = [None] * len(self.weight_names)
         self.weight_zero_point_names = [None] * len(self.weight_names)
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 c0fc3ff557..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
@@ -14,7 +14,7 @@ def _silu_and_mul_quant_fp8(
     input_stride: tl.constexpr,
     size_n: tl.constexpr,
     groups,
-    scale_stride_n: tl.constexpr,
+    scale_stride_n,
     limit,
     fp8_max: tl.constexpr,
     GROUP_SIZE: tl.constexpr,
diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py
index 92461be61f..c7c5b8d482 100644
--- a/lightllm/common/quantization/deepgemm.py
+++ b/lightllm/common/quantization/deepgemm.py
@@ -97,11 +97,9 @@ def apply(
     ) -> torch.Tensor:
         qweight = weight_pack.weight
         weight_scale = weight_pack.weight_scale
-        input_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]
+        input_scale = getattr(input_tensor, "scale", None)
         if input_scale is None:
+            alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty
             qinput_tensor, input_scale = per_token_group_quant_fp8(
                 input_tensor,
                 self.block_size,
@@ -112,32 +110,15 @@ 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)
-        _deepgemm_fp8_nt((qinput_tensor, input_scale), (qweight, weight_scale), out)
-        return out
-
-    def apply_swiglu(self, input_tensor, weight_pack, limit=None, out=None, use_custom_tensor_mananger=True, bias=None):
-        from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_quant_fp8
-
-        if self.use_packed_ue8m0 or not self.use_ue8m0_scales:
-            return super().apply_swiglu(input_tensor, weight_pack, limit, out, use_custom_tensor_mananger, bias)
-        alloc_func = self.cache_manager.empty if use_custom_tensor_mananger else torch.empty
-        quant, scale = silu_and_mul_quant_fp8(
-            input_tensor,
-            self.block_size,
-            weight_pack.weight.dtype,
-            limit=limit,
-            alloc_func=alloc_func,
-        )
         if out is None:
             out = alloc_func(
-                (input_tensor.shape[0], weight_pack.weight.shape[0]),
-                dtype=input_tensor.dtype,
-                device=input_tensor.device,
+                (input_tensor.shape[0], qweight.shape[0]), dtype=input_tensor.dtype, device=input_tensor.device
             )
-        _deepgemm_fp8_nt((quant, scale), (weight_pack.weight, weight_pack.weight_scale), out)
+        _deepgemm_fp8_nt((qinput_tensor, input_scale), (qweight, weight_scale), out)
         return out
 
     def _create_weight(
diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py
index c6dbdef935..95d8d806f9 100644
--- a/lightllm/common/quantization/quantize_method.py
+++ b/lightllm/common/quantization/quantize_method.py
@@ -63,21 +63,6 @@ def apply(
     def method_name(self):
         pass
 
-    def apply_swiglu(self, input_tensor, weight_pack, limit=None, out=None, use_custom_tensor_mananger=True, bias=None):
-        from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
-        from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
-
-        alloc_func = g_cache_manager.empty if use_custom_tensor_mananger else torch.empty
-        activation = alloc_func(
-            (input_tensor.shape[0], input_tensor.shape[1] // 2), dtype=input_tensor.dtype, device=input_tensor.device
-        )
-        silu_and_mul_fwd(
-            input_tensor, activation, limit=limit, alpha=1.0 if limit is not None else None, clamp_up_add_one=False
-        )
-        return self.apply(
-            activation, weight_pack, out, use_custom_tensor_mananger=use_custom_tensor_mananger, bias=bias
-        )
-
     def create_weight(
         self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int
     ) -> Tuple[WeightPack, List[WeightPack]]:
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 9ea1b70fe4..f6af9da283 100644
--- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py
@@ -4,6 +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,
+    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
@@ -533,7 +538,29 @@ 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)
-        return layer_weight.down_proj.mm_swiglu(gate_up, limit=self.swiglu_limit)
+        down_proj = layer_weight.down_proj
+        quant_method = down_proj.quant_method
+        if (
+            isinstance(quant_method, DeepGEMMFP8w8a8B128QuantizationMethod)
+            and quant_method.use_ue8m0_scales
+            and not quant_method.use_packed_ue8m0
+        ):
+            shared, scale = silu_and_mul_quant_fp8(
+                gate_up,
+                quant_method.block_size,
+                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), down_proj.mm_param.weight.shape[0]), input.dtype)
+            return 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, shared, alpha=1.0, limit=self.swiglu_limit, clamp_up_add_one=False)
+        del gate_up
+        return down_proj.mm(shared)
 
     def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
         x = x.view(-1, self.embed_dim_)

From 55b4db39821bd1c7bfd7045b80558b9e33e92c85 Mon Sep 17 00:00:00 2001
From: wanzihao <1060304770@qq.com>
Date: Tue, 3 Nov 2026 14:34:40 +0800
Subject: [PATCH 3/3] =?UTF-8?q?refactor(quant):=20=E7=BC=A9=E5=B0=8F?=
 =?UTF-8?q?=E9=A2=84=E9=87=8F=E5=8C=96=E6=8E=A5=E5=85=A5=E5=B7=AE=E5=88=86?=
 =?UTF-8?q?=E5=B9=B6=E4=BF=9D=E7=95=99=E5=8E=9F=E6=9C=89=E5=B1=9E=E6=80=A7?=
 =?UTF-8?q?=E8=AE=BF=E9=97=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 lightllm/common/quantization/deepgemm.py      |  8 +++---
 .../layer_infer/transformer_layer_infer.py    | 26 +++++++++++--------
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py
index c7c5b8d482..8683f1ec4e 100644
--- a/lightllm/common/quantization/deepgemm.py
+++ b/lightllm/common/quantization/deepgemm.py
@@ -98,8 +98,10 @@ def apply(
         qweight = weight_pack.weight
         weight_scale = weight_pack.weight_scale
         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]
         if input_scale is None:
-            alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty
             qinput_tensor, input_scale = per_token_group_quant_fp8(
                 input_tensor,
                 self.block_size,
@@ -115,9 +117,7 @@ def apply(
             qinput_tensor = input_tensor
 
         if out is None:
-            out = alloc_func(
-                (input_tensor.shape[0], qweight.shape[0]), dtype=input_tensor.dtype, device=input_tensor.device
-            )
+            out = alloc_func((m, n), dtype=input_tensor.dtype, device=input_tensor.device)
         _deepgemm_fp8_nt((qinput_tensor, input_scale), (qweight, weight_scale), out)
         return out
 
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 f6af9da283..0df3c618fe 100644
--- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py
@@ -538,29 +538,33 @@ 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)
-        down_proj = layer_weight.down_proj
-        quant_method = down_proj.quant_method
         if (
-            isinstance(quant_method, DeepGEMMFP8w8a8B128QuantizationMethod)
-            and quant_method.use_ue8m0_scales
-            and not quant_method.use_packed_ue8m0
+            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,
-                quant_method.block_size,
-                down_proj.mm_param.weight.dtype,
+                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), down_proj.mm_param.weight.shape[0]), input.dtype)
-            return down_proj.mm(shared, out=out)
+            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, shared, alpha=1.0, limit=self.swiglu_limit, clamp_up_add_one=False)
+        silu_and_mul_fwd(
+            gate_up,
+            shared,
+            alpha=1.0,
+            limit=self.swiglu_limit,
+            clamp_up_add_one=False,
+        )
         del gate_up
-        return down_proj.mm(shared)
+        return layer_weight.down_proj.mm(shared)
 
     def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV4TransformerLayerWeight):
         x = x.view(-1, self.embed_dim_)