« 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#
17# The state-summary decomposition follows the MIT-licensed context-parallel
18# implementation in flash-linear-attention/fla/ops/cp/chunk_delta_h.py.
20# pylint: disable=missing-public-type-hints,invalid-name
22"""Fixed-shape Triton-Ascend kernels for GDN state summaries."""
24import triton
25import triton.language as tl
27from .utils import get_autotune_config
30@triton.autotune(
31 configs=get_autotune_config(
32 multibuffer_list=(True, False),
33 set_workspace_multibuffer_list=(2, 4),
34 tile_mix_vector_loop_num_list=(2,),
35 tile_mix_cube_loop_num_list=(2,),
36 ),
37 key=["H", "K", "V", "BT", "BV"],
38)
39@triton.jit(do_not_specialize=["T"])
40def gdn_packed_state_summary_kernel(
41 k,
42 w,
43 u,
44 g,
45 packed_summary,
46 T,
47 H: tl.constexpr,
48 K: tl.constexpr,
49 V: tl.constexpr,
50 BT: tl.constexpr,
51 BV: tl.constexpr,
52 NT: tl.constexpr,
53):
54 """Build the local affine state transition and extension in one buffer."""
55 i_v = tl.program_id(0)
56 i_bh = tl.program_id(1)
57 i_b = i_bh // H
58 i_h = i_bh % H
60 stride_k = H * K
61 stride_v = H * V
62 k += (i_b * T * H + i_h) * K
63 w += (i_b * T * H + i_h) * K
64 u += (i_b * T * H + i_h) * V
65 g += i_b * T * H + i_h
66 packed_summary += i_bh * K * (V + K)
68 col = tl.arange(0, BV)
69 row1 = tl.arange(0, 64)
70 row2 = 64 + tl.arange(0, 64)
71 is_transition = i_v * BV >= V
72 transition_col = i_v * BV - V + col
73 b_h1 = tl.where(
74 is_transition & (row1[:, None] == transition_col[None, :]), 1.0, 0.0
75 ).to(tl.float32)
76 b_h2 = tl.where(
77 is_transition & (row2[:, None] == transition_col[None, :]), 1.0, 0.0
78 ).to(tl.float32)
80 for i_t in range(NT):
81 p_w1 = tl.make_block_ptr(
82 w, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)
83 )
84 p_w2 = tl.make_block_ptr(
85 w, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)
86 )
87 b_w1 = tl.load(p_w1, boundary_check=(0, 1))
88 b_w2 = tl.load(p_w2, boundary_check=(0, 1))
89 b_v = tl.dot(b_w1, b_h1.to(b_w1.dtype))
90 b_v += tl.dot(b_w2, b_h2.to(b_w2.dtype))
92 p_u = tl.make_block_ptr(
93 u,
94 (T, V),
95 (stride_v, 1),
96 (i_t * BT, i_v * BV),
97 (BT, BV),
98 (1, 0),
99 )
100 b_v = tl.load(p_u, boundary_check=(0, 1)) - b_v
102 last_idx = min((i_t + 1) * BT, T) - 1
103 token = i_t * BT + tl.arange(0, BT)
104 mask = token < T
105 b_g_last = tl.load(g + last_idx * H).to(tl.float32)
106 b_g = tl.load(g + token * H, mask=mask, other=0.0).to(tl.float32)
107 b_v *= tl.where(mask, tl.exp(b_g_last - b_g), 0.0)[:, None]
108 decay = tl.exp(b_g_last)
109 b_h1 *= decay
110 b_h2 *= decay
111 b_v = b_v.to(k.dtype.element_ty)
113 p_k1 = tl.make_block_ptr(
114 k, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)
115 )
116 p_k2 = tl.make_block_ptr(
117 k, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
118 )
119 b_h1 += tl.dot(tl.load(p_k1, boundary_check=(0, 1)), b_v)
120 b_h2 += tl.dot(tl.load(p_k2, boundary_check=(0, 1)), b_v)
122 p_out1 = tl.make_block_ptr(
123 packed_summary,
124 (K, V + K),
125 (V + K, 1),
126 (0, i_v * BV),
127 (64, BV),
128 (1, 0),
129 )
130 p_out2 = tl.make_block_ptr(
131 packed_summary,
132 (K, V + K),
133 (V + K, 1),
134 (64, i_v * BV),
135 (64, BV),
136 (1, 0),
137 )
138 tl.store(p_out1, b_h1, boundary_check=(0, 1))
139 tl.store(p_out2, b_h2, boundary_check=(0, 1))
142@triton.autotune(
143 configs=get_autotune_config(
144 multibuffer_list=(True, False),
145 set_workspace_multibuffer_list=(2, 4),
146 tile_mix_vector_loop_num_list=(2,),
147 tile_mix_cube_loop_num_list=(2,),
148 ),
149 key=["H", "K", "V", "BT"],
150)
151@triton.jit(do_not_specialize=["T"])
152def gdn_state_grad_ext_kernel(
153 q,
154 k,
155 w,
156 g,
157 do,
158 dv,
159 grad_state_ext,
160 scale,
161 T,
162 H: tl.constexpr,
163 K: tl.constexpr,
164 V: tl.constexpr,
165 BT: tl.constexpr,
166 BV: tl.constexpr,
167 NT: tl.constexpr,
168):
169 """Build the local-loss contribution to the incoming state gradient."""
170 i_v = tl.program_id(0)
171 i_bh = tl.program_id(1)
172 i_b = i_bh // H
173 i_h = i_bh % H
175 stride_k = H * K
176 stride_v = H * V
177 q += (i_b * T * H + i_h) * K
178 k += (i_b * T * H + i_h) * K
179 w += (i_b * T * H + i_h) * K
180 g += i_b * T * H + i_h
181 do += (i_b * T * H + i_h) * V
182 dv += (i_b * T * H + i_h) * V
183 grad_state_ext += i_bh * K * V
185 b_dh1 = tl.zeros([64, BV], dtype=tl.float32)
186 b_dh2 = tl.zeros([64, BV], dtype=tl.float32)
188 for reverse_idx in range(NT):
189 i_t = NT - 1 - reverse_idx
190 last_idx = min((i_t + 1) * BT, T) - 1
191 token = i_t * BT + tl.arange(0, BT)
192 mask = token < T
193 b_g_last = tl.load(g + last_idx * H).to(tl.float32)
194 b_g = tl.load(g + token * H, mask=mask, other=0.0).to(tl.float32)
196 p_k1 = tl.make_block_ptr(
197 k, (T, K), (stride_k, 1), (i_t * BT, 0), (BT, 64), (1, 0)
198 )
199 p_k2 = tl.make_block_ptr(
200 k, (T, K), (stride_k, 1), (i_t * BT, 64), (BT, 64), (1, 0)
201 )
202 b_k1 = tl.load(p_k1, boundary_check=(0, 1))
203 b_k2 = tl.load(p_k2, boundary_check=(0, 1))
204 b_dv = tl.dot(b_k1, b_dh1.to(b_k1.dtype))
205 b_dv += tl.dot(b_k2, b_dh2.to(b_k2.dtype))
206 b_dv *= tl.where(mask, tl.exp(b_g_last - b_g), 0.0)[:, None]
208 p_dv = tl.make_block_ptr(
209 dv,
210 (T, V),
211 (stride_v, 1),
212 (i_t * BT, i_v * BV),
213 (BT, BV),
214 (1, 0),
215 )
216 b_dv += tl.load(p_dv, boundary_check=(0, 1))
218 p_do = tl.make_block_ptr(
219 do,
220 (T, V),
221 (stride_v, 1),
222 (i_t * BT, i_v * BV),
223 (BT, BV),
224 (1, 0),
225 )
226 b_do = tl.load(p_do, boundary_check=(0, 1))
227 decay = tl.exp(b_g_last)
228 b_dh1 *= decay
229 b_dh2 *= decay
231 p_q1 = tl.make_block_ptr(
232 q, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)
233 )
234 p_q2 = tl.make_block_ptr(
235 q, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
236 )
237 p_w1 = tl.make_block_ptr(
238 w, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1)
239 )
240 p_w2 = tl.make_block_ptr(
241 w, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1)
242 )
243 b_q1 = tl.load(p_q1, boundary_check=(0, 1))
244 b_q2 = tl.load(p_q2, boundary_check=(0, 1))
245 b_w1 = tl.load(p_w1, boundary_check=(0, 1))
246 b_w2 = tl.load(p_w2, boundary_check=(0, 1))
247 gate = tl.exp(b_g)[None, :]
248 b_q1 *= gate
249 b_q2 *= gate
250 b_dh1 += tl.dot(b_q1, b_do.to(b_q1.dtype)) * scale
251 b_dh1 -= tl.dot(b_w1, b_dv.to(b_w1.dtype))
252 b_dh2 += tl.dot(b_q2, b_do.to(b_q2.dtype)) * scale
253 b_dh2 -= tl.dot(b_w2, b_dv.to(b_w2.dtype))
255 p_out1 = tl.make_block_ptr(
256 grad_state_ext, (K, V), (V, 1), (0, i_v * BV), (64, BV), (1, 0)
257 )
258 p_out2 = tl.make_block_ptr(
259 grad_state_ext, (K, V), (V, 1), (64, i_v * BV), (64, BV), (1, 0)
260 )
261 tl.store(p_out1, b_dh1, boundary_check=(0, 1))
262 tl.store(p_out2, b_dh2, boundary_check=(0, 1))
265__all__ = ["gdn_packed_state_summary_kernel", "gdn_state_grad_ext_kernel"]