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, 05:30 UTC
[GDN] Enable FI Blackwell GDN prefill kernel (#40717)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
May 20, 2026 at 11:46 UTC
1cb224430bea0d037b57e24cf91001f47b69ddf3
2 files changed
+58
-27
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+52
-27
@@ -83,6 +83,49 @@ if GDN_AITER_TRITON_AVAILABLE:
83
logger = init_logger(__name__)
84
85
86
+def _should_use_flashinfer_gdn_prefill(backend: str, head_k_dim: int | None) -> bool:
87
+ """Whether to use FlashInfer's GDN prefill kernel instead of the
88
+ Triton/FLA fallback.
89
+
90
+ Requirements:
91
+ * ``requested in ["flashinfer", "auto"]``;
92
+ * ``platform == cuda``;
93
+ * one of the following:
94
+ - Hopper (SM90) — no further constraints;
95
+ - Blackwell (SM10.x) with ``head_k_dim == 128`` and ``cuda_runtime >= 13``.
96
+ """
97
+ if backend not in ["flashinfer", "auto"]:
98
+ return False
99
+ if not current_platform.is_cuda():
100
+ return False
101
+ if current_platform.is_device_capability(90):
102
+ return True # Hopper — no further constraints.
103
+ if not current_platform.is_device_capability_family(100):
104
+ return False # Neither Hopper nor Blackwell.
105
+ if head_k_dim != 128:
106
+ return False
107
+ return current_platform.get_cuda_runtime_major() >= 13
108
+
109
+
110
+def _log_gdn_backend_decision(
111
+ backend: str, head_k_dim: int | None, use_flashinfer: bool
112
+) -> None:
113
+ """Log the GDN prefill backend choice in the attention-selector style."""
114
+ chosen = "FlashInfer" if use_flashinfer else "Triton/FLA"
115
+ logger.info_once(
116
+ "Using %s GDN prefill kernel (requested=%s, head_k_dim=%s).",
117
+ chosen,
118
+ backend,
119
+ head_k_dim,
120
+ )
121
+ # JIT-compiled cutlass path is only used on SM90 (Hopper).
122
+ if use_flashinfer and current_platform.is_device_capability(90):
123
+ logger.warning_once(
124
+ "FlashInfer GDN prefill is JIT-compiled; first run may take a "
125
+ "while. Set --gdn-prefill-backend triton to skip JIT.",
126
+ )
127
+
128
+
129
def fi_chunk_gated_delta_rule(
130
q: torch.Tensor,
131
k: torch.Tensor,
@@ -134,39 +177,21 @@ def fi_chunk_gated_delta_rule(
177
178
@CustomOp.register("chunk_gated_delta_rule")
179
class ChunkGatedDeltaRule(CustomOp):
137
- def __init__(self) -> None:
180
+ def __init__(self, head_k_dim: int | None = None) -> None:
181
super().__init__()
182
additional_config = get_current_vllm_config().additional_config
183
assert isinstance(additional_config, dict)
184
backend_cfg = additional_config.get("gdn_prefill_backend", "auto")
185
backend = str(backend_cfg).strip().lower()
186
144
- supports_flashinfer = (
145
- current_platform.is_cuda() and current_platform.is_device_capability(90)
146
- )
147
-
148
- if backend == "flashinfer":
149
- use_flashinfer = supports_flashinfer
150
- if not use_flashinfer:
151
- logger.warning_once(
152
- "GDN prefill backend 'flashinfer' is selected but "
153
- "cannot use this kernel on the current platform. "
154
- "Falling back to Triton/FLA."
155
- )
156
- elif backend == "triton":
157
- use_flashinfer = False
158
- else:
159
- use_flashinfer = supports_flashinfer
160
-
161
- if use_flashinfer:
162
- logger.info_once("Using FlashInfer GDN prefill kernel")
163
- logger.info_once(
164
- "FlashInfer GDN prefill kernel is JIT-compiled; first run may "
165
- "take a while to compile. Set `--gdn-prefill-backend triton` to "
166
- "avoid JIT compile time.",
187
+ use_flashinfer = _should_use_flashinfer_gdn_prefill(backend, head_k_dim)
188
+ if backend == "flashinfer" and not use_flashinfer:
189
+ logger.warning_once(
190
+ "GDN prefill backend 'flashinfer' is selected but "
191
+ "cannot use this kernel on the current platform. "
192
+ "Falling back to Triton/FLA."
193
)
168
- else:
169
- logger.info_once("Using Triton/FLA GDN prefill kernel")
194
+ _log_gdn_backend_decision(backend, head_k_dim, use_flashinfer)
195
196
self._forward_method = (
197
self.forward_cuda if use_flashinfer else self.forward_native
@@ -396,7 +421,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
421
prefix=f"{prefix}.out_proj",
422
)
423
399
- self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
424
+ self.chunk_gated_delta_rule = ChunkGatedDeltaRule(head_k_dim=self.head_k_dim)
425
self._prefill_kernels_warmed_up = False
426
self.enable_packed_recurrent_decode = (
427
envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
vllm/platforms/cuda.py
+6
@@ -202,6 +202,12 @@ class CudaPlatformBase(Platform):
202
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
203
raise NotImplementedError
204
205
+ @classmethod
206
+ def get_cuda_runtime_major(cls) -> int:
207
+ """Major ``torch.version.cuda`` version, or ``0`` if undetermined."""
208
+ major = (torch.version.cuda or "0").split(".", 1)[0]
209
+ return int(major) if major.isdigit() else 0
210
+
211
@classmethod
212
def get_device_name(cls, device_id: int = 0) -> str:
213
raise NotImplementedError