Initial release

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,234 @@
+---
+language:
+- dna
+library_name: transformers
+tags:
+- DNA
+- language-model
+- StripedHyena
+- Evo
+- Evo1.5
+license: apache-2.0
+---
+
+# Evo1-1.5-7B-8K
+
+A clean, minimal HuggingFace port of [Evo 1.5 (8k)](https://huggingface.co/evo-design/evo-1.5-8k-base), which extends Evo 1 (8k) by approximately 50% additional pretraining tokens. ~7B parameters. Native support for layer-by-layer hidden state extraction, attention-weight extraction, and a runtime-switchable attention backend.
+
+## Why this port?
+
+[evo-design/evo-1.5-8k-base](https://huggingface.co/evo-design/evo-1.5-8k-base) ships a `trust_remote_code` HF implementation but it has four gaps that force every downstream user to monkey-patch the model:
+
+1. `output_hidden_states=True` is hardcoded to `None` (intermediate embeddings require forward hooks).
+2. `output_attentions=True` is unsupported (flash-attn discards the `(B, H, T, T)` matrix; users must patch the attention module).
+3. `attn_implementation` cannot be switched at load time - flash_attn is mandatory at every attention layer.
+4. The bare backbone is not exposed via `AutoModel.from_pretrained`; only the LM-head wrapper exists.
+
+This repo fixes all four. The math is **bit-exact** with the evo-design reference (`max_abs_diff = 0.000e+00` at every layer; see [Parity Verification](#parity-verification)).
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Total parameters | ~7B |
+| Layers | 32 |
+| Attention heads | 32 |
+| Embedding dimension | 4096 |
+| Inner MLP size | 10928 |
+| Vocabulary size | 512 (UTF-8 byte-level) |
+| Attention layer indices | [8, 16, 24] |
+| Hyena layer indices | all others |
+| Hyena state size | 8 |
+| Positional encoding | RoPE (base = 10000) |
+| Architecture | StripedHyena (alternating Hyena / MHA blocks) |
+| Max sequence length | 8 192 |
+| Training dtype | bfloat16 (Hyena modal-form `poles` / `residues` kept in fp32) |
+
+Architecturally identical to `Evo1-1-7B-8K`; only the trained weights differ (Evo 1.5 = Evo 1 (8k) + ~50% more pretraining tokens).
+
+## Pretraining
+
+- **Objective:** causal byte-level next-token prediction.
+- **Data:** [OpenGenome](https://huggingface.co/datasets/LongSafari/open-genome), the prokaryotic whole-genome dataset assembled for Evo 1 (~80 000 prokaryotic genomes plus ~2 million phage and plasmid sequences), with training extended by an additional 50% - totaling **~450 billion tokens** (vs. ~300 B for Evo 1).
+- **Initialization:** continued pretraining from `evo-1-8k-base` at 8 192-token context (i.e. Evo 1.5 is not trained from scratch; the additional ~150 B tokens are appended to Evo 1's 300 B-token training run).
+- **Source checkpoint:** `evo-design/evo-1.5-8k-base@main`.
+
+### Downstream use
+
+Evo 1.5 is the model used in [Merchant et al., 2025 (Nature)](https://doi.org/10.1038/s41586-025-09749-7) for *semantic mining* of functional de novo genes - generating new candidate genes conditioned on genomic context. It was used to produce [SynGenome](https://evodesign.org/syngenome/), the first large-scale AI-generated genomic database (over 120 B base pairs of synthetic DNA, organised by predicted function, protein domain, and host species).
+
+## Parity Verification
+
+Hidden-state representations verified **bit-exact** (`max_abs_diff = 0.000e+00`) to the evo-design reference at all 33 representation levels (token embedding + each of the 32 transformer blocks + final RMSNorm), using `attn_implementation="flash_attention_2"` in bf16 (matches the reference's backend choice and the trained dtype). Logits from `Evo1ForCausalLM` were also verified bit-exact (top-1 agreement: 128/128 positions). Verified on H100 with PyTorch 2.7.1 / CUDA 12.9.
+
+### Numerical equivalence across attention backends
+
+`flash_attention_2` is bit-exact with the original togethercomputer / evo-design implementations (same CUDA kernel). The `sdpa` and `eager` backends use different kernels (PyTorch's bundled flash kernel and pure-PyTorch matmul, respectively); these compute mathematically equivalent attention but accumulate floating-point operations in slightly different orders, producing per-block diffs at the bf16 noise floor (relative error roughly `1e-4` to `1e-2`).
+
+Unlike a standard transformer, where attention is softmax-bounded and per-block diffs stay small through the stack, StripedHyena's Hyena layers use an unbounded-gain IIR filter (no softmax) - so any small per-attention-block diff gets amplified by Hyena's filter gain. Across 32 layers this compounds to ~1% relative error in the intermediate residual stream, though the final post-RMSNorm output is bounded. Use `flash_attention_2` if you need to match the reference's activations bit-for-bit.
+
+## Related Models
+
+See the full [Evo1 collection](https://huggingface.co/collections/Taykhoom/evo1-6a24ae4a98f04906482db8c2) on the Hub.
+
+| Model | Context | Notes |
+|---|---|---|
+| [Taykhoom/Evo1-1-7B-8K](https://huggingface.co/Taykhoom/Evo1-1-7B-8K) | 8 192 | Original Evo 1 base model (8k context). |
+| [Taykhoom/Evo1-1-7B-131K](https://huggingface.co/Taykhoom/Evo1-1-7B-131K) | 131 072 | Long-context Evo 1 with linearly-scaled RoPE (131k context). |
+| **[Taykhoom/Evo1-1.5-7B-8K](https://huggingface.co/Taykhoom/Evo1-1.5-7B-8K)** | 8 192 | Evo 1.5: Evo 1 (8k) further trained on ~50% more pretraining tokens. |
+
+## Usage
+
+> **Note on dtype.** Evo1 was trained in **bfloat16**, with the Hyena `poles` / `residues` (modal-form filter parameters) kept in fp32 for numerical stability. **Passing `dtype=...` to `from_pretrained` only affects the initial load precision** (peak memory during loading) **and does not change the inference dtype** - `Evo1Model.__init__` and `Evo1ForCausalLM.__init__` unconditionally call `to_bfloat16_except_poles_residues()`, so the model always runs in bf16 with poles/residues in fp32. This is intentional: the trained activations are bf16-stable and fp16-unstable, and the modal-form filter requires fp32 for numerical stability - a single mixed config is the only valid one.
+
+> **Note on attention backend.** By default, `from_pretrained` selects `attn_implementation="sdpa"` (PyTorch's bundled scaled-dot-product-attention) - this works out of the box without `flash_attn` installed. The original togethercomputer / evo-design implementations use `flash_attn` unconditionally; **for bit-exact reproduction of reference outputs, explicitly pass `attn_implementation="flash_attention_2"`** (and `pip install flash-attn`). See [Numerical equivalence across attention backends](#numerical-equivalence-across-attention-backends) for the magnitude of the difference.
+
+### Embedding generation (no LM head)
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1.5-7B-8K", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "Taykhoom/Evo1-1.5-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",  # bit-exact with reference; or omit to default to "sdpa"
+).cuda().eval()
+
+seqs = ["ACGTACGTACGT", "GGGTTTAAACCC"]
+inputs = tokenizer(seqs, return_tensors="pt", padding=True).to(model.device)
+
+with torch.no_grad():
+    out = model(**inputs, output_hidden_states=True)
+
+last_hidden  = out.last_hidden_state  # (B, T, 4096)
+all_layers   = out.hidden_states      # tuple of (B, T, 4096), len = 34 (embed + 32 blocks + post-norm)
+layer_12_emb = all_layers[12]         # often used as a "middle" representation
+```
+
+### LM logits
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1.5-7B-8K", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "Taykhoom/Evo1-1.5-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",
+).cuda().eval()
+
+inputs = tokenizer(["ACGT"], return_tensors="pt").to(model.device)
+with torch.no_grad():
+    logits = model(**inputs).logits   # (1, T, 512)
+```
+
+### Generation
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1.5-7B-8K", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "Taykhoom/Evo1-1.5-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",
+).cuda().eval()
+
+inputs = tokenizer(["ACGT"], return_tensors="pt").to(model.device)
+out = model.generate(**inputs, max_new_tokens=128, do_sample=True, top_k=4, temperature=1.0)
+print(tokenizer.decode(out[0]))
+```
+
+`generation_config.json` ships with `eos_token_id = 0` (the EOD byte) and `pad_token_id = 1` so `model.generate()` stops naturally at the trained end-of-document token without needing extra kwargs. Note that the tokenizer itself does **not** add an EOS at encoding time - this matches the original Evo1 inference pipeline (only generation stops on EOS; embedding/scoring uses raw byte input).
+
+### Attention weights
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1.5-7B-8K", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "Taykhoom/Evo1-1.5-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="eager",  # required for output_attentions to populate
+).cuda().eval()
+
+inputs = tokenizer(["ACGTACGT"], return_tensors="pt").to(model.device)
+with torch.no_grad():
+    out = model(**inputs, output_attentions=True)
+
+# out.attentions is a tuple of length 32. Entries at indices not in [8, 16, 24]
+# are None (Hyena blocks have no attention matrix). Entries at [8, 16, 24] are
+# (B, num_heads, T, T) tensors.
+attn_block_8 = out.attentions[8]
+```
+
+### Multi-GPU loading (optional)
+
+Loading via `accelerate`'s `device_map` is supported (`_no_split_modules` is set so each `AttentionBlock` / `ParallelGatedConvBlock` stays atomic on one device, with hidden state automatically transferred across device boundaries):
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo1-1.5-7B-8K", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "Taykhoom/Evo1-1.5-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",
+    device_map="auto",     # auto-shard across all visible GPUs; falls back to single GPU if only one is present
+).eval()
+```
+
+Requires `pip install accelerate`.
+
+### Fine-tuning
+
+Standard HuggingFace conventions. For sequence-level tasks, take the final `last_hidden_state` (or any intermediate `hidden_states[i]`) and feed it into a downstream head.
+
+## Implementation Notes
+
+- **Custom attention module (`attention.py`).** Replaces `flash_attn.modules.mha.MHA` with a small in-repo `MHA` class that supports `attn_implementation="eager"` / `"sdpa"` / `"flash_attention_2"`. Parameter names (`Wqkv`, `out_proj`, `rotary_emb.inv_freq`) are preserved so existing checkpoints load unchanged. When `output_attentions=True`, the sdpa and flash paths automatically fall back to eager so the attention matrix is materialized.
+- **Custom rotary embedding (`rotary.py`).** When `flash_attn` is installed we delegate to its Triton kernel (faster on long sequences). The pure-PyTorch fallback does the rotary multiply in fp32 internally (then casts back) so it produces bit-exactly identical results to the Triton kernel - a bf16 multiply here introduces ~3e-2 error per layer that compounds to ~1% relative across 32 layers.
+- **Hyena engine (`engine.py`).** Copied verbatim from the togethercomputer reference (FFT-based long convolution, modal-form prefill).
+- **Cache subclass (`cache.py`).** `Evo1Cache(transformers.cache_utils.Cache)` wraps the two block-type-specific inference param dataclasses (`InferenceParams` for attention KV cache, `RecurrentInferenceParams` for Hyena FIR window + IIR modal state). Exposes `get_seq_length()` / `get_max_cache_shape()` so HF's `model.generate()` can introspect cache state; falls through to `cache["mha"]` / `cache["hyena"]` for the model internals.
+- **Tokenizer (`tokenization_evo1.py`).** Byte-level UTF-8 with vocab_size = 512. Pad token is byte `\x01`. No CLS, no EOS appended at encoding time (matches original Evo1 inference). The `_decode` method is numpy-2.x compatible (the original `np.uint8.clip(min=32, max=512)` was an overflow on numpy 2).
+- **Dependencies.** `torch`, `transformers`, `numpy`, `safetensors`, `huggingface_hub` (only for `from_pretrained` downloads). `flash_attn` is **only** required if you pass `attn_implementation="flash_attention_2"`.
+
+## Citation
+
+```bibtex
+@article{merchant2025_evo_1_5,
+  title   = {Semantic design of functional de novo genes from a genomic language model},
+  author  = {Merchant, Aditi T. and King, Samuel H. and Nguyen, Eric and Hie, Brian L.},
+  journal = {Nature},
+  year    = {2025},
+  doi     = {10.1038/s41586-025-09749-7}
+}
+
+@article{nguyen2024_evo,
+  title   = {Sequence modeling and design from molecular to genome scale with {Evo}},
+  author  = {Nguyen, Eric and Poli, Michael and Durrant, Matthew G. and Kang, Brian and Katrekar, Dhruva and Li, David B. and Bartie, Liam J. and Thomas, Armin W. and King, Samuel H. and Brixi, Garyk and Sullivan, Jeremy and Ng, Madelena Y. and Lewis, Ashley and Lou, Aaron and Ermon, Stefano and Baccus, Stephen A. and Hernandez-Boussard, Tina and {R{\'e}}, Christopher and Hsu, Patrick D. and Hie, Brian L.},
+  journal = {Science},
+  volume  = {386},
+  number  = {6723},
+  pages   = {eado9336},
+  year    = {2024},
+  doi     = {10.1126/science.ado9336}
+}
+```
+
+## Credits
+
+Original Evo / Evo 1.5 model and code by Nguyen et al. and Merchant et al. Source repo: [evo-design/evo](https://github.com/evo-design/evo). Source checkpoint: [evo-design/evo-1.5-8k-base](https://huggingface.co/evo-design/evo-1.5-8k-base).
+
+The HuggingFace conversion code in this repo was authored primarily by [Claude](https://claude.ai/code) and reviewed manually by Taykhoom Dalal.
+
+## License
+
+Apache 2.0, following the original Evo1 release.

attention.py

@@ -0,0 +1,294 @@
+# Copyright (c) Together / Apache 2.0.
+#
+# Minimal multi-head attention block for the Evo1 HF port.
+#
+# Replaces flash_attn.modules.mha.MHA with a small, dependency-light
+# implementation that:
+#   - keeps the same parameter names (Wqkv, out_proj, rotary_emb.inv_freq)
+#     so existing checkpoints load directly,
+#   - supports attn_implementation in {"eager", "sdpa", "flash_attention_2"},
+#   - returns attention weights when output_attentions=True (eager path),
+#   - falls back to eager when output_attentions=True for sdpa/flash backends
+#     (per the standard HuggingFace dispatch convention),
+#   - keeps a one-method KV cache compatible with the existing
+#     InferenceParams dataclass for autoregressive generation.
+#
+# Math is causal, single-stream (no cross-attention), no ALiBi, no sliding
+# window. Evo1 only ever exercised the qkv-packed self-attention path.
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .rotary import RotaryEmbedding
+
+
+def _flash_attn_required():
+    try:
+        from flash_attn import flash_attn_func, flash_attn_varlen_func  # noqa: F401
+        from flash_attn.bert_padding import pad_input, unpad_input  # noqa: F401
+    except ImportError as exc:  # pragma: no cover - optional dep
+        raise ImportError(
+            "attn_implementation='flash_attention_2' requires the flash-attn "
+            "package. Install with `pip install flash-attn --no-build-isolation`."
+        ) from exc
+
+
+def _update_kv_cache(kv: torch.Tensor, inference_params, layer_idx: int) -> torch.Tensor:
+    """Append `kv` to inference_params.key_value_memory_dict[layer_idx].
+
+    kv: (B, S, 2, H_kv, D) where S is the new-token chunk length (may be 1).
+    Returns the cumulative kv up to the current sequence position.
+    """
+    num_heads, head_dim = kv.shape[-2:]
+    if layer_idx not in inference_params.key_value_memory_dict:
+        kv_cache = torch.empty(
+            inference_params.max_batch_size,
+            inference_params.max_seqlen,
+            2,
+            num_heads,
+            head_dim,
+            dtype=kv.dtype,
+            device=kv.device,
+        )
+        inference_params.key_value_memory_dict[layer_idx] = kv_cache
+    else:
+        kv_cache = inference_params.key_value_memory_dict[layer_idx]
+    batch_start = inference_params.batch_size_offset
+    batch_end = batch_start + kv.shape[0]
+    sequence_start = inference_params.seqlen_offset
+    sequence_end = sequence_start + kv.shape[1]
+    kv_cache[batch_start:batch_end, sequence_start:sequence_end, ...] = kv
+    return kv_cache[batch_start:batch_end, :sequence_end, ...]
+
+
+class MHA(nn.Module):
+    """Multi-head self-attention with backend-dispatch.
+
+    Constructor signature is a strict subset of flash_attn.modules.mha.MHA so
+    that the existing AttentionBlock instantiation site is left untouched.
+    Unsupported kwargs (cross_attn, dwconv, alibi, window_size, ...) are
+    accepted and ignored or hard-asserted: Evo1 never exercises them.
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        num_heads_kv: int | None = None,
+        cross_attn: bool = False,
+        qkv_proj_bias: bool = True,
+        out_proj_bias: bool = True,
+        dropout: float = 0.0,
+        softmax_scale: float | None = None,
+        causal: bool = False,
+        layer_idx: int | None = None,
+        rotary_emb_dim: int = 0,
+        rotary_emb_base: float = 10000.0,
+        rotary_emb_scale_base: float | None = None,
+        rotary_emb_interleaved: bool = False,
+        use_flash_attn: bool = False,  # legacy kwarg, kept for ctor compatibility
+        attn_implementation: str = "eager",
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__()
+        if cross_attn:
+            raise NotImplementedError("Cross-attention is not supported in this minimal MHA.")
+
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.num_heads_kv = num_heads_kv if num_heads_kv is not None else num_heads
+        if self.embed_dim % num_heads != 0:
+            raise ValueError("embed_dim must be divisible by num_heads")
+        if self.num_heads % self.num_heads_kv != 0:
+            raise ValueError("num_heads must be divisible by num_heads_kv")
+        self.head_dim = self.embed_dim // num_heads
+        self.causal = causal
+        self.softmax_scale = softmax_scale
+        self.layer_idx = layer_idx
+        self.rotary_emb_dim = rotary_emb_dim
+        self.attn_implementation = attn_implementation
+        self.dropout_p = dropout
+
+        if self.rotary_emb_dim > 0:
+            self.rotary_emb = RotaryEmbedding(
+                self.rotary_emb_dim,
+                base=rotary_emb_base,
+                interleaved=rotary_emb_interleaved,
+                scale_base=rotary_emb_scale_base,
+                device=device,
+            )
+
+        qkv_dim = self.head_dim * (self.num_heads + 2 * self.num_heads_kv)
+        self.Wqkv = nn.Linear(embed_dim, qkv_dim, bias=qkv_proj_bias, **factory_kwargs)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=out_proj_bias, **factory_kwargs)
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
+        dtype = self.out_proj.weight.dtype if dtype is None else dtype
+        device = self.out_proj.weight.device
+        return torch.empty(
+            batch_size, max_seqlen, 2, self.num_heads_kv, self.head_dim,
+            dtype=dtype, device=device,
+        )
+
+    def _project_qkv(self, x: torch.Tensor) -> torch.Tensor:
+        """Compute Wqkv(x) and reshape to (B, T, 3, H, D) when MHA, or
+        return (q, kv) tuple-like layout when GQA. Returns the packed qkv
+        tensor in either case (kv heads broadcast for SDPA/flash later).
+
+        For Evo1 we have num_heads_kv == num_heads (proj_groups=1), so the
+        common-case packed layout is fine; we keep a GQA branch for future
+        flexibility but assert MHA at construction time.
+        """
+        qkv = self.Wqkv(x)
+        if self.num_heads_kv == self.num_heads:
+            return qkv.view(*qkv.shape[:-1], 3, self.num_heads, self.head_dim)
+        # GQA path (unused by Evo1):
+        q = qkv[..., : self.num_heads * self.head_dim]
+        kv = qkv[..., self.num_heads * self.head_dim:]
+        q = q.view(*q.shape[:-1], self.num_heads, self.head_dim)
+        kv = kv.view(*kv.shape[:-1], 2, self.num_heads_kv, self.head_dim)
+        return q, kv  # type: ignore[return-value]
+
+    # ------------------------------------------------------------------ eager
+    def _forward_eager(
+        self,
+        qkv: torch.Tensor,
+        output_attentions: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        # qkv: (B, T, 3, H, D). Match flash_attn / sdpa numerical behaviour by
+        # doing the attention math in fp32 internally (q*scale, QK^T matmul,
+        # softmax, attn @ V). Without this, the bf16 matmuls accumulate
+        # ~1e-2 absolute error per attention block and diverge meaningfully
+        # from flash_attn (which always accumulates in fp32 inside its CUDA
+        # kernel). Output is cast back to the original dtype for the residual
+        # add.
+        orig_dtype = qkv.dtype
+        q, k, v = qkv.unbind(dim=2)
+        q = q.permute(0, 2, 1, 3).float()  # (B, H, T, D), fp32
+        k = k.permute(0, 2, 1, 3).float()
+        v = v.permute(0, 2, 1, 3).float()
+        scale = self.softmax_scale if self.softmax_scale is not None else 1.0 / math.sqrt(self.head_dim)
+
+        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
+        if self.causal:
+            T = q.shape[-2]
+            mask = torch.triu(
+                torch.ones(T, T, device=scores.device, dtype=torch.bool), diagonal=1
+            )
+            scores = scores.masked_fill(mask, float("-inf"))
+        attn = F.softmax(scores, dim=-1)
+        if self.training and self.dropout_p > 0:
+            attn = F.dropout(attn, p=self.dropout_p)
+        out = torch.matmul(attn, v).permute(0, 2, 1, 3)  # (B, T, H, D), fp32
+        out = out.to(orig_dtype)
+        return out, (attn.to(orig_dtype) if output_attentions else None)
+
+    # -------------------------------------------------------------------- sdpa
+    def _forward_sdpa(self, qkv: torch.Tensor) -> torch.Tensor:
+        q, k, v = qkv.unbind(dim=2)
+        q = q.permute(0, 2, 1, 3)  # (B, H, T, D)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+        scale = self.softmax_scale if self.softmax_scale is not None else None
+        out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=None,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            is_causal=self.causal,
+            scale=scale,
+        )
+        return out.permute(0, 2, 1, 3)  # (B, T, H, D)
+
+    # -------------------------------------------------------- flash_attention_2
+    def _forward_flash(self, qkv: torch.Tensor) -> torch.Tensor:
+        _flash_attn_required()
+        from flash_attn import flash_attn_qkvpacked_func
+        # flash_attn expects (B, T, 3, H, D) in fp16/bf16 already; Evo1 attn
+        # blocks already cast to bf16 in __init__.
+        out = flash_attn_qkvpacked_func(
+            qkv,
+            dropout_p=self.dropout_p if self.training else 0.0,
+            softmax_scale=self.softmax_scale,
+            causal=self.causal,
+        )
+        return out  # (B, T, H, D)
+
+    # ----------------------------------------------------------- KV-cache path
+    def _forward_with_cache(
+        self,
+        qkv: torch.Tensor,
+        inference_params,
+    ) -> torch.Tensor:
+        # qkv: (B, T, 3, H, D). Apply rotary at the current offset, append kv
+        # to cache, attend over the cumulative kv. For correctness we use SDPA
+        # which has stable behaviour at all sequence lengths.
+        if self.rotary_emb_dim > 0:
+            qkv = self.rotary_emb(
+                qkv,
+                seqlen_offset=inference_params.seqlen_offset,
+                max_seqlen=inference_params.max_seqlen,
+            )
+        q, k, v = qkv.unbind(dim=2)
+        kv = torch.stack((k, v), dim=2)  # (B, T, 2, H, D)
+        kv = _update_kv_cache(kv, inference_params, self.layer_idx)
+        k_full, v_full = kv.unbind(dim=2)  # (B, S_total, H, D)
+
+        q = q.permute(0, 2, 1, 3)
+        k_full = k_full.permute(0, 2, 1, 3)
+        v_full = v_full.permute(0, 2, 1, 3)
+        scale = self.softmax_scale if self.softmax_scale is not None else None
+        is_causal = self.causal and q.shape[-2] == k_full.shape[-2]
+        out = F.scaled_dot_product_attention(
+            q, k_full, v_full, is_causal=is_causal, scale=scale,
+        )
+        return out.permute(0, 2, 1, 3)  # (B, T, H, D)
+
+    # ---------------------------------------------------------------- forward
+    def forward(
+        self,
+        x: torch.Tensor,
+        inference_params=None,
+        output_attentions: bool = False,
+        **_unused,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Returns (out, attn_weights_or_None) where out is (B, T, embed_dim)."""
+        if self.num_heads_kv != self.num_heads:
+            raise NotImplementedError("GQA is not exercised by Evo1; please file an issue if needed.")
+
+        qkv = self._project_qkv(x)  # (B, T, 3, H, D)
+
+        if inference_params is not None:
+            out_btd = self._forward_with_cache(qkv, inference_params)
+            attn_weights = None
+        else:
+            if self.rotary_emb_dim > 0:
+                qkv = self.rotary_emb(qkv, seqlen_offset=0, max_seqlen=qkv.shape[1])
+
+            backend = self.attn_implementation
+            if output_attentions and backend != "eager":
+                # Standard HF behaviour: silently fall back to eager so we can
+                # actually compute and return the attention matrix.
+                backend = "eager"
+
+            if backend == "eager":
+                out_btd, attn_weights = self._forward_eager(qkv, output_attentions=output_attentions)
+            elif backend == "sdpa":
+                out_btd = self._forward_sdpa(qkv)
+                attn_weights = None
+            elif backend == "flash_attention_2":
+                out_btd = self._forward_flash(qkv)
+                attn_weights = None
+            else:
+                raise ValueError(f"Unknown attn_implementation: {backend!r}")
+
+        # (B, T, H, D) -> (B, T, embed_dim)
+        B, T, H, D = out_btd.shape
+        out_flat = out_btd.reshape(B, T, H * D)
+        return self.out_proj(out_flat), attn_weights

cache.py

@@ -0,0 +1,132 @@
+# 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)
+

config.json

@@ -0,0 +1,86 @@
+{
+  "architectures": [
+    "Evo1ForCausalLM"
+  ],
+  "attn_layer_idxs": [
+    8,
+    16,
+    24
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_evo1.Evo1Config",
+    "AutoModel": "modeling_evo1.Evo1Model",
+    "AutoModelForCausalLM": "modeling_evo1.Evo1ForCausalLM"
+  },
+  "column_split": false,
+  "column_split_hyena": true,
+  "dtype": "bfloat16",
+  "eps": 1e-06,
+  "final_norm": true,
+  "hidden_size": 4096,
+  "hyena_filter_groups": 1,
+  "hyena_layer_idxs": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22,
+    23,
+    25,
+    26,
+    27,
+    28,
+    29,
+    30,
+    31
+  ],
+  "inference_mode": false,
+  "inner_mlp_size": 10928,
+  "log_intermediate_values": false,
+  "make_vocab_size_divisible_by": 8,
+  "max_seqlen": 8192,
+  "mha_out_proj_bias": true,
+  "mlp_activation": "gelu",
+  "model_parallel_size": 1,
+  "model_type": "evo1",
+  "num_attention_heads": 32,
+  "num_filters": 4096,
+  "num_layers": 32,
+  "pipe_parallel_size": 1,
+  "prefill_style": "fft",
+  "proj_groups": 1,
+  "qkv_proj_bias": true,
+  "rotary_emb_base": 10000,
+  "rotary_emb_scaling_factor": 1.0,
+  "short_filter_bias": true,
+  "short_filter_length": 3,
+  "smeared_gqa": false,
+  "split_k0": true,
+  "state_size": 8,
+  "tie_embeddings": true,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "use_flash_attention_2": true,
+  "use_flash_attn": false,
+  "use_flash_depthwise": false,
+  "use_flash_rmsnorm": false,
+  "use_flashfft": false,
+  "use_interpolated_rotary_pos_emb": false,
+  "vocab_size": 512
+}

configuration_evo1.py

@@ -0,0 +1,153 @@
+# Apache 2.0 - port of togethercomputer's StripedHyenaConfig.
+"""Configuration for Evo1 (StripedHyena 7B family)."""
+
+from __future__ import annotations
+
+import json
+
+from transformers import PretrainedConfig
+
+
+class Evo1Config(PretrainedConfig):
+    """Configuration for the Evo1 family.
+
+    Defaults match the evo-1-8k-base / evo-1-131k-base / evo-1.5-8k-base
+    checkpoints (32 layers, 4096 hidden, 32 heads, attn at idx [8, 16, 24]
+    and Hyena everywhere else, byte-level vocab_size=512). The 131k variant
+    overrides ``use_interpolated_rotary_pos_emb`` and ``rotary_emb_scaling_factor``
+    plus a longer ``max_seqlen``.
+    """
+
+    model_type = "evo1"
+
+    def __init__(
+        self,
+        # Architecture
+        vocab_size: int = 512,
+        hidden_size: int = 4096,
+        num_filters: int = 4096,
+        inner_mlp_size: int = 10928,
+        attn_layer_idxs=None,
+        hyena_layer_idxs=None,
+        num_layers: int = 32,
+        num_attention_heads: int = 32,
+        proj_groups: int = 1,
+        hyena_filter_groups: int = 1,
+        short_filter_length: int = 3,
+        short_filter_bias: bool = True,
+        state_size: int = 8,
+        column_split: bool = False,
+        column_split_hyena: bool = True,
+        split_k0: bool = True,
+        smeared_gqa: bool = False,
+        # Norms
+        eps: float = 1e-6,
+        final_norm: bool = True,
+        # Linear biases
+        mha_out_proj_bias: bool = True,
+        qkv_proj_bias: bool = True,
+        # Embeddings
+        tie_embeddings: bool = True,
+        make_vocab_size_divisible_by: int = 8,
+        # Activations
+        mlp_activation: str = "gelu",
+        # Sequence length / RoPE
+        max_seqlen: int = 8192,
+        rotary_emb_base: float = 10000,
+        use_interpolated_rotary_pos_emb: bool = False,
+        rotary_emb_scaling_factor: float = 1.0,
+        # Inference engine
+        prefill_style: str = "fft",
+        inference_mode: bool = False,
+        # Backend toggles
+        use_cache: bool = True,
+        use_flash_attention_2: bool = True,
+        use_flash_rmsnorm: bool = False,
+        use_flash_depthwise: bool = False,
+        use_flashfft: bool = False,
+        use_flash_attn: bool = False,
+        # Misc
+        log_intermediate_values: bool = False,
+        model_parallel_size: int = 1,
+        pipe_parallel_size: int = 1,
+        **kwargs,
+    ):
+        if attn_layer_idxs is None:
+            attn_layer_idxs = [8, 16, 24]
+        if hyena_layer_idxs is None:
+            hyena_layer_idxs = [i for i in range(num_layers) if i not in attn_layer_idxs]
+
+        # Architecture
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_filters = num_filters
+        self.inner_mlp_size = inner_mlp_size
+        self.attn_layer_idxs = attn_layer_idxs
+        self.hyena_layer_idxs = hyena_layer_idxs
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+        self.proj_groups = proj_groups
+        self.hyena_filter_groups = hyena_filter_groups
+        self.short_filter_length = short_filter_length
+        self.short_filter_bias = short_filter_bias
+        self.state_size = state_size
+        self.column_split = column_split
+        self.column_split_hyena = column_split_hyena
+        self.split_k0 = split_k0
+        self.smeared_gqa = smeared_gqa
+        # Norms
+        self.eps = eps
+        self.final_norm = final_norm
+        # Biases
+        self.mha_out_proj_bias = mha_out_proj_bias
+        self.qkv_proj_bias = qkv_proj_bias
+        # Embeddings
+        self.tie_embeddings = tie_embeddings
+        self.make_vocab_size_divisible_by = make_vocab_size_divisible_by
+        # Activations
+        self.mlp_activation = mlp_activation
+        # Length / RoPE
+        self.max_seqlen = max_seqlen
+        self.rotary_emb_base = rotary_emb_base
+        self.use_interpolated_rotary_pos_emb = use_interpolated_rotary_pos_emb
+        self.rotary_emb_scaling_factor = rotary_emb_scaling_factor
+        # Engine
+        self.prefill_style = prefill_style
+        self.inference_mode = inference_mode
+        # Backend toggles
+        self.use_cache = use_cache
+        self.use_flash_attention_2 = use_flash_attention_2
+        self.use_flash_rmsnorm = use_flash_rmsnorm
+        self.use_flash_depthwise = use_flash_depthwise
+        self.use_flashfft = use_flashfft
+        self.use_flash_attn = use_flash_attn
+        # Misc
+        self.log_intermediate_values = log_intermediate_values
+        self.model_parallel_size = model_parallel_size
+        self.pipe_parallel_size = pipe_parallel_size
+        super().__init__(**kwargs)
+
+    # ------------------------------------------------------------------
+    # Backwards-compatible attribute access.
+    #
+    # The internal blocks (RMSNorm, ParallelGatedMLP, ...) call
+    # ``config.get(key, default)`` because they were originally written
+    # against a `dotdict`. PretrainedConfig has a different `.get`, so we
+    # provide a dict-like one that delegates to attribute access.
+    # ------------------------------------------------------------------
+    @property
+    def num_hidden_layers(self) -> int:
+        # HF generation utilities (DynamicCache, etc.) expect this name; we
+        # keep ``num_layers`` as the source of truth to match the upstream
+        # StripedHyena config.
+        return self.num_layers
+
+    def get(self, key, default=None):
+        # Dict-style access used by internal blocks (RMSNorm, MHA, ...).
+        return getattr(self, key, default)
+
+    @classmethod
+    def from_original_config(cls, config_path: str, **kwargs) -> "Evo1Config":
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        return cls(**config, **kwargs)

engine.py

@@ -0,0 +1,388 @@
+# Copyright (c) Together
+# This software is distributed under the terms of the Apache License, Version 2.0
+# Author: Michael Poli
+
+import gc
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    import conv1d_cpp  # noqa: F401
+except Exception:
+    pass
+
+IIR_PREFILL_MODES = [
+    "recurrence",
+    "modal-fft",
+    "hybrid-modal-recurrence",
+    "modal-scan",
+    "canonical-fft",
+    "iir-fir-caching",
+]
+
+
+def canonicalize_modal_system(poles, residues):
+    """Canonicalize a modal system.
+
+    Args:
+        poles (Tensor): The poles of the system.
+        residues (Tensor): The residues of the system.
+
+    Returns:
+        Tuple[Tensor, Tensor]: The canonicalized poles and residues.
+    """
+    raise NotImplementedError
+
+
+def list_tensors(idx):
+    for obj in gc.get_objects():
+        try:
+            if torch.is_tensor(obj) and isinstance(obj, torch.Tensor):
+                # dump to log
+                print(type(obj), obj.size())
+                el = obj[0]
+                with open(f"tensors_{idx}.txt", "a") as f:
+                    f.write(f"{type(obj)} {obj.size()} {el}\n")
+        except Exception as e:
+            pass
+
+
+class HyenaInferenceEngine:
+    def __init__(
+        self,
+        fir_fn=None,
+        iir_prefill_style="modal-fft",
+        layer_idx=None,
+    ) -> None:
+        self.fir_fn = fir_fn
+        assert iir_prefill_style in IIR_PREFILL_MODES, f"iir_prefill_style must be one of {IIR_PREFILL_MODES}"
+        self.iir_prefill_style = iir_prefill_style
+        self.layer_idx = layer_idx
+        self.low_mem_mode = False
+
+    def parallel_fir(
+        self,
+        fir_fn,
+        u,
+        weight,
+        bias,
+        L,
+        fir_length=3,
+        inference_params=None,
+        prefill_mode=None,
+        padding_mask=None,
+    ):
+        """Compute the output state of the long convolutional filter."""
+        # prepare input layout, dimensions and dispatch to fir kernel
+        if fir_fn != torch.nn.functional.conv1d:
+            z_pre = fir_fn(u)[:, :L]  # B, L, D
+            z_pre = z_pre.permute(0, 2, 1)
+        else:
+            u = u.permute(0, 2, 1)  # B, D, L
+            z_pre = fir_fn(
+                u,
+                weight,
+                bias=None,  # don't pass it here, add manually instead!  source of small error
+                stride=1,
+                padding=fir_length - 1,
+                groups=u.shape[1],
+            )[..., :L]
+
+            # add manually instead!  source of small error
+            z_pre = z_pre + bias[None, :, None]
+
+        # handle padding post fir, the only place with biases
+        if type(padding_mask) == torch.Tensor:
+            z_pre = z_pre * padding_mask[:, None]
+
+        if inference_params is not None:
+            # handle seqlen last and dim last cases for `u`
+            if fir_fn != torch.nn.functional.conv1d:
+                fir_state = u[:, -fir_length + 1 :].permute(0, 2, 1)
+            else:
+                fir_state = u[..., -fir_length + 1 :]
+        else:
+            fir_state = None
+
+        return z_pre, fir_state
+
+    def parallel_iir(
+        self,
+        z_pre,
+        h,
+        D,
+        L,
+        poles,
+        residues,
+        t,
+        dims,
+        layer_idx,
+        inference_params=None,
+        prefill_style="fft",
+        fftconv_fn=None,
+        padding_mask=None,
+        use_flashfft=False,
+        column_split_hyena=False,
+        long_fir_threshold=None,
+    ):
+        """Compute the output state of the short convolutional filter."""
+        fft_size = 2 * L
+        hidden_size, num_attention_heads, hidden_size_per_attention_head, _, _ = dims
+        # Compatibility with training infra that column splits the projections
+        if column_split_hyena:
+            z = z_pre.reshape(
+                z_pre.shape[0],
+                num_attention_heads,
+                3 * hidden_size_per_attention_head,
+                z_pre.shape[2],
+            )
+            x2, x1, v = (
+                z[:, :, :hidden_size_per_attention_head],
+                z[
+                    :,
+                    :,
+                    hidden_size_per_attention_head : 2 * hidden_size_per_attention_head,
+                ],
+                z[:, :, 2 * hidden_size_per_attention_head :],
+            )
+            x2, x1, v = (
+                x2.reshape(x2.shape[0], -1, x2.shape[-1]),
+                x1.reshape(x1.shape[0], -1, x1.shape[-1]),
+                v.reshape(v.shape[0], -1, v.shape[-1]),
+            )
+        else:
+            x2, x1, v = z_pre.split([hidden_size, hidden_size, hidden_size], dim=1)
+
+        x1v = x1 * v
+
+        if inference_params is not None and prefill_style == "recurrence":
+            y = self.prefill_via_direct_recurrence(
+                inference_params=inference_params,
+                x1v=x1v,
+                L=L,
+                poles=poles,
+                residues=residues,
+            )
+
+        else:
+            if use_flashfft and (L % 2) == 0:  # only works with even L
+                y = fftconv_fn(
+                    x1v.to(dtype=torch.bfloat16).contiguous(),
+                    h.to(dtype=torch.float32),
+                )
+                X_s = None
+
+            elif long_fir_threshold is None:
+                H = torch.fft.rfft(h.to(dtype=torch.float32), n=fft_size) / fft_size
+                X_s = torch.fft.fft(x1v.to(dtype=torch.float32), n=fft_size)
+                X = X_s[..., : H.shape[-1]]
+                if len(z_pre.shape) > 3:
+                    H = H.unsqueeze(1)
+                y = torch.fft.irfft(X * H, n=fft_size, norm="forward")[..., :L]
+
+            else:
+                assert h.shape[0] == 1, "batch size must be 1 for long_fir_threshold"
+                h = h[0][:, None]  # rearrange to d, 1, l for depthwise conv1d
+                h = h[..., :long_fir_threshold]
+                y = F.conv1d(
+                    x1v,
+                    h.to(dtype=x1v.dtype),
+                    stride=1,
+                    groups=x1v.shape[1],
+                    padding=h.shape[-1] - 1,
+                )[..., :L]
+
+        y = y.to(dtype=x1v.dtype)
+        y = (y + x1v * D.unsqueeze(-1)) * x2
+
+        if inference_params is not None:
+            if prefill_style == "fft":
+                self.prefill_via_modal_fft(
+                    inference_params=inference_params,
+                    x1v=x1v,
+                    X_s=X_s,
+                    L=L,
+                    t=t,
+                    poles=poles,
+                    dims=dims,
+                    layer_idx=layer_idx,
+                    use_flashfft=use_flashfft,
+                    fftconv_fn=fftconv_fn,
+                )
+
+            elif prefill_style == "recurrence":
+                # recurrent prefill is done before
+                pass
+            else:
+                raise NotImplementedError
+            if self.low_mem_mode:
+                # TODO: smarter gc
+                del z_pre, x2, x1, v, x1v, h, poles, residues
+                torch.cuda.empty_cache()
+
+        return y.permute(0, 2, 1)
+
+    def step_fir(self, u, fir_state, weight, bias=None):
+        """Step the FIR filter.
+
+        Note:
+        `fir_state` contains the last `short_filter_length - 1` elements of `u`: `u_(L-2), u_{L-1), ...`
+        We assume dimensions of `short_filter_weight` to be `[d, 1, short_filter_len]` (SISO / multi SISO layout).
+        """
+        h0, h = weight[..., 0, -1], weight[..., 0, :-1]
+        h0, h = h0[None], h[None]
+        y = h0 * u + torch.sum(fir_state * h, dim=-1) + bias
+
+        # update
+        fir_state = torch.roll(fir_state, -1, dims=2)
+        fir_state[..., -1] = u
+        return y, fir_state
+
+    def step_iir(self, x2, x1, v, D, residues, poles, iir_state, iir_groups=1):
+        x1v = x1 * v
+
+        residues, poles = (
+            torch.view_as_complex(residues.to(torch.float32)),
+            torch.view_as_complex(poles.to(torch.float32)),
+        )
+        # squeeze the dummy seqlen dimension
+        # D, state_dim, 1 -> 1, D, state_dim
+        residues, poles = residues[..., 0][None], poles[..., 0][None]
+        iir_state = poles * iir_state + x1v[..., None]
+
+        res_state = torch.sum(residues * iir_state, dim=-1).real
+
+        if iir_groups > 1:
+            raise NotImplementedError
+        y = x2 * (res_state + D * x1v)
+
+        return y, iir_state
+
+    def prefill_via_fir_caching(self, u, inference_params, L, *args, **kwargs):
+        """Turns the IIR filter into a FIR and uses a cache for decoding."""
+        raise NotImplementedError(":)")
+
+    def prefill_via_direct_recurrence(
+        self, inference_params, x1v, L, residues, poles, *args, **kwargs
+    ) -> torch.Tensor:
+        """
+        Compute the IIR state via explicit SSM recurrence (modal form)
+
+        This is the most memory efficient prefilling method for Hyena filters.
+
+        Note:
+            dtypes: [state: float32, poles: float32, x1v: bfloat16, output: bfloat16]
+        """
+        state_dim = poles.shape[1]
+        x1v_ = x1v[..., None, None]  # b, d, l, sdim, reim
+        x1v_ = x1v_.repeat(1, 1, 1, state_dim, 2)  # b, d, l, sdim, reim
+        x1v_[..., 1] = 0
+
+        state = 0 * x1v_[:, :, 0]
+        output = 0 * x1v_[:, :, :, 0, 0]  # b, d, l
+
+        # suppress dummy seqlen dimension
+        poles = poles[:, :, 0][None]
+        residues = residues[:, :, 0][None].repeat(x1v_.shape[0], 1, 1, 1)  # b, d, sdim, reim
+
+        # state: b, d, sdim, reim
+        # poles: 1, d, sdim, reim
+        # x1v_: b, d, l, sdim, reim
+        for i in range(L):
+            state[..., 0] = poles[..., 0] * state[..., 0] - poles[..., 1] * state[..., 1] + x1v_[:, :, i, :, 0]
+            state[..., 1] = poles[..., 0] * state[..., 1] + poles[..., 1] * state[..., 0] + x1v_[:, :, i, :, 1] 
+            output[:, :, i] = torch.sum(residues * state, dim=-2)[..., 0]  # .real
+            
+        inference_params.state_dict[self.layer_idx] = torch.view_as_complex(state.to(dtype=torch.float32))
+
+        return output
+
+    def prefill_via_hybrid_recurrence(self, inference_params, u, log_poles, x1v_f_a, L, *args, **kwargs):
+        """
+        Compute the IIR state via hybrid recurrence-convolution over blocks
+        """
+        raise NotImplementedError(":)")
+
+    def prefill_via_scan(self, u, inference_params=None, *args, **kwargs):
+        raise NotImplementedError
+
+    def prefill_via_canonical_fft(self, u, inference_params=None, *args, **kwargs):
+        """
+        Compute the IIR state via a single FFT with the denominator of the SSM in companion form.
+
+        This is the most memory efficient "parallelized" prefilling method for Hyena.
+
+        From: https://arxiv.org/abs/2310.18780
+        """
+        raise NotImplementedError(":)")
+
+    def prefill_via_modal_fft(
+        self,
+        inference_params,
+        x1v,
+        L,
+        poles,
+        t,
+        dims,
+        layer_idx,
+        X_s=None,
+        use_flashfft=False,
+        fftconv_fn=None,
+        state_dtype=torch.complex64,
+        *args,
+        **kwargs,
+    ):
+        """
+        Compute the IIR state via a single FFT, using the poles of the SSM in modal form.
+        """
+        # When the model has a long convolution derived from a SSM in modal form and prefill_style is "fft",
+        # we split the filter into poles and residues and reuse FFT computation on the input.
+        # This optimization is currently not supported when using flashfftconv.
+        hidden_size, _, _, state_size, hyena_filter_groups = dims
+
+        if use_flashfft:
+            # using real states
+            poles = poles.squeeze().reshape(poles.shape[0], -1)[..., None]
+
+            state_s = poles**t
+            if hyena_filter_groups > 1:
+                raise NotImplementedError
+
+            x1v = x1v[:, :, None].repeat(1, 1, 2 * state_size, 1)
+            x1v = x1v.reshape(x1v.shape[0], -1, x1v.shape[-1])
+            state_s = state_s[None]
+
+            state = fftconv_fn(
+                x1v.contiguous(),
+                state_s.to(dtype=torch.float32),
+            )
+            state = state[..., L - 1].reshape(x1v.shape[0], hidden_size, state_size, 2)
+            state = torch.view_as_complex(state.contiguous().to(dtype=torch.float32))
+            inference_params.state_dict[self.layer_idx] = state
+        else:
+            assert X_s is not None
+            bs = x1v.shape[0]
+            fft_size = 2 * L
+            poles = torch.view_as_complex(poles.to(torch.float32))
+            state_s = poles**t
+            state_S = torch.fft.fft(state_s, n=fft_size).repeat(bs, 1, 1, 1)  # B, D, state_dim, 2 * L
+            if hyena_filter_groups > 1:
+                state_S = state_S.repeat_interleave(hidden_size // hyena_filter_groups, 1)
+            state = torch.fft.ifft(X_s[..., None, :] * state_S, n=fft_size)
+            inference_params.state_dict[layer_idx] = state[..., L - 1].to(dtype=state_dtype)
+
+    def _compute_state(self, log_poles, u, t, L, *args, **kwargs):
+        """
+        Compute the IIR state given an input `u` and log_poles of the modal system.
+        """
+        bs = u.shape[0]
+        fft_size = 2 * L
+        U = torch.fft.rfft(u.to(torch.float32), n=fft_size)
+        fft_size = 2 * L
+        x = (log_poles * t).exp()
+        # [batch, hidden_size, state_dim, 2 * seqlen]
+        X = torch.fft.fft(x, n=fft_size).repeat(bs, 1, 1, 1)
+        state = torch.fft.ifft(U[..., None, :] * X, n=fft_size)[..., :L]
+        return state

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "eos_token_id": 0,
+  "pad_token_id": 1,
+  "_from_model_config": true
+}

layers.py

@@ -0,0 +1,82 @@
+# Copyright (c) Together
+# Apache 2.0 - Author: Michael Poli
+# Adapted for the minimal Evo1 HF port.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+def grab_first_if_tuple(x):
+    if x.__class__.__name__ == "tuple":
+        return x[0]
+    return x
+
+
+class RMSNorm(torch.nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.eps = config.eps
+        self.hidden_size = config.hidden_size
+        self.scale = torch.nn.Parameter(torch.ones(self.hidden_size))
+        self.register_parameter("scale", self.scale)
+        self.use_flash_rmsnorm = config.get("use_flash_rmsnorm", False)
+        if self.use_flash_rmsnorm:
+            from flash_attn.ops.rms_norm import rms_norm as rmsnorm_func
+            self.rmsnorm_func = rmsnorm_func
+
+    def forward(self, x):
+        if self.use_flash_rmsnorm:
+            return self.rmsnorm_func(x, self.scale, self.eps)
+        y = x / (x.norm(2, dim=-1, keepdim=True) * self.hidden_size ** (-1.0 / 2) + self.eps)
+        return self.scale * y
+
+
+class ParallelGatedMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        multiple_of = config.get("inner_size_multiple_of", 64)
+        self.act_type = config.get("mlp_activation", "silu")
+        if self.act_type == "gelu":
+            self.act = F.gelu
+        elif self.act_type == "silu":
+            self.act = F.silu
+        else:
+            raise NotImplementedError(f"Unknown mlp_activation: {self.act_type}")
+
+        self.multiple_of = multiple_of * config.model_parallel_size
+        inner_size = int(2 * config.hidden_size * 4 / 3)
+        inner_size = self.multiple_of * ((inner_size + self.multiple_of - 1) // self.multiple_of)
+        if config.get("inner_mlp_size", None) is not None:
+            inner_size = config.inner_mlp_size
+
+        self.l1 = nn.Linear(config.hidden_size, inner_size, bias=False)
+        self.l2 = nn.Linear(config.hidden_size, inner_size, bias=False)
+        self.l3 = nn.Linear(inner_size, config.hidden_size, bias=False)
+
+    def forward(self, z):
+        z1, z2 = self.l1(z), self.l2(z)
+        z1, z2 = grab_first_if_tuple(z1), grab_first_if_tuple(z2)
+        y = self.l3(self.act(z1) * z2)
+        return grab_first_if_tuple(y)
+
+
+class VocabParallelEmbedding(nn.Embedding):
+    """Single-process variant of the original VocabParallelEmbedding.
+
+    The original supports tensor-parallel embedding sharding. We keep the
+    naming so existing checkpoints load directly, but drop all distributed
+    paths since this minimal port runs on a single device.
+    """
+
+    def __init__(self, config):
+        vocab_size = config.vocab_size
+        padding_idx = config.get("padding_idx", None)
+        super().__init__(vocab_size, embedding_dim=config.hidden_size, padding_idx=padding_idx)
+
+    def embed(self, x: Tensor) -> Tensor:
+        return self.forward(x)
+
+    def unembed(self, u: Tensor) -> Tensor:
+        return u @ self.weight.T

model-00001-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bac7a0023292c7b07b32251b03ed1316c3e1a6d77d72ce34a813b8eee2593941
+size 4985304472

model-00002-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7f5ce86f31aaa6e14aee17cfb5da728d39cf7b2b2bf9083cab48e6216b2cb60a
+size 4928275752

model-00003-of-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:573e321dfe66a07a9fb2529be5ac9c406d081c7575405d5bb0620d742758a78d
+size 2999633368

model.safetensors.index.json

@@ -0,0 +1,446 @@
+{
+  "metadata": {
+    "total_parameters": 6452781056,
+    "total_size": 12913164672
+  },
+  "weight_map": {
+    "backbone.blocks.0.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.0.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.12.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.12.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.12.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.13.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.13.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.13.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.14.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.14.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.15.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.15.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.inner_mha_cls.Wqkv.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.inner_mha_cls.Wqkv.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.inner_mha_cls.out_proj.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.inner_mha_cls.out_proj.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.inner_mha_cls.rotary_emb.inv_freq": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.17.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.17.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.18.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.18.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.19.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.19.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.2.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.20.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.20.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.20.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.21.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.21.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.22.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.22.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.23.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.23.filter.short_filter_bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.filter.short_filter_weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.mlp.l2.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.mlp.l3.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.out_filter_dense.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.projections.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.inner_mha_cls.Wqkv.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.inner_mha_cls.Wqkv.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.inner_mha_cls.out_proj.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.inner_mha_cls.out_proj.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.inner_mha_cls.rotary_emb.inv_freq": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.mlp.l1.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.24.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.24.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.24.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.25.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.25.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.25.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.25.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.26.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.26.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.27.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.27.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.28.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.28.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.29.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.29.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.3.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.30.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.30.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.30.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.filter.poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.31.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.31.filter.short_filter_bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.filter.short_filter_weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.mlp.l1.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.mlp.l2.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.mlp.l3.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.out_filter_dense.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.out_filter_dense.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.projections.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.projections.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.4.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.4.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.inner_mha_cls.Wqkv.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.inner_mha_cls.Wqkv.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.inner_mha_cls.out_proj.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.inner_mha_cls.out_proj.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.inner_mha_cls.rotary_emb.inv_freq": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.residues": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.short_filter_bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.short_filter_weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.mlp.l1.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.mlp.l2.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.mlp.l3.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.out_filter_dense.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.projections.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.embedding_layer.weight": "model-00001-of-00003.safetensors",
+    "backbone.norm.scale": "model-00001-of-00003.safetensors"
+  }
+}

modeling_evo1.py

@@ -0,0 +1,764 @@
+"""Minimal Evo1 (StripedHyena) HuggingFace port.
+
+This module is a refactor of togethercomputer/evo-1-131k-base@1.1_fix's
+``modeling_hyena.py`` + ``model.py`` into a single self-contained file with:
+
+  * ``output_hidden_states`` and ``output_attentions`` plumbed end-to-end,
+  * ``attn_implementation`` switch (``eager`` / ``sdpa`` / ``flash_attention_2``),
+  * ``Evo1Model`` (no LM head, ``BaseModelOutputWithPast``) for ``AutoModel``,
+  * ``Evo1ForCausalLM`` (with logits, ``CausalLMOutputWithPast``)
+    for ``AutoModelForCausalLM``,
+  * minimal external imports (only ``torch`` + ``transformers``; ``flash-attn``
+    is loaded lazily and only when ``attn_implementation='flash_attention_2'``).
+
+Hyena blocks have no attention matrix by construction, so they always emit
+``None`` in the per-layer ``attentions`` tuple. Attention blocks (layers 8,
+16, 24 for Evo1) emit the (B, H, T, T) softmax matrix when
+``output_attentions=True`` (this triggers a one-time fallback from sdpa /
+flash_attention_2 to the eager backend).
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.utils import logging
+
+from .attention import MHA
+from .cache import Evo1Cache, InferenceParams, RecurrentInferenceParams
+from .configuration_evo1 import Evo1Config
+from .engine import HyenaInferenceEngine
+from .layers import ParallelGatedMLP, RMSNorm, VocabParallelEmbedding
+from .rotary import swap_mha_rope
+# dummy import so that trust_remote_code bundles the tokenizer file
+from .tokenization_evo1 import ByteTokenizer  # noqa: F401
+
+logger = logging.get_logger(__name__)
+
+
+# =============================================================================
+# Block: attention (used at layers config.attn_layer_idxs)
+# =============================================================================
+
+
+class AttentionBlock(nn.Module):
+    """Pre-norm Transformer block: norm -> MHA -> residual -> norm -> MLP -> residual."""
+
+    def __init__(self, config, layer_idx) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.pre_norm, self.post_norm = RMSNorm(config), RMSNorm(config)
+        self.proj_groups = config.get("proj_groups", 1)
+        dtype = config.get("attn_block_dtype", torch.bfloat16)
+        mlp_dtype = config.get("mlp_dtype", torch.bfloat16)
+        self.num_attention_heads = config.num_attention_heads
+        self.hidden_size_per_attention_head = (
+            config.hidden_size // config.num_attention_heads
+        )
+
+        attn_impl = getattr(config, "_attn_implementation", "eager")
+
+        self.inner_mha_cls = MHA(
+            embed_dim=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_heads_kv=config.num_attention_heads // self.proj_groups,
+            rotary_emb_dim=config.hidden_size // config.num_attention_heads,
+            qkv_proj_bias=config.get("qkv_proj_bias", True),
+            rotary_emb_base=config.get("rotary_emb_base", 10000),
+            causal=True,
+            layer_idx=layer_idx,
+            out_proj_bias=config.get("mha_out_proj_bias", True),
+            attn_implementation=attn_impl,
+        ).to(dtype=dtype)
+
+        if config.get("use_interpolated_rotary_pos_emb", False):
+            swap_mha_rope(
+                mha=self.inner_mha_cls,
+                kwargs_new_rope={
+                    "scaling_factor": config.get("rotary_emb_scaling_factor", 1.0)
+                },
+            )
+
+        if self.config.get("smeared_gqa", False):
+            self.inner_mha_cls.num_heads_kv = self.inner_mha_cls.num_heads
+        # Make sure the inv_freq buffer round-trips through to_bfloat16/state_dict.
+        self.inner_mha_cls.rotary_emb.register_buffer(
+            "inv_freq", self.inner_mha_cls.rotary_emb.inv_freq
+        )
+
+        self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)
+
+    def forward(
+        self,
+        u: torch.Tensor,
+        inference_params=None,
+        padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        *args,
+        **kwargs,
+    ):
+        if isinstance(padding_mask, torch.Tensor):
+            # Workaround for masking with no qkv bias: this zeros the attended
+            # values at pad positions so they don't leak via attention.
+            u = u * padding_mask[..., None]
+
+        attn_out, attn_weights = self.inner_mha_cls(
+            self.pre_norm(u),
+            inference_params=inference_params,
+            output_attentions=output_attentions,
+        )
+        u = attn_out + u
+
+        if isinstance(padding_mask, torch.Tensor):
+            u = u * padding_mask[..., None]
+        u = self.mlp(self.post_norm(u)) + u
+        return u, attn_weights
+
+
+# =============================================================================
+# Block: Hyena (used at all other layers)
+# =============================================================================
+
+
+class ParallelHyenaFilter(nn.Module):
+    def __init__(self, config, layer_idx) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hyena_filter_groups = config.get(
+            "hyena_filter_groups", self.config.hidden_size
+        )
+
+        self.use_flashfft = config.get("use_flashfft", False)
+        self.state_size = config.state_size
+        self.hidden_size = config.hidden_size
+        self.num_filters = config.num_filters
+        self.inference_mode = config.get("inference_mode", True)
+        self.column_split_hyena = config.get("column_split_hyena", True)
+
+        assert self.hidden_size % self.num_filters == 0
+        assert self.num_filters <= self.hidden_size
+
+        self.D = nn.Parameter(torch.zeros(self.hidden_size))
+
+        # heads only used to slice post-FIR projections like the checkpoint
+        self.num_attention_heads = config.num_attention_heads
+        self.hidden_size_per_attention_head = (
+            self.hidden_size // self.num_attention_heads
+        )
+
+        self.short_filter_length = config.short_filter_length
+        self.short_filter_weight = nn.Parameter(
+            torch.randn(3 * config.hidden_size, 1, config.short_filter_length)
+        )
+        self.short_filter_bias = (
+            nn.Parameter(torch.randn(3 * config.hidden_size))
+            if config.short_filter_bias
+            else None
+        )
+
+        self.engine = HyenaInferenceEngine(layer_idx=layer_idx)
+        self.use_flash_depthwise = config.get("use_flash_depthwise", False)
+        self.data_dtype = None
+
+        if self.use_flash_depthwise:
+            # importlib avoids the top-level static-import check that HF's
+            # dynamic_module_utils.check_imports performs against the file.
+            import importlib
+            FlashDepthwiseConv1d = importlib.import_module("flashfftconv").FlashDepthwiseConv1d
+            self.fir_fn = FlashDepthwiseConv1d(
+                channels=3 * self.hidden_size,
+                kernel_size=self.short_filter_length,
+                padding=self.short_filter_length - 1,
+                weights=self.short_filter_weight,
+                bias=self.short_filter_bias,
+                device=None,
+                dtype=self.config.get("depthwise_dtype", torch.bfloat16),
+            )
+        else:
+            self.fir_fn = F.conv1d
+
+        self.fftconv_fn = None
+        self.long_fir_threshold = config.get("long_fir_threshold", None)
+        if self.long_fir_threshold is not None:
+            assert self.use_flashfft is False, (
+                "long_fir_threshold not compatible with fused flashfft"
+            )
+
+        self.num_systems = self.hidden_size // self.hyena_filter_groups
+        poles = torch.randn(self.num_systems, self.state_size, 1, 2)
+        poles[..., 0] = 1e-2 * torch.randn(self.num_systems, self.state_size, 1)
+        poles[..., 1] = 1e-3 * torch.randn(self.num_systems, self.state_size, 1)
+        self.poles = nn.Parameter(poles)
+        self.residues = nn.Parameter(
+            torch.randn(self.num_systems, self.state_size, 1, 2)
+        )
+        self.h = None
+
+    def forward(self, u, inference_params=None, padding_mask=None, *args, **kwargs):
+        if (
+            inference_params is not None
+            and self.layer_idx in inference_params.fir_state_dict.keys()
+        ):
+            return self.sequential_forward(u, inference_params)
+        return self.parallel_forward(u, inference_params, padding_mask)
+
+    def parallel_forward(self, u, inference_params=None, padding_mask=None):
+        L = u.shape[1]
+        z_pre, fir_state = self.engine.parallel_fir(
+            self.fir_fn,
+            u,
+            self.short_filter_weight,
+            self.short_filter_bias,
+            L,
+            fir_length=self.short_filter_length,
+            inference_params=inference_params,
+            padding_mask=padding_mask,
+        )
+        if inference_params:
+            inference_params.fir_state_dict[self.layer_idx] = fir_state
+
+        if self.h is None:
+            h, _, _, _ = self.compute_filter(L, u.device)
+        else:
+            h = self.h
+
+        if self.hyena_filter_groups > 1:
+            h = h.repeat_interleave(self.hidden_size // self.hyena_filter_groups, 1)
+
+        dims = (
+            self.hidden_size,
+            self.num_attention_heads,
+            self.hidden_size_per_attention_head,
+            self.state_size,
+            self.hyena_filter_groups,
+        )
+        y = self.engine.parallel_iir(
+            z_pre,
+            h,
+            self.D,
+            L,
+            t=self.t,
+            poles=self.poles,
+            residues=self.residues,
+            dims=dims,
+            inference_params=inference_params,
+            layer_idx=self.layer_idx,
+            prefill_style=self.config.get("prefill_style", "fft"),
+            use_flashfft=self.use_flashfft,
+            fftconv_fn=self.fftconv_fn,
+            column_split_hyena=self.column_split_hyena,
+            long_fir_threshold=self.long_fir_threshold,
+            padding_mask=padding_mask,
+        )
+        return y, inference_params
+
+    def sequential_forward(self, u, inference_params):
+        if self.data_dtype is None:
+            self.data_dtype = u.dtype
+        if len(u.shape) > 2:
+            u = u[:, -1]
+
+        fir_state = inference_params.fir_state_dict[self.layer_idx]
+        iir_state = inference_params.state_dict[self.layer_idx]
+
+        z_pre, fir_state = self.engine.step_fir(
+            u, fir_state,
+            weight=self.short_filter_weight, bias=self.short_filter_bias,
+        )
+        if self.column_split_hyena:
+            x_reshaped = z_pre.reshape(
+                z_pre.shape[0],
+                self.num_attention_heads,
+                3 * self.hidden_size_per_attention_head,
+            )
+            head = self.hidden_size_per_attention_head
+            x2 = x_reshaped[:, :, :head].reshape(z_pre.shape[0], -1)
+            x1 = x_reshaped[:, :, head : 2 * head].reshape(z_pre.shape[0], -1)
+            v = x_reshaped[:, :, 2 * head:].reshape(z_pre.shape[0], -1)
+        else:
+            x2, x1, v = z_pre.split(
+                [self.hidden_size, self.hidden_size, self.hidden_size], dim=1
+            )
+
+        y, iir_state = self.engine.step_iir(
+            x2, x1, v, self.D, self.residues, self.poles, iir_state,
+            iir_groups=self.hyena_filter_groups,
+        )
+        inference_params.fir_state_dict[self.layer_idx] = fir_state
+        inference_params.state_dict[self.layer_idx] = iir_state
+        y = y.to(dtype=self.data_dtype)
+        return y[:, None], inference_params
+
+    def update_time(self, L, device):
+        if not hasattr(self, "t"):
+            self.t = torch.arange(L, device=device)[None, None]
+        elif self.t.shape[-1] < L:
+            self.t = torch.arange(L, device=device)[None, None]
+        else:
+            self.t = self.t[..., :L]
+
+    def compute_filter(self, L, device):
+        self.update_time(L, device)
+        filter_dtype = torch.float32
+        residues = torch.view_as_complex(self.residues.to(filter_dtype))
+        log_poles = torch.view_as_complex(self.poles.to(filter_dtype)).log()
+        h = (residues * (log_poles * self.t).exp()).real.sum(1)[None]
+        return h, filter_dtype, log_poles, residues
+
+
+class ParallelGatedConvBlock(nn.Module):
+    def __init__(self, config, layer_idx) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.low_mem_mode = config.get("low_mem_mode", False)
+        dtype = config.get("hyena_block_dtype", torch.float32)
+        mlp_dtype = config.get("mlp_dtype", torch.bfloat16)
+        self.pre_norm = RMSNorm(config).to(dtype=dtype)
+        self.post_norm = RMSNorm(config).to(dtype=dtype)
+        self.filter = ParallelHyenaFilter(config, layer_idx).to(dtype=dtype)
+        self.projections = nn.Linear(config.hidden_size, 3 * config.hidden_size)
+        self.out_filter_dense = nn.Linear(
+            config.hidden_size, config.hidden_size
+        ).to(dtype)
+        self.mlp = ParallelGatedMLP(config).to(dtype=mlp_dtype)
+
+    def forward(
+        self,
+        u,
+        inference_params=None,
+        padding_mask=None,
+        output_attentions: bool = False,
+        *args,
+        **kwargs,
+    ):
+        z = self.projections(self.pre_norm(u))
+
+        if isinstance(padding_mask, torch.Tensor):
+            z = z * padding_mask[..., None]
+
+        z, inference_params = self.filter(
+            z, inference_params=inference_params, padding_mask=padding_mask
+        )
+        z_in = self.out_filter_dense(z) + u
+
+        if isinstance(padding_mask, torch.Tensor):
+            z_in = z_in * padding_mask[..., None]
+
+        y = self.mlp(self.post_norm(z_in)) + z_in
+        # Hyena blocks have no attention matrix.
+        return y, None
+
+
+def get_block(config, layer_idx, flash_fft=None):
+    if layer_idx in config.attn_layer_idxs:
+        return AttentionBlock(config, layer_idx)
+    if layer_idx in config.hyena_layer_idxs:
+        block = ParallelGatedConvBlock(config, layer_idx)
+        if config.get("use_flashfft", False):
+            block.filter.fftconv_fn = flash_fft
+        return block
+    raise NotImplementedError(f"layer_idx {layer_idx} not in attn or hyena indices")
+
+
+# =============================================================================
+# Backbone (StripedHyena)
+# =============================================================================
+
+
+class StripedHyena(nn.Module):
+    """Pure backbone: token embedding -> N blocks -> RMSNorm.
+
+    The unembed step is owned by the LM head wrapper, not here, so that
+    ``Evo1Model`` (no LM head) can return the post-norm hidden state as
+    ``last_hidden_state`` cleanly.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embedding_layer = VocabParallelEmbedding(config)
+        self.norm = RMSNorm(config) if config.get("final_norm", True) else None
+
+        if config.get("use_flashfft", False):
+            import importlib
+            FlashFFTConv = importlib.import_module("flashfftconv").FlashFFTConv
+            # NOTE: the original togethercomputer reference had ``config.seqlen``
+            # here, which is a typo - that attribute doesn't exist on the
+            # config (it's ``max_seqlen``). The bug was unreachable upstream
+            # because ``use_flashfft`` defaults to False; we fix it so the
+            # path is at least loadable for users who do enable it.
+            # FlashFFTConv requires its build-time seqlen to be 2x the
+            # longest input it'll ever see (zero-padding for FFT).
+            self.flash_fft = FlashFFTConv(2 * config.max_seqlen, dtype=torch.bfloat16)
+        else:
+            self.flash_fft = None
+
+        self.blocks = nn.ModuleList(
+            get_block(config, i, flash_fft=self.flash_fft)
+            for i in range(config.num_layers)
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        inference_params_dict=None,
+        padding_mask: Optional[torch.Tensor] = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ):
+        x = self.embedding_layer.embed(x)
+
+        all_hidden_states: list[torch.Tensor] = []
+        all_attentions: list[Optional[torch.Tensor]] = []
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        if inference_params_dict is not None:
+            x, inference_params_dict_out = self._stateful_forward(
+                x, inference_params_dict,
+                all_hidden_states=all_hidden_states,
+                all_attentions=all_attentions,
+                output_hidden_states=output_hidden_states,
+                output_attentions=output_attentions,
+            )
+        else:
+            x, inference_params_dict_out = self._stateless_forward(
+                x, padding_mask=padding_mask,
+                all_hidden_states=all_hidden_states,
+                all_attentions=all_attentions,
+                output_hidden_states=output_hidden_states,
+                output_attentions=output_attentions,
+            )
+
+        if self.norm is not None:
+            x = self.norm(x)
+        if output_hidden_states:
+            all_hidden_states.append(x)
+
+        return x, inference_params_dict_out, all_hidden_states, all_attentions
+
+    def _stateful_forward(
+        self, x, inference_params_dict,
+        all_hidden_states, all_attentions,
+        output_hidden_states, output_attentions,
+    ):
+        for block_idx, block in enumerate(self.blocks):
+            block_name = (
+                "mha" if block_idx in self.config.attn_layer_idxs else "hyena"
+            )
+            inference_params = inference_params_dict[block_name]
+            x, attn = block(
+                x, inference_params=inference_params,
+                output_attentions=output_attentions,
+            )
+            if output_hidden_states:
+                all_hidden_states.append(x)
+            if output_attentions:
+                all_attentions.append(attn)
+        return x, inference_params_dict
+
+    def _stateless_forward(
+        self, x, padding_mask,
+        all_hidden_states, all_attentions,
+        output_hidden_states, output_attentions,
+    ):
+        if isinstance(padding_mask, torch.Tensor):
+            x = x * padding_mask[..., None]
+        for block in self.blocks:
+            x, attn = block(
+                x, inference_params=None, padding_mask=padding_mask,
+                output_attentions=output_attentions,
+            )
+            if output_hidden_states:
+                all_hidden_states.append(x)
+            if output_attentions:
+                all_attentions.append(attn)
+        return x, None
+
+    def initialize_inference_params(self, max_batch_size: int = 1) -> Evo1Cache:
+        return Evo1Cache(
+            max_seqlen=self.config.get("max_seqlen", 8192),
+            max_batch_size=max_batch_size,
+            short_filter_length=self.config.short_filter_length,
+            state_size=self.config.state_size,
+        )
+
+    def to_bfloat16_except_poles_residues(self):
+        """Cast all parameters to bfloat16 except Hyena poles/residues."""
+        for k, p in self.named_parameters():
+            if "poles" not in k and "residues" not in k:
+                p.data = p.data.to(torch.bfloat16)
+
+
+# =============================================================================
+# HuggingFace wrappers
+# =============================================================================
+
+
+class Evo1PreTrainedModel(PreTrainedModel):
+    config_class = Evo1Config
+    base_model_prefix = "backbone"
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["AttentionBlock", "ParallelGatedConvBlock"]
+    _skip_keys_device_placement = "past_key_values"
+    _keys_to_ignore_on_load_missing = [r"freq", r"\.t$"]
+    _keys_to_ignore_on_load_unexpected = [r"fftconv", r"twiddle_factors"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    # Hyena filter SSM parameters (poles / residues) MUST stay in fp32: they
+    # parametrize a long-range modal-form filter whose stability collapses
+    # in bf16. HF will keep these in fp32 even when the rest of the model is
+    # loaded in bf16 (or fp16) via the dtype= kwarg of from_pretrained.
+    _keep_in_fp32_modules = ["poles", "residues"]
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        # Evo1 was trained in bfloat16, with the modal-form filter parameters
+        # (Hyena poles / residues) kept in fp32 via _keep_in_fp32_modules.
+        # bf16 works correctly for all three attention backends (eager, sdpa,
+        # flash_attention_2). Default to bf16 so users don't have to pass it
+        # explicitly; this also silences HF's flash_attention_2 dtype warning
+        # (which inspects the model dtype before force_dtype() runs in __init__).
+        if "dtype" not in kwargs and "torch_dtype" not in kwargs:
+            kwargs["dtype"] = torch.bfloat16
+        return super().from_pretrained(*args, **kwargs)
+
+
+class Evo1Model(Evo1PreTrainedModel):
+    """Bare backbone: returns ``BaseModelOutputWithPast``.
+
+    ``last_hidden_state`` is the final (post-RMSNorm) representation, ready
+    to be fed into a downstream head or unembed projection.
+    """
+
+    def __init__(self, config: Evo1Config):
+        super().__init__(config)
+        self.backbone = StripedHyena(config)
+        self.config = config
+        self.post_init()
+        self.force_dtype()
+
+    def force_dtype(self):
+        # Cast everything except poles/residues to bf16 (the trained dtype).
+        # This runs at __init__ time so the model is usable even without an
+        # explicit ``dtype=torch.bfloat16`` kwarg to ``from_pretrained``.
+        self.backbone.to_bfloat16_except_poles_residues()
+
+    def get_input_embeddings(self):
+        return self.backbone.embedding_layer
+
+    def set_input_embeddings(self, value):
+        self.backbone.embedding_layer = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        # Evo1Model is the bare backbone (no LM head). Default to no caching:
+        # KV caches and Hyena recurrent state are only useful for autoregressive
+        # generation (Evo1ForCausalLM). For embedding extraction the caches
+        # have a large per-layer memory footprint with no benefit. The user
+        # can still opt-in by passing ``use_cache=True`` explicitly.
+        use_cache = use_cache if use_cache is not None else False
+        if use_cache and self.training:
+            use_cache = False
+
+        inputs = input_ids
+        if use_cache and past_key_values is None:
+            past_key_values = self.backbone.initialize_inference_params(
+                max_batch_size=input_ids.shape[0],
+            )
+
+        last_hidden, past_kv, hidden_states, attentions = self.backbone(
+            inputs,
+            padding_mask=attention_mask,
+            inference_params_dict=past_key_values if use_cache else None,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+
+        if not return_dict:
+            outputs = (last_hidden,)
+            if use_cache:
+                outputs += (past_kv,)
+            if output_hidden_states:
+                outputs += (tuple(hidden_states),)
+            if output_attentions:
+                outputs += (tuple(attentions),)
+            return outputs
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=last_hidden,
+            past_key_values=past_kv if use_cache else None,
+            hidden_states=tuple(hidden_states) if output_hidden_states else None,
+            attentions=tuple(attentions) if output_attentions else None,
+        )
+
+
+class Evo1ForCausalLM(Evo1PreTrainedModel, GenerationMixin):
+    """LM head wrapper. Tied to ``backbone.embedding_layer`` (Evo1 ties weights)."""
+
+    def __init__(self, config: Evo1Config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.backbone = StripedHyena(config)
+        self.config = config
+
+        # Pad-to-multiple-of for the vocab (matches togethercomputer config).
+        vocab_size = config.vocab_size
+        if vocab_size % config.make_vocab_size_divisible_by != 0:
+            vocab_size += config.make_vocab_size_divisible_by - (
+                vocab_size % config.make_vocab_size_divisible_by
+            )
+        self.vocab_size = vocab_size
+        self.post_init()
+        self.force_dtype()
+
+    def force_dtype(self):
+        self.backbone.to_bfloat16_except_poles_residues()
+
+    def get_input_embeddings(self):
+        return self.backbone.embedding_layer
+
+    def set_input_embeddings(self, value):
+        self.backbone.embedding_layer = value
+
+    def get_output_embeddings(self):
+        return self.backbone.embedding_layer
+
+    def set_output_embeddings(self, value):
+        self.backbone.embedding_layer = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        if use_cache and labels is not None:
+            logger.warning_once(
+                "use_cache=True is incompatible with loss computation; "
+                "disabling cache."
+            )
+            use_cache = False
+
+        inputs = input_ids
+        if use_cache:
+            # If the user (or HF generation) didn't pass our Evo1Cache,
+            # initialize a fresh one on the first call.
+            if not isinstance(past_key_values, Evo1Cache):
+                past_key_values = self.backbone.initialize_inference_params(
+                    max_batch_size=input_ids.shape[0],
+                )
+            else:
+                seqlen_offset = past_key_values.seqlen_offset
+                if seqlen_offset == 0:
+                    # Prefill done; set offset to prompt length minus the one
+                    # token we're about to consume (and that we'll keep
+                    # consuming one-at-a-time below).
+                    past_key_values.set_offset(input_ids.shape[-1] - 1)
+                else:
+                    past_key_values.advance(1)
+                inputs = input_ids[:, -1:]
+
+        last_hidden, past_kv, hidden_states, attentions = self.backbone(
+            inputs,
+            padding_mask=attention_mask,
+            inference_params_dict=past_key_values if use_cache else None,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+
+        # Tied unembed: matmul against embedding weights.
+        logits = last_hidden @ self.backbone.embedding_layer.weight.T
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = F.cross_entropy(shift_logits, shift_labels)
+
+        if not return_dict:
+            outputs = (logits,)
+            if use_cache:
+                outputs += (past_kv,)
+            if output_hidden_states:
+                outputs += (tuple(hidden_states),)
+            if output_attentions:
+                outputs += (tuple(attentions),)
+            if loss is not None:
+                outputs = (loss,) + outputs
+            return outputs
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_kv if use_cache else None,
+            hidden_states=tuple(hidden_states) if output_hidden_states else None,
+            attentions=tuple(attentions) if output_attentions else None,
+        )
+
+    @classmethod
+    def can_generate(cls) -> bool:
+        return True
+
+    def prepare_inputs_for_generation(
+        self, input_ids, attention_mask=None, past_key_values=None, **kwargs
+    ):
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "past_key_values": past_key_values,
+        }

rotary.py

@@ -0,0 +1,245 @@
+"""Rotary embeddings for the Evo1 HF port.
+
+Two modes:
+  * **Fast path** - when ``flash_attn`` is installed, we delegate to
+    ``flash_attn.layers.rotary.RotaryEmbedding``, whose Triton kernel does
+    the rotary multiply in fp32 internally (and is bit-exact with our
+    pure-PyTorch path below).
+  * **Fallback** - pure-PyTorch implementation, mathematically identical to
+    flash_attn's kernel (multiply done in fp32 then cast back to bf16). Used
+    when ``flash_attn`` isn't available.
+
+The ``LinearlyScaledRotaryEmbedding`` subclass (used for the 131k variant)
+overrides ``_update_cos_sin_cache`` to scale position indices, which works
+identically against either parent class.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+
+try:
+    from flash_attn.layers.rotary import RotaryEmbedding as _FlashRotaryEmbedding
+    _HAS_FLASH_ROTARY = True
+except ImportError:  # pragma: no cover - optional dep
+    _FlashRotaryEmbedding = None  # type: ignore[assignment]
+    _HAS_FLASH_ROTARY = False
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate the second half of the last dim into the first half (with sign).
+
+    [x1, x2] -> [-x2, x1]
+    """
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Apply non-interleaved RoPE to the last `2 * cos.shape[-1]` dims of x.
+
+    cos / sin shape: (T, rot_dim/2). x shape: (..., T, ..., D), where the rot
+    is applied along the last dim. We expand cos/sin to broadcast over the
+    leading dims.
+
+    The multiplication is performed in fp32 internally (then cast back to
+    x.dtype) to match flash_attn's Triton rotary kernel bit-exactly. Doing
+    the multiply in bf16 directly compounds rounding error of ~3e-2 per
+    layer, which becomes a ~1% relative error after 32 transformer blocks.
+    """
+    rot_dim = cos.shape[-1] * 2
+    x_rot = x[..., :rot_dim]
+    x_pass = x[..., rot_dim:]
+    orig_dtype = x.dtype
+    cos_full = torch.cat((cos, cos), dim=-1).float()
+    sin_full = torch.cat((sin, sin), dim=-1).float()
+    x_rot_f = x_rot.float()
+    rotated = (x_rot_f * cos_full) + (_rotate_half(x_rot_f) * sin_full)
+    rotated = rotated.to(orig_dtype)
+    return torch.cat((rotated, x_pass), dim=-1)
+
+
+class _PureRotaryEmbedding(nn.Module):
+    """Pure-PyTorch fallback RoPE (used when flash_attn is unavailable).
+
+    Mirrors the public surface of ``flash_attn.layers.rotary.RotaryEmbedding``
+    for the subset used by the Evo1 attention block: exposes ``inv_freq`` as
+    a buffer (so it serializes/deserializes the same way) and a forward(qkv)
+    -> qkv method that rotates Q and K.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base: float = 10000.0,
+        interleaved: bool = False,
+        scale_base: float | None = None,
+        pos_idx_in_fp32: bool = True,
+        device=None,
+    ):
+        super().__init__()
+        if interleaved:
+            raise NotImplementedError("Interleaved RoPE is not implemented.")
+        if scale_base is not None:
+            raise NotImplementedError("xPos scale_base is not implemented.")
+        self.dim = dim
+        self.base = float(base)
+        self.interleaved = interleaved
+        self.scale_base = scale_base
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+
+        inv_freq = self._compute_inv_freq(device=device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.scale = None  # xPos slot kept for swap_mha_rope compatibility
+
+        self._seq_len_cached = 0
+        self._cos_cached: torch.Tensor | None = None
+        self._sin_cached: torch.Tensor | None = None
+
+    def _compute_inv_freq(self, device=None) -> torch.Tensor:
+        return 1.0 / (
+            self.base
+            ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
+        )
+
+    def _update_cos_sin_cache(self, seqlen: int, device=None, dtype=None):
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                if self.inv_freq.dtype != torch.float32:
+                    inv_freq = self._compute_inv_freq(device=device)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                inv_freq = self.inv_freq
+
+            freqs = torch.outer(t, inv_freq)
+            self._cos_cached = torch.cos(freqs).to(dtype)
+            self._sin_cached = torch.sin(freqs).to(dtype)
+
+    def forward(
+        self,
+        qkv: torch.Tensor,
+        seqlen_offset: int | torch.Tensor = 0,
+        max_seqlen: int | None = None,
+    ) -> torch.Tensor:
+        """Rotate Q and K of a packed (B, T, 3, H, D) qkv tensor.
+
+        seqlen_offset is supported as int only (no per-sample offsets); for
+        the inference KV-cache fast path we fall back to int(seqlen_offset).
+        """
+        if isinstance(seqlen_offset, torch.Tensor):
+            seqlen_offset = int(seqlen_offset.max().item())
+        T = qkv.shape[1]
+        seqlen = max_seqlen if max_seqlen is not None else (T + seqlen_offset)
+        self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)
+
+        cos = self._cos_cached[seqlen_offset : seqlen_offset + T]
+        sin = self._sin_cached[seqlen_offset : seqlen_offset + T]
+        q, k, v = qkv.unbind(dim=2)
+        cos_b = cos[None, :, None, :]
+        sin_b = sin[None, :, None, :]
+        q = _apply_rotary(q, cos_b, sin_b)
+        k = _apply_rotary(k, cos_b, sin_b)
+        return torch.stack((q, k, v), dim=2)
+
+
+# Public ``RotaryEmbedding``: delegates to flash_attn's Triton kernel when
+# available, falls back to our pure-PyTorch implementation otherwise.
+RotaryEmbedding: type = (
+    _FlashRotaryEmbedding if _HAS_FLASH_ROTARY else _PureRotaryEmbedding
+)
+
+
+class LinearlyScaledRotaryEmbedding(RotaryEmbedding):
+    """RoPE with linear interpolation of position indices.
+
+    Used for evo-1-131k-base: positions are divided by ``scaling_factor``
+    before the cos/sin tables are computed, effectively stretching the
+    trained context. The override is the same shape regardless of whether
+    the parent class is flash_attn's RotaryEmbedding or our pure-PyTorch
+    fallback (both expose the same ``_update_cos_sin_cache`` hook).
+    """
+
+    def __init__(self, dim: int, scaling_factor: float = 1.0, **kwargs):
+        super().__init__(dim=dim, **kwargs)
+        self._linear_scaling_factor = float(scaling_factor)
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        # Mirrors the parent body but divides position indices by the linear
+        # scaling factor before computing the cos/sin tables.
+        if (
+            seqlen <= self._seq_len_cached
+            and self._cos_cached is not None
+            and self._cos_cached.device == device
+            and self._cos_cached.dtype == dtype
+            and not (self.training and self._cos_cached.is_inference())
+        ):
+            return
+
+        self._seq_len_cached = seqlen
+        if self.pos_idx_in_fp32:
+            t = torch.arange(seqlen, device=device, dtype=torch.float32)
+            t = t / self._linear_scaling_factor
+            if self.inv_freq.dtype != torch.float32:
+                inv_freq = self._compute_inv_freq(device=device) \
+                    if hasattr(self, "_compute_inv_freq") \
+                    else self.inv_freq.float()
+            else:
+                inv_freq = self.inv_freq
+        else:
+            t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+            t = t / self._linear_scaling_factor
+            inv_freq = self.inv_freq
+
+        freqs = torch.outer(t, inv_freq)
+        if self.scale is None:
+            self._cos_cached = torch.cos(freqs).to(dtype)
+            self._sin_cached = torch.sin(freqs).to(dtype)
+        else:  # pragma: no cover - xPos not used by Evo1
+            from einops import rearrange
+            power = (
+                torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device)
+                - seqlen // 2
+            ) / self.scale_base
+            scale = self.scale.to(device=power.device) ** rearrange(power, "s -> s 1")
+            self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+            self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+            self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+            self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+
+def swap_mha_rope(mha, new_rope=LinearlyScaledRotaryEmbedding, kwargs_new_rope=None):
+    """Replace ``mha.rotary_emb`` with a freshly-constructed scaled RoPE.
+
+    Mirrors ``stripedhyena.positional_embeddings.swap_mha_rope``: inherits
+    dim/base/interleaved/scale_base/pos_idx_in_fp32 from the existing rope,
+    deletes the old module, and attaches a new one of ``new_rope`` type
+    configured with ``kwargs_new_rope``.
+    """
+    weight_attr = "Wq" if getattr(mha, "cross_attn", False) else "Wqkv"
+    weight = getattr(mha, weight_attr).weight
+    dtype = weight.dtype
+    kwargs_old_rope = dict(
+        dim=mha.rotary_emb.dim,
+        base=mha.rotary_emb.base,
+        interleaved=mha.rotary_emb.interleaved,
+        scale_base=mha.rotary_emb.scale_base,
+        pos_idx_in_fp32=mha.rotary_emb.pos_idx_in_fp32,
+        device=mha.rotary_emb.inv_freq.device,
+    )
+    del mha.rotary_emb
+    kwargs_new_rope = kwargs_new_rope or {"scaling_factor": 1.0}
+    scaled = new_rope(**kwargs_new_rope, **kwargs_old_rope).to(dtype)
+    mha.rotary_emb = scaled
+    return mha

special_tokens_map.json

@@ -0,0 +1 @@
+{}

tokenization_evo1.py

@@ -0,0 +1,128 @@
+# Based on togethercomputer/evo-1-131k-base/tokenizer.py (Apache 2.0).
+# Adapted: minor fix (typing), Dict -> dict for forward-compatibility.
+"""ByteTokenizer for Evo1.
+
+Maps text to/from raw UTF-8 byte values, vocab_size = 512 (no real vocab,
+just the byte range padded out).
+"""
+from __future__ import annotations
+
+from os import PathLike
+from typing import List, Tuple
+
+import numpy as np
+
+from transformers.tokenization_utils import PreTrainedTokenizer
+from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
+from transformers.utils.generic import PaddingStrategy
+
+
+EMPTY: str = ""
+
+
+class ByteTokenizer(PreTrainedTokenizer):
+    """UTF-8 byte-level tokenizer (vocab_size = 512).
+
+    Special tokens follow the original ``CharLevelTokenizer`` convention:
+      * ``\x00`` (chr(0)) is end-of-document / end-of-sequence.
+      * ``\x01`` (chr(1)) is the padding token.
+
+    These are wired up so ``HF`` tokenization helpers (``padding=True``,
+    ``model_max_length``, etc.) work with this tokenizer.
+    """
+
+    def __init__(self, byte_level: bool = True, **kwargs):
+        # Defaults; ``tokenizer_config.json`` may override via kwargs.
+        # Padding token: byte 0x01 (matches the original CharLevelTokenizer
+        # pad convention). We deliberately do NOT set eos_token / cls_token:
+        # Evo1 is a pure byte-level model with no special tokens added at
+        # encoding time, so downstream pooling logic (e.g. mRNABench) should
+        # treat every non-pad position as a real token.
+        kwargs.setdefault("pad_token", chr(1))
+        super().__init__(byte_level=byte_level, **kwargs)
+        # The model only consumes input_ids and attention_mask (no segment ids).
+        self.model_input_names = ["input_ids", "attention_mask"]
+
+    @property
+    def vocab_size(self) -> int:
+        return 512
+
+    @property
+    def byte_level(self) -> bool:
+        return self.init_kwargs.get("byte_level", True)
+
+    def get_vocab(self) -> dict:
+        return {chr(i): i for i in range(self.vocab_size)}
+
+    def __len__(self) -> int:
+        return self.vocab_size
+
+    def clamp(self, n: int) -> int:
+        return max(32, min(n, self.vocab_size))
+
+    def _tokenize(self, text: str, **kwargs) -> List[str]:
+        return list(text)
+
+    def byte_tokenize(self, text: str) -> np.ndarray:
+        return np.frombuffer(text.encode("utf-8"), dtype=np.uint8)
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self.clamp(ord(token))
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return chr(self.clamp(index))
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return EMPTY.join(tokens)
+
+    def _decode(self, token_ids: List[int], **kwargs) -> str:
+        # numpy >= 2 disallows clipping a uint8 array to a value outside [0, 255]
+        # (vocab_size=512). Convert to int16 first, clip, then back to uint8.
+        indices = np.asarray(token_ids, dtype=np.int16)
+        indices = indices.clip(min=32, max=255).astype(np.uint8)
+        return indices.tobytes().decode("utf-8", errors="replace")
+
+    def _encode_plus(self, text: str, **kwargs) -> BatchEncoding:
+        first_ids = self.byte_tokenize(text).tolist()
+        return self.prepare_for_model(
+            first_ids,
+            pair_ids=None,
+            add_special_tokens=kwargs.get("add_special_tokens", False),
+            padding=kwargs.get("padding_strategy", PaddingStrategy.DO_NOT_PAD).value,
+            truncation=kwargs.get("truncation_strategy", TruncationStrategy.DO_NOT_TRUNCATE).value,
+            max_length=kwargs.get("max_length"),
+            stride=kwargs.get("stride", 0),
+            pad_to_multiple_of=kwargs.get("pad_to_multiple_of"),
+            return_tensors=kwargs.get("return_tensors"),
+            prepend_batch_axis=True,
+            return_attention_mask=kwargs.get("return_attention_mask"),
+            return_token_type_ids=kwargs.get("return_token_type_ids"),
+            return_overflowing_tokens=kwargs.get("return_overflowing_tokens", False),
+            return_special_tokens_mask=kwargs.get("return_special_tokens_mask", False),
+            return_length=kwargs.get("return_length", False),
+            verbose=kwargs.get("verbose", True),
+        )
+
+    def _batch_encode_plus(self, batch_text_or_text_pairs, **kwargs) -> BatchEncoding:
+        input_ids = [(self.byte_tokenize(t).tolist(), None) for t in batch_text_or_text_pairs]
+        return self._batch_prepare_for_model(
+            input_ids,
+            add_special_tokens=kwargs.get("add_special_tokens", False),
+            padding_strategy=kwargs.get("padding_strategy", PaddingStrategy.DO_NOT_PAD),
+            truncation_strategy=kwargs.get("truncation_strategy", TruncationStrategy.DO_NOT_TRUNCATE),
+            max_length=kwargs.get("max_length"),
+            stride=kwargs.get("stride", 0),
+            pad_to_multiple_of=kwargs.get("pad_to_multiple_of"),
+            return_attention_mask=kwargs.get("return_attention_mask"),
+            return_token_type_ids=kwargs.get("return_token_type_ids"),
+            return_overflowing_tokens=kwargs.get("return_overflowing_tokens", False),
+            return_special_tokens_mask=kwargs.get("return_special_tokens_mask", False),
+            return_length=kwargs.get("return_length", False),
+            return_tensors=kwargs.get("return_tensors"),
+            verbose=kwargs.get("verbose", True),
+        )
+
+    def _save_pretrained(
+        self, save_directory: str | PathLike, file_names: Tuple[str], **kwargs
+    ) -> Tuple[str]:
+        return file_names

tokenizer_config.json

@@ -0,0 +1,12 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_evo1.ByteTokenizer",
+      null
+    ]
+  },
+  "tokenizer_class": "ByteTokenizer",
+  "model_max_length": 8192,
+  "pad_token": "\u0001",
+  "byte_level": true
+}