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 13, 08:15 UTC
[Bugfix] Fix SM100 gpt-oss regression due to faulty attn sink support (#28561)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Nov 12, 2025 at 21:40 UTC
a543e678b45a08c6bd98a4e5ebcc244679003659
2 files changed
+36
-10
vllm/utils/flashinfer.py
+21
-10
@@ -35,9 +35,20 @@ FLASHINFER_CUBINS_REPOSITORY = os.environ.get(
35
)
36
37
38
+@functools.cache
39
+def has_flashinfer_cubin() -> bool:
40
+ """Return `True` if flashinfer-cubin package is available."""
41
+ if envs.VLLM_HAS_FLASHINFER_CUBIN:
42
+ return True
43
+ if importlib.util.find_spec("flashinfer_cubin") is not None:
44
+ return True
45
+ logger.debug_once("flashinfer-cubin package was not found")
46
+ return False
47
+
48
+
49
@functools.cache
50
def has_flashinfer() -> bool:
40
- """Return `True` if FlashInfer is available."""
51
+ """Return `True` if flashinfer-python package is available."""
52
# Use find_spec to check if the module exists without importing it
53
# This avoids potential CUDA initialization side effects
54
if importlib.util.find_spec("flashinfer") is None:
@@ -45,7 +56,7 @@ def has_flashinfer() -> bool:
56
return False
57
# When not using flashinfer cubin,
58
# Also check if nvcc is available since it's required to JIT compile flashinfer
48
- if not envs.VLLM_HAS_FLASHINFER_CUBIN and shutil.which("nvcc") is None:
59
+ if not has_flashinfer_cubin() and shutil.which("nvcc") is None:
60
logger.debug_once(
61
"FlashInfer unavailable since nvcc was not found "
62
"and not using pre-downloaded cubins"
@@ -183,9 +194,8 @@ def has_nvidia_artifactory() -> bool:
194
This checks connectivity to the kernel inference library artifactory
195
which is required for downloading certain cubin kernels like TRTLLM FHMA.
196
"""
186
- # Since FLASHINFER_CUBIN_DIR defines the pre-downloaded cubins path, when
187
- # it's true, we could assume the cubins are available.
188
- if envs.VLLM_HAS_FLASHINFER_CUBIN:
197
+ # If we have pre-downloaded cubins, we can assume the cubins are available.
198
+ if has_flashinfer_cubin():
199
return True
200
201
try:
@@ -208,9 +218,13 @@ def has_nvidia_artifactory() -> bool:
218
@functools.cache
219
def supports_trtllm_attention() -> bool:
220
"""
211
- TRTLLM attention is supported if the platform is SM100 and
212
- NVIDIA artifactory is accessible
221
+ TRTLLM attention is supported if the platform is SM100,
222
+ NVIDIA artifactory is accessible, and batch-invariant mode is not enabled.
223
"""
224
+ # Batch-invariant mode disables TRTLLM attention
225
+ if vllm_is_batch_invariant():
226
+ return False
227
+
228
# Requires SM100 and NVIDIA artifactory to be accessible to download cubins
229
return current_platform.is_device_capability(100) and has_nvidia_artifactory()
230
@@ -229,9 +243,6 @@ def force_use_trtllm_attention() -> bool | None:
243
return `True` if TRTLLM attention is forced to be used,
244
return `False` if TRTLLM attention is forced to be not used.
245
"""
232
- if vllm_is_batch_invariant():
233
- logger.info_once("VLLM_USE_TRTLLM_ATTENTION is disabled for batch-invariant")
234
- return False
246
return _force_use_trtllm_attention(envs.VLLM_USE_TRTLLM_ATTENTION)
247
248
vllm/v1/attention/backends/flashinfer.py
+15
@@ -229,6 +229,21 @@ class FlashInferBackend(AttentionBackend):
229
12, 1
230
)
231
232
+ @classmethod
233
+ def supports_sink(cls) -> bool:
234
+ """FlashInfer supports sinks when TRTLLM attention is available (SM100)."""
235
+ from vllm.utils.flashinfer import (
236
+ force_use_trtllm_attention,
237
+ supports_trtllm_attention,
238
+ )
239
+
240
+ # Respect explicit disable flag (e.g., VLLM_USE_TRTLLM_ATTENTION=0)
241
+ if force_use_trtllm_attention() is False:
242
+ return False
243
+
244
+ # Check if TRTLLM is supported on this platform
245
+ return supports_trtllm_attention()
246
+
247
@classmethod
248
def get_required_kv_cache_layout(cls) -> KVCacheLayoutType | None:
249
from vllm.platforms import current_platform