@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a8eab8f30d
Fork 0

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 5, 22:15 UTC

[Model] Extract GatedDeltaNetAttention into shared layer for Qwen3Next and Qwen3.5 (#37975)

Signed-off-by: wxsIcey <1790571317@qq.com> Signed-off-by: Icey <1790571317@qq.com>

Xiaoshuang Wang committed Mar 27, 2026 at 14:13 UTC a8eab8f30dda380a67c515c743365840b4347ccd
3 files changed +1053 -1126
vllm/model_executor/layers/mamba/gdn_linear_attn.py new
+1046
@@ -0,0 +1,1046 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Inference-only Qwen3-Next/Qwen3.5 model."""
4 +
5 +import torch
6 +from einops import rearrange
7 +from torch import nn
8 +from transformers.activations import ACT2FN
9 +
10 +from vllm import envs
11 +from vllm.config import (
12 + VllmConfig,
13 + get_current_vllm_config,
14 +)
15 +from vllm.distributed import (
16 + divide,
17 + get_tensor_model_parallel_rank,
18 + get_tensor_model_parallel_world_size,
19 +)
20 +from vllm.forward_context import ForwardContext, get_forward_context
21 +from vllm.logger import init_logger
22 +from vllm.model_executor.custom_op import CustomOp, PluggableLayer
23 +from vllm.model_executor.layers.fla.ops import (
24 + chunk_gated_delta_rule as fla_chunk_gated_delta_rule,
25 +)
26 +from vllm.model_executor.layers.fla.ops import (
27 + fused_recurrent_gated_delta_rule_packed_decode,
28 + fused_sigmoid_gating_delta_rule_update,
29 +)
30 +from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd
31 +from vllm.model_executor.layers.layernorm import RMSNormGated
32 +from vllm.model_executor.layers.linear import (
33 + ColumnParallelLinear,
34 + MergedColumnParallelLinear,
35 + RowParallelLinear,
36 +)
37 +from vllm.model_executor.layers.mamba.abstract import MambaBase
38 +from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weight_loader
39 +from vllm.model_executor.layers.mamba.mamba_utils import (
40 + MambaStateDtypeCalculator,
41 + MambaStateShapeCalculator,
42 +)
43 +from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
44 + causal_conv1d_fn,
45 + causal_conv1d_update,
46 +)
47 +from vllm.model_executor.layers.quantization import QuantizationConfig
48 +from vllm.model_executor.model_loader.weight_utils import (
49 + sharded_weight_loader,
50 +)
51 +from vllm.model_executor.models.utils import extract_layer_index
52 +from vllm.model_executor.utils import set_weight_attrs
53 +from vllm.platforms import current_platform
54 +from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
55 +from vllm.triton_utils import tl, triton
56 +from vllm.utils.torch_utils import direct_register_custom_op
57 +from vllm.v1.attention.backend import AttentionMetadata
58 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
59 +
60 +logger = init_logger(__name__)
61 +
62 +
63 +def fi_chunk_gated_delta_rule(
64 + q: torch.Tensor,
65 + k: torch.Tensor,
66 + v: torch.Tensor,
67 + g: torch.Tensor,
68 + beta: torch.Tensor,
69 + initial_state: torch.Tensor,
70 + output_final_state: bool,
71 + cu_seqlens: torch.Tensor | None = None,
72 + use_qk_l2norm_in_kernel: bool = True,
73 +):
74 + from flashinfer.gdn_prefill import (
75 + chunk_gated_delta_rule as chunk_gated_delta_rule_fi,
76 + )
77 +
78 + if use_qk_l2norm_in_kernel:
79 + q = l2norm_fwd(q)
80 + k = l2norm_fwd(k)
81 +
82 + # use flashinfer implementation
83 + q = q.squeeze(0).contiguous()
84 + k = k.squeeze(0).contiguous()
85 + v = v.squeeze(0).contiguous()
86 +
87 + g = g.squeeze(0).contiguous()
88 + beta = beta.squeeze(0).contiguous()
89 + fi_state = initial_state.to(torch.float32)
90 + fi_g = g.to(torch.float32)
91 + fi_beta = beta.to(torch.float32)
92 + result = chunk_gated_delta_rule_fi(
93 + q=q,
94 + k=k,
95 + v=v,
96 + g=torch.exp(fi_g),
97 + beta=fi_beta,
98 + initial_state=fi_state,
99 + output_final_state=output_final_state,
100 + cu_seqlens=cu_seqlens,
101 + )
102 + # FlashInfer returns (output, state) when output_final_state=True,
103 + # or just output when output_final_state=False.
104 + # Unsqueeze back to 4D (1, L, H, D) to match fla output format
105 + if output_final_state:
106 + output, final_state = result
107 + return output.unsqueeze(0), final_state
108 + else:
109 + return result.unsqueeze(0), None
110 +
111 +
112 +@CustomOp.register("chunk_gated_delta_rule")
113 +class ChunkGatedDeltaRule(CustomOp):
114 + def __init__(self) -> None:
115 + super().__init__()
116 + backend_cfg = get_current_vllm_config().additional_config.get(
117 + "gdn_prefill_backend", "auto"
118 + )
119 + backend = str(backend_cfg).strip().lower()
120 +
121 + supports_flashinfer = (
122 + current_platform.is_cuda() and current_platform.is_device_capability(90)
123 + )
124 +
125 + if backend == "flashinfer":
126 + use_flashinfer = supports_flashinfer
127 + if not use_flashinfer:
128 + logger.warning_once(
129 + "GDN prefill backend 'flashinfer' is selected but "
130 + "cannot use this kernel on the current platform. "
131 + "Falling back to Triton/FLA."
132 + )
133 + elif backend == "triton":
134 + use_flashinfer = False
135 + else:
136 + use_flashinfer = supports_flashinfer
137 +
138 + if use_flashinfer:
139 + logger.info_once("Using FlashInfer GDN prefill kernel", scope="local")
140 + logger.info_once(
141 + "FlashInfer GDN prefill kernel is JIT-compiled; first run may "
142 + "take a while to compile. Set `--gdn-prefill-backend triton` to "
143 + "avoid JIT compile time.",
144 + scope="local",
145 + )
146 + else:
147 + logger.info_once("Using Triton/FLA GDN prefill kernel", scope="local")
148 +
149 + self._forward_method = (
150 + self.forward_cuda if use_flashinfer else self.forward_native
151 + )
152 +
153 + def forward_cuda(
154 + self,
155 + q: torch.Tensor,
156 + k: torch.Tensor,
157 + v: torch.Tensor,
158 + g: torch.Tensor,
159 + beta: torch.Tensor,
160 + initial_state: torch.Tensor,
161 + output_final_state: bool,
162 + cu_seqlens: torch.Tensor | None = None,
163 + use_qk_l2norm_in_kernel: bool = True,
164 + ):
165 + return fi_chunk_gated_delta_rule(
166 + q=q,
167 + k=k,
168 + v=v,
169 + g=g,
170 + beta=beta,
171 + initial_state=initial_state,
172 + output_final_state=output_final_state,
173 + cu_seqlens=cu_seqlens,
174 + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
175 + )
176 +
177 + def forward_native(
178 + self,
179 + q: torch.Tensor,
180 + k: torch.Tensor,
181 + v: torch.Tensor,
182 + g: torch.Tensor,
183 + beta: torch.Tensor,
184 + initial_state: torch.Tensor,
185 + output_final_state: bool,
186 + cu_seqlens: torch.Tensor | None = None,
187 + use_qk_l2norm_in_kernel: bool = True,
188 + ):
189 + return fla_chunk_gated_delta_rule(
190 + q=q,
191 + k=k,
192 + v=v,
193 + g=g,
194 + beta=beta,
195 + initial_state=initial_state,
196 + output_final_state=output_final_state,
197 + cu_seqlens=cu_seqlens,
198 + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
199 + )
200 +
201 +
202 +@PluggableLayer.register("gated_delta_net_attention")
203 +class GatedDeltaNetAttention(PluggableLayer, MambaBase):
204 + @property
205 + def mamba_type(self) -> str:
206 + return "gdn_attention"
207 +
208 + def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
209 + return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
210 + self.model_config.dtype,
211 + self.cache_config.mamba_cache_dtype,
212 + self.cache_config.mamba_ssm_cache_dtype,
213 + )
214 +
215 + def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
216 + return MambaStateShapeCalculator.gated_delta_net_state_shape(
217 + self.tp_size,
218 + self.num_k_heads,
219 + self.num_v_heads,
220 + self.head_k_dim,
221 + self.head_v_dim,
222 + self.conv_kernel_size,
223 + self.num_spec,
224 + )
225 +
226 + def __init__(
227 + self,
228 + config: Qwen3NextConfig,
229 + vllm_config: VllmConfig,
230 + prefix: str = "",
231 + create_in_proj_qkvz: bool = True,
232 + gqa_interleaved_layout=False,
233 + ) -> None:
234 + super().__init__()
235 + self.tp_size = get_tensor_model_parallel_world_size()
236 + self.tp_rank = get_tensor_model_parallel_rank()
237 + self.hidden_size = config.hidden_size
238 + self.num_v_heads = config.linear_num_value_heads
239 + self.num_k_heads = config.linear_num_key_heads
240 + self.head_k_dim = config.linear_key_head_dim
241 + self.head_v_dim = config.linear_value_head_dim
242 + self.key_dim = self.head_k_dim * self.num_k_heads
243 + self.value_dim = self.head_v_dim * self.num_v_heads
244 +
245 + self.conv_kernel_size = config.linear_conv_kernel_dim
246 + self.layer_idx = extract_layer_index(prefix)
247 + self.activation = config.hidden_act
248 + self.act = ACT2FN[config.hidden_act]
249 + self.layer_norm_epsilon = config.rms_norm_eps
250 + self.prefix = prefix
251 + self.config = config
252 + self.model_config = vllm_config.model_config
253 + self.cache_config = vllm_config.cache_config
254 + quant_config = vllm_config.quant_config
255 + self.speculative_config = vllm_config.speculative_config
256 + self.num_spec = (
257 + self.speculative_config.num_speculative_tokens
258 + if self.speculative_config
259 + else 0
260 + )
261 + self.gqa_interleaved_layout = gqa_interleaved_layout
262 +
263 + # QKV
264 + self.conv_dim = self.key_dim * 2 + self.value_dim
265 + self.conv1d = ColumnParallelLinear(
266 + input_size=self.conv_kernel_size,
267 + output_size=self.conv_dim,
268 + bias=False,
269 + prefix=f"{prefix}.conv1d",
270 + )
271 + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
272 +
273 + # projection of the input hidden states
274 + # Qwen3-Next and Qwen3.5 has a different qkv_proj layout,
275 + # we need to create qkvz_proj adaptively here.
276 + # When create_in_proj_qkvz is False (e.g. LoRA enabled in Qwen3.5),
277 + # in_proj_qkv and in_proj_z are created separately instead.
278 + if create_in_proj_qkvz:
279 + self.in_proj_qkvz = self.create_qkvz_proj(
280 + hidden_size=self.hidden_size,
281 + key_dim=self.key_dim,
282 + value_dim=self.value_dim,
283 + quant_config=quant_config,
284 + prefix=f"{prefix}.in_proj_qkvz",
285 + )
286 + else:
287 + # LoRA case (Qwen3.5 only): keep q/k/v and z as separate modules
288 + # so that LoRA adapters can be applied independently.
289 + self.in_proj_qkv = MergedColumnParallelLinear(
290 + input_size=self.hidden_size,
291 + output_sizes=[self.key_dim, self.key_dim, self.value_dim],
292 + bias=False,
293 + quant_config=quant_config,
294 + prefix=f"{prefix}.in_proj_qkv",
295 + )
296 + self.in_proj_z = ColumnParallelLinear(
297 + input_size=self.hidden_size,
298 + output_size=self.value_dim,
299 + bias=False,
300 + quant_config=quant_config,
301 + prefix=f"{prefix}.in_proj_z",
302 + )
303 + # ba_proj doesn't support blockwise fp8 quantization.
304 + # Qwen3-Next and Qwen3.5 have different in_proj_ba checkpoint
305 + # layouts, so we use a factory method to create the projection.
306 + self.in_proj_ba = self.create_ba_proj(
307 + hidden_size=self.hidden_size,
308 + num_v_heads=self.num_v_heads,
309 + quant_config=quant_config,
310 + prefix=f"{prefix}.in_proj_ba",
311 + )
312 +
313 + query_key_settings = (self.key_dim, 0, False)
314 + value_settings = (self.value_dim, 0, False)
315 +
316 + self.conv1d.weight.weight_loader = mamba_v2_sharded_weight_loader(
317 + [
318 + query_key_settings,
319 + query_key_settings,
320 + value_settings,
321 + ],
322 + self.tp_size,
323 + self.tp_rank,
324 + )
325 +
326 + # selective projection used to make dt, B and C input dependent
327 +
328 + # time step projection (discretization)
329 + # instantiate once and copy inv_dt in init_weights of PretrainedModel
330 + self.dt_bias = nn.Parameter(
331 + torch.ones(self.num_v_heads // self.tp_size),
332 + )
333 + self.A_log = nn.Parameter(
334 + torch.empty(
335 + divide(self.num_v_heads, self.tp_size),
336 + dtype=torch.float32,
337 + )
338 + )
339 +
340 + set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
341 + set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
342 +
343 + self.norm = RMSNormGated(
344 + self.head_v_dim,
345 + eps=self.layer_norm_epsilon,
346 + group_size=None,
347 + norm_before_gate=True,
348 + device=current_platform.current_device(),
349 + )
350 +
351 + self.out_proj = RowParallelLinear(
352 + self.value_dim,
353 + self.hidden_size,
354 + bias=False,
355 + input_is_parallel=True,
356 + quant_config=quant_config,
357 + prefix=f"{prefix}.out_proj",
358 + )
359 +
360 + self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
361 + self.enable_packed_recurrent_decode = (
362 + envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
363 + )
364 +
365 + compilation_config = get_current_vllm_config().compilation_config
366 + if prefix in compilation_config.static_forward_context:
367 + raise ValueError(f"Duplicate layer name: {prefix}")
368 + compilation_config.static_forward_context[prefix] = self
369 +
370 + def create_qkvz_proj(
371 + self,
372 + hidden_size: int,
373 + key_dim: int,
374 + value_dim: int,
375 + quant_config: QuantizationConfig | None,
376 + prefix: str,
377 + ) -> MergedColumnParallelLinear:
378 + # When gqa_interleaved_layout=True (Qwen3-Next), qkvz weights are
379 + # stored as a single fused tensor with interleaved GQA layout, so we
380 + # use one output shard to preserve the interleaving across TP ranks.
381 + # When gqa_interleaved_layout=False (Qwen3.5), the checkpoint has
382 + # separate q, k, v, z weights, so we use 4 independent output sizes.
383 + output_sizes = (
384 + [sum((key_dim, key_dim, value_dim, value_dim))]
385 + if self.gqa_interleaved_layout
386 + else [key_dim, key_dim, value_dim, value_dim]
387 + )
388 + return MergedColumnParallelLinear(
389 + input_size=hidden_size,
390 + output_sizes=output_sizes,
391 + bias=False,
392 + quant_config=quant_config,
393 + prefix=prefix,
394 + )
395 +
396 + def create_ba_proj(
397 + self,
398 + hidden_size: int,
399 + num_v_heads: int,
400 + quant_config: QuantizationConfig | None,
401 + prefix: str,
402 + ) -> MergedColumnParallelLinear:
403 + # When gqa_interleaved_layout=True (Qwen3-Next), in_proj_ba is stored
404 + # as a single fused weight [b_g0, a_g0, b_g1, a_g1, ...] interleaved
405 + # by key-head group; a single output shard preserves this across TP.
406 + # When gqa_interleaved_layout=False (Qwen3.5), in_proj_b and in_proj_a
407 + # are separate checkpoint weights, so we use 2 independent output sizes.
408 + output_sizes = (
409 + [num_v_heads * 2] if self.gqa_interleaved_layout else [num_v_heads] * 2
410 + )
411 + return MergedColumnParallelLinear(
412 + input_size=hidden_size,
413 + output_sizes=output_sizes,
414 + bias=False,
415 + quant_config=quant_config,
416 + prefix=prefix,
417 + )
418 +
419 + def fix_query_key_value_ordering(
420 + self,
421 + mixed_qkvz: torch.Tensor,
422 + mixed_ba: torch.Tensor,
423 + ):
424 + """
425 + Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
426 + """
427 + new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
428 + self.num_k_heads // self.tp_size,
429 + (
430 + self.head_k_dim
431 + + self.head_k_dim
432 + + (self.head_v_dim + self.head_v_dim)
433 + * self.num_v_heads
434 + // self.num_k_heads
435 + ),
436 + )
437 + new_tensor_shape_ba = mixed_ba.size()[:-1] + (
438 + self.num_k_heads // self.tp_size,
439 + 2 * self.num_v_heads // self.num_k_heads,
440 + )
441 +
442 + mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
443 + mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
444 +
445 + split_arg_list_qkvz = [
446 + self.head_k_dim,
447 + self.head_k_dim,
448 + (self.num_v_heads // self.num_k_heads * self.head_v_dim),
449 + (self.num_v_heads // self.num_k_heads * self.head_v_dim),
450 + ]
451 + split_arg_list_ba = [
452 + self.num_v_heads // self.num_k_heads,
453 + self.num_v_heads // self.num_k_heads,
454 + ]
455 +
456 + # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)]
457 + # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn],
458 + # [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng]
459 + (query, key, value, z) = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2)
460 + (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=2)
461 +
462 + # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]
463 + value = value.reshape(value.size(0), -1, self.head_v_dim)
464 + z = z.reshape(z.size(0), -1, self.head_v_dim)
465 + b = b.reshape(b.size(0), self.num_v_heads // self.tp_size)
466 + a = a.reshape(a.size(0), self.num_v_heads // self.tp_size)
467 +
468 + return query, key, value, z, b, a
469 +
470 + def rearrange_mixed_qkv(self, mixed_qkv):
471 + if mixed_qkv is None:
472 + return None, None, None
473 + query, key, value = torch.split(
474 + mixed_qkv,
475 + [
476 + self.key_dim // self.tp_size,
477 + self.key_dim // self.tp_size,
478 + self.value_dim // self.tp_size,
479 + ],
480 + dim=-1,
481 + )
482 + query, key = map(
483 + lambda x: rearrange(x, "l (h d) -> 1 l h d", d=self.head_k_dim),
484 + (query, key),
485 + )
486 + value = rearrange(value, "l (h d) -> 1 l h d", d=self.head_v_dim)
487 + return query.contiguous(), key.contiguous(), value.contiguous()
488 +
489 + def forward(
490 + self,
491 + hidden_states: torch.Tensor,
492 + output: torch.Tensor,
493 + ):
494 + """
495 + Forward pass with three parts:
496 + 1. Input projection
497 + 2. Core attention (custom op)
498 + 3. Output projection
499 + """
500 + num_tokens = hidden_states.size(0)
501 + # ============================================================
502 + # Part 1: Input Projection
503 + # ============================================================
504 + if hasattr(self, "in_proj_qkv"):
505 + # LoRA path (Qwen3.5 only): separate in_proj_qkv and in_proj_z
506 + mixed_qkv, _ = self.in_proj_qkv(hidden_states)
507 + ba, _ = self.in_proj_ba(hidden_states)
508 + z, _ = self.in_proj_z(hidden_states)
509 + z = z.reshape(z.size(0), -1, self.head_v_dim)
510 + b, a = ba.chunk(2, dim=-1)
511 + b = b.contiguous()
512 + a = a.contiguous()
513 + else:
514 + mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
515 + ba, _ = self.in_proj_ba(hidden_states)
516 +
517 + if self.gqa_interleaved_layout:
518 + # Qwen3-Next: unpack the interleaved GQA layout
519 + query, key, value, z, b, a = self.fix_query_key_value_ordering(
520 + mixed_qkvz, ba
521 + )
522 + query, key, value = map(
523 + lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)
524 + )
525 + mixed_qkv = torch.cat((query, key, value), dim=-1)
526 + else:
527 + # Qwen3.5: weights are already in [q, k, v, z] and [b, a] order
528 + qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
529 + z_size = self.value_dim // self.tp_size
530 + mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
531 + z = z.reshape(z.size(0), -1, self.head_v_dim)
532 + b, a = ba.chunk(2, dim=-1)
533 + b = b.contiguous()
534 + a = a.contiguous()
535 +
536 + # ============================================================
537 + # Part 2: Core Attention (Custom Op)
538 + # ============================================================
539 + # Note: we should not use torch.empty here like other attention backends,
540 + # see discussions in https://github.com/vllm-project/vllm/pull/28182
541 + core_attn_out = torch.zeros(
542 + (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
543 + dtype=hidden_states.dtype,
544 + device=hidden_states.device,
545 + )
546 +
547 + torch.ops.vllm.gdn_attention_core(
548 + mixed_qkv,
549 + b,
550 + a,
551 + core_attn_out,
552 + self.prefix,
553 + )
554 +
555 + # ============================================================
556 + # Part 3: Output Projection
557 + # ============================================================
558 + z_shape_og = z.shape
559 + # Reshape input data into 2D tensor
560 + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
561 + z = z.reshape(-1, z.shape[-1])
562 + core_attn_out = self.norm(core_attn_out, z)
563 + core_attn_out = core_attn_out.reshape(z_shape_og)
564 + core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
565 + output[:num_tokens], _ = self.out_proj(core_attn_out)
566 +
567 + def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
568 + """Warm up GDN prefill kernels during V1 profiling.
569 +
570 + During V1 profile runs, ``_forward_core`` returns early because
571 + ``attn_metadata`` is ``None``, so the autotuned kernels used by
572 + ``chunk_gated_delta_rule`` (e.g. ``solve_tril``,
573 + ``chunk_scaled_dot_kkt``) are never invoked. After profiling,
574 + vLLM allocates KV cache using most of the remaining GPU memory.
575 + When the first real inference triggers the autotuner it OOMs
576 + because there is not enough memory left for benchmarking.
577 +
578 + This method runs minimal forward passes through
579 + ``chunk_gated_delta_rule`` with small dummy tensors to force
580 + autotuning while GPU memory is still plentiful. The autotuner
581 + results are cached globally, so only the first layer incurs
582 + actual benchmarking cost.
583 +
584 + Most kernels use a fixed ``BT = chunk_size`` (64), but
585 + ``chunk_fwd_kernel_o`` recomputes ``BT`` from the sequence
586 + length: ``min(64, max(16, next_power_of_2(T)))``. Since ``BT``
587 + is part of its autotune key, we run warmup passes with T = 16,
588 + 32, and 64 to cover all possible ``BT`` values.
589 +
590 + The decode path uses ``fused_sigmoid_gating_delta_rule_update``
591 + which has fixed kernel parameters (no autotuning), so only the
592 + prefill (chunked) path needs warming up.
593 + """
594 + if hasattr(self, "_prefill_kernels_warmed_up"):
595 + return
596 + self._prefill_kernels_warmed_up = True
597 +
598 + device = mixed_qkv.device
599 + dtype = mixed_qkv.dtype
600 + num_k_heads = self.num_k_heads // self.tp_size
601 + num_v_heads = self.num_v_heads // self.tp_size
602 + _, state_dtype = self.get_state_dtype()
603 +
604 + # Run warmup for each possible BT value of chunk_fwd_kernel_o:
605 + # T=16 → BT=16, T=32 → BT=32, T=64 → BT=64.
606 + # Other kernels always use BT=chunk_size(64), so their autotune
607 + # cache is populated on the first pass and reused thereafter.
608 + for T in (16, 32, 64):
609 + q = torch.randn(
610 + 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
611 + )
612 + k = torch.randn(
613 + 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
614 + )
615 + v = torch.randn(
616 + 1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype
617 + )
618 + # NOTE: g and beta must have the same dtypes as during
619 + # inference, so we construct them with the same function
620 + # (fused_gdn_gating). dummy_a and dummy_b are throwaway
621 + # inputs required by that function.
622 + dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)
623 + dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)
624 + g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias)
625 + state = torch.zeros(
626 + 1,
627 + num_v_heads,
628 + self.head_v_dim,
629 + self.head_k_dim,
630 + device=device,
631 + dtype=state_dtype,
632 + )
633 + cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)
634 +
635 + try:
636 + self.chunk_gated_delta_rule(
637 + q=q,
638 + k=k,
639 + v=v,
640 + g=g,
641 + beta=beta,
642 + initial_state=state,
643 + output_final_state=True,
644 + cu_seqlens=cu_seqlens,
645 + use_qk_l2norm_in_kernel=True,
646 + )
647 + except Exception:
648 + logger.warning(
649 + "GDN prefill kernel warmup (T=%d) failed for "
650 + "layer %s. First inference may OOM due to "
651 + "autotuner.",
652 + T,
653 + self.prefix,
654 + exc_info=True,
655 + )
656 + else:
657 + logger.debug(
658 + "GDN prefill kernel warmup (T=%d) completed for layer %s",
659 + T,
660 + self.prefix,
661 + )
662 + finally:
663 + del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
664 +
665 + torch.accelerator.empty_cache()
666 +
667 + def _forward_core(
668 + self,
669 + mixed_qkv: torch.Tensor,
670 + b: torch.Tensor,
671 + a: torch.Tensor,
672 + core_attn_out: torch.Tensor,
673 + ):
674 + forward_context = get_forward_context()
675 + attn_metadata: AttentionMetadata = forward_context.attn_metadata
676 +
677 + if attn_metadata is None:
678 + # V1 profile run — warm up prefill kernels so that
679 + # autotuning completes before KV cache allocation.
680 + self._warmup_prefill_kernels(mixed_qkv)
681 + return
682 +
683 + assert isinstance(attn_metadata, dict)
684 + attn_metadata = attn_metadata[self.prefix]
685 + assert isinstance(attn_metadata, GDNAttentionMetadata)
686 +
687 + if (
688 + self.enable_packed_recurrent_decode
689 + and attn_metadata.spec_sequence_masks is None
690 + and attn_metadata.num_prefills == 0
691 + and attn_metadata.num_decodes > 0
692 + ):
693 + return self._forward_core_decode_non_spec(
694 + mixed_qkv=mixed_qkv,
695 + b=b,
696 + a=a,
697 + core_attn_out=core_attn_out,
698 + attn_metadata=attn_metadata,
699 + )
700 +
701 + has_initial_state = attn_metadata.has_initial_state
702 + spec_query_start_loc = attn_metadata.spec_query_start_loc
703 + non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
704 + spec_sequence_masks = attn_metadata.spec_sequence_masks
705 + spec_token_indx = attn_metadata.spec_token_indx
706 + non_spec_token_indx = attn_metadata.non_spec_token_indx
707 + spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
708 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
709 + self_kv_cache = self.kv_cache
710 + conv_state = self_kv_cache[0].transpose(-1, -2)
711 + ssm_state = self_kv_cache[1]
712 + num_actual_tokens = attn_metadata.num_actual_tokens
713 + num_accepted_tokens = attn_metadata.num_accepted_tokens
714 +
715 + mixed_qkv = mixed_qkv[:num_actual_tokens]
716 + b = b[:num_actual_tokens]
717 + a = a[:num_actual_tokens]
718 +
719 + # 1. Convolution sequence transformation
720 + conv_weights = self.conv1d.weight.view(
721 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)
722 + )
723 +
724 + if spec_sequence_masks is not None:
725 + if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
726 + mixed_qkv_spec = mixed_qkv
727 + mixed_qkv_non_spec = None
728 + else:
729 + mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
730 + mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
731 + else:
732 + mixed_qkv_spec = None
733 + mixed_qkv_non_spec = mixed_qkv
734 +
735 + # 1.1: Process the multi-query part
736 + if spec_sequence_masks is not None:
737 + mixed_qkv_spec = causal_conv1d_update(
738 + mixed_qkv_spec,
739 + conv_state,
740 + conv_weights,
741 + self.conv1d.bias,
742 + self.activation,
743 + conv_state_indices=spec_state_indices_tensor[:, 0][
744 + : attn_metadata.num_spec_decodes
745 + ],
746 + num_accepted_tokens=num_accepted_tokens,
747 + query_start_loc=spec_query_start_loc,
748 + max_query_len=spec_state_indices_tensor.size(-1),
749 + validate_data=False,
750 + )
751 +
752 + # 1.2: Process the remaining part
753 + if attn_metadata.num_prefills > 0:
754 + assert mixed_qkv_non_spec is not None
755 + mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
756 + # - "cache_indices" updates the conv_state cache in positions
757 + # pointed to by "state_indices_tensor"
758 + mixed_qkv_non_spec = causal_conv1d_fn(
759 + mixed_qkv_non_spec_T,
760 + conv_weights,
761 + self.conv1d.bias,
762 + activation=self.activation,
763 + conv_states=conv_state,
764 + has_initial_state=has_initial_state,
765 + cache_indices=non_spec_state_indices_tensor,
766 + query_start_loc=non_spec_query_start_loc,
767 + metadata=attn_metadata,
768 + ).transpose(0, 1)
769 + elif attn_metadata.num_decodes > 0:
770 + assert mixed_qkv_non_spec is not None
771 + mixed_qkv_non_spec = causal_conv1d_update(
772 + mixed_qkv_non_spec,
773 + conv_state,
774 + conv_weights,
775 + self.conv1d.bias,
776 + self.activation,
777 + conv_state_indices=non_spec_state_indices_tensor[
778 + : attn_metadata.num_actual_tokens
779 + ],
780 + validate_data=True,
781 + )
782 + else:
783 + mixed_qkv_non_spec = None
784 +
785 + query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
786 + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
787 + mixed_qkv_non_spec
788 + )
789 +
790 + if attn_metadata.num_prefills > 0:
791 + g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias)
792 + if spec_sequence_masks is not None:
793 + g_non_spec = g.index_select(1, non_spec_token_indx)
794 + beta_non_spec = beta.index_select(1, non_spec_token_indx)
795 + else:
796 + g_non_spec = g
797 + beta_non_spec = beta
798 + else:
799 + g_non_spec = None
800 + beta_non_spec = None
801 +
802 + # 2. Recurrent attention
803 +
804 + # 2.1: Process the multi-query part
805 + if spec_sequence_masks is not None:
806 + core_attn_out_spec, last_recurrent_state = (
807 + fused_sigmoid_gating_delta_rule_update(
808 + A_log=self.A_log,
809 + a=a,
810 + b=b,
811 + dt_bias=self.dt_bias,
812 + q=query_spec,
813 + k=key_spec,
814 + v=value_spec,
815 + initial_state=ssm_state,
816 + inplace_final_state=True,
817 + cu_seqlens=spec_query_start_loc[
818 + : attn_metadata.num_spec_decodes + 1
819 + ],
820 + ssm_state_indices=spec_state_indices_tensor,
821 + num_accepted_tokens=num_accepted_tokens,
822 + use_qk_l2norm_in_kernel=True,
823 + )
824 + )
825 + else:
826 + core_attn_out_spec, last_recurrent_state = None, None
827 +
828 + # 2.2: Process the remaining part
829 + if attn_metadata.num_prefills > 0:
830 + initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
831 + initial_state[~has_initial_state, ...] = 0
832 + (
833 + core_attn_out_non_spec,
834 + last_recurrent_state,
835 + ) = self.chunk_gated_delta_rule(
836 + q=query_non_spec,
837 + k=key_non_spec,
838 + v=value_non_spec,
839 + g=g_non_spec,
840 + beta=beta_non_spec,
841 + initial_state=initial_state,
842 + output_final_state=True,
843 + cu_seqlens=non_spec_query_start_loc,
844 + use_qk_l2norm_in_kernel=True,
845 + )
846 + # Init cache
847 + ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
848 + ssm_state.dtype
849 + )
850 + elif attn_metadata.num_decodes > 0:
851 + core_attn_out_non_spec, last_recurrent_state = (
852 + fused_sigmoid_gating_delta_rule_update(
853 + A_log=self.A_log,
854 + a=a,
855 + b=b,
856 + dt_bias=self.dt_bias,
857 + q=query_non_spec,
858 + k=key_non_spec,
859 + v=value_non_spec,
860 + initial_state=ssm_state,
861 + inplace_final_state=True,
862 + cu_seqlens=non_spec_query_start_loc[
863 + : attn_metadata.num_decodes + 1
864 + ],
865 + ssm_state_indices=non_spec_state_indices_tensor,
866 + use_qk_l2norm_in_kernel=True,
867 + )
868 + )
869 + else:
870 + core_attn_out_non_spec, last_recurrent_state = None, None
871 +
872 + # 3. Merge core attention output
873 + if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
874 + merged_out = torch.empty(
875 + (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
876 + dtype=core_attn_out_non_spec.dtype,
877 + device=core_attn_out_non_spec.device,
878 + )
879 + merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
880 + merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
881 + core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
882 + elif spec_sequence_masks is not None:
883 + core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
884 + else:
885 + core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
886 +
887 + def _forward_core_decode_non_spec(
888 + self,
889 + mixed_qkv: torch.Tensor,
890 + b: torch.Tensor,
891 + a: torch.Tensor,
892 + core_attn_out: torch.Tensor,
893 + attn_metadata: GDNAttentionMetadata,
894 + ):
895 + """
896 + Core attention computation with a packed non-spec decode fast path.
897 + """
898 + non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
899 + self_kv_cache = self.kv_cache
900 + conv_state = self_kv_cache[0].transpose(-1, -2)
901 + ssm_state = self_kv_cache[1]
902 + num_actual_tokens = attn_metadata.num_actual_tokens
903 +
904 + mixed_qkv = mixed_qkv[:num_actual_tokens]
905 + b = b[:num_actual_tokens]
906 + a = a[:num_actual_tokens]
907 +
908 + conv_weights = self.conv1d.weight.view(
909 + self.conv1d.weight.size(0), self.conv1d.weight.size(2)
910 + )
911 + mixed_qkv_non_spec = causal_conv1d_update(
912 + mixed_qkv,
913 + conv_state,
914 + conv_weights,
915 + self.conv1d.bias,
916 + self.activation,
917 + conv_state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
918 + validate_data=False,
919 + )
920 + out_buf = core_attn_out[:num_actual_tokens].unsqueeze(1)
921 + fused_recurrent_gated_delta_rule_packed_decode(
922 + mixed_qkv=mixed_qkv_non_spec,
923 + a=a,
924 + b=b,
925 + A_log=self.A_log,
926 + dt_bias=self.dt_bias,
927 + scale=self.head_k_dim**-0.5,
928 + initial_state=ssm_state,
929 + out=out_buf,
930 + ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
931 + use_qk_l2norm_in_kernel=True,
932 + )
933 + return
934 +
935 +
936 +def gdn_attention_core(
937 + mixed_qkv: torch.Tensor,
938 + b: torch.Tensor,
939 + a: torch.Tensor,
940 + core_attn_out: torch.Tensor,
941 + layer_name: str,
942 +) -> None:
943 + """
944 + Custom op for the core attention computation.
945 + Only handles the convolution + recurrent attention part.
946 + Input/output projections are handled outside this op.
947 + """
948 + forward_context: ForwardContext = get_forward_context()
949 + self = forward_context.no_compile_layers[layer_name]
950 + self._forward_core(
951 + mixed_qkv=mixed_qkv,
952 + b=b,
953 + a=a,
954 + core_attn_out=core_attn_out,
955 + )
956 +
957 +
958 +def gdn_attention_core_fake(
959 + mixed_qkv: torch.Tensor,
960 + b: torch.Tensor,
961 + a: torch.Tensor,
962 + core_attn_out: torch.Tensor,
963 + layer_name: str,
964 +) -> None:
965 + """Fake implementation for torch.compile."""
966 + return
967 +
968 +
969 +direct_register_custom_op(
970 + op_name="gdn_attention_core",
971 + op_func=gdn_attention_core,
972 + mutates_args=["core_attn_out"],
973 + fake_impl=gdn_attention_core_fake,
974 +)
975 +
976 +
977 +@triton.jit
978 +def fused_gdn_gating_kernel(
979 + g,
980 + beta_output,
981 + A_log,
982 + a,
983 + b,
984 + dt_bias,
985 + seq_len,
986 + NUM_HEADS: tl.constexpr,
987 + beta: tl.constexpr,
988 + threshold: tl.constexpr,
989 + BLK_HEADS: tl.constexpr,
990 +):
991 + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
992 + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
993 + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
994 + mask = head_off < NUM_HEADS
995 + blk_A_log = tl.load(A_log + head_off, mask=mask)
996 + blk_a = tl.load(a + off, mask=mask)
997 + blk_b = tl.load(b + off, mask=mask)
998 + blk_bias = tl.load(dt_bias + head_off, mask=mask)
999 + # If the model is loaded in fp16, without the .float() here, A might be -inf
1000 + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
1001 + softplus_x = tl.where(
1002 + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
1003 + )
1004 + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
1005 + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
1006 + # compute beta_output = sigmoid(b)
1007 + blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
1008 + tl.store(
1009 + beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
1010 + )
1011 +
1012 +
1013 +def fused_gdn_gating(
1014 + A_log: torch.Tensor,
1015 + a: torch.Tensor,
1016 + b: torch.Tensor,
1017 + dt_bias: torch.Tensor,
1018 + beta: float = 1.0,
1019 + threshold: float = 20.0,
1020 +) -> tuple[torch.Tensor, torch.Tensor]:
1021 + """
1022 + Fused computation of g and beta for Gated Delta Net.
1023 + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
1024 + beta_output = b.sigmoid()
1025 + TODO maybe use torch.compile to replace this triton kernel
1026 + """
1027 + batch, num_heads = a.shape
1028 + seq_len = 1
1029 + grid = (batch, seq_len, triton.cdiv(num_heads, 8))
1030 + g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
1031 + beta_output = torch.empty(1, batch, num_heads, dtype=b.dtype, device=b.device)
1032 + fused_gdn_gating_kernel[grid](
1033 + g,
1034 + beta_output,
1035 + A_log,
1036 + a,
1037 + b,
1038 + dt_bias,
1039 + seq_len,
1040 + num_heads,
1041 + beta,
1042 + threshold,
1043 + 8,
1044 + num_warps=1,
1045 + )
1046 + return g, beta_output
vllm/model_executor/models/qwen3_5.py
+4 -151
@@ -28,7 +28,6 @@ import typing
28 from collections.abc import Callable, Iterable
29
30 import torch
31 -from einops import rearrange
31 from torch import nn
32
33 from vllm.compilation.decorators import support_torch_compile
@@ -40,18 +39,14 @@ from vllm.logger import init_logger
39 from vllm.model_executor.layers.layernorm import (
40 GemmaRMSNorm as Qwen3_5RMSNorm,
41 )
43 -from vllm.model_executor.layers.linear import (
44 - ColumnParallelLinear,
45 - MergedColumnParallelLinear,
46 -)
42 from vllm.model_executor.layers.logits_processor import LogitsProcessor
43 +from vllm.model_executor.layers.mamba.gdn_linear_attn import GatedDeltaNetAttention
44 from vllm.model_executor.layers.mamba.mamba_utils import (
45 MambaStateCopyFunc,
46 MambaStateCopyFuncCalculator,
47 MambaStateDtypeCalculator,
48 MambaStateShapeCalculator,
49 )
54 -from vllm.model_executor.layers.quantization import QuantizationConfig
50 from vllm.model_executor.layers.vocab_parallel_embedding import (
51 ParallelLMHead,
52 VocabParallelEmbedding,
@@ -85,7 +80,6 @@ from .qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
80 from .qwen3_next import (
81 Qwen3NextAttention,
82 Qwen3NextDecoderLayer,
88 - Qwen3NextGatedDeltaNet,
83 Qwen3NextModel,
84 Qwen3NextSparseMoeBlock,
85 QwenNextMixtureOfExperts,
@@ -121,149 +115,6 @@ class Qwen3_5MoeProcessingInfo(Qwen3VLProcessingInfo):
115 return self.ctx.get_hf_config(Qwen3_5MoeConfig)
116
117
124 -class Qwen3_5GatedDeltaNet(Qwen3NextGatedDeltaNet):
125 - def fix_query_key_value_ordering(
126 - self,
127 - mixed_qkvz: torch.Tensor,
128 - mixed_ba: torch.Tensor,
129 - ):
130 - raise NotImplementedError(
131 - "Qwen3.5 Series dont need to fix query key value ordering"
132 - )
133 -
134 - def __init__(
135 - self,
136 - config: Qwen3_5Config,
137 - vllm_config: VllmConfig,
138 - prefix: str = "",
139 - ) -> None:
140 - create_in_proj_qkvz = vllm_config.lora_config is None
141 - super().__init__(
142 - config,
143 - vllm_config=vllm_config,
144 - prefix=prefix,
145 - create_in_proj_qkvz=create_in_proj_qkvz,
146 - )
147 - if vllm_config.lora_config is not None:
148 - # Separate in_proj_qkv (Q,K,V) and in_proj_z for LoRA compatibility.
149 - # Use MergedColumnParallelLinear for in_proj_qkv because GDN can have
150 - # linear_num_key_heads != linear_num_value_heads (e.g. 16 vs 32), so
151 - # output sizes [key_dim, key_dim, value_dim] are not representable
152 - # with a single QKVParallelLinear (which ties K and V head counts).
153 - self.in_proj_qkv = MergedColumnParallelLinear(
154 - input_size=self.hidden_size,
155 - output_sizes=[self.key_dim, self.key_dim, self.value_dim],
156 - bias=False,
157 - quant_config=vllm_config.quant_config,
158 - prefix=f"{prefix}.in_proj_qkv",
159 - )
160 - self.in_proj_z = ColumnParallelLinear(
161 - input_size=self.hidden_size,
162 - output_size=self.value_dim,
163 - bias=False,
164 - quant_config=vllm_config.quant_config,
165 - prefix=f"{prefix}.in_proj_z",
166 - )
167 -
168 - def create_qkvz_proj(
169 - self,
170 - hidden_size: int,
171 - key_dim: int,
172 - value_dim: int,
173 - quant_config: QuantizationConfig | None,
174 - prefix: str,
175 - ) -> MergedColumnParallelLinear:
176 - return MergedColumnParallelLinear(
177 - input_size=hidden_size,
178 - output_sizes=[key_dim, key_dim, value_dim, value_dim],
179 - bias=False,
180 - quant_config=quant_config,
181 - prefix=prefix,
182 - )
183 -
184 - def create_ba_proj(
185 - self,
186 - hidden_size: int,
187 - num_v_heads: int,
188 - quant_config: QuantizationConfig | None,
189 - prefix: str,
190 - ) -> MergedColumnParallelLinear:
191 - # Qwen3.5 has separate in_proj_b and in_proj_a weights in the
192 - # checkpoint, which are loaded into the fused in_proj_ba parameter
193 - # via stacked_params_mapping with shard_id 0 and 1 respectively.
194 - return MergedColumnParallelLinear(
195 - input_size=hidden_size,
196 - output_sizes=[num_v_heads] * 2,
197 - bias=False,
198 - quant_config=quant_config,
199 - prefix=prefix,
200 - )
201 -
202 - def forward(
203 - self,
204 - hidden_states: torch.Tensor,
205 - output: torch.Tensor,
206 - ):
207 - """
208 - Forward pass with three parts:
209 - 1. Input projection
210 - 2. Core attention (custom op)
211 - 3. Output projection
212 - """
213 - num_tokens = hidden_states.size(0)
214 -
215 - # ============================================================
216 - # Part 1: Input Projection
217 - # ============================================================
218 - if hasattr(self, "in_proj_qkv"):
219 - # LoRA path: separate in_proj_qkv and in_proj_z
220 - mixed_qkv, _ = self.in_proj_qkv(hidden_states)
221 - ba, _ = self.in_proj_ba(hidden_states)
222 - z, _ = self.in_proj_z(hidden_states)
223 - else:
224 - mixed_qkvz, _ = self.in_proj_qkvz(hidden_states)
225 - ba, _ = self.in_proj_ba(hidden_states)
226 - qkv_size = (self.key_dim * 2 + self.value_dim) // self.tp_size
227 - z_size = self.value_dim // self.tp_size
228 - mixed_qkv, z = mixed_qkvz.split([qkv_size, z_size], dim=-1)
229 - z = z.reshape(z.size(0), -1, self.head_v_dim)
230 - b, a = ba.chunk(2, dim=-1)
231 -
232 - b = b.contiguous()
233 - a = a.contiguous()
234 -
235 - # ============================================================
236 - # Part 2: Core Attention (Custom Op)
237 - # ============================================================
238 - # Note: we should not use torch.empty here like other attention backends,
239 - # see discussions in https://github.com/vllm-project/vllm/pull/28182
240 - core_attn_out = torch.zeros(
241 - (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
242 - dtype=hidden_states.dtype,
243 - device=hidden_states.device,
244 - )
245 -
246 - torch.ops.vllm.gdn_attention_core(
247 - mixed_qkv,
248 - b,
249 - a,
250 - core_attn_out,
251 - self.prefix,
252 - )
253 -
254 - # ============================================================
255 - # Part 3: Output Projection
256 - # ============================================================
257 - z_shape_og = z.shape
258 - # Reshape input data into 2D tensor
259 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
260 - z = z.reshape(-1, z.shape[-1])
261 - core_attn_out = self.norm(core_attn_out, z)
262 - core_attn_out = core_attn_out.reshape(z_shape_og)
263 - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
264 - output[:num_tokens], _ = self.out_proj(core_attn_out)
265 -
266 -
118 class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
119 def __init__(
120 self,
@@ -282,10 +133,12 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
133 self.layer_idx = extract_layer_index(prefix)
134
135 if self.layer_type == "linear_attention":
285 - self.linear_attn = Qwen3_5GatedDeltaNet(
136 + self.linear_attn = GatedDeltaNetAttention(
137 config=config,
138 vllm_config=vllm_config,
139 prefix=f"{prefix}.linear_attn",
140 + gqa_interleaved_layout=False,
141 + create_in_proj_qkvz=vllm_config.lora_config is None,
142 )
143 elif self.layer_type == "full_attention":
144 self.self_attn = Qwen3NextAttention(
vllm/model_executor/models/qwen3_next.py
+3 -975
@@ -6,11 +6,8 @@ from collections.abc import Iterable
6 from itertools import islice
7
8 import torch
9 -from einops import rearrange
9 from torch import nn
11 -from transformers.activations import ACT2FN
10
13 -from vllm import envs
11 from vllm.compilation.decorators import support_torch_compile
12 from vllm.config import (
13 CacheConfig,
@@ -19,50 +16,30 @@ from vllm.config import (
16 get_current_vllm_config,
17 )
18 from vllm.distributed import (
22 - divide,
19 get_ep_group,
20 get_pp_group,
25 - get_tensor_model_parallel_rank,
21 get_tensor_model_parallel_world_size,
22 tensor_model_parallel_all_gather,
23 )
29 -from vllm.forward_context import ForwardContext, get_forward_context
24 from vllm.logger import init_logger
31 -from vllm.model_executor.custom_op import CustomOp
25 from vllm.model_executor.layers.attention import Attention
33 -from vllm.model_executor.layers.fla.ops import (
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_packed_decode,
38 - fused_sigmoid_gating_delta_rule_update,
39 -)
40 -from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd
26 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
27 from vllm.model_executor.layers.layernorm import (
28 GemmaRMSNorm as Qwen3NextRMSNorm,
29 )
45 -from vllm.model_executor.layers.layernorm import RMSNormGated
30 from vllm.model_executor.layers.linear import (
47 - ColumnParallelLinear,
48 - MergedColumnParallelLinear,
31 QKVParallelLinear,
32 ReplicatedLinear,
33 RowParallelLinear,
34 )
35 from vllm.model_executor.layers.logits_processor import LogitsProcessor
54 -from vllm.model_executor.layers.mamba.abstract import MambaBase
55 -from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weight_loader
36 +from vllm.model_executor.layers.mamba.gdn_linear_attn import GatedDeltaNetAttention
37 from vllm.model_executor.layers.mamba.mamba_utils import (
38 MambaStateCopyFunc,
39 MambaStateCopyFuncCalculator,
40 MambaStateDtypeCalculator,
41 MambaStateShapeCalculator,
42 )
62 -from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
63 - causal_conv1d_fn,
64 - causal_conv1d_update,
65 -)
43 from vllm.model_executor.layers.quantization import QuantizationConfig
44 from vllm.model_executor.layers.rotary_embedding import get_rope
45 from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -72,18 +49,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
49 from vllm.model_executor.model_loader.weight_utils import (
50 default_weight_loader,
51 maybe_remap_kv_scale_name,
75 - sharded_weight_loader,
52 )
53 from vllm.model_executor.models.qwen2_moe import Qwen2MoeMLP as Qwen3NextMLP
54 from vllm.model_executor.models.utils import sequence_parallel_chunk
79 -from vllm.model_executor.utils import set_weight_attrs
80 -from vllm.platforms import current_platform
55 from vllm.sequence import IntermediateTensors
56 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
83 -from vllm.triton_utils import tl, triton
84 -from vllm.utils.torch_utils import direct_register_custom_op
85 -from vllm.v1.attention.backend import AttentionMetadata
86 -from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
57
58 from .interfaces import (
59 HasInnerState,
@@ -107,149 +77,6 @@ logger = init_logger(__name__)
77 KVCache = tuple[torch.Tensor, torch.Tensor]
78
79
110 -def fi_chunk_gated_delta_rule(
111 - q: torch.Tensor,
112 - k: torch.Tensor,
113 - v: torch.Tensor,
114 - g: torch.Tensor,
115 - beta: torch.Tensor,
116 - initial_state: torch.Tensor,
117 - output_final_state: bool,
118 - cu_seqlens: torch.Tensor | None = None,
119 - use_qk_l2norm_in_kernel: bool = True,
120 -):
121 - from flashinfer.gdn_prefill import (
122 - chunk_gated_delta_rule as chunk_gated_delta_rule_fi,
123 - )
124 -
125 - if use_qk_l2norm_in_kernel:
126 - q = l2norm_fwd(q)
127 - k = l2norm_fwd(k)
128 -
129 - # use flashinfer implementation
130 - q = q.squeeze(0).contiguous()
131 - k = k.squeeze(0).contiguous()
132 - v = v.squeeze(0).contiguous()
133 -
134 - g = g.squeeze(0).contiguous()
135 - beta = beta.squeeze(0).contiguous()
136 - fi_state = initial_state.to(torch.float32)
137 - fi_g = g.to(torch.float32)
138 - fi_beta = beta.to(torch.float32)
139 - result = chunk_gated_delta_rule_fi(
140 - q=q,
141 - k=k,
142 - v=v,
143 - g=torch.exp(fi_g),
144 - beta=fi_beta,
145 - initial_state=fi_state,
146 - output_final_state=output_final_state,
147 - cu_seqlens=cu_seqlens,
148 - )
149 - # FlashInfer returns (output, state) when output_final_state=True,
150 - # or just output when output_final_state=False.
151 - # Unsqueeze back to 4D (1, L, H, D) to match fla output format
152 - if output_final_state:
153 - output, final_state = result
154 - return output.unsqueeze(0), final_state
155 - else:
156 - return result.unsqueeze(0), None
157 -
158 -
159 -@CustomOp.register("chunk_gated_delta_rule")
160 -class ChunkGatedDeltaRule(CustomOp):
161 - def __init__(self) -> None:
162 - super().__init__()
163 - backend = (
164 - str(
165 - get_current_vllm_config().additional_config.get(
166 - "gdn_prefill_backend", "auto"
167 - )
168 - )
169 - .strip()
170 - .lower()
171 - )
172 - supports_flashinfer = (
173 - current_platform.is_cuda() and current_platform.is_device_capability(90)
174 - )
175 -
176 - if backend == "flashinfer":
177 - use_flashinfer = supports_flashinfer
178 - if not use_flashinfer:
179 - logger.warning_once(
180 - "GDN prefill backend 'flashinfer' is selected but "
181 - "cannot use this kernel on the current platform. "
182 - "Falling back to Triton/FLA."
183 - )
184 - elif backend == "triton":
185 - use_flashinfer = False
186 - else:
187 - use_flashinfer = supports_flashinfer
188 -
189 - if use_flashinfer:
190 - logger.info_once("Using FlashInfer GDN prefill kernel", scope="local")
191 - logger.info_once(
192 - "FlashInfer GDN prefill kernel is JIT-compiled; first run may "
193 - "take a while to compile. Set `--gdn-prefill-backend triton` to "
194 - "avoid JIT compile time.",
195 - scope="local",
196 - )
197 - else:
198 - logger.info_once("Using Triton/FLA GDN prefill kernel", scope="local")
199 -
200 - self._forward_method = (
201 - self.forward_cuda if use_flashinfer else self.forward_native
202 - )
203 -
204 - def forward_cuda(
205 - self,
206 - q: torch.Tensor,
207 - k: torch.Tensor,
208 - v: torch.Tensor,
209 - g: torch.Tensor,
210 - beta: torch.Tensor,
211 - initial_state: torch.Tensor,
212 - output_final_state: bool,
213 - cu_seqlens: torch.Tensor | None = None,
214 - use_qk_l2norm_in_kernel: bool = True,
215 - ):
216 - return fi_chunk_gated_delta_rule(
217 - q=q,
218 - k=k,
219 - v=v,
220 - g=g,
221 - beta=beta,
222 - initial_state=initial_state,
223 - output_final_state=output_final_state,
224 - cu_seqlens=cu_seqlens,
225 - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
226 - )
227 -
228 - def forward_native(
229 - self,
230 - q: torch.Tensor,
231 - k: torch.Tensor,
232 - v: torch.Tensor,
233 - g: torch.Tensor,
234 - beta: torch.Tensor,
235 - initial_state: torch.Tensor,
236 - output_final_state: bool,
237 - cu_seqlens: torch.Tensor | None = None,
238 - use_qk_l2norm_in_kernel: bool = True,
239 - ):
240 - return fla_chunk_gated_delta_rule(
241 - q=q,
242 - k=k,
243 - v=v,
244 - g=g,
245 - beta=beta,
246 - initial_state=initial_state,
247 - output_final_state=output_final_state,
248 - cu_seqlens=cu_seqlens,
249 - use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
250 - )
251 -
252 -
80 class Qwen3NextSparseMoeBlock(nn.Module):
81 def __init__(self, vllm_config: VllmConfig, prefix: str = ""):
82 super().__init__()
@@ -370,693 +197,6 @@ class Qwen3NextSparseMoeBlock(nn.Module):
197 return final_hidden_states.view(orig_shape)
198
199
373 -class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
374 - @property
375 - def mamba_type(self) -> str:
376 - return "gdn_attention"
377 -
378 - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:
379 - return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
380 - self.model_config.dtype,
381 - self.cache_config.mamba_cache_dtype,
382 - self.cache_config.mamba_ssm_cache_dtype,
383 - )
384 -
385 - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:
386 - return MambaStateShapeCalculator.gated_delta_net_state_shape(
387 - self.tp_size,
388 - self.num_k_heads,
389 - self.num_v_heads,
390 - self.head_k_dim,
391 - self.head_v_dim,
392 - self.conv_kernel_size,
393 - self.num_spec,
394 - )
395 -
396 - def __init__(
397 - self,
398 - config: Qwen3NextConfig,
399 - vllm_config: VllmConfig,
400 - prefix: str = "",
401 - create_in_proj_qkvz: bool = True,
402 - ) -> None:
403 - super().__init__()
404 - self.tp_size = get_tensor_model_parallel_world_size()
405 - self.tp_rank = get_tensor_model_parallel_rank()
406 - self.hidden_size = config.hidden_size
407 - self.num_v_heads = config.linear_num_value_heads
408 - self.num_k_heads = config.linear_num_key_heads
409 - self.head_k_dim = config.linear_key_head_dim
410 - self.head_v_dim = config.linear_value_head_dim
411 - self.key_dim = self.head_k_dim * self.num_k_heads
412 - self.value_dim = self.head_v_dim * self.num_v_heads
413 -
414 - self.conv_kernel_size = config.linear_conv_kernel_dim
415 - self.layer_idx = extract_layer_index(prefix)
416 - self.activation = config.hidden_act
417 - self.act = ACT2FN[config.hidden_act]
418 - self.layer_norm_epsilon = config.rms_norm_eps
419 - self.prefix = prefix
420 -
421 - self.config = config
422 - self.model_config = vllm_config.model_config
423 - self.cache_config = vllm_config.cache_config
424 - quant_config = vllm_config.quant_config
425 - self.speculative_config = vllm_config.speculative_config
426 - self.num_spec = (
427 - self.speculative_config.num_speculative_tokens
428 - if self.speculative_config
429 - else 0
430 - )
431 -
432 - # QKV
433 - self.conv_dim = self.key_dim * 2 + self.value_dim
434 - self.conv1d = ColumnParallelLinear(
435 - input_size=self.conv_kernel_size,
436 - output_size=self.conv_dim,
437 - bias=False,
438 - prefix=f"{prefix}.conv1d",
439 - )
440 - self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
441 -
442 - # projection of the input hidden states
443 - # Qwen3-Next and Qwen3.5 has a different qkv_proj layout,
444 - # we need to create qkvz_proj adaptively here.
445 - # When create_in_proj_qkvz is False (e.g. LoRA enabled in Qwen3.5),
446 - # the subclass creates in_proj_qkv and in_proj_z separately.
447 - if create_in_proj_qkvz:
448 - self.in_proj_qkvz = self.create_qkvz_proj(
449 - hidden_size=self.hidden_size,
450 - key_dim=self.key_dim,
451 - value_dim=self.value_dim,
452 - quant_config=quant_config,
453 - prefix=f"{prefix}.in_proj_qkvz",
454 - )
455 - # ba_proj doesn't support blockwise fp8 quantization.
456 - # Qwen3-Next and Qwen3.5 have different in_proj_ba checkpoint
457 - # layouts, so we use a factory method to create the projection.
458 - self.in_proj_ba = self.create_ba_proj(
459 - hidden_size=self.hidden_size,
460 - num_v_heads=self.num_v_heads,
461 - quant_config=quant_config,
462 - prefix=f"{prefix}.in_proj_ba",
463 - )
464 -
465 - query_key_settings = (self.key_dim, 0, False)
466 - value_settings = (self.value_dim, 0, False)
467 -
468 - delattr(self.conv1d.weight, "weight_loader")
469 - set_weight_attrs(
470 - self.conv1d.weight,
471 - {
472 - "weight_loader": mamba_v2_sharded_weight_loader(
473 - [
474 - query_key_settings,
475 - query_key_settings,
476 - value_settings,
477 - ],
478 - self.tp_size,
479 - self.tp_rank,
480 - )
481 - },
482 - )
483 -
484 - # selective projection used to make dt, B and C input dependent
485 -
486 - # time step projection (discretization)
487 - # instantiate once and copy inv_dt in init_weights of PretrainedModel
488 - self.dt_bias = nn.Parameter(
489 - torch.ones(self.num_v_heads // self.tp_size),
490 - )
491 - self.A_log = nn.Parameter(
492 - torch.empty(
493 - divide(self.num_v_heads, self.tp_size),
494 - dtype=torch.float32,
495 - )
496 - )
497 -
498 - set_weight_attrs(self.A_log, {"weight_loader": sharded_weight_loader(0)})
499 - set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
500 -
501 - self.norm = RMSNormGated(
502 - self.head_v_dim,
503 - eps=self.layer_norm_epsilon,
504 - group_size=None,
505 - norm_before_gate=True,
506 - device=current_platform.current_device(),
507 - )
508 -
509 - self.out_proj = RowParallelLinear(
510 - self.value_dim,
511 - self.hidden_size,
512 - bias=False,
513 - input_is_parallel=True,
514 - quant_config=quant_config,
515 - prefix=f"{prefix}.out_proj",
516 - )
517 -
518 - self.chunk_gated_delta_rule = ChunkGatedDeltaRule()
519 - self.enable_packed_recurrent_decode = (
520 - envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE
521 - )
522 -
523 - compilation_config = get_current_vllm_config().compilation_config
524 - if prefix in compilation_config.static_forward_context:
525 - raise ValueError(f"Duplicate layer name: {prefix}")
526 - compilation_config.static_forward_context[prefix] = self
527 -
528 - def create_qkvz_proj(
529 - self,
530 - hidden_size: int,
531 - key_dim: int,
532 - value_dim: int,
533 - quant_config: QuantizationConfig | None,
534 - prefix: str,
535 - ) -> MergedColumnParallelLinear:
536 - return MergedColumnParallelLinear(
537 - input_size=hidden_size,
538 - output_sizes=[sum((key_dim, key_dim, value_dim, value_dim))],
539 - bias=False,
540 - quant_config=quant_config,
541 - prefix=prefix,
542 - )
543 -
544 - def create_ba_proj(
545 - self,
546 - hidden_size: int,
547 - num_v_heads: int,
548 - quant_config: QuantizationConfig | None,
549 - prefix: str,
550 - ) -> MergedColumnParallelLinear:
551 - # Qwen3-Next stores in_proj_ba as a single fused weight with an
552 - # interleaved GQA layout: [b_g0, a_g0, b_g1, a_g1, ...] where
553 - # each group corresponds to a key-head group. We must use a single
554 - # output shard so that ColumnParallel sharding preserves this
555 - # interleaved structure across TP ranks.
556 - # Qwen3.5 overrides this to use [num_v_heads, num_v_heads] since
557 - # its checkpoint has separate in_proj_b and in_proj_a weights.
558 - return MergedColumnParallelLinear(
559 - input_size=hidden_size,
560 - output_sizes=[num_v_heads * 2],
561 - bias=False,
562 - quant_config=quant_config,
563 - prefix=prefix,
564 - )
565 -
566 - def fix_query_key_value_ordering(
567 - self,
568 - mixed_qkvz: torch.Tensor,
569 - mixed_ba: torch.Tensor,
570 - ):
571 - """
572 - Derives `query`, `key` and `value` tensors from `mixed_qkvzba`.
573 - """
574 - new_tensor_shape_qkvz = mixed_qkvz.size()[:-1] + (
575 - self.num_k_heads // self.tp_size,
576 - (
577 - self.head_k_dim
578 - + self.head_k_dim
579 - + (self.head_v_dim + self.head_v_dim)
580 - * self.num_v_heads
581 - // self.num_k_heads
582 - ),
583 - )
584 - new_tensor_shape_ba = mixed_qkvz.size()[:-1] + (
585 - self.num_k_heads // self.tp_size,
586 - 2 * self.num_v_heads // self.num_k_heads,
587 - )
588 -
589 - mixed_qkvz = mixed_qkvz.view(*new_tensor_shape_qkvz)
590 - mixed_ba = mixed_ba.view(*new_tensor_shape_ba)
591 -
592 - split_arg_list_qkvz = [
593 - self.head_k_dim,
594 - self.head_k_dim,
595 - (self.num_v_heads // self.num_k_heads * self.head_v_dim),
596 - (self.num_v_heads // self.num_k_heads * self.head_v_dim),
597 - ]
598 - split_arg_list_ba = [
599 - self.num_v_heads // self.num_k_heads,
600 - self.num_v_heads // self.num_k_heads,
601 - ]
602 -
603 - # [b, sq, ng, (hn + hn + np/ng * hn + np/ng + np/ng)]
604 - # --> [b, sq, ng, hn], [b, sq, ng, hn], [b, sq, ng, np/ng * hn],
605 - # [b, sq, ng, np/ng * hn], [b, sq, ng, np/ng], [b, sq, ng, np/ng]
606 - (query, key, value, z) = torch.split(mixed_qkvz, split_arg_list_qkvz, dim=2)
607 - (b, a) = torch.split(mixed_ba, split_arg_list_ba, dim=2)
608 -
609 - # [b, sq, ng, np/ng * hn] -> [b, sq, np, hn]
610 - value = value.reshape(value.size(0), -1, self.head_v_dim)
611 - z = z.reshape(z.size(0), -1, self.head_v_dim)
612 - b = b.reshape(b.size(0), self.num_v_heads // self.tp_size)
613 - a = a.reshape(a.size(0), self.num_v_heads // self.tp_size)
614 -
615 - return query, key, value, z, b, a
616 -
617 - def rearrange_mixed_qkv(self, mixed_qkv):
618 - if mixed_qkv is None:
619 - return None, None, None
620 - query, key, value = torch.split(
621 - mixed_qkv,
622 - [
623 - self.key_dim // self.tp_size,
624 - self.key_dim // self.tp_size,
625 - self.value_dim // self.tp_size,
626 - ],
627 - dim=-1,
628 - )
629 - query, key = map(
630 - lambda x: rearrange(x, "l (h d) -> 1 l h d", d=self.head_k_dim),
631 - (query, key),
632 - )
633 - value = rearrange(value, "l (h d) -> 1 l h d", d=self.head_v_dim)
634 - return query.contiguous(), key.contiguous(), value.contiguous()
635 -
636 - def forward(
637 - self,
638 - hidden_states: torch.Tensor,
639 - output: torch.Tensor,
640 - ):
641 - """
642 - Forward pass with three parts:
643 - 1. Input projection
644 - 2. Core attention (custom op)
645 - 3. Output projection
646 - """
647 - num_tokens = hidden_states.size(0)
648 -
649 - # ============================================================
650 - # Part 1: Input Projection
651 - # ============================================================
652 - projected_states_qkvz, _ = self.in_proj_qkvz(hidden_states)
653 - projected_states_ba, _ = self.in_proj_ba(hidden_states)
654 - query, key, value, z, b, a = self.fix_query_key_value_ordering(
655 - projected_states_qkvz, projected_states_ba
656 - )
657 - query, key, value = map(
658 - lambda x: rearrange(x, "l p d -> l (p d)"), (query, key, value)
659 - )
660 - mixed_qkv = torch.cat((query, key, value), dim=-1)
661 -
662 - # ============================================================
663 - # Part 2: Core Attention (Custom Op)
664 - # ============================================================
665 - # Note: we should not use torch.empty here like other attention backends,
666 - # see discussions in https://github.com/vllm-project/vllm/pull/28182
667 - core_attn_out = torch.zeros(
668 - (num_tokens, self.num_v_heads // self.tp_size, self.head_v_dim),
669 - dtype=hidden_states.dtype,
670 - device=hidden_states.device,
671 - )
672 -
673 - torch.ops.vllm.gdn_attention_core(
674 - mixed_qkv,
675 - b,
676 - a,
677 - core_attn_out,
678 - self.prefix,
679 - )
680 -
681 - # ============================================================
682 - # Part 3: Output Projection
683 - # ============================================================
684 - z_shape_og = z.shape
685 - # Reshape input data into 2D tensor
686 - core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1])
687 - z = z.reshape(-1, z.shape[-1])
688 - core_attn_out = self.norm(core_attn_out, z)
689 - core_attn_out = core_attn_out.reshape(z_shape_og)
690 - core_attn_out = rearrange(core_attn_out, "... h d -> ... (h d)")
691 - output[:num_tokens], _ = self.out_proj(core_attn_out)
692 -
693 - def _warmup_prefill_kernels(self, mixed_qkv: torch.Tensor) -> None:
694 - """Warm up GDN prefill kernels during V1 profiling.
695 -
696 - During V1 profile runs, ``_forward_core`` returns early because
697 - ``attn_metadata`` is ``None``, so the autotuned kernels used by
698 - ``chunk_gated_delta_rule`` (e.g. ``solve_tril``,
699 - ``chunk_scaled_dot_kkt``) are never invoked. After profiling,
700 - vLLM allocates KV cache using most of the remaining GPU memory.
701 - When the first real inference triggers the autotuner it OOMs
702 - because there is not enough memory left for benchmarking.
703 -
704 - This method runs minimal forward passes through
705 - ``chunk_gated_delta_rule`` with small dummy tensors to force
706 - autotuning while GPU memory is still plentiful. The autotuner
707 - results are cached globally, so only the first layer incurs
708 - actual benchmarking cost.
709 -
710 - Most kernels use a fixed ``BT = chunk_size`` (64), but
711 - ``chunk_fwd_kernel_o`` recomputes ``BT`` from the sequence
712 - length: ``min(64, max(16, next_power_of_2(T)))``. Since ``BT``
713 - is part of its autotune key, we run warmup passes with T = 16,
714 - 32, and 64 to cover all possible ``BT`` values.
715 -
716 - The decode path uses ``fused_sigmoid_gating_delta_rule_update``
717 - which has fixed kernel parameters (no autotuning), so only the
718 - prefill (chunked) path needs warming up.
719 - """
720 - if hasattr(self, "_prefill_kernels_warmed_up"):
721 - return
722 - self._prefill_kernels_warmed_up = True
723 -
724 - device = mixed_qkv.device
725 - dtype = mixed_qkv.dtype
726 - num_k_heads = self.num_k_heads // self.tp_size
727 - num_v_heads = self.num_v_heads // self.tp_size
728 - _, state_dtype = self.get_state_dtype()
729 -
730 - # Run warmup for each possible BT value of chunk_fwd_kernel_o:
731 - # T=16 → BT=16, T=32 → BT=32, T=64 → BT=64.
732 - # Other kernels always use BT=chunk_size(64), so their autotune
733 - # cache is populated on the first pass and reused thereafter.
734 - for T in (16, 32, 64):
735 - q = torch.randn(
736 - 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
737 - )
738 - k = torch.randn(
739 - 1, T, num_k_heads, self.head_k_dim, device=device, dtype=dtype
740 - )
741 - v = torch.randn(
742 - 1, T, num_v_heads, self.head_v_dim, device=device, dtype=dtype
743 - )
744 - # NOTE: g and beta must have the same dtypes as during
745 - # inference, so we construct them with the same function
746 - # (fused_gdn_gating). dummy_a and dummy_b are throwaway
747 - # inputs required by that function.
748 - dummy_a = torch.randn(T, num_v_heads, device=device, dtype=dtype)
749 - dummy_b = torch.randn(T, num_v_heads, device=device, dtype=dtype)
750 - g, beta = fused_gdn_gating(self.A_log, dummy_a, dummy_b, self.dt_bias)
751 - state = torch.zeros(
752 - 1,
753 - num_v_heads,
754 - self.head_v_dim,
755 - self.head_k_dim,
756 - device=device,
757 - dtype=state_dtype,
758 - )
759 - cu_seqlens = torch.tensor([0, T], device=device, dtype=torch.int32)
760 -
761 - try:
762 - self.chunk_gated_delta_rule(
763 - q=q,
764 - k=k,
765 - v=v,
766 - g=g,
767 - beta=beta,
768 - initial_state=state,
769 - output_final_state=True,
770 - cu_seqlens=cu_seqlens,
771 - use_qk_l2norm_in_kernel=True,
772 - )
773 - except Exception:
774 - logger.warning(
775 - "GDN prefill kernel warmup (T=%d) failed for "
776 - "layer %s. First inference may OOM due to "
777 - "autotuner.",
778 - T,
779 - self.prefix,
780 - exc_info=True,
781 - )
782 - else:
783 - logger.debug(
784 - "GDN prefill kernel warmup (T=%d) completed for layer %s",
785 - T,
786 - self.prefix,
787 - )
788 - finally:
789 - del q, k, v, dummy_a, dummy_b, g, beta, state, cu_seqlens
790 -
791 - torch.accelerator.empty_cache()
792 -
793 - def _forward_core(
794 - self,
795 - mixed_qkv: torch.Tensor,
796 - b: torch.Tensor,
797 - a: torch.Tensor,
798 - core_attn_out: torch.Tensor,
799 - ):
800 - forward_context = get_forward_context()
801 - attn_metadata: AttentionMetadata = forward_context.attn_metadata
802 -
803 - if attn_metadata is None:
804 - # V1 profile run — warm up prefill kernels so that
805 - # autotuning completes before KV cache allocation.
806 - self._warmup_prefill_kernels(mixed_qkv)
807 - return
808 -
809 - assert isinstance(attn_metadata, dict)
810 - attn_metadata = attn_metadata[self.prefix]
811 - assert isinstance(attn_metadata, GDNAttentionMetadata)
812 -
813 - if (
814 - self.enable_packed_recurrent_decode
815 - and attn_metadata.spec_sequence_masks is None
816 - and attn_metadata.num_prefills == 0
817 - and attn_metadata.num_decodes > 0
818 - ):
819 - return self._forward_core_decode_non_spec(
820 - mixed_qkv=mixed_qkv,
821 - b=b,
822 - a=a,
823 - core_attn_out=core_attn_out,
824 - attn_metadata=attn_metadata,
825 - )
826 -
827 - has_initial_state = attn_metadata.has_initial_state
828 - spec_query_start_loc = attn_metadata.spec_query_start_loc
829 - non_spec_query_start_loc = attn_metadata.non_spec_query_start_loc
830 - spec_sequence_masks = attn_metadata.spec_sequence_masks
831 - spec_token_indx = attn_metadata.spec_token_indx
832 - non_spec_token_indx = attn_metadata.non_spec_token_indx
833 - spec_state_indices_tensor = attn_metadata.spec_state_indices_tensor # noqa: E501
834 - non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
835 - self_kv_cache = self.kv_cache
836 - conv_state = self_kv_cache[0].transpose(-1, -2)
837 - ssm_state = self_kv_cache[1]
838 - num_actual_tokens = attn_metadata.num_actual_tokens
839 - num_accepted_tokens = attn_metadata.num_accepted_tokens
840 -
841 - mixed_qkv = mixed_qkv[:num_actual_tokens]
842 - b = b[:num_actual_tokens]
843 - a = a[:num_actual_tokens]
844 -
845 - # 1. Convolution sequence transformation
846 - conv_weights = self.conv1d.weight.view(
847 - self.conv1d.weight.size(0), self.conv1d.weight.size(2)
848 - )
849 -
850 - if spec_sequence_masks is not None:
851 - if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
852 - mixed_qkv_spec = mixed_qkv
853 - mixed_qkv_non_spec = None
854 - else:
855 - mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
856 - mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
857 - else:
858 - mixed_qkv_spec = None
859 - mixed_qkv_non_spec = mixed_qkv
860 -
861 - # 1.1: Process the multi-query part
862 - if spec_sequence_masks is not None:
863 - mixed_qkv_spec = causal_conv1d_update(
864 - mixed_qkv_spec,
865 - conv_state,
866 - conv_weights,
867 - self.conv1d.bias,
868 - self.activation,
869 - conv_state_indices=spec_state_indices_tensor[:, 0][
870 - : attn_metadata.num_spec_decodes
871 - ],
872 - num_accepted_tokens=num_accepted_tokens,
873 - query_start_loc=spec_query_start_loc,
874 - max_query_len=spec_state_indices_tensor.size(-1),
875 - validate_data=False,
876 - )
877 -
878 - # 1.2: Process the remaining part
879 - if attn_metadata.num_prefills > 0:
880 - mixed_qkv_non_spec_T = mixed_qkv_non_spec.transpose(0, 1)
881 - # - "cache_indices" updates the conv_state cache in positions
882 - # pointed to by "state_indices_tensor"
883 - mixed_qkv_non_spec = causal_conv1d_fn(
884 - mixed_qkv_non_spec_T,
885 - conv_weights,
886 - self.conv1d.bias,
887 - activation=self.activation,
888 - conv_states=conv_state,
889 - has_initial_state=has_initial_state,
890 - cache_indices=non_spec_state_indices_tensor,
891 - query_start_loc=non_spec_query_start_loc,
892 - metadata=attn_metadata,
893 - ).transpose(0, 1)
894 - elif attn_metadata.num_decodes > 0:
895 - mixed_qkv_non_spec = causal_conv1d_update(
896 - mixed_qkv_non_spec,
897 - conv_state,
898 - conv_weights,
899 - self.conv1d.bias,
900 - self.activation,
901 - conv_state_indices=non_spec_state_indices_tensor[
902 - : attn_metadata.num_actual_tokens
903 - ],
904 - validate_data=True,
905 - )
906 - else:
907 - mixed_qkv_non_spec = None
908 -
909 - query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
910 - query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
911 - mixed_qkv_non_spec
912 - )
913 -
914 - if attn_metadata.num_prefills > 0:
915 - g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias)
916 - if spec_sequence_masks is not None:
917 - g_non_spec = g.index_select(1, non_spec_token_indx)
918 - beta_non_spec = beta.index_select(1, non_spec_token_indx)
919 - else:
920 - g_non_spec = g
921 - beta_non_spec = beta
922 - else:
923 - g_non_spec = None
924 - beta_non_spec = None
925 -
926 - # 2. Recurrent attention
927 -
928 - # 2.1: Process the multi-query part
929 - if spec_sequence_masks is not None:
930 - core_attn_out_spec, last_recurrent_state = (
931 - fused_sigmoid_gating_delta_rule_update(
932 - A_log=self.A_log,
933 - a=a,
934 - b=b,
935 - dt_bias=self.dt_bias,
936 - q=query_spec,
937 - k=key_spec,
938 - v=value_spec,
939 - initial_state=ssm_state,
940 - inplace_final_state=True,
941 - cu_seqlens=spec_query_start_loc[
942 - : attn_metadata.num_spec_decodes + 1
943 - ],
944 - ssm_state_indices=spec_state_indices_tensor,
945 - num_accepted_tokens=num_accepted_tokens,
946 - use_qk_l2norm_in_kernel=True,
947 - )
948 - )
949 - else:
950 - core_attn_out_spec, last_recurrent_state = None, None
951 -
952 - # 2.2: Process the remaining part
953 - if attn_metadata.num_prefills > 0:
954 - initial_state = ssm_state[non_spec_state_indices_tensor].contiguous()
955 - initial_state[~has_initial_state, ...] = 0
956 - (
957 - core_attn_out_non_spec,
958 - last_recurrent_state,
959 - ) = self.chunk_gated_delta_rule(
960 - q=query_non_spec,
961 - k=key_non_spec,
962 - v=value_non_spec,
963 - g=g_non_spec,
964 - beta=beta_non_spec,
965 - initial_state=initial_state,
966 - output_final_state=True,
967 - cu_seqlens=non_spec_query_start_loc,
968 - use_qk_l2norm_in_kernel=True,
969 - )
970 - # Init cache
971 - ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(
972 - ssm_state.dtype
973 - )
974 - elif attn_metadata.num_decodes > 0:
975 - core_attn_out_non_spec, last_recurrent_state = (
976 - fused_sigmoid_gating_delta_rule_update(
977 - A_log=self.A_log,
978 - a=a,
979 - b=b,
980 - dt_bias=self.dt_bias,
981 - q=query_non_spec,
982 - k=key_non_spec,
983 - v=value_non_spec,
984 - initial_state=ssm_state,
985 - inplace_final_state=True,
986 - cu_seqlens=non_spec_query_start_loc[
987 - : attn_metadata.num_decodes + 1
988 - ],
989 - ssm_state_indices=non_spec_state_indices_tensor,
990 - use_qk_l2norm_in_kernel=True,
991 - )
992 - )
993 - else:
994 - core_attn_out_non_spec, last_recurrent_state = None, None
995 -
996 - # 3. Merge core attention output
997 - if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
998 - merged_out = torch.empty(
999 - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
1000 - dtype=core_attn_out_non_spec.dtype,
1001 - device=core_attn_out_non_spec.device,
1002 - )
1003 - merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
1004 - merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
1005 - core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
1006 - elif spec_sequence_masks is not None:
1007 - core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
1008 - else:
1009 - core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
1010 -
1011 - def _forward_core_decode_non_spec(
1012 - self,
1013 - mixed_qkv: torch.Tensor,
1014 - b: torch.Tensor,
1015 - a: torch.Tensor,
1016 - core_attn_out: torch.Tensor,
1017 - attn_metadata: GDNAttentionMetadata,
1018 - ):
1019 - """
1020 - Core attention computation with a packed non-spec decode fast path.
1021 - """
1022 - non_spec_state_indices_tensor = attn_metadata.non_spec_state_indices_tensor # noqa: E501
1023 - self_kv_cache = self.kv_cache
1024 - conv_state = self_kv_cache[0].transpose(-1, -2)
1025 - ssm_state = self_kv_cache[1]
1026 - num_actual_tokens = attn_metadata.num_actual_tokens
1027 -
1028 - mixed_qkv = mixed_qkv[:num_actual_tokens]
1029 - b = b[:num_actual_tokens]
1030 - a = a[:num_actual_tokens]
1031 -
1032 - conv_weights = self.conv1d.weight.view(
1033 - self.conv1d.weight.size(0), self.conv1d.weight.size(2)
1034 - )
1035 - mixed_qkv_non_spec = causal_conv1d_update(
1036 - mixed_qkv,
1037 - conv_state,
1038 - conv_weights,
1039 - self.conv1d.bias,
1040 - self.activation,
1041 - conv_state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
1042 - validate_data=False,
1043 - )
1044 - out_buf = core_attn_out[:num_actual_tokens].unsqueeze(1)
1045 - fused_recurrent_gated_delta_rule_packed_decode(
1046 - mixed_qkv=mixed_qkv_non_spec,
1047 - a=a,
1048 - b=b,
1049 - A_log=self.A_log,
1050 - dt_bias=self.dt_bias,
1051 - scale=self.head_k_dim**-0.5,
1052 - initial_state=ssm_state,
1053 - out=out_buf,
1054 - ssm_state_indices=non_spec_state_indices_tensor[:num_actual_tokens],
1055 - use_qk_l2norm_in_kernel=True,
1056 - )
1057 - return
1058 -
1059 -
200 class Qwen3NextAttention(nn.Module):
201 def __init__(
202 self,
@@ -1192,10 +332,11 @@ class Qwen3NextDecoderLayer(nn.Module):
332 self.layer_idx = extract_layer_index(prefix)
333
334 if self.layer_type == "linear_attention":
1195 - self.linear_attn = Qwen3NextGatedDeltaNet(
335 + self.linear_attn = GatedDeltaNetAttention(
336 config,
337 vllm_config=vllm_config,
338 prefix=f"{prefix}.linear_attn",
339 + gqa_interleaved_layout=True,
340 )
341 elif self.layer_type == "full_attention":
342 self.self_attn = Qwen3NextAttention(
@@ -1669,116 +810,3 @@ class Qwen3NextForCausalLM(
810
811 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
812 return self.model.get_expert_mapping()
1672 -
1673 -
1674 -def gdn_attention_core(
1675 - mixed_qkv: torch.Tensor,
1676 - b: torch.Tensor,
1677 - a: torch.Tensor,
1678 - core_attn_out: torch.Tensor,
1679 - layer_name: str,
1680 -) -> None:
1681 - """
1682 - Custom op for the core attention computation.
1683 - Only handles the convolution + recurrent attention part.
1684 - Input/output projections are handled outside this op.
1685 - """
1686 - forward_context: ForwardContext = get_forward_context()
1687 - self = forward_context.no_compile_layers[layer_name]
1688 - self._forward_core(
1689 - mixed_qkv=mixed_qkv,
1690 - b=b,
1691 - a=a,
1692 - core_attn_out=core_attn_out,
1693 - )
1694 -
1695 -
1696 -def gdn_attention_core_fake(
1697 - mixed_qkv: torch.Tensor,
1698 - b: torch.Tensor,
1699 - a: torch.Tensor,
1700 - core_attn_out: torch.Tensor,
1701 - layer_name: str,
1702 -) -> None:
1703 - """Fake implementation for torch.compile."""
1704 - return
1705 -
1706 -
1707 -direct_register_custom_op(
1708 - op_name="gdn_attention_core",
1709 - op_func=gdn_attention_core,
1710 - mutates_args=["core_attn_out"],
1711 - fake_impl=gdn_attention_core_fake,
1712 -)
1713 -
1714 -
1715 -@triton.jit
1716 -def fused_gdn_gating_kernel(
1717 - g,
1718 - beta_output,
1719 - A_log,
1720 - a,
1721 - b,
1722 - dt_bias,
1723 - seq_len,
1724 - NUM_HEADS: tl.constexpr,
1725 - beta: tl.constexpr,
1726 - threshold: tl.constexpr,
1727 - BLK_HEADS: tl.constexpr,
1728 -):
1729 - i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
1730 - head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS)
1731 - off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off
1732 - mask = head_off < NUM_HEADS
1733 - blk_A_log = tl.load(A_log + head_off, mask=mask)
1734 - blk_a = tl.load(a + off, mask=mask)
1735 - blk_b = tl.load(b + off, mask=mask)
1736 - blk_bias = tl.load(dt_bias + head_off, mask=mask)
1737 - # If the model is loaded in fp16, without the .float() here, A might be -inf
1738 - x = blk_a.to(tl.float32) + blk_bias.to(tl.float32)
1739 - softplus_x = tl.where(
1740 - beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
1741 - )
1742 - blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x
1743 - tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask)
1744 - # compute beta_output = sigmoid(b)
1745 - blk_beta_output = tl.sigmoid(blk_b.to(tl.float32))
1746 - tl.store(
1747 - beta_output + off, blk_beta_output.to(beta_output.dtype.element_ty), mask=mask
1748 - )
1749 -
1750 -
1751 -def fused_gdn_gating(
1752 - A_log: torch.Tensor,
1753 - a: torch.Tensor,
1754 - b: torch.Tensor,
1755 - dt_bias: torch.Tensor,
1756 - beta: float = 1.0,
1757 - threshold: float = 20.0,
1758 -) -> tuple[torch.Tensor, torch.Tensor]:
1759 - """
1760 - Fused computation of g and beta for Gated Delta Net.
1761 - g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias)
1762 - beta_output = b.sigmoid()
1763 - TODO maybe use torch.compile to replace this triton kernel
1764 - """
1765 - batch, num_heads = a.shape
1766 - seq_len = 1
1767 - grid = (batch, seq_len, triton.cdiv(num_heads, 8))
1768 - g = torch.empty(1, batch, num_heads, dtype=torch.float32, device=a.device)
1769 - beta_output = torch.empty(1, batch, num_heads, dtype=b.dtype, device=b.device)
1770 - fused_gdn_gating_kernel[grid](
1771 - g,
1772 - beta_output,
1773 - A_log,
1774 - a,
1775 - b,
1776 - dt_bias,
1777 - seq_len,
1778 - num_heads,
1779 - beta,
1780 - threshold,
1781 - 8,
1782 - num_warps=1,
1783 - )
1784 - return g, beta_output