« 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# The affine state decomposition follows the MIT-licensed implementation in
16# flash-linear-attention/fla/ops/cp/chunk_delta_h.py.
18# pylint: disable=invalid-name,missing-public-type-hints
20"""Fixed-shape Triton-Ascend kernels for KDA state and transfer summaries."""
22__all__ = [
23 "kda_split_state_summary_kernel",
24 "kda_state_transition_kernel",
25 "kda_state_grad_ext_kernel",
26]
28import triton
29import triton.language as tl
32@triton.jit
33def _exp2(value):
34 """Evaluate base-two exponentiation in FP32."""
35 return tl.math.exp2(value.to(tl.float32))
38@triton.jit(do_not_specialize=["T"])
39def kda_split_state_summary_kernel(
40 key: tl.tensor,
41 w: tl.tensor,
42 u: tl.tensor,
43 gate: tl.tensor,
44 state_ext: tl.tensor,
45 transition: tl.tensor,
46 T: tl.tensor,
47 H: tl.constexpr,
48 H_TOTAL: tl.constexpr,
49 HEAD_OFFSET: tl.constexpr,
50 K: tl.constexpr,
51 V: tl.constexpr,
52 BT: tl.constexpr,
53 BV: tl.constexpr,
54 OUTPUT_MODE: tl.constexpr = 0,
55 W_HEAD_FIRST: tl.constexpr = False,
56) -> None:
57 """Build ``S_ext`` and ``M`` directly without a packed-output copy.
59 Args:
60 key: Prepared gated-key pointer.
61 w: Prepared feedback-weight pointer.
62 u: Prepared value pointer.
63 gate: Per-chunk cumulative log2-gate pointer.
64 state_ext: FP32 local-state output pointer.
65 transition: FP32 transfer output pointer, unused in state-only mode.
66 T: Local sequence length.
67 H: Number of heads processed by this launch.
68 H_TOTAL: Head extent in the input layout.
69 HEAD_OFFSET: First input head processed by this launch.
70 K: Key width, currently 128.
71 V: Value width, currently 128.
72 BT: Tokens per chunk, currently 64.
73 BV: Value-column tile width.
74 OUTPUT_MODE: Select combined, state-only, or transfer-only output.
75 W_HEAD_FIRST: Whether feedback weights use head-first storage.
76 """
77 i_col = tl.program_id(0)
78 i_bh = tl.program_id(1)
79 i_b = i_bh // H
80 i_h = i_bh % H + HEAD_OFFSET
82 stride_k = H_TOTAL * K
83 stride_w = K if W_HEAD_FIRST else stride_k
84 stride_v = H_TOTAL * V
85 key += (i_b * T * H_TOTAL + i_h) * K
86 if W_HEAD_FIRST:
87 w += (i_b * H_TOTAL + i_h) * T * K
88 else:
89 w += (i_b * T * H_TOTAL + i_h) * K
90 u += (i_b * T * H_TOTAL + i_h) * V
91 gate += (i_b * T * H_TOTAL + i_h) * K
92 state_ext += i_bh * K * V
93 transition += i_bh * K * K
95 col = i_col * BV + tl.arange(0, BV)
96 row1 = tl.arange(0, 64)
97 row2 = 64 + tl.arange(0, 64)
98 if OUTPUT_MODE == 1:
99 is_transition = False
100 transition_col = col
101 elif OUTPUT_MODE == 2:
102 is_transition = True
103 transition_col = col
104 else:
105 is_transition = i_col * BV >= V
106 transition_col = col - V
107 state1 = tl.where(is_transition & (row1[:, None] == transition_col[None, :]), 1.0, 0.0).to(tl.float32)
108 state2 = tl.where(is_transition & (row2[:, None] == transition_col[None, :]), 1.0, 0.0).to(tl.float32)
110 num_chunks = tl.cdiv(T, BT)
111 for chunk_idx in range(num_chunks):
112 w1_ptr = tl.make_block_ptr(w, (T, K), (stride_w, 1), (chunk_idx * BT, 0), (BT, 64), (1, 0))
113 w2_ptr = tl.make_block_ptr(w, (T, K), (stride_w, 1), (chunk_idx * BT, 64), (BT, 64), (1, 0))
114 w1 = tl.load(w1_ptr, boundary_check=(0, 1))
115 w2 = tl.load(w2_ptr, boundary_check=(0, 1))
116 if OUTPUT_MODE == 1:
117 # Match Ascend FLA h: state and residual feedback stay in FP32.
118 w1 = w1.to(tl.float32)
119 w2 = w2.to(tl.float32)
120 value_new = tl.dot(w1, state1.to(w1.dtype))
121 value_new += tl.dot(w2, state2.to(w2.dtype))
123 if OUTPUT_MODE == 2:
124 value_new = -value_new
125 else:
126 u_col_offset = tl.where(is_transition, 0, i_col * BV)
127 u_ptr = tl.make_block_ptr(u, (T, V), (stride_v, 1), (chunk_idx * BT, u_col_offset), (BT, BV), (1, 0))
128 u_value = tl.load(u_ptr, boundary_check=(0, 1), padding_option='zero')
129 value_new = tl.where(is_transition, 0.0, u_value) - value_new
130 if OUTPUT_MODE != 1:
131 value_new = value_new.to(key.dtype.element_ty)
133 last_idx = min((chunk_idx + 1) * BT, T) - 1
134 gate_last_ptr = gate + last_idx * H_TOTAL * K
135 gate_last1 = tl.load(gate_last_ptr + row1, mask=row1 < K, other=0.0)
136 gate_last2 = tl.load(gate_last_ptr + row2, mask=row2 < K, other=0.0)
137 state1 *= _exp2(gate_last1)[:, None]
138 state2 *= _exp2(gate_last2)[:, None]
140 key1_ptr = tl.make_block_ptr(key, (K, T), (1, stride_k), (0, chunk_idx * BT), (64, BT), (0, 1))
141 key2_ptr = tl.make_block_ptr(key, (K, T), (1, stride_k), (64, chunk_idx * BT), (64, BT), (0, 1))
142 key1 = tl.load(key1_ptr, boundary_check=(0, 1))
143 key2 = tl.load(key2_ptr, boundary_check=(0, 1))
144 if OUTPUT_MODE == 1:
145 key1 = key1.to(tl.float32)
146 key2 = key2.to(tl.float32)
147 state1 += tl.dot(key1, value_new)
148 state2 += tl.dot(key2, value_new)
150 state_mask = (~is_transition) & (col[None, :] < V)
151 transition_mask = is_transition & (transition_col[None, :] < K)
152 tl.store(state_ext + row1[:, None] * V + col[None, :], state1, mask=state_mask)
153 tl.store(state_ext + row2[:, None] * V + col[None, :], state2, mask=state_mask)
154 tl.store(transition + row1[:, None] * K + transition_col[None, :], state1, mask=transition_mask)
155 tl.store(transition + row2[:, None] * K + transition_col[None, :], state2, mask=transition_mask)
158@triton.jit(do_not_specialize=["T"])
159def kda_state_transition_kernel(
160 key: tl.tensor,
161 w: tl.tensor,
162 gate: tl.tensor,
163 transition: tl.tensor,
164 T: tl.tensor,
165 H: tl.constexpr,
166 K: tl.constexpr,
167 BT: tl.constexpr,
168 BV: tl.constexpr,
169) -> None:
170 """Compose explicit chunk transfers with FP32 feedback, as in FLA CP.
172 Args:
173 key: Contiguous prepared gated-key pointer, shaped [B, T, H, K].
174 w: Contiguous prepared feedback-weight pointer with the same shape.
175 gate: Contiguous per-chunk cumulative log2-gate pointer.
176 transition: FP32 output pointer, shaped [B, H, K, K].
177 T: Local sequence length, a multiple of BT.
178 H: Number of local heads.
179 K: Key width, currently 128.
180 BT: Tokens per chunk, currently 64.
181 BV: Number of output columns owned by each program.
182 """
183 column = tl.program_id(0)
184 head = tl.program_id(1)
185 batch = head // H
186 local_head = head % H
187 rows = tl.arange(0, K)
188 columns = column * BV + tl.arange(0, BV)
189 tokens = tl.arange(0, BT)
190 output_offsets = head * K * K + rows[:, None] * K + columns[None, :]
191 matrix = (rows[:, None] == columns[None, :]).to(tl.float32)
192 # Ascend cannot carry the Cube result directly into the next Cube input.
193 # Each program owns disjoint output columns, so this FP32 feedback needs
194 # neither an inter-program barrier nor a per-chunk matrix workspace.
195 tl.store(transition + output_offsets, matrix)
196 for chunk in range(T // BT):
197 offsets = ((batch * T + chunk * BT + tokens[:, None]) * H + local_head) * K + rows[None, :]
198 key_tile = tl.load(key + offsets)
199 weight_tile = tl.load(w + offsets)
200 last = ((batch * T + chunk * BT + BT - 1) * H + local_head) * K + rows
201 decay = tl.exp2(tl.load(gate + last).to(tl.float32))
202 diagonal = tl.where(rows[:, None] == rows[None, :], decay[:, None], 0.0)
203 interaction = tl.dot(tl.trans(key_tile), weight_tile)
204 chunk_transition = diagonal - interaction
205 previous = tl.load(transition + output_offsets)
206 updated = tl.dot(chunk_transition, previous, input_precision="ieee")
207 tl.store(transition + output_offsets, updated)
210@triton.jit(do_not_specialize=["T"])
211def kda_state_grad_ext_kernel(
212 query: tl.tensor,
213 key: tl.tensor,
214 w: tl.tensor,
215 gate: tl.tensor,
216 grad_output: tl.tensor,
217 grad_value: tl.tensor,
218 grad_state_ext: tl.tensor,
219 scale: tl.tensor,
220 T: tl.tensor,
221 H: tl.constexpr,
222 H_TOTAL: tl.constexpr,
223 HEAD_OFFSET: tl.constexpr,
224 K: tl.constexpr,
225 V: tl.constexpr,
226 BT: tl.constexpr,
227 BV: tl.constexpr,
228) -> None:
229 """Build the local-output contribution to the incoming state gradient.
231 Args:
232 query: Prepared gated-query pointer.
233 key: Prepared gated-key pointer.
234 w: Prepared feedback-weight pointer.
235 gate: Per-chunk cumulative log2-gate pointer.
236 grad_output: Original token-output gradient pointer.
237 grad_value: Local value-gradient pointer before recurrent feedback.
238 grad_state_ext: FP32 local state-gradient output pointer.
239 scale: Attention scale applied before the dot-input cast.
240 T: Local sequence length.
241 H: Number of heads processed by this launch.
242 H_TOTAL: Head extent in the input layout.
243 HEAD_OFFSET: First input head processed by this launch.
244 K: Key width, currently 128.
245 V: Value width, currently 128.
246 BT: Tokens per chunk, currently 64.
247 BV: Value-column tile width.
248 """
249 i_col = tl.program_id(0)
250 i_bh = tl.program_id(1)
251 i_b = i_bh // H
252 i_h = i_bh % H + HEAD_OFFSET
254 stride_k = H_TOTAL * K
255 stride_v = H_TOTAL * V
256 query += (i_b * T * H_TOTAL + i_h) * K
257 key += (i_b * T * H_TOTAL + i_h) * K
258 w += (i_b * T * H_TOTAL + i_h) * K
259 gate += (i_b * T * H_TOTAL + i_h) * K
260 grad_output += (i_b * T * H_TOTAL + i_h) * V
261 grad_value += (i_b * T * H_TOTAL + i_h) * V
262 grad_state_ext += i_bh * K * V
264 row1 = tl.arange(0, 64)
265 row2 = 64 + tl.arange(0, 64)
266 state1 = tl.zeros([64, BV], dtype=tl.float32)
267 state2 = tl.zeros([64, BV], dtype=tl.float32)
269 num_chunks = tl.cdiv(T, BT)
270 for reverse_idx in range(num_chunks):
271 chunk_idx = num_chunks - 1 - reverse_idx
273 key1_ptr = tl.make_block_ptr(key, (T, K), (stride_k, 1), (chunk_idx * BT, 0), (BT, 64), (1, 0))
274 key2_ptr = tl.make_block_ptr(key, (T, K), (stride_k, 1), (chunk_idx * BT, 64), (BT, 64), (1, 0))
275 key1 = tl.load(key1_ptr, boundary_check=(0, 1))
276 key2 = tl.load(key2_ptr, boundary_check=(0, 1))
277 value_grad = tl.dot(key1, state1.to(key1.dtype))
278 value_grad += tl.dot(key2, state2.to(key2.dtype))
280 grad_value_ptr = tl.make_block_ptr(
281 grad_value,
282 (T, V),
283 (stride_v, 1),
284 (chunk_idx * BT, i_col * BV),
285 (BT, BV),
286 (1, 0),
287 )
288 value_grad += tl.load(grad_value_ptr, boundary_check=(0, 1))
290 last_idx = min((chunk_idx + 1) * BT, T) - 1
291 gate_last_ptr = gate + last_idx * H_TOTAL * K
292 gate_last1 = tl.load(gate_last_ptr + row1, mask=row1 < K, other=0.0)
293 gate_last2 = tl.load(gate_last_ptr + row2, mask=row2 < K, other=0.0)
294 state1 *= _exp2(gate_last1)[:, None]
295 state2 *= _exp2(gate_last2)[:, None]
297 output_grad_ptr = tl.make_block_ptr(
298 grad_output,
299 (T, V),
300 (stride_v, 1),
301 (chunk_idx * BT, i_col * BV),
302 (BT, BV),
303 (1, 0),
304 )
305 output_grad = tl.load(output_grad_ptr, boundary_check=(0, 1))
306 # Match FLA dhu's scale-before-cast contract for the Q contribution.
307 output_grad = output_grad * scale
308 output_grad_second = output_grad + 0.0
310 query1_ptr = tl.make_block_ptr(query, (K, T), (1, stride_k), (0, chunk_idx * BT), (64, BT), (0, 1))
311 query2_ptr = tl.make_block_ptr(query, (K, T), (1, stride_k), (64, chunk_idx * BT), (64, BT), (0, 1))
312 w1_ptr = tl.make_block_ptr(w, (K, T), (1, stride_k), (0, chunk_idx * BT), (64, BT), (0, 1))
313 w2_ptr = tl.make_block_ptr(w, (K, T), (1, stride_k), (64, chunk_idx * BT), (64, BT), (0, 1))
314 query1 = tl.load(query1_ptr, boundary_check=(0, 1))
315 query2 = tl.load(query2_ptr, boundary_check=(0, 1))
316 w1 = tl.load(w1_ptr, boundary_check=(0, 1))
317 w2 = tl.load(w2_ptr, boundary_check=(0, 1))
318 state1 += tl.dot(query1, output_grad.to(query1.dtype)) - tl.dot(w1, value_grad.to(w1.dtype))
319 state2 += tl.dot(query2, output_grad_second.to(query2.dtype)) - tl.dot(w2, value_grad.to(w2.dtype))
321 output1_ptr = tl.make_block_ptr(grad_state_ext, (K, V), (V, 1), (0, i_col * BV), (64, BV), (1, 0))
322 output2_ptr = tl.make_block_ptr(grad_state_ext, (K, V), (V, 1), (64, i_col * BV), (64, BV), (1, 0))
323 tl.store(output1_ptr, state1.to(output1_ptr.dtype.element_ty), boundary_check=(0, 1))
324 tl.store(output2_ptr, state2.to(output2_ptr.dtype.element_ty), boundary_check=(0, 1))