From dbe23136fd0f679094c218bcf725357d39f6a58d Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Sat, 4 Jul 2026 08:06:02 +0000 Subject: [PATCH 1/5] feat: opt mega moe perf --- .../fused_moe/grouped_fused_moe_ep.py | 170 +++++++++++++++--- lightllm/common/quantization/deepgemm.py | 23 ++- 2 files changed, 170 insertions(+), 23 deletions(-) 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 58d4d4551..2c660c723 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 @@ -51,6 +51,154 @@ def use_sm100_mega_moe(quant_method: Any) -> bool: return is_sm100_gpu() and quant_method.method_name == "fp4fp8-b32-deepgemm" +def _per_token_cast_to_fp8_packed_ue8m0(hidden_states: torch.Tensor, gran_k: int): + from deep_gemm.utils import per_token_cast_to_fp8 + + hidden_states, scale = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=gran_k, + use_packed_ue8m0=False, + ) + assert scale.size(-1) % 4 == 0, "packed UE8M0 scale requires scale groups divisible by 4" + scale = (scale.view(torch.int32) >> 23).to(torch.uint8).view(torch.int32) + return hidden_states, scale + + +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True), exp + + +@triton.jit +def _mega_moe_quant_topk_to_buffer_kernel( + x_ptr, + x_out_ptr, + x_sf_out_ptr, + topk_idx_ptr, + topk_idx_out_ptr, + topk_weights_ptr, + topk_weights_out_ptr, + stride_x_m: tl.constexpr, + stride_x_k: tl.constexpr, + stride_x_out_m: tl.constexpr, + stride_x_out_k: tl.constexpr, + stride_x_sf_out_m: tl.constexpr, + stride_x_sf_out_k: tl.constexpr, + stride_topk_idx_m: tl.constexpr, + stride_topk_idx_k: tl.constexpr, + stride_topk_idx_out_m: tl.constexpr, + stride_topk_idx_out_k: tl.constexpr, + stride_topk_weights_m: tl.constexpr, + stride_topk_weights_k: tl.constexpr, + stride_topk_weights_out_m: tl.constexpr, + stride_topk_weights_out_k: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + TOPK: tl.constexpr, + GROUP_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + token_id = tl.program_id(0) + pack_id = tl.program_id(1) + offsets = tl.arange(0, BLOCK) + cols = pack_id * BLOCK + offsets + + x = tl.load(x_ptr + token_id * stride_x_m + cols * stride_x_k).to(tl.float32) + abs_x = tl.abs(x) + group_id = offsets // GROUP_SIZE + + amax0 = tl.max(tl.where(group_id == 0, abs_x, 0.0)) + amax1 = tl.max(tl.where(group_id == 1, abs_x, 0.0)) + amax2 = tl.max(tl.where(group_id == 2, abs_x, 0.0)) + amax3 = tl.max(tl.where(group_id == 3, abs_x, 0.0)) + + scale0, exp0 = _ceil_to_ue8m0(tl.maximum(amax0, 1.0e-4) / FP8_MAX) + scale1, exp1 = _ceil_to_ue8m0(tl.maximum(amax1, 1.0e-4) / FP8_MAX) + scale2, exp2 = _ceil_to_ue8m0(tl.maximum(amax2, 1.0e-4) / FP8_MAX) + scale3, exp3 = _ceil_to_ue8m0(tl.maximum(amax3, 1.0e-4) / FP8_MAX) + + scale = tl.where( + group_id == 0, + scale0, + tl.where(group_id == 1, scale1, tl.where(group_id == 2, scale2, scale3)), + ) + x_q = tl.clamp(x / scale, FP8_MIN, FP8_MAX).to(x_out_ptr.dtype.element_ty) + tl.store(x_out_ptr + token_id * stride_x_out_m + cols * stride_x_out_k, x_q) + + packed_scale = exp0 | (exp1 << 8) | (exp2 << 16) | (exp3 << 24) + tl.store(x_sf_out_ptr + token_id * stride_x_sf_out_m + pack_id * stride_x_sf_out_k, packed_scale) + + if pack_id == 0: + topk_offsets = tl.arange(0, TOPK) + topk_idx = tl.load(topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k) + topk_weights = tl.load( + topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k + ) + tl.store( + topk_idx_out_ptr + token_id * stride_topk_idx_out_m + topk_offsets * stride_topk_idx_out_k, + topk_idx.to(topk_idx_out_ptr.dtype.element_ty), + ) + tl.store( + topk_weights_out_ptr + + token_id * stride_topk_weights_out_m + + topk_offsets * stride_topk_weights_out_k, + topk_weights.to(topk_weights_out_ptr.dtype.element_ty), + ) + + +def _prepare_mega_moe_buffer( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + buffer: Any, + group_size: int, +): + num_tokens, hidden_size = hidden_states.shape + if num_tokens == 0: + return + assert hidden_size % (group_size * 4) == 0, "packed UE8M0 scale requires four FP8 groups per int32" + assert hidden_states.is_contiguous(), "hidden_states must be contiguous" + assert buffer.x.shape[0] >= num_tokens and buffer.x.shape[1] == hidden_size + assert buffer.x_sf.shape[0] >= num_tokens and buffer.x_sf.shape[1] == hidden_size // group_size // 4 + + block = group_size * 4 + finfo = torch.finfo(buffer.x.dtype) + _mega_moe_quant_topk_to_buffer_kernel[(num_tokens, hidden_size // block)]( + hidden_states, + buffer.x, + buffer.x_sf, + topk_ids, + buffer.topk_idx, + topk_weights, + buffer.topk_weights, + hidden_states.stride(0), + hidden_states.stride(1), + buffer.x.stride(0), + buffer.x.stride(1), + buffer.x_sf.stride(0), + buffer.x_sf.stride(1), + topk_ids.stride(0), + topk_ids.stride(1), + buffer.topk_idx.stride(0), + buffer.topk_idx.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + buffer.topk_weights.stride(0), + buffer.topk_weights.stride(1), + FP8_MIN=finfo.min, + FP8_MAX=finfo.max, + TOPK=topk_ids.shape[1], + GROUP_SIZE=group_size, + BLOCK=block, + num_warps=4, + num_stages=4, + ) + + def check_ep_expert_dtype(quant_method: Any): expert_dtype = getattr(quant_method, "method_name", None) if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES: @@ -131,8 +279,6 @@ def mega_moe_impl( if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") - from deep_gemm.utils import per_token_cast_to_fp8 - buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) if buffer is None: raise RuntimeError("SM100 Mega MoE requires dist_group_manager.ep_mega_moe_buffer to be initialized") @@ -143,19 +289,10 @@ def mega_moe_impl( f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" ) - qinput_tensor = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) state = _get_mega_moe_cache_state(w13, w2) l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) - buffer.x[:num_tokens].copy_(qinput_tensor[0]) - buffer.x_sf[:num_tokens].copy_(qinput_tensor[1]) - buffer.topk_idx[:num_tokens].copy_(topk_ids) - buffer.topk_weights[:num_tokens].copy_(topk_weights) + _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) deep_gemm.fp8_fp4_mega_moe( @@ -175,14 +312,7 @@ def quantize_fused_experts_input( ): check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): - from deep_gemm.utils import per_token_cast_to_fp8 - - return per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) + return _per_token_cast_to_fp8_packed_ue8m0(hidden_states, quant_method.block_size) block_size_k = 0 if w13.weight.ndim == 3: diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index 3c3ee30bb..3645f6ef6 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -182,11 +182,28 @@ def _create_weight( out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims assert in_dim % 2 == 0, "FP4 packed weight requires even input dimension" assert in_dim % self.block_size == 0, "FP4 scale dimension must be divisible by block_size" + scales_per_int32 = 4 + scale_layout_k = self.block_size * scales_per_int32 + assert in_dim % scale_layout_k == 0, ( + f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" + ) expert_prefix = (num_experts,) if num_experts > 1 else () weight = torch.empty(expert_prefix + (out_dim, in_dim // 2), dtype=torch.int8).cuda(device_id) - weight_scale = torch.empty(expert_prefix + (out_dim, in_dim // self.block_size), dtype=torch.int32).cuda( - device_id - ) + scale_dim = in_dim // scale_layout_k + if num_experts > 1: + weight_scale = torch.empty_strided( + (num_experts, out_dim, scale_dim), + (out_dim * scale_dim, 1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) + else: + weight_scale = torch.empty_strided( + (out_dim, scale_dim), + (1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) mm_param_list = self._split_weight_pack( mm_param, From c896d3cf53061abad4f38f16084cef29ab745252 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 6 Jul 2026 12:19:58 +0000 Subject: [PATCH 2/5] feat: support UE8M0 --- .../quantization/fp8act_quant_kernel.py | 77 +++++++++++++------ .../fp8w8a8_block_quant_kernel.py | 35 +++++++-- lightllm/common/quantization/deepgemm.py | 3 +- 3 files changed, 84 insertions(+), 31 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py index 0a6837288..af4aff5c8 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py @@ -14,6 +14,14 @@ pass +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py @triton.jit def _per_token_group_quant_fp8( @@ -25,21 +33,19 @@ def _per_token_group_quant_fp8( eps, fp8_min, fp8_max, - xs_m, xs_n, - xs_row_major: tl.constexpr, + xs_stride_m, + xs_stride_n, BLOCK: tl.constexpr, NEED_MASK: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, ): g_id = tl.program_id(0) y_ptr += g_id * y_stride y_q_ptr += g_id * y_stride - if xs_row_major: - y_s_ptr += g_id - else: - row_id = g_id // xs_n - col_id = g_id % xs_n - y_s_ptr += col_id * xs_m + row_id # col major + row_id = g_id // xs_n + col_id = g_id % xs_n + y_s_ptr += row_id * xs_stride_m + col_id * xs_stride_n cols = tl.arange(0, BLOCK) # N <= BLOCK @@ -52,8 +58,11 @@ def _per_token_group_quant_fp8( y = tl.load(y_ptr + cols, mask=mask, other=other).to(tl.float32) # Quant - _absmax = tl.maximum(tl.max(tl.abs(y)), eps) - y_s = _absmax / fp8_max + _absmax = tl.max(tl.abs(y)) + if USE_UE8M0_SCALE: + y_s = _ceil_to_ue8m0(tl.maximum(_absmax, 1.0e-4) / fp8_max) + else: + y_s = tl.maximum(_absmax, eps) / fp8_max y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) tl.store(y_q_ptr + cols, y_q, mask=mask) @@ -67,6 +76,7 @@ def lightllm_per_token_group_quant_fp8( x_s: torch.Tensor, eps: float = 1e-10, dtype: torch.dtype = torch.float8_e4m3fn, + use_ue8m0_scales: bool = False, ): """group-wise, per-token quantization on input tensor `x`. Args: @@ -80,8 +90,8 @@ def lightllm_per_token_group_quant_fp8( assert x.shape[-1] % group_size == 0, "the last dimension of `x` cannot be divisible by `group_size`" assert x.is_contiguous(), "`x` is not contiguous" - xs_row_major = x_s.is_contiguous() - xs_m, xs_n = x_s.shape + xs_n = x_s.shape[-1] + xs_stride_m, xs_stride_n = x_s.stride() finfo = torch.finfo(dtype) fp8_max = finfo.max @@ -102,11 +112,12 @@ def lightllm_per_token_group_quant_fp8( eps, fp8_min=fp8_min, fp8_max=fp8_max, - xs_m=xs_m, xs_n=xs_n, - xs_row_major=xs_row_major, + xs_stride_m=xs_stride_m, + xs_stride_n=xs_stride_n, BLOCK=BLOCK, NEED_MASK=BLOCK != group_size, + USE_UE8M0_SCALE=use_ue8m0_scales, num_warps=num_warps, num_stages=num_stages, ) @@ -121,12 +132,13 @@ def per_token_group_quant_fp8( column_major_scales: bool = False, scale_tma_aligned: bool = False, alloc_func: Callable = torch.empty, + use_ue8m0_scales: bool = False, ): x_q = alloc_func(x.shape, dtype=dtype, device=x.device) x_s = None # Adapted from # https://github.com/sgl-project/sglang/blob/7e257cd666c0d639626487987ea8e590da1e9395/python/sglang/srt/layers/quantization/fp8_kernel.py#L290 - if HAS_SGL_KERNEL: + if HAS_SGL_KERNEL and not use_ue8m0_scales: finfo = torch.finfo(dtype) fp8_max, fp8_min = finfo.max, finfo.min @@ -157,14 +169,35 @@ def per_token_group_quant_fp8( sgl_ops.sgl_per_token_group_quant_fp8(x, x_q, x_s, group_size, 1e-10, fp8_min, fp8_max, False, enable_v2=True) else: # 使用LightLLM kernel进行量化 - x_s = alloc_func( - x.shape[:-1] + (x.shape[-1] // group_size,), - device=x.device, - dtype=torch.float32, + if column_major_scales: + if scale_tma_aligned: + aligned_size = (x.shape[-2] + 3) // 4 * 4 + x_s = alloc_func( + x.shape[:-2] + (x.shape[-1] // group_size, aligned_size), + device=x.device, + dtype=torch.float32, + ).permute(-1, -2)[: x.shape[-2], :] + else: + x_s = alloc_func( + (x.shape[-1] // group_size,) + x.shape[:-1], + device=x.device, + dtype=torch.float32, + ).permute(-1, -2) + else: + x_s = alloc_func( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + lightllm_per_token_group_quant_fp8( + x, + group_size, + x_q, + x_s, + eps=eps, + dtype=dtype, + use_ue8m0_scales=use_ue8m0_scales, ) - lightllm_per_token_group_quant_fp8(x, group_size, x_q, x_s, eps=1e-10, dtype=torch.float8_e4m3fn) - if column_major_scales and scale_tma_aligned: - x_s = tma_align_input_scale(x_s) return x_q, x_s diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py index 3881cfe4b..1c2caa967 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py @@ -5,7 +5,15 @@ @triton.jit -def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + +@triton.jit +def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr, USE_UE8M0_SCALE: tl.constexpr): pid_m = tl.program_id(axis=0) pid_n = tl.program_id(axis=1) n_blocks = tl.cdiv(N, BLOCK_SIZE) @@ -20,14 +28,21 @@ def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): amax = tl.max(tl.abs(x)) max_fp8e4m3_val = 448.0 - scale = amax / max_fp8e4m3_val - y = (x / (scale + 1e-6)).to(y_ptr.dtype.element_ty) + if USE_UE8M0_SCALE: + scale = _ceil_to_ue8m0(tl.maximum(amax, 1.0e-4) / max_fp8e4m3_val) + denom = scale + else: + scale = amax / max_fp8e4m3_val + denom = scale + 1e-6 + y = (x / denom).to(y_ptr.dtype.element_ty) tl.store(y_ptr + offs, y, mask=mask) tl.store(s_ptr + pid_m * n_blocks + pid_n, scale) -def mm_weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: +def mm_weight_quant( + x: torch.Tensor, block_size: int = 128, use_ue8m0_scales: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous(), "Input tensor must be contiguous" M, N = x.size() @@ -38,11 +53,15 @@ def mm_weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tenso s_scales = torch.empty((num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) - weight_quant_kernel[grid](x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size) + weight_quant_kernel[grid]( + x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales + ) return y_quant, s_scales -def weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: +def weight_quant( + x: torch.Tensor, block_size: int = 128, use_ue8m0_scales: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous(), "Input tensor must be contiguous" x = x.cuda(get_current_device_id()) if x.dim() == 3: @@ -51,8 +70,8 @@ def weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, num_blocks_n = triton.cdiv(x.shape[2], block_size) s_scales = torch.empty((x.shape[0], num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) for i in range(x.shape[0]): - y_quant[i], s_scales[i] = mm_weight_quant(x[i], block_size) + y_quant[i], s_scales[i] = mm_weight_quant(x[i], block_size, use_ue8m0_scales=use_ue8m0_scales) return y_quant, s_scales else: - y_quant, s_scales = mm_weight_quant(x, block_size) + y_quant, s_scales = mm_weight_quant(x, block_size, use_ue8m0_scales=use_ue8m0_scales) return y_quant, s_scales diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index 3645f6ef6..ce4d43b30 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -62,7 +62,7 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant device = output.weight.device - weight, scale = weight_quant(weight.cuda(device), self.block_size) + weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=True) output.weight.copy_(weight) output.weight_scale.copy_(scale) return @@ -90,6 +90,7 @@ def apply( column_major_scales=True, scale_tma_aligned=True, alloc_func=alloc_func, + use_ue8m0_scales=True, ) if out is None: From 49a2e5007dec88b2a6e77ae18dc9c33bca0b60b3 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 13 Jul 2026 12:20:16 +0800 Subject: [PATCH 3/5] feat: mega moe weight in place --- .../layer_weights/base_layer_weight.py | 1 + .../layer_weights/meta_weights/base_weight.py | 3 ++ .../fused_moe/fused_moe_weight.py | 10 +++++ .../fused_moe/grouped_fused_moe_ep.py | 40 +++++++++---------- 4 files changed, 33 insertions(+), 21 deletions(-) diff --git a/lightllm/common/basemodel/layer_weights/base_layer_weight.py b/lightllm/common/basemodel/layer_weights/base_layer_weight.py index b1d992a7c..7a49787e2 100644 --- a/lightllm/common/basemodel/layer_weights/base_layer_weight.py +++ b/lightllm/common/basemodel/layer_weights/base_layer_weight.py @@ -38,6 +38,7 @@ def verify_load(self): else: layer_num = None assert attr.verify_load(), f"Loading {attr_name} of layers {layer_num} fails." + attr.finalize_load() def _cuda(self, cpu_tensor): return cpu_tensor.contiguous().to(self.data_type_).cuda(get_current_device_id()) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py index 714e7acf4..d9f2b76ac 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py @@ -21,6 +21,9 @@ def _create_weight(self): def verify_load(self) -> bool: pass + def finalize_load(self) -> None: + pass + class BaseWeightTpl(BaseWeight): def __init__(self, tp_rank: int = None, tp_world_size: int = None, data_type: torch.dtype = None): 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 7f369c4fd..db0a27acc 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 @@ -294,6 +294,16 @@ def verify_load(self): ) return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok + def finalize_load(self): + if self.enable_ep_moe: + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( + transform_mega_moe_weights_in_place, + use_sm100_mega_moe, + ) + + if use_sm100_mega_moe(self.quant_method): + transform_mega_moe_weights_in_place(self.w13, self.w2) + def _create_weight(self): intermediate_size = self.split_inter_size self.e_score_correction_bias = 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 2c660c723..cd919922e 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 @@ -29,7 +29,7 @@ from lightllm.utils.tensor_buffer_manager import TensorBufferManager logger = init_logger(__name__) -_MEGA_MOE_STATES: Dict[Tuple[int, int, int, int], Dict[str, Any]] = {} +_MEGA_MOE_STATS: Dict[Tuple[int, int, int, int], torch.Tensor] = {} SUPPORTED_EP_EXPERT_DTYPES = ("fp8w8a8-b128-deepgemm", "fp4fp8-b32-deepgemm") @@ -241,31 +241,29 @@ def masked_group_gemm( return gemm_out_b -def _get_mega_moe_cache_state(w13: Any, w2: Any): +def _get_mega_moe_cumulative_stats(w13: Any, w2: Any): state_key = ( w13.weight.data_ptr(), w13.weight_scale.data_ptr(), w2.weight.data_ptr(), w2.weight_scale.data_ptr(), ) - return _MEGA_MOE_STATES.setdefault(state_key, {}) - - -def _get_mega_moe_weights(w13: Any, w2: Any, state: Dict[str, Any]): - if "weight_cache" not in state: - state["weight_cache"] = deep_gemm.transform_weights_for_mega_moe( - (w13.weight, w13.weight_scale), - (w2.weight, w2.weight_scale), - ) - return state["weight_cache"] + stats = _MEGA_MOE_STATS.get(state_key) + if stats is None: + stats = torch.zeros((w13.weight.shape[0],), device=w13.weight.device, dtype=torch.int32) + _MEGA_MOE_STATS[state_key] = stats + return stats -def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, state: Dict[str, Any]): - stats = state.get("stats") - if stats is None or stats.numel() != num_local_experts or stats.device != device: - stats = torch.zeros((num_local_experts,), device=device, dtype=torch.int32) - state["stats"] = stats - return stats +def transform_mega_moe_weights_in_place(w13: Any, w2: Any): + """Convert to Mega MoE layout without retaining a second weight copy.""" + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( + (w13.weight, w13.weight_scale), + (w2.weight, w2.weight_scale), + ) + w13.weight.copy_(transformed_l1[0]) + w13.weight_scale.copy_(transformed_l1[1]) + w2.weight_scale.copy_(transformed_l2[1]) def mega_moe_impl( @@ -289,9 +287,9 @@ def mega_moe_impl( f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" ) - state = _get_mega_moe_cache_state(w13, w2) - l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) - stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) + l1_weights = (w13.weight, w13.weight_scale) + l2_weights = (w2.weight, w2.weight_scale) + stats = _get_mega_moe_cumulative_stats(w13, w2) _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) From 47799e26ba24658caaa620bfd6e9fcba58abde0c Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 14 Jul 2026 14:53:53 +0800 Subject: [PATCH 4/5] refine --- .../fused_moe/fused_moe_weight.py | 3 +- .../fused_moe/grouped_fused_moe_ep.py | 38 +++++++++---------- .../fp8w8a8_block_quant_kernel.py | 4 +- lightllm/common/quantization/deepgemm.py | 36 +++++++----------- 4 files changed, 33 insertions(+), 48 deletions(-) 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 db0a27acc..8420f5260 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 @@ -295,7 +295,7 @@ def verify_load(self): return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok def finalize_load(self): - if self.enable_ep_moe: + if self.enable_ep_moe and not getattr(self, "_mega_moe_weights_transformed", False): from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( transform_mega_moe_weights_in_place, use_sm100_mega_moe, @@ -303,6 +303,7 @@ def finalize_load(self): if use_sm100_mega_moe(self.quant_method): transform_mega_moe_weights_in_place(self.w13, self.w2) + self._mega_moe_weights_transformed = True def _create_weight(self): intermediate_size = self.split_inter_size 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 cd919922e..59412cd30 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 @@ -3,7 +3,7 @@ import torch import triton import triton.language as tl -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, List, Optional, Tuple from lightllm.distributed import dist_group_manager from lightllm.utils.log_utils import init_logger from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd @@ -29,7 +29,6 @@ from lightllm.utils.tensor_buffer_manager import TensorBufferManager logger = init_logger(__name__) -_MEGA_MOE_STATS: Dict[Tuple[int, int, int, int], torch.Tensor] = {} SUPPORTED_EP_EXPERT_DTYPES = ("fp8w8a8-b128-deepgemm", "fp4fp8-b32-deepgemm") @@ -99,6 +98,7 @@ def _mega_moe_quant_topk_to_buffer_kernel( FP8_MIN: tl.constexpr, FP8_MAX: tl.constexpr, TOPK: tl.constexpr, + TOPK_BLOCK: tl.constexpr, GROUP_SIZE: tl.constexpr, BLOCK: tl.constexpr, ): @@ -133,20 +133,27 @@ def _mega_moe_quant_topk_to_buffer_kernel( tl.store(x_sf_out_ptr + token_id * stride_x_sf_out_m + pack_id * stride_x_sf_out_k, packed_scale) if pack_id == 0: - topk_offsets = tl.arange(0, TOPK) - topk_idx = tl.load(topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k) + topk_offsets = tl.arange(0, TOPK_BLOCK) + topk_mask = topk_offsets < TOPK + topk_idx = tl.load( + topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k, + mask=topk_mask, + ) topk_weights = tl.load( - topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k + topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k, + mask=topk_mask, ) tl.store( topk_idx_out_ptr + token_id * stride_topk_idx_out_m + topk_offsets * stride_topk_idx_out_k, topk_idx.to(topk_idx_out_ptr.dtype.element_ty), + mask=topk_mask, ) tl.store( topk_weights_out_ptr + token_id * stride_topk_weights_out_m + topk_offsets * stride_topk_weights_out_k, topk_weights.to(topk_weights_out_ptr.dtype.element_ty), + mask=topk_mask, ) @@ -162,8 +169,12 @@ def _prepare_mega_moe_buffer( return assert hidden_size % (group_size * 4) == 0, "packed UE8M0 scale requires four FP8 groups per int32" assert hidden_states.is_contiguous(), "hidden_states must be contiguous" + assert topk_ids.shape == topk_weights.shape and topk_ids.shape[0] == num_tokens + assert topk_ids.shape[1] > 0 assert buffer.x.shape[0] >= num_tokens and buffer.x.shape[1] == hidden_size assert buffer.x_sf.shape[0] >= num_tokens and buffer.x_sf.shape[1] == hidden_size // group_size // 4 + assert buffer.topk_idx.shape[0] >= num_tokens and buffer.topk_idx.shape[1] == topk_ids.shape[1] + assert buffer.topk_weights.shape[0] >= num_tokens and buffer.topk_weights.shape[1] == topk_ids.shape[1] block = group_size * 4 finfo = torch.finfo(buffer.x.dtype) @@ -192,6 +203,7 @@ def _prepare_mega_moe_buffer( FP8_MIN=finfo.min, FP8_MAX=finfo.max, TOPK=topk_ids.shape[1], + TOPK_BLOCK=triton.next_power_of_2(topk_ids.shape[1]), GROUP_SIZE=group_size, BLOCK=block, num_warps=4, @@ -241,20 +253,6 @@ def masked_group_gemm( return gemm_out_b -def _get_mega_moe_cumulative_stats(w13: Any, w2: Any): - state_key = ( - w13.weight.data_ptr(), - w13.weight_scale.data_ptr(), - w2.weight.data_ptr(), - w2.weight_scale.data_ptr(), - ) - stats = _MEGA_MOE_STATS.get(state_key) - if stats is None: - stats = torch.zeros((w13.weight.shape[0],), device=w13.weight.device, dtype=torch.int32) - _MEGA_MOE_STATS[state_key] = stats - return stats - - def transform_mega_moe_weights_in_place(w13: Any, w2: Any): """Convert to Mega MoE layout without retaining a second weight copy.""" transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( @@ -289,7 +287,6 @@ def mega_moe_impl( l1_weights = (w13.weight, w13.weight_scale) l2_weights = (w2.weight, w2.weight_scale) - stats = _get_mega_moe_cumulative_stats(w13, w2) _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) @@ -298,7 +295,6 @@ def mega_moe_impl( l1_weights, l2_weights, buffer, - cumulative_local_expert_recv_stats=stats, ) return output diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py index 1c2caa967..c711aae85 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py @@ -53,9 +53,7 @@ def mm_weight_quant( s_scales = torch.empty((num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) - weight_quant_kernel[grid]( - x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales - ) + weight_quant_kernel[grid](x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales) return y_quant, s_scales diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index ce4d43b30..8f55db2ff 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -5,6 +5,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.log_utils import init_logger +from lightllm.utils.device_utils import is_sm100_gpu logger = init_logger(__name__) @@ -62,7 +63,7 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant device = output.weight.device - weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=True) + weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=is_sm100_gpu()) output.weight.copy_(weight) output.weight_scale.copy_(scale) return @@ -90,7 +91,7 @@ def apply( column_major_scales=True, scale_tma_aligned=True, alloc_func=alloc_func, - use_ue8m0_scales=True, + use_ue8m0_scales=is_sm100_gpu(), ) if out is None: @@ -181,30 +182,19 @@ def _create_weight( self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims - assert in_dim % 2 == 0, "FP4 packed weight requires even input dimension" - assert in_dim % self.block_size == 0, "FP4 scale dimension must be divisible by block_size" - scales_per_int32 = 4 - scale_layout_k = self.block_size * scales_per_int32 - assert in_dim % scale_layout_k == 0, ( - f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" - ) + scale_layout_k = self.block_size * 4 # Each int32 packs four UE8M0 scales. + assert ( + in_dim % scale_layout_k == 0 + ), f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" expert_prefix = (num_experts,) if num_experts > 1 else () weight = torch.empty(expert_prefix + (out_dim, in_dim // 2), dtype=torch.int8).cuda(device_id) scale_dim = in_dim // scale_layout_k - if num_experts > 1: - weight_scale = torch.empty_strided( - (num_experts, out_dim, scale_dim), - (out_dim * scale_dim, 1, out_dim), - dtype=torch.int32, - device=f"cuda:{device_id}", - ) - else: - weight_scale = torch.empty_strided( - (out_dim, scale_dim), - (1, out_dim), - dtype=torch.int32, - device=f"cuda:{device_id}", - ) + weight_scale = torch.empty_strided( + expert_prefix + (out_dim, scale_dim), + (out_dim * scale_dim, 1, out_dim) if num_experts > 1 else (1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) mm_param_list = self._split_weight_pack( mm_param, From 54b1a7fb1a6060b89d4229d4b369b6b2e9bf8013 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 14 Jul 2026 18:29:26 +0800 Subject: [PATCH 5/5] fix: fix bugs for fp8 quant in sm100 --- .../fused_moe/impl/deepgemm_impl.py | 3 +++ .../fused_moe/grouped_fused_moe_ep.py | 20 ++++++++++++++++--- .../moe_silu_and_mul_mix_quant_ep.py | 16 ++++++++++++++- .../test_moe_silu_and_mul_mix_quant_ep.py | 13 ++++++++++-- 4 files changed, 46 insertions(+), 6 deletions(-) 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 a5ba656c9..c9cc25fac 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 @@ -17,6 +17,7 @@ from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.triton_utils.autotuner import Autotuner from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair +from lightllm.utils.device_utils import is_sm100_gpu class FuseMoeDeepGEMM(FuseMoeTriton): @@ -123,6 +124,8 @@ 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=is_sm100_gpu(), + use_ue8m0=is_sm100_gpu(), async_finish=False, return_recv_hook=True, ) 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 59412cd30..ef2563ff6 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 @@ -246,7 +246,14 @@ def masked_group_gemm( qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) - silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + gemm_out_a, + qsilu_out, + qsilu_out_scale, + block_size, + masked_m, + use_ue8m0_scales=is_sm100_gpu(), + ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) @@ -312,7 +319,9 @@ def quantize_fused_experts_input( if w13.weight.ndim == 3: block_size_k = w13.weight.shape[2] // w13.weight_scale.shape[2] assert block_size_k == 128, "block_size_k must be 128" - return per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w13.weight.dtype) + return per_token_group_quant_fp8( + hidden_states, block_size_k, dtype=w13.weight.dtype, use_ue8m0_scales=is_sm100_gpu() + ) def fused_experts( @@ -385,7 +394,9 @@ 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) + qinput_tensor, input_scale = per_token_group_quant_fp8( + hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=is_sm100_gpu() + ) allocate_on_comm_stream = previous_event is not None # Expanded dispatch directly produces expert-contiguous, alignment-padded inputs: # recv_x[0]: [num_expanded_tokens, hidden] @@ -470,6 +481,8 @@ def fused_experts_impl( num_max_dispatch_tokens_per_rank, num_experts, use_fp8=use_fp8_w8a8, + round_scale=is_sm100_gpu(), + use_ue8m0=is_sm100_gpu(), async_finish=False, return_recv_hook=False, ) @@ -679,6 +692,7 @@ def workspace_quant_alloc(shape, dtype, device): column_major_scales=True, scale_tma_aligned=True, alloc_func=workspace_quant_alloc, + use_ue8m0_scales=is_sm100_gpu(), ) 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 aa91f15ed..7380cd1b7 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 @@ -6,6 +6,14 @@ from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + @triton.jit def _silu_and_mul_post_quant_kernel( input_ptr, @@ -26,6 +34,7 @@ def _silu_and_mul_post_quant_kernel( fp8_min, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, ): expert_id = tl.program_id(2) @@ -61,7 +70,10 @@ def _silu_and_mul_post_quant_kernel( gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) - output_s = _absmax / fp8_max + if USE_UE8M0_SCALE: + output_s = _ceil_to_ue8m0(tl.maximum(_absmax, 1.0e-4) / fp8_max) + else: + output_s = _absmax / fp8_max output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to(output_ptr.dtype.element_ty) tl.store( output_ptr_offs + token_index * stride_output_1, @@ -80,6 +92,7 @@ def silu_and_mul_masked_post_quant_fwd( output_scale: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, + use_ue8m0_scales: bool = False, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -135,6 +148,7 @@ def silu_and_mul_masked_post_quant_fwd( fp8_min, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, + USE_UE8M0_SCALE=use_ue8m0_scales, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), num_warps=num_warps, ) diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py index 8783f35a4..3bebd2919 100644 --- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py +++ b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py @@ -36,7 +36,8 @@ def is_fp8_native_supported(): for token_num in range(1, 7, 2) ], ) -def test_silu_and_mul_masked(expert_num, token_num, hidden_dim): +@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) +def test_silu_and_mul_masked(expert_num, token_num, hidden_dim, use_ue8m0_scales): quant_group_size = 128 in_tensor = torch.randn((expert_num, token_num, hidden_dim), dtype=torch.bfloat16, device="cuda") out_tensor = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=torch.float8_e4m3fn, device="cuda") @@ -53,9 +54,17 @@ def test_silu_and_mul_masked(expert_num, token_num, hidden_dim): true_out_tensor_mid.view(-1, hidden_dim // 2), quant_group_size, alloc_func=torch.empty, + use_ue8m0_scales=use_ue8m0_scales, ) - silu_and_mul_masked_post_quant_fwd(in_tensor, out_tensor, out_scale_tensor, quant_group_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + in_tensor, + out_tensor, + out_scale_tensor, + quant_group_size, + masked_m, + use_ue8m0_scales=use_ue8m0_scales, + ) true_out_tensor = true_out_tensor.view(out_tensor.shape) true_out_scale_tensor = true_out_scale_tensor.view(out_scale_tensor.shape)