[AMD] [GLM5] use tuned MXFP4 kernels for MLA absorbed BMM

[Raiden-Makoto]

Motivation

The absorbed K/V projections remain the largest dense-attention bottleneck in GLM-5.2 prefill because the existing path does not use the checkpoint’s MXFP4 weights. This change preserves those weights in MXFP4 and runs the absorbed BMMs on the matching gfx950 kernels.

Modifications

• Preserve GLM MLA absorbed w_kc / w_vc weights and scales in packed MXFP4 behind SGLANG_USE_MXFP4_MLA_BMM=1.
• Dispatch A16WFP4 K/V BMMs with NUM_KSPLIT=1.
• Use BLOCK_SIZE_K=64 for GLM K-up (K=192).
• Use BLOCK_SIZE_M=128, BLOCK_SIZE_K=128 for V-up when M>256; preserve AITER's small-M decode buckets.
• Route the current MLA_ROCM K/V path through the tuned dispatch helpers.
• Write V-up output directly in batch-major layout.

Accuracy Tests

Test	Pre-tuning	This PR
GSM8K, 20 examples	0.950	0.950
Focused unit tests	—	14/14

Speed Benchmarks

GLM-5.2-MXFP4, TP4 MI355X, 8192 input / 1024 output, REPS=3.

A. Pre-tuning A16WFP4:

concurrency	TTFT (ms)	ITL (ms)	E2EL (ms)	output tok/s
4	996.05	11.825	13386.47	305.49
8	1610.50	13.583	16449.63	497.61
16	2871.90	15.943	21346.01	765.34
32	5405.12	19.817	30273.44	1081.76
64	10497.72	26.803	47538.85	1379.12

B. Tuned A16WFP4 (delta vs A):

concurrency	TTFT (ms)	delta	ITL (ms)	delta	E2EL (ms)	delta	output tok/s	delta
4	1001.84	+0.58%	11.820	-0.04%	13357.40	-0.22%	306.13	+0.21%
8	1578.35	-2.00%	13.577	-0.04%	16397.56	-0.32%	499.28	+0.34%
16	2805.58	-2.31%	15.953	+0.06%	21287.72	-0.27%	769.49	+0.54%
32	5291.87	-2.10%	19.827	+0.05%	30107.75	-0.55%	1087.52	+0.53%
64	10344.29	-1.46%	26.823	+0.07%	47212.48	-0.69%	1388.12	+0.65%

Baseline = rocm/sgl-dev:v0.5.16-rocm720-mi35x-20260804, plus these 7 PRs
PR #30575: Triton sparse MLA prefill/decode
PR #30715: fused indexer query Hadamard + FP8 quant
PR #30808: gfx950 dense-MHA short-context fallback
PR #31323: fused shared-expert append
PR #31324: skip unnecessary dense-regime decode indexer work
PR #33602: FP8 dense GEMMs
PR #33620: pre-tuning A16WFP4 absorbed K/V BMM

This PR = Baseline + M128/K128 V-BMM tuning

Checklist

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

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users.
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ❌ Run #33441160421
Latest PR Test (Extra): ❌ Run #33441159899
Latest PR Test (AMD ROCm 7.2): ❌ Run #33441160098

[Raiden-Makoto]

Doesn't work on image rocm/sgl-dev:v0.5.17-rocm720-mi35x-20260810 because AiTER commit changed

[Raiden-Makoto]

Resolved: now works on new image as well.

[HaiShaw]

/tag-and-rerun-ci

[Raiden-Makoto]

MI35x, AMD stage-A, AMD nondeterministic, AMD large 1-GPU, AMD small 1-GPU shards 1–13, and AMD 2-GPU shard 0 passed.

Failures appear unrelated:

• AMD: Qwen MoE HIPBLAS failure during CUDA-graph capture, ending with server exit -9
• AMD: AITER custom-allreduce IPC metadata timeout during 2-GPU graph capture
• AMD disaggregation: decode endpoint unavailable; simulated-retract GSM8K received Internal Server Error Xeon: sgl_kernel.silu_and_mul_cpu missing from the CPU extension Remaining short failures were gate/finish cascades

Failed AMD root causes:

• stage-a-test-1-gpu-small-amd-rocm720: passed — does not touch PR code
• stage-b-test-1-gpu-small-amd-mi35x-rocm720: passed — exercises the PR’s target hardware
• stage-b-test-1-gpu-small-amd-rocm720 (linux-mi300-1gpu, shard 0): test_torch_compile_moe.py — HIPBLAS internal error during CUDA-graph capture; server exited -9 — does not touch PR code
• stage-b-test-2-gpu-large-amd-rocm720 (linux-mi300-2gpu, shard 1): test_bench_one_batch_2gpu.py — AITER IPC metadata store timed out during graph capture — does not touch PR code
• stage-b-test-large-8-gpu-mi35x-disaggregation-amd-rocm720: test_disaggregation_basic.py — decode endpoint refused connections during simulated retract; GSM8K returned Internal Server Error — does not touch PR code
• stage-c AMD: skipped because stage-b failed.

cc @HaiShaw

[Raiden-Makoto]

/rerun-failed-ci