modeling_nested_mamba.py

"""Inference-only architecture for the three-level nested byte Mamba-2 model."""
from __future__ import annotations

import copy
import math
import threading
from collections import deque
from typing import Dict, List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F

try:
    from mamba_ssm.utils.generation import InferenceParams
except Exception:
    InferenceParams = None


def _new_mamba_inference_params(max_batch_size: int, max_seqlen: int):
    if InferenceParams is None:
        return None
    try:
        return InferenceParams(
            max_seqlen=max_seqlen,
            max_batch_size=max_batch_size,
        )
    except TypeError:
        try:
            return InferenceParams(
                max_batch_size=max_batch_size,
                max_seqlen=max_seqlen,
            )
        except TypeError:
            return InferenceParams(max_seqlen, max_batch_size)


def _set_inference_seqlen_offset(inference_params, offset: int) -> None:
    if inference_params is not None and hasattr(inference_params, "seqlen_offset"):
        inference_params.seqlen_offset = int(offset)

MAMBA_IMPORT_ERROR: Optional[Exception] = None
try:
    from mamba_ssm import Mamba, Mamba2
except Exception as exc:
    Mamba = None
    Mamba2 = None
    MAMBA_IMPORT_ERROR = exc


class BLTGlobalBlock(nn.Module):
    def __init__(
        self,
        dim: int,
        use_mamba: bool = True,
        ff_mult: int = 2,
        layer_idx: int = 0,
        mamba_version: int = 2,
        mamba_d_state: int = 64,
        mamba2_headdim: int = 0,
    ):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.mamba_version = int(mamba_version)
        self.mamba_d_state = int(mamba_d_state)
        self.mamba2_headdim = int(mamba2_headdim)
        if self.mamba_version not in (1, 2):
            raise ValueError(f"mamba_version must be 1 or 2, got {self.mamba_version}")
        if self.mamba_d_state < 1:
            raise ValueError("mamba_d_state must be positive")
        if use_mamba and self.mamba_version == 2:
            if Mamba2 is None:
                raise RuntimeError(
                    "Mamba-2 was requested but mamba_ssm.Mamba2 is unavailable"
                ) from MAMBA_IMPORT_ERROR
            inner_dim = int(dim) * 2
            if self.mamba2_headdim <= 0:
                self.mamba2_headdim = next(
                    (candidate for candidate in (64, 128, 100, 80, 50, 40, 32, 25, 20, 16, 10, 8, 5, 4, 2, 1)
                     if inner_dim % candidate == 0),
                    1,
                )
            if inner_dim % self.mamba2_headdim:
                raise ValueError(
                    f"Mamba-2 inner dimension {inner_dim} (2*dim) must be divisible by "
                    f"mamba2_headdim={self.mamba2_headdim}"
                )
            self.mixer = Mamba2(
                d_model=dim,
                d_state=self.mamba_d_state,
                d_conv=4,
                expand=2,
                headdim=self.mamba2_headdim,
                layer_idx=layer_idx,
            )
            self.kind = "mamba"
        elif use_mamba and Mamba is not None:
            try:
                self.mixer = Mamba(
                    d_model=dim,
                    d_state=self.mamba_d_state,
                    d_conv=4,
                    expand=2,
                    layer_idx=layer_idx,
                )
            except TypeError:
                self.mixer = Mamba(
                    d_model=dim, d_state=self.mamba_d_state, d_conv=4, expand=2
                )
                self.mixer.layer_idx = layer_idx
            self.kind = "mamba"
        else:
            self.mixer = nn.GRU(dim, dim, batch_first=True)
            self.kind = "gru"
        self.ff = nn.Sequential(
            nn.LayerNorm(dim),
            nn.Linear(dim, dim * ff_mult),
            nn.GELU(),
            nn.Linear(dim * ff_mult, dim),
        )

    def forward(self, x: torch.Tensor, inference_params=None) -> torch.Tensor:
        h = self.norm(x)
        if self.kind == "gru":
            h, _ = self.mixer(h)
        elif inference_params is None:
            h = self.mixer(h)
        else:
            h = self.mixer(h, inference_params=inference_params)
        x = x + h
        return x + self.ff(x)


class CausalConv1d(nn.Module):
    def __init__(self, dim: int, kernel_size: int):
        super().__init__()
        self.left_pad = max(0, int(kernel_size) - 1)
        self.conv = nn.Conv1d(dim, dim, kernel_size=kernel_size, padding=0, groups=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.left_pad:
            x = F.pad(x, (self.left_pad, 0))
        return self.conv(x)


class LearnedPoolEncoder(nn.Module):
    """Causal, learned replacement for the old mean patch autoencoder.

    ``close_probability`` is used to decide whether an eligible pool ends; a
    maximum length still forces a close, so the global stream has a bounded
    memory footprint.  The same probabilities also participate in the pooled
    representation, which gives the close head a learning signal from the LM
    objective (rather than making it a detached routing heuristic).
    """

    def __init__(self, dim: int, *, detach_close_content_gradient: bool = True):
        super().__init__()
        self.detach_close_content_gradient = bool(detach_close_content_gradient)
        # A recurrent encoder is impractical for the 660k-token training
        # windows this trainer supports (and cuDNN rejects some such shapes).
        # This is still a learned causal encoder, but its convolutional form
        # is safe for very long windows and compatible with AMP/cuDNN.
        self.context_norm = nn.LayerNorm(dim)
        self.context = nn.Conv1d(dim, dim, kernel_size=3, padding=0)
        self.value = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, dim))
        self.weight = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, 1))
        self.close = nn.Sequential(nn.LayerNorm(dim), nn.Linear(dim, dim), nn.GELU(), nn.Linear(dim, 1))
        self.close_value = nn.Linear(1, dim, bias=False)
        # Start with stable max-sized pools; the head can learn earlier closes
        # without immediately multiplying the global sequence length at init.
        nn.init.constant_(self.close[-1].bias, -4.0)

    def contextualize(self, local_h: torch.Tensor) -> torch.Tensor:
        h = self.context_norm(local_h).transpose(1, 2)
        h = self.context(F.pad(h, (2, 0))).transpose(1, 2)
        return local_h + F.gelu(h)

    def close_probability(self, contextual_h: torch.Tensor) -> torch.Tensor:
        return torch.sigmoid(self.close_logits(contextual_h))

    def close_logits(self, contextual_h: torch.Tensor) -> torch.Tensor:
        return self.close(contextual_h).squeeze(-1)

    def pool(
        self,
        contextual_h: torch.Tensor,
        patch_ids: torch.Tensor,
        patch_counts: torch.Tensor,
        token_mask: torch.Tensor,
        max_patches: int,
    ) -> torch.Tensor:
        b, _, c = contextual_h.shape
        value = self.value(contextual_h)
        # Preserve existing checkpoint forward values while preventing the LM
        # objective from training the boundary probability as a content side
        # channel. Boundary heads receive their own causal routing objective.
        close_prob = self.close_probability(contextual_h)
        if self.detach_close_content_gradient:
            close_prob = close_prob.detach()
        close_prob = close_prob.unsqueeze(-1)
        weight = torch.sigmoid(self.weight(contextual_h)) * token_mask.unsqueeze(-1).to(contextual_h.dtype)
        encoded = (value + self.close_value(close_prob)) * weight
        pooled = contextual_h.new_zeros((b, max_patches, c))
        normalizer = contextual_h.new_zeros((b, max_patches, 1))
        index = patch_ids.unsqueeze(-1)
        pooled.scatter_add_(1, index.expand(-1, -1, c), encoded)
        normalizer.scatter_add_(1, index, weight)
        return pooled / normalizer.clamp_min(1e-6)


class ByteLatentMambaCore(nn.Module):
    """Causal hierarchical BLT with adaptive, bounded level transitions."""

    def __init__(
        self,
        vocab_size: int,
        dim: int = 256,
        layers: int = 4,
        position_bins: int = 8192,
        use_mamba: bool = True,
        mamba_version: int = 2,
        mamba_d_state: int = 64,
        mamba2_headdim: int = 0,
        num_sections: int = 16,
        min_patch_bytes: int = 4,
        max_patch_bytes: int = 16,
        patch_change_threshold: int = 48,
        close_threshold: float = 0.90,
        mid_close_bonus: float = 0.05,
        nested_pool_factor: int = 16,
        nested_min_pool_factor: int = 0,
        nested_close_threshold: Optional[float] = None,
        nested_layers: int = 0,
        tertiary_pool_factor: int = 0,
        tertiary_min_pool_factor: int = 0,
        tertiary_close_threshold: Optional[float] = None,
        tertiary_layers: int = 0,
        decoder_dim: int = 0,
        detach_inactive_coarse_gradients: bool = True,
        legacy_pool_closure_training: bool = False,
        decoder_pool_controller: bool = False,
        pool_controller_alpha: float = 1.0,
        pool_controller_beta: float = 0.5,
        pool_controller_gamma: float = 0.5,
        short_pool_budget: int = 0,
        short_pool_window: int = 0,
        secondary_min_patch_bytes: int = 16,
    ):
        super().__init__()
        self.vocab_size = int(vocab_size)
        self.mamba_version = int(mamba_version)
        self.mamba_d_state = int(mamba_d_state)
        inner_dim = int(dim) * 2
        requested_headdim = int(mamba2_headdim)
        self.mamba2_headdim = requested_headdim
        if self.mamba_version == 2 and self.mamba2_headdim <= 0:
            self.mamba2_headdim = next(
                (candidate for candidate in (64, 128, 100, 80, 50, 40, 32, 25, 20, 16, 10, 8, 5, 4, 2, 1)
                 if inner_dim % candidate == 0),
                1,
            )
        self.min_patch_bytes = max(1, int(min_patch_bytes))
        self.max_patch_bytes = max(self.min_patch_bytes, int(max_patch_bytes))
        self.patch_change_threshold = max(0, int(patch_change_threshold))
        self.close_threshold = max(0.0, min(1.0, float(close_threshold)))
        self.mid_close_bonus = max(0.0, float(mid_close_bonus))
        self.nested_pool_factor = max(2, int(nested_pool_factor))
        self.nested_min_pool_factor = (
            self.nested_pool_factor
            if int(nested_min_pool_factor) <= 0
            else max(1, min(int(nested_min_pool_factor), self.nested_pool_factor))
        )
        self.nested_close_threshold = self.close_threshold if nested_close_threshold is None else max(
            0.0, min(1.0, float(nested_close_threshold))
        )
        self.tertiary_pool_factor = int(tertiary_pool_factor)
        self.tertiary_min_pool_factor = (
            self.tertiary_pool_factor
            if int(tertiary_min_pool_factor) <= 0
            else max(1, min(int(tertiary_min_pool_factor), self.tertiary_pool_factor))
        )
        self.tertiary_close_threshold = self.close_threshold if tertiary_close_threshold is None else max(
            0.0, min(1.0, float(tertiary_close_threshold))
        )
        self.tertiary_layers = int(tertiary_layers)
        self.tertiary_enabled = self.tertiary_pool_factor >= 2 and self.tertiary_layers > 0
        if (self.tertiary_pool_factor > 0 or self.tertiary_layers > 0) and not self.tertiary_enabled:
            raise ValueError("tertiary_pool_factor must be at least 2 and tertiary_layers must be positive together")
        self.nested_layers = max(1, int(nested_layers) if int(nested_layers) > 0 else max(1, int(layers) // 2))
        self.fine_layers = max(0, int(layers) - self.nested_layers)
        self.decoder_dim = int(decoder_dim) if int(decoder_dim) > 0 else int(dim) * 2
        if self.decoder_dim < int(dim) * 2:
            raise ValueError("decoder_dim must be 0 (the 2*dim default) or at least 2*dim")
        self.detach_inactive_coarse_gradients = bool(detach_inactive_coarse_gradients)
        self.legacy_pool_closure_training = bool(legacy_pool_closure_training)
        self.decoder_pool_controller = bool(decoder_pool_controller)
        self.pool_controller_alpha = float(pool_controller_alpha)
        self.pool_controller_beta = float(pool_controller_beta)
        self.pool_controller_gamma = float(pool_controller_gamma)
        self.short_pool_budget = max(0, int(short_pool_budget))
        self.short_pool_window = max(0, int(short_pool_window))
        self.secondary_min_patch_bytes = min(
            self.max_patch_bytes,
            max(self.min_patch_bytes, int(secondary_min_patch_bytes)),
        )
        self.byte_offset = 4
        self.byte_emb = nn.Embedding(vocab_size, dim)
        self.local_norm = nn.LayerNorm(dim)
        self.local_conv = nn.Sequential(
            CausalConv1d(dim, kernel_size=5),
            nn.GELU(),
            CausalConv1d(dim, kernel_size=3),
            nn.GELU(),
        )
        pool_encoder_kwargs = {
            "detach_close_content_gradient": not self.legacy_pool_closure_training
        }
        self.pool_encoder = LearnedPoolEncoder(dim, **pool_encoder_kwargs)
        # The second pooler operates on completed first-level pool states.
        # Fixed-size grouping keeps this hierarchy vectorized and bounded;
        # its representation/weights remain learned through this encoder.
        self.nested_pool_encoder = LearnedPoolEncoder(dim, **pool_encoder_kwargs)
        self.tertiary_pool_encoder = (
            LearnedPoolEncoder(dim, **pool_encoder_kwargs)
            if self.tertiary_enabled else None
        )
        self.boe_patch = nn.Parameter(torch.zeros(1, 1, dim))
        self.patch_pos_emb = nn.Embedding(position_bins, dim)
        self.patch_len_emb = nn.Embedding(self.max_patch_bytes + 1, dim)
        self.nested_len_emb = nn.Embedding(self.nested_pool_factor + 1, dim)
        self.tertiary_len_emb = nn.Embedding(self.tertiary_pool_factor + 1, dim) if self.tertiary_enabled else None
        block_kwargs = dict(
            dim=dim,
            use_mamba=use_mamba,
            mamba_version=self.mamba_version,
            mamba_d_state=self.mamba_d_state,
            mamba2_headdim=self.mamba2_headdim,
        )
        self.global_blocks = nn.ModuleList([
            BLTGlobalBlock(**block_kwargs, layer_idx=i) for i in range(self.fine_layers)
        ])
        self.nested_global_blocks = nn.ModuleList([
            BLTGlobalBlock(**block_kwargs, layer_idx=self.fine_layers + i)
            for i in range(self.nested_layers)
        ])
        self.tertiary_global_blocks = nn.ModuleList([
            BLTGlobalBlock(
                **block_kwargs,
                layer_idx=self.fine_layers + self.nested_layers + i,
            )
            for i in range(self.tertiary_layers)
        ])
        self.decoder = nn.Sequential(
            nn.LayerNorm(dim * (5 if self.tertiary_enabled else 4)),
            nn.Linear(dim * (5 if self.tertiary_enabled else 4), self.decoder_dim),
            nn.GELU(),
            nn.Linear(self.decoder_dim, dim),
            nn.GELU(),
        )
        # Three causal readouts (hold, refresh, compress) for each hierarchy.
        # The zero initialization produces 0.5 for every readout.  With the
        # default alpha=beta+gamma this is an exact no-op, which lets an older
        # checkpoint acquire the controller without changing its first
        # generated byte or its initial routing policy.
        self.pool_controller = nn.Linear(dim, 9) if self.decoder_pool_controller else None
        if self.pool_controller is not None:
            nn.init.zeros_(self.pool_controller.weight)
            nn.init.zeros_(self.pool_controller.bias)
        self.lm_head = nn.Linear(dim, vocab_size, bias=False)
        self.position_bins = int(position_bins)
        # Placement is deliberately runtime-only: checkpoints contain the
        # ordinary architecture state dict and can still be loaded on CPU, a
        # single GPU, or a browser-export host.
        self.model_parallel_enabled = False
        # Standalone evaluation may request byte-aligned hierarchy details.
        # Keep this disabled during training so large activations are not
        # retained after forward/backward.
        self.capture_evaluation_details = False
        self._evaluation_details_by_thread: Dict[int, object] = {}
        self.last_forward_evaluation_details = None
        self.last_pool_controller_scores = None
        self.last_pool_controller_close_signals = None
        self.last_pool_routing_training = None
        self.fine_device = None
        self.nested_devices: List[torch.device] = []

    @property
    def last_forward_evaluation_details(self):
        return self._evaluation_details_by_thread.get(threading.get_ident())

    @last_forward_evaluation_details.setter
    def last_forward_evaluation_details(self, value) -> None:
        thread_id = threading.get_ident()
        if value is None:
            self._evaluation_details_by_thread.pop(thread_id, None)
        else:
            self._evaluation_details_by_thread[thread_id] = value

    def clear_forward_evaluation_details(self) -> None:
        self._evaluation_details_by_thread.clear()

    def _pool_controller_scores(self, decoded: torch.Tensor) -> Optional[torch.Tensor]:
        """Return [batch, time, level, (hold, refresh, compress)].

        Routing is discrete, so controller supervision intentionally does not
        backpropagate into the already-trained decoder.  Recovery therefore
        teaches only the newly introduced readout and cannot damage the
        decoder while its policy is being calibrated.
        """
        if self.pool_controller is None:
            return None
        return torch.sigmoid(self.pool_controller(decoded.detach())).view(
            *decoded.shape[:-1], 3, 3
        )

    def _controller_close_adjustment(self, scores, level: int) -> float:
        if scores is None or not self.decoder_pool_controller:
            return 0.0
        level_scores = scores[level]
        hold, refresh, compress = (
            float(value.detach().item()) if isinstance(value, torch.Tensor) else float(value)
            for value in level_scores
        )
        # Centering makes the zero-initialized 0.5/0.5/0.5 controller neutral
        # for every coefficient combination, not only the defaults.
        return (
            self.pool_controller_alpha * (refresh - 0.5)
            - self.pool_controller_beta * (hold - 0.5)
            - self.pool_controller_gamma * (compress - 0.5)
        )

    def pool_controller_auxiliary_loss(
        self, logits: torch.Tensor, targets: torch.Tensor
    ) -> torch.Tensor:
        """Supervise controller semantics without differentiating hard routing.

        The existing close head supplies the refresh target.  Remaining mass
        is split between hold and compress using detached next-byte surprise:
        difficult bytes request more local evidence (hold), while predictable
        bytes permit compression.  Each level receives its own aligned close
        signal from the encoder at that level.
        """
        scores = self.last_pool_controller_scores
        signals = self.last_pool_controller_close_signals
        if scores is None or signals is None:
            return logits.new_zeros(())
        valid = targets.ne(-100)
        with torch.no_grad():
            nll = F.cross_entropy(
                logits.detach().transpose(1, 2), targets, ignore_index=-100,
                reduction="none",
            )
            difficulty = nll / (nll + 2.0)
            refresh = signals.detach().clamp(0.0, 1.0)
            unresolved = 1.0 - refresh
            target_scores = torch.stack(
                [
                    unresolved * difficulty.unsqueeze(-1),
                    refresh,
                    unresolved * (1.0 - difficulty.unsqueeze(-1)),
                ],
                dim=-1,
            )
        # C[t] controls the boundary evaluated at the beginning of step t+1.
        # Train it against that future step rather than the byte that emitted
        # it. The final controller output has no within-window target.
        future_scores = scores[:, :-1]
        target_scores = target_scores[:, 1:]
        future_valid = valid[:, 1:]
        expanded_valid = future_valid.unsqueeze(-1).unsqueeze(-1).expand_as(
            future_scores
        )
        if not bool(expanded_valid.any()):
            return logits.new_zeros(())
        # PyTorch deliberately rejects probability-space BCE inside CUDA
        # autocast because its sigmoid gradient can underflow in FP16/BF16.
        # Keep the checkpoint-compatible sigmoid controller, but evaluate this
        # small auxiliary objective in FP32 outside the surrounding training
        # autocast region. Gradients still flow through the float conversion
        # into pool_controller; decoded was intentionally detached above.
        with torch.autocast(device_type=future_scores.device.type, enabled=False):
            return F.binary_cross_entropy(
                future_scores[expanded_valid].float(),
                target_scores[expanded_valid].float(),
            )

    def configure_model_parallel(
        self,
        fine_device: torch.device,
        nested_devices: List[torch.device],
        tertiary_device: Optional[torch.device] = None,
    ) -> None:
        """Place fine BLT work on one device and coarse levels on others.

        The byte embedding, local convolution, first-level pooling, fine
        blocks, and decoder stay together because they carry the long sequence.
        Only completed fine-pool states cross to the coarse hierarchy.  Coarse
        blocks are assigned in contiguous chunks to avoid an interconnect hop
        after every layer.
        """
        if not nested_devices:
            raise ValueError("nested model parallelism requires at least one coarse device")
        fine_device = torch.device(fine_device)
        nested_devices = [torch.device(item) for item in nested_devices]
        if fine_device.type != "cuda" or any(item.type != "cuda" for item in nested_devices):
            raise ValueError("nested model parallelism requires CUDA devices")

        # Start with a coherent root copy, then move just the coarse stage.
        self.to(fine_device)
        coarse_root = nested_devices[0]
        self.nested_pool_encoder.to(coarse_root)
        self.nested_len_emb.to(coarse_root)
        if self.tertiary_enabled:
            if tertiary_device is None:
                raise ValueError("three-level nested model parallelism requires a tertiary device")
            tertiary_device = torch.device(tertiary_device)
            if tertiary_device.type != "cuda":
                raise ValueError("tertiary model-parallel device must be CUDA")
            assert self.tertiary_pool_encoder is not None and self.tertiary_len_emb is not None
            self.tertiary_pool_encoder.to(tertiary_device)
            self.tertiary_len_emb.to(tertiary_device)
            for block in self.tertiary_global_blocks:
                block.to(tertiary_device)
        block_count = len(self.nested_global_blocks)
        for index, block in enumerate(self.nested_global_blocks):
            # Contiguous placement gives one transfer per participating GPU,
            # rather than alternating CUDA devices every coarse layer.
            device_index = min(len(nested_devices) - 1, index * len(nested_devices) // max(1, block_count))
            block.to(nested_devices[device_index])
        self.model_parallel_enabled = True
        self.fine_device = fine_device
        self.nested_devices = nested_devices

    @staticmethod
    def _module_device(module: nn.Module) -> torch.device:
        return next(module.parameters()).device

    @classmethod
    def _run_block_on_own_device(
        cls,
        block: nn.Module,
        h: torch.Tensor,
        inference_params=None,
    ) -> torch.Tensor:
        """Launch Triton-backed mixers under their parameter device context.

        Ordinary PyTorch operators dispatch from tensor placement, but
        Mamba-2's Triton SSD kernels use CUDA's current device when launching.
        In a model-parallel process that current device can remain cuda:0 even
        while a level and its activations live on cuda:1 or cuda:2, producing
        Triton's misleading "cpu tensor?" pointer error.
        """
        block_device = cls._module_device(block)
        if h.device != block_device:
            h = h.to(block_device, non_blocking=True)
        if block_device.type == "cuda":
            with torch.cuda.device(block_device):
                return block(h, inference_params=inference_params)
        return block(h, inference_params=inference_params)

    def _run_nested_blocks(self, nested_global_h: torch.Tensor, nested_global_mask: torch.Tensor) -> torch.Tensor:
        """Run coarse blocks, transferring only at device-stage boundaries."""
        h = nested_global_h
        mask = nested_global_mask
        for block in self.nested_global_blocks:
            block_device = self._module_device(block)
            if h.device != block_device:
                h = h.to(block_device, non_blocking=True)
                mask = mask.to(block_device, non_blocking=True)
            h = self._run_block_on_own_device(block, h, inference_params=None)
            h = h.masked_fill(~mask.unsqueeze(-1), 0.0)
        return h

    def _run_tertiary_blocks(self, h: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        for block in self.tertiary_global_blocks:
            block_device = self._module_device(block)
            if h.device != block_device:
                h, mask = h.to(block_device, non_blocking=True), mask.to(block_device, non_blocking=True)
            h = self._run_block_on_own_device(
                block, h, inference_params=None
            ).masked_fill(~mask.unsqueeze(-1), 0.0)
        return h

    def _local_features(self, x: torch.Tensor) -> torch.Tensor:
        h = self.byte_emb(x.clamp(0, self.vocab_size - 1))
        conv = self.local_conv(self.local_norm(h).transpose(1, 2)).transpose(1, 2)
        return h + conv

    @staticmethod
    def _stream_conv_last(
        conv: nn.Conv1d,
        current: torch.Tensor,
        state: Dict[str, object],
        cache_key: str,
    ) -> torch.Tensor:
        """Evaluate one causal Conv1d position and retain only its left context."""
        kernel_size = int(conv.kernel_size[0])
        history = state.get(cache_key)
        values = current if history is None else torch.cat([history, current], dim=1)
        conv_input = values
        if values.shape[1] < kernel_size:
            conv_input = F.pad(values.transpose(1, 2), (kernel_size - values.shape[1], 0))
        else:
            conv_input = values[:, -kernel_size:].transpose(1, 2)
        output = conv(conv_input).transpose(1, 2)[:, -1:, :]
        keep = max(0, kernel_size - 1)
        state[cache_key] = values[:, -keep:].detach() if keep else None
        return output

    def _local_stream_context(self, token: torch.Tensor, state: Dict[str, object]) -> torch.Tensor:
        """Incrementally compute the exact causal byte and pool-context feature."""
        embedded = self.byte_emb(token.clamp(0, self.vocab_size - 1))
        conv1 = self.local_conv[0]
        conv2 = self.local_conv[2]
        first = self._stream_conv_last(
            conv1.conv, self.local_norm(embedded), state, "local_conv1_tail"
        )
        first = F.gelu(first)
        second = self._stream_conv_last(
            conv2.conv, first, state, "local_conv2_tail"
        )
        local_h = embedded + F.gelu(second)
        contextual_delta = self._stream_conv_last(
            self.pool_encoder.context,
            self.pool_encoder.context_norm(local_h),
            state,
            "pool_context_tail",
        )
        return local_h + F.gelu(contextual_delta)

    def pool_close_probabilities(self, x: torch.Tensor) -> torch.Tensor:
        """Run only the learned encoder, for greedy-policy bootstrap training."""
        local_h = self._local_features(x)
        return self.pool_encoder.close_probability(self.pool_encoder.contextualize(local_h))

    def _close_score(
        self, probability: float, patch_len: int, minimum: Optional[int] = None
    ) -> float:
        """Positive-only midpoint incentive; it never lowers a close score."""
        minimum = self.min_patch_bytes if minimum is None else int(minimum)
        span = max(1, self.max_patch_bytes - minimum)
        progress = max(0.0, min(1.0, (float(patch_len) - minimum) / span))
        midpoint = max(0.0, 1.0 - abs(2.0 * progress - 1.0))
        return float(probability) + self.mid_close_bonus * midpoint

    def _runtime_patch_ids(
        self,
        contextual_h: torch.Tensor,
        token_mask: torch.Tensor,
        close_probabilities: Optional[torch.Tensor] = None,
        minimum: Optional[int] = None,
        minima_by_token: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Causally route tokens into learned pools, closing the final pool too.

        Boundary decisions are only allowed after ``min_patch_bytes`` and are
        always forced at ``max_patch_bytes``.  Padding never opens a pool.
        """
        minimum = self.min_patch_bytes if minimum is None else max(
            1, min(int(minimum), self.max_patch_bytes)
        )
        probabilities = (
            self.pool_encoder.close_probability(contextual_h)
            if close_probabilities is None
            else close_probabilities
        ).detach()
        positions = torch.arange(contextual_h.shape[1], device=contextual_h.device).view(1, -1)
        candidate_close = token_mask & (positions >= minimum) & (
            probabilities + self.mid_close_bonus >= self.close_threshold
        )
        if minima_by_token is None and not bool(candidate_close.any()):
            # The initialized close bias intentionally takes this fast path:
            # bounded, vectorized max-sized pools with no Python per-token
            # loop or CUDA synchronization for long training windows.
            self.last_short_pool_budget_metrics = {
                "short_pool_count": 0.0,
                "budget_activated_rows": 0.0,
                "batch_rows": float(contextual_h.shape[0]),
            }
            return (
                (positions // self.max_patch_bytes)
                .expand_as(token_mask)
                .long()
                .masked_fill(~token_mask, 0)
            )

        # Learned closes are sparse.  Move the small routing decision to CPU
        # once and scan ordinary Python values, avoiding one CUDA sync per byte.
        probability_rows = probabilities.float().cpu().tolist()
        mask_rows = token_mask.cpu().tolist()
        minima_rows = (
            minima_by_token.detach().to(device="cpu", dtype=torch.long).tolist()
            if minima_by_token is not None else None
        )
        rows: List[torch.Tensor] = []
        short_pool_counts: List[int] = []
        budget_activated_rows = 0
        for row_index, (row_prob, row_mask) in enumerate(zip(probability_rows, mask_rows)):
            patch_id = 0
            patch_len = 0
            short_pools = 0
            secondary_active = False
            recent_short_pools = deque()
            recent_short_count = 0
            budget_was_activated = False
            ids: List[int] = []
            previous_policy = None
            row_minima = minima_rows[row_index] if minima_rows is not None else None
            for token_index, (prob, valid) in enumerate(zip(row_prob, row_mask)):
                if not valid:
                    ids.append(patch_id)
                    continue
                policy = (
                    tuple(int(value) for value in row_minima[token_index])
                    if row_minima is not None else (minimum, 0, 0)
                )
                policy_minimum = max(1, min(int(policy[0]), self.max_patch_bytes))
                runtime_minimum = max(
                    policy_minimum,
                    self.secondary_min_patch_bytes if secondary_active else 1,
                )
                policy_transition = (
                    patch_len > 0 and previous_policy is not None and policy != previous_policy
                )
                learned_close = patch_len >= self.max_patch_bytes or (
                    patch_len >= runtime_minimum
                    and self._close_score(prob, patch_len, policy_minimum) >= self.close_threshold
                )
                if policy_transition or learned_close:
                    # A text<->codec transition is a structural boundary, not
                    # evidence that the learned router emitted a short pool.
                    completed_short = (
                        not policy_transition and patch_len < self.secondary_min_patch_bytes
                    )
                    if completed_short:
                        short_pools += 1
                    if self.short_pool_budget > 0 and self.short_pool_window > 0:
                        recent_short_pools.append(completed_short)
                        recent_short_count += int(completed_short)
                        if len(recent_short_pools) > self.short_pool_window:
                            recent_short_count -= int(recent_short_pools.popleft())
                        secondary_active = recent_short_count >= self.short_pool_budget
                    elif self.short_pool_budget > 0 and short_pools >= self.short_pool_budget:
                        secondary_active = True
                    budget_was_activated = budget_was_activated or secondary_active
                    patch_id += 1
                    patch_len = 0
                ids.append(patch_id)
                patch_len += 1
                previous_policy = policy
            rows.append(torch.tensor(ids, dtype=torch.long, device=contextual_h.device))
            short_pool_counts.append(short_pools)
            budget_activated_rows += int(budget_was_activated)
        self.last_short_pool_budget_metrics = {
            "short_pool_count": float(sum(short_pool_counts)),
            "budget_activated_rows": float(budget_activated_rows),
            "batch_rows": float(len(rows)),
        }
        return torch.stack(rows, dim=0).masked_fill(~token_mask, 0)

    def _runtime_hierarchy_ids(
        self,
        contextual_h: torch.Tensor,
        token_mask: torch.Tensor,
        encoder: LearnedPoolEncoder,
        minimum: int,
        maximum: int,
        threshold: float,
        minima_by_token: Optional[torch.Tensor] = None,
        policy_by_token: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Causally group completed lower-level states with learned closures.

        A lower-level state first joins its current group.  Its close score may
        then close that completed group for the *next* lower-level state.  This
        is important: closing the previous group from the current state would
        let full-window inference expose a coarse result before cached
        inference had actually completed the deciding state.
        """
        positions = torch.arange(contextual_h.shape[1], device=contextual_h.device).view(1, -1)
        if minimum >= maximum and minima_by_token is None and policy_by_token is None:
            return (
                (positions // maximum)
                .expand_as(token_mask)
                .long()
                .masked_fill(~token_mask, 0)
            )
        probabilities = encoder.close_probability(contextual_h).detach()
        probability_rows = probabilities.float().cpu().tolist()
        mask_rows = token_mask.cpu().tolist()
        minima_rows = (
            minima_by_token.detach().to(device="cpu", dtype=torch.long).tolist()
            if minima_by_token is not None else None
        )
        policy_rows = (
            policy_by_token.detach().to(device="cpu", dtype=torch.long).tolist()
            if policy_by_token is not None else None
        )
        rows: List[torch.Tensor] = []
        span = max(1, maximum - minimum)
        for row_index, (row_probabilities, row_mask) in enumerate(zip(probability_rows, mask_rows)):
            pool_id = 0
            pool_len = 0
            ids: List[int] = []
            previous_policy = None
            for token_index, (probability, valid) in enumerate(zip(row_probabilities, row_mask)):
                if not valid:
                    ids.append(pool_id)
                    continue
                runtime_minimum = (
                    max(1, min(int(minima_rows[row_index][token_index]), maximum))
                    if minima_rows is not None else minimum
                )
                policy = (
                    tuple(int(value) for value in policy_rows[row_index][token_index])
                    if policy_rows is not None else (runtime_minimum,)
                )
                if pool_len > 0 and previous_policy is not None and policy != previous_policy:
                    pool_id += 1
                    pool_len = 0
                ids.append(pool_id)
                pool_len += 1
                runtime_span = max(1, maximum - runtime_minimum)
                progress = max(0.0, min(1.0, (pool_len - runtime_minimum) / runtime_span))
                midpoint = max(0.0, 1.0 - abs(2.0 * progress - 1.0))
                close_score = float(probability) + self.mid_close_bonus * midpoint
                if pool_len >= maximum or (
                    pool_len >= runtime_minimum and close_score >= threshold
                ):
                    pool_id += 1
                    pool_len = 0
                previous_policy = policy
            rows.append(torch.tensor(ids, dtype=torch.long, device=contextual_h.device))
        # scatter_add validates every index, including entries whose weight is
        # zero.  A group that closes on a row's final valid state increments
        # pool_id for the following padding positions; force those invalid
        # positions to the always-valid zero bucket before learned pooling.
        return torch.stack(rows, dim=0).masked_fill(~token_mask, 0)

    @staticmethod
    def _pool_runtime_minima(
        minima: torch.Tensor,
        ids: torch.Tensor,
        mask: torch.Tensor,
        group_count: int,
    ) -> torch.Tensor:
        """Carry a causal span policy into its forced-aligned pooled states."""
        output = torch.zeros(
            (minima.shape[0], int(group_count), 3),
            dtype=torch.long,
            device=minima.device,
        )
        safe_ids = ids.masked_fill(~mask, 0).unsqueeze(-1).expand(-1, -1, 3)
        output.scatter_reduce_(
            1,
            safe_ids,
            minima.masked_fill(~mask.unsqueeze(-1), 0),
            reduce="amax",
            include_self=True,
        )
        return output

    def _pool_patches(
        self,
        local_h: torch.Tensor,
        patch_ids: torch.Tensor,
        token_mask: torch.Tensor,
        patch_counts: Optional[torch.Tensor] = None,
        max_patches: Optional[int] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        b, t, c = local_h.shape
        if patch_counts is None:
            patch_counts = patch_ids.amax(dim=1) + 1
        else:
            patch_counts = patch_counts.to(device=local_h.device, dtype=torch.long).reshape(-1)
            if int(patch_counts.numel()) != b:
                raise ValueError(f"patch_counts has {patch_counts.numel()} values for batch size {b}")
        if max_patches is None:
            max_patches = int(patch_counts.max().detach().cpu().item())
        else:
            max_patches = int(max_patches)
        if max_patches < 1:
            raise ValueError("max_patches must be positive")
        pooled = self.pool_encoder.pool(local_h, patch_ids, patch_counts, token_mask, max_patches)
        counts = local_h.new_zeros((b, max_patches, 1))
        counts.scatter_add_(1, patch_ids.unsqueeze(-1), token_mask.unsqueeze(-1).to(dtype=local_h.dtype))
        patch_mask = torch.arange(max_patches, device=local_h.device).view(1, -1) < patch_counts.view(-1, 1)
        patch_lens = counts.squeeze(-1).round().long().clamp(0, self.max_patch_bytes)
        return pooled, patch_lens, patch_mask

    @staticmethod
    def _first_valid_token_mask(token_mask: torch.Tensor) -> torch.Tensor:
        """Select exactly the first non-padding token in each batch row."""
        return token_mask & token_mask.long().cumsum(dim=1).eq(1)

    def _gate_inactive_coarse_gradient(
        self,
        latent: torch.Tensor,
        active_mask: torch.Tensor,
        token_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Keep forward values while limiting inactive coarse BOE gradients.

        The bootstrap state remains trainable on the first valid token. After
        that, an inactive state is detached until a completed causal pool is
        available. ``torch.where`` changes only autograd routing; inference
        values and checkpoint shapes are unchanged.
        """
        if not self.detach_inactive_coarse_gradients:
            return latent
        allow_gradient = active_mask | self._first_valid_token_mask(token_mask)
        return torch.where(allow_gradient.unsqueeze(-1), latent, latent.detach())

    def _global_stream_step(self, patch_embed: torch.Tensor, inference_params, offset: int) -> torch.Tensor:
        _set_inference_seqlen_offset(inference_params, int(offset))
        if patch_embed.device.type == "cuda":
            # Mamba's inference conv/SSM cache is dtype-strict. Keep recurrent
            # state updates in fp32 even when streaming validation runs under AMP.
            with torch.amp.autocast("cuda", enabled=False):
                stream_dtype = next(self.global_blocks[0].parameters()).dtype
                h = patch_embed.to(dtype=stream_dtype)
                for block in self.global_blocks:
                    h = self._run_block_on_own_device(
                        block, h, inference_params=inference_params
                    )
        else:
            h = patch_embed
            for block in self.global_blocks:
                h = self._run_block_on_own_device(
                    block, h, inference_params=inference_params
                )
        return h

    def _nested_global_stream_step(self, pool_embed: torch.Tensor, inference_params, offset: int) -> torch.Tensor:
        """Advance the coarse Mamba cache by one completed coarse pool."""
        _set_inference_seqlen_offset(inference_params, int(offset))
        if pool_embed.device.type == "cuda":
            with torch.amp.autocast("cuda", enabled=False):
                stream_dtype = next(self.nested_global_blocks[0].parameters()).dtype
                h = pool_embed.to(dtype=stream_dtype)
                for block in self.nested_global_blocks:
                    block_device = self._module_device(block)
                    if h.device != block_device:
                        h = h.to(block_device, non_blocking=True)
                    h = self._run_block_on_own_device(
                        block, h, inference_params=inference_params
                    )
        else:
            h = pool_embed
            for block in self.nested_global_blocks:
                h = self._run_block_on_own_device(
                    block, h, inference_params=inference_params
                )
        return h

    def _tertiary_global_stream_step(self, pool_embed: torch.Tensor, inference_params, offset: int) -> torch.Tensor:
        """Advance level 3 by one completed level-2 group."""
        _set_inference_seqlen_offset(inference_params, int(offset))
        h = pool_embed
        if h.device.type == "cuda":
            with torch.amp.autocast("cuda", enabled=False):
                stream_dtype = next(self.tertiary_global_blocks[0].parameters()).dtype
                h = h.to(dtype=stream_dtype)
                for block in self.tertiary_global_blocks:
                    block_device = self._module_device(block)
                    if h.device != block_device:
                        h = h.to(block_device, non_blocking=True)
                    h = self._run_block_on_own_device(
                        block, h, inference_params=inference_params
                    )
        else:
            for block in self.tertiary_global_blocks:
                h = self._run_block_on_own_device(
                    block, h, inference_params=inference_params
                )
        return h

    def _initialize_stream_mamba_caches(
        self,
        blocks: nn.ModuleList,
        inference_params,
        max_seqlen: int,
    ) -> None:
        """Preallocate numerically stable Mamba-2 recurrent states.

        Casting a checkpoint to FP16 should not also reduce the accumulator
        range of its long-lived SSD state. Mamba-2's public cache allocator
        uses one dtype for both caches, so retain the short convolution cache
        in the model dtype and promote only the SSM accumulator to FP32.
        The Triton selective-state-update kernel supports this mixed layout.
        """
        if inference_params is None:
            return
        for block in blocks:
            if (
                getattr(block, "kind", None) != "mamba"
                or int(getattr(block, "mamba_version", 1)) != 2
            ):
                continue
            mixer = block.mixer
            layer_idx = getattr(mixer, "layer_idx", None)
            if layer_idx is None:
                raise RuntimeError("cached Mamba-2 inference requires a unique layer_idx")
            block_device = self._module_device(block)
            if block_device.type == "cuda":
                with torch.cuda.device(block_device):
                    conv_state, ssm_state = mixer.allocate_inference_cache(
                        1, int(max_seqlen)
                    )
            else:
                conv_state, ssm_state = mixer.allocate_inference_cache(
                    1, int(max_seqlen)
                )
            inference_params.key_value_memory_dict[int(layer_idx)] = (
                conv_state,
                ssm_state.float(),
            )

    def _patch_conditioning(
        self,
        patch_h: torch.Tensor,
        patch_len: int,
        patch_index: int,
    ) -> torch.Tensor:
        b, _, _ = patch_h.shape
        device = patch_h.device
        patch_pos = torch.full((b, 1), int(patch_index), dtype=torch.long, device=device)
        patch_lens = torch.full((b, 1), max(1, min(int(patch_len), self.max_patch_bytes)), dtype=torch.long, device=device)
        return patch_h + self.patch_pos_emb(patch_pos.remainder(self.position_bins)) + self.patch_len_emb(patch_lens)

    def _concat_patch_parts(self, parts: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
        keys = [k for k, v in parts[0].items() if isinstance(v, torch.Tensor) and v.dim() >= 2 and int(v.shape[1]) == 1]
        return {k: torch.cat([p[k] for p in parts], dim=1) for k in keys}

    def _hierarchy_should_close(
        self,
        probability: float,
        pool_len: int,
        minimum: int,
        maximum: int,
        threshold: float,
        controller_adjustment: float = 0.0,
    ) -> bool:
        if pool_len >= maximum:
            return True
        if pool_len < minimum:
            return False
        span = max(1, maximum - minimum)
        progress = max(0.0, min(1.0, (pool_len - minimum) / span))
        midpoint = max(0.0, 1.0 - abs(2.0 * progress - 1.0))
        return (
            float(probability)
            + self.mid_close_bonus * midpoint
            + float(controller_adjustment)
            >= threshold
        )

    def _append_nested_stream_part(
        self, state: Dict[str, object], fine_completed_h: torch.Tensor
    ) -> None:
        """Route one completed fine state into the adaptive level-2 pool."""
        nested_device = self._module_device(self.nested_pool_encoder)
        current = fine_completed_h
        if current.device != nested_device:
            current = current.to(nested_device, non_blocking=True)
        routing_tail = state.get("nested_routing_tail")
        routing_input = current if routing_tail is None else torch.cat([routing_tail, current], dim=1)
        contextual_current = self.nested_pool_encoder.contextualize(routing_input)[:, -1:, :]
        probability = float(
            self.nested_pool_encoder.close_probability(contextual_current)[0, 0]
            .detach()
            .cpu()
            .item()
        )
        state.setdefault("nested_fine_parts", []).append(current)
        state["nested_routing_tail"] = routing_input[:, -2:].detach()
        state["pending_nested_close_probability"] = probability
        controller_adjustment = self._controller_close_adjustment(
            state.get("pending_pool_controller_scores"), 1
        )
        if self._hierarchy_should_close(
            probability,
            len(state["nested_fine_parts"]),
            self.nested_min_pool_factor,
            self.nested_pool_factor,
            self.nested_close_threshold,
            controller_adjustment,
        ):
            self._close_nested_stream_pool(state, force=True)

    def _close_nested_stream_pool(self, state: Dict[str, object], force: bool = False) -> None:
        """Commit a group of completed fine pools to the coarse cached stream."""
        fine_parts = state.get("nested_fine_parts", [])
        if not fine_parts or (not force and len(fine_parts) < self.nested_pool_factor):
            return
        fine_h = torch.cat(fine_parts, dim=1)
        nested_device = self._module_device(self.nested_pool_encoder)
        if fine_h.device != nested_device:
            fine_h = fine_h.to(nested_device, non_blocking=True)
        tail = state.get("nested_context_tail")
        contextual_input = fine_h if tail is None else torch.cat([tail, fine_h], dim=1)
        contextual_h = self.nested_pool_encoder.contextualize(contextual_input)[:, -fine_h.shape[1]:]
        token_mask = torch.ones(fine_h.shape[:2], dtype=torch.bool, device=fine_h.device)
        nested_ids = torch.zeros(fine_h.shape[:2], dtype=torch.long, device=fine_h.device)
        nested_h = self.nested_pool_encoder.pool(
            contextual_h, nested_ids, torch.ones(1, dtype=torch.long, device=fine_h.device), token_mask, 1
        )
        nested_index = int(state["completed_nested_patches"])
        position_device = self.patch_pos_emb.weight.device
        nested_pos = torch.full((1, 1), nested_index, dtype=torch.long, device=position_device)
        nested_len = torch.full((1, 1), int(fine_h.shape[1]), dtype=torch.long, device=fine_h.device)
        nested_pos_h = self.patch_pos_emb(nested_pos.remainder(self.position_bins))
        if nested_pos_h.device != nested_device:
            nested_pos_h = nested_pos_h.to(nested_device, non_blocking=True)
        nested_h = nested_h + nested_pos_h + self.nested_len_emb(nested_len)
        state["current_nested_global"] = self._nested_global_stream_step(
            nested_h, state["nested_inference_params"], nested_index + 1
        )
        state["completed_nested_patches"] = nested_index + 1
        state["nested_context_tail"] = contextual_input[:, -2:].detach()
        state.setdefault("closed_nested_pools", []).append({
            "length": int(fine_h.shape[1]),
            "close_probability": float(
                state.get("pending_nested_close_probability", 0.0)
            ),
        })
        state["nested_fine_parts"] = []
        if self.tertiary_enabled:
            self._append_tertiary_stream_part(state, state["current_nested_global"])

    def _append_tertiary_stream_part(
        self, state: Dict[str, object], nested_completed_h: torch.Tensor
    ) -> None:
        """Route one completed level-2 state into the adaptive level-3 pool."""
        if not self.tertiary_enabled:
            return
        assert self.tertiary_pool_encoder is not None
        tertiary_device = self._module_device(self.tertiary_pool_encoder)
        current = nested_completed_h
        if current.device != tertiary_device:
            current = current.to(tertiary_device, non_blocking=True)
        routing_tail = state.get("tertiary_routing_tail")
        routing_input = current if routing_tail is None else torch.cat([routing_tail, current], dim=1)
        contextual_current = self.tertiary_pool_encoder.contextualize(routing_input)[:, -1:, :]
        probability = float(
            self.tertiary_pool_encoder.close_probability(contextual_current)[0, 0]
            .detach()
            .cpu()
            .item()
        )
        state.setdefault("tertiary_nested_parts", []).append(current)
        state["tertiary_routing_tail"] = routing_input[:, -2:].detach()
        state["pending_tertiary_close_probability"] = probability
        controller_adjustment = self._controller_close_adjustment(
            state.get("pending_pool_controller_scores"), 2
        )
        if self._hierarchy_should_close(
            probability,
            len(state["tertiary_nested_parts"]),
            self.tertiary_min_pool_factor,
            self.tertiary_pool_factor,
            self.tertiary_close_threshold,
            controller_adjustment,
        ):
            self._close_tertiary_stream_pool(state, force=True)

    def _close_tertiary_stream_pool(self, state: Dict[str, object], force: bool = False) -> None:
        """Commit completed level-2 states to the cached level-3 stream."""
        if not self.tertiary_enabled:
            return
        nested_parts = state.get("tertiary_nested_parts", [])
        if not nested_parts or (not force and len(nested_parts) < self.tertiary_pool_factor):
            return
        assert self.tertiary_pool_encoder is not None and self.tertiary_len_emb is not None
        tertiary_device = self._module_device(self.tertiary_pool_encoder)
        nested_h = torch.cat(nested_parts, dim=1)
        if nested_h.device != tertiary_device:
            nested_h = nested_h.to(tertiary_device, non_blocking=True)
        tail = state.get("tertiary_context_tail")
        contextual_input = nested_h if tail is None else torch.cat([tail, nested_h], dim=1)
        contextual_h = self.tertiary_pool_encoder.contextualize(contextual_input)[:, -nested_h.shape[1]:]
        token_mask = torch.ones(nested_h.shape[:2], dtype=torch.bool, device=tertiary_device)
        tertiary_ids = torch.zeros(nested_h.shape[:2], dtype=torch.long, device=tertiary_device)
        tertiary_h = self.tertiary_pool_encoder.pool(
            contextual_h,
            tertiary_ids,
            torch.ones(1, dtype=torch.long, device=tertiary_device),
            token_mask,
            1,
        )
        tertiary_index = int(state["completed_tertiary_patches"])
        position_device = self.patch_pos_emb.weight.device
        tertiary_pos = torch.full((1, 1), tertiary_index, dtype=torch.long, device=position_device)
        tertiary_pos_h = self.patch_pos_emb(tertiary_pos.remainder(self.position_bins))
        if tertiary_pos_h.device != tertiary_device:
            tertiary_pos_h = tertiary_pos_h.to(tertiary_device, non_blocking=True)
        tertiary_len = torch.full(
            (1, 1), int(nested_h.shape[1]), dtype=torch.long, device=tertiary_device
        )
        tertiary_h = tertiary_h + tertiary_pos_h + self.tertiary_len_emb(tertiary_len)
        state["current_tertiary_global"] = self._tertiary_global_stream_step(
            tertiary_h, state["tertiary_inference_params"], tertiary_index + 1
        )
        state["completed_tertiary_patches"] = tertiary_index + 1
        state["tertiary_context_tail"] = contextual_input[:, -2:].detach()
        state.setdefault("closed_tertiary_pools", []).append({
            "length": int(nested_h.shape[1]),
            "close_probability": float(
                state.get("pending_tertiary_close_probability", 0.0)
            ),
        })
        state["tertiary_nested_parts"] = []

    def _close_stream_patch(self, state: Dict[str, object]) -> None:
        patch_tokens = state.get("patch_tokens", [])
        if not patch_tokens:
            return
        contextual_h = torch.cat(state["patch_contextual_parts"], dim=1)
        patch_ids = torch.zeros(contextual_h.shape[:2], dtype=torch.long, device=contextual_h.device)
        token_mask = torch.ones(contextual_h.shape[:2], dtype=torch.bool, device=contextual_h.device)
        patch_h = self.pool_encoder.pool(
            contextual_h,
            patch_ids,
            torch.ones(1, dtype=torch.long, device=contextual_h.device),
            token_mask,
            1,
        )
        patch_index = int(state["completed_patches"])
        patch_h = self._patch_conditioning(
            patch_h,
            int(contextual_h.shape[1]),
            patch_index,
        )
        fine_completed_h = self._global_stream_step(
            patch_h,
            state["inference_params"],
            int(state["completed_patches"]) + 1,
        )
        state["current_global"] = fine_completed_h
        state["completed_patches"] = int(state["completed_patches"]) + 1
        self._append_nested_stream_part(state, fine_completed_h)
        state.setdefault("closed_pools", []).append({
            "length": int(contextual_h.shape[1]),
            "close_probability": float(state.get("pending_close_probability", 0.0)),
            "reason": str(state.get("pending_close_reason", "end")),
            "byte_pos": int(state.get("patch_end_byte_pos", 0)),
        })
        if int(contextual_h.shape[1]) < self.secondary_min_patch_bytes:
            state["short_pool_count"] = int(state.get("short_pool_count", 0)) + 1
        if self.short_pool_budget > 0 and self.short_pool_window > 0:
            history = state.setdefault("short_pool_history", deque())
            completed_short = int(contextual_h.shape[1]) < self.secondary_min_patch_bytes
            history.append(completed_short)
            state["short_pool_window_count"] = int(
                state.get("short_pool_window_count", 0)
            ) + int(completed_short)
            if len(history) > self.short_pool_window:
                state["short_pool_window_count"] -= int(history.popleft())
            state["secondary_min_active"] = (
                int(state["short_pool_window_count"]) >= self.short_pool_budget
            )
        elif (
            self.short_pool_budget > 0
            and int(state.get("short_pool_count", 0)) >= self.short_pool_budget
        ):
            state["secondary_min_active"] = True
        state["patch_tokens"] = []
        state["patch_contextual_parts"] = []
        state["patch_prev_raw"] = None
        state["pending_close_probability"] = 0.0
        state["pending_close_reason"] = "end"

    def new_stream_state(self, batch: Dict[str, torch.Tensor], max_patches: int) -> Dict[str, object]:
        inference_params = _new_mamba_inference_params(max_batch_size=1, max_seqlen=max(2, int(max_patches) + 1))
        # Adaptive groups can close at their minimum, so cache capacity must be
        # based on the shortest possible group rather than the hard maximum.
        nested_max_patches = max(
            2, int(math.ceil(int(max_patches) / self.nested_min_pool_factor)) + 1
        )
        nested_inference_params = _new_mamba_inference_params(max_batch_size=1, max_seqlen=nested_max_patches)
        tertiary_max_patches = max(
            2,
            int(
                math.ceil(
                    nested_max_patches / max(1, self.tertiary_min_pool_factor)
                )
            )
            + 1,
        )
        tertiary_inference_params = (
            _new_mamba_inference_params(max_batch_size=1, max_seqlen=tertiary_max_patches)
            if self.tertiary_enabled
            else None
        )
        if inference_params is None or nested_inference_params is None or (
            self.tertiary_enabled and tertiary_inference_params is None
        ):
            raise RuntimeError("BLT streaming evaluation requires mamba_ssm InferenceParams.")
        self._initialize_stream_mamba_caches(
            self.global_blocks, inference_params, max(2, int(max_patches) + 1)
        )
        self._initialize_stream_mamba_caches(
            self.nested_global_blocks,
            nested_inference_params,
            nested_max_patches,
        )
        if self.tertiary_enabled:
            self._initialize_stream_mamba_caches(
                self.tertiary_global_blocks,
                tertiary_inference_params,
                tertiary_max_patches,
            )
        boe = self.boe_patch.expand(1, 1, self.boe_patch.shape[-1])
        current_global = self._global_stream_step(boe, inference_params, 0)
        nested_device = self._module_device(self.nested_global_blocks[0])
        current_nested_global = self._nested_global_stream_step(
            boe.to(nested_device, non_blocking=True), nested_inference_params, 0
        )
        current_tertiary_global = None
        if self.tertiary_enabled:
            tertiary_device = self._module_device(self.tertiary_global_blocks[0])
            current_tertiary_global = self._tertiary_global_stream_step(
                boe.to(tertiary_device, non_blocking=True), tertiary_inference_params, 0
            )
        return {
            "inference_params": inference_params,
            "nested_inference_params": nested_inference_params,
            "tertiary_inference_params": tertiary_inference_params,
            "current_global": current_global,
            "current_nested_global": current_nested_global,
            "current_tertiary_global": current_tertiary_global,
            "initial_nested_global": current_nested_global.detach().clone(),
            "initial_tertiary_global": (
                current_tertiary_global.detach().clone()
                if current_tertiary_global is not None
                else None
            ),
            "prev_byte_latent": current_global,
            "completed_patches": 0,
            "completed_nested_patches": 0,
            "completed_tertiary_patches": 0,
            "patch_tokens": [],
            "patch_contextual_parts": [],
            "patch_prev_raw": None,
            "pending_close_probability": 0.0,
            "pending_close_reason": "end",
            "closed_pools": [],
            "nested_fine_parts": [],
            "closed_nested_pools": [],
            "nested_context_tail": None,
            "nested_routing_tail": None,
            "pending_nested_close_probability": 0.0,
            "tertiary_nested_parts": [],
            "closed_tertiary_pools": [],
            "tertiary_context_tail": None,
            "tertiary_routing_tail": None,
            "pending_tertiary_close_probability": 0.0,
            "pending_pool_controller_scores": None,
            "short_pool_count": 0,
            "short_pool_history": deque(),
            "short_pool_window_count": 0,
            "secondary_min_active": False,
            "local_conv1_tail": None,
            "local_conv2_tail": None,
            "pool_context_tail": None,
        }

    def stream_step(self, state: Dict[str, object], one: Dict[str, torch.Tensor]) -> torch.Tensor:
        # A controller emitted at t is consumed only when step t+1 begins.
        # Closing here makes any newly completed fine/L2/L3 state visible to
        # the next decode, while never allowing that decode to alter its own
        # routing decision.
        token = one["x"]
        contextual_current = self._local_stream_context(token, state)
        close_prob = float(
            self.pool_encoder.close_probability(contextual_current)[0, 0].detach().cpu().item()
        )
        patch_len = len(state.get("patch_tokens", []))
        runtime_minimum = (
            self.secondary_min_patch_bytes
            if bool(state.get("secondary_min_active", False))
            else self.min_patch_bytes
        )
        if patch_len:
            controller_adjustment = self._controller_close_adjustment(
                state.get("pending_pool_controller_scores"), 0
            )
            close_reason = None
            if patch_len >= self.max_patch_bytes:
                close_reason = "max"
            elif (
                patch_len >= runtime_minimum
                and self._close_score(close_prob, patch_len)
                + controller_adjustment
                >= self.close_threshold
            ):
                close_reason = (
                    "decoder_controller"
                    if self.decoder_pool_controller
                    else "learned"
                )
            if close_reason is not None:
                state["pending_close_probability"] = close_prob
                state["pending_close_reason"] = close_reason
                self._close_stream_patch(state)

        state["patch_tokens"].append(token)
        state["patch_contextual_parts"].append(contextual_current)

        local_last = contextual_current
        byte_latent = state["current_global"]
        prev_latent = state.get("prev_byte_latent", byte_latent)
        nested_latent = state["current_nested_global"]
        decoder_device = local_last.device
        if nested_latent.device != decoder_device:
            nested_latent = nested_latent.to(decoder_device, non_blocking=True)
        decode_parts = [local_last, byte_latent, prev_latent, nested_latent]
        if self.tertiary_enabled:
            tertiary_latent = state["current_tertiary_global"]
            if tertiary_latent.device != decoder_device:
                tertiary_latent = tertiary_latent.to(decoder_device, non_blocking=True)
            decode_parts.append(tertiary_latent)
        self.last_stream_decode_parts = tuple(part.detach() for part in decode_parts)
        decoded = self.decoder(torch.cat(decode_parts, dim=-1))
        logits = self.lm_head(decoded)
        controller_scores = self._pool_controller_scores(decoded)
        if controller_scores is not None:
            current_scores = controller_scores[0, 0].detach().float().cpu().tolist()
            state["pending_pool_controller_scores"] = current_scores
        else:
            state["pending_pool_controller_scores"] = None
        state["prev_byte_latent"] = byte_latent
        return logits

    def finish_stream(self, state: Dict[str, object]) -> None:
        """Flush an under-length pool at an explicit end-of-stream boundary."""
        self._close_stream_patch(state)
        self._close_nested_stream_pool(state, force=True)
        self._close_tertiary_stream_pool(state, force=True)

    def forward(
        self,
        x,
        patch_ids=None,
        patch_counts=None,
        max_patches=None,
        inference_params=None,
        runtime_pool_minima=None,
        **_unused,
    ):
        if self.model_parallel_enabled and inference_params is not None:
            raise RuntimeError("stateful nested streaming is not supported with --nested-devices; use full-window training/validation or a single device.")
        fine_minimum = self.min_patch_bytes
        nested_minimum = self.nested_min_pool_factor
        tertiary_minimum = self.tertiary_min_pool_factor
        per_token_minima = None
        if runtime_pool_minima is not None:
            if runtime_pool_minima.numel() == 3:
                minima = runtime_pool_minima.detach().to(device="cpu").reshape(-1).tolist()
                fine_minimum = max(1, min(int(minima[0]), self.max_patch_bytes))
                nested_minimum = max(1, min(int(minima[1]), self.nested_pool_factor))
                tertiary_minimum = max(1, min(int(minima[2]), self.tertiary_pool_factor))
            else:
                if runtime_pool_minima.ndim != 3 or runtime_pool_minima.shape[-1] != 3:
                    raise ValueError(
                        "runtime_pool_minima must be [3] or [batch, bytes, 3]"
                    )
                if tuple(runtime_pool_minima.shape[:2]) != tuple(x.shape):
                    raise ValueError(
                        "per-byte runtime_pool_minima must match x batch and length"
                    )
                per_token_minima = runtime_pool_minima.to(
                    device=x.device, dtype=torch.long, non_blocking=True
                ).clone()
                defaults = per_token_minima.new_tensor([
                    fine_minimum, nested_minimum, tertiary_minimum,
                ]).view(1, 1, 3)
                per_token_minima = torch.where(
                    per_token_minima.gt(0), per_token_minima, defaults
                )
                per_token_minima[..., 0].clamp_(1, self.max_patch_bytes)
                per_token_minima[..., 1].clamp_(1, self.nested_pool_factor)
                per_token_minima[..., 2].clamp_(1, self.tertiary_pool_factor)
        local_h = self._local_features(x)
        contextual_h = self.pool_encoder.contextualize(local_h)
        # Pool boundaries are model predictions.  Ignore legacy precomputed
        # greedy IDs so cached batches/checkpoints remain usable during the
        # transition to learned pooling.
        token_mask = x.ne(0)
        close_logits = self.pool_encoder.close_logits(contextual_h)
        close_probabilities = torch.sigmoid(close_logits)
        self.last_pool_close_probabilities = close_probabilities
        patch_ids = self._runtime_patch_ids(
            contextual_h, token_mask, close_probabilities, minimum=fine_minimum,
            minima_by_token=per_token_minima,
        )
        patch_counts = patch_ids.masked_fill(~token_mask, 0).amax(dim=1) + 1
        max_patches = int(patch_counts.max().detach().cpu().item())
        patch_h, patch_lens, patch_mask = self._pool_patches(
            contextual_h,
            patch_ids,
            token_mask,
            patch_counts=patch_counts,
            max_patches=max_patches,
        )
        patch_minima = (
            self._pool_runtime_minima(
                per_token_minima, patch_ids, token_mask, max_patches
            )
            if per_token_minima is not None else None
        )
        valid_pool_lens = patch_lens[patch_mask]
        self.last_pool_metrics = {
            "pool_length_sum": valid_pool_lens.detach().float().sum(),
            "pool_count": valid_pool_lens.detach().new_tensor(float(valid_pool_lens.numel())),
            "pool_length_mean": valid_pool_lens.detach().float().mean() if valid_pool_lens.numel() else patch_h.new_zeros(()),
            "pools_per_window": patch_counts.detach().float().mean(),
        }
        budget_metrics = getattr(self, "last_short_pool_budget_metrics", {})
        self.last_pool_metrics.update({
            "short_pool_count": patch_h.new_tensor(
                float(budget_metrics.get("short_pool_count", 0.0))
            ),
            "short_pool_budget_activated_rows": patch_h.new_tensor(
                float(budget_metrics.get("budget_activated_rows", 0.0))
            ),
            "short_pool_budget_batch_rows": patch_h.new_tensor(
                float(budget_metrics.get("batch_rows", contextual_h.shape[0]))
            ),
        })
        b, p, c = patch_h.shape
        patch_pos = torch.arange(p, device=x.device, dtype=torch.long).view(1, p).expand(b, p)
        patch_h = patch_h + self.patch_pos_emb(patch_pos.remainder(self.position_bins)) + self.patch_len_emb(patch_lens)

        patch_h = patch_h.masked_fill(~patch_mask.unsqueeze(-1), 0.0)
        boe_patch = self.boe_patch.expand(b, 1, c)
        global_h = torch.cat([boe_patch, patch_h], dim=1)
        global_mask = torch.cat(
            [torch.ones((b, 1), dtype=torch.bool, device=x.device), patch_mask],
            dim=1,
        )
        for block in self.global_blocks:
            global_h = self._run_block_on_own_device(
                block, global_h, inference_params=inference_params
            )
            global_h = global_h.masked_fill(~global_mask.unsqueeze(-1), 0.0)

        # Causal BLT alignment: bytes in patch p decode from global state p,
        # which is BOE for patch 0 and completed patch p-1 thereafter.
        byte_latents = global_h.gather(1, patch_ids.unsqueeze(-1).expand(-1, -1, c))
        prev_latents = torch.cat([byte_latents[:, :1], byte_latents[:, :-1]], dim=1)

        # Nest completed fine-pool states into a coarser causal stream.  Each
        # coarse state is only read by a *later* group of fine pools, so the
        # multi-scale decoder never receives information from its current
        # target pool.  This lets later layers operate on roughly
        # ``nested_pool_factor`` fewer positions without losing fine readouts.
        fine_h = global_h[:, 1:]
        # The nested hierarchy is substantially shorter than the byte path.
        # In model-parallel mode, make that single fine->coarse transfer here.
        nested_device = self._module_device(self.nested_pool_encoder)
        fine_h_nested = fine_h if fine_h.device == nested_device else fine_h.to(nested_device, non_blocking=True)
        patch_mask_nested = patch_mask if patch_mask.device == nested_device else patch_mask.to(nested_device, non_blocking=True)
        patch_counts_nested = patch_counts if patch_counts.device == nested_device else patch_counts.to(nested_device, non_blocking=True)
        # All IDs used by coarse pooling must live with the coarse activations.
        fine_contextual_h = self.nested_pool_encoder.contextualize(fine_h_nested)
        nested_close_logits = self.nested_pool_encoder.close_logits(fine_contextual_h)
        nested_close_probabilities = torch.sigmoid(nested_close_logits)
        nested_ids = self._runtime_hierarchy_ids(
            fine_contextual_h,
            patch_mask_nested,
            self.nested_pool_encoder,
            nested_minimum,
            self.nested_pool_factor,
            self.nested_close_threshold,
            minima_by_token=(
                patch_minima[..., 1].to(nested_device, non_blocking=True)
                if patch_minima is not None else None
            ),
            policy_by_token=(
                patch_minima.to(nested_device, non_blocking=True)
                if patch_minima is not None else None
            ),
        )
        nested_counts = nested_ids.masked_fill(~patch_mask_nested, 0).amax(dim=1) + 1
        max_nested_patches = int(nested_counts.max().detach().cpu().item())
        nested_h = self.nested_pool_encoder.pool(
            fine_contextual_h,
            nested_ids,
            nested_counts,
            patch_mask_nested,
            max_nested_patches,
        )
        nested_minima = (
            self._pool_runtime_minima(
                patch_minima.to(nested_device, non_blocking=True),
                nested_ids,
                patch_mask_nested,
                max_nested_patches,
            )
            if patch_minima is not None else None
        )
        nested_token_counts = fine_h_nested.new_zeros((b, max_nested_patches, 1))
        nested_token_counts.scatter_add_(
            1,
            nested_ids.unsqueeze(-1),
            patch_mask_nested.unsqueeze(-1).to(dtype=fine_h_nested.dtype),
        )
        nested_lens = nested_token_counts.squeeze(-1).round().long().clamp(0, self.nested_pool_factor)
        nested_mask = torch.arange(max_nested_patches, device=nested_device).view(1, -1) < nested_counts.view(-1, 1)
        valid_nested_lens = nested_lens[nested_mask]
        self.last_pool_metrics.update({
            "nested_pool_length_sum": valid_nested_lens.detach().float().sum(),
            "nested_pool_count": valid_nested_lens.detach().new_tensor(float(valid_nested_lens.numel())),
            "nested_pool_length_mean": (
                valid_nested_lens.detach().float().mean()
                if valid_nested_lens.numel()
                else nested_h.new_zeros(())
            ),
            "nested_pools_per_window": nested_counts.detach().float().mean(),
        })
        nested_pos = torch.arange(max_nested_patches, device=x.device, dtype=torch.long).view(1, -1).expand(b, -1)
        # patch_pos_emb is intentionally shared with the fine stage, so look
        # it up on the fine device and transfer its small coarse sequence.
        nested_pos_emb = self.patch_pos_emb(nested_pos.remainder(self.position_bins)).to(nested_device, non_blocking=True)
        nested_h = (
            nested_h
            + nested_pos_emb
            + self.nested_len_emb(nested_lens)
        ).masked_fill(~nested_mask.unsqueeze(-1), 0.0)
        nested_boe = boe_patch if boe_patch.device == nested_device else boe_patch.to(nested_device, non_blocking=True)
        nested_global_h = torch.cat([nested_boe, nested_h], dim=1)
        nested_global_mask = torch.cat(
            [torch.ones((b, 1), dtype=torch.bool, device=nested_device), nested_mask],
            dim=1,
        )
        nested_global_h = self._run_nested_blocks(nested_global_h, nested_global_mask)
        tertiary_byte_latents = None
        tertiary_initial_byte_latents = None
        tertiary_ids_fine = None
        tertiary_lens = None
        tertiary_mask = None
        tertiary_close_probabilities = None
        tertiary_close_logits = None
        tertiary_ids = None
        level2_mask = None
        if self.tertiary_enabled:
            assert self.tertiary_pool_encoder is not None and self.tertiary_len_emb is not None
            tertiary_device = self._module_device(self.tertiary_pool_encoder)
            level2_h = nested_global_h[:, 1:]
            if level2_h.device != tertiary_device:
                level2_h = level2_h.to(tertiary_device, non_blocking=True)
            level2_mask = nested_mask if nested_mask.device == tertiary_device else nested_mask.to(tertiary_device, non_blocking=True)
            level2_counts = nested_counts if nested_counts.device == tertiary_device else nested_counts.to(tertiary_device, non_blocking=True)
            level2_contextual = self.tertiary_pool_encoder.contextualize(level2_h)
            tertiary_close_logits = self.tertiary_pool_encoder.close_logits(level2_contextual)
            tertiary_close_probabilities = torch.sigmoid(tertiary_close_logits)
            tertiary_ids = self._runtime_hierarchy_ids(
                level2_contextual,
                level2_mask,
                self.tertiary_pool_encoder,
                tertiary_minimum,
                self.tertiary_pool_factor,
                self.tertiary_close_threshold,
                minima_by_token=(
                    nested_minima[..., 2].to(tertiary_device, non_blocking=True)
                    if nested_minima is not None else None
                ),
                policy_by_token=(
                    nested_minima.to(tertiary_device, non_blocking=True)
                    if nested_minima is not None else None
                ),
            )
            tertiary_counts = tertiary_ids.masked_fill(~level2_mask, 0).amax(dim=1) + 1
            max_tertiary = int(tertiary_counts.max().detach().cpu().item())
            tertiary_h = self.tertiary_pool_encoder.pool(level2_contextual, tertiary_ids, tertiary_counts, level2_mask, max_tertiary)
            tertiary_token_counts = level2_h.new_zeros((b, max_tertiary, 1))
            tertiary_token_counts.scatter_add_(1, tertiary_ids.unsqueeze(-1), level2_mask.unsqueeze(-1).to(level2_h.dtype))
            tertiary_lens = tertiary_token_counts.squeeze(-1).round().long().clamp(0, self.tertiary_pool_factor)
            tertiary_mask = torch.arange(max_tertiary, device=tertiary_device).view(1, -1) < tertiary_counts.view(-1, 1)
            valid_tertiary_lens = tertiary_lens[tertiary_mask]
            self.last_pool_metrics.update({
                "tertiary_pool_length_sum": valid_tertiary_lens.detach().float().sum(),
                "tertiary_pool_count": valid_tertiary_lens.detach().new_tensor(float(valid_tertiary_lens.numel())),
                "tertiary_pool_length_mean": (
                    valid_tertiary_lens.detach().float().mean()
                    if valid_tertiary_lens.numel()
                    else tertiary_h.new_zeros(())
                ),
                "tertiary_pools_per_window": tertiary_counts.detach().float().mean(),
            })
            tertiary_pos = torch.arange(max_tertiary, device=x.device, dtype=torch.long).view(1, -1).expand(b, -1)
            tertiary_pos_emb = self.patch_pos_emb(tertiary_pos.remainder(self.position_bins)).to(tertiary_device, non_blocking=True)
            tertiary_h = (tertiary_h + tertiary_pos_emb + self.tertiary_len_emb(tertiary_lens)).masked_fill(~tertiary_mask.unsqueeze(-1), 0.0)
            tertiary_boe = boe_patch.to(tertiary_device, non_blocking=True)
            tertiary_global_h = self._run_tertiary_blocks(
                torch.cat([tertiary_boe, tertiary_h], dim=1),
                torch.cat([torch.ones((b, 1), dtype=torch.bool, device=tertiary_device), tertiary_mask], dim=1),
            )
            # Map level-3 causal readouts back to level 2, then fine patches.
            tertiary_global_h = tertiary_global_h.to(x.device, non_blocking=True)
            tertiary_initial_byte_latents = tertiary_global_h[:, :1].expand(
                -1, x.shape[1], -1
            )
            tertiary_ids_fine = tertiary_ids.to(x.device, non_blocking=True)
            level2_tertiary = tertiary_global_h.gather(1, tertiary_ids_fine.unsqueeze(-1).expand(-1, -1, c))
            nested_ids_fine_for_tertiary = nested_ids.to(x.device, non_blocking=True)
            tertiary_byte_latents = level2_tertiary.gather(1, nested_ids_fine_for_tertiary.unsqueeze(-1).expand(-1, -1, c)).gather(1, patch_ids.unsqueeze(-1).expand(-1, -1, c))
        # Return only the coarse result; decoder and all byte-level gathers
        # remain on the fine device.
        if nested_global_h.device != x.device:
            nested_global_h = nested_global_h.to(x.device, non_blocking=True)
        nested_ids_fine = nested_ids if nested_ids.device == x.device else nested_ids.to(x.device, non_blocking=True)
        nested_readouts = nested_global_h.gather(1, nested_ids_fine.unsqueeze(-1).expand(-1, -1, c))
        nested_byte_latents = nested_readouts.gather(1, patch_ids.unsqueeze(-1).expand(-1, -1, c))
        nested_byte_pool_ids = nested_ids_fine.gather(1, patch_ids)
        nested_byte_latents = self._gate_inactive_coarse_gradient(
            nested_byte_latents,
            nested_byte_pool_ids.gt(0),
            token_mask,
        )
        nested_initial_byte_latents = nested_global_h[:, :1].expand(-1, x.shape[1], -1)
        if tertiary_byte_latents is not None and tertiary_ids_fine is not None:
            tertiary_fine_pool_ids = tertiary_ids_fine.gather(1, nested_ids_fine)
            tertiary_byte_pool_ids = tertiary_fine_pool_ids.gather(1, patch_ids)
            tertiary_byte_latents = self._gate_inactive_coarse_gradient(
                tertiary_byte_latents,
                tertiary_byte_pool_ids.gt(0),
                token_mask,
            )
        decode_parts = [contextual_h, byte_latents, prev_latents, nested_byte_latents]
        if tertiary_byte_latents is not None:
            decode_parts.append(tertiary_byte_latents)
        if self.capture_evaluation_details:
            self.last_forward_evaluation_details = {
                "token_mask": token_mask.detach(),
                "patch_ids": patch_ids.detach(),
                "patch_lens": patch_lens.detach(),
                "patch_mask": patch_mask.detach(),
                "nested_ids": nested_ids_fine.detach(),
                "nested_lens": nested_lens.detach(),
                "nested_mask": nested_mask.detach(),
                "tertiary_ids": (
                    tertiary_ids_fine.detach() if tertiary_ids_fine is not None else None
                ),
                "tertiary_lens": (
                    tertiary_lens.detach() if tertiary_lens is not None else None
                ),
                "tertiary_mask": (
                    tertiary_mask.detach() if tertiary_mask is not None else None
                ),
                "decode_parts": tuple(part.detach() for part in decode_parts),
                "initial_nested_byte_latents": nested_initial_byte_latents.detach(),
                "initial_tertiary_byte_latents": (
                    tertiary_initial_byte_latents.detach()
                    if tertiary_initial_byte_latents is not None
                    else None
                ),
            }
        self.last_pool_routing_training = {
            "fine_logits": close_logits,
            "fine_ids": patch_ids.detach(),
            "fine_mask": token_mask.detach(),
            "fine_minimum": fine_minimum,
            "fine_maximum": self.max_patch_bytes,
            "fine_threshold": self.close_threshold,
            "nested_logits": nested_close_logits,
            "nested_ids": nested_ids.detach(),
            "nested_mask": patch_mask_nested.detach(),
            "nested_minimum": nested_minimum,
            "nested_maximum": self.nested_pool_factor,
            "nested_threshold": self.nested_close_threshold,
            "tertiary_logits": tertiary_close_logits,
            "tertiary_ids": tertiary_ids.detach() if tertiary_ids is not None else None,
            "tertiary_mask": level2_mask.detach() if level2_mask is not None else None,
            "tertiary_minimum": tertiary_minimum,
            "tertiary_maximum": self.tertiary_pool_factor,
            "tertiary_threshold": self.tertiary_close_threshold,
        }
        decoded = self.decoder(torch.cat(decode_parts, dim=-1))
        controller_scores = self._pool_controller_scores(decoded)
        self.last_pool_controller_scores = controller_scores
        if controller_scores is not None:
            nested_signal_fine = nested_close_probabilities.to(
                x.device, non_blocking=True
            ).gather(1, patch_ids)
            if tertiary_close_probabilities is not None and tertiary_ids_fine is not None:
                tertiary_signal_level2 = tertiary_close_probabilities.to(
                    x.device, non_blocking=True
                ).gather(1, nested_ids_fine)
                tertiary_signal_fine = tertiary_signal_level2.gather(1, patch_ids)
            else:
                tertiary_signal_fine = nested_signal_fine
            self.last_pool_controller_close_signals = torch.stack(
                [close_probabilities, nested_signal_fine, tertiary_signal_fine],
                dim=-1,
            )
        else:
            self.last_pool_controller_close_signals = None
        return self.lm_head(decoded)

    @staticmethod
    def _routing_decision_ages(
        ids: torch.Tensor, mask: torch.Tensor, *, closes_before_token: bool
    ) -> torch.Tensor:
        """Age seen by a boundary decision, stable across the resulting close."""
        batch, length = ids.shape
        positions = torch.arange(length, device=ids.device).view(1, -1).expand(batch, -1)
        group_count = max(
            1, int(ids.masked_fill(~mask, 0).amax().detach().cpu().item()) + 1
        )
        safe_ids = ids.masked_fill(~mask, 0).clamp(0, group_count - 1)
        starts = torch.full(
            (batch, group_count), length, dtype=torch.long, device=ids.device
        )
        starts.scatter_reduce_(
            1, safe_ids, positions.masked_fill(~mask, length),
            reduce="amin", include_self=True,
        )
        ages = (positions - starts.gather(1, safe_ids) + 1).masked_fill(~mask, 0)
        if not closes_before_token:
            return ages
        # Fine routing closes the previous patch immediately before the current
        # byte. On a boundary byte, supervise from the completed previous
        # length so the positive label remains stable after the close occurs.
        counts = torch.zeros(
            (batch, group_count), dtype=torch.long, device=ids.device
        )
        counts.scatter_add_(1, safe_ids, mask.long())
        previous_lengths = counts.gather(1, (safe_ids - 1).clamp_min(0))
        starts_group = mask & positions.gt(0) & ages.eq(1)
        return torch.where(starts_group, previous_lengths, (ages - 1).clamp_min(0))

    @staticmethod
    def _routing_aggregate(
        values: torch.Tensor, ids: torch.Tensor, mask: torch.Tensor
    ) -> torch.Tensor:
        values = values.to(ids.device)
        mask = mask.to(ids.device)
        groups = max(
            1, int(ids.masked_fill(~mask, 0).amax().detach().cpu().item()) + 1
        )
        safe_ids = ids.masked_fill(~mask, 0).clamp(0, groups - 1)
        sums = values.new_zeros((values.shape[0], groups))
        counts = values.new_zeros((values.shape[0], groups))
        sums.scatter_add_(1, safe_ids, values.masked_fill(~mask, 0.0))
        counts.scatter_add_(1, safe_ids, mask.to(values.dtype))
        return sums / counts.clamp_min(1.0)

    @staticmethod
    def _routing_future_surprise(
        values: torch.Tensor, valid: torch.Tensor
    ) -> torch.Tensor:
        future = torch.cat([values[:, 1:], values[:, -1:]], dim=1).detach().float()
        valid = valid.to(future.device)
        selected = future[valid]
        if not selected.numel():
            return future.new_zeros(future.shape)
        return torch.tanh(
            (future - selected.mean()) / selected.std(unbiased=False).clamp_min(1e-4)
        ).masked_fill(~valid, 0.0)

    @staticmethod
    def _routing_level_objective(
        logits: torch.Tensor,
        ids: torch.Tensor,
        mask: torch.Tensor,
        surprise: torch.Tensor,
        *,
        minimum: int,
        maximum: int,
        target_length: int,
        semantic_span: float,
        closes_before_token: bool,
        threshold: float,
        collect_metrics: bool = True,
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        mask = mask.to(logits.device)
        ids = ids.to(logits.device)
        surprise = surprise.to(logits.device)
        minimum = max(1, min(int(minimum), int(maximum)))
        target_length = max(minimum, min(int(target_length), int(maximum)))
        valid_logits = logits.float()[mask]
        if not valid_logits.numel():
            return logits.float().sum() * 0.0, {}
        valid_surprise = surprise.float()[mask]
        # Select a fixed causal refresh budget, assigning it to positions whose
        # following region was hardest to predict. This makes the close head a
        # semantic selector rather than asking a content-only network to infer
        # an invisible current-pool age.
        target_rate = min(1.0, 1.0 / max(1.0, float(target_length)))
        positive_count = max(1, int(round(valid_logits.numel() * target_rate)))
        positive_count = min(positive_count, valid_logits.numel())
        selected = torch.topk(valid_surprise, positive_count, sorted=False).indices
        labels = torch.zeros_like(valid_logits)
        labels[selected] = 1.0
        threshold = max(1e-4, min(1.0 - 1e-4, float(threshold)))
        threshold_logit = math.log(threshold / (1.0 - threshold))
        decision_logits = (valid_logits - threshold_logit) / 0.5
        semantic_loss = F.binary_cross_entropy_with_logits(decision_logits, labels)
        soft_rate = torch.sigmoid(decision_logits).mean()
        rate_loss = (soft_rate - target_rate) ** 2
        objective = max(0.0, float(semantic_span)) * semantic_loss + 4.0 * rate_loss
        if not collect_metrics:
            return objective, {}
        probabilities = torch.sigmoid(valid_logits.detach())
        return objective, {
            "prob_mean": float(probabilities.mean().cpu()),
            "prob_std": float(probabilities.std(unbiased=False).cpu()),
            "above_threshold": float(probabilities.ge(float(threshold)).float().mean().cpu()),
            "target_length": float(target_length),
            "target_close_rate": float(target_rate),
            "soft_close_rate": float(soft_rate.detach().cpu()),
        }

    def pool_routing_auxiliary_loss(
        self,
        logits: torch.Tensor,
        targets: torch.Tensor,
        *,
        fine_target: int = 0,
        nested_target: int = 0,
        tertiary_target: int = 0,
        semantic_span: float = 2.0,
        collect_metrics: bool = True,
    ) -> Tuple[torch.Tensor, Dict[str, float]]:
        """Explicitly train causal close heads against rate and surprise labels."""
        state = self.last_pool_routing_training
        if not state:
            return logits.float().sum() * 0.0, {}
        with torch.no_grad():
            byte_loss = F.cross_entropy(
                logits.detach().float().transpose(1, 2), targets,
                ignore_index=-100, reduction="none",
            )
            fine_mask = state["fine_mask"]
            scored_bytes = targets.ne(-100).to(fine_mask.device) & fine_mask
            fine_surprise = self._routing_future_surprise(
                byte_loss.to(fine_mask.device), scored_bytes
            )
            fine_group_loss = self._routing_aggregate(
                byte_loss, state["fine_ids"], scored_bytes
            )
            nested_mask = state["nested_mask"]
            nested_surprise = self._routing_future_surprise(
                fine_group_loss.to(nested_mask.device), nested_mask
            )
            nested_group_loss = self._routing_aggregate(
                fine_group_loss, state["nested_ids"], nested_mask
            )
            tertiary_mask = state["tertiary_mask"]
            tertiary_surprise = (
                self._routing_future_surprise(
                    nested_group_loss.to(tertiary_mask.device), tertiary_mask
                )
                if tertiary_mask is not None else None
            )

        def target(requested: int, prefix: str) -> int:
            minimum = int(state[f"{prefix}_minimum"])
            maximum = int(state[f"{prefix}_maximum"])
            return (
                max(minimum, min(int(requested), maximum))
                if int(requested) > 0
                else max(minimum, min(maximum, minimum * 4))
            )

        levels = [
            ("fine", "fine", fine_target, fine_surprise, True),
            ("level2", "nested", nested_target, nested_surprise, False),
        ]
        if state["tertiary_logits"] is not None:
            levels.append((
                "level3", "tertiary", tertiary_target, tertiary_surprise, False
            ))
        objectives: List[torch.Tensor] = []
        metrics: Dict[str, float] = {}
        destination = logits.device
        for name, prefix, requested, surprise, before in levels:
            level_objective, level_metrics = self._routing_level_objective(
                state[f"{prefix}_logits"], state[f"{prefix}_ids"],
                state[f"{prefix}_mask"], surprise,
                minimum=int(state[f"{prefix}_minimum"]),
                maximum=int(state[f"{prefix}_maximum"]),
                target_length=target(requested, prefix),
                semantic_span=semantic_span,
                closes_before_token=before,
                threshold=float(state[f"{prefix}_threshold"]),
                collect_metrics=collect_metrics,
            )
            objectives.append(level_objective.to(destination, non_blocking=True))
            metrics.update({f"{name}_{key}": value for key, value in level_metrics.items()})
        return torch.stack(objectives).mean(), metrics


class ForwardBackwardRepairModel(nn.Module):
    def __init__(
        self,
        vocab_size: int,
        dim: int = 256,
        layers: int = 4,
        position_bins: int = 8192,
        use_mamba: bool = True,
        mamba_version: int = 2,
        mamba_d_state: int = 64,
        mamba2_headdim: int = 0,
        bidirectional: bool = False,
        num_sections: int = 16,
        min_patch_bytes: int = 4,
        max_patch_bytes: int = 16,
        patch_change_threshold: int = 48,
        close_threshold: float = 0.90,
        mid_close_bonus: float = 0.05,
        nested_pool_factor: int = 16,
        nested_min_pool_factor: int = 0,
        nested_close_threshold: Optional[float] = None,
        nested_layers: int = 0,
        tertiary_pool_factor: int = 0,
        tertiary_min_pool_factor: int = 0,
        tertiary_close_threshold: Optional[float] = None,
        tertiary_layers: int = 0,
        decoder_dim: int = 0,
        detach_inactive_coarse_gradients: bool = True,
        legacy_pool_closure_training: bool = False,
        decoder_pool_controller: bool = False,
        pool_controller_alpha: float = 1.0,
        pool_controller_beta: float = 0.5,
        pool_controller_gamma: float = 0.5,
        short_pool_budget: int = 0,
        short_pool_window: int = 0,
        secondary_min_patch_bytes: int = 16,
    ):
        super().__init__()
        if bidirectional:
            raise ValueError("BLT Mamba trainer is forward-only; disable --bidirectional-repair.")
        self.bidirectional = False
        self.forward_model = ByteLatentMambaCore(
            vocab_size,
            dim=dim,
            layers=layers,
            position_bins=position_bins,
            use_mamba=use_mamba,
            mamba_version=mamba_version,
            mamba_d_state=mamba_d_state,
            mamba2_headdim=mamba2_headdim,
            num_sections=num_sections,
            min_patch_bytes=min_patch_bytes,
            max_patch_bytes=max_patch_bytes,
            patch_change_threshold=patch_change_threshold,
            close_threshold=close_threshold,
            mid_close_bonus=mid_close_bonus,
            nested_pool_factor=nested_pool_factor,
            nested_min_pool_factor=nested_min_pool_factor,
            nested_close_threshold=nested_close_threshold,
            nested_layers=nested_layers,
            tertiary_pool_factor=tertiary_pool_factor,
            tertiary_min_pool_factor=tertiary_min_pool_factor,
            tertiary_close_threshold=tertiary_close_threshold,
            tertiary_layers=tertiary_layers,
            decoder_dim=decoder_dim,
            detach_inactive_coarse_gradients=detach_inactive_coarse_gradients,
            legacy_pool_closure_training=legacy_pool_closure_training,
            decoder_pool_controller=decoder_pool_controller,
            pool_controller_alpha=pool_controller_alpha,
            pool_controller_beta=pool_controller_beta,
            pool_controller_gamma=pool_controller_gamma,
            short_pool_budget=short_pool_budget,
            short_pool_window=short_pool_window,
            secondary_min_patch_bytes=secondary_min_patch_bytes,
        )

    def configure_model_parallel(self, fine_device: torch.device, nested_devices: List[torch.device], tertiary_device: Optional[torch.device] = None) -> None:
        """Enable nested-level model parallelism without changing checkpoint layout."""
        self.forward_model.configure_model_parallel(fine_device, nested_devices, tertiary_device=tertiary_device)

    def forward(self, batch, inference_params=None):
        return self.forward_model(**batch, inference_params=inference_params)