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..5652279702 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 @@ -306,9 +306,10 @@ def _create_weight(self): self.per_expert_scale = None # Create e_score_correction_bias if self.e_score_correction_bias_name: + # Preserve FP32 correction scores at expert-selection boundaries. self.e_score_correction_bias = torch.empty( (self.n_routed_experts,), - dtype=self.data_type_, + dtype=torch.float32, device=f"cuda:{self.device_id_}", ) self.e_score_correction_bias.load_ok = False 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..3aa9b8b32f 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, @@ -130,7 +130,7 @@ def low_latency_dispatch( num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, num_experts=self.total_expert_num_contain_redundancy, use_fp8=use_fp8_w8a8, - round_scale=get_env_use_ue8m0_scales(), + round_scale=get_env_use_ue8m0_activation_scales(), async_finish=False, return_recv_hook=True, ) 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..d2a56c188c 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 @@ -391,7 +392,9 @@ def moe_align_fused(expert_to_token_index, expert_to_weight, expert_token_num, t else: # Expert rows may be unordered, but grouped matmul reuses this same # mapping for up/down projections and writes back to original topk slots. - _moe_align_fused_atomic_token(expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights) + _moe_align_fused_atomic_token( + expert_to_token_index, expert_to_weight, expert_token_num, topk_ids, topk_weights + ) return expert_to_token_index, expert_to_weight, expert_token_num @@ -878,8 +881,12 @@ def grouped_matmul( # input 使用 per group quant 量化 _m, _k = token_inputs.shape assert _k % block_size_k == 0 + # Triton consumes FP32 activation scales even when weight scales are packed UE8M0. token_inputs, token_input_scale = per_token_group_quant_fp8( - token_inputs, block_size_k, dtype=expert_weights.dtype + token_inputs, + block_size_k, + dtype=expert_weights.dtype, + use_ue8m0_scales=get_env_use_ue8m0_activation_scales(get_env_use_ue8m0_scales() or use_packed_ue8m0), ) if reused_mblock_infos is None: 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..da37626171 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(), ) @@ -288,7 +288,7 @@ def fused_experts_impl( combined_x = None if is_prefill: qinput_tensor, input_scale = per_token_group_quant_fp8( - hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_scales() + hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=get_env_use_ue8m0_activation_scales() ) allocate_on_comm_stream = previous_event is not None # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: @@ -383,7 +383,7 @@ def fused_experts_impl( num_max_dispatch_tokens_per_rank, num_experts, use_fp8=use_fp8_w8a8, - round_scale=get_env_use_ue8m0_scales(), + round_scale=get_env_use_ue8m0_activation_scales(), async_finish=False, return_recv_hook=False, ) @@ -608,7 +608,7 @@ def workspace_quant_alloc(shape, dtype, device): column_major_scales=True, scale_tma_aligned=True, alloc_func=workspace_quant_alloc, - use_ue8m0_scales=get_env_use_ue8m0_scales(), + use_ue8m0_scales=get_env_use_ue8m0_activation_scales(), ) workspace_manager.free(silu_out) del silu_out 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..9a631c8ea7 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 @@ -55,7 +55,9 @@ def _silu_and_mul_post_quant_kernel( ) for token_index in tl.range(token_id, token_num_cur_expert, block_num_per_expert, num_stages=NUM_STAGE): - gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to(tl.float32) + gate = tl.load(input_ptr_offs + token_index * stride_input_1, mask=offs_in_d < size_n, other=0.0).to( + tl.float32 + ) up = tl.load(input_ptr_offs + token_index * stride_input_1 + size_n, mask=offs_in_d < size_n, other=0.0) if USE_LIMIT_AND_ALPHA: gate = tl.minimum(gate, limit) @@ -161,7 +163,7 @@ def silu_and_mul_masked_post_quant_fwd( alpha=alpha, limit=limit, CLAMP_UP_ADD_ONE=clamp_up_add_one, - USE_UE8M0_SCALE=get_env_use_ue8m0_scales(), + USE_UE8M0_SCALE=get_env_use_ue8m0_activation_scales(), num_warps=num_warps, ) return 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/kv_cache_mem_manager/glm5_next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py index 45a3be667c..1779fabf0c 100644 --- a/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/glm5_next_mem_manager.py @@ -39,6 +39,10 @@ def get_att_input_params(self, layer_index): return self._layer_buffer(layer_index)[:, :, : self.mla_head_dim] def get_indexer_k_buffer(self, layer_index): + if self.linear_config.indexer_bf16: + return self._layer_buffer(layer_index)[ + :, :, self.mla_head_dim : self.mla_head_dim + self.linear_config.index_head_dim + ] return self._layer_buffer(layer_index).view(torch.uint8)[:, :, -132:] def write_to_shm(self, req_manager): diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index c953afd27a..b0221f3304 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, ) @@ -187,7 +190,9 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): else: num_groups, n, k = weight.shape packed_weight = torch.empty((num_groups, n, k // 2), device=weight.device, dtype=torch.int8) - weight_scale = torch.empty((num_groups, n, k // self.block_size), device=weight.device, dtype=torch.float32) + weight_scale = torch.empty( + (num_groups, n, k // self.block_size), device=weight.device, dtype=torch.float32 + ) for i in range(num_groups): packed_weight[i], weight_scale[i] = per_token_cast_to_fp4( weight[i], use_ue8m0=True, gran_k=self.block_size 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/common/state_cache_manager/glm5_next.py b/lightllm/common/state_cache_manager/glm5_next.py index 13195a74ac..4c833dd15d 100644 --- a/lightllm/common/state_cache_manager/glm5_next.py +++ b/lightllm/common/state_cache_manager/glm5_next.py @@ -1,6 +1,6 @@ import dataclasses -from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args +from lightllm.utils.envs_utils import get_added_mtp_kv_layer_num, get_env_start_args, get_env_glm5_indexer_bf16 from lightllm.utils.torch_dtype_utils import get_torch_dtype from .linear_att import LinearAttCacheConfig @@ -12,6 +12,7 @@ class Glm5NextCacheConfig(LinearAttCacheConfig): index_kpool: int = 4 index_head_dim: int = 128 + indexer_bf16: bool = False INDEX_PADDING_BYTES = 144 @@ -26,8 +27,12 @@ def from_model_config(cls, config, args): "deepseek_sparse_attention" if i % 4 == 3 else "linear_attention" for i in range(layers) ] dtype = get_torch_dtype(args.data_type) + indexer_bf16 = get_env_glm5_indexer_bf16() + if indexer_bf16: + assert args.data_type in ("bf16", "bfloat16"), "BF16 indexer requires --data_type bfloat16" # Raw index keys and compression scores only live in request tails. - packed_dim = config["kv_lora_rank"] + cls.INDEX_PADDING_BYTES // dtype.itemsize + index_bytes = config["index_head_dim"] * dtype.itemsize if indexer_bf16 else cls.INDEX_PADDING_BYTES + packed_dim = config["kv_lora_rank"] + index_bytes // dtype.itemsize return cls( tp_world_size=tp, full_att_all_num_kv_heads=1, @@ -49,6 +54,7 @@ def from_model_config(cls, config, args): draft_full_att_kv_layer_num=get_added_mtp_kv_layer_num() if args.mtp_mode is not None else 0, index_kpool=config["index_kpool"], index_head_dim=config["index_head_dim"], + indexer_bf16=indexer_bf16, ) def get_cpu_cache_full_att_bytes(self): diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index cf13ab83de..652e848c1b 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -2,9 +2,18 @@ import triton from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops +from lightllm.utils.envs_utils import get_env_glm5_indexer_bf16 +from .triton_kernel.index_bf16 import bf16_mqa_logits from .triton_kernel.index_quant import hadamard_transform_quant_fp8 -from .triton_kernel.kpool import compress_pools, gather_pools, gather_paged_pools, get_pool_ranges, expand_topk +from .triton_kernel.kpool import ( + compress_pools, + gather_pools, + gather_bf16_pools, + gather_paged_pools, + get_pool_ranges, + expand_topk, +) class Glm5NextNsaInfer: @@ -15,7 +24,9 @@ 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 + self.indexer_bf16 = get_env_glm5_indexer_bf16() def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" @@ -38,6 +49,7 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig raw = torch.cat((k, gate), -1) tail = infer_state.req_manager.get_indexer_tail_buffer(self.layer_idx) packed_buffer = infer_state.mem_manager.get_indexer_k_buffer(self.layer_idx) + assert (packed_buffer.dtype == torch.bfloat16) == self.indexer_bf16, "Indexer/cache precision mismatch" compress_pools( raw=raw, tail=tail, @@ -54,15 +66,24 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig ) if infer_state.max_kv_seq_len <= self.topk: - return expand_topk(None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True) + return expand_topk( + None, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk, dense=True + ) # The small indexer is replicated: no all-gather of query heads and # identical pool selection on every TP rank. q = layer_weight.wq_b_proj_.mm(q_lora).view(-1, self.heads, self.dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim ** -0.5) - weights = layer_weight.weights_proj_.mm(hidden_states.float()) - weights = weights * (self.heads ** -0.5 * self.dim ** -0.5) * q_scale.squeeze(-1) max_pools = triton.cdiv(infer_state.max_kv_seq_len, 4 * 128) * 128 + if self.indexer_bf16: + # Keep the existing projections and pooling arithmetic: only remove + # the rotated FP8 Q/K representation and FP8 scoring from this A/B. + weights = layer_weight.weights_proj_.mm(hidden_states.float()) * self.heads**-0.5 + groups = self._get_bf16_indices(q, weights, packed_buffer, infer_state, att_state, max_pools) + return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) + + q_fp8, q_scale = hadamard_transform_quant_fp8(q, scale=self.dim**-0.5) + weights = layer_weight.weights_proj_.mm(hidden_states.float()) + weights = weights * (self.heads**-0.5 * self.dim**-0.5) * q_scale.squeeze(-1) # Prefill usually has many queries: parallelize over Q and reuse K within each query tile. # Decode has few queries, so paged MQA also splits K across SMs for parallelism. if infer_state.is_prefill: @@ -71,6 +92,26 @@ def _get_indices(self, hidden_states, q_lora, infer_state, att_state, layer_weig groups = self._get_decode_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) return expand_topk(groups, att_state.lengths, att_state.ks, att_state.ragged_mem_index, self.topk) + def _get_bf16_indices(self, q, weights, packed_buffer, infer_state, att_state, max_pools): + # Decode uses one request entry per MTP position; its visible length is + # per row, so future verify pools and HOLD rows remain masked. + seq_lens = infer_state.b_seq_len if infer_state.is_prefill else att_state.lengths + keys = gather_bf16_pools( + packed_buffer, infer_state.req_manager.req_to_token_indexs, infer_state.b_req_idx, seq_lens, max_pools + ) + starts, ends, lengths = get_pool_ranges( + att_state.lengths, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools + ) + groups = torch.empty((q.shape[0], self.topk // 4), dtype=torch.int32, device=q.device) + chunk_size = max(1, min(q.shape[0], 16 * 1024 * 1024 // max_pools)) + for start in range(0, q.shape[0], chunk_size): + end = min(start + chunk_size, q.shape[0]) + logits = bf16_mqa_logits( + q[start:end], keys, weights[start:end], starts[start:end], ends[start:end], max_pools + ) + self.select_topk_indices(logits, lengths[start:end], groups[start:end]) + return groups + def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): keys = gather_pools( packed_buffer, 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..2032e50c96 100644 --- a/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/glm5_next/layer_infer/transformer_layer_infer.py @@ -428,6 +428,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 +439,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_bf16.py b/lightllm/models/glm5_next/triton_kernel/index_bf16.py new file mode 100644 index 0000000000..d17ff0e50c --- /dev/null +++ b/lightllm/models/glm5_next/triton_kernel/index_bf16.py @@ -0,0 +1,76 @@ +"""Unrotated BF16 Q/K with IEEE FP32 scoring for the indexer precision ablation.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _bf16_mqa_logits( + Q, + K, + Weights, + Starts, + Ends, + Out, + Q_STRIDE: tl.constexpr, + Q_HEAD_STRIDE: tl.constexpr, + W_STRIDE: tl.constexpr, + POOLS: tl.constexpr, + HEADS: tl.constexpr, + DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_P: tl.constexpr, +): + row = tl.program_id(1) + pools = tl.program_id(0) * BLOCK_P + tl.arange(0, BLOCK_P) + heads = tl.arange(0, BLOCK_H) + dims = tl.arange(0, 32) + start = tl.load(Starts + row).to(tl.int64) + length = tl.load(Ends + row).to(tl.int64) - start + valid = (pools < POOLS) & (pools < length) + # One query per program preserves request boundaries in ragged/MTP batches. + # IEEE FP32 avoids TF32 settings changing the precision experiment. + dots = tl.full((BLOCK_H, BLOCK_P), 0, tl.float32) + for block in range(tl.cdiv(DIM, 32)): + cols = block * 32 + dims + q = tl.load( + Q + row * Q_STRIDE + heads[:, None] * Q_HEAD_STRIDE + cols[None, :], + (heads[:, None] < HEADS) & (cols[None, :] < DIM), + 0, + ).to(tl.float32) + k = tl.load( + K + (start + pools[None, :]) * DIM + cols[:, None], + valid[None, :] & (cols[:, None] < DIM), + 0, + ).to(tl.float32) + dots = tl.dot(q, k, dots, input_precision="ieee") + weights = tl.load(Weights + row * W_STRIDE + heads, heads < HEADS, 0).to(tl.float32) + scores = tl.sum(tl.maximum(dots * (DIM**-0.5), 0) * weights[:, None], 0) + tl.store(Out + row * POOLS + pools, tl.where(valid, scores, -float("inf")), pools < POOLS) + + +def bf16_mqa_logits(q, keys, weights, starts, ends, max_pools): + """Return row-relative logits; weights already include heads**-0.5.""" + assert q.dtype == keys.dtype == torch.bfloat16 + assert keys.is_contiguous() and weights.stride(1) == 1 + logits = torch.empty((q.shape[0], max_pools), device=q.device, dtype=torch.float32) + _bf16_mqa_logits[(triton.cdiv(max_pools, 32), q.shape[0])]( + q, + keys, + weights, + starts, + ends, + logits, + q.stride(0), + q.stride(1), + weights.stride(0), + max_pools, + q.shape[1], + q.shape[2], + max(16, triton.next_power_of_2(q.shape[1])), + 32, + num_warps=4, + enable_fp_fusion=False, + ) + return logits diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index 0e126aea6a..0bf42af58b 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -45,6 +45,7 @@ def _compress_pools( SINGLE_QUERY: tl.constexpr, TAIL_SIZE: tl.constexpr, HAS_MTP_INDEX: tl.constexpr, + QUANTIZED: tl.constexpr, ): if SINGLE_QUERY: batch = tl.program_id(0) @@ -85,15 +86,18 @@ def _compress_pools( score = tl.exp(score - tl.max(score, 0)[None, :]) weights = score / tl.sum(score, 0)[None, :] key = tl.sum(raw * weights, 0).to(tl.bfloat16).to(tl.float32).reshape(1, 128) - 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))) - 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 - tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) - tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) + if QUANTIZED: + 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))) + key = tl.minimum(tl.maximum(key / scale[:, None], -448.0), 448.0).to(tl.float8e4nv) + tl.store(dest + cols, key.reshape(128).to(tl.uint8, bitcast=True)) + tl.store((dest + 128).to(tl.pointer_type(tl.float32)), tl.sum(scale, 0)) + else: + tl.store(dest + cols, key.reshape(128)) @triton.jit @@ -156,6 +160,7 @@ def compress_pools( single_query, tail.shape[1], mtp_index is not None, + packed_buffer.dtype == torch.uint8, num_warps=4, ) # Complete all boundary pools before replacing the previous chunk's tail. @@ -226,6 +231,7 @@ def _gather_pools( PACKED_STRIDE: tl.constexpr, REQ_STRIDE: tl.constexpr, POOLS: tl.constexpr, + QUANTIZED: tl.constexpr, ): pool, batch = tl.program_id(0), tl.program_id(1) length = tl.load(SeqLen + batch) @@ -234,10 +240,13 @@ def _gather_pools( loc = tl.load(ReqTable + req * REQ_STRIDE + pool * 4 + 3, valid, 0).to(tl.int64) cols = tl.arange(0, 128) packed = tl.load(Packed + loc * PACKED_STRIDE + cols, valid, 0) - scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) row = batch.to(tl.int64) * POOLS + pool - tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) - tl.store(Scale + row, scale) + if QUANTIZED: + scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) + tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) + tl.store(Scale + row, scale) + else: + tl.store(K + row * 128 + cols, packed) def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): @@ -254,11 +263,30 @@ def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): packed_buffer.stride(0), req_table.stride(0), max_pools, + True, num_warps=4, ) return keys, scales +def gather_bf16_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): + keys = torch.empty((req_idx.numel() * max_pools, 128), device=packed_buffer.device, dtype=torch.bfloat16) + _gather_pools[(max_pools, req_idx.numel())]( + packed_buffer, + req_table, + req_idx, + seq_len, + keys, + None, + packed_buffer.stride(0), + req_table.stride(0), + max_pools, + False, + num_warps=4, + ) + return keys + + @triton.jit def _gather_paged_pools( Packed, diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 748f1ad2f7..ba8f690b36 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -77,6 +77,30 @@ 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. + Indexer Q/K quantization has its own precision switch. + """ + 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 + + +def get_env_glm5_indexer_bf16() -> bool: + """Precision ablation; set identically on all GLM5 prefill/decode workers.""" + return os.getenv("LIGHTLLM_GLM5_INDEXER_BF16", "0").upper() in ["ON", "TRUE", "1"] + + @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..d126c65ca8 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..f8565027b5 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -61,14 +61,14 @@ def record_call(*args, **kwargs): assert scales.dtype == (torch.int32 if use_packed_ue8m0 else torch.float32) if use_ue8m0_scales: actual = _unpack_scales(scales, 4) if use_packed_ue8m0 else scales - torch.testing.assert_close(actual, torch.full_like(actual, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(actual, torch.full_like(actual, 2.0**-8), rtol=0, atol=0) @pytest.mark.parametrize("use_packed_ue8m0", [False, True]) def test_ue8m0_rounding_at_power_of_two_boundaries(monkeypatch, use_packed_ue8m0): monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr(activation, "HAS_SGL_KERNEL", False) - powers = torch.tensor([2.0 ** exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") + powers = torch.tensor([2.0**exponent for exponent in (-16, -8, 0, 8, 16, 118)], device="cuda") above = torch.nextafter(powers, torch.full_like(powers, float("inf"))) below = torch.nextafter(powers, torch.zeros_like(powers)) magnitudes = torch.cat((powers, above, below, powers.new_zeros(1))) * 448.0 @@ -134,8 +134,10 @@ def test_packed_ue8m0_scales(monkeypatch, rows, groups, group_size, use_sgl): x[0].zero_() if rows > 1: x[1].fill_(1e-12) - x[-1, -group_size:].fill_(2.0 ** 16) - q, packed_scales = activation.per_token_group_quant_fp8(x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True) + x[-1, -group_size:].fill_(2.0**16) + q, packed_scales = activation.per_token_group_quant_fp8( + x, group_size, use_ue8m0_scales=True, use_packed_ue8m0=True + ) assert packed_scales.dtype == torch.int32 assert packed_scales.shape == (rows, (groups + 3) // 4) @@ -194,9 +196,9 @@ def check_gemm_inputs(a, b, out): seen.append(a[1].dtype) assert b[1].dtype == (torch.int32 if sm100 else torch.float32) scales = _unpack_scales(a[1], 4) if sm100 else a[1] - torch.testing.assert_close(scales, torch.full_like(scales, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(scales, torch.full_like(scales, 2.0**-8), rtol=0, atol=0) weight_scales = _unpack_scales(b[1], 4) if sm100 else b[1] - torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0 ** -8), rtol=0, atol=0) + torch.testing.assert_close(weight_scales, torch.full_like(weight_scales, 2.0**-8), rtol=0, atol=0) out.zero_() # Exercise format selection and real quantization on Hopper; SM100 GEMM needs a separate machine. @@ -266,13 +268,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 +342,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 +498,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 +514,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 +525,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_cache.py b/unit_tests/models/glm5_next/test_cache.py index ae40702d75..d6646a5aea 100644 --- a/unit_tests/models/glm5_next/test_cache.py +++ b/unit_tests/models/glm5_next/test_cache.py @@ -14,7 +14,9 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_nope_cache_config_uses_native_mla_width(): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_nope_cache_config_uses_native_mla_width(monkeypatch, indexer_bf16): + monkeypatch.setenv("LIGHTLLM_GLM5_INDEXER_BF16", str(int(indexer_bf16))) config = Glm5NextCacheConfig.from_model_config( { "kv_lora_rank": 512, @@ -31,14 +33,16 @@ def test_nope_cache_config_uses_native_mla_width(): }, StartArgs(tp=4, data_type="bfloat16"), ) - assert config.full_att_head_dim == 584 - assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + 144 + assert config.indexer_bf16 == indexer_bf16 + assert config.full_att_head_dim == (640 if indexer_bf16 else 584) + assert config.full_att_head_dim * config.full_att_dtype.itemsize == 512 * 2 + (256 if indexer_bf16 else 144) @pytest.mark.parametrize("small_page", [False, True]) @pytest.mark.parametrize("tp_world_size", [1, 4]) @pytest.mark.parametrize("mtp_step", [0, 2]) -def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, tp_world_size, mtp_step, indexer_bf16): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) @@ -54,12 +58,14 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t set_unique_server_name(args) get_env_start_args.cache_clear() set_env_start_args(dataclasses.asdict(args)) + head_dim = 640 if indexer_bf16 else 584 config = Glm5NextCacheConfig( tp_world_size=tp_world_size, full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=head_dim, + indexer_bf16=indexer_bf16, global_linear_k_heads=2 * tp_world_size, global_linear_v_heads=2 * tp_world_size, num_linear_k_heads=2, @@ -75,12 +81,14 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t draft_full_att_kv_layer_num=int(mtp_step > 0), ) full_layers = config.get_full_att_kv_layer_num_with_draft_model() - mem = Glm5NextMemManager(16, torch.bfloat16, 1, 584, full_layers, config) + mem = Glm5NextMemManager(16, torch.bfloat16, 1, head_dim, full_layers, config) req = Glm5NextReqManager(3, 16, mem, config) att_kv = mem.get_att_input_params(3) assert att_kv.shape == (17, 1, 512) - assert att_kv.stride(0) == 584 + assert att_kv.stride(0) == head_dim index_bytes = mem.get_indexer_k_buffer(3) + assert index_bytes.dtype == (torch.bfloat16 if indexer_bf16 else torch.uint8) + assert index_bytes.shape[-1] == (128 if indexer_bf16 else 132) index_bytes.random_(0, 256) expected_index_bytes = index_bytes.clone() new_kv = torch.randn(2, 1, 512, dtype=torch.bfloat16, device="cuda") @@ -133,7 +141,7 @@ def test_hybrid_checkpoint_restore_and_packed_kv_copy(monkeypatch, small_page, t packed_bytes[:, 0].random_(0, 256) mem.operator.copy_mem_to_mem(torch.tensor([0]), torch.tensor([7])) assert torch.equal(packed_bytes[:, 0], packed_bytes[:, 7]) - assert mem.get_cell_size() == 584 * 2 * full_layers + assert mem.get_cell_size() == head_dim * 2 * full_layers assert config.get_cpu_cache_full_att_bytes() == mem.get_cell_size() * 8 from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( diff --git a/unit_tests/models/glm5_next/test_indexer_bf16.py b/unit_tests/models/glm5_next/test_indexer_bf16.py new file mode 100644 index 0000000000..048a77d647 --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_bf16.py @@ -0,0 +1,153 @@ +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer +from lightllm.models.glm5_next.triton_kernel.index_bf16 import bf16_mqa_logits +from lightllm.models.glm5_next.triton_kernel.kpool import compress_pools + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +@pytest.mark.parametrize("heads", [4, 32]) +@pytest.mark.parametrize("max_pools", [640, 8192, 262144]) +def test_bf16_scores_match_fp32_and_replay_changing_bounds(heads, max_pools): + torch.manual_seed(53) + q = torch.randn(6, heads, 128, device="cuda", dtype=torch.bfloat16) + keys = torch.randn(2 * max_pools, 128, device="cuda", dtype=torch.bfloat16) + weights = torch.randn(6, heads, device="cuda") * heads**-0.5 + starts = torch.tensor([0, max_pools, 0, max_pools, 0, max_pools], device="cuda", dtype=torch.int32) + lengths = torch.tensor([0, 1, 511, 512, 639, max_pools], device="cuda", dtype=torch.int32) + ends = starts + lengths + bf16_mqa_logits(q, keys, weights, starts, ends, max_pools) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = bf16_mqa_logits(q, keys, weights, starts, ends, max_pools) + # Explicit FP64 oracle avoids any global TF32 setting affecting the check. + for shift in range(2): + lengths.copy_(lengths.roll(shift)) + ends.copy_(starts + lengths) + graph.replay() + for row, (start, length) in enumerate(zip(starts.tolist(), lengths.tolist())): + reference = (q[row].double() @ keys[start : start + length].double().T * 128**-0.5).relu() + reference = (reference * weights[row, :, None].double()).sum(0).float() + torch.testing.assert_close(actual[row, :length], reference, atol=3e-6, rtol=2e-5) + assert torch.isneginf(actual[row, length:]).all() + + +@pytest.mark.parametrize("is_prefill", [True, False]) +@pytest.mark.parametrize("cuda_graph", [False, True]) +def test_bf16_indexer_forward_causal_pools_mtp_and_hold(monkeypatch, is_prefill, cuda_graph): + monkeypatch.setenv("LIGHTLLM_GLM5_INDEXER_BF16", "1") + monkeypatch.setattr("lightllm.models.glm5_next.indexer.HAS_VLLM", False) + + def unexpected_fp8(*args, **kwargs): + raise AssertionError("BF16 indexer must not quantize Q or call the FP8 scorer") + + monkeypatch.setattr("lightllm.models.glm5_next.indexer.hadamard_transform_quant_fp8", unexpected_fp8) + torch.manual_seed(57) + heads, dim = 4, 128 + prefixes, counts, requests = [2300, 2049, 7], [5, 3, 3], [2, 0, 1] + lengths = [prefix + count for prefix, count in zip(prefixes, counts)] + table = torch.randperm(4 * 2500, device="cuda", dtype=torch.int32).view(4, 2500) + storage = torch.zeros(4 * 2500, 1, 640, device="cuda", dtype=torch.bfloat16) + packed = storage[:, :, 512:] + tail = torch.zeros(4, 8, 256, device="cuda", dtype=torch.bfloat16) + ape = torch.randn(4, dim, device="cuda") + tensor = lambda values: torch.tensor(values, device="cuda", dtype=torch.int32) + sources = [torch.randn(length, 256, device="cuda", dtype=torch.bfloat16) for length in lengths] + for req, prefix, source in zip(requests, prefixes, sources): + compress_pools( + source[:prefix], + tail, + packed, + ape, + tensor(range(1, prefix + 1)), + tensor([0] * prefix), + table[req], + tensor([req]), + tensor([0, prefix]), + tensor([prefix]), + prefix, + ) + hidden = torch.cat([source[prefix:] for source, prefix in zip(sources, prefixes)]) + q = torch.randn(hidden.shape[0], heads, dim, device="cuda", dtype=torch.bfloat16) + visible, req_rows, mtp_rows, ragged, starts, cu = [], [], [], [], [], [0] + for req, prefix, count, length in zip(requests, prefixes, counts, lengths): + visible.extend(range(prefix + 1, length + 1)) + req_rows.extend([req] * count) + mtp_rows.extend(range(count)) + starts.extend([len(ragged)] * count) + ragged.extend(table[req, :length].tolist()) + cu.append(cu[-1] + count) + if not is_prefill: + # Each verify position is an independent row with its own visible pool + # length and ragged start; graph padding uses the HOLD request slot. + hidden = torch.cat((hidden, torch.zeros(1, 256, device="cuda", dtype=torch.bfloat16))) + q = torch.cat((q, torch.zeros(1, heads, dim, device="cuda", dtype=torch.bfloat16))) + req_rows.append(3) + mtp_rows.append(0) + visible.append(0) + starts.append(len(ragged)) + infer = SimpleNamespace( + is_prefill=is_prefill, + max_kv_seq_len=max(lengths), + max_q_seq_len=max(counts) if is_prefill else 1, + b_req_idx=tensor(requests if is_prefill else req_rows), + b_seq_len=tensor(lengths if is_prefill else visible), + b1_cu_q_seq_len=tensor(cu if is_prefill else range(hidden.shape[0] + 1)), + b_mtp_index=tensor(mtp_rows), + mem_manager=SimpleNamespace(get_indexer_k_buffer=lambda _: packed), + req_manager=SimpleNamespace(req_to_token_indexs=table, get_indexer_tail_buffer=lambda _: tail), + ) + att = SimpleNamespace(lengths=tensor(visible), ks=tensor(starts), ragged_mem_index=tensor(ragged)) + weight = SimpleNamespace( + wk_proj_=SimpleNamespace(mm=lambda x: x[:, :dim].contiguous()), + k_norm_=lambda x, **kwargs: x, + index_kpool_compress_gate=SimpleNamespace(mm=lambda x: x[:, dim:].contiguous()), + index_kpool_compress_ape=SimpleNamespace(weight=ape), + wq_b_proj_=SimpleNamespace(mm=lambda x: x), + weights_proj_=SimpleNamespace(mm=lambda x: x[:, :heads].contiguous() * 0.1), + ) + indexer = Glm5NextNsaInfer(3, {"index_topk": 2048, "index_n_heads": heads, "index_head_dim": dim}, 8) + run = lambda: indexer._get_indices(hidden, q.flatten(1), infer, att, weight) + before = tail.clone() + run() + tail.copy_(before) + if cuda_graph: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual, relative = run() + tail.copy_(before) + graph.replay() + else: + actual, relative = run() + row = 0 + for req, prefix, count, source in zip(requests, prefixes, counts, sources): + for step in range(count): + length = prefix + step + 1 + pools = length // 4 + values = source[: pools * 4].view(pools, 4, 256) + pooled = (values[:, :, :dim].float() * (values[:, :, dim:].float() + ape).softmax(1)).sum(1).bfloat16() + scores = (q[row].double() @ pooled.double().T * dim**-0.5).relu() + weights = hidden[row, :heads].float() * 0.1 * heads**-0.5 + scores = (scores * weights[:, None].double()).sum(0) + selected = relative[row][relative[row] >= 0] + # Dense requests include all tokens; sparse requests include full + # selected pools followed by the uncompressed partial tail. + if length <= 2048: + assert sorted(selected.tolist()) == list(range(length)) + else: + groups = selected[:2048:4] // 4 + assert groups.unique().numel() == 512 + expected = scores.topk(512).values.sort().values + got = scores[groups.long()].sort().values + torch.testing.assert_close(got, expected, atol=2e-5, rtol=2e-5) + assert selected[2048:].tolist() == list(range(pools * 4, length)) + assert torch.equal(actual[row, : selected.numel()], table[req, selected.long()]) + row += 1 + if not is_prefill: + assert (actual[-1] == -1).all() + assert not tail[3].any() diff --git a/unit_tests/models/glm5_next/test_kernels.py b/unit_tests/models/glm5_next/test_kernels.py index 1a5e32567b..8a8ec45afd 100644 --- a/unit_tests/models/glm5_next/test_kernels.py +++ b/unit_tests/models/glm5_next/test_kernels.py @@ -182,7 +182,8 @@ 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("indexer_bf16", [False, True]) +def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph, indexer_bf16): # Unaligned chunks, reordered requests, and pools spanning old/new tails. seqs = [9, 7] ragged = torch.randperm(40, device="cuda", dtype=torch.int32)[: sum(seqs)] @@ -191,8 +192,10 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): table[1, :9], table[3, :7] = ragged[:9], ragged[9:] source = torch.randn(sum(seqs), 256, device="cuda", dtype=torch.bfloat16) ape = torch.randn(4, 128, device="cuda") - packed_storage = torch.zeros(40, 1, 584, device="cuda", dtype=torch.bfloat16) - packed = packed_storage.view(torch.uint8)[:, :, -132:] + head_dim = 640 if indexer_bf16 else 584 + packed_storage = torch.zeros(40, 1, head_dim, device="cuda", dtype=torch.bfloat16) + get_packed = lambda storage: storage[:, :, 512:] if indexer_bf16 else storage.view(torch.uint8)[:, :, -132:] + packed = get_packed(packed_storage) tail = torch.randn(5, 4, 256, device="cuda", dtype=torch.bfloat16) unchanged = tail[[0, 2, 4]].clone() tensor = lambda x: torch.tensor(x, device="cuda", dtype=torch.int32) @@ -237,13 +240,22 @@ def test_kpool_chunk_boundaries_and_fragmented_token_kv(cuda_graph): assert torch.equal(tail[[0, 2, 4]], unchanged) # Model the full KV copy performed by cache offload/load or request move. packed_storage = packed_storage.clone() - packed = packed_storage.view(torch.uint8)[:, :, -132:] - keys, scales = gather_pools(packed, table, req_idx, torch.tensor(seqs, device="cuda", dtype=torch.int32), 3) + packed = get_packed(packed_storage) + seq_lens = torch.tensor(seqs, device="cuda", dtype=torch.int32) + if indexer_bf16: + from lightllm.models.glm5_next.triton_kernel.kpool import gather_bf16_pools + + keys = gather_bf16_pools(packed, table, req_idx, seq_lens, 3) + else: + keys, scales = gather_pools(packed, table, req_idx, seq_lens, 3) for batch, start in enumerate([0, 9]): for group in range(seqs[batch] // 4): values = source[start + group * 4 : start + group * 4 + 4] expected = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128 ** -0.5) + if indexer_bf16: + torch.testing.assert_close(keys[batch * 3 + group], expected, atol=0, rtol=0) + continue + expected_key, expected_scale = hadamard_transform_quant_fp8(expected[None], scale=128**-0.5) torch.testing.assert_close(keys[batch * 3 + group].float(), expected_key[0].float(), atol=0, rtol=0) torch.testing.assert_close(scales[batch * 3 + group], expected_scale[0, 0], atol=0, rtol=0) lengths = torch.tensor(seqs, device="cuda", dtype=torch.int32) @@ -300,7 +312,7 @@ def run(): if pending == 0: values = source[batch, pos - 3 : pos + 1] pooled = (values[:, :128].float() * (values[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - key, scale = hadamard_transform_quant_fp8(pooled[None], 128 ** -0.5) + key, scale = hadamard_transform_quant_fp8(pooled[None], 128**-0.5) actual = packed[table[batch, pos].long(), 0] assert torch.equal(actual[:128], key.view(torch.uint8)[0]) assert torch.equal(actual[128:].view(torch.float32), scale.flatten()) @@ -392,7 +404,7 @@ def check(keys, scales): def test_kpool_packed_kv_above_two_gib(): # Automatic KV sizing on H200 can put pool keys beyond a 32-bit byte offset. - loc = 2 ** 31 // (584 * 2) + 1 + loc = 2**31 // (584 * 2) + 1 storage = torch.empty(loc + 1, 1, 584, device="cuda", dtype=torch.bfloat16) packed = storage.view(torch.uint8)[:, :, -132:] raw = torch.randn(4, 256, device="cuda", dtype=torch.bfloat16) @@ -405,7 +417,7 @@ def test_kpool_packed_kv_above_two_gib(): compress_pools(raw, tail, packed, ape, lengths, zero, ragged, zero[:1], cu_q_lens, lengths[-1:], max_q_len=4) keys, scales = gather_pools(packed, ragged[None], zero[:1], lengths[-1:], 1) pooled = (raw[:, :128].float() * (raw[:, 128:].float() + ape).softmax(0)).sum(0).bfloat16() - expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128 ** -0.5) + expected_key, expected_scale = hadamard_transform_quant_fp8(pooled[None], scale=128**-0.5) assert torch.equal(keys.view(torch.uint8), expected_key.view(torch.uint8)) assert torch.equal(scales, expected_scale.flatten()) @@ -482,7 +494,7 @@ def test_shared_chunk_kernel_preserves_gdn_natural_log_decay(): state *= gate[:, i, :, None, None].exp() delta = (v[:, i].float() - torch.einsum("bhkv,bhk->bhv", state, k[:, i].float())) * beta[:, i, :, None] state += k[:, i, :, :, None].float() * delta[:, :, None, :] - expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim ** 0.5) + expected.append(torch.einsum("bhkv,bhk->bhv", state, q[:, i].float()) / dim**0.5) out, final = chunk_gated_delta_rule( q, k, @@ -637,10 +649,10 @@ def test_kpool_indexer_long_prefill_and_cached_decode(max_kv_seq_len): pool_count = tokens // 4 pool_values = hidden[: pool_count * 4].view(pool_count, 4, dim) pooled = (pool_values.float() * (pool_values.float() * 0.25 + ape).softmax(1)).sum(1).bfloat16() - k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim ** -0.5) + k_fp8, k_scale = hadamard_transform_quant_fp8(pooled, dim**-0.5) query = (hidden[-1:] @ q_weight).view(heads, dim) - q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim ** -0.5) + q_fp8, q_scale = hadamard_transform_quant_fp8(query, dim**-0.5) logits = (q_fp8.float() @ k_fp8.float().T * k_scale.flatten()).clamp_min(0) - scores = (logits * q_scale * (heads ** -0.5 * dim ** -0.5)).sum(0) + scores = (logits * q_scale * (heads**-0.5 * dim**-0.5)).sum(0) expected_groups = set(scores.topk(512).indices.tolist()) assert set((decoded[0, :2048:4] // 4).tolist()) == expected_groups diff --git a/unit_tests/models/glm5_next/test_mtp.py b/unit_tests/models/glm5_next/test_mtp.py index 00d4e326f0..e863ab91a7 100644 --- a/unit_tests/models/glm5_next/test_mtp.py +++ b/unit_tests/models/glm5_next/test_mtp.py @@ -72,13 +72,16 @@ def test_kda_mtp_matches_single_token_decode_and_restores_accepted_state(mtp_ste @pytest.mark.parametrize("mtp_step", [1, 2, 3, 5]) @pytest.mark.parametrize("prefix", [1, 2, 3, 4, 7, 15]) -def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix): +@pytest.mark.parametrize("indexer_bf16", [False, True]) +def test_kpool_verify_and_draft_rewind_match_full_prefill(mtp_step, prefix, indexer_bf16): torch.manual_seed(53) width, capacity = mtp_step + 1, 128 raw = torch.randn(capacity, 256, device="cuda", dtype=torch.bfloat16) ape = torch.randn(4, 128, device="cuda", dtype=torch.bfloat16) ring = torch.zeros(2, 4 + mtp_step, 256, device="cuda", dtype=torch.bfloat16) - packed = torch.zeros(capacity, 132, device="cuda", dtype=torch.uint8) + packed = torch.zeros( + capacity, 128 if indexer_bf16 else 132, device="cuda", dtype=torch.bfloat16 if indexer_bf16 else torch.uint8 + ) ragged = torch.arange(capacity, device="cuda", dtype=torch.int32) def prefill(values, destination, tail): diff --git a/unit_tests/models/glm5_next/test_pd_cache.py b/unit_tests/models/glm5_next/test_pd_cache.py index 57665cbbe3..c561ea5da9 100644 --- a/unit_tests/models/glm5_next/test_pd_cache.py +++ b/unit_tests/models/glm5_next/test_pd_cache.py @@ -14,12 +14,13 @@ pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.fixture(params=[0, 2]) +@pytest.fixture(params=[(mtp, bf16) for mtp in (0, 2) for bf16 in (False, True)]) def make_mems(monkeypatch, request): monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_NODE", "0") monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setattr("lightllm.common.req_manager.req_sampling_params.get_vocab_size", lambda _: 128) - mtp_step = request.param + mtp_step, indexer_bf16 = request.param + head_dim = 640 if indexer_bf16 else 584 args = StartArgs( data_type="bfloat16", linear_att_hash_page_size=4, @@ -45,7 +46,8 @@ def create(tp): full_att_all_num_kv_heads=1, full_att_dtype=torch.bfloat16, full_att_num_kv_heads=1, - full_att_head_dim=584, + full_att_head_dim=head_dim, + indexer_bf16=indexer_bf16, global_linear_k_heads=8, global_linear_v_heads=8, num_linear_k_heads=8 // tp, @@ -62,7 +64,7 @@ def create(tp): ) mems = [] for _ in range(tp): - mem = Glm5NextMemManager(32, torch.bfloat16, 1, 584, 2 + int(mtp_step > 0), config) + mem = Glm5NextMemManager(32, torch.bfloat16, 1, head_dim, 2 + int(mtp_step > 0), config) req = Glm5NextReqManager(3, 32, mem, config) mem.write_to_shm(req) assert mem.big_page_buffers is not None @@ -73,7 +75,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) @@ -83,7 +85,7 @@ def test_pd_kv_and_runtime_state_roundtrip(make_mems, prefill_tp, decode_tp, rem length = 8 + remainder src_indexes = torch.randperm(32)[:length].tolist() dst_indexes = torch.randperm(32)[:length].tolist() - packed = torch.randint(0, 256, (layers, length, 1, 1168), dtype=torch.uint8, device="cuda") + packed = torch.randint(0, 256, (layers, length, 1, source[0].head_dim * 2), dtype=torch.uint8, device="cuda") conv = torch.randn(6, 3, 8, 128, 3, dtype=torch.bfloat16, device="cuda") ssm = torch.randn(6, 8, 128, 128, dtype=torch.float32, device="cuda") # Non-live slots deliberately contain stale data. The sequence length, not @@ -170,4 +172,4 @@ def test_pd_page_capacity_includes_tail_without_resizing(make_mems): with pytest.raises(AssertionError, match="smaller than global linear att state"): helper.assert_page_size() shape = mem.get_paged_kv_move_buffer_shape(2, 2048) - assert shape == (2, 2048, mem.layer_num, 1, 584) + assert shape == (2, 2048, mem.layer_num, 1, mem.head_dim) 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..746397c385 --- /dev/null +++ b/unit_tests/models/glm5_next/test_precision_alignment.py @@ -0,0 +1,200 @@ +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)