| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Compatibility wrapper for FlashInfer API changes. |
| 4 | |
| 5 | Users of vLLM should always import **only** these wrappers. |
| 6 | """ |
| 7 | |
| 8 | import contextlib |
| 9 | import functools |
| 10 | import importlib |
| 11 | import importlib.util |
| 12 | import os |
| 13 | import shutil |
| 14 | from collections.abc import Callable, Iterator |
| 15 | from contextvars import ContextVar |
| 16 | from typing import Any, NoReturn |
| 17 | |
| 18 | import requests |
| 19 | import torch |
| 20 | |
| 21 | import vllm.envs as envs |
| 22 | from vllm.logger import init_logger |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.utils.math_utils import cdiv |
| 25 | from vllm.utils.torch_utils import PIN_MEMORY |
| 26 | |
| 27 | logger = init_logger(__name__) |
| 28 | |
| 29 | |
| 30 | _bf16_autotune_buckets: ContextVar[tuple[int, ...] | None] = ContextVar( |
| 31 | "flashinfer_bf16_autotune_buckets", default=None |
| 32 | ) |
| 33 | |
| 34 | |
| 35 | @contextlib.contextmanager |
| 36 | def autotune_bf16_only( |
| 37 | tuning_buckets: tuple[int, ...], *, skip_ops: set[str] | None = None |
| 38 | ) -> Iterator[None]: |
| 39 | """Tune BF16 calls with bounded buckets, outside full-model autotuning.""" |
| 40 | token = _bf16_autotune_buckets.set(tuning_buckets) |
| 41 | try: |
| 42 | with autotune(tune_mode=False, skip_ops=skip_ops): |
| 43 | yield |
| 44 | finally: |
| 45 | _bf16_autotune_buckets.reset(token) |
| 46 | |
| 47 | |
| 48 | def flashinfer_bf16_mm( |
| 49 | a: torch.Tensor, |
| 50 | b: torch.Tensor, |
| 51 | bias: torch.Tensor | None, |
| 52 | pdl: bool, |
| 53 | backend: str, |
| 54 | ) -> torch.Tensor: |
| 55 | from flashinfer import mm_bf16 |
| 56 | |
| 57 | tuning_buckets = _bf16_autotune_buckets.get() |
| 58 | tuning = ( |
| 59 | autotune(tune_mode=True, tuning_buckets=tuning_buckets) |
| 60 | if tuning_buckets is not None |
| 61 | else contextlib.nullcontext() |
| 62 | ) |
| 63 | with tuning: |
| 64 | return mm_bf16( |
| 65 | a, |
| 66 | b, |
| 67 | bias=bias, |
| 68 | pdl=pdl, |
| 69 | out_dtype=torch.bfloat16, |
| 70 | backend=backend, |
| 71 | ) |
| 72 | |
| 73 | |
| 74 | # This is the storage path for the cubins, it can be replaced |
| 75 | # with a local path for testing. |
| 76 | # Referenced from https://github.com/flashinfer-ai/flashinfer/blob/0c9a92c3d9a7e043ab6f3f7b2273269caf6ab044/flashinfer/jit/cubin_loader.py#L35 # noqa: E501 |
| 77 | FLASHINFER_CUBINS_REPOSITORY = os.environ.get( |
| 78 | "FLASHINFER_CUBINS_REPOSITORY", |
| 79 | "https://edge.urm.nvidia.com/artifactory/sw-kernelinferencelibrary-public-generic-local/", # noqa: E501 |
| 80 | ) |
| 81 | |
| 82 | _DEFAULT_CUDA_HOME = "/usr/local/cuda" |
| 83 | |
| 84 | |
| 85 | def _flashinfer_nvcc_path() -> str | None: |
| 86 | """Return the nvcc FlashInfer's JIT would run, or None if it is missing. |
| 87 | |
| 88 | Mirrors ``flashinfer.jit.cpp_ext.get_cuda_path()`` without importing |
| 89 | FlashInfer, whose import initializes CUDA. |
| 90 | """ |
| 91 | cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") |
| 92 | if not cuda_home: |
| 93 | nvcc = shutil.which("nvcc") |
| 94 | cuda_home = ( |
| 95 | os.path.dirname(os.path.dirname(nvcc)) if nvcc else _DEFAULT_CUDA_HOME |
| 96 | ) |
| 97 | return shutil.which(os.path.join(cuda_home, "bin", "nvcc")) |
| 98 | |
| 99 | |
| 100 | @functools.cache |
| 101 | def has_flashinfer_cubin() -> bool: |
| 102 | """Return `True` if flashinfer-cubin package is available.""" |
| 103 | if envs.VLLM_HAS_FLASHINFER_CUBIN: |
| 104 | return True |
| 105 | if importlib.util.find_spec("flashinfer_cubin") is not None: |
| 106 | return True |
| 107 | logger.debug_once("flashinfer-cubin package was not found") |
| 108 | return False |
| 109 | |
| 110 | |
| 111 | @functools.cache |
| 112 | def has_flashinfer() -> bool: |
| 113 | """Return `True` if flashinfer-python package is available.""" |
| 114 | # Use find_spec to check if the module exists without importing it |
| 115 | # This avoids potential CUDA initialization side effects |
| 116 | if importlib.util.find_spec("flashinfer") is None: |
| 117 | logger.debug_once("FlashInfer unavailable since package was not found") |
| 118 | return False |
| 119 | # FlashInfer's JIT runs nvcc and `ninja` (from PATH). |
| 120 | if not has_flashinfer_cubin() and ( |
| 121 | _flashinfer_nvcc_path() is None or shutil.which("ninja") is None |
| 122 | ): |
| 123 | logger.warning_once( |
| 124 | "FlashInfer kernels are disabled: flashinfer-cubin is not installed " |
| 125 | "and nvcc (CUDA_HOME, CUDA_PATH, PATH or /usr/local/cuda) or ninja " |
| 126 | "(PATH) is missing. Set CUDA_HOME to a CUDA toolkit and put ninja on " |
| 127 | "PATH, or run `flashinfer download-kernels`." |
| 128 | ) |
| 129 | return False |
| 130 | return True |
| 131 | |
| 132 | |
| 133 | @functools.cache |
| 134 | def has_flashinfer_bf16_gemm() -> bool: |
| 135 | """Return whether FlashInfer exposes the BF16 dense GEMM API.""" |
| 136 | if not has_flashinfer(): |
| 137 | return False |
| 138 | mod = _get_submodule("flashinfer") |
| 139 | return mod is not None and callable(getattr(mod, "mm_bf16", None)) |
| 140 | |
| 141 | |
| 142 | @functools.cache |
| 143 | def is_flashinfer_bf16_gemm_supported( |
| 144 | backend: str, |
| 145 | compute_capability: int | None = None, |
| 146 | ) -> bool: |
| 147 | """Return whether an exact FlashInfer BF16 backend is available.""" |
| 148 | if not current_platform.is_cuda() or not has_flashinfer_bf16_gemm(): |
| 149 | return False |
| 150 | |
| 151 | mod = _get_submodule("flashinfer") |
| 152 | mm_bf16 = getattr(mod, "mm_bf16", None) if mod is not None else None |
| 153 | backend_supported = getattr(mm_bf16, "is_backend_supported", None) |
| 154 | if not callable(backend_supported): |
| 155 | return False |
| 156 | |
| 157 | if compute_capability is None: |
| 158 | device_capability = current_platform.get_device_capability() |
| 159 | if device_capability is None: |
| 160 | return False |
| 161 | compute_capability = device_capability.to_int() |
| 162 | |
| 163 | try: |
| 164 | return bool(backend_supported(backend, compute_capability)) |
| 165 | except (AttributeError, ImportError, RuntimeError, TypeError, ValueError): |
| 166 | return False |
| 167 | |
| 168 | |
| 169 | @functools.cache |
| 170 | def is_flashinfer_cutedsl_bf16_gemm_supported() -> bool: |
| 171 | """Return whether the CuTeDSL BF16 dense GEMM backend is available.""" |
| 172 | if not is_flashinfer_bf16_gemm_supported("cute-dsl"): |
| 173 | return False |
| 174 | try: |
| 175 | from flashinfer.cute_dsl.utils import is_cute_dsl_available |
| 176 | from flashinfer.utils import is_sm100a_supported |
| 177 | except (ImportError, ModuleNotFoundError): |
| 178 | return False |
| 179 | try: |
| 180 | return is_cute_dsl_available() and is_sm100a_supported(torch.device("cuda")) |
| 181 | except (RuntimeError, TypeError, ValueError): |
| 182 | return False |
| 183 | |
| 184 | |
| 185 | def _missing(*_: Any, **__: Any) -> NoReturn: |
| 186 | """Placeholder for unavailable FlashInfer backend.""" |
| 187 | raise RuntimeError( |
| 188 | "FlashInfer backend is not available. Please install the package " |
| 189 | "to enable FlashInfer kernels: " |
| 190 | "https://github.com/flashinfer-ai/flashinfer" |
| 191 | ) |
| 192 | |
| 193 | |
| 194 | def _missing_sparse_mla(*_: Any, **__: Any) -> NoReturn: |
| 195 | raise RuntimeError( |
| 196 | "FlashInfer sparse MLA decode APIs are not available. " |
| 197 | "Install a FlashInfer build that includes sparse MLA decode support." |
| 198 | ) |
| 199 | |
| 200 | |
| 201 | def _get_submodule(module_name: str) -> Any | None: |
| 202 | """Safely import a submodule and return it, or None if not available.""" |
| 203 | try: |
| 204 | return importlib.import_module(module_name) |
| 205 | except (ImportError, ModuleNotFoundError): |
| 206 | return None |
| 207 | |
| 208 | |
| 209 | # General lazy import wrapper |
| 210 | def _lazy_import_wrapper( |
| 211 | module_name: str, attr_name: str, fallback_fn: Callable[..., Any] = _missing |
| 212 | ): |
| 213 | """Create a lazy import wrapper for a specific function.""" |
| 214 | |
| 215 | @functools.cache |
| 216 | def _get_impl(): |
| 217 | if not has_flashinfer(): |
| 218 | return None |
| 219 | mod = _get_submodule(module_name) |
| 220 | return getattr(mod, attr_name, None) if mod else None |
| 221 | |
| 222 | def wrapper(*args, **kwargs): |
| 223 | impl = _get_impl() |
| 224 | if impl is None: |
| 225 | return fallback_fn(*args, **kwargs) |
| 226 | return impl(*args, **kwargs) |
| 227 | |
| 228 | return wrapper |
| 229 | |
| 230 | |
| 231 | def pin_host_range_buf(length: int) -> None: |
| 232 | """Pin flashinfer's cached host-side arange buffer covering `length`. |
| 233 | |
| 234 | flashinfer's decode plan() builds its host qo_indptr from the cached |
| 235 | `_get_range_buf` entry and copies it to the GPU with non_blocking=True on |
| 236 | every call; the cache entry is unpinned, so the copy silently blocks the |
| 237 | host. Replace the entry with a pinned copy so the copies stay async. |
| 238 | """ |
| 239 | if not PIN_MEMORY: |
| 240 | return |
| 241 | mod = _get_submodule("flashinfer.utils") |
| 242 | if mod is None: |
| 243 | return |
| 244 | get_range_buf = getattr(mod, "_get_range_buf", None) |
| 245 | cache_buf = getattr(mod, "_cache_buf", None) |
| 246 | ceil_pow2 = getattr(mod, "_ceil_pow2", None) |
| 247 | if get_range_buf is None or cache_buf is None or ceil_pow2 is None: |
| 248 | return |
| 249 | get_range_buf(length, "cpu") # Ensure the cache entry exists. |
| 250 | key = (f"range_{ceil_pow2(length)}", "cpu") |
| 251 | buf = cache_buf[key] |
| 252 | if not buf.is_pinned(): |
| 253 | cache_buf[key] = buf.pin_memory() |
| 254 | |
| 255 | |
| 256 | # Create lazy wrappers for each function |
| 257 | flashinfer_trtllm_bf16_moe = _lazy_import_wrapper( |
| 258 | "flashinfer.fused_moe", "trtllm_bf16_moe" |
| 259 | ) |
| 260 | flashinfer_trtllm_fp8_block_scale_moe = _lazy_import_wrapper( |
| 261 | "flashinfer.fused_moe", "trtllm_fp8_block_scale_moe" |
| 262 | ) |
| 263 | flashinfer_trtllm_fp8_per_tensor_scale_moe = _lazy_import_wrapper( |
| 264 | "flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe" |
| 265 | ) |
| 266 | flashinfer_cutlass_fused_moe = _lazy_import_wrapper( |
| 267 | "flashinfer.fused_moe", "cutlass_fused_moe" |
| 268 | ) |
| 269 | flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper( |
| 270 | "flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked" |
| 271 | ) |
| 272 | flashinfer_prepare_bf16_fp4_weights = _lazy_import_wrapper( |
| 273 | "flashinfer.gemm", "prepare_bf16_fp4_weights" |
| 274 | ) |
| 275 | flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize") |
| 276 | flashinfer_mxfp4_quantize = _lazy_import_wrapper("flashinfer", "mxfp4_quantize") |
| 277 | nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize") |
| 278 | silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper( |
| 279 | "flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize" |
| 280 | ) |
| 281 | scaled_fp4_grouped_quantize = _lazy_import_wrapper( |
| 282 | "flashinfer", "scaled_fp4_grouped_quantize" |
| 283 | ) |
| 284 | nvfp4_block_scale_interleave = _lazy_import_wrapper( |
| 285 | "flashinfer.fp4_quantization", "block_scale_interleave" |
| 286 | ) |
| 287 | flashinfer_cute_dsl_fused_moe_nvfp4 = _lazy_import_wrapper( |
| 288 | "flashinfer", "cute_dsl_fused_moe_nvfp4" |
| 289 | ) |
| 290 | flashinfer_convert_sf_to_mma_layout = _lazy_import_wrapper( |
| 291 | "flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout" |
| 292 | ) |
| 293 | flashinfer_b12x_fused_moe = _lazy_import_wrapper( |
| 294 | "flashinfer.fused_moe", "b12x_fused_moe" |
| 295 | ) |
| 296 | flashinfer_get_hybrid_num_tokens_buckets = _lazy_import_wrapper( |
| 297 | "flashinfer.fused_moe.utils", "get_hybrid_num_tokens_buckets" |
| 298 | ) |
| 299 | trtllm_fp4_block_scale_moe = _lazy_import_wrapper( |
| 300 | "flashinfer", "trtllm_fp4_block_scale_moe" |
| 301 | ) |
| 302 | flashinfer_trtllm_batch_decode_with_kv_cache_mla = _lazy_import_wrapper( |
| 303 | "flashinfer.decode", |
| 304 | "trtllm_batch_decode_with_kv_cache_mla", |
| 305 | fallback_fn=_missing_sparse_mla, |
| 306 | ) |
| 307 | flashinfer_trtllm_batch_decode_sparse_mla_dsv4 = _lazy_import_wrapper( |
| 308 | "flashinfer.decode", |
| 309 | "trtllm_batch_decode_sparse_mla_dsv4", |
| 310 | fallback_fn=_missing_sparse_mla, |
| 311 | ) |
| 312 | flashinfer_xqa_batch_decode_with_kv_cache = _lazy_import_wrapper( |
| 313 | "flashinfer.decode", |
| 314 | "xqa_batch_decode_with_kv_cache", |
| 315 | ) |
| 316 | flashinfer_packed_fused_kda_decode = _lazy_import_wrapper( |
| 317 | "flashinfer", "packed_fused_kda_decode" |
| 318 | ) |
| 319 | flashinfer_recurrent_kda = _lazy_import_wrapper( |
| 320 | "flashinfer.kda", |
| 321 | "recurrent_kda", |
| 322 | ) |
| 323 | flashinfer_fused_kda_decode = _lazy_import_wrapper( |
| 324 | "flashinfer.kda_decode", |
| 325 | "fused_kda_decode", |
| 326 | ) |
| 327 | |
| 328 | |
| 329 | # Special case for autotune since it returns a context manager |
| 330 | autotune = _lazy_import_wrapper( |
| 331 | "flashinfer.autotuner", |
| 332 | "autotune", |
| 333 | fallback_fn=lambda *args, **kwargs: contextlib.nullcontext(), |
| 334 | ) |
| 335 | |
| 336 | |
| 337 | @functools.cache |
| 338 | def has_flashinfer_comm() -> bool: |
| 339 | """Return `True` if FlashInfer comm module is available.""" |
| 340 | return has_flashinfer() and importlib.util.find_spec("flashinfer.comm") is not None |
| 341 | |
| 342 | |
| 343 | @functools.cache |
| 344 | def has_flashinfer_packed_fused_kda_decode() -> bool: |
| 345 | """Return whether FlashInfer's packed fused KDA decode API is available.""" |
| 346 | if not has_flashinfer(): |
| 347 | return False |
| 348 | module = _get_submodule("flashinfer") |
| 349 | return bool(module and callable(getattr(module, "packed_fused_kda_decode", None))) |
| 350 | |
| 351 | |
| 352 | @functools.cache |
| 353 | def has_flashinfer_nvlink_two_sided() -> bool: |
| 354 | """Return `True` if FlashInfer mnnvl all2all is available.""" |
| 355 | if not has_flashinfer_comm(): |
| 356 | return False |
| 357 | |
| 358 | # Check if all required functions are available |
| 359 | required_functions = [ |
| 360 | ("flashinfer.comm", "Mapping"), |
| 361 | ("flashinfer.comm.mnnvl", "MnnvlMemory"), |
| 362 | ("flashinfer.comm.trtllm_alltoall", "MnnvlMoe"), |
| 363 | ("flashinfer.comm.trtllm_alltoall", "MoEAlltoallInfo"), |
| 364 | ] |
| 365 | |
| 366 | for module_name, attr_name in required_functions: |
| 367 | mod = _get_submodule(module_name) |
| 368 | if not mod or not hasattr(mod, attr_name): |
| 369 | return False |
| 370 | return True |
| 371 | |
| 372 | |
| 373 | @functools.cache |
| 374 | def has_flashinfer_nvlink_one_sided() -> bool: |
| 375 | """Return `True` if FlashInfer trtllm_moe_alltoall module is available.""" |
| 376 | if not has_flashinfer_comm(): |
| 377 | return False |
| 378 | return importlib.util.find_spec("flashinfer.comm.trtllm_moe_alltoall") is not None |
| 379 | |
| 380 | |
| 381 | @functools.cache |
| 382 | def has_flashinfer_moe() -> bool: |
| 383 | """Return `True` if FlashInfer MoE module is available.""" |
| 384 | return ( |
| 385 | has_flashinfer() |
| 386 | and importlib.util.find_spec("flashinfer.fused_moe") is not None |
| 387 | ) |
| 388 | |
| 389 | |
| 390 | @functools.cache |
| 391 | def has_flashinfer_sm90_nope_mla() -> bool: |
| 392 | """FlashInfer SM90 NoPE MLA (FP8 KV with in-kernel dequant, kpe=0). |
| 393 | |
| 394 | Feature-detected via the ``ckv_scale_arr`` run() kwarg introduced with |
| 395 | the SM90 NoPE support (FlashInfer >= 0.6.18), so dev builds carry the |
| 396 | gate without a version parse. |
| 397 | """ |
| 398 | if not has_flashinfer(): |
| 399 | return False |
| 400 | try: |
| 401 | import inspect |
| 402 | |
| 403 | from flashinfer.mla import BatchMLAPagedAttentionWrapper |
| 404 | except ImportError: |
| 405 | return False |
| 406 | try: |
| 407 | params = inspect.signature(BatchMLAPagedAttentionWrapper.run).parameters |
| 408 | except (TypeError, ValueError): |
| 409 | return False |
| 410 | return ( |
| 411 | "ckv_scale_arr" in params |
| 412 | and params["ckv_scale_arr"].kind is inspect.Parameter.KEYWORD_ONLY |
| 413 | ) |
| 414 | |
| 415 | |
| 416 | @functools.cache |
| 417 | def has_flashinfer_sparse_mla_sm120() -> bool: |
| 418 | """Return ``True`` if FlashInfer sparse MLA decode support is available.""" |
| 419 | if not has_flashinfer(): |
| 420 | return False |
| 421 | try: |
| 422 | from flashinfer.autotuner import autotune |
| 423 | from flashinfer.decode import ( |
| 424 | trtllm_batch_decode_sparse_mla_dsv4, |
| 425 | trtllm_batch_decode_with_kv_cache_mla, |
| 426 | ) |
| 427 | except ImportError: |
| 428 | return False |
| 429 | return ( |
| 430 | callable(trtllm_batch_decode_sparse_mla_dsv4) |
| 431 | and callable(trtllm_batch_decode_with_kv_cache_mla) |
| 432 | and callable(autotune) |
| 433 | ) |
| 434 | |
| 435 | |
| 436 | @functools.cache |
| 437 | def has_flashinfer_sparse_mla_sm120_config(num_q_heads: int, top_k: int) -> bool: |
| 438 | """Return whether FlashInfer ships an SM120 DSV4 decode specialization. |
| 439 | |
| 440 | The public sparse MLA API predates some DSV4 shapes, so checking only that |
| 441 | the callable exists can select a package that later aborts or rejects a |
| 442 | valid vLLM configuration. Inspect FlashInfer's dispatch table until it |
| 443 | exposes a public capability query. |
| 444 | """ |
| 445 | if not has_flashinfer_sparse_mla_sm120(): |
| 446 | return False |
| 447 | mod = _get_submodule("flashinfer.mla._sparse_mla_sm120") |
| 448 | dispatch = getattr(mod, "_DECODE_DSV4_DISPATCH", None) if mod else None |
| 449 | return dispatch is not None and (int(num_q_heads), int(top_k)) in dispatch |
| 450 | |
| 451 | |
| 452 | @functools.cache |
| 453 | def has_flashinfer_cutedsl() -> bool: |
| 454 | """Return ``True`` if FlashInfer cutedsl module is available.""" |
| 455 | return ( |
| 456 | has_flashinfer() and importlib.util.find_spec("flashinfer.cute_dsl") is not None |
| 457 | ) |
| 458 | |
| 459 | |
| 460 | @functools.cache |
| 461 | def has_flashinfer_bf16_fp4() -> bool: |
| 462 | """Return ``True`` if FlashInfer's CuTe-DSL W4A16 GEMM is available.""" |
| 463 | if not has_flashinfer_cutedsl(): |
| 464 | return False |
| 465 | mod = _get_submodule("flashinfer.gemm") |
| 466 | return mod is not None and all( |
| 467 | hasattr(mod, name) for name in ("mm_bf16_fp4", "prepare_bf16_fp4_weights") |
| 468 | ) |
| 469 | |
| 470 | |
| 471 | @functools.cache |
| 472 | def has_flashinfer_recurrent_kda() -> bool: |
| 473 | """Return whether FlashInfer recurrent KDA prefill is available.""" |
| 474 | if not has_flashinfer(): |
| 475 | return False |
| 476 | mod = _get_submodule("flashinfer.kda") |
| 477 | return mod is not None and callable(getattr(mod, "recurrent_kda", None)) |
| 478 | |
| 479 | |
| 480 | @functools.cache |
| 481 | def has_flashinfer_fused_kda_decode() -> bool: |
| 482 | """Return whether FlashInfer fused KDA decode is available.""" |
| 483 | if not has_flashinfer(): |
| 484 | return False |
| 485 | mod = _get_submodule("flashinfer.kda_decode") |
| 486 | return ( |
| 487 | mod is not None |
| 488 | and bool(getattr(mod, "_FUSED_KDA_DECODE_AVAILABLE", False)) |
| 489 | and callable(getattr(mod, "fused_kda_decode", None)) |
| 490 | ) |
| 491 | |
| 492 | |
| 493 | @functools.cache |
| 494 | def has_flashinfer_trtllm_fused_moe() -> bool: |
| 495 | """Return `True` if FlashInfer TRTLLM fused MoE is available.""" |
| 496 | if not has_flashinfer_moe(): |
| 497 | return False |
| 498 | required_functions = [ |
| 499 | ("flashinfer.fused_moe", "trtllm_fp8_block_scale_moe"), |
| 500 | ("flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe"), |
| 501 | ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"), |
| 502 | ("flashinfer.fused_moe", "trtllm_mxint4_block_scale_moe"), |
| 503 | ("flashinfer.fused_moe", "trtllm_bf16_moe"), |
| 504 | ] |
| 505 | for module_name, attr_name in required_functions: |
| 506 | mod = _get_submodule(module_name) |
| 507 | if not mod or not hasattr(mod, attr_name): |
| 508 | return False |
| 509 | return True |
| 510 | |
| 511 | |
| 512 | @functools.cache |
| 513 | def has_flashinfer_cutlass_fused_moe() -> bool: |
| 514 | """Return `True` if FlashInfer CUTLASS fused MoE is available.""" |
| 515 | if not has_flashinfer_moe(): |
| 516 | return False |
| 517 | |
| 518 | # Check if all required functions are available |
| 519 | required_functions = [ |
| 520 | ("flashinfer.fused_moe", "cutlass_fused_moe"), |
| 521 | ("flashinfer", "fp4_quantize"), |
| 522 | ("flashinfer", "nvfp4_block_scale_interleave"), |
| 523 | ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"), |
| 524 | ] |
| 525 | |
| 526 | for module_name, attr_name in required_functions: |
| 527 | mod = _get_submodule(module_name) |
| 528 | if not mod or not hasattr(mod, attr_name): |
| 529 | return False |
| 530 | return True |
| 531 | |
| 532 | |
| 533 | @functools.cache |
| 534 | def has_flashinfer_cutedsl_grouped_gemm_nt_masked() -> bool: |
| 535 | """Return ``True`` if FlashInfer CUTLASS fused MoE is available.""" |
| 536 | if not has_flashinfer_cutedsl(): |
| 537 | return False |
| 538 | |
| 539 | # Check if all required functions are available |
| 540 | required_functions = [ |
| 541 | ("flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"), |
| 542 | ("flashinfer", "scaled_fp4_grouped_quantize"), |
| 543 | ("flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize"), |
| 544 | ] |
| 545 | |
| 546 | for module_name, attr_name in required_functions: |
| 547 | mod = _get_submodule(module_name) |
| 548 | if not mod or not hasattr(mod, attr_name): |
| 549 | return False |
| 550 | return True |
| 551 | |
| 552 | |
| 553 | @functools.cache |
| 554 | def has_flashinfer_cutedsl_moe_nvfp4() -> bool: |
| 555 | """Return ``True`` if FlashInfer cute_dsl_fused_moe_nvfp4 is available.""" |
| 556 | if not has_flashinfer_cutedsl(): |
| 557 | return False |
| 558 | mod = _get_submodule("flashinfer") |
| 559 | return mod is not None and hasattr(mod, "cute_dsl_fused_moe_nvfp4") |
| 560 | |
| 561 | |
| 562 | @functools.cache |
| 563 | def has_flashinfer_b12x_gemm() -> bool: |
| 564 | """Return True if FlashInfer b12x FP4 GEMM backend is available (SM120+).""" |
| 565 | if not has_flashinfer_cutedsl(): |
| 566 | return False |
| 567 | mod = _get_submodule("flashinfer.gemm") |
| 568 | if mod is None: |
| 569 | return False |
| 570 | # FlashInfer 0.6.11 renamed Sm120BlockScaledDenseGemmKernel -> |
| 571 | # Sm120B12xBlockScaledDenseGemmKernel (commit 223f2a49). Accept either. |
| 572 | return hasattr(mod, "Sm120B12xBlockScaledDenseGemmKernel") or hasattr( |
| 573 | mod, "Sm120BlockScaledDenseGemmKernel" |
| 574 | ) |
| 575 | |
| 576 | |
| 577 | @functools.cache |
| 578 | def has_flashinfer_b12x_moe() -> bool: |
| 579 | """Return ``True`` if FlashInfer CuteDSL SM12x fused MoE is available.""" |
| 580 | if not has_flashinfer_moe(): |
| 581 | return False |
| 582 | |
| 583 | required_functions = [ |
| 584 | ("flashinfer.fused_moe", "b12x_fused_moe"), |
| 585 | ("flashinfer.cute_dsl.utils", "convert_sf_to_mma_layout"), |
| 586 | ] |
| 587 | |
| 588 | for module_name, attr_name in required_functions: |
| 589 | mod = _get_submodule(module_name) |
| 590 | if not mod or not hasattr(mod, attr_name): |
| 591 | return False |
| 592 | return True |
| 593 | |
| 594 | |
| 595 | @functools.cache |
| 596 | def has_nvidia_artifactory() -> bool: |
| 597 | """Return `True` if NVIDIA's artifactory is accessible. |
| 598 | |
| 599 | This checks connectivity to the kernel inference library artifactory |
| 600 | which is required for downloading certain cubin kernels like TRTLLM FHMA. |
| 601 | """ |
| 602 | # If we have pre-downloaded cubins, we can assume the cubins are available. |
| 603 | if has_flashinfer_cubin(): |
| 604 | return True |
| 605 | |
| 606 | try: |
| 607 | # Use a short timeout to avoid blocking for too long |
| 608 | response = requests.get(FLASHINFER_CUBINS_REPOSITORY, timeout=5) |
| 609 | accessible = response.status_code == 200 |
| 610 | if accessible: |
| 611 | logger.debug_once("NVIDIA artifactory is accessible") |
| 612 | else: |
| 613 | logger.warning_once( |
| 614 | "NVIDIA artifactory returned failed status code: %d", |
| 615 | response.status_code, |
| 616 | ) |
| 617 | return accessible |
| 618 | except Exception as e: |
| 619 | logger.warning_once("Failed to connect to NVIDIA artifactory: %s", e) |
| 620 | return False |
| 621 | |
| 622 | |
| 623 | @functools.cache |
| 624 | def supports_trtllm_attention(is_prefill: bool = False) -> bool: |
| 625 | """Return whether TRTLLM attention is available on the current platform |
| 626 | for the given attention phase. |
| 627 | |
| 628 | SM90 (Hopper) and SM12x support the XQA decode kernel but not TRTLLM |
| 629 | prefill. SM100+ supports TRTLLM for both phases. All others are unsupported. |
| 630 | """ |
| 631 | # Batch-invariant mode disables TRTLLM attention |
| 632 | if envs.VLLM_BATCH_INVARIANT: |
| 633 | return False |
| 634 | |
| 635 | # Requires NVIDIA artifactory to be accessible to download cubins |
| 636 | if not has_nvidia_artifactory(): |
| 637 | return False |
| 638 | |
| 639 | # SM90 and SM12x have XQA decode only. |
| 640 | if current_platform.is_device_capability( |
| 641 | 90 |
| 642 | ) or current_platform.is_device_capability_family(120): |
| 643 | return not is_prefill |
| 644 | |
| 645 | # SM100/SM103 has both prefill and decode TRTLLM kernels. |
| 646 | return current_platform.is_device_capability_family(100) |
| 647 | |
| 648 | |
| 649 | def force_use_trtllm_attention() -> bool | None: |
| 650 | """This function should only be called during initialization stage when vllm config |
| 651 | is set. |
| 652 | Return `None` if --attention-config.use_trtllm_attention is not set, |
| 653 | return `True` if TRTLLM attention is forced to be used, |
| 654 | return `False` if TRTLLM attention is forced to be not used. |
| 655 | """ |
| 656 | from vllm.config import get_current_vllm_config |
| 657 | |
| 658 | vllm_config = get_current_vllm_config() |
| 659 | return vllm_config.attention_config.use_trtllm_attention |
| 660 | |
| 661 | |
| 662 | def can_use_trtllm_attention( |
| 663 | num_qo_heads: int, num_kv_heads: int, is_prefill: bool = False |
| 664 | ) -> bool: |
| 665 | """Check if the current configuration supports TRTLLM attention.""" |
| 666 | if force_use_trtllm_attention() is False: |
| 667 | return False |
| 668 | return supports_trtllm_attention(is_prefill=is_prefill) and ( |
| 669 | num_qo_heads % num_kv_heads == 0 |
| 670 | ) |
| 671 | |
| 672 | |
| 673 | def use_trtllm_attention( |
| 674 | num_qo_heads: int, |
| 675 | num_kv_heads: int, |
| 676 | num_tokens: int, |
| 677 | max_seq_len: int, |
| 678 | dcp_world_size: int, |
| 679 | kv_cache_dtype: str, |
| 680 | q_dtype: torch.dtype, |
| 681 | is_prefill: bool, |
| 682 | # None means auto-detection, True means force on, False means force off |
| 683 | force_use_trtllm: bool | None = None, |
| 684 | has_sinks: bool = False, |
| 685 | has_spec: bool = False, |
| 686 | ) -> bool: |
| 687 | """Return `True` if TRTLLM attention is used.""" |
| 688 | # CLI argument is set to 0 - respect it |
| 689 | if force_use_trtllm is not None and not force_use_trtllm: |
| 690 | return False |
| 691 | |
| 692 | # TRTLLM prefill attends only the DCP-local KV shard and has no |
| 693 | # cross-rank LSE combine, so it cannot be used with DCP; fall back to |
| 694 | # FlashInfer's DCP prefill path. TRTLLM decode under DCP is selected |
| 695 | # separately (all-gathered query heads + LSE combine in forward). |
| 696 | if dcp_world_size > 1: |
| 697 | logger.warning_once( |
| 698 | "TRTLLM prefill does not support DCP, reverting to FlashInfer" |
| 699 | ) |
| 700 | return False |
| 701 | |
| 702 | # The platform is not supported |
| 703 | if not supports_trtllm_attention(is_prefill=is_prefill): |
| 704 | if force_use_trtllm: |
| 705 | logger.warning_once( |
| 706 | "TRTLLM attention is not supported on this platform for %s, " |
| 707 | "but --attention-config.use_trtllm_attention is set to 1", |
| 708 | "prefill" if is_prefill else "decode", |
| 709 | ) |
| 710 | return False |
| 711 | |
| 712 | # The combination of query and key heads is not supported |
| 713 | if num_qo_heads % num_kv_heads != 0: |
| 714 | if force_use_trtllm: |
| 715 | logger.warning_once( |
| 716 | "TRTLLM attention is not supported for this combination of " |
| 717 | "query and key heads, but --attention-config.use_trtllm_attention is " |
| 718 | "set to 1" |
| 719 | ) |
| 720 | return False |
| 721 | |
| 722 | if has_spec and not is_prefill: |
| 723 | # Speculative decoding requires TRTLLM attention for decodes |
| 724 | logger.info_once("Using TRTLLM attention (enabled for speculative decoding).") |
| 725 | return True |
| 726 | |
| 727 | # Must use TRTLLM attention if query is FP8 quantized |
| 728 | if q_dtype == current_platform.fp8_dtype(): |
| 729 | logger.info_once("Using TRTLLM attention (query is quantized).") |
| 730 | return True |
| 731 | |
| 732 | # If sinks are being used, we must use TRTLLM attention as it's |
| 733 | # the only backend that supports them |
| 734 | if has_sinks: |
| 735 | logger.info_once("Using TRTLLM attention (required for attention sinks).") |
| 736 | return True |
| 737 | |
| 738 | if force_use_trtllm is None: |
| 739 | # CLI argument not set - use auto-detection |
| 740 | if is_prefill: |
| 741 | # Prefill auto-detection |
| 742 | use_trtllm = kv_cache_dtype == "auto" |
| 743 | elif ( |
| 744 | current_platform.is_device_capability(90) |
| 745 | or current_platform.is_device_capability_family(120) |
| 746 | ) and kv_cache_dtype.startswith("fp8"): |
| 747 | # SM90/SM12x + FP8 KV cache: prefer the XQA decode kernel. |
| 748 | use_trtllm = True |
| 749 | else: |
| 750 | # Decode auto-detection |
| 751 | use_trtllm = num_tokens <= 256 and kv_cache_dtype == "auto" |
| 752 | if use_trtllm: |
| 753 | logger.warning_once( |
| 754 | "Using TRTLLM %s attention (auto-detected).", |
| 755 | "prefill" if is_prefill else "decode", |
| 756 | ) |
| 757 | return use_trtllm |
| 758 | |
| 759 | # CLI argument is set to 1 - respect it |
| 760 | logger.info_once( |
| 761 | "Using TRTLLM attention (--attention-config.use_trtllm_attention is set to 1)" |
| 762 | ) |
| 763 | return True |
| 764 | |
| 765 | |
| 766 | if has_flashinfer(): |
| 767 | from vllm.utils.torch_utils import direct_register_custom_op |
| 768 | |
| 769 | def _flashinfer_concat_mla_k( |
| 770 | k: torch.Tensor, |
| 771 | k_nope: torch.Tensor, |
| 772 | k_pe: torch.Tensor, |
| 773 | ) -> None: |
| 774 | """Custom op wrapper for flashinfer's concat_mla_k. |
| 775 | |
| 776 | This is an in-place operation that concatenates k_nope and k_pe into k. |
| 777 | |
| 778 | The kernel is optimized for DeepSeek V3 dimensions: |
| 779 | - num_heads=128 |
| 780 | - nope_dim=128 |
| 781 | - rope_dim=64 |
| 782 | |
| 783 | Key optimizations: |
| 784 | - Warp-based processing with software pipelining |
| 785 | - Vectorized memory access (int2 for nope, int for rope) |
| 786 | - L2 prefetching for next row while processing current |
| 787 | - Register reuse for rope values across all heads |
| 788 | |
| 789 | Args: |
| 790 | k: Output tensor, shape [num_tokens, num_heads, nope_dim + rope_dim]. |
| 791 | Modified in-place. |
| 792 | k_nope: The nope part of k, shape [num_tokens, num_heads, nope_dim]. |
| 793 | k_pe: The rope part of k (shared), shape [num_tokens, 1, rope_dim]. |
| 794 | This is broadcast to all heads. |
| 795 | |
| 796 | """ |
| 797 | from flashinfer.concat_ops import concat_mla_k |
| 798 | |
| 799 | concat_mla_k(k, k_nope, k_pe) |
| 800 | |
| 801 | # Register flashinfer concat_mla_k custom op |
| 802 | direct_register_custom_op( |
| 803 | op_name="flashinfer_concat_mla_k", |
| 804 | op_func=_flashinfer_concat_mla_k, |
| 805 | mutates_args=["k"], # k tensor is modified in-place |
| 806 | ) |
| 807 | |
| 808 | @torch.library.custom_op( |
| 809 | "vllm::flashinfer_mm_fp4", |
| 810 | mutates_args=[], |
| 811 | device_types="cuda", |
| 812 | ) |
| 813 | def flashinfer_mm_fp4( |
| 814 | A: torch.Tensor, |
| 815 | B: torch.Tensor, |
| 816 | A_scale: torch.Tensor, |
| 817 | B_scale: torch.Tensor, |
| 818 | g_scale: torch.Tensor, |
| 819 | dtype: torch.dtype, |
| 820 | use_8x4_sf_layout: bool, |
| 821 | backend: str, |
| 822 | block_size: int = 16, |
| 823 | use_nvfp4: bool = True, |
| 824 | ) -> torch.Tensor: |
| 825 | from flashinfer import mm_fp4 as flashinfer_mm_fp4_ |
| 826 | |
| 827 | return flashinfer_mm_fp4_( |
| 828 | A, |
| 829 | B, |
| 830 | A_scale, |
| 831 | B_scale, |
| 832 | g_scale, |
| 833 | dtype, |
| 834 | block_size=block_size, |
| 835 | use_8x4_sf_layout=use_8x4_sf_layout, |
| 836 | use_nvfp4=use_nvfp4, |
| 837 | backend=backend, |
| 838 | ) |
| 839 | |
| 840 | @torch.library.register_fake( |
| 841 | "vllm::flashinfer_mm_fp4", |
| 842 | ) |
| 843 | def flashinfer_mm_fp4_fake( |
| 844 | A: torch.Tensor, |
| 845 | B: torch.Tensor, |
| 846 | A_scale: torch.Tensor, |
| 847 | B_scale: torch.Tensor, |
| 848 | g_scale: torch.Tensor, |
| 849 | dtype: torch.dtype, |
| 850 | use_8x4_sf_layout: bool, |
| 851 | backend: str, |
| 852 | block_size: int = 16, |
| 853 | use_nvfp4: bool = True, |
| 854 | ) -> torch.Tensor: |
| 855 | return torch.empty(A.shape[0], B.shape[1], dtype=dtype, device=A.device) |
| 856 | |
| 857 | @torch.library.custom_op( |
| 858 | "vllm::flashinfer_mm_bf16_fp4", |
| 859 | mutates_args=[], |
| 860 | device_types="cuda", |
| 861 | ) |
| 862 | def flashinfer_mm_bf16_fp4( |
| 863 | A: torch.Tensor, |
| 864 | B: torch.Tensor, |
| 865 | B_scale: torch.Tensor, |
| 866 | global_scale: torch.Tensor, |
| 867 | ) -> torch.Tensor: |
| 868 | from flashinfer.gemm import mm_bf16_fp4 |
| 869 | |
| 870 | return mm_bf16_fp4( |
| 871 | A, |
| 872 | B, |
| 873 | B_scale, |
| 874 | global_scale, |
| 875 | backend="cute-dsl", |
| 876 | ) |
| 877 | |
| 878 | @torch.library.register_fake( |
| 879 | "vllm::flashinfer_mm_bf16_fp4", |
| 880 | ) |
| 881 | def flashinfer_mm_bf16_fp4_fake( |
| 882 | A: torch.Tensor, |
| 883 | B: torch.Tensor, |
| 884 | B_scale: torch.Tensor, |
| 885 | global_scale: torch.Tensor, |
| 886 | ) -> torch.Tensor: |
| 887 | output_size = B.shape[0] if B.dtype == torch.uint8 else B.shape[1] // 2 |
| 888 | return torch.empty(A.shape[0], output_size, dtype=A.dtype, device=A.device) |
| 889 | |
| 890 | @torch.library.custom_op( |
| 891 | "vllm::flashinfer_mxfp4_quantize", |
| 892 | mutates_args=[], |
| 893 | device_types="cuda", |
| 894 | ) |
| 895 | def flashinfer_mxfp4_quantize( |
| 896 | a: torch.Tensor, |
| 897 | backend: str, |
| 898 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 899 | from flashinfer import mxfp4_quantize as _mxfp4_quantize |
| 900 | |
| 901 | return _mxfp4_quantize(a, backend=backend) |
| 902 | |
| 903 | @torch.library.register_fake("vllm::flashinfer_mxfp4_quantize") |
| 904 | def flashinfer_mxfp4_quantize_fake( |
| 905 | a: torch.Tensor, |
| 906 | backend: str, |
| 907 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 908 | m, k = a.shape |
| 909 | sf_vec_size = 32 |
| 910 | padded_m = cdiv(m, 128) * 128 |
| 911 | sf_cols = cdiv(k // sf_vec_size, 4) * 4 |
| 912 | return ( |
| 913 | torch.empty(m, k // 2, dtype=torch.uint8, device=a.device), |
| 914 | torch.empty(padded_m, sf_cols, dtype=torch.uint8, device=a.device), |
| 915 | ) |
| 916 | |
| 917 | @torch.library.custom_op( |
| 918 | "vllm::bmm_fp8", |
| 919 | mutates_args=[], |
| 920 | device_types="cuda", |
| 921 | ) |
| 922 | def bmm_fp8( |
| 923 | A: torch.Tensor, |
| 924 | B: torch.Tensor, |
| 925 | A_scale: torch.Tensor, |
| 926 | B_scale: torch.Tensor, |
| 927 | dtype: torch.dtype, |
| 928 | backend: str, |
| 929 | ) -> torch.Tensor: |
| 930 | from flashinfer import bmm_fp8 as bmm_fp8_ |
| 931 | |
| 932 | return bmm_fp8_(A, B, A_scale, B_scale, dtype, None, backend) |
| 933 | |
| 934 | @torch.library.register_fake( |
| 935 | "vllm::bmm_fp8", |
| 936 | ) |
| 937 | def bmm_fp8_fake( |
| 938 | A: torch.Tensor, |
| 939 | B: torch.Tensor, |
| 940 | A_scale: torch.Tensor, |
| 941 | B_scale: torch.Tensor, |
| 942 | dtype: torch.dtype, |
| 943 | backend: str, |
| 944 | ) -> torch.Tensor: |
| 945 | return torch.empty( |
| 946 | A.shape[0], A.shape[1], B.shape[2], dtype=dtype, device=A.device |
| 947 | ) |
| 948 | |
| 949 | @torch.library.custom_op( |
| 950 | "vllm::flashinfer_nvfp4_quantize", |
| 951 | mutates_args=[], |
| 952 | device_types="cuda", |
| 953 | ) |
| 954 | def flashinfer_nvfp4_quantize( |
| 955 | a: torch.Tensor, a_global_sf: torch.Tensor |
| 956 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 957 | from flashinfer import SfLayout |
| 958 | from flashinfer import nvfp4_quantize as nvfp4_quantize_ |
| 959 | |
| 960 | return nvfp4_quantize_( |
| 961 | a, a_global_sf, sfLayout=SfLayout.layout_8x4, do_shuffle=False |
| 962 | ) |
| 963 | |
| 964 | @torch.library.register_fake( |
| 965 | "vllm::flashinfer_nvfp4_quantize", |
| 966 | ) |
| 967 | def flashinfer_nvfp4_quantize_fake( |
| 968 | a: torch.Tensor, a_global_sf: torch.Tensor |
| 969 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 970 | m, n = a.shape |
| 971 | |
| 972 | round_up = lambda x, y: (x + y - 1) // y * y |
| 973 | |
| 974 | rounded_m = round_up(m, 8) |
| 975 | scale_n = n // 16 |
| 976 | rounded_n = round_up(scale_n, 4) |
| 977 | |
| 978 | return torch.empty(m, n // 2, dtype=torch.uint8, device=a.device), torch.empty( |
| 979 | rounded_m, rounded_n, dtype=torch.uint8, device=a.device |
| 980 | ) |
| 981 | |
| 982 | @torch.library.custom_op( |
| 983 | "vllm::flashinfer_mxfp8_quantize_8x4", |
| 984 | mutates_args=[], |
| 985 | device_types="cuda", |
| 986 | ) |
| 987 | def flashinfer_mxfp8_quantize_8x4( |
| 988 | a: torch.Tensor, |
| 989 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 990 | from flashinfer import SfLayout |
| 991 | from flashinfer import mxfp8_quantize as mxfp8_quantize_ |
| 992 | |
| 993 | return mxfp8_quantize_( |
| 994 | a, |
| 995 | backend="cuda", |
| 996 | sf_swizzle_layout=SfLayout.layout_8x4, |
| 997 | ) |
| 998 | |
| 999 | @torch.library.register_fake( |
| 1000 | "vllm::flashinfer_mxfp8_quantize_8x4", |
| 1001 | ) |
| 1002 | def flashinfer_mxfp8_quantize_8x4_fake( |
| 1003 | a: torch.Tensor, |
| 1004 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1005 | m, k = a.shape |
| 1006 | scale_size = cdiv(m, 8) * 8 * cdiv(k // 32, 4) * 4 |
| 1007 | return ( |
| 1008 | torch.empty(m, k, dtype=torch.float8_e4m3fn, device=a.device), |
| 1009 | torch.empty(scale_size, dtype=torch.uint8, device=a.device), |
| 1010 | ) |
| 1011 | |
| 1012 | @torch.library.custom_op( |
| 1013 | "vllm::mm_mxfp8", |
| 1014 | mutates_args=[], |
| 1015 | device_types="cuda", |
| 1016 | ) |
| 1017 | def mm_mxfp8( |
| 1018 | A: torch.Tensor, |
| 1019 | B: torch.Tensor, |
| 1020 | A_scale: torch.Tensor, |
| 1021 | B_scale: torch.Tensor, |
| 1022 | out_dtype: torch.dtype, |
| 1023 | backend: str = "cutlass", |
| 1024 | use_8x4_sf_layout: bool = False, |
| 1025 | ) -> torch.Tensor: |
| 1026 | from flashinfer import mm_mxfp8 as mm_mxfp8_ |
| 1027 | |
| 1028 | return mm_mxfp8_( |
| 1029 | A, |
| 1030 | B, |
| 1031 | A_scale, |
| 1032 | B_scale, |
| 1033 | out=None, |
| 1034 | out_dtype=out_dtype, |
| 1035 | backend=backend, |
| 1036 | use_8x4_sf_layout=use_8x4_sf_layout, |
| 1037 | ) |
| 1038 | |
| 1039 | @torch.library.register_fake( |
| 1040 | "vllm::mm_mxfp8", |
| 1041 | ) |
| 1042 | def mm_mxfp8_fake( |
| 1043 | A: torch.Tensor, |
| 1044 | B: torch.Tensor, |
| 1045 | A_scale: torch.Tensor, |
| 1046 | B_scale: torch.Tensor, |
| 1047 | out_dtype: torch.dtype, |
| 1048 | backend: str = "cutlass", |
| 1049 | use_8x4_sf_layout: bool = False, |
| 1050 | ) -> torch.Tensor: |
| 1051 | # A is [m, k], B is [k, n] -> output [m, n] |
| 1052 | return torch.empty(A.shape[0], B.shape[1], dtype=out_dtype, device=A.device) |
| 1053 | |
| 1054 | |
| 1055 | def flashinfer_scaled_fp4_mm( |
| 1056 | a: torch.Tensor, |
| 1057 | b: torch.Tensor, |
| 1058 | block_scale_a: torch.Tensor, |
| 1059 | block_scale_b: torch.Tensor, |
| 1060 | alpha: torch.Tensor | None, |
| 1061 | out_dtype: torch.dtype, |
| 1062 | backend: str, |
| 1063 | block_size: int = 16, |
| 1064 | use_nvfp4: bool = True, |
| 1065 | ) -> torch.Tensor: |
| 1066 | assert a.ndim == 2 and b.ndim == 2 |
| 1067 | assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2 |
| 1068 | assert a.stride(-1) == 1 and b.stride(-1) == 1 |
| 1069 | assert a.shape[1] == b.shape[1] |
| 1070 | |
| 1071 | if alpha is None: |
| 1072 | alpha = torch.ones(1, dtype=torch.float32, device=a.device) |
| 1073 | |
| 1074 | if backend in ("cutlass", "cudnn"): |
| 1075 | block_scale_a = block_scale_a.view(torch.uint8) |
| 1076 | block_scale_b = block_scale_b.view(torch.uint8) |
| 1077 | |
| 1078 | use_8x4_sf_layout = True if backend == "trtllm" and a.shape[0] <= 32 else False # noqa: SIM210 |
| 1079 | |
| 1080 | return flashinfer_mm_fp4( |
| 1081 | a, |
| 1082 | b.t(), |
| 1083 | block_scale_a, |
| 1084 | block_scale_b.t(), |
| 1085 | alpha, |
| 1086 | out_dtype, |
| 1087 | use_8x4_sf_layout=use_8x4_sf_layout, |
| 1088 | backend=backend, |
| 1089 | block_size=block_size, |
| 1090 | use_nvfp4=use_nvfp4, |
| 1091 | ) |
| 1092 | |
| 1093 | |
| 1094 | def flashinfer_scaled_fp4_mm_out( |
| 1095 | a: torch.Tensor, |
| 1096 | b: torch.Tensor, |
| 1097 | block_scale_a: torch.Tensor, |
| 1098 | block_scale_b: torch.Tensor, |
| 1099 | alpha: torch.Tensor, |
| 1100 | out: torch.Tensor, |
| 1101 | out_dtype: torch.dtype | None, |
| 1102 | use_8x4_sf_layout: bool, |
| 1103 | backend: str, |
| 1104 | ) -> torch.Tensor: |
| 1105 | assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 |
| 1106 | assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2 |
| 1107 | assert a.stride(-1) == 1 |
| 1108 | assert a.shape[1] == b.shape[0] |
| 1109 | assert out.shape == (a.shape[0], b.shape[1]) |
| 1110 | assert out.device.type == "cuda" |
| 1111 | |
| 1112 | if backend in ("cutlass", "cudnn"): |
| 1113 | if block_scale_a.dtype != torch.uint8: |
| 1114 | block_scale_a = block_scale_a.view(torch.uint8) |
| 1115 | if block_scale_b.dtype != torch.uint8: |
| 1116 | block_scale_b = block_scale_b.view(torch.uint8) |
| 1117 | |
| 1118 | from flashinfer import mm_fp4 as flashinfer_mm_fp4_ |
| 1119 | |
| 1120 | flashinfer_mm_fp4_( |
| 1121 | a, |
| 1122 | b, |
| 1123 | block_scale_a, |
| 1124 | block_scale_b, |
| 1125 | alpha, |
| 1126 | out_dtype or out.dtype, |
| 1127 | out=out, |
| 1128 | block_size=16, |
| 1129 | use_8x4_sf_layout=use_8x4_sf_layout, |
| 1130 | backend=backend, |
| 1131 | ) |
| 1132 | return out |
| 1133 | |
| 1134 | |
| 1135 | def flashinfer_scaled_fp8_mm( |
| 1136 | a: torch.Tensor, |
| 1137 | b: torch.Tensor, |
| 1138 | scale_a: torch.Tensor, |
| 1139 | scale_b: torch.Tensor, |
| 1140 | out_dtype: torch.dtype, |
| 1141 | bias: torch.Tensor | None = None, |
| 1142 | ) -> torch.Tensor: |
| 1143 | assert a.ndim == 2 and b.ndim == 2 |
| 1144 | assert a.shape[1] == b.shape[0] |
| 1145 | assert scale_a.numel() == 1 and scale_b.numel() == 1 |
| 1146 | assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn |
| 1147 | assert a.device.type == "cuda" and b.device.type == "cuda" |
| 1148 | assert scale_a.dtype == torch.float32 and scale_b.dtype == torch.float32 |
| 1149 | assert scale_a.device.type == "cuda" and scale_b.device.type == "cuda" |
| 1150 | |
| 1151 | output = bmm_fp8( |
| 1152 | a.unsqueeze(0), |
| 1153 | b.unsqueeze(0), |
| 1154 | scale_a, |
| 1155 | scale_b, |
| 1156 | out_dtype, |
| 1157 | "auto", |
| 1158 | ).view(a.shape[0], b.shape[1]) |
| 1159 | |
| 1160 | if bias is not None: |
| 1161 | output = output + bias |
| 1162 | return output |
| 1163 | |
| 1164 | |
| 1165 | def flashinfer_scaled_fp8_mm_out( |
| 1166 | a: torch.Tensor, |
| 1167 | b: torch.Tensor, |
| 1168 | scale_a: torch.Tensor, |
| 1169 | scale_b: torch.Tensor, |
| 1170 | out: torch.Tensor, |
| 1171 | out_dtype: torch.dtype | None = None, |
| 1172 | ) -> torch.Tensor: |
| 1173 | assert a.ndim == 2 and b.ndim == 2 and out.ndim == 2 |
| 1174 | assert a.shape[1] == b.shape[0] |
| 1175 | assert out.shape == (a.shape[0], b.shape[1]) |
| 1176 | assert scale_a.numel() == 1 and scale_b.numel() == 1 |
| 1177 | assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn |
| 1178 | assert out.device.type == "cuda" |
| 1179 | assert a.is_contiguous() |
| 1180 | |
| 1181 | from flashinfer import bmm_fp8 as bmm_fp8_ |
| 1182 | |
| 1183 | bmm_fp8_( |
| 1184 | a.unsqueeze(0), |
| 1185 | # FlashInfer expects the weight in the same column-major view layout |
| 1186 | # consumed by flashinfer_scaled_fp8_mm, so keep the transposed view. |
| 1187 | b.unsqueeze(0), |
| 1188 | scale_a, |
| 1189 | scale_b, |
| 1190 | out_dtype or out.dtype, |
| 1191 | out.unsqueeze(0), |
| 1192 | "auto", |
| 1193 | ) |
| 1194 | return out |
| 1195 | |
| 1196 | |
| 1197 | def flashinfer_quant_nvfp4_8x4_sf_layout( |
| 1198 | a: torch.Tensor, a_global_sf: torch.Tensor |
| 1199 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 1200 | return flashinfer_nvfp4_quantize(a, a_global_sf) |
| 1201 | |
| 1202 | |
| 1203 | flashinfer_fp8_blockscale_gemm = _lazy_import_wrapper( |
| 1204 | "flashinfer.gemm", "fp8_blockscale_gemm_sm90" |
| 1205 | ) |
| 1206 | |
| 1207 | |
| 1208 | @functools.cache |
| 1209 | def has_flashinfer_fp8_blockscale_gemm() -> bool: |
| 1210 | """Return `True` if FlashInfer block-scale FP8 GEMM is available.""" |
| 1211 | return ( |
| 1212 | has_flashinfer() |
| 1213 | and current_platform.is_device_capability(90) |
| 1214 | and hasattr(_get_submodule("flashinfer.gemm"), "fp8_blockscale_gemm_sm90") |
| 1215 | ) |
| 1216 | |
| 1217 | |
| 1218 | @functools.cache |
| 1219 | def is_flashinfer_fp8_blockscale_gemm_supported() -> bool: |
| 1220 | """Return `True` if FlashInfer block-scale FP8 GEMM is supported.""" |
| 1221 | return ( |
| 1222 | envs.VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER |
| 1223 | and has_flashinfer_fp8_blockscale_gemm() |
| 1224 | ) |
| 1225 | |
| 1226 | |
| 1227 | def should_use_flashinfer_for_blockscale_fp8_gemm( |
| 1228 | is_flashinfer_supported: bool, |
| 1229 | output_dtype: torch.dtype, |
| 1230 | input_dtype: torch.dtype, |
| 1231 | weight_dtype: torch.dtype, |
| 1232 | weight_shape: tuple[int, int], |
| 1233 | ): |
| 1234 | if not is_flashinfer_supported: |
| 1235 | return False |
| 1236 | |
| 1237 | # Verify DeepGEMM N/K dims requirements |
| 1238 | # NOTE: Also synchronized with test_w8a8_block_fp8_deep_gemm_matmul |
| 1239 | # test inside kernels/quantization/test_block_fp8.py |
| 1240 | N_MULTIPLE = 64 |
| 1241 | K_MULTIPLE = 128 |
| 1242 | |
| 1243 | should_use_flashinfer = ( |
| 1244 | output_dtype == torch.bfloat16 |
| 1245 | and input_dtype == torch.bfloat16 |
| 1246 | and weight_dtype == torch.float8_e4m3fn |
| 1247 | and weight_shape[0] % N_MULTIPLE == 0 |
| 1248 | and weight_shape[1] % K_MULTIPLE == 0 |
| 1249 | ) |
| 1250 | |
| 1251 | return should_use_flashinfer |
| 1252 | |
| 1253 | |
| 1254 | _MIN_CUDNN_FP8 = 91701 # cuDNN >= 9.17.1 required for FP8 ViT attention |
| 1255 | |
| 1256 | |
| 1257 | @functools.cache |
| 1258 | def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool: |
| 1259 | """Check if FP8 ViT attention is supported on this platform. |
| 1260 | |
| 1261 | Requires Blackwell (SM 100) or newer, the FlashInfer cuDNN backend, |
| 1262 | and cuDNN >= 9.17.1. |
| 1263 | |
| 1264 | cuDNN's FP8 SDPA forward path with bf16/fp16 output (used by |
| 1265 | ``MMEncoderAttention._forward_flashinfer``) gates internally on |
| 1266 | ``prop.major >= 10``; on Hopper it raises a misleading |
| 1267 | ``cudnnGraphNotSupportedError: ... cuDNN version 9.13.0 and newer`` |
| 1268 | even when the installed cuDNN is new enough. See PR #38065 for the |
| 1269 | original Blackwell-only design intent. |
| 1270 | """ |
| 1271 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 1272 | |
| 1273 | # cuDNN SDPA FP8 with bf16/fp16 output requires Blackwell (SM 100) or newer. |
| 1274 | if not current_platform.has_device_capability(100): |
| 1275 | return False |
| 1276 | |
| 1277 | try: |
| 1278 | supported = current_platform.get_supported_vit_attn_backends() |
| 1279 | if AttentionBackendEnum.FLASHINFER not in supported: |
| 1280 | return False |
| 1281 | except (ImportError, AttributeError): |
| 1282 | return False |
| 1283 | |
| 1284 | try: |
| 1285 | import torch.backends.cudnn as cudnn |
| 1286 | |
| 1287 | if cudnn.is_available() and cudnn.version() < _MIN_CUDNN_FP8: |
| 1288 | return False |
| 1289 | except (ImportError, AttributeError): |
| 1290 | pass |
| 1291 | |
| 1292 | return True |
| 1293 | |
| 1294 | |
| 1295 | __all__ = [ |
| 1296 | "has_flashinfer", |
| 1297 | "flashinfer_bf16_mm", |
| 1298 | "autotune_bf16_only", |
| 1299 | "has_flashinfer_bf16_gemm", |
| 1300 | "is_flashinfer_bf16_gemm_supported", |
| 1301 | "is_flashinfer_cutedsl_bf16_gemm_supported", |
| 1302 | "flashinfer_trtllm_fp8_block_scale_moe", |
| 1303 | "flashinfer_cutlass_fused_moe", |
| 1304 | "flashinfer_cutedsl_grouped_gemm_nt_masked", |
| 1305 | "flashinfer_prepare_bf16_fp4_weights", |
| 1306 | "flashinfer_fp4_quantize", |
| 1307 | "silu_and_mul_scaled_nvfp4_experts_quantize", |
| 1308 | "scaled_fp4_grouped_quantize", |
| 1309 | "nvfp4_block_scale_interleave", |
| 1310 | "flashinfer_cute_dsl_fused_moe_nvfp4", |
| 1311 | "flashinfer_b12x_fused_moe", |
| 1312 | "flashinfer_get_hybrid_num_tokens_buckets", |
| 1313 | "flashinfer_convert_sf_to_mma_layout", |
| 1314 | "trtllm_fp4_block_scale_moe", |
| 1315 | "flashinfer_trtllm_batch_decode_with_kv_cache_mla", |
| 1316 | "flashinfer_trtllm_batch_decode_sparse_mla_dsv4", |
| 1317 | "flashinfer_xqa_batch_decode_with_kv_cache", |
| 1318 | "flashinfer_recurrent_kda", |
| 1319 | "flashinfer_fused_kda_decode", |
| 1320 | "autotune", |
| 1321 | "has_flashinfer_moe", |
| 1322 | "has_flashinfer_comm", |
| 1323 | "has_flashinfer_nvlink_two_sided", |
| 1324 | "has_flashinfer_nvlink_one_sided", |
| 1325 | "has_flashinfer_cutlass_fused_moe", |
| 1326 | "has_flashinfer_cutedsl_grouped_gemm_nt_masked", |
| 1327 | "has_flashinfer_recurrent_kda", |
| 1328 | "has_flashinfer_fused_kda_decode", |
| 1329 | "has_flashinfer_cutedsl_moe_nvfp4", |
| 1330 | "has_flashinfer_bf16_fp4", |
| 1331 | "has_flashinfer_b12x_moe", |
| 1332 | "has_flashinfer_b12x_gemm", |
| 1333 | "has_flashinfer_fp8_blockscale_gemm", |
| 1334 | "has_nvidia_artifactory", |
| 1335 | "supports_trtllm_attention", |
| 1336 | "can_use_trtllm_attention", |
| 1337 | "use_trtllm_attention", |
| 1338 | "flashinfer_mxfp4_quantize", |
| 1339 | "flashinfer_scaled_fp4_mm", |
| 1340 | "flashinfer_scaled_fp4_mm_out", |
| 1341 | "flashinfer_scaled_fp8_mm", |
| 1342 | "flashinfer_scaled_fp8_mm_out", |
| 1343 | "flashinfer_quant_nvfp4_8x4_sf_layout", |
| 1344 | "flashinfer_fp8_blockscale_gemm", |
| 1345 | "should_use_flashinfer_for_blockscale_fp8_gemm", |
| 1346 | "is_flashinfer_fp8_blockscale_gemm_supported", |
| 1347 | "is_flashinfer_cudnn_fp8_prefill_attn_supported", |
| 1348 | ] |