@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e88bdd60d9

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 14, 05:30 UTC

[FLA] Introduce Kimi Delta Attention(KDA) to VLLM (#27654)

Signed-off-by: lizhiyuan <lizhiyuan@moonshot.cn>

Zhiyuan Li committed Oct 28, 2025 at 22:56 UTC e88bdd60d9a25d985168c9f4a60ab10095236d7c

5 files changed +1451 -58

vllm/model_executor/layers/fla/ops/chunk.py

+1 -1

@@ -36,7 +36,7 @@ def chunk_gated_delta_rule_fwd(

36 g = chunk_local_cumsum(g, chunk_size=64, cu_seqlens=cu_seqlens)

37 # obtain WY representation. u is actually the new v.

38 A = chunk_scaled_dot_kkt_fwd(

39 - k=k, beta=beta, g_cumsum=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32

39 + k=k, beta=beta, g=g, cu_seqlens=cu_seqlens, output_dtype=torch.float32

40 )

41 A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)

42 w, u = recompute_w_u_fwd(

vllm/model_executor/layers/fla/ops/chunk_delta_h.py

+70 -40

@@ -14,14 +14,15 @@ from vllm.triton_utils import tl, triton

15 from .index import prepare_chunk_indices, prepare_chunk_offsets

16 from .op import exp

17 -from .utils import is_nvidia_hopper, use_cuda_graph

17 +from .utils import use_cuda_graph

19 -NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16]

19 +NUM_WARPS = [2, 4, 8, 16]

22 @triton.heuristics(

23 {

24 "USE_G": lambda args: args["g"] is not None,

25 + "USE_GK": lambda args: args["gk"] is not None,

26 "USE_INITIAL_STATE": lambda args: args["h0"] is not None,

27 "STORE_FINAL_STATE": lambda args: args["ht"] is not None,

28 "SAVE_NEW_VALUE": lambda args: args["v_new"] is not None,

@@ -35,7 +36,7 @@ NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16]

36 for num_stages in [2, 3, 4]

37 for BV in [32, 64]

38 ],

38 - key=["H", "K", "V", "BT", "USE_G"],

39 + key=["H", "K", "V", "BT"],

40 use_cuda_graph=use_cuda_graph,

41 )

42 @triton.jit(do_not_specialize=["T"])

@@ -45,6 +46,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

46 w,

47 v_new,

48 g,

49 + gk,

50 h,

51 h0,

52 ht,

@@ -58,6 +60,7 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

60 BT: tl.constexpr,

61 BV: tl.constexpr,

62 USE_G: tl.constexpr,

63 + USE_GK: tl.constexpr,

64 USE_INITIAL_STATE: tl.constexpr,

65 STORE_FINAL_STATE: tl.constexpr,

66 SAVE_NEW_VALUE: tl.constexpr,

@@ -88,12 +91,12 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

91 b_h4 = tl.zeros([64, BV], dtype=tl.float32)

93 # calculate offset

91 - h += (boh * H + i_h) * K * V

92 - v += (bos * H + i_h) * V

93 - k += (bos * Hg + i_h // (H // Hg)) * K

94 - w += (bos * H + i_h) * K

94 + h += ((boh * H + i_h) * K * V).to(tl.int64)

95 + v += ((bos * H + i_h) * V).to(tl.int64)

96 + k += ((bos * Hg + i_h // (H // Hg)) * K).to(tl.int64)

97 + w += ((bos * H + i_h) * K).to(tl.int64)

98 if SAVE_NEW_VALUE:

96 - v_new += (bos * H + i_h) * V

99 + v_new += ((bos * H + i_h) * V).to(tl.int64)

100 stride_v = H * V

101 stride_h = H * K * V

102 stride_k = Hg * K

@@ -145,92 +148,115 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(

148 )

149 tl.store(p_h4, b_h4.to(p_h4.dtype.element_ty), boundary_check=(0, 1))

150

148 - p_v = tl.make_block_ptr(

149 - v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)

150 - )

151 - p_v_new = (

152 - tl.make_block_ptr(

153 - v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)

154 - )

155 - if SAVE_NEW_VALUE

156 - else None

157 - )

158 - b_v_new = tl.zeros([BT, BV], dtype=tl.float32)

151 p_w = tl.make_block_ptr(

152 w, (T, K), (stride_w, 1), (i_t * BT, 0), (BT, 64), (1, 0)

153 )

154 b_w = tl.load(p_w, boundary_check=(0, 1))

163 - b_v_new += tl.dot(b_w, b_h1.to(b_w.dtype))

155 + b_v = tl.dot(b_w, b_h1.to(b_w.dtype))

156 if K > 64:

157 p_w = tl.make_block_ptr(

158 w, (T, K), (stride_w, 1), (i_t * BT, 64), (BT, 64), (1, 0)

159 )

160 b_w = tl.load(p_w, boundary_check=(0, 1))

169 - b_v_new += tl.dot(b_w, b_h2.to(b_w.dtype))

161 + b_v += tl.dot(b_w, b_h2.to(b_w.dtype))

162 if K > 128:

163 p_w = tl.make_block_ptr(

164 w, (T, K), (stride_w, 1), (i_t * BT, 128), (BT, 64), (1, 0)

165 )

166 b_w = tl.load(p_w, boundary_check=(0, 1))

175 - b_v_new += tl.dot(b_w, b_h3.to(b_w.dtype))

167 + b_v += tl.dot(b_w, b_h3.to(b_w.dtype))

168 if K > 192:

169 p_w = tl.make_block_ptr(

170 w, (T, K), (stride_w, 1), (i_t * BT, 192), (BT, 64), (1, 0)

171 )

172 b_w = tl.load(p_w, boundary_check=(0, 1))

181 - b_v_new += tl.dot(b_w, b_h4.to(b_w.dtype))

182 - b_v_new = -b_v_new + tl.load(p_v, boundary_check=(0, 1))

173 + b_v += tl.dot(b_w, b_h4.to(b_w.dtype))

174 + p_v = tl.make_block_ptr(

175 + v, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)

176 + )

177 + b_v = tl.load(p_v, boundary_check=(0, 1)) - b_v

178

179 if SAVE_NEW_VALUE:

185 - p_v_new = tl.make_block_ptr(

180 + p_v = tl.make_block_ptr(

181 v_new, (T, V), (stride_v, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0)

182 )

188 - tl.store(

189 - p_v_new, b_v_new.to(p_v_new.dtype.element_ty), boundary_check=(0, 1)

190 - )

183 + tl.store(p_v, b_v.to(p_v.dtype.element_ty), boundary_check=(0, 1))

184

185 + last_idx = min((i_t + 1) * BT, T) - 1

186 if USE_G:

187 m_t = (i_t * BT + tl.arange(0, BT)) < T

194 - last_idx = min((i_t + 1) * BT, T) - 1

188 b_g_last = tl.load(g + bos * H + last_idx * H + i_h)

189 p_g = tl.make_block_ptr(

190 g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)

191 )

192 b_g = tl.load(p_g, boundary_check=(0,))

200 - b_v_new = b_v_new * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None]

193 + b_v = b_v * tl.where(m_t, exp(b_g_last - b_g), 0)[:, None]

194 b_g_last = exp(b_g_last)

202 - b_h1 = b_h1 * b_g_last

195 + b_h1 *= b_g_last

196 + if K > 64:

197 + b_h2 *= b_g_last

198 + if K > 128:

199 + b_h3 *= b_g_last

200 + if K > 192:

201 + b_h4 *= b_g_last

202 +

203 + if USE_GK:

204 + o_k1 = tl.arange(0, 64)

205 + b_gk_last1 = tl.load(

206 + gk + (bos + last_idx) * H * K + i_h * K + o_k1,

207 + mask=(o_k1 < K),

208 + other=0.0,

209 + )

210 + b_h1 *= exp(b_gk_last1)[:, None]

211 if K > 64:

204 - b_h2 = b_h2 * b_g_last

212 + o_k2 = 64 + o_k1

213 + b_gk_last2 = tl.load(

214 + gk + (bos + last_idx) * H * K + i_h * K + o_k2,

215 + mask=(o_k2 < K),

216 + other=0.0,

217 + )

218 + b_h2 *= exp(b_gk_last2)[:, None]

219 if K > 128:

206 - b_h3 = b_h3 * b_g_last

220 + o_k3 = 128 + o_k1

221 + b_gk_last3 = tl.load(

222 + gk + (bos + last_idx) * H * K + i_h * K + o_k3,

223 + mask=(o_k3 < K),

224 + other=0.0,

225 + )

226 + b_h3 *= exp(b_gk_last3)[:, None]

227 if K > 192:

208 - b_h4 = b_h4 * b_g_last

209 - b_v_new = b_v_new.to(k.dtype.element_ty)

228 + o_k4 = 192 + o_k1

229 + b_gk_last4 = tl.load(

230 + gk + (bos + last_idx) * H * K + i_h * K + o_k4,

231 + mask=(o_k4 < K),

232 + other=0.0,

233 + )

234 + b_h4 *= exp(b_gk_last4)[:, None]

235 + b_v = b_v.to(k.dtype.element_ty)

236 +

237 p_k = tl.make_block_ptr(

238 k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)

239 )

240 b_k = tl.load(p_k, boundary_check=(0, 1))

214 - b_h1 += tl.dot(b_k, b_v_new)

241 + b_h1 += tl.dot(b_k, b_v)

242 if K > 64:

243 p_k = tl.make_block_ptr(

244 k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)

245 )

246 b_k = tl.load(p_k, boundary_check=(0, 1))

220 - b_h2 += tl.dot(b_k, b_v_new)

247 + b_h2 += tl.dot(b_k, b_v)

248 if K > 128:

249 p_k = tl.make_block_ptr(

250 k, (K, T), (1, stride_k), (128, i_t * BT), (64, BT), (0, 1)

251 )

252 b_k = tl.load(p_k, boundary_check=(0, 1))

226 - b_h3 += tl.dot(b_k, b_v_new)

253 + b_h3 += tl.dot(b_k, b_v)

254 if K > 192:

255 p_k = tl.make_block_ptr(

256 k, (K, T), (1, stride_k), (192, i_t * BT), (64, BT), (0, 1)

257 )

258 b_k = tl.load(p_k, boundary_check=(0, 1))

232 - b_h4 += tl.dot(b_k, b_v_new)

233 -

259 + b_h4 += tl.dot(b_k, b_v)

260 # epilogue

261 if STORE_FINAL_STATE:

262 p_ht = tl.make_block_ptr(ht, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0))

@@ -257,12 +283,15 @@ def chunk_gated_delta_rule_fwd_h(

283 w: torch.Tensor,

284 u: torch.Tensor,

285 g: torch.Tensor | None = None,

286 + gk: torch.Tensor | None = None,

287 initial_state: torch.Tensor | None = None,

288 output_final_state: bool = False,

289 chunk_size: int = 64, # SY: remove this argument and force chunk size 64?

290 save_new_value: bool = True,

291 cu_seqlens: torch.LongTensor | None = None,

292 ) -> tuple[torch.Tensor, torch.Tensor]:

293 + # This kernel is slightly different from fla to support Q/K with different head numbers.

294 + # In fla, Q/K always have the same head number, so Hg is always equal to H.

295 B, T, Hg, K, V = *k.shape, u.shape[-1]

296 H = u.shape[-2]

297 BT = chunk_size

@@ -299,6 +328,7 @@ def chunk_gated_delta_rule_fwd_h(

328 w=w,

329 v_new=v_new,

330 g=g,

331 + gk=gk,

332 h=h,

333 h0=initial_state,

334 ht=final_state,

vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py

+11 -13

@@ -18,8 +18,8 @@ from .op import exp

19 @triton.heuristics(

20 {

21 + "USE_G": lambda args: args["g"] is not None,

22 "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,

22 - "USE_G": lambda args: args["g_cumsum"] is not None,

23 }

24 )

25 @triton.autotune(

@@ -35,7 +35,7 @@ from .op import exp

35 def chunk_scaled_dot_kkt_fwd_kernel(

36 k,

37 beta,

38 - g_cumsum,

38 + g,

39 A,

40 cu_seqlens,

41 chunk_indices,

@@ -85,9 +85,7 @@ def chunk_scaled_dot_kkt_fwd_kernel(

85 b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k))

87 if USE_G:

88 - p_g = tl.make_block_ptr(

89 - g_cumsum + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,)

90 - )

88 + p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))

89 b_g = tl.load(p_g, boundary_check=(0,))

90 b_g_diff = b_g[:, None] - b_g[None, :]

91 b_A = b_A * exp(b_g_diff)

@@ -102,8 +100,8 @@ def chunk_scaled_dot_kkt_fwd_kernel(

100

101 def chunk_scaled_dot_kkt_fwd(

102 k: torch.Tensor,

105 - beta: torch.Tensor,

106 - g_cumsum: torch.Tensor | None = None,

103 + g: torch.Tensor | None = None,

104 + beta: torch.Tensor | None = None,

105 cu_seqlens: torch.LongTensor | None = None,

106 chunk_size: int = 64,

107 output_dtype: torch.dtype = torch.float32,

@@ -116,9 +114,8 @@ def chunk_scaled_dot_kkt_fwd(

114 The key tensor of shape `[B, T, H, K]`.

115 beta (torch.Tensor):

116 The beta tensor of shape `[B, T, H]`.

119 - g_cumsum (torch.Tensor):

120 - The cumulative sum of the gate tensor of shape `[B, T, H]`.

121 - Default: None

117 + g (torch.Tensor):

118 + The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.

119 cu_seqlens (torch.LongTensor):

120 The cumulative sequence lengths of the input tensor.

121 Default: None

@@ -130,20 +127,21 @@ def chunk_scaled_dot_kkt_fwd(

127 Returns:

128 beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.

129 """

133 -

130 + # This kernel is slightly different from fla to support Q/K with different head numbers.

131 + # In fla, Q/K always have the same head number, so Hg is always equal to H.

132 B, T, Hg, K = k.shape

135 -

133 H = beta.shape[-1]

134 BT = chunk_size

135 chunk_indices = (

136 prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

137 )

138 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

139 +

140 A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)

141 chunk_scaled_dot_kkt_fwd_kernel[(NT, B * H)](

142 k=k,

143 + g=g,

144 beta=beta,

146 - g_cumsum=g_cumsum,

145 A=A,

146 cu_seqlens=cu_seqlens,

147 chunk_indices=chunk_indices,

vllm/model_executor/layers/fla/ops/fused_recurrent.py

+18 -4

@@ -57,6 +57,7 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(

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

@@ -86,7 +87,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(

87 p_beta = beta + (bos * HV + i_hv) * V + o_v

88 else:

89 p_beta = beta + bos * HV + i_hv

89 - p_g = g + bos * HV + i_hv

90 +

91 + if not IS_KDA:

92 + p_g = g + bos * HV + i_hv

93 + else:

94 + p_gk = g + (bos * HV + i_hv) * K + o_k

95 +

96 p_o = o + ((i_k * all + bos) * HV + i_hv) * V + o_v

98 mask_k = o_k < K

@@ -116,14 +122,18 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(

122 b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)

123 b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)

124 b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)

119 - b_g = tl.load(p_g).to(tl.float32)

125

126 if USE_QK_L2NORM_IN_KERNEL:

127 b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)

128 b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)

129 b_q = b_q * scale

130 # [BK, BV]

126 - b_h *= exp(b_g)

131 + if not IS_KDA:

132 + b_g = tl.load(p_g).to(tl.float32)

133 + b_h *= exp(b_g)

134 + else:

135 + b_gk = tl.load(p_gk).to(tl.float32)

136 + b_h *= exp(b_gk[:, None])

137 # [BV]

138 b_v -= tl.sum(b_h * b_k[:, None], 0)

139 if IS_BETA_HEADWISE:

@@ -155,7 +165,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(

165 p_k += H * K

166 p_o += HV * V

167 p_v += HV * V

158 - p_g += HV

168 + if not IS_KDA:

169 + p_g += HV

170 + else:

171 + p_gk += HV * K

172 p_beta += HV * (V if IS_BETA_HEADWISE else 1)

173

174

@@ -228,6 +241,7 @@ def fused_recurrent_gated_delta_rule_fwd(

241 IS_BETA_HEADWISE=beta.ndim == v.ndim,

242 USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,

243 INPLACE_FINAL_STATE=inplace_final_state,

244 + IS_KDA=False,

245 num_warps=num_warps,

246 num_stages=num_stages,

247 )

vllm/model_executor/layers/fla/ops/kda.py new

+1351

@@ -0,0 +1,1351 @@

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:

9 +# ruff: noqa: E501

10 +

11 +

12 +import torch

13 +import torch.nn as nn

14 +

15 +from vllm.triton_utils import tl, triton

16 +from vllm.utils.math_utils import cdiv, next_power_of_2

17 +

18 +from .chunk_delta_h import chunk_gated_delta_rule_fwd_h

19 +from .cumsum import chunk_local_cumsum

20 +from .fused_recurrent import fused_recurrent_gated_delta_rule_fwd_kernel

21 +from .index import prepare_chunk_indices

22 +from .l2norm import l2norm_fwd

23 +from .op import exp, log

24 +from .solve_tril import solve_tril

25 +from .utils import is_amd

26 +

27 +BT_LIST_AUTOTUNE = [32, 64, 128]

28 +NUM_WARPS_AUTOTUNE = [2, 4, 8, 16] if is_amd else [4, 8, 16, 32]

29 +

30 +

31 +def fused_recurrent_kda_fwd(

32 + q: torch.Tensor,

33 + k: torch.Tensor,

34 + v: torch.Tensor,

35 + g: torch.Tensor,

36 + beta: torch.Tensor,

37 + scale: float,

38 + initial_state: torch.Tensor,

39 + inplace_final_state: bool = True,

40 + cu_seqlens: torch.LongTensor | None = None,

41 + ssm_state_indices: torch.Tensor | None = None,

42 + num_accepted_tokens: torch.Tensor | None = None,

43 + use_qk_l2norm_in_kernel: bool = False,

44 +) -> tuple[torch.Tensor, torch.Tensor]:

45 + B, T, H, K, V = *k.shape, v.shape[-1]

46 + HV = v.shape[2]

47 + N = B if cu_seqlens is None else len(cu_seqlens) - 1

48 + BK, BV = next_power_of_2(K), min(next_power_of_2(V), 8)

49 + NK, NV = cdiv(K, BK), cdiv(V, BV)

50 + assert NK == 1, "NK > 1 is not supported yet"

51 + num_stages = 3

52 + num_warps = 1

53 +

54 + o = torch.empty_like(k)

55 + if inplace_final_state:

56 + final_state = initial_state

57 + else:

58 + final_state = q.new_empty(T, HV, K, V, dtype=initial_state.dtype)

59 +

60 + stride_init_state_token = initial_state.stride(0)

61 + stride_final_state_token = final_state.stride(0)

62 +

63 + if ssm_state_indices is None:

64 + stride_indices_seq, stride_indices_tok = 1, 1

65 + elif ssm_state_indices.ndim == 1:

66 + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride(0), 1

67 + else:

68 + stride_indices_seq, stride_indices_tok = ssm_state_indices.stride()

69 +

70 + grid = (NK, NV, N * HV)

71 + fused_recurrent_gated_delta_rule_fwd_kernel[grid](

72 + q=q,

73 + k=k,

74 + v=v,

75 + g=g,

76 + beta=beta,

77 + o=o,

78 + h0=initial_state,

79 + ht=final_state,

80 + cu_seqlens=cu_seqlens,

81 + ssm_state_indices=ssm_state_indices,

82 + num_accepted_tokens=num_accepted_tokens,

83 + scale=scale,

84 + N=N,

85 + T=T,

86 + B=B,

87 + H=H,

88 + HV=HV,

89 + K=K,

90 + V=V,

91 + BK=BK,

92 + BV=BV,

93 + stride_init_state_token=stride_init_state_token,

94 + stride_final_state_token=stride_final_state_token,

95 + stride_indices_seq=stride_indices_seq,

96 + stride_indices_tok=stride_indices_tok,

97 + IS_BETA_HEADWISE=beta.ndim == v.ndim,

98 + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,

99 + INPLACE_FINAL_STATE=inplace_final_state,

100 + IS_KDA=True,

101 + num_warps=num_warps,

102 + num_stages=num_stages,

103 + )

104 +

105 + return o, final_state

106 +

107 +

108 +def fused_recurrent_kda(

109 + q: torch.Tensor,

110 + k: torch.Tensor,

111 + v: torch.Tensor,

112 + g: torch.Tensor,

113 + beta: torch.Tensor = None,

114 + scale: float = None,

115 + initial_state: torch.Tensor = None,

116 + inplace_final_state: bool = True,

117 + use_qk_l2norm_in_kernel: bool = True,

118 + cu_seqlens: torch.LongTensor | None = None,

119 + ssm_state_indices: torch.LongTensor | None = None,

120 + **kwargs,

121 +) -> tuple[torch.Tensor, torch.Tensor]:

122 + if cu_seqlens is not None and q.shape[0] != 1:

123 + raise ValueError(

124 + f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."

125 + f"Please flatten variable-length inputs before processing."

126 + )

127 + if scale is None:

128 + scale = k.shape[-1] ** -0.5

129 +

130 + o, final_state = fused_recurrent_kda_fwd(

131 + q=q.contiguous(),

132 + k=k.contiguous(),

133 + v=v.contiguous(),

134 + g=g.contiguous(),

135 + beta=beta.contiguous(),

136 + scale=scale,

137 + initial_state=initial_state,

138 + inplace_final_state=inplace_final_state,

139 + cu_seqlens=cu_seqlens,

140 + ssm_state_indices=ssm_state_indices,

141 + num_accepted_tokens=None,

142 + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,

143 + )

144 + return o, final_state

145 +

146 +

147 +@triton.heuristics(

148 + {

149 + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,

150 + "HAS_RESIDUAL": lambda args: args["residual"] is not None,

151 + "HAS_WEIGHT": lambda args: args["w"] is not None,

152 + "HAS_BIAS": lambda args: args["b"] is not None,

153 + }

154 +)

155 +@triton.jit

156 +def layer_norm_gated_fwd_kernel(

157 + x, # pointer to the input

158 + g, # pointer to the gate

159 + y, # pointer to the output

160 + w, # pointer to the weights

161 + b, # pointer to the biases

162 + residual, # pointer to the residual

163 + residual_out, # pointer to the residual

164 + mean, # pointer to the mean

165 + rstd, # pointer to the 1/std

166 + eps, # epsilon to avoid division by zero

167 + T, # number of rows in x

168 + D: tl.constexpr, # number of columns in x

169 + BT: tl.constexpr,

170 + BD: tl.constexpr,

171 + ACTIVATION: tl.constexpr,

172 + IS_RMS_NORM: tl.constexpr,

173 + STORE_RESIDUAL_OUT: tl.constexpr,

174 + HAS_RESIDUAL: tl.constexpr,

175 + HAS_WEIGHT: tl.constexpr,

176 + HAS_BIAS: tl.constexpr,

177 +):

178 + i_t = tl.program_id(0)

179 +

180 + o_d = tl.arange(0, BD)

181 + m_d = o_d < D

182 +

183 + p_x = tl.make_block_ptr(x, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))

184 + b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)

185 + if HAS_RESIDUAL:

186 + p_res = tl.make_block_ptr(

187 + residual, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)

188 + )

189 + b_x += tl.load(p_res, boundary_check=(0, 1)).to(tl.float32)

190 + if STORE_RESIDUAL_OUT:

191 + p_res_out = tl.make_block_ptr(

192 + residual_out, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0)

193 + )

194 + tl.store(p_res_out, b_x.to(p_res_out.dtype.element_ty), boundary_check=(0, 1))

195 + if not IS_RMS_NORM:

196 + b_mean = tl.sum(b_x, axis=1) / D

197 + p_mean = tl.make_block_ptr(mean, (T,), (1,), (i_t * BT,), (BT,), (0,))

198 + tl.store(p_mean, b_mean.to(p_mean.dtype.element_ty), boundary_check=(0,))

199 + b_xbar = tl.where(m_d[None, :], b_x - b_mean[:, None], 0.0)

200 + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D

201 + else:

202 + b_xbar = tl.where(m_d[None, :], b_x, 0.0)

203 + b_var = tl.sum(b_xbar * b_xbar, axis=1) / D

204 + b_rstd = 1 / tl.sqrt(b_var + eps)

205 +

206 + p_rstd = tl.make_block_ptr(rstd, (T,), (1,), (i_t * BT,), (BT,), (0,))

207 + tl.store(p_rstd, b_rstd.to(p_rstd.dtype.element_ty), boundary_check=(0,))

208 +

209 + if HAS_WEIGHT:

210 + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)

211 + if HAS_BIAS:

212 + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)

213 + b_x_hat = (

214 + (b_x - b_mean[:, None]) * b_rstd[:, None]

215 + if not IS_RMS_NORM

216 + else b_x * b_rstd[:, None]

217 + )

218 + b_y = b_x_hat * b_w[None, :] if HAS_WEIGHT else b_x_hat

219 + if HAS_BIAS:

220 + b_y = b_y + b_b[None, :]

221 +

222 + # swish/sigmoid output gate

223 + p_g = tl.make_block_ptr(g, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))

224 + b_g = tl.load(p_g, boundary_check=(0, 1)).to(tl.float32)

225 + if ACTIVATION == "swish" or ACTIVATION == "silu":

226 + b_y = b_y * b_g * tl.sigmoid(b_g)

227 + elif ACTIVATION == "sigmoid":

228 + b_y = b_y * tl.sigmoid(b_g)

229 +

230 + # Write output

231 + p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))

232 + tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))

233 +

234 +

235 +@triton.heuristics(

236 + {

237 + "STORE_RESIDUAL_OUT": lambda args: args["residual_out"] is not None,

238 + "HAS_RESIDUAL": lambda args: args["residual"] is not None,

239 + "HAS_WEIGHT": lambda args: args["w"] is not None,

240 + "HAS_BIAS": lambda args: args["b"] is not None,

241 + }

242 +)

243 +@triton.jit

244 +def layer_norm_gated_fwd_kernel1(

245 + x, # pointer to the input

246 + g, # pointer to the gate

247 + y, # pointer to the output

248 + w, # pointer to the weights

249 + b, # pointer to the biases

250 + residual, # pointer to the residual

251 + residual_out, # pointer to the residual

252 + mean, # pointer to the mean

253 + rstd, # pointer to the 1/std

254 + eps, # epsilon to avoid division by zero

255 + D: tl.constexpr, # number of columns in x

256 + BD: tl.constexpr,

257 + ACTIVATION: tl.constexpr,

258 + IS_RMS_NORM: tl.constexpr,

259 + STORE_RESIDUAL_OUT: tl.constexpr,

260 + HAS_RESIDUAL: tl.constexpr,

261 + HAS_WEIGHT: tl.constexpr,

262 + HAS_BIAS: tl.constexpr,

263 +):

264 + i_t = tl.program_id(0)

265 + x += i_t * D

266 + y += i_t * D

267 + g += i_t * D

268 + if HAS_RESIDUAL:

269 + residual += i_t * D

270 + if STORE_RESIDUAL_OUT:

271 + residual_out += i_t * D

272 +

273 + o_d = tl.arange(0, BD)

274 + m_d = o_d < D

275 + b_x = tl.load(x + o_d, mask=m_d, other=0.0).to(tl.float32)

276 + if HAS_RESIDUAL:

277 + b_x += tl.load(residual + o_d, mask=m_d, other=0.0).to(tl.float32)

278 + if STORE_RESIDUAL_OUT:

279 + tl.store(residual_out + o_d, b_x, mask=m_d)

280 + if not IS_RMS_NORM:

281 + b_mean = tl.sum(b_x, axis=0) / D

282 + tl.store(mean + i_t, b_mean)

283 + b_xbar = tl.where(m_d, b_x - b_mean, 0.0)

284 + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D

285 + else:

286 + b_xbar = tl.where(m_d, b_x, 0.0)

287 + b_var = tl.sum(b_xbar * b_xbar, axis=0) / D

288 + b_rstd = 1 / tl.sqrt(b_var + eps)

289 + tl.store(rstd + i_t, b_rstd)

290 +

291 + if HAS_WEIGHT:

292 + b_w = tl.load(w + o_d, mask=m_d).to(tl.float32)

293 + if HAS_BIAS:

294 + b_b = tl.load(b + o_d, mask=m_d).to(tl.float32)

295 + b_x_hat = (b_x - b_mean) * b_rstd if not IS_RMS_NORM else b_x * b_rstd

296 + b_y = b_x_hat * b_w if HAS_WEIGHT else b_x_hat

297 + if HAS_BIAS:

298 + b_y = b_y + b_b

299 +

300 + # swish/sigmoid output gate

301 + b_g = tl.load(g + o_d, mask=m_d, other=0.0).to(tl.float32)

302 + if ACTIVATION == "swish" or ACTIVATION == "silu":

303 + b_y = b_y * b_g * tl.sigmoid(b_g)

304 + elif ACTIVATION == "sigmoid":

305 + b_y = b_y * tl.sigmoid(b_g)

306 +

307 + # Write output

308 + tl.store(y + o_d, b_y, mask=m_d)

309 +

310 +

311 +def layer_norm_gated_fwd(

312 + x: torch.Tensor,

313 + g: torch.Tensor,

314 + weight: torch.Tensor,

315 + bias: torch.Tensor,

316 + activation: str = "swish",

317 + eps: float = 1e-5,

318 + residual: torch.Tensor = None,

319 + out_dtype: torch.dtype = None,

320 + residual_dtype: torch.dtype = None,

321 + is_rms_norm: bool = False,

322 +):

323 + if residual is not None:

324 + residual_dtype = residual.dtype

325 + T, D = x.shape

326 + if residual is not None:

327 + assert residual.shape == (T, D)

328 + if weight is not None:

329 + assert weight.shape == (D,)

330 + if bias is not None:

331 + assert bias.shape == (D,)

332 + # allocate output

333 + y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype)

334 + if residual is not None or (

335 + residual_dtype is not None and residual_dtype != x.dtype

336 + ):

337 + residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)

338 + else:

339 + residual_out = None

340 + mean = (

341 + torch.empty((T,), dtype=torch.float, device=x.device)

342 + if not is_rms_norm

343 + else None

344 + )

345 + rstd = torch.empty((T,), dtype=torch.float, device=x.device)

346 + # Less than 64KB per feature: enqueue fused kernel

347 + MAX_FUSED_SIZE = 65536 // x.element_size()

348 + BD = min(MAX_FUSED_SIZE, next_power_of_2(D))

349 + if D > BD:

350 + raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")

351 + # heuristics for number of warps

352 +

353 + if D <= 512:

354 + BT = 32

355 + layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](

356 + x=x,

357 + g=g,

358 + y=y,

359 + w=weight,

360 + b=bias,

361 + residual=residual,

362 + residual_out=residual_out,

363 + mean=mean,

364 + rstd=rstd,

365 + eps=eps,

366 + T=T,

367 + D=D,

368 + BD=BD,

369 + BT=BT,

370 + ACTIVATION=activation,

371 + IS_RMS_NORM=is_rms_norm,

372 + num_warps=4,

373 + )

374 + else:

375 + layer_norm_gated_fwd_kernel1[(T,)](

376 + x=x,

377 + g=g,

378 + y=y,

379 + w=weight,

380 + b=bias,

381 + residual=residual,

382 + residual_out=residual_out,

383 + mean=mean,

384 + rstd=rstd,

385 + eps=eps,

386 + D=D,

387 + BD=BD,

388 + ACTIVATION=activation,

389 + IS_RMS_NORM=is_rms_norm,

390 + num_warps=4,

391 + )

392 + # residual_out is None if residual is None and residual_dtype == input_dtype

393 + return y, mean, rstd, residual_out if residual_out is not None else x

394 +

395 +

396 +def rms_norm_gated(

397 + x: torch.Tensor,

398 + g: torch.Tensor,

399 + weight: torch.Tensor,

400 + bias: torch.Tensor,

401 + activation: str = "swish",

402 + residual: torch.Tensor | None = None,

403 + prenorm: bool = False,

404 + residual_in_fp32: bool = False,

405 + eps: float = 1e-6,

406 +):

407 + x_shape_og = x.shape

408 + # reshape input data into 2D tensor

409 + x = x.contiguous().reshape(-1, x.shape[-1])

410 + g = g.contiguous().reshape(-1, g.shape[-1])

411 + if residual is not None:

412 + assert residual.shape == x_shape_og

413 + residual = residual.contiguous().reshape(-1, residual.shape[-1])

414 + residual_dtype = (

415 + residual.dtype

416 + if residual is not None

417 + else (torch.float if residual_in_fp32 else None)

418 + )

419 + y, _, _, residual_out = layer_norm_gated_fwd(

420 + x=x,

421 + g=g,

422 + weight=weight,

423 + bias=bias,

424 + activation=activation,

425 + eps=eps,

426 + residual=residual,

427 + residual_dtype=residual_dtype,

428 + is_rms_norm=True,

429 + )

430 + y = y.reshape(x_shape_og)

431 + return y if not prenorm else (y, residual_out.reshape(x_shape_og))

432 +

433 +

434 +class FusedRMSNormGated(nn.Module):

435 + def __init__(

436 + self,

437 + hidden_size: int,

438 + elementwise_affine: bool = True,

439 + eps: float = 1e-5,

440 + activation: str = "swish",

441 + device: torch.device | None = None,

442 + dtype: torch.dtype | None = None,

443 + ) -> None:

444 + factory_kwargs = {"device": device, "dtype": dtype}

445 + super().__init__()

446 +

447 + self.hidden_size = hidden_size

448 + self.elementwise_affine = elementwise_affine

449 + self.eps = eps

450 + self.activation = activation

451 +

452 + if self.activation not in ["swish", "silu", "sigmoid"]:

453 + raise ValueError(f"Unsupported activation: {self.activation}")

454 +

455 + if elementwise_affine:

456 + self.weight = nn.Parameter(torch.empty(hidden_size, **factory_kwargs))

457 + else:

458 + self.register_parameter("weight", None)

459 + self.register_parameter("bias", None)

460 +

461 + def forward(

462 + self,

463 + x: torch.Tensor,

464 + g: torch.Tensor,

465 + residual: torch.Tensor | None = None,

466 + prenorm: bool = False,

467 + residual_in_fp32: bool = False,

468 + ) -> torch.Tensor:

469 + return rms_norm_gated(

470 + x,

471 + g,

472 + self.weight,

473 + self.bias,

474 + self.activation,

475 + residual=residual,

476 + eps=self.eps,

477 + prenorm=prenorm,

478 + residual_in_fp32=residual_in_fp32,

479 + )

480 +

481 +

482 +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})

483 +@triton.autotune(

484 + configs=[

485 + triton.Config({"BK": BK}, num_warps=num_warps, num_stages=num_stages)

486 + for BK in [32, 64]

487 + for num_warps in [1, 2, 4, 8]

488 + for num_stages in [2, 3, 4]

489 + ],

490 + key=["BC"],

491 +)

492 +@triton.jit(do_not_specialize=["T"])

493 +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter(

494 + q,

495 + k,

496 + g,

497 + beta,

498 + A,

499 + Aqk,

500 + scale,

501 + cu_seqlens,

502 + chunk_indices,

503 + T,

504 + H: tl.constexpr,

505 + K: tl.constexpr,

506 + BT: tl.constexpr,

507 + BC: tl.constexpr,

508 + BK: tl.constexpr,

509 + NC: tl.constexpr,

510 + IS_VARLEN: tl.constexpr,

511 +):

512 + i_t, i_c, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)

513 + i_b, i_h = i_bh // H, i_bh % H

514 + i_i, i_j = i_c // NC, i_c % NC

515 + if IS_VARLEN:

516 + i_n, i_t = (

517 + tl.load(chunk_indices + i_t * 2).to(tl.int32),

518 + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),

519 + )

520 + bos, eos = (

521 + tl.load(cu_seqlens + i_n).to(tl.int32),

522 + tl.load(cu_seqlens + i_n + 1).to(tl.int32),

523 + )

524 + T = eos - bos

525 + else:

526 + bos, eos = i_b * T, i_b * T + T

527 +

528 + if i_t * BT + i_i * BC >= T:

529 + return

530 + if i_i <= i_j:

531 + return

532 +

533 + q += (bos * H + i_h) * K

534 + k += (bos * H + i_h) * K

535 + g += (bos * H + i_h) * K

536 + A += (bos * H + i_h) * BT

537 + Aqk += (bos * H + i_h) * BT

538 +

539 + p_b = tl.make_block_ptr(

540 + beta + bos * H + i_h, (T,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,)

541 + )

542 + b_b = tl.load(p_b, boundary_check=(0,))

543 +

544 + b_A = tl.zeros([BC, BC], dtype=tl.float32)

545 + b_Aqk = tl.zeros([BC, BC], dtype=tl.float32)

546 + for i_k in range(tl.cdiv(K, BK)):

547 + p_q = tl.make_block_ptr(

548 + q, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)

549 + )

550 + p_k = tl.make_block_ptr(

551 + k, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)

552 + )

553 + p_g = tl.make_block_ptr(

554 + g, (T, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK), (1, 0)

555 + )

556 + b_kt = tl.make_block_ptr(

557 + k, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)

558 + )

559 + p_gk = tl.make_block_ptr(

560 + g, (K, T), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC), (0, 1)

561 + )

562 +

563 + o_k = i_k * BK + tl.arange(0, BK)

564 + m_k = o_k < K

565 + # [BK,]

566 + b_gn = tl.load(g + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0)

567 + # [BC, BK]

568 + b_g = tl.load(p_g, boundary_check=(0, 1))

569 + b_k = tl.load(p_k, boundary_check=(0, 1)) * exp(b_g - b_gn[None, :])

570 + # [BK, BC]

571 + b_gk = tl.load(p_gk, boundary_check=(0, 1))

572 + b_kt = tl.load(b_kt, boundary_check=(0, 1))

573 + # [BC, BC]

574 + b_ktg = b_kt * exp(b_gn[:, None] - b_gk)

575 + b_A += tl.dot(b_k, b_ktg)

576 +

577 + b_q = tl.load(p_q, boundary_check=(0, 1))

578 + b_qg = b_q * exp(b_g - b_gn[None, :]) * scale

579 + b_Aqk += tl.dot(b_qg, b_ktg)

580 +

581 + b_A *= b_b[:, None]

582 +

583 + p_A = tl.make_block_ptr(

584 + A, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)

585 + )

586 + tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))

587 + p_Aqk = tl.make_block_ptr(

588 + Aqk, (T, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0)

589 + )

590 + tl.store(p_Aqk, b_Aqk.to(Aqk.dtype.element_ty), boundary_check=(0, 1))

591 +

592 +

593 +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})

594 +@triton.autotune(

595 + configs=[triton.Config({}, num_warps=num_warps) for num_warps in [1, 2, 4, 8]],

596 + key=["BK", "BT"],

597 +)

598 +@triton.jit(do_not_specialize=["T"])

599 +def chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra(

600 + q,

601 + k,

602 + g,

603 + beta,

604 + A,

605 + Aqk,

606 + scale,

607 + cu_seqlens,

608 + chunk_indices,

609 + T,

610 + H: tl.constexpr,

611 + K: tl.constexpr,

612 + BT: tl.constexpr,

613 + BC: tl.constexpr,

614 + BK: tl.constexpr,

615 + IS_VARLEN: tl.constexpr,

616 +):

617 + i_t, i_i, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)

618 + i_b, i_h = i_bh // H, i_bh % H

619 + if IS_VARLEN:

620 + i_n, i_t = (

621 + tl.load(chunk_indices + i_t * 2).to(tl.int32),

622 + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),

623 + )

624 + bos, eos = (

625 + tl.load(cu_seqlens + i_n).to(tl.int32),

626 + tl.load(cu_seqlens + i_n + 1).to(tl.int32),

627 + )

628 + T = eos - bos

629 + else:

630 + bos, eos = i_b * T, i_b * T + T

631 +

632 + if i_t * BT + i_i * BC >= T:

633 + return

634 +

635 + o_i = tl.arange(0, BC)

636 + o_k = tl.arange(0, BK)

637 + m_k = o_k < K

638 + m_A = (i_t * BT + i_i * BC + o_i) < T

639 + o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC

640 +

641 + p_q = tl.make_block_ptr(

642 + q + (bos * H + i_h) * K,

643 + (T, K),

644 + (H * K, 1),

645 + (i_t * BT + i_i * BC, 0),

646 + (BC, BK),

647 + (1, 0),

648 + )

649 + p_k = tl.make_block_ptr(

650 + k + (bos * H + i_h) * K,

651 + (T, K),

652 + (H * K, 1),

653 + (i_t * BT + i_i * BC, 0),

654 + (BC, BK),

655 + (1, 0),

656 + )

657 + p_g = tl.make_block_ptr(

658 + g + (bos * H + i_h) * K,

659 + (T, K),

660 + (H * K, 1),

661 + (i_t * BT + i_i * BC, 0),

662 + (BC, BK),

663 + (1, 0),

664 + )

665 + b_q = tl.load(p_q, boundary_check=(0, 1))

666 + b_k = tl.load(p_k, boundary_check=(0, 1))

667 + b_g = tl.load(p_g, boundary_check=(0, 1))

668 +

669 + p_b = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h

670 + b_k = b_k * tl.load(p_b, mask=m_A, other=0)[:, None]

671 +

672 + p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k

673 + p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k

674 +

675 + for j in range(0, min(BC, T - i_t * BT - i_i * BC)):

676 + b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32)

677 + b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32)

678 + b_ktg = b_kt[None, :] * exp(b_g - b_gk[None, :])

679 + b_A = tl.sum(b_k * b_ktg, 1)

680 + b_A = tl.where(o_i > j, b_A, 0.0)

681 + b_Aqk = tl.sum(b_q * b_ktg, 1)

682 + b_Aqk = tl.where(o_i >= j, b_Aqk * scale, 0.0)

683 + tl.store(A + o_A + j, b_A, mask=m_A)

684 + tl.store(Aqk + o_A + j, b_Aqk, mask=m_A)

685 + p_kt += H * K

686 + p_gk += H * K

687 +

688 +

689 +def chunk_kda_scaled_dot_kkt_fwd(

690 + q: torch.Tensor,

691 + k: torch.Tensor,

692 + gk: torch.Tensor | None = None,

693 + beta: torch.Tensor | None = None,

694 + scale: float | None = None,

695 + cu_seqlens: torch.LongTensor | None = None,

696 + chunk_size: int = 64,

697 + output_dtype: torch.dtype = torch.float32,

698 +) -> tuple[torch.Tensor, torch.Tensor]:

699 + r"""

700 + Compute beta * K * K^T.

701 +

702 + Args:

703 + k (torch.Tensor):

704 + The key tensor of shape `[B, T, H, K]`.

705 + beta (torch.Tensor):

706 + The beta tensor of shape `[B, T, H]`.

707 + gk (torch.Tensor):

708 + The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.

709 + cu_seqlens (torch.LongTensor):

710 + The cumulative sequence lengths of the input tensor.

711 + Default: None

712 + chunk_size (int):

713 + The chunk size. Default: 64.

714 + output_dtype (torch.dtype):

715 + The dtype of the output tensor. Default: `torch.float32`

716 +

717 + Returns:

718 + beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.

719 + """

720 + B, T, H, K = k.shape

721 + assert K <= 256

722 + BT = chunk_size

723 + chunk_indices = (

724 + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

725 + )

726 + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

727 +

728 + BC = min(16, BT)

729 + NC = cdiv(BT, BC)

730 + BK = max(next_power_of_2(K), 16)

731 + A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)

732 + Aqk = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)

733 + grid = (NT, NC * NC, B * H)

734 + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid](

735 + q=q,

736 + k=k,

737 + g=gk,

738 + beta=beta,

739 + A=A,

740 + Aqk=Aqk,

741 + scale=scale,

742 + cu_seqlens=cu_seqlens,

743 + chunk_indices=chunk_indices,

744 + T=T,

745 + H=H,

746 + K=K,

747 + BT=BT,

748 + BC=BC,

749 + NC=NC,

750 + )

751 +

752 + grid = (NT, NC, B * H)

753 + chunk_kda_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid](

754 + q=q,

755 + k=k,

756 + g=gk,

757 + beta=beta,

758 + A=A,

759 + Aqk=Aqk,

760 + scale=scale,

761 + cu_seqlens=cu_seqlens,

762 + chunk_indices=chunk_indices,

763 + T=T,

764 + H=H,

765 + K=K,

766 + BT=BT,

767 + BC=BC,

768 + BK=BK,

769 + )

770 + return A, Aqk

771 +

772 +

773 +@triton.heuristics(

774 + {

775 + "STORE_QG": lambda args: args["qg"] is not None,

776 + "STORE_KG": lambda args: args["kg"] is not None,

777 + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None,

778 + }

779 +)

780 +@triton.autotune(

781 + configs=[

782 + triton.Config({}, num_warps=num_warps, num_stages=num_stages)

783 + for num_warps in [2, 4, 8]

784 + for num_stages in [2, 3, 4]

785 + ],

786 + key=["H", "K", "V", "BT", "BK", "BV", "IS_VARLEN"],

787 +)

788 +@triton.jit(do_not_specialize=["T"])

789 +def recompute_w_u_fwd_kernel(

790 + q,

791 + k,

792 + qg,

793 + kg,

794 + v,

795 + beta,

796 + w,

797 + u,

798 + A,

799 + gk,

800 + cu_seqlens,

801 + chunk_indices,

802 + T,

803 + H: tl.constexpr,

804 + K: tl.constexpr,

805 + V: tl.constexpr,

806 + BT: tl.constexpr,

807 + BK: tl.constexpr,

808 + BV: tl.constexpr,

809 + STORE_QG: tl.constexpr,

810 + STORE_KG: tl.constexpr,

811 + IS_VARLEN: tl.constexpr,

812 + DOT_PRECISION: tl.constexpr,

813 +):

814 + i_t, i_bh = tl.program_id(0), tl.program_id(1)

815 + i_b, i_h = i_bh // H, i_bh % H

816 + if IS_VARLEN:

817 + i_n, i_t = (

818 + tl.load(chunk_indices + i_t * 2).to(tl.int32),

819 + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),

820 + )

821 + bos, eos = (

822 + tl.load(cu_seqlens + i_n).to(tl.int32),

823 + tl.load(cu_seqlens + i_n + 1).to(tl.int32),

824 + )

825 + T = eos - bos

826 + else:

827 + bos, eos = i_b * T, i_b * T + T

828 + p_b = tl.make_block_ptr(beta + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))

829 + b_b = tl.load(p_b, boundary_check=(0,))

830 +

831 + p_A = tl.make_block_ptr(

832 + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)

833 + )

834 + b_A = tl.load(p_A, boundary_check=(0, 1))

835 +

836 + for i_v in range(tl.cdiv(V, BV)):

837 + p_v = tl.make_block_ptr(

838 + v + (bos * H + i_h) * V,

839 + (T, V),

840 + (H * V, 1),

841 + (i_t * BT, i_v * BV),

842 + (BT, BV),

843 + (1, 0),

844 + )

845 + p_u = tl.make_block_ptr(

846 + u + (bos * H + i_h) * V,

847 + (T, V),

848 + (H * V, 1),

849 + (i_t * BT, i_v * BV),

850 + (BT, BV),

851 + (1, 0),

852 + )

853 + b_v = tl.load(p_v, boundary_check=(0, 1))

854 + b_vb = (b_v * b_b[:, None]).to(b_v.dtype)

855 + b_u = tl.dot(b_A, b_vb, input_precision=DOT_PRECISION)

856 + tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))

857 +

858 + for i_k in range(tl.cdiv(K, BK)):

859 + p_w = tl.make_block_ptr(

860 + w + (bos * H + i_h) * K,

861 + (T, K),

862 + (H * K, 1),

863 + (i_t * BT, i_k * BK),

864 + (BT, BK),

865 + (1, 0),

866 + )

867 + p_k = tl.make_block_ptr(

868 + k + (bos * H + i_h) * K,

869 + (T, K),

870 + (H * K, 1),

871 + (i_t * BT, i_k * BK),

872 + (BT, BK),

873 + (1, 0),

874 + )

875 + b_k = tl.load(p_k, boundary_check=(0, 1))

876 + b_kb = b_k * b_b[:, None]

877 +

878 + p_gk = tl.make_block_ptr(

879 + gk + (bos * H + i_h) * K,

880 + (T, K),

881 + (H * K, 1),

882 + (i_t * BT, i_k * BK),

883 + (BT, BK),

884 + (1, 0),

885 + )

886 + b_gk = tl.load(p_gk, boundary_check=(0, 1))

887 + b_kb *= exp(b_gk)

888 + if STORE_QG:

889 + p_q = tl.make_block_ptr(

890 + q + (bos * H + i_h) * K,

891 + (T, K),

892 + (H * K, 1),

893 + (i_t * BT, i_k * BK),

894 + (BT, BK),

895 + (1, 0),

896 + )

897 + p_qg = tl.make_block_ptr(

898 + qg + (bos * H + i_h) * K,

899 + (T, K),

900 + (H * K, 1),

901 + (i_t * BT, i_k * BK),

902 + (BT, BK),

903 + (1, 0),

904 + )

905 + b_q = tl.load(p_q, boundary_check=(0, 1))

906 + b_qg = b_q * exp(b_gk)

907 + tl.store(p_qg, b_qg.to(p_qg.dtype.element_ty), boundary_check=(0, 1))

908 + if STORE_KG:

909 + last_idx = min(i_t * BT + BT, T) - 1

910 +

911 + o_k = i_k * BK + tl.arange(0, BK)

912 + m_k = o_k < K

913 + b_gn = tl.load(

914 + gk + ((bos + last_idx) * H + i_h) * K + o_k, mask=m_k, other=0.0

915 + )

916 + b_kg = b_k * exp(b_gn - b_gk)

917 +

918 + p_kg = tl.make_block_ptr(

919 + kg + (bos * H + i_h) * K,

920 + (T, K),

921 + (H * K, 1),

922 + (i_t * BT, i_k * BK),

923 + (BT, BK),

924 + (1, 0),

925 + )

926 + tl.store(p_kg, b_kg.to(p_kg.dtype.element_ty), boundary_check=(0, 1))

927 +

928 + b_w = tl.dot(b_A, b_kb.to(b_k.dtype))

929 + tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))

930 +

931 +

932 +def recompute_w_u_fwd(

933 + k: torch.Tensor,

934 + v: torch.Tensor,

935 + beta: torch.Tensor,

936 + A: torch.Tensor,

937 + q: torch.Tensor | None = None,

938 + gk: torch.Tensor | None = None,

939 + cu_seqlens: torch.LongTensor | None = None,

940 +) -> tuple[torch.Tensor, torch.Tensor]:

941 + B, T, H, K, V = *k.shape, v.shape[-1]

942 + BT = A.shape[-1]

943 + BK = 64

944 + BV = 64

945 +

946 + chunk_indices = (

947 + prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None

948 + )

949 + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

950 +

951 + w = torch.empty_like(k)

952 + u = torch.empty_like(v)

953 + kg = torch.empty_like(k) if gk is not None else None

954 + recompute_w_u_fwd_kernel[(NT, B * H)](

955 + q=q,

956 + k=k,

957 + qg=None,

958 + kg=kg,

959 + v=v,

960 + beta=beta,

961 + w=w,

962 + u=u,

963 + A=A,

964 + gk=gk,

965 + cu_seqlens=cu_seqlens,

966 + chunk_indices=chunk_indices,

967 + T=T,

968 + H=H,

969 + K=K,

970 + V=V,

971 + BT=BT,

972 + BK=BK,

973 + BV=BV,

974 + DOT_PRECISION="ieee",

975 + )

976 + return w, u, None, kg

977 +

978 +

979 +@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})

980 +@triton.autotune(

981 + configs=[

982 + triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)

983 + for BK in [32, 64]

984 + for BV in [64, 128]

985 + for num_warps in [2, 4, 8]

986 + for num_stages in [2, 3, 4]

987 + ],

988 + key=["BT"],

989 +)

990 +@triton.jit(do_not_specialize=["T"])

991 +def chunk_gla_fwd_kernel_o(

992 + q,

993 + v,

994 + g,

995 + h,

996 + o,

997 + A,

998 + cu_seqlens,

999 + chunk_indices,

1000 + scale,

1001 + T,

1002 + H: tl.constexpr,

1003 + K: tl.constexpr,

1004 + V: tl.constexpr,

1005 + BT: tl.constexpr,

1006 + BK: tl.constexpr,

1007 + BV: tl.constexpr,

1008 + IS_VARLEN: tl.constexpr,

1009 +):

1010 + i_v, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)

1011 + i_b, i_h = i_bh // H, i_bh % H

1012 + if IS_VARLEN:

1013 + i_tg = i_t

1014 + i_n, i_t = (

1015 + tl.load(chunk_indices + i_t * 2).to(tl.int32),

1016 + tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),

1017 + )

1018 + bos, eos = (

1019 + tl.load(cu_seqlens + i_n).to(tl.int32),

1020 + tl.load(cu_seqlens + i_n + 1).to(tl.int32),

1021 + )

1022 + T = eos - bos

1023 + NT = tl.cdiv(T, BT)

1024 + else:

1025 + NT = tl.cdiv(T, BT)

1026 + i_tg = i_b * NT + i_t

1027 + bos, eos = i_b * T, i_b * T + T

1028 +

1029 + m_s = tl.arange(0, BT)[:, None] >= tl.arange(0, BT)[None, :]

1030 +

1031 + b_o = tl.zeros([BT, BV], dtype=tl.float32)

1032 + for i_k in range(tl.cdiv(K, BK)):

1033 + p_q = tl.make_block_ptr(

1034 + q + (bos * H + i_h) * K,

1035 + (T, K),

1036 + (H * K, 1),

1037 + (i_t * BT, i_k * BK),

1038 + (BT, BK),

1039 + (1, 0),

1040 + )

1041 + p_g = tl.make_block_ptr(

1042 + g + (bos * H + i_h) * K,

1043 + (T, K),

1044 + (H * K, 1),

1045 + (i_t * BT, i_k * BK),

1046 + (BT, BK),

1047 + (1, 0),

1048 + )

1049 + p_h = tl.make_block_ptr(

1050 + h + (i_tg * H + i_h) * K * V,

1051 + (K, V),

1052 + (V, 1),

1053 + (i_k * BK, i_v * BV),

1054 + (BK, BV),

1055 + (1, 0),

1056 + )

1057 +

1058 + # [BT, BK]

1059 + b_q = tl.load(p_q, boundary_check=(0, 1))

1060 + b_q = (b_q * scale).to(b_q.dtype)

1061 + # [BT, BK]

1062 + b_g = tl.load(p_g, boundary_check=(0, 1))

1063 + # [BT, BK]

1064 + b_qg = (b_q * exp(b_g)).to(b_q.dtype)

1065 + # [BK, BV]

1066 + b_h = tl.load(p_h, boundary_check=(0, 1))

1067 + # works but dkw, owing to divine benevolence

1068 + # [BT, BV]

1069 + if i_k >= 0:

1070 + b_o += tl.dot(b_qg, b_h.to(b_qg.dtype))

1071 + p_v = tl.make_block_ptr(

1072 + v + (bos * H + i_h) * V,

1073 + (T, V),

1074 + (H * V, 1),

1075 + (i_t * BT, i_v * BV),

1076 + (BT, BV),

1077 + (1, 0),

1078 + )

1079 + p_o = tl.make_block_ptr(

1080 + o + (bos * H + i_h) * V,

1081 + (T, V),

1082 + (H * V, 1),

1083 + (i_t * BT, i_v * BV),

1084 + (BT, BV),

1085 + (1, 0),

1086 + )

1087 + p_A = tl.make_block_ptr(

1088 + A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0)

1089 + )

1090 + # [BT, BV]

1091 + b_v = tl.load(p_v, boundary_check=(0, 1))

1092 + # [BT, BT]

1093 + b_A = tl.load(p_A, boundary_check=(0, 1))

1094 + b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)

1095 + b_o += tl.dot(b_A, b_v, allow_tf32=False)

1096 + tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))

1097 +

1098 +

1099 +def chunk_gla_fwd_o_gk(

1100 + q: torch.Tensor,

1101 + v: torch.Tensor,

1102 + g: torch.Tensor,

1103 + A: torch.Tensor,

1104 + h: torch.Tensor,

1105 + o: torch.Tensor,

1106 + scale: float,

1107 + cu_seqlens: torch.LongTensor | None = None,

1108 + chunk_size: int = 64,

1109 +):

1110 + B, T, H, K, V = *q.shape, v.shape[-1]

1111 + BT = chunk_size

1112 +

1113 + chunk_indices = (

1114 + prepare_chunk_indices(cu_seqlens, chunk_size)

1115 + if cu_seqlens is not None

1116 + else None

1117 + )

1118 + NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)

1119 +

1120 + def grid(meta):

1121 + return (cdiv(V, meta["BV"]), NT, B * H)

1122 +

1123 + chunk_gla_fwd_kernel_o[grid](

1124 + q=q,

1125 + v=v,

1126 + g=g,

1127 + h=h,

1128 + o=o,

1129 + A=A,

1130 + cu_seqlens=cu_seqlens,

1131 + chunk_indices=chunk_indices,

1132 + scale=scale,

1133 + T=T,

1134 + H=H,

1135 + K=K,

1136 + V=V,

1137 + BT=BT,

1138 + )

1139 + return o

1140 +

1141 +

1142 +def chunk_kda_fwd(

1143 + q: torch.Tensor,

1144 + k: torch.Tensor,

1145 + v: torch.Tensor,

1146 + g: torch.Tensor,

1147 + beta: torch.Tensor,

1148 + scale: float,

1149 + initial_state: torch.Tensor,

1150 + output_final_state: bool,

1151 + cu_seqlens: torch.LongTensor | None = None,

1152 +):

1153 + chunk_size = 64

1154 + g = chunk_local_cumsum(g, chunk_size=chunk_size, cu_seqlens=cu_seqlens)

1155 + # the intra Aqk is kept in fp32

1156 + # the computation has very marginal effect on the entire throughput

1157 + A, Aqk = chunk_kda_scaled_dot_kkt_fwd(

1158 + q=q,

1159 + k=k,

1160 + gk=g,

1161 + beta=beta,

1162 + scale=scale,

1163 + cu_seqlens=cu_seqlens,

1164 + output_dtype=torch.float32,

1165 + )

1166 + A = solve_tril(A=A, cu_seqlens=cu_seqlens, output_dtype=k.dtype)

1167 + w, u, _, kg = recompute_w_u_fwd(

1168 + k=k,

1169 + v=v,

1170 + beta=beta,

1171 + A=A,

1172 + gk=g,

1173 + cu_seqlens=cu_seqlens,

1174 + )

1175 + del A

1176 + h, v_new, final_state = chunk_gated_delta_rule_fwd_h(

1177 + k=kg,

1178 + w=w,

1179 + u=u,

1180 + gk=g,

1181 + initial_state=initial_state,

1182 + output_final_state=output_final_state,

1183 + cu_seqlens=cu_seqlens,

1184 + )

1185 + del w, u, kg

1186 + o = chunk_gla_fwd_o_gk(

1187 + q=q,

1188 + v=v_new,

1189 + g=g,

1190 + A=Aqk,

1191 + h=h,

1192 + o=v,

1193 + scale=scale,

1194 + cu_seqlens=cu_seqlens,

1195 + chunk_size=chunk_size,

1196 + )

1197 + del Aqk, v_new, h

1198 + return o, final_state

1199 +

1200 +

1201 +def chunk_kda(

1202 + q: torch.Tensor,

1203 + k: torch.Tensor,

1204 + v: torch.Tensor,

1205 + g: torch.Tensor,

1206 + beta: torch.Tensor,

1207 + scale: float = None,

1208 + initial_state: torch.Tensor = None,

1209 + output_final_state: bool = False,

1210 + use_qk_l2norm_in_kernel: bool = False,

1211 + cu_seqlens: torch.LongTensor | None = None,

1212 + **kwargs,

1213 +):

1214 + if scale is None:

1215 + scale = k.shape[-1] ** -0.5

1216 +

1217 + if use_qk_l2norm_in_kernel:

1218 + q = l2norm_fwd(q.contiguous())

1219 + k = l2norm_fwd(k.contiguous())

1220 +

1221 + o, final_state = chunk_kda_fwd(

1222 + q=q,

1223 + k=k,

1224 + v=v.contiguous(),

1225 + g=g.contiguous(),

1226 + beta=beta.contiguous(),

1227 + scale=scale,

1228 + initial_state=initial_state.contiguous(),

1229 + output_final_state=output_final_state,

1230 + cu_seqlens=cu_seqlens,

1231 + )

1232 + return o, final_state

1233 +

1234 +

1235 +@triton.autotune(

1236 + configs=[

1237 + triton.Config({"BT": bt}, num_warps=nw, num_stages=ns)

1238 + for bt in BT_LIST_AUTOTUNE

1239 + for nw in NUM_WARPS_AUTOTUNE

1240 + for ns in [2, 3]

1241 + ],

1242 + key=["H", "D"],

1243 +)

1244 +@triton.jit

1245 +def kda_gate_fwd_kernel(

1246 + g,

1247 + A,

1248 + y,

1249 + g_bias,

1250 + beta: tl.constexpr,

1251 + threshold: tl.constexpr,

1252 + T,

1253 + H,

1254 + D: tl.constexpr,

1255 + BT: tl.constexpr,

1256 + BD: tl.constexpr,

1257 + HAS_BIAS: tl.constexpr,

1258 +):

1259 + i_t, i_h = tl.program_id(0), tl.program_id(1)

1260 + n_t = i_t * BT

1261 +

1262 + b_a = tl.load(A + i_h).to(tl.float32)

1263 + b_a = -tl.exp(b_a)

1264 +

1265 + stride_row = H * D

1266 + stride_col = 1

1267 +

1268 + g_ptr = tl.make_block_ptr(

1269 + base=g + i_h * D,

1270 + shape=(T, D),

1271 + strides=(stride_row, stride_col),

1272 + offsets=(n_t, 0),

1273 + block_shape=(BT, BD),

1274 + order=(1, 0),

1275 + )

1276 +

1277 + y_ptr = tl.make_block_ptr(

1278 + base=y + i_h * D,

1279 + shape=(T, D),

1280 + strides=(stride_row, stride_col),

1281 + offsets=(n_t, 0),

1282 + block_shape=(BT, BD),

1283 + order=(1, 0),

1284 + )

1285 +

1286 + b_g = tl.load(g_ptr, boundary_check=(0, 1)).to(tl.float32)

1287 +

1288 + if HAS_BIAS:

1289 + n_d = tl.arange(0, BD)

1290 + bias_mask = n_d < D

1291 + b_bias = tl.load(g_bias + i_h * D + n_d, mask=bias_mask, other=0.0).to(

1292 + tl.float32

1293 + )

1294 + b_g = b_g + b_bias[None, :]

1295 +

1296 + # softplus(x, beta) = (1/beta) * log(1 + exp(beta * x))

1297 + # When beta * x > threshold, use linear approximation x

1298 + # Use threshold to switch to linear when beta*x > threshold

1299 + g_scaled = b_g * beta

1300 + use_linear = g_scaled > threshold

1301 + sp = tl.where(use_linear, b_g, (1.0 / beta) * log(1.0 + tl.exp(g_scaled)))

1302 + b_y = b_a * sp

1303 +

1304 + tl.store(y_ptr, b_y.to(y.dtype.element_ty), boundary_check=(0, 1))

1305 +

1306 +

1307 +def kda_gate_fwd(

1308 + g: torch.Tensor,

1309 + A: torch.Tensor,

1310 + head_k_dim: int,

1311 + g_bias: torch.Tensor | None = None,

1312 + beta: float = 1.0,

1313 + threshold: float = 20.0,

1314 +) -> torch.Tensor:

1315 + """

1316 + Forward pass for KDA gate:

1317 + input g: [..., H*D]

1318 + param A: [H] or [1, 1, H, 1]

1319 + beta: softplus beta parameter

1320 + threshold: softplus threshold parameter

1321 + return : [..., H, D]

1322 + """

1323 + orig_shape = g.shape[:-1]

1324 +

1325 + g = g.view(-1, g.shape[-1])

1326 + T = g.shape[0]

1327 + HD = g.shape[1]

1328 + H = A.numel()

1329 + assert H * head_k_dim == HD

1330 +

1331 + y = torch.empty_like(g, dtype=torch.float32)

1332 +

1333 + def grid(meta):

1334 + return (cdiv(T, meta["BT"]), H)

1335 +

1336 + kda_gate_fwd_kernel[grid](

1337 + g,

1338 + A,

1339 + y,

1340 + g_bias,

1341 + beta,

1342 + threshold,

1343 + T,

1344 + H,

1345 + head_k_dim,

1346 + BD=next_power_of_2(head_k_dim),

1347 + HAS_BIAS=g_bias is not None,

1348 + )

1349 +

1350 + y = y.view(*orig_shape, H, head_k_dim)

1351 + return y