Add model_max_length=8192 and max_seqlen=8192 (training context); note StripedHyena extension in README

.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,247 @@
+---
+language:
+- dna
+library_name: transformers
+tags:
+- DNA
+- language-model
+- StripedHyena2
+- Evo2
+license: apache-2.0
+---
+
+# Evo2-7B-8K
+
+A clean, minimal HuggingFace port of [Evo 2 7B base](https://huggingface.co/arcinstitute/evo2_7b_base), the 8K-context StripedHyena2 DNA foundation model. Provides native support for layer-by-layer hidden state extraction, attention-weight extraction, and a runtime-switchable attention backend.
+
+> **No Transformer Engine required.** This variant runs in pure bf16 on any
+> CUDA-capable GPU (no FP8 / TE dependency).
+
+
+
+## Why this port?
+
+[arcinstitute/evo2_7b_base](https://huggingface.co/arcinstitute/evo2_7b_base) ships a `.pt` checkpoint that requires the [evo2](https://github.com/ArcInstitute/evo2) and [vortex](https://github.com/Zymrael/vortex) Python packages just to instantiate the model. Even with both installed, common pain points remain:
+
+1. **Not a HuggingFace model.** No `from_pretrained`, no `AutoModel`, no `AutoModelForCausalLM` - the original ships a thin Python wrapper around a custom `nn.Module`.
+2. **No way to extract attention weights.** The reference uses flash-attn unconditionally and discards the `(B, H, T, T)` attention matrix; there is no official path to read it back.
+3. **`evo2` + `vortex` packages mandatory** even for inference.
+
+This repo fixes all three. The math is **bit-exact** with the vortex reference (`max_abs_diff = 0.000e+00` at every layer; see [Parity Verification](#parity-verification)). Loads with `from_pretrained` and `trust_remote_code=True` - no `evo2` / `vortex` install needed.
+
+## Architecture
+
+| Parameter | Value |
+|---|---|
+| Total parameters | ~6.5B |
+| Architecture | StripedHyena 2 (interleaved Hyena cascade + MHA blocks) |
+| Layers | 32 |
+| Attention heads | 32 |
+| Embedding dimension | 4096 |
+| Inner MLP size | 11 008 |
+| Vocabulary size | 512 (UTF-8 byte-level) |
+| Attention block indices | 3, 10, 17, 24, 31 (5 blocks total) |
+| Hyena block indices | all others (27 blocks: hcs / hcm / hcl pattern) |
+| Positional encoding | RoPE (base = 10 000) |
+| Max sequence length | 8 192 (training context; StripedHyena supports extension beyond this with adjusted FFT buffers) |
+| Training dtype | bfloat16 (Hyena modal-form `log_poles` / `residues` and rotary `inv_freq` kept in fp32) |
+| FP8 input projections | no |
+| Weight format | `model.safetensors` (6.5B params, 3 files) |
+
+## Pretraining
+
+- **Objective:** causal byte-level next-token prediction.
+- **Data:** [OpenGenome2](https://huggingface.co/datasets/arcinstitute/opengenome2), 8.8 trillion tokens spanning all domains of life.
+- **Source checkpoint:** `arcinstitute/evo2_7b_base` ([evo2_7b_base.pt](https://huggingface.co/arcinstitute/evo2_7b_base/blob/main/evo2_7b_base.pt)).
+
+## Parity Verification
+
+Hidden-state representations verified **bit-exact** (`max_abs_diff = 0.000e+00`) to the vortex reference at every block output, using `attn_implementation="sdpa"` in bf16 (the same backend vortex's `SelfAttention` calls when `use_flash_attn=False`). Logits from `Evo2ForCausalLM` were also verified bit-exact (top-1 agreement: 128/128 positions on a 128-byte ACGT input). Verified on H100 with PyTorch 2.7 / CUDA 12.
+
+Two non-obvious correctness fixes were required versus a naive port (see Implementation Notes for details):
+
+1. **`inv_freq` recomputation.** `from_pretrained(dtype=bf16)` casts buffers to bf16, which loses ~7 bits of mantissa in the rotary `inv_freq = 1 / base^(2i/dim)`. Our `to_bfloat16_except_poles_residues()` recomputes `inv_freq` in fp32 from `self.base` to match vortex's `to_bfloat16_except_pr_lc(to_float32=True)`.
+2. **SDPA backend used for parity.** Vortex's reference `SelfAttention` (`use_flash_attn=False`) calls `F.scaled_dot_product_attention`, not a textbook softmax loop. Parity is measured with `attn_implementation="sdpa"` on our side. Using `"eager"` (textbook einsum + softmax) is mathematically equivalent but not bit-exact in bf16; using `"flash_attention_2"` (the recommended runtime backend) is also not bit-exact but agrees within bf16 noise.
+
+## Related Models
+
+See the full Evo 2 collection on the [Arc Institute HF org](https://huggingface.co/arcinstitute) for the original weights, or the [`Taykhoom/Evo2-*` collection](https://huggingface.co/collections/Taykhoom/evo2-6a24b06c05955a295025a006) for our minimal HF ports.
+
+| Model | Size | Context | Notes |
+|---|---|---|---|
+| [Taykhoom/Evo2-1B-8K](https://huggingface.co/Taykhoom/Evo2-1B-8K) | 1B | 8 192 |  |
+| **[Taykhoom/Evo2-7B-8K](https://huggingface.co/Taykhoom/Evo2-7B-8K)** | 7B | 8 192 | <- this model |
+| [Taykhoom/Evo2-7B-262K](https://huggingface.co/Taykhoom/Evo2-7B-262K) | 7B | 262 144 |  |
+| [Taykhoom/Evo2-7B-1M](https://huggingface.co/Taykhoom/Evo2-7B-1M) | 7B | 1 048 576 |  |
+| [Taykhoom/Evo2-20B-1M](https://huggingface.co/Taykhoom/Evo2-20B-1M) | 20B | 1 048 576 |  |
+| [Taykhoom/Evo2-40B-8K](https://huggingface.co/Taykhoom/Evo2-40B-8K) | 40B | 8 192 |  |
+| [Taykhoom/Evo2-40B-1M](https://huggingface.co/Taykhoom/Evo2-40B-1M) | 40B | 1 048 576 |  |
+
+## Usage
+
+> **Note on dtype.** Evo 2 was trained in bfloat16, with the Hyena `log_poles` / `residues` (modal-form filter parameters) and the rotary `inv_freq` kept in fp32 for numerical stability. Passing `dtype=...` to `from_pretrained` only affects the initial load precision - `Evo2Model.__init__` and `Evo2ForCausalLM.__init__` always call `to_bfloat16_except_poles_residues()`, so the model runs in bf16 with these fp32 invariants regardless. This is intentional: bf16 is the trained precision, and the fp32 islands are required for stability.
+
+> **Note on attention backend.** By HuggingFace convention this model defaults to `attn_implementation="sdpa"` (`F.scaled_dot_product_attention`) since SDPA needs only `torch` and runs on any GPU. The original Arc Institute Evo 2 inference path uses **flash_attention_2**, which is faster on long sequences but requires a separate `flash-attn` install. All usage examples below opt in to `flash_attention_2` explicitly because most real users will want it. Drop the kwarg (or pass `"sdpa"` / `"eager"`) if you don't have `flash-attn` installed.
+
+### Embedding generation (no LM head)
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-7B-8K", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "Taykhoom/Evo2-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",  # or "sdpa" (default) or "eager"
+).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
+middle_layer = all_layers[16]          # input to block 16 (= output of block 15)
+```
+
+### Recommended embedding: pre-norm of the middle block
+
+The Evo 2 paper reports that intermediate representations work better than the final layer for downstream tasks - specifically the **pre-norm of a middle block**. For this variant the middle block is `blocks[16]`, so the recommended embedding is `blocks[16].pre_norm(hidden_states[16])`:
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-7B-8K", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "Taykhoom/Evo2-7B-8K",
+    trust_remote_code=True,
+    attn_implementation="flash_attention_2",
+).cuda().eval()
+
+inputs = tokenizer(["ACGTACGTACGT"], return_tensors="pt").to(model.device)
+with torch.no_grad():
+    out = model(**inputs, output_hidden_states=True)
+    pre_norm_middle = model.backbone.blocks[16].pre_norm(
+        out.hidden_states[16]
+    )                                      # (B, T, 4096)
+```
+
+HF has no built-in API for sub-block intermediates like pre-norm outputs (only block outputs via `output_hidden_states`). The pattern above applies the block's `pre_norm` submodule directly to the corresponding `hidden_states` entry; this gives a bit-identical result to registering a forward hook on `backbone.blocks[i].pre_norm` and is simpler than using PyTorch hooks. Note that it does require running the full forward pass and then re-applying `pre_norm`, so a forward hook is more efficient if you only need this single intermediate.
+
+### LM logits
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-7B-8K", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "Taykhoom/Evo2-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/Evo2-7B-8K", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "Taykhoom/Evo2-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.
+
+### Attention weights
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+tokenizer = AutoTokenizer.from_pretrained("Taykhoom/Evo2-7B-8K", trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    "Taykhoom/Evo2-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
+# [3, 10, 17, 24, 31] are None (Hyena blocks have no attention matrix).
+# The 5 attention block(s) at those indices return a (B, num_heads, T, T) tensor.
+attn_block_3 = out.attentions[3]
+```
+
+### Multi-GPU loading (optional)
+
+For sharding across multiple GPUs (required for 40B, optional for smaller
+variants), install [accelerate](https://github.com/huggingface/accelerate)
+and pass `device_map="auto"`:
+
+```python
+from transformers import AutoModelForCausalLM
+# pip install accelerate
+model = AutoModelForCausalLM.from_pretrained(
+    "Taykhoom/Evo2-7B-8K", trust_remote_code=True,
+    device_map="auto",  # accelerate will shard across all visible GPUs
+)
+```
+
+### Fine-tuning
+
+This HuggingFace port has **not** been tested for fine-tuning - it's verified only for inference parity. For fine-tuning, follow the original Arc Institute guidance and use either [Savanna](https://github.com/Zymrael/savanna) (the framework Evo 2 was pretrained in) or [Nvidia BioNeMo](https://github.com/NVIDIA/bionemo-framework), which provides an official [Evo 2 fine-tuning recipe](https://github.com/NVIDIA-BioNeMo/bionemo-framework/tree/main/bionemo-recipes/recipes/evo2_megatron#fine-tuning-from-an-existing-checkpoint).
+
+## Implementation Notes
+
+- **`inv_freq` kept in fp32 (critical for parity).** HF's `from_pretrained(dtype=bf16)` casts all buffers, including the rotary `inv_freq`, to bf16. The geometric series `inv_freq[i] = 1 / base^(2i/dim)` loses ~7 bits of mantissa in bf16, which shifts the cos/sin tables by ~5e-2 per cell at higher positions and contributes ~3e-2 of Q/K noise per attention layer. Our `to_bfloat16_except_poles_residues()` (called in each `__init__`) recomputes `inv_freq` in fp32 from the stored `self.base` and invalidates the cos/sin cache, mirroring vortex's `to_bfloat16_except_pr_lc(to_float32=True)`.
+- **`log_poles` / `residues` kept in fp32 (critical for stability).** The Hyena cascade long (hcl) blocks parameterize an IIR filter via `log_poles` and `residues`; bf16 quantisation makes the recurrence numerically unstable. Both are stored as fp32 in the safetensors and restored to fp32 by `force_dtype()` after load.
+- **`attn_implementation` switching (`attention.py`).** Three backends, selected via the standard HF `attn_implementation` kwarg to `from_pretrained` (default chosen by HF auto-detection - typically `"sdpa"`):
+  - `"sdpa"`: calls `F.scaled_dot_product_attention`. Bit-exact with vortex's reference path (when vortex uses `use_flash_attn=False`).
+  - `"flash_attention_2"`: calls `flash_attn.flash_attn_qkvpacked_func`. Matches the original Arc Institute inference path; faster on long sequences; requires `flash-attn` installed.
+  - `"eager"`: textbook einsum + softmax(QK^T) + einsum. Slowest, used internally when `output_attentions=True` so the attention matrix is materialized.
+- **Block dispatch (`hyena.py`).** StripedHyena 2 has 4 block types, dispatched by `layer_idx` membership in four config lists: `attn_layer_idxs` (MHA + RoPE), `hcl_layer_idxs` (modal-form IIR via FFT), `hcm_layer_idxs` (medium FIR cascade, inner length 128), `hcs_layer_idxs` (short FIR cascade, inner length 7). The disjoint union must equal `range(num_layers)`.
+- **`TELinear` with pure-PyTorch fallback (`layers.py`).** Hyena cascade blocks use a TransformerEngine-backed input projection (3x hidden_size output) that supports FP8 quantisation. When TE is not installed, a `TELinear` fallback class with the same state_dict layout (`weight`, `bias`) is used - checkpoints are cross-loadable.
+- **Custom cache (`cache.py`).** `Evo2Cache` wraps four block-type-specific dataclasses: `InferenceParams` for MHA KV cache, `HyenaCascadeIIRInferenceParams` for hcl, and two `HyenaCascadeFIRInferenceParams` for hcm / hcs. Passed through `model.generate()` as `past_key_values` (we set `_supports_cache_class = False` so HF treats it as an opaque dict rather than wrapping it in a `DynamicCache`).
+- **Tokenizer (`tokenization_evo2.py`).** Byte-level UTF-8, vocab_size = 512. Pad token = byte `\x01`. EOS = byte `\x00` (set as `eos_token_id` in `generation_config.json`). Tokenizer does not add EOS at encoding time - matches the original Evo 2 inference pipeline.
+- **Dependencies.** `torch`, `transformers`, `numpy`, `safetensors`, `huggingface_hub`. `accelerate` is optional but recommended if you want to load with `device_map="auto"` for multi-GPU sharding. `flash_attn` is optional (only needed if you pass `attn_implementation="flash_attention_2"`).
+
+## Citation
+
+```bibtex
+@article{brixi2026_evo2,
+  title   = {Genome modelling and design across all domains of life with {Evo} 2},
+  author  = {Brixi, Garyk and Durrant, Matthew G. and Ku, Jerome and Naghipourfar, Mohsen and Poli, Michael and Sun, Gwanggyu and Brockman, Greg and Chang, Daniel and Fanton, Alison and Gonzalez, Gabriel A. and King, Samuel H. and Li, David B. and Merchant, Aditi T. and Nguyen, Eric and Ricci-Tam, Chiara and Romero, David W. and Schmok, Jonathan C. and Taghibakhshi, Ali and Vorontsov, Anton and Yang, Brandon and Deng, Myra and Gorton, Liv and Nguyen, Nam and Wang, Nicholas K. and Pearce, Michael T. and Simon, Elana and Adams, Etowah and Amador, Zachary J. and Ashley, Euan A. and Baccus, Stephen A. and Dai, Haoyu and Dillmann, Steven and Ermon, Stefano and Guo, Daniel and Herschl, Michael H. and Ilango, Rajesh and Janik, Ken and Lu, Amy X. and Mehta, Reshma and Mofrad, Mohammad R. K. and Ng, Madelena Y. and Pannu, Jaspreet and {R{\'e}}, Christopher and St. John, John and Sullivan, Jeremy and Tey, Joseph and Viggiano, Ben and Zhu, Kevin and Zynda, Greg and Balsam, Daniel and Collison, Patrick and Costa, Anthony B. and Hernandez-Boussard, Tina and Ho, Eric and Liu, Ming-Yu and McGrath, Thomas and Powell, Kimberly and Pinglay, Sudarshan and Burke, Dave P. and Goodarzi, Hani and Hsu, Patrick D. and Hie, Brian L.},
+  journal = {Nature},
+  year    = {2026},
+  doi     = {10.1038/s41586-026-10176-5}
+}
+```
+
+## Credits
+
+Original Evo 2 model and code by Brixi et al. ([arcinstitute/evo2](https://github.com/arcinstitute/evo2), [Zymrael/vortex](https://github.com/Zymrael/vortex)). Source checkpoint: [arcinstitute/evo2_7b_base](https://huggingface.co/arcinstitute/evo2_7b_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 Evo 2 release.

attention.py

@@ -0,0 +1,261 @@
+"""Multi-head attention for Evo2.
+
+Causal self-attention with RoPE. Three backends dispatchable via the standard
+HF attn_implementation kwarg:
+
+  * "eager"              -- pure F.softmax(QK^T)V; returns weights on demand
+  * "sdpa"               -- F.scaled_dot_product_attention; no weights
+  * "flash_attention_2"  -- flash_attn.flash_attn_qkvpacked_func; no weights
+
+Constructor signature is a strict subset of vortex's MHA so the
+AttentionBlock instantiation site only changes the attn_implementation kwarg.
+KV cache (used by Evo2 generation) goes through an SDPA path that handles
+arbitrary sequence lengths.
+"""
+
+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_qkvpacked_func  # noqa: F401
+    except ImportError as exc:
+        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:
+    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."""
+
+    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 compat
+        attn_implementation: str = "eager",
+        device=None,
+        dtype=None,
+    ) -> None:
+        super().__init__()
+        if cross_attn:
+            raise NotImplementedError("Cross-attention is not supported.")
+
+        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
+        self.cross_attn = cross_attn
+
+        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:
+        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)
+        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
+
+    def _forward_eager(
+        self,
+        qkv: torch.Tensor,
+        output_attentions: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Bit-identical port of flash_attn.modules.mha.SelfAttention.forward.
+
+        Key choices that must match vortex (which delegates to this class
+        when use_flash_attn=False):
+
+          * einsum (not matmul) for QK^T and PV: bf16 matmul vs bf16 einsum
+            use different accumulator orders, ~3e-2 per attention layer.
+          * -10000.0 causal mask (not -inf): vortex uses a finite mask value.
+          * softmax(scores, dim=-1, dtype=v.dtype) (not F.softmax(.float())):
+            PyTorch promotes to fp32 internally then casts to v.dtype.
+            Explicit .float() then .to(v.dtype) cast rounds differently.
+          * scale applied to K (k * scale), not Q. Same math, different
+            bf16 multiply target.
+        """
+        q, k, v = qkv.unbind(dim=2)  # each (B, T, H, D)
+        softmax_scale = (
+            self.softmax_scale if self.softmax_scale is not None
+            else 1.0 / math.sqrt(self.head_dim)
+        )
+        scores = torch.einsum("bthd,bshd->bhts", q, k * softmax_scale)
+        T = q.shape[1]
+        if self.causal:
+            causal_mask = torch.triu(
+                torch.full((T, T), -10000.0, device=scores.device), 1
+            )
+            scores = scores + causal_mask.to(dtype=scores.dtype)
+        attention = torch.softmax(scores, dim=-1, dtype=v.dtype)
+        if self.training and self.dropout_p > 0:
+            attention = F.dropout(attention, p=self.dropout_p)
+        out = torch.einsum("bhts,bshd->bthd", attention, v)
+        return out, (attention if output_attentions else None)
+
+    def _forward_sdpa(self, qkv: torch.Tensor) -> torch.Tensor:
+        q, k, v = qkv.unbind(dim=2)
+        q = q.permute(0, 2, 1, 3)
+        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)
+
+    def _forward_flash(self, qkv: torch.Tensor) -> torch.Tensor:
+        _flash_attn_required()
+        from flash_attn import flash_attn_qkvpacked_func
+        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
+
+    def _forward_with_cache(
+        self,
+        qkv: torch.Tensor,
+        inference_params,
+    ) -> torch.Tensor:
+        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)
+        kv = _update_kv_cache(kv, inference_params, self.layer_idx)
+        k_full, v_full = kv.unbind(dim=2)
+
+        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)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        inference_params=None,
+        output_attentions: bool = False,
+        **_unused,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if self.num_heads_kv != self.num_heads:
+            raise NotImplementedError("GQA is not exercised by Evo2.")
+
+        qkv = self._project_qkv(x)
+
+        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":
+                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 = 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,141 @@
+"""Inference-time caches for Evo2 blocks.
+
+StripedHyena2 has four block types with different caching needs:
+
+  * `attn` blocks   -> InferenceParams (standard KV cache)
+  * `hcl` blocks    -> HyenaCascadeIIRInferenceParams (FIR window + IIR state)
+  * `hcm` blocks    -> HyenaCascadeFIRInferenceParams (outer FIR + inner FIR)
+  * `hcs` blocks    -> HyenaCascadeFIRInferenceParams (outer FIR + inner FIR)
+
+Layer outputs of these caches are wrapped together inside an HF Cache subclass
+(`Evo2Cache`) so model.generate() can drive autoregressive decoding without
+the user having to instantiate four separate caches by hand.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class InferenceParams:
+    """Standard KV cache for attention blocks."""
+
+    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 HyenaCascadeIIRInferenceParams:
+    """Cache for `hcl` blocks: short FIR window + IIR modal state."""
+
+    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.seqlen_offset = 0
+
+
+@dataclass
+class HyenaCascadeFIRInferenceParams:
+    """Cache for `hcm` and `hcs` blocks: outer short FIR + inner FIR cascade."""
+
+    fir_filter_length: int = 3
+    fir_inner_filter_length: int = 4
+    seqlen_offset: int = 0
+    fir_inner_state_dict: dict = field(default_factory=dict)
+    fir_state_dict: dict = field(default_factory=dict)
+    state_dict: dict = field(default_factory=dict)
+
+    def reset(self):
+        self.seqlen_offset = 0
+
+
+class Evo2Cache:
+    """Container for per-block-type inference params.
+
+    Not a transformers.Cache subclass (the new Cache API requires per-layer
+    dataclasses, which doesn't fit StripedHyena 2's 4 block-type-specific
+    state structures). Instead we set Evo2PreTrainedModel._supports_cache_class
+    = False so HF's generate() treats this as an opaque past_key_values dict.
+    """
+
+    is_compileable = False
+
+    def __init__(
+        self,
+        max_seqlen: int,
+        max_batch_size: int,
+        short_filter_length: int,
+        hcm_filter_length: int,
+        hcs_filter_length: int,
+        state_size: int,
+    ):
+        self.mha = InferenceParams(
+            max_seqlen=max_seqlen,
+            max_batch_size=max_batch_size,
+        )
+        self.hcl = HyenaCascadeIIRInferenceParams(
+            fir_filter_length=short_filter_length,
+            state_dim=state_size,
+        )
+        self.hcm = HyenaCascadeFIRInferenceParams(
+            fir_filter_length=short_filter_length,
+            fir_inner_filter_length=hcm_filter_length,
+        )
+        self.hcs = HyenaCascadeFIRInferenceParams(
+            fir_filter_length=short_filter_length,
+            fir_inner_filter_length=hcs_filter_length,
+        )
+
+    @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:
+        return self.mha.max_seqlen
+
+    def advance(self, n: int = 1) -> None:
+        self.mha.seqlen_offset += n
+        self.hcl.seqlen_offset += n
+        self.hcm.seqlen_offset += n
+        self.hcs.seqlen_offset += n
+
+    def set_offset(self, offset: int) -> None:
+        self.mha.seqlen_offset = offset
+        self.hcl.seqlen_offset = offset
+        self.hcm.seqlen_offset = offset
+        self.hcs.seqlen_offset = offset
+
+    def reset(self) -> None:
+        self.mha.reset(self.mha.max_seqlen, self.mha.max_batch_size)
+        self.hcl.reset()
+        self.hcm.reset()
+        self.hcs.reset()
+
+    def by_block_name(self, name: str):
+        return getattr(self, name)

config.json

@@ -0,0 +1,108 @@
+{
+  "architectures": [
+    "Evo2ForCausalLM"
+  ],
+  "attn_block_dtype": "bfloat16",
+  "attn_layer_idxs": [
+    3,
+    10,
+    17,
+    24,
+    31
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_evo2.Evo2Config",
+    "AutoModel": "modeling_evo2.Evo2Model",
+    "AutoModelForCausalLM": "modeling_evo2.Evo2ForCausalLM"
+  },
+  "column_split": true,
+  "column_split_hyena": false,
+  "dtype": "bfloat16",
+  "eps": 1e-06,
+  "evo2_style_activations": true,
+  "final_norm": true,
+  "hcl_filter_groups": 4096,
+  "hcl_layer_idxs": [
+    2,
+    6,
+    9,
+    13,
+    16,
+    20,
+    23,
+    27,
+    30
+  ],
+  "hcm_filter_groups": 256,
+  "hcm_filter_length": 128,
+  "hcm_layer_idxs": [
+    1,
+    5,
+    8,
+    12,
+    15,
+    19,
+    22,
+    26,
+    29
+  ],
+  "hcs_filter_groups": 256,
+  "hcs_filter_length": 7,
+  "hcs_layer_idxs": [
+    0,
+    4,
+    7,
+    11,
+    14,
+    18,
+    21,
+    25,
+    28
+  ],
+  "hidden_size": 4096,
+  "hyena_block_dtype": "bfloat16",
+  "hyena_filter_groups": 1,
+  "hyena_flip_x1x2": false,
+  "hyena_out_proj_bias": true,
+  "inference_mode": true,
+  "inner_mlp_size": 11008,
+  "inner_size_multiple_of": 16,
+  "interleave": true,
+  "log_intermediate_values": false,
+  "make_vocab_size_divisible_by": 8,
+  "max_batch_size": 1,
+  "max_seqlen": 8192,
+  "mha_out_proj_bias": true,
+  "mlp_activation": "gelu",
+  "mlp_dtype": "bfloat16",
+  "mlp_init_method": "torch.nn.init.zeros_",
+  "mlp_output_init_method": "torch.nn.init.zeros_",
+  "model_parallel_size": 1,
+  "model_type": "evo2",
+  "num_attention_heads": 32,
+  "num_filters": 4096,
+  "num_layers": 32,
+  "pipe_parallel_size": 1,
+  "prefill_style": "fft",
+  "print_activations": false,
+  "proj_groups": 1,
+  "qkv_proj_bias": false,
+  "rotary_emb_base": 10000,
+  "rotary_emb_scaling_factor": 1.0,
+  "short_filter_bias": false,
+  "short_filter_length": 3,
+  "smeared_gqa": false,
+  "state_size": 16,
+  "tie_embeddings": true,
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "use_flash_attn": false,
+  "use_flash_depthwise": false,
+  "use_flash_rmsnorm": false,
+  "use_flashfft": false,
+  "use_fp8_input_projections": false,
+  "use_interpolated_rotary_pos_emb": false,
+  "use_laughing_hyena": false,
+  "vocab_size": 512,
+  "model_max_length": 8192
+}

configuration_evo2.py

@@ -0,0 +1,197 @@
+"""Configuration for Evo2 (StripedHyena2)."""
+
+from __future__ import annotations
+
+import json
+from typing import List
+
+from transformers import PretrainedConfig
+
+
+class Evo2Config(PretrainedConfig):
+    """Evo2 config.
+
+    Defaults match evo2_7b_base (32-layer, 4096-hidden). Per-variant overrides
+    are written into the config.json of each repo by convert_checkpoint.py.
+
+    Block dispatch is driven by four index lists:
+      attn_layer_idxs : transformer (MHA + RoPE) blocks
+      hcl_layer_idxs  : Hyena Cascade Long  (IIR via log_poles + residues)
+      hcm_layer_idxs  : Hyena Cascade Medium (FIR, fir_inner_filter_length=128)
+      hcs_layer_idxs  : Hyena Cascade Short  (FIR, fir_inner_filter_length=7)
+    Their disjoint union must equal range(num_layers).
+    """
+
+    model_type = "evo2"
+
+    def __init__(
+        self,
+        # Architecture core
+        vocab_size: int = 512,
+        hidden_size: int = 4096,
+        num_filters: int = 4096,
+        inner_mlp_size: int = 11008,
+        num_layers: int = 32,
+        num_attention_heads: int = 32,
+        # Block dispatch
+        attn_layer_idxs: List[int] | None = None,
+        hcl_layer_idxs: List[int] | None = None,
+        hcm_layer_idxs: List[int] | None = None,
+        hcs_layer_idxs: List[int] | None = None,
+        # Filter geometry
+        hcm_filter_length: int = 128,
+        hcs_filter_length: int = 7,
+        hcl_filter_groups: int = 4096,
+        hcm_filter_groups: int = 256,
+        hcs_filter_groups: int = 256,
+        hyena_filter_groups: int = 1,
+        short_filter_length: int = 3,
+        short_filter_bias: bool = False,
+        state_size: int = 16,
+        # Channel-split conventions
+        column_split: bool = True,
+        column_split_hyena: bool = False,
+        interleave: bool = True,
+        hyena_flip_x1x2: bool = False,
+        # Norms / activations
+        eps: float = 1e-6,
+        final_norm: bool = True,
+        mlp_activation: str = "gelu",
+        evo2_style_activations: bool = True,
+        # Linear biases
+        mha_out_proj_bias: bool = True,
+        hyena_out_proj_bias: bool = True,
+        qkv_proj_bias: bool = False,
+        # MLP geometry
+        inner_size_multiple_of: int = 16,
+        make_vocab_size_divisible_by: int = 8,
+        # Embeddings
+        tie_embeddings: bool = True,
+        # Rotary / sequence
+        max_seqlen: int = 8192,
+        max_batch_size: int = 1,
+        rotary_emb_base: float = 10000.0,
+        use_interpolated_rotary_pos_emb: bool = False,
+        rotary_emb_scaling_factor: float = 1.0,
+        # Inference engine
+        prefill_style: str = "fft",
+        inference_mode: bool = True,
+        # Projection precision (TE FP8 - required for 1b / 20b / 40b)
+        use_fp8_input_projections: bool = False,
+        # Backend toggles
+        use_cache: bool = True,
+        # Vortex used this flag at runtime. Our HF port ignores it entirely
+        # (attention dispatch is driven by ``config._attn_implementation``,
+        # the standard HF mechanism). Default is False here so the config
+        # reflects what actually runs: SDPA by default unless the user passes
+        # ``attn_implementation="flash_attention_2"`` to from_pretrained.
+        use_flash_attn: bool = False,
+        use_flash_rmsnorm: bool = False,
+        use_flash_depthwise: bool = False,
+        use_flashfft: bool = False,
+        # Per-block dtypes (cast at module init)
+        attn_block_dtype: str = "bfloat16",
+        hyena_block_dtype: str = "bfloat16",
+        mlp_dtype: str = "bfloat16",
+        # Multi-tensor parallel knobs (kept for ParallelGatedMLP._compute_inner_size)
+        model_parallel_size: int = 1,
+        pipe_parallel_size: int = 1,
+        # GQA
+        proj_groups: int = 1,
+        smeared_gqa: bool = False,
+        **kwargs,
+    ):
+        if attn_layer_idxs is None:
+            attn_layer_idxs = [3, 10, 17, 24, 31]
+        if hcl_layer_idxs is None:
+            hcl_layer_idxs = [2, 6, 9, 13, 16, 20, 23, 27, 30]
+        if hcm_layer_idxs is None:
+            hcm_layer_idxs = [1, 5, 8, 12, 15, 19, 22, 26, 29]
+        if hcs_layer_idxs is None:
+            hcs_layer_idxs = [0, 4, 7, 11, 14, 18, 21, 25, 28]
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_filters = num_filters
+        self.inner_mlp_size = inner_mlp_size
+        self.num_layers = num_layers
+        self.num_attention_heads = num_attention_heads
+
+        self.attn_layer_idxs = attn_layer_idxs
+        self.hcl_layer_idxs = hcl_layer_idxs
+        self.hcm_layer_idxs = hcm_layer_idxs
+        self.hcs_layer_idxs = hcs_layer_idxs
+
+        self.hcm_filter_length = hcm_filter_length
+        self.hcs_filter_length = hcs_filter_length
+        self.hcl_filter_groups = hcl_filter_groups
+        self.hcm_filter_groups = hcm_filter_groups
+        self.hcs_filter_groups = hcs_filter_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.interleave = interleave
+        self.hyena_flip_x1x2 = hyena_flip_x1x2
+
+        self.eps = eps
+        self.final_norm = final_norm
+        self.mlp_activation = mlp_activation
+        self.evo2_style_activations = evo2_style_activations
+
+        self.mha_out_proj_bias = mha_out_proj_bias
+        self.hyena_out_proj_bias = hyena_out_proj_bias
+        self.qkv_proj_bias = qkv_proj_bias
+
+        self.inner_size_multiple_of = inner_size_multiple_of
+        self.make_vocab_size_divisible_by = make_vocab_size_divisible_by
+
+        self.tie_embeddings = tie_embeddings
+
+        self.max_seqlen = max_seqlen
+        self.max_batch_size = max_batch_size
+        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
+
+        self.prefill_style = prefill_style
+        self.inference_mode = inference_mode
+
+        self.use_fp8_input_projections = use_fp8_input_projections
+
+        self.use_cache = use_cache
+        self.use_flash_attn = use_flash_attn
+        self.use_flash_rmsnorm = use_flash_rmsnorm
+        self.use_flash_depthwise = use_flash_depthwise
+        self.use_flashfft = use_flashfft
+
+        self.attn_block_dtype = attn_block_dtype
+        self.hyena_block_dtype = hyena_block_dtype
+        self.mlp_dtype = mlp_dtype
+
+        self.model_parallel_size = model_parallel_size
+        self.pipe_parallel_size = pipe_parallel_size
+
+        self.proj_groups = proj_groups
+        self.smeared_gqa = smeared_gqa
+
+        super().__init__(**kwargs)
+
+    # HF generation helpers expect `num_hidden_layers`.
+    @property
+    def num_hidden_layers(self) -> int:
+        return self.num_layers
+
+    # Internal blocks were originally written against a dotdict and call
+    # `config.get(key, default)` extensively; provide a dict-like getter.
+    def get(self, key, default=None):
+        return getattr(self, key, default)
+
+    @classmethod
+    def from_original_config(cls, config_path: str, **kwargs) -> "Evo2Config":
+        with open(config_path, "r") as f:
+            config = json.load(f)
+        return cls(**config, **kwargs)

engine.py

@@ -0,0 +1,327 @@
+"""Hyena inference engine for Evo2 (StripedHyena2).
+
+Three operator families are exercised by Evo2 blocks:
+
+  * parallel_fir(gate=False)  -- outer FIR used by all hyena blocks before any
+                                 channel split (input projection convolution).
+  * parallel_fir(gate=True)   -- inner FIR cascade used by hcm/hcs blocks
+                                 (x1 * v gated, then convolved by `h`,
+                                 then multiplied by x2 postgate).
+  * parallel_iir              -- modal-form IIR (long convolution via FFT)
+                                 used by hcl blocks; poles + residues parameterize
+                                 a stable, long-range linear filter.
+
+Sequential step paths (step_fir / step_iir) are used during generation.
+
+Layout conventions match vortex exactly so checkpoints are bit-identical.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+IIR_PREFILL_MODES = ["recurrence", "modal-fft"]
+
+
+def adjust_filter_shape_for_broadcast(u, h):
+    h = h.squeeze()
+    if len(u.shape) > len(h.shape):
+        h = h.unsqueeze(0)
+    if len(u.shape) > 3:
+        h = h.unsqueeze(1)
+    return h
+
+
+def fftconv_func(u, k, D, dropout_mask, gelu=True, k_rev=None, bidirectional=False, **kwargs):
+    """FFT convolution for long FIR filters (length >= 128 path)."""
+    seqlen = u.shape[-1]
+    fft_size = 2 * seqlen
+
+    k_f = torch.fft.rfft(k, n=fft_size) / fft_size
+    k_f = adjust_filter_shape_for_broadcast(u, k_f)
+    k = k.squeeze()
+
+    if bidirectional:
+        u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)
+        k, k2 = k.split(k.shape[1] // 2, dim=1)
+        k2_f = torch.fft.rfft(k2, n=fft_size) / fft_size
+        y1 = u_f * k_f
+        y2 = u_f.conj() * k2_f.conj()
+        y = torch.fft.irfft(y1 + y2, n=fft_size, norm="forward")[..., :seqlen]
+    else:
+        if k_rev is not None:
+            k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size
+            k_f = k_f + k_rev_f.conj()
+        u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)
+        y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen]
+
+    out = y + u * D.unsqueeze(-1)
+    return out.to(dtype=u.dtype)
+
+
+def _column_split(x, num_heads, head_size):
+    """Compatibility helper for column_split_hyena=True (not used by Evo2)."""
+    x = x.reshape(x.shape[0], num_heads, 3 * head_size, x.shape[2])
+    x2 = x[:, :, :head_size].reshape(x.shape[0], -1, x.shape[-1])
+    x1 = x[:, :, head_size : 2 * head_size].reshape(x.shape[0], -1, x.shape[-1])
+    v = x[:, :, 2 * head_size :].reshape(x.shape[0], -1, x.shape[-1])
+    return x2, x1, v
+
+
+class HyenaInferenceEngine:
+    def __init__(
+        self,
+        layer_idx: int | None = None,
+        iir_prefill_style: str = "modal-fft",
+        hyena_flip_x1x2: bool = False,
+    ) -> None:
+        assert iir_prefill_style in IIR_PREFILL_MODES, iir_prefill_style
+        self.iir_prefill_style = iir_prefill_style
+        self.layer_idx = layer_idx
+        self.low_mem_mode = False
+        self.hyena_flip_x1x2 = hyena_flip_x1x2
+
+    # ---------------------------------------------------------------- FIR
+    def parallel_fir(
+        self,
+        fir_fn,
+        u,
+        weight,
+        bias,
+        L,
+        dims,
+        groups=None,
+        gated_bias=False,
+        column_split_hyena=False,
+        dim_last=True,
+        fir_length=3,
+        gate=False,
+        inference_params=None,
+        padding_mask=None,
+    ):
+        L = u.shape[1] if dim_last else u.shape[2]
+        if gate:
+            hidden_size, num_attention_heads, hidden_size_per_attention_head, _, _ = dims
+            if column_split_hyena:
+                x2, x1, v = _column_split(u, num_attention_heads, hidden_size_per_attention_head)
+            else:
+                x2, x1, v = u.split([hidden_size, hidden_size, hidden_size], dim=1)
+            if self.hyena_flip_x1x2:
+                x1, x2 = x2, x1
+            u = x1 * v
+
+        if fir_length >= 128:
+            with torch.autocast("cuda"):
+                z = fftconv_func(
+                    u.to(torch.float32),
+                    weight[:, :, :L].to(torch.float32),
+                    bias,
+                    None,
+                    gelu=False,
+                    bidirectional=False,
+                    groups=groups,
+                )
+                z = z.to(u.dtype)
+        else:
+            if dim_last:
+                u = u.permute(0, 2, 1)  # B, D, L
+
+            z = fir_fn(
+                u.to(torch.float32),
+                weight.to(torch.float32),
+                bias=None,
+                stride=1,
+                padding=fir_length - 1,
+                groups=u.shape[1],
+            )[..., :L]
+
+            z = z.to(u.dtype)
+
+            if bias is not None:
+                if gated_bias:
+                    z = z + bias[None, :, None] * u
+                else:
+                    z = z + bias[None, :, None]
+
+        if isinstance(padding_mask, torch.Tensor):
+            z = z * padding_mask[:, None]
+
+        if gate:
+            z = x2 * z
+
+        if inference_params is not None:
+            fir_state = u[..., -fir_length + 1 :]
+        else:
+            fir_state = None
+
+        return z, fir_state
+
+    # ---------------------------------------------------------------- IIR
+    def parallel_iir(
+        self,
+        z_pre,
+        h,
+        D,
+        L,
+        poles,
+        residues,
+        t,
+        dims,
+        layer_idx,
+        inference_params=None,
+        prefill_style: str = "fft",
+        fftconv_fn=None,
+        padding_mask=None,
+        use_flashfft: bool = False,
+        column_split_hyena: bool = False,
+        long_fir_threshold: int | None = None,
+    ):
+        fft_size = 2 * L
+        hidden_size, num_attention_heads, hidden_size_per_attention_head, _, _ = dims
+        if column_split_hyena:
+            x2, x1, v = _column_split(z_pre, num_attention_heads, hidden_size_per_attention_head)
+        else:
+            x2, x1, v = z_pre.split([hidden_size, hidden_size, hidden_size], dim=1)
+
+        if self.hyena_flip_x1x2:
+            x1, x2 = x2, x1
+
+        x1v = x1 * v
+
+        X_s = None
+        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:
+                y = fftconv_fn(
+                    x1v.to(dtype=torch.bfloat16).contiguous(),
+                    h.to(dtype=torch.float32),
+                )
+            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]
+                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 and 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,
+            )
+
+        return y.permute(0, 2, 1)
+
+    # --------------------------------------------------------- step (decode)
+    def step_fir(self, u, fir_state, weight, bias=None, gated_bias=False, flip_filter=False):
+        """Single-step FIR. fir_state holds the last (filter_len - 1) inputs."""
+        weight = weight.squeeze()
+        cache_size = fir_state.shape[-1]
+        filter_length = weight.shape[-1]
+        if flip_filter:
+            weight = weight.flip(-1)
+            weight = weight[..., -cache_size - 1 :].unsqueeze(0)
+        else:
+            weight = weight[..., : cache_size + 1].unsqueeze(0)
+
+        input_dtype = u.dtype
+        weight = weight.to(torch.float32)
+        u = u.to(torch.float32)
+        fir_state = fir_state.to(torch.float32)
+        bias = bias.to(torch.float32) if bias is not None else None
+
+        h0, h = weight[..., -1], weight[..., :-1]
+        y = h0 * u + torch.sum(fir_state * h, dim=-1)
+
+        if bias is not None:
+            if gated_bias:
+                y = y + bias * u
+            else:
+                y = y + bias
+
+        if cache_size < filter_length - 1:
+            fir_state = torch.cat([fir_state, u[..., None]], dim=-1)
+        else:
+            fir_state = torch.roll(fir_state, -1, dims=2)
+            fir_state[..., -1] = u
+
+        return y.to(input_dtype), fir_state
+
+    def step_iir(self, x2, x1, v, D, residues, poles, iir_state, iir_groups=1):
+        x1v = x1 * v
+        # `poles` arg contains log_poles (real, in modal form for evo2)
+        poles = torch.exp(poles)
+        poles = poles[..., 0][None]
+        residues = residues[None]
+        iir_state = poles * iir_state + x1v[..., None]
+        res_state = torch.sum(residues * iir_state, dim=-1)
+
+        if iir_groups > 1:
+            raise NotImplementedError
+        y = x2 * (res_state + D * x1v)
+        return y, iir_state
+
+    def prefill_via_direct_recurrence(self, inference_params, x1v, L, residues, poles, *args, **kwargs):
+        state_dim = poles.shape[1]
+        x1v_ = x1v[..., None, None]
+        x1v_ = x1v_.repeat(1, 1, 1, state_dim, 2)
+        x1v_[..., 1] = 0
+
+        state = 0 * x1v_[:, :, 0]
+        output = 0 * x1v_[:, :, :, 0, 0]
+
+        poles = poles[:, :, 0][None]
+        residues = residues[:, :, 0][None].repeat(x1v_.shape[0], 1, 1, 1)
+
+        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]
+
+        inference_params.state_dict[self.layer_idx] = state.to(dtype=torch.float32)
+        return output
+
+    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.float32, *args, **kwargs,
+    ):
+        """Compute IIR state via a single FFT.
+
+        Evo2 uses *real* `log_poles` (not the complex view-as-real layout
+        that Evo1 uses), so the impulse-response IFFT is mathematically real;
+        the imaginary component is FFT round-off. We take ``.real`` explicitly
+        instead of relying on torch's lossy complex->real cast, which avoids
+        the "Casting complex values to real discards the imaginary part"
+        UserWarning at every decoded token.
+        """
+        hidden_size, _, _, state_size, hyena_filter_groups = dims
+
+        assert X_s is not None
+        bs = x1v.shape[0]
+        fft_size = 2 * L
+        state_s = (poles.to(torch.float32) * t).exp()
+        state_S = torch.fft.fft(state_s, n=fft_size).repeat(bs, 1, 1, 1)
+        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].real.to(dtype=state_dtype)

generation_config.json

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

hyena.py

@@ -0,0 +1,517 @@
+"""StripedHyena2 block definitions for Evo2.
+
+Block types:
+
+  * AttentionBlock       -- standard pre-norm transformer block with MHA + RoPE.
+  * ParallelGatedConvBlock -- pre-norm hyena block; the inner `filter` is a
+                              HyenaCascade configured per block-type idx:
+                                hcl: IIR via log_poles + residues
+                                hcm: FIR cascade, fir_inner_filter_length=128
+                                hcs: FIR cascade, fir_inner_filter_length=7
+
+Inside ParallelGatedConvBlock, `projections` is a TELinear (3x hidden_size
+output) -- the FP8-capable input projection. For 7B variants without TE,
+the pure-PyTorch fallback in layers.py is used.
+
+The `interleave` config flag (True for all Evo2 variants) reorders the
+channel dim after the outer FIR before the channel-split / cascade.
+"""
+
+from __future__ import annotations
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .attention import MHA
+from .engine import HyenaInferenceEngine
+from .layers import ParallelGatedMLP, RMSNorm, TELinear
+from .rotary import swap_mha_rope
+
+
+def _interleave(z_pre: torch.Tensor) -> torch.Tensor:
+    """Deinterleave a [x1_0,x2_0,v_0,x1_1,x2_1,v_1,...] channel layout into
+    blocks of [x1_all | x2_all | v_all]. Matches vortex.model.utils.interleave."""
+    if len(z_pre.shape) == 3:  # (B, C, L) non-cached path
+        x1 = z_pre[:, 0::3, :]
+        x2 = z_pre[:, 1::3, :]
+        v = z_pre[:, 2::3, :]
+        return torch.cat([x1, x2, v], dim=1)
+    x1 = z_pre[..., 0::3]
+    x2 = z_pre[..., 1::3]
+    v = z_pre[..., 2::3]
+    return torch.cat([x1, x2, v], dim=-1)
+
+
+def _column_split_step(z_pre, num_heads, head_size):
+    """Per-step variant of column_split for sequential_forward (cache step)."""
+    x = z_pre.reshape(z_pre.shape[0], num_heads, 3 * head_size)
+    x2 = x[..., :head_size].reshape(z_pre.shape[0], -1)
+    x1 = x[..., head_size : 2 * head_size].reshape(z_pre.shape[0], -1)
+    v = x[..., 2 * head_size :].reshape(z_pre.shape[0], -1)
+    return x2, x1, v
+
+
+# =============================================================================
+# Attention block
+# =============================================================================
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, config, layer_idx: int) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.pre_norm = RMSNorm(config)
+        self.post_norm = RMSNorm(config)
+        self.proj_groups = config.get("proj_groups", 1)
+        dtype = self._resolve_dtype(config.get("attn_block_dtype", "bfloat16"))
+        mlp_dtype = self._resolve_dtype(config.get("mlp_dtype", "bfloat16"))
+        self.num_attention_heads = config.num_attention_heads
+        self.hidden_size = config.hidden_size
+        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", False),
+            rotary_emb_base=config.get("rotary_emb_base", 1000000),
+            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
+        # Round-trip inv_freq through state_dict for safety.
+        self.inner_mha_cls.rotary_emb.register_buffer(
+            "inv_freq", self.inner_mha_cls.rotary_emb.inv_freq
+        )
+
+        self.mlp = ParallelGatedMLP(config, layer_idx).to(dtype=mlp_dtype)
+
+    @staticmethod
+    def _resolve_dtype(name) -> torch.dtype:
+        if isinstance(name, torch.dtype):
+            return name
+        mapping = {
+            "bfloat16": torch.bfloat16,
+            "float16": torch.float16,
+            "float32": torch.float32,
+        }
+        return mapping.get(name, torch.bfloat16)
+
+    def forward(
+        self,
+        u: torch.Tensor,
+        inference_params=None,
+        padding_mask: Optional[torch.Tensor] = None,
+        output_attentions: bool = False,
+        *args,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        if isinstance(padding_mask, torch.Tensor):
+            # Zero attended values at pad positions (no qkv bias here).
+            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
+
+
+# =============================================================================
+# Hyena cascade (hcl / hcm / hcs)
+# =============================================================================
+
+
+class HyenaCascade(nn.Module):
+    """The inner mixer of a hyena block.
+
+    Two execution modes selected by `fir_inner_filter_length`:
+      * None        -> hcl: outer FIR -> channel split -> parallel_iir
+                              (modal-form long convolution)
+      * int (e.g. 7 or 128) -> hcm/hcs: outer FIR -> interleave -> inner FIR
+                              cascade via parallel_fir(gate=True)
+    """
+
+    def __init__(
+        self,
+        config,
+        layer_idx: int,
+        hyena_filter_groups: int,
+        fir_inner_filter_length: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hyena_filter_groups = hyena_filter_groups
+
+        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", False)
+        self.hyena_flip_x1x2 = config.get("hyena_flip_x1x2", False)
+
+        assert self.hidden_size % self.num_filters == 0
+        assert self.num_filters <= self.hidden_size
+
+        self.num_attention_heads = config.num_attention_heads
+        self.hidden_size_per_attention_head = (
+            self.hidden_size // self.num_attention_heads
+        )
+
+        self.fir_inner_filter_length = fir_inner_filter_length
+        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,
+            hyena_flip_x1x2=config.get("hyena_flip_x1x2", False),
+        )
+        self.fir_fn = F.conv1d
+        self.fir_inner_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 not self.use_flashfft, "long_fir_threshold incompatible with flashfft"
+
+        self.num_systems = self.hyena_filter_groups
+        self.channels_per_group = self.hidden_size // self.hyena_filter_groups
+
+        if self.fir_inner_filter_length:
+            # hcm / hcs: explicit FIR filter `h` of shape (groups, 1, inner_len)
+            self.h = nn.Parameter(
+                torch.randn(self.hyena_filter_groups, 1, fir_inner_filter_length)
+            )
+            if fir_inner_filter_length >= 128:
+                self.D = nn.Parameter(torch.zeros(self.hidden_size))
+            else:
+                self.D = None
+        else:
+            # hcl: modal-form IIR, log_poles in (num_systems, state_size, 1)
+            log_poles = torch.randn(self.num_systems, self.state_size, 1, dtype=torch.float32)
+            self.log_poles = nn.Parameter(log_poles)
+            self.residues = nn.Parameter(
+                torch.randn(self.num_systems, self.state_size, dtype=torch.float32)
+            )
+            self.D = nn.Parameter(torch.zeros(self.hidden_size))
+            self.h = None
+
+        self.t = 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
+        ):
+            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]
+        dims = (
+            self.hidden_size,
+            self.num_attention_heads,
+            self.hidden_size_per_attention_head,
+            self.state_size,
+            self.hyena_filter_groups,
+        )
+        z_pre, fir_state = self.engine.parallel_fir(
+            self.fir_fn, u,
+            self.short_filter_weight, self.short_filter_bias,
+            L, dims=dims, gate=False,
+            column_split_hyena=self.column_split_hyena,
+            fir_length=self.short_filter_length,
+            inference_params=inference_params,
+            padding_mask=padding_mask, dim_last=True,
+        )
+        if inference_params:
+            inference_params.fir_state_dict[self.layer_idx] = fir_state
+
+        if self.config.interleave:
+            z_pre = _interleave(z_pre)
+
+        if self.h is None:
+            # hcl path: compute IIR impulse response then convolve via FFT
+            h, _, _, _ = self.compute_filter(L, u.device)
+        else:
+            h = self.h
+
+        D = self.D
+        if self.hyena_filter_groups > 1:
+            h = h.repeat_interleave(self.hidden_size // self.hyena_filter_groups, 0)
+
+        if self.fir_inner_filter_length is not None:
+            # hcm / hcs: inner FIR cascade
+            y, fir_inner_state = self.engine.parallel_fir(
+                self.fir_inner_fn, z_pre, h, D, L,
+                dims=dims, gate=True,
+                gated_bias=self.fir_inner_filter_length >= 128,
+                dim_last=False,
+                column_split_hyena=self.column_split_hyena,
+                fir_length=self.fir_inner_filter_length,
+                inference_params=inference_params,
+                padding_mask=padding_mask,
+                groups=self.hyena_filter_groups,
+            )
+            y = y.permute(0, 2, 1)
+            if inference_params:
+                inference_params.fir_inner_state_dict[self.layer_idx] = fir_inner_state
+        else:
+            # hcl: parallel IIR via FFT
+            y = self.engine.parallel_iir(
+                z_pre, h, D, L, t=self.t,
+                poles=self.log_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 len(u.shape) > 2:
+            u = u[:, -1]
+        # Track input dtype so we can cast back at the end. step_iir produces
+        # fp32 output (log_poles / residues are fp32 for stability) which would
+        # otherwise dtype-mismatch the downstream bf16 Linear (out_filter_dense).
+        input_dtype = u.dtype
+
+        z_pre, fir_state = self.engine.step_fir(
+            u, inference_params.fir_state_dict[self.layer_idx],
+            weight=self.short_filter_weight, bias=self.short_filter_bias,
+        )
+        inference_params.fir_state_dict[self.layer_idx] = fir_state
+
+        if self.config.interleave:
+            # For a single step, dim-1 length is C and `_interleave` operates
+            # along the channel dim; reshape to (B, C, 1) so the dim==3 branch
+            # applies, then squeeze the last dim back.
+            z_pre = _interleave(z_pre.unsqueeze(-1)).squeeze(-1)
+
+        if self.column_split_hyena:
+            x2, x1, v = _column_split_step(
+                z_pre, self.num_attention_heads, self.hidden_size_per_attention_head
+            )
+        else:
+            x2, x1, v = z_pre.split([self.hidden_size, self.hidden_size, self.hidden_size], dim=1)
+
+        if self.hyena_flip_x1x2:
+            x1, x2 = x2, x1
+
+        if self.fir_inner_filter_length is not None:
+            if self.hyena_filter_groups > 1:
+                h = self.h.repeat_interleave(self.hidden_size // self.hyena_filter_groups, 0)
+            else:
+                h = self.h
+            y, fir_inner_state = self.engine.step_fir(
+                x1 * v,
+                inference_params.fir_inner_state_dict[self.layer_idx],
+                weight=h, bias=self.D,
+                flip_filter=self.fir_inner_filter_length >= 128,
+                gated_bias=self.fir_inner_filter_length >= 128,
+            )
+            y = y * x2
+            inference_params.fir_inner_state_dict[self.layer_idx] = fir_inner_state
+        else:
+            y, iir_state = self.engine.step_iir(
+                x2, x1, v, self.D,
+                self.residues, self.log_poles,
+                inference_params.state_dict[self.layer_idx],
+                iir_groups=1,
+            )
+            inference_params.state_dict[self.layer_idx] = iir_state
+
+        y = y.to(input_dtype)
+        return y[:, None], inference_params
+
+    def update_time(self, L, device):
+        if self.t is None:
+            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 = self.residues.to(filter_dtype)
+        log_poles = self.log_poles.to(filter_dtype)
+        h = (residues[..., None] * (log_poles * self.t).exp()).sum(1)[None]  # B, D, L
+        return h, filter_dtype, log_poles, residues
+
+
+# =============================================================================
+# Hyena block (wraps HyenaCascade with norms / projections / mlp)
+# =============================================================================
+
+
+class ParallelGatedConvBlock(nn.Module):
+    def __init__(
+        self,
+        config,
+        layer_idx: int,
+        hyena_filter_groups: int | None = None,
+        fir_inner_filter_length: int | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.fir_inner_filter_length = fir_inner_filter_length
+        self.hyena_filter_groups = (
+            hyena_filter_groups if hyena_filter_groups is not None else config.hidden_size
+        )
+        dtype = AttentionBlock._resolve_dtype(config.get("hyena_block_dtype", "bfloat16"))
+        mlp_dtype = AttentionBlock._resolve_dtype(config.get("mlp_dtype", "bfloat16"))
+        self.pre_norm = RMSNorm(config).to(dtype=dtype)
+        self.post_norm = RMSNorm(config).to(dtype=dtype)
+        self.filter = HyenaCascade(
+            config, layer_idx,
+            hyena_filter_groups=self.hyena_filter_groups,
+            fir_inner_filter_length=fir_inner_filter_length,
+        ).to(dtype=dtype)
+
+        self.projections = TELinear(
+            config.hidden_size,
+            3 * config.hidden_size,
+            bias=config.qkv_proj_bias,
+            init_method=torch.nn.init.xavier_uniform_,
+            use_fp8=config.get("use_fp8_input_projections", False),
+        )
+
+        self.out_filter_dense = nn.Linear(
+            config.hidden_size, config.hidden_size,
+            bias=config.hyena_out_proj_bias,
+        ).to(dtype)
+        self.mlp = ParallelGatedMLP(config, layer_idx).to(dtype=mlp_dtype)
+
+    def _pad_to_multiple(self, x: torch.Tensor, multiple: int = 16) -> torch.Tensor:
+        """Right-pad along seq dim to a multiple of `multiple` when FP8 input
+        projections are enabled. TE's FP8 path requires the product of all
+        dims except the last to be divisible by 8 (we use 16 to be safe).
+        No-op when FP8 is off."""
+        if not self.config.get("use_fp8_input_projections", False):
+            return x
+        seq_len = x.size(1)
+        pad_len = (multiple - (seq_len % multiple)) % multiple
+        if pad_len == 0:
+            return x
+        return F.pad(x, (0, 0, 0, pad_len))
+
+    def _proj_norm(self, x):
+        original_seq_len = x.size(1)
+        normalized = self.pre_norm(x)
+        normalized = self._pad_to_multiple(normalized)
+        with torch.cuda.device(x.device) if x.is_cuda else _nullctx():
+            projected = self.projections(normalized)
+        if isinstance(projected, tuple):
+            projected = projected[0]
+        # Slice back to original seq length if padding was added.
+        if projected.size(1) > original_seq_len:
+            projected = projected[:, :original_seq_len, :]
+        return projected
+
+    def forward(
+        self,
+        u: torch.Tensor,
+        inference_params=None,
+        padding_mask=None,
+        output_attentions: bool = False,
+        *args,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, None]:
+        z = self._proj_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
+
+
+class _nullctx:
+    def __enter__(self): return None
+    def __exit__(self, *a): return False
+
+
+# =============================================================================
+# Block dispatch
+# =============================================================================
+
+
+def get_block(config, layer_idx: int):
+    if layer_idx in config.attn_layer_idxs:
+        return AttentionBlock(config, layer_idx)
+    if layer_idx in config.hcl_layer_idxs:
+        return ParallelGatedConvBlock(
+            config, layer_idx,
+            hyena_filter_groups=config.hcl_filter_groups,
+            fir_inner_filter_length=None,
+        )
+    if layer_idx in config.hcm_layer_idxs:
+        return ParallelGatedConvBlock(
+            config, layer_idx,
+            hyena_filter_groups=config.hcm_filter_groups,
+            fir_inner_filter_length=config.hcm_filter_length,
+        )
+    if layer_idx in config.hcs_layer_idxs:
+        return ParallelGatedConvBlock(
+            config, layer_idx,
+            hyena_filter_groups=config.hcs_filter_groups,
+            fir_inner_filter_length=config.hcs_filter_length,
+        )
+    raise NotImplementedError(f"layer_idx {layer_idx} not in any block-type idxs")
+
+
+def block_type_for_idx(config, layer_idx: int) -> str:
+    if layer_idx in config.attn_layer_idxs: return "mha"
+    if layer_idx in config.hcl_layer_idxs:  return "hcl"
+    if layer_idx in config.hcm_layer_idxs:  return "hcm"
+    if layer_idx in config.hcs_layer_idxs:  return "hcs"
+    raise ValueError(f"block idx {layer_idx} not classified")

layers.py

@@ -0,0 +1,252 @@
+"""Basic layers for Evo2 (StripedHyena2).
+
+  * RMSNorm                -- pre/post norm in each block
+  * ParallelGatedMLP       -- GLU(act(l1(x)) * l2(x)) -> l3(...); evo2_style
+                              activations replace `act` with Identity for layer > 0
+  * VocabParallelEmbedding -- single-process variant of vortex's
+                              VocabParallelEmbedding so checkpoint keys match
+  * TELinear               -- TransformerEngine FP8-capable Linear, with a
+                              pure-PyTorch fallback when TE is not installed.
+                              Used for the input QKV-like projections of
+                              Hyena cascade blocks; required for 1B/20B/40B.
+"""
+
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+
+try:
+    from transformer_engine.pytorch import Linear as _TELinearBase
+    from transformer_engine.common.recipe import Format, DelayedScaling
+    import transformer_engine.pytorch as te
+    HAS_TE = True
+except ImportError:
+    HAS_TE = False
+
+
+def grab_first_if_tuple(x):
+    if x.__class__.__name__ == "tuple":
+        return x[0]
+    return x
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.eps = config.eps
+        self.hidden_size = config.hidden_size
+        self.scale = 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):
+    """GLU MLP. With evo2_style_activations=True, layer_idx > 0 uses Identity
+    in place of the gating activation; layer 0 keeps gelu."""
+
+    def __init__(self, config, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        multiple_of = config.get("inner_size_multiple_of", 64)
+        self.act_type = config.get("mlp_activation", "gelu")
+        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}")
+
+        if self.layer_idx > 0 and config.get("evo2_style_activations", False):
+            self.act = nn.Identity()
+
+        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 vortex's VocabParallelEmbedding.
+
+    Drops distributed sharding so this minimal port runs on a single device
+    (HF/accelerate handles cross-device placement via device_map="auto").
+    """
+
+    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
+
+
+if HAS_TE:
+
+    _TE_WORKSPACE_FIXUP_APPLIED = False
+
+    def fixup_te_workspace():
+        """Patch TE's Linear module to use per-device cuBLAS workspaces.
+
+        Vortex helper, ported verbatim. Without this, TE's Linear uses a single
+        workspace tensor allocated on whatever CUDA device was current at first
+        call, then fails ("cuBLAS Error: the function failed to launch on the
+        GPU") when called from layers on other devices -- which happens with
+        device_map="auto" sharding. Idempotent; safe to call multiple times.
+        No-op when TE is not installed (the import guard above gates it).
+        """
+        global _TE_WORKSPACE_FIXUP_APPLIED
+        if _TE_WORKSPACE_FIXUP_APPLIED:
+            return
+        from functools import lru_cache
+
+        @lru_cache
+        def te_cublas_get_workspace_per_device(device):
+            import transformer_engine.pytorch.module.base as tebase
+            with torch.cuda.device(device):
+                tebase._cublas_workspace = None  # force get_workspace() to reallocate
+                return tebase.get_workspace()
+
+        def get_workspace():
+            return te_cublas_get_workspace_per_device(torch.cuda.current_device())
+
+        import transformer_engine.pytorch.module.linear as telinear
+        telinear.get_workspace = get_workspace
+        _TE_WORKSPACE_FIXUP_APPLIED = True
+
+    def set_format_recipe():
+        fp8_format = Format.HYBRID
+        fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max")
+        return fp8_format, fp8_recipe
+
+    class TELinear(_TELinearBase):
+        """Wrapper for Transformer-Engine's Linear, matching vortex's signature.
+
+        Returns (out, None) so callers can grab_first_if_tuple, matching vortex.
+        """
+
+        def __init__(
+            self,
+            input_size: int,
+            output_size: int,
+            init_method: Callable | None = None,
+            bias: bool = True,
+            skip_bias_add: bool = False,
+            use_fp8: bool = False,
+            **kwargs,
+        ):
+            params_dtype = torch.bfloat16
+            self.te_return_bias = skip_bias_add and bias
+            self.use_fp8_input_projections = use_fp8
+            if use_fp8:
+                self.fp8_format, self.fp8_recipe = set_format_recipe()
+
+            super().__init__(
+                in_features=input_size,
+                out_features=output_size,
+                sequence_parallel=False,
+                fuse_wgrad_accumulation=False,
+                tp_group=None,
+                tp_size=1,
+                init_method=init_method,
+                params_dtype=params_dtype,
+                parallel_mode=None,
+                bias=bias,
+                return_bias=self.te_return_bias,
+                **kwargs,
+            )
+
+        def forward(self, x):
+            if self.use_fp8_input_projections:
+                with te.fp8_autocast(enabled=True, fp8_recipe=self.fp8_recipe):
+                    out = super().forward(x)
+            else:
+                out = super().forward(x)
+
+            if self.te_return_bias:
+                return out
+            return out, None
+
+else:
+
+    def fixup_te_workspace():
+        """No-op when TransformerEngine is not installed."""
+        return
+
+    class TELinear(nn.Module):
+        """Pure-PyTorch fallback for TELinear (no FP8). Used by 7B variants
+        when TransformerEngine isn't installed.
+
+        Parameters are registered with names that match TE's state_dict layout
+        (`weight`, `bias`), so checkpoints saved from either path are
+        cross-loadable.
+        """
+
+        def __init__(
+            self,
+            input_size: int,
+            output_size: int,
+            init_method: Callable | None = None,
+            bias: bool = True,
+            skip_bias_add: bool = False,
+            use_fp8: bool = False,
+            **kwargs,
+        ):
+            super().__init__()
+            if use_fp8:
+                raise RuntimeError(
+                    "FP8 requires Transformer Engine, which is not installed. "
+                    "Install it with: pip install transformer_engine>=2.3.0"
+                )
+
+            self.te_return_bias = skip_bias_add and bias
+            self.use_fp8_input_projections = False
+            self.in_features = input_size
+            self.out_features = output_size
+            self.has_bias = bias
+
+            self.weight = nn.Parameter(torch.empty(output_size, input_size, dtype=torch.bfloat16))
+            if bias:
+                self.bias = nn.Parameter(torch.zeros(output_size, dtype=torch.bfloat16))
+            else:
+                self.register_parameter("bias", None)
+
+            if init_method is not None:
+                init_method(self.weight)
+            else:
+                nn.init.xavier_uniform_(self.weight)
+
+        def forward(self, x):
+            out = F.linear(x.to(self.weight.dtype), self.weight, self.bias)
+            if self.te_return_bias:
+                return out, self.bias
+            return out, None

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,352 @@
+{
+  "metadata": {
+    "total_parameters": 6481059584,
+    "total_size": 12964479104
+  },
+  "weight_map": {
+    "backbone.blocks.0.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.1.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.inner_mha_cls.Wqkv.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.inner_mha_cls.out_proj.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.inner_mha_cls.out_proj.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.inner_mha_cls.rotary_emb.inv_freq": "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.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.10.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.11.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.12.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.13.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.filter.log_poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.13.filter.residues": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.14.filter.h": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.15.filter.h": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.filter.log_poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.filter.residues": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.filter.short_filter_weight": "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.out_filter_dense.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.16.out_filter_dense.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.16.projections.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.inner_mha_cls.Wqkv.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.inner_mha_cls.out_proj.bias": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.inner_mha_cls.out_proj.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.inner_mha_cls.rotary_emb.inv_freq": "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.post_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.17.pre_norm.scale": "model-00002-of-00003.safetensors",
+    "backbone.blocks.18.filter.h": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.19.filter.h": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.2.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.filter.log_poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.2.filter.residues": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.20.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.filter.log_poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.20.filter.residues": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.21.filter.h": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.22.filter.h": "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.weight": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.filter.D": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.filter.log_poles": "model-00002-of-00003.safetensors",
+    "backbone.blocks.23.filter.residues": "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.weight": "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.h": "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.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.26.filter.h": "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.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.filter.log_poles": "model-00003-of-00003.safetensors",
+    "backbone.blocks.27.filter.residues": "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.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.28.filter.h": "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.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.29.filter.h": "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.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.3.inner_mha_cls.Wqkv.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.inner_mha_cls.out_proj.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.inner_mha_cls.out_proj.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.inner_mha_cls.rotary_emb.inv_freq": "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.post_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.3.pre_norm.scale": "model-00001-of-00003.safetensors",
+    "backbone.blocks.30.filter.D": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.filter.log_poles": "model-00003-of-00003.safetensors",
+    "backbone.blocks.30.filter.residues": "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.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.inner_mha_cls.Wqkv.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.inner_mha_cls.out_proj.bias": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.inner_mha_cls.out_proj.weight": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.inner_mha_cls.rotary_emb.inv_freq": "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.post_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.31.pre_norm.scale": "model-00003-of-00003.safetensors",
+    "backbone.blocks.4.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.5.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.log_poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.6.filter.residues": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.7.filter.h": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.filter.h": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.filter.short_filter_weight": "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.out_filter_dense.bias": "model-00001-of-00003.safetensors",
+    "backbone.blocks.8.out_filter_dense.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.8.projections.weight": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.D": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.log_poles": "model-00001-of-00003.safetensors",
+    "backbone.blocks.9.filter.residues": "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.weight": "model-00001-of-00003.safetensors",
+    "backbone.embedding_layer.weight": "model-00001-of-00003.safetensors",
+    "backbone.norm.scale": "model-00003-of-00003.safetensors"
+  }
+}

modeling_evo2.py

@@ -0,0 +1,425 @@
+"""HuggingFace wrappers for Evo2 (StripedHyena2).
+
+Two top-level model classes:
+
+  * Evo2Model        -- bare backbone returning BaseModelOutputWithPast
+                        (no LM head); the post-RMSNorm hidden state is the
+                        last_hidden_state.
+  * Evo2ForCausalLM  -- with the tied LM head + GenerationMixin so
+                        model.generate() works out of the box.
+
+Caching for autoregressive generation uses the custom Evo2Cache (see
+cache.py): a dict-like container of four block-type-specific param objects
+(mha, hcl, hcm, hcs) so all four StripedHyena2 block types can decode in
+constant time per new token.
+
+Per-block dtype: bfloat16 for everything except the modal-form filter
+parameters (`log_poles`, `residues`) which must stay fp32 for stability.
+This is enforced both at convert time (the safetensors store these as fp32)
+and at runtime via `force_dtype()` (called from each model's __init__).
+"""
+
+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 .cache import Evo2Cache
+from .configuration_evo2 import Evo2Config
+from .hyena import (
+    AttentionBlock,
+    ParallelGatedConvBlock,
+    block_type_for_idx,
+    get_block,
+)
+from .layers import RMSNorm, VocabParallelEmbedding
+# Bundle the tokenizer file via trust_remote_code:
+from .tokenization_evo2 import ByteTokenizer  # noqa: F401
+# Force HF's trust_remote_code loader to copy these transitive deps into its
+# dynamic-module cache (it only walks top-level `from .X import Y` patterns
+# of modeling_evo2.py, not `from . import X`, so we use the explicit form).
+from .attention import MHA as _MHA  # noqa: F401
+from .engine import HyenaInferenceEngine as _HyenaInferenceEngine  # noqa: F401
+from .rotary import RotaryEmbedding as _RotaryEmbedding  # noqa: F401
+
+logger = logging.get_logger(__name__)
+
+
+# =============================================================================
+# Backbone
+# =============================================================================
+
+
+class StripedHyena2(nn.Module):
+    """Pure backbone: token embedding -> N blocks -> final RMSNorm."""
+
+    def __init__(self, config: Evo2Config):
+        super().__init__()
+        # Patch TE's cuBLAS workspace handling so multi-GPU device_map="auto"
+        # works (TE otherwise allocates a single workspace on the first device
+        # it sees, then fails on layers placed on other devices). No-op when
+        # TE isn't installed.
+        from .layers import fixup_te_workspace
+        fixup_te_workspace()
+
+        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
+            self.flash_fft = FlashFFTConv(2 * config.max_seqlen, dtype=torch.bfloat16)
+        else:
+            self.flash_fft = None
+
+        self.blocks = nn.ModuleList(
+            get_block(config, i) for i in range(config.num_layers)
+        )
+        # Wire fftconv_fn into the hcl filters when flashfft is on.
+        if self.flash_fft is not None:
+            for block in self.blocks:
+                if isinstance(block, ParallelGatedConvBlock) and block.filter.h is None:
+                    block.filter.fftconv_fn = self.flash_fft
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        inference_params_dict: Optional[Evo2Cache] = 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, params_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, params_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, params_out, all_hidden_states, all_attentions
+
+    def _stateful_forward(
+        self, x, cache: Evo2Cache,
+        all_hidden_states, all_attentions,
+        output_hidden_states, output_attentions,
+    ):
+        for block_idx, block in enumerate(self.blocks):
+            block_name = block_type_for_idx(self.config, block_idx)
+            inference_params = cache.by_block_name(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, cache
+
+    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) -> Evo2Cache:
+        return Evo2Cache(
+            max_seqlen=self.config.get("max_seqlen", 8192),
+            max_batch_size=max_batch_size,
+            short_filter_length=self.config.short_filter_length,
+            hcm_filter_length=self.config.hcm_filter_length,
+            hcs_filter_length=self.config.hcs_filter_length,
+            state_size=self.config.state_size,
+        )
+
+    def to_bfloat16_except_poles_residues(self):
+        """Cast params to bf16, restore fp32 invariants:
+
+          * log_poles / residues -- bf16 collapses the modal-form IIR (the
+            poles get rounded too aggressively, the recurrence blows up).
+            Mirrors vortex's ``to_bfloat16_except_pr_lc(to_float32=True)``
+            first-pass behaviour.
+
+          * rotary_emb.inv_freq -- HF ``from_pretrained(dtype=bf16)`` casts
+            ALL buffers to bf16. ``inv_freq = 1 / base^(2i/dim)`` loses ~7
+            bits of mantissa in bf16, which shifts cos/sin tables by ~5e-2
+            per cell at position 64+, causing each attention layer to add
+            ~4e-2 of noise on Q/K. We recompute inv_freq in fp32 here so
+            the cos/sin builder uses the precise stored value.
+        """
+        for k, p in self.named_parameters():
+            if "log_poles" in k or "residues" in k:
+                p.data = p.data.to(torch.float32)
+            else:
+                p.data = p.data.to(torch.bfloat16)
+        for module in self.modules():
+            if hasattr(module, "_compute_inv_freq") and hasattr(module, "inv_freq"):
+                fresh = module._compute_inv_freq(device=module.inv_freq.device)
+                module.inv_freq.data = fresh
+                # Invalidate any cached cos/sin so the next forward rebuilds
+                # them from the precise fp32 inv_freq.
+                module._seq_len_cached = 0
+                module._cos_cached = None
+                module._sin_cached = None
+
+
+# =============================================================================
+# HF wrappers
+# =============================================================================
+
+
+class Evo2PreTrainedModel(PreTrainedModel):
+    config_class = Evo2Config
+    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", r"_extra_state$"]
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    # Evo2 cache is a custom Evo2Cache (not a DynamicCache); tell HF not to wrap it.
+    _supports_cache_class = False
+    # log_poles / residues parameterize a modal long-range filter; bf16 collapses
+    # them. HF will keep these in fp32 even with dtype=bf16 at load time.
+    _keep_in_fp32_modules = ["log_poles", "residues"]
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        if "dtype" not in kwargs and "torch_dtype" not in kwargs:
+            kwargs["dtype"] = torch.bfloat16
+        return super().from_pretrained(*args, **kwargs)
+
+
+class Evo2Model(Evo2PreTrainedModel):
+    """Bare backbone returning BaseModelOutputWithPast."""
+
+    def __init__(self, config: Evo2Config):
+        super().__init__(config)
+        self.backbone = StripedHyena2(config)
+        self.config = config
+        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 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
+        )
+        # Default to no caching for the bare backbone: caches have large
+        # per-layer memory cost with no benefit for embedding extraction.
+        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 Evo2ForCausalLM(Evo2PreTrainedModel, GenerationMixin):
+    """LM head wrapper. Tied to backbone.embedding_layer (Evo2 ties weights)."""
+
+    def __init__(self, config: Evo2Config, **kwargs):
+        super().__init__(config, **kwargs)
+        self.backbone = StripedHyena2(config)
+        self.config = 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.")
+            use_cache = False
+
+        inputs = input_ids
+        if use_cache:
+            if not isinstance(past_key_values, Evo2Cache):
+                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:
+                    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,
+        )
+
+        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,201 @@
+"""Rotary embeddings for Evo2.
+
+Mirrors models/evo/hf/rotary.py (same RoPE math): flash_attn fast path when
+available, pure-PyTorch fallback otherwise. LinearlyScaledRotaryEmbedding
+implements linear position-index interpolation (`t = t / scaling_factor`)
+used by all the 1M / 262K context variants.
+"""
+
+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:
+    _FlashRotaryEmbedding = None
+    _HAS_FLASH_ROTARY = False
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    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:
+    # Compute the multiply in fp32 then cast back to x.dtype to match
+    # flash_attn's Triton kernel bit-exactly. Doing the multiply in bf16
+    # directly compounds rounding error of ~3e-2 per layer.
+    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. Mirrors flash_attn.layers.rotary.RotaryEmbedding's surface."""
+
+    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:
+        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)
+
+
+RotaryEmbedding: type = (
+    _FlashRotaryEmbedding if _HAS_FLASH_ROTARY else _PureRotaryEmbedding
+)
+
+
+class LinearlyScaledRotaryEmbedding(RotaryEmbedding):
+    """RoPE with linear interpolation of position indices.
+
+    Used for evo2_7b (1M), evo2_7b_262k, evo2_20b, evo2_40b. Positions are
+    divided by ``scaling_factor`` before cos/sin tables are computed.
+    """
+
+    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):
+        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:
+            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."""
+    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_evo2.py

@@ -0,0 +1,119 @@
+"""ByteTokenizer for Evo2 (StripedHyena2).
+
+Wraps raw UTF-8 byte values into a vocab of size 512 (only [0, 255] are real
+bytes; the upper range pads out to match the model's vocab embedding size).
+Mirrors vortex.model.tokenizer.CharLevelTokenizer exactly:
+
+  * eod / eos = byte 0 (chr(0))
+  * pad       = byte 1 (chr(1))
+
+No real special tokens are added at encoding time -- Evo2 is a pure byte-level
+model and downstream pooling should treat every non-pad position as a real
+token.
+"""
+
+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 for Evo2 (vocab_size = 512)."""
+
+    def __init__(self, byte_level: bool = True, **kwargs):
+        kwargs.setdefault("pad_token", chr(1))
+        super().__init__(byte_level=byte_level, **kwargs)
+        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:
+        # Matches vortex CharLevelTokenizer.clamp.
+        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:
+        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_evo2.ByteTokenizer",
+      null
+    ]
+  },
+  "tokenizer_class": "ByteTokenizer",
+  "model_max_length": 8192,
+  "pad_token": "\u0001",
+  "byte_level": true
+}