Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / kimi_delta_attention.py: 14%
224 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"""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 hyper_parallel.components.functional.linear_attention_packed import PackedSequenceInfo, resolve_packed_sequence
29from .kimi_delta_attention_fla_adapter import get_fla_kda_staged_ops, run_fla_chunk_kda
30from .kimi_delta_attention_state_summary import (
31 apply_kda_state_gradient_summary,
32 apply_kda_state_summary,
33 kda_state_gradient_summary_from_prepared,
34 kda_state_summary_forward_from_prepared,
35)
38def _validate_local_shapes(
39 query: torch.Tensor,
40 key: torch.Tensor,
41 value: torch.Tensor,
42 gate: torch.Tensor,
43 beta: torch.Tensor,
44) -> tuple[int, int, int, int]:
45 """Validate local KDA tensor shapes and return dimensions used by later checks."""
46 if (query.ndim, key.ndim, value.ndim, gate.ndim) != (4, 4, 4, 4):
47 raise ValueError("Fused KDA expects rank-4 query/key/value/gate tensors.")
48 if beta.ndim != 3:
49 raise ValueError("Fused KDA expects a rank-3 beta tensor.")
50 if query.shape != key.shape:
51 raise ValueError("Fused KDA query and key must have identical shapes.")
52 batch, sequence_length, num_query_heads, key_dim = query.shape
53 num_value_heads, value_dim = value.shape[2:]
54 if value.shape[:2] != (batch, sequence_length):
55 raise ValueError("Fused KDA value must match query batch and sequence dimensions.")
56 if gate.shape != (batch, sequence_length, num_value_heads, key_dim):
57 raise ValueError("Fused KDA gate has an incompatible shape.")
58 if beta.shape != (batch, sequence_length, num_value_heads):
59 raise ValueError("Fused KDA beta has an incompatible shape.")
60 if num_value_heads % num_query_heads:
61 raise ValueError("Fused KDA value heads must be divisible by query heads.")
62 return sequence_length, key_dim, value_dim, num_value_heads
65def _validate_local_inputs(
66 query: torch.Tensor,
67 key: torch.Tensor,
68 value: torch.Tensor,
69 gate: torch.Tensor,
70 beta: torch.Tensor,
71 a_log: torch.Tensor,
72 dt_bias: torch.Tensor,
73 *,
74 chunk_size: int,
75 lower_bound: float,
76) -> None:
77 """Validate the fixed-shape local KDA backend before compiling kernels."""
78 _, key_dim, value_dim, num_value_heads = _validate_local_shapes(
79 query, key, value, gate, beta
80 )
81 if not (
82 query.dtype == key.dtype == value.dtype == gate.dtype == beta.dtype
83 == torch.bfloat16
84 ):
85 raise TypeError("Fused KDA requires q/k/v/gate/beta tensors in bfloat16.")
86 if a_log.dtype != torch.float32 or dt_bias.dtype != torch.float32:
87 raise TypeError("Fused KDA requires a_log and dt_bias tensors in float32.")
88 if key_dim != 128 or value_dim != 128 or chunk_size != 64:
89 raise NotImplementedError(
90 "Fused KDA currently requires key_dim=value_dim=128 and chunk_size=64."
91 )
92 if a_log.numel() != num_value_heads:
93 raise ValueError("Fused KDA a_log must contain one value per value head.")
94 if dt_bias.numel() != num_value_heads * key_dim:
95 raise ValueError("Fused KDA dt_bias must contain one value per gate channel.")
96 if not -5.0 <= lower_bound < 0:
97 raise ValueError("Fused KDA lower_bound must lie in [-5, 0).")
98 if not query.is_npu:
99 raise RuntimeError("The fused KDA backend requires Ascend NPU tensors.")
100 if any(
101 tensor.device != query.device
102 for tensor in (key, value, gate, beta, a_log, dt_bias)
103 ):
104 raise ValueError("Fused KDA inputs must reside on the same NPU device.")
105 get_fla_kda_staged_ops()
108def _validate_p2p_inputs(
109 query: torch.Tensor,
110 key: torch.Tensor,
111 value: torch.Tensor,
112 gate: torch.Tensor,
113 beta: torch.Tensor,
114 a_log: torch.Tensor,
115 dt_bias: torch.Tensor,
116 *,
117 chunk_size: int,
118 lower_bound: float,
119 cp_rank: int,
120 cp_size: int,
121) -> None:
122 """Validate the fixed-shape KDA P2P backend before communication starts."""
123 _validate_local_inputs(
124 query,
125 key,
126 value,
127 gate,
128 beta,
129 a_log,
130 dt_bias,
131 chunk_size=chunk_size,
132 lower_bound=lower_bound,
133 )
134 if cp_size <= 0 or not 0 <= cp_rank < cp_size:
135 raise ValueError(f"Invalid KDA P2P rank {cp_rank} for cp_size {cp_size}.")
138def fused_chunk_kda(
139 query: torch.Tensor,
140 key: torch.Tensor,
141 value: torch.Tensor,
142 gate: torch.Tensor,
143 beta: torch.Tensor,
144 *,
145 a_log: torch.Tensor,
146 dt_bias: torch.Tensor,
147 scale: Optional[float] = None,
148 lower_bound: float = -5.0,
149 chunk_size: int = 64,
150 safe_gate: bool = True,
151 cu_seqlens: Optional[torch.Tensor] = None,
152) -> torch.Tensor:
153 """Run the local Triton-Ascend KDA backend without CP communication.
155 Args:
156 query: Local query tensor [B,T,H,K].
157 key: Matching local key tensor.
158 value: Local value tensor [B,T,HV,V].
159 gate: Per-token gate logits [B,T,HV,K].
160 beta: Per-token beta logits [B,T,HV].
161 a_log: Decay parameter with one value per value head.
162 dt_bias: Learned bias with one value per gate channel.
163 scale: Optional attention scaling factor.
164 lower_bound: Lower bound for the activated gate.
165 chunk_size: Tokens per local recurrence chunk.
166 safe_gate: Use the staged lower-bounded gate implementation.
167 """
168 _validate_local_inputs(
169 query,
170 key,
171 value,
172 gate,
173 beta,
174 a_log,
175 dt_bias,
176 chunk_size=chunk_size,
177 lower_bound=lower_bound,
178 )
179 output, _ = run_fla_chunk_kda(
180 query,
181 key,
182 value,
183 gate,
184 beta,
185 a_log=a_log,
186 dt_bias=dt_bias.reshape(-1),
187 scale=scale,
188 lower_bound=lower_bound,
189 chunk_size=chunk_size,
190 safe_gate=safe_gate,
191 cu_seqlens=cu_seqlens,
192 )
193 return output
196def _prepare_packed_kda_summary(kg, w, u, gate, sequence_info, state_cp, send_next, chunk_size):
197 """Skip summaries with no consumer when the invocation supports absent tensors."""
198 if state_cp is None and not send_next:
199 return None, None
200 state_shape = (kg.shape[0], u.shape[2], kg.shape[-1], u.shape[-1])
201 transition = None
202 if getattr(state_cp, "needs_state_summary", True):
203 tail = sequence_info.tail
204 state, transition = kda_state_summary_forward_from_prepared(
205 kg[:, tail], w[:, tail], u[:, tail], gate[:, tail], chunk_size=chunk_size,
206 output_transition=(sequence_info.transfer_active
207 and getattr(state_cp, "needs_transition_summary", True)),
208 )
209 elif getattr(state_cp, "allows_absent_summaries", False):
210 return None, None
211 else:
212 state = kg.new_zeros(state_shape, dtype=torch.float32)
213 if state_cp is not None and transition is None:
214 # Shape placeholders preserve existing collective validation; no M kernel runs.
215 transition = kg.new_zeros((*state_shape[:-1], kg.shape[-1]), dtype=torch.float32)
216 return state, transition
219class _KDAStateP2PFunction(torch.autograd.Function):
220 """Run shared local KDA and autograd with the selected state exchange."""
222 @staticmethod
223 def forward( # pylint: disable=arguments-differ,too-many-arguments,too-many-locals
224 ctx: Any,
225 query: torch.Tensor,
226 key: torch.Tensor,
227 value: torch.Tensor,
228 gate_raw: torch.Tensor,
229 beta_raw: torch.Tensor,
230 a_log: torch.Tensor,
231 dt_bias: torch.Tensor,
232 scale: float,
233 lower_bound: float,
234 chunk_size: int,
235 safe_gate: bool,
236 cp_group: Any,
237 prev_rank: int,
238 next_rank: int,
239 cp_rank: int,
240 cp_size: int,
241 state_cp: Any,
242 sequence_info: Optional[PackedSequenceInfo],
243 ) -> torch.Tensor:
244 """Prepare local KDA, propagate its state, and produce token outputs.
246 Args:
247 ctx: Autograd context holding tensors and communication metadata.
248 query: Local query tensor [B,T,H,K].
249 key: Matching local key tensor.
250 value: Local value tensor [B,T,HV,V].
251 gate_raw: Unactivated per-token gate logits.
252 beta_raw: Unactivated per-token beta logits.
253 a_log: Decay parameter with one value per value head.
254 dt_bias: Learned bias with one value per gate channel.
255 scale: Optional attention scaling factor.
256 lower_bound: Lower bound for the activated gate.
257 chunk_size: Tokens per local recurrence chunk.
258 safe_gate: Use the staged lower-bounded gate implementation.
259 cp_group: Precreated chronological CP process group.
260 prev_rank: Global rank of the preceding sequence partition.
261 next_rank: Global rank of the following sequence partition.
262 cp_rank: Chronological index within the CP group.
263 cp_size: Number of state partitions.
264 state_cp: Executor returning an invocation cache; None selects P2P.
265 """
266 ops = get_fla_kda_staged_ops()
267 bind_sequence = getattr(state_cp, "bind_sequence", None)
268 if sequence_info is not None and bind_sequence is not None:
269 state_cp = bind_sequence(
270 sequence_info, state_shape=(query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]))
271 plan = None if sequence_info is None else sequence_info.materialize(query.device, chunk_size)
272 varlen = {} if plan is None else plan.kwargs()
273 recurrent_varlen = dict(varlen)
274 if plan is not None:
275 recurrent_varlen["chunk_offsets"] = plan.chunk_offsets
277 # FLA consumes base-2 log gates: exp(x) = exp2(x / ln(2)).
278 # Use the same conversion when recomputing gates in backward.
279 rcp_ln2 = 1.4426950408889634
281 query, key, value, gate_raw, beta_raw = (
282 tensor.contiguous()
283 for tensor in (query, key, value, gate_raw, beta_raw)
284 )
285 query, query_rstd = ops.l2norm_fwd(query)
286 key, key_rstd = ops.l2norm_fwd(key)
287 beta = ops.fused_beta_sigmoid(beta_raw)
288 gate = ops.kda_gate_chunk_cumsum(
289 g=gate_raw, A_log=a_log, dt_bias=dt_bias, scale=rcp_ln2, chunk_size=chunk_size,
290 lower_bound=lower_bound,
291 **varlen,
292 )
294 state_shape = (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1])
295 receive_previous = cp_rank > 0 and (sequence_info is None or sequence_info.continues_previous)
296 send_next = cp_rank < cp_size - 1 and (sequence_info is None or sequence_info.continues_next)
297 initial_state = None
298 recv_handle = None
299 if state_cp is None and receive_previous:
300 initial_state = torch.empty(state_shape, device=query.device, dtype=torch.float32)
301 recv_handle = dist.irecv(initial_state, src=prev_rank, group=cp_group)
303 w, u, _, kg, attention_qk, attention_kk = ops.chunk_kda_fwd_intra(
304 q=query, k=key, v=value, gk=gate, beta=beta, scale=scale, chunk_size=chunk_size, safe_gate=safe_gate,
305 disable_recompute=True,
306 **varlen,
307 )
308 # The gated query is unused here and is recomputed in backward.
309 del _
310 state_ext = None
311 transition = None
312 if sequence_info is not None:
313 state_ext, transition = _prepare_packed_kda_summary(
314 kg, w, u, gate, sequence_info, state_cp, send_next, chunk_size,
315 )
316 else:
317 if state_cp is None and cp_rank == 0 and cp_size > 1:
318 # Zero entry uses S directly, and rank 0 has no reverse-state recipient.
319 state_ext, transition = kda_state_summary_forward_from_prepared(
320 kg, w, u, gate, chunk_size=chunk_size, output_transition=False,
321 )
322 elif state_cp is not None or cp_rank < cp_size - 1:
323 state_ext, transition = kda_state_summary_forward_from_prepared(kg, w, u, gate, chunk_size=chunk_size)
324 # A reset makes the outgoing summary independent of the incoming state.
325 # Wait before apply only when it actually reads that state; local h still waits below.
326 if recv_handle is not None and (sequence_info is None or sequence_info.transfer_active):
327 recv_handle.wait()
328 recv_handle = None
330 send_handle = None
331 send_state = None
332 state_cache = (transition,) if transition is not None else ()
333 if state_cp is not None:
334 initial_state, state_cache = state_cp.forward(state_ext, transition)
335 transition = None
336 elif send_next:
337 summary_initial = initial_state if transition is not None else None
338 final_state = apply_kda_state_summary(state_ext, transition, summary_initial)
339 send_state = final_state.contiguous()
340 send_handle = dist.isend(send_state, dst=next_rank, group=cp_group)
342 if recv_handle is not None:
343 recv_handle.wait()
344 kernel_initial_state = initial_state if sequence_info is None else sequence_info.first_state(initial_state)
345 states, value_new, _ = ops.chunk_gated_delta_rule_fwd_h(
346 k=kg, w=w, u=u, gk=gate, initial_state=kernel_initial_state,
347 output_final_state=False, chunk_size=chunk_size,
348 **recurrent_varlen,
349 )
350 # Local h was their last forward consumer on this stream.
351 del w, u, kg, _
352 output = ops.chunk_gla_fwd_o_gk(
353 q=query, v=value_new, g=gate, A=attention_qk, h=states, scale=scale, chunk_size=chunk_size,
354 **varlen,
355 )
356 if send_handle is not None:
357 send_handle.wait()
359 saved_initial_state = initial_state
360 if saved_initial_state is None:
361 saved_initial_state = query.new_empty(0, dtype=torch.float32)
362 packed_tensors = () if plan is None else (
363 plan.cu_seqlens, plan.chunk_indices, plan.chunk_offsets)
364 ctx.save_for_backward(
365 query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias, attention_qk,
366 attention_kk, saved_initial_state, *state_cache, *packed_tensors,
367 )
368 ctx.scale = scale
369 ctx.lower_bound = lower_bound
370 ctx.chunk_size = chunk_size
371 ctx.safe_gate = safe_gate
372 ctx.cp_group = cp_group
373 ctx.prev_rank = prev_rank
374 ctx.next_rank = next_rank
375 ctx.cp_rank = cp_rank
376 ctx.cp_size = cp_size
377 ctx.state_cp = state_cp
378 ctx.state_cache_count = len(state_cache)
379 ctx.receive_previous = receive_previous
380 ctx.send_next = send_next
381 ctx.sequence_info = sequence_info
382 return output.to(value.dtype)
384 @staticmethod
385 def backward( # pylint: disable=arguments-differ,too-many-locals
386 ctx: Any,
387 grad_output: torch.Tensor,
388 ) -> tuple[Any, ...]:
389 """Reverse the state wavefront, then finish the fused local backward.
391 Args:
392 ctx: Autograd context holding tensors and communication metadata.
393 grad_output: Upstream gradient for the local token output.
394 """
395 ops = get_fla_kda_staged_ops()
397 # FLA consumes base-2 log gates: exp(x) = exp2(x / ln(2)).
398 # Use the same conversion when recomputing gates in backward.
399 rcp_ln2 = 1.4426950408889634
401 (
402 query, query_rstd, key, key_rstd, value, gate_raw, beta_raw, a_log, dt_bias,
403 attention_qk, attention_kk, saved_initial_state, *saved_tail,
404 ) = ctx.saved_tensors
405 state_cache = tuple(saved_tail[:ctx.state_cache_count])
406 packed_tensors = saved_tail[ctx.state_cache_count:]
407 sequence_info = ctx.sequence_info
408 cu_seqlens = chunk_indices = chunk_offsets = None
409 if sequence_info is not None:
410 cu_seqlens, chunk_indices, chunk_offsets = packed_tensors
411 varlen = {} if sequence_info is None else {"cu_seqlens": cu_seqlens, "chunk_indices": chunk_indices}
412 recurrent_varlen = dict(varlen)
413 if sequence_info is not None:
414 recurrent_varlen["chunk_offsets"] = chunk_offsets
415 initial_state = saved_initial_state if saved_initial_state.numel() else None
416 kernel_initial_state = initial_state if sequence_info is None else sequence_info.first_state(initial_state)
417 grad_output = grad_output.contiguous()
419 grad_final_state = None
420 recv_handle = None
421 if ctx.state_cp is None and ctx.send_next:
422 grad_final_state = torch.empty(
423 (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]), device=query.device,
424 dtype=torch.float32,
425 )
426 recv_handle = dist.irecv(grad_final_state, src=ctx.next_rank, group=ctx.cp_group)
428 beta = ops.fused_beta_sigmoid(beta_raw)
429 gate = ops.kda_gate_chunk_cumsum(
430 g=gate_raw, A_log=a_log, dt_bias=dt_bias, scale=rcp_ln2, chunk_size=ctx.chunk_size,
431 lower_bound=ctx.lower_bound,
432 **varlen,
433 )
434 w, u, query_gated, key_gated = ops.recompute_w_u_fwd(
435 q=query, k=key, v=value, beta=beta, A=attention_kk, gk=gate,
436 **varlen,
437 )
438 states, value_new, _ = ops.chunk_gated_delta_rule_fwd_h(
439 k=key_gated, w=w, u=u, gk=gate, initial_state=kernel_initial_state, output_final_state=False,
440 chunk_size=ctx.chunk_size,
441 **recurrent_varlen,
442 )
443 grad_attention_qk, grad_value_local = ops.chunk_kda_bwd_dav(
444 q=query, k=key, v=value_new, do=grad_output, A=attention_qk, scale=ctx.scale,
445 chunk_size=ctx.chunk_size,
446 **varlen,
447 )
448 grad_state_ext = None
449 if ctx.state_cp is not None or ctx.receive_previous:
450 if sequence_info is not None and not sequence_info.continues_previous:
451 if not getattr(ctx.state_cp, "allows_absent_summaries", False):
452 grad_state_ext = query.new_zeros(
453 (query.shape[0], value.shape[2], query.shape[-1], value.shape[-1]), dtype=torch.float32)
454 else:
455 head = slice(None) if sequence_info is None else sequence_info.head
456 grad_state_ext = kda_state_gradient_summary_from_prepared(
457 query_gated[:, head], key_gated[:, head], w[:, head], gate[:, head],
458 grad_output[:, head], grad_value_local[:, head], ctx.scale, chunk_size=ctx.chunk_size,
459 )
460 if recv_handle is not None and (sequence_info is None or sequence_info.transfer_active):
461 recv_handle.wait()
462 recv_handle = None
463 send_handle = None
464 send_state_gradient = None
465 if ctx.state_cp is not None:
466 grad_final_state = ctx.state_cp.backward(grad_state_ext, state_cache)
467 elif ctx.receive_previous:
468 transition = state_cache[0] if state_cache else None
469 summary_final_gradient = grad_final_state if transition is not None else None
470 grad_initial_state = apply_kda_state_gradient_summary(grad_state_ext, transition, summary_final_gradient)
471 send_state_gradient = grad_initial_state.contiguous()
472 send_handle = dist.isend(send_state_gradient, dst=ctx.prev_rank, group=ctx.cp_group)
474 if recv_handle is not None:
475 recv_handle.wait()
476 kernel_grad_final = grad_final_state if sequence_info is None else sequence_info.last_state(grad_final_state)
477 grad_states, _, grad_value = ops.chunk_gated_delta_rule_bwd_dhu(
478 q=query_gated, k=key_gated, w=w, gk=gate, h0=kernel_initial_state, dht=kernel_grad_final, do=grad_output,
479 dv=grad_value_local, scale=ctx.scale, chunk_size=ctx.chunk_size,
480 **recurrent_varlen,
481 )
482 # These recomputed tensors have no consumers after dhu. Dropping local
483 # references lets the allocator reuse their storage on this stream.
484 del query_gated, key_gated, w, u, grad_value_local, _
485 (
486 grad_query, grad_key, grad_value, grad_beta, grad_gate, grad_attention_kk,
487 ) = ops.chunk_kda_bwd_wy_dqkg_fused(
488 q=query, k=key, v=value, v_new=value_new, g=gate, beta=beta, A=attention_kk, h=states, do=grad_output,
489 dh=grad_states, dv=grad_value, scale=ctx.scale, chunk_size=ctx.chunk_size,
490 **varlen,
491 )
492 # Release recomputed states before intra backward allocates its outputs.
493 del states, grad_states, value_new
494 grad_query, grad_key, grad_beta, grad_gate = ops.chunk_kda_bwd_intra(
495 q=query, k=key, g=gate, beta=beta, dAqk=grad_attention_qk, dAkk=grad_attention_kk, dq=grad_query,
496 dk=grad_key, db=grad_beta, dg=grad_gate, chunk_size=ctx.chunk_size, safe_gate=ctx.safe_gate,
497 **varlen,
498 )
500 num_query_heads = query.shape[2]
501 num_value_heads = value.shape[2]
502 if num_value_heads > num_query_heads:
503 groups = num_value_heads // num_query_heads
504 grad_query = grad_query.view(*grad_query.shape[:2], num_query_heads, groups, grad_query.shape[-1]).sum(
505 dim=3,
506 )
507 grad_key = grad_key.view(*grad_key.shape[:2], num_query_heads, groups, grad_key.shape[-1]).sum(dim=3)
509 grad_gate = ops.chunk_local_cumsum(grad_gate, chunk_size=ctx.chunk_size, reverse=True, **varlen)
510 grad_gate, grad_a_log, grad_dt_bias = ops.kda_gate_bwd(
511 g=gate_raw, A_log=a_log, dt_bias=dt_bias, dyg=grad_gate, lower_bound=ctx.lower_bound,
512 )
513 grad_beta = ops.fused_beta_sigmoid_bwd(beta_raw, grad_beta)
514 grad_query = ops.l2norm_bwd(query, query_rstd, grad_query)
515 grad_key = ops.l2norm_bwd(key, key_rstd, grad_key)
516 if send_handle is not None:
517 send_handle.wait()
519 return (
520 grad_query, grad_key, grad_value, grad_gate, grad_beta, grad_a_log, grad_dt_bias, None, None, None,
521 None, None, None, None, None, None, None, None,
522 )
525@torch.compiler.disable
526def fused_chunk_kda_p2p(
527 query: torch.Tensor,
528 key: torch.Tensor,
529 value: torch.Tensor,
530 gate: torch.Tensor,
531 beta: torch.Tensor,
532 *,
533 a_log: torch.Tensor,
534 dt_bias: torch.Tensor,
535 cp_group: Any,
536 prev_rank: int,
537 next_rank: int,
538 cp_rank: int,
539 cp_size: int,
540 scale: Optional[float] = None,
541 lower_bound: float = -5.0,
542 chunk_size: int = 64,
543 safe_gate: bool = True,
544 state_cp: Any = None,
545 sequence_info: Optional[PackedSequenceInfo] = None,
546) -> torch.Tensor:
547 """Run one sequence-sharded KDA segment with a selected state CP executor.
549 The historical P2P entry name is retained for existing callers; passing
550 state_cp selects gathered or recursive state exchange with the same local ops.
552 Args:
553 query: Query input [B,T,H,K].
554 key: Matching key input.
555 value: Value input [B,T,HV,V].
556 gate: Per-token gate logits.
557 beta: Per-token beta logits.
558 a_log: FP32 decay parameter per value head.
559 dt_bias: FP32 gate-channel bias.
560 cp_group: Existing chronological state process group.
561 prev_rank: Global predecessor rank for P2P.
562 next_rank: Global successor rank for P2P.
563 cp_rank: Chronological index within the state group.
564 cp_size: Number of state partitions.
565 scale: Optional attention scale.
566 lower_bound: Lower bound for the activated gate.
567 chunk_size: Tokens per local chunk.
568 safe_gate: Enable the staged safe-gate path.
569 state_cp: State executor returning an invocation cache; None selects existing P2P.
571 Returns:
572 Output tokens in the local sequence layout.
573 """
574 _validate_p2p_inputs(
575 query,
576 key,
577 value,
578 gate,
579 beta,
580 a_log,
581 dt_bias,
582 chunk_size=chunk_size,
583 lower_bound=lower_bound,
584 cp_rank=cp_rank,
585 cp_size=cp_size,
586 )
587 if sequence_info is not None:
588 if query.shape[0] != 1:
589 raise ValueError("Packed state CP requires batch dimension one.")
590 resolve_packed_sequence({"packed_sequence_info": sequence_info}, query.shape[1],
591 cp_rank=cp_rank, cp_size=cp_size)
592 effective_scale = query.shape[-1] ** -0.5 if scale is None else float(scale)
593 return _KDAStateP2PFunction.apply(
594 query,
595 key,
596 value,
597 gate,
598 beta,
599 a_log,
600 dt_bias.reshape(-1),
601 effective_scale,
602 float(lower_bound),
603 chunk_size,
604 safe_gate,
605 cp_group,
606 prev_rank,
607 next_rank,
608 cp_rank,
609 cp_size,
610 state_cp,
611 sequence_info,
612 )