fix(glm5): add independent activation and indexer precision controls

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py

@@ -306,9 +306,10 @@ def _create_weight(self):
         self.per_expert_scale = None
         # Create e_score_correction_bias
         if self.e_score_correction_bias_name:
+            # Preserve FP32 correction scores at expert-selection boundaries.
             self.e_score_correction_bias = torch.empty(
                 (self.n_routed_experts,),
-                dtype=self.data_type_,
+                dtype=torch.float32,
                 device=f"cuda:{self.device_id_}",
             )
             self.e_score_correction_bias.load_ok = False

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py

@@ -6,7 +6,7 @@
 from lightllm.utils.envs_utils import (
     get_deepep_num_max_dispatch_tokens_per_rank_prefill,
     get_deepep_num_max_dispatch_tokens_per_rank_decode,
-    get_env_use_ue8m0_scales,
+    get_env_use_ue8m0_activation_scales,
 )
 from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import (
     fused_experts,
@@ -130,7 +130,7 @@ def low_latency_dispatch(
             num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank,
             num_experts=self.total_expert_num_contain_redundancy,
             use_fp8=use_fp8_w8a8,
-            round_scale=get_env_use_ue8m0_scales(),
+            round_scale=get_env_use_ue8m0_activation_scales(),
             async_finish=False,
             return_recv_hook=True,
         )

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

@@ -25,6 +25,7 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.vllm_utils import vllm_ops
 from lightllm.utils.device_utils import triton_support_tensor_descriptor
+from lightllm.utils.envs_utils import get_env_use_ue8m0_scales, get_env_use_ue8m0_activation_scales
 from .moe_silu_and_mul import silu_and_mul_fwd
 from .moe_sum_reduce import moe_sum_reduce
 from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
@@ -391,7 +392,9 @@ def moe_align_fused(expert_to_token_index, expert_to_weight, expert_token_num, t
     else:
         # Expert rows may be unordered, but grouped matmul reuses this same
         # mapping for up/down projections and writes back to original topk slots.
-        _moe_align_fused_atomic_token(expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights)
+        _moe_align_fused_atomic_token(
+            expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights
+        )
     return expert_to_token_index, expert_to_weight, expert_token_num
 
 
@@ -878,8 +881,12 @@ def grouped_matmul(
             # input 使用 per group quant 量化
             _m, _k = token_inputs.shape
             assert _k % block_size_k == 0
+            # Triton consumes FP32 activation scales even when weight scales are packed UE8M0.
             token_inputs, token_input_scale = per_token_group_quant_fp8(
-                token_inputs, block_size_k, dtype=expert_weights.dtype
+                token_inputs,
+                block_size_k,
+                dtype=expert_weights.dtype,
+                use_ue8m0_scales=get_env_use_ue8m0_activation_scales(get_env_use_ue8m0_scales() or use_packed_ue8m0),
             )
 
     if reused_mblock_infos is None:

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

@@ -22,7 +22,7 @@
 from lightllm.utils.envs_utils import (
     get_deepep_num_max_dispatch_tokens_per_rank_prefill,
     get_deepep_num_max_dispatch_tokens_per_rank_decode,
-    get_env_use_ue8m0_scales,
+    get_env_use_ue8m0_activation_scales,
 )
 from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType
 from lightllm.utils.device_utils import is_sm100_gpu
@@ -158,7 +158,7 @@ def mega_moe_impl(
 
     qinput_tensor = per_token_cast_to_fp8(
         hidden_states,
-        use_ue8m0=True,
+        use_ue8m0=get_env_use_ue8m0_activation_scales(default=True, require_ue8m0=True),
         gran_k=quant_method.block_size,
         use_packed_ue8m0=True,
     )
@@ -192,7 +192,7 @@ def quantize_fused_experts_input(
 
         return per_token_cast_to_fp8(
             hidden_states,
-            use_ue8m0=True,
+            use_ue8m0=get_env_use_ue8m0_activation_scales(default=True, require_ue8m0=True),
             gran_k=quant_method.block_size,
             use_packed_ue8m0=True,
         )
@@ -204,7 +204,7 @@ def quantize_fused_experts_input(
         hidden_states,
         block_size_k,
         dtype=w13.weight.dtype,
-        use_ue8m0_scales=get_env_use_ue8m0_scales(),
+        use_ue8m0_scales=get_env_use_ue8m0_activation_scales(),
     )
 
 
@@ -288,7 +288,7 @@ def fused_experts_impl(
     combined_x = None
     if is_prefill:
         qinput_tensor, input_scale = per_token_group_quant_fp8(
-            hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_scales()
+            hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_activation_scales()
         )
         allocate_on_comm_stream = previous_event is not None
         # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs:
@@ -383,7 +383,7 @@ def fused_experts_impl(
             num_max_dispatch_tokens_per_rank,
             num_experts,
             use_fp8=use_fp8_w8a8,
-            round_scale=get_env_use_ue8m0_scales(),
+            round_scale=get_env_use_ue8m0_activation_scales(),
             async_finish=False,
             return_recv_hook=False,
         )
@@ -608,7 +608,7 @@ def workspace_quant_alloc(shape, dtype, device):
                 column_major_scales=True,
                 scale_tma_aligned=True,
                 alloc_func=workspace_quant_alloc,
-                use_ue8m0_scales=get_env_use_ue8m0_scales(),
+                use_ue8m0_scales=get_env_use_ue8m0_activation_scales(),
             )
             workspace_manager.free(silu_out)
             del silu_out

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

@@ -4,7 +4,7 @@
 import triton.language as tl
 
 from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu
-from lightllm.utils.envs_utils import get_env_use_ue8m0_scales
+from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales
 
 
 @triton.jit
@@ -55,7 +55,9 @@ def _silu_and_mul_post_quant_kernel(
     )
 
     for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE):
-        gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32)
+        gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(
+            tl.float32
+        )
         up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0)
         if USE_LIMIT_AND_ALPHA:
             gate = tl.minimum(gate, limit)
@@ -161,7 +163,7 @@ def silu_and_mul_masked_post_quant_fwd(
         alpha=alpha,
         limit=limit,
         CLAMP_UP_ADD_ONE=clamp_up_add_one,
-        USE_UE8M0_SCALE=get_env_use_ue8m0_scales(),
+        USE_UE8M0_SCALE=get_env_use_ue8m0_activation_scales(),
         num_warps=num_warps,
     )
     return

lightllm/common/basemodel/triton_kernel/mhc/post.py

@@ -13,8 +13,8 @@
 def _hc_post_4stream_kernel(
     layer_output,  # 输入 [T, H]，激活 dtype。
     residual,  # 输入 [T, 4 * H]，激活 dtype；当前子层 pre-mixing 前的残差。
-    residual_mix,  # 输入 [T, 4, 4]，fp32；最后两维为输入、输出 stream。
-    post_mix,  # 输入 [T, 4]，fp32。
+    residual_mix,  # 输入 [T, 4, 4]，fp32 或激活 dtype；最后两维为输入、输出 stream。
+    post_mix,  # 输入 [T, 4]，fp32 或激活 dtype。
     output,  # 输出 [T, 4 * H]，与 layer_output 相同 dtype。
     hidden: tl.constexpr,
     layer_stride_m: tl.constexpr,
@@ -23,6 +23,7 @@ def _hc_post_4stream_kernel(
     post_stride_m: tl.constexpr,
     out_stride_m: tl.constexpr,
     BLOCK_H: tl.constexpr,
+    GLM5_NUMERICS: tl.constexpr,
 ):
     """将子层输出和原始残差混合为四条新 stream：[T, H] + [T, 4, H] -> [T, 4, H]。
 
@@ -50,10 +51,10 @@ def _hc_post_4stream_kernel(
 
     post_base = post_mix + token * post_stride_m
     mix_base = residual_mix + token * mix_stride_m
-    accumulator_0 = layer_value * tl.load(post_base)
-    accumulator_1 = layer_value * tl.load(post_base + 1)
-    accumulator_2 = layer_value * tl.load(post_base + 2)
-    accumulator_3 = layer_value * tl.load(post_base + 3)
+    accumulator_0 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base)
+    accumulator_1 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base + 1)
+    accumulator_2 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base + 2)
+    accumulator_3 = tl.full((BLOCK_H,), 0.0, tl.float32) if GLM5_NUMERICS else layer_value * tl.load(post_base + 3)
 
     accumulator_0 += residual_0 * tl.load(mix_base)
     accumulator_0 += residual_1 * tl.load(mix_base + 4)
@@ -72,6 +73,18 @@ def _hc_post_4stream_kernel(
     accumulator_3 += residual_2 * tl.load(mix_base + 11)
     accumulator_3 += residual_3 * tl.load(mix_base + 15)
 
+    if GLM5_NUMERICS:
+        # With TORCH_COMPILE_DISABLE=1, Megatron materializes both the
+        # residual bmm and the post product in the activation dtype.
+        accumulator_0 = accumulator_0.to(residual.dtype.element_ty).to(tl.float32)
+        accumulator_1 = accumulator_1.to(residual.dtype.element_ty).to(tl.float32)
+        accumulator_2 = accumulator_2.to(residual.dtype.element_ty).to(tl.float32)
+        accumulator_3 = accumulator_3.to(residual.dtype.element_ty).to(tl.float32)
+        accumulator_0 += (layer_value * tl.load(post_base)).to(layer_output.dtype.element_ty).to(tl.float32)
+        accumulator_1 += (layer_value * tl.load(post_base + 1)).to(layer_output.dtype.element_ty).to(tl.float32)
+        accumulator_2 += (layer_value * tl.load(post_base + 2)).to(layer_output.dtype.element_ty).to(tl.float32)
+        accumulator_3 += (layer_value * tl.load(post_base + 3)).to(layer_output.dtype.element_ty).to(tl.float32)
+
     output_base = output + token * out_stride_m + hidden_offsets
     tl.store(output_base, accumulator_0, mask=hidden_mask)
     tl.store(output_base + hidden, accumulator_1, mask=hidden_mask)
@@ -85,6 +98,8 @@ def hc_post(
     residual_mix: torch.Tensor,
     post_mix: torch.Tensor,
     streams: int,
+    *,
+    glm5_numerics: bool = False,
 ) -> torch.Tensor:
     """用一次 Triton launch 将子层输出混回残差 streams。
 
@@ -93,9 +108,10 @@ def hc_post(
     Args:
         layer_output: [T, H]，连续的 attention / FFN 子层输出。
         residual: [T, S * H]，连续的、该子层 pre-mixing 前保存的原始残差。
-        residual_mix: [T, S, S]，fp32；来自 hc_pre_norm。
-        post_mix: [T, S]，fp32；来自同一次 pre 调用。
+        residual_mix: [T, S, S]，fp32 或激活 dtype；来自 hc_pre_norm。
+        post_mix: [T, S]，与 residual_mix 相同 dtype；来自同一次 pre 调用。
         streams: 残差 stream 数 S；当前 Triton 入口要求 S = 4。
+        glm5_numerics: 匹配 GLM 训练默认 native post，在加子层输出前将残差混合结果和 post 乘积分别舍入到激活 dtype。
 
     Returns:
         新残差 [T, S * H]，dtype 与 layer_output 相同；输入张量保持不变。
@@ -127,6 +143,8 @@ def hc_post(
         post_stride_m=post_mix.stride(0),
         out_stride_m=output.stride(0),
         BLOCK_H=block_h,
+        GLM5_NUMERICS=glm5_numerics,
+        enable_fp_fusion=not glm5_numerics,
         num_warps=8,
     )
     return output