File size: 4,312 Bytes
0eca2aa | 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 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | # Copyright (c) Together
# Apache 2.0 - Author: Michael Poli
# Adapted for the minimal Evo1 HF port.
"""Inference-time caches for Evo1 blocks.
Evo1 has two block types with different caching needs:
* `mha` blocks -> InferenceParams (standard KV cache)
* `hyena` blocks -> RecurrentInferenceParams (FIR window + IIR modal state)
Per-block dataclasses are wrapped in an HF ``Cache`` subclass (``Evo1Cache``)
so ``model.generate()`` can drive autoregressive decoding without the user
having to instantiate the two caches by hand, and so HF generation helpers
can introspect cache state (``get_seq_length``, ``get_max_cache_shape``).
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
from torch import Tensor
from transformers.cache_utils import Cache
@dataclass
class InferenceParams:
"""KV-cache parameters for the attention blocks (mha branch)."""
max_seqlen: int
max_batch_size: int
seqlen_offset: int = 0
batch_size_offset: int = 0
key_value_memory_dict: dict = field(default_factory=dict)
lengths_per_sample: Optional[Tensor] = None
def reset(self, max_seqlen, max_batch_size):
self.max_seqlen = max_seqlen
self.max_batch_size = max_batch_size
self.seqlen_offset = 0
if self.lengths_per_sample is not None:
self.lengths_per_sample.zero_()
@dataclass
class RecurrentInferenceParams:
"""SSM-cache parameters for the Hyena blocks (hyena branch)."""
fir_filter_length: int = 3
state_dim: int = 16
seqlen_offset: int = 0
fir_state_dict: dict = field(default_factory=dict)
state_dict: dict = field(default_factory=dict)
def reset(self):
self.fir_filter_length = 3
self.state_dim = 16
self.seqlen_offset = 0
class Evo1Cache(Cache):
"""HF-compatible wrapper around the per-block inference params.
Internally holds two dataclasses keyed by block type. Exposes
``seqlen_offset`` so HF generation helpers can read the current decoded
length, and implements ``get_seq_length()`` / ``get_max_cache_shape()``
per the transformers ``Cache`` interface.
The model internals (``StripedHyena.stateful_forward``) look up caches
via ``cache["mha"]`` and ``cache["hyena"]``; ``__getitem__`` is delegated
to attribute access so the original dict-keyed API keeps working.
"""
is_compileable = False
def __init__(
self,
max_seqlen: int,
max_batch_size: int,
short_filter_length: int = 3,
state_size: int = 8,
):
# transformers >= 4.55 Cache.__init__ requires either ``layers`` or
# ``layer_class_to_replicate``. We don't use HF's per-layer cache
# model (our two block-type-specific caches handle storage), so we
# pass an empty layers list.
super().__init__(layers=[])
self.mha = InferenceParams(
max_seqlen=max_seqlen,
max_batch_size=max_batch_size,
)
self.hyena = RecurrentInferenceParams(
fir_filter_length=short_filter_length,
state_dim=state_size,
)
# --- HF Cache interface ------------------------------------------------
@property
def seqlen_offset(self) -> int:
return self.mha.seqlen_offset
def get_seq_length(self, layer_idx: int = 0) -> int:
return self.mha.seqlen_offset
def get_max_cache_shape(self) -> int:
return self.mha.max_seqlen
def get_max_length(self) -> int:
# deprecated alias kept for older transformers versions
return self.mha.max_seqlen
# --- our convenience helpers ------------------------------------------
def advance(self, n: int = 1) -> None:
self.mha.seqlen_offset += n
self.hyena.seqlen_offset += n
def set_offset(self, offset: int) -> None:
self.mha.seqlen_offset = offset
self.hyena.seqlen_offset = offset
def reset(self) -> None:
self.mha.reset(self.mha.max_seqlen, self.mha.max_batch_size)
self.hyena.reset()
# --- dict-like access so existing call sites keep working --------------
def __getitem__(self, name: str):
return getattr(self, name)
def by_block_name(self, name: str):
return getattr(self, name)
|