From fa040aec4c6a04704519d3923a20c20a07886dad Mon Sep 17 00:00:00 2001 From: Xiake Sun Date: Thu, 3 Sep 2026 10:41:07 +0000 Subject: [PATCH] feat (qwen3.8): support qwen3.8 quark mxfp4 model on mi308x - Add a TP8/EP1 AITER A16W4 MOE execution path for serialized Qwen3.8 Quark MXFP4 checkpoints on gfx942. - Preserve checkpoint-native packed weights and E8M0 scales, use zero-copy column-major views, and run the two routed GEMMs with the required gate/up activation. - Add extra-buffer Mamba radix caching support - Add focused layout, dispatch, and dequantized-reference coverage Signed-off-by: Xiake Sun --- .../launch_qwen3.8-2.4T-quark-mxfp4_tp8.sh | 48 ++++ python/sglang/srt/arg_groups/overrides.py | 4 + .../quark/schemes/quark_w4a4_mxfp4_moe.py | 106 ++++++++- .../test_quark_mxfp4_a16w4_gfx942.py | 212 ++++++++++++++++++ test/registered/unit/test_model_overrides.py | 15 ++ 5 files changed, 381 insertions(+), 4 deletions(-) create mode 100755 evaluation/launch_qwen3.8-2.4T-quark-mxfp4_tp8.sh create mode 100644 test/registered/unit/layers/quantization/test_quark_mxfp4_a16w4_gfx942.py diff --git a/evaluation/launch_qwen3.8-2.4T-quark-mxfp4_tp8.sh b/evaluation/launch_qwen3.8-2.4T-quark-mxfp4_tp8.sh new file mode 100755 index 000000000000..552b07eabce7 --- /dev/null +++ b/evaluation/launch_qwen3.8-2.4T-quark-mxfp4_tp8.sh @@ -0,0 +1,48 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL_PATH=${MODEL_PATH:-/models/Qwen3.8-2.4T-A95B-Quark-MXFP4/} +SERVED_MODEL_NAME=${SERVED_MODEL_NAME:-Qwen3.8-2.4T-A95B-Quark-MXFP4} +HOST=${HOST:-0.0.0.0} +PORT=${PORT:-30000} + +export SGLANG_USE_AITER=1 +export USE_AITER_COMM=1 +export HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7} +export PYTHONPATH=/opt/sglang/python${PYTHONPATH:+:${PYTHONPATH}} + +if [[ -n "${SGLANG_TORCH_PROFILER_DIR:-}" ]]; then + mkdir -p "${SGLANG_TORCH_PROFILER_DIR}" +fi + +cd /opt/sglang +exec python3 -m sglang.launch_server \ + --model-path "${MODEL_PATH}" \ + --served-model-name "${SERVED_MODEL_NAME}" \ + --tp-size 8 \ + --moe-runner-backend aiter \ + --attention-backend aiter \ + --linear-attn-backend aiter \ + --linear-attn-decode-backend aiter \ + --linear-attn-prefill-backend aiter \ + --page-size 64 \ + --kv-cache-dtype fp8_e4m3 \ + --chunked-prefill-size 8192 \ + --max-prefill-tokens 8192 \ + --max-total-tokens 273536 \ + --max-running-requests 9 \ + --max-mamba-cache-size 9 \ + --mamba-ssm-dtype bfloat16 \ + --mem-fraction-static 1.05 \ + --disable-radix-cache \ + --disable-custom-all-reduce \ + --cuda-graph-max-bs-decode 8 \ + --cuda-graph-bs-decode 1 2 4 8 \ + --cuda-graph-backend-prefill disabled \ + --watchdog-timeout 1200 \ + --reasoning-parser qwen3 \ + --tool-call-parser qwen3_coder \ + --trust-remote-code \ + --host "${HOST}" \ + --port "${PORT}" \ + "$@" diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index dc5443bfa617..561b6c162ab1 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -962,6 +962,8 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: "KimiLinearForCausalLM", "BailingMoeV2_5ForCausalLM", "Qwen3NextForCausalLM", + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForCausalLM", "Qwen3_5MoeForConditionalGeneration", "InternS2PreviewForConditionalGeneration", "Qwen3_5ForConditionalGeneration", @@ -981,6 +983,8 @@ def _step3p_overrides(server_args: Any, hf_config: Any) -> dict: # delegates here. _MAMBA_EXTRA_BUFFER_ARCHS = frozenset( { + "Qwen3_5ForCausalLM", + "Qwen3_5MoeForCausalLM", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", "Qwen3NextForCausalLM", diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index c6890dfd5014..0ad04075fc26 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -12,6 +12,7 @@ from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme from sglang.srt.utils import ( get_bool_env_var, + is_gfx942_supported, is_gfx95_supported, is_hip, set_weight_attrs, @@ -27,16 +28,15 @@ logger = logging.getLogger(__name__) _is_shuffle_moe_mxfp4 = is_gfx95_supported() +_is_gfx942 = is_gfx942_supported() __all__ = ["QuarkW4A4MXFp4MoE"] _is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip -if _use_aiter: +if _is_hip: from aiter.ops.shuffle import shuffle_weight from aiter.utility.fp4_utils import e8m0_shuffle - -if _is_hip: from aiter.ops.triton.quant import dynamic_mxfp4_quant else: dynamic_mxfp4_quant = None @@ -45,7 +45,6 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): - def __init__( self, weight_config: dict[str, Any], @@ -67,6 +66,7 @@ def __init__( self.static_input_scales = not self.input_quant.get("is_dynamic") self.with_bias = False + self.use_aiter_a16w4 = False if not self.is_checkpoint_mxfp4_serialized: if not mxfp_supported(): @@ -93,6 +93,19 @@ def create_weights( params_dtype: torch.dtype, **extra_weight_attrs, ): + self.use_aiter_a16w4 = ( + _use_aiter + and _is_gfx942 + and self.is_checkpoint_mxfp4_serialized + and num_experts == 512 + and hidden_size == 8192 + and intermediate_size_per_partition == 256 + ) + if self.use_aiter_a16w4: + logger.info_once( + "Using the gfx942 A16W4 MoE path for Qwen3.8 TP8 " + "(512 experts, top-k 10, hidden 8192, local intermediate 256)." + ) from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported @@ -218,6 +231,17 @@ def online_mxfp4_moe_weight_loader( return online_mxfp4_moe_weight_loader def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if self.use_aiter_a16w4: + # The gfx942 A16W4 kernel consumes the checkpoint's original E8M0 + # scale order and column-major transposed views. Do not apply the + # CDNA4-oriented scale/weight shuffles used by the existing A4W4 + # path. + if hasattr(layer, "dispatcher"): + layer.dispatcher.set_quant_config( + {"weight_dtype": torch.float4_e2m1fn_x2} + ) + return + # Pre-shuffle weight scales s0, s1, _ = layer.w13_weight_scale.shape w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) @@ -268,6 +292,9 @@ def apply_weights( layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: + if self.use_aiter_a16w4: + return self._apply_aiter_a16w4(layer, dispatch_output) + from sglang.srt.layers.moe.moe_runner.aiter import ( AiterMoeQuantInfo, AiterQuantType, @@ -293,3 +320,74 @@ def apply_weights( expert_mask=layer.dispatcher.expert_mask_gpu, ) return self.runner.run(dispatch_output, quant_info) + + def _apply_aiter_a16w4( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + from aiter import silu_and_mul + from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_a16w4 + from aiter.ops.triton.moe.moe_routing.routing import routing + + from sglang.srt.layers.moe.token_dispatcher.standard import ( + StandardCombineInput, + ) + + hidden_states = dispatch_output.hidden_states + if hidden_states.shape[0] == 0: + return StandardCombineInput(hidden_states=hidden_states) + + if layer.moe_ep_size != 1 or layer.moe_tp_size != 8: + raise NotImplementedError( + "The gfx942 Qwen3.8 A16W4 path currently supports TP8/EP1 only. " + f"Got TP{layer.moe_tp_size}/EP{layer.moe_ep_size}." + ) + if self.moe_runner_config.top_k != 10: + raise ValueError( + "The gfx942 Qwen3.8 A16W4 path requires top-k 10, got " + f"{self.moe_runner_config.top_k}." + ) + + router_logits = dispatch_output.topk_output.router_logits + routing_data, gather_idx, scatter_idx = routing( + router_logits, + self.moe_runner_config.top_k, + sm_first=False, + ) + + # Checkpoint layout is [E, N, K/2] / [E, N, K/32]. AITER A16W4 + # consumes logical column-major [E, K/2, N] / [E, K/32, N] views. + w13_weight = layer.w13_weight.view(torch.uint8).transpose(1, 2) + w13_scale = layer.w13_weight_scale.view(torch.uint8).transpose(1, 2) + w2_weight = layer.w2_weight.view(torch.uint8).transpose(1, 2) + w2_scale = layer.w2_weight_scale.view(torch.uint8).transpose(1, 2) + + stage1 = moe_gemm_a16w4( + hidden_states, + w13_weight, + None, + w13_scale, + routing_data=routing_data, + gather_indx=gather_idx, + out_dtype=hidden_states.dtype, + ) + + intermediate = torch.empty( + (stage1.shape[0], stage1.shape[1] // 2), + dtype=stage1.dtype, + device=stage1.device, + ) + silu_and_mul(intermediate, stage1) + + output = moe_gemm_a16w4( + intermediate, + w2_weight, + None, + w2_scale, + routing_data=routing_data, + scatter_indx=scatter_idx, + gammas=routing_data.gate_scal, + out_dtype=hidden_states.dtype, + ) + return StandardCombineInput(hidden_states=output) diff --git a/test/registered/unit/layers/quantization/test_quark_mxfp4_a16w4_gfx942.py b/test/registered/unit/layers/quantization/test_quark_mxfp4_a16w4_gfx942.py new file mode 100644 index 000000000000..4cf5d37fdd95 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_quark_mxfp4_a16w4_gfx942.py @@ -0,0 +1,212 @@ +"""Focused integration coverage for Quark MXFP4 MoE on gfx942.""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from aiter.ops.triton.moe.moe_op_gemm_a16w4 import moe_gemm_torch +from aiter.ops.triton.moe.moe_routing.routing import routing +from aiter.ops.triton.moe.quant_moe import downcast_to_mxfp, upcast_from_mxfp +from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput +from sglang.srt.layers.moe.topk import StandardTopKOutput +from sglang.srt.layers.quantization.quark.schemes.quark_w4a4_mxfp4_moe import ( + QuarkW4A4MXFp4MoE, +) +from sglang.srt.utils.common import is_gfx942_supported + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or not is_gfx942_supported(), + reason="MI308X/gfx942-specific test", +) + + +class _Dispatcher: + expert_mask_gpu = None + + def __init__(self): + self.quant_config = None + + def set_quant_config(self, quant_config): + self.quant_config = quant_config + + +def _scheme() -> QuarkW4A4MXFp4MoE: + scheme = QuarkW4A4MXFp4MoE( + weight_config={"qscheme": "per_group"}, + input_config={"qscheme": "per_group", "is_dynamic": True}, + ) + scheme.use_aiter_a16w4 = True + scheme.moe_runner_config = SimpleNamespace(top_k=10) + return scheme + + +def _assert_close(ref: torch.Tensor, actual: torch.Tensor) -> None: + ref = ref.float() + actual = actual.float() + diff = (ref - actual).abs() + scale = ref.abs().amax().clamp_min(1.0e-30) + ref_normalized = ref / scale + actual_normalized = actual / scale + ref_rms = torch.sqrt(torch.square(ref_normalized).mean()).clamp_min(1.0e-30) + rel = (ref_normalized - actual_normalized).abs() / torch.maximum( + ref_normalized.abs(), ref_rms + ) + + assert torch.isfinite(actual).all() + assert diff.max().item() <= 5.0e-2 + assert rel.max().item() <= 4.0e-1 + assert torch.sqrt(torch.square(rel).mean()).item() <= 4.0e-2 + + +def test_create_weights_selects_qwen3_8_tp8_a16w4() -> None: + scheme = QuarkW4A4MXFp4MoE( + weight_config={"qscheme": "per_group"}, + input_config={"qscheme": "per_group", "is_dynamic": True}, + ) + layer = torch.nn.Module() + + with torch.device("meta"): + scheme.create_weights( + layer=layer, + num_experts=512, + hidden_size=8192, + intermediate_size_per_partition=256, + params_dtype=torch.bfloat16, + weight_loader=lambda *args: None, + ) + + assert scheme.use_aiter_a16w4 + assert tuple(layer.w13_weight.shape) == (512, 512, 4096) + assert tuple(layer.w2_weight.shape) == (512, 8192, 128) + assert tuple(layer.w13_weight_scale.shape) == (512, 512, 256) + assert tuple(layer.w2_weight_scale.shape) == (512, 8192, 8) + + +def test_process_weights_preserves_checkpoint_layout() -> None: + scheme = _scheme() + layer = SimpleNamespace( + w13_weight=torch.nn.Parameter( + torch.arange(64, dtype=torch.uint8).view(2, 4, 8), requires_grad=False + ), + w2_weight=torch.nn.Parameter( + torch.arange(64, dtype=torch.uint8).view(2, 8, 4), requires_grad=False + ), + w13_weight_scale=torch.nn.Parameter( + torch.arange(16, dtype=torch.uint8).view(2, 4, 2), requires_grad=False + ), + w2_weight_scale=torch.nn.Parameter( + torch.arange(16, dtype=torch.uint8).view(2, 8, 1), requires_grad=False + ), + dispatcher=_Dispatcher(), + ) + before = [ + tensor.detach().clone() + for tensor in ( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + ) + ] + + scheme.process_weights_after_loading(layer) + + after = ( + layer.w13_weight, + layer.w2_weight, + layer.w13_weight_scale, + layer.w2_weight_scale, + ) + assert all(torch.equal(expected, actual) for expected, actual in zip(before, after)) + assert layer.dispatcher.quant_config == {"weight_dtype": torch.float4_e2m1fn_x2} + + +def test_quark_a16w4_apply_matches_dequantized_reference() -> None: + torch.manual_seed(3807) + m = 4 + hidden_size = 8192 + intermediate_size_per_rank = 256 + num_experts = 16 + topk = 10 + + # Build logical AITER weights, then transpose back to SGLang's checkpoint + # parameter layout so the integration's zero-copy views are exercised. + w13 = ( + torch.randn( + (num_experts, hidden_size, 2 * intermediate_size_per_rank), + device="cuda", + dtype=torch.bfloat16, + ) + * 0.02 + ) + w13_quant, w13_scale = downcast_to_mxfp(w13, torch.uint8, axis=1) + w2 = ( + torch.randn( + (num_experts, intermediate_size_per_rank, hidden_size), + device="cuda", + dtype=torch.bfloat16, + ) + * 0.02 + ) + w2_quant, w2_scale = downcast_to_mxfp(w2, torch.uint8, axis=1) + + layer = SimpleNamespace( + w13_weight=torch.nn.Parameter(w13_quant.transpose(1, 2), requires_grad=False), + w13_weight_scale=torch.nn.Parameter( + w13_scale.transpose(1, 2), requires_grad=False + ), + w2_weight=torch.nn.Parameter(w2_quant.transpose(1, 2), requires_grad=False), + w2_weight_scale=torch.nn.Parameter( + w2_scale.transpose(1, 2), requires_grad=False + ), + moe_ep_size=1, + moe_tp_size=8, + ) + hidden_states = ( + torch.randn((m, hidden_size), device="cuda", dtype=torch.bfloat16) * 0.02 + ) + router_logits = torch.randn((m, num_experts), device="cuda", dtype=torch.float16) + dispatch_output = StandardDispatchOutput( + hidden_states=hidden_states, + hidden_states_scale=None, + topk_output=StandardTopKOutput( + topk_weights=torch.empty((m, topk), device="cuda"), + topk_ids=torch.empty((m, topk), dtype=torch.int32, device="cuda"), + router_logits=router_logits, + ), + ) + + actual = _scheme()._apply_aiter_a16w4(layer, dispatch_output).hidden_states + + routing_data, gather_idx, scatter_idx = routing(router_logits, topk) + w13_dequant = upcast_from_mxfp(w13_quant, w13_scale, torch.bfloat16, axis=1) + stage1_ref = moe_gemm_torch( + hidden_states, + w13_dequant, + None, + routing_data, + gather_idx, + None, + None, + False, + ) + intermediate_ref = ( + F.silu(stage1_ref[:, :intermediate_size_per_rank].float()) + * stage1_ref[:, intermediate_size_per_rank:].float() + ).to(torch.bfloat16) + w2_dequant = upcast_from_mxfp(w2_quant, w2_scale, torch.bfloat16, axis=1) + expected = moe_gemm_torch( + intermediate_ref, + w2_dequant, + None, + routing_data, + None, + scatter_idx, + routing_data.gate_scal, + False, + ) + + _assert_close(expected, actual) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 5af3874ce278..048a3c281840 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1580,6 +1580,21 @@ def _view(arch, layer_types=None, **kw): "Qwen3_5ForConditionalGeneration", ) ) + for arch in ("Qwen3_5ForCausalLM", "Qwen3_5MoeForCausalLM"): + self.assertEqual( + _mamba_radix_cache_resolution( + _view(arch, linear_attn_backend="aiter", page_size=16) + ), + { + "uses_mamba_radix_cache": True, + "mamba_radix_cache_strategy": "extra_buffer", + }, + ) + self.assertTrue( + supports_mamba_cache_extra_buffer( + SimpleNamespace(linear_attn_backend="aiter"), arch + ) + ) self.assertFalse( supports_mamba_cache_extra_buffer( SimpleNamespace(linear_attn_backend="aiter"),