@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4eefbf9609
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, 16:00 UTC

[Perf] fuse kernels in gdn (#37813)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

Jiangyun Zhu committed Apr 2, 2026 at 19:52 UTC 4eefbf9609e5ddb996e3ac37e192e92466ec35cc
4 files changed +495 -10
tests/kernels/test_fused_gdn_post_conv.py new
+209
@@ -0,0 +1,209 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for fused_gdn_prefill_post_conv kernel.
4 +
5 +Verifies that the fused kernel matches the reference:
6 + split → rearrange → contiguous → l2norm → gating
7 +"""
8 +
9 +import pytest
10 +import torch
11 +import torch.nn.functional as F
12 +
13 +from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (
14 + fused_post_conv_prep,
15 +)
16 +
17 +
18 +def reference_post_conv(
19 + conv_output: torch.Tensor,
20 + a: torch.Tensor,
21 + b: torch.Tensor,
22 + A_log: torch.Tensor,
23 + dt_bias: torch.Tensor,
24 + H: int,
25 + K: int,
26 + V: int,
27 + apply_l2norm: bool = True,
28 + output_g_exp: bool = False,
29 +):
30 + """Reference implementation using individual ops."""
31 + L = conv_output.shape[0]
32 + HV = A_log.shape[0]
33 +
34 + # Split
35 + q_flat, k_flat, v_flat = torch.split(conv_output, [H * K, H * K, HV * V], dim=-1)
36 +
37 + # Rearrange + contiguous
38 + q = q_flat.view(L, H, K).contiguous()
39 + k = k_flat.view(L, H, K).contiguous()
40 + v = v_flat.view(L, HV, V).contiguous()
41 +
42 + # L2 norm
43 + if apply_l2norm:
44 + q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6).to(conv_output.dtype)
45 + k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6).to(conv_output.dtype)
46 +
47 + # Gating
48 + x = a.float() + dt_bias.float()
49 + sp = F.softplus(x, beta=1.0, threshold=20.0)
50 + g = -torch.exp(A_log.float()) * sp
51 +
52 + if output_g_exp:
53 + g = torch.exp(g)
54 +
55 + beta_out = torch.sigmoid(b.float())
56 +
57 + return q, k, v, g, beta_out
58 +
59 +
60 +# Qwen3.5-35B config: H=16, HV=32, K=128, V=128
61 +# Qwen3.5-397B config: H=16, HV=64, K=128, V=128
62 +@pytest.mark.parametrize(
63 + "H, HV, K, V",
64 + [
65 + (16, 32, 128, 128), # 35B
66 + (16, 64, 128, 128), # 397B
67 + (4, 8, 64, 64), # small
68 + ],
69 +)
70 +@pytest.mark.parametrize("L", [1, 16, 128, 512, 2048])
71 +@pytest.mark.parametrize("apply_l2norm", [True, False])
72 +@pytest.mark.parametrize("output_g_exp", [True, False])
73 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
74 +def test_fused_post_conv_correctness(H, HV, K, V, L, apply_l2norm, output_g_exp, dtype):
75 + """Test fused kernel matches reference for all configs."""
76 + torch.manual_seed(42)
77 + device = "cuda"
78 + qkv_dim = 2 * H * K + HV * V
79 +
80 + conv_output = torch.randn(L, qkv_dim, dtype=dtype, device=device)
81 + a = torch.randn(L, HV, dtype=dtype, device=device)
82 + b = torch.randn(L, HV, dtype=dtype, device=device)
83 + A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0
84 + dt_bias = torch.randn(HV, dtype=torch.float32, device=device) * 0.1
85 +
86 + # Reference
87 + ref_q, ref_k, ref_v, ref_g, ref_beta = reference_post_conv(
88 + conv_output,
89 + a,
90 + b,
91 + A_log,
92 + dt_bias,
93 + H,
94 + K,
95 + V,
96 + apply_l2norm,
97 + output_g_exp,
98 + )
99 +
100 + # Fused kernel
101 + fused_q, fused_k, fused_v, fused_g, fused_beta = fused_post_conv_prep(
102 + conv_output,
103 + a,
104 + b,
105 + A_log,
106 + dt_bias,
107 + num_k_heads=H,
108 + head_k_dim=K,
109 + head_v_dim=V,
110 + apply_l2norm=apply_l2norm,
111 + output_g_exp=output_g_exp,
112 + )
113 +
114 + # Check shapes
115 + assert fused_q.shape == (L, H, K), f"q shape: {fused_q.shape}"
116 + assert fused_k.shape == (L, H, K), f"k shape: {fused_k.shape}"
117 + assert fused_v.shape == (L, HV, V), f"v shape: {fused_v.shape}"
118 + assert fused_g.shape == (L, HV), f"g shape: {fused_g.shape}"
119 + assert fused_beta.shape == (L, HV), f"beta shape: {fused_beta.shape}"
120 +
121 + # Check dtypes
122 + assert fused_q.dtype == dtype
123 + assert fused_k.dtype == dtype
124 + assert fused_v.dtype == dtype
125 + assert fused_g.dtype == torch.float32
126 + assert fused_beta.dtype == torch.float32
127 +
128 + # Check contiguity
129 + assert fused_q.is_contiguous()
130 + assert fused_k.is_contiguous()
131 + assert fused_v.is_contiguous()
132 +
133 + # Check values
134 + atol_qkv = 1e-2 if apply_l2norm else 1e-3
135 + rtol_qkv = 1e-2 if apply_l2norm else 1e-3
136 +
137 + torch.testing.assert_close(fused_q, ref_q, atol=atol_qkv, rtol=rtol_qkv)
138 + torch.testing.assert_close(fused_k, ref_k, atol=atol_qkv, rtol=rtol_qkv)
139 + torch.testing.assert_close(fused_v, ref_v, atol=1e-3, rtol=1e-3)
140 + torch.testing.assert_close(fused_g, ref_g, atol=1e-4, rtol=1e-4)
141 + torch.testing.assert_close(fused_beta, ref_beta, atol=1e-4, rtol=1e-4)
142 +
143 +
144 +@pytest.mark.parametrize("L", [1, 64, 256])
145 +def test_fused_post_conv_sanity(L):
146 + """Sanity checks: no NaN, unit-norm q/k, beta in (0,1)."""
147 + torch.manual_seed(0)
148 + device = "cuda"
149 + H, HV, K, V = 16, 32, 128, 128
150 + qkv_dim = 2 * H * K + HV * V
151 +
152 + conv_output = torch.randn(L, qkv_dim, dtype=torch.bfloat16, device=device)
153 + a = torch.randn(L, HV, dtype=torch.bfloat16, device=device)
154 + b = torch.randn(L, HV, dtype=torch.bfloat16, device=device)
155 + A_log = torch.randn(HV, dtype=torch.float32, device=device) - 2.0
156 + dt_bias = torch.randn(HV, dtype=torch.float32, device=device)
157 +
158 + q, k, v, g, beta = fused_post_conv_prep(
159 + conv_output,
160 + a,
161 + b,
162 + A_log,
163 + dt_bias,
164 + num_k_heads=H,
165 + head_k_dim=K,
166 + head_v_dim=V,
167 + )
168 +
169 + # Basic sanity
170 + assert not torch.isnan(q).any(), "NaN in q"
171 + assert not torch.isnan(k).any(), "NaN in k"
172 + assert not torch.isnan(v).any(), "NaN in v"
173 + assert not torch.isnan(g).any(), "NaN in g"
174 + assert not torch.isnan(beta).any(), "NaN in beta"
175 +
176 + # L2 norm check: each head vector should have unit norm
177 + q_norms = torch.norm(q.float(), dim=-1)
178 + k_norms = torch.norm(k.float(), dim=-1)
179 + torch.testing.assert_close(q_norms, torch.ones_like(q_norms), atol=1e-3, rtol=1e-3)
180 + torch.testing.assert_close(k_norms, torch.ones_like(k_norms), atol=1e-3, rtol=1e-3)
181 +
182 + # Beta should be in (0, 1)
183 + assert (beta >= 0).all() and (beta <= 1).all(), "beta out of range"
184 +
185 +
186 +def test_fused_post_conv_l0():
187 + """Test L=0 edge case."""
188 + device = "cuda"
189 + H, HV, K, V = 16, 32, 128, 128
190 + qkv_dim = 2 * H * K + HV * V
191 +
192 + conv_output = torch.empty(0, qkv_dim, dtype=torch.bfloat16, device=device)
193 + a = torch.empty(0, HV, dtype=torch.bfloat16, device=device)
194 + b = torch.empty(0, HV, dtype=torch.bfloat16, device=device)
195 + A_log = torch.randn(HV, dtype=torch.float32, device=device)
196 + dt_bias = torch.randn(HV, dtype=torch.float32, device=device)
197 +
198 + q, k, v, g, beta = fused_post_conv_prep(
199 + conv_output,
200 + a,
201 + b,
202 + A_log,
203 + dt_bias,
204 + num_k_heads=H,
205 + head_k_dim=K,
206 + head_v_dim=V,
207 + )
208 + assert q.shape == (0, H, K)
209 + assert g.shape == (0, HV)
vllm/model_executor/layers/fla/ops/__init__.py
+2
@@ -7,6 +7,7 @@
7 # the following copyright notice:
8 # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
9 from .chunk import chunk_gated_delta_rule
10 +from .fused_gdn_prefill_post_conv import fused_post_conv_prep
11 from .fused_recurrent import (
12 fused_recurrent_gated_delta_rule,
13 fused_recurrent_gated_delta_rule_packed_decode,
@@ -19,5 +20,6 @@ __all__ = [
20 "chunk_gated_delta_rule",
21 "fused_recurrent_gated_delta_rule",
22 "fused_recurrent_gated_delta_rule_packed_decode",
23 + "fused_post_conv_prep",
24 "fused_sigmoid_gating_delta_rule_update",
25 ]
vllm/model_executor/layers/fla/ops/fused_gdn_prefill_post_conv.py new
+248
@@ -0,0 +1,248 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Fused post-conv1d preparation for GDN prefill.
4 +
5 +Replaces the chain:
6 + split → rearrange → contiguous * 3 → l2norm * 2 → gating
7 +with a **single Triton kernel** that reads the conv'd mixed_qkv output
8 +and writes directly to q/k/v/g/beta in the target contiguous layout.
9 +
10 +"""
11 +
12 +from __future__ import annotations
13 +
14 +import torch
15 +
16 +from vllm.triton_utils import tl, triton
17 +
18 +
19 +@triton.jit
20 +def _fused_post_conv_kernel(
21 + # ---- inputs ----
22 + mixed_qkv_ptr, # [L, qkv_dim] conv'd output (contiguous)
23 + a_ptr, # [L, HV]
24 + b_ptr, # [L, HV]
25 + # ---- params ----
26 + A_log_ptr, # [HV]
27 + dt_bias_ptr, # [HV]
28 + # ---- outputs ----
29 + q_ptr, # [L, H, K] contiguous
30 + k_ptr, # [L, H, K] contiguous
31 + v_ptr, # [L, HV, V] contiguous
32 + g_ptr, # [L, HV] float32
33 + beta_ptr, # [L, HV] float32
34 + # ---- strides ----
35 + stride_x_tok, # qkv_dim
36 + stride_a_tok, # HV
37 + stride_b_tok, # HV
38 + stride_q_tok, # H * K
39 + stride_k_tok, # H * K
40 + stride_v_tok, # HV * V
41 + # ---- dims ----
42 + L,
43 + H: tl.constexpr,
44 + HV: tl.constexpr,
45 + K: tl.constexpr,
46 + V: tl.constexpr,
47 + APPLY_L2NORM: tl.constexpr,
48 + L2NORM_EPS: tl.constexpr,
49 + OUTPUT_G_EXP: tl.constexpr,
50 + SOFTPLUS_THRESHOLD: tl.constexpr,
51 + BLOCK_T: tl.constexpr,
52 + BK: tl.constexpr,
53 + BV: tl.constexpr,
54 +):
55 + """Single fused kernel for post-conv1d preparation.
56 +
57 + Grid: (ceil(L, BLOCK_T), H + HV)
58 + - program_id(1) in [0, H): Q/K head processing + l2norm
59 + - program_id(1) in [H, H+HV): V head processing + gating
60 + """
61 + i_tb = tl.program_id(0)
62 + i_head = tl.program_id(1)
63 +
64 + HK: tl.constexpr = H * K
65 +
66 + offs_t = i_tb * BLOCK_T + tl.arange(0, BLOCK_T) # [BLOCK_T]
67 + mask_t = offs_t < L
68 +
69 + if i_head < H:
70 + # ============ Q/K head processing ============
71 + i_h = i_head
72 + offs_k = tl.arange(0, BK) # [BK]
73 + mask_k = offs_k < K
74 + mask_2d = mask_t[:, None] & mask_k[None, :] # [BLOCK_T, BK]
75 +
76 + # Load Q features: mixed_qkv[t, i_h*K + k]
77 + q_offsets = offs_t[:, None] * stride_x_tok + i_h * K + offs_k[None, :]
78 + q_f32 = tl.load(mixed_qkv_ptr + q_offsets, mask=mask_2d, other=0).to(tl.float32)
79 +
80 + # Load K features: mixed_qkv[t, HK + i_h*K + k]
81 + k_offsets = offs_t[:, None] * stride_x_tok + HK + i_h * K + offs_k[None, :]
82 + k_f32 = tl.load(mixed_qkv_ptr + k_offsets, mask=mask_2d, other=0).to(tl.float32)
83 +
84 + if APPLY_L2NORM:
85 + q_sq_sum = tl.sum(q_f32 * q_f32, axis=1) # [BLOCK_T]
86 + q_inv = 1.0 / tl.sqrt(q_sq_sum + L2NORM_EPS)
87 + q_f32 = q_f32 * q_inv[:, None]
88 +
89 + k_sq_sum = tl.sum(k_f32 * k_f32, axis=1)
90 + k_inv = 1.0 / tl.sqrt(k_sq_sum + L2NORM_EPS)
91 + k_f32 = k_f32 * k_inv[:, None]
92 +
93 + # Store Q
94 + q_out = offs_t[:, None] * stride_q_tok + i_h * K + offs_k[None, :]
95 + tl.store(
96 + q_ptr + q_out,
97 + q_f32.to(q_ptr.dtype.element_ty),
98 + mask=mask_2d,
99 + )
100 +
101 + # Store K
102 + k_out = offs_t[:, None] * stride_k_tok + i_h * K + offs_k[None, :]
103 + tl.store(
104 + k_ptr + k_out,
105 + k_f32.to(k_ptr.dtype.element_ty),
106 + mask=mask_2d,
107 + )
108 + else:
109 + # ============ V head + gating processing ============
110 + i_hv = i_head - H
111 + offs_v = tl.arange(0, BV) # [BV]
112 + mask_v = offs_v < V
113 + mask_2d = mask_t[:, None] & mask_v[None, :] # [BLOCK_T, BV]
114 +
115 + V_OFFSET: tl.constexpr = 2 * H * K
116 +
117 + # Load V features: mixed_qkv[t, 2*H*K + i_hv*V + v]
118 + v_offsets = (
119 + offs_t[:, None] * stride_x_tok + V_OFFSET + i_hv * V + offs_v[None, :]
120 + )
121 + v_vals = tl.load(mixed_qkv_ptr + v_offsets, mask=mask_2d, other=0)
122 +
123 + # Store V
124 + v_out = offs_t[:, None] * stride_v_tok + i_hv * V + offs_v[None, :]
125 + tl.store(v_ptr + v_out, v_vals, mask=mask_2d)
126 +
127 + # Gating: one scalar per (token, v-head)
128 + A_log_val = tl.load(A_log_ptr + i_hv).to(tl.float32)
129 + dt_bias_val = tl.load(dt_bias_ptr + i_hv).to(tl.float32)
130 +
131 + a_offsets = offs_t * stride_a_tok + i_hv
132 + b_offsets = offs_t * stride_b_tok + i_hv
133 + a_vals = tl.load(a_ptr + a_offsets, mask=mask_t, other=0).to(tl.float32)
134 + b_vals = tl.load(b_ptr + b_offsets, mask=mask_t, other=0).to(tl.float32)
135 +
136 + # g = -exp(A_log) * softplus(a + dt_bias)
137 + x = a_vals + dt_bias_val
138 + sp = tl.where(x > 0, x + tl.log(1.0 + tl.exp(-x)), tl.log(1.0 + tl.exp(x)))
139 + sp = tl.where(x <= SOFTPLUS_THRESHOLD, sp, x)
140 + g_vals = -tl.exp(A_log_val) * sp
141 +
142 + if OUTPUT_G_EXP:
143 + g_vals = tl.exp(g_vals)
144 +
145 + beta_vals = tl.sigmoid(b_vals)
146 +
147 + gb_offsets = offs_t * HV + i_hv
148 + tl.store(g_ptr + gb_offsets, g_vals, mask=mask_t)
149 + tl.store(beta_ptr + gb_offsets, beta_vals, mask=mask_t)
150 +
151 +
152 +def fused_post_conv_prep(
153 + conv_output: torch.Tensor, # [L, qkv_dim] conv'd mixed_qkv
154 + a: torch.Tensor, # [L, HV]
155 + b: torch.Tensor, # [L, HV]
156 + A_log: torch.Tensor, # [HV]
157 + dt_bias: torch.Tensor, # [HV]
158 + num_k_heads: int,
159 + head_k_dim: int,
160 + head_v_dim: int,
161 + apply_l2norm: bool = True,
162 + output_g_exp: bool = False,
163 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
164 + """Fused post-conv1d prep: split + l2norm + gating in one kernel.
165 +
166 + Args:
167 + conv_output: [L, qkv_dim] contiguous conv'd mixed_qkv
168 + a: [L, HV] gating input
169 + b: [L, HV] gating input
170 + A_log: [HV] log decay parameter
171 + dt_bias: [HV] dt bias parameter
172 + num_k_heads: number of K heads (H)
173 + head_k_dim: dimension per K head (K)
174 + head_v_dim: dimension per V head (V)
175 + apply_l2norm: whether to L2-normalize q and k
176 + output_g_exp: if True, output exp(g) instead of g (for FlashInfer)
177 +
178 + Returns:
179 + q: [L, H, K] contiguous, optionally l2-normalized
180 + k: [L, H, K] contiguous, optionally l2-normalized
181 + v: [L, HV, V] contiguous
182 + g: [L, HV] float32
183 + beta: [L, HV] float32
184 + """
185 + L = conv_output.shape[0]
186 + qkv_dim = conv_output.shape[1]
187 + H = num_k_heads
188 + K = head_k_dim
189 + V = head_v_dim
190 + HV = A_log.shape[0]
191 + dtype = conv_output.dtype
192 + device = conv_output.device
193 +
194 + assert qkv_dim == 2 * H * K + HV * V, (
195 + f"qkv_dim={qkv_dim} != 2*H*K + HV*V = {2 * H * K + HV * V}"
196 + )
197 +
198 + # Allocate outputs in target contiguous layout
199 + q = torch.empty(L, H, K, dtype=dtype, device=device)
200 + k = torch.empty(L, H, K, dtype=dtype, device=device)
201 + v = torch.empty(L, HV, V, dtype=dtype, device=device)
202 + g = torch.empty(L, HV, dtype=torch.float32, device=device)
203 + beta = torch.empty(L, HV, dtype=torch.float32, device=device)
204 +
205 + if L == 0:
206 + return q, k, v, g, beta
207 +
208 + # ---- Kernel config ----
209 + BK = triton.next_power_of_2(K)
210 + BV = triton.next_power_of_2(V)
211 + BLOCK_T = 16 # tokens per block
212 +
213 + # Single kernel: blocks [0,H) do Q/K, blocks [H, H+HV) do V+gating
214 + grid = (triton.cdiv(L, BLOCK_T), H + HV)
215 + _fused_post_conv_kernel[grid](
216 + mixed_qkv_ptr=conv_output,
217 + a_ptr=a,
218 + b_ptr=b,
219 + A_log_ptr=A_log,
220 + dt_bias_ptr=dt_bias,
221 + q_ptr=q,
222 + k_ptr=k,
223 + v_ptr=v,
224 + g_ptr=g,
225 + beta_ptr=beta,
226 + stride_x_tok=conv_output.stride(0),
227 + stride_a_tok=a.stride(0),
228 + stride_b_tok=b.stride(0),
229 + stride_q_tok=q.stride(0),
230 + stride_k_tok=k.stride(0),
231 + stride_v_tok=v.stride(0),
232 + L=L,
233 + H=H,
234 + HV=HV,
235 + K=K,
236 + V=V,
237 + APPLY_L2NORM=apply_l2norm,
238 + L2NORM_EPS=1e-6,
239 + OUTPUT_G_EXP=output_g_exp,
240 + SOFTPLUS_THRESHOLD=20.0,
241 + BLOCK_T=BLOCK_T,
242 + BK=BK,
243 + BV=BV,
244 + num_warps=4,
245 + num_stages=2,
246 + )
247 +
248 + return q, k, v, g, beta
vllm/model_executor/layers/mamba/gdn_linear_attn.py
+36 -10
@@ -24,6 +24,7 @@ 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_post_conv_prep,
28 fused_recurrent_gated_delta_rule_packed_decode,
29 fused_sigmoid_gating_delta_rule_update,
30 )
@@ -774,19 +775,44 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
775 mixed_qkv_non_spec = None
776
777 query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
777 - query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
778 - mixed_qkv_non_spec
779 - )
780 -
778 if attn_metadata.num_prefills > 0:
782 - g, beta = fused_gdn_gating(self.A_log, a, b, self.dt_bias)
779 + assert mixed_qkv_non_spec is not None, (
780 + "mixed_qkv_non_spec must be provided for prefill path"
781 + )
782 if spec_sequence_masks is not None:
784 - g_non_spec = g.index_select(1, non_spec_token_indx)
785 - beta_non_spec = beta.index_select(1, non_spec_token_indx)
783 + a_non_spec = a.index_select(0, non_spec_token_indx)
784 + b_non_spec = b.index_select(0, non_spec_token_indx)
785 else:
787 - g_non_spec = g
788 - beta_non_spec = beta
786 + a_non_spec = a
787 + b_non_spec = b
788 +
789 + (
790 + query_non_spec,
791 + key_non_spec,
792 + value_non_spec,
793 + g_non_spec,
794 + beta_non_spec,
795 + ) = fused_post_conv_prep(
796 + conv_output=mixed_qkv_non_spec,
797 + a=a_non_spec,
798 + b=b_non_spec,
799 + A_log=self.A_log,
800 + dt_bias=self.dt_bias,
801 + num_k_heads=self.num_k_heads // self.tp_size,
802 + head_k_dim=self.head_k_dim,
803 + head_v_dim=self.head_v_dim,
804 + apply_l2norm=True,
805 + output_g_exp=False,
806 + )
807 + query_non_spec = query_non_spec.unsqueeze(0)
808 + key_non_spec = key_non_spec.unsqueeze(0)
809 + value_non_spec = value_non_spec.unsqueeze(0)
810 + g_non_spec = g_non_spec.unsqueeze(0)
811 + beta_non_spec = beta_non_spec.unsqueeze(0)
812 else:
813 + query_non_spec, key_non_spec, value_non_spec = self.rearrange_mixed_qkv(
814 + mixed_qkv_non_spec
815 + )
816 g_non_spec = None
817 beta_non_spec = None
818
@@ -832,7 +858,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase):
858 initial_state=initial_state,
859 output_final_state=True,
860 cu_seqlens=non_spec_query_start_loc,
835 - use_qk_l2norm_in_kernel=True,
861 + use_qk_l2norm_in_kernel=False,
862 )
863 # Init cache
864 ssm_state[non_spec_state_indices_tensor] = last_recurrent_state.to(