[Attention] Fix FlashInfer SM12x prefill with sinks (vllm-project#52148)

Signed-off-by: Andrii Skliar <askliar@nvidia.com>
Co-authored-by: Andrii Skliar <askliar@nvidia.com>
Co-authored-by: Michael Goin <mgoin64@gmail.com>
Signed-off-by: Wyett <wyettzeng@gmail.com>

Author: 3 people

tests/v1/attention/test_attention_backends.py

@@ -848,7 +848,7 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype):
     AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
     reason="FlashInfer is not available.",
 )
-def test_flashinfer_native_prefill_with_sinks():
+def test_flashinfer_native_prefill_with_sinks(default_vllm_config):
     if not (
         current_platform.is_cuda() and current_platform.is_device_capability_family(120)
     ):

vllm/v1/attention/backends/flashinfer.py

@@ -10,6 +10,7 @@
 import numpy as np
 import torch
 from flashinfer import (
+    BatchAttentionWithAttentionSinkWrapper,
     BatchDecodeWithPagedKVCacheWrapper,
     BatchPrefillWithPagedKVCacheWrapper,
     BatchPrefillWithRaggedKVCacheWrapper,
@@ -1113,11 +1114,27 @@ def _get_prefill_wrapper(
                     "NVFP4 KV cache."
                 )
             if self._noncausal_prefill_wrapper is None:
-                self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
-                    self._get_workspace_buffer(),
-                    get_kv_cache_layout(),
-                    backend="auto",
-                )
+                if self.has_sinks and current_platform.is_device_capability_family(120):
+                    self._noncausal_prefill_wrapper = (
+                        BatchAttentionWithAttentionSinkWrapper(
+                            self._get_workspace_buffer(),
+                            get_kv_cache_layout(),
+                            backend="auto",
+                            q_data_type=self.q_data_type_prefill,
+                            kv_data_type=self.kv_cache_dtype,
+                            head_dim_qk=self.head_dim,
+                            head_dim_vo=self.head_dim,
+                            window_left=self.window_left,
+                        )
+                    )
+                else:
+                    self._noncausal_prefill_wrapper = (
+                        BatchPrefillWithPagedKVCacheWrapper(
+                            self._get_workspace_buffer(),
+                            get_kv_cache_layout(),
+                            backend="auto",
+                        )
+                    )
             return self._noncausal_prefill_wrapper
 
         if self._prefill_wrapper is None:
@@ -1127,14 +1144,27 @@ def _get_prefill_wrapper(
                     dcp_a2a=self.dcp_a2a,
                 )
             else:
-                # NVFP4 KV cache requires the trtllm-gen backend inside
-                # the wrapper; fa2/fa3 do not support nvfp4.
-                backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
-                self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
-                    self._get_workspace_buffer(),
-                    get_kv_cache_layout(),
-                    backend=backend,
-                )
+                if self.has_sinks and current_platform.is_device_capability_family(120):
+                    assert not self.is_kvcache_nvfp4
+                    self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper(
+                        self._get_workspace_buffer(),
+                        get_kv_cache_layout(),
+                        backend="auto",
+                        q_data_type=self.q_data_type_prefill,
+                        kv_data_type=self.kv_cache_dtype,
+                        head_dim_qk=self.head_dim,
+                        head_dim_vo=self.head_dim,
+                        window_left=self.window_left,
+                    )
+                else:
+                    # NVFP4 KV cache requires the trtllm-gen backend inside
+                    # the wrapper; fa2/fa3 do not support nvfp4.
+                    backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
+                    self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
+                        self._get_workspace_buffer(),
+                        get_kv_cache_layout(),
+                        backend=backend,
+                    )
         assert self._prefill_wrapper is not None
         return self._prefill_wrapper
 
@@ -2086,16 +2116,28 @@ def forward(
                     else:
                         out_prefill = output[num_decode_tokens:]
 
-                    prefill_wrapper.run(
-                        prefill_query,
-                        kv_cache_for_fi,
-                        q_scale=layer._q_scale_float,
-                        k_scale=layer._k_scale_float,
-                        v_scale=layer._v_scale_float,
-                        out=out_prefill,
-                        kv_cache_sf=kv_cache_sf,
-                        sinks=self.sinks,
-                    )
+                    if isinstance(
+                        prefill_wrapper, BatchAttentionWithAttentionSinkWrapper
+                    ):
+                        assert self.sinks is not None
+                        prefill_wrapper.run(
+                            prefill_query,
+                            kv_cache_for_fi,
+                            self.sinks,
+                            self.scale * layer._q_scale_float * layer._k_scale_float,
+                            v_scale=layer._v_scale_float,
+                            out=out_prefill,
+                        )
+                    else:
+                        prefill_wrapper.run(
+                            prefill_query,
+                            kv_cache_for_fi,
+                            q_scale=layer._q_scale_float,
+                            k_scale=layer._k_scale_float,
+                            v_scale=layer._v_scale_float,
+                            out=out_prefill,
+                            kv_cache_sf=kv_cache_sf,
+                        )
 
                     if needs_fp8_out_prefill:
                         output[