[feat] support aiter pa for Qwen3.5 GA module in decode phase

[At1a8]

Motivation

Support page attention ragged in decode phase

Modifications

• fix weight loader issue in aiter attn backend to make it compatible with hybrid linear models (self attention module + gated delta net)
• load aiter backend when specify SGLANG_USE_AITER_PA=1, it will help Qwen3.5 use aiter page attention in decode stage

Accuracy Tests

root@smci355-ccs-aus-n12-13:/mnt/raid0/fangyuan/qwen3.5# SGLANG_USE_AITER_PA=1 python3 /opt/sglang/benchmark/gsm8k/bench_sglang.py --port 8000
Downloading from https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl to /tmp/test.jsonl
/tmp/test.jsonl: 732kB [00:00, 31.0MB/s]
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 200/200 [00:26<00:00, 7.57it/s]
Accuracy: 0.945
Invalid: 0.005
Latency: 26.430 s
Output throughput: 1195.773 token/s

Benchmarking 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.

[Copilot]

Pull request overview

Adds support for using the AMD aiter paged-attention backend specifically in the decode phase for hybrid-linear (full-attn + linear-attn) models (notably Qwen3.5), and fixes an aiter KV loader incompatibility for certain hybrid layouts.

Changes:

• Extend HybridLinearAttnBackend to optionally use a separate full-attention backend for decode.
• In attn_backend_wrapper, select AiterAttnBackend for full-attention decode when SGLANG_USE_AITER_PA=1.
• Adjust AiterAttnBackend v-head-dim detection for some hybrid-linear model runners.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 4 comments.

File	Description
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py	Adds an optional decode-only full-attn backend and routes decode calls accordingly.
python/sglang/srt/layers/attention/attention_registry.py	Enables env-var-based selection of aiter paged attention for decode in hybrid-linear models.
python/sglang/srt/layers/attention/aiter_backend.py	Tweaks v-head-dim retrieval to avoid hybrid KV pool layer-id assumptions in some cases.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

full_attn_decode_backend.forward_decode(..., **kwargs) can raise TypeError when the decode backend doesn’t accept extra keyword args. For example, AiterAttnBackend.forward_decode currently has no **kwargs parameter, so any forwarded args like q_rope/k_rope/sinks will break decode. Consider making the decode backend interface-compatible (accept and ignore **kwargs), or filter kwargs before calling the decode backend.

[Copilot]

The new SGLANG_USE_AITER_PA path can select AiterAttnBackend on systems where aiter isn’t installed / supported, leading to a later runtime failure (the aiter_backend module swallows ImportError and continues). Add an explicit availability check here (e.g., import aiter in a try/except and raise a clear error) and/or hardware gating so enabling this env var fails fast with an actionable message.

[Copilot]

This v_head_dim selection is based on specific model configs (hybrid_gdn_config/kimi_linear_config), but the underlying issue is whether token_to_kv_pool.get_value_buffer(0) is valid (it can raise for HybridLinearKVPool when layer 0 isn’t a full-attention layer). To avoid reintroducing the same loader failure for other hybrid models (e.g., NemotronH / other mambaish configs), prefer a capability-based check like hasattr(token_to_kv_pool, "get_v_head_dim") (or isinstance(..., HybridLinearKVPool)) and use get_v_head_dim() when available.