Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / kimi_delta_attention.py: 0%
160 statements
« 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"""Dense Triton-Ascend KDA execution helpers for Torch training."""
16from __future__ import annotations
18# pylint: disable=forbidden-backend-import
20__all__ = ["fused_chunk_kda", "fused_chunk_kda_p2p"]
22from typing import Any, Optional
24import torch
25import torch.distributed as dist
27from .kimi_delta_attention_fla_adapter import get_fla_kda_staged_ops, run_fla_chunk_kda
28from .kimi_delta_attention_state_summary import (
29 apply_kda_state_gradient_summary,
30 apply_kda_state_summary,
31 kda_state_gradient_summary_from_prepared,
32 kda_state_summary_forward_from_prepared,
33)
36def _validate_local_shapes(
37 query: torch.Tensor,
38 key: torch.Tensor,
39 value: torch.Tensor,
40 gate: torch.Tensor,
41 beta: torch.Tensor,
42) -> tuple[int, int, int, int]:
43 """Validate local KDA tensor shapes and return dimensions used by later checks."""
44 if (query.ndim, key.ndim, value.ndim, gate.ndim) != (4, 4, 4, 4):
45 raise ValueError("Fused KDA expects rank-4 query/key/value/gate tensors.")
46 if beta.ndim != 3:
47 raise ValueError("Fused KDA expects a rank-3 beta tensor.")
48 if query.shape != key.shape:
49 raise ValueError("Fused KDA query and key must have identical shapes.")
50 batch, sequence_length, num_query_heads, key_dim = query.shape
51 num_value_heads, value_dim = value.shape[2:]
52 if value.shape[:2] != (batch, sequence_length):
53 raise ValueError("Fused KDA value must match query batch and sequence dimensions.")
54 if gate.shape != (batch, sequence_length, num_value_heads, key_dim):
55 raise ValueError("Fused KDA gate has an incompatible shape.")
56 if beta.shape != (batch, sequence_length, num_value_heads):
57 raise ValueError("Fused KDA beta has an incompatible shape.")
58 if num_value_heads % num_query_heads:
59 raise ValueError("Fused KDA value heads must be divisible by query heads.")
60 return sequence_length, key_dim, value_dim, num_value_heads
63def _validate_local_inputs(
64 query: torch.Tensor,
65 key: torch.Tensor,
66 value: torch.Tensor,
67 gate: torch.Tensor,
68 beta: torch.Tensor,
69 a_log: torch.Tensor,
70 dt_bias: torch.Tensor,
71 *,
72 chunk_size: int,
73 lower_bound: float,
74) -> None:
75 """Validate the fixed-shape local KDA backend before compiling kernels."""
76 sequence_length, key_dim, value_dim, num_value_heads = _validate_local_shapes(
77 query, key, value, gate, beta
78 )
79 if not (
80 query.dtype == key.dtype == value.dtype == gate.dtype == beta.dtype
81 == torch.bfloat16
82 ):
83 raise TypeError("Fused KDA requires q/k/v/gate/beta tensors in bfloat16.")
84 if a_log.dtype != torch.float32 or dt_bias.dtype != torch.float32:
85 raise TypeError("Fused KDA requires a_log and dt_bias tensors in float32.")
86 if key_dim != 128 or value_dim != 128 or chunk_size != 64:
87 raise NotImplementedError(
88 "Fused KDA currently requires key_dim=value_dim=128 and chunk_size=64."
89 )
90 if sequence_length % chunk_size:
91 raise ValueError(
92 f"Fused KDA sequence length {sequence_length} must be divisible "
93 f"by chunk_size {chunk_size}."
94 )
95 if a_log.numel() != num_value_heads:
96 raise ValueError("Fused KDA a_log must contain one value per value head.")
97 if dt_bias.numel() != num_value_heads * key_dim:
98 raise ValueError("Fused KDA dt_bias must contain one value per gate channel.")
99 if not -5.0 <= lower_bound < 0:
100 raise ValueError("Fused KDA lower_bound must lie in [-5, 0).")
101 if not query.is_npu:
102 raise RuntimeError("The fused KDA backend requires Ascend NPU tensors.")
103 if any(
104 tensor.device != query.device
105 for tensor in (key, value, gate, beta, a_log, dt_bias)
106 ):
107 raise ValueError("Fused KDA inputs must reside on the same NPU device.")
108 get_fla_kda_staged_ops()
111def _validate_p2p_inputs(
112 query: torch.Tensor,
113 key: torch.Tensor,
114 value: torch.Tensor,
115 gate: torch.Tensor,
116 beta: torch.Tensor,
117 a_log: torch.Tensor,
118 dt_bias: torch.Tensor,
119 *,
120 chunk_size: int,
121 lower_bound: float,
122 cp_rank: int,
123 cp_size: int,
124) -> None:
125 """Validate the fixed-shape KDA P2P backend before communication starts."""
126 _validate_local_inputs(
127 query,
128 key,
129 value,
130 gate,
131 beta,
132 a_log,
133 dt_bias,
134 chunk_size=chunk_size,
135 lower_bound=lower_bound,
136 )
137 if cp_size <= 0 or not 0 <= cp_rank < cp_size:
138 raise ValueError(f"Invalid KDA P2P rank {cp_rank} for cp_size {cp_size}.")
141def fused_chunk_kda(
142 query: torch.Tensor,
143 key: torch.Tensor,
144 value: torch.Tensor,
145 gate: torch.Tensor,
146 beta: torch.Tensor,
147 *,
148 a_log: torch.Tensor,
149 dt_bias: torch.Tensor,
150 scale: Optional[float] = None,
151 lower_bound: float = -5.0,
152 chunk_size: int = 64,
153 safe_gate: bool = True,
154) -> torch.Tensor:
155 """Run the local Triton-Ascend KDA backend without CP communication.
157 Args:
158 query: Local query tensor [B,T,H,K].
159 key: Matching local key tensor.
160 value: Local value tensor [B,T,HV,V].
161 gate: Per-token gate logits [B,T,HV,K].
162 beta: Per-token beta logits [B,T,HV].
163 a_log: Decay parameter with one value per value head.
164 dt_bias: Learned bias with one value per gate channel.
165 scale: Optional attention scaling factor.
166 lower_bound: Lower bound for the activated gate.
167 chunk_size: Tokens per local recurrence chunk.
168 safe_gate: Use the staged lower-bounded gate implementation.
169 """
170 _validate_local_inputs(
171 query,
172 key,
173 value,
174 gate,
175 beta,
176 a_log,
177 dt_bias,
178 chunk_size=chunk_size,
179 lower_bound=lower_bound,
180 )
181 output, _ = run_fla_chunk_kda(
182 query,
183 key,
184 value,
185 gate,
186 beta,
187 a_log=a_log,
188 dt_bias=dt_bias.reshape(-1),
189 scale=scale,
190 lower_bound=lower_bound,
191 chunk_size=chunk_size,
192 safe_gate=safe_gate,
193 )
194 return output
197class _KDAStateP2PFunction(torch.autograd.Function):
198 """Run shared local KDA and autograd for P2P or gathered state exchange."""
200 @staticmethod
201 def forward( # pylint: disable=arguments-differ,too-many-arguments,too-many-locals
202 ctx: Any,
203 query: torch.Tensor,
204 key: torch.Tensor,
205 value: torch.Tensor,
206 gate_raw: torch.Tensor,
207 beta_raw: torch.Tensor,
208 a_log: torch.Tensor,
209 dt_bias: torch.Tensor,
210 scale: float,
211 lower_bound: float,
212 chunk_size: int,
213 safe_gate: bool,
214 cp_group: Any,
215 prev_rank: int,
216 next_rank: int,
217 cp_rank: int,
218 cp_size: int,
219 state_cp: Any,
220 ) -> torch.Tensor:
221 """Prepare local KDA, propagate its state, and produce token outputs.
223 Args:
224 ctx: Autograd context holding tensors and communication metadata.
225 query: Local query tensor [B,T,H,K].
226 key: Matching local key tensor.
227 value: Local value tensor [B,T,HV,V].
228 gate_raw: Unactivated per-token gate logits.
229 beta_raw: Unactivated per-token beta logits.
230 a_log: Decay parameter with one value per value head.
231 dt_bias: Learned bias with one value per gate channel.
232 scale: Optional attention scaling factor.
233 lower_bound: Lower bound for the activated gate.
234 chunk_size: Tokens per local recurrence chunk.
235 safe_gate: Use the staged lower-bounded gate implementation.
236 cp_group: Precreated chronological CP process group.
237 prev_rank: Global rank of the preceding sequence partition.
238 next_rank: Global rank of the following sequence partition.
239 cp_rank: Chronological index within the CP group.
240 cp_size: Number of state partitions.
241 state_cp: Optional gather executor; None selects P2P.
242 """
243 ops = get_fla_kda_staged_ops()
245 # FLA consumes base-2 log gates: exp(x) = exp2(x / ln(2)).
246 # Use the same conversion when recomputing gates in backward.
247 rcp_ln2 = 1.4426950408889634
249 query, key, value, gate_raw, beta_raw = (
250 tensor.contiguous()
251 for tensor in (query, key, value, gate_raw, beta_raw)
252 )
253 query, query_rstd = ops.l2norm_fwd(query)
254 key, key_rstd = ops.l2norm_fwd(key)
255 beta = ops.fused_beta_sigmoid(beta_raw)
256 gate = ops.kda_gate_chunk_cumsum(
257 g=gate_raw, A_log=a_log, dt_bias=dt_bias, scale=rcp_ln2, chunk_size=chunk_size,
258 lower_bound=lower_bound,
259 )
261 state_shape = (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1])
262 initial_state = None
263 recv_handle = None
264 if state_cp is None and cp_rank > 0:
265 initial_state = torch.empty(state_shape, device=query.device, dtype=torch.float32)
266 recv_handle = dist.irecv(initial_state, src=prev_rank, group=cp_group)
268 w, u, _, kg, attention_qk, attention_kk = ops.chunk_kda_fwd_intra(
269 q=query, k=key, v=value, gk=gate, beta=beta, scale=scale, chunk_size=chunk_size, safe_gate=safe_gate,
270 disable_recompute=True,
271 )
272 state_ext = None
273 transition = None
274 if state_cp is None and cp_rank == 0 and cp_size > 1:
275 # Zero entry uses S directly, and rank 0 has no reverse-state recipient.
276 state_ext, transition = kda_state_summary_forward_from_prepared(
277 kg, w, u, gate, chunk_size=chunk_size, output_transition=False,
278 )
279 elif state_cp is not None or cp_rank < cp_size - 1:
280 state_ext, transition = kda_state_summary_forward_from_prepared(kg, w, u, gate, chunk_size=chunk_size)
281 if recv_handle is not None:
282 recv_handle.wait()
284 send_handle = None
285 send_state = None
286 if state_cp is not None:
287 initial_state = state_cp.forward(state_ext, transition)
288 elif cp_rank < cp_size - 1:
289 final_state = apply_kda_state_summary(state_ext, transition, initial_state)
290 send_state = final_state.contiguous()
291 send_handle = dist.isend(send_state, dst=next_rank, group=cp_group)
293 states, value_new, _ = ops.chunk_gated_delta_rule_fwd_h(
294 k=kg, w=w, u=u, gk=gate, initial_state=initial_state, output_final_state=False, chunk_size=chunk_size,
295 )
296 output = ops.chunk_gla_fwd_o_gk(
297 q=query, v=value_new, g=gate, A=attention_qk, h=states, scale=scale, chunk_size=chunk_size,
298 )
299 if send_handle is not None:
300 send_handle.wait()
302 saved_initial_state = initial_state
303 if saved_initial_state is None:
304 saved_initial_state = query.new_empty(0, dtype=torch.float32)
305 ctx.save_for_backward(
306 query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias, attention_qk,
307 attention_kk, transition if transition is not None else query.new_empty(0), saved_initial_state,
308 )
309 ctx.scale = scale
310 ctx.lower_bound = lower_bound
311 ctx.chunk_size = chunk_size
312 ctx.safe_gate = safe_gate
313 ctx.cp_group = cp_group
314 ctx.prev_rank = prev_rank
315 ctx.next_rank = next_rank
316 ctx.cp_rank = cp_rank
317 ctx.cp_size = cp_size
318 ctx.state_cp = state_cp
319 return output.to(value.dtype)
321 @staticmethod
322 def backward( # pylint: disable=arguments-differ,too-many-locals
323 ctx: Any,
324 grad_output: torch.Tensor,
325 ) -> tuple[Any, ...]:
326 """Reverse the state wavefront, then finish the fused local backward.
328 Args:
329 ctx: Autograd context holding tensors and communication metadata.
330 grad_output: Upstream gradient for the local token output.
331 """
332 ops = get_fla_kda_staged_ops()
334 # FLA consumes base-2 log gates: exp(x) = exp2(x / ln(2)).
335 # Use the same conversion when recomputing gates in backward.
336 rcp_ln2 = 1.4426950408889634
338 (
339 query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias,
340 attention_qk, attention_kk, transition, saved_initial_state,
341 ) = ctx.saved_tensors
342 initial_state = saved_initial_state if saved_initial_state.numel() else None
343 grad_output = grad_output.contiguous()
345 grad_final_state = None
346 recv_handle = None
347 if ctx.state_cp is None and ctx.cp_rank < ctx.cp_size - 1:
348 grad_final_state = torch.empty(
349 (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]), device=query.device,
350 dtype=torch.float32,
351 )
352 recv_handle = dist.irecv(grad_final_state, src=ctx.next_rank, group=ctx.cp_group)
354 beta = ops.fused_beta_sigmoid(beta_raw)
355 gate = ops.kda_gate_chunk_cumsum(
356 g=gate_raw, A_log=a_log, dt_bias=dt_bias, scale=rcp_ln2, chunk_size=ctx.chunk_size,
357 lower_bound=ctx.lower_bound,
358 )
359 w, u, query_gated, key_gated = ops.recompute_w_u_fwd(
360 q=query, k=key, v=value, beta=beta, A=attention_kk, gk=gate,
361 )
362 states, value_new, _ = ops.chunk_gated_delta_rule_fwd_h(
363 k=key_gated, w=w, u=u, gk=gate, initial_state=initial_state, output_final_state=False,
364 chunk_size=ctx.chunk_size,
365 )
366 grad_attention_qk, grad_value_local = ops.chunk_kda_bwd_dav(
367 q=query, k=key, v=value_new, do=grad_output, A=attention_qk, scale=ctx.scale,
368 chunk_size=ctx.chunk_size,
369 )
370 grad_state_ext = None
371 if ctx.state_cp is not None or ctx.cp_rank > 0:
372 grad_state_ext = kda_state_gradient_summary_from_prepared(
373 query_gated, key_gated, w, gate, grad_output, grad_value_local, ctx.scale,
374 chunk_size=ctx.chunk_size,
375 )
376 if recv_handle is not None:
377 recv_handle.wait()
378 send_handle = None
379 send_state_gradient = None
380 if ctx.state_cp is not None:
381 grad_final_state = ctx.state_cp.backward(grad_state_ext, transition)
382 elif ctx.cp_rank > 0:
383 grad_initial_state = apply_kda_state_gradient_summary(grad_state_ext, transition, grad_final_state)
384 send_state_gradient = grad_initial_state.contiguous()
385 send_handle = dist.isend(send_state_gradient, dst=ctx.prev_rank, group=ctx.cp_group)
387 grad_states, _, grad_value = ops.chunk_gated_delta_rule_bwd_dhu(
388 q=query_gated, k=key_gated, w=w, gk=gate, h0=initial_state, dht=grad_final_state, do=grad_output,
389 dv=grad_value_local, scale=ctx.scale, chunk_size=ctx.chunk_size,
390 )
391 # These recomputed tensors have no consumers after dhu. Dropping local
392 # references lets the allocator reuse their storage on this stream.
393 del query_gated, key_gated, w, u, grad_value_local, _
394 (
395 grad_query, grad_key, grad_value, grad_beta, grad_gate, grad_attention_kk,
396 ) = ops.chunk_kda_bwd_wy_dqkg_fused(
397 q=query, k=key, v=value, v_new=value_new, g=gate, beta=beta, A=attention_kk, h=states, do=grad_output,
398 dh=grad_states, dv=grad_value, scale=ctx.scale, chunk_size=ctx.chunk_size,
399 )
400 # Release recomputed states before intra backward allocates its outputs.
401 del states, grad_states, value_new
402 grad_query, grad_key, grad_beta, grad_gate = ops.chunk_kda_bwd_intra(
403 q=query, k=key, g=gate, beta=beta, dAqk=grad_attention_qk, dAkk=grad_attention_kk, dq=grad_query,
404 dk=grad_key, db=grad_beta, dg=grad_gate, chunk_size=ctx.chunk_size, safe_gate=ctx.safe_gate,
405 )
407 num_query_heads = query.shape[2]
408 num_value_heads = value.shape[2]
409 if num_value_heads > num_query_heads:
410 groups = num_value_heads // num_query_heads
411 grad_query = grad_query.view(*grad_query.shape[:2], num_query_heads, groups, grad_query.shape[-1]).sum(
412 dim=3,
413 )
414 grad_key = grad_key.view(*grad_key.shape[:2], num_query_heads, groups, grad_key.shape[-1]).sum(dim=3)
416 grad_gate = ops.chunk_local_cumsum(grad_gate, chunk_size=ctx.chunk_size, reverse=True)
417 grad_gate, grad_a_log, grad_dt_bias = ops.kda_gate_bwd(
418 g=gate_raw, A_log=a_log, dt_bias=dt_bias, dyg=grad_gate, lower_bound=ctx.lower_bound,
419 )
420 grad_beta = ops.fused_beta_sigmoid_bwd(beta_raw, grad_beta)
421 grad_query = ops.l2norm_bwd(query, query_rstd, grad_query)
422 grad_key = ops.l2norm_bwd(key, key_rstd, grad_key)
423 if send_handle is not None:
424 send_handle.wait()
426 return (
427 grad_query, grad_key, grad_value, grad_gate, grad_beta, grad_a_log, grad_dt_bias, None, None, None,
428 None, None, None, None, None, None, None,
429 )
432@torch.compiler.disable
433def fused_chunk_kda_p2p(
434 query: torch.Tensor,
435 key: torch.Tensor,
436 value: torch.Tensor,
437 gate: torch.Tensor,
438 beta: torch.Tensor,
439 *,
440 a_log: torch.Tensor,
441 dt_bias: torch.Tensor,
442 cp_group: Any,
443 prev_rank: int,
444 next_rank: int,
445 cp_rank: int,
446 cp_size: int,
447 scale: Optional[float] = None,
448 lower_bound: float = -5.0,
449 chunk_size: int = 64,
450 safe_gate: bool = True,
451 state_cp: Any = None,
452) -> torch.Tensor:
453 """Run one sequence-sharded KDA segment with a selected state CP executor.
455 The historical P2P entry name is retained for existing callers; passing
456 state_cp selects AllGather or grouped AllGather/P2P with the same local ops.
458 Args:
459 query: Query input [B,T,H,K].
460 key: Matching key input.
461 value: Value input [B,T,HV,V].
462 gate: Per-token gate logits.
463 beta: Per-token beta logits.
464 a_log: FP32 decay parameter per value head.
465 dt_bias: FP32 gate-channel bias.
466 cp_group: Existing chronological state process group.
467 prev_rank: Global predecessor rank for P2P.
468 next_rank: Global successor rank for P2P.
469 cp_rank: Chronological index within the state group.
470 cp_size: Number of state partitions.
471 scale: Optional attention scale.
472 lower_bound: Lower bound for the activated gate.
473 chunk_size: Tokens per local chunk.
474 safe_gate: Enable the staged safe-gate path.
475 state_cp: Optional stateless AG executor; None selects existing P2P.
477 Returns:
478 Output tokens in the local sequence layout.
479 """
480 _validate_p2p_inputs(
481 query,
482 key,
483 value,
484 gate,
485 beta,
486 a_log,
487 dt_bias,
488 chunk_size=chunk_size,
489 lower_bound=lower_bound,
490 cp_rank=cp_rank,
491 cp_size=cp_size,
492 )
493 effective_scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
494 return _KDAStateP2PFunction.apply(
495 query,
496 key,
497 value,
498 gate,
499 beta,
500 a_log,
501 dt_bias.reshape(-1),
502 effective_scale,
503 float(lower_bound),
504 chunk_size,
505 safe_gate,
506 cp_group,
507 prev_rank,
508 next_rank,
509 cp_rank,
510 cp_size,
511 state_cp,
512 )