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 8, 16:30 UTC
[Kernel] use flashinfer for gdn prefill (#32846)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Feb 10, 2026 at 01:17 UTC
285bab47526cbc4d4e26c61d831eaeb17b253d0f
1 file changed
+115
-2
vllm/model_executor/models/qwen3_next.py
+115
-2
@@ -28,11 +28,15 @@ from vllm.distributed import (
28
)
29
from vllm.forward_context import ForwardContext, get_forward_context
30
from vllm.logger import init_logger
31
+from vllm.model_executor.custom_op import CustomOp
32
from vllm.model_executor.layers.attention import Attention
33
from vllm.model_executor.layers.fla.ops import (
33
- chunk_gated_delta_rule,
34
+ chunk_gated_delta_rule as fla_chunk_gated_delta_rule,
35
+)
36
+from vllm.model_executor.layers.fla.ops import (
37
fused_recurrent_gated_delta_rule,
38
)
39
+from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd
40
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
41
from vllm.model_executor.layers.layernorm import (
42
GemmaRMSNorm as Qwen3NextRMSNorm,
@@ -101,6 +105,113 @@ logger = init_logger(__name__)
105
KVCache = tuple[torch.Tensor, torch.Tensor]
106
107
108
+def fi_chunk_gated_delta_rule(
109
+ q: torch.Tensor,
110
+ k: torch.Tensor,
111
+ v: torch.Tensor,
112
+ g: torch.Tensor,
113
+ beta: torch.Tensor,
114
+ initial_state: torch.Tensor,
115
+ output_final_state: bool,
116
+ cu_seqlens: torch.LongTensor | None = None,
117
+ head_first: bool = False,
118
+ use_qk_l2norm_in_kernel: bool = True,
119
+):
120
+ from flashinfer.gdn_prefill import (
121
+ chunk_gated_delta_rule as chunk_gated_delta_rule_fi,
122
+ )
123
+
124
+ if use_qk_l2norm_in_kernel:
125
+ q = l2norm_fwd(q)
126
+ k = l2norm_fwd(k)
127
+
128
+ # use flashinfer implementation
129
+ q = q.squeeze(0).contiguous()
130
+ k = k.squeeze(0).contiguous()
131
+ v = v.squeeze(0).contiguous()
132
+
133
+ g = g.squeeze(0).contiguous()
134
+ beta = beta.squeeze(0).contiguous()
135
+ fi_state = initial_state.to(torch.float32)
136
+ fi_g = g.to(torch.float32)
137
+ fi_beta = beta.to(torch.float32)
138
+ return chunk_gated_delta_rule_fi(
139
+ q=q,
140
+ k=k,
141
+ v=v,
142
+ g=torch.exp(fi_g),
143
+ beta=fi_beta,
144
+ initial_state=fi_state,
145
+ output_final_state=output_final_state,
146
+ cu_seqlens=cu_seqlens,
147
+ )
148
+
149
+
150
+@CustomOp.register("chunk_gated_delta_rule")
151
+class ChunkGatedDeltaRule(CustomOp):
152
+ def __init__(self) -> None:
153
+ super().__init__()
154
+ if current_platform.is_cuda() and current_platform.is_device_capability(90):
155
+ logger.info_once(
156
+ "Using FlashInfer GDN prefill kernel on CUDA compute capability 90"
157
+ )
158
+ self._forward_method = self.forward_cuda
159
+ else:
160
+ self._forward_method = self.forward_native
161
+
162
+ def forward_cuda(
163
+ self,
164
+ q: torch.Tensor,
165
+ k: torch.Tensor,
166
+ v: torch.Tensor,
167
+ g: torch.Tensor,
168
+ beta: torch.Tensor,
169
+ initial_state: torch.Tensor,
170
+ output_final_state: bool,
171
+ cu_seqlens: torch.LongTensor | None = None,
172
+ head_first: bool = False,
173
+ use_qk_l2norm_in_kernel: bool = True,
174
+ ):
175
+ return fi_chunk_gated_delta_rule(
176
+ q=q,
177
+ k=k,
178
+ v=v,
179
+ g=g,
180
+ beta=beta,
181
+ initial_state=initial_state,
182
+ output_final_state=output_final_state,
183
+ cu_seqlens=cu_seqlens,
184
+ head_first=head_first,
185
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
186
+ )
187
+
188
+ def forward_native(
189
+ self,
190
+ q: torch.Tensor,
191
+ k: torch.Tensor,
192
+ v: torch.Tensor,
193
+ g: torch.Tensor,
194
+ beta: torch.Tensor,
195
+ initial_state: torch.Tensor,
196
+ output_final_state: bool,
197
+ cu_seqlens: torch.LongTensor | None = None,
198
+ head_first: bool = False,
199
+ use_qk_l2norm_in_kernel: bool = True,
200
+ ):
201
+ return fla_chunk_gated_delta_rule(
202
+ q=q,
203
+ k=k,
204
+ v=v,
205
+ g=g,
206
+ beta=beta,
207
+ initial_state=initial_state,
208
+ output_final_state=output_final_state,
209
+ cu_seqlens=cu_seqlens,
210
+ head_first=head_first,
211
+ use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
212
+ )
213
+
214
+
215
class Qwen3NextSparseMoeBlock(nn.Module):
216
def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
217
super().__init__()
@@ -362,6 +473,8 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
473
prefix=f"{prefix}.out_proj",
474
)
475
476
+ self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
477
+
478
compilation_config = get_current_vllm_config().compilation_config
479
if prefix in compilation_config.static_forward_context:
480
raise ValueError(f"Duplicate layer name: {prefix}")
@@ -647,7 +760,7 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
760
(
761
core_attn_out_non_spec,
762
last_recurrent_state,
650
- ) = chunk_gated_delta_rule(
763
+ ) = self.chunk_gated_delta_rule(
764
q=query_non_spec,
765
k=key_non_spec,
766
v=value_non_spec,