File size: 3,139 Bytes
431ffa3 | 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 | # -*- coding: utf-8 -*-
"""Incremental decode caches for bit-exact autoregressive generation.
The Daffodils 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` -- the constant-size ``(conv_window, ssm_state)`` of one
NaiveSSM layer, so the selective scan can be stepped one token at a time.
: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-NaiveSSM constant-size recurrent state.
``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.
"""
conv_window: torch.Tensor | None = None
ssm_state: torch.Tensor | None = None
@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
|