minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[AMD][WA] force to use gate_mode interleaved to fix tp2/tp4/tp8 acc issue (#27201)

Browse Source 
Co-authored-by: wunhuang <wunhuang@amd.com>
This commit is contained in:
kk
2026-06-05 20:18:11 -07:00
committed by GitHub
co-authored by wunhuang
parent bf66b7b6da
commit aa55657e9e
4 changed files with 73 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+10 -3
View File
@@ -902,10 +902,12 @@ class AiterAttnBackend(AttentionBackend):
)
if self.use_sliding_window_kv_pool:
# AITER attention kernels require int32 page indices;
# full_to_swa_index_mapping is stored as int64.
swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
kv_indices
)
).to(torch.int32)
)
kv_indices = self._transform_table_1_to_real(kv_indices)
@@ -1381,10 +1383,13 @@ class AiterAttnBackend(AttentionBackend):
)
if self.use_sliding_window_kv_pool:
# AITER attention kernels (e.g. mha_batch_prefill_func)
# require int32 page indices; full_to_swa_index_mapping is
# stored as int64.
swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.indices_updater_prefill.kv_indices
)
).to(torch.int32)
)
self.forward_metadata = ForwardMetadata(
@@ -1580,10 +1585,12 @@ class AiterAttnBackend(AttentionBackend):
]
if self.use_sliding_window_kv_pool:
# AITER attention kernels require int32 page indices;
# full_to_swa_index_mapping is stored as int64.
swa_page_indices = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
page_indices
)
).to(torch.int32)
)
page_indices = self._transform_table_1_to_real(page_indices)
python/sglang/srt/layers/quantization/mxfp4.py
+54 -43
View File
@@ -28,11 +28,11 @@ from torch.nn.parameter import Parameter
# cutlass_fused_moe. Its C++ logger reads TLLM_LOG_LEVEL on first kernel launch;
# setdefault preserves any explicit user override.
os.environ.setdefault("TLLM_LOG_LEVEL", "INFO")
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.amx_utils import (
CPUQuantMethod,
_amx_process_weight_after_loading,
@@ -155,7 +155,9 @@ if _is_hip:
try:
from aiter.ops.shuffle import (
shuffle_scale,
shuffle_scale_a16w4,
shuffle_weight,
shuffle_weight_a16w4,
)
from aiter.ops.triton.quant import dynamic_mxfp4_quant
from aiter.utility.fp4_utils import e8m0_shuffle
@@ -773,7 +775,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
)
return
if _use_aiter:
# Bias must be fp32 for the AITER kernels.
if layer.w13_weight_bias is not None:
layer.w13_weight_bias.data = layer.w13_weight_bias.data.to(
torch.float32
@@ -781,12 +782,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
if layer.w2_weight_bias is not None:
layer.w2_weight_bias.data = layer.w2_weight_bias.data.to(torch.float32)
# HF GPT-OSS stores w13 as gate/up *interleaved* row pairs
# [(g0, u0), (g1, u1), ...]. The AITER MXFP4 fused MoE kernels
# (FlyDSL `gate_mode="separated"` and CK `preshuffle_on`) expect
# the *separated* layout [gate_0..gate_{N-1}, up_0..up_{N-1}].
# De-interleave weights, scales, and bias before the tile shuffle
# so the post-shuffle bytes land in the layout the kernel reads.
e, n, k = layer.w13_weight.shape
layer.w13_weight.view(torch.uint8).copy_(
layer.w13_weight.data.view(torch.uint8)
@@ -808,32 +803,46 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
.view(-1, n)
)
# ATOM-aligned MXFP4 preshuffle (is_guinterleave=False). Both
# `gptoss_fp4_tuned_fmoe.csv` flydsl entries (e.g.
# flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w2) and the CK
# `module_moe_ck2stages_*_preshuffle_on_*` fallback use the
# standard (16,16) tile layout produced here. The previous
# shuffle_weight_a16w4 path produced a gate/up-interleaved tile
# layout that does not match the separated-gate kernels and
# caused silent accuracy loss.
layer.w13_weight.data = shuffle_weight(
layer.w13_weight, is_guinterleave=False, gate_up=True
)
shuffled_w13_scale = shuffle_scale(
layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
experts_cnt=self.num_experts,
is_guinterleave=False,
gate_up=True,
)
layer.w2_weight.data = shuffle_weight(
layer.w2_weight, is_guinterleave=False, gate_up=False
)
shuffled_w2_scale = shuffle_scale(
layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
experts_cnt=self.num_experts,
is_guinterleave=False,
gate_up=False,
)
if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get():
layer.w13_weight.data = shuffle_weight_a16w4(layer.w13_weight, 16, True)
shuffled_w13_scale = shuffle_scale_a16w4(
layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
self.num_experts,
True,
)
layer.w2_weight.data = shuffle_weight_a16w4(layer.w2_weight, 16, False)
shuffled_w2_scale = shuffle_scale_a16w4(
layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
self.num_experts,
False,
)
else:
layer.w13_weight.data = shuffle_weight(
layer.w13_weight, is_guinterleave=False, gate_up=True
)
shuffled_w13_scale = shuffle_scale(
layer.w13_weight_scale.view(-1, layer.w13_weight_scale.shape[-1]),
experts_cnt=self.num_experts,
is_guinterleave=False,
gate_up=True,
)
layer.w2_weight.data = shuffle_weight(
layer.w2_weight, is_guinterleave=False, gate_up=False
)
shuffled_w2_scale = shuffle_scale(
layer.w2_weight_scale.view(-1, layer.w2_weight_scale.shape[-1]),
experts_cnt=self.num_experts,
is_guinterleave=False,
gate_up=False,
)
# shuffle_weight_a16w4(gate_up=True) above preshuffles w13 into aiter's
# preshuffle + gate/up-interleaved layout. Tag the Parameter so apply()
# can carry the metadata across .view(float4_e2m1fn_x2) and aiter's
# fused_moe selects the preshuffle_on kernel family.
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
layer.w13_weight_scale = torch.nn.Parameter(
shuffled_w13_scale, requires_grad=False
@@ -842,14 +851,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
shuffled_w2_scale, requires_grad=False
)
# Tell aiter.fused_moe these weights are already preshuffled so it
# picks the preshuffle_on CK / FlyDSL kernels (which match the
# actual layout) instead of falling back to preshuffle_off kernels
# that interpret the shuffled bytes as a non-shuffled tensor and
# produce garbage / OOB accesses.
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
return
if self.use_triton_kernels:
@@ -1281,7 +1282,17 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
doweight_stage1=self.moe_runner_config.apply_router_weight_on_input,
hidden_pad=self.hidden_pad,
intermediate_pad=self.intermediate_pad,
swiglu_limit=self.moe_runner_config.swiglu_limit or 0.0,
# Triggers aiter's INTERLEAVE gate_mode dispatch (required for our
# preshuffled gate/up-interleaved weight layout) and applies the
# model's swiglu clamp. Models populate the same scalar under
# different MoeRunnerConfig fields: gpt-oss uses `gemm1_clamp_limit`
# (renamed in `models/gpt_oss.py` from `config.swiglu_limit`); DSv4
# / FP8 uses `swiglu_limit` directly. Accept either.
swiglu_limit=(
self.moe_runner_config.gemm1_clamp_limit
or self.moe_runner_config.swiglu_limit
or 0.0
),
)
return self.runner.run(
dispatch_output._replace(hidden_states=x_padded), quant_info
python/sglang/srt/server_args.py
+7 -7
View File
@@ -2175,13 +2175,13 @@ class ServerArgs:
logger.warning(
"Detected ROCm and MXFP4 quantization format for GPT-OSS model, enabling aiter MXFP4 MOE kernel."
)
# The AITER MXFP4 fused-MoE path for GPT-OSS expects the
# SEPARATED gate/up tile layout (matches the
# `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
# Mxfp4MoEMethod weight shuffle). Other AITER MXFP4
# callers default to INTERLEAVE; opt this path out
# unless the user explicitly overrode it.
envs.SGLANG_USE_AITER_MOE_GU_ITLV.set(False)
## The AITER MXFP4 fused-MoE path for GPT-OSS expects the
## SEPARATED gate/up tile layout (matches the
## `gptoss_fp4_tuned_fmoe.csv` flydsl entries and the
## Mxfp4MoEMethod weight shuffle). Other AITER MXFP4
## callers default to INTERLEAVE; opt this path out
## unless the user explicitly overrode it.
# envs.SGLANG_USE_AITER_MOE_GU_ITLV.set(False)
elif is_hip() and envs.SGLANG_USE_AITER.get():
# For GPT-OSS bf16 on ROCm with aiter, use triton backend
# because aiter CK kernel doesn't support all GEMM dimensions
test/registered/amd/accuracy/mi35x/test_gpt_oss_eval_mi35x.py
+2 -2
View File
@@ -76,7 +76,7 @@ MI35X_GPT_OSS_MODELS = [
# to INTERLEAVE, so opt out explicitly here.
env_vars={
"SGLANG_USE_AITER": "1",
"SGLANG_USE_AITER_MOE_GU_ITLV": "0",
"SGLANG_USE_AITER_MOE_GU_ITLV": "1",
},
),
ModelConfig(
@@ -97,7 +97,7 @@ MI35X_GPT_OSS_MODELS = [
],
env_vars={
"SGLANG_USE_AITER": "1",
"SGLANG_USE_AITER_MOE_GU_ITLV": "0",
"SGLANG_USE_AITER_MOE_GU_ITLV": "1",
},
),
]