« 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 state exchange through AllGather or grouped AllGather/P2P."""
16# pylint: disable=forbidden-backend-import
17from __future__ import annotations
19from typing import Any
21import torch
22import torch.distributed as dist
24from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
27def _validate_mesh_order(mesh: DeviceMesh) -> None:
28 """Prevent collective slots from being mistaken for chronological positions."""
29 group_ranks = tuple(dist.get_process_group_ranks(mesh.get_group()))
30 if tuple(mesh.rank_list) != group_ranks:
31 raise ValueError("KDA mesh sequence order must match its communication group order.")
34def split_kda_mesh(mesh: DeviceMesh, degree: int, names: tuple[str, str]) -> DeviceMesh:
35 """Split a CP axis without losing the root rank map or sibling DP groups.
37 Args:
38 mesh: One-dimensional chronological CP mesh.
39 degree: Size of the inner split axis.
40 names: Names of the outer and inner axes.
42 Returns:
43 Two-dimensional submesh attached to the original root.
44 """
45 if mesh.ndim != 1 or not isinstance(degree, int) or isinstance(degree, bool):
46 raise ValueError("KDA mesh splitting requires a 1-D mesh and integer degree.")
47 if degree < 1 or mesh.size() % degree:
48 raise ValueError("KDA split degree must be positive and divide the CP size.")
49 _validate_mesh_order(mesh)
50 # Group creation uses WORLD; every rank must derive the same sibling groups.
51 root = mesh._get_root_mesh() # pylint: disable=protected-access
52 if root.size() != dist.get_world_size():
53 raise ValueError("Hybrid KDA requires a CP submesh retaining its WORLD root mesh.")
54 # Like the existing hybrid CP helper, preserve all sibling groups through the root layout.
55 return mesh._unflatten(0, (mesh.size() // degree, degree), names) # pylint: disable=protected-access
58def build_kda_state_cp(mesh: DeviceMesh, protocol: str, allgather_group_size: int) -> object | None:
59 """Construct a state-CP executor holding only communication metadata.
61 Args:
62 mesh: One-dimensional mesh ordered along the state sequence.
63 protocol: p2p, allgather, or grouped_allgather_p2p.
64 allgather_group_size: Consecutive ranks per local gather, only for the grouped protocol.
66 Returns:
67 Gather executor, or None to select the existing P2P path.
68 """
69 if protocol not in ("p2p", "allgather", "grouped_allgather_p2p"):
70 raise ValueError(f"Unknown KDA state_cp_method: {protocol!r}.")
71 if not isinstance(allgather_group_size, int) or isinstance(allgather_group_size, bool) or allgather_group_size < 1:
72 raise ValueError("KDA allgather_group_size must be a positive integer.")
73 if protocol != "grouped_allgather_p2p" and allgather_group_size != 1:
74 raise ValueError("KDA allgather_group_size is only valid for grouped_allgather_p2p.")
75 if protocol == "p2p":
76 return None
77 ranks = tuple(mesh.rank_list)
78 rank, size = mesh.get_local_rank(), mesh.size()
79 if protocol == "allgather":
80 _validate_mesh_order(mesh)
81 return KDAStateAllGather(mesh.get_group(), rank, size)
82 if size % allgather_group_size:
83 raise ValueError("KDA allgather_group_size must divide the state CP size.")
84 if allgather_group_size == 1:
85 return None
86 if allgather_group_size == size:
87 _validate_mesh_order(mesh)
88 return KDAStateAllGather(mesh.get_group(), rank, size)
89 split = split_kda_mesh(mesh, allgather_group_size, ("kda_outer", "kda_group"))
90 return KDAStateAllGatherP2P(
91 mesh.get_group(), split["kda_group"].get_group(), rank, ranks, allgather_group_size)
94def _validate_summary(state: torch.Tensor, matrix: torch.Tensor) -> None:
95 """Validate the dense affine contract before entering communication."""
96 if (state.ndim != 4 or state.shape[-2:] != (128, 128) or matrix.shape != state.shape
97 or state.shape[0] == 0 or state.shape[1] == 0):
98 raise ValueError("KDA state CP requires nonempty matching [B,H,128,128] summaries.")
99 if state.dtype != torch.float32 or matrix.dtype != torch.float32 or state.device != matrix.device:
100 raise ValueError("KDA state CP requires FP32 summaries on the same device.")
101 if state.device.type not in ("cpu", "npu"):
102 raise ValueError("KDA state CP supports CPU protocol checks and Ascend NPU execution only.")
105def merge_affine_summaries(
106 states: torch.Tensor, matrices: torch.Tensor, rank: int, forward: bool,
107) -> torch.Tensor:
108 """Merge the exclusive prefix or transposed suffix in chronological order.
110 Args:
111 states: Affine offsets shaped [P,B,H,128,128].
112 matrices: Matching transfer matrices, possibly views of packed storage.
113 rank: Chronological rank whose input state or final-state gradient is requested.
114 forward: Select forward prefix instead of backward suffix.
116 Returns:
117 Contiguous FP32 recurrent state or state gradient shaped [B,H,128,128].
118 """
119 if states.ndim != 5 or states.shape != matrices.shape or states.shape[0] == 0:
120 raise ValueError("KDA merge requires matching rank-stacked summaries.")
121 size, batch, heads = states.shape[:3]
122 _validate_summary(states[0], matrices[0])
123 if not 0 <= rank < size:
124 raise ValueError("KDA merge rank is outside the summary group.")
125 if any(t.stride(1) != heads * t.stride(2) or t.stride(-1) != 1 for t in (states, matrices)):
126 raise ValueError("KDA merge requires consecutive batch/head slabs and columns.")
127 if states.device.type != "npu":
128 raise ValueError("KDA fused state merge requires Ascend NPU tensors.")
129 count = rank if forward else size - rank - 1
130 output = states.new_empty((batch, heads, 128, 128))
131 if count == 0:
132 return output.zero_()
133 # Triton is optional for imports; production merging requires Ascend.
134 # pylint: disable=import-outside-toplevel
135 from hyper_parallel.components.functional._triton.kimi_delta_attention.cp_merge import (
136 kda_cached_affine_merge_kernel,
137 )
138 # pylint: enable=import-outside-toplevel
140 kda_cached_affine_merge_kernel[(2, batch * heads)](
141 output, states, matrices, count, rank,
142 state_rank_stride=states.stride(0), state_head_stride=states.stride(2), state_row_stride=states.stride(3),
143 matrix_rank_stride=matrices.stride(0), matrix_head_stride=matrices.stride(2),
144 matrix_row_stride=matrices.stride(3),
145 forward=forward, block_v=64, num_warps=2, num_stages=2)
146 return output
149def _gather_maps(state: torch.Tensor, matrix: torch.Tensor, group: Any, size: int) -> torch.Tensor:
150 """Pack validated local summaries for one invocation; the caller owns the gathered storage."""
151 packed = torch.cat((state, matrix), dim=-1)
152 output = packed.new_empty((size * packed.shape[0], *packed.shape[1:]))
153 dist.all_gather_into_tensor(output, packed, group=group)
154 return output.view(size, *packed.shape)
157class KDAStateAllGather:
158 """Gather S/M forward and G/M backward; only local M survives forward."""
160 def __init__(self, group: Any, rank: int, size: int) -> None:
161 """Bind precreated communication metadata, without tensor workspaces."""
162 if not isinstance(size, int) or isinstance(size, bool) or size < 1 or not 0 <= rank < size:
163 raise ValueError("AllGather requires a positive size and valid local rank.")
164 self.group, self.rank, self.size = group, rank, size
166 def forward(self, state: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
167 """Return the incoming state; the caller saves its original local M.
169 Args:
170 state: Local FP32 offset S [B,H,128,128].
171 matrix: Local forward transfer M of the same shape.
173 Returns:
174 This rank's incoming recurrent state.
175 """
176 _validate_summary(state, matrix)
177 maps = _gather_maps(state, matrix, self.group, self.size)
178 return merge_affine_summaries(maps[..., :128], maps[..., 128:], self.rank, True)
180 def backward(self, gradient: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
181 """Collect new G with the saved local M and recover the outgoing gradient.
183 Args:
184 gradient: Local FP32 offset G [B,H,128,128].
185 matrix: This invocation's cached local forward M.
187 Returns:
188 This rank's outgoing recurrent-state gradient.
189 """
190 _validate_summary(gradient, matrix)
191 maps = _gather_maps(gradient, matrix, self.group, self.size)
192 return merge_affine_summaries(maps[..., :128], maps[..., 128:], self.rank, False)
195def _validate_scan_range(indices: range, size: int, forward: bool, prefix_count: int | None) -> None:
196 """Check the requested chronological interval and optional saved prefix."""
197 count = len(indices)
198 if prefix_count is not None and not 0 <= prefix_count <= count:
199 raise ValueError("The saved prefix must lie within the scan interval.")
200 if count:
201 endpoints_valid = all(0 <= index < size for index in (indices[0], indices[-1]))
202 if indices.step != (1 if forward else -1) or not endpoints_valid:
203 raise ValueError("KDA scans require an ordered interval inside the group.")
206def _scan_maps(maps: torch.Tensor, incoming: torch.Tensor,
207 indices: range, forward: bool, prefix_count: int | None = None,
208 ) -> tuple[torch.Tensor, torch.Tensor | None]:
209 """Scan validated group maps once, optionally returning an intermediate prefix."""
210 count = len(indices)
211 _validate_scan_range(indices, maps.shape[0], forward, prefix_count)
212 if not count:
213 return incoming, incoming if prefix_count == 0 else None
214 if incoming.device.type != "npu" or not incoming.is_contiguous():
215 raise ValueError("KDA fused scans require a contiguous Ascend input state.")
216 output = torch.empty_like(incoming)
217 prefix = None
218 if prefix_count is not None:
219 prefix = incoming if prefix_count == 0 else output if prefix_count == count else torch.empty_like(incoming)
220 save_prefix = prefix is not None and prefix_count not in (0, count)
221 # Keep the optional Triton dependency out of CPU-only model imports.
222 # pylint: disable=import-outside-toplevel
223 from hyper_parallel.components.functional._triton.kimi_delta_attention.cp_merge import kda_affine_scan_kernel
224 # pylint: enable=import-outside-toplevel
226 batch, heads = incoming.shape[:2]
227 kda_affine_scan_kernel[(2, batch * heads)](
228 output, prefix, maps, incoming, count, indices.start, prefix_count or 0,
229 rank_stride=maps.stride(0), head_stride=maps.stride(2), row_stride=maps.stride(3),
230 forward=forward, save_prefix=save_prefix, block_v=64, num_warps=2, num_stages=2)
231 return output, prefix
234class KDAStateAllGatherP2P:
235 """Use group-local gathers and an owner chain, retaining only each local M."""
237 def __init__(self, group: Any, intra_group: Any, rank: int,
238 ranks: tuple[int, ...], width: int) -> None:
239 """Bind ranks already validated against the process-group order by the factory."""
240 size = len(ranks)
241 if not isinstance(width, int) or isinstance(width, bool) or width < 2 or size % width:
242 raise ValueError("Grouped AllGather width must be >= 2 and divide the CP size.")
243 if not 0 <= rank < size or len(set(ranks)) != size:
244 raise ValueError("Grouped AllGather requires unique global ranks and a valid local rank.")
245 self.group, self.intra, self.rank, self.width = group, intra_group, rank, width
246 self.ranks = tuple(ranks)
247 self.block, self.local = divmod(rank, width)
248 self.blocks = size // width
249 self.owner = self.block * width + width - 1
251 def _receive(self, state: torch.Tensor, forward: bool) -> tuple[torch.Tensor, Any]:
252 """Post the owner receive before gathering local summaries."""
253 incoming = torch.zeros_like(state)
254 recv_handle = None
255 peer = self.owner - self.width if forward else self.owner + self.width
256 if self.rank == self.owner and 0 <= peer < len(self.ranks):
257 recv_handle = dist.irecv(incoming, src=self.ranks[peer], group=self.group)
258 return incoming, recv_handle
260 def forward(self, state: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
261 """Relay ordered group scans and restore this rank's incoming state.
263 Args:
264 state: Local FP32 offset S [B,H,128,128].
265 matrix: Local forward transfer M of the same shape.
267 Returns:
268 This rank's incoming recurrent state.
269 """
270 _validate_summary(state, matrix)
271 incoming, recv_handle = self._receive(state, True)
272 maps = _gather_maps(state, matrix, self.intra, self.width)
273 send_handle, send_state = None, None
274 result = None
275 if self.rank == self.owner:
276 if recv_handle is not None:
277 recv_handle.wait()
278 if self.block + 1 < self.blocks:
279 send_state, result = _scan_maps(maps, incoming, range(self.width), True,
280 prefix_count=self.local)
281 send_handle = dist.isend(send_state, dst=self.ranks[self.owner + self.width], group=self.group)
282 if self.block > 0:
283 dist.broadcast(incoming, src=self.ranks[self.owner], group=self.intra)
284 if result is None:
285 result, _ = _scan_maps(maps, incoming, range(self.local), True)
286 if send_handle is not None:
287 send_handle.wait()
288 return result
290 def backward(self, gradient: torch.Tensor, matrix: torch.Tensor) -> torch.Tensor:
291 """Regather G/M and reverse the owner chain with ordered state feedback.
293 Args:
294 gradient: Local FP32 offset G [B,H,128,128].
295 matrix: This invocation's cached local forward M.
297 Returns:
298 This rank's outgoing recurrent-state gradient.
299 """
300 _validate_summary(gradient, matrix)
301 incoming, recv_handle = self._receive(gradient, False)
302 maps = _gather_maps(gradient, matrix, self.intra, self.width)
303 send_handle, send_state = None, None
304 if self.rank == self.owner:
305 if recv_handle is not None:
306 recv_handle.wait()
307 if self.block > 0:
308 send_state, _ = _scan_maps(maps, incoming, range(self.width - 1, -1, -1), False)
309 send_handle = dist.isend(send_state, dst=self.ranks[self.owner - self.width], group=self.group)
310 if self.block + 1 < self.blocks:
311 dist.broadcast(incoming, src=self.ranks[self.owner], group=self.intra)
312 result, _ = _scan_maps(maps, incoming, range(self.width - 1, self.local, -1), False)
313 if send_handle is not None:
314 send_handle.wait()
315 return result