diff --git a/docker/Dockerfile b/docker/Dockerfile index b47853a06c7..1263bff436f 100644 --- a/docker/Dockerfile +++ b/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 '.') diff --git a/docker/versions.json b/docker/versions.json index 4dffa00985c..e6839bbb05c 100644 --- a/docker/versions.json +++ b/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" diff --git a/requirements/cuda.txt b/requirements/cuda.txt index 91a57997684..ab1df3ffe60 100644 --- a/requirements/cuda.txt +++ b/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 diff --git a/setup.py b/setup.py index e8f52970184..6419e3073ea 100644 --- a/setup.py +++ b/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") diff --git a/tests/test_jit_monitor.py b/tests/test_jit_monitor.py index 9f3285ddec0..50261a479d9 100644 --- a/tests/test_jit_monitor.py +++ b/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): diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index 8228e24c1c3..4c3327283a3 100644 --- a/vllm/utils/jit_monitor.py +++ b/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 "" + ) + _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 "" - ) - _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