A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Add missing sink tensor into flash attn cascade attn implementation (#26325)
Pei-Lun Liao committed
Oct 7, 2025 at 11:56 UTC
eb577e465589bf4394bc4621228483d0ef34704c
1 file changed
+5
vllm/v1/attention/backends/flash_attn.py
+5
@@ -607,6 +607,7 @@ class FlashAttentionImpl(AttentionImpl):
607
q_descale=layer._q_scale,
608
k_descale=layer._k_scale,
609
v_descale=layer._v_scale,
610
+ s_aux=self.sinks,
611
)
612
return output
613
@@ -767,6 +768,7 @@ def cascade_attention(
768
q_descale: Optional[torch.Tensor] = None,
769
k_descale: Optional[torch.Tensor] = None,
770
v_descale: Optional[torch.Tensor] = None,
771
+ s_aux: Optional[torch.Tensor] = None,
772
) -> torch.Tensor:
773
assert alibi_slopes is None, "Cascade attention does not support ALiBi."
774
# TODO: Support sliding window.
@@ -801,6 +803,9 @@ def cascade_attention(
803
q_descale=q_descale.expand(descale_shape) if q_descale is not None else None,
804
k_descale=k_descale.expand(descale_shape) if k_descale is not None else None,
805
v_descale=v_descale.expand(descale_shape) if v_descale is not None else None,
806
+ # s_aux is incorporated into prefix_lse inside the GPU kernel,
807
+ # enabling its effect during the final attention merge.
808
+ s_aux=s_aux,
809
)
810
811
descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])