[AMD] support qlen>1 for aiter gluon path for Kimi K3

[billishyahao]

Motivation

This patch is to add qlen>1 support for aiter gluon path. #34647 only support qlen=1, for other cases, it will fall back to mla asm ps. If you specify the following configs,

SGLANG_AITER_HONOR_EXPLICIT_MEM_FRACTION=1 \
SGLANG_USE_AITER=1 \
SGLANG_AITER_K3_OPT=1 \
AITER_FLYDSL_FORCE=1 \
AITER_SITUV2_A8W4=1 \
SGLANG_AITER_MLA_GLUON=1 \
python3 -m sglang.launch_server --model-path /data-models/Kimi-K3 --trust-remote-code \
    --tp-size 8  \
	--prefill-attention-backend aiter \
	--decode-attention-backend aiter \
	--dtype bfloat16 \
	--kv-cache-dtype fp8_e4m3 \
    --cuda-graph-max-bs-decode 32 \
	--max-mamba-cache-size 160 \
	--mem-fraction-static 0.85 \
    --reasoning-parser kimi_k3 \
	--tool-call-parser kimi_k3 \
	--speculative-algorithm DSPARK \
        --speculative-draft-model-path /data-models/RadixArk-Kimi-K3-DSpark \
        --speculative-dspark-block-size 7 \
         --enable-linear-replayssm-spec \
       --speculative-attention-mode decode

then will get the server crash:

[AITER] /sgl-workspace/aiter/aiter_meta/csrc/py_itfs_cu/asm_mla.cu:193 get_heuristic_kernel_mla: cannot get heuristic kernel! q_type:bf16 kv_type:fp8 gqa:16 ps:0 prefill:0 causal:1 qseqlen:8 lse:0 cprr:0
Fatal Python error: Aborted

Modifications

layers/attention/aiter_mla_gluon.py — mla_gluon_decode() introduces qlen to decide whether q_nope and q_pe needs to be viewed

layers/attention/aiter_backend.py — remove max_q_len == 1 gate to allow qlen >1 path being activated for dspark

distributed/parallel_state.py — DSPARK's in-graph sampler use broadcast inside a CUDA graph, which is crashed by using torch api during cuda graph capture. Change it to pynccl.

test/registered/attention/test_aiter_gluon_h12_fp8.py — renamed from test_mla_gluon_h12_fp8.py to make the name more clear and rewritten (17 cases, mocked, no kernel invoked)

Accuracy Tests

python benchmark/gsm8k/bench_sglang.py --num-shots 5 --num-questions 1300 --parallel 64 --port 8100

Accuracy: 0.953
Invalid: 0.001
Latency: 100.714 s
Output throughput: 1296.612 token/s

Speed Tests and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the 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 to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
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 #34074956295
Latest PR Test (Extra): ❌ Run #34074956107
Latest PR Test (AMD ROCm 7.2): ⏳ Run #34074956212

[billishyahao]

re-trigger ci due to
https://github.com/sgl-project/sglang/actions/runs/33819121460/job/100859975831

  Error: A major update has landed on `main`. All PRs must rebase onto the latest `main` before CI will run.
  Error: Required base commit: `a5f07b1241fc` (your PR is `diverged` relative to this commit).

[billishyahao]

NV and CPU test are passed. https://github.com/sgl-project/sglang/actions/runs/33998322033/job/101393878232?pr=37601

[HaiShaw]

HIP/AITER gated

[mmangkad]

Why was base-a failure https://github.com/sgl-project/sglang/actions/runs/34074956295/job/101599269921#step:15:1084 ignored here?

[billishyahao]

Hi @mmangkad This base-a failure

======================================================================
ERROR: test_env_disable_not_ready (__main__.TestMlaGluonCapability)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/home/runner/work/sglang/sglang/test/registered/attention/test_mla_gluon_h12_fp8.py", line 20, in setUp
    from sglang.srt.layers.attention.aiter_mla_gluon import (
ImportError: cannot import name 'reset_mla_gluon_state_for_test' from 'sglang.srt.layers.attention.aiter_mla_gluon' (/home/runner/work/sglang/sglang/python/sglang/srt/layers/attention/aiter_mla_gluon.py)

comes from stale unittest test/registered/attention/test_mla_gluon_h12_fp8.py introduced by #34647 which is renamed and replaced by test/registered/attention/test_aiter_gluon_h12_fp8.py (I made this name to let it more specific to aiter)
through initial commits
df4ada8

but then rename behaviour tracking is lost after several rebasing to main. It was reverted by 511e75f Merge remote-tracking branch 'sgl/main' into gluon_dspark. This is unexpected.

Thanks for dropping it #38259