[AMD] Add opt-in AITER ASM FP4 GEMM for Quark dense linear

[vorapolsiloai]

Motivation

On ROCm, Quark W4A4 MXFP4 dense linear layers always run AITER's Triton gemm_afp4wfp4. That kernel does not read the tuned a4w4_blockscale_tuned_gemm.csv tables, which AITER's gemm_a4w4 and vLLM both use. So SGLang cannot use the per-model CK/ASM tuning that vLLM gets (for example ROCm/aiter#5350), and we would have to keep a separate set of Triton tunings for SGLang.

This PR adds an opt-in path that runs these layers through gemm_a4w4, so SGLang uses the same tuned kernels.

Modifications

• Add SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM to Envs (default off). It only takes effect on gfx95.
• After weight loading, preshuffle the FP4 weights and the E8M0 weight scales once into the layout that gemm_a4w4 expects.
• At run time, quantize activations with AITER's per_1x32_f4_quant (shuffled scales) and call gemm_a4w4 with bpreshuffle=True. AITER then picks the model-specific CSV entry if there is one.
• Both AITER calls are registered as custom ops with fake implementations, so torch.compile can trace them.
• If a layer's scale shape cannot be swizzled (rows must be a multiple of 32, scale columns a multiple of 8), that layer keeps the existing Triton path and a warning is logged once.

Limitation: Quark's tuple-input fusion projections are not supported on this path. The weights are preshuffled at load time, so we cannot fall back to Triton per call. These projections raise an error that tells the user to turn the flag off.

Open question for maintainers

Should this stay behind its own env var, or should gfx95 with SGLANG_USE_AITER=1 use CK/ASM by default (as vLLM does) and keep Triton as the fallback? Details: #39598 (comment)

Accuracy Tests

MI355X, synthetic GEMM (M, N, K) = (16, 1024, 5120):

• AITER loaded the Qwen3.8-27B tuned A4W4 CSV entry and selected f4gemm_bf16_per1x32Fp4_BpreShuffle_32x128.
• process_weights_after_loading + apply_weights gave exactly the same output as a separately dequantized FP4 reference (rtol=0, atol=0).
• CPU unit tests (test/registered/unit/layers/quantization/test_quark_mxfp4_asm.py) cover which scale shapes are supported and check the scale permutation against AITER's tile order (all 3 tests pass cleanly).
• Re-validated post-rebase on AMD Instinct MI355X (ROCm 10, AITER): env var registration in Envs and QuarkW4A4MXFP4 module imports pass cleanly.

Speed Tests and Profiling

TODO: add a serving comparison on one image, flag off (Triton) vs flag on (CK/ASM).

Checklist

• Format code with pre-commit.
• Add unit tests.
• Update documentation.
• Provide accuracy and speed benchmark results.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #37922164076
Latest PR Test (Extra): ❌ Run #37922163642
Latest PR Test (AMD ROCm 10): ❌ Run #37922164006

[vorapolsiloai]

Open question for maintainers: default CK vs extra env

This PR currently keeps SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM opt-in, default off. That is intentional for the first landing, but we would like maintainers to decide the longer-term policy.

What this PR does today

• SGLANG_USE_AITER=1 still uses Triton gemm_afp4wfp4 for Quark dense (existing behavior).
• SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM=1 on gfx95 switches eligible dense layers to AITER CK/ASM gemm_a4w4 (shared A4W4 CSV path used by AITER/vLLM).

Why we did not copy vLLM blindly
vLLM does not add a second env. On gfx950 + AITER it prefers CK, and only falls back to Triton gemm_afp4wfp4_preshuffle for M <= 64 when (N,K) is on a hardcoded vLLM Triton-tuned list. That list is not a gfx950 policy and does not cover SGLang Quark shapes (e.g. Qwen3.8-27B 5120 / 14336 / 34816) or other SGLang Triton-tuned models (Kimi, MiniMax, …). CK and Triton also consume different CSVs and different weight layouts, so defaulting CK without a real fallback would ignore existing Triton tune tables.

Ask for maintainers
Please choose one of:

1. Keep this opt-in env for now (current PR). Lowest risk for Triton-CSV workloads; extra flag to track in serving recipes.
2. No extra env: SGLANG_USE_AITER=1 + gfx95 implies CK for Quark dense. Matches vLLM operationally, but needs an explicit Triton fallback story or it will regress Triton-only tables.
3. Follow-up hybrid: still no extra env, but pick CK vs Triton from which AITER CSV actually has the shape (and add the preshuffle Triton kernel so shuffled weights can take both paths). Closest to vLLM’s intent without copying their (N,K) list.

We can keep (1) in this PR and do (2)/(3) as a follow-up if that is preferred.

[BowenBao]

Thanks @vorapolsiloai, looks good overall.

Could you share some end to end performance and accuracy data comparison for your model use case?

[vorapolsiloai]

Thanks @BowenBao!

Here is the end-to-end performance and accuracy comparison for Qwen3.8-27B (Quark-AWQ-MXFP4) on AMD Instinct MI355X (TP2), demonstrating the uplift when enabling this PR's opt-in ASM FP4 GEMM path (SGLANG_ROCM_USE_AITER_FP4_ASM_GEMM=1).

1. Setup & Environment

• Hardware: AMD Instinct MI355X (TP2)
• Base Environment: SGLang ROCm nightly (v0.5.21-rocm724-mi35x)
• Companion PRs / Patches used in this test stack:
  • SGLang:
    • [AMD] Add opt-in AITER ASM FP4 GEMM for Quark dense linear #39598 (this PR: opt-in AITER ASM FP4 GEMM)
    • [AMD] Fix AITER unified decode page table for page size one #40410 (Fix AITER unified decode page table for page size 1)
    • [AMD] Quantize AITER unified prefill queries to FP8 #39998 (Quantize AITER unified prefill queries to FP8)
    • [AMD] Fuse Qwen3.8 GDN prefill prologue and MXFP4 output epilogue #39912 (Fuse Qwen3.8 GDN prefill prologue and MXFP4 output epilogue)
  • AITER:
    • [Triton/Gluon] [GFX950] Add Unified Attention Gluon Kernel ROCm/aiter#4614 (Gluon Unified Attention, merged)
    • FlyDSL A4W4 decode GEMM for M<=16 on gfx950 ROCm/aiter#6218 (FlyDSL decode GEMM backport for small-M)
• Model: amd/Qwen3.8-27B-Quark-AWQ-MXFP4
• Serving Flags: --tp-size 2 --quantization quark --attention-backend aiter --page-size 1 --kv-cache-dtype fp8_e4m3 --chunked-prefill-size 16384

2. End-to-End Performance Comparison

Workload	Baseline Thr (tok/s)	ASM FP4 Thr (tok/s)	Throughput %	TTFT %	TPOT %
text_isl8192_osl1024_c4	493.8	578.1	+17.08%	-10.49%	-15.36%
text_isl8192_osl1024_c32	1969.0	2034.2	+3.31%	-1.48%	-3.29%
text_isl128000_osl1024_c4	169.5	177.3	+4.58%	-2.49%	-5.04%

• Geomean: +8.15% throughput uplift across the text workloads, with -4.91% faster TTFT and -8.05% lower inter-token latency (TPOT).

3. Accuracy Verification

• GSM8K (200 questions, parallel 64):
  • Baseline: 0.955 accuracy (0 invalid)
  • Candidate (with ASM FP4): 0.955 accuracy (0 invalid)
  • Exact parity with zero regressions.