Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / kimi_delta_attention.py: 25%
489 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"""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.core.context_parallel.context_parallel import _ensure_1d
36from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
37from hyper_parallel.core.dtensor.dtensor import DTensor
38from hyper_parallel.components.modules.kimi_delta_attention import (
39 chunk_kda,
40 torch_apply_kda_state_summary,
41 torch_kda_state_summary,
42)
43from hyper_parallel.core.utils import communication
44from hyper_parallel.distributed.context_parallel.kimi_delta_attention_state import (
45 build_kda_state_cp,
46 split_kda_mesh,
47)
50_KDA_BACKENDS = frozenset({"eager", "triton"})
53def _run_local_kda(
54 query: torch.Tensor,
55 key: torch.Tensor,
56 value: torch.Tensor,
57 gate: torch.Tensor,
58 beta: torch.Tensor,
59 *,
60 a_log: torch.Tensor,
61 dt_bias: torch.Tensor,
62 scale: Optional[float],
63 lower_bound: float,
64 chunk_size: int,
65 safe_gate: bool,
66 backend: str,
67) -> torch.Tensor:
68 """Run the explicitly selected local KDA implementation."""
69 if backend == "triton" and not safe_gate:
70 raise RuntimeError("The fused KDA backend requires the lower-bounded gate.")
71 output, _ = chunk_kda(
72 query,
73 key,
74 value,
75 gate,
76 beta,
77 a_log=a_log,
78 dt_bias=dt_bias,
79 scale=scale,
80 lower_bound=lower_bound,
81 chunk_size=chunk_size,
82 initial_state=None,
83 output_final_state=False,
84 use_qk_l2norm_in_kernel=True,
85 use_gate_in_kernel=True,
86 use_beta_sigmoid_in_kernel=True,
87 backend=backend,
88 )
89 return output
92def _is_official_kimi_layer(module: nn.Module) -> bool:
93 """Whether ``module`` follows the official Kimi K3 projection contract."""
94 return hasattr(module, "f_a_proj") and hasattr(module, "f_b_proj")
97def _num_value_heads(module: nn.Module) -> int:
98 """Return the number of KDA value heads for either supported layer API."""
99 if hasattr(module, "num_v_heads"):
100 return int(module.num_v_heads)
101 return int(module.num_heads)
104def _value_head_dim(module: nn.Module) -> int:
105 """Return the KDA value head width for either supported layer API."""
106 if hasattr(module, "head_v_dim"):
107 return int(module.head_v_dim)
108 return int(module.head_dim)
111def _key_projection_dim(module: nn.Module) -> int:
112 """Return the flattened Q/K projection width."""
113 if hasattr(module, "key_dim"):
114 return int(module.key_dim)
115 return int(module.num_heads * module.head_k_dim)
118def _value_projection_dim(module: nn.Module) -> int:
119 """Return the flattened value projection width."""
120 if hasattr(module, "value_dim"):
121 return int(module.value_dim)
122 return int(_num_value_heads(module) * _value_head_dim(module))
125def _gate_lower_bound(module: nn.Module) -> Optional[float]:
126 """Return the lower bound used by the safe KDA gate parameterization."""
127 return getattr(module, "lower_bound", getattr(module, "gate_lower_bound", None))
130def _uses_safe_gate(module: nn.Module) -> bool:
131 """Return whether the layer enables the safe lower-bounded gate."""
132 return bool(getattr(module, "safe_gate", _gate_lower_bound(module) is not None))
135def _uses_short_conv(module: nn.Module) -> bool:
136 """Return whether Q/K/V pass through the layer's ShortConvolution modules."""
137 return bool(
138 getattr(
139 module,
140 "use_short_conv",
141 all(hasattr(module, name) for name in ("q_conv1d", "k_conv1d", "v_conv1d")),
142 )
143 )
146def _forget_gate_projection(module: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
147 """Run the module-specific state-decay gate projection."""
148 if hasattr(module, "f_proj"):
149 return module.f_proj(hidden_states)
150 return module.f_b_proj(module.f_a_proj(hidden_states))
153def _output_gate_projection(module: nn.Module, hidden_states: torch.Tensor) -> torch.Tensor:
154 """Run the module-specific output-gate projection."""
155 if hasattr(module, "g_proj"):
156 return module.g_proj(hidden_states)
157 return module.g_b_proj(module.g_a_proj(hidden_states))
160def _format_layer_output(module: nn.Module, output: torch.Tensor):
161 """Preserve the original layer's tensor or three-tuple output contract."""
162 if _is_official_kimi_layer(module):
163 return output
164 return output, None, None
167def _global_peer_rank(cp_mesh: DeviceMesh, local_rank: int) -> int:
168 """Map one CP-local rank index to its global distributed rank."""
169 return int(cp_mesh.rank_list[local_rank])
172def _local_cp_tensor(tensor: torch.Tensor) -> torch.Tensor:
173 """Return the local sequence shard of a CP input tensor."""
174 if isinstance(tensor, DTensor):
175 return tensor.to_local()
176 return tensor
179def _all_to_all_previous_rank_halo(
180 tail: torch.Tensor,
181 cp_mesh: DeviceMesh,
182 cp_rank: int,
183 cp_size: int,
184) -> torch.Tensor:
185 """Send a causal-convolution halo only to the next CP rank."""
186 if cp_size == 1:
187 return torch.zeros_like(tail)
189 cp_group = cp_mesh.get_group()
190 group_ranks = tuple(int(rank) for rank in dist.get_process_group_ranks(cp_group))
191 rank_list = tuple(int(rank) for rank in cp_mesh.rank_list)
192 rank_to_group_index = {rank: index for index, rank in enumerate(group_ranks)}
193 halo_width = tail.shape[1]
195 input_splits = [0] * cp_size
196 exchange_input = tail.permute(1, 0, 2).contiguous()
197 if cp_rank < cp_size - 1:
198 input_splits[rank_to_group_index[rank_list[cp_rank + 1]]] = halo_width
199 else:
200 exchange_input = exchange_input[:0]
202 output_splits = [0] * cp_size
203 if cp_rank > 0:
204 output_splits[rank_to_group_index[rank_list[cp_rank - 1]]] = halo_width
206 exchange_output = communication.differentiable_all_to_all_single(
207 exchange_input,
208 input_splits,
209 output_splits,
210 group=cp_group,
211 )
212 if cp_rank == 0:
213 return torch.zeros_like(tail) + exchange_output.sum().to(tail.dtype) * 0
214 return exchange_output.permute(1, 0, 2).contiguous()
217def _run_causal_short_conv(
218 tensor: torch.Tensor,
219 convolution: nn.Conv1d,
220 halo: Optional[torch.Tensor] = None,
221) -> torch.Tensor:
222 """Apply one local ShortConv, optionally with a preceding-rank halo."""
223 if halo is None:
224 conv_input = tensor.transpose(1, 2)
225 padding = convolution.padding
226 else:
227 conv_input = torch.cat((halo, tensor), dim=1).transpose(1, 2)
228 padding = 0
229 output = F.conv1d( # pylint: disable=not-callable
230 input=conv_input,
231 weight=convolution.weight,
232 bias=convolution.bias,
233 stride=convolution.stride,
234 padding=padding,
235 dilation=convolution.dilation,
236 groups=convolution.groups,
237 )
238 if halo is None:
239 output = output[:, :, : tensor.shape[1]]
240 return F.silu(output).transpose(1, 2)
243def _causal_short_convs_with_cp_halo(
244 projected: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
245 convolutions: tuple[nn.Conv1d, nn.Conv1d, nn.Conv1d],
246 cp_mesh: DeviceMesh,
247 cp_rank: int,
248 cp_size: int,
249) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
250 """Run three causal depthwise ShortConvs with one packed halo exchange."""
251 halo_widths = {
252 (conv.kernel_size[0] - 1) * conv.dilation[0] for conv in convolutions
253 }
254 if len(halo_widths) != 1:
255 raise ValueError("KDA P2P requires Q/K/V ShortConv halo widths to match.")
256 halo_width = halo_widths.pop()
257 if halo_width == 0 or cp_size == 1:
258 return tuple(
259 _run_causal_short_conv(tensor, convolution)
260 for tensor, convolution in zip(projected, convolutions)
261 )
262 if projected[0].shape[1] < halo_width:
263 raise ValueError(
264 "KDA P2P ShortConv requires local_seq_len >= halo width, got "
265 f"{projected[0].shape[1]} < {halo_width}."
266 )
268 channel_sizes = [tensor.shape[-1] for tensor in projected]
269 packed_tail = torch.cat(
270 [tensor[:, -halo_width:, :] for tensor in projected],
271 dim=-1,
272 ).contiguous()
273 packed_halo = _all_to_all_previous_rank_halo(
274 packed_tail,
275 cp_mesh,
276 cp_rank,
277 cp_size,
278 )
279 halos = torch.split(packed_halo, channel_sizes, dim=-1)
280 return tuple(
281 _run_causal_short_conv(tensor, convolution, halo)
282 for tensor, halo, convolution in zip(projected, halos, convolutions)
283 )
286def _slice_local_heads(
287 tensor: torch.Tensor,
288 cp_rank: int,
289 cp_size: int,
290) -> torch.Tensor:
291 """Return the current CP rank's contiguous head slice."""
292 if tensor.shape[0] % cp_size != 0:
293 raise ValueError(
294 f"KDA head dimension {tensor.shape[0]} must be divisible by "
295 f"cp_size {cp_size}."
296 )
297 local_heads = tensor.shape[0] // cp_size
298 return tensor.narrow(0, cp_rank * local_heads, local_heads)
301def _differentiable_all_to_all_shard(
302 tensor: torch.Tensor,
303 device_mesh: DeviceMesh,
304 *,
305 split_dim: int,
306 concat_dim: int,
307) -> torch.Tensor:
308 """Redistribute a local tensor from concat-dim shard to split-dim shard."""
309 split_count = device_mesh.size()
310 if split_count == 1:
311 return tensor
313 original_shape = tuple(tensor.shape)
314 if original_shape[split_dim] % split_count != 0:
315 raise ValueError(
316 f"KDA all-to-all split dimension {split_dim} with size "
317 f"{original_shape[split_dim]} must be divisible by cp_size "
318 f"{split_count}."
319 )
320 split_size = original_shape[split_dim] // split_count
321 final_shape = list(original_shape)
322 final_shape[split_dim] = split_size
323 final_shape[concat_dim] *= split_count
325 reshape_dims = list(original_shape)
326 reshape_dims[split_dim] = split_count
327 reshape_dims.insert(split_dim + 1, split_size)
328 permutation = list(range(len(reshape_dims)))
329 permutation.remove(split_dim)
330 permutation.insert(0, split_dim)
331 all_to_all_input = tensor.reshape(reshape_dims).permute(permutation).contiguous()
332 all_to_all_input = all_to_all_input.flatten(0, 1)
334 split_length = all_to_all_input.shape[0] // split_count
335 splits = [split_length] * split_count
336 output = communication.differentiable_all_to_all_single(
337 all_to_all_input,
338 splits,
339 splits,
340 group=device_mesh.get_group(),
341 )
343 output_reshape = list(output.shape)
344 output_reshape[0] = split_count
345 output_reshape.insert(1, output.shape[0] // split_count)
346 permutation = list(range(len(output_reshape)))
347 rank_dimension = permutation.pop(0)
348 permutation.insert(concat_dim, rank_dimension)
349 output = output.reshape(output_reshape).permute(permutation).contiguous()
351 output_shape = list(output.shape)
352 output_shape[concat_dim] *= output_shape[concat_dim + 1]
353 output_shape.pop(concat_dim + 1)
354 return output.reshape(output_shape).view(final_shape)
357class _RecvKDAInitialState(torch.autograd.Function):
358 """Receive a KDA initial state and return its gradient to the sender."""
360 @staticmethod
361 def forward( # pylint: disable=arguments-differ
362 ctx: Any,
363 anchor: torch.Tensor,
364 cp_group: Any,
365 prev_rank: int,
366 state_shape: tuple[int, ...],
367 ) -> torch.Tensor:
368 """Receive the initial recurrent state from the previous CP rank.
370 Args:
371 ctx: Autograd context holding tensors and communication metadata.
372 anchor: Tensor supplying the receive device and autograd dependency.
373 cp_group: Precreated chronological CP process group.
374 prev_rank: Global rank of the preceding sequence partition.
375 state_shape: Shape of the FP32 recurrent state.
376 """
377 state = torch.empty(state_shape, device=anchor.device, dtype=torch.float32)
378 dist.recv(state, src=prev_rank, group=cp_group)
379 ctx.cp_group = cp_group
380 ctx.prev_rank = prev_rank
381 return state
383 @staticmethod
384 def backward(
385 ctx: Any,
386 grad_state: Optional[torch.Tensor],
387 ) -> tuple[None, None, None, None]:
388 """Send the accumulated initial-state gradient to the previous rank.
390 Args:
391 ctx: Autograd context holding tensors and communication metadata.
392 grad_state: Gradient of the received initial state.
393 """
394 if grad_state is None:
395 raise RuntimeError("KDA P2P backward is missing the initial-state gradient.")
396 dist.send(grad_state.contiguous(), dst=ctx.prev_rank, group=ctx.cp_group)
397 return None, None, None, None
400class _SendKDAFinalState(torch.autograd.Function):
401 """Send a KDA final state and receive its gradient during backward."""
403 @staticmethod
404 def forward( # pylint: disable=arguments-differ
405 ctx: Any,
406 final_state: torch.Tensor,
407 cp_group: Any,
408 next_rank: int,
409 ) -> torch.Tensor:
410 """Send the final recurrent state to the next CP rank.
412 Args:
413 ctx: Autograd context holding tensors and communication metadata.
414 final_state: Outgoing recurrent state.
415 cp_group: Precreated chronological CP process group.
416 next_rank: Global rank of the following sequence partition.
417 """
418 dist.send(final_state.contiguous(), dst=next_rank, group=cp_group)
419 ctx.cp_group = cp_group
420 ctx.next_rank = next_rank
421 ctx.state_shape = tuple(final_state.shape)
422 ctx.state_dtype = final_state.dtype
423 return final_state.new_zeros(())
425 @staticmethod
426 def backward(
427 ctx: Any,
428 grad_token: torch.Tensor,
429 ) -> tuple[torch.Tensor, None, None]:
430 """Receive the final-state gradient from the next rank.
432 Args:
433 ctx: Autograd context holding tensors and communication metadata.
434 grad_token: Gradient token supplying the receive device.
435 """
436 grad_state = torch.empty(
437 ctx.state_shape,
438 device=grad_token.device,
439 dtype=ctx.state_dtype,
440 )
441 dist.recv(grad_state, src=ctx.next_rank, group=ctx.cp_group)
442 return grad_state, None, None
445def _eager_kda_state_p2p(
446 query: torch.Tensor,
447 key: torch.Tensor,
448 value: torch.Tensor,
449 gate: torch.Tensor,
450 beta: torch.Tensor,
451 *,
452 a_log: torch.Tensor,
453 dt_bias: torch.Tensor,
454 scale: Optional[float],
455 lower_bound: float,
456 chunk_size: int,
457 cp_group: Any,
458 prev_rank: int,
459 next_rank: int,
460 cp_rank: int,
461 cp_size: int,
462) -> torch.Tensor:
463 """Run the differentiable Torch reference for affine state-P2P KDA."""
464 if cp_size == 1:
465 return _run_local_kda(
466 query,
467 key,
468 value,
469 gate,
470 beta,
471 a_log=a_log,
472 dt_bias=dt_bias,
473 scale=scale,
474 lower_bound=lower_bound,
475 chunk_size=chunk_size,
476 safe_gate=True,
477 backend="eager",
478 )
480 state_ext = None
481 transition = None
482 if cp_rank < cp_size - 1:
483 state_ext, transition = torch_kda_state_summary(
484 query,
485 key,
486 value,
487 gate,
488 beta,
489 a_log=a_log,
490 dt_bias=dt_bias,
491 scale=scale,
492 use_qk_l2norm_in_kernel=True,
493 use_gate_in_kernel=True,
494 use_beta_sigmoid_in_kernel=True,
495 lower_bound=lower_bound,
496 chunk_size=chunk_size,
497 )
499 initial_state = None
500 if cp_rank > 0:
501 state_shape = (
502 query.shape[0],
503 value.shape[2],
504 query.shape[-1],
505 value.shape[-1],
506 )
507 initial_state = _RecvKDAInitialState.apply(
508 query,
509 cp_group,
510 prev_rank,
511 state_shape,
512 )
514 send_token = None
515 if cp_rank < cp_size - 1:
516 final_state = torch_apply_kda_state_summary(
517 state_ext,
518 transition,
519 initial_state,
520 )
521 send_token = _SendKDAFinalState.apply(final_state, cp_group, next_rank)
523 output, _ = chunk_kda(
524 query,
525 key,
526 value,
527 gate,
528 beta,
529 a_log=a_log,
530 dt_bias=dt_bias,
531 scale=scale,
532 initial_state=initial_state,
533 output_final_state=False,
534 use_qk_l2norm_in_kernel=True,
535 use_gate_in_kernel=True,
536 use_beta_sigmoid_in_kernel=True,
537 lower_bound=lower_bound,
538 chunk_size=chunk_size,
539 backend="eager",
540 )
541 if send_token is not None:
542 output = output + send_token.to(output.dtype) * 0
543 return output
546def _validate_projected_inputs(
547 ulysses_degree: int,
548 query: torch.Tensor,
549 key: torch.Tensor,
550 value: torch.Tensor,
551 gate: torch.Tensor,
552 beta: torch.Tensor,
553 a_log: torch.Tensor,
554 dt_bias: torch.Tensor,
555) -> None:
556 """Validate the projected-tensor contract and Ulysses divisibility."""
557 if (query.dim(), key.dim(), value.dim(), gate.dim()) != (4, 4, 4, 4):
558 raise ValueError("query, key, value, and gate must be rank-4 tensors.")
559 if beta.dim() != 3:
560 raise ValueError("beta must be a rank-3 tensor.")
561 if query.shape != key.shape:
562 raise ValueError("query and key must have identical shapes.")
564 batch_size, local_seq_len, num_k_heads, k_head_dim = query.shape
565 expected_prefix = (batch_size, local_seq_len)
566 if value.shape[:2] != expected_prefix:
567 raise ValueError("value batch and sequence dimensions must match query.")
568 num_v_heads = value.shape[2]
569 if gate.shape != (batch_size, local_seq_len, num_v_heads, k_head_dim):
570 raise ValueError(
571 "gate must have shape [B, S_local, num_v_heads, k_head_dim]."
572 )
573 if beta.shape != (batch_size, local_seq_len, num_v_heads):
574 raise ValueError("beta must have shape [B, S_local, num_v_heads].")
575 if num_v_heads % num_k_heads != 0:
576 raise ValueError("num_v_heads must be divisible by num_k_heads.")
577 if num_k_heads % ulysses_degree != 0:
578 raise ValueError(
579 f"KDA num_k_heads ({num_k_heads}) must be divisible by "
580 f"cp_size ({ulysses_degree}) for Ulysses CP."
581 )
582 if num_v_heads % ulysses_degree != 0:
583 raise ValueError(
584 f"KDA num_v_heads ({num_v_heads}) must be divisible by "
585 f"cp_size ({ulysses_degree}) for Ulysses CP."
586 )
587 if a_log.numel() != num_v_heads:
588 raise ValueError(
589 f"a_log must contain {num_v_heads} elements, got {a_log.numel()}."
590 )
591 if dt_bias.numel() != num_v_heads * k_head_dim:
592 raise ValueError(
593 f"dt_bias must contain {num_v_heads * k_head_dim} elements, "
594 f"got {dt_bias.numel()}."
595 )
598class KimiDeltaAttentionUlyssesCP(nn.Module):
599 """Run KDA on sequence-sharded projected tensors with Ulysses all-to-all.
601 The input and output layout is token-first with a local sequence shard:
603 ``[B, S/P, H, D] -> [B, S, H/P, D] -> [B, S/P, H, D]``.
605 This first-stage interface intentionally starts after projections and
606 ShortConv and ends before the output norm/gate/projection. It is therefore
607 independent of one particular Kimi K3 model class. Training uses a zero
608 initial recurrent state and does not materialize a final state.
609 """
611 def __init__(
612 self,
613 device_mesh: DeviceMesh,
614 *,
615 chunk_size: int = 64,
616 lower_bound: float = -5.0,
617 safe_gate: bool = True,
618 backend: str = "eager",
619 ) -> None:
620 """Initialize the KDA Ulysses executor for one 1-D CP mesh."""
621 super().__init__()
622 if chunk_size <= 0:
623 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
624 if lower_bound >= 0:
625 raise ValueError(f"lower_bound must be negative, got {lower_bound}.")
626 backend = backend.lower()
627 if backend not in _KDA_BACKENDS:
628 raise ValueError(
629 f"unsupported KDA backend {backend!r}; "
630 f"expected one of {sorted(_KDA_BACKENDS)}."
631 )
632 self.cp_mesh = _ensure_1d(device_mesh)
633 self.cp_size = self.cp_mesh.size()
634 self.cp_rank = self.cp_mesh.get_local_rank()
635 self.chunk_size = chunk_size
636 self.lower_bound = lower_bound
637 self.safe_gate = safe_gate
638 self.backend = backend
639 self.seq_dim = 1
640 self.head_dim = 2
642 def _validate_inputs(
643 self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
644 gate: torch.Tensor, beta: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor,
645 ) -> None:
646 """Validate the projected contract for this Ulysses degree."""
647 _validate_projected_inputs(self.cp_size, query, key, value, gate, beta, a_log, dt_bias)
649 def _seq_to_head(self, tensor: torch.Tensor) -> torch.Tensor:
650 """Redistribute a local sequence shard into a local head shard."""
651 return _differentiable_all_to_all_shard(
652 tensor,
653 self.cp_mesh,
654 split_dim=self.head_dim,
655 concat_dim=self.seq_dim,
656 )
658 def _head_to_seq(self, tensor: torch.Tensor) -> torch.Tensor:
659 """Redistribute a local head shard back into a local sequence shard."""
660 return _differentiable_all_to_all_shard(
661 tensor,
662 self.cp_mesh,
663 split_dim=self.seq_dim,
664 concat_dim=self.head_dim,
665 )
667 def forward(
668 self,
669 query: torch.Tensor,
670 key: torch.Tensor,
671 value: torch.Tensor,
672 gate: torch.Tensor,
673 beta: torch.Tensor,
674 *,
675 a_log: torch.Tensor,
676 dt_bias: torch.Tensor,
677 scale: Optional[float] = None,
678 ) -> torch.Tensor:
679 """Run chunkwise KDA using full sequence and local heads.
681 Args:
682 query: Local query tensor [B,T,H,K].
683 key: Matching local key tensor.
684 value: Local value tensor [B,T,HV,V].
685 gate: Per-token gate logits [B,T,HV,K].
686 beta: Per-token beta logits [B,T,HV].
687 a_log: Decay parameter with one value per value head.
688 dt_bias: Learned bias with one value per gate channel.
689 scale: Optional attention scaling factor.
690 """
691 self._validate_inputs(query, key, value, gate, beta, a_log, dt_bias)
692 num_v_heads = value.shape[2]
693 k_head_dim = query.shape[-1]
695 query = self._seq_to_head(query)
696 key = self._seq_to_head(key)
697 value = self._seq_to_head(value)
698 gate = self._seq_to_head(gate)
699 beta = self._seq_to_head(beta)
701 local_a_log = _slice_local_heads(
702 a_log.reshape(num_v_heads),
703 self.cp_rank,
704 self.cp_size,
705 )
706 local_dt_bias = _slice_local_heads(
707 dt_bias.reshape(num_v_heads, k_head_dim),
708 self.cp_rank,
709 self.cp_size,
710 )
711 output = _run_local_kda(
712 query,
713 key,
714 value,
715 gate,
716 beta,
717 a_log=local_a_log,
718 dt_bias=local_dt_bias,
719 scale=scale,
720 lower_bound=self.lower_bound,
721 chunk_size=self.chunk_size,
722 safe_gate=self.safe_gate,
723 backend=self.backend,
724 )
725 return self._head_to_seq(output)
728class _KDAStateCP(nn.Module):
729 """Run local KDA with the configured P2P or gathered state exchange."""
731 def __init__(
732 self,
733 device_mesh: DeviceMesh,
734 *,
735 chunk_size: int = 64,
736 lower_bound: float = -5.0,
737 safe_gate: bool = True,
738 backend: str = "eager",
739 state_cp_method: str = "p2p",
740 allgather_group_size: int = 1,
741 ) -> None:
742 """Initialize local KDA and its state exchange on one chronological axis."""
743 super().__init__()
744 if chunk_size <= 0:
745 raise ValueError(f"chunk_size must be positive, got {chunk_size}.")
746 backend = backend.lower()
747 if backend not in _KDA_BACKENDS:
748 raise ValueError(
749 f"unsupported KDA backend {backend!r}; "
750 f"expected one of {sorted(_KDA_BACKENDS)}."
751 )
752 self.cp_mesh = _ensure_1d(device_mesh)
753 self.cp_size = self.cp_mesh.size()
754 self.cp_rank = self.cp_mesh.get_local_rank()
755 self.cp_group = self.cp_mesh.get_group()
756 if backend != "triton" and state_cp_method != "p2p":
757 raise ValueError("KDA AllGather state exchange requires backend='triton'.")
758 self.state_cp = build_kda_state_cp(self.cp_mesh, state_cp_method, allgather_group_size)
759 self.prev_rank = _global_peer_rank(
760 self.cp_mesh,
761 max(self.cp_rank - 1, 0),
762 )
763 self.next_rank = _global_peer_rank(
764 self.cp_mesh,
765 min(self.cp_rank + 1, self.cp_size - 1),
766 )
767 self.chunk_size = chunk_size
768 self.lower_bound = lower_bound
769 self.safe_gate = safe_gate
770 self.backend = backend
772 def forward(
773 self,
774 query: torch.Tensor,
775 key: torch.Tensor,
776 value: torch.Tensor,
777 gate: torch.Tensor,
778 beta: torch.Tensor,
779 *,
780 a_log: torch.Tensor,
781 dt_bias: torch.Tensor,
782 scale: Optional[float] = None,
783 ) -> torch.Tensor:
784 """Run one local KDA segment and propagate its recurrent state.
786 Args:
787 query: Local query tensor [B,T,H,K].
788 key: Matching local key tensor.
789 value: Local value tensor [B,T,HV,V].
790 gate: Per-token gate logits [B,T,HV,K].
791 beta: Per-token beta logits [B,T,HV].
792 a_log: Decay parameter with one value per value head.
793 dt_bias: Learned bias with one value per gate channel.
794 scale: Optional attention scaling factor.
795 """
796 if self.cp_size == 1:
797 return _run_local_kda(
798 query,
799 key,
800 value,
801 gate,
802 beta,
803 a_log=a_log,
804 dt_bias=dt_bias,
805 scale=scale,
806 lower_bound=self.lower_bound,
807 chunk_size=self.chunk_size,
808 safe_gate=self.safe_gate,
809 backend=self.backend,
810 )
811 if self.backend == "eager":
812 return _eager_kda_state_p2p(
813 query,
814 key,
815 value,
816 gate,
817 beta,
818 a_log=a_log,
819 dt_bias=dt_bias,
820 scale=scale,
821 lower_bound=self.lower_bound,
822 chunk_size=self.chunk_size,
823 cp_group=self.cp_group,
824 prev_rank=self.prev_rank,
825 next_rank=self.next_rank,
826 cp_rank=self.cp_rank,
827 cp_size=self.cp_size,
828 )
830 from hyper_parallel.components.functional.kimi_delta_attention import ( # pylint: disable=import-outside-toplevel
831 fused_chunk_kda_p2p,
832 )
833 return fused_chunk_kda_p2p(
834 query,
835 key,
836 value,
837 gate,
838 beta,
839 a_log=a_log.float(),
840 dt_bias=dt_bias.float(),
841 cp_group=self.cp_group,
842 prev_rank=self.prev_rank,
843 next_rank=self.next_rank,
844 cp_rank=self.cp_rank,
845 cp_size=self.cp_size,
846 scale=scale,
847 lower_bound=self.lower_bound,
848 chunk_size=self.chunk_size,
849 safe_gate=self.safe_gate,
850 state_cp=self.state_cp,
851 )
854class KimiDeltaAttentionP2PCP(_KDAStateCP):
855 """Preserve the public projected-P2P type while sharing state execution.
857 Inherited defaults select P2P. Keeping a named subclass preserves the old
858 import and type name without duplicating the common implementation.
859 New hybrid configurations use KimiDeltaAttentionCoreCP.
860 """
863def _validate_layer_short_conv(module: nn.Module, mode: str) -> None:
864 """Check the shared depthwise-convolution contract before building execution."""
865 if _uses_short_conv(module):
866 for name in ("q_conv1d", "k_conv1d", "v_conv1d"):
867 if not hasattr(module, name):
868 raise TypeError(f"Kimi K3 layer is missing {name}.")
869 convolution = getattr(module, name)
870 if convolution.stride != (1,) or convolution.dilation != (1,):
871 raise ValueError(
872 f"Kimi K3 {mode} CP supports ShortConv stride=1 and "
873 "dilation=1 only."
874 )
875 if convolution.groups != convolution.in_channels or convolution.out_channels != convolution.in_channels:
876 raise ValueError(f"Kimi K3 {mode} CP expects depthwise ShortConv with equal input/output channels.")
877 if convolution.padding != (convolution.kernel_size[0] - 1,) or convolution.padding_mode != "zeros":
878 raise ValueError(f"Kimi K3 {mode} CP expects causal zero-padded ShortConv.")
879 if getattr(convolution, "activation", None) not in ("silu", "swish"):
880 raise ValueError(
881 f"Kimi K3 {mode} CP expects ShortConv SiLU activation."
882 )
885class KimiDeltaAttentionLayerUlyssesCP(KimiDeltaAttentionUlyssesCP):
886 """Wrap a training-time Kimi K3 attention layer with Ulysses CP.
888 The wrapped module keeps ownership of all model parameters. Its local
889 sequence shard is projected first, then Q/K/V, per-dimension gate logits,
890 and beta logits are packed into one sequence-to-head all-to-all. Each CP
891 rank runs the wrapped layer's ShortConv weights and fused KDA core for its
892 local head shard and the full sequence. The reverse all-to-all restores a
893 local sequence shard before the original output gate, norm, and projection.
895 This first integration targets dense training batches. Inference cache and
896 padded or packed variable-length sequences require separate CP metadata
897 handling and are rejected explicitly.
898 """
900 def __init__(
901 self,
902 module: nn.Module,
903 device_mesh: DeviceMesh,
904 *,
905 chunk_size: int = 64,
906 backend: str = "eager",
907 ) -> None:
908 """Initialize the full-layer adapter without replacing its parameters."""
909 lower_bound = _gate_lower_bound(module)
910 if lower_bound is None:
911 raise ValueError(
912 "Kimi K3 Ulysses CP requires the lower-bounded gate parameterization."
913 )
914 super().__init__(
915 device_mesh,
916 chunk_size=chunk_size,
917 lower_bound=float(lower_bound),
918 safe_gate=_uses_safe_gate(module),
919 backend=backend,
920 )
921 self.module = module
922 self._validate_module()
924 def _validate_module(self) -> None:
925 """Validate the first-stage KimiDeltaAttention module contract."""
926 required_attributes = (
927 "q_proj",
928 "k_proj",
929 "v_proj",
930 "b_proj",
931 "o_norm",
932 "o_proj",
933 "A_log",
934 "dt_bias",
935 "num_heads",
936 "head_k_dim",
937 )
938 missing = [name for name in required_attributes if not hasattr(self.module, name)]
939 if missing:
940 raise TypeError(
941 "Kimi K3 layer is missing required attributes: "
942 + ", ".join(missing)
943 )
944 if self.module.num_heads % self.cp_size != 0:
945 raise ValueError(
946 f"KDA num_heads ({self.module.num_heads}) must be divisible by "
947 f"cp_size ({self.cp_size})."
948 )
949 num_v_heads = _num_value_heads(self.module)
950 if num_v_heads % self.cp_size != 0:
951 raise ValueError(
952 f"KDA num_v_heads ({num_v_heads}) must be divisible "
953 f"by cp_size ({self.cp_size})."
954 )
955 if num_v_heads % self.module.num_heads != 0:
956 raise ValueError("KDA num_v_heads must be divisible by num_heads.")
957 if not (hasattr(self.module, "f_proj") or _is_official_kimi_layer(self.module)):
958 raise TypeError("Kimi K3 layer is missing its forget-gate projection.")
959 if not (
960 hasattr(self.module, "g_proj")
961 or all(hasattr(self.module, name) for name in ("g_a_proj", "g_b_proj"))
962 ):
963 raise TypeError("Kimi K3 layer is missing its output-gate projection.")
964 if getattr(self.module, "allow_neg_eigval", False):
965 raise NotImplementedError(
966 "Kimi K3 Ulysses CP does not yet support allow_neg_eigval=True."
967 )
969 _validate_layer_short_conv(self.module, "Ulysses")
971 def _pack_projected_inputs(
972 self,
973 query: torch.Tensor,
974 key: torch.Tensor,
975 value: torch.Tensor,
976 gate: torch.Tensor,
977 beta: torch.Tensor,
978 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
979 """Pack all projected paths into one sequence-to-head all-to-all."""
980 if self.cp_size == 1:
981 return query, key, value, gate, beta
983 base = self.module
984 split_sizes = (
985 _key_projection_dim(base) // self.cp_size,
986 _key_projection_dim(base) // self.cp_size,
987 _value_projection_dim(base) // self.cp_size,
988 _num_value_heads(base) * base.head_k_dim // self.cp_size,
989 _num_value_heads(base) // self.cp_size,
990 )
991 projected = (query, key, value, gate, beta)
992 chunks = [torch.split(tensor, size, dim=-1) for tensor, size in zip(projected, split_sizes)]
993 rank_major = [
994 torch.cat(parts, dim=-1)
995 for parts in zip(*chunks)
996 ]
997 packed = self._seq_to_head(torch.cat(rank_major, dim=-1).contiguous())
998 return torch.split(packed, split_sizes, dim=-1)
1000 def _local_short_conv(
1001 self,
1002 projected: torch.Tensor,
1003 convolution: nn.Conv1d,
1004 ) -> torch.Tensor:
1005 """Run one wrapped ShortConv on this rank's contiguous channel slice."""
1006 local_channels = projected.shape[-1]
1007 channel_start = self.cp_rank * local_channels
1008 weight = convolution.weight.narrow(0, channel_start, local_channels)
1009 bias = convolution.bias
1010 if bias is not None:
1011 bias = bias.narrow(0, channel_start, local_channels)
1012 output = F.conv1d( # pylint: disable=not-callable
1013 projected.transpose(1, 2),
1014 weight,
1015 bias,
1016 stride=convolution.stride,
1017 padding=convolution.padding,
1018 dilation=convolution.dilation,
1019 groups=local_channels,
1020 )
1021 output = output[:, :, : projected.shape[1]].transpose(1, 2)
1022 return F.silu(output)
1024 def forward( # pylint: disable=arguments-renamed
1025 self,
1026 hidden_states: torch.Tensor,
1027 attention_mask: Optional[torch.Tensor] = None,
1028 past_key_values: Optional[Any] = None,
1029 use_cache: bool = False,
1030 output_attentions: bool = False,
1031 **kwargs: Any,
1032 ) -> tuple[torch.Tensor, None, None]:
1033 """Run the real Kimi layer using KDA and Ulysses CP.
1035 Args:
1036 hidden_states: Local sequence shard [B,T,hidden_size].
1037 attention_mask: Optional dense mask; padding is unsupported.
1038 past_key_values: Recurrent inference cache, currently unsupported.
1039 use_cache: Must be False for this training adapter.
1040 output_attentions: Must be False; KDA does not materialize attention weights.
1041 kwargs: Additional model arguments; packed sequences are unsupported.
1042 """
1043 if not self.module.training:
1044 raise NotImplementedError("Kimi K3 Ulysses CP currently supports training only.")
1045 if past_key_values is not None or use_cache:
1046 raise NotImplementedError("Kimi K3 Ulysses CP does not support recurrent cache.")
1047 if output_attentions:
1048 raise NotImplementedError("KDA does not materialize attention weights.")
1049 if kwargs.get("cu_seqlens") is not None:
1050 raise NotImplementedError("Packed variable-length KDA CP is not implemented.")
1051 if attention_mask is not None and not bool(attention_mask.bool().all().item()):
1052 raise NotImplementedError("Padded KDA CP batches are not implemented.")
1054 hidden_states = _local_cp_tensor(hidden_states)
1055 base = self.module
1056 batch_size, local_seq_len, _ = hidden_states.shape
1057 query = base.q_proj(hidden_states)
1058 key = base.k_proj(hidden_states)
1059 value = base.v_proj(hidden_states)
1060 gate = _forget_gate_projection(base, hidden_states)
1061 beta = base.b_proj(hidden_states)
1062 output_gate = _output_gate_projection(base, hidden_states).reshape(
1063 batch_size,
1064 local_seq_len,
1065 _num_value_heads(base),
1066 _value_head_dim(base),
1067 )
1069 query, key, value, gate, beta = self._pack_projected_inputs(query, key, value, gate, beta)
1070 if _uses_short_conv(base):
1071 query = self._local_short_conv(query, base.q_conv1d)
1072 key = self._local_short_conv(key, base.k_conv1d)
1073 value = self._local_short_conv(value, base.v_conv1d)
1074 else:
1075 query, key, value = (F.silu(tensor) for tensor in (query, key, value))
1077 full_seq_len = query.shape[1]
1078 local_k_heads = base.num_heads // self.cp_size
1079 local_v_heads = _num_value_heads(base) // self.cp_size
1080 query = query.reshape(
1081 batch_size, full_seq_len, local_k_heads, base.head_k_dim
1082 )
1083 key = key.reshape(batch_size, full_seq_len, local_k_heads, base.head_k_dim)
1084 value = value.reshape(
1085 batch_size, full_seq_len, local_v_heads, _value_head_dim(base)
1086 )
1087 gate = gate.reshape(
1088 batch_size, full_seq_len, local_v_heads, base.head_k_dim
1089 )
1090 beta = beta.reshape(batch_size, full_seq_len, local_v_heads)
1092 local_a_log = _slice_local_heads(
1093 base.A_log.reshape(_num_value_heads(base)),
1094 self.cp_rank,
1095 self.cp_size,
1096 )
1097 local_dt_bias = _slice_local_heads(
1098 base.dt_bias.reshape(_num_value_heads(base), base.head_k_dim),
1099 self.cp_rank,
1100 self.cp_size,
1101 )
1102 output = _run_local_kda(
1103 query,
1104 key,
1105 value,
1106 gate,
1107 beta,
1108 a_log=local_a_log,
1109 dt_bias=local_dt_bias,
1110 scale=None,
1111 lower_bound=self.lower_bound,
1112 chunk_size=self.chunk_size,
1113 safe_gate=self.safe_gate,
1114 backend=self.backend,
1115 )
1116 output = self._head_to_seq(output)
1117 output = base.o_norm(output, output_gate)
1118 output = output.reshape(
1119 batch_size, local_seq_len, _value_projection_dim(base)
1120 )
1121 output = base.o_proj(output)
1122 return _format_layer_output(base, output)
1125def _validate_cp_degrees(cp_size: int, ulysses_degree: int, allgather_group_size: int) -> None:
1126 """Validate both split extents before head division or process-group creation."""
1127 if (not isinstance(ulysses_degree, int) or isinstance(ulysses_degree, bool)
1128 or ulysses_degree < 1 or cp_size % ulysses_degree):
1129 raise ValueError("ulysses_degree must be a positive integer dividing the CP size.")
1130 if (not isinstance(allgather_group_size, int) or isinstance(allgather_group_size, bool)
1131 or allgather_group_size < 1 or (cp_size // ulysses_degree) % allgather_group_size):
1132 raise ValueError("allgather_group_size must be a positive integer dividing the state CP size.")
1135class KimiDeltaAttentionLayerCP(nn.Module):
1136 """Project local tokens, exchange convolution halos, and invoke the configured CP core."""
1138 def __init__(
1139 self,
1140 module: nn.Module,
1141 device_mesh: DeviceMesh,
1142 *,
1143 chunk_size: int = 64,
1144 backend: str = "eager",
1145 ulysses_degree: int = 1,
1146 state_cp_method: str = "p2p",
1147 allgather_group_size: int = 1,
1148 ) -> None:
1149 """Keep layer ownership and full-CP halo metadata separate from the core."""
1150 super().__init__()
1151 lower_bound = _gate_lower_bound(module)
1152 if lower_bound is None:
1153 raise ValueError("Kimi K3 CP requires the lower-bounded gate parameterization.")
1154 self.cp_mesh = _ensure_1d(device_mesh)
1155 self.cp_rank, self.cp_size = self.cp_mesh.get_local_rank(), self.cp_mesh.size()
1156 self.module, self.backend = module, backend.lower()
1157 self._validate_module()
1158 _validate_cp_degrees(self.cp_size, ulysses_degree, allgather_group_size)
1159 if module.num_heads % ulysses_degree or _num_value_heads(module) % ulysses_degree:
1160 raise ValueError("KDA query and value heads must both be divisible by ulysses_degree.")
1161 self.core = KimiDeltaAttentionCoreCP(
1162 device_mesh, ulysses_degree=ulysses_degree, chunk_size=chunk_size,
1163 lower_bound=float(lower_bound), safe_gate=_uses_safe_gate(module), backend=backend,
1164 state_cp_method=state_cp_method, allgather_group_size=allgather_group_size)
1166 def _validate_module(self) -> None:
1167 """Validate the KimiDeltaAttention contract used by state CP."""
1168 required_attributes = (
1169 "q_proj",
1170 "k_proj",
1171 "v_proj",
1172 "b_proj",
1173 "o_norm",
1174 "o_proj",
1175 "A_log",
1176 "dt_bias",
1177 "num_heads",
1178 "head_k_dim",
1179 )
1180 missing = [name for name in required_attributes if not hasattr(self.module, name)]
1181 if missing:
1182 raise TypeError(
1183 "Kimi K3 layer is missing required attributes: "
1184 + ", ".join(missing)
1185 )
1186 num_v_heads = _num_value_heads(self.module)
1187 if num_v_heads % self.module.num_heads:
1188 raise ValueError("KDA num_v_heads must be divisible by num_heads.")
1189 if self.backend == "triton" and (
1190 self.module.head_k_dim != 128 or _value_head_dim(self.module) != 128
1191 ):
1192 raise NotImplementedError(
1193 "Kimi K3 CP currently requires key/value head dimensions 128."
1194 )
1195 if getattr(self.module, "allow_neg_eigval", False):
1196 raise NotImplementedError(
1197 "Kimi K3 CP does not support allow_neg_eigval=True."
1198 )
1199 if not (hasattr(self.module, "f_proj") or _is_official_kimi_layer(self.module)):
1200 raise TypeError("Kimi K3 layer is missing its forget-gate projection.")
1201 if not (
1202 hasattr(self.module, "g_proj")
1203 or all(hasattr(self.module, name) for name in ("g_a_proj", "g_b_proj"))
1204 ):
1205 raise TypeError("Kimi K3 layer is missing its output-gate projection.")
1206 _validate_layer_short_conv(self.module, "state")
1208 def forward( # pylint: disable=arguments-renamed
1209 self,
1210 hidden_states: torch.Tensor,
1211 attention_mask: Optional[torch.Tensor] = None,
1212 past_key_values: Optional[Any] = None,
1213 use_cache: bool = False,
1214 output_attentions: bool = False,
1215 **kwargs: Any,
1216 ) -> tuple[torch.Tensor, None, None]:
1217 """Run projections, full-CP halos, the configured CP core, and output transforms.
1219 Args:
1220 hidden_states: Local sequence shard [B,T,hidden_size].
1221 attention_mask: Optional dense mask; padding is unsupported.
1222 past_key_values: Recurrent inference cache, currently unsupported.
1223 use_cache: Must be False for this training adapter.
1224 output_attentions: Must be False; KDA does not materialize attention weights.
1225 kwargs: Additional model arguments; packed sequences are unsupported.
1226 """
1227 if not self.module.training:
1228 raise NotImplementedError("Kimi K3 CP currently supports training only.")
1229 if past_key_values is not None or use_cache:
1230 raise NotImplementedError("Kimi K3 CP does not support recurrent cache.")
1231 if output_attentions:
1232 raise NotImplementedError("KDA does not materialize attention weights.")
1233 if kwargs.get("cu_seqlens") is not None:
1234 raise NotImplementedError("Packed variable-length KDA CP is not implemented.")
1235 if attention_mask is not None and not bool(attention_mask.bool().all().item()):
1236 raise NotImplementedError("Padded KDA CP batches are not implemented.")
1238 hidden_states = _local_cp_tensor(hidden_states)
1239 base = self.module
1240 batch_size, local_seq_len, _ = hidden_states.shape
1241 query = base.q_proj(hidden_states)
1242 key = base.k_proj(hidden_states)
1243 value = base.v_proj(hidden_states)
1244 if _uses_short_conv(base):
1245 query, key, value = _causal_short_convs_with_cp_halo( # pylint: disable=unbalanced-tuple-unpacking
1246 (query, key, value),
1247 (base.q_conv1d, base.k_conv1d, base.v_conv1d),
1248 self.cp_mesh,
1249 self.cp_rank,
1250 self.cp_size,
1251 )
1252 else:
1253 query, key, value = (F.silu(tensor) for tensor in (query, key, value))
1255 query = query.reshape(
1256 batch_size,
1257 local_seq_len,
1258 base.num_heads,
1259 base.head_k_dim,
1260 )
1261 key = key.reshape(
1262 batch_size,
1263 local_seq_len,
1264 base.num_heads,
1265 base.head_k_dim,
1266 )
1267 value = value.reshape(
1268 batch_size,
1269 local_seq_len,
1270 _num_value_heads(base),
1271 _value_head_dim(base),
1272 )
1273 gate = _forget_gate_projection(base, hidden_states).reshape(
1274 batch_size,
1275 local_seq_len,
1276 _num_value_heads(base),
1277 base.head_k_dim,
1278 )
1279 beta = base.b_proj(hidden_states).reshape(
1280 batch_size,
1281 local_seq_len,
1282 _num_value_heads(base),
1283 )
1284 output_gate = _output_gate_projection(base, hidden_states).reshape(
1285 batch_size,
1286 local_seq_len,
1287 _num_value_heads(base),
1288 _value_head_dim(base),
1289 )
1291 output = self.core(
1292 query,
1293 key,
1294 value,
1295 gate,
1296 beta,
1297 a_log=base.A_log,
1298 dt_bias=base.dt_bias,
1299 )
1300 output = base.o_norm(output, output_gate)
1301 output = output.reshape(
1302 batch_size, local_seq_len, _value_projection_dim(base)
1303 )
1304 output = base.o_proj(output)
1305 return _format_layer_output(base, output)
1308class KimiDeltaAttentionCoreCP(nn.Module):
1309 """Compose optional Ulysses redistribution with local KDA and state exchange.
1311 Inputs are projected Q/K/V/gate/beta after ShortConv. The layer adapter
1312 handles projections and full-sequence halo neighbors before calling here.
1313 """
1315 def __init__(self, device_mesh: DeviceMesh, *, ulysses_degree: int = 1,
1316 chunk_size: int = 64, lower_bound: float = -5.0, safe_gate: bool = True,
1317 backend: str = "triton", state_cp_method: str = "p2p",
1318 allgather_group_size: int = 1) -> None:
1319 """Build all communication groups once, before forward or checkpoint replay."""
1320 super().__init__()
1321 device_mesh = _ensure_1d(device_mesh)
1322 if state_cp_method not in ("p2p", "allgather", "grouped_allgather_p2p"):
1323 raise ValueError("KDA requires a supported state_cp_method.")
1324 if backend.lower() != "triton" and state_cp_method != "p2p":
1325 raise ValueError("KDA AllGather states require backend='triton'.")
1326 _validate_cp_degrees(device_mesh.size(), ulysses_degree, allgather_group_size)
1327 if state_cp_method != "grouped_allgather_p2p" and allgather_group_size != 1:
1328 raise ValueError("allgather_group_size is only valid for grouped_allgather_p2p.")
1329 if lower_bound >= 0:
1330 raise ValueError("lower_bound must be negative.")
1331 self.ulysses_degree = ulysses_degree
1332 self.ulysses_mesh = None
1333 self.ulysses_rank = 0
1334 state_mesh = device_mesh
1335 if ulysses_degree > 1:
1336 hybrid = split_kda_mesh(device_mesh, ulysses_degree, ("kda_state", "kda_ulysses"))
1337 self.ulysses_mesh = hybrid["kda_ulysses"]
1338 self.ulysses_rank = self.ulysses_mesh.get_local_rank()
1339 state_mesh = hybrid["kda_state"]
1340 self.state = _KDAStateCP(
1341 state_mesh, chunk_size=chunk_size, lower_bound=lower_bound, safe_gate=safe_gate,
1342 backend=backend, state_cp_method=state_cp_method, allgather_group_size=allgather_group_size)
1344 def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
1345 gate: torch.Tensor, beta: torch.Tensor, *, a_log: torch.Tensor,
1346 dt_bias: torch.Tensor, scale: float | None = None) -> torch.Tensor:
1347 """Run state CP with optional token/head redistribution.
1349 Args:
1350 query: Projected query tensor [B,T,H,K].
1351 key: Matching local key tensor.
1352 value: Projected value tensor [B,T,HV,V].
1353 gate: Per-token gate logits [B,T,HV,K].
1354 beta: Per-token beta logits [B,T,HV].
1355 a_log: Learned decay parameter with one entry per value head.
1356 dt_bias: Learned gate bias with HV times K entries.
1357 scale: Attention scaling factor; None selects the local default.
1359 Returns:
1360 Output restored to this rank's original token/head layout.
1361 """
1362 if self.ulysses_mesh is None:
1363 return self.state(query, key, value, gate, beta, a_log=a_log, dt_bias=dt_bias, scale=scale)
1364 _validate_projected_inputs(self.ulysses_degree, query, key, value, gate, beta, a_log, dt_bias)
1365 heads, key_dim = value.shape[2], query.shape[-1]
1366 inputs = self._exchange_inputs((query, key, value, gate, beta))
1367 local_a = _slice_local_heads(a_log.reshape(heads), self.ulysses_rank, self.ulysses_degree)
1368 local_bias = _slice_local_heads(dt_bias.reshape(heads, key_dim), self.ulysses_rank, self.ulysses_degree)
1369 output = self.state(*inputs, a_log=local_a, dt_bias=local_bias, scale=scale)
1370 return _differentiable_all_to_all_shard(output, self.ulysses_mesh, split_dim=1, concat_dim=2)
1372 def _exchange_inputs(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]:
1373 """Redistribute projected inputs using the established synchronous schedule."""
1374 return tuple(_differentiable_all_to_all_shard(tensor, self.ulysses_mesh, split_dim=2, concat_dim=1)
1375 for tensor in inputs)
1378class KimiDeltaAttentionLayerP2PCP(KimiDeltaAttentionLayerCP):
1379 """Preserve the public layer-P2P type while sharing layer execution.
1381 Inherited defaults select P2P with no Ulysses. The subclass retains the old
1382 import and type name; new hybrid configurations use KimiDeltaAttentionLayerCP.
1383 """