Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / context_parallel / kimi_delta_attention_rd.py: 72%
232 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"""Cached eager recursive doubling using public batched point-to-point transport."""
16from __future__ import annotations
18# pylint: disable=forbidden-backend-import
20from dataclasses import dataclass
21from functools import lru_cache
22import importlib
23from typing import Any
25import torch
26import torch.distributed as dist
28from hyper_parallel.distributed.context_parallel.kimi_delta_attention_rd_plan import (
29 RDResetPlan, build_reset_plan, build_scan_rounds,
30)
33_RD_EPILOGUE_BLOCK_ROWS = 64
36@lru_cache(maxsize=1)
37def _rd_kernels() -> Any:
38 """Load optional Ascend kernels only when an NPU epilogue is executed."""
39 return importlib.import_module(
40 "hyper_parallel.components.functional._triton.kimi_delta_attention.rd_pointwise"
41 )
44def _finish_round(product: torch.Tensor, previous: torch.Tensor, save_m: bool = True) -> torch.Tensor | None:
45 """Add the previous S and copy the next M in one launch, without fusing GEMM.
47 Inputs must be distinct contiguous FP32 [B,H,128,256] buffers. Only the
48 newly computed product's S half is modified; input summaries are immutable.
49 The validated scan inputs and fresh matrix product establish this layout.
51 Args:
52 product: Fresh contiguous FP32 packed matrix product [B, H, 128, 256].
53 previous: Previous immutable packed S/M tensor with the same layout.
54 save_m: Whether a future round consumes the new contiguous M.
55 """
56 state_width, packed_width = product.shape[-2:]
57 matrix = (torch.empty((*product.shape[:-1], state_width), dtype=product.dtype, device=product.device)
58 if save_m else None)
59 kernels = _rd_kernels()
61 rows = product.numel() // packed_width
62 block_rows = _RD_EPILOGUE_BLOCK_ROWS
63 kernels.kda_rd_finish_round_kernel[((rows + block_rows - 1) // block_rows,)](
64 product, previous, matrix if save_m else product,
65 row_count=rows, state_width=state_width, packed_width=packed_width,
66 block_rows=block_rows, save_matrix=save_m, num_warps=4)
67 return matrix
70def _finish_to_state(
71 product: torch.Tensor, previous: torch.Tensor, save_m: bool,
72) -> tuple[torch.Tensor | None, torch.Tensor]:
73 """Emit contiguous final-stage S alongside the unchanged packed epilogue.
75 Args:
76 product: Fresh contiguous FP32 packed matrix product [B, H, 128, 256].
77 previous: Previous immutable packed S/M tensor with the same layout.
78 save_m: Whether a future round consumes the new contiguous M.
79 """
80 state_width, packed_width = product.shape[-2:]
81 state = torch.empty((*product.shape[:-1], state_width), dtype=product.dtype, device=product.device)
82 matrix = torch.empty_like(state) if save_m else None
83 kernels = _rd_kernels()
85 rows = product.numel() // packed_width
86 block_rows = _RD_EPILOGUE_BLOCK_ROWS
87 kernels.kda_rd_finish_to_state_kernel[((rows + block_rows - 1) // block_rows,)](
88 product, previous, matrix if save_m else product, state,
89 row_count=rows, state_width=state_width, packed_width=packed_width,
90 block_rows=block_rows, save_matrix=save_m, num_warps=4)
91 return matrix, state
94class RoundExchange:
95 """Keep immutable payload alive until send completion; outputs belong to the call."""
97 def __init__(self, group: Any, rank: int, ranks: tuple[int, ...]) -> None:
98 """Resolve ordered peers and the process group's protocol device."""
99 if not ranks or len(set(ranks)) != len(ranks) or not 0 <= rank < len(ranks):
100 raise ValueError("RD requires unique global peers and a valid local rank")
101 if tuple(dist.get_process_group_ranks(group)) != ranks:
102 raise ValueError("RD mesh sequence order must match its communication group order")
103 self.group, self.rank, self.size, self.peers = group, rank, len(ranks), ranks
104 if dist.get_backend(group) != "hccl":
105 raise ValueError("KDA recursive doubling requires an HCCL process group")
106 # Receive-only TND ranks have no local summary to infer a device from.
107 # Load the optional extension here and bind buffers to the caller's NPU.
108 torch_npu = importlib.import_module("torch_npu")
109 self.device = torch.device("npu", torch_npu.npu.current_device())
111 def __call__(
112 self, payload: torch.Tensor | None, destination: int | None, source: int | None,
113 receive_shape: tuple[int, ...] | None = None,
114 ) -> torch.Tensor | None:
115 """Exchange peers, allocating only a live receive independently of the send.
117 Args:
118 payload: Immutable send tensor; may be absent when there is no destination.
119 destination: Group-local send peer, or None.
120 source: Group-local receive peer, or None.
121 receive_shape: Explicit receive shape, required for a receive without a payload.
122 """
123 output = self._prepare_receive(payload, destination, source, receive_shape)
124 if source is None and destination is None:
125 return output
126 operations = []
127 if source is not None:
128 operations.append(dist.P2POp(dist.irecv, output, self.peers[source], self.group))
129 if destination is not None:
130 operations.append(dist.P2POp(dist.isend, payload, self.peers[destination], self.group))
131 for handle in dist.batch_isend_irecv(operations):
132 handle.wait()
133 return output
135 def _prepare_receive(
136 self, payload: torch.Tensor | None, destination: int | None, source: int | None,
137 receive_shape: tuple[int, ...] | None,
138 ) -> torch.Tensor | None:
139 """Validate one exchange and allocate only its live receive buffer."""
140 if payload is None:
141 if destination is not None:
142 raise ValueError('RD send requires a payload')
143 if source is not None and receive_shape is None:
144 raise ValueError('RD receive without a payload requires a shape')
145 elif payload.device != self.device or not payload.is_contiguous():
146 raise ValueError('RD payload must be contiguous on the protocol device')
147 if any(peer is not None and not 0 <= peer < self.size for peer in (source, destination)):
148 raise ValueError('RD peer is outside the process group')
149 output = None
150 if source is not None:
151 output = (torch.empty_like(payload) if receive_shape is None else torch.empty(
152 receive_shape, dtype=torch.float32 if payload is None else payload.dtype, device=self.device))
153 return output
156def _validate_summary(
157 tensor: torch.Tensor, device: torch.device, state_shape: tuple[int, ...] | None = None,
158) -> None:
159 """Validate caller-owned S/M/G and cached transfers before any peer exchange."""
160 if (not isinstance(tensor, torch.Tensor) or tensor.ndim != 4
161 or tensor.shape[-2:] != (128, 128) or tensor.dtype != torch.float32
162 or tensor.device != device or not tensor.is_contiguous()):
163 raise ValueError("RD requires contiguous FP32 B/H/K128/V128 summaries on the protocol device")
164 if state_shape is not None and tensor.shape != state_shape:
165 raise ValueError("RD summary shape differs from the invocation")
168def _validate_backward_summary(
169 gradient: torch.Tensor | None, cache: tuple, expected: int,
170 device: torch.device, state_shape: tuple[int, ...] | None = None,
171) -> None:
172 """Reject incompatible invocation cache layouts before any peer exchange."""
173 if gradient is None:
174 if state_shape is None:
175 raise ValueError("RD gradient may be absent only with invocation shape metadata")
176 else:
177 _validate_summary(gradient, device, state_shape)
178 state_shape = gradient.shape
179 if not isinstance(cache, tuple) or len(cache) != expected:
180 raise ValueError("RD received an incompatible invocation cache")
181 for matrix in cache:
182 _validate_summary(matrix, device, state_shape)
185def _compact_transfer(matrix: torch.Tensor) -> torch.Tensor:
186 """Keep a contiguous M view from retaining a larger producer allocation."""
187 if matrix.untyped_storage().nbytes() != matrix.numel() * matrix.element_size():
188 return matrix.clone()
189 return matrix
192class KDAStateRecursiveDoubling:
193 """Exclusive forward scan and its cached adjoint, preserving the established rounding.
195 Caller creates/warm-ups the HCCL group collectively. Inputs are immutable,
196 contiguous FP32 B/H/K128/V128 summaries. Keep the returned invocation cache
197 with its own autograd context. No shared graph buffers or mutable layer state.
198 """
200 requires_contiguous_summaries = True
202 def __init__(self, group: Any, rank: int, ranks: tuple[int, ...]) -> None:
203 """Prepare communication metadata and logarithmic scan distances."""
204 self.exchange = RoundExchange(group, rank, ranks)
205 self.rank, self.size = self.exchange.rank, self.exchange.size
206 self.device = self.exchange.device
207 self.distances = tuple(1 << level for level in range((self.size - 1).bit_length()))
208 self.rounds = build_scan_rounds(self.rank, self.size, self.distances)
210 def bind_sequence(
211 self, sequence_info: Any, state_shape: tuple[int, ...] | None = None,
212 ) -> _ResetRDInvocation:
213 """Bind immutable packed metadata, optionally allowing absent dead summaries.
215 Args:
216 sequence_info: This call's packed document offsets and local interval.
217 state_shape: B/H/K/V dimensions, without retaining a tensor allocation.
218 """
219 if sequence_info.start != self.rank * sequence_info.length:
220 raise ValueError("RD packed interval must follow the state process-group order")
221 if state_shape is not None:
222 state_shape = tuple(state_shape)
223 if (len(state_shape) != 4 or state_shape[-2:] != (128, 128)
224 or any(not isinstance(dim, int) or dim < 1 for dim in state_shape)):
225 raise ValueError("RD invocation shape requires positive B/H and K=V128")
226 plan = build_reset_plan(sequence_info.documents, sequence_info.length, self.size, self.rank)
227 return _ResetRDInvocation(self, plan, state_shape)
229 def directional_plan(self, sequence_info: Any, local_length: int, reverse: bool = False) -> RDResetPlan:
230 """Reuse reset planning for independent forward or reverse affine maps.
232 Args:
233 sequence_info: Packed metadata, or None for one continuous document.
234 local_length: Equal token interval length in this state group.
235 reverse: Whether to scan fresh G/N in reverse chronological order.
236 """
237 if sequence_info is not None and sequence_info.start != self.rank * local_length:
238 raise ValueError("RD packed interval must follow the state process-group order")
239 documents = (0, self.size * local_length) if sequence_info is None else sequence_info.documents
240 return build_reset_plan(documents, local_length, self.size, self.rank, reverse=reverse)
242 def scan_affine(
243 self, packed: torch.Tensor | None, plan: RDResetPlan,
244 state_shape: tuple[int, ...], reverse: bool = False,
245 ) -> torch.Tensor | None:
246 """Scan independently supplied affine maps without a forward-adjoint cache.
248 Args:
249 packed: Immutable FP32 [B,H,128,256] S/M or G/N; absent for a dead tail.
250 plan: Directional reset plan from this invocation's metadata.
251 state_shape: B/H/K/V dimensions for receive-only fragments.
252 reverse: Map mirrored logical ranks onto the original process group.
254 Returns:
255 Exclusive input state, or None at a real document boundary.
256 N is already a reverse transfer and is never transposed here.
257 """
258 rank = self.size - self.rank - 1 if reverse else self.rank
259 if plan.size != self.size or plan.local_rank != rank:
260 raise ValueError("RD affine scan received another rank's directional plan")
261 rounds = plan.execution_rounds(rank)
262 destination, source = plan.shift_peers(rank)
263 if packed is None:
264 if rounds or destination is not None:
265 raise ValueError("RD affine map may be absent only without an exit consumer")
266 state = None
267 else:
268 if (packed.shape != (*state_shape[:-1], 256) or packed.dtype != torch.float32
269 or packed.device != self.device or not packed.is_contiguous()):
270 raise ValueError("RD affine maps must be contiguous FP32 [B,H,128,256] on the protocol device")
271 with torch.autocast(device_type=packed.device.type, enabled=False):
272 state = self._scan_independent_maps(packed, plan, rounds, reverse)
273 return self._exchange_direction(state, destination, source, reverse, state_shape)
275 def _exchange_direction(self, payload, destination, source, reverse, receive_shape=None):
276 """Reverse only logical peers; retain the original communicator and global ranks."""
277 if reverse:
278 destination = None if destination is None else self.size - destination - 1
279 source = None if source is None else self.size - source - 1
280 return self.exchange(payload, destination, source, receive_shape=receive_shape)
282 def _scan_independent_maps(self, packed, plan, rounds, reverse):
283 """Use FP32 Torch GEMM/add while sharing RD transport and pruning plans."""
284 state = packed[..., :128]
285 matrix = packed[..., 128:].contiguous() if rounds else None
286 for step in rounds:
287 final = step.stage == len(plan.distances) - 1
288 payload = state.contiguous() if final else packed
289 received = self._exchange_direction(payload, step.destination, step.source, reverse)
290 if received is None:
291 continue
292 if final:
293 state = matrix @ received + state
294 else:
295 updated = matrix @ received
296 updated[..., :128].add_(state)
297 packed, state = updated, updated[..., :128]
298 matrix = updated[..., 128:].contiguous() if step.save_matrix else None
299 return state.contiguous()
301 def _shift_peers(self, plan: RDResetPlan | None) -> tuple[int | None, int | None]:
302 """Only continued documents consume the preceding rank's exit."""
303 if plan is not None:
304 return plan.shift_peers(self.rank)
305 return (self.rank + 1 if self.rank + 1 < self.size else None,
306 self.rank - 1 if self.rank else None)
308 def _compose_packed_round(self, step, matrix, received, previous):
309 """Keep packed arithmetic, materializing a transfer only for a later receive."""
310 aggregate = matrix @ received
311 if 2 * self.distances[step.stage] == self.distances[-1]:
312 next_matrix, state = _finish_to_state(aggregate, previous, step.save_matrix)
313 else:
314 next_matrix = _finish_round(aggregate, previous, save_m=step.save_matrix)
315 state = None
316 return aggregate, next_matrix, state
318 def forward(
319 self, state: torch.Tensor | None, matrix: torch.Tensor | None,
320 plan: RDResetPlan | None = None, state_shape: tuple[int, ...] | None = None,
321 ) -> tuple[torch.Tensor | None, tuple[torch.Tensor, ...]]:
322 """Return the exclusive incoming state and immutable pre-round transfer cache.
324 Args:
325 state: Local summary S tensor or kernel output.
326 matrix: Contiguous FP32 local transfer M.
327 plan: Internal immutable plan supplied by the packed invocation wrapper.
328 state_shape: Optional per-call dimensions enabling absent dead summaries.
329 """
330 rounds = self.rounds if plan is None else plan.execution_rounds(self.rank)
331 destination, source = self._shift_peers(plan)
332 if state_shape is not None and plan is None:
333 raise ValueError("RD optional summaries require a packed invocation plan")
334 if state is None:
335 if state_shape is None or matrix is not None or rounds or destination is not None:
336 raise ValueError("RD state may be absent only when no exit-state consumer exists")
337 return self.exchange(None, None, source, receive_shape=state_shape), ()
338 _validate_summary(state, self.device, state_shape)
339 _validate_summary(matrix, self.device, state.shape)
340 narrow_state, cache = self._scan_forward(state, matrix, rounds)
341 received = self.exchange(narrow_state, destination, source)
342 if received is None and state_shape is None:
343 received = torch.zeros_like(state)
344 return received, cache
346 def _scan_forward(
347 self, state: torch.Tensor, matrix: torch.Tensor, rounds: tuple,
348 ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...]]:
349 """Compose live inclusive-scan rounds, preserving per-round adjoint inputs."""
350 matrix = _compact_transfer(matrix) if rounds else matrix
351 # Reset-only shifts never consume a packed S/M payload.
352 packed = (torch.cat((state, matrix), dim=-1)
353 if any(step.stage < len(self.distances) - 1 for step in rounds) else None)
354 before_m, narrow_state = matrix, state
355 cache = []
356 for step in rounds:
357 final = step.stage == len(self.distances) - 1
358 if final and narrow_state is None:
359 narrow_state = packed[..., :128].contiguous()
360 destination, source = step.destination, step.source
361 received = self.exchange(narrow_state if final else packed, destination, source)
362 if source is not None:
363 cache.append(before_m)
364 if final:
365 narrow_state = before_m @ received + narrow_state
366 else:
367 packed, before_m, narrow_state = self._compose_packed_round(
368 step, before_m, received, packed)
369 if narrow_state is None:
370 narrow_state = packed[..., :128].contiguous()
371 return narrow_state, tuple(cache)
373 def backward(
374 self, gradient: torch.Tensor | None, cache: tuple[torch.Tensor, ...],
375 plan: RDResetPlan | None = None, state_shape: tuple[int, ...] | None = None,
376 ) -> torch.Tensor | None:
377 """Transpose the exact recorded forward composition order.
379 Args:
380 gradient: Contiguous FP32 local state-gradient summary G.
381 cache: Invocation-owned transfers returned by this protocol forward.
382 plan: The same internal plan used by this invocation's forward.
383 state_shape: The same optional per-call shape used by forward.
384 """
385 expected = (sum(self.rank >= distance for distance in self.distances)
386 if plan is None else plan.cache_count(self.rank))
387 source, destination = self._shift_peers(plan)
388 if state_shape is not None and plan is None:
389 raise ValueError("RD optional summaries require a packed invocation plan")
390 if gradient is None and destination is not None:
391 raise ValueError("RD gradient may be absent only when no head-gradient consumer exists")
392 _validate_backward_summary(gradient, cache, expected, self.device, state_shape)
393 received = (self.exchange(gradient, destination, source) if gradient is not None
394 else self.exchange(None, destination, source, receive_shape=state_shape))
395 if received is None and state_shape is not None:
396 # A tail without a consumer has no live scan rounds in a reset plan.
397 return None
398 adjoint = torch.zeros_like(gradient) if received is None else received
399 index = len(cache)
400 rounds = self.rounds if plan is None else plan.execution_rounds(self.rank)
401 for step in reversed(rounds):
402 source, destination = step.destination, step.source
403 payload = adjoint
404 if destination is not None:
405 index -= 1
406 payload = cache[index].transpose(-2, -1) @ adjoint
407 received = self.exchange(payload, destination, source)
408 if received is not None:
409 # Adjoint is invocation-owned and every send has completed.
410 adjoint.add_(received)
411 return adjoint
414@dataclass(frozen=True)
415class _ResetRDInvocation:
416 """Share static transport while keeping the exact plan with one autograd call."""
418 requires_contiguous_summaries = True
420 protocol: KDAStateRecursiveDoubling
421 plan: RDResetPlan
422 state_shape: tuple[int, ...] | None = None
424 @property
425 def allows_absent_summaries(self) -> bool:
426 """Whether shape metadata replaces unused full-size placeholders."""
427 return self.state_shape is not None
429 @property
430 def needs_state_summary(self) -> bool:
431 """Only a continuing tail document has an exit-state consumer."""
432 return self.plan.shift_peers(self.protocol.rank)[0] is not None
434 @property
435 def needs_transition_summary(self) -> bool:
436 """Produce M only when a consumed exit depends on this rank's entry."""
437 return self.needs_state_summary and self.plan.is_transfer_active(self.protocol.rank)
439 def forward(
440 self, state: torch.Tensor | None, matrix: torch.Tensor | None,
441 ) -> tuple[torch.Tensor | None, tuple[torch.Tensor, ...]]:
442 """Execute the pruned graph and return its tensor-only cache."""
443 return self.protocol.forward(state, matrix, self.plan, self.state_shape)
445 def backward(self, gradient: torch.Tensor | None, cache: tuple[torch.Tensor, ...]) -> torch.Tensor | None:
446 """Transpose this invocation's graph, independent of later forwards."""
447 return self.protocol.backward(gradient, cache, self.plan, self.state_shape)