File size: 3,677 Bytes
af0f593 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | # -*- coding: utf-8 -*-
"""Incremental decode caches for bit-exact autoregressive generation.
The DoL-OCR RDT has exactly one cross-position operation: causal MLA
attention. Every other sub-layer is per-position (mHC stream mixing, SwiGLU
FFN, RMSNorm) or a sequential recurrence (the NaiveSSM Mamba scan). Therefore
incremental decoding -- appending one token at a time while reusing frozen
earlier state -- is *numerically identical* to re-running the full forward on
the growing prefix, and these caches make that explicit.
Two primitive caches:
* :class:`MLACache` -- post-RoPE ``K``/``V`` of one attention call, grown by
``append``. Because each refinement pass of the Stage-2 loop is itself a
fresh causal attention over the frozen earlier positions, every
``(step, layer)`` pair owns its own :class:`MLACache`.
* :class:`MambaCache` -- either the constant-size NaiveSSM recurrent state or
one call site's isolated upstream ``InferenceParams`` object.
:class:`DecodeCache` is a flat keyed container the orchestrators
(``RDTForCausalLM`` / ``TwoStageCore``) populate lazily; keys encode the call
site (e.g. ``"prelude.0"``, ``"stage1.2"``, ``"stage2.s3.l0"``, ``"coda.1"``).
"""
from __future__ import annotations
from dataclasses import dataclass, field
import torch
@dataclass
class MLACache:
"""Per-attention-call key/value cache (post-RoPE, ``[B, H, T, head_dim]``)."""
k: torch.Tensor | None = None
v: torch.Tensor | None = None
@property
def length(self) -> int:
return 0 if self.k is None else self.k.shape[-2]
def append(
self, k: torch.Tensor, v: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""Append the new tokens' ``k``/``v`` and return the full cached tensors."""
if self.k is None:
self.k, self.v = k, v
else:
self.k = torch.cat([self.k, k], dim=-2)
self.v = torch.cat([self.v, v], dim=-2)
return self.k, self.v
@dataclass
class MambaCache:
"""Per-call-site recurrent state for exactly one Mamba backend.
``conv_window`` holds the last ``d_conv - 1`` *pre-convolution* columns
(shape ``[B, d_inner, d_conv - 1]``); ``ssm_state`` is the selective-scan
state ``[B, nheads, headdim, d_state]`` in float32. The ``official_*``
fields hold an isolated upstream cache; they must never be shared between
logical call sites because upstream indexes state by ``layer_idx``.
"""
conv_window: torch.Tensor | None = None
ssm_state: torch.Tensor | None = None
backend: str | None = None
official_inference_params: object | None = None
official_owner: object | None = field(default=None, repr=False)
official_batch_size: int | None = None
official_max_seqlen: int | None = None
official_device: torch.device | None = None
official_dtype: torch.dtype | None = None
official_poisoned: bool = False
@dataclass
class DecodeCache:
"""Flat keyed container of per-layer caches for one decode stream."""
mla: dict[str, MLACache] = field(default_factory=dict)
mamba: dict[str, MambaCache] = field(default_factory=dict)
seq_len: int = 0
def mla_cache(self, key: str) -> MLACache:
cache = self.mla.get(key)
if cache is None:
cache = self.mla[key] = MLACache()
return cache
def mamba_cache(self, key: str) -> MambaCache:
cache = self.mamba.get(key)
if cache is None:
cache = self.mamba[key] = MambaCache()
return cache
@property
def past_len(self) -> int:
"""Number of tokens already cached (0 before the prefill call)."""
return self.seq_len
|