runtime/src/resynthesis/sequence_parallel.py

"""Exact tensor boundaries for sequence-parallel Resynthesis operators.

KDA is a recurrent affine state transform.  A Python loop that runs shard 0,
then passes its final state to shard 1, is exact but is not context parallel.
This module therefore exposes the actual associative segment algebra needed by
KDA context parallelism and keeps cross-rank transport outside the model-owned
math boundary.  No single-device path claims that it exercised multiple GPUs.

The older generic LASP+ runner remains as a compatibility boundary, but it now
executes one full native kernel launch.  It never changes model routing from an
environment flag.  USP remains an explicitly diagnostic softmax canary.
"""

from __future__ import annotations

import os
from collections.abc import Callable
from dataclasses import dataclass
from typing import TypeVar, cast

import torch

_T = TypeVar("_T", bound=torch.Tensor)


@dataclass(frozen=True)
class ResynthesisKDAAffineSegmentTensorPacket:
    """One or more exact affine KDA segment transforms.

    For each segment, ``transition_t`` and ``source_t`` encode
    ``state_out = transition_t @ state_in + source_t``.  Both fields retain a
    leading segment axis and are ordinary differentiable tensors, so a real
    transport owner can all-gather/scan them without serializing Python state.
    """

    transition_t: torch.Tensor
    source_t: torch.Tensor


@dataclass(frozen=True)
class ResynthesisKDAContextParallelTensorPacket:
    """Tensor-only exact KDA segment and inclusive-prefix authority."""

    segment_transition_t: torch.Tensor
    segment_source_t: torch.Tensor
    prefix_transition_t: torch.Tensor
    prefix_source_t: torch.Tensor
    initial_states_t: torch.Tensor


def lasp_plus_enabled_boundary() -> bool:
    """Observe a legacy request for receipts; never route model compute."""

    raw = os.environ.get("NNF_RESYNTHESIS_LASP_PLUS", "0")
    return raw in {"1", "true", "True", "yes", "on"}


def usp_enabled_boundary() -> bool:
    raw = os.environ.get("NNF_RESYNTHESIS_USP", "0")
    return raw in {"1", "true", "True", "yes", "on"}


def sequence_parallel_world_size_boundary() -> int:
    raw = os.environ.get("NNF_RESYNTHESIS_SEQUENCE_PARALLEL_WORLD_SIZE", "2")
    try:
        world_size = int(raw)
    except ValueError as error:
        raise RuntimeError(
            "NNF_RESYNTHESIS_SEQUENCE_PARALLEL_WORLD_SIZE must be an integer"
        ) from error
    if world_size < 1:
        raise RuntimeError(
            "NNF_RESYNTHESIS_SEQUENCE_PARALLEL_WORLD_SIZE must be positive"
        )
    return world_size


def sequence_parallel_audit_boundary() -> dict[str, object]:
    return {
        "schema": "nnf.resynthesis.sequence_parallel_audit.v2",
        "legacyLaspPlusRequested": lasp_plus_enabled_boundary(),
        "legacySequentialKdaCanaryActive": False,
        "exactKdaAssociativeSegmentComposition": True,
        "crossRankTransportClaimed": False,
        "uspEnabled": usp_enabled_boundary(),
        "worldSize": sequence_parallel_world_size_boundary(),
        "productionDefault": False,
    }


def _validate_kda_segment_geometry(
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
) -> None:
    if k.ndim != 5 or v.ndim != 5 or g.ndim != 5 or beta.ndim != 4:
        raise ValueError("KDA affine segments require a leading segment axis")
    if any(
        width < 1
        for width in (
            k.shape[0],
            k.shape[1],
            k.shape[2],
            k.shape[3],
            k.shape[4],
            v.shape[4],
        )
    ):
        raise ValueError("KDA affine segments require nonempty tensor geometry")
    if k.shape != g.shape or k.shape[:-1] != v.shape[:-1]:
        raise ValueError("KDA affine segment K/V/decay geometry differs")
    if beta.shape != k.shape[:-1]:
        raise ValueError("KDA affine segment write-gate geometry differs")
    if k.device != v.device or k.device != g.device or k.device != beta.device:
        raise ValueError("KDA affine segment tensors occupy different devices")
    if not k.is_floating_point() or not v.is_floating_point():
        raise TypeError("KDA affine segments require floating-point K/V tensors")
    if not g.is_floating_point() or not beta.is_floating_point():
        raise TypeError("KDA affine segments require floating-point gates")


def kda_affine_segments_boundary(
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
) -> ResynthesisKDAAffineSegmentTensorPacket:
    """Build exact differentiable affine transforms for gathered KDA segments.

    Inputs use ``[segments,batch,tokens,heads,width]`` except ``beta``, whose
    shape is ``[segments,batch,tokens,heads]``.  State algebra is accumulated in
    FP32 (FP64 when the inputs are FP64), matching KDA's stable recurrent-state
    contract while preserving gradients to every input tensor.
    """

    _validate_kda_segment_geometry(k, v, g, beta)
    segments, batch, tokens, heads, key_dim = k.shape
    value_dim = v.shape[-1]
    state_dtype = torch.float64 if k.dtype == torch.float64 else torch.float32
    k_state_t = k.to(dtype=state_dtype)
    v_state_t = v.to(dtype=state_dtype)
    g_state_t = g.to(dtype=state_dtype)
    beta_state_t = beta.to(dtype=state_dtype)
    identity_t = torch.eye(
        key_dim,
        device=k.device,
        dtype=state_dtype,
    ).view(1, 1, 1, key_dim, key_dim)
    transition_t = identity_t.expand(
        segments,
        batch,
        heads,
        key_dim,
        key_dim,
    )
    source_t = k_state_t.new_zeros(
        segments,
        batch,
        heads,
        key_dim,
        value_dim,
    )
    for token_index in range(tokens):
        key_t = k_state_t[:, :, token_index]
        value_t = v_state_t[:, :, token_index]
        decay_t = g_state_t[:, :, token_index].exp()
        write_t = beta_state_t[:, :, token_index]
        key_outer_t = key_t.unsqueeze(-1) * key_t.unsqueeze(-2)
        token_transition_t = (
            identity_t - write_t.unsqueeze(-1).unsqueeze(-1) * key_outer_t
        ) * decay_t.unsqueeze(-2)
        token_source_t = (
            write_t.unsqueeze(-1).unsqueeze(-1)
            * key_t.unsqueeze(-1)
            * value_t.unsqueeze(-2)
        )
        source_t = torch.matmul(token_transition_t, source_t) + token_source_t
        transition_t = torch.matmul(token_transition_t, transition_t)
    return ResynthesisKDAAffineSegmentTensorPacket(
        transition_t=transition_t,
        source_t=source_t,
    )


def kda_compose_affine_segments_boundary(
    upstream: ResynthesisKDAAffineSegmentTensorPacket,
    downstream: ResynthesisKDAAffineSegmentTensorPacket,
) -> ResynthesisKDAAffineSegmentTensorPacket:
    """Compose exact KDA transforms in sequence order.

    ``upstream`` executes first and ``downstream`` second.  The operation is
    associative, which is the mathematical property required by a real
    all-gather plus parallel-prefix KDA context-parallel implementation.
    """

    if (
        upstream.transition_t.shape != downstream.transition_t.shape
        or upstream.source_t.shape != downstream.source_t.shape
        or upstream.transition_t.device != downstream.transition_t.device
        or upstream.source_t.device != downstream.source_t.device
        or upstream.transition_t.dtype != downstream.transition_t.dtype
        or upstream.source_t.dtype != downstream.source_t.dtype
    ):
        raise ValueError("KDA affine composition geometry differs")
    transition_t = torch.matmul(
        downstream.transition_t,
        upstream.transition_t,
    )
    source_t = (
        torch.matmul(downstream.transition_t, upstream.source_t)
        + downstream.source_t
    )
    return ResynthesisKDAAffineSegmentTensorPacket(
        transition_t=transition_t,
        source_t=source_t,
    )


def kda_associative_prefix_boundary(
    segments: ResynthesisKDAAffineSegmentTensorPacket,
) -> ResynthesisKDAAffineSegmentTensorPacket:
    """Return the inclusive KDA prefix using an exact doubling scan."""

    transition_t = segments.transition_t
    source_t = segments.source_t
    if transition_t.ndim != 5 or source_t.ndim != 5:
        raise ValueError("KDA affine prefix requires a segment axis")
    if (
        transition_t.shape[0] != source_t.shape[0]
        or transition_t.shape[1:3] != source_t.shape[1:3]
        or transition_t.shape[-1] != transition_t.shape[-2]
        or transition_t.shape[-1] != source_t.shape[-2]
    ):
        raise ValueError("KDA affine prefix geometry differs")
    stride = 1
    segment_count = transition_t.shape[0]
    while stride < segment_count:
        downstream_t = transition_t[stride:]
        downstream_source_t = source_t[stride:]
        composed_transition_t = torch.matmul(
            downstream_t,
            transition_t[:-stride],
        )
        composed_source_t = (
            torch.matmul(downstream_t, source_t[:-stride])
            + downstream_source_t
        )
        transition_t = torch.cat(
            (transition_t[:stride], composed_transition_t),
            dim=0,
        )
        source_t = torch.cat(
            (source_t[:stride], composed_source_t),
            dim=0,
        )
        stride *= 2
    return ResynthesisKDAAffineSegmentTensorPacket(
        transition_t=transition_t,
        source_t=source_t,
    )


def kda_context_parallel_packet_boundary(
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    *,
    initial_state: torch.Tensor | None = None,
) -> ResynthesisKDAContextParallelTensorPacket:
    """Build exact incoming states for already-gathered KDA segments.

    This boundary deliberately does not create a process group or infer that
    data-parallel ranks are context-parallel ranks.  A real transport owner
    must provide gathered segment tensors in sequence order; this function then
    supplies the exact, differentiable prefix math with no sequential carry.
    """

    segments = kda_affine_segments_boundary(k, v, g, beta)
    prefix = kda_associative_prefix_boundary(segments)
    segment_count, batch, heads, key_dim, value_dim = (
        segments.source_t.shape
    )
    if initial_state is None:
        first_state_t = segments.source_t.new_zeros(
            batch,
            heads,
            key_dim,
            value_dim,
        )
    else:
        expected_shape = (batch, heads, key_dim, value_dim)
        if initial_state.shape != expected_shape:
            raise ValueError("KDA context-parallel initial-state geometry differs")
        first_state_t = initial_state.to(
            device=segments.source_t.device,
            dtype=segments.source_t.dtype,
        )
    if segment_count == 1:
        initial_states_t = first_state_t.unsqueeze(0)
    else:
        later_states_t = (
            torch.matmul(prefix.transition_t[:-1], first_state_t.unsqueeze(0))
            + prefix.source_t[:-1]
        )
        initial_states_t = torch.cat(
            (first_state_t.unsqueeze(0), later_states_t),
            dim=0,
        )
    return ResynthesisKDAContextParallelTensorPacket(
        segment_transition_t=segments.transition_t,
        segment_source_t=segments.source_t,
        prefix_transition_t=prefix.transition_t,
        prefix_source_t=prefix.source_t,
        initial_states_t=initial_states_t,
    )


def _active_world_size(*, feature_enabled: bool) -> int:
    if not feature_enabled:
        return 1
    return sequence_parallel_world_size_boundary()


def lasp_plus_shard_sequence_boundary(
    tensor: _T,
    *,
    dim: int = 1,
) -> list[_T]:
    """Split one ``[batch, seq, ...]`` tensor into ring shards."""

    world_size = _active_world_size(feature_enabled=lasp_plus_enabled_boundary())
    if world_size <= 1 or tensor.shape[dim] < world_size:
        return [tensor]
    return [
        cast(_T, shard)
        for shard in tensor.tensor_split(world_size, dim=dim)
    ]


def lasp_plus_gather_sequence_boundary(
    shards: list[torch.Tensor],
    *,
    dim: int = 1,
) -> torch.Tensor:
    """Merge ring shards back along ``dim``."""

    if len(shards) == 1:
        return shards[0]
    return torch.cat(shards, dim=dim)


def lasp_plus_run_sequence_chunks_boundary(
    tensors: tuple[_T, ...],
    runner: Callable[
        ...,
        torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    ],
    *,
    dim: int = 1,
    carry_state: bool = True,
) -> _T:
    """Compatibility boundary that executes one complete native kernel.

    The old implementation split tensors according to host environment flags
    and passed recurrent state through a Python loop.  Since that is not LASP+
    or context parallelism, the compatibility surface now preserves the native
    full-sequence launch regardless of those diagnostic settings.
    """

    def output_tensor(
        result: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
    ) -> torch.Tensor:
        return result[0] if isinstance(result, tuple) else result

    return cast(_T, output_tensor(runner(*tensors)))


def usp_softmax_boundary(
    scores: torch.Tensor,
    *,
    dim: int = -1,
) -> torch.Tensor:
    """Ulysses×Ring-style stable softmax canary on one device."""

    if not usp_enabled_boundary():
        return torch.softmax(scores, dim=dim)
    world_size = sequence_parallel_world_size_boundary()
    if world_size <= 1 or scores.shape[dim] < world_size:
        return torch.softmax(scores, dim=dim)
    parts = scores.tensor_split(world_size, dim=dim)
    local_max = torch.stack(
        [part.amax(dim=dim, keepdim=True) for part in parts],
        dim=0,
    ).amax(dim=0)
    exp_parts = [(part - local_max).exp() for part in parts]
    local_sum = torch.stack(
        [part.sum(dim=dim, keepdim=True) for part in exp_parts],
        dim=0,
    ).sum(dim=0)
    tiny = torch.finfo(scores.dtype).tiny
    return torch.cat(
        [part / local_sum.clamp_min(tiny) for part in exp_parts],
        dim=dim,
    )