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