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
[gpt-oss] flashinfer attention sink init (#22330)
Signed-off-by: simon-mo <xmo@berkeley.edu> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Minseok Lee <47620120+minseokl@users.noreply.github.com>
Yongye Zhu committed
Aug 5, 2025 at 23:48 UTC
90ec006937c4bcb33b4c0423285fd72502659cfe
1 file changed
+10
vllm/v1/attention/backends/flashinfer.py
+10
@@ -611,6 +611,7 @@ class FlashInferImpl(AttentionImpl):
611
logits_soft_cap: Optional[float] = None,
612
attn_type: AttentionType = AttentionType.DECODER,
613
kv_sharing_target_layer_name: Optional[int] = None,
614
+ sinks: Optional[torch.Tensor] = None,
615
) -> None:
616
self.num_heads = num_heads
617
self.head_size = head_size
@@ -635,6 +636,15 @@ class FlashInferImpl(AttentionImpl):
636
"are not implemented for "
637
"FlashInferImpl")
638
639
+ self.sinks: Optional[torch.Tensor] = None
640
+ if sinks is not None:
641
+ assert sinks.shape[0] == num_heads, (
642
+ "Sinks must have the same number of heads "
643
+ "as the number of heads in the layer"
644
+ )
645
+ assert sinks.dtype == torch.float32, "Sinks must be of type float32"
646
+ self.sinks = sinks
647
+
648
def forward(
649
self,
650
layer: torch.nn.Module,