Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / kimi_delta_attention.py: 23%
543 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"""KDA context-parallel executors: state CP, Ulysses and their combinations."""
16# pylint: disable=forbidden-backend-import
17from __future__ import annotations
19__all__ = [
20 "KimiDeltaAttentionCoreCP",
21 "KimiDeltaAttentionLayerCP",
22 "KimiDeltaAttentionLayerP2PCP",
23 "KimiDeltaAttentionLayerUlyssesCP",
24 "KimiDeltaAttentionP2PCP",
25 "KimiDeltaAttentionUlyssesCP",
26]
28from typing import Any, Optional
30import torch
31import torch.distributed as dist
32from torch import nn
33from torch.nn import functional as F
35from hyper_parallel.components.functional.linear_attention_convolution import (
36 fused_causal_conv1d,
37)
38from hyper_parallel.components.functional.linear_attention_packed import (
39 PackedSequenceInfo, packed_causal_conv1d, resolve_packed_sequence,
40)
42from hyper_parallel.core.context_parallel.context_parallel import _ensure_1d
43from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
44from hyper_parallel.core.dtensor.dtensor import DTensor
45from hyper_parallel.components.modules.kimi_delta_attention import (
46 chunk_kda,
47 torch_apply_kda_state_summary,
48 torch_kda_state_summary,
49)
50from hyper_parallel.core.utils import communication
51from hyper_parallel.distributed.context_parallel.kimi_delta_attention_state import (
52 build_kda_state_cp,
53 split_kda_mesh,
54)
57_KDA_BACKENDS = frozenset({"eager", "triton", "ascendc"})
60def _run_local_kda(
61 query: torch.Tensor,
62 key: torch.Tensor,
63 value: torch.Tensor,
64 gate: torch.Tensor,
65 beta: torch.Tensor,
66 *,
67 a_log: torch.Tensor,
68 dt_bias: torch.Tensor,
69 scale: Optional[float],
70 lower_bound: float,
71 chunk_size: int,
72 safe_gate: bool,
73 backend: str,
74 sequence_info: Optional[PackedSequenceInfo] = None,
75) -> torch.Tensor:
76 """Run the explicitly selected local KDA implementation."""
77 if backend in ("triton", "ascendc") and not safe_gate:
78 raise RuntimeError("The fused KDA backend requires the lower-bounded gate.")
79 cu_seqlens = None
80 if sequence_info is not None:
81 # Native APIs consume host metadata; do not materialize an NPU plan
82 # only to copy its offsets back to the host at every layer.
83 cu_seqlens = (sequence_info.local_offsets if backend == "ascendc" else
84 sequence_info.materialize(query.device, chunk_size).cu_seqlens)
85 output, _ = chunk_kda(
86 query,
87 key,
88 value,
89 gate,
90 beta,
91 a_log=a_log,
92 dt_bias=dt_bias,
93 scale=scale,
94 lower_bound=lower_bound,
95 chunk_size=chunk_size,
96 initial_state=None,
97 output_final_state=False,
98 use_qk_l2norm_in_kernel=True,
99 use_gate_in_kernel=True,
100 use_beta_sigmoid_in_kernel=True,
101 backend=backend,
102 cu_seqlens=cu_seqlens,
103 )
104 return output
107def _is_official_kimi_layer(module: nn.Module) -> bool:
108 """Whether ``module`` follows the official Kimi K3 projection contract."""
109 return hasattr(module, "f_a_proj") and hasattr(module, "f_b_proj")
112def _num_value_heads(module: nn.Module) -> int:
113 """Return the number of KDA value heads for either supported layer API."""
114 if hasattr(module, "num_v_heads"):
115 return int(module.num_v_heads)
116 return int(module.num_heads)
119def _value_head_dim(module: nn.Module) -> int:
120 """Return the KDA value head width for either supported layer API."""
121 if hasattr(module, "head_v_dim"):
122 return int(module.head_v_dim)
123 return int(module.head_dim)
126def _key_projection_dim(module: nn.Module) -> int:
127 """Return the flattened Q/K projection width."""
128 if hasattr(module, "key_dim"):
129 return int(module.key_dim)
130 return int(module.num_heads * module.head_k_dim)
133def _value_projection_dim(module: nn.Module) -> int:
134 """Return the flattened value projection width."""
135 if hasattr(module, "value_dim"):
136 return int(module.value_dim)
137 return int(_num_value_heads(module) * _value_head_dim(module))
140def _gate_lower_bound(module: nn.Module) -> Optional[float]:
141 """Return the lower bound used by the safe KDA gate parameterization."""
142 return getattr(module, "lower_bound", getattr(module, "gate_lower_bound", None))
145def _uses_safe_gate(module: nn.Module) -> bool:
146 """Return whether the layer enables the safe lower-bounded gate."""
147 return bool(getattr(module, "safe_gate", _gate_lower_bound(module) is not None))
150def _uses_short_conv(module: nn.Module) -> bool:
151 """Return whether Q/K/V pass through the layer's ShortConvolution modules."""
152 return bool(
153 getattr(
154 module,
155 "use_short_conv",
156 all(hasattr(module, name) for name in ("q_conv1d", "k_conv1d", "v_conv1d")),
157 )
158 )
161def _forget_gate_projection(module: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
162 """Run the module-specific state-decay gate projection."""
163 if hasattr(module, "f_proj"):
164 return module.f_proj(hidden_states)
165 return module.f_b_proj(module.f_a_proj(hidden_states))
168def _output_gate_projection(module: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
169 """Run the module-specific output-gate projection."""
170 if hasattr(module, "g_proj"):
171 return module.g_proj(hidden_states)
172 return module.g_b_proj(module.g_a_proj(hidden_states))
175def _format_layer_output(module: nn.Module, output: torch.Tensor):
176 """Preserve the original layer's tensor or three-tuple output contract."""
177 if _is_official_kimi_layer(module):
178 return output
179 return output, None, None
182def _prepare_layer_input(
183 hidden_states: torch.Tensor,
184 kwargs: dict[str, Any],
185 cp_rank: int,
186 cp_size: int,
187) -> tuple[torch.Tensor, Optional[PackedSequenceInfo], bool]:
188 """Normalize a local layer input and resolve this invocation's packed metadata."""
189 hidden_states = _local_cp_tensor(hidden_states)
190 # Restore the caller's layout on return; offsets define document boundaries.
191 is_flat_input = hidden_states.ndim == 2
192 if is_flat_input:
193 hidden_states = hidden_states.unsqueeze(0)
194 sequence_info = resolve_packed_sequence(
195 kwargs, hidden_states.shape[1], cp_rank=cp_rank, cp_size=cp_size)
196 if sequence_info is not None and hidden_states.shape[0] != 1:
197 raise ValueError("Packed KDA requires batch dimension one.")
198 return hidden_states, sequence_info, is_flat_input
201def _global_peer_rank(cp_mesh: DeviceMesh, local_rank: int) -> int:
202 """Map one CP-local rank index to its global distributed rank."""
203 return int(cp_mesh.rank_list[local_rank])
206def _local_cp_tensor(tensor: torch.Tensor) -> torch.Tensor:
207 """Return the local sequence shard of a CP input tensor."""
208 if isinstance(tensor, DTensor):
209 return tensor.to_local()
210 return tensor
213def _all_to_all_previous_rank_halo(
214 tail: torch.Tensor,
215 cp_mesh: DeviceMesh,
216 cp_rank: int,
217 cp_size: int,
218) -> torch.Tensor:
219 """Send a causal-convolution halo only to the next CP rank."""
220 if cp_size == 1:
221 return torch.zeros_like(tail)
223 cp_group = cp_mesh.get_group()
224 group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
225 rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
226 rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
227 halo_width = tail.shape[1]
229 input_splits = [0] * cp_size
230 exchange_input = tail.permute(1, 0, 2).contiguous()
231 if cp_rank < cp_size - 1:
232 input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
233 else:
234 exchange_input = exchange_input[:0]
236 output_splits = [0] * cp_size
237 if cp_rank > 0:
238 output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width
240 exchange_output = communication.differentiable_all_to_all_single(
241 exchange_input,
242 input_splits,
243 output_splits,
244 group=cp_group,
245 )
246 if cp_rank == 0:
247 return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
248 return exchange_output.permute(1, 0, 2).contiguous()
251def _run_causal_short_conv(
252 tensor: torch.Tensor,
253 convolution: nn.Conv1d,
254 halo: Optional[torch.Tensor] = None,
255 sequence_info: Optional[PackedSequenceInfo] = None,
256 backend: str = "eager",
257) -> torch.Tensor:
258 """Apply one local ShortConv, optionally with a preceding-rank halo."""
259 if backend != "eager":
260 return fused_causal_conv1d(
261 tensor, convolution.weight, convolution.bias, sequence_info, halo,
262 backend=backend, dilation=convolution.dilation[0])
263 if sequence_info is not None:
264 return packed_causal_conv1d(tensor, convolution, sequence_info, halo)
265 if halo is None:
266 conv_input = tensor.transpose(1, 2)
267 padding = convolution.padding
268 else:
269 conv_input = torch.cat((halo, tensor), dim=1).transpose(1, 2)
270 padding = 0
271 output = F.conv1d( # pylint: disable=not-callable
272 input=conv_input,
273 weight=convolution.weight,
274 bias=convolution.bias,
275 stride=convolution.stride,
276 padding=padding,
277 dilation=convolution.dilation,
278 groups=convolution.groups,
279 )
280 if halo is None:
281 output = output[:, :, : tensor.shape[1]]
282 return F.silu(output).transpose(1, 2)
285def _causal_short_convs_with_cp_halo(
286 projected: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
287 convolutions: tuple[nn.Conv1d, nn.Conv1d, nn.Conv1d],
288 cp_mesh: DeviceMesh,
289 cp_rank: int,
290 cp_size: int,
291 sequence_info: Optional[PackedSequenceInfo] = None,
292 backend: str = "eager",
293) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
294 """Run three causal depthwise ShortConvs with one packed halo exchange."""
295 halo_widths = {
296 (conv.kernel_size[0] - 1) * conv.dilation[0] for conv in convolutions
297 }
298 if len(halo_widths) != 1:
299 raise ValueError("KDA P2P requires Q/K/V ShortConv halo widths to match.")
300 halo_width = halo_widths.pop()
301 if halo_width == 0 or cp_size == 1:
302 return tuple(
303 _run_causal_short_conv(tensor, convolution, sequence_info=sequence_info, backend=backend)
304 for tensor, convolution in zip(projected, convolutions)
305 )
306 if projected[0].shape[1] < halo_width:
307 raise ValueError(
308 "KDA P2P ShortConv requires local_seq_len >= halo width, got "
309 f"{projected[0].shape[1]} < {halo_width}."
310 )
312 channel_sizes = [tensor.shape[-1] for tensor in projected]
313 packed_tail = torch.cat(
314 [tensor[:, -halo_width:, :] for tensor in projected],
315 dim=-1,
316 ).contiguous()
317 packed_halo = _all_to_all_previous_rank_halo(
318 packed_tail,
319 cp_mesh,
320 cp_rank,
321 cp_size,
322 )
323 halos = torch.split(packed_halo, channel_sizes, dim=-1)
324 return tuple(
325 _run_causal_short_conv(tensor, convolution, halo, sequence_info=sequence_info, backend=backend)
326 for tensor, halo, convolution in zip(projected, halos, convolutions)
327 )
330def _slice_local_heads(
331 tensor: torch.Tensor,
332 cp_rank: int,
333 cp_size: int,
334) -> torch.Tensor:
335 """Return the current CP rank's contiguous head slice."""
336 if tensor.shape[0] % cp_size != 0:
337 raise ValueError(
338 f"KDA head dimension {tensor.shape[0]} must be divisible by "
339 f"cp_size {cp_size}."
340 )
341 local_heads = tensor.shape[0] // cp_size
342 return tensor.narrow(0, cp_rank * local_heads, local_heads)
345def _differentiable_all_to_all_shard(
346 tensor: torch.Tensor,
347 device_mesh: DeviceMesh,
348 *,
349 split_dim: int,
350 concat_dim: int,
351) -> torch.Tensor:
352 """Redistribute a local tensor from concat-dim shard to split-dim shard."""
353 split_count = device_mesh.size()
354 if split_count == 1:
355 return tensor
357 original_shape = tuple(tensor.shape)
358 if original_shape[split_dim] % split_count != 0:
359 raise ValueError(
360 f"KDA all-to-all split dimension {split_dim} with size "
361 f"{original_shape[split_dim]} must be divisible by cp_size "
362 f"{split_count}."
363 )
364 split_size = original_shape[split_dim] // split_count
365 final_shape = list(original_shape)
366 final_shape[split_dim] = split_size
367 final_shape[concat_dim] *= split_count
369 reshape_dims = list(original_shape)
370 reshape_dims[split_dim] = split_count
371 reshape_dims.insert(split_dim + 1, split_size)
372 permutation = list(range(len(reshape_dims)))
373 permutation.remove(split_dim)
374 permutation.insert(0, split_dim)
375 all_to_all_input = tensor.reshape(reshape_dims).permute(permutation).contiguous()
376 all_to_all_input = all_to_all_input.flatten(0, 1)
378 split_length = all_to_all_input.shape[0] // split_count
379 splits = [split_length] * split_count
380 output = communication.differentiable_all_to_all_single(
381 all_to_all_input,
382 splits,
383 splits,
384 group=device_mesh.get_group(),
385 )
387 output_reshape = list(output.shape)
388 output_reshape[0] = split_count
389 output_reshape.insert(1, output.shape[0] // split_count)
390 permutation = list(range(len(output_reshape)))
391 rank_dimension = permutation.pop(0)
392 permutation.insert(concat_dim, rank_dimension)
393 output = output.reshape(output_reshape).permute(permutation).contiguous()
395 output_shape = list(output.shape)
396 output_shape[concat_dim] *= output_shape[concat_dim + 1]
397 output_shape.pop(concat_dim + 1)
398 return output.reshape(output_shape).view(final_shape)
401class _RecvKDAInitialState(torch.autograd.Function):
402 """Receive a KDA initial state and return its gradient to the sender."""
404 @staticmethod
405 def forward( # pylint: disable=arguments-differ
406 ctx: Any,
407 anchor: torch.Tensor,
408 cp_group: Any,
409 prev_rank: int,
410 state_shape: tuple[int, ...],
411 ) -> torch.Tensor:
412 """Receive the initial recurrent state from the previous CP rank.
414 Args:
415 ctx: Autograd context holding tensors and communication metadata.
416 anchor: Tensor supplying the receive device and autograd dependency.
417 cp_group: Precreated chronological CP process group.
418 prev_rank: Global rank of the preceding sequence partition.
419 state_shape: Shape of the FP32 recurrent state.
420 """
421 state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
422 dist.recv(state, src=prev_rank, group=cp_group)
423 ctx.cp_group = cp_group
424 ctx.prev_rank = prev_rank
425 return state
427 @staticmethod
428 def backward(
429 ctx: Any,
430 grad_state: Optional[torch.Tensor],
431 ) -> tuple[None, None, None, None]:
432 """Send the accumulated initial-state gradient to the previous rank.
434 Args:
435 ctx: Autograd context holding tensors and communication metadata.
436 grad_state: Gradient of the received initial state.
437 """
438 if grad_state is None:
439 raise RuntimeError("KDA P2P backward is missing the initial-state gradient.")
440 dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
441 return None, None, None, None
444class _SendKDAFinalState(torch.autograd.Function):
445 """Send a KDA final state and receive its gradient during backward."""
447 @staticmethod
448 def forward( # pylint: disable=arguments-differ
449 ctx: Any,
450 final_state: torch.Tensor,
451 cp_group: Any,
452 next_rank: int,
453 ) -> torch.Tensor:
454 """Send the final recurrent state to the next CP rank.
456 Args:
457 ctx: Autograd context holding tensors and communication metadata.
458 final_state: Outgoing recurrent state.
459 cp_group: Precreated chronological CP process group.
460 next_rank: Global rank of the following sequence partition.
461 """
462 dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
463 ctx.cp_group = cp_group
464 ctx.next_rank = next_rank
465 ctx.state_shape = tuple(final_state.shape)
466 ctx.state_dtype = final_state.dtype
467 return final_state.new_zeros(())
469 @staticmethod
470 def backward(
471 ctx: Any,
472 grad_token: torch.Tensor,
473 ) -> tuple[torch.Tensor, None, None]:
474 """Receive the final-state gradient from the next rank.
476 Args:
477 ctx: Autograd context holding tensors and communication metadata.
478 grad_token: Gradient token supplying the receive device.
479 """
480 grad_state = torch.empty(
481 ctx.state_shape,
482 device=grad_token.device,
483 dtype=ctx.state_dtype,
484 )
485 dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
486 return grad_state, None, None
489def _eager_kda_packed_p2p(
490 query: torch.Tensor,
491 key: torch.Tensor,
492 value: torch.Tensor,
493 gate: torch.Tensor,
494 beta: torch.Tensor,
495 *,
496 a_log: torch.Tensor,
497 dt_bias: torch.Tensor,
498 scale: Optional[float],
499 lower_bound: float,
500 chunk_size: int,
501 cp_group: Any,
502 prev_rank: int,
503 next_rank: int,
504 cp_rank: int,
505 cp_size: int,
506 sequence_info: Optional[PackedSequenceInfo] = None,
507) -> torch.Tensor:
508 """Run native per-document recurrence with differentiable neighbor states."""
509 # Portable semantic reference: native fragment recurrence, fixed sends.
510 incoming = None
511 if cp_rank > 0:
512 shape = (1, value.shape[2], query.shape[-1], value.shape[-1])
513 incoming = _RecvKDAInitialState.apply(query, cp_group, prev_rank, shape)
514 output, finals = chunk_kda(
515 query, key, value, gate, beta, a_log=a_log, dt_bias=dt_bias, scale=scale,
516 initial_state=sequence_info.first_state(incoming), output_final_state=True,
517 use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, use_beta_sigmoid_in_kernel=True,
518 lower_bound=lower_bound, chunk_size=chunk_size, backend="eager",
519 cu_seqlens=sequence_info.local_offsets,
520 )
521 if incoming is not None:
522 output = output + incoming.sum().to(output.dtype) * 0
523 if cp_rank < cp_size - 1:
524 token = _SendKDAFinalState.apply(finals[-1:].contiguous(), cp_group, next_rank)
525 output = output + token.to(output.dtype) * 0
526 return output
529def _eager_kda_state_p2p(
530 query: torch.Tensor,
531 key: torch.Tensor,
532 value: torch.Tensor,
533 gate: torch.Tensor,
534 beta: torch.Tensor,
535 *,
536 a_log: torch.Tensor,
537 dt_bias: torch.Tensor,
538 scale: Optional[float],
539 lower_bound: float,
540 chunk_size: int,
541 cp_group: Any,
542 prev_rank: int,
543 next_rank: int,
544 cp_rank: int,
545 cp_size: int,
546 sequence_info: Optional[PackedSequenceInfo] = None,
547) -> torch.Tensor:
548 """Run the differentiable Torch reference for affine state-P2P KDA."""
549 if sequence_info is not None:
550 return _eager_kda_packed_p2p(
551 query, key, value, gate, beta, a_log=a_log, dt_bias=dt_bias, scale=scale,
552 lower_bound=lower_bound, chunk_size=chunk_size, cp_group=cp_group,
553 prev_rank=prev_rank, next_rank=next_rank, cp_rank=cp_rank, cp_size=cp_size,
554 sequence_info=sequence_info,
555 )
556 if cp_size == 1:
557 return _run_local_kda(
558 query,
559 key,
560 value,
561 gate,
562 beta,
563 a_log=a_log,
564 dt_bias=dt_bias,
565 scale=scale,
566 lower_bound=lower_bound,
567 chunk_size=chunk_size,
568 safe_gate=True,
569 backend="eager",
570 )
572 state_ext = None
573 transition = None
574 if cp_rank < cp_size - 1:
575 state_ext, transition = torch_kda_state_summary(
576 query,
577 key,
578 value,
579 gate,
580 beta,
581 a_log=a_log,
582 dt_bias=dt_bias,
583 scale=scale,
584 use_qk_l2norm_in_kernel=True,
585 use_gate_in_kernel=True,
586 use_beta_sigmoid_in_kernel=True,
587 lower_bound=lower_bound,
588 chunk_size=chunk_size,
589 )
591 initial_state = None
592 if cp_rank > 0:
593 state_shape = (
594 query.shape[0],
595 value.shape[2],
596 query.shape[-1],
597 value.shape[-1],
598 )
599 initial_state = _RecvKDAInitialState.apply(
600 query,
601 cp_group,
602 prev_rank,
603 state_shape,
604 )
606 send_token = None
607 if cp_rank < cp_size - 1:
608 final_state = torch_apply_kda_state_summary(state_ext, transition, initial_state)
609 send_token = _SendKDAFinalState.apply(final_state, cp_group, next_rank)
611 output, _ = chunk_kda(
612 query,
613 key,
614 value,
615 gate,
616 beta,
617 a_log=a_log,
618 dt_bias=dt_bias,
619 scale=scale,
620 initial_state=initial_state,
621 output_final_state=False,
622 use_qk_l2norm_in_kernel=True,
623 use_gate_in_kernel=True,
624 use_beta_sigmoid_in_kernel=True,
625 lower_bound=lower_bound,
626 chunk_size=chunk_size,
627 backend="eager",
628 )
629 if send_token is not None:
630 output = output + send_token.to(output.dtype) * 0
631 return output
634def _validate_projected_inputs(
635 ulysses_degree: int,
636 query: torch.Tensor,
637 key: torch.Tensor,
638 value: torch.Tensor,
639 gate: torch.Tensor,
640 beta: torch.Tensor,
641 a_log: torch.Tensor,
642 dt_bias: torch.Tensor,
643) -> None:
644 """Validate the projected-tensor contract and Ulysses divisibility."""
645 if (query.dim(), key.dim(), value.dim(), gate.dim()) != (4, 4, 4, 4):
646 raise ValueError("query, key, value, and gate must be rank-4 tensors.")
647 if beta.dim() != 3:
648 raise ValueError("beta must be a rank-3 tensor.")
649 if query.shape != key.shape:
650 raise ValueError("query and key must have identical shapes.")
652 batch_size, local_seq_len, num_k_heads, k_head_dim = query.shape
653 expected_prefix = (batch_size, local_seq_len)
654 if value.shape[:2] != expected_prefix:
655 raise ValueError("value batch and sequence dimensions must match query.")
656 num_v_heads = value.shape[2]
657 if gate.shape != (batch_size, local_seq_len, num_v_heads, k_head_dim):
658 raise ValueError(
659 "gate must have shape [B, S_local, num_v_heads, k_head_dim]."
660 )
661 if beta.shape != (batch_size, local_seq_len, num_v_heads):
662 raise ValueError("beta must have shape [B, S_local, num_v_heads].")
663 if num_v_heads % num_k_heads != 0:
664 raise ValueError("num_v_heads must be divisible by num_k_heads.")
665 if num_k_heads % ulysses_degree != 0:
666 raise ValueError(
667 f"KDA num_k_heads ({num_k_heads}) must be divisible by "
668 f"cp_size ({ulysses_degree}) for Ulysses CP."
669 )
670 if num_v_heads % ulysses_degree != 0:
671 raise ValueError(
672 f"KDA num_v_heads ({num_v_heads}) must be divisible by "
673 f"cp_size ({ulysses_degree}) for Ulysses CP."
674 )
675 if a_log.numel() != num_v_heads:
676 raise ValueError(
677 f"a_log must contain {num_v_heads} elements, got {a_log.numel()}."
678 )
679 if dt_bias.numel() != num_v_heads * k_head_dim:
680 raise ValueError(
681 f"dt_bias must contain {num_v_heads * k_head_dim} elements, "
682 f"got {dt_bias.numel()}."
683 )
686class KimiDeltaAttentionUlyssesCP(nn.Module):
687 """Run KDA on sequence-sharded projected tensors with Ulysses all-to-all.
689 The input and output layout is token-first with a local sequence shard:
691 ``[B, S/P, H, D] -> [B, S, H/P, D] -> [B, S/P, H, D]``.
693 This first-stage interface intentionally starts after projections and
694 ShortConv and ends before the output norm/gate/projection. It is therefore
695 independent of one particular Kimi K3 model class. Training uses a zero
696 initial recurrent state and does not materialize a final state.
697 """
699 def __init__(
700 self,
701 device_mesh: DeviceMesh,
702 *,
703 chunk_size: int = 64,
704 lower_bound: float = -5.0,
705 safe_gate: bool = True,
706 backend: str = "eager",
707 ) -> None:
708 """Initialize the KDA Ulysses executor for one 1-D CP mesh."""
709 super().__init__()
710 if chunk_size <= 0:
711 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
712 if lower_bound >= 0:
713 raise ValueError(f"lower_bound must be negative, got {lower_bound}.")
714 backend = backend.lower()
715 if backend not in _KDA_BACKENDS:
716 raise ValueError(
717 f"unsupported KDA backend {backend!r}; "
718 f"expected one of {sorted(_KDA_BACKENDS)}."
719 )
720 self.cp_mesh = _ensure_1d(device_mesh)
721 self.cp_size = self.cp_mesh.size()
722 self.cp_rank = self.cp_mesh.get_local_rank()
723 self.chunk_size = chunk_size
724 self.lower_bound = lower_bound
725 self.safe_gate = safe_gate
726 self.backend = backend
727 self.seq_dim = 1
728 self.head_dim = 2
730 def _validate_inputs(
731 self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
732 gate: torch.Tensor, beta: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor,
733 ) -> None:
734 """Validate the projected contract for this Ulysses degree."""
735 _validate_projected_inputs(self.cp_size, query, key, value, gate, beta, a_log, dt_bias)
737 def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
738 """Redistribute a local sequence shard into a local head shard."""
739 return _differentiable_all_to_all_shard(
740 tensor,
741 self.cp_mesh,
742 split_dim=self.head_dim,
743 concat_dim=self.seq_dim,
744 )
746 def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
747 """Redistribute a local head shard back into a local sequence shard."""
748 return _differentiable_all_to_all_shard(
749 tensor,
750 self.cp_mesh,
751 split_dim=self.seq_dim,
752 concat_dim=self.head_dim,
753 )
755 def forward(
756 self,
757 query: torch.Tensor,
758 key: torch.Tensor,
759 value: torch.Tensor,
760 gate: torch.Tensor,
761 beta: torch.Tensor,
762 *,
763 a_log: torch.Tensor,
764 dt_bias: torch.Tensor,
765 scale: Optional[float] = None,
766 sequence_info: Optional[PackedSequenceInfo] = None,
767 ) -> torch.Tensor:
768 """Run chunkwise KDA using full sequence and local heads.
770 Args:
771 query: Local query tensor [B,T,H,K].
772 key: Matching local key tensor.
773 value: Local value tensor [B,T,HV,V].
774 gate: Per-token gate logits [B,T,HV,K].
775 beta: Per-token beta logits [B,T,HV].
776 a_log: Decay parameter with one value per value head.
777 dt_bias: Learned bias with one value per gate channel.
778 scale: Optional attention scaling factor.
779 """
780 is_flat_input = query.ndim == 3
781 if is_flat_input:
782 output = self.forward(*(tensor.unsqueeze(0) for tensor in (query, key, value, gate, beta)),
783 a_log=a_log, dt_bias=dt_bias, scale=scale, sequence_info=sequence_info)
784 return output.squeeze(0)
785 self._validate_inputs(query, key, value, gate, beta, a_log, dt_bias)
786 exchanged_info = None if sequence_info is None else sequence_info.ulysses(self.cp_size, self.cp_rank)
787 num_v_heads = value.shape[2]
788 k_head_dim = query.shape[-1]
790 query = self._seq_to_head(query)
791 key = self._seq_to_head(key)
792 value = self._seq_to_head(value)
793 gate = self._seq_to_head(gate)
794 beta = self._seq_to_head(beta)
796 local_a_log = _slice_local_heads(
797 a_log.reshape(num_v_heads),
798 self.cp_rank,
799 self.cp_size,
800 )
801 local_dt_bias = _slice_local_heads(
802 dt_bias.reshape(num_v_heads, k_head_dim),
803 self.cp_rank,
804 self.cp_size,
805 )
806 output = _run_local_kda(
807 query,
808 key,
809 value,
810 gate,
811 beta,
812 a_log=local_a_log,
813 dt_bias=local_dt_bias,
814 scale=scale,
815 lower_bound=self.lower_bound,
816 chunk_size=self.chunk_size,
817 safe_gate=self.safe_gate,
818 backend=self.backend,
819 sequence_info=exchanged_info,
820 )
821 return self._head_to_seq(output)
824class _KDAStateCP(nn.Module):
825 """Run local KDA with the configured state exchange."""
827 def __init__(
828 self,
829 device_mesh: DeviceMesh,
830 *,
831 chunk_size: int = 64,
832 lower_bound: float = -5.0,
833 safe_gate: bool = True,
834 backend: str = "eager",
835 state_cp_method: str = "p2p",
836 allgather_group_size: int = 1,
837 ) -> None:
838 """Initialize local KDA and its state exchange on one chronological axis."""
839 super().__init__()
840 if chunk_size <= 0:
841 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
842 backend = backend.lower()
843 if backend not in _KDA_BACKENDS:
844 raise ValueError(
845 f"unsupported KDA backend {backend!r}; "
846 f"expected one of {sorted(_KDA_BACKENDS)}."
847 )
848 self.cp_mesh = _ensure_1d(device_mesh)
849 self.cp_size = self.cp_mesh.size()
850 self.cp_rank = self.cp_mesh.get_local_rank()
851 self.cp_group = self.cp_mesh.get_group()
852 if backend == "ascendc" and (chunk_size != 64 or not safe_gate):
853 raise ValueError("Ascend C KDA requires safe_gate=True and chunk_size=64.")
854 if backend == "eager" and state_cp_method != "p2p":
855 raise ValueError("Eager KDA state exchange supports P2P only.")
856 self.state_cp = build_kda_state_cp(self.cp_mesh, state_cp_method, allgather_group_size)
857 self.prev_rank = _global_peer_rank(
858 self.cp_mesh,
859 max(self.cp_rank - 1, 0),
860 )
861 self.next_rank = _global_peer_rank(
862 self.cp_mesh,
863 min(self.cp_rank + 1, self.cp_size - 1),
864 )
865 self.chunk_size = chunk_size
866 self.lower_bound = lower_bound
867 self.safe_gate = safe_gate
868 self.backend = backend
870 def forward(
871 self,
872 query: torch.Tensor,
873 key: torch.Tensor,
874 value: torch.Tensor,
875 gate: torch.Tensor,
876 beta: torch.Tensor,
877 *,
878 a_log: torch.Tensor,
879 dt_bias: torch.Tensor,
880 scale: Optional[float] = None,
881 sequence_info: Optional[PackedSequenceInfo] = None,
882 ) -> torch.Tensor:
883 """Run one local KDA segment and propagate its recurrent state.
885 Args:
886 query: Local query tensor [B,T,H,K].
887 key: Matching local key tensor.
888 value: Local value tensor [B,T,HV,V].
889 gate: Per-token gate logits [B,T,HV,K].
890 beta: Per-token beta logits [B,T,HV].
891 a_log: Decay parameter with one value per value head.
892 dt_bias: Learned bias with one value per gate channel.
893 scale: Optional attention scaling factor.
894 """
895 if self.cp_size == 1:
896 return _run_local_kda(
897 query,
898 key,
899 value,
900 gate,
901 beta,
902 a_log=a_log,
903 dt_bias=dt_bias,
904 scale=scale,
905 lower_bound=self.lower_bound,
906 chunk_size=self.chunk_size,
907 safe_gate=self.safe_gate,
908 backend=self.backend,
909 sequence_info=sequence_info,
910 )
911 if self.backend == "eager":
912 return _eager_kda_state_p2p(
913 query,
914 key,
915 value,
916 gate,
917 beta,
918 a_log=a_log,
919 dt_bias=dt_bias,
920 scale=scale,
921 lower_bound=self.lower_bound,
922 chunk_size=self.chunk_size,
923 cp_group=self.cp_group,
924 prev_rank=self.prev_rank,
925 next_rank=self.next_rank,
926 cp_rank=self.cp_rank,
927 cp_size=self.cp_size,
928 sequence_info=sequence_info,
929 )
931 if self.backend == "ascendc":
932 # The native package remains optional; loading it belongs to the selected backend.
933 from hyper_parallel.components.functional.kimi_delta_attention_ascendc import ( # pylint: disable=import-outside-toplevel
934 ascendc_state_kda,
935 )
936 return ascendc_state_kda(
937 query, key, value, gate, beta, a_log=a_log, dt_bias=dt_bias, scale=scale,
938 lower_bound=self.lower_bound, state_cp=self.state_cp, cp_mesh=self.cp_mesh,
939 sequence_info=sequence_info)
941 from hyper_parallel.components.functional.kimi_delta_attention import ( # pylint: disable=import-outside-toplevel
942 fused_chunk_kda_p2p,
943 )
944 return fused_chunk_kda_p2p(
945 query,
946 key,
947 value,
948 gate,
949 beta,
950 a_log=a_log.float(),
951 dt_bias=dt_bias.float(),
952 cp_group=self.cp_group,
953 prev_rank=self.prev_rank,
954 next_rank=self.next_rank,
955 cp_rank=self.cp_rank,
956 cp_size=self.cp_size,
957 scale=scale,
958 lower_bound=self.lower_bound,
959 chunk_size=self.chunk_size,
960 safe_gate=self.safe_gate,
961 state_cp=self.state_cp,
962 sequence_info=sequence_info,
963 )
966class KimiDeltaAttentionP2PCP(_KDAStateCP):
967 """Preserve the public projected-P2P type while sharing state execution.
969 Inherited defaults select P2P. Keeping a named subclass preserves the old
970 import and type name without duplicating the common implementation.
971 New hybrid configurations use KimiDeltaAttentionCoreCP.
972 """
975def _validate_layer_short_conv(module: nn.Module, mode: str) -> None:
976 """Check the shared depthwise-convolution contract before building execution."""
977 if _uses_short_conv(module):
978 for name in ("q_conv1d", "k_conv1d", "v_conv1d"):
979 if not hasattr(module, name):
980 raise TypeError(f"Kimi K3 layer is missing {name}.")
981 convolution = getattr(module, name)
982 if convolution.stride != (1,) or convolution.dilation != (1,):
983 raise ValueError(
984 f"Kimi K3 {mode} CP supports ShortConv stride=1 and "
985 "dilation=1 only."
986 )
987 if convolution.groups != convolution.in_channels or convolution.out_channels != convolution.in_channels:
988 raise ValueError(f"Kimi K3 {mode} CP expects depthwise ShortConv with equal input/output channels.")
989 if convolution.padding != (convolution.kernel_size[0] - 1,) or convolution.padding_mode != "zeros":
990 raise ValueError(f"Kimi K3 {mode} CP expects causal zero-padded ShortConv.")
991 if getattr(convolution, "activation", None) not in ("silu", "swish"):
992 raise ValueError(
993 f"Kimi K3 {mode} CP expects ShortConv SiLU activation."
994 )
997class KimiDeltaAttentionLayerUlyssesCP(KimiDeltaAttentionUlyssesCP):
998 """Wrap a training-time Kimi K3 attention layer with Ulysses CP.
1000 The wrapped module keeps ownership of all model parameters. Its local
1001 sequence shard is projected first, then Q/K/V, per-dimension gate logits,
1002 and beta logits are packed into one sequence-to-head all-to-all. Each CP
1003 rank runs the wrapped layer's ShortConv weights and fused KDA core for its
1004 local head shard and the full sequence. The reverse all-to-all restores a
1005 local sequence shard before the original output gate, norm, and projection.
1007 Packed samples use global document offsets. ``backend`` selects both the
1008 attention core and local convolution for dense and packed inputs.
1009 Inference cache and padded sequences are unsupported.
1010 """
1012 def __init__(
1013 self,
1014 module: nn.Module,
1015 device_mesh: DeviceMesh,
1016 *,
1017 chunk_size: int = 64,
1018 backend: str = "eager",
1019 ) -> None:
1020 """Initialize the full-layer adapter without replacing its parameters.
1022 Args:
1023 module: Original Kimi attention layer that owns the parameters.
1024 device_mesh: One-dimensional Ulysses mesh.
1025 chunk_size: Chunk length for the attention core.
1026 backend: Attention and convolution backend: ``eager`` or ``triton``.
1028 Raises:
1029 ValueError: If a backend or the wrapped layer contract is unsupported.
1030 """
1031 lower_bound = _gate_lower_bound(module)
1032 if lower_bound is None:
1033 raise ValueError(
1034 "Kimi K3 Ulysses CP requires the lower-bounded gate parameterization."
1035 )
1036 super().__init__(
1037 device_mesh,
1038 chunk_size=chunk_size,
1039 lower_bound=float(lower_bound),
1040 safe_gate=_uses_safe_gate(module),
1041 backend=backend,
1042 )
1043 self.module = module
1044 self._validate_module()
1046 def _validate_module(self) -> None:
1047 """Validate the first-stage KimiDeltaAttention module contract."""
1048 required_attributes = (
1049 "q_proj",
1050 "k_proj",
1051 "v_proj",
1052 "b_proj",
1053 "o_norm",
1054 "o_proj",
1055 "A_log",
1056 "dt_bias",
1057 "num_heads",
1058 "head_k_dim",
1059 )
1060 missing = [name for name in required_attributes if not hasattr(self.module, name)]
1061 if missing:
1062 raise TypeError(
1063 "Kimi K3 layer is missing required attributes: "
1064 + ", ".join(missing)
1065 )
1066 if self.module.num_heads % self.cp_size != 0:
1067 raise ValueError(
1068 f"KDA num_heads ({self.module.num_heads}) must be divisible by "
1069 f"cp_size ({self.cp_size})."
1070 )
1071 num_v_heads = _num_value_heads(self.module)
1072 if num_v_heads % self.cp_size != 0:
1073 raise ValueError(
1074 f"KDA num_v_heads ({num_v_heads}) must be divisible "
1075 f"by cp_size ({self.cp_size})."
1076 )
1077 if num_v_heads % self.module.num_heads != 0:
1078 raise ValueError("KDA num_v_heads must be divisible by num_heads.")
1079 if not (hasattr(self.module, "f_proj") or _is_official_kimi_layer(self.module)):
1080 raise TypeError("Kimi K3 layer is missing its forget-gate projection.")
1081 if not (
1082 hasattr(self.module, "g_proj")
1083 or all(hasattr(self.module, name) for name in ("g_a_proj", "g_b_proj"))
1084 ):
1085 raise TypeError("Kimi K3 layer is missing its output-gate projection.")
1086 if getattr(self.module, "allow_neg_eigval", False):
1087 raise NotImplementedError(
1088 "Kimi K3 Ulysses CP does not yet support allow_neg_eigval=True."
1089 )
1091 _validate_layer_short_conv(self.module, "Ulysses")
1093 def _pack_projected_inputs(
1094 self,
1095 query: torch.Tensor,
1096 key: torch.Tensor,
1097 value: torch.Tensor,
1098 gate: torch.Tensor,
1099 beta: torch.Tensor,
1100 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
1101 """Pack all projected paths into one sequence-to-head all-to-all."""
1102 if self.cp_size == 1:
1103 return query, key, value, gate, beta
1105 base = self.module
1106 split_sizes = (
1107 _key_projection_dim(base) // self.cp_size,
1108 _key_projection_dim(base) // self.cp_size,
1109 _value_projection_dim(base) // self.cp_size,
1110 _num_value_heads(base) * base.head_k_dim // self.cp_size,
1111 _num_value_heads(base) // self.cp_size,
1112 )
1113 projected = (query, key, value, gate, beta)
1114 chunks = [torch.split(tensor, size, dim=-1) for tensor, size in zip(projected, split_sizes)]
1115 rank_major = [
1116 torch.cat(parts, dim=-1)
1117 for parts in zip(*chunks)
1118 ]
1119 packed = self._seq_to_head(torch.cat(rank_major, dim=-1).contiguous())
1120 return torch.split(packed, split_sizes, dim=-1)
1122 def _local_short_conv(
1123 self,
1124 projected: torch.Tensor,
1125 convolution: nn.Conv1d,
1126 sequence_info: Optional[PackedSequenceInfo] = None,
1127 ) -> torch.Tensor:
1128 """Run one wrapped ShortConv on this rank's contiguous channel slice."""
1129 local_channels = projected.shape[-1]
1130 channel_start = self.cp_rank * local_channels
1131 weight = convolution.weight.narrow(0, channel_start, local_channels)
1132 bias = convolution.bias
1133 if bias is not None:
1134 bias = bias.narrow(0, channel_start, local_channels)
1135 if sequence_info is not None:
1136 return packed_causal_conv1d(
1137 projected, convolution, sequence_info, weight=weight, bias=bias,
1138 backend=self.backend,
1139 )
1140 if self.backend != "eager":
1141 return fused_causal_conv1d(
1142 projected, weight, bias, backend=self.backend,
1143 dilation=convolution.dilation[0],
1144 )
1145 output = F.conv1d( # pylint: disable=not-callable
1146 projected.transpose(1, 2),
1147 weight,
1148 bias,
1149 stride=convolution.stride,
1150 padding=convolution.padding,
1151 dilation=convolution.dilation,
1152 groups=local_channels,
1153 )
1154 output = output[:, :, : projected.shape[1]].transpose(1, 2)
1155 return F.silu(output)
1157 def forward( # pylint: disable=arguments-renamed
1158 self,
1159 hidden_states: torch.Tensor,
1160 attention_mask: Optional[torch.Tensor] = None,
1161 past_key_values: Optional[Any] = None,
1162 use_cache: bool = False,
1163 output_attentions: bool = False,
1164 **kwargs: Any,
1165 ) -> tuple[torch.Tensor, None, None]:
1166 """Run the real Kimi layer using KDA and Ulysses CP.
1168 Args:
1169 hidden_states: Local sequence shard [B,T,hidden_size].
1170 attention_mask: Optional dense mask; padding is unsupported.
1171 past_key_values: Recurrent inference cache, currently unsupported.
1172 use_cache: Must be False for this training adapter.
1173 output_attentions: Must be False; KDA does not materialize attention weights.
1174 kwargs: Optional global sample offsets or prebuilt packed_sequence_info.
1175 """
1176 if not self.module.training:
1177 raise NotImplementedError("Kimi K3 Ulysses CP currently supports training only.")
1178 if past_key_values is not None or use_cache:
1179 raise NotImplementedError("Kimi K3 Ulysses CP does not support recurrent cache.")
1180 if output_attentions:
1181 raise NotImplementedError("KDA does not materialize attention weights.")
1182 if attention_mask is not None and not bool(attention_mask.bool().all().item()):
1183 raise NotImplementedError("Padded KDA CP batches are not implemented.")
1185 hidden_states, sequence_info, is_flat_input = _prepare_layer_input(
1186 hidden_states, kwargs, self.cp_rank, self.cp_size)
1187 exchanged_info = None if sequence_info is None else sequence_info.ulysses(self.cp_size, self.cp_rank)
1188 base = self.module
1189 batch_size, local_seq_len, _ = hidden_states.shape
1190 query = base.q_proj(hidden_states)
1191 key = base.k_proj(hidden_states)
1192 value = base.v_proj(hidden_states)
1193 gate = _forget_gate_projection(base, hidden_states)
1194 beta = base.b_proj(hidden_states)
1195 output_gate = _output_gate_projection(base, hidden_states).reshape(
1196 batch_size,
1197 local_seq_len,
1198 _num_value_heads(base),
1199 _value_head_dim(base),
1200 )
1202 query, key, value, gate, beta = self._pack_projected_inputs(query, key, value, gate, beta)
1203 if _uses_short_conv(base):
1204 query = self._local_short_conv(query, base.q_conv1d, sequence_info=exchanged_info)
1205 key = self._local_short_conv(key, base.k_conv1d, sequence_info=exchanged_info)
1206 value = self._local_short_conv(value, base.v_conv1d, sequence_info=exchanged_info)
1207 else:
1208 query, key, value = (F.silu(tensor) for tensor in (query, key, value))
1210 full_seq_len = query.shape[1]
1211 local_k_heads = base.num_heads // self.cp_size
1212 local_v_heads = _num_value_heads(base) // self.cp_size
1213 query = query.reshape(
1214 batch_size, full_seq_len, local_k_heads, base.head_k_dim
1215 )
1216 key = key.reshape(batch_size, full_seq_len, local_k_heads, base.head_k_dim)
1217 value = value.reshape(
1218 batch_size, full_seq_len, local_v_heads, _value_head_dim(base)
1219 )
1220 gate = gate.reshape(
1221 batch_size, full_seq_len, local_v_heads, base.head_k_dim
1222 )
1223 beta = beta.reshape(batch_size, full_seq_len, local_v_heads)
1225 local_a_log = _slice_local_heads(
1226 base.A_log.reshape(_num_value_heads(base)),
1227 self.cp_rank,
1228 self.cp_size,
1229 )
1230 local_dt_bias = _slice_local_heads(
1231 base.dt_bias.reshape(_num_value_heads(base), base.head_k_dim),
1232 self.cp_rank,
1233 self.cp_size,
1234 )
1235 output = _run_local_kda(
1236 query,
1237 key,
1238 value,
1239 gate,
1240 beta,
1241 a_log=local_a_log,
1242 dt_bias=local_dt_bias,
1243 scale=None,
1244 lower_bound=self.lower_bound,
1245 chunk_size=self.chunk_size,
1246 safe_gate=self.safe_gate,
1247 backend=self.backend,
1248 sequence_info=exchanged_info,
1249 )
1250 output = self._head_to_seq(output)
1251 output = base.o_norm(output, output_gate)
1252 output = output.reshape(
1253 batch_size, local_seq_len, _value_projection_dim(base)
1254 )
1255 output = base.o_proj(output)
1256 return _format_layer_output(base, output.squeeze(0) if is_flat_input else output)
1259def _validate_cp_degrees(cp_size: int, ulysses_degree: int, allgather_group_size: int) -> None:
1260 """Validate both split extents before head division or process-group creation."""
1261 if (not isinstance(ulysses_degree, int) or isinstance(ulysses_degree, bool)
1262 or ulysses_degree < 1 or cp_size % ulysses_degree):
1263 raise ValueError("ulysses_degree must be a positive integer dividing the CP size.")
1264 if (not isinstance(allgather_group_size, int) or isinstance(allgather_group_size, bool)
1265 or allgather_group_size < 1 or (cp_size // ulysses_degree) % allgather_group_size):
1266 raise ValueError("allgather_group_size must be a positive integer dividing the state CP size.")
1269class KimiDeltaAttentionLayerCP(nn.Module):
1270 """Project local tokens, exchange convolution halos, and invoke the configured CP core."""
1272 def __init__(
1273 self,
1274 module: nn.Module,
1275 device_mesh: DeviceMesh,
1276 *,
1277 chunk_size: int = 64,
1278 backend: str = "eager",
1279 ulysses_degree: int = 1,
1280 state_cp_method: str = "p2p",
1281 allgather_group_size: int = 1,
1282 ) -> None:
1283 """Keep layer ownership and full-CP halo metadata separate from the core."""
1284 super().__init__()
1285 lower_bound = _gate_lower_bound(module)
1286 if lower_bound is None:
1287 raise ValueError("Kimi K3 CP requires the lower-bounded gate parameterization.")
1288 self.cp_mesh = _ensure_1d(device_mesh)
1289 self.cp_rank, self.cp_size = self.cp_mesh.get_local_rank(), self.cp_mesh.size()
1290 self.module, self.backend = module, backend.lower()
1291 self._validate_module()
1292 _validate_cp_degrees(self.cp_size, ulysses_degree, allgather_group_size)
1293 if module.num_heads % ulysses_degree or _num_value_heads(module) % ulysses_degree:
1294 raise ValueError("KDA query and value heads must both be divisible by ulysses_degree.")
1295 self.core = KimiDeltaAttentionCoreCP(
1296 device_mesh, ulysses_degree=ulysses_degree, chunk_size=chunk_size,
1297 lower_bound=float(lower_bound), safe_gate=_uses_safe_gate(module), backend=backend,
1298 state_cp_method=state_cp_method, allgather_group_size=allgather_group_size)
1300 def _validate_module(self) -> None:
1301 """Validate the KimiDeltaAttention contract used by state CP."""
1302 required_attributes = (
1303 "q_proj",
1304 "k_proj",
1305 "v_proj",
1306 "b_proj",
1307 "o_norm",
1308 "o_proj",
1309 "A_log",
1310 "dt_bias",
1311 "num_heads",
1312 "head_k_dim",
1313 )
1314 missing = [name for name in required_attributes if not hasattr(self.module, name)]
1315 if missing:
1316 raise TypeError(
1317 "Kimi K3 layer is missing required attributes: "
1318 + ", ".join(missing)
1319 )
1320 num_v_heads = _num_value_heads(self.module)
1321 if num_v_heads % self.module.num_heads:
1322 raise ValueError("KDA num_v_heads must be divisible by num_heads.")
1323 if self.backend in ("triton", "ascendc") and (
1324 self.module.head_k_dim != 128 or _value_head_dim(self.module) != 128
1325 ):
1326 raise NotImplementedError(
1327 "Kimi K3 CP currently requires key/value head dimensions 128."
1328 )
1329 if getattr(self.module, "allow_neg_eigval", False):
1330 raise NotImplementedError(
1331 "Kimi K3 CP does not support allow_neg_eigval=True."
1332 )
1333 if not (hasattr(self.module, "f_proj") or _is_official_kimi_layer(self.module)):
1334 raise TypeError("Kimi K3 layer is missing its forget-gate projection.")
1335 if not (
1336 hasattr(self.module, "g_proj")
1337 or all(hasattr(self.module, name) for name in ("g_a_proj", "g_b_proj"))
1338 ):
1339 raise TypeError("Kimi K3 layer is missing its output-gate projection.")
1340 _validate_layer_short_conv(self.module, "state")
1342 def forward( # pylint: disable=arguments-renamed
1343 self,
1344 hidden_states: torch.Tensor,
1345 attention_mask: Optional[torch.Tensor] = None,
1346 past_key_values: Optional[Any] = None,
1347 use_cache: bool = False,
1348 output_attentions: bool = False,
1349 **kwargs: Any,
1350 ) -> tuple[torch.Tensor, None, None]:
1351 """Run projections, full-CP halos, the configured CP core, and output transforms.
1353 Args:
1354 hidden_states: Local sequence shard [B,T,hidden_size].
1355 attention_mask: Optional dense mask; padding is unsupported.
1356 past_key_values: Recurrent inference cache, currently unsupported.
1357 use_cache: Must be False for this training adapter.
1358 output_attentions: Must be False; KDA does not materialize attention weights.
1359 kwargs: Optional global sample offsets or prebuilt packed_sequence_info.
1360 """
1361 if not self.module.training:
1362 raise NotImplementedError("Kimi K3 CP currently supports training only.")
1363 if past_key_values is not None or use_cache:
1364 raise NotImplementedError("Kimi K3 CP does not support recurrent cache.")
1365 if output_attentions:
1366 raise NotImplementedError("KDA does not materialize attention weights.")
1367 if attention_mask is not None and not bool(attention_mask.bool().all().item()):
1368 raise NotImplementedError("Padded KDA CP batches are not implemented.")
1370 hidden_states, sequence_info, is_flat_input = _prepare_layer_input(
1371 hidden_states, kwargs, self.cp_rank, self.cp_size)
1372 base = self.module
1373 batch_size, local_seq_len, _ = hidden_states.shape
1374 query = base.q_proj(hidden_states)
1375 key = base.k_proj(hidden_states)
1376 value = base.v_proj(hidden_states)
1377 if _uses_short_conv(base):
1378 query, key, value = _causal_short_convs_with_cp_halo( # pylint: disable=unbalanced-tuple-unpacking
1379 (query, key, value),
1380 (base.q_conv1d, base.k_conv1d, base.v_conv1d),
1381 self.cp_mesh,
1382 self.cp_rank,
1383 self.cp_size,
1384 sequence_info=sequence_info,
1385 backend="ascendc" if self.backend == "ascendc" else "eager",
1386 )
1387 else:
1388 query, key, value = (F.silu(tensor) for tensor in (query, key, value))
1390 query = query.reshape(
1391 batch_size,
1392 local_seq_len,
1393 base.num_heads,
1394 base.head_k_dim,
1395 )
1396 key = key.reshape(
1397 batch_size,
1398 local_seq_len,
1399 base.num_heads,
1400 base.head_k_dim,
1401 )
1402 value = value.reshape(
1403 batch_size,
1404 local_seq_len,
1405 _num_value_heads(base),
1406 _value_head_dim(base),
1407 )
1408 gate = _forget_gate_projection(base, hidden_states).reshape(
1409 batch_size,
1410 local_seq_len,
1411 _num_value_heads(base),
1412 base.head_k_dim,
1413 )
1414 beta = base.b_proj(hidden_states).reshape(
1415 batch_size,
1416 local_seq_len,
1417 _num_value_heads(base),
1418 )
1419 output_gate = _output_gate_projection(base, hidden_states).reshape(
1420 batch_size,
1421 local_seq_len,
1422 _num_value_heads(base),
1423 _value_head_dim(base),
1424 )
1426 output = self.core(
1427 query,
1428 key,
1429 value,
1430 gate,
1431 beta,
1432 a_log=base.A_log,
1433 dt_bias=base.dt_bias,
1434 sequence_info=sequence_info,
1435 )
1436 output = base.o_norm(output, output_gate)
1437 output = output.reshape(
1438 batch_size, local_seq_len, _value_projection_dim(base)
1439 )
1440 output = base.o_proj(output)
1441 return _format_layer_output(base, output.squeeze(0) if is_flat_input else output)
1444class KimiDeltaAttentionCoreCP(nn.Module):
1445 """Compose optional Ulysses redistribution with local KDA and state exchange.
1447 Inputs are projected Q/K/V/gate/beta after ShortConv. The layer adapter
1448 handles projections and full-sequence halo neighbors before calling here.
1449 """
1451 def __init__(self, device_mesh: DeviceMesh, *, ulysses_degree: int = 1,
1452 chunk_size: int = 64, lower_bound: float = -5.0, safe_gate: bool = True,
1453 backend: str = "triton", state_cp_method: str = "p2p",
1454 allgather_group_size: int = 1) -> None:
1455 """Build all communication groups once, before forward or checkpoint replay."""
1456 super().__init__()
1457 device_mesh = _ensure_1d(device_mesh)
1458 if state_cp_method not in ("p2p", "allgather", "grouped_allgather_p2p", "recursive_doubling"):
1459 raise ValueError("KDA requires a supported state_cp_method.")
1460 if backend.lower() == "eager" and state_cp_method != "p2p":
1461 raise ValueError("Eager KDA state exchange supports P2P only.")
1462 _validate_cp_degrees(device_mesh.size(), ulysses_degree, allgather_group_size)
1463 if state_cp_method != "grouped_allgather_p2p" and allgather_group_size != 1:
1464 raise ValueError("allgather_group_size is only valid for grouped_allgather_p2p.")
1465 if lower_bound >= 0:
1466 raise ValueError("lower_bound must be negative.")
1467 self.cp_size = device_mesh.size()
1468 self.cp_rank = device_mesh.get_local_rank()
1469 self.ulysses_degree = ulysses_degree
1470 self.ulysses_mesh = None
1471 self.ulysses_rank = 0
1472 state_mesh = device_mesh
1473 if ulysses_degree > 1:
1474 hybrid = split_kda_mesh(device_mesh, ulysses_degree, ("kda_state", "kda_ulysses"))
1475 self.ulysses_mesh = hybrid["kda_ulysses"]
1476 self.ulysses_rank = self.ulysses_mesh.get_local_rank()
1477 state_mesh = hybrid["kda_state"]
1478 self.state = _KDAStateCP(
1479 state_mesh, chunk_size=chunk_size, lower_bound=lower_bound, safe_gate=safe_gate,
1480 backend=backend, state_cp_method=state_cp_method, allgather_group_size=allgather_group_size)
1482 def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
1483 gate: torch.Tensor, beta: torch.Tensor, *, a_log: torch.Tensor,
1484 dt_bias: torch.Tensor, scale: float | None = None,
1485 sequence_info: Optional[PackedSequenceInfo] = None) -> torch.Tensor:
1486 """Run state CP with optional token/head redistribution.
1488 Args:
1489 query: Projected query tensor [B,T,H,K].
1490 key: Matching local key tensor.
1491 value: Projected value tensor [B,T,HV,V].
1492 gate: Per-token gate logits [B,T,HV,K].
1493 beta: Per-token beta logits [B,T,HV].
1494 a_log: Learned decay parameter with one entry per value head.
1495 dt_bias: Learned gate bias with HV times K entries.
1496 scale: Attention scaling factor; None selects the local default.
1498 Returns:
1499 Output restored to this rank's original token/head layout.
1500 """
1501 is_flat_input = query.ndim == 3
1502 if is_flat_input:
1503 output = self.forward(*(tensor.unsqueeze(0) for tensor in (query, key, value, gate, beta)),
1504 a_log=a_log, dt_bias=dt_bias, scale=scale, sequence_info=sequence_info)
1505 return output.squeeze(0)
1506 if sequence_info is not None:
1507 if query.shape[0] != 1:
1508 raise ValueError("Packed KDA core requires batch dimension one.")
1509 resolve_packed_sequence({"packed_sequence_info": sequence_info}, query.shape[1],
1510 cp_rank=self.cp_rank, cp_size=self.cp_size)
1511 if self.ulysses_mesh is None:
1512 return self.state(query, key, value, gate, beta, a_log=a_log, dt_bias=dt_bias, scale=scale,
1513 sequence_info=sequence_info)
1514 _validate_projected_inputs(self.ulysses_degree, query, key, value, gate, beta, a_log, dt_bias)
1515 heads, key_dim = value.shape[2], query.shape[-1]
1516 inputs = self._exchange_inputs((query, key, value, gate, beta))
1517 local_a = _slice_local_heads(a_log.reshape(heads), self.ulysses_rank, self.ulysses_degree)
1518 local_bias = _slice_local_heads(dt_bias.reshape(heads, key_dim), self.ulysses_rank, self.ulysses_degree)
1519 exchanged_info = None if sequence_info is None else sequence_info.ulysses(
1520 self.ulysses_degree, self.ulysses_rank)
1521 output = self.state(*inputs, a_log=local_a, dt_bias=local_bias, scale=scale, sequence_info=exchanged_info)
1522 return _differentiable_all_to_all_shard(output, self.ulysses_mesh, split_dim=1, concat_dim=2)
1524 def _exchange_inputs(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
1525 """Redistribute projected inputs using the established synchronous schedule."""
1526 return tuple(_differentiable_all_to_all_shard(tensor, self.ulysses_mesh, split_dim=2, concat_dim=1)
1527 for tensor in inputs)
1530class KimiDeltaAttentionLayerP2PCP(KimiDeltaAttentionLayerCP):
1531 """Preserve the public layer-P2P type while sharing layer execution.
1533 Inherited defaults select P2P with no Ulysses. The subclass retains the old
1534 import and type name; new hybrid configurations use KimiDeltaAttentionLayerCP.
1535 """