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 6, 17:00 UTC
[GDN] add a config for gdn kernel selection (#36647)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Jiangyun Zhu committed
Mar 16, 2026 at 00:40 UTC
697e4ff3528c72806a4d00ed9b7581332b9efd43
2 files changed
+47
-4
vllm/engine/arg_utils.py
+11
@@ -614,6 +614,7 @@ class EngineArgs:
614
)
615
616
fail_on_environ_validation: bool = False
617
+ gdn_prefill_backend: Literal["flashinfer", "triton"] | None = None
618
619
def __post_init__(self):
620
# support `EngineArgs(compilation_config={...})`
@@ -1318,6 +1319,13 @@ class EngineArgs:
1319
help="Shutdown timeout in seconds. 0 = abort, >0 = wait.",
1320
)
1321
1322
+ parser.add_argument(
1323
+ "--gdn-prefill-backend",
1324
+ dest="gdn_prefill_backend",
1325
+ choices=["flashinfer", "triton"],
1326
+ default=None,
1327
+ help="Select GDN prefill backend.",
1328
+ )
1329
return parser
1330
1331
@classmethod
@@ -1903,6 +1911,9 @@ class EngineArgs:
1911
),
1912
)
1913
1914
+ if self.gdn_prefill_backend is not None:
1915
+ self.additional_config["gdn_prefill_backend"] = self.gdn_prefill_backend
1916
+
1917
config = VllmConfig(
1918
model_config=model_config,
1919
cache_config=cache_config,
vllm/model_executor/models/qwen3_next.py
+36
-4
@@ -161,13 +161,45 @@ def fi_chunk_gated_delta_rule(
161
class ChunkGatedDeltaRule(CustomOp):
162
def __init__(self) -> None:
163
super().__init__()
164
- if current_platform.is_cuda() and current_platform.is_device_capability(90):
164
+ backend = (
165
+ str(
166
+ get_current_vllm_config().additional_config.get(
167
+ "gdn_prefill_backend", "auto"
168
+ )
169
+ )
170
+ .strip()
171
+ .lower()
172
+ )
173
+ supports_flashinfer = (
174
+ current_platform.is_cuda() and current_platform.is_device_capability(90)
175
+ )
176
+
177
+ if backend == "flashinfer":
178
+ use_flashinfer = supports_flashinfer
179
+ if not use_flashinfer:
180
+ logger.warning_once(
181
+ "GDN prefill backend 'flashinfer' is selected but "
182
+ "cannot use this kernel on the current platform. "
183
+ "Falling back to Triton/FLA."
184
+ )
185
+ elif backend == "triton":
186
+ use_flashinfer = False
187
+ else:
188
+ use_flashinfer = supports_flashinfer
189
+
190
+ if use_flashinfer:
191
+ logger.info_once("Using FlashInfer GDN prefill kernel")
192
logger.info_once(
166
- "Using FlashInfer GDN prefill kernel on CUDA compute capability 90"
193
+ "FlashInfer GDN prefill kernel is JIT-compiled; first run may "
194
+ "take a while to compile. Set `--gdn-prefill-backend triton` to "
195
+ "avoid JIT compile time."
196
)
168
- self._forward_method = self.forward_cuda
197
else:
170
- self._forward_method = self.forward_native
198
+ logger.info_once("Using Triton/FLA GDN prefill kernel")
199
+
200
+ self._forward_method = (
201
+ self.forward_cuda if use_flashinfer else self.forward_native
202
+ )
203
204
def forward_cuda(
205
self,