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,