Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / runtime_input.py: 72%
110 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"""Resolve sequence metadata and model-owned runtime inputs."""
17from __future__ import annotations
19from abc import ABC, abstractmethod
20from collections.abc import Mapping, Sequence
21from dataclasses import dataclass
22from typing import Any
24import torch
27class OnlineBoundaryResolver:
28 """Read global cumulative sequence boundaries emitted by Online packing."""
30 @staticmethod
31 def resolve(canonical_batch: Mapping[str, Any]) -> Any:
32 """Read leading-zero ``cu_seq_lens`` from an Online batch."""
33 raw_cu_seq_lens = canonical_batch["cu_seq_lens"]
34 cu_seq_lens = raw_cu_seq_lens.to(torch.int32)
35 return cu_seq_lens
38class IndexedBoundaryResolver:
39 """Recover global cumulative sequence boundaries from Indexed input IDs."""
41 def __init__(self, eod_token_id: int | None) -> None:
42 """Store the token separating packed Indexed sequences."""
43 self.eod_token_id = eod_token_id
45 def resolve(self, canonical_batch: Mapping[str, Any]) -> Any:
46 """Recover leading-zero ``cu_seq_lens`` from Indexed tokens."""
47 input_ids = canonical_batch["input_ids"]
48 batch_size, seq_len = input_ids.shape
49 if "cu_seq_lens" in canonical_batch:
50 rows = canonical_batch["cu_seq_lens"]
51 if isinstance(rows, torch.Tensor) and rows.ndim == 1:
52 return rows.to(torch.int32)
53 if len(rows) != batch_size:
54 raise ValueError("cu_seq_lens must have one boundary array per batch row")
55 # Only metadata is flattened; tokens and labels remain [B, S].
56 ends = [end + row_idx * seq_len for row_idx, row in enumerate(rows) for end in row[1:]]
57 return torch.tensor([0, *ends], dtype=torch.int32)
58 token_indices = torch.arange(seq_len, dtype=input_ids.dtype, device=input_ids.device)
60 seq_ends = []
61 for batch_idx in range(batch_size):
62 if self.eod_token_id is None:
63 eod_indices = token_indices[:0]
64 else:
65 eod_indices = token_indices[input_ids[batch_idx] == self.eod_token_id]
67 prev_eod_idx = -1
68 for eod_idx in eod_indices:
69 eod_idx = int(eod_idx.item())
70 if eod_idx == prev_eod_idx:
71 break
73 seq_end = batch_idx * seq_len + eod_idx + 1
74 seq_ends.append(seq_end)
75 prev_eod_idx = eod_idx + 1
77 row_end = (batch_idx + 1) * seq_len
78 if not seq_ends or seq_ends[-1] != row_end:
79 seq_ends.append(row_end)
81 cu_seq_lens = torch.tensor([0, *seq_ends], dtype=torch.int32)
82 return cu_seq_lens
85@dataclass(frozen=True)
86class RuntimeInputContext:
87 """Framework execution facts available to a model runtime adapter.
89 ``options`` deliberately carries feature-specific policy. The framework
90 owns batch movement and parallel coordinates, while a model adapter may
91 translate those facts into attention metadata, modality routing inputs,
92 cache descriptors, or another forward-only contract.
93 """
95 local_input_shape: Sequence[int]
96 parallel_ranks: Mapping[str, int]
97 parallel_sizes: Mapping[str, int]
98 options: Mapping[str, Any]
100 @classmethod
101 def from_batch(
102 cls,
103 *,
104 batch: Mapping[str, Any],
105 parallel_context: Any,
106 options: Mapping[str, Any] | None = None,
107 ) -> RuntimeInputContext:
108 """Build runtime context from one parallel-local batch."""
109 runtime_context = cls(
110 local_input_shape=batch["input_ids"].shape,
111 parallel_ranks={
112 "tp": parallel_context.tp_rank,
113 "cp": parallel_context.cp_rank,
114 },
115 parallel_sizes={
116 "tp": parallel_context.tp_world_size,
117 "cp": parallel_context.cp_world_size,
118 },
119 options={} if options is None else options,
120 )
121 return runtime_context
124class RuntimeInputAdapter(ABC):
125 """Extend a generic batch with model-owned forward inputs."""
127 def build(
128 self,
129 *,
130 batch: Mapping[str, Any],
131 parallel_context: Any,
132 options: Mapping[str, Any] | None = None,
133 ) -> dict[str, Any]:
134 """Build runtime context and produce one runtime's model inputs."""
135 runtime_context = RuntimeInputContext.from_batch(
136 batch=batch,
137 parallel_context=parallel_context,
138 options=options,
139 )
140 runtime_inputs = self.build_runtime_inputs(
141 batch=batch,
142 context=runtime_context,
143 )
144 if not isinstance(runtime_inputs, Mapping):
145 raise TypeError(
146 "RuntimeInputAdapter.build_runtime_inputs must return a mapping, "
147 f"got {type(runtime_inputs).__name__}"
148 )
149 return dict(runtime_inputs)
151 def prepare_graph_inputs(self, model_inputs: dict[str, Any]) -> dict[str, Any]:
152 """Return inputs unchanged unless the model needs graph-only metadata.
154 Args:
155 model_inputs: Real model inputs immediately before graph execution.
157 Returns:
158 Inputs used for both graph capture and replay. Eager never calls this hook.
159 """
160 return model_inputs
162 def runtime_input_fields(self) -> tuple[str, ...]:
163 """Return the forward fields exclusively owned by this adapter.
165 Dynamic values may vary by batch, but field ownership must be stable so
166 diagnostics can reject collisions before interpreting precision runs.
167 """
168 return ()
170 @abstractmethod
171 def build_runtime_inputs(
172 self,
173 *,
174 batch: Mapping[str, Any],
175 context: RuntimeInputContext,
176 ) -> Mapping[str, Any]:
177 """Build additional model inputs without mutating ``batch``.
179 Args:
180 batch: Device-resident, parallel-local batch. Global metadata such
181 as ``cu_seq_lens`` remains available when required.
182 context: Local geometry, parallel coordinates, and feature-specific
183 recipe options.
185 Returns:
186 Mapping merged into the model forward inputs. Keys must not replace
187 fields already owned by the generic batch path.
188 """
189 raise NotImplementedError
192class AttentionRuntime(RuntimeInputAdapter):
193 """Build dense attention fields from common runtime context."""
195 def __init__(
196 self,
197 *,
198 mode: str,
199 create_mask: bool,
200 reset_mask: bool,
201 sliding_window: int | None,
202 ) -> None:
203 """Store the attention representation and mask policy."""
204 self.mode = mode
205 self.create_mask = create_mask
206 self.reset_mask = reset_mask
207 self.sliding_window = sliding_window
209 def build_runtime_inputs(
210 self,
211 *,
212 batch: Mapping[str, Any],
213 context: RuntimeInputContext,
214 ) -> Mapping[str, Any]:
215 """Build attention fields without owning model-specific packed metadata."""
216 attention_inputs = {"attention_mask": None}
217 if not self.create_mask or self.mode == "compressed":
218 return attention_inputs
220 micro_batch_size, local_sequence_length = batch["input_ids"].shape
221 cp_size = context.parallel_sizes["cp"]
222 global_sequence_length = local_sequence_length * cp_size
223 attention_mask, sliding_window_mask = self._build_dense_attention_masks(
224 cu_seq_lens=batch["cu_seq_lens"],
225 micro_batch_size=micro_batch_size,
226 seq_length=global_sequence_length,
227 device=batch["input_ids"].device,
228 )
229 attention_inputs["attention_mask"] = attention_mask
230 if sliding_window_mask is not None:
231 attention_inputs["swa_mask"] = sliding_window_mask
233 return attention_inputs
235 def _build_dense_attention_masks(
236 self,
237 *,
238 cu_seq_lens: Any,
239 micro_batch_size: int,
240 seq_length: int,
241 device: Any,
242 ) -> tuple[Any, Any | None]:
243 """Build global dense attention and optional sliding-window masks."""
244 boundaries = [int(boundary) for boundary in cu_seq_lens.tolist()]
245 if len(boundaries) < 2 or boundaries[0] != 0:
246 raise ValueError("cu_seq_lens must contain a leading zero and at least one sequence")
248 if any(end <= start for start, end in zip(boundaries[:-1], boundaries[1:])):
249 raise ValueError("cu_seq_lens must be strictly increasing")
251 physical_sequence_length = micro_batch_size * seq_length
252 if boundaries[-1] != physical_sequence_length:
253 raise ValueError(
254 "cu_seq_lens must cover the physical batch length "
255 f"({physical_sequence_length}), but got {boundaries[-1]}"
256 )
258 attention_batch_size = micro_batch_size if self.reset_mask else 1
259 attention_mask = torch.ones(
260 (attention_batch_size, seq_length, seq_length),
261 dtype=torch.bool,
262 device=device,
263 ).tril()
264 attention_mask = attention_mask.view(attention_batch_size, 1, seq_length, seq_length)
266 if self.reset_mask:
267 for sequence_start in boundaries[:-1]:
268 batch_index, row_sequence_start = divmod(sequence_start, seq_length)
269 attention_mask[batch_index, 0, row_sequence_start:, :row_sequence_start] = False
271 sliding_window_mask = None
272 if self.sliding_window is not None:
273 positions = torch.arange(seq_length, dtype=torch.int64, device=device)
274 outside_window = positions.unsqueeze(1) - positions.unsqueeze(0) > self.sliding_window
275 sliding_window_mask = attention_mask & ~outside_window.unsqueeze(0).unsqueeze(0)
277 return attention_mask, sliding_window_mask
280__all__ = [
281 "AttentionRuntime",
282 "IndexedBoundaryResolver",
283 "OnlineBoundaryResolver",
284 "RuntimeInputAdapter",
285 "RuntimeInputContext",
286]