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 3, 00:30 UTC
[Bugfix] Detect wrong libcute_dsl_runtime.so variant in FlashInfer GDN (#43427)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
May 22, 2026 at 22:33 UTC
23f7b11bf4b715ee6084ea368c929789184a8555
1 file changed
+84
-2
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+84
-2
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Inference-only Qwen3-Next/Qwen3.5 model."""
4
5
+import functools
6
+
7
import torch
8
from einops import rearrange
9
from torch import nn
@@ -80,6 +82,70 @@ if GDN_AITER_TRITON_AVAILABLE:
82
logger = init_logger(__name__)
83
84
85
+# TODO(arpera): remove ``_is_libs_cu13_install_intact`` and its caller in
86
+# ``_should_use_flashinfer_gdn_prefill`` once the upstream packaging bug is
87
+# fixed and the broken wheels are yanked / superseded on PyPI:
88
+# https://github.com/NVIDIA/cutlass/issues/3170
89
+# https://github.com/NVIDIA/cutlass/issues/3259
90
+@functools.cache
91
+def _is_libs_cu13_install_intact() -> bool:
92
+ """Return True if every file installed by ``nvidia-cutlass-dsl-libs-cu13``
93
+ matches the SHA-256 declared in its wheel ``RECORD``.
94
+
95
+ ``nvidia-cutlass-dsl-libs-base`` and ``nvidia-cutlass-dsl-libs-cu13``
96
+ both ship into the shared ``nvidia_cutlass_dsl/`` namespace and
97
+ write many of the same on-disk paths (the runtime ``.so``, the MLIR
98
+ Python bindings, cuTe-DSL Python sources, ...) with different
99
+ content. Whichever wheel extracts last wins; with a parallel
100
+ installer (e.g. ``uv``) the order is racy and the resulting venv
101
+ can end up with a mix of files from both variants. The
102
+ ``-libs-base`` variant fails MLIR legalization when JIT-compiling
103
+ the FlashInfer Blackwell GDN prefill kernel, and any other
104
+ cuTe-DSL-based kernel can break too if on-disk files diverge from
105
+ what ``-libs-cu13``'s wheel expects. Tracked upstream at:
106
+
107
+ * https://github.com/NVIDIA/cutlass/issues/3170
108
+ * https://github.com/NVIDIA/cutlass/issues/3259
109
+
110
+ This helper re-hashes every file the ``-libs-cu13`` wheel claims to
111
+ own and compares against its declared SHA-256. Returns False on any
112
+ error (uninstalled, missing RECORD, missing file, hash mismatch).
113
+ Result is cached per-process.
114
+ """
115
+ import hashlib
116
+ import importlib.metadata
117
+
118
+ import pybase64 as base64
119
+
120
+ try:
121
+ dist = importlib.metadata.distribution("nvidia-cutlass-dsl-libs-cu13")
122
+ except importlib.metadata.PackageNotFoundError:
123
+ return False
124
+
125
+ files = dist.files
126
+ if not files:
127
+ return False
128
+
129
+ for pkg_path in files:
130
+ file_hash = pkg_path.hash
131
+ # Skip RECORD rows without a hash (RECORD itself, generated
132
+ # ``.pyc`` files, ...) and any non-SHA-256 hash modes.
133
+ if file_hash is None or not file_hash.value:
134
+ continue
135
+ if file_hash.mode != "sha256":
136
+ continue
137
+ try:
138
+ with open(pkg_path.locate(), "rb") as f:
139
+ digest = hashlib.sha256(f.read()).digest()
140
+ except OSError:
141
+ return False
142
+ actual = base64.urlsafe_b64encode(digest).decode().rstrip("=")
143
+ if actual != file_hash.value:
144
+ return False
145
+
146
+ return True
147
+
148
+
149
def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) -> bool:
150
"""Whether to use FlashInfer's GDN prefill kernel instead of the
151
Triton/FLA fallback.
@@ -89,7 +155,9 @@ def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) ->
155
* ``platform == cuda``;
156
* one of the following:
157
- Hopper (SM90) — no further constraints;
92
- - Blackwell (SM10.x) with ``head_k_dim == 128`` and ``cuda_runtime >= 13``.
158
+ - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``,
159
+ and an intact ``nvidia-cutlass-dsl-libs-cu13`` install on disk
160
+ (see :func:`_is_libs_cu13_install_intact`).
161
"""
162
if backend not in ["flashinfer", "auto"]:
163
return False
@@ -101,7 +169,21 @@ def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) ->
169
return False # Neither Hopper nor Blackwell.
170
if head_k_dim != 128:
171
return False
104
- return current_platform.get_cuda_runtime_major() >= 13
172
+ if current_platform.get_cuda_runtime_major() < 13:
173
+ return False
174
+ if not _is_libs_cu13_install_intact():
175
+ logger.warning_once(
176
+ "FlashInfer Blackwell GDN requires an intact nvidia-cutlass-dsl"
177
+ "-libs-cu13 install, but some on-disk files do not match the "
178
+ "SHA-256 declared in its RECORD (install-order race in "
179
+ "nvidia-cutlass-dsl packaging — see "
180
+ "https://github.com/NVIDIA/cutlass/issues/3170 and "
181
+ "https://github.com/NVIDIA/cutlass/issues/3259). Falling back "
182
+ "to Triton/FLA. Repair with: pip install --force-reinstall "
183
+ "--no-deps nvidia-cutlass-dsl-libs-cu13"
184
+ )
185
+ return False
186
+ return True
187
188
189
def _log_gdn_backend_decision(