@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dc6b578466
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 7, 01:30 UTC

[Kernel] Add fused_sigmoid_gating_delta_rule_update kernel for Qwen3 Next (#35777)

Signed-off-by: Xin Yang <xyangx@amazon.com>

Xin Yang committed Mar 8, 2026 at 23:41 UTC dc6b57846686206d6d77fe788f71ab7fe8e568ab
4 files changed +509 -31
tests/kernels/test_fused_sigmoid_gating_delta_rule.py new
+196
@@ -0,0 +1,196 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +import torch.nn.functional as F
7 +
8 +from vllm.model_executor.layers.fla.ops import (
9 + fused_recurrent_gated_delta_rule,
10 + fused_sigmoid_gating_delta_rule_update,
11 +)
12 +from vllm.platforms import current_platform
13 +from vllm.utils.torch_utils import set_random_seed
14 +
15 +DEVICE = current_platform.device_type
16 +
17 +
18 +@pytest.mark.parametrize("tp_size", [1])
19 +@pytest.mark.parametrize("num_reqs", [1, 2, 4])
20 +@pytest.mark.parametrize("num_k_heads", [16])
21 +@pytest.mark.parametrize("num_v_heads", [32])
22 +@pytest.mark.parametrize("head_k_dim", [128])
23 +@pytest.mark.parametrize("head_v_dim", [128])
24 +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
25 +def test_fused_sigmoid_gating_delta_rule_update_non_spec(
26 + tp_size: int,
27 + num_reqs: int,
28 + num_k_heads: int,
29 + num_v_heads: int,
30 + head_k_dim: int,
31 + head_v_dim: int,
32 + dtype: torch.dtype,
33 +) -> None:
34 + torch.set_default_device(DEVICE)
35 + set_random_seed(0)
36 + key_dim = head_k_dim * num_k_heads
37 + value_dim = head_v_dim * num_v_heads
38 + mixed_qkv_dim = (key_dim * 2 + value_dim) // tp_size
39 + seq_len = 1 # seq_len is 1 for decode
40 + num_tokens = num_reqs * seq_len
41 + total_entries = num_tokens * 2
42 +
43 + mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype)
44 + query, key, value = torch.split(
45 + mixed_qkv,
46 + [
47 + key_dim // tp_size,
48 + key_dim // tp_size,
49 + value_dim // tp_size,
50 + ],
51 + dim=-1,
52 + )
53 + query = query.view(1, num_tokens, num_k_heads, head_k_dim)
54 + key = key.view(1, num_tokens, num_k_heads, head_k_dim)
55 + value = value.view(1, num_tokens, num_v_heads, head_v_dim)
56 +
57 + A_log = torch.rand(num_v_heads // tp_size, dtype=dtype)
58 + dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
59 + a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
60 + b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
61 + ssm_state = torch.rand(
62 + total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
63 + )
64 + state_indices = torch.randperm(total_entries, dtype=torch.int32)[:num_tokens]
65 + cu_seqlens = torch.arange(0, num_tokens + 1, dtype=torch.int32)
66 +
67 + beta = b.sigmoid()
68 + g = -A_log.float().exp() * F.softplus(a.float() + dt_bias)
69 + core_attn_out_ref, last_recurrent_state_ref = fused_recurrent_gated_delta_rule(
70 + q=query,
71 + k=key,
72 + v=value,
73 + g=g.unsqueeze(0),
74 + beta=beta.unsqueeze(0),
75 + initial_state=ssm_state.clone(),
76 + inplace_final_state=True,
77 + ssm_state_indices=state_indices,
78 + cu_seqlens=cu_seqlens,
79 + use_qk_l2norm_in_kernel=True,
80 + )
81 +
82 + core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
83 + A_log=A_log,
84 + a=a,
85 + b=b,
86 + dt_bias=dt_bias,
87 + q=query,
88 + k=key,
89 + v=value,
90 + initial_state=ssm_state,
91 + inplace_final_state=True,
92 + ssm_state_indices=state_indices,
93 + cu_seqlens=cu_seqlens,
94 + use_qk_l2norm_in_kernel=True,
95 + )
96 +
97 + torch.testing.assert_close(core_attn_out, core_attn_out_ref, atol=1e-2, rtol=1e-2)
98 + torch.testing.assert_close(
99 + last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2
100 + )
101 +
102 +
103 +@pytest.mark.parametrize("tp_size", [1])
104 +@pytest.mark.parametrize("num_reqs", [1, 2, 4])
105 +@pytest.mark.parametrize("num_k_heads", [16])
106 +@pytest.mark.parametrize("num_v_heads", [32])
107 +@pytest.mark.parametrize("head_k_dim", [128])
108 +@pytest.mark.parametrize("head_v_dim", [128])
109 +@pytest.mark.parametrize("num_speculative_tokens", [1, 3])
110 +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
111 +def test_fused_sigmoid_gating_delta_rule_update_spec(
112 + tp_size: int,
113 + num_reqs: int,
114 + num_k_heads: int,
115 + num_v_heads: int,
116 + head_k_dim: int,
117 + head_v_dim: int,
118 + num_speculative_tokens: int,
119 + dtype: torch.dtype,
120 +) -> None:
121 + torch.set_default_device(DEVICE)
122 + set_random_seed(0)
123 + key_dim = head_k_dim * num_k_heads
124 + value_dim = head_v_dim * num_v_heads
125 + mixed_qkv_dim = (key_dim * 2 + value_dim) // tp_size
126 + num_tokens = num_reqs * (num_speculative_tokens + 1)
127 + total_entries = num_tokens * 2
128 +
129 + mixed_qkv = torch.rand(num_tokens, mixed_qkv_dim, dtype=dtype)
130 + query, key, value = torch.split(
131 + mixed_qkv,
132 + [
133 + key_dim // tp_size,
134 + key_dim // tp_size,
135 + value_dim // tp_size,
136 + ],
137 + dim=-1,
138 + )
139 + query = query.view(1, num_tokens, num_k_heads, head_k_dim)
140 + key = key.view(1, num_tokens, num_k_heads, head_k_dim)
141 + value = value.view(1, num_tokens, num_v_heads, head_v_dim)
142 +
143 + A_log = torch.rand(num_v_heads // tp_size, dtype=dtype)
144 + dt_bias = torch.rand(num_v_heads // tp_size, dtype=dtype)
145 + a = torch.rand(num_tokens, num_v_heads, dtype=dtype)
146 + b = torch.rand(num_tokens, num_v_heads, dtype=dtype)
147 + ssm_state = torch.rand(
148 + total_entries, num_v_heads, head_k_dim, head_v_dim, dtype=dtype
149 + )
150 + state_indices = torch.randperm(
151 + total_entries,
152 + dtype=torch.int32,
153 + )[:num_tokens].view(num_reqs, num_speculative_tokens + 1)
154 + num_accepted_tokens = torch.randint(
155 + 1, num_speculative_tokens + 1, (num_reqs,), dtype=torch.int32
156 + )
157 + cu_seqlens = torch.arange(
158 + 0, num_tokens + 1, num_speculative_tokens + 1, dtype=torch.int32
159 + )
160 +
161 + beta = b.sigmoid()
162 + g = -A_log.float().exp() * F.softplus(a.float() + dt_bias)
163 + core_attn_out_ref, last_recurrent_state_ref = fused_recurrent_gated_delta_rule(
164 + q=query,
165 + k=key,
166 + v=value,
167 + g=g.unsqueeze(0),
168 + beta=beta.unsqueeze(0),
169 + initial_state=ssm_state.clone(),
170 + inplace_final_state=True,
171 + ssm_state_indices=state_indices,
172 + cu_seqlens=cu_seqlens,
173 + num_accepted_tokens=num_accepted_tokens,
174 + use_qk_l2norm_in_kernel=True,
175 + )
176 +
177 + core_attn_out, last_recurrent_state = fused_sigmoid_gating_delta_rule_update(
178 + A_log=A_log,
179 + a=a,
180 + b=b,
181 + dt_bias=dt_bias,
182 + q=query,
183 + k=key,
184 + v=value,
185 + initial_state=ssm_state,
186 + inplace_final_state=True,
187 + ssm_state_indices=state_indices,
188 + cu_seqlens=cu_seqlens,
189 + num_accepted_tokens=num_accepted_tokens,
190 + use_qk_l2norm_in_kernel=True,
191 + )
192 +
193 + torch.testing.assert_close(core_attn_out, core_attn_out_ref, atol=1e-2, rtol=1e-2)
194 + torch.testing.assert_close(
195 + last_recurrent_state, last_recurrent_state_ref, atol=1e-2, rtol=1e-2
196 + )
vllm/model_executor/layers/fla/ops/__init__.py
+2
@@ -8,10 +8,12 @@
8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9 from .chunk import chunk_gated_delta_rule
10 from .fused_recurrent import fused_recurrent_gated_delta_rule
11 +from .fused_sigmoid_gating import fused_sigmoid_gating_delta_rule_update
12 from .layernorm_guard import RMSNormGated
13
14 __all__ = [
15 "RMSNormGated",
16 "chunk_gated_delta_rule",
17 "fused_recurrent_gated_delta_rule",
18 + "fused_sigmoid_gating_delta_rule_update",
19 ]
vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py new
+279
@@ -0,0 +1,279 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +# SPDX-FileCopyrightText: Songlin Yang, Yu Zhang
4 +#
5 +# This file contains code copied from the flash-linear-attention project.
6 +# The original source code was licensed under the MIT license and included
7 +# the following copyright notice:
8 +# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9 +
10 +import torch
11 +
12 +from vllm.triton_utils import tl, triton
13 +
14 +
15 +@triton.heuristics(
16 + {
17 + "USE_INITIAL_STATE": lambda args: args["h0"] is not None,
18 + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,
19 + "IS_CONTINUOUS_BATCHING": lambda args: args["ssm_state_indices"] is not None,
20 + "IS_SPEC_DECODING": lambda args: args["num_accepted_tokens"] is not None,
21 + }
22 +)
23 +@triton.jit(do_not_specialize=["N", "T"])
24 +def fused_sigmoid_gating_delta_rule_update_kernel(
25 + A_log,
26 + a,
27 + b,
28 + dt_bias,
29 + beta,
30 + threshold,
31 + q,
32 + k,
33 + v,
34 + o,
35 + h0,
36 + ht,
37 + cu_seqlens,
38 + ssm_state_indices,
39 + num_accepted_tokens,
40 + scale,
41 + N: tl.int64, # num of sequences
42 + T: tl.int64, # num of tokens
43 + B: tl.constexpr,
44 + H: tl.constexpr,
45 + HV: tl.constexpr,
46 + K: tl.constexpr,
47 + V: tl.constexpr,
48 + BK: tl.constexpr,
49 + BV: tl.constexpr,
50 + stride_init_state_token: tl.constexpr,
51 + stride_final_state_token: tl.constexpr,
52 + stride_indices_seq: tl.constexpr,
53 + stride_indices_tok: tl.constexpr,
54 + USE_INITIAL_STATE: tl.constexpr, # whether to use initial state
55 + INPLACE_FINAL_STATE: tl.constexpr, # whether to store final state inplace
56 + USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
57 + IS_VARLEN: tl.constexpr,
58 + IS_CONTINUOUS_BATCHING: tl.constexpr,
59 + IS_SPEC_DECODING: tl.constexpr,
60 + IS_KDA: tl.constexpr,
61 +):
62 + i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
63 + i_n, i_hv = i_nh // HV, i_nh % HV
64 + i_h = i_hv // (HV // H)
65 + if IS_VARLEN:
66 + bos, eos = (
67 + tl.load(cu_seqlens + i_n).to(tl.int64),
68 + tl.load(cu_seqlens + i_n + 1).to(tl.int64),
69 + )
70 + all = T
71 + T = eos - bos
72 + else:
73 + bos, eos = i_n * T, i_n * T + T
74 + all = B * T
75 +
76 + if T == 0:
77 + # no tokens to process for this sequence
78 + return
79 +
80 + o_k = i_k * BK + tl.arange(0, BK)
81 + o_v = i_v * BV + tl.arange(0, BV)
82 +
83 + p_q = q + (bos * H + i_h) * K + o_k
84 + p_k = k + (bos * H + i_h) * K + o_k
85 + p_v = v + (bos * HV + i_hv) * V + o_v
86 +
87 + p_A_log = A_log + i_hv
88 + if not IS_KDA:
89 + p_a = a + bos * HV + i_hv
90 + p_dt_bias = dt_bias + i_hv
91 + else:
92 + p_a = a + (bos * HV + i_hv) * K + o_k
93 + p_dt_bias = dt_bias + i_hv * K + o_k
94 +
95 + p_b = b + bos * HV + i_hv
96 + p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v
97 +
98 + mask_k = o_k < K
99 + mask_v = o_v < V
100 + mask_h = mask_v[:, None] & mask_k[None, :]
101 +
102 + b_h = tl.zeros([BV, BK], dtype=tl.float32)
103 + if USE_INITIAL_STATE:
104 + if IS_CONTINUOUS_BATCHING:
105 + if IS_SPEC_DECODING:
106 + i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
107 + else:
108 + i_t = 0
109 + # Load state index and check for PAD_SLOT_ID (-1)
110 + state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
111 + tl.int64
112 + )
113 + # Skip if state index is invalid (PAD_SLOT_ID = -1)
114 + if state_idx < 0:
115 + return
116 + p_h0 = h0 + state_idx * stride_init_state_token
117 + else:
118 + p_h0 = h0 + bos * HV * V * K
119 + p_h0 = p_h0 + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
120 + b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
121 +
122 + for i_t in range(0, T):
123 + b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
124 + b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
125 + b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
126 + b_b = tl.load(p_b).to(tl.float32)
127 +
128 + # If the model is loaded in fp16, without the .float() here, A might be -inf
129 + x = tl.load(p_a).to(tl.float32) + tl.load(p_dt_bias).to(tl.float32)
130 + softplus_x = tl.where(
131 + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x
132 + )
133 + b_g = -tl.exp(tl.load(p_A_log).to(tl.float32)) * softplus_x
134 +
135 + # compute beta_output = sigmoid(b)
136 + b_beta = tl.sigmoid(b_b.to(tl.float32))
137 +
138 + if USE_QK_L2NORM_IN_KERNEL:
139 + b_q = b_q * (tl.rsqrt(tl.sum(b_q * b_q) + 1e-6))
140 + b_k = b_k * (tl.rsqrt(tl.sum(b_k * b_k) + 1e-6))
141 + b_q = b_q * scale
142 + # [BV, BK]
143 + if not IS_KDA:
144 + b_h *= tl.exp(b_g)
145 + else:
146 + b_h *= tl.exp(b_g[None, :])
147 + # [BV]
148 + b_v -= tl.sum(b_h * b_k[None, :], 1)
149 + b_v *= b_beta
150 + # [BV, BK]
151 + b_h += b_v[:, None] * b_k[None, :]
152 + # [BV]
153 + b_o = tl.sum(b_h * b_q[None, :], 1)
154 + tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
155 +
156 + # keep the states for multi-query tokens
157 + if INPLACE_FINAL_STATE:
158 + # Load state index and check for PAD_SLOT_ID (-1)
159 + final_state_idx = tl.load(
160 + ssm_state_indices + i_n * stride_indices_seq + i_t
161 + ).to(tl.int64)
162 + # Only store if state index is valid (not PAD_SLOT_ID)
163 + if final_state_idx >= 0:
164 + p_ht = ht + final_state_idx * stride_final_state_token
165 + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
166 + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
167 + else:
168 + p_ht = ht + (bos + i_t) * stride_final_state_token
169 + p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
170 + tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
171 +
172 + # Update pointers for next timestep
173 + p_q += H * K
174 + p_k += H * K
175 + p_o += HV * V
176 + p_v += HV * V
177 + p_b += HV
178 + p_a += HV
179 +
180 +
181 +def fused_sigmoid_gating_delta_rule_update(
182 + A_log: torch.Tensor,
183 + a: torch.Tensor,
184 + b: torch.Tensor,
185 + dt_bias: torch.Tensor,
186 + q: torch.Tensor,
187 + k: torch.Tensor,
188 + v: torch.Tensor,
189 + beta: float = 1.0,
190 + threshold: float = 20.0,
191 + scale: float = None,
192 + initial_state: torch.Tensor = None,
193 + inplace_final_state: bool = True,
194 + cu_seqlens: torch.LongTensor | None = None,
195 + ssm_state_indices: torch.Tensor | None = None,
196 + num_accepted_tokens: torch.Tensor | None = None,
197 + use_qk_l2norm_in_kernel: bool = False,
198 + is_kda: bool = False,
199 +):
200 + """
201 + Fused triton implementation of sigmoid gating delta rule update.
202 + This function uses a single fused kernel that combines both sigmoid gating
203 + computation and the recurrent delta rule update for better performance.
204 + """
205 + B, T, H, K, V = *k.shape, v.shape[-1]
206 + HV = v.shape[2]
207 + N = B if cu_seqlens is None else len(cu_seqlens) - 1
208 + BK, BV = triton.next_power_of_2(K), min(triton.next_power_of_2(V), 32)
209 + NK, NV = triton.cdiv(K, BK), triton.cdiv(V, BV)
210 + assert NK == 1, "NK > 1 is not supported yet"
211 + num_stages = 3
212 + num_warps = 4
213 +
214 + if cu_seqlens is not None and q.shape[0] != 1:
215 + raise ValueError(
216 + f"The batch size is expected to be 1 rather than {q.shape[0]}"
217 + f" when using `cu_seqlens`. Please flatten variable-length"
218 + f" inputs before processing."
219 + )
220 + if scale is None:
221 + scale = k.shape[-1] ** -0.5
222 + else:
223 + assert scale > 0, "scale must be positive"
224 +
225 + o = q.new_empty(NK, *v.shape)
226 + if inplace_final_state:
227 + final_state = initial_state
228 + else:
229 + final_state = q.new_empty(T, HV, V, K, dtype=initial_state.dtype)
230 +
231 + stride_init_state_token = initial_state.stride(0)
232 + stride_final_state_token = final_state.stride(0)
233 +
234 + if ssm_state_indices is None:
235 + stride_indices_seq, stride_indices_tok = 1, 1
236 + elif ssm_state_indices.ndim == 1:
237 + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1
238 + else:
239 + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()
240 +
241 + grid = (NK, NV, N * HV)
242 + fused_sigmoid_gating_delta_rule_update_kernel[grid](
243 + A_log=A_log,
244 + a=a.contiguous(),
245 + b=b.contiguous(),
246 + dt_bias=dt_bias,
247 + beta=beta,
248 + threshold=threshold,
249 + q=q.contiguous(),
250 + k=k.contiguous(),
251 + v=v.contiguous(),
252 + o=o,
253 + h0=initial_state,
254 + ht=final_state,
255 + cu_seqlens=cu_seqlens,
256 + ssm_state_indices=ssm_state_indices,
257 + num_accepted_tokens=num_accepted_tokens,
258 + scale=scale,
259 + N=N,
260 + T=T,
261 + B=B,
262 + H=H,
263 + HV=HV,
264 + K=K,
265 + V=V,
266 + BK=BK,
267 + BV=BV,
268 + stride_init_state_token=stride_init_state_token,
269 + stride_final_state_token=stride_final_state_token,
270 + stride_indices_seq=stride_indices_seq,
271 + stride_indices_tok=stride_indices_tok,
272 + INPLACE_FINAL_STATE=inplace_final_state,
273 + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
274 + IS_KDA=is_kda,
275 + num_warps=num_warps,
276 + num_stages=num_stages,
277 + )
278 + o = o.squeeze(0)
279 + return o, final_state
vllm/model_executor/models/qwen3_next.py
+32 -31
@@ -34,7 +34,7 @@ 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,
37 + fused_sigmoid_gating_delta_rule_update,
38 )
39 from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd
40 from vllm.model_executor.layers.fused_moe import SharedFusedMoE
@@ -731,41 +731,40 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
731 mixed_qkv_non_spec
732 )
733
734 - g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias)
735 -
736 - if spec_sequence_masks is not None:
737 - if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
738 - g_spec = g
739 - beta_spec = beta
740 - g_non_spec = None
741 - beta_non_spec = None
742 - else:
743 - g_spec = g.index_select(1, spec_token_indx)
744 - beta_spec = beta.index_select(1, spec_token_indx)
734 + if attn_metadata.num_prefills > 0:
735 + g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias)
736 + if spec_sequence_masks is not None:
737 g_non_spec = g.index_select(1, non_spec_token_indx)
738 beta_non_spec = beta.index_select(1, non_spec_token_indx)
739 + else:
740 + g_non_spec = g
741 + beta_non_spec = beta
742 else:
748 - g_spec = None
749 - beta_spec = None
750 - g_non_spec = g
751 - beta_non_spec = beta
743 + g_non_spec = None
744 + beta_non_spec = None
745
746 # 2. Recurrent attention
747
748 # 2.1: Process the multi-query part
749 if spec_sequence_masks is not None:
757 - core_attn_out_spec, last_recurrent_state = fused_recurrent_gated_delta_rule(
758 - q=query_spec,
759 - k=key_spec,
760 - v=value_spec,
761 - g=g_spec,
762 - beta=beta_spec,
763 - initial_state=ssm_state,
764 - inplace_final_state=True,
765 - cu_seqlens=spec_query_start_loc[: attn_metadata.num_spec_decodes + 1],
766 - ssm_state_indices=spec_state_indices_tensor,
767 - num_accepted_tokens=num_accepted_tokens,
768 - use_qk_l2norm_in_kernel=True,
750 + core_attn_out_spec, last_recurrent_state = (
751 + fused_sigmoid_gating_delta_rule_update(
752 + A_log=self.A_log,
753 + a=a,
754 + b=b,
755 + dt_bias=self.dt_bias,
756 + q=query_spec,
757 + k=key_spec,
758 + v=value_spec,
759 + initial_state=ssm_state,
760 + inplace_final_state=True,
761 + cu_seqlens=spec_query_start_loc[
762 + : attn_metadata.num_spec_decodes + 1
763 + ],
764 + ssm_state_indices=spec_state_indices_tensor,
765 + num_accepted_tokens=num_accepted_tokens,
766 + use_qk_l2norm_in_kernel=True,
767 + )
768 )
769 else:
770 core_attn_out_spec, last_recurrent_state = None, None
@@ -794,12 +793,14 @@ class Qwen3NextGatedDeltaNet(nn.Module, MambaBase):
793 )
794 elif attn_metadata.num_decodes > 0:
795 core_attn_out_non_spec, last_recurrent_state = (
797 - fused_recurrent_gated_delta_rule(
796 + fused_sigmoid_gating_delta_rule_update(
797 + A_log=self.A_log,
798 + a=a,
799 + b=b,
800 + dt_bias=self.dt_bias,
801 q=query_non_spec,
802 k=key_non_spec,
803 v=value_non_spec,
801 - g=g_non_spec,
802 - beta=beta_non_spec,
804 initial_state=ssm_state,
805 inplace_final_state=True,
806 cu_seqlens=non_spec_query_start_loc[