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