diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index d851040cfc37..dcdd0328ca8c 100644 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -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 + ] self.num_kv_head = model_runner.model_config.get_num_kv_heads( get_attention_tp_size() ) diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 25ffc3f0b412..e3177ae261e5 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -1,5 +1,6 @@ import logging from typing import TYPE_CHECKING +from sglang.srt.utils import get_bool_env_var logger = logging.getLogger(__name__) @@ -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), + ) return HybridLinearAttnBackend( full_attn_backend, linear_attn_backend, full_attn_layers ) diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 47a74066cd0d..a405b1df959f 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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: @@ -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 + ) + 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,