from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint

from .attention import GQAAttentionCache, GatedGroupedQueryAttention
from .attn_res import (
    SlicedDepthAttentionRouter,
    SublayerDepthRouters,
    sliced_block_accumulate,
    sliced_phase1_from_logits,
    sliced_phase2_merge,
    sliced_source_logits,
)
from .config import YunmoNextConfig
from .kda_runtime import configure_kda_mixer_runtime
from .packed import document_position_ids
from .training_engine import TrainingEngineConfig


ATTENTION_RESIDUAL_QUERY_INIT_SEED = 17
ATTENTION_RESIDUAL_QUERY_INIT_OFFSET = 7_900_031
ATTENTION_RESIDUAL_QUERY_INIT_STD = 0.02
ATTENTION_RESIDUAL_ORTHOGONAL_QUERY_DEPTH = 24
ACTIVATION_CHECKPOINT_POLICY_FULL = "full"
ACTIVATION_CHECKPOINT_POLICIES = frozenset({ACTIVATION_CHECKPOINT_POLICY_FULL})


def orthogonal_attention_residual_query_bank(
    query_count: int,
    key_dim: int,
    *,
    device: torch.device | str = "cpu",
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """Return the adopted private-seed orthogonal 24-layer query bank.

    The first row is generated for lineage stability but remains zero in the
    model: the first pre-mixer read has one completed source, so that query is
    structurally unidentifiable.  A private CPU generator prevents this
    initialization from advancing the common model RNG stream.
    """

    if query_count <= 0 or key_dim <= 0 or query_count > key_dim:
        raise ValueError(
            "orthogonal Attention Residual query bank requires "
            "0 < queries <= key_dim"
        )
    generator = torch.Generator(device="cpu")
    generator.manual_seed(
        (
            ATTENTION_RESIDUAL_QUERY_INIT_SEED
            + ATTENTION_RESIDUAL_QUERY_INIT_OFFSET
        )
        % (2**63 - 1)
    )
    basis = torch.randn(
        key_dim,
        query_count,
        generator=generator,
        dtype=torch.float32,
    )
    rows = torch.linalg.qr(basis, mode="reduced").Q.T.contiguous()
    rows.mul_((key_dim**0.5) * ATTENTION_RESIDUAL_QUERY_INIT_STD)
    return rows.to(device=device, dtype=dtype)


def _build_recurrent_mixer(
    config: YunmoNextConfig,
    engine: TrainingEngineConfig,
    layer_idx: int,
) -> nn.Module:
    try:
        from fla.layers import KimiDeltaAttention
    except ImportError as exc:
        if not torch.cuda.is_available():
            if engine.recurrent_segment_tokens:
                raise RuntimeError(
                    "serial recurrent execution requires the CUDA KDA runtime"
                ) from exc
            from .reference_recurrent import ReferenceKimiDeltaAttention

            return ReferenceKimiDeltaAttention(
                hidden_size=config.hidden_size,
                head_dim=config.recurrent_head_dim,
                num_heads=config.recurrent_num_heads,
                expand_v=config.recurrent_expand_v,
                use_short_conv=True,
                conv_size=config.short_conv_size,
                layer_idx=layer_idx,
                norm_eps=config.rms_norm_eps,
            )
        raise ImportError(
            "Yunmo-Next recurrent layers require the FLA CUDA environment"
        ) from exc

    mixer = KimiDeltaAttention(
        hidden_size=config.hidden_size,
        head_dim=config.recurrent_head_dim,
        num_heads=config.recurrent_num_heads,
        expand_v=config.recurrent_expand_v,
        use_short_conv=True,
        conv_size=config.short_conv_size,
        mode="chunk",
        layer_idx=layer_idx,
        norm_eps=config.rms_norm_eps,
    )
    return configure_kda_mixer_runtime(mixer, engine)


class _SerialChunkedSwiGLU(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        hidden_states: torch.Tensor,
        gate_weight: torch.Tensor,
        up_weight: torch.Tensor,
        down_weight: torch.Tensor,
        chunk_size: int,
    ) -> torch.Tensor:
        ctx.chunk_size = chunk_size
        ctx.save_for_backward(hidden_states, gate_weight, up_weight, down_weight)
        return torch.cat(
            [
                F.linear(
                    F.silu(F.linear(chunk, gate_weight))
                    * F.linear(chunk, up_weight),
                    down_weight,
                )
                for chunk in hidden_states.split(chunk_size, dim=-2)
            ],
            dim=-2,
        )

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor):
        hidden_states, gate_weight, up_weight, down_weight = ctx.saved_tensors
        grad_hidden_chunks: list[torch.Tensor] = []
        accumulator_dtype = (
            torch.float64 if gate_weight.dtype == torch.float64 else torch.float32
        )
        grad_gate = torch.zeros_like(gate_weight, dtype=accumulator_dtype)
        grad_up = torch.zeros_like(up_weight, dtype=accumulator_dtype)
        grad_down = torch.zeros_like(down_weight, dtype=accumulator_dtype)
        hidden_chunks = hidden_states.split(ctx.chunk_size, dim=-2)
        output_grad_chunks = grad_output.split(ctx.chunk_size, dim=-2)
        create_graph = torch.is_grad_enabled()
        for hidden_chunk, output_grad_chunk in zip(
            hidden_chunks, output_grad_chunks, strict=True
        ):
            with torch.enable_grad():
                chunk = hidden_chunk.detach().requires_grad_(True)
                gate = gate_weight.detach().requires_grad_(True)
                up = up_weight.detach().requires_grad_(True)
                down = down_weight.detach().requires_grad_(True)
                output = F.linear(
                    F.silu(F.linear(chunk, gate)) * F.linear(chunk, up),
                    down,
                )
                chunk_grad, gate_grad, up_grad, down_grad = torch.autograd.grad(
                    output,
                    (chunk, gate, up, down),
                    output_grad_chunk,
                    create_graph=create_graph,
                )
            grad_hidden_chunks.append(chunk_grad)
            grad_gate.add_(gate_grad.to(accumulator_dtype))
            grad_up.add_(up_grad.to(accumulator_dtype))
            grad_down.add_(down_grad.to(accumulator_dtype))
        return (
            torch.cat(grad_hidden_chunks, dim=-2),
            grad_gate.to(gate_weight.dtype),
            grad_up.to(up_weight.dtype),
            grad_down.to(down_weight.dtype),
            None,
        )


class SwiGLU(nn.Module):
    def __init__(
        self,
        hidden_size: int,
        intermediate_size: int,
        engine: TrainingEngineConfig,
    ) -> None:
        super().__init__()
        self.sequence_chunk_size = engine.feed_forward_chunk_tokens
        self.sequence_serial_backward = engine.feed_forward_serial_backward
        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)

    def _project(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gate = self.gate_proj(hidden_states)
        up = self.up_proj(hidden_states)
        activated = F.silu(gate) * up
        return self.down_proj(activated)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        chunk_size = self.sequence_chunk_size
        if chunk_size <= 0 or hidden_states.shape[-2] <= chunk_size:
            return self._project(hidden_states)
        if self.sequence_serial_backward:
            return _SerialChunkedSwiGLU.apply(
                hidden_states,
                self.gate_proj.weight,
                self.up_proj.weight,
                self.down_proj.weight,
                chunk_size,
            )
        return torch.cat(
            [
                self._project(chunk)
                for chunk in hidden_states.split(chunk_size, dim=-2)
            ],
            dim=-2,
        )


class YunmoBlock(nn.Module):
    def __init__(
        self,
        config: YunmoNextConfig,
        engine: TrainingEngineConfig,
        layer_idx: int,
    ) -> None:
        super().__init__()
        self.layer_idx = layer_idx
        self.input_norm = nn.RMSNorm(
            config.hidden_size,
            eps=config.rms_norm_eps,
        )
        self.post_mixer_norm = nn.RMSNorm(
            config.hidden_size,
            eps=config.rms_norm_eps,
        )
        self.mixer = (
            GatedGroupedQueryAttention(config, layer_idx)
            if config.is_attention_layer(layer_idx)
            else _build_recurrent_mixer(config, engine, layer_idx)
        )
        self.mlp = SwiGLU(
            config.hidden_size,
            config.intermediate_size,
            engine,
        )
        self.depth_routers = SublayerDepthRouters(
            config.hidden_size,
            config.rms_norm_eps,
        )

    def mixer_output(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Return one mixer transformation before its residual write."""

        normalized = self.input_norm(hidden_states)
        if isinstance(self.mixer, GatedGroupedQueryAttention):
            return self.mixer(normalized, position_ids, cu_seqlens)
        return self.mixer(normalized, cu_seqlens=cu_seqlens)[0]

    def mlp_output(self, hidden_states: torch.Tensor) -> torch.Tensor:
        """Return one feed-forward transformation before its residual write."""

        return self.mlp(self.post_mixer_norm(hidden_states))

    def mixer_output_cached(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
        *,
        fla_cache: object,
        attention_cache: GQAAttentionCache | None,
    ) -> tuple[torch.Tensor, GQAAttentionCache | None]:
        normalized = self.input_norm(hidden_states)
        if isinstance(self.mixer, GatedGroupedQueryAttention):
            if attention_cache is not None and not isinstance(
                attention_cache,
                GQAAttentionCache,
            ):
                raise TypeError("wrong cache type for grouped-query attention")
            return self.mixer.forward_cached(
                normalized,
                position_ids,
                attention_cache,
            )
        mixed, _, _ = self.mixer(
            normalized,
            past_key_values=fla_cache,
            use_cache=True,
        )
        return mixed, attention_cache


@dataclass
class YunmoNextOutput:
    loss: torch.Tensor | None = None
    logits: torch.Tensor | None = None
    hidden_states: torch.Tensor | None = None


@dataclass
class YunmoNextCUDAState:
    """Bounded KDA/local-attention state plus sparse global KV."""

    fla_cache: object
    attention: dict[int, GQAAttentionCache]
    seen_tokens: int = 0


class YunmoNextForCausalLM(nn.Module):
    def __init__(
        self,
        config: YunmoNextConfig,
        training_engine: TrainingEngineConfig | None = None,
    ) -> None:
        super().__init__()
        self.config = config
        self.training_engine = (
            TrainingEngineConfig()
            if training_engine is None
            else training_engine
        )
        if not isinstance(self.training_engine, TrainingEngineConfig):
            raise TypeError("training_engine must be TrainingEngineConfig")
        self.activation_checkpointing_runtime_policy = (
            ACTIVATION_CHECKPOINT_POLICY_FULL
        )
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(
            [
                YunmoBlock(config, self.training_engine, layer_idx)
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
        self._enable_cell_echo_biases()
        self.final_depth_router = SlicedDepthAttentionRouter(
            config.hidden_size,
            config.rms_norm_eps,
        )
        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.reset_parameters()
        self._install_cell_phase_mlp_geometry()

    def set_activation_checkpointing_runtime_policy(
        self,
        policy: str,
    ) -> None:
        if policy not in ACTIVATION_CHECKPOINT_POLICIES:
            raise ValueError(
                f"unsupported activation-checkpointing runtime policy: {policy}"
            )
        self.activation_checkpointing_runtime_policy = policy

    def _activation_checkpoint(
        self,
        function,
        *args: torch.Tensor,
    ):
        return checkpoint(function, *args, use_reentrant=False)

    @torch.no_grad()
    def reset_parameters(self) -> None:
        nn.init.normal_(
            self.embed_tokens.weight,
            mean=0.0,
            std=self.config.initializer_range,
        )
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(
                    module.weight,
                    mean=0.0,
                    std=self.config.initializer_range,
                )

        residual_scale = self.config.residual_write_initialization_scale
        for block in self.layers:
            block.mlp.down_proj.weight.mul_(residual_scale)
            output_projection = getattr(block.mixer, "o_proj", None)
            if output_projection is not None:
                output_projection.weight.mul_(residual_scale)

        self._reset_attention_residual_queries()

    @torch.no_grad()
    def _install_cell_phase_mlp_geometry(self) -> None:
        """Install layer-local widths without perturbing non-MLP RNG mapping.

        Every layer first constructs and initializes the common baseline width.
        Differently shaped MLPs are then installed from private layer-local
        random streams without moving the rest of the model's seed mapping.
        """

        seed = int(torch.initial_seed())
        modulus = 2**63 - 1
        for layer_idx, (layer, width) in enumerate(
            zip(
                self.layers,
                self.config.mlp_intermediate_sizes,
                strict=True,
            )
        ):
            mlp = SwiGLU(
                self.config.hidden_size,
                width,
                self.training_engine,
            )
            if mlp.gate_proj.weight.device.type == "meta":
                layer.mlp = mlp
                continue
            generator = torch.Generator(device=mlp.gate_proj.weight.device)
            generator.manual_seed(
                (seed * 1_000_000 + 3 + layer_idx) % modulus
            )
            for projection in (
                mlp.gate_proj,
                mlp.up_proj,
                mlp.down_proj,
            ):
                nn.init.normal_(
                    projection.weight,
                    mean=0.0,
                    std=self.config.initializer_range,
                    generator=generator,
                )
            mlp.down_proj.weight.mul_(
                self.config.residual_write_initialization_scale
            )
            layer.mlp = mlp

    @torch.no_grad()
    def _reset_attention_residual_queries(self) -> None:
        if (
            self.config.num_hidden_layers
            != ATTENTION_RESIDUAL_ORTHOGONAL_QUERY_DEPTH
        ):
            return
        routers = self._sliced_router_bank()
        expected = 2 * ATTENTION_RESIDUAL_ORTHOGONAL_QUERY_DEPTH + 1
        if len(routers) != expected or any(
            router.key_dim != 64 for router in routers
        ):
            raise RuntimeError(
                "orthogonal Attention Residual initialization requires "
                "49 rank-64 queries"
            )
        bank = orthogonal_attention_residual_query_bank(
            len(routers),
            routers[0].key_dim,
            device=routers[0].query.device,
            dtype=routers[0].query.dtype,
        )
        routers[0].query.zero_()
        for router, row in zip(routers[1:], bank[1:], strict=True):
            router.query.copy_(row)

    def parameter_count(self) -> int:
        return sum(parameter.numel() for parameter in self.parameters())

    def _apply_depth_router(
        self,
        router: SlicedDepthAttentionRouter,
        sources: list[torch.Tensor],
        routing_trace: list[torch.Tensor] | None,
    ) -> torch.Tensor:
        if (
            self.training_engine.activation_checkpointing
            and self.training
            and routing_trace is None
        ):
            def route_checkpointed(*inputs: torch.Tensor) -> torch.Tensor:
                routed = router(list(inputs))
                assert isinstance(routed, torch.Tensor)
                return routed

            return self._activation_checkpoint(route_checkpointed, *sources)
        if routing_trace is None:
            output = router(sources)
            assert isinstance(output, torch.Tensor)
            return output
        output, weights = router(sources, return_weights=True)
        routing_trace.append(weights.detach())
        return output

    def _mixer_output(
        self,
        layer: YunmoBlock,
        hidden_states: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor | None,
    ) -> torch.Tensor:
        if self.training_engine.activation_checkpointing and self.training:
            return self._activation_checkpoint(
                layer.mixer_output,
                hidden_states,
                position_ids,
                cu_seqlens,
            )
        return layer.mixer_output(hidden_states, position_ids, cu_seqlens)

    def _mlp_output(
        self,
        layer: YunmoBlock,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        if self.training_engine.activation_checkpointing and self.training:
            return self._activation_checkpoint(layer.mlp_output, hidden_states)
        return layer.mlp_output(hidden_states)

    def _sliced_router_bank(self) -> list[SlicedDepthAttentionRouter]:
        routers: list[SlicedDepthAttentionRouter] = []
        for layer in self.layers:
            routers.extend(
                (
                    layer.depth_routers.pre_mixer,
                    layer.depth_routers.pre_mlp,
                )
            )
        routers.append(self.final_depth_router)
        return routers

    def _enable_cell_echo_biases(self) -> None:
        """Address each attention-first cell's completed GQA source."""

        block_size = self.config.attn_res_sublayers_per_block
        period = self.config.recurrent_period
        for cell_start in range(0, len(self.layers), period):
            target_layer = cell_start + period - 1
            router = self.layers[target_layer].depth_routers.pre_mixer
            gqa_sublayer = 2 * cell_start
            preferred_source = 1 + gqa_sublayer // block_size
            target_read = 2 * target_layer
            if preferred_source * block_size > target_read:
                raise RuntimeError(
                    "cell GQA source is not completed before Cell-Echo read"
                )
            router.enable_source_bias(preferred_source)

    def _forward_hidden_checkpointed_blocks(
        self,
        embedding: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor | None,
        routers: list[SlicedDepthAttentionRouter],
    ) -> torch.Tensor:
        """Checkpoint each complete two-phase Attention Residual block."""

        block_size = self.config.attn_res_sublayers_per_block
        total_sublayers = 2 * len(self.layers)
        if block_size % 2 != 0 or total_sublayers % block_size != 0:
            raise ValueError(
                "block-checkpointed Attention Residual requires complete "
                "even sublayer blocks"
            )

        layers_per_block = block_size // 2
        completed_values = [embedding]
        for layer_first in range(0, len(self.layers), layers_per_block):
            layer_end = layer_first + layers_per_block
            sublayer_first = 2 * layer_first
            block_layers = tuple(self.layers[layer_first:layer_end])
            block_routers = tuple(
                routers[sublayer_first : sublayer_first + block_size]
            )
            source_count = len(completed_values)

            def run_block(
                *inputs: torch.Tensor,
                block_layers: tuple[YunmoBlock, ...] = block_layers,
                block_routers: tuple[SlicedDepthAttentionRouter, ...] = block_routers,
                source_count: int = source_count,
            ) -> torch.Tensor:
                checkpoint_sources = list(inputs[:source_count])
                logits = torch.stack(
                    [
                        sliced_source_logits(
                            source,
                            block_routers,
                            source_index=source_index,
                        )
                        for source_index, source in enumerate(checkpoint_sources)
                    ],
                    dim=1,
                )
                outputs, lses, _ = sliced_phase1_from_logits(
                    checkpoint_sources,
                    logits,
                    routers=block_routers,
                )
                phase_outputs = outputs.unbind(0)
                phase_lses = lses.unbind(0)

                partial_block: torch.Tensor | None = None
                local_sublayer = 0
                for layer in block_layers:
                    mixer_input = phase_outputs[local_sublayer]
                    if partial_block is not None:
                        mixer_input, _ = sliced_phase2_merge(
                            block_routers[local_sublayer],
                            partial_block,
                            mixer_input,
                            phase_lses[local_sublayer],
                        )
                    mixer_delta = layer.mixer_output(
                        mixer_input,
                        position_ids,
                        cu_seqlens,
                    )
                    partial_block = sliced_block_accumulate(
                        partial_block,
                        mixer_delta,
                    )
                    local_sublayer += 1

                    mlp_input, _ = sliced_phase2_merge(
                        block_routers[local_sublayer],
                        partial_block,
                        phase_outputs[local_sublayer],
                        phase_lses[local_sublayer],
                    )
                    partial_block = sliced_block_accumulate(
                        partial_block,
                        layer.mlp_output(mlp_input),
                    )
                    local_sublayer += 1

                if partial_block is None or local_sublayer != len(block_routers):
                    raise RuntimeError(
                        "Attention Residual block did not close exactly"
                    )
                return partial_block

            completed_values.append(
                self._activation_checkpoint(run_block, *completed_values)
            )

        hidden_states = self._apply_depth_router(
            self.final_depth_router,
            completed_values,
            None,
        )
        return self.norm(hidden_states)

    def _forward_hidden_two_phase_attention_residual(
        self,
        embedding: torch.Tensor,
        position_ids: torch.Tensor,
        cu_seqlens: torch.Tensor | None,
        layer_trace: list[torch.Tensor] | None,
        routing_trace: list[torch.Tensor] | None,
    ) -> torch.Tensor:
        routers = self._sliced_router_bank()
        total_sublayers = 2 * len(self.layers)
        block_size = self.config.attn_res_sublayers_per_block
        if (
            self.training_engine.activation_checkpointing
            and self.training
            and layer_trace is None
            and routing_trace is None
            and block_size % 2 == 0
            and total_sublayers % block_size == 0
        ):
            return self._forward_hidden_checkpointed_blocks(
                embedding,
                position_ids,
                cu_seqlens,
                routers,
            )

        completed_values = [embedding]
        completed_logits = [
            sliced_source_logits(embedding, routers, source_index=0)
        ]
        completed_first_reads = [0]
        partial_block: torch.Tensor | None = None
        phase_first = -1
        phase_end = -1
        phase_outputs: tuple[torch.Tensor, ...] = ()
        phase_lses: tuple[torch.Tensor, ...] = ()
        phase_weights: tuple[torch.Tensor, ...] | None = None

        def invalidate_phase() -> None:
            nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights
            phase_first = -1
            phase_end = -1
            phase_outputs = ()
            phase_lses = ()
            phase_weights = None

        def append_completed(source: torch.Tensor, first_read: int) -> None:
            source_index = len(completed_values)
            completed_values.append(source)
            completed_logits.append(
                sliced_source_logits(
                    source,
                    routers[first_read:],
                    source_index=source_index,
                )
            )
            completed_first_reads.append(first_read)
            invalidate_phase()

        def ensure_phase(read_idx: int) -> None:
            nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights
            if phase_first <= read_idx < phase_end:
                return
            phase_first = read_idx
            phase_end = (
                min(
                    ((read_idx // block_size) + 1) * block_size,
                    total_sublayers,
                )
                if read_idx < total_sublayers
                else total_sublayers + 1
            )
            logits = torch.stack(
                [
                    source_logits[
                        phase_first - first_read : phase_end - first_read
                    ]
                    for source_logits, first_read in zip(
                        completed_logits,
                        completed_first_reads,
                        strict=True,
                    )
                ],
                dim=1,
            )
            outputs, lses, weights = sliced_phase1_from_logits(
                completed_values,
                logits,
                routers=routers[phase_first:phase_end],
                return_weights=routing_trace is not None,
            )
            phase_outputs = outputs.unbind(0)
            phase_lses = lses.unbind(0)
            phase_weights = weights.unbind(0) if weights is not None else None

        def read_residual(read_idx: int) -> torch.Tensor:
            ensure_phase(read_idx)
            phase_idx = read_idx - phase_first
            output = phase_outputs[phase_idx]
            weights = (
                phase_weights[phase_idx] if phase_weights is not None else None
            )
            if partial_block is not None:
                output, partial_probability = sliced_phase2_merge(
                    routers[read_idx],
                    partial_block,
                    output,
                    phase_lses[phase_idx],
                )
                if weights is not None:
                    completed_mass = 1.0 - partial_probability
                    weights = torch.cat(
                        [
                            weights * completed_mass.unsqueeze(0),
                            partial_probability.unsqueeze(0),
                        ],
                        dim=0,
                    )
            if routing_trace is not None:
                assert weights is not None
                routing_trace.append(weights.detach())
            return output

        sublayer_idx = 0
        for layer in self.layers:
            mixer_input = read_residual(sublayer_idx)
            mixer_delta = self._mixer_output(
                layer,
                mixer_input,
                position_ids,
                cu_seqlens,
            )
            partial_block = sliced_block_accumulate(
                partial_block,
                mixer_delta,
            )
            sublayer_idx += 1
            if sublayer_idx % block_size == 0:
                append_completed(partial_block, sublayer_idx)
                partial_block = None

            mlp_input = read_residual(sublayer_idx)
            mlp_delta = self._mlp_output(layer, mlp_input)
            partial_block = sliced_block_accumulate(partial_block, mlp_delta)
            sublayer_idx += 1
            if layer_trace is not None:
                layer_trace.append(partial_block)
            if sublayer_idx % block_size == 0:
                append_completed(partial_block, sublayer_idx)
                partial_block = None

        if partial_block is not None:
            append_completed(partial_block, total_sublayers)
        hidden_states = read_residual(total_sublayers)
        return self.norm(hidden_states)

    def forward_hidden(
        self,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        *,
        layer_trace: list[torch.Tensor] | None = None,
        routing_trace: list[torch.Tensor] | None = None,
    ) -> torch.Tensor:
        return self._forward_hidden_from_embeddings(
            self.embed_tokens(input_ids),
            position_ids,
            cu_seqlens,
            layer_trace=layer_trace,
            routing_trace=routing_trace,
        )

    def _forward_hidden_from_embeddings(
        self,
        embeddings: torch.Tensor,
        position_ids: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        *,
        layer_trace: list[torch.Tensor] | None = None,
        routing_trace: list[torch.Tensor] | None = None,
    ) -> torch.Tensor:
        """Run the adopted text stack from already-constructed embeddings.

        This is the only non-token entry seam.  It adds no module or parameter
        and leaves creation or transformation of embeddings to the caller,
        keeping the ordinary token path independent of those concerns.
        """

        if embeddings.ndim != 3:
            raise ValueError(
                "embeddings must have shape [batch, sequence, hidden]"
            )
        batch, sequence, hidden_size = embeddings.shape
        if hidden_size != self.config.hidden_size:
            raise ValueError("embedding hidden size disagrees with model config")
        if position_ids is None:
            position_ids = (
                document_position_ids(cu_seqlens, sequence)
                if cu_seqlens is not None
                else torch.arange(
                    sequence,
                    device=embeddings.device,
                ).expand(batch, -1)
            )
        elif position_ids.shape != (batch, sequence):
            raise ValueError(
                "position IDs must match embedding batch and sequence"
            )
        return self._forward_hidden_two_phase_attention_residual(
            embeddings,
            position_ids,
            cu_seqlens,
            layer_trace,
            routing_trace,
        )

    def _forward_hidden_two_phase_attention_residual_cached(
        self,
        embedding: torch.Tensor,
        position_ids: torch.Tensor,
        cache: YunmoNextCUDAState,
        layer_trace: list[torch.Tensor] | None,
        routing_trace: list[torch.Tensor] | None,
    ) -> torch.Tensor:
        routers = self._sliced_router_bank()
        total_sublayers = 2 * len(self.layers)
        block_size = self.config.attn_res_sublayers_per_block
        completed_values = [embedding]
        completed_logits = [
            sliced_source_logits(embedding, routers, source_index=0)
        ]
        completed_first_reads = [0]
        partial_block: torch.Tensor | None = None
        phase_first = -1
        phase_end = -1
        phase_outputs: tuple[torch.Tensor, ...] = ()
        phase_lses: tuple[torch.Tensor, ...] = ()
        phase_weights: tuple[torch.Tensor, ...] | None = None

        def invalidate_phase() -> None:
            nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights
            phase_first = -1
            phase_end = -1
            phase_outputs = ()
            phase_lses = ()
            phase_weights = None

        def append_completed(source: torch.Tensor, first_read: int) -> None:
            source_index = len(completed_values)
            completed_values.append(source)
            completed_logits.append(
                sliced_source_logits(
                    source,
                    routers[first_read:],
                    source_index=source_index,
                )
            )
            completed_first_reads.append(first_read)
            invalidate_phase()

        def ensure_phase(read_idx: int) -> None:
            nonlocal phase_first, phase_end, phase_outputs, phase_lses, phase_weights
            if phase_first <= read_idx < phase_end:
                return
            phase_first = read_idx
            phase_end = (
                min(
                    ((read_idx // block_size) + 1) * block_size,
                    total_sublayers,
                )
                if read_idx < total_sublayers
                else total_sublayers + 1
            )
            logits = torch.stack(
                [
                    source_logits[
                        phase_first - first_read : phase_end - first_read
                    ]
                    for source_logits, first_read in zip(
                        completed_logits,
                        completed_first_reads,
                        strict=True,
                    )
                ],
                dim=1,
            )
            outputs, lses, weights = sliced_phase1_from_logits(
                completed_values,
                logits,
                routers=routers[phase_first:phase_end],
                return_weights=routing_trace is not None,
            )
            phase_outputs = outputs.unbind(0)
            phase_lses = lses.unbind(0)
            phase_weights = weights.unbind(0) if weights is not None else None

        def read_residual(read_idx: int) -> torch.Tensor:
            ensure_phase(read_idx)
            phase_idx = read_idx - phase_first
            output = phase_outputs[phase_idx]
            weights = (
                phase_weights[phase_idx] if phase_weights is not None else None
            )
            if partial_block is not None:
                output, partial_probability = sliced_phase2_merge(
                    routers[read_idx],
                    partial_block,
                    output,
                    phase_lses[phase_idx],
                )
                if weights is not None:
                    weights = torch.cat(
                        [
                            weights
                            * (1.0 - partial_probability).unsqueeze(0),
                            partial_probability.unsqueeze(0),
                        ],
                        dim=0,
                    )
            if routing_trace is not None:
                assert weights is not None
                routing_trace.append(weights.detach())
            return output

        sublayer_idx = 0
        for layer_idx, layer in enumerate(self.layers):
            mixer_input = read_residual(sublayer_idx)
            mixer_delta, layer_attention_cache = layer.mixer_output_cached(
                mixer_input,
                position_ids,
                fla_cache=cache.fla_cache,
                attention_cache=cache.attention.get(layer_idx),
            )
            if layer_attention_cache is not None:
                cache.attention[layer_idx] = layer_attention_cache
            partial_block = sliced_block_accumulate(
                partial_block,
                mixer_delta,
            )
            sublayer_idx += 1
            if sublayer_idx % block_size == 0:
                append_completed(partial_block, sublayer_idx)
                partial_block = None

            mlp_input = read_residual(sublayer_idx)
            mlp_delta = layer.mlp_output(mlp_input)
            partial_block = sliced_block_accumulate(partial_block, mlp_delta)
            sublayer_idx += 1
            if layer_trace is not None:
                layer_trace.append(partial_block)
            if sublayer_idx % block_size == 0:
                append_completed(partial_block, sublayer_idx)
                partial_block = None

        if partial_block is not None:
            append_completed(partial_block, total_sublayers)
        return self.norm(read_residual(total_sublayers))

    def forward_hidden_cached(
        self,
        input_ids: torch.Tensor,
        cache: YunmoNextCUDAState | None = None,
        position_ids: torch.Tensor | None = None,
        *,
        layer_trace: list[torch.Tensor] | None = None,
        routing_trace: list[torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, YunmoNextCUDAState]:
        """Exact single-document chunked prefill for the adopted model."""

        if self.training:
            raise RuntimeError("cached inference requires model.eval()")
        return self._forward_hidden_cached_from_embeddings(
            self.embed_tokens(input_ids),
            cache,
            position_ids,
            layer_trace=layer_trace,
            routing_trace=routing_trace,
        )

    def _forward_hidden_cached_from_embeddings(
        self,
        embeddings: torch.Tensor,
        cache: YunmoNextCUDAState | None = None,
        position_ids: torch.Tensor | None = None,
        *,
        layer_trace: list[torch.Tensor] | None = None,
        routing_trace: list[torch.Tensor] | None = None,
    ) -> tuple[torch.Tensor, YunmoNextCUDAState]:
        if self.training:
            raise RuntimeError("cached inference requires model.eval()")
        if embeddings.ndim != 3:
            raise ValueError(
                "embeddings must have shape [batch, sequence, hidden]"
            )
        batch, sequence, hidden_size = embeddings.shape
        if hidden_size != self.config.hidden_size:
            raise ValueError("embedding hidden size disagrees with model config")
        if cache is None:
            try:
                from fla.models.utils import Cache
            except ImportError as exc:
                raise ImportError(
                    "cached recurrent inference requires FLA"
                ) from exc
            cache = YunmoNextCUDAState(
                fla_cache=Cache(),
                attention={},
                seen_tokens=0,
            )
        if position_ids is None:
            position_ids = torch.arange(
                cache.seen_tokens,
                cache.seen_tokens + sequence,
                device=embeddings.device,
            ).expand(batch, -1)
        elif position_ids.shape != (batch, sequence):
            raise ValueError("cached position IDs must match embeddings")

        hidden_states = self._forward_hidden_two_phase_attention_residual_cached(
            embeddings,
            position_ids,
            cache,
            layer_trace,
            routing_trace,
        )
        cache.seen_tokens += sequence
        return hidden_states, cache

    def forward_cached(
        self,
        input_ids: torch.Tensor,
        cache: YunmoNextCUDAState | None = None,
        position_ids: torch.Tensor | None = None,
        *,
        materialize_logits: bool = True,
    ) -> tuple[YunmoNextOutput, YunmoNextCUDAState]:
        hidden_states, cache = self.forward_hidden_cached(
            input_ids,
            cache,
            position_ids,
        )
        logits = (
            F.linear(hidden_states, self.embed_tokens.weight)
            if materialize_logits
            else None
        )
        return YunmoNextOutput(
            logits=logits,
            hidden_states=hidden_states,
        ), cache

    def _forward_cached_from_embeddings(
        self,
        embeddings: torch.Tensor,
        cache: YunmoNextCUDAState | None = None,
        position_ids: torch.Tensor | None = None,
        *,
        materialize_logits: bool = True,
    ) -> tuple[YunmoNextOutput, YunmoNextCUDAState]:
        hidden_states, cache = self._forward_hidden_cached_from_embeddings(
            embeddings,
            cache,
            position_ids,
        )
        logits = (
            F.linear(hidden_states, self.embed_tokens.weight)
            if materialize_logits
            else None
        )
        return YunmoNextOutput(
            logits=logits,
            hidden_states=hidden_states,
        ), cache

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        materialize_logits: bool = False,
    ) -> YunmoNextOutput:
        return self._forward_from_embeddings(
            self.embed_tokens(input_ids),
            labels=labels,
            position_ids=position_ids,
            cu_seqlens=cu_seqlens,
            materialize_logits=materialize_logits,
        )

    def _forward_from_embeddings(
        self,
        embeddings: torch.Tensor,
        labels: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        cu_seqlens: torch.Tensor | None = None,
        materialize_logits: bool = False,
    ) -> YunmoNextOutput:
        hidden_states = self._forward_hidden_from_embeddings(
            embeddings,
            position_ids,
            cu_seqlens,
        )

        loss = None
        logits = None
        if (
            labels is not None
            and self.training_engine.cross_entropy_backend == "liger"
        ):
            try:
                from liger_kernel.transformers import (
                    LigerFusedLinearCrossEntropyLoss,
                )
            except ImportError as exc:
                raise ImportError(
                    "the configured cross-entropy backend requires liger-kernel"
                ) from exc
            loss_fn = LigerFusedLinearCrossEntropyLoss(
                accum_dtype=torch.float32
            )
            loss = loss_fn(
                self.embed_tokens.weight,
                hidden_states.reshape(-1, self.config.hidden_size),
                labels.reshape(-1),
            )
        elif labels is not None or materialize_logits:
            logits = F.linear(hidden_states, self.embed_tokens.weight)
            if labels is not None:
                loss = F.cross_entropy(
                    logits.float().reshape(-1, self.config.vocab_size),
                    labels.reshape(-1),
                )

        return YunmoNextOutput(
            loss=loss,
            logits=logits,
            hidden_states=hidden_states,
        )