Merge branch 'main' into wentao-fix-flashinfer-layout

Signed-off-by: yewentao256 <zhyanwentao@126.com>

vllm/_custom_ops.py

@@ -12,6 +12,7 @@ from vllm.platforms import current_platform
 from vllm.scalar_type import ScalarType
 from vllm.utils.flashinfer import (
     flashinfer_quant_nvfp4_8x4_sf_layout,
+    flashinfer_trtllm_fp4_8x4_is_safe,
 )
 from vllm.utils.math_utils import cdiv
 
@@ -1642,7 +1643,9 @@ def scaled_fp4_quant(
             f"padded_n has to be a multiple of {block_size}, but got {padded_n}."
         )
 
-    use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False  # noqa: SIM210
+    use_8x4_sf_layout = (
+        "trtllm" in backend and m <= 32 and flashinfer_trtllm_fp4_8x4_is_safe()
+    )
     if use_8x4_sf_layout and padded_n is not None and padded_n != n:
         # TODO: support this case
         raise ValueError("padded_n is not supported with TRTLLM 8x4 scale layout.")

vllm/utils/flashinfer.py

@@ -8,6 +8,7 @@ Users of vLLM should always import **only** these wrappers.
 import contextlib
 import functools
 import importlib
+import importlib.metadata
 import importlib.util
 import os
 import shutil
@@ -16,6 +17,7 @@ from typing import Any, NoReturn
 
 import requests
 import torch
+from packaging.version import Version
 
 import vllm.envs as envs
 from vllm.logger import init_logger
@@ -63,6 +65,43 @@ def has_flashinfer() -> bool:
     return True
 
 
+@functools.cache
+def get_flashinfer_version() -> Version | None:
+    if not has_flashinfer():
+        return None
+
+    return Version(importlib.metadata.version("flashinfer-python"))
+
+
+@functools.cache
+def flashinfer_trtllm_fp4_8x4_is_safe() -> bool:
+    """Return whether FlashInfer TRTLLM FP4 8x4 is safe to use.
+
+    FlashInfer documents TRTLLM FP4 as allowing 8x4 only for the activation
+    (A) scale layout while keeping the weight (B) path in 128x4. However,
+    released FlashInfer versions up to and including 0.6.11.* are known to
+    have correctness bugs on this path; see flashinfer-ai/flashinfer#2861.
+    """
+    version = get_flashinfer_version()
+    if version is None:
+        logger.warning_once(
+            "Disabling FlashInfer TRTLLM FP4 8x4 path because the installed "
+            "FlashInfer version could not be determined."
+        )
+        return False
+
+    if Version(version.base_version) <= Version("0.6.11"):
+        logger.warning_once(
+            "Disabling FlashInfer TRTLLM FP4 8x4 path for flashinfer==%s due "
+            "to known upstream correctness bugs in released versions up to "
+            "0.6.11.* (see flashinfer-ai/flashinfer#2861).",
+            version,
+        )
+        return False
+
+    return True
+
+
 def _missing(*_: Any, **__: Any) -> NoReturn:
     """Placeholder for unavailable FlashInfer backend."""
     raise RuntimeError(
@@ -742,7 +781,9 @@ def flashinfer_scaled_fp4_mm(
         block_scale_a = block_scale_a.view(torch.uint8)
         block_scale_b = block_scale_b.view(torch.uint8)
 
-    use_8x4_sf_layout = True if backend == "trtllm" and a.shape[0] <= 32 else False  # noqa: SIM210
+    use_8x4_sf_layout = (
+        backend == "trtllm" and a.shape[0] <= 32 and flashinfer_trtllm_fp4_8x4_is_safe()
+    )
 
     return flashinfer_mm_fp4(
         a,
@@ -954,6 +995,7 @@ def is_flashinfer_cudnn_fp8_prefill_attn_supported() -> bool:
 
 __all__ = [
     "has_flashinfer",
+    "get_flashinfer_version",
     "flashinfer_trtllm_fp8_block_scale_moe",
     "flashinfer_cutlass_fused_moe",
     "flashinfer_cutedsl_grouped_gemm_nt_masked",
@@ -982,6 +1024,7 @@ __all__ = [
     "use_trtllm_attention",
     "flashinfer_mxfp4_quantize",
     "flashinfer_scaled_fp4_mm",
+    "flashinfer_trtllm_fp4_8x4_is_safe",
     "flashinfer_scaled_fp4_mm_out",
     "flashinfer_scaled_fp8_mm",
     "flashinfer_scaled_fp8_mm_out",