uni_engine.py

from __future__ import annotations

"""UNI MAX: unified exact-prefill / quantized-decode recurrent inference.

The GDN24 cache topology is identical between the exact BF16 model and a model
whose **state-safe tail** (final MLP + LM head) is quantized. UNI MAX v1.1 exploits that causal invariant:

1. prefill the prompt with the exact BF16 model;
2. copy only the fixed recurrent cache state into a state-compatible persistent decode engine;
3. decode with the fastest CUDA-Graph candidate (typically FP8 on NVIDIA L4);
4. optionally verify FP8 draft blocks with exact BF16 chunk forwards to recover
   exact greedy-token semantics.

No KV sequence is copied: the bridge is O(1) in context length. Full-MLP quantization is deliberately excluded from this bridge because it changes the hidden trajectory that writes future recurrent states.
"""

from dataclasses import dataclass
import time
from typing import Any

import torch

from .engine import MaxTurboGraphDecoder, snapshot_cache, restore_cache_


@dataclass(frozen=True)
class BridgeReport:
    seconds: float
    mib: float
    seen_tokens: int


@dataclass(frozen=True)
class VerifyReport:
    generated_tokens: int
    drafted_tokens: int
    accepted_draft_tokens: int
    rejected_blocks: int
    verifier_blocks: int

    @property
    def acceptance_rate(self) -> float:
        if self.drafted_tokens <= 0:
            return 1.0
        return self.accepted_draft_tokens / self.drafted_tokens


def cache_payload_bytes(cache) -> int:
    total = 0
    for layer in cache.layers:
        for table_name in ("conv_states", "recurrent_states"):
            table = getattr(layer, table_name, {})
            for tensor in table.values():
                if tensor is not None:
                    total += tensor.numel() * tensor.element_size()
    return int(total)


@torch.inference_mode()
def copy_cache_(dst, src) -> None:
    """Copy recurrent state without reallocating destination tensors.

    Destination storage must already be materialized (CUDA-Graph capture does
    this).  Tensor addresses are preserved, so captured graphs remain valid.
    """
    if len(dst.layers) != len(src.layers):
        raise ValueError("cache topology mismatch")
    for d_layer, s_layer in zip(dst.layers, src.layers):
        for table_name in ("conv_states", "recurrent_states"):
            d_table = getattr(d_layer, table_name)
            s_table = getattr(s_layer, table_name)
            for idx, s_tensor in s_table.items():
                if s_tensor is None:
                    continue
                d_tensor = d_table.get(idx)
                if d_tensor is None:
                    raise RuntimeError(
                        f"destination cache storage is not materialized: {table_name}[{idx}]"
                    )
                if d_tensor.shape != s_tensor.shape or d_tensor.dtype != s_tensor.dtype:
                    raise RuntimeError(
                        f"cache state mismatch for {table_name}[{idx}]: "
                        f"dst={tuple(d_tensor.shape)}/{d_tensor.dtype}, "
                        f"src={tuple(s_tensor.shape)}/{s_tensor.dtype}"
                    )
                d_tensor.copy_(s_tensor)
        d_layer.has_previous_state.clear()
        d_layer.has_previous_state.update(dict(s_layer.has_previous_state))
        d_layer.is_conv_states_initialized.clear()
        d_layer.is_conv_states_initialized.update(dict(s_layer.is_conv_states_initialized))
        d_layer.is_recurrent_states_initialized.clear()
        d_layer.is_recurrent_states_initialized.update(dict(s_layer.is_recurrent_states_initialized))
    dst.seen_tokens = int(src.seen_tokens)


class UniMaxEngine:
    """Phase-specialized recurrent inference engine.

    ``exact_model`` always handles prefill.  ``decode_model`` can be the same model (UNI-SAFE) or a state-safe quantized clone (UNI-SPEED / UNI-EXACT).
    """

    def __init__(
        self,
        exact_model,
        decode_model,
        decode_graph: MaxTurboGraphDecoder,
    ):
        self.exact_model = exact_model.eval()
        self.decode_model = decode_model.eval()
        self.decode_graph = decode_graph
        self.device = self.exact_model.get_input_embeddings().weight.device
        if self.device != self.decode_graph.device:
            raise ValueError("exact and decode engines must live on the same CUDA device")
        self.exact_cache = self.exact_model.make_recurrent_cache()
        self.last_bridge: BridgeReport | None = None

    @torch.inference_mode()
    def reset(self) -> None:
        self.exact_cache.reset()
        self.decode_graph.reset()
        self.last_bridge = None

    @torch.inference_mode()
    def _bridge(self) -> BridgeReport:
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        t0 = time.perf_counter()
        copy_cache_(self.decode_graph.cache, self.exact_cache)
        if self.device.type == "cuda":
            torch.cuda.synchronize(self.device)
        dt = time.perf_counter() - t0
        report = BridgeReport(
            seconds=float(dt),
            mib=cache_payload_bytes(self.exact_cache) / 2**20,
            seen_tokens=int(self.exact_cache.seen_tokens),
        )
        self.last_bridge = report
        return report

    @torch.inference_mode()
    def prefill_exact(self, input_ids: torch.Tensor) -> tuple[torch.Tensor, Any, BridgeReport]:
        """Exact BF16 prefill, then O(1)-context state-compatible recurrent bridge."""
        self.exact_cache.reset()
        out = self.exact_model(
            input_ids=input_ids,
            past_key_values=self.exact_cache,
            use_cache=True,
            logits_to_keep=1,
        )
        first = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
        bridge = self._bridge()
        self.decode_graph.static_token.copy_(first)
        return first, out, bridge

    @torch.inference_mode()
    def decode_fast(self, first_token: torch.Tensor, recurrent_forwards: int) -> torch.Tensor:
        return self.decode_graph.decode_forwards(first_token, int(recurrent_forwards))

    @torch.inference_mode()
    def generate_fast(self, input_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
        """UNI-SPEED generation: exact first token, quantized graph thereafter."""
        n = int(max_new_tokens)
        if n <= 0:
            return torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
        first, _, _ = self.prefill_exact(input_ids)
        if n == 1:
            return first
        tail = self.decode_graph.decode_tokens(first, n - 1)
        return torch.cat([first, tail], dim=1)

    @torch.inference_mode()
    def decode_verified(
        self,
        first_token: torch.Tensor,
        recurrent_forwards: int,
        *,
        draft_block: int | None = None,
    ) -> tuple[torch.Tensor, VerifyReport]:
        """Verify ``recurrent_forwards`` tokens after ``first_token``.

        ``prefill_exact`` must have been called immediately before this method so
        both exact and decode caches represent the same prompt state.
        """
        remaining = int(recurrent_forwards)
        if remaining <= 0:
            empty = torch.empty((first_token.shape[0], 0), dtype=torch.long, device=first_token.device)
            return empty, VerifyReport(0, 0, 0, 0, 0)
        if first_token.shape[0] != 1:
            raise ValueError("UNI-EXACT currently supports batch size 1")

        current = first_token
        generated: list[torch.Tensor] = []
        k_default = max(1, int(draft_block or self.decode_graph.block_size))
        drafted = accepted_drafts = rejected_blocks = verifier_blocks = 0

        while remaining > 0:
            k = min(k_default, remaining)
            exact_before = snapshot_cache(self.exact_cache)
            draft = self.decode_graph.decode_tokens(current, k)
            drafted += k

            verify_input = current if k == 1 else torch.cat([current, draft[:, :-1]], dim=1)
            verify_out = self.exact_model(
                input_ids=verify_input,
                past_key_values=self.exact_cache,
                use_cache=True,
                logits_to_keep=k,
            )
            exact_pred = torch.argmax(verify_out.logits[:, -k:, :], dim=-1)
            verifier_blocks += 1
            mismatch_positions = (~exact_pred.eq(draft)[0]).nonzero(as_tuple=False)

            if mismatch_positions.numel() == 0:
                generated.append(draft.detach().clone())
                accepted_drafts += k
                current = draft[:, -1:]
                remaining -= k
                continue

            m = int(mismatch_positions[0, 0].item())
            accepted_drafts += m
            rejected_blocks += 1
            restore_cache_(self.exact_cache, exact_before)
            replay_input = current if m == 0 else torch.cat([current, draft[:, :m]], dim=1)
            replay = self.exact_model(
                input_ids=replay_input,
                past_key_values=self.exact_cache,
                use_cache=True,
                logits_to_keep=1,
            )
            corrected = torch.argmax(replay.logits[:, -1, :], dim=-1, keepdim=True)
            if m:
                generated.append(draft[:, :m].detach().clone())
            generated.append(corrected.detach().clone())
            current = corrected
            emitted = m + 1
            remaining -= emitted
            copy_cache_(self.decode_graph.cache, self.exact_cache)
            self.decode_graph.static_token.copy_(current)

        out = torch.cat(generated, dim=1)
        return out, VerifyReport(
            generated_tokens=int(out.shape[1]),
            drafted_tokens=int(drafted),
            accepted_draft_tokens=int(accepted_drafts),
            rejected_blocks=int(rejected_blocks),
            verifier_blocks=int(verifier_blocks),
        )

    @torch.inference_mode()
    def generate_verified(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int,
        *,
        draft_block: int | None = None,
    ) -> tuple[torch.Tensor, VerifyReport]:
        """UNI-EXACT generation with exact greedy-token semantics."""
        n = int(max_new_tokens)
        if n <= 0:
            empty = torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
            return empty, VerifyReport(0, 0, 0, 0, 0)
        first, _, _ = self.prefill_exact(input_ids)
        if n == 1:
            return first, VerifyReport(1, 0, 0, 0, 0)
        tail, report = self.decode_verified(first, n - 1, draft_block=draft_block)
        full = torch.cat([first, tail], dim=1)
        return full, VerifyReport(
            generated_tokens=int(full.shape[1]),
            drafted_tokens=report.drafted_tokens,
            accepted_draft_tokens=report.accepted_draft_tokens,
            rejected_blocks=report.rejected_blocks,
            verifier_blocks=report.verifier_blocks,
        )



@torch.inference_mode()
def exact_greedy_tokens(model, input_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
    n = int(max_new_tokens)
    if n <= 0:
        return torch.empty((input_ids.shape[0], 0), dtype=torch.long, device=input_ids.device)
    cache = model.make_recurrent_cache()
    out = model(input_ids=input_ids, past_key_values=cache, use_cache=True, logits_to_keep=1)
    token = torch.argmax(out.logits[:, -1, :], dim=-1, keepdim=True)
    pieces = [token]
    for _ in range(n - 1):
        token = model.greedy_step(token, cache)
        pieces.append(token)
    return torch.cat(pieces, dim=1)