« 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"""Immutable reset-aware RD plans derived from packed document metadata."""
16from __future__ import annotations
18from bisect import bisect_left
19from dataclasses import dataclass
20from functools import lru_cache
22from hyper_parallel.components.functional.linear_attention_packed import PackedDocumentInfo
25@dataclass(frozen=True)
26class RDScanRound:
27 """One immutable exchange and its subsequent transfer-matrix consumer."""
29 stage: int
30 destination: int | None
31 source: int | None
32 save_matrix: bool
35def build_scan_rounds(
36 rank: int, size: int, distances: tuple[int, ...],
37) -> tuple[RDScanRound, ...]:
38 """Precompute dense inclusive-scan peers and transfer consumers for one rank.
40 Args:
41 rank: Rank in the state process group.
42 size: Number of chronological ranks in the state process group.
43 distances: Increasing powers of two used by the inclusive scan.
44 """
45 # Dense execution retains the final rank's inclusive output. Use size as
46 # an exclusive scan bound, instead of pruning a document's final rank.
47 return _interval_rounds(rank, distances, _RankInterval(0, size, rank > 0))
50@dataclass(frozen=True)
51class _RankInterval:
52 """Rank span of the tail document and independent head continuation."""
54 first: int
55 last: int
56 continues_previous: bool
58 def peers(self, rank: int, distance: int) -> tuple[int | None, int | None]:
59 """Return the exact live scan edge at one doubling distance.
61 A document's last rank has no exit-state consumer. Earlier ranks reach
62 a predecessor precisely when that predecessor belongs to this document.
63 """
64 destination = rank + distance if rank + distance < self.last else None
65 source = rank - distance if rank < self.last and rank - distance >= self.first else None
66 return destination, source
68 def cache_count(self, rank: int) -> int:
69 """Count powers of two that reach a preceding rank of this document."""
70 return (rank - self.first).bit_length() if rank < self.last else 0
73def _rank_interval(offsets: tuple[int, ...], local_length: int, rank: int) -> _RankInterval:
74 """Find the tail document without treating a rank cut as a reset."""
75 tail = bisect_left(offsets, (rank + 1) * local_length) - 1
76 start = rank * local_length
77 head = bisect_left(offsets, start)
78 return _RankInterval(offsets[tail] // local_length, (offsets[tail + 1] - 1) // local_length,
79 offsets[head] != start)
82def _interval_rounds(
83 rank: int, distances: tuple[int, ...], interval: _RankInterval,
84) -> tuple[RDScanRound, ...]:
85 """Keep stage order and transfer lifetimes from the inclusive scan.
87 Each effective receive caches the matrix used by its adjoint. Retain its
88 composed output only when a subsequent receive consumes that output.
89 """
90 last_receive = interval.cache_count(rank) - 1
91 rounds = []
92 for stage, distance in enumerate(distances):
93 destination, source = interval.peers(rank, distance)
94 if destination is not None or source is not None:
95 save_matrix = stage < last_receive
96 rounds.append(RDScanRound(stage, destination, source, save_matrix))
97 return tuple(rounds)
100@dataclass(frozen=True)
101class RDResetPlan:
102 """An immutable local schedule derived from true document boundaries.
104 Cached-adjoint backward transposes the same plan; fresh reverse maps use
105 a mirrored plan. Matrices are never inspected on device: only true document
106 boundaries can remove an edge. Binding one rank
107 retains its tail interval and rounds, without building the all-rank graph.
108 """
110 size: int
111 distances: tuple[int, ...]
112 offsets: tuple[int, ...]
113 local_length: int
114 local_rank: int | None
115 local_rounds: tuple[RDScanRound, ...]
116 _local_interval: _RankInterval | None
118 def _interval(self, rank: int) -> _RankInterval:
119 """Reuse local bounds or locate another rank for an audit query."""
120 if rank == self.local_rank and self._local_interval is not None:
121 return self._local_interval
122 if not 0 <= rank < self.size:
123 raise ValueError("RD execution-plan rank is outside the state process group")
124 return _rank_interval(self.offsets, self.local_length, rank)
126 def peers(self, stage: int, rank: int) -> tuple[int | None, int | None]:
127 """Return forward destination/source, both in process-group order."""
128 return self._interval(rank).peers(rank, self.distances[stage])
130 def shift_peers(self, rank: int) -> tuple[int | None, int | None]:
131 """Return the tail's consumer and the head's producer independently."""
132 interval = self._interval(rank)
133 return (rank + 1 if rank < interval.last else None,
134 rank - 1 if interval.continues_previous else None)
136 def is_transfer_active(self, rank: int) -> bool:
137 """Whether a rank's entry can reach its exit without a document reset."""
138 return rank > self._interval(rank).first
140 def execution_rounds(self, rank: int) -> tuple[RDScanRound, ...]:
141 """Reuse the caller's plan without materializing schedules for every peer."""
142 if rank == self.local_rank:
143 return self.local_rounds
144 return _interval_rounds(rank, self.distances, self._interval(rank))
146 def cache_count(self, rank: int) -> int:
147 """Count invocation-owned pre-round transfers on one rank."""
148 return self._interval(rank).cache_count(rank)
151@lru_cache(maxsize=32)
152def _validate_intervals(offsets: tuple[int, ...], local_length: int, size: int) -> None:
153 """Share document validation across rank-specific plans with bounded storage."""
154 if (size < 1 or local_length < 1 or len(offsets) < 2 or offsets[0] != 0
155 or offsets[-1] != size * local_length
156 or any(left >= right for left, right in zip(offsets, offsets[1:]))):
157 raise ValueError("RD reset planning requires ordered documents covering equal state-rank intervals")
160@lru_cache(maxsize=32)
161def _reverse_documents(offsets: tuple[int, ...] | PackedDocumentInfo) -> PackedDocumentInfo:
162 """Cache the reversed immutable snapshot across layers of the same batch."""
163 values = offsets.offsets if isinstance(offsets, PackedDocumentInfo) else offsets
164 return PackedDocumentInfo(tuple(values[-1] - value for value in reversed(values)))
167@lru_cache(maxsize=32)
168def build_reset_plan(
169 offsets: tuple[int, ...] | PackedDocumentInfo, local_length: int, size: int, rank: int | None = None,
170 *, reverse: bool = False,
171) -> RDResetPlan:
172 """Derive one rank's exact live edges from its tail document interval.
174 Args:
175 offsets: Global document offsets, or a batch's already validated snapshot.
176 Reusing the snapshot avoids full-offset validation and hashing.
177 local_length: Equal token interval length after any Ulysses exchange.
178 size: Number of chronological state ranks, independent of global rank IDs.
179 rank: Optional physical group rank whose execution schedule should be cached.
180 reverse: Plan a fresh reverse-map scan in mirrored token/rank order.
181 The executor maps logical peers back to the unchanged process group.
183 Returns:
184 An immutable CPU plan reusable across layers with identical metadata.
185 Omitting rank retains only metadata for subsequent diagnostic queries.
186 """
187 if rank is not None and not 0 <= rank < size:
188 raise ValueError("RD execution-plan rank is outside the state process group")
189 if reverse:
190 if not isinstance(offsets, PackedDocumentInfo):
191 _validate_intervals(offsets, local_length, size)
192 offsets = _reverse_documents(offsets)
193 rank = None if rank is None else size - rank - 1
194 if isinstance(offsets, PackedDocumentInfo):
195 offsets = offsets.offsets
196 if size < 1 or local_length < 1 or offsets[-1] != size * local_length:
197 raise ValueError("RD reset planning requires ordered documents covering equal state-rank intervals")
198 else:
199 _validate_intervals(offsets, local_length, size)
200 distances = tuple(1 << level for level in range((size - 1).bit_length()))
201 interval = None if rank is None else _rank_interval(offsets, local_length, rank)
202 rounds = () if rank is None else _interval_rounds(rank, distances, interval)
203 return RDResetPlan(size, distances, offsets, local_length, rank, rounds, interval)