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

[AMD] [GLM-5.3-Flash Day 0] Enable FP8 and Quark MXFP4 MoE on gfx950 (#38546)

Browse Source 
Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: Kevin Mi <mikevin920@yahoo.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Raiden Makoto
2026-09-21 21:07:37 -07:00
committed by GitHub
co-authored by Raiden-Makoto Thomas Wang andyluo7 Kevin Mi Claude Opus 5
parent 15ba54bd5d
commit b44e248682
8 changed files with 756 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/aiter.py
+7 -4
View File
@@ -282,10 +282,13 @@ class AiterRunnerCore(MoeRunnerCore):
# `SGLANG_USE_AITER_MOE_GU_ITLV=0` to switch to SEPARATED, which
# matches the layout produced by `Mxfp4MoEMethod` (gpt-oss
# MXFP4) and the gptoss_fp4 tuned FlyDSL kernels.
extra["gate_mode"] = (
GateMode.INTERLEAVE.value
if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
else GateMode.SEPARATED.value
extra.setdefault(
"gate_mode",
(
GateMode.INTERLEAVE.value
if envs.SGLANG_USE_AITER_MOE_GU_ITLV.get()
else GateMode.SEPARATED.value
),
)
extra["swiglu_limit"] = quant_info.swiglu_limit
if self.config.no_combine:
python/sglang/srt/layers/quantization/fp8.py
+22
View File
@@ -1886,6 +1886,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
)
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
layer._aiter_gate_up_interleaved = False
return
elif self.use_mxfp8 and get_moe_a2a_backend().is_flashinfer_megamoe():
from sglang.srt.layers.moe.flashinfer_megamoe import (
@@ -1933,6 +1934,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
)
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
layer._aiter_gate_up_interleaved = False
elif _use_aiter:
# Pre-shuffle weights
t = shuffle_weight(layer.w13_weight, (16, 16))
@@ -1943,6 +1945,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
del t
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
layer._aiter_gate_up_interleaved = False
elif _is_cpu:
assert _is_cpu_amx_available, (
"Fp8MoEMethod on CPU requires that CPU has AMX support"
@@ -2650,6 +2653,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
requires_grad=False,
)
torch.cuda.empty_cache()
layer._aiter_gate_up_interleaved = False
# ROCm (_use_aiter): using column-wise scaling
layer.w13_weight_scale1 *= layer.w13_weight_scale.unsqueeze(-1)
@@ -3158,6 +3162,23 @@ class Fp8MoEMethod(FusedMoEMethodBase):
quant_type = AiterQuantType.PER_TOKEN
w13_scale = layer.w13_weight_scale1
w2_scale = layer.w2_weight_scale1
fused_moe_kwargs = None
gate_up_interleaved = getattr(layer, "_aiter_gate_up_interleaved", None)
if (
gate_up_interleaved is not None
and (self.moe_runner_config.swiglu_limit or 0.0) > 0
):
from aiter.ops.flydsl.moe_common import GateMode
fused_moe_kwargs = {
"gate_mode": (
GateMode.INTERLEAVE.value
if gate_up_interleaved
else GateMode.SEPARATED.value
)
}
return AiterMoeQuantInfo(
w13_weight=w13_weight,
w2_weight=w2_weight,
@@ -3168,6 +3189,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
swiglu_limit=self.moe_runner_config.swiglu_limit or 0.0,
hidden_pad=getattr(layer, "hidden_pad", 0),
intermediate_pad=getattr(layer, "intermediate_pad", 0),
fused_moe_kwargs=fused_moe_kwargs,
)
python/sglang/srt/layers/quantization/quark/quark.py
+56 -3
View File
@@ -15,7 +15,11 @@ from sglang.srt.layers.quantization.base_config import ( # noqa: E501
QuantizationConfig,
QuantizeMethodBase,
)
from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8LinearMethod
from sglang.srt.layers.quantization.fp8 import (
Fp8Config,
Fp8LinearMethod,
Fp8MoEMethod,
)
from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
from sglang.srt.layers.quantization.quark.schemes import (
QuarkLinearScheme,
@@ -375,6 +379,46 @@ class QuarkConfig(QuantizationConfig):
expanded.append(name.removeprefix("language_model."))
self.exclude_layers = list(dict.fromkeys(expanded))
layer_quant_config = self.quant_config.get("layer_quant_config")
if layer_quant_config:
self.quant_config["layer_quant_config"] = hf_to_sglang_mapper.apply_dict(
layer_quant_config
)
if self.kv_cache_group:
self.kv_cache_group = hf_to_sglang_mapper.apply_list(self.kv_cache_group)
@staticmethod
def _get_block_fp8_config(
layer_quant_config: Optional[dict[str, Any]],
packed_modules_mapping: dict[str, list[str]],
) -> Optional[Fp8Config]:
if layer_quant_config is None:
return None
weight_config = layer_quant_config.get("weight") or {}
input_config = layer_quant_config.get("input_tensors") or {}
block_size = weight_config.get("block_size")
if not (
not layer_quant_config.get("output_tensors")
and not layer_quant_config.get("bias")
and weight_config.get("dtype") in {"fp8_e4m3", "fp8_e4m3fn"}
and weight_config.get("qscheme") == "per_block"
and weight_config.get("is_dynamic") is False
and isinstance(block_size, list)
and len(block_size) == 2
and input_config.get("dtype") in {"fp8_e4m3", "fp8_e4m3fn"}
and input_config.get("is_dynamic") is True
):
return None
return Fp8Config(
is_checkpoint_fp8_serialized=True,
activation_scheme="dynamic",
weight_block_size=block_size,
packed_modules_mapping=packed_modules_mapping,
)
def get_quant_method(
self, layer: torch.nn.Module, prefix: str
) -> Optional["QuantizeMethodBase"]:
@@ -396,6 +440,17 @@ class QuarkConfig(QuantizationConfig):
return QuarkKVCacheMethod(self)
return None
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
block_fp8_config = self._get_block_fp8_config(
self._find_matched_config(prefix, layer), self.packed_modules_mapping
)
if block_fp8_config is not None:
if isinstance(layer, LinearBase):
return Fp8LinearMethod(block_fp8_config)
if isinstance(layer, FusedMoE):
return Fp8MoEMethod(block_fp8_config)
if isinstance(layer, LinearBase):
scheme = self.get_linear_scheme(layer=layer, layer_name=prefix)
layer.scheme = scheme
@@ -406,8 +461,6 @@ class QuarkConfig(QuantizationConfig):
self._online_quantized_layers.add(prefix)
return QuarkKVCacheMethod(self)
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
if isinstance(layer, FusedMoE):
self._online_quantized_layers.add(prefix)
layer.scheme = self.get_moe_scheme(layer, prefix)
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py
+24 -4
View File
@@ -38,12 +38,12 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_is_shuffle_moe_mxfp4 = is_gfx95_supported()
__all__ = ["QuarkW4A4MXFp4MoE"]
_is_hip = is_hip()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_is_gfx95 = is_gfx95_supported()
_is_shuffle_moe_mxfp4 = _use_aiter and _is_gfx95
if _use_aiter:
from aiter.ops.shuffle import moe_shuffle_scale, moe_shuffle_weight, shuffle_weight
from aiter.utility.fp4_utils import e8m0_shuffle
@@ -202,6 +202,8 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
is_concat=True,
is_packed=True,
)
layer.hidden_pad = 0
layer.intermediate_pad = w13_up_dim // 2 - intermediate_size_per_partition
# Add the quantization method used (per tensor/grouped/channel)
# to ensure the weight scales are loaded in properly
@@ -819,6 +821,11 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
layer.w2_weight = torch.nn.Parameter(qw2_weight, requires_grad=False)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
if not getattr(self, "_owns_moe_runner", False):
raise RuntimeError(
"Quark MXFP4 weight preshuffling requires an owned AITER runner."
)
if (
not self.is_checkpoint_mxfp4_serialized
or self.dequantization_config is not None
@@ -889,15 +896,19 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
)
self.moe_runner_config = moe_runner_config
self._owns_moe_runner = False
moe_runner_backend = get_moe_runner_backend()
if moe_runner_backend.is_auto() and get_moe_a2a_backend().supports_aiter():
moe_runner_backend = MoeRunnerBackend.AITER
if moe_runner_backend.is_aiter():
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
self._owns_moe_runner = True
else:
# TODO(cwan): refactor other backends
pass
raise NotImplementedError(
"Quark MXFP4 MoE currently requires the AITER runner; "
f"got {moe_runner_backend.value!r}."
)
def apply_weights(
self,
@@ -924,6 +935,12 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
from aiter.ops.flydsl.moe_common import GateMode
_fused_moe_kwargs = {"gate_mode": GateMode.INTERLEAVE.value}
elif _is_gfx95:
from aiter.ops.flydsl.moe_common import GateMode
# Quark checkpoints store gate and up projections as separate
# contiguous row ranges. Keep that ordering for correctness.
_fused_moe_kwargs = {"gate_mode": GateMode.SEPARATED.value}
else:
_fused_moe_kwargs = None
@@ -934,6 +951,9 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
w13_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
expert_mask=layer.dispatcher.expert_mask_gpu,
hidden_pad=getattr(layer, "hidden_pad", 0),
intermediate_pad=getattr(layer, "intermediate_pad", 0),
swiglu_limit=self.moe_runner_config.swiglu_limit or 0.0,
fused_moe_kwargs=_fused_moe_kwargs,
)
return self.runner.run(dispatch_output, quant_info)
python/sglang/srt/models/glm5_next.py
+17 -3
View File
@@ -1233,11 +1233,13 @@ class Glm5NextModel(nn.Module):
class Glm5NextForConditionalGeneration(nn.Module):
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_substr={
"model.language_model.": "model.",
"model.visual": "visual",
}
},
orig_to_new_prefix={
"model.language_model.": "model.",
},
orig_to_new_suffix={".attn.qkv": ".attn.qkv_proj"},
)
packed_modules_mapping = {
"fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"],
**Glm5NextLinearAttention._PACKED_MODULES_MAPPING,
@@ -1570,6 +1572,14 @@ class Glm5NextForConditionalGeneration(nn.Module):
fused_cat_dim = 0
params_dict = dict(self.named_parameters())
def maybe_map_fp8_block_scale_name(name: str) -> str:
if name.endswith("weight_scale"):
candidate = name.removesuffix("weight_scale") + "weight_scale_inv"
if candidate in params_dict:
return candidate
return name
weight_names = []
for name, loaded_weight in weights:
is_visual_weight = "visual" in name
@@ -1633,6 +1643,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
if "mlp.experts" in name:
continue
candidate = name.replace(weight_name, param_name)
candidate = maybe_map_fp8_block_scale_name(candidate)
if (
param_name
in {
@@ -1662,6 +1673,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
continue
is_expert_weight = True
name = name.replace(weight_name, param_name)
name = maybe_map_fp8_block_scale_name(name)
if name not in params_dict:
continue
param = params_dict[name]
@@ -1713,6 +1725,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
"fused_qkv_a_proj_with_mqa",
)
)
target = maybe_map_fp8_block_scale_name(target)
if target in params_dict:
param = params_dict[target]
weight_loader = getattr(
@@ -1723,6 +1736,7 @@ class Glm5NextForConditionalGeneration(nn.Module):
cached_a_proj.pop(kv_a_proj_name, None)
continue
name = maybe_map_fp8_block_scale_name(name)
if name not in params_dict:
continue