Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
At1a8 wants to merge 2 commits into
base: Qwen3.5_v0.5.9
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
60ab005
support aiter pa for Qwen3.5 GA module in decode phase
Apr 3, 2026
c46ff84
Update python/sglang/srt/layers/attention/attention_registry.py
At1a8 Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion python/sglang/srt/layers/attention/aiter_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,17 @@ def __init__(
model_runner.model_config.num_attention_heads // get_attention_tp_size()
)
self.head_dim = model_runner.model_config.head_dim
self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]
# self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[-1]
if (
model_runner.hybrid_gdn_config is not None
or model_runner.kimi_linear_config is not None
):
# For hybrid linear models, layer_id = 0 may not be full attention
self.v_head_dim = model_runner.token_to_kv_pool.get_v_head_dim()
else:
self.v_head_dim = model_runner.token_to_kv_pool.get_value_buffer(0).shape[
-1
]
Comment on lines +127 to +136

Copilot AI Apr 9, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
self.num_kv_head = model_runner.model_config.get_num_kv_heads(
get_attention_tp_size()
)
Expand Down
9 changes: 9 additions & 0 deletions python/sglang/srt/layers/attention/attention_registry.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
from typing import TYPE_CHECKING
from sglang.srt.utils import get_bool_env_var

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -221,6 +222,14 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
"Expected hybrid GDN or NemotronH models, but got unknown model."
)
full_attn_layers = cfg.full_attention_layer_ids
if get_bool_env_var("SGLANG_USE_AITER_PA"):
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
return HybridLinearAttnBackend(
full_attn_backend,
linear_attn_backend,
full_attn_layers,
AiterAttnBackend(runner),
)
Comment on lines +225 to +232

Copilot AI Apr 9, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
return HybridLinearAttnBackend(
full_attn_backend, linear_attn_backend, full_attn_layers
)
Expand Down
16 changes: 13 additions & 3 deletions python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1567,11 +1567,16 @@ def __init__(
full_attn_backend: AttentionBackend,
linear_attn_backend: MambaAttnBackendBase,
full_attn_layers: list[int],
full_attn_decode_backend: Optional[AttentionBackend] = None,
):
self.full_attn_layers = full_attn_layers
self.full_attn_backend = full_attn_backend
self.linear_attn_backend = linear_attn_backend
self.attn_backend_list = [full_attn_backend, linear_attn_backend]
self.full_attn_decode_backend = None
if full_attn_decode_backend is not None:
self.full_attn_decode_backend = full_attn_decode_backend
self.attn_backend_list += [full_attn_decode_backend]

def init_forward_metadata(self, forward_batch: ForwardBatch):
for attn_backend in self.attn_backend_list:
Expand Down Expand Up @@ -1643,9 +1648,14 @@ def forward_decode(
):
layer_id = layer.layer_id if layer else kwargs["layer_id"]
if layer_id in self.full_attn_layers:
return self.full_attn_backend.forward_decode(
q, k, v, layer, forward_batch, save_kv_cache, **kwargs
)
if self.full_attn_decode_backend is not None:
return self.full_attn_decode_backend.forward_decode(
q, k, v, layer, forward_batch, save_kv_cache, **kwargs
)
Comment on lines +1651 to +1654

Copilot AI Apr 9, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 uses AI. Check for mistakes.
else:
return self.full_attn_backend.forward_decode(
q, k, v, layer, forward_batch, save_kv_cache, **kwargs
)
# Linear attention backend
return self.linear_attn_backend.forward_decode(
q=q,
Expand Down