fix(glm5): align FP8 precision and preserve reasoning tool calls

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

@@ -140,6 +140,7 @@ def experts(
         alpha: Optional[float] = None,
         limit: Optional[float] = None,
         clamp_up_add_one: bool = True,
+        shared_expert_output: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         # Captures MoE topk expert ids for routed-experts metadata when enabled.
         moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_)
@@ -162,6 +163,7 @@ def experts(
             alpha=alpha,
             limit=limit,
             clamp_up_add_one=clamp_up_add_one,
+            **({"shared_expert_output": shared_expert_output} if shared_expert_output is not None else {}),
         )
 
     def low_latency_dispatch(
@@ -306,9 +308,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,
@@ -80,6 +80,7 @@ def _fused_experts(
         alpha: Optional[float] = None,
         limit: Optional[float] = None,
         clamp_up_add_one: bool = True,
+        shared_expert_output: Optional[torch.Tensor] = None,
     ):
         output = fused_experts(
             hidden_states=input_tensor,
@@ -94,6 +95,7 @@ def _fused_experts(
             alpha=alpha,
             limit=limit,
             clamp_up_add_one=clamp_up_add_one,
+            shared_expert_output=shared_expert_output,
         )
         return output
 
@@ -130,7 +132,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/layer_weights/meta_weights/fused_moe/impl/triton_impl.py

@@ -111,6 +111,7 @@ def __call__(
         alpha: Optional[float] = None,
         limit: Optional[float] = None,
         clamp_up_add_one: bool = True,
+        shared_expert_output: Optional[torch.Tensor] = None,
     ):
         topk_weights, topk_ids, origin_topk_ids = self._select_experts(
             input_tensor=input_tensor,
@@ -140,5 +141,6 @@ def __call__(
             alpha=alpha,
             limit=limit,
             clamp_up_add_one=clamp_up_add_one,
+            **({"shared_expert_output": shared_expert_output} if shared_expert_output is not None else {}),
         )
         return output

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
@@ -878,8 +879,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(),
     )
 
 
@@ -221,10 +221,15 @@ def fused_experts(
     alpha: Optional[float] = None,
     limit: Optional[float] = None,
     clamp_up_add_one: bool = True,
+    shared_expert_output: Optional[torch.Tensor] = None,
 ):
     assert (limit is None and alpha is None) or (limit is not None and alpha is not None)
+    if shared_expert_output is not None and not is_prefill:
+        raise ValueError("Shared combine bias requires normal prefill combine")
     check_ep_expert_dtype(quant_method)
     if use_sm100_mega_moe(quant_method):
+        if shared_expert_output is not None:
+            raise NotImplementedError("Shared combine bias requires DeepEP normal combine")
         if limit is not None:
             raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU")
         return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method)
@@ -249,6 +254,7 @@ def fused_experts(
         alpha=alpha,
         limit=limit,
         clamp_up_add_one=clamp_up_add_one,
+        shared_expert_output=shared_expert_output,
     )
 
 
@@ -271,7 +277,11 @@ def fused_experts_impl(
     limit: Optional[float] = None,
     clamp_up_add_one: bool = True,
     block_size_k: int = 128,
+    shared_expert_output: Optional[torch.Tensor] = None,
 ):
+    # Shared output joins normal combine before its final BF16 cast.
+    if shared_expert_output is not None and not is_prefill:
+        raise ValueError("Shared combine bias requires normal prefill combine")
     # Check constraints.
     assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
     assert topk_weights.shape == topk_idx.shape, "topk shape mismatch"
@@ -288,7 +298,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:
@@ -372,6 +382,7 @@ def fused_experts_impl(
             topk_weights=None,
             previous_event=previous_event,
             allocate_on_comm_stream=allocate_on_comm_stream,
+            **({"bias": shared_expert_output} if shared_expert_output is not None else {}),
         )
     else:
         # low latency dispatch
@@ -383,7 +394,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 +619,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
@@ -161,7 +161,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