diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
index fd3de5c9b4..ef970735d1 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py
+++ b/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
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
index 74e457426b..c22a5c595a 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py
+++ b/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,
)
diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
index f748aea467..222e209374 100644
--- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py
+++ b/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
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py
index 90eeb5a4ca..19df68e041 100644
--- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe.py
+++ b/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:
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
index 61cd81556a..fda0ab09f9 100644
--- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py
+++ b/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
diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
index 8c94f088c2..ef058407dd 100644
--- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py
+++ b/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
diff --git a/lightllm/common/basemodel/triton_kernel/mhc/post.py b/lightllm/common/basemodel/triton_kernel/mhc/post.py
index b4050e0765..b420bb2edf 100644
--- a/lightllm/common/basemodel/triton_kernel/mhc/post.py
+++ b/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
diff --git a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py
index 2710040673..3191ce0716 100644
--- a/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py
+++ b/lightllm/common/basemodel/triton_kernel/mhc/pre_norm.py
@@ -14,6 +14,63 @@
LIGHTLLM_DISABLE_DEEPGEMM_MHC = os.getenv("LIGHTLLM_DISABLE_DEEPGEMM_MHC", "False").upper() in ["ON", "TRUE", "1"]
+@triton.jit
+def _hc_input_inv_rms_kernel(
+ x,
+ inv_rms,
+ FLATTENED_HIDDEN: tl.constexpr,
+ RMS_EPS: tl.constexpr,
+ BLOCK_K: tl.constexpr,
+):
+ token = tl.program_id(0)
+ offsets = tl.arange(0, BLOCK_K)
+ values = tl.load(x + token * FLATTENED_HIDDEN + offsets, offsets < FLATTENED_HIDDEN, 0).to(tl.float32)
+ variance = tl.sum(values * values, axis=0) / FLATTENED_HIDDEN
+ tl.store(inv_rms + token, tl.rsqrt(variance + RMS_EPS))
+
+
+@triton.jit
+def _hc_normalized_projection_kernel(
+ x,
+ fn,
+ inv_rms,
+ gemm_partial,
+ tokens,
+ FLATTENED_HIDDEN: tl.constexpr,
+ MIX_SIZE: tl.constexpr,
+ SPLIT_K_BLOCKS: tl.constexpr,
+ BLOCK_T: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+ BLOCK_K: tl.constexpr,
+):
+ token_offsets = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T)
+ split = tl.program_id(1)
+ mix_offsets = tl.arange(0, BLOCK_M)
+ k_offsets = split * SPLIT_K_BLOCKS * BLOCK_K + tl.arange(0, BLOCK_K)
+ rms = tl.load(inv_rms + token_offsets, token_offsets < tokens, 0)
+ accumulator = tl.zeros((BLOCK_T, BLOCK_M), dtype=tl.float32)
+ for _ in range(SPLIT_K_BLOCKS):
+ values = tl.load(
+ x + token_offsets[:, None] * FLATTENED_HIDDEN + k_offsets[None, :],
+ (token_offsets[:, None] < tokens) & (k_offsets[None, :] < FLATTENED_HIDDEN),
+ 0,
+ ).to(tl.float32)
+ weights = tl.load(
+ fn + mix_offsets[None, :] * FLATTENED_HIDDEN + k_offsets[:, None],
+ (mix_offsets[None, :] < MIX_SIZE) & (k_offsets[:, None] < FLATTENED_HIDDEN),
+ 0,
+ )
+ # Normalize in FP32 before the IEEE projection, independent of TF32 settings.
+ normalized = values * rms[:, None]
+ accumulator = tl.dot(normalized, weights, accumulator, input_precision="ieee")
+ k_offsets += BLOCK_K
+ tl.store(
+ gemm_partial + (split * tokens + token_offsets[:, None]) * MIX_SIZE + mix_offsets[None, :],
+ accumulator,
+ (token_offsets[:, None] < tokens) & (mix_offsets[None, :] < MIX_SIZE),
+ )
+
+
@triton.jit
def _hc_prepare_prenorm_kernel(
gemm_partial, # 输入 [P, T, M]，fp32；投影的 split-K 部分和。
@@ -21,8 +78,8 @@ def _hc_prepare_prenorm_kernel(
scale, # 输入 [3]，fp32；pre、post、residual 三组缩放系数。
base, # 输入 [M]，fp32；M = 2 * S + S * S。
pre, # 输出 [T, S]，fp32。
- post, # 输出 [T, S]，fp32。
- residual_mix, # 输出 [T, S, S]，fp32；最后两维为输入、输出 stream。
+ post, # 输出 [T, S]，fp32 或激活 dtype。
+ residual_mix, # 输出 [T, S, S]，fp32 或激活 dtype；最后两维为输入、输出 stream。
gemm_stride_s,
gemm_stride_m: tl.constexpr,
sqrsum_stride_s,
@@ -36,6 +93,7 @@ def _hc_prepare_prenorm_kernel(
POST_MULTIPLIER: tl.constexpr,
SINKHORN_ITERS: tl.constexpr,
N_SPLITS: tl.constexpr,
+ PROJECTION_IS_NORMALIZED: tl.constexpr,
):
"""归并分片并准备混合权重：[P, T, M] / [P, T] -> [T, S] / [T, S, S]。
@@ -55,11 +113,13 @@ def _hc_prepare_prenorm_kernel(
pre_raw += tl.load(partial_base + stream_offsets)
post_raw += tl.load(partial_base + STREAMS + stream_offsets)
matrix_raw += tl.load(partial_base + 2 * STREAMS + matrix_offsets)
- sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m)
- inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS)
- pre_raw *= inv_rms
- post_raw *= inv_rms
- matrix_raw *= inv_rms
+ if not PROJECTION_IS_NORMALIZED:
+ sqrsum += tl.load(sqrsum_partial + split * sqrsum_stride_s + token * sqrsum_stride_m)
+ if not PROJECTION_IS_NORMALIZED:
+ inv_rms = tl.rsqrt(sqrsum / FLATTENED_HIDDEN + RMS_EPS)
+ pre_raw *= inv_rms
+ post_raw *= inv_rms
+ matrix_raw *= inv_rms
pre_values = tl.sigmoid(pre_raw * tl.load(scale) + tl.load(base + stream_offsets)) + HC_EPS
post_values = POST_MULTIPLIER * tl.sigmoid(post_raw * tl.load(scale + 1) + tl.load(base + STREAMS + stream_offsets))
@@ -141,6 +201,38 @@ def _compute_prenorm_splits(tokens: int, flattened_hidden: int, device: torch.de
return max(1, min(sms // max(grid_size, 1), k_blocks // 4))
+def _hc_normalized_projection(x: torch.Tensor, fn: torch.Tensor, rms_eps: float) -> torch.Tensor:
+ """FP32 normalize-before-project, without materializing a [T, S * H] FP32 copy."""
+ tokens, flattened_hidden = x.shape
+ mix_size = fn.shape[0]
+ n_splits = _compute_prenorm_splits(tokens, flattened_hidden, x.device)
+ inv_rms = torch.empty((tokens,), dtype=torch.float32, device=x.device)
+ gemm_partial = torch.empty((n_splits, tokens, mix_size), dtype=torch.float32, device=x.device)
+ _hc_input_inv_rms_kernel[(tokens,)](
+ x,
+ inv_rms,
+ FLATTENED_HIDDEN=flattened_hidden,
+ RMS_EPS=rms_eps,
+ BLOCK_K=triton.next_power_of_2(flattened_hidden),
+ num_warps=8,
+ )
+ _hc_normalized_projection_kernel[(triton.cdiv(tokens, 16), n_splits)](
+ x,
+ fn,
+ inv_rms,
+ gemm_partial,
+ tokens,
+ FLATTENED_HIDDEN=flattened_hidden,
+ MIX_SIZE=mix_size,
+ SPLIT_K_BLOCKS=triton.cdiv(triton.cdiv(flattened_hidden, 64), n_splits),
+ BLOCK_T=16,
+ BLOCK_M=triton.next_power_of_2(mix_size),
+ BLOCK_K=64,
+ num_warps=4,
+ )
+ return gemm_partial
+
+
def hc_pre_norm(
x: torch.Tensor,
fn: torch.Tensor,
@@ -153,6 +245,8 @@ def hc_pre_norm(
hc_eps: float,
sinkhorn_iters: int,
post_multiplier: float = 2.0,
+ *,
+ glm5_numerics: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""将 mHC pre-mixing 与紧随其后的子层 RMSNorm 融合。
@@ -170,15 +264,18 @@ def hc_pre_norm(
hc_eps: pre sigmoid 后的偏移量及 Sinkhorn 的稳定项。
sinkhorn_iters: Sinkhorn 迭代次数。
post_multiplier: post sigmoid 的乘数，默认 2.0。
+ glm5_numerics: FP32 先归一化后投影；pre 系数保持 fp32，post/residual 舍入到激活 dtype。
Returns:
- (layer_input, residual_mix, post_mix)：分别为 bf16 [T, H]、fp32 [T, S, S]、fp32 [T, S]。
+ (layer_input, residual_mix, post_mix)：形状为 [T, H]、[T, S, S]、[T, S]。
+ layer_input 为 bf16；两个混合系数默认 fp32，glm5_numerics=True 时为激活 dtype。
layer_input 已完成子层 RMSNorm；residual_mix[t, i, j] 对应输入 i 到输出 j。
x 保持不变，留给对应 hc_post 使用。
默认用 DeepGEMM 融合投影与平方和计算，Triton 完成 Sinkhorn、stream 合并和 RMSNorm。
合并结果先舍入到 bf16，再做 RMSNorm。
LIGHTLLM_DISABLE_DEEPGEMM_MHC=1 禁用 DeepGEMM，改用 PyTorch 计算投影与平方和。
+ glm5_numerics=True 则使用独立的 IEEE FP32 Triton 投影，不依赖 DeepGEMM 或全局 TF32 设置。
"""
assert x.ndim == 2 and x.shape[-1] % streams == 0
@@ -187,8 +284,13 @@ def hc_pre_norm(
assert x.is_contiguous() and fn.is_contiguous() and norm_weight.is_contiguous()
tokens, flattened_hidden = x.shape
hidden = flattened_hidden // streams
+ assert fn.shape == ((2 + streams) * streams, flattened_hidden)
- if LIGHTLLM_DISABLE_DEEPGEMM_MHC:
+ if glm5_numerics:
+ gemm_partial = _hc_normalized_projection(x, fn, rms_eps)
+ sqrsum_partial = None
+ n_splits = gemm_partial.shape[0]
+ elif LIGHTLLM_DISABLE_DEEPGEMM_MHC:
# PyTorch produces one complete partition: [1, T, M] and [1, T].
x_fp32 = x.float()
gemm_partial = (x_fp32 @ fn.T).unsqueeze(0)
@@ -204,8 +306,9 @@ def hc_pre_norm(
tf32_hc_prenorm_gemm(x, fn, gemm_partial, sqrsum_partial, n_splits)
pre = torch.empty((tokens, streams), dtype=torch.float32, device=x.device)
- post = torch.empty_like(pre)
- residual_mix = torch.empty((tokens, streams, streams), dtype=torch.float32, device=x.device)
+ mix_dtype = x.dtype if glm5_numerics else torch.float32
+ post = torch.empty((tokens, streams), dtype=mix_dtype, device=x.device)
+ residual_mix = torch.empty((tokens, streams, streams), dtype=mix_dtype, device=x.device)
_hc_prepare_prenorm_kernel[(tokens,)](
gemm_partial,
sqrsum_partial,
@@ -216,8 +319,8 @@ def hc_pre_norm(
residual_mix,
gemm_partial.stride(0),
gemm_partial.stride(1),
- sqrsum_partial.stride(0),
- sqrsum_partial.stride(1),
+ sqrsum_partial.stride(0) if sqrsum_partial is not None else 0,
+ sqrsum_partial.stride(1) if sqrsum_partial is not None else 0,
pre.stride(0),
residual_mix.stride(0),
FLATTENED_HIDDEN=flattened_hidden,
@@ -227,6 +330,7 @@ def hc_pre_norm(
POST_MULTIPLIER=post_multiplier,
SINKHORN_ITERS=sinkhorn_iters,
N_SPLITS=n_splits,
+ PROJECTION_IS_NORMALIZED=glm5_numerics,
num_warps=1,
)
diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py
index c953afd27a..5aca76bd5d 100644
--- a/lightllm/common/quantization/deepgemm.py
+++ b/lightllm/common/quantization/deepgemm.py
@@ -6,6 +6,7 @@
from lightllm.common.quantization.registry import QUANTMETHODS
from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8
from lightllm.utils.device_utils import is_sm100_gpu
+from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales
from lightllm.utils.log_utils import init_logger
logger = init_logger(__name__)
@@ -109,7 +110,9 @@ def apply(
column_major_scales=True,
scale_tma_aligned=True,
alloc_func=alloc_func,
- use_ue8m0_scales=self.use_ue8m0_scales,
+ use_ue8m0_scales=get_env_use_ue8m0_activation_scales(
+ self.use_ue8m0_scales, require_ue8m0=self.use_packed_ue8m0
+ ),
use_packed_ue8m0=self.use_packed_ue8m0,
)
diff --git a/lightllm/common/quantization/w8a8.py b/lightllm/common/quantization/w8a8.py
index 472db0d543..92075b8282 100644
--- a/lightllm/common/quantization/w8a8.py
+++ b/lightllm/common/quantization/w8a8.py
@@ -13,6 +13,7 @@
from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_gemm_kernel import w8a8_block_fp8_matmul
from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops, cutlass_scaled_mm
from lightllm.utils.sgl_utils import HAS_SGL_KERNEL, sgl_ops
+from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales
_HAS_SGL_FP8 = HAS_SGL_KERNEL and sgl_ops is not None and hasattr(sgl_ops, "fp8_scaled_mm")
@@ -447,7 +448,11 @@ def _dynamic_quant_input(
alloc_func = torch.empty if not use_custom_tensor_mananger else self.cache_manager.empty
if input_scale is None:
qinput_tensor, input_scale = per_token_group_quant_fp8(
- input_tensor, self.block_size, dtype=qweight.dtype, alloc_func=alloc_func
+ input_tensor,
+ self.block_size,
+ dtype=qweight.dtype,
+ alloc_func=alloc_func,
+ use_ue8m0_scales=get_env_use_ue8m0_activation_scales(default=False),
)
if out is None:
out = alloc_func((m, n), dtype=input_tensor.dtype, device=input_tensor.device)
diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py
index cf13ab83de..30a9883881 100644
--- a/lightllm/models/glm5_next/indexer.py
+++ b/lightllm/models/glm5_next/indexer.py
@@ -15,7 +15,8 @@ def __init__(self, layer_idx, network_config, tp_world_size):
self.topk = network_config["index_topk"]
self.heads = network_config["index_n_heads"]
self.dim = network_config["index_head_dim"]
- self.eps = network_config["rms_norm_eps"]
+ # The indexer LayerNorm has a separate epsilon from the model RMSNorm.
+ self.eps = 1e-6
def select_topk_indices(self, logits, lengths, indices):
"""Select row-relative indices, with valid entries before -1 padding."""
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 620554bc06..e748701967 100644
--- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py
+++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py
@@ -17,7 +17,7 @@
)
from lightllm.common.triton_utils.autotuner import Autotuner
from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer
-from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.envs_utils import enable_env_vars, get_env_start_args
from lightllm.utils.tensor_utils import tensor_to_no_ref_tensor
@@ -37,6 +37,7 @@ def __init__(self, layer_num, network_config):
self.hc_sinkhorn_iters = network_config.get("hc_sinkhorn_iters", 20)
self.swiglu_limit = network_config["swiglu_limit"]
self.enable_ep_moe = get_env_start_args().enable_ep_moe
+ self.ep_shared_combine_bias = self.enable_ep_moe and enable_env_vars("LIGHTLLM_EP_SHARED_COMBINE_BIAS")
self.is_moe = (
network_config["n_routed_experts"] is not None
and layer_num >= network_config["first_k_dense_replace"]
@@ -101,10 +102,12 @@ def _moe_ffn(self, input, infer_state, layer_weight) -> torch.Tensor:
hidden_states = input.view(-1, self.embed_dim_)
num_tokens, hidden_dim = hidden_states.shape
- # if fused_shared_experts is not enabled, compute shared_output
+ shared_output = None
if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0:
shared_output = self._ffn_tp(hidden_states, infer_state, layer_weight)
+ combine_shared = self.ep_shared_combine_bias and infer_state.is_prefill and shared_output is not None
+
moe_gate_dtype = layer_weight.moe_gate.data_type_
router_logits = layer_weight.moe_gate.mm(hidden_states.to(moe_gate_dtype))
output = layer_weight.experts.experts(
@@ -120,9 +123,10 @@ def _moe_ffn(self, input, infer_state, layer_weight) -> torch.Tensor:
alpha=1.0,
limit=self.swiglu_limit,
clamp_up_add_one=False,
+ **({"shared_expert_output": shared_output} if combine_shared else {}),
)
- if self.n_shared_experts is not None and layer_weight.num_fused_shared_experts == 0:
+ if shared_output is not None and not combine_shared:
output.add_(shared_output)
return output.view(num_tokens, hidden_dim)
@@ -428,6 +432,7 @@ def _hc_pre(self, streams, layer_weight, prefix, norm_weight):
norm_eps=self.eps_,
hc_eps=self.hc_eps,
sinkhorn_iters=self.hc_sinkhorn_iters,
+ glm5_numerics=True,
)
def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill):
@@ -438,11 +443,11 @@ def _forward_mhc(self, input_embeddings, infer_state, layer_weight, *, prefill):
layer_output = self.context_attention_forward(layer_input, infer_state, layer_weight)
else:
layer_output = self.token_attention_forward(layer_input, infer_state, layer_weight)
- streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams)
+ streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams, glm5_numerics=True)
layer_input, residual_mix, post_mix = self._hc_pre(streams, layer_weight, "ffn", layer_weight.ffn_norm_weight_)
layer_output = self._ffn(layer_input, infer_state, layer_weight)
- streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams)
+ streams = hc_post(layer_output, streams, residual_mix, post_mix, self.mhc_streams, glm5_numerics=True)
# Only prefill autotuning truncates the model. Decode autotuning still
# executes every layer and must keep all residual streams until the end.
is_autotune_last_layer = (
diff --git a/lightllm/models/glm5_next/triton_kernel/index_quant.py b/lightllm/models/glm5_next/triton_kernel/index_quant.py
index 67485d9b52..25c09fd908 100644
--- a/lightllm/models/glm5_next/triton_kernel/index_quant.py
+++ b/lightllm/models/glm5_next/triton_kernel/index_quant.py
@@ -2,6 +2,7 @@
import triton
import triton.language as tl
from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage
+from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales
@triton.jit
@@ -13,6 +14,7 @@ def _hadamard_transform_quant_fp8_kernel(
scale: tl.constexpr,
BLOCK_R: tl.constexpr,
BLOCK_N: tl.constexpr,
+ USE_UE8M0_SCALES: tl.constexpr,
):
pid = tl.program_id(0)
rows = pid * BLOCK_R + tl.arange(0, BLOCK_R)
@@ -32,7 +34,9 @@ def _hadamard_transform_quant_fp8_kernel(
# Match the unfused path's bf16 Hadamard output before FP8 quantization.
x = (x * scale).to(tl.bfloat16).to(tl.float32)
absmax = tl.maximum(tl.max(tl.abs(x), axis=1), 1e-4)
- quant_scale = tl.exp2(tl.ceil(tl.log2(absmax * (1.0 / 448.0))))
+ quant_scale = absmax * (1.0 / 448.0)
+ if USE_UE8M0_SCALES:
+ quant_scale = tl.exp2(tl.ceil(tl.log2(quant_scale)))
y = tl.minimum(tl.maximum(x / quant_scale[:, None], -448.0), 448.0)
tl.store(Y + offsets, y, mask=row_mask[:, None])
@@ -40,7 +44,7 @@ def _hadamard_transform_quant_fp8_kernel(
def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[torch.Tensor, torch.Tensor]:
- """Fuse Hadamard-128 with the following ue8m0 FP8 quantization."""
+ """Fuse Hadamard-128 and FP8 quantization with FP32 scale storage."""
assert x.is_cuda, "hadamard_transform_quant_fp8 only supports CUDA tensors"
assert x.dtype == torch.bfloat16, "Hadamard transform expects bfloat16 input"
@@ -64,6 +68,7 @@ def hadamard_transform_quant_fp8(x: torch.Tensor, scale: float = 1.0) -> tuple[t
scale,
BLOCK_R=block_r,
BLOCK_N=128,
+ USE_UE8M0_SCALES=get_env_use_ue8m0_activation_scales(default=True),
num_warps=2,
)
return output.view(original_shape), output_scale
diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py
index 0e126aea6a..c40ab42abd 100644
--- a/lightllm/models/glm5_next/triton_kernel/kpool.py
+++ b/lightllm/models/glm5_next/triton_kernel/kpool.py
@@ -4,6 +4,7 @@
from lightllm.models.deepseek3_2.triton_kernel.hadamard_transform import _butterfly_stage
from lightllm.utils.device_utils import get_device_sm_count
+from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales
@triton.jit
@@ -45,6 +46,7 @@ def _compress_pools(
SINGLE_QUERY: tl.constexpr,
TAIL_SIZE: tl.constexpr,
HAS_MTP_INDEX: tl.constexpr,
+ USE_UE8M0_SCALES: tl.constexpr,
):
if SINGLE_QUERY:
batch = tl.program_id(0)
@@ -88,7 +90,9 @@ def _compress_pools(
for step in tl.static_range(7):
key = _butterfly_stage(key, 64 >> step, 1 << step, 1, 128)
key = (key * (128 ** -0.5)).to(tl.bfloat16).to(tl.float32)
- scale = tl.exp2(tl.ceil(tl.log2(tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0)))
+ scale = tl.maximum(tl.max(tl.abs(key), 1), 1e-4) / 448.0
+ if USE_UE8M0_SCALES:
+ scale = tl.exp2(tl.ceil(tl.log2(scale)))
key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv)
loc = tl.load(Ragged + start + length - 1).to(tl.int64)
dest = Packed + loc * PACKED_STRIDE
@@ -156,6 +160,7 @@ def compress_pools(
single_query,
tail.shape[1],
mtp_index is not None,
+ USE_UE8M0_SCALES=get_env_use_ue8m0_activation_scales(default=True),
num_warps=4,
)
# Complete all boundary pools before replacing the previous chunk's tail.
diff --git a/lightllm/server/api_openai.py b/lightllm/server/api_openai.py
index fa61a311f4..fee264748a 100644
--- a/lightllm/server/api_openai.py
+++ b/lightllm/server/api_openai.py
@@ -181,7 +181,7 @@ def _get_history_tool_calls_cnt(request: ChatCompletionRequest) -> int:
def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool:
- """Whether this request uses forced thinking / reasoning (parser + template)."""
+ """Resolve the request's thinking switch for rendering the chat template."""
from .build_prompt import tokenizer_supports_force_thinking
reasoning_parser = get_env_start_args().reasoning_parser
@@ -201,11 +201,10 @@ def _is_force_thinking_mode(request: ChatCompletionRequest) -> bool:
return True # default
-def _get_guided_reasoning_end(request: ChatCompletionRequest, tokenizer) -> List[int]:
- """Use the same thinking mode as the chat template for grammar activation."""
- if not _is_force_thinking_mode(request):
+def _get_guided_reasoning_end(tokenizer, parser_name: Optional[str], initial_reasoning_state: bool) -> List[int]:
+ """Delay answer constraints while generation starts inside reasoning."""
+ if not parser_name or not initial_reasoning_state:
return []
- parser_name = get_env_start_args().reasoning_parser
detector = ReasoningParser.DetectorMap[parser_name.lower()]()
token_ids = tokenizer.encode(detector.grammar_start_token, add_special_tokens=False)
return token_ids
@@ -217,13 +216,14 @@ def _process_reasoning_stream(
reasoning_parser_dict: Dict[int, ReasoningParser],
metadata: Dict[str, Any],
request: ChatCompletionRequest,
-) -> tuple[Optional[str], str]:
+ force_reasoning: bool,
+) -> tuple[str, str]:
"""Process reasoning content and update its token usage."""
if index not in reasoning_parser_dict:
reasoning_parser_dict[index] = ReasoningParser(
get_env_start_args().reasoning_parser,
request.stream_reasoning,
- _is_force_thinking_mode(request),
+ force_reasoning,
)
parser = reasoning_parser_dict[index]
token_id = metadata.get("id")
@@ -353,6 +353,10 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
multimodal_params = _build_multimodal_params(request)
tools = _select_chat_template_tools(request)
prompt = await build_prompt(request, tools)
+ reasoning_parser = get_env_start_args().reasoning_parser
+ initial_reasoning_state = ReasoningParser.get_initial_reasoning_state(
+ reasoning_parser, prompt, default=_is_force_thinking_mode(request)
+ )
sampling_params_dict = {
"do_sample": request.do_sample,
"presence_penalty": request.presence_penalty,
@@ -385,9 +389,9 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
sampling_params_dict["max_new_tokens"] = request.max_tokens
if request.stop is not None:
sampling_params_dict["stop_sequences"] = request.stop
- if get_env_start_args().reasoning_parser:
+ if reasoning_parser:
# The prompt may already open reasoning; generated tokens may omit the opening marker.
- sampling_params_dict["_initial_reasoning_state"] = int(_is_force_thinking_mode(request))
+ sampling_params_dict["_initial_reasoning_state"] = int(initial_reasoning_state)
# Chat 接口提取 reasoning，按请求设置 stop 范围；通用 generate 接口不区分阶段。
sampling_params_dict["stop_str_match_mode"] = (
StopStrMatchMode.ALL_TEXT if is_stop_matching_in_reasoning_enabled() else StopStrMatchMode.CONTENT_ONLY
@@ -405,7 +409,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
if sampling_params_dict.get("guided_json") or sampling_params_dict.get("guided_grammar"):
sampling_params_dict["guided_reasoning_end"] = _get_guided_reasoning_end(
- request, g_objs.httpserver_manager.tokenizer
+ g_objs.httpserver_manager.tokenizer, reasoning_parser, initial_reasoning_state
)
sampling_params = SamplingParams()
@@ -418,8 +422,6 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
# Non-streaming case
if not request.stream:
- reasoning_parser = get_env_start_args().reasoning_parser
- request_enable_reasoning = _is_force_thinking_mode(request) if reasoning_parser else False
final_output_dict = collections.defaultdict(list)
count_output_tokens_dict = collections.defaultdict(lambda: 0)
finish_reason_dict = {}
@@ -440,7 +442,7 @@ async def chat_completions_impl(request: ChatCompletionRequest, raw_request: Req
parser = ReasoningParser(
reasoning_parser,
stream_reasoning=False,
- force_reasoning=request_enable_reasoning,
+ force_reasoning=initial_reasoning_state,
)
reasoning_parser_dict[sub_req_id] = parser
token_id = metadata.get("id")
@@ -590,10 +592,21 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
yield f"data: {_serialize_sse_chunk(first_chunk, _first_choice_nulls, _first_resp_nulls)}\n\n"
# Handle reasoning content
- if get_env_start_args().reasoning_parser:
+ if reasoning_parser:
reasoning_text, delta = _process_reasoning_stream(
- choice_index, delta, reasoning_parser_dict, metadata, request
+ choice_index,
+ delta,
+ reasoning_parser_dict,
+ metadata,
+ request,
+ force_reasoning=initial_reasoning_state,
)
+ if current_finish_reason is not None:
+ # Flush before tool parsing so buffered argument text follows
+ # the same path as every other chunk, including on truncation.
+ flush_reasoning, flush_text = reasoning_parser_dict[choice_index].flush()
+ reasoning_text += flush_reasoning
+ delta += flush_text
if reasoning_text:
if request.separate_reasoning:
choice_data = ChatCompletionStreamResponseChoice(
@@ -609,7 +622,7 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
)
yield f"data: {_serialize_sse_chunk(chunk, _choice_nulls)}\n\n"
else:
- delta = reasoning_text + (delta or "")
+ delta = reasoning_text + delta
if request.tool_choice != "none" and request.tools:
# parse_increment => returns (normal_text, calls)
@@ -773,46 +786,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
# Emit a per-choice final chunk with finish_reason (for tool_calls path
# or when no delta was emitted alongside finish_reason).
if current_finish_reason is not None:
- # Flush any buffered reasoning content that was never released
- # (e.g., max_completion_tokens hit before was seen).
- if get_env_start_args().reasoning_parser:
- parser = reasoning_parser_dict.get(choice_index)
- if parser is not None:
- flush_reasoning, flush_text = parser.flush()
- if flush_reasoning:
- if request.separate_reasoning:
- flush_choice = ChatCompletionStreamResponseChoice(
- index=choice_index,
- delta=DeltaMessage(reasoning=flush_reasoning),
- finish_reason=None,
- )
- else:
- # vLLM compat: emit buffered thinking as content
- flush_choice = ChatCompletionStreamResponseChoice(
- index=choice_index,
- delta=DeltaMessage(content=flush_reasoning),
- finish_reason=None,
- )
- flush_chunk = ChatCompletionStreamResponse(
- id=chat_completion_id,
- created=created_time,
- model=request.model,
- choices=[flush_choice],
- )
- yield f"data: {_serialize_sse_chunk(flush_chunk, _choice_nulls)}\n\n"
- if flush_text:
- flush_choice = ChatCompletionStreamResponseChoice(
- index=choice_index,
- delta=DeltaMessage(content=flush_text),
- finish_reason=None,
- )
- flush_chunk = ChatCompletionStreamResponse(
- id=chat_completion_id,
- created=created_time,
- model=request.model,
- choices=[flush_choice],
- )
- yield f"data: {_serialize_sse_chunk(flush_chunk, _choice_nulls)}\n\n"
if has_emitted_tool_calls[sub_req_id] and current_finish_reason == "stop":
current_finish_reason = "tool_calls"
final_choice = ChatCompletionStreamResponseChoice(
@@ -828,7 +801,6 @@ async def stream_results() -> AsyncGenerator[bytes, None]:
)
yield f"data: {_serialize_sse_chunk(final_chunk, _final_choice_nulls)}\n\n"
- reasoning_parser = get_env_start_args().reasoning_parser
completion_tokens_details = None
if reasoning_parser:
reasoning_tokens = sum(parser.reasoning_tokens for parser in reasoning_parser_dict.values())
diff --git a/lightllm/server/function_call_parser.py b/lightllm/server/function_call_parser.py
index e5d7564081..e2b6d2ab84 100644
--- a/lightllm/server/function_call_parser.py
+++ b/lightllm/server/function_call_parser.py
@@ -13,6 +13,7 @@
import ast
import json
+import math
import os
import orjson
import logging
@@ -20,7 +21,7 @@
from abc import ABC, abstractmethod
from json import JSONDecodeError, JSONDecoder
from json.decoder import WHITESPACE
-from typing import Any, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, Iterator, List, Optional, Tuple, Type
import partial_json_parser
from partial_json_parser.core.exceptions import MalformedJSON
@@ -28,6 +29,7 @@
from pydantic import BaseModel, Field
from .api_models import Tool
+from .parser_utils import GlmToolCallScanner
logger = logging.getLogger(__name__)
ENABLE_TOOL_NAME_CHECK = os.getenv("LIGHTLLM_ENABLE_TOOL_NAME_CHECK", "False").upper() in ["ON", "TRUE", "1"]
@@ -1187,64 +1189,152 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
class Glm47Detector(BaseFormatDetector):
- """
- Detector for GLM-4.7/GLM-4.7-Flash model function call format.
-
- The GLM-4.7 format uses an XML-style envelope with arg_key/arg_value pairs
- instead of JSON arguments.
-
- Format Structure:
- ```
- function_name
- param1
- value1
- param2
- value2
-
- ```
-
- Example:
- ```
- tool_brave_web_search_post
- query
- test search
- count
- 5
-
- ```
-
- Key Components:
- - Tool Call Tags: `` and `` wrap each individual call
- - Function Name: Appears on the first line after ``
- - Arguments: Pairs of `name` and `value`
+ """Parse GLM's XML-style tool calls.
- Reference: https://github.com/vllm-project/vllm/blob/main/vllm/tool_parsers/glm4_moe_tool_parser.py
+ Format: namekeyvalue
+ Arguments stay literal until the call is complete, then convert according
+ to the function's parameter schema.
"""
+ _TYPE_ALIASES = {
+ "str": "string",
+ "text": "string",
+ "varchar": "string",
+ "char": "string",
+ "enum": "string",
+ "int": "integer",
+ "int32": "integer",
+ "int64": "integer",
+ "uint": "integer",
+ "uint32": "integer",
+ "uint64": "integer",
+ "long": "integer",
+ "short": "integer",
+ "unsigned": "integer",
+ "float": "number",
+ "float32": "number",
+ "float64": "number",
+ "double": "number",
+ "bool": "boolean",
+ "dict": "object",
+ "arr": "array",
+ "list": "array",
+ "sequence": "array",
+ }
+ _TYPE_PRIORITY = ("null", "integer", "number", "boolean", "object", "array", "string")
+
def __init__(self):
super().__init__()
self.bot_token = ""
self.eot_token = ""
self.tool_call_separator = "\n"
- # Regex patterns for parsing GLM-4.7 tool calls
- # Match complete tool call blocks
- self.func_call_regex = re.compile(r".*?", re.DOTALL)
- # Extract function name and arguments from a tool call block
- # Function name can be followed by newline OR directly by
- # Pattern: function_name(\n|)...
- self.func_detail_regex = re.compile(
- r"([^<\n]+?)(?:\n|(?=)|(?=))(.*?)", re.DOTALL
- )
+ # A function name ends at a newline, the first argument, or the closing tag.
self.func_name_regex = re.compile(r"([^<\n]+?)(?:\n|(?=)|(?=))")
- self._streaming_tool_name: Optional[str] = None
- # Extract arg_key/arg_value pairs
self.func_arg_regex = re.compile(r"(.*?)\s*(.*?)", re.DOTALL)
+ self._tool_scanner = GlmToolCallScanner()
+ self._tool_scan_offset = len(self.bot_token)
def has_tool_call(self, text: str) -> bool:
"""Check if the text contains a GLM-4.7 format tool call."""
return self.bot_token in text
+ @classmethod
+ def _schema_types(cls, schema: Any) -> set[str]:
+ if not isinstance(schema, dict):
+ return {"string"}
+ types = set()
+ declared = schema.get("type")
+ if isinstance(declared, str):
+ types.add(declared)
+ elif isinstance(declared, list):
+ types.update(item for item in declared if isinstance(item, str))
+ enum_values = schema.get("enum", [])
+ if not isinstance(enum_values, list):
+ enum_values = []
+ for value in enum_values:
+ if value is None:
+ types.add("null")
+ elif isinstance(value, bool):
+ types.add("boolean")
+ elif isinstance(value, int):
+ types.add("integer")
+ elif isinstance(value, float):
+ types.add("number")
+ elif isinstance(value, str):
+ types.add("string")
+ elif isinstance(value, list):
+ types.add("array")
+ elif isinstance(value, dict):
+ types.add("object")
+ for field in ("anyOf", "oneOf", "allOf"):
+ choices = schema.get(field)
+ if isinstance(choices, list):
+ for choice in choices:
+ types.update(cls._schema_types(choice))
+ return types or {"string"}
+
+ @classmethod
+ def _is_finite_json(cls, value: Any) -> bool:
+ if isinstance(value, float):
+ return math.isfinite(value)
+ if isinstance(value, dict):
+ return all(cls._is_finite_json(item) for item in value.values())
+ if isinstance(value, list):
+ return all(cls._is_finite_json(item) for item in value)
+ return True
+
+ @staticmethod
+ def _parse_typed_value(value: str, value_type: str) -> Any:
+ match value_type:
+ case "string":
+ return value
+ case "null":
+ if value.lower() == "null":
+ return None
+ case "integer":
+ return int(value)
+ case "number":
+ number = float(value)
+ return int(number) if number.is_integer() else number
+ case "boolean":
+ normalized = value.strip().lower()
+ if normalized in ("true", "1"):
+ return True
+ if normalized in ("false", "0"):
+ return False
+ case "object" | "array":
+ return json.loads(value)
+ raise ValueError(f"Cannot parse value as {value_type}")
+
+ @classmethod
+ def _convert_param_value(cls, value: str, schema: Any) -> Any:
+ """Match vLLM's schema-based coercion without changing literal strings.
+
+ Union types use null/integer/number/boolean/object/array/string priority.
+ Unspecified fields remain strings. Non-finite numbers remain raw text so
+ the resulting arguments are always valid JSON.
+ """
+ schema_types = set()
+ for type_name in cls._schema_types(schema):
+ type_name = type_name.strip().lower()
+ schema_types.add(cls._TYPE_ALIASES.get(type_name, type_name))
+
+ for value_type in cls._TYPE_PRIORITY:
+ if value_type not in schema_types:
+ continue
+ try:
+ parsed = cls._parse_typed_value(value, value_type)
+ except (ValueError, TypeError):
+ continue
+ if cls._is_finite_json(parsed):
+ return parsed
+ try:
+ parsed = json.loads(value)
+ return parsed if cls._is_finite_json(parsed) else value
+ except ValueError:
+ return value
+
def _parse_xml_arguments(self, arg_text: str, tool: Tool) -> dict:
"""
Parse XML-style arguments into a dictionary.
@@ -1260,23 +1350,25 @@ def _parse_xml_arguments(self, arg_text: str, tool: Tool) -> dict:
args = {}
properties = (tool.function.parameters or {}).get("properties", {})
+ if not isinstance(properties, dict):
+ properties = {}
matches = self.func_arg_regex.findall(arg_text)
for key, value in matches:
key = key.strip()
- # XML string values are literal: file content and edit targets may
- # contain meaningful whitespace or text that happens to be JSON.
- if properties.get(key, {}).get("type") == "string":
- args[key] = value
- continue
- # Try to parse value as JSON for complex types (arrays, objects, numbers, booleans)
- try:
- parsed_value = json.loads(value)
- args[key] = parsed_value
- except (json.JSONDecodeError, ValueError):
- # Keep as string if not valid JSON
- args[key] = value
+ args[key] = self._convert_param_value(value, properties.get(key))
return args
+ def _iter_tool_call_spans(self, text: str) -> Iterator[Tuple[int, int]]:
+ start = text.find(self.bot_token)
+ while start != -1:
+ body_start = start + len(self.bot_token)
+ relative_end = GlmToolCallScanner().find_end(text[body_start:])
+ if relative_end is None:
+ return
+ end = body_start + relative_end
+ yield start, end
+ start = text.find(self.bot_token, end)
+
def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
"""
One-time parsing: Detects and parses tool calls in the provided text.
@@ -1285,28 +1377,26 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult
:param tools: List of available tools.
:return: StreamingParseResult with normal_text and parsed calls.
"""
- idx = text.find(self.bot_token)
- normal_text = text[:idx].strip() if idx != -1 else text
-
if self.bot_token not in text:
- return StreamingParseResult(normal_text=normal_text, calls=[])
+ return StreamingParseResult(normal_text=text, calls=[])
tool_indices = self._get_tool_indices(tools)
calls = []
+ normal_parts = []
+ cursor = 0
- # Find all ... blocks
- match_result_list = self.func_call_regex.findall(text)
-
- for match_result in match_result_list:
+ for start, end in self._iter_tool_call_spans(text):
+ normal_parts.append(text[cursor:start])
+ cursor = end
+ tool_text = text[start:end]
try:
- # Extract function name and arguments
- func_detail = self.func_detail_regex.search(match_result)
- if not func_detail:
- logger.warning(f"Failed to parse GLM-4.7 tool call: {match_result}")
+ header = self.func_name_regex.match(tool_text)
+ if header is None:
+ logger.warning(f"Failed to parse GLM-4.7 tool call: {tool_text}")
continue
- func_name = func_detail.group(1).strip()
- arg_text = func_detail.group(2) if func_detail.group(2) else ""
+ func_name = header.group(1).strip()
+ arg_text = tool_text[header.end() : -len(self.eot_token)]
# Validate function name
if func_name not in tool_indices:
@@ -1324,10 +1414,32 @@ def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult
)
)
except Exception as e:
- logger.warning(f"Failed to parse GLM-4.7 tool call: {match_result}, error: {str(e)}")
+ logger.warning(f"Failed to parse GLM-4.7 tool call: {tool_text}, error: {str(e)}")
continue
- return StreamingParseResult(normal_text=normal_text, calls=calls)
+ # Keep ordinary text between and after calls. An unfinished envelope
+ # stays buffered in streaming mode and must not leak in non-streaming.
+ content_end = text.find(self.bot_token, cursor)
+ if content_end == -1:
+ content_end = len(text)
+ normal_parts.append(text[cursor:content_end])
+ return StreamingParseResult(normal_text="".join(normal_parts).strip(), calls=calls)
+
+ def _stream_incomplete_call(self, tools: List[Tool]) -> Optional[ToolCallItem]:
+ """Announce the name once, then emit empty argument deltas while buffering."""
+ if self.current_tool_name_sent:
+ return ToolCallItem(tool_index=self.current_tool_id, parameters="")
+
+ header = self.func_name_regex.match(self._buffer)
+ if header is None:
+ return None
+ name = header.group(1).strip()
+ if name not in self._get_tool_indices(tools):
+ return None
+
+ self.current_tool_id += 1
+ self.current_tool_name_sent = True
+ return ToolCallItem(tool_index=self.current_tool_id, name=name, parameters="")
def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult:
"""Buffer XML arguments while emitting tool deltas to keep the stream alive."""
@@ -1345,35 +1457,29 @@ def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> Streami
normal_text += self._buffer[:start]
self._buffer = self._buffer[start:]
- end = self._buffer.find(self.eot_token)
- if end == -1:
- # Match the buffered Qwen3-Coder flow: announce the name once,
- # then send empty argument deltas until the XML call is complete.
- # Long Write/Bash arguments must not leave the HTTP stream idle.
- name = None
- if self._streaming_tool_name is None:
- match = self.func_name_regex.match(self._buffer)
- if match and match.group(1).strip() in self._get_tool_indices(tools):
- name = match.group(1).strip()
- self._streaming_tool_name = name
- self.current_tool_id += 1
- if self._streaming_tool_name is not None:
- calls.append(ToolCallItem(tool_index=self.current_tool_id, name=name, parameters=""))
+ relative_end = self._tool_scanner.find_end(self._buffer[self._tool_scan_offset :])
+ if relative_end is None:
+ self._tool_scan_offset = len(self._buffer)
+ pending_call = self._stream_incomplete_call(tools)
+ if pending_call is not None:
+ calls.append(pending_call)
break
- end += len(self.eot_token)
+ end = self._tool_scan_offset + relative_end
# Waiting for the closing tag avoids re-emitting partial JSON or
# rewriting already sent arguments when another XML key arrives.
result = self.detect_and_parse(self._buffer[:end], tools)
for call in result.calls:
- if self._streaming_tool_name is None:
- self.current_tool_id += 1
- else:
+ if self.current_tool_name_sent:
call.name = None
+ else:
+ self.current_tool_id += 1
call.tool_index = self.current_tool_id
calls.append(call)
- self._streaming_tool_name = None
+ self.current_tool_name_sent = False
self._buffer = self._buffer[end:]
+ self._tool_scanner = GlmToolCallScanner()
+ self._tool_scan_offset = len(self.bot_token)
return StreamingParseResult(normal_text=normal_text, calls=calls)
diff --git a/lightllm/server/parser_utils.py b/lightllm/server/parser_utils.py
new file mode 100644
index 0000000000..8299357554
--- /dev/null
+++ b/lightllm/server/parser_utils.py
@@ -0,0 +1,62 @@
+# Copyright 2025 ModelTC Team
+# Copyright 2023-2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from typing import Optional, Sequence, Tuple
+
+
+def prefix_hold(text: str, markers: Sequence[str]) -> Tuple[str, str]:
+ """Split off the longest suffix that could become a complete marker."""
+ max_length = min(len(text), max((len(marker) - 1 for marker in markers), default=0))
+ for length in range(max_length, 0, -1):
+ suffix = text[-length:]
+ if any(len(suffix) < len(marker) and marker.startswith(suffix) for marker in markers):
+ return text[:-length], suffix
+ return text, ""
+
+
+class GlmToolCallScanner:
+ """Find the outer GLM tool-call end while treating argument values as literal text."""
+
+ _MARKERS = ("", "", "")
+ _MARKER_PATTERN = re.compile(r"|")
+
+ def __init__(self):
+ self._in_arg_value = False
+ self._pending_marker = ""
+
+ def find_end(self, new_text: str) -> Optional[int]:
+ """Return the closing tag's end offset in this chunk, or None if incomplete.
+
+ Callers may forward every character immediately. The retained marker
+ prefix is only for recognizing a delimiter split across chunks.
+ """
+ pending_length = len(self._pending_marker)
+ text = self._pending_marker + new_text
+ self._pending_marker = ""
+ offset = 0
+
+ for match in self._MARKER_PATTERN.finditer(text):
+ offset = match.end()
+ marker = match.group()
+ if self._in_arg_value:
+ if marker == "":
+ self._in_arg_value = False
+ elif marker == "":
+ self._in_arg_value = True
+ elif marker == "":
+ return offset - pending_length
+
+ _, self._pending_marker = prefix_hold(text[offset:], self._MARKERS)
+ return None
diff --git a/lightllm/server/reasoning_parser.py b/lightllm/server/reasoning_parser.py
index 1e7c2f1315..b0036ae288 100644
--- a/lightllm/server/reasoning_parser.py
+++ b/lightllm/server/reasoning_parser.py
@@ -16,6 +16,7 @@
from typing import Iterator, List, Tuple, Dict, Optional, Type
from lightllm.utils.config_utils import get_token_id
+from lightllm.server.parser_utils import GlmToolCallScanner, prefix_hold
@dataclass
@@ -36,28 +37,6 @@ class Token:
end: int
-def prefix_hold(text: str, tokens: List[str]) -> Tuple[str, str]:
- """
- Holds back the longest suffix of `text` that could be a prefix of any token.
- Returns (emit_now, keep_for_later).
- """
- if not text:
- return "", ""
- max_hold = 0
- for tok in tokens:
- if not tok:
- continue
- # Check for prefixes of tok in the suffix of text
- L = min(len(tok) - 1, len(text))
- for k in range(L, 0, -1):
- if tok.startswith(text[-k:]):
- max_hold = max(max_hold, k)
- break
- if max_hold == 0:
- return text, ""
- return text[:-max_hold], text[-max_hold:]
-
-
def iter_tokens(text: str, start_pos: int = 0) -> Iterator[Token]:
"""Iterate over structural tokens in left-to-right order."""
TOKENS = {
@@ -600,6 +579,11 @@ def __init__(
self._buffer = ""
self.stripped_think_start = False
+ @classmethod
+ def initial_reasoning_state(cls, prompt: str) -> Optional[bool]:
+ """Infer the generation phase from the prompt, or defer to the request."""
+ return None
+
@property
def grammar_start_token(self) -> str:
"""Generated delimiter after which the final answer's grammar applies."""
@@ -756,6 +740,83 @@ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False)
)
+class Glm45Detector(BaseReasoningFormatDetector):
+ """Route GLM tool envelopes to content, even before an explicit ````."""
+
+ tool_start_token = ""
+ _MARKERS = ("", "", tool_start_token)
+ _MARKER_PATTERN = re.compile("|".join(re.escape(marker) for marker in _MARKERS))
+
+ def __init__(self, stream_reasoning: bool = True, force_reasoning: bool = False):
+ super().__init__("", "", stream_reasoning=stream_reasoning, force_reasoning=force_reasoning)
+ self._tool_scanner: Optional[GlmToolCallScanner] = None
+
+ @classmethod
+ def initial_reasoning_state(cls, prompt: str) -> Optional[bool]:
+ # GLM-5.3 always opens thinking, even if enable_thinking=False was
+ # requested. Inspect only the final assistant turn, never past turns.
+ _, assistant_marker, assistant_prompt = prompt.rpartition("<|assistant|>")
+ if not assistant_marker:
+ return None
+ if any(marker in assistant_prompt for marker in ("<|system|>", "<|user|>", "<|observation|>")):
+ return None
+ think_start = assistant_prompt.rfind("")
+ think_end = assistant_prompt.rfind("")
+ if think_start == think_end == -1:
+ return None
+ return think_start > think_end
+
+ def detect_and_parse(self, text: str) -> StreamingParseResult:
+ detector = type(self)(stream_reasoning=self.stream_reasoning, force_reasoning=self._in_reasoning)
+ result = detector.parse_streaming_increment(text)
+ remaining = detector.flush()
+ return StreamingParseResult(
+ reasoning_text=(result.reasoning_text + remaining.reasoning_text).strip(),
+ normal_text=result.normal_text + remaining.normal_text,
+ )
+
+ def parse_streaming_increment(self, new_text: str) -> StreamingParseResult:
+ self._buffer += new_text
+ reasoning_parts = []
+ normal_parts = []
+
+ while self._buffer:
+ if self._tool_scanner is not None:
+ # Tool arguments are literal, including any reasoning markers.
+ end = self._tool_scanner.find_end(self._buffer)
+ if end is None:
+ end = len(self._buffer)
+ else:
+ self._tool_scanner = None
+ normal_parts.append(self._buffer[:end])
+ self._buffer = self._buffer[end:]
+ continue
+
+ match = self._MARKER_PATTERN.search(self._buffer)
+ output_parts = reasoning_parts if self._in_reasoning else normal_parts
+ if match is None:
+ # A delimiter may share a chunk with reasoning or span chunks.
+ text, self._buffer = prefix_hold(self._buffer, self._MARKERS)
+ output_parts.append(text)
+ break
+
+ output_parts.append(self._buffer[: match.start()])
+ marker = match.group()
+ self._buffer = self._buffer[match.end() :]
+ if marker == self.tool_start_token:
+ self._in_reasoning = False
+ self._tool_scanner = GlmToolCallScanner()
+ # Preserve the opener and literal thinking tags in tool arguments.
+ normal_parts.append(marker)
+ else:
+ self._in_reasoning = marker == self.think_start_token
+
+ return StreamingParseResult(
+ reasoning_text="".join(reasoning_parts),
+ normal_text="".join(normal_parts),
+ )
+
+
class KimiDetector(BaseReasoningFormatDetector):
"""
Detector for Kimi Thinking model.
@@ -919,7 +980,7 @@ class ReasoningParser:
DetectorMap: Dict[str, Type[BaseReasoningFormatDetector]] = {
"deepseek-r1": DeepSeekR1Detector,
"deepseek-v3": Qwen3Detector,
- "glm45": Qwen3Detector,
+ "glm45": Glm45Detector,
"gpt-oss": GptOssDetector,
"kimi": KimiDetector,
"kimi_k2": DeepSeekR1Detector,
@@ -933,6 +994,17 @@ class ReasoningParser:
"gemma4": Gemma4Detector,
}
+ @classmethod
+ def get_initial_reasoning_state(cls, model_type: Optional[str], prompt: str, default: bool) -> bool:
+ """Resolve the starting phase once for parsing, stops and answer constraints."""
+ if not model_type:
+ return False
+ detector_class = cls.DetectorMap.get(model_type.lower())
+ if detector_class is None:
+ raise ValueError(f"Unsupported model type: {model_type}")
+ state = detector_class.initial_reasoning_state(prompt)
+ return default if state is None else state
+
def __init__(
self,
model_type: Optional[str] = None,
@@ -975,17 +1047,17 @@ def update_reasoning_token_count(self, token_id: int) -> None:
else:
self.reasoning_tokens += 1
- def parse_non_stream(self, full_text: str) -> Tuple[Optional[str], Optional[str]]:
+ def parse_non_stream(self, full_text: str) -> Tuple[str, str]:
"""Non-streaming call: one-time parsing"""
ret = self.detector.detect_and_parse(full_text)
return ret.reasoning_text, ret.normal_text
- def parse_stream_chunk(self, chunk_text: str) -> Tuple[Optional[str], Optional[str]]:
+ def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, str]:
"""Streaming call: incremental parsing"""
ret = self.detector.parse_streaming_increment(chunk_text)
return ret.reasoning_text, ret.normal_text
- def flush(self) -> Tuple[Optional[str], Optional[str]]:
+ def flush(self) -> Tuple[str, str]:
"""Flush remaining buffered content when generation ends prematurely."""
ret = self.detector.flush()
return ret.reasoning_text, ret.normal_text
diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py
index 748f1ad2f7..7ad9d68cce 100644
--- a/lightllm/utils/envs_utils.py
+++ b/lightllm/utils/envs_utils.py
@@ -77,6 +77,25 @@ def get_env_use_ue8m0_scales() -> bool:
return os.getenv("LIGHTLLM_USE_UE8M0_SCALES", "0").upper() in ["ON", "TRUE", "1"]
+def get_env_use_ue8m0_activation_scales(default=None, *, require_ue8m0=False) -> bool:
+ """Select block-FP8 activation scales without changing checkpoint weights.
+
+ LIGHTLLM_ACTIVATION_SCALE_FORMAT: auto keeps the caller's existing policy;
+ fp32 uses continuous scales; ue8m0 rounds scales upward to powers of two.
+ GLM5 indexer Q/K use this policy too; their payload and cache remain FP8.
+ """
+ scale_format = os.getenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "auto").strip().lower()
+ if scale_format == "auto":
+ use_ue8m0 = get_env_use_ue8m0_scales() if default is None else default
+ elif scale_format in ("fp32", "ue8m0"):
+ use_ue8m0 = scale_format == "ue8m0"
+ else:
+ raise ValueError("LIGHTLLM_ACTIVATION_SCALE_FORMAT must be auto, fp32, or ue8m0")
+ if require_ue8m0 and not use_ue8m0:
+ raise ValueError("This packed DeepGEMM backend requires UE8M0 activation scales; fp32 is supported on Hopper")
+ return use_ue8m0
+
+
@lru_cache(maxsize=None)
def is_stop_matching_in_reasoning_enabled() -> bool:
"""是否允许用户配置的 stop 在 reasoning 阶段参与匹配。"""
diff --git a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py
index 03ae1f3c33..2fea058ea6 100644
--- a/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py
+++ b/unit_tests/common/basemodel/triton_kernel/mhc/test_mhc.py
@@ -1,7 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
+import math
import sys
-from types import ModuleType
+from types import ModuleType, SimpleNamespace
import pytest
import torch
@@ -16,7 +17,19 @@
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_iters, post_multiplier=2.0):
+def _hc_pre_reference(
+ x,
+ fn,
+ scale,
+ base,
+ streams,
+ rms_eps,
+ hc_eps,
+ sinkhorn_iters,
+ post_multiplier=2.0,
+ *,
+ normalize_before_projection=False,
+):
"""Unfused FP32 reference; returns the merged input before sublayer RMSNorm."""
assert x.ndim == 2 and x.shape[-1] % streams == 0
tokens, flattened_hidden = x.shape
@@ -25,7 +38,10 @@ def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_ite
x_fp32 = x.float()
inv_rms = torch.rsqrt(x_fp32.square().mean(dim=-1, keepdim=True) + rms_eps)
- mixes = F.linear(x_fp32, fn) * inv_rms
+ if normalize_before_projection:
+ mixes = F.linear(x_fp32 * inv_rms, fn)
+ else:
+ mixes = F.linear(x_fp32, fn) * inv_rms
pre_raw = mixes[:, :streams]
post_raw = mixes[:, streams : 2 * streams]
@@ -44,12 +60,17 @@ def _hc_pre_reference(x, fn, scale, base, streams, rms_eps, hc_eps, sinkhorn_ite
return layer_input, residual_mix, post
-def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams):
+def _hc_post_reference(layer_output, residual, residual_mix, post_mix, streams, *, round_residual=False):
"""Unfused FP32 reference; residual_mix maps input streams to output streams."""
tokens, hidden = layer_output.shape
residual_3d = residual.view(tokens, streams, hidden)
mixed_residual = (residual_mix.unsqueeze(-1) * residual_3d.float().unsqueeze(2)).sum(dim=1)
- out = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1) + mixed_residual
+ if round_residual:
+ mixed_residual = mixed_residual.to(layer_output.dtype).float()
+ expanded = post_mix.unsqueeze(-1) * layer_output.float().unsqueeze(1)
+ if round_residual:
+ expanded = expanded.to(layer_output.dtype).float()
+ out = expanded + mixed_residual
return out.to(layer_output.dtype).reshape(tokens, streams * hidden)
@@ -98,3 +119,148 @@ def test_mhc_reports_missing_deepgemm(monkeypatch, missing):
with pytest.raises(ImportError, match="deep_gemm"):
hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20)
+
+
+@pytest.mark.parametrize("tokens,hidden,prefix", [(1, 4096, "attn"), (19, 780, "ffn"), (3073, 4096, "attn")])
+@pytest.mark.parametrize("seed", [1525, 2026])
+def test_glm_mhc_caller_matches_normalize_first_reference(tokens, hidden, prefix, seed, monkeypatch):
+ from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer
+
+ torch.manual_seed(seed)
+ monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False)
+ monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
+ streams = 4
+ x = torch.randn(tokens, streams * hidden, device="cuda", dtype=torch.bfloat16)
+ fn = torch.randn(24, streams * hidden, device="cuda") * 0.005
+ scale = torch.randn(3, device="cuda")
+ base = torch.randn(24, device="cuda")
+ norm = torch.randn(hidden, device="cuda", dtype=torch.bfloat16)
+ merged, expected_mix, expected_post = _hc_pre_reference(
+ x, fn, scale, base, streams, 1e-5, 1e-6, 20, normalize_before_projection=True
+ )
+ merged = merged.float()
+ expected_input = (merged * torch.rsqrt(merged.square().mean(-1, keepdim=True) + 1e-5) * norm.float()).to(x.dtype)
+ expected_mix, expected_post = expected_mix.to(x.dtype), expected_post.to(x.dtype)
+
+ layer = SimpleNamespace(mhc_streams=streams, eps_=1e-5, hc_eps=1e-6, hc_sinkhorn_iters=20)
+ weights = SimpleNamespace(
+ **{
+ f"hc_{prefix}_{name}": SimpleNamespace(weight=value)
+ for name, value in (("fn", fn), ("scale", scale), ("base", base))
+ }
+ )
+ monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", True)
+ actual_input, actual_mix, actual_post = Glm5NextTransformerLayerInfer._hc_pre(
+ layer, x, weights, prefix, SimpleNamespace(weight=norm)
+ )
+ assert torch.backends.cuda.matmul.allow_tf32
+ assert actual_mix.dtype == actual_post.dtype == x.dtype
+ # FP32 reduction differences can cross the merged BF16 midpoint; RMSNorm
+ # rounds to BF16 again. Check FP32 projection separately at strict tolerance.
+ torch.testing.assert_close(actual_input, expected_input, atol=0.004, rtol=2 * torch.finfo(x.dtype).eps)
+ torch.testing.assert_close(actual_mix, expected_mix, atol=1e-5, rtol=0.008)
+ torch.testing.assert_close(actual_post, expected_post, atol=1e-5, rtol=0.008)
+ layer_out = torch.randn(tokens, hidden, device="cuda", dtype=x.dtype)
+ torch.testing.assert_close(
+ hc_post(layer_out, x, actual_mix, actual_post, streams, glm5_numerics=True),
+ _hc_post_reference(layer_out, x, actual_mix, actual_post, streams, round_residual=True),
+ atol=0.004,
+ rtol=0.008,
+ )
+
+
+@pytest.mark.parametrize("tokens,hidden", [(1, 4096), (19, 780), (3073, 4096)])
+def test_glm_mhc_projection_retains_fp32_inputs(tokens, hidden, monkeypatch):
+ monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
+ x = torch.randn(tokens, 4 * hidden, device="cuda", dtype=torch.bfloat16)
+ fn = torch.randn(24, 4 * hidden, device="cuda") * 0.005
+ x_fp32 = x.float()
+ expected = F.linear(x_fp32 * torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5), fn)
+
+ monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", True)
+ partials = pre_norm._hc_normalized_projection(x, fn, 1e-5)
+ assert torch.backends.cuda.matmul.allow_tf32
+ torch.testing.assert_close(partials.sum(0), expected, atol=2e-5, rtol=2e-5)
+
+
+def test_glm_mhc_normalizes_before_cancelling_large_products(monkeypatch):
+ monkeypatch.setattr(pre_norm, "LIGHTLLM_DISABLE_DEEPGEMM_MHC", False)
+ monkeypatch.setitem(sys.modules, "deep_gemm", None)
+ monkeypatch.setattr(torch.backends.cuda.matmul, "allow_tf32", False)
+ # Project-first overflows; normalize-first has finite cancelling products.
+ x = torch.full((1, 4 * 17), 2.0 ** 32, device="cuda", dtype=torch.bfloat16)
+ fn = torch.zeros(24, x.shape[-1], device="cuda")
+ fn[:, 0], fn[:, 1] = 2.0 ** 110, -(2.0 ** 110)
+ x_fp32 = x.float()
+ inv_rms = torch.rsqrt(x_fp32.square().mean(-1, keepdim=True) + 1e-5)
+ expected = F.linear(x_fp32 * inv_rms, fn)
+ old_order = F.linear(x_fp32, fn) * inv_rms
+ assert torch.isfinite(expected).all()
+ assert not torch.isfinite(old_order).all()
+ torch.testing.assert_close(pre_norm._hc_normalized_projection(x, fn, 1e-5).sum(0), expected, atol=0, rtol=0)
+
+
+def test_glm_mhc_rounds_post_coefficients_before_mixing():
+ x = torch.full((1, 4 * 17), 1 / 256, device="cuda", dtype=torch.bfloat16)
+ fn = torch.zeros(24, x.shape[-1], device="cuda")
+ scale = torch.zeros(3, device="cuda")
+ base = torch.zeros(24, device="cuda")
+ base[4:8] = math.log(1.003 / 0.997)
+ norm = torch.ones(17, device="cuda", dtype=x.dtype)
+ _, mix, post = hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20, glm5_numerics=True)
+ assert mix.dtype == post.dtype == x.dtype
+ torch.testing.assert_close(post, torch.ones_like(post), atol=0, rtol=0)
+ torch.testing.assert_close(mix, torch.full_like(mix, 0.25), atol=0, rtol=0)
+ layer_out = torch.ones(1, 17, device="cuda", dtype=x.dtype)
+ actual = hc_post(layer_out, x, mix, post, 4, glm5_numerics=True)
+ torch.testing.assert_close(actual, torch.ones_like(actual), atol=0, rtol=0)
+ _, unrounded_mix, unrounded_post = _hc_pre_reference(
+ x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True
+ )
+ old_result = _hc_post_reference(layer_out, x, unrounded_mix, unrounded_post, 4)
+ assert (old_result != actual).all()
+
+
+def test_glm_mhc_post_rounds_residual_like_training():
+ x = torch.zeros(1, 4, 17, device="cuda", dtype=torch.bfloat16)
+ x[:, 0], x[:, 1] = 1, 1 / 256
+ mix = torch.full((1, 4, 4), 0.25, device="cuda", dtype=x.dtype)
+ post = torch.ones(1, 4, device="cuda", dtype=x.dtype)
+ layer_out = torch.full((1, 17), 3 / 512, device="cuda", dtype=x.dtype)
+ # The native training bmm rounds 0.2509765625 to 0.25 before expansion.
+ mixed = torch.bmm(mix.transpose(1, 2), x)
+ expected = torch.addcmul(mixed.float(), layer_out[:, None].float(), post[:, :, None].float()).to(x.dtype)
+ actual = hc_post(layer_out, x.flatten(1), mix, post, 4, glm5_numerics=True)
+ torch.testing.assert_close(actual, expected.flatten(1), atol=0, rtol=0)
+ assert (hc_post(layer_out, x.flatten(1), mix, post, 4) != actual).all()
+
+
+def test_glm_mhc_keeps_pre_coefficients_in_fp32():
+ x = torch.zeros(1, 4, 17, device="cuda", dtype=torch.bfloat16)
+ x[:, 0], x[:, 1] = 1, -1
+ x = x.flatten(1)
+ fn = torch.zeros(24, x.shape[-1], device="cuda")
+ scale = torch.zeros(3, device="cuda")
+ base = torch.zeros(24, device="cuda")
+ # Both coefficients round to BF16 0.5, which would erase this input entirely.
+ pre_values = torch.tensor([0.5005, 0.5], device="cuda") - 1e-6
+ base[:2] = torch.logit(pre_values)
+ norm = torch.ones(17, device="cuda", dtype=x.dtype)
+ actual, _, _ = hc_pre_norm(x, fn, scale, base, norm, 4, 1e-5, 1e-5, 1e-6, 20, glm5_numerics=True)
+ merged, _, _ = _hc_pre_reference(x, fn, scale, base, 4, 1e-5, 1e-6, 20, normalize_before_projection=True)
+ merged = merged.float()
+ expected = (merged * torch.rsqrt(merged.square().mean(-1, keepdim=True) + 1e-5)).to(x.dtype)
+ assert (expected > 0.1).all()
+ torch.testing.assert_close(actual, expected, atol=0.001, rtol=0.008)
+
+
+def test_glm_mhc_rounds_post_product_before_residual_add():
+ x = torch.full((1, 4 * 17), -1.015625, device="cuda", dtype=torch.bfloat16)
+ layer = torch.full((1, 17), 1.0078125, device="cuda", dtype=x.dtype)
+ mix = torch.eye(4, device="cuda", dtype=x.dtype).unsqueeze(0)
+ post = torch.full((1, 4), 1.0078125, device="cuda", dtype=x.dtype)
+ expected = _hc_post_reference(layer, x, mix, post, 4, round_residual=True)
+ actual = hc_post(layer, x, mix, post, 4, glm5_numerics=True)
+ torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+ assert torch.count_nonzero(actual) == 0
+ assert torch.count_nonzero(hc_post(layer, x, mix, post, 4)) > 0
diff --git a/unit_tests/common/fused_moe/test_deepgemm_activation.py b/unit_tests/common/fused_moe/test_deepgemm_activation.py
index e9ba2ab26e..2df337f6b2 100644
--- a/unit_tests/common/fused_moe/test_deepgemm_activation.py
+++ b/unit_tests/common/fused_moe/test_deepgemm_activation.py
@@ -14,7 +14,12 @@
@pytest.mark.parametrize("is_prefill", [False, True])
@pytest.mark.parametrize("use_ue8m0_scales", [False, True])
@pytest.mark.parametrize("alpha,limit,add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)])
-def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one):
+@pytest.mark.parametrize("activation_format", ["auto", "fp32", "ue8m0"])
+def test_grouped_gemm_activation_matches_reference(
+ monkeypatch, is_prefill, use_ue8m0_scales, alpha, limit, add_one, activation_format
+):
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format)
+ activation_ue8m0 = use_ue8m0_scales if activation_format == "auto" else activation_format == "ue8m0"
monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0")
monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES")
config = {"quant_method": "fp8", "weight_block_size": [128, 128]}
@@ -31,6 +36,24 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_
w2_scale = torch.ones(2, 1, 1, device="cuda")
counts = torch.tensor([3, 2], device="cuda", dtype=torch.int32)
activation_args = dict(alpha=alpha, limit=limit, clamp_up_add_one=add_one)
+ captured_scales = []
+ quantize = ep.per_token_group_quant_fp8
+ masked_quantize = ep.silu_and_mul_masked_post_quant_fwd
+
+ def capture_contiguous(*args, **kwargs):
+ assert kwargs["use_ue8m0_scales"] == activation_ue8m0
+ result = quantize(*args, **kwargs)
+ count = (3, 2)[len(captured_scales)]
+ captured_scales.append(result[1][:count].clone())
+ return result
+
+ def capture_masked(*args, **kwargs):
+ result = masked_quantize(*args, **kwargs)
+ captured_scales.append(torch.cat([args[2][0, :3], args[2][1, :2]]).clone())
+ return result
+
+ monkeypatch.setattr(ep, "per_token_group_quant_fp8", capture_contiguous)
+ monkeypatch.setattr(ep, "silu_and_mul_masked_post_quant_fwd", capture_masked)
recv = torch.zeros(2, padded, dim, device="cuda", dtype=torch.float8_e4m3fn)
recv[0, :3] = x[:3]
recv[1, :2] = x[3:]
@@ -77,7 +100,8 @@ def test_grouped_gemm_activation_matches_reference(monkeypatch, is_prefill, use_
gate = (gate * torch.sigmoid(gate * (alpha or 1.0))).bfloat16().float()
activation = (gate * up).bfloat16().float()
scales = activation.abs().amax(-1, keepdim=True).clamp(min=1e-10) / 448
- if use_ue8m0_scales:
+ if activation_ue8m0:
scales = torch.exp2(torch.ceil(torch.log2(scales)))
+ torch.testing.assert_close(torch.cat(captured_scales), scales, atol=1e-6, rtol=1e-5)
expected = ((activation / scales).to(torch.float8_e4m3fn).float() * scales).bfloat16()
torch.testing.assert_close(actual, expected, atol=0.05, rtol=0.02)
diff --git a/unit_tests/common/quantization/test_activation_scale_override.py b/unit_tests/common/quantization/test_activation_scale_override.py
new file mode 100644
index 0000000000..21fc9e522a
--- /dev/null
+++ b/unit_tests/common/quantization/test_activation_scale_override.py
@@ -0,0 +1,149 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from lightllm.common.quantization import Quantcfg, deepgemm
+from lightllm.common.quantization.quantize_method import WeightPack
+from lightllm.utils.envs_utils import get_env_use_ue8m0_activation_scales, get_env_use_ue8m0_scales
+
+
+@pytest.mark.parametrize("scale_format", [None, "auto", "fp32", "ue8m0", " FP32 "])
+@pytest.mark.parametrize("legacy", [False, True])
+@pytest.mark.parametrize("default", [None, False, True])
+def test_activation_override_preserves_weight_policy(monkeypatch, scale_format, legacy, default):
+ monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(legacy)))
+ if scale_format is None:
+ monkeypatch.delenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", raising=False)
+ else:
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format)
+ expected = legacy if default is None else default
+ if scale_format is not None and scale_format.strip().lower() != "auto":
+ expected = scale_format.strip().lower() == "ue8m0"
+ assert get_env_use_ue8m0_activation_scales(default) == expected
+ assert get_env_use_ue8m0_scales() == legacy
+
+
+def test_invalid_or_unsupported_activation_format_fails_explicitly(monkeypatch):
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "typo")
+ with pytest.raises(ValueError, match="LIGHTLLM_ACTIVATION_SCALE_FORMAT"):
+ get_env_use_ue8m0_activation_scales()
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", "fp32")
+ with pytest.raises(ValueError, match="requires UE8M0"):
+ get_env_use_ue8m0_activation_scales(True, require_ue8m0=True)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or not deepgemm.HAS_DEEPGEMM, reason="CUDA + DeepGEMM required")
+@pytest.mark.parametrize("weight_ue8m0", [False, True])
+def test_dense_activation_override_reuses_exact_weight_pack(monkeypatch, weight_ue8m0):
+ if torch.cuda.get_device_capability()[0] != 9:
+ pytest.skip("Continuous FP32 activation scales target Hopper")
+ monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0")
+ config = {
+ "n_layer": 1,
+ "quantization_config": {
+ "quant_method": "fp8",
+ "weight_block_size": [128, 128],
+ "scale_fmt": "ue8m0" if weight_ue8m0 else "float32",
+ },
+ }
+ method = Quantcfg(config).get_quant_method(0, "q_proj")
+ pack, _ = method.create_weight([128], 512, torch.bfloat16, 0)
+ torch.manual_seed(54)
+ method.load_weight(torch.randn(128, 512, device="cuda", dtype=torch.bfloat16), pack)
+ before_weight, before_scale = pack.weight.view(torch.uint8).clone(), pack.weight_scale.clone()
+ x = torch.randn(17, 512, device="cuda", dtype=torch.bfloat16)
+ quantize = deepgemm.per_token_group_quant_fp8
+ calls = []
+
+ def capture(values, group_size, **kwargs):
+ result = quantize(values, group_size, **kwargs)
+ expected = values.float().reshape(17, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448
+ if kwargs["use_ue8m0_scales"]:
+ expected = expected.log2().ceil().exp2()
+ assert group_size == 128 and result[1].dtype == torch.float32
+ torch.testing.assert_close(result[1], expected, rtol=1e-6, atol=0)
+ calls.append(kwargs["use_ue8m0_scales"])
+ return result
+
+ monkeypatch.setattr(deepgemm, "per_token_group_quant_fp8", capture)
+ for scale_format in ("auto", "fp32", "ue8m0"):
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format)
+ result = method.apply(x, pack, use_custom_tensor_mananger=False)
+ assert torch.isfinite(result).all()
+ assert torch.equal(pack.weight.view(torch.uint8), before_weight)
+ assert torch.equal(pack.weight_scale, before_scale)
+ assert method.use_ue8m0_scales == weight_ue8m0
+ assert calls == [weight_ue8m0, False, True]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"])
+def test_fallback_block_quantization_honors_override(monkeypatch, scale_format):
+ from lightllm.common.quantization.w8a8 import FP8w8a8B128QuantizationMethod
+
+ monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1")
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format)
+ method = object.__new__(FP8w8a8B128QuantizationMethod)
+ method.block_size = 128
+ x = torch.randn(3, 512, device="cuda", dtype=torch.bfloat16)
+ pack = WeightPack(
+ weight=torch.ones(128, 512, device="cuda", dtype=torch.float8_e4m3fn),
+ weight_scale=torch.full((1, 4), 0.125, device="cuda"),
+ )
+ result = method._dynamic_quant_input(x, pack, use_custom_tensor_mananger=False)
+ expected = x.float().reshape(3, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448
+ if scale_format == "ue8m0":
+ expected = expected.log2().ceil().exp2()
+ torch.testing.assert_close(result[2], expected, rtol=1e-6, atol=0)
+ assert torch.all(pack.weight_scale == 0.125)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"])
+@pytest.mark.parametrize("legacy", [False, True])
+@pytest.mark.parametrize("overlap", [False, True])
+def test_deepep_dispatch_uses_activation_override(monkeypatch, scale_format, legacy, overlap):
+ from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep
+ from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import FuseMoeDeepGEMM
+
+ monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", str(int(legacy)))
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format)
+ expected = legacy if scale_format == "auto" else scale_format == "ue8m0"
+ hidden = torch.zeros(3, 128, device="cuda", dtype=torch.bfloat16)
+ ids = torch.zeros(3, 1, device="cuda", dtype=torch.int64)
+ weights = torch.ones(3, 1, device="cuda")
+ calls = []
+
+ def dispatch(*args, **kwargs):
+ calls.append(kwargs["round_scale"])
+ return object(), object(), object(), None, None
+
+ buffer = SimpleNamespace(
+ group_size=1,
+ low_latency_dispatch=dispatch,
+ low_latency_combine=lambda *args, **kwargs: (hidden, None, None),
+ )
+ monkeypatch.setattr(ep.dist_group_manager, "ep_low_latency_buffer", buffer, raising=False)
+ if overlap:
+ impl = object.__new__(FuseMoeDeepGEMM)
+ impl.quant_method = SimpleNamespace(method_name="fp8w8a8-b128-deepgemm")
+ impl.total_expert_num_contain_redundancy = 1
+ impl._select_experts = lambda **kwargs: (weights, ids, ids)
+ impl.low_latency_dispatch(hidden, hidden, None, False, 1, True, 1, 1, "sigmoid")
+ else:
+ monkeypatch.setattr(ep, "masked_group_gemm", lambda *args, **kwargs: hidden)
+ ep.fused_experts_impl(
+ hidden,
+ torch.zeros(1, 256, 128, device="cuda", dtype=torch.float8_e4m3fn),
+ torch.zeros(1, 128, 128, device="cuda", dtype=torch.float8_e4m3fn),
+ weights,
+ ids,
+ buffer=buffer,
+ is_prefill=False,
+ num_experts=1,
+ use_fp8_w8a8=True,
+ use_fp8_all2all=True,
+ block_size_k=128,
+ )
+ assert calls == [expected]
diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py
index cf87ac3ffb..0bdfbcd63f 100644
--- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py
+++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py
@@ -266,13 +266,66 @@ def test_weight_scale_loading_preserves_fused_views(
assert len(packed_weight_layout) == (2 * (experts or 1) if packed else 0)
+@pytest.mark.parametrize("use_ue8m0", [False, True])
+@pytest.mark.parametrize("shared_expert", [False, True])
+@pytest.mark.parametrize("activation_format", ["auto", "fp32", "ue8m0"])
+def test_tp_moe_activation_scale_rule(monkeypatch, use_ue8m0, shared_expert, activation_format):
+ from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe as moe
+ from lightllm.common.basemodel.triton_kernel.fused_moe.append_shared_expert_topk import append_fused_shared_experts
+
+ monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "0")
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format)
+ activation_ue8m0 = use_ue8m0 if activation_format == "auto" else activation_format == "ue8m0"
+ Quantcfg(
+ {
+ "n_layer": 1,
+ "quantization_config": {
+ "quant_method": "fp8",
+ "weight_block_size": [128, 128],
+ "scale_fmt": "ue8m0" if use_ue8m0 else "float32",
+ },
+ }
+ )
+ torch.manual_seed(20261023)
+ x = torch.randn((17, 512), device="cuda", dtype=torch.bfloat16)
+ w1, s1 = weight.weight_quant(torch.randn((3, 512, 512), device="cuda", dtype=torch.bfloat16) * 0.02)
+ w2, s2 = weight.weight_quant(torch.randn((3, 512, 256), device="cuda", dtype=torch.bfloat16) * 0.02)
+ ids = torch.tensor([0, 1], device="cuda", dtype=torch.int32).repeat(17, 1)
+ scores = torch.tensor([0.3, 0.7], device="cuda").repeat(17, 1)
+ if shared_expert:
+ scores, ids = append_fused_shared_experts(scores, ids, shared_expert_start_id=2, num_fused_shared_experts=1)
+ calls = []
+ quantize = moe.per_token_group_quant_fp8
+
+ def check_activation(values, group_size, **kwargs):
+ q, scales = quantize(values, group_size, **kwargs)
+ assert group_size == 128
+ assert scales.dtype == torch.float32
+ expected = values.float().reshape(values.shape[0], -1, 128).abs().amax(-1).clamp_min(1e-10) / 448.0
+ if activation_ue8m0:
+ expected = torch.exp2(torch.ceil(torch.log2(expected.double()))).float()
+ torch.testing.assert_close(scales, expected, rtol=1e-6, atol=0)
+ calls.append(values.shape[1])
+ return q, scales
+
+ monkeypatch.setattr(moe, "per_token_group_quant_fp8", check_activation)
+ out = moe.fused_experts_impl(x, w1, w2, scores, ids, use_fp8_w8a8=True, w1_scale=s1, w2_scale=s2)
+ assert calls == [512, 256] # W13 and W2, including the optional shared-expert slot.
+ assert torch.isfinite(out).all()
+
+
@pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM")
+@pytest.mark.parametrize("packed_env", [False, True])
@pytest.mark.parametrize("need_trans", [False, True])
@pytest.mark.parametrize("k", [512, 640])
-def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout, need_trans, k):
+@pytest.mark.parametrize("activation_format", ["auto", "fp32"])
+def test_tp_moe_consumes_packed_weight_scales(
+ monkeypatch, packed_weight_layout, need_trans, k, packed_env, activation_format
+):
from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe import grouped_matmul
monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1")
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format)
monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: True)
method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj")
full, _ = method.create_moe_weight([256], k, torch.bfloat16, 0, 2)
@@ -287,6 +340,7 @@ def test_tp_moe_consumes_packed_weight_scales(monkeypatch, packed_weight_layout,
indices[1, :8] = torch.arange(9, 17, device="cuda")
outputs = []
for sf in (scales, full.weight_scale):
+ monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1" if sf.dtype != torch.int32 or packed_env else "0")
out = torch.empty((17, 256), device="cuda", dtype=torch.bfloat16)
grouped_matmul(
17,
@@ -442,11 +496,15 @@ def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout
@pytest.mark.parametrize("is_prefill", [False, True])
@pytest.mark.parametrize("sm100", [False, True])
@pytest.mark.parametrize("scale_source", ["none", "env", "config"])
-def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_prefill, sm100, scale_source):
+@pytest.mark.parametrize("activation_format", ["auto", "fp32", "ue8m0"])
+def test_ep_moe_rounds_scales_and_preserves_logical_block_size(
+ monkeypatch, is_prefill, sm100, scale_source, activation_format
+):
from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep
from lightllm.utils.envs_utils import get_env_use_ue8m0_scales
monkeypatch.delenv("LIGHTLLM_USE_UE8M0_SCALES", raising=False)
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", activation_format)
assert not get_env_use_ue8m0_scales()
config = {"n_layer": 1}
if scale_source == "env":
@@ -454,6 +512,7 @@ def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_p
elif scale_source == "config":
config["quantization_config"] = {"quant_method": "fp8", "weight_block_size": [128, 128], "scale_fmt": "ue8m0"}
use_ue8m0_scales = scale_source != "none"
+ activation_ue8m0 = use_ue8m0_scales if activation_format == "auto" else activation_format == "ue8m0"
monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100)
method = Quantcfg(config, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj")
assert get_env_use_ue8m0_scales() == use_ue8m0_scales
@@ -464,7 +523,7 @@ def test_ep_moe_rounds_scales_and_preserves_logical_block_size(monkeypatch, is_p
assert q.shape == x.shape and sf.shape == (3, 4)
assert sf.dtype == torch.float32
expected_scales = x.float().reshape(3, 4, 128).abs().amax(-1).clamp_min(1e-10) / 448
- if use_ue8m0_scales:
+ if activation_ue8m0:
expected_scales = torch.exp2(torch.ceil(torch.log2(expected_scales)))
expected_q = (x.float().reshape(3, 4, 128) / expected_scales.unsqueeze(-1)).to(q.dtype).reshape_as(q)
torch.testing.assert_close(q.float(), expected_q.float(), atol=0, rtol=0)
diff --git a/unit_tests/models/glm5_next/test_indexer_scale.py b/unit_tests/models/glm5_next/test_indexer_scale.py
new file mode 100644
index 0000000000..a6bcd30cd5
--- /dev/null
+++ b/unit_tests/models/glm5_next/test_indexer_scale.py
@@ -0,0 +1,30 @@
+import pytest
+import torch
+
+from lightllm.models.glm5_next.triton_kernel.index_quant import hadamard_transform_quant_fp8
+
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"])
+@pytest.mark.parametrize("rows", [0, 1, 33])
+def test_indexer_scale_matches_hadamard_reference(monkeypatch, scale_format, rows):
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format)
+ # FP32 matrix products are exact for these power-of-two input fractions.
+ x = (torch.randint(-64, 64, (rows, 128), device="cuda").float() / 64).bfloat16()
+ actual, scales = hadamard_transform_quant_fp8(x, 128 ** -0.5)
+ assert actual.dtype == torch.float8_e4m3fn and scales.dtype == torch.float32
+ assert actual.shape == x.shape and scales.shape == (rows, 1)
+ if not rows:
+ return
+ matrix = torch.ones((1, 1), device="cuda")
+ for _ in range(7):
+ matrix = torch.cat((torch.cat((matrix, matrix), 1), torch.cat((matrix, -matrix), 1)), 0)
+ transformed = ((x.float() @ matrix) * 128 ** -0.5).bfloat16().float()
+ expected_scale = transformed.abs().amax(-1, keepdim=True).clamp_min(1e-4) / 448
+ if scale_format != "fp32":
+ expected_scale = expected_scale.log2().ceil().exp2()
+ torch.testing.assert_close(scales, expected_scale, rtol=1e-6, atol=0)
+ expected = (transformed / expected_scale).clamp(-448, 448).to(torch.float8_e4m3fn)
+ assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py
index 1a5e32567b..c87afd3f2a 100644
--- a/unit_tests/models/glm5_next/test_kernels.py
+++ b/unit_tests/models/glm5_next/test_kernels.py
@@ -182,7 +182,9 @@ def set_inputs(counts, phase):
@pytest.mark.parametrize("cuda_graph", [False, True])
-def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph):
+@pytest.mark.parametrize("scale_format", ["auto", "fp32", "ue8m0"])
+def test_kpool_chunk_boundaries_and_fragmented_token_kv(monkeypatch, cuda_graph, scale_format):
+ monkeypatch.setenv("LIGHTLLM_ACTIVATION_SCALE_FORMAT", scale_format)
# Unaligned chunks, reordered requests, and pools spanning old/new tails.
seqs = [9, 7]
ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)]
diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py
index 57665cbbe3..39ee6c039c 100644
--- a/unit_tests/models/glm5_next/test_pd_cache.py
+++ b/unit_tests/models/glm5_next/test_pd_cache.py
@@ -73,7 +73,7 @@ def create(tp):
return create
-@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1)])
+@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (4, 4), (1, 4), (4, 1), (8, 1)])
@pytest.mark.parametrize("remainder", range(4))
def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, remainder):
source, dest = make_mems(prefill_tp), make_mems(decode_tp)
diff --git a/unit_tests/models/glm5_next/test_precision_alignment.py b/unit_tests/models/glm5_next/test_precision_alignment.py
new file mode 100644
index 0000000000..309e56bf05
--- /dev/null
+++ b/unit_tests/models/glm5_next/test_precision_alignment.py
@@ -0,0 +1,198 @@
+import dataclasses
+import json
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import FuseMoeDeepGEMM
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.triton_impl import FuseMoeTriton
+from lightllm.common.quantization import Quantcfg
+from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer
+from lightllm.models.glm5_next.layer_weights.transformer_layer_weight import Glm5NextTransformerLayerWeight
+from lightllm.server.core.objs.start_args_type import StartArgs
+from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_ids, get_redundancy_expert_num
+
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.fixture(autouse=True)
+def runtime(monkeypatch):
+ monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(StartArgs())))
+ for name, value in {
+ "GLOBAL_RANK": 0,
+ "GLOBAL_WORLD_SIZE": 1,
+ "DP_WORLD_SIZE": 1,
+ "CURRENT_RANK_IN_DP": 0,
+ "CURRENT_RANK_IN_NODE": 0,
+ "CURRENT_DEVICE_ID": 0,
+ }.items():
+ monkeypatch.setenv("LIGHTLLM_" + name, str(value))
+ for getter in (get_env_start_args, get_redundancy_expert_ids, get_redundancy_expert_num):
+ getter.cache_clear()
+ yield
+ for getter in (get_env_start_args, get_redundancy_expert_ids, get_redundancy_expert_num):
+ getter.cache_clear()
+
+
+@pytest.fixture
+def config():
+ return {
+ "hidden_size": 128,
+ "n_embed": 128,
+ "intermediate_size": 128,
+ "moe_intermediate_size": 128,
+ "num_hidden_layers": 4,
+ "n_layer": 4,
+ "num_attention_heads": 4,
+ "num_key_value_heads": 4,
+ "rms_norm_eps": 1e-5,
+ "q_lora_rank": 128,
+ "kv_lora_rank": 512,
+ "qk_nope_head_dim": 128,
+ "qk_rope_head_dim": 0,
+ "v_head_dim": 128,
+ "index_n_heads": 32,
+ "index_head_dim": 128,
+ "index_topk": 2048,
+ "index_kpool": 4,
+ "index_kpool_always_select_tail": True,
+ "linear_attn_config": {"num_heads": 4, "head_dim": 128, "short_conv_kernel_size": 4},
+ "layer_types": ["linear_attention"] * 3 + ["deepseek_sparse_attention"],
+ "mhc": False,
+ "n_routed_experts": 4,
+ "n_shared_experts": 1,
+ "first_k_dense_replace": 4,
+ "num_experts_per_tok": 2,
+ "norm_topk_prob": True,
+ "scoring_func": "sigmoid",
+ "n_group": 1,
+ "topk_group": 1,
+ }
+
+
+def test_indexer_layernorm_epsilon_preserves_pooled_topk_boundary(config, monkeypatch):
+ pytest.importorskip("deep_gemm")
+ if torch.cuda.get_device_capability()[0] < 9:
+ pytest.skip("FP8 indexer requires Hopper or newer")
+ monkeypatch.setattr("lightllm.models.glm5_next.indexer.HAS_VLLM", False)
+ dim, heads, topk = config["index_head_dim"], config["index_n_heads"], config["index_topk"]
+ weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config))
+ indexer = Glm5NextNsaInfer(3, config, 1)
+ columns = torch.arange(dim)
+ a = torch.where(columns % 2 == 0, -1.0, 1.0)
+ b = torch.where((columns // 2) % 2 == 0, -1.0, 1.0)
+ query = 2 * a + b
+ prefix = "model.language_model.layers.3.self_attn.indexer"
+ eye = torch.eye(dim, dtype=torch.bfloat16)
+ weight.load_hf_weights(
+ {
+ f"{prefix}.wk.weight": eye,
+ f"{prefix}.k_norm.weight": torch.ones(dim, dtype=torch.bfloat16),
+ f"{prefix}.k_norm.bias": torch.zeros(dim, dtype=torch.bfloat16),
+ f"{prefix}.wq_b.weight": eye.repeat(heads, 1),
+ f"{prefix}.weights_proj.weight": ((a + b) / dim).repeat(heads, 1),
+ f"{prefix}.index_kpool_compress_gate": torch.zeros(dim, dim, dtype=torch.bfloat16),
+ f"{prefix}.index_kpool_compress_ape": torch.zeros(4, dim),
+ }
+ )
+
+ # 513 pools compete for 512 slots: epsilon changes the final selected pool.
+ hidden = torch.cat(
+ (
+ (query * 2 ** -4).repeat(topk - 4, 1),
+ (a * 2 ** -10).repeat(4, 1),
+ (b * 2 ** -4).repeat(4, 1),
+ )
+ ).to(device="cuda", dtype=torch.bfloat16)
+ tokens = hidden.shape[0]
+ q_lora = query.to(device="cuda", dtype=torch.bfloat16).repeat(tokens, 1)
+ projected_k = weight.wk_proj_.mm(hidden)
+ actual_k = weight.k_norm_(projected_k, eps=indexer.eps)
+ expected_k = F.layer_norm(projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=1e-6)
+ torch.testing.assert_close(actual_k, expected_k, atol=0, rtol=0)
+ wrong_k = F.layer_norm(projected_k, (dim,), weight.k_norm_.weight, weight.k_norm_.bias, eps=config["rms_norm_eps"])
+ assert not torch.equal(expected_k, wrong_k)
+
+ storage = torch.zeros(tokens + 1, 1, 584, dtype=torch.bfloat16, device="cuda")
+ tail = torch.zeros(2, 4, 2 * dim, dtype=torch.bfloat16, device="cuda")
+ locations = torch.arange(tokens, dtype=torch.int32, device="cuda")
+ infer = SimpleNamespace(
+ is_prefill=True,
+ max_kv_seq_len=tokens,
+ max_q_seq_len=tokens,
+ b_req_idx=torch.zeros(1, dtype=torch.int32, device="cuda"),
+ b_seq_len=torch.tensor([tokens], dtype=torch.int32, device="cuda"),
+ b1_cu_q_seq_len=torch.tensor([0, tokens], dtype=torch.int32, device="cuda"),
+ req_manager=SimpleNamespace(
+ req_to_token_indexs=locations[None],
+ get_indexer_tail_buffer=lambda _: tail,
+ ),
+ mem_manager=SimpleNamespace(get_indexer_k_buffer=lambda _: storage.view(torch.uint8)[:, :, -132:]),
+ )
+ state = SimpleNamespace(lengths=locations + 1, ks=torch.zeros_like(locations), ragged_mem_index=locations)
+ _, actual = indexer._get_indices(hidden, q_lora, infer, state, weight)
+ torch.testing.assert_close(actual[-1, :topk].sort().values, locations[:topk], atol=0, rtol=0)
+ assert (actual[-1, topk:] == -1).all()
+
+ # Negative control: the previous epsilon selects the other boundary pool.
+ indexer.eps = config["rms_norm_eps"]
+ _, wrong = indexer._get_indices(hidden, q_lora, infer, state, weight)
+ wrong_selection = torch.cat((locations[: topk - 4], locations[topk:]))
+ torch.testing.assert_close(wrong[-1, :topk].sort().values, wrong_selection, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("enable_ep_moe", [False, True], ids=["tp", "ep"])
+def test_checkpoint_correction_bias_preserves_fp32_routing(config, monkeypatch, enable_ep_moe):
+ args = StartArgs(tp=2, enable_ep_moe=enable_ep_moe, enable_fused_shared_experts=False)
+ monkeypatch.setenv("LIGHTLLM_START_ARGS", json.dumps(dataclasses.asdict(args)))
+ # Rank 1 must retain every expert's global correction bias.
+ for name, value in {"GLOBAL_RANK": 1, "GLOBAL_WORLD_SIZE": 2, "DP_WORLD_SIZE": 2, "CURRENT_RANK_IN_DP": 1}.items():
+ monkeypatch.setenv("LIGHTLLM_" + name, str(value))
+ get_env_start_args.cache_clear()
+ config.update(n_routed_experts=288, num_experts_per_tok=8, first_k_dense_replace=1)
+ weight = Glm5NextTransformerLayerWeight(3, torch.bfloat16, config, Quantcfg(config))
+ experts = weight.experts
+ expected_impl = FuseMoeDeepGEMM if enable_ep_moe else FuseMoeTriton
+ assert isinstance(experts.fuse_moe_impl, expected_impl)
+ source_bias = torch.zeros(config["n_routed_experts"], dtype=torch.float32)
+ source_bias[7:9] = torch.tensor([0.1002, 0.1001], dtype=torch.float32)
+ weight.load_hf_weights({"model.language_model.layers.3.mlp.gate.e_score_correction_bias": source_bias})
+ assert experts.e_score_correction_bias.dtype == torch.float32
+ assert experts.e_score_correction_bias.load_ok
+ torch.testing.assert_close(experts.e_score_correction_bias.cpu(), source_bias, atol=0, rtol=0)
+ assert experts.w13.weight.dtype == experts.w2.weight.dtype == torch.bfloat16
+
+ hidden = torch.zeros(1, config["hidden_size"], dtype=torch.bfloat16, device="cuda")
+ logits = torch.full((1, config["n_routed_experts"]), -10.0, dtype=torch.float32, device="cuda")
+ logits[:, :7] = 4.0
+ logits[0, 7:9] = torch.tensor([0.0, 4e-5], dtype=torch.float32, device="cuda")
+
+ def route(bias):
+ return experts.fuse_moe_impl._select_experts(
+ input_tensor=hidden,
+ router_logits=logits,
+ correction_bias=bias,
+ top_k=8,
+ renormalize=True,
+ use_grouped_topk=True,
+ topk_group=1,
+ num_expert_group=1,
+ scoring_func="sigmoid",
+ )
+
+ weights, ids, origin_ids = route(experts.e_score_correction_bias)
+ expected_ids = torch.arange(8, dtype=ids.dtype, device="cuda")[None]
+ torch.testing.assert_close(ids.sort(dim=-1).values, expected_ids, atol=0, rtol=0)
+ torch.testing.assert_close(origin_ids, ids, atol=0, rtol=0)
+ # Correction bias affects selection; weights use the original sigmoid scores.
+ expected_weights = logits.sigmoid().gather(1, ids.long())
+ expected_weights /= expected_weights.sum(-1, keepdim=True)
+ torch.testing.assert_close(weights, expected_weights, atol=1e-7, rtol=1e-6)
+
+ # Negative control: BF16 rounding makes expert 8 displace expert 7.
+ _, rounded_ids, _ = route(experts.e_score_correction_bias.to(torch.bfloat16))
+ rounded_expected = torch.tensor([list(range(7)) + [8]], dtype=rounded_ids.dtype, device="cuda")
+ torch.testing.assert_close(rounded_ids.sort(dim=-1).values, rounded_expected, atol=0, rtol=0)
diff --git a/unit_tests/models/glm5_next/test_shared_combine.py b/unit_tests/models/glm5_next/test_shared_combine.py
new file mode 100644
index 0000000000..67c11f657e
--- /dev/null
+++ b/unit_tests/models/glm5_next/test_shared_combine.py
@@ -0,0 +1,148 @@
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from lightllm.models.glm5_next.layer_infer.transformer_layer_infer import Glm5NextTransformerLayerInfer
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as weight_module
+from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import deepgemm_impl
+from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep
+
+
+@pytest.mark.parametrize("prefill", [False, True])
+@pytest.mark.parametrize("enabled", [False, True])
+@pytest.mark.parametrize("shared_mode", ["absent", "separate", "native_fused"])
+def test_shared_output_is_added_once(prefill, enabled, shared_mode):
+ hidden = torch.zeros(2, 128, dtype=torch.bfloat16)
+ shared = torch.full_like(hidden, 0.25)
+ calls = []
+
+ def experts(input_tensor, **kwargs):
+ calls.append(kwargs)
+ routed = torch.ones_like(input_tensor)
+ return routed if "shared_expert_output" not in kwargs else routed + kwargs["shared_expert_output"]
+
+ layer = SimpleNamespace(
+ embed_dim_=128,
+ n_shared_experts=None if shared_mode == "absent" else 1,
+ ep_shared_combine_bias=enabled,
+ swiglu_limit=7.0,
+ num_experts_per_tok=2,
+ norm_topk_prob=True,
+ n_group=1,
+ topk_group=1,
+ _ffn_tp=lambda *args: shared,
+ )
+ weights = SimpleNamespace(
+ num_fused_shared_experts=int(shared_mode == "native_fused"),
+ moe_gate=SimpleNamespace(data_type_=torch.float32, mm=lambda x: torch.zeros(2, 4)),
+ experts=SimpleNamespace(experts=experts),
+ )
+ result = Glm5NextTransformerLayerInfer._moe_ffn(layer, hidden, SimpleNamespace(is_prefill=prefill), weights)
+ expected = torch.ones_like(hidden) + (shared if shared_mode == "separate" else 0)
+ torch.testing.assert_close(result, expected, rtol=0, atol=0)
+ if enabled and prefill and shared_mode == "separate":
+ assert calls[0]["shared_expert_output"] is shared
+ else:
+ assert "shared_expert_output" not in calls[0]
+
+
+@pytest.mark.parametrize("with_shared", [False, True])
+def test_shared_bias_reaches_normal_combine_backend(monkeypatch, with_shared):
+ hidden = torch.zeros(2, 128, dtype=torch.bfloat16)
+ shared = torch.ones_like(hidden) if with_shared else None
+ ids = torch.zeros(2, 1, dtype=torch.int32)
+ impl = object.__new__(deepgemm_impl.FuseMoeDeepGEMM)
+ impl._select_experts = lambda **kwargs: (torch.ones(2, 1), ids, ids)
+ impl.total_expert_num_contain_redundancy = 8
+ impl.quant_method = SimpleNamespace(block_size=128)
+ pack = SimpleNamespace(weight=hidden, weight_scale=None)
+ calls = []
+
+ def capture(**kwargs):
+ calls.append(kwargs)
+ return hidden
+
+ monkeypatch.setattr(weight_module, "get_moe_capture_callback", lambda *args: None)
+ monkeypatch.setattr(ep, "check_ep_expert_dtype", lambda method: None)
+ monkeypatch.setattr(ep, "use_sm100_mega_moe", lambda method: False)
+ monkeypatch.setattr(ep, "fused_experts_impl", capture)
+ monkeypatch.setattr(ep.dist_group_manager, "ep_buffer", object(), raising=False)
+ weight = SimpleNamespace(
+ layer_num_=3,
+ fuse_moe_impl=impl,
+ w13=pack,
+ w2=pack,
+ e_score_correction_bias=None,
+ scoring_func="sigmoid",
+ per_expert_scale=None,
+ )
+ result = weight_module.FusedMoeWeight.experts(
+ weight,
+ hidden,
+ hidden,
+ top_k=1,
+ renormalize=True,
+ use_grouped_topk=False,
+ topk_group=1,
+ num_expert_group=1,
+ is_prefill=True,
+ shared_expert_output=shared,
+ )
+ assert result is hidden and calls[0]["shared_expert_output"] is shared
+
+
+def test_unsupported_shared_bias_fails_before_dispatch(monkeypatch):
+ x = torch.zeros(1, 128, dtype=torch.bfloat16)
+ args = dict(
+ hidden_states=x,
+ w13=None,
+ w2=None,
+ topk_weights=None,
+ topk_idx=None,
+ num_experts=8,
+ quant_method=None,
+ shared_expert_output=x,
+ )
+ with pytest.raises(ValueError, match="normal prefill combine"):
+ ep.fused_experts(**args, is_prefill=False)
+ monkeypatch.setattr(ep, "check_ep_expert_dtype", lambda method: None)
+ monkeypatch.setattr(ep, "use_sm100_mega_moe", lambda method: True)
+ with pytest.raises(NotImplementedError, match="DeepEP normal combine"):
+ ep.fused_experts(**args, is_prefill=True)
+
+
+@pytest.mark.parametrize("with_shared", [False, True])
+def test_normal_combine_preserves_default_call_contract(monkeypatch, with_shared):
+ x = torch.zeros(2, 128, dtype=torch.bfloat16)
+ bias = torch.ones_like(x) if with_shared else None
+ calls = []
+ handle = SimpleNamespace(num_recv_tokens_per_expert_list=[0], recv_src_metadata=torch.empty(0, 3))
+
+ def combine(*args, **kwargs):
+ calls.append(kwargs)
+ return x, None, None
+
+ buffer = SimpleNamespace(dispatch=lambda *args, **kwargs: ((x, x), None, None, handle, None), combine=combine)
+ monkeypatch.setattr(ep, "per_token_group_quant_fp8", lambda *args, **kwargs: (x, x))
+ monkeypatch.setattr(ep, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 256)
+ monkeypatch.setattr(ep.Autotuner, "is_kernel_autotune_warmup", lambda *args: False)
+ ids = torch.zeros(2, 1, dtype=torch.int64)
+ result = ep.fused_experts_impl(
+ x,
+ torch.empty(1, 256, 128),
+ torch.empty(1, 128, 128),
+ torch.ones(2, 1),
+ ids,
+ num_experts=1,
+ buffer=buffer,
+ is_prefill=True,
+ use_fp8_w8a8=True,
+ use_fp8_all2all=True,
+ shared_expert_output=bias,
+ )
+ assert result is x
+ if with_shared:
+ assert calls[0]["bias"] is bias
+ else:
+ assert "bias" not in calls[0]
diff --git a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py
index 8c265847db..e4f08421ad 100644
--- a/unit_tests/server/router/model_infer/structured_output/test_reasoning.py
+++ b/unit_tests/server/router/model_infer/structured_output/test_reasoning.py
@@ -49,7 +49,8 @@ def encode(text, add_special_tokens):
chat_template_kwargs=template_kwargs,
reasoning_effort=effort,
)
- result = api_openai._get_guided_reasoning_end(request, SimpleNamespace(encode=encode))
+ thinking = api_openai._is_force_thinking_mode(request)
+ result = api_openai._get_guided_reasoning_end(SimpleNamespace(encode=encode), parser, thinking)
assert encoded == ([marker] if marker else [])
assert result == ([7, 8] if marker else [])
@@ -134,8 +135,11 @@ def test_fixed_reasoning_does_not_require_a_template_thinking_switch(monkeypatch
monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: False)
tokenizer = SimpleNamespace(encode=lambda text, **kwargs: list(text.encode()))
request = ChatCompletionRequest(model="test", messages=[])
- assert api_openai._is_force_thinking_mode(request) == (marker is not None)
- assert api_openai._get_guided_reasoning_end(request, tokenizer) == (list(marker.encode()) if marker else [])
+ thinking = api_openai._is_force_thinking_mode(request)
+ assert thinking == (marker is not None)
+ assert api_openai._get_guided_reasoning_end(tokenizer, parser, thinking) == (
+ list(marker.encode()) if marker else []
+ )
@pytest.mark.parametrize("stream", [False, True])
diff --git a/unit_tests/server/test_glm47_tool_parser.py b/unit_tests/server/test_glm47_tool_parser.py
index 8f626d6343..eb57175ac6 100644
--- a/unit_tests/server/test_glm47_tool_parser.py
+++ b/unit_tests/server/test_glm47_tool_parser.py
@@ -84,6 +84,9 @@ def test_text_around_tool_calls_survives_chunk_boundaries(chunk_size):
content, calls = collect_stream(text, chunk_size)
assert content == "Before between after"
assert list(calls) == [0, 1]
+ normal, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text)
+ assert normal == content
+ assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()]
def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name():
@@ -95,7 +98,7 @@ def test_buffered_arguments_keep_the_stream_alive_without_repeating_the_name():
assert calls[0].name == "Write"
assert calls[0].parameters == ""
- value = ' line 1\n{"literal": true}\n' + "x" * 100_000
+ value = ' line 1\n{"literal": true}\nRead\n' + "x" * 100_000
for chunk in (value[:13], value[13:27], value[27:]):
normal, calls = parser.parse_stream_chunk(chunk)
assert normal == ""
@@ -125,3 +128,79 @@ def test_undefined_buffered_tool_does_not_consume_a_call_index():
assert len(calls) == 1
assert calls[0].name == "Read"
assert calls[0].tool_index == 0
+
+
+@pytest.mark.parametrize("chunk_size", [1, 7, 4096])
+def test_tool_delimiters_inside_values_do_not_close_or_create_calls(chunk_size):
+ value = "Read\nliteralliteral"
+ text = (
+ f"Writecontent{value}"
+ "Readfile_path/tmp/a"
+ )
+ normal, calls = collect_stream(text, chunk_size)
+ assert normal == ""
+ assert list(calls) == [0, 1]
+ assert [call["name"] for call in calls.values()] == ["Write", "Read"]
+ assert [json.loads(call["arguments"]) for call in calls.values()] == [{"content": value}, {"file_path": "/tmp/a"}]
+ normal, nonstream = FunctionCallParser(TOOLS, "glm47").parse_non_stream(text)
+ assert normal == ""
+ assert [call.parameters for call in nonstream] == [call["arguments"] for call in calls.values()]
+
+
+@pytest.mark.parametrize(
+ "schema,value,expected",
+ [
+ ({"type": ["string"]}, "123", "123"),
+ ({"anyOf": [{"type": "string"}]}, "false", "false"),
+ ({"oneOf": [{"type": "string"}]}, "null", "null"),
+ ({"allOf": [{"type": "string"}]}, "123", "123"),
+ ({"enum": ["123", "false"]}, "123", "123"),
+ ({"type": ["string", "null"]}, "null", None),
+ ({"anyOf": [{"type": "integer"}, {"type": "string"}]}, "0012", 12),
+ ({"type": "boolean"}, "False", False),
+ ({"type": "boolean"}, "1", True),
+ ({"type": "boolean"}, " 0 ", False),
+ ({"type": "integer"}, "0012", 12),
+ ({"type": "number"}, ".5", 0.5),
+ ({"type": "number"}, "1.0", 1),
+ ({"type": "array"}, '[1, "中", false]', [1, "中", False]),
+ ({"type": "object"}, '{"nested": {"ok": true}}', {"nested": {"ok": True}}),
+ ({}, "123", "123"),
+ (None, "false", "false"),
+ ({"type": "number"}, "NaN", "NaN"),
+ ({"type": "number"}, "Infinity", "Infinity"),
+ ({"type": "number"}, "1e999", "1e999"),
+ ({"type": "array"}, "[1e999]", "[1e999]"),
+ ({"type": "object"}, '{"x": NaN}', '{"x": NaN}'),
+ ],
+)
+def test_schema_coercion_preserves_declared_types_and_finite_json(schema, value, expected):
+ tools = [
+ Tool.model_validate(
+ {
+ "type": "function",
+ "function": {
+ "name": "Probe",
+ "parameters": {
+ "type": "object",
+ "properties": {"value": schema} if schema is not None else {},
+ },
+ },
+ }
+ )
+ ]
+ text = f"Probevalue{value}"
+ parser = FunctionCallParser(tools, "glm47")
+ normal, calls = parser.parse_non_stream(text)
+ assert normal == ""
+ assert len(calls) == 1
+ parsed = json.loads(calls[0].parameters)["value"]
+ assert parsed == expected
+ assert type(parsed) is type(expected)
+ json.dumps(parsed, allow_nan=False)
+ parameters = ""
+ for char in text:
+ normal, deltas = parser.parse_stream_chunk(char)
+ assert normal == ""
+ parameters += "".join(delta.parameters for delta in deltas)
+ assert parameters == calls[0].parameters
diff --git a/unit_tests/server/test_glm_reasoning_tool_calls.py b/unit_tests/server/test_glm_reasoning_tool_calls.py
new file mode 100644
index 0000000000..92bc3e8496
--- /dev/null
+++ b/unit_tests/server/test_glm_reasoning_tool_calls.py
@@ -0,0 +1,328 @@
+import asyncio
+import json
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+
+from lightllm.server import api_http, api_openai, build_prompt, reasoning_parser
+from lightllm.server.api_models import ChatCompletionRequest, Tool
+from lightllm.server.core.objs import FinishStatus
+from lightllm.server.core.objs.start_args_type import StartArgs
+from lightllm.server.function_call_parser import FunctionCallParser
+from lightllm.server.httpserver.stop_seq_match import ReasoningStateMachine
+from lightllm.server.reasoning_parser import ReasoningParser
+
+
+TOOLS = [
+ Tool.model_validate(
+ {
+ "type": "function",
+ "function": {
+ "name": "Write",
+ "parameters": {"type": "object", "properties": {"content": {"type": "string"}}},
+ },
+ }
+ )
+]
+VALUE = ' 中文 "quoted" \\path\nliteralRead\n '
+TOOL_CALL = f"Writecontent{VALUE}"
+
+
+@pytest.fixture(autouse=True)
+def mock_token_ids(monkeypatch):
+ monkeypatch.setattr(
+ reasoning_parser, "get_token_id", {"": 98, "": 99, "": 77}.__getitem__
+ )
+
+
+@pytest.fixture
+def complete_chat(monkeypatch):
+ args = StartArgs(reasoning_parser="glm45", tool_call_parser="glm47")
+ monkeypatch.setattr(api_openai, "get_env_start_args", lambda: args)
+ monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True)
+
+ def complete(
+ chunks,
+ finish_status=FinishStatus.FINISHED_STOP,
+ prompt="<|assistant|>",
+ parser_name="glm45",
+ compiler=None,
+ **request_kwargs,
+ ):
+ args.reasoning_parser = parser_name
+ monkeypatch.setattr(api_openai, "build_prompt", AsyncMock(return_value=prompt))
+ captured_params = []
+
+ async def generate(_prompt, sampling_params, *_args, **_kwargs):
+ captured_params.append(sampling_params)
+ for index, (text, token_id) in enumerate(chunks):
+ finish = FinishStatus(finish_status) if index == len(chunks) - 1 else FinishStatus()
+ yield 0, text, {"id": token_id, "prompt_tokens": 1}, finish
+
+ tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [99]}[text])
+ manager = SimpleNamespace(tokenizer=tokenizer, generate=generate, output_grammar_compiler=compiler)
+ monkeypatch.setattr(api_http, "g_objs", SimpleNamespace(args=args, httpserver_manager=manager))
+ request = ChatCompletionRequest(model="glm", messages=[], tools=TOOLS, **request_kwargs)
+
+ async def run():
+ response = await api_openai.chat_completions_impl(request, None)
+ result = SimpleNamespace(reasoning="", content="", calls=[], finish_reason=None, sampling_params=None)
+ if not request.stream:
+ choice = response.choices[0]
+ result.reasoning = choice.message.reasoning or ""
+ result.content = choice.message.content or ""
+ result.calls = [
+ {"id": call.id, "name": call.function.name, "arguments": call.function.arguments}
+ for call in choice.message.tool_calls or []
+ ]
+ result.finish_reason = choice.finish_reason
+ result.sampling_params = captured_params[0]
+ return result
+
+ calls = {}
+ done = False
+ async for chunk in response.body_iterator:
+ chunk = chunk.decode() if isinstance(chunk, bytes) else chunk
+ if chunk.strip() == "data: [DONE]":
+ done = True
+ continue
+ for choice in json.loads(chunk.removeprefix("data: "))["choices"]:
+ delta = choice["delta"]
+ result.reasoning += delta.get("reasoning") or ""
+ result.content += delta.get("content") or ""
+ for call in delta.get("tool_calls") or []:
+ index = call["index"]
+ name = call["function"].get("name")
+ if name is not None:
+ assert index not in calls # Send each name and ID once.
+ assert call["id"]
+ calls[index] = {"id": call["id"], "name": name, "arguments": ""}
+ assert index in calls
+ if call.get("id") is not None:
+ assert call["id"] == calls[index]["id"]
+ calls[index]["arguments"] += call["function"].get("arguments") or ""
+ result.finish_reason = choice.get("finish_reason") or result.finish_reason
+ assert done
+ assert list(calls) == list(range(len(calls)))
+ result.calls = list(calls.values())
+ result.sampling_params = captured_params[0]
+ return result
+
+ return asyncio.run(run())
+
+ return complete
+
+
+@pytest.mark.parametrize(
+ "prefix,force_reasoning",
+ [
+ ("Need a file.", False),
+ ("Need a file.", True),
+ ("Need a file.", False),
+ ("Need a file.", True),
+ ],
+)
+@pytest.mark.parametrize("chunk_size", [1, 7, 4096])
+def test_glm_reasoning_and_tools_agree_across_stream_boundaries(prefix, force_reasoning, chunk_size):
+ text = prefix + TOOL_CALL + TOOL_CALL
+ parser = ReasoningParser("glm45", force_reasoning=force_reasoning)
+ reasoning, content = parser.parse_non_stream(text)
+ assert reasoning == "Need a file."
+ assert content == TOOL_CALL + TOOL_CALL
+
+ parser = ReasoningParser("glm45", force_reasoning=force_reasoning)
+ tool_parser = FunctionCallParser(TOOLS, "glm47")
+ reasoning_parts, content_parts, calls = [], [], {}
+ for start in range(0, len(text), chunk_size):
+ reasoning, content = parser.parse_stream_chunk(text[start : start + chunk_size])
+ reasoning_parts.append(reasoning)
+ content, deltas = tool_parser.parse_stream_chunk(content)
+ content_parts.append(content)
+ for delta in deltas:
+ if delta.name is not None:
+ assert delta.tool_index not in calls
+ calls[delta.tool_index] = {"name": delta.name, "arguments": ""}
+ calls[delta.tool_index]["arguments"] += delta.parameters
+
+ assert parser.flush() == ("", "")
+ assert "".join(reasoning_parts) == "Need a file."
+ assert "".join(content_parts) == ""
+ assert list(calls) == [0, 1]
+ assert all(call["name"] == "Write" for call in calls.values())
+ assert all(json.loads(call["arguments"]) == {"content": VALUE} for call in calls.values())
+
+
+@pytest.mark.parametrize("suffix", [""])
+def test_truncated_or_similar_markers_remain_reasoning(suffix):
+ parser = ReasoningParser("glm45", force_reasoning=True)
+ reasoning, content = parser.parse_stream_chunk("Keep thinking." + suffix)
+ tail_reasoning, tail_content = parser.flush()
+ assert reasoning + tail_reasoning == "Keep thinking." + suffix
+ assert content + tail_content == ""
+
+
+def test_non_thinking_glm_preserves_tool_payload():
+ parser = ReasoningParser("glm45", force_reasoning=False)
+ assert parser.parse_non_stream(TOOL_CALL) == ("", TOOL_CALL)
+ assert parser.parse_stream_chunk(TOOL_CALL) == ("", TOOL_CALL)
+
+
+def test_other_reasoning_parsers_keep_their_existing_boundary():
+ parser = ReasoningParser("qwen3", force_reasoning=True)
+ reasoning, content = parser.parse_non_stream("Still thinking." + TOOL_CALL.replace(VALUE, "literal"))
+ assert reasoning.startswith("Still thinking.")
+ assert content == ""
+
+
+def test_glm_tool_parsing_keeps_existing_reasoning_token_accounting():
+ parser = ReasoningParser("glm45", force_reasoning=True)
+ for token_id in [1, 2, 77, 3, 99, 4]:
+ parser.update_reasoning_token_count(token_id)
+ assert parser.reasoning_tokens == 4
+
+
+def test_glm_tools_inside_reasoning_do_not_enable_string_stops():
+ tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [98], "": [99]}[text])
+ state = ReasoningStateMachine("glm45", tokenizer, 1, [])
+ assert not any(state.handle(token_id) for token_id in [1, 77, 2, 3])
+ assert state.in_reasoning
+ assert state.handle(99) is False
+ assert not state.in_reasoning
+ assert state.handle(4) is True
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("separate_reasoning", [False, True])
+@pytest.mark.parametrize("reasoning_end", ["missing", "before_tool", "after_tool"])
+@pytest.mark.parametrize("chunk_size", [1, 4096])
+def test_chat_completion_exposes_glm_tool_calls(complete_chat, stream, separate_reasoning, reasoning_end, chunk_size):
+ segments = [("Need a file.", 1)]
+ if reasoning_end == "before_tool":
+ segments.append(("", 99))
+ segments.extend([("", 77), (TOOL_CALL[len("") :], 2)] * 2)
+ if reasoning_end == "after_tool":
+ segments.append(("", 99))
+ chunks = [
+ (text[start : start + chunk_size], token_id if start + chunk_size >= len(text) else 1)
+ for text, token_id in segments
+ for start in range(0, len(text), chunk_size)
+ ]
+ result = complete_chat(
+ chunks,
+ stream=stream,
+ separate_reasoning=separate_reasoning,
+ # GLM-5.3's template opens thinking even with this toggle disabled.
+ chat_template_kwargs={"enable_thinking": False},
+ )
+ assert result.reasoning == ("Need a file." if separate_reasoning else "")
+ assert result.content == ("" if separate_reasoning else "Need a file.")
+ assert [call["name"] for call in result.calls] == ["Write", "Write"]
+ assert len({call["id"] for call in result.calls}) == 2
+ assert [json.loads(call["arguments"]) for call in result.calls] == [{"content": VALUE}] * 2
+ assert result.finish_reason == "tool_calls"
+
+
+@pytest.mark.parametrize(
+ "prompt,state",
+ [
+ ("<|assistant|>", True),
+ ("<|assistant|>reasoning directive", True),
+ ("<|assistant|>", False),
+ ("<|assistant|>", False),
+ ("<|assistant|>firstsecond", True),
+ ("<|assistant|>old<|user|>question<|assistant|>", True),
+ ("<|assistant|>old<|user|>question<|assistant|>", None),
+ ("<|assistant|>old<|observation|>tool result", None),
+ ("<|user|>quoted", None),
+ ],
+)
+def test_glm_initial_reasoning_state_comes_from_current_prompt_turn(prompt, state):
+ assert ReasoningParser.DetectorMap["glm45"].initial_reasoning_state(prompt) is state
+
+
+def test_glm_prompt_override_does_not_change_other_models(monkeypatch):
+ monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="qwen3"))
+ monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True)
+ request = ChatCompletionRequest(model="qwen", messages=[], chat_template_kwargs={"enable_thinking": False})
+ default = api_openai._is_force_thinking_mode(request)
+ assert ReasoningParser.get_initial_reasoning_state("qwen3", "<|assistant|>", default) is False
+
+
+@pytest.mark.parametrize("default", [False, True])
+@pytest.mark.parametrize("prompt", ["custom prompt", "<|assistant|>", "<|assistant|>old<|user|>question"])
+def test_glm_unrecognized_prompt_preserves_request_mode(prompt, default):
+ assert ReasoningParser.get_initial_reasoning_state("glm45", prompt, default) is default
+
+
+@pytest.mark.parametrize("thinking", [False, True])
+def test_glm_grammar_activation_follows_rendered_prompt(monkeypatch, thinking):
+ monkeypatch.setattr(api_openai, "get_env_start_args", lambda: StartArgs(reasoning_parser="glm45"))
+ monkeypatch.setattr(build_prompt, "tokenizer_supports_force_thinking", lambda: True)
+ request = ChatCompletionRequest(model="glm", messages=[], chat_template_kwargs={"enable_thinking": not thinking})
+ prompt = "<|assistant|>" + ("" if thinking else "")
+ tokenizer = SimpleNamespace(encode=lambda text, **kwargs: {"": [99]}[text])
+ default = api_openai._is_force_thinking_mode(request)
+ initial_reasoning_state = ReasoningParser.get_initial_reasoning_state("glm45", prompt, default)
+ assert default is not thinking
+ assert initial_reasoning_state is thinking
+ assert api_openai._get_guided_reasoning_end(tokenizer, "glm45", initial_reasoning_state) == (
+ [99] if thinking else []
+ )
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("thinking", [False, True])
+def test_glm_prompt_phase_is_shared_by_output_stop_and_grammar(complete_chat, compiler, stream, thinking):
+ chunks = [("Need a file.", 1), ("", 99)] if thinking else []
+ chunks.append(('{"ok": true}', 2))
+ result = complete_chat(
+ chunks,
+ prompt="<|assistant|>" + ("" if thinking else ""),
+ compiler=compiler,
+ stream=stream,
+ separate_reasoning=True,
+ tool_choice="none",
+ stop=["STOP"],
+ response_format={"type": "json_object"},
+ chat_template_kwargs={"enable_thinking": not thinking},
+ )
+ assert result.reasoning == ("Need a file." if thinking else "")
+ assert result.content == '{"ok": true}'
+ assert result.sampling_params._initial_reasoning_state == int(thinking)
+ assert result.sampling_params.constraint_out_params.guided_reasoning_end.to_list() == ([99] if thinking else [])
+
+
+@pytest.mark.parametrize("parser_name", ["glm45", "qwen3"])
+@pytest.mark.parametrize("thinking", [False, True])
+@pytest.mark.parametrize("separate_reasoning", [False, True])
+def test_stream_flush_preserves_truncated_text_in_each_phase(complete_chat, parser_name, thinking, separate_reasoning):
+ result = complete_chat(
+ [("Body", 1), ("<", 2)],
+ finish_status=FinishStatus.FINISHED_LENGTH,
+ parser_name=parser_name,
+ prompt="<|assistant|>" + ("" if thinking else ""),
+ stream=True,
+ separate_reasoning=separate_reasoning,
+ tool_choice="none",
+ chat_template_kwargs={"enable_thinking": thinking},
+ )
+ assert result.reasoning == ("Body<" if thinking and separate_reasoning else "")
+ assert result.content == ("" if thinking and separate_reasoning else "Body<")
+ assert result.finish_reason == "length"
+
+
+@pytest.mark.parametrize("suffix", ["<", "Writecontentbody", 2), (suffix, 3)],
+ finish_status=FinishStatus.FINISHED_LENGTH,
+ stream=True,
+ separate_reasoning=separate_reasoning,
+ )
+ assert result.reasoning == ("Need a file." if separate_reasoning else "")
+ assert result.content == ("" if separate_reasoning else "Need a file.")
+ assert [call["name"] for call in result.calls] == ["Write"]
+ assert result.calls[0]["arguments"] == ""
+ assert result.finish_reason == "length"