Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / modules / kimi_delta_attention.py: 82%
372 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"""Pure PyTorch Kimi Delta Attention operators and training layer.
17The implementations in this module intentionally favor readable equations over
18performance. They are the numerical oracle shared by the standalone layer and
19context-parallel wrappers.
20"""
21# pylint: disable=invalid-name
22# This module mirrors the existing Torch model building blocks in this package.
23# pylint: disable=forbidden-backend-import
24from __future__ import annotations
26__all__ = [
27 "KimiDeltaAttention",
28 "KimiRMSNormGated",
29 "chunk_kda",
30 "is_triton_kda_available",
31 "torch_apply_kda_state_summary",
32 "torch_chunk_kda",
33 "torch_compose_kda_state_summaries",
34 "torch_kda_gate",
35 "torch_kda_state_summary",
36 "torch_recurrent_kda",
37]
39from collections.abc import Sequence
40from typing import Any, Optional
42import torch
43from torch import nn
44from torch.nn import functional as F
46from hyper_parallel.components.functional.kimi_delta_attention_ascendc import ascendc_chunk_kda
47from hyper_parallel.components.functional.linear_attention_packed import (
48 packed_causal_conv1d, packed_reference, resolve_packed_sequence,
49)
51_KDA_BACKENDS = frozenset({"eager", "triton", "ascendc"})
52_TRITON_KDA_HEAD_DIM = 128
53_TRITON_KDA_CHUNK_SIZE = 64
56def _has_triton_kda_shapes(
57 query: torch.Tensor,
58 value: torch.Tensor,
59 gate: torch.Tensor,
60 beta: torch.Tensor,
61 a_log: torch.Tensor,
62 dt_bias: torch.Tensor,
63 chunk_size: int,
64) -> bool:
65 """Check the fixed dense shapes accepted by the Triton KDA path."""
66 batch_size, sequence_length, num_query_heads, key_dim = query.shape
67 num_value_heads, value_dim = value.shape[2:]
68 return all((
69 value.shape[:2] == (batch_size, sequence_length),
70 gate.shape == (batch_size, sequence_length, num_value_heads, key_dim),
71 beta.shape == (batch_size, sequence_length, num_value_heads),
72 num_value_heads % num_query_heads == 0,
73 a_log.numel() == num_value_heads,
74 dt_bias.numel() == num_value_heads * key_dim,
75 key_dim == value_dim == _TRITON_KDA_HEAD_DIM,
76 chunk_size == _TRITON_KDA_CHUNK_SIZE,
77 ))
80def _is_triton_kda_input_supported(
81 query: torch.Tensor,
82 key: Optional[torch.Tensor],
83 value: Optional[torch.Tensor],
84 gate: Optional[torch.Tensor],
85 beta: Optional[torch.Tensor],
86 a_log: Optional[torch.Tensor],
87 dt_bias: Optional[torch.Tensor],
88 *,
89 chunk_size: int,
90 lower_bound: Optional[float],
91) -> bool:
92 """Check the fixed Kimi K3 contract validated by the fused backend."""
93 operands = (key, value, gate, beta, a_log, dt_bias)
94 if any(tensor is None for tensor in operands):
95 return False
96 if lower_bound is None:
97 return False
99 basic_contract = (
100 query.device.type == "npu",
101 all(tensor.dtype == torch.bfloat16 for tensor in (query, key, value, gate, beta)),
102 all(tensor.dtype == torch.float32 for tensor in (a_log, dt_bias)),
103 all(tensor.ndim == 4 for tensor in (query, key, value, gate)),
104 beta.ndim == 3,
105 query.shape == key.shape,
106 )
107 if not all(basic_contract):
108 return False
110 shape_supported = _has_triton_kda_shapes(
111 query, value, gate, beta, a_log, dt_bias, chunk_size,
112 )
113 lower_bound_supported = -5.0 <= lower_bound < 0
114 if not (shape_supported and lower_bound_supported):
115 return False
116 return all(tensor.device == query.device for tensor in operands)
119def is_triton_kda_available(
120 query: Optional[torch.Tensor] = None,
121 key: Optional[torch.Tensor] = None,
122 value: Optional[torch.Tensor] = None,
123 gate: Optional[torch.Tensor] = None,
124 beta: Optional[torch.Tensor] = None,
125 a_log: Optional[torch.Tensor] = None,
126 dt_bias: Optional[torch.Tensor] = None,
127 *,
128 chunk_size: int = _TRITON_KDA_CHUNK_SIZE,
129 lower_bound: Optional[float] = -5.0,
130) -> bool:
131 """Return whether the validated Triton-Ascend KDA backend is available."""
132 if chunk_size != _TRITON_KDA_CHUNK_SIZE:
133 return False
134 if query is not None and not _is_triton_kda_input_supported(
135 query,
136 key,
137 value,
138 gate,
139 beta,
140 a_log,
141 dt_bias,
142 chunk_size=chunk_size,
143 lower_bound=lower_bound,
144 ):
145 return False
146 from hyper_parallel.components.functional.kimi_delta_attention_fla_adapter import ( # pylint: disable=import-outside-toplevel
147 is_fla_triton_kda_available,
148 )
149 return is_fla_triton_kda_available()
152def _validate_kda_inputs(
153 query: torch.Tensor,
154 key: torch.Tensor,
155 value: torch.Tensor,
156 gate: torch.Tensor,
157 beta: torch.Tensor,
158) -> tuple[int, int, int, int, int, int]:
159 """Validate dense token-first KDA inputs and return their dimensions."""
160 if (query.dim(), key.dim(), value.dim(), gate.dim()) != (4, 4, 4, 4):
161 raise ValueError("query, key, value, and gate must be rank-4 tensors.")
162 if beta.dim() != 3:
163 raise ValueError("beta must be a rank-3 tensor.")
164 if query.shape != key.shape:
165 raise ValueError("query and key must have identical shapes.")
167 batch_size, sequence_length, num_k_heads, k_head_dim = query.shape
168 value_batch, value_sequence, num_v_heads, v_head_dim = value.shape
169 if (value_batch, value_sequence) != (batch_size, sequence_length):
170 raise ValueError("value batch and sequence dimensions must match query.")
171 if gate.shape != (batch_size, sequence_length, num_v_heads, k_head_dim):
172 raise ValueError(
173 "gate must have shape [B, S, num_v_heads, k_head_dim], got "
174 f"{tuple(gate.shape)}."
175 )
176 if beta.shape != (batch_size, sequence_length, num_v_heads):
177 raise ValueError(
178 "beta must have shape [B, S, num_v_heads], got "
179 f"{tuple(beta.shape)}."
180 )
181 if num_v_heads % num_k_heads != 0:
182 raise ValueError("num_v_heads must be divisible by num_k_heads.")
183 return (
184 batch_size,
185 sequence_length,
186 num_k_heads,
187 num_v_heads,
188 k_head_dim,
189 v_head_dim,
190 )
193def _reshape_gate_bias(
194 parameter: torch.Tensor,
195 num_v_heads: int,
196 k_head_dim: int,
197 name: str,
198) -> torch.Tensor:
199 """Return a gate parameter as ``[num_v_heads, k_head_dim]``."""
200 if parameter.numel() != num_v_heads * k_head_dim:
201 raise ValueError(
202 f"{name} must contain {num_v_heads * k_head_dim} elements, "
203 f"got {parameter.numel()}."
204 )
205 return parameter.reshape(num_v_heads, k_head_dim)
208def torch_kda_gate(
209 gate_logits: torch.Tensor,
210 a_log: torch.Tensor,
211 dt_bias: torch.Tensor,
212 *,
213 lower_bound: Optional[float] = None,
214) -> torch.Tensor:
215 """Activate KDA log-decay logits.
217 ``lower_bound`` selects the Kimi K3 parameterization
218 ``g_min * sigmoid(exp(a_log) * (logits + bias))``. With no lower bound,
219 this function implements the earlier Kimi Linear negative-Softplus gate.
220 """
221 if gate_logits.dim() != 4:
222 raise ValueError("gate_logits must have shape [B, S, H, K].")
223 num_v_heads, k_head_dim = gate_logits.shape[-2:]
224 if a_log.numel() != num_v_heads:
225 raise ValueError(
226 f"a_log must contain {num_v_heads} elements, got {a_log.numel()}."
227 )
228 bias = _reshape_gate_bias(dt_bias, num_v_heads, k_head_dim, "dt_bias")
229 scale = a_log.float().exp().reshape(1, 1, num_v_heads, 1)
230 preactivation = gate_logits.float() + bias.float().reshape(
231 1, 1, num_v_heads, k_head_dim
232 )
233 if lower_bound is None:
234 return -scale * F.softplus(preactivation) # pylint: disable=not-callable
235 if lower_bound >= 0:
236 raise ValueError(f"lower_bound must be negative, got {lower_bound}.")
237 return float(lower_bound) * torch.sigmoid(scale * preactivation)
240def _prepare_kda_inputs(
241 query: torch.Tensor,
242 key: torch.Tensor,
243 value: torch.Tensor,
244 gate: torch.Tensor,
245 beta: torch.Tensor,
246 *,
247 a_log: Optional[torch.Tensor],
248 dt_bias: Optional[torch.Tensor],
249 scale: Optional[float],
250 use_qk_l2norm_in_kernel: bool,
251 use_gate_in_kernel: bool,
252 use_beta_sigmoid_in_kernel: bool,
253 lower_bound: Optional[float],
254) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, float]:
255 """Apply the common KDA input transformations in FP32."""
256 _, _, num_k_heads, num_v_heads, k_head_dim, _ = _validate_kda_inputs(
257 query, key, value, gate, beta
258 )
259 if use_qk_l2norm_in_kernel:
260 query = query * torch.rsqrt(
261 (query.float() * query.float()).sum(dim=-1, keepdim=True) + 1e-6
262 ).to(query.dtype)
263 key = key * torch.rsqrt(
264 (key.float() * key.float()).sum(dim=-1, keepdim=True) + 1e-6
265 ).to(key.dtype)
266 if use_gate_in_kernel:
267 if a_log is None or dt_bias is None:
268 raise ValueError("a_log and dt_bias are required when gate activation is enabled.")
269 gate = torch_kda_gate(
270 gate,
271 a_log,
272 dt_bias,
273 lower_bound=lower_bound,
274 )
275 if use_beta_sigmoid_in_kernel:
276 beta = beta.float().sigmoid()
278 value_groups = num_v_heads // num_k_heads
279 query = query.float().repeat_interleave(value_groups, dim=2)
280 key = key.float().repeat_interleave(value_groups, dim=2)
281 value = value.float()
282 gate = gate.float()
283 beta = beta.float()
284 effective_scale = k_head_dim ** -0.5 if scale is None else scale
285 return query, key, value, gate, beta, float(effective_scale)
288def _initial_kda_state(
289 value: torch.Tensor,
290 k_head_dim: int,
291 initial_state: Optional[torch.Tensor],
292) -> torch.Tensor:
293 """Build or validate a KDA state in FP32."""
294 expected_shape = (value.shape[0], value.shape[2], k_head_dim, value.shape[3])
295 if initial_state is None:
296 return value.new_zeros(expected_shape, dtype=torch.float32)
297 if tuple(initial_state.shape) != expected_shape:
298 raise ValueError(
299 f"initial_state must have shape {expected_shape}, got "
300 f"{tuple(initial_state.shape)}."
301 )
302 return initial_state.float()
305def torch_recurrent_kda(
306 query: torch.Tensor,
307 key: torch.Tensor,
308 value: torch.Tensor,
309 gate: torch.Tensor,
310 beta: torch.Tensor,
311 *,
312 a_log: Optional[torch.Tensor] = None,
313 dt_bias: Optional[torch.Tensor] = None,
314 scale: Optional[float] = None,
315 initial_state: Optional[torch.Tensor] = None,
316 output_final_state: bool = False,
317 use_qk_l2norm_in_kernel: bool = False,
318 use_gate_in_kernel: bool = False,
319 use_beta_sigmoid_in_kernel: bool = False,
320 lower_bound: Optional[float] = None,
321) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
322 """Evaluate the exact token-recurrent KDA equations with PyTorch ops."""
323 output_dtype = value.dtype
324 query, key, value, gate, beta, scale = _prepare_kda_inputs(
325 query,
326 key,
327 value,
328 gate,
329 beta,
330 a_log=a_log,
331 dt_bias=dt_bias,
332 scale=scale,
333 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
334 use_gate_in_kernel=use_gate_in_kernel,
335 use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
336 lower_bound=lower_bound,
337 )
338 state = _initial_kda_state(value, key.shape[-1], initial_state)
339 outputs = []
340 for token_idx in range(query.shape[1]):
341 query_i = query[:, token_idx]
342 key_i = key[:, token_idx]
343 value_i = value[:, token_idx]
344 gate_i = gate[:, token_idx]
345 beta_i = beta[:, token_idx]
347 state = state * gate_i.exp().unsqueeze(-1)
348 state_value = torch.einsum("bhk,bhkv->bhv", key_i, state)
349 residual = value_i - state_value
350 state = state + torch.einsum(
351 "bhk,bhv->bhkv",
352 key_i * beta_i.unsqueeze(-1),
353 residual,
354 )
355 outputs.append(torch.einsum("bhk,bhkv->bhv", query_i * scale, state))
357 output = torch.stack(outputs, dim=1).to(output_dtype)
358 return output, state if output_final_state else None
361def _inverse_unit_lower_triangular(strict_lower: torch.Tensor) -> torch.Tensor:
362 """Invert ``I + strict_lower`` using differentiable forward substitution."""
363 chunk_size = strict_lower.shape[-1]
364 rows = []
365 for row_idx in range(chunk_size):
366 base = -strict_lower[..., row_idx, :]
367 if rows:
368 previous = torch.stack(rows, dim=-2)
369 row = torch.einsum(
370 "...i,...ij->...j",
371 base[..., :row_idx],
372 previous,
373 )
374 else:
375 row = torch.zeros_like(base)
376 diagonal = F.one_hot( # pylint: disable=not-callable
377 torch.tensor(row_idx, device=strict_lower.device),
378 num_classes=chunk_size,
379 ).to(strict_lower.dtype)
380 rows.append(row + diagonal)
381 return torch.stack(rows, dim=-2)
384def _prepare_kda_wy(
385 query: torch.Tensor,
386 key: torch.Tensor,
387 value: torch.Tensor,
388 gate: torch.Tensor,
389 beta: torch.Tensor,
390 chunk_size: int,
391) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
392 """Build the cumulative gate, causal mask, and chunkwise WY factors."""
393 gate = gate.cumsum(dim=-2)
394 causal = torch.ones(
395 chunk_size,
396 chunk_size,
397 dtype=torch.bool,
398 device=query.device,
399 ).tril()
400 strict_causal = causal.logical_xor(
401 torch.eye(chunk_size, dtype=torch.bool, device=query.device)
402 )
404 gate_pair = gate.unsqueeze(-2) - gate.unsqueeze(-3)
405 key_pair = torch.einsum(
406 "...ik,...jk,...ijk->...ij",
407 key,
408 key,
409 gate_pair.exp(),
410 )
411 strict_lower = (
412 key_pair * beta.unsqueeze(-1)
413 ).masked_fill(~strict_causal, 0)
414 inverse = _inverse_unit_lower_triangular(strict_lower)
415 transform = inverse * beta.unsqueeze(-2)
416 w = transform @ (key * gate.exp())
417 u = transform @ value
419 return gate, causal, w, u
422def _run_kda_chunks(
423 query: torch.Tensor,
424 key: torch.Tensor,
425 gate: torch.Tensor,
426 w: torch.Tensor,
427 u: torch.Tensor,
428 state: torch.Tensor,
429 causal: torch.Tensor,
430 scale: float,
431 num_chunks: int,
432) -> tuple[list[torch.Tensor], torch.Tensor]:
433 """Apply chunk WY factors to the recurrent state in sequence order."""
434 output_chunks = []
435 for chunk_idx in range(num_chunks):
436 query_i = query[:, :, chunk_idx]
437 key_i = key[:, :, chunk_idx]
438 gate_i = gate[:, :, chunk_idx]
439 w_i = w[:, :, chunk_idx]
440 u_i = u[:, :, chunk_idx]
442 relative_gate = gate_i.unsqueeze(-2) - gate_i.unsqueeze(-3)
443 qk = torch.einsum(
444 "bhik,bhjk,bhijk->bhij",
445 query_i,
446 key_i,
447 relative_gate.exp(),
448 ).masked_fill(~causal, 0)
449 value_new = u_i - w_i @ state
450 output_inter = (query_i * gate_i.exp()) @ state
451 output_chunks.append(output_inter * scale + qk @ value_new * scale)
453 last_gate = gate_i[:, :, -1]
454 key_decay = key_i * (last_gate.unsqueeze(-2) - gate_i).exp()
455 state = state * last_gate.exp().unsqueeze(-1)
456 state = state + key_decay.transpose(-1, -2) @ value_new
458 return output_chunks, state
461def torch_chunk_kda(
462 query: torch.Tensor,
463 key: torch.Tensor,
464 value: torch.Tensor,
465 gate: torch.Tensor,
466 beta: torch.Tensor,
467 *,
468 a_log: Optional[torch.Tensor] = None,
469 dt_bias: Optional[torch.Tensor] = None,
470 scale: Optional[float] = None,
471 initial_state: Optional[torch.Tensor] = None,
472 output_final_state: bool = False,
473 use_qk_l2norm_in_kernel: bool = False,
474 use_gate_in_kernel: bool = False,
475 use_beta_sigmoid_in_kernel: bool = False,
476 lower_bound: Optional[float] = None,
477 chunk_size: int = 64,
478 cu_seqlens: Optional[torch.Tensor] = None,
479) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
480 """Evaluate chunkwise KDA using the UT/WY representation.
482 This implementation is algebraically independent from
483 :func:`torch_recurrent_kda`; comparing the two catches errors in the
484 chunk transform that a second token loop would hide.
485 """
486 if cu_seqlens is not None:
487 return packed_reference(
488 torch_chunk_kda, (query, key, value, gate, beta), cu_seqlens,
489 initial_state=initial_state, output_final_state=output_final_state,
490 a_log=a_log, dt_bias=dt_bias, scale=scale,
491 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, use_gate_in_kernel=use_gate_in_kernel,
492 use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel, lower_bound=lower_bound, chunk_size=chunk_size,
493 )
494 if chunk_size <= 0:
495 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
496 output_dtype = value.dtype
497 sequence_length = query.shape[1]
498 query, key, value, gate, beta, scale = _prepare_kda_inputs(
499 query,
500 key,
501 value,
502 gate,
503 beta,
504 a_log=a_log,
505 dt_bias=dt_bias,
506 scale=scale,
507 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
508 use_gate_in_kernel=use_gate_in_kernel,
509 use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
510 lower_bound=lower_bound,
511 )
512 batch_size, _, num_heads, k_head_dim = query.shape
513 v_head_dim = value.shape[-1]
514 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
515 query = F.pad(query, (0, 0, 0, 0, 0, pad_size))
516 key = F.pad(key, (0, 0, 0, 0, 0, pad_size))
517 value = F.pad(value, (0, 0, 0, 0, 0, pad_size))
518 gate = F.pad(gate, (0, 0, 0, 0, 0, pad_size))
519 beta = F.pad(beta, (0, 0, 0, pad_size))
520 num_chunks = query.shape[1] // chunk_size
522 def chunk_view(tensor: torch.Tensor) -> torch.Tensor:
523 """Rearrange a token-first tensor to head-first chunks."""
524 shape = tensor.shape
525 return tensor.reshape(
526 shape[0], num_chunks, chunk_size, shape[2], *shape[3:]
527 ).permute(0, 3, 1, 2, *range(4, tensor.dim() + 1))
529 query, key, value, gate, beta = [
530 chunk_view(tensor) for tensor in (query, key, value, gate, beta)
531 ]
532 gate, causal, w, u = _prepare_kda_wy(query, key, value, gate, beta, chunk_size)
534 state = _initial_kda_state(
535 value.permute(0, 2, 3, 1, 4).reshape(
536 batch_size, num_chunks * chunk_size, num_heads, v_head_dim
537 ),
538 k_head_dim,
539 initial_state,
540 )
541 output_chunks, state = _run_kda_chunks(query, key, gate, w, u, state, causal, scale, num_chunks)
543 output = torch.stack(output_chunks, dim=2)
544 output = output.permute(0, 2, 3, 1, 4).reshape(
545 batch_size,
546 num_chunks * chunk_size,
547 num_heads,
548 v_head_dim,
549 )
550 output = output[:, :sequence_length].to(output_dtype)
551 return output, state if output_final_state else None
554def _normalize_kda_backend(backend: str) -> str:
555 """Normalize the selected backend and reject unsupported names."""
556 backend = backend.lower()
557 if backend not in _KDA_BACKENDS:
558 raise ValueError(
559 f"unsupported KDA backend {backend!r}; "
560 f"expected one of {sorted(_KDA_BACKENDS)}."
561 )
562 return backend
565def _validate_fused_state_contract(initial_state, output_final_state, qk_norm, gate, beta, backend) -> None:
566 """Keep raw-input dispatch separate from the stateful CP backend contract."""
567 if initial_state is not None or output_final_state:
568 raise NotImplementedError(
569 f"KDA backend={backend!r} currently supports zero initial state without "
570 "returning a final state; stateful execution is owned by the CP backend."
571 )
572 if not (qk_norm and gate and beta):
573 raise NotImplementedError(
574 f"KDA backend={backend!r} requires fused Q/K normalization, gate activation, "
575 "and beta sigmoid."
576 )
579def chunk_kda(
580 query: torch.Tensor,
581 key: torch.Tensor,
582 value: torch.Tensor,
583 gate: torch.Tensor,
584 beta: torch.Tensor,
585 *,
586 a_log: Optional[torch.Tensor] = None,
587 dt_bias: Optional[torch.Tensor] = None,
588 scale: Optional[float] = None,
589 initial_state: Optional[torch.Tensor] = None,
590 output_final_state: bool = False,
591 use_qk_l2norm_in_kernel: bool = False,
592 use_gate_in_kernel: bool = False,
593 use_beta_sigmoid_in_kernel: bool = False,
594 lower_bound: Optional[float] = None,
595 chunk_size: int = 64,
596 backend: str = "eager",
597 cu_seqlens: Optional[torch.Tensor | Sequence[int]] = None,
598) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
599 """Dispatch dense or packed KDA to the explicitly selected local backend."""
600 if query.ndim == 3:
601 output, final = chunk_kda(
602 *(tensor.unsqueeze(0) for tensor in (query, key, value, gate, beta)),
603 a_log=a_log, dt_bias=dt_bias, scale=scale, initial_state=initial_state,
604 output_final_state=output_final_state, use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
605 use_gate_in_kernel=use_gate_in_kernel, use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
606 lower_bound=lower_bound, chunk_size=chunk_size, backend=backend, cu_seqlens=cu_seqlens)
607 return output.squeeze(0), final
608 backend = _normalize_kda_backend(backend)
609 if backend == "eager":
610 return torch_chunk_kda(
611 query,
612 key,
613 value,
614 gate,
615 beta,
616 a_log=a_log,
617 dt_bias=dt_bias,
618 scale=scale,
619 initial_state=initial_state,
620 output_final_state=output_final_state,
621 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
622 use_gate_in_kernel=use_gate_in_kernel,
623 use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
624 lower_bound=lower_bound,
625 chunk_size=chunk_size,
626 cu_seqlens=cu_seqlens,
627 )
628 _validate_fused_state_contract(initial_state, output_final_state, use_qk_l2norm_in_kernel,
629 use_gate_in_kernel, use_beta_sigmoid_in_kernel, backend)
630 if backend == "ascendc":
631 output = ascendc_chunk_kda(
632 query, key, value, gate, beta,
633 a_log=a_log.float() if a_log is not None else None,
634 dt_bias=dt_bias.float() if dt_bias is not None else None,
635 scale=scale, lower_bound=lower_bound, chunk_size=chunk_size, cu_seqlens=cu_seqlens,
636 )
637 return output, None
638 triton_a_log = a_log.float() if a_log is not None else None
639 triton_dt_bias = dt_bias.float() if dt_bias is not None else None
640 if not _is_triton_kda_input_supported(
641 query,
642 key,
643 value,
644 gate,
645 beta,
646 triton_a_log,
647 triton_dt_bias,
648 chunk_size=chunk_size,
649 lower_bound=lower_bound,
650 ):
651 raise RuntimeError(
652 "KDA backend='triton' received an unsupported input contract; it requires "
653 "q/k/v/gate/beta=bf16, a_log/dt_bias=fp32, "
654 "head_k_dim=head_v_dim=128, chunk_size=64, and a supported safe gate."
655 )
657 from hyper_parallel.components.functional.kimi_delta_attention import ( # pylint: disable=import-outside-toplevel
658 fused_chunk_kda,
659 )
661 output = fused_chunk_kda(
662 query,
663 key,
664 value,
665 gate,
666 beta,
667 a_log=triton_a_log,
668 dt_bias=triton_dt_bias,
669 scale=scale,
670 lower_bound=lower_bound,
671 chunk_size=chunk_size,
672 safe_gate=True,
673 cu_seqlens=cu_seqlens,
674 )
675 return output, None
678def _chunk_kda_summary_inputs(
679 key: torch.Tensor,
680 value: torch.Tensor,
681 gate: torch.Tensor,
682 beta: torch.Tensor,
683 chunk_size: int,
684) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
685 """Pad and convert token-first KDA tensors into head-first chunks."""
686 sequence_length = key.shape[1]
687 pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size
688 key = F.pad(key, (0, 0, 0, 0, 0, pad_size))
689 value = F.pad(value, (0, 0, 0, 0, 0, pad_size))
690 gate = F.pad(gate, (0, 0, 0, 0, 0, pad_size))
691 beta = F.pad(beta, (0, 0, 0, pad_size))
692 num_chunks = key.shape[1] // chunk_size
694 def chunk_view(tensor: torch.Tensor) -> torch.Tensor:
695 """Rearrange one token-first tensor to head-first chunks."""
696 shape = tensor.shape
697 return tensor.reshape(
698 shape[0], num_chunks, chunk_size, shape[2], *shape[3:]
699 ).permute(0, 3, 1, 2, *range(4, tensor.dim() + 1))
701 key, value, gate, beta = (
702 chunk_view(tensor) for tensor in (key, value, gate, beta)
703 )
704 return key, value, gate, beta, num_chunks
707def torch_kda_state_summary(
708 query: torch.Tensor,
709 key: torch.Tensor,
710 value: torch.Tensor,
711 gate: torch.Tensor,
712 beta: torch.Tensor,
713 *,
714 a_log: Optional[torch.Tensor] = None,
715 dt_bias: Optional[torch.Tensor] = None,
716 scale: Optional[float] = None,
717 use_qk_l2norm_in_kernel: bool = False,
718 use_gate_in_kernel: bool = False,
719 use_beta_sigmoid_in_kernel: bool = False,
720 lower_bound: Optional[float] = None,
721 chunk_size: int = 64,
722) -> tuple[torch.Tensor, torch.Tensor]:
723 """Summarize a local KDA segment as ``H_out = M @ H_in + S``.
725 The implementation follows the same chunkwise WY equations as
726 :func:`torch_chunk_kda`, but computes only the state-independent affine
727 transition. ``S`` has shape ``[B, HV, K, V]`` and ``M`` has shape
728 ``[B, HV, K, K]``. Query is accepted to share KDA input preprocessing and
729 grouped-value validation; the state transition itself is independent of
730 query and output scale.
731 """
732 if chunk_size <= 0:
733 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
734 _, key, value, gate, beta, _ = _prepare_kda_inputs(
735 query,
736 key,
737 value,
738 gate,
739 beta,
740 a_log=a_log,
741 dt_bias=dt_bias,
742 scale=scale,
743 use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
744 use_gate_in_kernel=use_gate_in_kernel,
745 use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
746 lower_bound=lower_bound,
747 )
748 key, value, gate, beta, num_chunks = _chunk_kda_summary_inputs(
749 key,
750 value,
751 gate,
752 beta,
753 chunk_size,
754 )
755 gate = gate.cumsum(dim=-2)
756 strict_causal = torch.ones(
757 chunk_size,
758 chunk_size,
759 dtype=torch.bool,
760 device=key.device,
761 ).tril(diagonal=-1)
762 gate_pair = gate.unsqueeze(-2) - gate.unsqueeze(-3)
763 key_pair = torch.einsum(
764 "...ik,...jk,...ijk->...ij",
765 key,
766 key,
767 gate_pair.exp(),
768 )
769 strict_lower = (key_pair * beta.unsqueeze(-1)).masked_fill(
770 ~strict_causal,
771 0,
772 )
773 inverse = _inverse_unit_lower_triangular(strict_lower)
774 transform = inverse * beta.unsqueeze(-2)
775 w = transform @ (key * gate.exp())
776 u = transform @ value
778 batch_size, num_heads, _, _, k_head_dim = key.shape
779 v_head_dim = value.shape[-1]
780 eye = torch.eye(
781 k_head_dim,
782 dtype=torch.float32,
783 device=key.device,
784 ).reshape(1, 1, k_head_dim, k_head_dim)
785 state_ext = torch.zeros(
786 batch_size,
787 num_heads,
788 k_head_dim,
789 v_head_dim,
790 dtype=torch.float32,
791 device=key.device,
792 )
793 transition = eye.expand(batch_size, num_heads, -1, -1).clone()
795 for chunk_idx in range(num_chunks):
796 key_i = key[:, :, chunk_idx]
797 gate_i = gate[:, :, chunk_idx]
798 w_i = w[:, :, chunk_idx]
799 u_i = u[:, :, chunk_idx]
800 last_gate = gate_i[:, :, -1]
801 key_decay = key_i * (last_gate.unsqueeze(-2) - gate_i).exp()
802 transition_i = torch.diag_embed(last_gate.exp())
803 transition_i = transition_i - key_decay.transpose(-1, -2) @ w_i
804 state_ext_i = key_decay.transpose(-1, -2) @ u_i
805 state_ext = transition_i @ state_ext + state_ext_i
806 transition = transition_i @ transition
808 return state_ext.float(), transition.float()
811def torch_apply_kda_state_summary(
812 state_ext: torch.Tensor,
813 transition: torch.Tensor,
814 initial_state: Optional[torch.Tensor],
815) -> torch.Tensor:
816 """Apply a KDA affine state summary to an incoming state."""
817 if initial_state is None:
818 return state_ext.float()
819 return (transition @ initial_state.to(transition) + state_ext).float()
822def torch_compose_kda_state_summaries(
823 first_state_ext: torch.Tensor,
824 first_transition: torch.Tensor,
825 second_state_ext: torch.Tensor,
826 second_transition: torch.Tensor,
827) -> tuple[torch.Tensor, torch.Tensor]:
828 """Compose two chronological KDA summaries, first then second."""
829 state_ext = second_transition @ first_state_ext + second_state_ext
830 transition = second_transition @ first_transition
831 return state_ext, transition
834class KimiRMSNormGated(nn.Module):
835 """Per-head RMSNorm followed by the sigmoid output gate used by KDA."""
837 def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
838 """Initialize the per-head normalization weight and epsilon."""
839 super().__init__()
840 self.weight = nn.Parameter(torch.ones(hidden_size))
841 self.variance_epsilon = eps
843 def forward(self, hidden_states: torch.Tensor, gate: torch.Tensor) -> torch.Tensor:
844 """Normalize in FP32 and apply a sigmoid gate in the input dtype."""
845 input_dtype = hidden_states.dtype
846 normalized = hidden_states.float()
847 variance = normalized.square().mean(dim=-1, keepdim=True)
848 normalized = normalized * torch.rsqrt(variance + self.variance_epsilon)
849 normalized = normalized * self.weight.float()
850 return (normalized * gate.float().sigmoid()).to(input_dtype)
853def _make_short_convolution(channels: int, kernel_size: int) -> nn.Conv1d:
854 """Create the depthwise causal convolution used by the training layer."""
855 convolution = nn.Conv1d(
856 channels,
857 channels,
858 kernel_size=kernel_size,
859 groups=channels,
860 padding=kernel_size - 1,
861 bias=False,
862 )
863 convolution.activation = "silu"
864 return convolution
867class KimiDeltaAttention(nn.Module):
868 """Training-only Kimi Delta Attention layer implemented with eager Torch.
870 The layer owns the projection, ShortConv, gate, normalization, and output
871 parameters used by Kimi K3. Accelerator selection is intentionally outside
872 this model module; context-parallel policies may replace ``forward`` while
873 preserving these parameters and state-dict names.
874 """
876 def __init__(
877 self,
878 hidden_size: int,
879 num_heads: int = 32,
880 num_v_heads: Optional[int] = None,
881 head_k_dim: int = 128,
882 head_v_dim: int = 128,
883 conv_kernel_size: int = 4,
884 chunk_size: int = 64,
885 gate_lower_bound: float = -5.0,
886 use_full_rank_gate: bool = True,
887 rms_norm_eps: float = 1e-6,
888 layer_idx: int = 0,
889 ) -> None:
890 """Initialize a configurable, training-only KDA attention layer."""
891 super().__init__()
892 if hidden_size <= 0 or num_heads <= 0:
893 raise ValueError("hidden_size and num_heads must be positive.")
894 if num_v_heads is None:
895 num_v_heads = num_heads
896 if num_v_heads <= 0 or num_v_heads % num_heads:
897 raise ValueError("num_v_heads must be positive and divisible by num_heads.")
898 if head_k_dim <= 0 or head_v_dim <= 0:
899 raise ValueError("KDA head dimensions must be positive.")
900 if conv_kernel_size <= 0 or chunk_size <= 0:
901 raise ValueError("conv_kernel_size and chunk_size must be positive.")
902 if gate_lower_bound >= 0:
903 raise ValueError("gate_lower_bound must be negative.")
905 self.hidden_size = hidden_size
906 self.num_heads = num_heads
907 self.num_k_heads = num_heads
908 self.num_v_heads = num_v_heads
909 self.head_k_dim = head_k_dim
910 self.head_v_dim = head_v_dim
911 self.head_dim = head_v_dim
912 self.key_dim = num_heads * head_k_dim
913 self.value_dim = num_v_heads * head_v_dim
914 self.conv_kernel_size = conv_kernel_size
915 self.conv_size = conv_kernel_size
916 self.chunk_size = chunk_size
917 self.layer_idx = layer_idx
918 self.gate_lower_bound = gate_lower_bound
919 self.lower_bound = gate_lower_bound
920 self.safe_gate = True
921 self.allow_neg_eigval = False
922 self.use_short_conv = True
923 self.use_full_rank_gate = use_full_rank_gate
925 self.q_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
926 self.k_proj = nn.Linear(hidden_size, self.key_dim, bias=False)
927 self.v_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
928 self.q_conv1d = _make_short_convolution(self.key_dim, conv_kernel_size)
929 self.k_conv1d = _make_short_convolution(self.key_dim, conv_kernel_size)
930 self.v_conv1d = _make_short_convolution(self.value_dim, conv_kernel_size)
932 self.A_log = nn.Parameter(torch.empty(num_v_heads, dtype=torch.float32))
933 self.dt_bias = nn.Parameter(
934 torch.zeros(num_v_heads * head_k_dim, dtype=torch.float32)
935 )
936 with torch.no_grad():
937 self.A_log.copy_(torch.empty_like(self.A_log).uniform_(1, 16).log())
939 self.f_a_proj = nn.Linear(hidden_size, head_k_dim, bias=False)
940 self.f_b_proj = nn.Linear(head_k_dim, num_v_heads * head_k_dim, bias=False)
941 self.b_proj = nn.Linear(hidden_size, num_v_heads, bias=False)
942 if use_full_rank_gate:
943 self.g_proj = nn.Linear(hidden_size, self.value_dim, bias=False)
944 else:
945 self.g_a_proj = nn.Linear(hidden_size, head_v_dim, bias=False)
946 self.g_b_proj = nn.Linear(head_v_dim, self.value_dim, bias=False)
947 self.o_norm = KimiRMSNormGated(head_v_dim, eps=rms_norm_eps)
948 self.o_proj = nn.Linear(self.value_dim, hidden_size, bias=False)
950 @staticmethod
951 def _run_short_convolution(
952 projected: torch.Tensor,
953 convolution: nn.Conv1d,
954 sequence_info=None,
955 ) -> torch.Tensor:
956 """Apply one causal depthwise convolution and trim its right padding."""
957 if sequence_info is not None:
958 return packed_causal_conv1d(projected, convolution, sequence_info)
959 sequence_length = projected.shape[1]
960 output = convolution(projected.transpose(1, 2))[:, :, :sequence_length]
961 return F.silu(output.transpose(1, 2))
963 def _output_gate(self, hidden_states: torch.Tensor) -> torch.Tensor:
964 """Project the layer input to the per-head output gate."""
965 if self.use_full_rank_gate:
966 return self.g_proj(hidden_states)
967 return self.g_b_proj(self.g_a_proj(hidden_states))
969 def forward(
970 self,
971 hidden_states: torch.Tensor,
972 attention_mask: Optional[torch.Tensor] = None,
973 *,
974 use_cache: bool = False,
975 output_attentions: bool = False,
976 **kwargs: Any,
977 ) -> torch.Tensor:
978 """Run dense or packed training-time KDA without recurrent cache."""
979 if use_cache or kwargs.get("cache_params") is not None:
980 raise NotImplementedError("KimiDeltaAttention does not support inference cache.")
981 if output_attentions:
982 raise NotImplementedError("KDA does not materialize attention weights.")
983 if attention_mask is not None and not bool(attention_mask.bool().all().item()):
984 raise NotImplementedError("Padded KDA batches are not implemented.")
986 is_flat_input = hidden_states.ndim == 2
987 if is_flat_input:
988 hidden_states = hidden_states.unsqueeze(0)
989 batch_size, sequence_length, _ = hidden_states.shape
990 sequence_info = resolve_packed_sequence(kwargs, sequence_length)
991 if sequence_info is not None and batch_size != 1:
992 raise ValueError("Packed KDA requires a flattened batch dimension of one.")
993 query = self._run_short_convolution(
994 self.q_proj(hidden_states), self.q_conv1d, sequence_info
995 ).reshape(batch_size, sequence_length, self.num_heads, self.head_k_dim)
996 key = self._run_short_convolution(
997 self.k_proj(hidden_states), self.k_conv1d, sequence_info
998 ).reshape(batch_size, sequence_length, self.num_heads, self.head_k_dim)
999 value = self._run_short_convolution(
1000 self.v_proj(hidden_states), self.v_conv1d, sequence_info
1001 ).reshape(batch_size, sequence_length, self.num_v_heads, self.head_v_dim)
1002 gate = self.f_b_proj(self.f_a_proj(hidden_states)).reshape(
1003 batch_size, sequence_length, self.num_v_heads, self.head_k_dim
1004 )
1005 beta = self.b_proj(hidden_states).reshape(
1006 batch_size, sequence_length, self.num_v_heads
1007 )
1008 output, _ = torch_chunk_kda(
1009 query,
1010 key,
1011 value,
1012 gate,
1013 beta,
1014 a_log=self.A_log,
1015 dt_bias=self.dt_bias,
1016 initial_state=None,
1017 output_final_state=False,
1018 use_qk_l2norm_in_kernel=True,
1019 use_gate_in_kernel=True,
1020 use_beta_sigmoid_in_kernel=True,
1021 lower_bound=self.gate_lower_bound,
1022 chunk_size=self.chunk_size,
1023 cu_seqlens=None if sequence_info is None else sequence_info.local_offsets,
1024 )
1025 output_gate = self._output_gate(hidden_states).reshape(
1026 batch_size, sequence_length, self.num_v_heads, self.head_v_dim
1027 )
1028 output = self.o_norm(output, output_gate)
1029 output = self.o_proj(output.reshape(batch_size, sequence_length, self.value_dim))
1030 return output.squeeze(0) if is_flat_input else output