« 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"""Affine KDA state summary built from prepared WY intermediates."""
17# pylint: disable=forbidden-backend-import
19__all__ = [
20 "apply_kda_state_gradient_summary",
21 "apply_kda_state_summary",
22 "kda_state_gradient_summary_from_prepared",
23 "kda_state_summary_forward_from_prepared",
24]
26from typing import Optional
28import torch
29import triton
32def _validate_prepared_summary_inputs(
33 key: torch.Tensor,
34 w: torch.Tensor,
35 u: torch.Tensor,
36 gate: torch.Tensor,
37 chunk_size: int,
38) -> tuple[int, int, int, int, int]:
39 """Validate the fixed Kimi K3 summary contract and return dimensions."""
40 if (key.ndim, w.ndim, u.ndim, gate.ndim) != (4, 4, 4, 4):
41 raise ValueError("KDA prepared summary expects rank-4 key/w/u/gate tensors.")
42 batch, sequence_length, heads, key_dim = key.shape
43 value_dim = u.shape[-1]
44 if key_dim != 128 or value_dim != 128 or chunk_size != 64:
45 raise NotImplementedError(
46 "Triton KDA state summary requires key_dim=value_dim=128 and "
47 "chunk_size=64."
48 )
49 if sequence_length % chunk_size:
50 raise ValueError(
51 f"KDA summary sequence length {sequence_length} must be divisible "
52 f"by chunk_size {chunk_size}."
53 )
54 if w.shape != key.shape or gate.shape != key.shape:
55 raise ValueError(
56 "KDA prepared key, w, and gate must have identical shapes, got "
57 f"key={tuple(key.shape)}, w={tuple(w.shape)}, gate={tuple(gate.shape)}."
58 )
59 if u.shape[:3] != key.shape[:3]:
60 raise ValueError(
61 f"KDA prepared u prefix {tuple(u.shape[:3])} must match "
62 f"key prefix {tuple(key.shape[:3])}."
63 )
64 if not key.is_npu:
65 raise RuntimeError("Triton KDA state summary requires Ascend NPU tensors.")
66 return batch, sequence_length, heads, key_dim, value_dim
69def _validate_gradient_summary_inputs(
70 query: torch.Tensor,
71 key: torch.Tensor,
72 w: torch.Tensor,
73 gate: torch.Tensor,
74 grad_output: torch.Tensor,
75 grad_value: torch.Tensor,
76 chunk_size: int,
77) -> tuple[int, int, int, int, int]:
78 """Validate prepared tensors used by the reverse state wavefront."""
79 tensors = (query, key, w, gate, grad_output, grad_value)
80 if any(tensor.ndim != 4 for tensor in tensors):
81 raise ValueError(
82 "KDA gradient summary expects rank-4 query/key/w/gate/do/dv tensors."
83 )
84 batch, sequence_length, heads, key_dim = query.shape
85 value_dim = grad_output.shape[-1]
86 if key_dim != 128 or value_dim != 128 or chunk_size != 64:
87 raise NotImplementedError(
88 "Triton KDA state-gradient summary requires key_dim=value_dim=128 "
89 "and chunk_size=64."
90 )
91 if sequence_length % chunk_size:
92 raise ValueError(
93 f"KDA gradient-summary sequence length {sequence_length} must be "
94 f"divisible by chunk_size {chunk_size}."
95 )
96 prepared_shape = (batch, sequence_length, heads, key_dim)
97 value_shape = (batch, sequence_length, heads, value_dim)
98 if key.shape != prepared_shape or w.shape != prepared_shape:
99 raise ValueError(
100 "KDA prepared query, key, and w must have identical shapes, got "
101 f"query={tuple(query.shape)}, key={tuple(key.shape)}, "
102 f"w={tuple(w.shape)}."
103 )
104 if gate.shape != prepared_shape:
105 raise ValueError(
106 f"KDA prepared gate shape {tuple(gate.shape)} must match "
107 f"query shape {prepared_shape}."
108 )
109 if grad_output.shape != value_shape or grad_value.shape != value_shape:
110 raise ValueError(
111 "KDA grad_output and grad_value must have shape "
112 f"{value_shape}, got do={tuple(grad_output.shape)}, "
113 f"dv={tuple(grad_value.shape)}."
114 )
115 if not query.is_npu:
116 raise RuntimeError(
117 "Triton KDA state-gradient summary requires Ascend NPU tensors."
118 )
119 return batch, sequence_length, heads, key_dim, value_dim
122def _launch_kda_state_summary(
123 key: torch.Tensor,
124 w: torch.Tensor,
125 u: torch.Tensor,
126 gate: torch.Tensor,
127 *,
128 chunk_size: int,
129 output_transition: bool = True,
130) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
131 """Launch the shared S/M summary used by the state CP protocols."""
132 return _launch_kda_paired_state_summary(
133 key,
134 w,
135 u,
136 gate,
137 chunk_size=chunk_size,
138 output_transition=output_transition,
139 )
142def _launch_kda_state_gradient_summary(
143 query: torch.Tensor,
144 key: torch.Tensor,
145 w: torch.Tensor,
146 gate: torch.Tensor,
147 grad_output: torch.Tensor,
148 grad_value: torch.Tensor,
149 scale: float,
150 *,
151 chunk_size: int,
152 block_size: int = 64,
153) -> torch.Tensor:
154 """Launch the reverse-wavefront Triton-Ascend summary kernel."""
155 from ._triton.kimi_delta_attention.state_summary import ( # pylint: disable=import-outside-toplevel
156 kda_state_grad_ext_kernel,
157 )
159 batch, sequence_length, heads, key_dim, value_dim = (
160 _validate_gradient_summary_inputs(
161 query,
162 key,
163 w,
164 gate,
165 grad_output,
166 grad_value,
167 chunk_size,
168 )
169 )
170 query, key, w, gate, grad_output, grad_value = (
171 tensor.contiguous()
172 for tensor in (query, key, w, gate, grad_output, grad_value)
173 )
174 grad_state_ext = torch.empty(
175 batch,
176 heads,
177 key_dim,
178 value_dim,
179 dtype=torch.float32,
180 device=query.device,
181 )
182 kda_state_grad_ext_kernel[
183 (triton.cdiv(value_dim, block_size), batch * heads)
184 ](
185 query,
186 key,
187 w,
188 gate,
189 grad_output,
190 grad_value,
191 grad_state_ext,
192 float(scale),
193 sequence_length,
194 H=heads,
195 H_TOTAL=heads,
196 HEAD_OFFSET=0,
197 K=key_dim,
198 V=value_dim,
199 BT=chunk_size,
200 BV=block_size,
201 num_warps=2,
202 )
203 return grad_state_ext
206@torch.compiler.disable
207def _launch_kda_paired_state_summary(
208 key: torch.Tensor,
209 w: torch.Tensor,
210 u: torch.Tensor,
211 gate: torch.Tensor,
212 *,
213 chunk_size: int = 64,
214 output_transition: bool = True,
215) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
216 """Build Ascend FLA-paired ``S_ext`` and an explicit FP32 transfer ``M``."""
217 from ._triton.kimi_delta_attention.state_summary import ( # pylint: disable=import-outside-toplevel
218 kda_split_state_summary_kernel,
219 kda_state_transition_kernel,
220 )
222 batch, sequence_length, heads, key_dim, value_dim = (
223 _validate_prepared_summary_inputs(key, w, u, gate, chunk_size)
224 )
225 key, w, u, gate = (
226 tensor.contiguous() for tensor in (key, w, u, gate)
227 )
228 state_ext = torch.empty(
229 batch,
230 heads,
231 key_dim,
232 value_dim,
233 dtype=torch.float32,
234 device=key.device,
235 )
236 transition = None
237 if output_transition:
238 transition = torch.empty(
239 batch,
240 heads,
241 key_dim,
242 key_dim,
243 dtype=torch.float32,
244 device=key.device,
245 )
246 kda_split_state_summary_kernel[(1, batch * heads)](
247 key,
248 w,
249 u,
250 gate,
251 state_ext,
252 transition if transition is not None else state_ext,
253 sequence_length,
254 H=heads,
255 H_TOTAL=heads,
256 HEAD_OFFSET=0,
257 K=key_dim,
258 V=value_dim,
259 BT=chunk_size,
260 BV=128,
261 OUTPUT_MODE=1,
262 num_warps=2,
263 )
264 if output_transition:
265 kda_state_transition_kernel[(triton.cdiv(key_dim, 64), batch * heads)](
266 key,
267 w,
268 gate,
269 transition,
270 sequence_length,
271 H=heads,
272 K=key_dim,
273 BT=chunk_size,
274 BV=64,
275 num_warps=4,
276 num_stages=2,
277 )
278 return state_ext, transition
281@torch.compiler.disable
282def kda_state_summary_forward_from_prepared(
283 key: torch.Tensor,
284 w: torch.Tensor,
285 u: torch.Tensor,
286 gate: torch.Tensor,
287 *,
288 chunk_size: int = 64,
289 output_transition: bool = True,
290) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
291 """Build the local affine summary for state context parallelism.
293 ``S_ext`` follows Ascend FLA h's FP32 feedback. ``M`` composes explicit
294 chunk transfers with FP32 feedback, following FLA CP.
295 ``output_transition=False`` returns only S and skips M allocation/computation.
296 It is valid for a zero-entry P2P endpoint whose backward never applies M.
298 Args:
299 key: Prepared keys, shaped [batch, tokens, heads, key_dim].
300 w: Prepared feedback weights with the same shape as key.
301 u: Prepared values, shaped [batch, tokens, heads, value_dim].
302 gate: Per-chunk cumulative log2 gates matching key.
303 chunk_size: Number of tokens per local chunk; the kernel requires 64.
304 output_transition: Whether the caller needs the local transition M.
306 Returns:
307 The FP32 state S and the FP32 transition M, or None when M is omitted.
308 """
309 return _launch_kda_state_summary(
310 key,
311 w,
312 u,
313 gate,
314 chunk_size=chunk_size,
315 output_transition=output_transition,
316 )
319@torch.compiler.disable
320def kda_state_gradient_summary_from_prepared(
321 query: torch.Tensor,
322 key: torch.Tensor,
323 w: torch.Tensor,
324 gate: torch.Tensor,
325 grad_output: torch.Tensor,
326 grad_value: torch.Tensor,
327 scale: float,
328 *,
329 chunk_size: int = 64,
330) -> torch.Tensor:
331 """Return the local-loss contribution to the incoming KDA state gradient.
333 Match Ascend FLA ``dhu`` at zero terminal state: scale ``grad_output``
334 before its dot-input cast and preserve the two K64 slabs and subtraction
335 grouping. This summary omits the full ``dh``/``dv`` outputs. All inputs
336 are read-only; the affine application to a nonzero state is a separate step.
338 Args:
339 query: Prepared gated queries [B,T,H,K].
340 key: Matching prepared gated keys.
341 w: Prepared feedback weights matching key.
342 gate: Per-chunk cumulative log2 gates matching key.
343 grad_output: Original token-output gradients [B,T,H,V].
344 grad_value: Local value gradients before recurrent feedback.
345 scale: Attention scaling factor.
346 chunk_size: Tokens per local chunk; the kernel requires 64.
348 Returns:
349 FP32 local contribution G, shaped [B,H,K,V].
350 """
351 return _launch_kda_state_gradient_summary(
352 query,
353 key,
354 w,
355 gate,
356 grad_output,
357 grad_value,
358 scale,
359 chunk_size=chunk_size,
360 block_size=128,
361 )
364def apply_kda_state_summary(
365 state_ext: torch.Tensor,
366 transition: Optional[torch.Tensor],
367 initial_state: Optional[torch.Tensor],
368) -> torch.Tensor:
369 """Apply a prepared KDA state summary in FP32.
371 Args:
372 state_ext: Local zero-entry state S.
373 transition: Local transfer M; may be omitted for a zero entry.
374 initial_state: Incoming FP32 state, or None for a zero entry.
376 Returns:
377 The outgoing FP32 state.
378 """
379 if initial_state is None:
380 return state_ext
381 return torch.matmul(transition, initial_state.float()) + state_ext
384def apply_kda_state_gradient_summary(
385 grad_state_ext: torch.Tensor,
386 transition: torch.Tensor,
387 grad_final_state: Optional[torch.Tensor],
388) -> torch.Tensor:
389 """Apply the adjoint summary for the reverse P2P state wavefront."""
390 if grad_final_state is None:
391 return grad_state_ext
392 return (
393 torch.matmul(transition.transpose(-2, -1), grad_final_state.float())
394 + grad_state_ext
395 )