From f0153ad225bdcbdd652245e07a101379430cc0c9 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Tue, 10 Mar 2026 03:52:05 +0800 Subject: [PATCH] [AMD][Feature] support fp4 dispatch and fp8 combine in moriep (#19757) Co-authored-by: Duyi-Wang --- docs/references/environment_variables.md | 2 + python/sglang/srt/layers/moe/ep_moe/layer.py | 42 +++++- .../sglang/srt/layers/moe/rocm_moe_utils.py | 140 ++++++++++++++++++ .../srt/layers/moe/token_dispatcher/moriep.py | 104 +++++++++++-- 4 files changed, 265 insertions(+), 23 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 5702f196f..872e7697f 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -75,6 +75,8 @@ SGLang supports various environment variables that can be used to configure its | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_MORI_FP8_DISP` | Use FP8 for dispatch | `"false"` | +| `SGLANG_MORI_FP4_DISP` | Use MXFP4 for dispatch | `"false"` | +| `SGLANG_MORI_FP8_COMB` | Use FP8 for combine | `"false"` | | `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | Maximum number of dispatch tokens per rank for MORI-EP buffer allocation | `4096` | | `SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD` | Threshold for switching between `InterNodeV1` and `InterNodeV1LL` kernel types. `InterNodeV1LL` is used if `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is less than or equal to this threshold; otherwise, `InterNodeV1` is used. | `256` | | `SGLANG_MORI_QP_PER_TRANSFER` | Number of RDMA Queue Pairs (QPs) used per transfer operation | `1` | diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 54908a8f2..5fb874efc 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -18,7 +18,7 @@ from sglang.srt.layers.moe.fused_moe_triton.layer import ( FusedMoE, moe_forward_piecewise_cuda_graph_impl, ) -from sglang.srt.layers.moe.rocm_moe_utils import upscale +from sglang.srt.layers.moe.rocm_moe_utils import upscale, upscale_mxfp4 from sglang.srt.layers.moe.token_dispatcher.deepep import ( DeepEPLLCombineInput, DeepEPNormalCombineInput, @@ -653,11 +653,36 @@ class MoriEPMoE(DeepEPMoE): quant_type = QuantType.No - if not is_fp8_quant and dispatch_scale is not None: - dispatch_a1 = upscale( - dispatch_a1, dispatch_scale, dispatch_recv_token_num, output_dtype - ) - dispatch_scale = None + if ( + not is_fp8_quant + and dispatch_scale is not None + and dispatch_a1.dtype != torch.float4_e2m1fn_x2 + ): + if is_quark_w4a4: + # W4A4 model with FP8 dispatch: must dequant FP8->BF16 first, + # because the FP4 per_1x32 quantization path needs BF16 input + dispatch_a1 = upscale( + dispatch_a1, dispatch_scale, dispatch_recv_token_num, output_dtype + ) + dispatch_scale = None + else: + # Non-W4A4 model with FP8 dispatch: pass FP8 hidden_states + scale + # directly to fused_moe, avoiding unnecessary dequant->requant round-trip + quant_type = QuantType.per_128x128 + + if dispatch_a1.dtype == torch.float4_e2m1fn_x2 and dispatch_scale is not None: + if is_fp8_quant: + # FP8 weights + FP4 dispatch is not supported by fused_moe kernels + # (no kernel for q_dtype_a=fp4x2, q_dtype_w=fp8). + # Must dequant FP4->BF16 first; fused_moe will re-quant to FP8 internally. + dispatch_a1 = upscale_mxfp4( + dispatch_a1, dispatch_scale, dispatch_recv_token_num, output_dtype + ) + dispatch_scale = None + elif quant_type == QuantType.No: + # Skip upscale_mxfp4: pass FP4 hidden_states + scale directly to fused_moe + # fused_moe with QuantType.per_1x32 can accept pre-quantized fp4x2 input + quant_type = QuantType.per_1x32 if is_quark_w4a4: if hasattr(torch, "float4_e2m1fn_x2"): @@ -677,7 +702,10 @@ class MoriEPMoE(DeepEPMoE): if hasattr(self, "w2_weight_scale_inv"): w2_scale = self.w2_weight_scale_inv - quant_type = QuantType.per_128x128 + # Only set per_128x128 if quant_type was not already set by + # a prior dispatch path (e.g. FP4 dispatch sets per_1x32) + if quant_type == QuantType.No: + quant_type = QuantType.per_128x128 # [KK TODO] should to call the apply of quant method to handle fused moe hidden_states = fused_moe( diff --git a/python/sglang/srt/layers/moe/rocm_moe_utils.py b/python/sglang/srt/layers/moe/rocm_moe_utils.py index 77a58a2f4..43382610f 100644 --- a/python/sglang/srt/layers/moe/rocm_moe_utils.py +++ b/python/sglang/srt/layers/moe/rocm_moe_utils.py @@ -187,3 +187,143 @@ def upscale(hidden_state, hidden_state_scale, recv_token_num, output_dtype): ) return Out + + +@triton.jit +def upscale_fp4x2_block32_kernel( + A_u8_ptr, # *uint8 (view from float4_e2m1fn_x2) + S_u8_ptr, # *uint8 (view from float8_e8m0fnu), shape (M, N_fp4/32) + Out_ptr, # *fp16/fp32/bf16, shape (M, N_fp4) + N_FP4: tl.constexpr, + recv_token_num, + stride_am, + stride_an, # A strides (in uint8 elements) for (M, packed_N) + stride_sm, + stride_sn, # S strides (in uint8 elements) for (M, N_FP4/32) + stride_om, + stride_on, # Out strides (in output elements) for (M, N_FP4) + BLOCK_N: tl.constexpr, + OUT_DTYPE: tl.constexpr, # tl.float16 / tl.float32 / tl.bfloat16 +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + recv_token_num_val = tl.load(recv_token_num) + if pid_m >= recv_token_num_val: + return + + offs = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask = offs < N_FP4 + + # -------------------------- + # Load packed fp4x2 byte + # -------------------------- + byte_idx = offs >> 1 # offs // 2 + is_hi = (offs & 1) != 0 # select high nibble? + + a_ptrs = A_u8_ptr + pid_m * stride_am + byte_idx * stride_an + a_byte = tl.load(a_ptrs, mask=mask, other=0).to(tl.int32) + + lo = a_byte & 0xF + hi = (a_byte >> 4) & 0xF + code = tl.where(is_hi, hi, lo).to(tl.int32) # 0..15 + + # -------------------------- + # Decode float4_e2m1fn + # layout: [sign|exp(2)|mant(1)] + # bias=1, finite-only + # -------------------------- + sign = (code >> 3) & 0x1 + exp = (code >> 1) & 0x3 + mant = code & 0x1 + + mant_f = mant.to(tl.float32) * 0.5 + is_sub = exp == 0 + + # normal: 2^(exp-bias) * (1 + mant/2), bias=1 + e_norm = (exp - 1).to(tl.float32) + val_norm = tl.exp2(e_norm) * (1.0 + mant_f) + + # subnorm/zero: mant/2 * 2^(1-bias) = mant/2 + val_sub = mant_f + + val = tl.where(is_sub, val_sub, val_norm) + val = tl.where(sign != 0, -val, val) # apply sign + + # -------------------------- + # Per-token block32 scale: scale_idx = offs // 32 + # scale dtype: float8_e8m0fnu stored in uint8 + # decode: e==0 -> 0 + # e in [1..254] -> 2^(e-127) + # e==255 -> clamp to 254 + # -------------------------- + scale_idx = offs >> 5 # offs // 32 + + s_ptrs = S_u8_ptr + pid_m * stride_sm + scale_idx * stride_sn + e = tl.load(s_ptrs, mask=mask, other=0).to(tl.int32) + + e = tl.minimum(e, 254) # clamp 255->254 + is_zero = e == 0 + exp_s = (e - 127).to(tl.float32) + s = tl.exp2(exp_s) + s = tl.where(is_zero, 0.0, s) + + out = (val * s).to(OUT_DTYPE) + + out_ptrs = Out_ptr + pid_m * stride_om + offs * stride_on + tl.store(out_ptrs, out, mask=mask) + + +def upscale_mxfp4(hidden_state, hidden_state_scale, recv_token_num, output_dtype): + """ + hidden_state: (M, packed_N) torch.float4_e2m1fn_x2 + hidden_state_scale: (M, packed_N*2/32) = (M, N_fp4/32) torch.float8_e8m0fnu + output: (M, N_fp4) output_dtype + """ + assert hidden_state.dtype == torch.float4_e2m1fn_x2, hidden_state.dtype + assert hidden_state_scale.dtype == torch.float8_e8m0fnu, hidden_state_scale.dtype + assert hidden_state.is_contiguous() or True # stride-based load OK + + M, packed_N = hidden_state.shape + N_fp4 = packed_N * 2 + + # scale second dim must be N_fp4/32 + assert hidden_state_scale.shape[0] == M + assert hidden_state_scale.shape[1] == (N_fp4 // 32), ( + hidden_state_scale.shape, + N_fp4, + ) + + # Triton doesn't (reliably) accept torch.float4/float8 pointers directly. + # Use raw uint8 views. + A_u8 = hidden_state.view(torch.uint8) + S_u8 = hidden_state_scale.view(torch.uint8) + + Out = torch.empty((M, N_fp4), dtype=output_dtype, device=hidden_state.device) + + BLOCK_N = 256 + grid = (M, triton.cdiv(N_fp4, BLOCK_N)) + + OUT_TL = ( + tl.float16 + if output_dtype == torch.float16 + else tl.bfloat16 if output_dtype == torch.bfloat16 else tl.float32 + ) + + upscale_fp4x2_block32_kernel[grid]( + A_u8, + S_u8, + Out, + N_FP4=N_fp4, + recv_token_num=recv_token_num, + stride_am=A_u8.stride(0), + stride_an=A_u8.stride(1), + stride_sm=S_u8.stride(0), + stride_sn=S_u8.stride(1), + stride_om=Out.stride(0), + stride_on=Out.stride(1), + BLOCK_N=BLOCK_N, + OUT_DTYPE=OUT_TL, + num_warps=4, + ) + return Out diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 6c189d0f7..85fe8653e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -170,6 +170,19 @@ def get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank: int = 4096): } +@lru_cache(maxsize=2) +def _get_mori_dispatch_quant_flags(): + fp8_dispatch = get_bool_env_var("SGLANG_MORI_FP8_DISP", "False") + fp4_dispatch = get_bool_env_var("SGLANG_MORI_FP4_DISP", "False") + if fp8_dispatch and fp4_dispatch: + logger.warning( + "Both SGLANG_MORI_FP8_DISP and SGLANG_MORI_FP4_DISP are set to True. " + "Using SGLANG_MORI_FP4_DISP and ignoring SGLANG_MORI_FP8_DISP." + ) + fp8_dispatch = False + return fp8_dispatch, fp4_dispatch + + # init_mori_op only needs do once in model initial stage # use lru_cache to reuse the same mori_op instance to avoid the init overhead for mori @lru_cache(maxsize=2) @@ -211,17 +224,41 @@ def init_mori_op( block_num = cfg.block_num rdma_block_num = cfg.rdma_block_num + hidden_dim = hidden_size + scale_dim = 1 + data_type = fp8_dtype + scale_type_size = torch.float32.itemsize + + fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags() + + if fp8_dispatch: + scale_dim = hidden_size // 128 + elif fp4_dispatch: + # FP4 kernel still takes the original hidden size and do quantization internally, so hidden_dim is not reduced. The reason is that for FP4 quantization, we need to keep the original hidden size to calculate the quantization scale correctly. don't use packed hidden size for FP4 kernel. + hidden_dim = hidden_size + scale_dim = hidden_size // 32 + data_type = torch.float4_e2m1fn_x2 + scale_type_size = torch.float8_e8m0fnu.itemsize + + if mode == EpMode.INTRA_NODE: + if num_max_dispatch_tokens_per_rank < 128: + block_num = 225 + warp_num_per_block = 5 + else: + block_num = 256 + warp_num_per_block = 16 + + combine_quant_type = "none" + if get_bool_env_var("SGLANG_MORI_FP8_COMB", "False"): + combine_quant_type = "fp8_direct_cast" + mori_config = mori.ops.EpDispatchCombineConfig( rank=rank, world_size=world_size, - data_type=fp8_dtype, - hidden_dim=hidden_size, - scale_dim=( - hidden_size // 128 - if get_bool_env_var("SGLANG_MORI_FP8_DISP", "False") - else 1 - ), - scale_type_size=torch.float32.itemsize, + data_type=data_type, + hidden_dim=hidden_dim, + scale_dim=scale_dim, + scale_type_size=scale_type_size, max_token_type_size=params_dtype.itemsize, max_num_inp_token_per_rank=num_max_dispatch_tokens_per_rank, num_experts_per_rank=num_local_experts, @@ -232,6 +269,7 @@ def init_mori_op( gpu_per_node=gpu_per_node, rdma_block_num=rdma_block_num, num_qp_per_pe=2, + quant_type=combine_quant_type, ) mori_op = mori.ops.EpDispatchCombineOp(mori_config) return mori_op @@ -346,7 +384,8 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase): self.async_finish = async_finish self.quant_config = {} # [kk TODO] need to support mxfp4 type - self.quant_func = get_hip_quant(QuantType.per_1x128) + self.fp8_quant_func = get_hip_quant(QuantType.per_1x128) + self.fp4_quant_func = get_hip_quant(QuantType.per_1x32) self.enable_dual_stream = is_tbo_enabled() self._comm_stream = None if self.enable_dual_stream: @@ -378,17 +417,17 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase): topk_ids, previous_event, ): - num_token = hidden_states.shape[0] + num_tokens = hidden_states.shape[0] output_dtype = hidden_states.dtype scale = None - fp8_dispatch = get_bool_env_var("SGLANG_MORI_FP8_DISP", "False") + fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags() if fp8_dispatch: # FP8 quant - if num_token > 0: + if num_tokens > 0: # NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD. - hidden_states, scale = self.quant_func( + hidden_states, scale = self.fp8_quant_func( hidden_states, quant_dtype=fp8_dtype ) else: @@ -401,6 +440,22 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase): device=hidden_states.device, ) + elif fp4_dispatch: + # FP4 quant + if num_tokens > 0: + hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False) + else: + hidden_states = torch.empty( + (0, self.hidden_size // 2), + dtype=torch.float4_e2m1fn_x2, + device=hidden_states.device, + ) + scale = torch.empty( + (0, self.hidden_size // 32), + dtype=torch.float8_e8m0fnu, + device=hidden_states.device, + ) + ( packed_recv_hidden, recv_topk_weights, @@ -571,7 +626,8 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase): def __init__(self, **kwargs): super().__init__(**kwargs) self.quant_config = {} - self.quant_func = get_hip_quant(QuantType.per_1x128) + self.fp8_quant_func = get_hip_quant(QuantType.per_1x128) + self.fp4_quant_func = get_hip_quant(QuantType.per_1x32) def dispatch_a( self, @@ -589,13 +645,13 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase): output_dtype = hidden_states.dtype scale = None - fp8_dispatch = get_bool_env_var("SGLANG_MORI_FP8_DISP", "False") + fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags() if fp8_dispatch: # FP8 quant if num_tokens > 0: # NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD. - hidden_states, scale = self.quant_func( + hidden_states, scale = self.fp8_quant_func( hidden_states, quant_dtype=fp8_dtype ) else: @@ -608,6 +664,22 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase): device=hidden_states.device, ) + elif fp4_dispatch: + # FP4 quant + if num_tokens > 0: + hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False) + else: + hidden_states = torch.empty( + (0, self.hidden_size // 2), + dtype=torch.float4_e2m1fn_x2, + device=hidden_states.device, + ) + scale = torch.empty( + (0, self.hidden_size // 32), + dtype=torch.float8_e8m0fnu, + device=hidden_states.device, + ) + topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids (