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