diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index e4e543cb431..b87b87f136d 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -256,11 +256,6 @@ class FusedMoEQuantConfig: mx_alignment: int = 0 - # When True, MXFP8 (1x32) activations are quantized into DeepGEMM's packed - # UE8M0 scale layout in the prepare phase. Set only for the DeepGEMM mxfp8 - # backend; other backends (FlashInfer/Marlin) want the unpacked layout. - use_deep_gemm_packed_mxfp8: bool = False - def __post_init__(self): assert not self.per_act_token_quant or self.block_shape is None, ( "illegal quantization" @@ -511,7 +506,6 @@ class FusedMoEQuantConfig: gemm1_alpha: float | None = None, gemm1_beta: float | None = None, gemm1_clamp_limit: float | None = None, - use_deep_gemm_packed_mxfp8: bool = False, ) -> "FusedMoEQuantConfig": """ General builder function for a FusedMoEQuantConfig. @@ -583,7 +577,6 @@ class FusedMoEQuantConfig: gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=gemm1_clamp_limit, - use_deep_gemm_packed_mxfp8=use_deep_gemm_packed_mxfp8, ) assert quant_config.per_act_token_quant == per_act_token_quant assert quant_config.per_out_ch_quant == per_out_ch_quant diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py index 7634d236e77..c74cb2d9a7b 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_utils.py @@ -130,6 +130,9 @@ def _fwd_kernel_ep_scatter_2( HIDDEN_SIZE_PAD: tl.constexpr, SCALE_HIDDEN_SIZE: tl.constexpr, SCALE_HIDDEN_SIZE_PAD: tl.constexpr, + PACK_UE8M0: tl.constexpr, + SCALE_PACKED_SIZE: tl.constexpr, + SCALE_PACKED_SIZE_PAD: tl.constexpr, ): start_token_id = tl.program_id(0) grid_num = tl.num_programs(0) @@ -137,16 +140,47 @@ def _fwd_kernel_ep_scatter_2( offset_in = tl.arange(0, HIDDEN_SIZE_PAD) mask = offset_in < HIDDEN_SIZE - offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) - mask_s = offset_in_s < SCALE_HIDDEN_SIZE - output_tensor_stride0 = output_tensor_stride0.to(tl.int64) + if PACK_UE8M0: + # One int32 per 4 consecutive 32-wide UE8M0 groups, stored MN-major. + offs_pk = tl.arange(0, SCALE_PACKED_SIZE_PAD) + mask_pk = offs_pk < SCALE_PACKED_SIZE + else: + offset_in_s = tl.arange(0, SCALE_HIDDEN_SIZE_PAD) + mask_s = offset_in_s < SCALE_HIDDEN_SIZE + for token_id in range(start_token_id, total_token_num, grid_num): to_copy = tl.load(recv_x + token_id * recv_x_stride0 + offset_in, mask=mask) - to_copy_s = tl.load( - recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, mask=mask_s - ) + + if PACK_UE8M0: + # Pack 4 UE8M0 bytes into one int32 (byte j = group 4*pk+j). + base_s = recv_x_scale + token_id * recv_x_scale_stride0 + g0, g1 = offs_pk * 4, offs_pk * 4 + 1 + g2, g3 = offs_pk * 4 + 2, offs_pk * 4 + 3 + b0 = tl.load( + base_s + g0 * recv_x_scale_stride1, mask=g0 < SCALE_HIDDEN_SIZE + ) + b1 = tl.load( + base_s + g1 * recv_x_scale_stride1, mask=g1 < SCALE_HIDDEN_SIZE + ) + b2 = tl.load( + base_s + g2 * recv_x_scale_stride1, mask=g2 < SCALE_HIDDEN_SIZE + ) + b3 = tl.load( + base_s + g3 * recv_x_scale_stride1, mask=g3 < SCALE_HIDDEN_SIZE + ) + packed_s = ( + b0.to(tl.int32) + | (b1.to(tl.int32) << 8) + | (b2.to(tl.int32) << 16) + | (b3.to(tl.int32) << 24) + ) + else: + to_copy_s = tl.load( + recv_x_scale + token_id * recv_x_scale_stride0 + offset_in_s, + mask=mask_s, + ) for topk_index in tl.range(0, topk_num, 1, num_stages=4): expert_id = tl.load(recv_topk + token_id * recv_topk_stride0 + topk_index) @@ -164,11 +198,21 @@ def _fwd_kernel_ep_scatter_2( output_tensor_ptr = ( output_tensor + dest_token_index_i64 * output_tensor_stride0 ) + tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) + output_tensor_scale_ptr = ( output_tensor_scale + dest_token_index * output_tensor_scale_stride0 ) - tl.store(output_tensor_ptr + offset_in, to_copy, mask=mask) - tl.store(output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s) + if PACK_UE8M0: + tl.store( + output_tensor_scale_ptr + offs_pk * output_tensor_scale_stride1, + packed_s, + mask=mask_pk, + ) + else: + tl.store( + output_tensor_scale_ptr + offset_in_s, to_copy_s, mask=mask_s + ) @torch.no_grad() @@ -184,6 +228,7 @@ def ep_scatter( m_indices: torch.Tensor, output_index: torch.Tensor, block_size: int = 128, + pack_ue8m0: bool = False, ): BLOCK_E = 128 # token num of per expert is aligned to 128 BLOCK_D = block_size # block size of activation-scale quantization @@ -196,6 +241,10 @@ def ep_scatter( assert m_indices.shape[0] % BLOCK_E == 0 assert expert_start_loc.shape[0] == num_experts + # pack_ue8m0: scatter packs 4 UE8M0 bytes per int32; else copies scales as-is. + scale_hidden_size = hidden_size // BLOCK_D + scale_packed_size = (scale_hidden_size + 3) // 4 if pack_ue8m0 else 1 + _fwd_kernel_ep_scatter_1[(grid,)]( num_recv_tokens_per_expert, expert_start_loc, @@ -235,8 +284,11 @@ def ep_scatter( num_warps=num_warps, HIDDEN_SIZE=hidden_size, HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size), - SCALE_HIDDEN_SIZE=hidden_size // BLOCK_D, - SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(hidden_size // BLOCK_D), + SCALE_HIDDEN_SIZE=scale_hidden_size, + SCALE_HIDDEN_SIZE_PAD=triton.next_power_of_2(scale_hidden_size), + PACK_UE8M0=pack_ue8m0, + SCALE_PACKED_SIZE=scale_packed_size, + SCALE_PACKED_SIZE_PAD=triton.next_power_of_2(scale_packed_size), ) return @@ -382,9 +434,21 @@ def deepgemm_moe_permute( if aq_out is None: aq_out = torch.empty((M_sum, H), device=device, dtype=aq.dtype) - aq_scale_out = torch.empty( - (M_sum, H // block_k), device=device, dtype=torch.float32 - ) + # uint8 UE8M0 (MXFP8) -> scatter packs into DeepGEMM's int32 MN-major + # TMA-aligned layout; float32 (FP8/FP4) scattered row-major as-is. + pack_ue8m0 = aq_scale.dtype == torch.uint8 + sf_k = H // block_k + if pack_ue8m0: + packed_sf_k = (sf_k + 3) // 4 + tma_aligned_mn = round_up(M_sum, 4) + aq_scale_out = torch.empty_strided( + (M_sum, packed_sf_k), + (1, tma_aligned_mn), + device=device, + dtype=torch.int32, + ) + else: + aq_scale_out = torch.empty((M_sum, sf_k), device=device, dtype=torch.float32) # DeepGEMM uses negative values in m_indices (here expert_ids) to mark # completely invalid / padded blocks that should be skipped. We always @@ -419,6 +483,7 @@ def deepgemm_moe_permute( m_indices=expert_ids, output_index=inv_perm, block_size=block_k, + pack_ue8m0=pack_ue8m0, ) return aq_out, aq_scale_out, expert_ids, inv_perm diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index b640cf61bce..b720b6c87f1 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -560,11 +560,9 @@ def make_fp8_moe_quant_config( g2_alphas=(w2_scale * a2_scale).squeeze(), gemm1_clamp_limit=swiglu_limit, ) - # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to the - # mxfp8 activation quant rather than standard FP8 block quantization. - # Non-swizzled layout is required since the TRTLLM kernel expects scales in - # (num_tokens, hidden_dim // 32) format. DeepGEMM instead needs its packed - # UE8M0 scale layout, selected via use_deep_gemm_packed_mxfp8. + # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are + # the non-swizzled (num_tokens, hidden_dim // 32) uint8 UE8M0 layout for all + # backends; the DeepGEMM expert permute repacks them for the grouped GEMM. if block_shape == [1, 32]: return FusedMoEQuantConfig.make( "mxfp8", @@ -579,7 +577,6 @@ def make_fp8_moe_quant_config( gemm1_alpha=gemm1_alpha, gemm1_beta=gemm1_beta, gemm1_clamp_limit=swiglu_limit, - use_deep_gemm_packed_mxfp8=(fp8_backend == Fp8MoeBackend.DEEPGEMM), ) # All other backends use normal config. diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py index f8f0e990363..69587770c40 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py @@ -32,7 +32,6 @@ def _quantize_input( block_shape=quant_config.block_shape, is_scale_swizzled=quant_config.is_scale_swizzled, mx_alignment=quant_config.mx_alignment, - use_deep_gemm_packed_mxfp8=quant_config.use_deep_gemm_packed_mxfp8, ) return a1q, a1q_scale diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index 55098c0ea09..b8c84ad2af2 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -258,7 +258,6 @@ def moe_kernel_quantize_input( ocp_mx_scheme: str | None = None, quantization_emulation: bool = False, mx_alignment: int = 0, - use_deep_gemm_packed_mxfp8: bool = False, ) -> tuple[torch.Tensor, torch.Tensor | None]: # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation if ocp_mx_scheme is not None: @@ -314,13 +313,8 @@ def moe_kernel_quantize_input( "moe_kernel_quantize_input does not support quant_dtype='mxfp8' MOE " "quantization emulation. Please open an issue." ) - # DeepGEMM consumes per-(1, group) float32 scales (cast to UE8M0 - # internally), exactly like the FP8 128-block path but with group=32. - # The expert permute (deepgemm_moe_permute) is given the same group so - # the scale width matches. - if use_deep_gemm_packed_mxfp8: - assert block_shape is not None - return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape) + # Non-swizzled (M, K/32) uint8 UE8M0 scales; deepgemm_moe_permute packs + # them for DeepGEMM, TRTLLM takes them as-is. return _mxfp8_e4m3_quantize( A, A_scale,