Bump flashinfer version to 0.6.14 (#47669)

Signed-off-by: AmeenP <[email protected]>
Signed-off-by: AmeenP <[email protected]>
Co-authored-by: AmeenP <[email protected]>
Co-authored-by: Jee Jee Li <[email protected]>
Co-authored-by: Claude <[email protected]>
Co-authored-by: Pavani Majety <[email protected]>

docker/Dockerfile

@@ -815,7 +815,7 @@ RUN --mount=type=cache,target=/opt/uv/cache \
 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)
 # https://docs.flashinfer.ai/installation.html
 # From versions.json: .flashinfer.version
-ARG FLASHINFER_VERSION=0.6.13
+ARG FLASHINFER_VERSION=0.6.14
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \
         --index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')

docker/versions.json

@@ -68,7 +68,7 @@
       "default": "true"
     },
     "FLASHINFER_VERSION": {
-      "default": "0.6.13"
+      "default": "0.6.14"
     },
     "GDRCOPY_CUDA_VERSION": {
       "default": "12.8"

requirements/cuda.txt

@@ -11,8 +11,11 @@ torchvision==0.26.0 # Required for phi3v processor. See https://github.com/pytor
 torchcodec >= 0.14
 PyNvVideoCodec==2.0.4
 # FlashInfer should be updated together with the Dockerfile
-flashinfer-python==0.6.13
-flashinfer-cubin==0.6.13
+# flashinfer-cubin is not on PyPI since 0.6.14; setup.py excludes it from
+# install_requires so the published wheel does not carry an unresolvable pin
+--extra-index-url https://flashinfer.ai/whl/
+flashinfer-python==0.6.14
+flashinfer-cubin==0.6.14
 apache-tvm-ffi==0.1.9
 tilelang==0.1.9
 nvidia-cudnn-frontend>=1.19.1

setup.py

@@ -1075,6 +1075,11 @@ def get_requirements() -> list[str]:
                 # vllm-flash-attn is built only for CUDA 12.x.
                 # Skip for other versions.
                 continue
+            if "flashinfer-cubin" in req:
+                # Not on PyPI since 0.6.14 (only https://flashinfer.ai/whl), so
+                # it cannot be a wheel dependency; flashinfer falls back to
+                # fetching cubins at runtime when the package is absent.
+                continue
             if "nvidia-cutlass-dsl[cu13]" in req and cuda_major == "12":
                 # [cu13] extra is the default; strip it on CUDA 12 builds.
                 req = req.replace("nvidia-cutlass-dsl[cu13]", "nvidia-cutlass-dsl")

tests/test_jit_monitor.py

@@ -320,6 +320,26 @@ class TestCuTeDSLHook:
             with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"):
                 cute.compile(lambda: None, "arg", option=True)
 
+    def test_subscripted_compile_is_monitored(self):
+        """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work."""
+
+        class FakeCompileCallable:
+            def __getitem__(self, options):
+                return self
+
+            def __call__(self, *args, **kwargs):
+                return "compiled"
+
+        with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()):
+            import cutlass.cute as cute
+
+            jit_monitor.activate()
+            with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
+                result = cute.compile[("opt_level", 3)](lambda: None, "arg")
+
+        assert result == "compiled"
+        warning_once.assert_called_once()
+
 
 class TestTileLangHook:
     def test_jit_kernel_logs_warning(self):

vllm/utils/jit_monitor.py

@@ -267,6 +267,26 @@ def _log_cutedsl_jit_compile(fn_name: str) -> None:
     )
 
 
+class _MonitoredCuteCompile:
+    """Logs JIT compilations; a plain function would break ``cute.compile[opts]``."""
+
+    def __init__(self, inner):
+        self._inner = inner
+
+    def __getitem__(self, options) -> "_MonitoredCuteCompile":
+        return _MonitoredCuteCompile(self._inner[options])
+
+    def __call__(self, *args, **kwargs):
+        kernel = args[0] if args else kwargs.get("function")
+        kernel_name = getattr(kernel, "__name__", None)
+        if kernel_name is None:
+            kernel_name = (
+                kernel.__class__.__name__ if kernel is not None else "<unknown>"
+            )
+        _log_cutedsl_jit_compile(kernel_name)
+        return self._inner(*args, **kwargs)
+
+
 def _setup_cutedsl_jit_hook() -> None:
     """Wrap ``cutlass.cute.compile`` to warn on compilation."""
     global _cutedsl_hook_installed
@@ -279,20 +299,7 @@ def _setup_cutedsl_jit_hook() -> None:
         logger.debug("CuTeDSL is not available; skipping CuTeDSL JIT monitor.")
         return
 
-    original_compile = cute.compile
-
-    @functools.wraps(original_compile)
-    def _compile_with_monitor(*args, **kwargs):
-        kernel = args[0] if args else kwargs.get("function")
-        kernel_name = getattr(kernel, "__name__", None)
-        if kernel_name is None:
-            kernel_name = (
-                kernel.__class__.__name__ if kernel is not None else "<unknown>"
-            )
-        _log_cutedsl_jit_compile(kernel_name)
-        return original_compile(*args, **kwargs)
-
-    cute.compile = _compile_with_monitor
+    cute.compile = _MonitoredCuteCompile(cute.compile)
     _cutedsl_hook_installed = True