[Perf][GDN] Enable the FlashInfer GDN prefill kernel on SM12x (#55715)
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
Stefano Castagnetta committed
Sep 8, 2026 at 14:26 UTC
f6326f53bda46898a331c2d24500332c285d9a2b
1 file changed
+10
-2
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+10
-2
@@ -100,8 +100,9 @@ def _resolve_gdn_prefill_backend(
100
* ``requested in ["flashinfer", "auto"]``;
101
* ``platform == cuda``;
102
* one of the following:
103
- - Hopper (SM90) — no further constraints;
104
- - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``.
103
+ - Hopper (SM90) - no further constraints;
104
+ - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``;
105
+ - Blackwell (SM12.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``.
106
107
In-tree CuteDSL GDN prefill kernel is chosen when:
108
* "cutedsl" is requested; (opt-in only)
@@ -134,6 +135,13 @@ def _resolve_gdn_prefill_backend(
135
):
136
supports_flashinfer = True
137
supports_cutedsl = True
138
+ elif (
139
+ current_platform.is_device_capability_family(120)
140
+ and head_k_dim == 128
141
+ and current_platform.get_cuda_runtime_major() >= 13
142
+ ):
143
+ # The in-tree CuteDSL kernel targets SM100 only, so it stays off here.
144
+ supports_flashinfer = True
145
146
if backend in ["flashinfer", "auto"] and supports_flashinfer:
147
return backend, "flashinfer"