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