From 73d272bddb6f23d4ce98e19e698494d07b76af9c Mon Sep 17 00:00:00 2001 From: akhilg-nv <165961486+akhilg-nv@users.noreply.github.com> Date: Thu, 5 Mar 2026 08:05:35 -0800 Subject: [PATCH] Revised fix for HybridAttnBackend forward for linear attn (#19369) --- .../layers/attention/hybrid_attn_backend.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 88d81ab59..57e10daa6 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_attn_backend.py @@ -111,6 +111,34 @@ class HybridAttnBackend(AttentionBackend): def get_cuda_graph_seq_len_fill_value(self): return self.decode_backend.get_cuda_graph_seq_len_fill_value() + def forward( + self, + q: Optional[torch.Tensor] = None, # For full attention + k: Optional[torch.Tensor] = None, # For full attention + v: Optional[torch.Tensor] = None, # For full attention + layer: Optional[RadixAttention] = None, + forward_batch: Optional[ForwardBatch] = None, + save_kv_cache: bool = True, + *, + mixed_qkv: Optional[torch.Tensor] = None, # For linear attention + a: Optional[torch.Tensor] = None, # For linear attention + b: Optional[torch.Tensor] = None, # For linear attention + **kwargs, + ): + """Forward method that supports both regular attention (q, k, v) and linear attention (mixed_qkv, a, b).""" + backend = self._select_backend(forward_batch.forward_mode) + if mixed_qkv is not None: + return backend.forward( + layer=layer, + forward_batch=forward_batch, + save_kv_cache=save_kv_cache, + mixed_qkv=mixed_qkv, + a=a, + b=b, + **kwargs, + ) + return backend.forward(q, k, v, layer, forward_batch, save_kv_cache, **kwargs) + def forward_decode( self, q: torch.Tensor,