runtime/src/resynthesis/quantile_balancing.py · namenotfoundai/Nucleus-Resynthesis at main

Nucleus-Resynthesis / runtime /src /resynthesis /quantile_balancing.py
Add portable Release 188 generation runtime
919fd68 verified 3 months ago
40.9 kB
"""Quantile balancing for tensor-native expert allocation.
Expert biases beta are updated from router-score quantiles so allocation stays
near-uniform without a dead auxiliary load-balance loss. Routing uses the
previous-step bias, making the state update causal.
Reference: https://kexue.fm/archives/11619
"""
from __future__ import annotations

import math
from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING, cast

import torch
import torch.distributed as dist
import torch.nn as nn
import torch.nn.functional as F

if TYPE_CHECKING:
    from resynthesis.anti_systems_bridge import TensorAntiThompsonRegistry


ANTI_THOMPSON_FAIL_COUNTS_BUFFER_SUFFIX = (
    ".quantile_router.anti_thompson_fail_counts_t"
)
ANTI_THOMPSON_FAIL_COUNTS_INITIALIZATION_SCHEME = (
    "zero_persistent_quantile_router_anti_thompson_fail_counts_v1"
)
QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT = 1000


@dataclass(frozen=True)
class QuantileBalancingStepSnapshot:
    """Tensor-native replay state for one uncommitted optimizer transaction."""

    histogram_counts_t: torch.Tensor
    selected_count_t: torch.Tensor
    token_count_t: torch.Tensor


def _apply_fp32_control_buffer_without_narrowing(
    buffer_t: torch.Tensor,
    fn: Callable[[torch.Tensor], torch.Tensor],
) -> torch.Tensor:
    """Move one FP32 learning-control buffer without casting its values.
    ``nn.Module.to(dtype=...)`` applies its dtype conversion to every floating
    buffer, even when a buffer is durable optimizer-like routing state rather
    than activation compute.  Use an integer probe to discover the destination
    device, then move the original FP32 values directly to that device.  This
    avoids the lossy FP32 -> BF16 -> FP32 round trip that would otherwise
    corrupt exact cold-resume authority.
    """

    if buffer_t.dtype != torch.float32:
        return fn(buffer_t)
    if buffer_t.device.type == "meta":
        # ``to_empty`` is the only supported way out of meta state.  There are
        # no resident values to preserve there, but the FP32 ownership remains.
        return fn(buffer_t).to(dtype=torch.float32)
    target_probe_t = fn(
        torch.empty(
            0,
            device=buffer_t.device,
            dtype=torch.uint8,
        )
    )
    return buffer_t.to(
        device=target_probe_t.device,
        dtype=torch.float32,
    )


def quantile_threshold(
    values: torch.Tensor,
    q_t: torch.Tensor,
    dim: int,
) -> torch.Tensor:
    """Return a compile-safe linear quantile with a retained ``q`` gradient.
    ``torch.quantile`` asks for a concrete ``numel`` while Dynamo traces a
    dynamic sequence dimension.  NoNE validation intentionally uses variable
    sequence widths, so perform the same sorted linear interpolation with
    tensor indices.  This avoids a graph break while preserving gradients for
    both the values and the learned quantile level.
    """

    if q_t.numel() != 1:
        raise ValueError("quantile level must be scalar")
    active_q_t = q_t.reshape(()).to(device=values.device, dtype=torch.float32)
    # A diverged learned routing fraction must not become a NaN tensor index
    # below.  Keep this repair tensor-native: the balancing state is diagnostic
    # routing pressure, not an answer or page-proof authority.
    active_q_t = torch.where(
        torch.isfinite(active_q_t),
        active_q_t,
        torch.zeros_like(active_q_t),
    ).clamp(0.0, 1.0)
    ordered = values.float().sort(dim=dim).values
    rank_t = active_q_t * (values.shape[dim] - 1)
    lower_rank_t = rank_t.floor()
    lower_index_t = lower_rank_t.to(dtype=torch.long).reshape(1)
    upper_index_t = rank_t.ceil().to(dtype=torch.long).reshape(1)
    lower_t = torch.index_select(ordered, dim, lower_index_t)
    upper_t = torch.index_select(ordered, dim, upper_index_t)
    return torch.lerp(lower_t, upper_t, rank_t - lower_rank_t).to(
        dtype=values.dtype
    )


def _exact_quantile_frontier_mask(
    adjusted_scores_t: torch.Tensor,
    threshold_t: torch.Tensor,
    frontier_count_t: torch.Tensor,
) -> torch.Tensor:
    """Select the exact stable quantile frontier without a second full sort.
    ``frontier_count_t`` may be widened independently for each uncertain row,
    beyond the learned quantile cut.  Build an exact stable descending rank
    from tensor comparisons: greater scores precede a candidate, and equal
    scores precede it only when their expert index is lower.  This preserves
    model score identity without a host ``topk`` width or a second sort.
    """

    if threshold_t.shape != adjusted_scores_t.shape[:-1] + (1,):
        raise ValueError("quantile frontier threshold geometry differs")
    if (
        frontier_count_t.numel() != 1
        and frontier_count_t.shape != adjusted_scores_t.shape[:-1]
    ):
        raise ValueError("quantile frontier count geometry differs")
    detached_scores_t = adjusted_scores_t.detach()
    score_i_t = detached_scores_t.unsqueeze(-1)
    score_j_t = detached_scores_t.unsqueeze(-2)
    index_t = torch.arange(
        adjusted_scores_t.shape[-1],
        device=adjusted_scores_t.device,
        dtype=torch.long,
    )
    index_i_t = index_t.unsqueeze(-1)
    index_j_t = index_t.unsqueeze(0)
    precedes_t = score_j_t.gt(score_i_t) | (
        score_j_t.eq(score_i_t)
        & index_j_t.lt(index_i_t)
    )
    stable_rank_t = precedes_t.sum(dim=-1)
    active_frontier_count_t = frontier_count_t.to(
        device=adjusted_scores_t.device,
        dtype=torch.long,
    )
    active_frontier_count_t = (
        active_frontier_count_t.reshape(
            *([1] * adjusted_scores_t.dim())
        )
        if active_frontier_count_t.numel() == 1
        else active_frontier_count_t.unsqueeze(-1)
    )
    return stable_rank_t.lt(active_frontier_count_t)


class QuantileBalancingRouter(nn.Module):
    """Persistent expert bias β with differentiable soft quantile routing.
    ``activation_logit`` learns the active share directly from task gradients.
    The share is bounded below by one expert's mass so every token has a valid
    path, but there is no host top-k decision or fixed frontier width.
    """

    expert_bias_t: torch.Tensor
    anti_thompson_fail_counts_t: torch.Tensor
    _quantile_histogram_counts_t: torch.Tensor
    _quantile_histogram_token_count_t: torch.Tensor
    _quantile_histogram_selected_count_t: torch.Tensor

    def __init__(
        self,
        num_experts: int,
        *,
        activation_fraction: float = 0.25,
        temperature: float = 1.0,
        solve_steps: int = 5,
    ) -> None:
        super().__init__()
        experts = max(1, int(num_experts))
        self.num_experts = experts
        frac = min(max(float(activation_fraction), 1.0 / float(experts)), 1.0)
        minimum_fraction = 1.0 / float(experts)
        unit_fraction = (
            (frac - minimum_fraction) / max(1.0 - minimum_fraction, 1.0e-6)
        )
        self.activation_logit = nn.Parameter(
            torch.logit(torch.tensor(unit_fraction).clamp(1.0e-4, 1.0 - 1.0e-4))
        )
        self.temperature_logit = nn.Parameter(
            torch.log(torch.tensor(max(float(temperature), 1.0e-3)))
        )
        self.solve_steps = max(1, int(solve_steps))
        self.register_buffer(
            "expert_bias_t",
            torch.zeros(experts, dtype=torch.float32),
            persistent=True,
        )
        self.register_buffer(
            "anti_thompson_fail_counts_t",
            torch.zeros(experts, dtype=torch.float32),
            persistent=True,
        )
        # The global quantile-balancing update needs only additive histogram
        # counts during an open optimizer step.  These tensors follow the
        # module device but are intentionally absent from checkpoints: an
        # interrupted, uncommitted step is replayed from the durable training
        # cursor, while the causally committed expert bias above is durable.
        self._quantile_histogram_counts_t = torch.zeros(
            experts,
            QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT,
            dtype=torch.long,
        )
        self._quantile_histogram_token_count_t = torch.zeros(
            (),
            dtype=torch.long,
        )
        self._quantile_histogram_selected_count_t = torch.zeros(
            (),
            dtype=torch.long,
        )
        # TRAUMA/MILT state belongs to this exact router/catalog and is part of
        # the model checkpoint.  A zero bank is inert, so route ownership does
        # not depend on a host flag or module-global dictionary.
        from resynthesis.trauma_system import TensorTraumaState

        self.trauma_state = TensorTraumaState(num_arms=experts)
        self.last_hard_mask: torch.Tensor | None = None
        self.last_utilization_t: torch.Tensor | None = None
        self._last_soft_for_loss: torch.Tensor | None = None
        self._last_frontier_for_loss: torch.Tensor | None = None
        self._last_activation_fraction_for_loss: torch.Tensor | None = None

    def rebuild_nonpersistent_buffers(self) -> None:
        """Rebuild step-local histogram tensors after meta-device construction."""

        device = self.expert_bias_t.device
        experts = int(self.expert_bias_t.shape[0])
        self._quantile_histogram_counts_t = torch.zeros(
            experts,
            QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT,
            dtype=torch.long,
            device=device,
        )
        self._quantile_histogram_token_count_t = torch.zeros(
            (),
            dtype=torch.long,
            device=device,
        )
        self._quantile_histogram_selected_count_t = torch.zeros(
            (),
            dtype=torch.long,
            device=device,
        )

    def _apply(
        self,
        fn: Callable[[torch.Tensor], torch.Tensor],
        recurse: bool = True,
    ) -> "QuantileBalancingRouter":
        """Keep durable route posteriors/counts FP32 across model casts.
        Expert activations and trainable matrix weights may run in BF16, but
        quantile bias and anti-Thompson failure counts are cumulative
        learning-control state.  Narrowing them changes accepted checkpoint
        values and eventually starves small decay/update increments.
        """

        transient_step_tensors = {
            name: getattr(self, name)
            for name in (
                "_quantile_histogram_counts_t",
                "_quantile_histogram_token_count_t",
                "_quantile_histogram_selected_count_t",
            )
        }
        fp32_control_buffers = {
            name: buffer_t
            for name in (
                "expert_bias_t",
                "anti_thompson_fail_counts_t",
            )
            if (
                (buffer_t := self._buffers.get(name)) is not None
                and buffer_t.dtype == torch.float32
            )
        }
        result = cast(
            "QuantileBalancingRouter",
            super()._apply(  # type: ignore[no-untyped-call]
                fn,
                recurse=recurse,
            ),
        )
        for name, buffer_t in fp32_control_buffers.items():
            self._buffers[name] = _apply_fp32_control_buffer_without_narrowing(
                buffer_t,
                fn,
            )
        for name, tensor_t in transient_step_tensors.items():
            moved_t = (
                _apply_fp32_control_buffer_without_narrowing(tensor_t, fn)
                if tensor_t.dtype == torch.float32
                else fn(tensor_t)
            )
            setattr(self, name, moved_t)
        # Moving/casting a model is a transaction fence.  A partial histogram
        # cannot cross that boundary because its previous rank/device cohort
        # no longer exists; the durable cursor will replay the uncommitted
        # step with the causally committed beta intact.
        self.begin_expert_bias_step_boundary()
        registry = getattr(self, "_anti_thompson_registry", None)
        if registry is not None:
            self.bind_anti_thompson_registry_boundary(registry)
        return result

    def bind_anti_thompson_registry_boundary(
        self,
        registry: TensorAntiThompsonRegistry | None = None,
    ) -> TensorAntiThompsonRegistry:
        """Bind one runtime registry to this router's persistent fail bank."""

        from resynthesis.anti_systems_bridge import (
            TensorAntiThompsonRegistry,
        )

        candidate = (
            registry
            if registry is not None
            else getattr(self, "_anti_thompson_registry", None)
        )
        if not isinstance(candidate, TensorAntiThompsonRegistry):
            raise RuntimeError(
                "quantile router anti-thompson registry is absent"
            )
        candidate.bind_fail_counts_t(self.anti_thompson_fail_counts_t)
        setattr(self, "_anti_thompson_registry", candidate)
        return candidate

    def set_mitm_branch(self, branch: int) -> None:
        """Compatibility boundary; branch identity never owns route width.
        Older paged routers call this method while constructing a branch-local
        learner.  Keeping the method avoids a migration-only API break, but the
        value is intentionally not retained: route width is learned from
        ``activation_logit`` and the tensor-native uncertainty floor below.
        """

        del branch

    def activation_fraction(self) -> torch.Tensor:
        # Keep the learned routing fraction entirely on the parameter device.
        # ``new_tensor(Python_scalar)`` enters the CUDA copy path and forces a
        # stream synchronization each time the router updates its persistent
        # quantile bias.  A device-native scalar has identical values and
        # gradients without a host-to-device scalar transfer in the hot path.
        minimum = torch.ones_like(self.activation_logit) / self.num_experts
        fraction = minimum + (1.0 - minimum) * torch.sigmoid(
            self.activation_logit
        )
        finite_fraction = torch.where(
            torch.isfinite(fraction),
            fraction,
            minimum,
        )
        return torch.minimum(
            torch.maximum(finite_fraction, minimum),
            torch.ones_like(finite_fraction),
        )

    def frontier_count_t(
        self,
        *,
        uncertain: torch.Tensor | bool | None = None,
    ) -> torch.Tensor:
        """Return the learned finite frontier width as a scalar tensor."""

        learned = (
            1
            + torch.floor(
                self.activation_fraction() * (self.num_experts - 1)
            )
        ).clamp(1, self.num_experts).to(dtype=torch.long)
        if isinstance(uncertain, torch.Tensor):
            uncertain_t = uncertain.to(
                device=learned.device,
                dtype=torch.bool,
            )
        else:
            uncertainty_active = self.training if uncertain is None else uncertain
            uncertain_t = (
                torch.ones_like(learned, dtype=torch.bool)
                if uncertainty_active
                else torch.zeros_like(learned, dtype=torch.bool)
            )
        available_t = torch.ones_like(learned) * self.num_experts
        fractional_floor_t = torch.div(
            available_t + 9,
            10,
            rounding_mode="floor",
        )
        uncertainty_floor_t = torch.minimum(
            available_t,
            torch.maximum(
                fractional_floor_t,
                torch.ones_like(learned) * 16,
            ),
        )
        return torch.where(
            uncertain_t,
            torch.maximum(learned, uncertainty_floor_t),
            learned,
        )

    def temperature(self) -> torch.Tensor:
        temp = self.temperature_logit.exp()
        # Guard against divergence: replace non-finite values with 1.0
        temp = torch.where(torch.isfinite(temp), temp, torch.ones_like(temp))
        return temp.clamp_min(1.0e-3)

    def begin_route_arm_boundary(self) -> torch.Tensor:
        """Release differentiable diagnostics from the preceding route arm.
        A paged training cohort deliberately keeps its tensor-owned hard page
        identity across multiple CUDA waves.  Later waves refine that identity
        without calling ``route`` again, so the prior wave's auxiliary-loss
        tensors must not remain eligible for a second backward traversal.
        Persistent bias, utilization telemetry, and the latched route itself
        are independent state and remain unchanged.
        """

        self._last_soft_for_loss = None
        self._last_frontier_for_loss = None
        self._last_activation_fraction_for_loss = None
        return self.expert_bias_t.new_ones((), dtype=torch.bool)

    def biased_scores(self, scores: torch.Tensor) -> torch.Tensor:
        # The forward owns an immutable view of beta. ``route`` updates the
        # persistent buffer only after gates have been computed for this step.
        bias = self.expert_bias_t.detach().to(
            device=scores.device,
            dtype=scores.dtype,
        ).clone()
        view = bias.view(*([1] * (scores.dim() - 1)), -1)
        return scores - view

    def soft_gates(self, scores: torch.Tensor) -> torch.Tensor:
        """Sparse forward gates with dense gradients through the quantile frontier."""

        if scores.shape[-1] != self.num_experts:
            raise ValueError("quantile balancing score geometry differs")
        score_limit = math.sqrt(torch.finfo(torch.float32).max)
        safe_scores = torch.nan_to_num(
            scores,
            nan=0.0,
            posinf=score_limit,
            neginf=-score_limit,
        ).clamp(-score_limit, score_limit)
        # ``expert_bias_t`` owns dispatch balance only.  Keep its immutable
        # step snapshot separate so the mixture weights below can omit that
        # bias while still retaining model-owned Trauma/MILT/Anti-Thompson
        # score pressure.
        balance_bias_t = self.expert_bias_t.detach().to(
            device=safe_scores.device,
            dtype=safe_scores.dtype,
        ).clone()
        balance_bias_view_t = balance_bias_t.view(
            *([1] * (safe_scores.dim() - 1)),
            -1,
        )
        adj = safe_scores - balance_bias_view_t
        anti_registry = getattr(self, "_anti_thompson_registry", None)
        if anti_registry is not None:
            anti_registry = self.bind_anti_thompson_registry_boundary(
                anti_registry
            )
            anti_t = anti_registry.anti_bias_for_quantile_router().to(
                device=adj.device,
                dtype=adj.dtype,
            )
            if anti_t.shape == (self.num_experts,):
                adj = adj - anti_t.reshape(
                    *([1] * (adj.dim() - 1)),
                    -1,
                )
        from resynthesis.hard_knowledge_router_boundary import (
            mitm_scaffold_logits_delta_t,
        )
        from resynthesis.hard_knowledge_surface import (
            HardKnowledgeSurfacePacket,
            hard_knowledge_surface_packet_t,
        )
        from resynthesis.trauma_system import TensorTraumaState

        packet: HardKnowledgeSurfacePacket | None
        trauma_state = self.trauma_state
        if isinstance(trauma_state, TensorTraumaState):
            packet = hard_knowledge_surface_packet_t(
                trauma_state,
                top_k=self.num_experts,
                device=adj.device,
            )
            setattr(self, "_hard_knowledge_packet", packet)
        else:
            resident_packet = getattr(self, "_hard_knowledge_packet", None)
            packet = (
                resident_packet
                if isinstance(resident_packet, HardKnowledgeSurfacePacket)
                else None
            )
        mitm_delta = mitm_scaffold_logits_delta_t(
            self,
            num_candidates=self.num_experts,
            device=adj.device,
            dtype=adj.dtype,
        )
        if mitm_delta.numel() == adj.shape[-1]:
            adj = adj + mitm_delta.reshape(*([1] * (adj.dim() - 1)), -1)
        if packet is not None:
            surface = packet.definitive_surface_t.to(
                device=adj.device,
                dtype=adj.dtype,
            )
            if surface.numel() == adj.shape[-1]:
                adj = adj + surface.reshape(*([1] * (adj.dim() - 1)), -1)
        adj = torch.nan_to_num(
            adj,
            nan=0.0,
            posinf=score_limit,
            neginf=-score_limit,
        ).clamp(-score_limit, score_limit)
        frac = self.activation_fraction().to(device=scores.device, dtype=scores.dtype)
        # The auxiliary loss belongs to this exact route arm.  Retain the
        # differentiable fraction already used to derive its frontier instead
        # of rebuilding the sigmoid/divide scalar graph for every science and
        # paged router after the CUDA wave completes.
        self._last_activation_fraction_for_loss = frac
        q = (1.0 - frac).clamp(0.0, 1.0)
        threshold = quantile_threshold(adj, q, dim=-1)
        temp = self.temperature().to(device=scores.device, dtype=scores.dtype)
        # Quantile beta selects the sparse frontier but must not change the
        # mixture weights or router gradients inside that frontier.  The
        # remaining score deltas are model-owned routing knowledge and stay in
        # the mixture.  Adding the immutable beta snapshot back removes only
        # load-balancing pressure.
        mixture_scores_t = adj + balance_bias_view_t
        model_probability_t = F.softmax(mixture_scores_t / temp, dim=-1)
        route_uncertainty_t = model_probability_t.detach().amax(dim=-1).lt(0.35)
        frontier_count = self.frontier_count_t(uncertain=route_uncertainty_t)
        # A quantile comparison can admit too many experts when scores tie at
        # the threshold. The threshold sort already establishes every strict
        # winner; tensor-native tie ranks preserve stable expert-index order
        # without redundantly sorting the complete axis a second time.
        if self.num_experts <= 64:
            hard_mask_t = _exact_quantile_frontier_mask(
                adj,
                threshold,
                frontier_count,
            )
        else:
            # Wide banks amortize the CUDA/CPU sort kernel better than the
            # cumulative tie scan. Keep the previous exact stable-rank
            # implementation there; r152's 8-way science and 43-way paged
            # routers take the lower-overhead quantile path above.
            frontier_order_t = torch.argsort(
                adj.detach(),
                dim=-1,
                descending=True,
                stable=True,
            )
            frontier_rank_t = torch.empty_like(frontier_order_t)
            frontier_rank_t.scatter_(
                -1,
                frontier_order_t,
                torch.arange(
                    self.num_experts,
                    device=adj.device,
                    dtype=torch.long,
                )
                .reshape(*([1] * (adj.dim() - 1)), -1)
                .expand_as(frontier_order_t),
            )
            active_frontier_count_t = frontier_count.to(device=adj.device)
            if active_frontier_count_t.numel() != 1:
                active_frontier_count_t = active_frontier_count_t.unsqueeze(-1)
            hard_mask_t = frontier_rank_t.lt(active_frontier_count_t)
        hard = hard_mask_t.to(dtype=adj.dtype)
        self.last_hard_mask = hard.detach()
        util = hard.reshape(-1, self.num_experts).mean(dim=0)
        self.last_utilization_t = util.detach()
        soft_frontier = torch.sigmoid((adj - threshold) / temp)
        # Forward values are exactly sparse while backward derivatives flow
        # through scores, temperature, and the learned activation quantile.
        frontier = hard + soft_frontier - soft_frontier.detach()
        weighted = model_probability_t * frontier
        gates = weighted / weighted.sum(dim=-1, keepdim=True).clamp_min(
            torch.finfo(weighted.dtype).tiny
        )
        self._last_frontier_for_loss = soft_frontier
        return gates

    @torch.no_grad()
    def begin_expert_bias_step_boundary(self) -> torch.Tensor:
        """Open one additive whole-step QB histogram.
        Gradient-accumulation microbatches call
        :meth:`accumulate_expert_bias_histogram` against the same immutable
        previous-step beta, then the optimizer-step boundary calls
        :meth:`commit_expert_bias_step_boundary` exactly once.  The histogram
        is transaction-local: durable resume replays an interrupted step from
        its committed cursor rather than checkpointing partial counts.
        """

        self._quantile_histogram_counts_t.zero_()
        self._quantile_histogram_token_count_t.zero_()
        self._quantile_histogram_selected_count_t.zero_()
        return self.expert_bias_t.new_ones((), dtype=torch.bool)

    @torch.no_grad()
    def expert_bias_step_snapshot_boundary(
        self,
    ) -> QuantileBalancingStepSnapshot:
        """Snapshot in-flight histogram state for exact checkpoint replay."""

        return QuantileBalancingStepSnapshot(
            histogram_counts_t=(
                self._quantile_histogram_counts_t.detach().clone()
            ),
            selected_count_t=(
                self._quantile_histogram_selected_count_t.detach().clone()
            ),
            token_count_t=(
                self._quantile_histogram_token_count_t.detach().clone()
            ),
        )

    @torch.no_grad()
    def restore_expert_bias_step_snapshot_boundary(
        self,
        snapshot: QuantileBalancingStepSnapshot,
    ) -> torch.Tensor:
        """Restore an exact pre/post-forward accumulator without committing."""

        if (
            snapshot.histogram_counts_t.shape
            != self._quantile_histogram_counts_t.shape
            or snapshot.histogram_counts_t.dtype
            != self._quantile_histogram_counts_t.dtype
            or snapshot.selected_count_t.shape
            != self._quantile_histogram_selected_count_t.shape
            or snapshot.selected_count_t.dtype
            != self._quantile_histogram_selected_count_t.dtype
            or snapshot.token_count_t.shape
            != self._quantile_histogram_token_count_t.shape
            or snapshot.token_count_t.dtype
            != self._quantile_histogram_token_count_t.dtype
        ):
            raise RuntimeError(
                "quantile balancing transaction snapshot geometry differs"
            )
        torch._foreach_copy_(
            (
                self._quantile_histogram_counts_t,
                self._quantile_histogram_selected_count_t,
                self._quantile_histogram_token_count_t,
            ),
            (
                snapshot.histogram_counts_t,
                snapshot.selected_count_t,
                snapshot.token_count_t,
            ),
        )
        return self.expert_bias_t.new_ones((), dtype=torch.bool)

    @torch.no_grad()
    def accumulate_expert_bias_histogram(
        self,
        scores: torch.Tensor,
    ) -> torch.Tensor:
        """Accumulate Top-(k+1) required-bias margins without communication.
        The model-owned adaptive frontier may differ per row.  Rows whose
        learned frontier already spans every expert are balanced by
        construction and therefore contribute neither a cutoff nor histogram
        pressure.  All other rows contribute one required-bias observation per
        expert and their exact learned route width to the target load.
        """

        if scores.shape[-1] != self.num_experts:
            raise ValueError("quantile balancing update geometry differs")
        score_rows_t = scores.detach().float().reshape(-1, self.num_experts)
        if score_rows_t.shape[0] < 1:
            raise ValueError("quantile balancing update has no token rows")
        score_limit = math.sqrt(torch.finfo(score_rows_t.dtype).max)
        score_rows_t = torch.nan_to_num(
            score_rows_t,
            nan=0.0,
            posinf=score_limit,
            neginf=-score_limit,
        ).clamp(-score_limit, score_limit)
        beta_t = self.expert_bias_t.detach().to(
            device=score_rows_t.device,
            dtype=score_rows_t.dtype,
        )
        adjusted_rows_t = score_rows_t - beta_t.reshape(1, -1)
        temperature_t = self.temperature().detach().to(
            device=score_rows_t.device,
            dtype=score_rows_t.dtype,
        )
        route_probability_t = F.softmax(
            adjusted_rows_t / temperature_t,
            dim=-1,
        )
        frontier_count_t = self.frontier_count_t(
            uncertain=route_probability_t.amax(dim=-1).lt(0.35),
        ).to(device=score_rows_t.device, dtype=torch.long)
        if frontier_count_t.numel() == 1:
            frontier_count_t = frontier_count_t.expand(score_rows_t.shape[0])
        frontier_order_t = torch.argsort(
            adjusted_rows_t,
            dim=-1,
            descending=True,
            stable=True,
        )
        cutoff_rank_t = frontier_count_t.clamp_max(
            self.num_experts - 1
        ).unsqueeze(-1)
        cutoff_expert_t = frontier_order_t.gather(
            -1,
            cutoff_rank_t,
        )
        cutoff_t = adjusted_rows_t.gather(
            -1,
            cutoff_expert_t,
        )
        # In the additive-bias sign convention this is the bias that would
        # place expert j exactly at row i's current Top-k cutoff.  Resynthesis
        # stores subtractive beta, so commit negates the recovered quantile.
        required_additive_bias_t = cutoff_t - score_rows_t
        active_row_t = frontier_count_t.lt(self.num_experts)
        active_observation_t = active_row_t.unsqueeze(-1).expand_as(
            required_additive_bias_t
        )
        # Obtain a rank-common raw range from sigmoid router scores.
        # Resynthesis routing logits are intentionally unbounded, so use a
        # strictly monotonic arctangent compander instead.  Quantile order is
        # unchanged, every rank shares the model-owned temperature/beta scale,
        # and all microbatches can add directly into one fixed [0, 1]
        # histogram without range collectives or raw-margin gathering.
        companding_scale_t = (
            temperature_t
            + beta_t.detach().abs().amax()
        ).clamp_min(torch.finfo(torch.float32).eps)
        companded_required_bias_t = (
            torch.atan(
                required_additive_bias_t / companding_scale_t
            )
            / math.pi
            + 0.5
        )
        bin_index_t = (
            (
                companded_required_bias_t
                * QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT
            )
            .floor()
            .to(dtype=torch.long)
            .clamp(0, QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT - 1)
        )
        observation_count_t = active_observation_t.to(dtype=torch.long)
        self._quantile_histogram_counts_t.scatter_add_(
            1,
            bin_index_t.transpose(0, 1),
            observation_count_t.transpose(0, 1),
        )
        self._quantile_histogram_token_count_t.add_(
            active_row_t.to(dtype=torch.long).sum()
        )
        self._quantile_histogram_selected_count_t.add_(
            torch.where(
                active_row_t,
                frontier_count_t,
                torch.zeros_like(frontier_count_t),
            ).sum()
        )
        return self._quantile_histogram_counts_t

    @staticmethod
    def _all_reduce_quantile_boundary_(
        tensor_t: torch.Tensor,
        operation: dist.ReduceOp.RedOpType,
    ) -> torch.Tensor:
        """Reduce one tensor at the explicit distributed-step boundary."""

        if dist.is_available() and dist.is_initialized():
            dist.all_reduce(tensor_t, op=operation)
        return tensor_t

    @torch.no_grad()
    def commit_expert_bias_step_boundary(self) -> torch.Tensor:
        """Pool whole-step histograms globally and commit next-step beta."""

        reduction_payload_t = torch.cat(
            (
                self._quantile_histogram_counts_t.reshape(-1),
                self._quantile_histogram_selected_count_t.reshape(1),
                self._quantile_histogram_token_count_t.reshape(1),
            ),
        )
        self._all_reduce_quantile_boundary_(
            reduction_payload_t,
            dist.ReduceOp.SUM,
        )
        histogram_value_count = (
            self.num_experts * QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT
        )
        global_histogram_t = reduction_payload_t.narrow(
            0,
            0,
            histogram_value_count,
        ).reshape(
            self.num_experts,
            QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT,
        )
        global_selected_count_t = reduction_payload_t.narrow(
            0,
            histogram_value_count,
            1,
        ).reshape(())
        global_token_count_t = reduction_payload_t.narrow(
            0,
            histogram_value_count + 1,
            1,
        ).reshape(())
        global_has_data_t = global_token_count_t.gt(0)
        target_load_t = (
            global_selected_count_t.to(dtype=torch.float32)
            / self.num_experts
        )
        target_rank_t = target_load_t.ceil().to(dtype=torch.long)
        cumulative_t = global_histogram_t.cumsum(dim=-1)
        reached_target_t = cumulative_t.ge(target_rank_t)
        selected_bin_t = reached_target_t.to(dtype=torch.long).argmax(
            dim=-1
        )
        selected_bin_count_t = global_histogram_t.gather(
            1,
            selected_bin_t.unsqueeze(-1),
        ).squeeze(-1)
        cumulative_at_bin_t = cumulative_t.gather(
            1,
            selected_bin_t.unsqueeze(-1),
        ).squeeze(-1)
        cumulative_before_t = (
            cumulative_at_bin_t - selected_bin_count_t
        )
        within_bin_fraction_t = (
            (
                target_load_t
                - cumulative_before_t.to(dtype=target_load_t.dtype)
            )
            / selected_bin_count_t.to(
                dtype=target_load_t.dtype
            ).clamp_min(1.0)
        ).clamp(0.0, 1.0)
        companded_quantile_t = (
            selected_bin_t.to(dtype=torch.float32)
            + within_bin_fraction_t
        ) / QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT
        half_bin_t = torch.ones_like(companded_quantile_t) * (
            0.5 / QUANTILE_BALANCING_HISTOGRAM_BIN_COUNT
        )
        companded_quantile_t = torch.minimum(
            torch.maximum(companded_quantile_t, half_bin_t),
            torch.ones_like(companded_quantile_t) - half_bin_t,
        )
        companding_scale_t = (
            self.temperature().detach().float()
            + self.expert_bias_t.detach().float().abs().amax()
        ).clamp_min(torch.finfo(torch.float32).eps)
        additive_bias_t = companding_scale_t * torch.tan(
            (companded_quantile_t - 0.5) * math.pi
        )
        beta_t = -additive_bias_t
        beta_t = beta_t - beta_t.mean()
        beta_t = torch.nan_to_num(
            beta_t,
            nan=0.0,
            posinf=0.0,
            neginf=0.0,
        )
        next_beta_t = torch.where(
            global_has_data_t,
            beta_t,
            self.expert_bias_t.detach().to(
                device=beta_t.device,
                dtype=beta_t.dtype,
            ),
        )
        # Quantile beta owns load balancing only.  TRAUMA/MILT/anti pressure is
        # composed once in ``soft_gates`` with the explicit logit equation
        # learned + pro + exploration - anti; persisting it into beta would
        # reverse the pro sign because ``biased_scores`` subtracts beta.
        self.expert_bias_t.copy_(
            next_beta_t.to(
                device=self.expert_bias_t.device,
                dtype=self.expert_bias_t.dtype,
            )
        )
        self.begin_expert_bias_step_boundary()
        return self.expert_bias_t

    @torch.no_grad()
    def update_expert_bias(self, scores: torch.Tensor) -> torch.Tensor:
        """Run one complete route arm as a causal global histogram step."""

        self.begin_expert_bias_step_boundary()
        self.accumulate_expert_bias_histogram(scores)
        return self.commit_expert_bias_step_boundary()

    def utilization_balance_loss(self) -> torch.Tensor:
        """Switch-style load×importance from hard quantile masks + soft gates."""

        util = self.last_utilization_t
        soft = self._last_soft_for_loss
        frontier = self._last_frontier_for_loss
        target_share = self._last_activation_fraction_for_loss
        if (
            util is None
            or soft is None
            or frontier is None
            or target_share is None
        ):
            return self.expert_bias_t.new_zeros(())
        importance = soft.reshape(-1, self.num_experts).mean(dim=0)
        load = util.to(device=importance.device, dtype=importance.dtype)
        balance = (
            torch.ones_like(importance.sum())
            * self.num_experts
            * (importance * load).sum()
        )
        active_share = frontier.mean()
        active_target_share = target_share.to(
            device=active_share.device,
            dtype=active_share.dtype,
        )
        return balance + (active_share - active_target_share).square()

    def route(
        self,
        scores: torch.Tensor,
        *,
        balance_scores_t: torch.Tensor | None = None,
    ) -> torch.Tensor:
        gates = self.soft_gates(scores)
        self._last_soft_for_loss = gates
        if self.training:
            active_balance_scores_t = (
                scores if balance_scores_t is None else balance_scores_t
            )
            if active_balance_scores_t.shape != scores.shape:
                raise ValueError("quantile balancing update geometry differs")
            # Production training owns the optimizer transaction explicitly:
            # every microbatch contributes additive counts, while the outer
            # group boundary commits beta only after the optimizer step
            # succeeds.  This keeps every forward in the group on the same
            # immutable previous-step beta and prevents checkpoint
            # recomputation from advancing routing state twice.
            self.accumulate_expert_bias_histogram(
                active_balance_scores_t
            )
        return gates

    def route_coherent_rows(
        self,
        scores_t: torch.Tensor,
        *,
        balance_scores_t: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Route an already coherent rank-two row bank exactly once.
        The caller owns the model-side proof that every row is the same
        coherent route.  Keeping that proof at the caller avoids a device
        synchronization or a redundant all-row comparison here.  Quantile
        selection and the causal beta update therefore consume one row, while
        the returned gates and loss diagnostics retain the complete public
        ``[rows, experts]`` geometry through differentiable tensor expansion.
        """

        if (
            scores_t.dim() != 2
            or scores_t.shape[0] < 1
            or scores_t.shape[-1] != self.num_experts
        ):
            raise ValueError("quantile coherent row geometry differs")
        if (
            balance_scores_t is not None
            and balance_scores_t.shape != scores_t.shape
        ):
            raise ValueError("quantile balancing update geometry differs")

        coherent_scores_t = scores_t.narrow(0, 0, 1)
        coherent_balance_scores_t = (
            None
            if balance_scores_t is None
            else balance_scores_t.narrow(0, 0, 1)
        )
        coherent_gates_t = self.route(
            coherent_scores_t,
            balance_scores_t=coherent_balance_scores_t,
        )
        gates_t = coherent_gates_t.expand_as(scores_t)

        # Public diagnostics describe every routed row even though the
        # quantile frontier is evaluated once.  ``expand_as`` preserves the
        # exact differentiable route graph without copying the repeated bank.
        self._last_soft_for_loss = gates_t
        if isinstance(self.last_hard_mask, torch.Tensor):
            self.last_hard_mask = self.last_hard_mask.expand_as(scores_t)
        if isinstance(self._last_frontier_for_loss, torch.Tensor):
            self._last_frontier_for_loss = (
                self._last_frontier_for_loss.expand_as(scores_t)
            )
        return gates_t