"""Ragged Ulysses sequence-parallel transport for H3 inference.""" from __future__ import annotations from dataclasses import dataclass import torch import torch.distributed as dist def balanced_ranges(total: int, parts: int) -> tuple[tuple[int, int], ...]: """Split ``total`` ordered items into balanced contiguous non-empty ranges.""" if parts < 1: raise ValueError("parts must be positive") if total < parts: raise ValueError(f"cannot split {total} items into {parts} non-empty ranges") base, extra = divmod(total, parts) lengths = [base + (rank < extra) for rank in range(parts)] ranges = [] start = 0 for length in lengths: stop = start + int(length) ranges.append((start, stop)) start = stop return tuple(ranges) def range_lengths(ranges: tuple[tuple[int, int], ...]) -> tuple[int, ...]: return tuple(stop - start for start, stop in ranges) def localize_segments( segments: list[tuple[int, int, int]], shard_start: int, shard_stop: int, ) -> list[tuple[int, int, int]]: """Clip global H3 AdaLN segments to one contiguous token shard.""" localized = [] for start, stop, row in segments: local_start = max(start, shard_start) local_stop = min(stop, shard_stop) if local_start < local_stop: localized.append((local_start - shard_start, local_stop - shard_start, row)) return localized @dataclass(frozen=True) class SequenceParallelContext: """One rank's ragged token and attention-head ownership.""" group: dist.ProcessGroup | None rank: int world_size: int token_ranges: tuple[tuple[int, int], ...] head_ranges: tuple[tuple[int, int], ...] head_dim: int @classmethod def create( cls, sequence_length: int, heads: int, head_dim: int, *, group: dist.ProcessGroup | None = None, ) -> "SequenceParallelContext": if not dist.is_initialized(): raise RuntimeError("torch.distributed process group is not initialized") world_size = dist.get_world_size(group) rank = dist.get_rank(group) return cls( group=group, rank=rank, world_size=world_size, token_ranges=balanced_ranges(sequence_length, world_size), head_ranges=balanced_ranges(heads, world_size), head_dim=head_dim, ) @property def sequence_length(self) -> int: return self.token_ranges[-1][1] @property def heads(self) -> int: return self.head_ranges[-1][1] @property def token_lengths(self) -> tuple[int, ...]: return range_lengths(self.token_ranges) @property def head_lengths(self) -> tuple[int, ...]: return range_lengths(self.head_ranges) @property def local_token_range(self) -> tuple[int, int]: return self.token_ranges[self.rank] @property def local_head_range(self) -> tuple[int, int]: return self.head_ranges[self.rank] @property def local_token_length(self) -> int: start, stop = self.local_token_range return stop - start @property def local_head_count(self) -> int: start, stop = self.local_head_range return stop - start def localize_segments(self, segments: list[tuple[int, int, int]]) -> list[tuple[int, int, int]]: return localize_segments(segments, *self.local_token_range) def seq_to_heads( self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Exchange local tokens for full-sequence Q/K/V over locally owned heads. Inputs use BSHD layout ``[1, local_tokens, all_heads, head_dim]``. Outputs use ``[1, all_tokens, local_heads, head_dim]``. """ expected = (1, self.local_token_length, self.heads, self.head_dim) if tuple(q.shape) != expected or tuple(k.shape) != expected or tuple(v.shape) != expected: raise ValueError( f"sequence-parallel Q/K/V must each have shape {expected}; " f"got {tuple(q.shape)}, {tuple(k.shape)}, {tuple(v.shape)}" ) if q.dtype != k.dtype or q.dtype != v.dtype or q.device != k.device or q.device != v.device: raise ValueError("sequence-parallel Q/K/V must share dtype and device") if self.world_size == 1: return q, k, v local_qkv = torch.stack((q[0], k[0], v[0]), dim=1) send_chunks = [ local_qkv[:, :, start:stop, :].contiguous().view(-1) for start, stop in self.head_ranges ] input_splits = [chunk.numel() for chunk in send_chunks] send = torch.cat(send_chunks) output_splits = [ token_length * 3 * self.local_head_count * self.head_dim for token_length in self.token_lengths ] receive = torch.empty(sum(output_splits), dtype=q.dtype, device=q.device) dist.all_to_all_single( receive, send, output_split_sizes=output_splits, input_split_sizes=input_splits, group=self.group, ) source_chunks = [] offset = 0 for token_length, count in zip(self.token_lengths, output_splits, strict=True): source_chunks.append( receive[offset : offset + count].view( token_length, 3, self.local_head_count, self.head_dim, ) ) offset += count full_qkv = torch.cat(source_chunks, dim=0) full_q, full_k, full_v = full_qkv.unbind(dim=1) return full_q.unsqueeze(0), full_k.unsqueeze(0), full_v.unsqueeze(0) def heads_to_seq(self, output: torch.Tensor) -> torch.Tensor: """Exchange full-sequence local-head output back to local-token all-head output. Input is BSHD ``[1, all_tokens, local_heads, head_dim]``. The return value is ``[local_tokens, all_heads, head_dim]``. """ expected = (1, self.sequence_length, self.local_head_count, self.head_dim) if tuple(output.shape) != expected: raise ValueError(f"sequence-parallel output must have shape {expected}, got {tuple(output.shape)}") if self.world_size == 1: return output[0] output = output[0] send_chunks = [] input_splits = [] token_offset = 0 for token_length in self.token_lengths: chunk = output[token_offset : token_offset + token_length].contiguous().view(-1) send_chunks.append(chunk) input_splits.append(chunk.numel()) token_offset += token_length send = torch.cat(send_chunks) output_splits = [ self.local_token_length * head_length * self.head_dim for head_length in self.head_lengths ] receive = torch.empty(sum(output_splits), dtype=output.dtype, device=output.device) dist.all_to_all_single( receive, send, output_split_sizes=output_splits, input_split_sizes=input_splits, group=self.group, ) head_chunks = [] offset = 0 for head_length, count in zip(self.head_lengths, output_splits, strict=True): head_chunks.append( receive[offset : offset + count].view( self.local_token_length, head_length, self.head_dim, ) ) offset += count return torch.cat(head_chunks, dim=1).contiguous() def target_intersection(self, target_start: int, target_stop: int) -> tuple[int, int]: """Return one global target span's bounds relative to this token shard.""" shard_start, shard_stop = self.local_token_range start = max(target_start, shard_start) stop = min(target_stop, shard_stop) if start >= stop: return (0, 0) return (start - shard_start, stop - shard_start) def target_counts(self, target_start: int, target_stop: int) -> tuple[int, ...]: """Return ordered target-row counts contributed by every token rank.""" counts = [] for shard_start, shard_stop in self.token_ranges: counts.append(max(0, min(target_stop, shard_stop) - max(target_start, shard_start))) return tuple(counts) def all_gather_target_rows( self, local_rows: torch.Tensor, target_start: int, target_stop: int, ) -> torch.Tensor: """Gather a global target span's projected rows onto every rank. Padding is transport-only and is removed before concatenation; it is never exposed to attention or model semantics. """ counts = self.target_counts(target_start, target_stop) if local_rows.ndim != 2: raise ValueError("target rows must be rank-2 [rows, features]") if local_rows.shape[0] != counts[self.rank]: raise ValueError( f"rank {self.rank} must contribute {counts[self.rank]} target rows, " f"got {local_rows.shape[0]}" ) if self.world_size == 1: return local_rows max_rows = max(counts) padded = torch.zeros( max_rows, local_rows.shape[1], dtype=local_rows.dtype, device=local_rows.device, ) if local_rows.shape[0]: padded[: local_rows.shape[0]].copy_(local_rows) gathered = [torch.empty_like(padded) for _ in range(self.world_size)] dist.all_gather(gathered, padded, group=self.group) return torch.cat( [rows[:count] for rows, count in zip(gathered, counts, strict=True) if count], dim=0, ) def all_gather_rows(self, local_rows: torch.Tensor) -> torch.Tensor: """Gather ragged token rows on every rank without exposing padding to the model.""" if local_rows.shape[0] != self.local_token_length: raise ValueError( f"rank {self.rank} must contribute {self.local_token_length} rows, " f"got {local_rows.shape[0]}" ) if self.world_size == 1: return local_rows max_rows = max(self.token_lengths) padded = torch.zeros( (max_rows, *local_rows.shape[1:]), dtype=local_rows.dtype, device=local_rows.device, ) padded[: local_rows.shape[0]].copy_(local_rows) gathered = [torch.empty_like(padded) for _ in range(self.world_size)] dist.all_gather(gathered, padded, group=self.group) return torch.cat( [rows[:count] for rows, count in zip(gathered, self.token_lengths, strict=True)], dim=0, ) def reduce_scatter_rows(self, partial_full_rows: torch.Tensor) -> torch.Tensor: """Sum tensor-parallel partials and return this rank's ragged token rows.""" if partial_full_rows.shape[0] != self.sequence_length: raise ValueError( f"partial rows must cover sequence length {self.sequence_length}, " f"got {partial_full_rows.shape[0]}" ) if self.world_size == 1: return partial_full_rows trailing_shape = partial_full_rows.shape[1:] row_width = partial_full_rows[0].numel() send_chunks = [] input_splits = [] offset = 0 for token_length in self.token_lengths: chunk = partial_full_rows[offset : offset + token_length].contiguous().view(-1) send_chunks.append(chunk) input_splits.append(chunk.numel()) offset += token_length send = torch.cat(send_chunks) output_splits = [self.local_token_length * row_width] * self.world_size receive = torch.empty(sum(output_splits), dtype=send.dtype, device=send.device) dist.all_to_all_single( receive, send, output_split_sizes=output_splits, input_split_sizes=input_splits, group=self.group, ) contributions = receive.view(self.world_size, self.local_token_length, *trailing_shape) return contributions.sum(dim=0)