Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / batching / runtime_input.py: 36%
100 statements
« 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"""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 token_indices = torch.arange(seq_len, dtype=input_ids.dtype, device=input_ids.device)
51 seq_ends = []
52 for batch_idx in range(batch_size):
53 if self.eod_token_id is None:
54 eod_indices = token_indices[:0]
55 else:
56 eod_indices = token_indices[input_ids[batch_idx] == self.eod_token_id]
58 prev_eod_idx = -1
59 for eod_idx in eod_indices:
60 eod_idx = int(eod_idx.item())
61 if eod_idx == prev_eod_idx:
62 break
64 seq_end = batch_idx * seq_len + eod_idx + 1
65 seq_ends.append(seq_end)
66 prev_eod_idx = eod_idx + 1
68 row_end = (batch_idx + 1) * seq_len
69 if not seq_ends or seq_ends[-1] != row_end:
70 seq_ends.append(row_end)
72 cu_seq_lens = torch.tensor([0, *seq_ends], dtype=torch.int32)
73 return cu_seq_lens
76@dataclass(frozen=True)
77class RuntimeInputContext:
78 """Framework execution facts available to a model runtime adapter.
80 ``options`` deliberately carries feature-specific policy. The framework
81 owns batch movement and parallel coordinates, while a model adapter may
82 translate those facts into attention metadata, modality routing inputs,
83 cache descriptors, or another forward-only contract.
84 """
86 local_input_shape: Sequence[int]
87 parallel_ranks: Mapping[str, int]
88 parallel_sizes: Mapping[str, int]
89 options: Mapping[str, Any]
91 @classmethod
92 def from_batch(
93 cls,
94 *,
95 batch: Mapping[str, Any],
96 parallel_context: Any,
97 options: Mapping[str, Any] | None = None,
98 ) -> RuntimeInputContext:
99 """Build runtime context from one parallel-local batch."""
100 runtime_context = cls(
101 local_input_shape=batch["input_ids"].shape,
102 parallel_ranks={
103 "tp": parallel_context.tp_rank,
104 "cp": parallel_context.cp_rank,
105 },
106 parallel_sizes={
107 "tp": parallel_context.tp_world_size,
108 "cp": parallel_context.cp_world_size,
109 },
110 options={} if options is None else options,
111 )
112 return runtime_context
115class RuntimeInputAdapter(ABC):
116 """Extend a generic batch with model-owned forward inputs."""
118 def build(
119 self,
120 *,
121 batch: Mapping[str, Any],
122 parallel_context: Any,
123 options: Mapping[str, Any] | None = None,
124 ) -> dict[str, Any]:
125 """Build runtime context and produce one runtime's model inputs."""
126 runtime_context = RuntimeInputContext.from_batch(
127 batch=batch,
128 parallel_context=parallel_context,
129 options=options,
130 )
131 runtime_inputs = self.build_runtime_inputs(
132 batch=batch,
133 context=runtime_context,
134 )
135 if not isinstance(runtime_inputs, Mapping):
136 raise TypeError(
137 "RuntimeInputAdapter.build_runtime_inputs must return a mapping, "
138 f"got {type(runtime_inputs).__name__}"
139 )
140 return dict(runtime_inputs)
142 def runtime_input_fields(self) -> tuple[str, ...]:
143 """Return the forward fields exclusively owned by this adapter.
145 Dynamic values may vary by batch, but field ownership must be stable so
146 diagnostics can reject collisions before interpreting precision runs.
147 """
148 return ()
150 @abstractmethod
151 def build_runtime_inputs(
152 self,
153 *,
154 batch: Mapping[str, Any],
155 context: RuntimeInputContext,
156 ) -> Mapping[str, Any]:
157 """Build additional model inputs without mutating ``batch``.
159 Args:
160 batch: Device-resident, parallel-local batch. Global metadata such
161 as ``cu_seq_lens`` remains available when required.
162 context: Local geometry, parallel coordinates, and feature-specific
163 recipe options.
165 Returns:
166 Mapping merged into the model forward inputs. Keys must not replace
167 fields already owned by the generic batch path.
168 """
169 raise NotImplementedError
172class AttentionRuntime(RuntimeInputAdapter):
173 """Build dense attention fields from common runtime context."""
175 def __init__(
176 self,
177 *,
178 mode: str,
179 create_mask: bool,
180 reset_mask: bool,
181 sliding_window: int | None,
182 ) -> None:
183 """Store the attention representation and mask policy."""
184 self.mode = mode
185 self.create_mask = create_mask
186 self.reset_mask = reset_mask
187 self.sliding_window = sliding_window
189 def build_runtime_inputs(
190 self,
191 *,
192 batch: Mapping[str, Any],
193 context: RuntimeInputContext,
194 ) -> Mapping[str, Any]:
195 """Build attention fields without owning model-specific packed metadata."""
196 attention_inputs = {"attention_mask": None}
197 if not self.create_mask or self.mode == "compressed":
198 return attention_inputs
200 micro_batch_size, local_sequence_length = batch["input_ids"].shape
201 cp_size = context.parallel_sizes["cp"]
202 global_sequence_length = local_sequence_length * cp_size
203 attention_mask, sliding_window_mask = self._build_dense_attention_masks(
204 cu_seq_lens=batch["cu_seq_lens"],
205 micro_batch_size=micro_batch_size,
206 seq_length=global_sequence_length,
207 device=batch["input_ids"].device,
208 )
209 attention_inputs["attention_mask"] = attention_mask
210 if sliding_window_mask is not None:
211 attention_inputs["swa_mask"] = sliding_window_mask
213 return attention_inputs
215 def _build_dense_attention_masks(
216 self,
217 *,
218 cu_seq_lens: Any,
219 micro_batch_size: int,
220 seq_length: int,
221 device: Any,
222 ) -> tuple[Any, Any | None]:
223 """Build global dense attention and optional sliding-window masks."""
224 boundaries = [int(boundary) for boundary in cu_seq_lens.tolist()]
225 if len(boundaries) < 2 or boundaries[0] != 0:
226 raise ValueError("cu_seq_lens must contain a leading zero and at least one sequence")
228 if any(end <= start for start, end in zip(boundaries[:-1], boundaries[1:])):
229 raise ValueError("cu_seq_lens must be strictly increasing")
231 physical_sequence_length = micro_batch_size * seq_length
232 if boundaries[-1] != physical_sequence_length:
233 raise ValueError(
234 "cu_seq_lens must cover the physical batch length "
235 f"({physical_sequence_length}), but got {boundaries[-1]}"
236 )
238 attention_batch_size = micro_batch_size if self.reset_mask else 1
239 attention_mask = torch.ones(
240 (attention_batch_size, seq_length, seq_length),
241 dtype=torch.bool,
242 device=device,
243 ).tril()
244 attention_mask = attention_mask.view(attention_batch_size, 1, seq_length, seq_length)
246 if self.reset_mask:
247 for sequence_start in boundaries[:-1]:
248 batch_index, row_sequence_start = divmod(sequence_start, seq_length)
249 attention_mask[batch_index, 0, row_sequence_start:, :row_sequence_start] = False
251 sliding_window_mask = None
252 if self.sliding_window is not None:
253 positions = torch.arange(seq_length, dtype=torch.int64, device=device)
254 outside_window = positions.unsqueeze(1) - positions.unsqueeze(0) > self.sliding_window
255 sliding_window_mask = attention_mask & ~outside_window.unsqueeze(0).unsqueeze(0)
257 return attention_mask, sliding_window_mask
260__all__ = [
261 "AttentionRuntime",
262 "IndexedBoundaryResolver",
263 "OnlineBoundaryResolver",
264 "RuntimeInputAdapter",
265 "RuntimeInputContext",
266]