[Attention] Fix FlashInfer SM12x prefill with sinks (#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>
Andrii Skliar committed
Aug 13, 2026 at 23:25 UTC
b652dedd0c505b6d300788f9b4f6c494d2d12d5c
2 files changed
+66
-24
tests/v1/attention/test_attention_backends.py
+1
-1
@@ -848,7 +848,7 @@ def test_flashinfer_attention_sinks_refreshed_after_reload(dtype):
848
AttentionBackendEnum.FLASHINFER not in BACKENDS_TO_TEST,
849
reason="FlashInfer is not available.",
850
)
851
-def test_flashinfer_native_prefill_with_sinks():
851
+def test_flashinfer_native_prefill_with_sinks(default_vllm_config):
852
if not (
853
current_platform.is_cuda() and current_platform.is_device_capability_family(120)
854
):
vllm/v1/attention/backends/flashinfer.py
+65
-23
@@ -10,6 +10,7 @@ from typing import ClassVar
10
import numpy as np
11
import torch
12
from flashinfer import (
13
+ BatchAttentionWithAttentionSinkWrapper,
14
BatchDecodeWithPagedKVCacheWrapper,
15
BatchPrefillWithPagedKVCacheWrapper,
16
BatchPrefillWithRaggedKVCacheWrapper,
@@ -1113,11 +1114,27 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1114
"NVFP4 KV cache."
1115
)
1116
if self._noncausal_prefill_wrapper is None:
1116
- self._noncausal_prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
1117
- self._get_workspace_buffer(),
1118
- get_kv_cache_layout(),
1119
- backend="auto",
1120
- )
1117
+ if self.has_sinks and current_platform.is_device_capability_family(120):
1118
+ self._noncausal_prefill_wrapper = (
1119
+ BatchAttentionWithAttentionSinkWrapper(
1120
+ self._get_workspace_buffer(),
1121
+ get_kv_cache_layout(),
1122
+ backend="auto",
1123
+ q_data_type=self.q_data_type_prefill,
1124
+ kv_data_type=self.kv_cache_dtype,
1125
+ head_dim_qk=self.head_dim,
1126
+ head_dim_vo=self.head_dim,
1127
+ window_left=self.window_left,
1128
+ )
1129
+ )
1130
+ else:
1131
+ self._noncausal_prefill_wrapper = (
1132
+ BatchPrefillWithPagedKVCacheWrapper(
1133
+ self._get_workspace_buffer(),
1134
+ get_kv_cache_layout(),
1135
+ backend="auto",
1136
+ )
1137
+ )
1138
return self._noncausal_prefill_wrapper
1139
1140
if self._prefill_wrapper is None:
@@ -1127,14 +1144,27 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
1144
dcp_a2a=self.dcp_a2a,
1145
)
1146
else:
1130
- # NVFP4 KV cache requires the trtllm-gen backend inside
1131
- # the wrapper; fa2/fa3 do not support nvfp4.
1132
- backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
1133
- self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
1134
- self._get_workspace_buffer(),
1135
- get_kv_cache_layout(),
1136
- backend=backend,
1137
- )
1147
+ if self.has_sinks and current_platform.is_device_capability_family(120):
1148
+ assert not self.is_kvcache_nvfp4
1149
+ self._prefill_wrapper = BatchAttentionWithAttentionSinkWrapper(
1150
+ self._get_workspace_buffer(),
1151
+ get_kv_cache_layout(),
1152
+ backend="auto",
1153
+ q_data_type=self.q_data_type_prefill,
1154
+ kv_data_type=self.kv_cache_dtype,
1155
+ head_dim_qk=self.head_dim,
1156
+ head_dim_vo=self.head_dim,
1157
+ window_left=self.window_left,
1158
+ )
1159
+ else:
1160
+ # NVFP4 KV cache requires the trtllm-gen backend inside
1161
+ # the wrapper; fa2/fa3 do not support nvfp4.
1162
+ backend = "trtllm-gen" if self.is_kvcache_nvfp4 else "auto"
1163
+ self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
1164
+ self._get_workspace_buffer(),
1165
+ get_kv_cache_layout(),
1166
+ backend=backend,
1167
+ )
1168
assert self._prefill_wrapper is not None
1169
return self._prefill_wrapper
1170
@@ -2086,16 +2116,28 @@ class FlashInferImpl(AttentionImpl):
2116
else:
2117
out_prefill = output[num_decode_tokens:]
2118
2089
- prefill_wrapper.run(
2090
- prefill_query,
2091
- kv_cache_for_fi,
2092
- q_scale=layer._q_scale_float,
2093
- k_scale=layer._k_scale_float,
2094
- v_scale=layer._v_scale_float,
2095
- out=out_prefill,
2096
- kv_cache_sf=kv_cache_sf,
2097
- sinks=self.sinks,
2098
- )
2119
+ if isinstance(
2120
+ prefill_wrapper, BatchAttentionWithAttentionSinkWrapper
2121
+ ):
2122
+ assert self.sinks is not None
2123
+ prefill_wrapper.run(
2124
+ prefill_query,
2125
+ kv_cache_for_fi,
2126
+ self.sinks,
2127
+ self.scale * layer._q_scale_float * layer._k_scale_float,
2128
+ v_scale=layer._v_scale_float,
2129
+ out=out_prefill,
2130
+ )
2131
+ else:
2132
+ prefill_wrapper.run(
2133
+ prefill_query,
2134
+ kv_cache_for_fi,
2135
+ q_scale=layer._q_scale_float,
2136
+ k_scale=layer._k_scale_float,
2137
+ v_scale=layer._v_scale_float,
2138
+ out=out_prefill,
2139
+ kv_cache_sf=kv_cache_sf,
2140
+ )
2141
2142
if needs_fp8_out_prefill:
2143
output[