Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / functional / linear_attention_packed.py: 88%
213 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"""Packed token metadata and causal convolution shared by linear attention."""
17from __future__ import annotations
19from bisect import bisect_left, bisect_right
20from dataclasses import dataclass, field
21from typing import Any, Callable, Mapping, Sequence
23import torch
24from torch import nn
25from torch.nn import functional as F
27from hyper_parallel.components.functional.linear_attention_convolution import (
28 fused_causal_conv1d, validate_conv_backend,
29)
32# Bound retained interval views and their materialized chunk indices. Evicted
33# views can be rebuilt; this cache budget does not limit CP degree or documents.
34_MAX_CACHED_INTERVALS = 32
37@dataclass(frozen=True)
38class PackedChunkPlan:
39 """Read-only kernel indices for one local packed interval and chunk size."""
41 cu_seqlens: torch.Tensor
42 cu_seqlens_cpu: torch.Tensor
43 chunk_indices: torch.Tensor
44 chunk_offsets: torch.Tensor
46 def kwargs(self) -> dict[str, torch.Tensor]:
47 """Return the common FLA varlen arguments without copying device tensors."""
48 return {"cu_seqlens": self.cu_seqlens, "chunk_indices": self.chunk_indices}
51@dataclass(frozen=True, eq=False)
52class PackedDocumentInfo:
53 """Validated immutable document snapshot shared by a batch's interval views.
55 Identity hashing lets bounded CPU-plan caches reuse this snapshot without
56 hashing every document offset at every layer. Cache keys retain the object,
57 so distinct batches cannot collide through reuse of a Python object ID.
58 """
60 offsets: tuple[int, ...]
62 def __post_init__(self) -> None:
63 offsets = tuple(self.offsets)
64 if (
65 len(offsets) < 2
66 or offsets[0] != 0
67 or any(not isinstance(value, int) or isinstance(value, bool) for value in offsets)
68 or any(left >= right for left, right in zip(offsets, offsets[1:]))
69 ):
70 raise ValueError("Packed sample offsets must be strictly increasing integers starting at zero.")
71 object.__setattr__(self, "offsets", offsets)
74@dataclass(frozen=True)
75class PackedSequenceInfo:
76 """Own immutable sample offsets and a contiguous physical token interval.
78 Offsets are global within one CP group, including any explicitly declared
79 padding sample. Local CP cuts are fragments, never new sample resets.
80 Device plans contain read-only tensors; each invocation saves the tensors
81 it consumes so that ordinary autograd version checks remain effective.
82 """
84 global_offsets: tuple[int, ...] | PackedDocumentInfo
85 start: int
86 end: int
87 local_offsets: tuple[int, ...] = field(init=False)
88 documents: PackedDocumentInfo = field(init=False, repr=False, compare=False, hash=False)
89 _head_index: int = field(init=False, repr=False, compare=False, hash=False)
90 _tail_index: int = field(init=False, repr=False, compare=False, hash=False)
91 _plans: dict[tuple[torch.device, int], PackedChunkPlan] = field(
92 default_factory=dict, init=False, repr=False, compare=False, hash=False
93 )
94 _intervals: dict[tuple[int, int], PackedSequenceInfo] = field(
95 default_factory=dict, init=False, repr=False, compare=False, hash=False
96 )
98 def __post_init__(self) -> None:
99 documents = (self.global_offsets if isinstance(self.global_offsets, PackedDocumentInfo)
100 else PackedDocumentInfo(self.global_offsets))
101 offsets = documents.offsets
102 if (
103 not isinstance(self.start, int) or isinstance(self.start, bool)
104 or not isinstance(self.end, int) or isinstance(self.end, bool)
105 or not 0 <= self.start < self.end <= offsets[-1]
106 ):
107 raise ValueError("Packed token interval must be nonempty and covered by global sample offsets.")
108 head = bisect_right(offsets, self.start) - 1
109 tail = bisect_left(offsets, self.end) - 1
110 local = (0, *(value - self.start for value in offsets[head + 1:tail + 1]), self.end - self.start)
111 object.__setattr__(self, "global_offsets", offsets)
112 object.__setattr__(self, "local_offsets", local)
113 object.__setattr__(self, "documents", documents)
114 object.__setattr__(self, "_head_index", head)
115 object.__setattr__(self, "_tail_index", tail)
117 @property
118 def length(self) -> int:
119 """Return this interval's physical token count."""
120 return self.end - self.start
122 @property
123 def num_sequences(self) -> int:
124 """Return the number of nonempty local sample fragments."""
125 return len(self.local_offsets) - 1
127 @property
128 def continues_previous(self) -> bool:
129 """Whether the first fragment consumes the preceding interval's state."""
130 return self.start != self.global_offsets[self._head_index]
132 @property
133 def continues_next(self) -> bool:
134 """Whether the last fragment feeds the following interval's state."""
135 return self.end != self.global_offsets[self._tail_index + 1]
137 @property
138 def transfer_active(self) -> bool:
139 """Whether the preceding interval can affect this interval's exit."""
140 return self.continues_previous and self.num_sequences == 1
142 @property
143 def head(self) -> slice:
144 """Return the first fragment's local token slice."""
145 return slice(0, self.local_offsets[1])
147 @property
148 def tail(self) -> slice:
149 """Return the last fragment's local token slice."""
150 return slice(self.local_offsets[-2], self.length)
152 @property
153 def preceding_sample_tokens(self) -> int:
154 """Return how many earlier tokens belong to the first local sample."""
155 return self.start - self.global_offsets[self._head_index]
157 def interval(self, start: int, end: int) -> PackedSequenceInfo:
158 """Describe another contiguous interval without retaining artificial cuts."""
159 key = (start, end)
160 if key == (self.start, self.end):
161 return self
162 if key not in self._intervals:
163 if len(self._intervals) == _MAX_CACHED_INTERVALS:
164 self._intervals.pop(next(iter(self._intervals)))
165 self._intervals[key] = PackedSequenceInfo(self.documents, start, end)
166 return self._intervals[key]
168 def ulysses(self, degree: int, local_rank: int) -> PackedSequenceInfo:
169 """Describe the token interval after consecutive-rank AllToAll."""
170 if degree < 1 or not 0 <= local_rank < degree:
171 raise ValueError("Invalid Ulysses degree or local rank.")
172 start = self.start - local_rank * self.length
173 return self.interval(start, start + degree * self.length)
175 def materialize(self, device: torch.device, chunk_size: int) -> PackedChunkPlan:
176 """Prepare immutable indices once per batch, layout, device and chunk size."""
177 device = torch.device(device)
178 cache_key = (device, chunk_size)
179 if cache_key in self._plans:
180 return self._plans[cache_key]
181 if chunk_size < 1:
182 raise ValueError("chunk_size must be positive.")
183 counts = [
184 (end - start + chunk_size - 1) // chunk_size
185 for start, end in zip(self.local_offsets, self.local_offsets[1:])
186 ]
187 indices = [(sequence, chunk) for sequence, count in enumerate(counts) for chunk in range(count)]
188 offsets = [0]
189 for count in counts:
190 offsets.append(offsets[-1] + count)
191 cpu = torch.tensor(self.local_offsets, dtype=torch.int64)
192 plan = PackedChunkPlan(
193 cpu.to(device),
194 cpu,
195 torch.tensor(indices, device=device, dtype=torch.int64),
196 torch.tensor(offsets, device=device, dtype=torch.int64),
197 )
198 self._plans[cache_key] = plan
199 return plan
201 def first_state(self, state: torch.Tensor | None) -> torch.Tensor | None:
202 """Expand only a continued first fragment's state for a varlen kernel."""
203 if state is None or not self.continues_previous:
204 return None
205 if self.num_sequences == 1:
206 return state
207 return torch.cat((state, state.new_zeros((self.num_sequences - 1, *state.shape[1:]))))
209 def last_state(self, state: torch.Tensor | None) -> torch.Tensor | None:
210 """Expand a terminal state gradient into the last fragment slot."""
211 if state is None or not self.continues_next:
212 return None
213 if self.num_sequences == 1:
214 return state
215 return torch.cat((state.new_zeros((self.num_sequences - 1, *state.shape[1:])), state))
218def _offset_tuple(value: torch.Tensor | Sequence[int]) -> tuple[int, ...]:
219 """Snapshot user metadata so a later batch cannot mutate saved semantics."""
220 if isinstance(value, torch.Tensor):
221 if value.ndim != 1 or value.dtype not in (torch.int32, torch.int64):
222 raise ValueError("Cumulative sample offsets must be a one-dimensional integer tensor.")
223 return tuple(value.detach().cpu().tolist())
224 return tuple(value)
227def _global_offsets(info: PackedSequenceInfo | None, aliases: list) -> tuple[int, ...]:
228 """Check alias agreement before constructing or reusing local metadata."""
229 if info is not None and not isinstance(info, PackedSequenceInfo):
230 raise TypeError("packed_sequence_info must be PackedSequenceInfo.")
231 offsets = info.global_offsets if info is not None else _offset_tuple(aliases[0])
232 # Runtime q/k aliases share this immutable tuple; comparing the same tuple
233 # still walks its elements in Python, so skip only exact object identity.
234 if any(_offset_tuple(value) != offsets for value in aliases if value is not offsets):
235 raise ValueError("Packed sequence aliases disagree; provide one consistent set of global sample offsets.")
236 return offsets
239def resolve_packed_sequence(
240 arguments: Mapping[str, Any],
241 local_length: int,
242 *,
243 cp_rank: int = 0,
244 cp_size: int = 1,
245) -> PackedSequenceInfo | None:
246 """Normalize model-level aliases, all of which describe GLOBAL sample offsets.
248 Kernel-level cu_seqlens is derived from the returned local interval. Prefer
249 a prebuilt packed_sequence_info from the runtime adapter to avoid host reads
250 or per-layer validation of device metadata.
251 """
252 info = arguments.get("packed_sequence_info")
253 aliases = [
254 arguments[name]
255 for name in ("cu_seq_lens", "cu_seqlens", "cu_seq_lens_q", "cu_seq_lens_k")
256 if arguments.get(name) is not None
257 ]
258 if arguments.get("packed_seq_params") is not None or arguments.get("seq_idx") is not None:
259 raise ValueError(
260 "Use global cu_seq_lens or PackedSequenceInfo; implicit packed/seq_idx conversion is unsupported."
261 )
262 if info is None and not aliases:
263 return None
264 if not 0 <= cp_rank < cp_size:
265 raise ValueError("Invalid CP rank or size for packed sequence metadata.")
266 offsets = _global_offsets(info, aliases)
267 if info is None:
268 info = PackedSequenceInfo(offsets, cp_rank * local_length, (cp_rank + 1) * local_length)
269 if (
270 info.length != local_length
271 or info.start != cp_rank * local_length
272 or info.global_offsets[-1] != cp_size * local_length
273 ):
274 raise ValueError("Packed metadata must cover the chronological, equal-length CP token intervals.")
275 return info
278class _PackedConvTap(torch.autograd.Function):
279 """Recreate a masked tap in backward instead of retaining its FP32 copy."""
281 @staticmethod
282 def forward(
283 ctx: Any, extended: torch.Tensor, weight: torch.Tensor, valid: torch.Tensor, offset: int, length: int,
284 ) -> torch.Tensor:
285 """Save shared input storage while preserving the eager product precision."""
286 ctx.save_for_backward(extended, weight, valid)
287 ctx.offset, ctx.length = offset, length
288 values = extended[:, offset : offset + length].to(weight.dtype)
289 values = torch.where(valid[None, :, None], values, 0.0)
290 return values * weight[None, None, :]
292 @staticmethod
293 def backward(
294 ctx: Any, grad_output: torch.Tensor,
295 ) -> tuple[torch.Tensor | None, torch.Tensor | None, None, None, None]:
296 """Rebuild only this tap; differentiable Torch operations retain higher derivatives."""
297 extended, weight, valid = ctx.saved_tensors
298 grad_extended = grad_weight = None
299 if ctx.needs_input_grad[0]:
300 grad_values = grad_output * weight[None, None, :]
301 grad_values = torch.where(valid[None, :, None], grad_values, 0.0).to(extended.dtype)
302 grad_extended = F.pad(grad_values, (0, 0, ctx.offset, extended.shape[1] - ctx.offset - ctx.length))
303 if ctx.needs_input_grad[1]:
304 values = extended[:, ctx.offset : ctx.offset + ctx.length].to(weight.dtype)
305 values = torch.where(valid[None, :, None], values, 0.0)
306 grad_weight = (grad_output * values).sum(dim=(0, 1))
307 return grad_extended, grad_weight, None, None, None
310def packed_causal_conv1d(
311 inputs: torch.Tensor,
312 convolution: nn.Conv1d,
313 info: PackedSequenceInfo,
314 halo: torch.Tensor | None = None,
315 *,
316 weight: torch.Tensor | None = None,
317 bias: torch.Tensor | None = None,
318 backend: str = "eager",
319) -> torch.Tensor:
320 """Apply depthwise causal SiLU convolution with exact sample isolation.
322 Vectorize over tokens rather than loop over samples. Accumulate low-precision
323 products in FP32; keep FP32/FP64 reference inputs at their original precision.
324 Masking taps also masks their autograd contribution to an exchanged halo.
325 Backward reconstructs each masked tap from the shared input rather than
326 retaining a full accumulation-precision input for every window position.
328 Args:
329 inputs: Local projected tokens with shape ``[1, T, C]``.
330 convolution: Source depthwise, stride-one convolution attributes and parameters.
331 info: True document boundaries and the current local token interval.
332 halo: Optional already exchanged causal history; this function adds no communication.
333 weight: Optional local parameter view with shape ``[C, 1, W]``.
334 bias: Optional local bias; used with an explicit weight override.
335 backend: ``eager`` by default, or explicitly ``triton`` / ``ascendc``
336 for optional NPU kernels. The selected backend never silently falls back.
338 Returns:
339 SiLU output with the input shape and dtype.
341 Raises:
342 ValueError: If metadata, depthwise shape, or the selected kernel contract is invalid.
343 ImportError: If an explicitly selected optional kernel package is unavailable.
344 """
345 validate_conv_backend(backend)
346 if inputs.ndim != 3 or inputs.shape[:2] != (1, info.length):
347 raise ValueError("Packed convolution expects [1,T,C] matching its token interval.")
348 channels = inputs.shape[-1]
349 if (
350 convolution.groups != convolution.in_channels
351 or convolution.out_channels != convolution.in_channels
352 or convolution.stride != (1,)
353 ):
354 raise ValueError("Packed convolution requires a stride-one depthwise Conv1d.")
355 if weight is None:
356 weight, bias = convolution.weight, convolution.bias
357 width, dilation = convolution.kernel_size[0], convolution.dilation[0]
358 if weight.shape != (channels, 1, width):
359 raise ValueError("Packed convolution weight must match the local head channels.")
360 halo_width = (width - 1) * dilation
361 if halo is not None and halo.shape != (1, halo_width, channels):
362 raise ValueError("Packed convolution halo has an incompatible shape.")
363 if backend != "eager":
364 return fused_causal_conv1d(
365 inputs, weight, bias, info, halo, backend=backend, dilation=dilation,
366 )
367 if halo is None:
368 halo = inputs.new_zeros((1, halo_width, channels))
369 extended = torch.cat((halo, inputs), dim=1)
370 positions = torch.arange(info.start, info.end, device=inputs.device)
371 boundaries = torch.tensor(info.global_offsets, device=inputs.device)
372 sample_indices = torch.bucketize(positions, boundaries[1:], right=True)
373 starts = boundaries[sample_indices]
374 dtype = torch.float32 if inputs.dtype in (torch.bfloat16, torch.float16) else inputs.dtype
375 output = torch.zeros_like(inputs, dtype=dtype)
376 weights = weight[:, 0].to(dtype)
377 for tap in range(width):
378 valid = positions - (width - 1 - tap) * dilation >= starts
379 output = output + _PackedConvTap.apply(extended, weights[:, tap], valid, tap * dilation, info.length)
380 if bias is not None:
381 output = output + bias.to(dtype)
382 return F.silu(output.to(inputs.dtype))
385def packed_reference(
386 operator: Callable[..., tuple[torch.Tensor, torch.Tensor | None]],
387 inputs: tuple[torch.Tensor, ...],
388 cu_seqlens: torch.Tensor | Sequence[int],
389 *,
390 initial_state: torch.Tensor | None = None,
391 output_final_state: bool = False,
392 **kwargs: Any,
393) -> tuple[torch.Tensor, torch.Tensor | None]:
394 """Evaluate the portable reference independently per sample fragment."""
395 if inputs[0].shape[0] != 1:
396 raise ValueError("Packed reference requires batch dimension one.")
397 info = PackedSequenceInfo(_offset_tuple(cu_seqlens), 0, inputs[0].shape[1])
398 if info.global_offsets[-1] != info.length:
399 raise ValueError("Local cumulative offsets must end at the input token count.")
400 if initial_state is not None and initial_state.shape[0] != info.num_sequences:
401 raise ValueError("Packed initial_state requires one slot per local sample fragment.")
402 outputs, finals = [], []
403 for index, (start, end) in enumerate(zip(info.local_offsets, info.local_offsets[1:])):
404 state = None if initial_state is None else initial_state[index : index + 1]
405 output, final = operator(
406 *(tensor[:, start:end] for tensor in inputs),
407 initial_state=state,
408 output_final_state=output_final_state,
409 **kwargs,
410 )
411 outputs.append(output)
412 if output_final_state:
413 finals.append(final)
414 return torch.cat(outputs, dim=1), torch.cat(finals) if output_final_state else None