Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / platform / torch / custom_ops / gdn / triton / wy_fast.py: 0%
177 statements
« 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
17# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved.
19# pylint: disable=line-too-long,missing-public-type-hints,missing-public-docstring
20# pylint: disable=invalid-name,missing-module-docstring,missing-function-docstring
22from typing import Optional, Tuple
24import torch
25import triton
26import triton.language as tl
28from .utils import prepare_chunk_indices, exp
31@triton.heuristics({
32 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
33})
34@triton.jit(do_not_specialize=['T'])
35def prepare_wy_repr_bwd_kernel(
36 k,
37 v,
38 beta,
39 g,
40 A,
41 dw,
42 du,
43 dk,
44 dv,
45 dbeta,
46 dg,
47 cu_seqlens,
48 chunk_indices,
49 T,
50 B,
51 H: tl.constexpr,
52 K: tl.constexpr,
53 V: tl.constexpr,
54 NT: tl.constexpr,
55 BT: tl.constexpr,
56 BK: tl.constexpr,
57 BV: tl.constexpr,
58 IS_VARLEN: tl.constexpr
59):
60 core_id = tl.program_id(0)
61 total_cores = tl.num_programs(0)
62 T_max = T
64 base_chunks_per_pid = NT // total_cores
65 remainder_chunks = NT % total_cores
67 if core_id < remainder_chunks:
68 chunks_this_pid = base_chunks_per_pid + 1
69 start_idx = core_id * chunks_this_pid
70 else:
71 chunks_this_pid = base_chunks_per_pid
72 start_idx = core_id * chunks_this_pid + remainder_chunks
74 for idx in range(start_idx, start_idx + chunks_this_pid):
75 for i_b in range(B):
76 if IS_VARLEN:
77 i_n, i_t = tl.load(chunk_indices + idx * 2).to(tl.int32), tl.load(chunk_indices + idx * 2 + 1).to(tl.int32)
78 bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
79 T = eos - bos
80 else:
81 i_t = idx
82 bos, eos = i_b * T, i_b * T + T
84 o_t = i_t * BT + tl.arange(0, BT)
85 m_t = o_t < T
86 m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
87 for i_h in range(0, H):
88 if IS_VARLEN:
89 offset = bos + i_h * T_max
90 else:
91 offset = bos * H + i_h * T_max
93 p_beta = tl.make_block_ptr(beta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
94 p_g = tl.make_block_ptr(g + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
95 p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (BT, T), (1, H * BT), (0, i_t * BT), (BT, BT), (0, 1))
97 b_A = tl.load(p_A, boundary_check=(0, 1))
98 b_beta = tl.load(p_beta, boundary_check=(0,))
99 b_g = tl.load(p_g, boundary_check=(0,))
100 b_g_exp = tl.exp(b_g)
102 b_dbeta = tl.zeros([BT], dtype=tl.float32)
103 b_dA = tl.zeros([BT, BT], dtype=tl.float32)
104 b_dg = tl.zeros([BT], dtype=tl.float32)
106 for i_k in range(tl.cdiv(K, BK)):
107 p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
108 p_dk = tl.make_block_ptr(dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
109 p_dw = tl.make_block_ptr(dw + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
110 b_k = tl.load(p_k, boundary_check=(0, 1))
111 b_k_beta_g = (b_k * b_beta[:, None] * b_g_exp[:, None]).to(b_k.dtype)
112 b_dw = tl.load(p_dw, boundary_check=(0, 1))
113 b_dA += tl.dot(b_dw, tl.trans(b_k_beta_g))
114 b_dk_beta_g = tl.dot(b_A, b_dw)
115 b_dk = b_dk_beta_g * b_beta[:, None] * b_g_exp[:, None]
116 b_dbeta += tl.sum(b_dk_beta_g * b_k * b_g_exp[:, None], 1)
117 b_dg += tl.sum(b_dk_beta_g * b_k * b_g_exp[:, None] * b_beta[:, None], 1)
118 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
120 for i_v in range(tl.cdiv(V, BV)):
121 p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
122 p_dv = tl.make_block_ptr(dv + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
123 p_du = tl.make_block_ptr(du + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
124 b_v = tl.load(p_v, boundary_check=(0, 1))
125 b_v_beta = (b_v * b_beta[:, None]).to(b_v.dtype)
126 b_du = tl.load(p_du, boundary_check=(0, 1))
127 b_dA += tl.dot(b_du, tl.trans(b_v_beta))
128 b_dv_beta = tl.dot(b_A, b_du)
129 b_dv = b_dv_beta * b_beta[:, None]
130 b_dbeta += tl.sum(b_dv_beta * b_v, 1)
131 tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), boundary_check=(0, 1))
133 b_dA = tl.where(m_A, b_dA, 0)
134 b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
135 b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
136 b_dA = tl.where(m_A, -b_dA * exp(b_g[:, None] - b_g[None, :]), 0)
137 b_dA = b_dA.to(k.dtype.element_ty)
138 b_A = tl.zeros([BT, BT], dtype=tl.float32)
140 for i_k in range(tl.cdiv(K, BK)):
141 p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
142 p_dk = tl.make_block_ptr(dk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
143 b_k = tl.load(p_k, boundary_check=(0, 1))
144 b_dk = tl.load(p_dk, boundary_check=(0, 1))
145 b_k_beta = (b_k * b_beta[:, None]).to(b_k.dtype)
146 b_A += tl.dot(b_k_beta, tl.trans(b_k))
147 b_dk_beta = tl.dot(b_dA, b_k)
148 b_dbeta += tl.sum(b_dk_beta * b_k, 1)
149 b_dk += tl.dot(tl.trans(b_dA), b_k_beta)
150 b_dk += b_dk_beta * b_beta[:, None]
151 tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), boundary_check=(0, 1))
153 b_dA_A = b_dA * b_A
154 b_dg += tl.sum(b_dA_A, axis=1) - tl.sum(b_dA_A, axis=0)
155 p_dg = tl.make_block_ptr(dg + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
156 p_dbeta = tl.make_block_ptr(dbeta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
157 tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), boundary_check=(0,))
158 tl.store(p_dbeta, b_dbeta.to(p_dbeta.dtype.element_ty), boundary_check=(0,))
161@triton.heuristics({
162 'USE_G': lambda args: args['g'] is not None,
163 'USE_GK': lambda args: args['gk'] is not None,
164 'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
165})
166@triton.jit(do_not_specialize=['T'])
167def recompute_w_u_fwd_kernel(
168 k,
169 v,
170 beta,
171 w,
172 u,
173 A,
174 g,
175 gk,
176 cu_seqlens,
177 chunk_indices,
178 T_tmp,
179 B,
180 H: tl.constexpr,
181 K: tl.constexpr,
182 V: tl.constexpr,
183 NT: tl.constexpr,
184 BT: tl.constexpr,
185 BK: tl.constexpr,
186 BV: tl.constexpr,
187 USE_G: tl.constexpr,
188 USE_GK: tl.constexpr,
189 IS_VARLEN: tl.constexpr
190):
191 core_id = tl.program_id(0)
192 total_cores = tl.num_programs(0)
193 T_max = T_tmp
195 base_chunks_per_pid = NT // total_cores
196 remainder_chunks = NT % total_cores
198 if core_id < remainder_chunks:
199 chunks_this_pid = base_chunks_per_pid + 1
200 start_idx = core_id * chunks_this_pid
201 else:
202 chunks_this_pid = base_chunks_per_pid
203 start_idx = core_id * chunks_this_pid + remainder_chunks
205 for idx in range(start_idx, start_idx + chunks_this_pid):
206 for i_b in range(B):
207 for i_h in range(0, H):
209 if IS_VARLEN:
210 i_n, i_t = tl.load(chunk_indices + idx * 2).to(tl.int32), tl.load(chunk_indices + idx * 2 + 1).to(tl.int32)
211 bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
212 offset = bos + i_h * T_max
213 T = eos - bos
214 else:
215 T = T_tmp
216 i_t = idx
217 bos, eos = i_b * T, i_b * T + T
218 offset = bos * H + i_h * T_max
220 p_beta = tl.make_block_ptr(beta + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
221 b_beta = tl.load(p_beta, boundary_check=(0,))
223 p_A = tl.make_block_ptr(A + (bos * H + i_h) * BT, (T, BT), (H * BT, 1), (i_t * BT, 0), (BT, BT), (1, 0))
224 b_A = tl.load(p_A, boundary_check=(0, 1))
226 for i_v in range(tl.cdiv(V, BV)):
227 p_v = tl.make_block_ptr(v + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
228 p_u = tl.make_block_ptr(u + (bos * H + i_h) * V, (T, V), (H * V, 1), (i_t * BT, i_v * BV), (BT, BV), (1, 0))
229 b_v = tl.load(p_v, boundary_check=(0, 1))
230 b_vb = (b_v * b_beta[:, None]).to(b_v.dtype)
231 b_u = tl.dot(b_A, b_vb, allow_tf32=False)
232 tl.store(p_u, b_u.to(p_u.dtype.element_ty), boundary_check=(0, 1))
234 if USE_G:
235 p_g = tl.make_block_ptr(g + offset, (T,), (1,), (i_t * BT,), (BT,), (0,))
236 b_g = tl.exp(tl.load(p_g, boundary_check=(0,)))
238 for i_k in range(tl.cdiv(K, BK)):
239 p_k = tl.make_block_ptr(k + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
240 p_w = tl.make_block_ptr(w + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
241 b_k = tl.load(p_k, boundary_check=(0, 1))
242 b_kb = b_k * b_beta[:, None]
243 if USE_G:
244 b_kb *= b_g[:, None]
245 if USE_GK:
246 p_gk = tl.make_block_ptr(gk + (bos * H + i_h) * K, (T, K), (H * K, 1), (i_t * BT, i_k * BK), (BT, BK), (1, 0))
247 b_kb *= tl.exp(tl.load(p_gk, boundary_check=(0, 1)))
248 b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
249 tl.store(p_w, b_w.to(p_w.dtype.element_ty), boundary_check=(0, 1))
252def recompute_w_u_fwd(
253 k: torch.Tensor,
254 v: torch.Tensor,
255 beta: torch.Tensor,
256 A: torch.Tensor,
257 g: Optional[torch.Tensor] = None,
258 gk: Optional[torch.Tensor] = None,
259 cu_seqlens: Optional[torch.LongTensor] = None,
260) -> Tuple[torch.Tensor, torch.Tensor]:
261 B, T, H, K, V = *k.shape, v.shape[-1]
262 BT = A.shape[-1]
263 BK = 128
264 BV = 128
266 chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
267 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
268 g = g.transpose(1, 2).contiguous() if g is not None else None
269 beta = beta.transpose(1, 2).contiguous()
271 w = torch.empty_like(k)
272 u = torch.empty_like(v)
273 cv_kernel_num = 24
274 recompute_w_u_fwd_kernel[(cv_kernel_num,)](
275 k=k,
276 v=v,
277 beta=beta,
278 w=w,
279 u=u,
280 A=A,
281 g=g,
282 gk=gk,
283 cu_seqlens=cu_seqlens,
284 chunk_indices=chunk_indices,
285 T_tmp=T,
286 B=B,
287 H=H,
288 K=K,
289 V=V,
290 NT=NT,
291 BT=BT,
292 BK=BK,
293 BV=BV,
294 )
295 return w, u
298def prepare_wy_repr_bwd(
299 k: torch.Tensor,
300 v: torch.Tensor,
301 g: torch.Tensor,
302 beta: torch.Tensor,
303 A: torch.Tensor,
304 dw: torch.Tensor,
305 du: torch.Tensor,
306 cu_seqlens: Optional[torch.LongTensor],
307 chunk_size: int = 64,
308) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
309 B, T, H, K, V = *k.shape, v.shape[-1]
310 BT = chunk_size
311 chunk_indices = prepare_chunk_indices(cu_seqlens, BT) if cu_seqlens is not None else None
312 NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
313 BK = 128
314 BV = 128
315 beta = beta.transpose(1, 2).contiguous()
316 g = g.transpose(1, 2).contiguous()
318 dk = torch.empty_like(k)
319 dv = torch.empty_like(v)
320 dbeta = torch.empty_like(beta)
321 dg = torch.empty_like(g)
323 cv_kernel_num = 24
324 prepare_wy_repr_bwd_kernel[(cv_kernel_num,)](
325 k=k,
326 v=v,
327 beta=beta,
328 g=g,
329 A=A,
330 dw=dw,
331 du=du,
332 dk=dk,
333 dv=dv,
334 dbeta=dbeta,
335 dg=dg,
336 cu_seqlens=cu_seqlens,
337 chunk_indices=chunk_indices,
338 T=T,
339 B=B,
340 H=H,
341 K=K,
342 V=V,
343 NT=NT,
344 BT=BT,
345 BK=BK,
346 BV=BV,
347 )
349 dbeta = dbeta.transpose(1, 2).contiguous()
350 dg = dg.transpose(1, 2).contiguous()
352 return dk, dv, dbeta, dg
355bwd_prepare_wy_repr = prepare_wy_repr_bwd
357fwd_recompute_w_u = recompute_w_u_fwd