Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / kimi_delta_attention_ascendc.py: 20%
273 statements
« 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"""Optional native KDA stages for dense/packed Ulysses and summary state exchange."""
17from __future__ import annotations
19import importlib
20import inspect
21from bisect import bisect_left
22from collections.abc import Sequence
23from functools import lru_cache
24from typing import Any, NamedTuple
26import torch
27import torch.distributed as dist
28from torch.autograd.function import once_differentiable
30from hyper_parallel.components.functional.kimi_delta_attention_state_summary import apply_kda_state_summary
31from hyper_parallel.distributed.context_parallel.kimi_delta_attention_rd import KDAStateRecursiveDoubling
34@lru_cache(maxsize=1)
35def _load_ops():
36 """Keep the optional native runtime out of ordinary framework imports."""
37 try:
38 ops = importlib.import_module("fla_npu.ops.ascendc")
39 except ImportError as error:
40 raise RuntimeError("backend='ascendc' requires a compatible fla-npu installation.") from error
41 required = {
42 "chunk_kda_fwd_prepare": ("backward_mode", "use_qk_l2norm_in_kernel", "use_beta_sigmoid_in_kernel"),
43 "chunk_fwd_h": ("output_final_state", "cu_seqlens", "chunk_indices"),
44 "chunk_kda_fwd_finalize": ("output_layout", "cu_seqlens", "chunk_indices"),
45 "chunk_kda_bwd_prepare": ("scale", "layout", "cu_seqlens", "chunk_indices"),
46 "chunk_gated_delta_rule_bwd_dhu": ("gK", "dht", "cu_seqlens", "chunk_indices"),
47 "chunk_kda_bwd_finalize": ("q_rstd", "k_rstd", "use_beta_sigmoid_in_kernel", "output_layout"),
48 }
49 for name, parameters in required.items():
50 operation = getattr(ops, name, None)
51 try:
52 signature = inspect.signature(operation)
53 except (TypeError, ValueError) as error:
54 raise RuntimeError(f"fla-npu lacks a compatible public {name} API.") from error
55 missing = set(parameters).difference(signature.parameters)
56 if missing:
57 raise RuntimeError(
58 f"fla-npu {name} lacks required parameters {sorted(missing)}; "
59 "install a fla-npu version providing these public KDA parameters."
60 )
61 return ops
64class _KDAPrepared(NamedTuple):
65 """Public Prepare save-mode slots, in the native ABI's fixed order."""
67 gk: torch.Tensor # FP32 chunk-local cumulative log2 decay.
68 aqk: torch.Tensor # Intra-chunk QK coefficients.
69 akk: torch.Tensor # Triangular KK coefficients for backward.
70 w: torch.Tensor # WY state-feedback coefficients.
71 u: torch.Tensor | None # Forward-only residual input to FwdH.
72 qg: torch.Tensor # Gated queries consumed by backward.
73 kg: torch.Tensor # Gated keys consumed by FwdH and backward.
74 qg_scaled: torch.Tensor | None # Forward-only queries including attention scale.
75 q_hat: torch.Tensor # Normalized head-major queries.
76 k_hat: torch.Tensor # Normalized head-major keys.
77 q_rstd: torch.Tensor # FP32 Q normalization auxiliary.
78 k_rstd: torch.Tensor # FP32 K normalization auxiliary.
79 beta_eff: torch.Tensor # FP32 head-major sigmoid(beta) cache.
82def _prepare_kda(inputs, scale, lower_bound, metadata):
83 """Prepare once for native state summaries and the local state recurrence."""
84 ops = _load_ops()
85 query, key, value, gate, beta, a_log, dt_bias = inputs
86 result = ops.chunk_kda_fwd_prepare(
87 query, key, value, gate.transpose(1, 2).contiguous(), beta.transpose(1, 2).contiguous(),
88 scale, layout="BNSD", chunk_size=64,
89 safe_gate=True, lower_bound=lower_bound, use_gate_in_kernel=True,
90 A_log=a_log, dt_bias=dt_bias.flatten(), epsilon=1e-6,
91 use_qk_l2norm_in_kernel=True, use_beta_sigmoid_in_kernel=True,
92 use_exp2=True, backward_mode="save", **metadata,
93 )
94 if len(result) != 13 or any(not isinstance(tensor, torch.Tensor) for tensor in result):
95 raise RuntimeError("fla-npu Prepare must return all 13 tensors for native-norm save mode.")
96 return _KDAPrepared(*result)
99def _native_backward_view(tensor: torch.Tensor, packed: bool) -> torch.Tensor:
100 """Expose packed NTD views without copying the B=1 Prepare/FwdH storage."""
101 return tensor.squeeze(0) if packed else tensor
104def _packed_offsets(cu_seqlens, tokens):
105 """Snapshot host offsets before deriving native chunk metadata."""
106 if isinstance(cu_seqlens, torch.Tensor):
107 if (cu_seqlens.device.type != "cpu" or cu_seqlens.ndim != 1
108 or cu_seqlens.dtype not in (torch.int32, torch.int64)):
109 raise ValueError("Ascend C cu_seqlens must be a one-dimensional CPU integer tensor or sequence.")
110 cu_seqlens = cu_seqlens.tolist()
111 offsets = tuple(cu_seqlens)
112 if (len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != tokens
113 or any(not isinstance(x, int) or isinstance(x, bool) for x in offsets)
114 or any(begin >= end for begin, end in zip(offsets, offsets[1:]))):
115 raise ValueError("Ascend C cu_seqlens must strictly increase from zero to the packed token count.")
116 return offsets
119def _packed_metadata(cu_seqlens, tokens):
120 """Snapshot validated host offsets and derive the canonical chunk order."""
121 if cu_seqlens is None:
122 return {}
123 offsets = _packed_offsets(cu_seqlens, tokens)
124 indices = tuple(value for seq, (begin, end) in enumerate(zip(offsets, offsets[1:]))
125 for chunk in range((end - begin + 63) // 64) for value in (seq, chunk))
126 return {"cu_seqlens": offsets, "chunk_indices": indices}
129def ascendc_chunk_kda(
130 query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
131 gate: torch.Tensor, beta: torch.Tensor, *, a_log: torch.Tensor,
132 dt_bias: torch.Tensor, scale: float | None = None, lower_bound: float = -5.0,
133 chunk_size: int = 64, cu_seqlens: Sequence[int] | torch.Tensor | None = None,
134) -> torch.Tensor:
135 """Run dense or packed KDA with native attention and paired raw-input gradients.
137 Args:
138 query: Raw BF16 queries [B,T,H,128] or packed [T,H,128].
139 key: Raw BF16 keys with the same shape.
140 value: BF16 values with the same shape.
141 gate: BF16 or FP32 gate logits with the same shape.
142 beta: BF16 or FP32 beta logits [B,T,H] or packed [T,H].
143 a_log: FP32 decay parameters with H elements.
144 dt_bias: FP32 gate bias with H*128 elements.
145 scale: Positive attention scale, default 1/sqrt(128).
146 lower_bound: Safe gate lower bound in [-5,0).
147 chunk_size: Must be 64 for this native kernel family.
148 cu_seqlens: Strictly increasing host offsets for packed inputs, including
149 zero and T. Packed rank-four inputs require B=1. Documents need not
150 be multiples of 64; zero-length documents are rejected.
152 Returns:
153 BF16 output with the query shape. External recurrent states and double
154 backward are not exposed. Native Q/K norm backward uses save mode;
155 native capability checks select the supported device implementation.
156 """
157 if query.ndim == 3:
158 if cu_seqlens is None:
159 raise ValueError("Packed TND inputs require cu_seqlens.")
160 return ascendc_chunk_kda(
161 *(tensor.unsqueeze(0) for tensor in (query, key, value, gate, beta)),
162 a_log=a_log, dt_bias=dt_bias, scale=scale, lower_bound=lower_bound,
163 chunk_size=chunk_size, cu_seqlens=cu_seqlens).squeeze(0)
164 if query.ndim != 4:
165 raise ValueError("KDA inputs must have shape [B,T,H,K] or packed [T,H,K].")
166 if chunk_size != 64:
167 raise ValueError("Ascend C KDA uses chunk_size=64.")
168 metadata = _packed_metadata(cu_seqlens, query.shape[1])
169 if metadata and query.shape[0] != 1:
170 raise ValueError("Packed Ascend C KDA requires B=1.")
171 scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
172 tokens = tuple(tensor.transpose(1, 2).contiguous() for tensor in (query, key, value))
173 # Gate/beta stay token-major for backward. Native operators validate their
174 # dtype, shape and device; only adapter-owned layout/metadata checks live here.
175 # Clone sliced parameters to give native tiling this shard's storage descriptor.
176 return _AscendCKDAStages.apply(
177 *tokens, gate.contiguous(), beta.contiguous(), a_log.reshape(-1).clone(),
178 dt_bias.reshape(-1, query.shape[-1]).clone(), scale, lower_bound, metadata, None)
181@lru_cache(maxsize=1)
182def _load_state_ops():
183 """Check every native stage before entering the first state collective."""
184 ops = _load_ops()
185 required = (
186 "pre_process_fwd_kernel_merged", "pre_process_bwd_kernel_merged",
187 "merge_fwd_bwd_kernel",
188 )
189 missing = [name for name in required if not callable(getattr(ops, name, None))]
190 if missing:
191 raise RuntimeError(f"Native KDA state CP requires these fla-npu APIs: {', '.join(missing)}.")
192 return ops
195def _native_ag_intervals(info, tokens, rank, size):
196 """Return tail/head windows and document-bounded exclusive scan lengths."""
197 if info is None:
198 return (0, tokens), (0, tokens), rank, size - rank - 1
199 first_rank = (info.start - info.preceding_sample_tokens) // tokens
200 last_rank = (info.global_offsets[bisect_left(info.global_offsets, info.end)] - 1) // tokens
201 return ((info.tail.start, info.tail.stop), (info.head.start, info.head.stop),
202 rank - first_rank if info.continues_previous else 0,
203 last_rank - rank if info.continues_next else 0)
206class _NativeStateExchange:
207 """Bind one invocation's sequence plan to precreated communication groups.
209 Only metadata survives forward. Receive buffers and send handles are owned
210 by the active direction and drained before that invocation returns.
211 """
213 def __init__(self, mesh: Any, executor: Any, info: Any, local_length: int, heads: int) -> None:
214 """Copy communication metadata and retain only this call's sequence plan."""
215 self.executor, self.info, self.heads = executor, info, heads
216 self.rank = mesh.get_local_rank()
217 self.size = mesh.size()
218 self.ranks = tuple(mesh.rank_list)
219 self.group = mesh.get_group()
220 self.width = getattr(executor, "width", self.size) if executor is not None else 1
221 self.block, self.local = divmod(self.rank, self.width)
222 self.owner = self.block * self.width + self.width - 1
223 self.is_rd = isinstance(executor, KDAStateRecursiveDoubling)
224 self.plans = (tuple(executor.directional_plan(info, local_length, reverse=reverse) for reverse in (False, True))
225 if self.is_rd else ())
227 def peer(self, forward: bool, receiving: bool) -> int | None:
228 """Return a chronological peer, pruning disconnected ordinary P2P edges."""
229 direction = 1 if forward else -1
230 peer = self.owner + (-direction if receiving else direction) * self.width
231 if not 0 <= peer < self.size:
232 return None
233 if self.width == 1 and self.info is not None:
234 needs_previous = forward == receiving
235 connected = self.info.continues_previous if needs_previous else self.info.continues_next
236 if not connected:
237 return None
238 return peer
240 def needs_summary(self, forward: bool) -> bool:
241 """Avoid producing a summary without an exit consumer.
243 Args:
244 forward: Select the forward S/M or backward G/N direction.
245 """
246 if self.is_rd:
247 plan = self.plans[not forward]
248 return plan.shift_peers(plan.local_rank)[0] is not None
249 return self.width > 1 or self.peer(forward, receiving=False) is not None
251 def receive(self, reference: torch.Tensor, forward: bool) -> tuple[torch.Tensor | None, Any]:
252 """Post reception before native prepare, independently of summary production.
254 Args:
255 reference: Head-major input supplying the head count and device.
256 forward: Select the chronological communication direction.
257 """
258 if self.is_rd or self.width == self.size:
259 return None, None
260 incoming = reference.new_zeros((1, reference.shape[1], 128, 128), dtype=torch.float32)
261 peer = self.peer(forward, receiving=True)
262 handle = None
263 if self.rank == self.owner and peer is not None:
264 handle = dist.irecv(incoming, src=self.ranks[peer], group=self.group)
265 return incoming, handle
267 @staticmethod
268 def apply(packed: torch.Tensor, incoming: torch.Tensor | None) -> torch.Tensor:
269 """Apply forward M or native backward N directly, without transposing N."""
270 return apply_kda_state_summary(packed[..., :128], packed[..., 128:], incoming)
272 def exchange(self, ops: Any, packed: torch.Tensor | None,
273 reception: tuple[torch.Tensor | None, Any], forward: bool,
274 count: int) -> tuple[torch.Tensor, Any, torch.Tensor | None]:
275 """Produce the local input state and retain an asynchronous outgoing payload.
277 Args:
278 ops: Native summary-merge operations for AllGather.
279 packed: Local S/M or G/N, absent when no exit consumes it.
280 reception: Previously posted receive buffer and handle.
281 forward: Select the chronological communication direction.
282 count: Number of maps in the relevant AllGather prefix or suffix.
283 """
284 if self.is_rd:
285 state_shape = (1, self.heads, 128, 128)
286 maps = None if packed is None else packed.reshape(1, self.heads, 128, 256)
287 result = self.executor.scan_affine(maps, self.plans[not forward], state_shape, reverse=not forward)
288 return result, None, None
289 if self.width == self.size:
290 maps = self.executor.gather_packed(packed.reshape(1, *packed.shape[-3:]))
291 incoming = maps.new_zeros((maps.shape[1], 128, 128))
292 if count:
293 ops.merge_fwd_bwd_kernel(incoming, maps, count, self.rank, forward=forward)
294 return incoming.unsqueeze(0), None, None
295 incoming, recv_handle = reception
296 if self.width > 1:
297 return self._grouped_exchange(packed, incoming, recv_handle, forward)
298 transfer = self.info is None or self.info.num_sequences == 1
299 if recv_handle is not None and transfer:
300 recv_handle.wait()
301 recv_handle = None
302 send_handle, outgoing = None, None
303 peer = self.peer(forward, receiving=False)
304 if peer is not None:
305 outgoing = self.apply(packed, incoming if transfer else None).contiguous()
306 send_handle = dist.isend(outgoing, dst=self.ranks[peer], group=self.group)
307 # A reset permits sending the independent tail/head before this wait,
308 # but local h/dhu still consumes the received first/last fragment state.
309 if recv_handle is not None:
310 recv_handle.wait()
311 return incoming, send_handle, outgoing
313 def _grouped_exchange(self, packed, incoming, recv_handle, forward):
314 """Gather native maps within a group and relay only its exit state."""
315 if self.info is not None:
316 if not self.info.transfer_active or not self.info.continues_next:
317 packed[..., 128:].zero_()
318 connected = self.info.continues_next if forward else self.info.continues_previous
319 if not connected:
320 packed[..., :128].zero_()
321 maps = packed.new_empty((self.width, *packed.shape[-3:]))
322 dist.all_gather_into_tensor(maps, packed.reshape(1, *packed.shape[-3:]), group=self.executor.intra)
323 send_handle, outgoing, result = None, None, None
324 indices = range(self.width) if forward else range(self.width - 1, -1, -1)
325 prefix_count = self.local if forward else self.width - self.local - 1
326 if self.rank == self.owner:
327 if recv_handle is not None:
328 recv_handle.wait()
329 peer = self.peer(forward, receiving=False)
330 if peer is not None:
331 outgoing, result = self._scan(maps, incoming, indices, prefix_count)
332 send_handle = dist.isend(outgoing, dst=self.ranks[peer], group=self.group)
333 if self.peer(forward, receiving=True) is not None:
334 dist.broadcast(incoming, src=self.ranks[self.owner], group=self.executor.intra)
335 if result is None:
336 local_indices = range(self.local) if forward else range(self.width - 1, self.local, -1)
337 result, _ = self._scan(maps, incoming, local_indices)
338 return result, send_handle, outgoing
340 @staticmethod
341 def _scan(maps, incoming, indices, prefix_count=None):
342 """Keep native M/N orientation while applying a nonzero group entry."""
343 state = incoming
344 prefix = incoming if prefix_count == 0 else None
345 for step, index in enumerate(indices, start=1):
346 state = _NativeStateExchange.apply(maps[index:index + 1], state)
347 if step == prefix_count:
348 prefix = state
349 return state, prefix
352def _exchange_native_gradient(ctx, ops, inputs, reception):
353 """Own G/N production, reset injection and the stock dHT dtype boundary."""
354 if ctx.exchange is None:
355 return None, None, None
356 gk, qg, kg, weights, grad, dv_local_head = inputs
357 maps = None
358 if ctx.exchange.needs_summary(False):
359 # Only the stock G/N ABI rounds gk; local dhu/finalize retain FP32.
360 maps = ops.pre_process_bwd_kernel_merged(
361 qg, kg, weights, grad, dv_local_head, ctx.scale, 64,
362 gk=gk.to(qg.dtype), cu_seqlens=ctx.head)
363 terminal, send_handle, outgoing = ctx.exchange.exchange(ops, maps, reception, False, ctx.suffix)
364 if ctx.exchange.info is not None:
365 terminal = ctx.exchange.info.last_state(terminal)
366 # Communication and summary apply retain FP32; only stock dhu requires BF16.
367 terminal = None if terminal is None else terminal.to(qg.dtype)
368 return terminal, send_handle, outgoing
371class _AscendCKDAStages(torch.autograd.Function):
372 """Share native local stages; state protocols optionally inject H and D."""
374 @staticmethod
375 def forward(ctx: Any, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
376 gate: torch.Tensor, beta: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor,
377 scale: float, lower_bound: float, metadata: dict,
378 exchange: _NativeStateExchange | None) -> torch.Tensor:
379 """Share prepare with S/M and local FwdH, then release gathered maps.
381 Args:
382 ctx: Per-invocation autograd context.
383 query: Raw head-major queries.
384 key: Raw head-major keys.
385 value: Head-major values.
386 gate: Sequence-major gate logits.
387 beta: Sequence-major beta logits.
388 a_log: FP32 decay parameters.
389 dt_bias: FP32 gate bias.
390 scale: Attention scale.
391 lower_bound: Safe-gate lower bound.
392 metadata: Native packed offsets and chunk indices.
393 exchange: Optional state protocol for this invocation.
394 """
395 ops = _load_state_ops() if exchange is not None else _load_ops()
396 inputs = (query, key, value, gate, beta, a_log, dt_bias)
397 reception = None if exchange is None else exchange.receive(query, True)
398 prepared = _prepare_kda(inputs, scale, lower_bound, metadata)
399 initial = send_handle = outgoing = head = suffix = None
400 if exchange is not None:
401 tail, head, prefix, suffix = _native_ag_intervals(
402 exchange.info, query.shape[2], exchange.rank, exchange.size)
403 maps = None
404 if exchange.needs_summary(True):
405 maps = ops.pre_process_fwd_kernel_merged(
406 prepared.kg, prepared.w, prepared.u, gk=prepared.gk, cu_seqlens=tail)
407 initial, send_handle, outgoing = exchange.exchange(ops, maps, reception, True, prefix)
408 del maps, reception
409 if exchange.info is not None:
410 initial = exchange.info.first_state(initial)
411 hidden, v_new, _ = ops.chunk_fwd_h(
412 prepared.kg, prepared.w, prepared.u, gk=prepared.gk, initial_state=initial,
413 output_final_state=False, use_exp2=True, **metadata)
414 output = ops.chunk_kda_fwd_finalize(
415 prepared.qg_scaled, prepared.aqk, v_new, hidden, output_layout="BSND", **metadata)
416 if send_handle is not None:
417 send_handle.wait()
418 del outgoing
419 prepared = prepared._replace(u=None, qg_scaled=None)
420 # The saved hidden tensor contains entry states. Summary protocols
421 # produce the outgoing gradient themselves, so native dh0 is unused.
422 ctx.save_for_backward(value, gate, a_log, dt_bias, *prepared, hidden, v_new)
423 ctx.dtypes = tuple(tensor.dtype for tensor in inputs)
424 ctx.scale, ctx.lower_bound, ctx.native_metadata = scale, lower_bound, metadata
425 ctx.exchange, ctx.head, ctx.suffix = exchange, head, suffix
426 return output
428 @staticmethod
429 @once_differentiable
430 def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
431 """Collect G/N, inject nonzero dHT and finish raw-input gradients.
433 Args:
434 ctx: Forward-owned activations and communication metadata.
435 grad_output: Sequence-major output gradient.
436 """
437 ops = _load_state_ops() if ctx.exchange is not None else _load_ops()
438 value, gate, a_log, dt_bias, *saved = ctx.saved_tensors
439 (gk, aqk, akk, weights, _, qg, kg, _, query, key,
440 query_rstd, key_rstd, beta, hidden, v_new) = saved
441 packed = bool(ctx.native_metadata)
443 reception = None if ctx.exchange is None else ctx.exchange.receive(query, False)
444 grad = grad_output.transpose(1, 2).contiguous()
445 d_aqk, dv_local, dq_raw = ops.chunk_kda_bwd_prepare(
446 *(_native_backward_view(tensor, packed) for tensor in (aqk, v_new, grad, hidden)), scale=ctx.scale,
447 layout="NTD" if packed else "BNSD", chunk_size=64, **ctx.native_metadata)
448 dv_local_head = dv_local.unsqueeze(0) if packed else dv_local
449 terminal, send_handle, outgoing = _exchange_native_gradient(
450 ctx, ops, (gk, qg, kg, weights, grad, dv_local_head), reception)
451 del reception
452 hidden_grad, _, dv_scan = ops.chunk_gated_delta_rule_bwd_dhu(
453 qg, kg, weights, grad, dv_local_head, ctx.scale, 64,
454 gK=gk, h0=None, dht=terminal, use_exp2=True, **ctx.native_metadata)
455 dq, dk, dv, dbeta, dgate, da_log, ddt_bias = ops.chunk_kda_bwd_finalize(
456 *(_native_backward_view(tensor, packed)
457 for tensor in (query, key, value, gk, gate.transpose(1, 2).contiguous(), beta)),
458 a_log, dt_bias,
459 *(_native_backward_view(tensor, packed) for tensor in (akk, v_new, hidden, hidden_grad, dv_scan)),
460 d_aqk, dq_raw,
461 scale=ctx.scale, lower_bound=ctx.lower_bound,
462 q_rstd=_native_backward_view(query_rstd, packed),
463 k_rstd=_native_backward_view(key_rstd, packed),
464 use_beta_sigmoid_in_kernel=False,
465 output_layout="NTD" if packed else "BNSD", **ctx.native_metadata)
466 if packed:
467 dq, dk, dv, dbeta, dgate = (tensor.unsqueeze(0) for tensor in (dq, dk, dv, dbeta, dgate))
468 # Prepare already saves FP32 sigmoid(beta); reuse it without device-specific
469 # native epilogue selection, recomputing sigmoid, or retaining raw beta.
470 dbeta = torch.ops.aten.sigmoid_backward.default(dbeta, beta)
471 dbeta, dgate = dbeta.transpose(1, 2), dgate.transpose(1, 2)
472 if send_handle is not None:
473 send_handle.wait()
474 del outgoing
475 gradients = (dq, dk, dv, dgate, dbeta, da_log, ddt_bias.reshape_as(dt_bias))
476 return (*(grad.to(dtype) for grad, dtype in zip(gradients, ctx.dtypes)), None, None, None, None)
479def ascendc_state_kda(
480 query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
481 gate: torch.Tensor, beta: torch.Tensor, *, a_log: torch.Tensor,
482 dt_bias: torch.Tensor, state_cp: Any, cp_mesh: Any, sequence_info: Any = None,
483 scale: float | None = None, lower_bound: float = -5.0,
484) -> torch.Tensor:
485 """Run native KDA with P2P, AllGather, grouped AllGather/P2P, or RD.
487 Args:
488 query: Raw BF16 queries [1,T,H,128].
489 key: Matching raw keys.
490 value: BF16 values with the query shape.
491 gate: Raw gate logits with the query shape.
492 beta: Raw beta logits [1,T,H].
493 a_log: Per-head FP32 decay parameters.
494 dt_bias: FP32 gate bias [H,128].
495 state_cp: Existing gather/RD executor, or None for ordinary P2P.
496 cp_mesh: Ordered state mesh after any Ulysses redistribution.
497 sequence_info: Packed interval metadata after any Ulysses redistribution.
498 scale: Attention scale, default 1/sqrt(128).
499 lower_bound: Safe gate lower bound.
501 Returns:
502 Output with the query shape. Uses the unmodified native stage APIs:
503 backward summary gk and local dhu dHT are converted to model dtype.
504 Prepared gates, communicated states and summary apply remain FP32.
505 These ABI conversions introduce rounding relative to FP32-input kernels.
506 """
507 if query.device.type != "npu" or query.dtype != torch.bfloat16:
508 raise ValueError("Native KDA state CP requires BF16 NPU inputs.")
509 if query.ndim != 4 or query.shape[0] != 1 or query.shape[-1] != 128:
510 raise ValueError("Native KDA state CP requires [1,T,H,128] inputs.")
511 if any(tensor.shape != query.shape for tensor in (key, value, gate)) or beta.shape != query.shape[:-1]:
512 raise ValueError("Native KDA state CP requires matching Q/K/V/gate heads and beta [1,T,H].")
513 metadata = _packed_metadata(None if sequence_info is None else sequence_info.local_offsets, query.shape[1])
514 scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
515 tokens = tuple(tensor.transpose(1, 2).contiguous() for tensor in (query, key, value))
516 exchange = _NativeStateExchange(cp_mesh, state_cp, sequence_info, query.shape[1], query.shape[2])
517 return _AscendCKDAStages.apply(
518 *tokens, gate.contiguous(), beta.contiguous(), a_log.reshape(-1).clone(),
519 dt_bias.reshape(-1, query.shape[-1]).clone(), scale, lower_bound, metadata, exchange)