From f6326f53bda46898a331c2d24500332c285d9a2b Mon Sep 17 00:00:00 2001 From: Stefano Castagnetta Date: Tue, 8 Sep 2026 14:26:03 +0200 Subject: [PATCH] [Perf][GDN] Enable the FlashInfer GDN prefill kernel on SM12x (#55715) Signed-off-by: Stefano Castagnetta --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index b7bd5b7b0b..bbbb1c8bd5 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -100,8 +100,9 @@ def _resolve_gdn_prefill_backend( * ``requested in ["flashinfer", "auto"]``; * ``platform == cuda``; * one of the following: - - Hopper (SM90) — no further constraints; - - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``. + - Hopper (SM90) - no further constraints; + - Blackwell (SM10.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``; + - Blackwell (SM12.x) with ``head_k_dim == 128``, ``cuda_runtime >= 13``. In-tree CuteDSL GDN prefill kernel is chosen when: * "cutedsl" is requested; (opt-in only) @@ -134,6 +135,13 @@ def _resolve_gdn_prefill_backend( ): supports_flashinfer = True supports_cutedsl = True + elif ( + current_platform.is_device_capability_family(120) + and head_k_dim == 128 + and current_platform.get_cuda_runtime_major() >= 13 + ): + # The in-tree CuteDSL kernel targets SM100 only, so it stays off here. + supports_flashinfer = True if backend in ["flashinfer", "auto"] and supports_flashinfer: return backend, "flashinfer"