« prev ^ index » next coverage.py v7.13.1, created at 2026-08-21 04:29 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15# -*- coding: utf-8 -*-
16# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
18# pylint: disable=line-too-long,missing-public-type-hints,missing-public-docstring
19# pylint: disable=unused-argument,invalid-name,missing-module-docstring
20# pylint: disable=missing-function-docstring
22from typing import Optional
24import torch
25import triton
26import triton.language as tl
28from .utils import prepare_chunk_indices
31@triton.heuristics({
32 'USE_G': lambda args: args['g'] is not None,
33 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
34})
35@triton.jit(do_not_specialize=['T', 'NT', 'TOTAL_TASKS'])
36def chunk_scaled_dot_kkt_fwd_kernel(
37 k,
38 g,
39 beta,
40 A,
41 cu_seqlens,
42 chunk_indices,
43 T,
44 H: tl.constexpr,
45 K: tl.constexpr,
46 BT: tl.constexpr,
47 BK: tl.constexpr,
48 IS_VARLEN: tl.constexpr,
49 USE_G: tl.constexpr,
50 NT,
51 B,
52 TOTAL_TASKS,
53):
54 core_id = tl.program_id(0)
55 num_blocks = tl.num_programs(0)
56 T_max = T
58 base_tasks_per_block = TOTAL_TASKS // num_blocks
59 remainder_tasks = TOTAL_TASKS % num_blocks
61 if core_id < remainder_tasks:
62 tasks_this_core = base_tasks_per_block + 1
63 start_idx = core_id * tasks_this_core
64 else:
65 tasks_this_core = base_tasks_per_block
66 start_idx = core_id * base_tasks_per_block + remainder_tasks
68 for idx in range(start_idx, start_idx + tasks_this_core):
69 i_b = idx // NT
70 local_idx = idx % NT
72 if IS_VARLEN:
73 i_n = tl.load(chunk_indices + local_idx * 2).to(tl.int32)
74 i_t = tl.load(chunk_indices + local_idx * 2 + 1).to(tl.int32)
75 bos = tl.load(cu_seqlens + i_n).to(tl.int32)
76 eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32)
77 T_local = eos - bos
78 else:
79 bos, eos = 0, T
80 i_t = local_idx
81 T_local = T
83 for i_h in range(H):
84 k_batch_off = i_b * T_max * H * K
85 beta_batch_off = i_b * H * T_max
86 g_batch_off = i_b * H * T_max
87 A_batch_off = i_b * T_max * H * BT
89 p_beta = tl.make_block_ptr(beta + beta_batch_off + bos + i_h * T_max, (T_local,), (1,), (i_t * BT,), (BT,), (0,))
90 b_beta = tl.load(p_beta, boundary_check=(0,))
92 b_A = tl.zeros([BT, BT], dtype=tl.float32)
93 for i_k in range(tl.cdiv(K, BK)):
94 p_k = tl.make_block_ptr(k + k_batch_off + (bos * H + i_h) * K, (T_local, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
95 b_k = tl.load(p_k, boundary_check=(0, 1))
96 dot_product = tl.dot(b_k, tl.trans(b_k))
98 o_t = i_t * BT + tl.arange(0, BT)
99 o_t = o_t.to(tl.float32)
100 T_mask = (o_t < T_local).to(tl.float32)
102 row_indices = tl.arange(0, BT)[:, None]
103 col_indices = tl.arange(0, BT)[None, :]
104 tril_mask = (row_indices > col_indices).to(tl.float32)
105 tril_mask = tril_mask * T_mask[:, None]
106 masked_dot = dot_product * tril_mask
107 b_A += masked_dot
109 if USE_G:
110 p_g = tl.make_block_ptr(g + g_batch_off + bos + i_h * T_max, (T_local,), (1,), (i_t * BT,), (BT,), (0,))
111 b_g = tl.load(p_g, boundary_check=(0,))
112 b_g_diff = b_g[:, None] - b_g[None, :]
113 b_g_diff = tl.minimum(tl.maximum(b_g_diff, -50.0), 50.0)
114 b_A *= tl.exp(b_g_diff)
115 b_A *= b_beta[:, None]
117 p_A = tl.make_block_ptr(A + A_batch_off + (bos * H + i_h) * BT, (T_local, BT), (BT * H, 1), (i_t * BT, 0), (BT, BT), (1, 0))
118 tl.store(p_A, b_A.to(p_A.dtype.element_ty), boundary_check=(0, 1))
121@triton.heuristics({
122 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
123})
124@triton.autotune(
125 configs=[
126 triton.Config({'BK': BK})
127 for BK in [32, 64]
128 ],
129 key=["BC"]
130)
131@triton.jit(do_not_specialize=['T'])
132def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter(
133 k,
134 g,
135 beta,
136 A,
137 cu_seqlens,
138 chunk_indices,
139 T,
140 H: tl.constexpr,
141 K: tl.constexpr,
142 BT: tl.constexpr,
143 BC: tl.constexpr,
144 BK: tl.constexpr,
145 NC: tl.constexpr,
146 IS_VARLEN: tl.constexpr,
147):
148 i_t, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
149 i_i, i_j = i_c // NC, i_c % NC
151 for i_h in range(H):
152 if IS_VARLEN:
153 i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
154 bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
155 T_val = eos - bos
156 else:
157 bos, eos = i_b * T, i_b * T + T
158 T_val = T
160 should_compute = (i_t * BT + i_i * BC < T_val) and (i_i > i_j)
162 if should_compute:
163 k_ptr = k + (bos * H + i_h) * K
164 g_ptr = g + (bos * H + i_h) * K
165 A_ptr = A + (bos * H + i_h) * BT
167 p_beta = tl.make_block_ptr(beta + bos * H + i_h, (T_val,), (H,), (i_t * BT + i_i * BC,), (BC,), (0,))
168 b_beta = tl.load(p_beta, boundary_check=(0,))
170 b_A = tl.zeros([BC, BC], dtype=tl.float32)
171 for i_k in range(tl.cdiv(K, BK)):
172 p_k = tl.make_block_ptr(k_ptr, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK),
173 (1, 0))
174 p_g = tl.make_block_ptr(g_ptr, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, i_k * BK), (BC, BK),
175 (1, 0))
176 b_kt = tl.make_block_ptr(k_ptr, (K, T_val), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC),
177 (0, 1))
178 p_gk = tl.make_block_ptr(g_ptr, (K, T_val), (1, H * K), (i_k * BK, i_t * BT + i_j * BC), (BK, BC),
179 (0, 1))
181 o_k = i_k * BK + tl.arange(0, BK)
182 m_k = o_k < K
183 b_gn = tl.load(g_ptr + (i_t * BT + i_i * BC) * H * K + o_k, mask=m_k, other=0)
184 b_g = tl.load(p_g, boundary_check=(0, 1))
185 b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.exp(b_g - b_gn[None, :])
186 b_gk = tl.load(p_gk, boundary_check=(0, 1))
187 b_kt = tl.load(b_kt, boundary_check=(0, 1)) * tl.exp(b_gn[:, None] - b_gk)
188 b_A += tl.dot(b_k, b_kt)
189 b_A *= b_beta[:, None]
191 p_A = tl.make_block_ptr(A_ptr, (T_val, BT), (H * BT, 1), (i_t * BT + i_i * BC, i_j * BC), (BC, BC), (1, 0))
192 tl.store(p_A, b_A.to(A.dtype.element_ty), boundary_check=(0, 1))
195@triton.heuristics({
196 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
197})
198@triton.jit(do_not_specialize=['T'])
199def chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra(
200 k,
201 g,
202 beta,
203 A,
204 cu_seqlens,
205 chunk_indices,
206 T,
207 H: tl.constexpr,
208 K: tl.constexpr,
209 BT: tl.constexpr,
210 BC: tl.constexpr,
211 BK: tl.constexpr,
212 IS_VARLEN: tl.constexpr,
213):
214 i_t, i_i, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
216 for i_h in range(H):
217 if IS_VARLEN:
218 i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
219 bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
220 T_val = eos - bos
221 else:
222 bos, eos = i_b * T, i_b * T + T
223 T_val = T
225 should_compute = i_t * BT + i_i * BC < T_val
227 if should_compute:
228 o_i = tl.arange(0, BC)
229 o_k = tl.arange(0, BK)
230 m_k = o_k < K
231 m_A = (i_t * BT + i_i * BC + o_i) < T_val
232 o_A = (bos + i_t * BT + i_i * BC + o_i) * H * BT + i_h * BT + i_i * BC
234 p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, 0), (BC, BK),
235 (1, 0))
236 p_g = tl.make_block_ptr(g + (bos * H + i_h) * K, (T_val, K), (H * K, 1), (i_t * BT + i_i * BC, 0), (BC, BK),
237 (1, 0))
238 p_beta = beta + (bos + i_t * BT + i_i * BC + o_i) * H + i_h
240 b_k = tl.load(p_k, boundary_check=(0, 1)) * tl.load(p_beta, mask=m_A, other=0)[:, None]
241 b_g = tl.load(p_g, boundary_check=(0, 1))
243 p_kt = k + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
244 p_gk = g + (bos + i_t * BT + i_i * BC) * H * K + i_h * K + o_k
246 for j in range(0, min(BC, T_val - i_t * BT - i_i * BC)):
247 b_kt = tl.load(p_kt, mask=m_k, other=0).to(tl.float32)
248 b_gk = tl.load(p_gk, mask=m_k, other=0).to(tl.float32)
249 b_A = tl.sum(b_k * b_kt[None, :] * tl.exp(b_g - b_gk[None, :]), 1)
250 # 转化成f32
251 o_i_tmp = o_i.to(tl.float32)
252 b_A = tl.where(o_i_tmp > j, b_A, 0.)
254 tl.store(A + o_A + j, b_A, mask=m_A)
255 p_kt += H * K
256 p_gk += H * K
259def chunk_scaled_dot_kkt_fwd(
260 k: torch.Tensor,
261 g: Optional[torch.Tensor] = None,
262 gk: Optional[torch.Tensor] = None,
263 beta: Optional[torch.Tensor] = None,
264 cu_seqlens: Optional[torch.LongTensor] = None,
265 chunk_size: int = 64,
266 output_dtype: torch.dtype = torch.float32
267) -> torch.Tensor:
268 r"""
269 Compute beta * K * K^T.
271 Args:
272 k (torch.Tensor):
273 The key tensor of shape `[B, T, H, K]`.
274 beta (torch.Tensor):
275 The beta tensor of shape `[B, T, H]`.
276 g (torch.Tensor):
277 The cumulative sum of the gate tensor of shape `[B, T, H]`. Default: `None`.
278 gk (torch.Tensor):
279 The cumulative sum of the gate tensor of shape `[B, T, H, K]` applied to the key tensor. Default: `None`.
280 cu_seqlens (torch.LongTensor):
281 The cumulative sequence lengths of the input tensor.
282 Default: None
283 chunk_size (int):
284 The chunk size. Default: 64.
285 output_dtype (torch.dtype):
286 The dtype of the output tensor. Default: `torch.float32`
288 Returns:
289 beta * K * K^T of shape `[B, T, H, BT]` where `BT` is the chunk size.
290 """
291 B, T, H, K = k.shape
292 BT = chunk_size
293 chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
294 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
295 beta = beta.transpose(1, 2).contiguous()
296 g = g.transpose(1, 2).contiguous()
297 BK = 128
298 kernel_num = 24
300 if gk is None:
301 A = torch.empty(B, T, H, BT, device=k.device, dtype=output_dtype)
302 chunk_scaled_dot_kkt_fwd_kernel[(kernel_num,)](
303 k=k,
304 g=g,
305 beta=beta,
306 A=A,
307 cu_seqlens=cu_seqlens,
308 chunk_indices=chunk_indices,
309 T=T,
310 H=H,
311 K=K,
312 BT=BT,
313 BK=BK,
314 NT=NT,
315 B=B,
316 TOTAL_TASKS=B * NT,
317 )
318 return A
320 BC = min(16, BT)
321 NC = triton.cdiv(BT, BC)
322 BK = max(triton.next_power_of_2(K), 16)
323 A = torch.zeros(B, T, H, BT, device=k.device, dtype=output_dtype)
324 grid = (NT, NC * NC, B)
325 chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter[grid](
326 k=k,
327 g=gk,
328 beta=beta,
329 A=A,
330 cu_seqlens=cu_seqlens,
331 chunk_indices=chunk_indices,
332 T=T,
333 H=H,
334 K=K,
335 BT=BT,
336 BC=BC,
337 NC=NC,
338 )
340 grid = (NT, NC, B)
341 chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra[grid](
342 k=k,
343 g=gk,
344 beta=beta,
345 A=A,
346 cu_seqlens=cu_seqlens,
347 chunk_indices=chunk_indices,
348 T=T,
349 H=H,
350 K=K,
351 BT=BT,
352 BC=BC,
353 BK=BK,
354 )
355 return A