chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD][Feature] support fp4 dispatch and fp8 combine in moriep (#19757)

Browse Source 
Co-authored-by: Duyi-Wang <duyi.wang@amd.com>
This commit is contained in:
AMD-yanfeiwang
2026-03-10 03:52:05 +08:00
committed by GitHub
parent ffb4b6f4c1
commit f0153ad225
4 changed files with 265 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -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(

140
python/sglang/srt/layers/moe/rocm_moe_utils.py
View File

@@ -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

104
python/sglang/srt/layers/moe/token_dispatcher/moriep.py
View File

@@ -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
(