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 15, 07:15 UTC
[V1] Use FlashInfer by default on Blackwell GPUs (#19118)
Michael Goin committed
Jun 5, 2025 at 15:40 UTC
87360308b7ee9b0e529ba9146c93a9ef5526b968
2 files changed
+39
vllm/platforms/cuda.py
+15
@@ -229,6 +229,21 @@ class CudaPlatformBase(Platform):
229
logger.info_once("Using Triton backend on V1 engine.")
230
return ("vllm.v1.attention.backends."
231
"triton_attn.TritonAttentionBackend")
232
+ if cls.is_device_capability(100):
233
+ # Prefer FlashInfer for V1 on Blackwell GPUs if installed
234
+ try:
235
+ import flashinfer # noqa: F401
236
+ logger.info_once(
237
+ "Using FlashInfer backend on V1 engine by default for "
238
+ "Blackwell (SM 10.0) GPUs.")
239
+ return ("vllm.v1.attention.backends."
240
+ "flashinfer.FlashInferBackend")
241
+ except ImportError:
242
+ logger.info_once(
243
+ "FlashInfer failed to import for V1 engine on "
244
+ "Blackwell (SM 10.0) GPUs; it is recommended to "
245
+ "install FlashInfer for better performance.")
246
+ pass
247
if cls.has_device_capability(80):
248
logger.info_once("Using Flash Attention backend on V1 engine.")
249
return ("vllm.v1.attention.backends."
vllm/platforms/interface.py
+24
@@ -228,6 +228,30 @@ class Platform:
228
229
return current_capability.to_int() >= capability
230
231
+ @classmethod
232
+ def is_device_capability(
233
+ cls,
234
+ capability: Union[tuple[int, int], int],
235
+ device_id: int = 0,
236
+ ) -> bool:
237
+ """
238
+ Test whether this platform has exactly the specified device capability.
239
+
240
+ The `capability` argument can either be:
241
+
242
+ - A tuple `(major, minor)`.
243
+ - An integer `<major><minor>`. (See
244
+ [`DeviceCapability.to_int`][vllm.platforms.interface.DeviceCapability.to_int])
245
+ """
246
+ current_capability = cls.get_device_capability(device_id=device_id)
247
+ if current_capability is None:
248
+ return False
249
+
250
+ if isinstance(capability, tuple):
251
+ return current_capability == capability
252
+
253
+ return current_capability.to_int() == capability
254
+
255
@classmethod
256
def get_device_name(cls, device_id: int = 0) -> str:
257
"""Get the name of a device."""