[gpt-oss] flashinfer attention sink init (vllm-project#22330)

Signed-off-by: simon-mo <[email protected]>
Co-authored-by: LiuXiaoxuanPKU <[email protected]>
Co-authored-by: simon-mo <[email protected]>
Co-authored-by: Chen Zhang <[email protected]>
Co-authored-by: Woosuk Kwon <[email protected]>
Co-authored-by: Hongxia Yang <[email protected]>
Co-authored-by: Minseok Lee <[email protected]>

Author: 7 people

vllm/v1/attention/backends/flashinfer.py

@@ -611,6 +611,7 @@ def __init__(
         logits_soft_cap: Optional[float] = None,
         attn_type: AttentionType = AttentionType.DECODER,
         kv_sharing_target_layer_name: Optional[int] = None,
+        sinks: Optional[torch.Tensor] = None,
     ) -> None:
         self.num_heads = num_heads
         self.head_size = head_size
@@ -635,6 +636,15 @@ def __init__(
                                       "are not implemented for "
                                       "FlashInferImpl")
 
+        self.sinks: Optional[torch.Tensor] = None
+        if sinks is not None:
+            assert sinks.shape[0] == num_heads, (
+                "Sinks must have the same number of heads "
+                "as the number of heads in the layer"
+            )
+            assert sinks.dtype == torch.float32, "Sinks must be of type float32"
+            self.sinks = sinks
+
     def forward(
         self,
         layer: torch.nn.Module,