Bit-exact with vortex: checkpoint RoPE frequencies, trained FP8 scales, fp8_recipe option

README.md

@@ -31,7 +31,7 @@ A clean, minimal HuggingFace port of [Evo 2 20B](https://huggingface.co/arcinsti
 2. **No public attention-weight API.** The default reference path uses Flash Attention and discards the `(B, H, T, T)` attention matrix.
 3. **`evo2` + `vortex` packages mandatory** even for inference.
 
-This repo fixes all three. Its converted **bf16 SDPA path** is bit-exact with a matching bf16 SDPA vortex reference configuration (`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.
+This repo fixes all three. With `attn_implementation="flash_attention_2"` its outputs are bit-exact with the original vortex model (see [Parity Verification](#parity-verification)). Loads with `from_pretrained` and `trust_remote_code=True` - no `evo2` / `vortex` install needed.
 
 ## Architecture
 
@@ -60,12 +60,12 @@ This repo fixes all three. Its converted **bf16 SDPA path** is bit-exact with a
 
 ## Parity Verification
 
-Hidden-state representations verified **bit-exact** (`max_abs_diff = 0.000e+00`) to the vortex reference at every block output after reconstructing the published checkpoint's bf16-rounded `inv_freq` from `base` and `dim` in fp32 on both sides. Parity uses `attn_implementation="sdpa"` in bf16, with FP8 input projections disabled on both sides where applicable. Logits from `Evo2ForCausalLM` were also bit-exact (top-1 agreement: 128/128 positions on a 128-byte ACGT input). This verifies conversion fidelity under that controlled configuration; it does not claim bit-exact output against the upstream default Flash Attention / FP8 execution path. Verified on H100 with PyTorch 2.7 / CUDA 12.
+With `attn_implementation="flash_attention_2"`, logits and every block output are **bit-exact** (`max_abs_diff = 0.000e+00`) against the original vortex model, with the FP8 projections both on and off. We checked this on Evo 2 1B on H100, with PyTorch 2.7 and 2.12 (Transformer Engine 2.5 and 2.17); the other sizes use the same code and checkpoint format. `"sdpa"` is also bit-exact with vortex's own SDPA path (`use_flash_attn=False`) when both use the same SDPA backend. PyTorch 2.7 picks its flash backend on H100, while 2.12 picks cuDNN; on newer PyTorch, wrap the forward in `torch.nn.attention.sdpa_kernel(SDPBackend.FLASH_ATTENTION)` to match. Against the default flash path, `"sdpa"` and `"eager"` agree closely but not bit for bit, since they use different attention kernels and the FP8 rounding amplifies the difference.
 
-Two non-obvious correctness fixes were required versus a naive port (see Implementation Notes for details):
+Two details matter for matching the original exactly (see Implementation Notes):
 
-1. **`inv_freq` recomputation.** A checkpoint-loaded rotary `inv_freq` can contain bf16-rounded values even when its destination buffer is fp32. The remote model code treats `inv_freq` as non-persistent runtime state and reconstructs it in fp32 from `base` and `dim` during construction, after checkpoint loading, and whenever the RoPE cache is rebuilt.
-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.
+1. **RoPE frequencies.** The checkpoint stores `inv_freq` in bf16, and vortex runs RoPE on those rounded values. The port rebuilds the same rounded values from `base` and `dim`. Using the exact values instead changes the frequencies by up to ~0.4%, which is small but enough to break parity.
+2. **Trained FP8 scales.** The original checkpoint carries per-layer FP8 scaling factors. The port loads them and keeps them fixed at inference, as vortex does. Without them Transformer Engine starts every scale at 1.0 and the outputs drift.
 
 ## Related Models
 
@@ -215,6 +215,12 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 ```
 
+> **Note on training across GPUs.** The FP8 projections use the same
+> Transformer Engine recipe as the original Evo 2, which can't be sharded
+> across GPUs when gradients are on. Inference is unaffected. To fine-tune on
+> multiple GPUs, load with `fp8_recipe="current"`; its outputs differ slightly
+> from the original's, so keep the default for inference.
+
 ### 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-recipes/tree/main/recipes/evo2_megatron#fine-tuning-from-an-existing-checkpoint).
@@ -222,17 +228,17 @@ This HuggingFace port has **not** been tested for fine-tuning - it's verified on
 ## Implementation Notes
 
 - **HF base-model contract.** `base_model_prefix = "backbone"` exposes the raw `StripedHyena2` module through `.base_model`. Its public `forward` accepts standard HF arguments (`input_ids`, `attention_mask`, `past_key_values`, `use_cache`, `output_hidden_states`, `output_attentions`, `return_dict`) and returns `BaseModelOutputWithPast`; `AutoModel` delegates to this same path.
-- **`inv_freq` kept exact in fp32 (critical for parity).** The geometric series `inv_freq[i] = 1 / base^(2i/dim)` loses ~7 bits of mantissa when rounded to bf16, which shifts the cos/sin tables and adds Q/K error at every attention layer. `inv_freq` is non-persistent runtime state: the remote code reconstructs it from `base` and `dim` in fp32 after loading and before rebuilding the RoPE cache. Rotary multiplication is also performed in fp32 before casting back to the activation dtype, matching Flash Attention 2.
+- **RoPE `inv_freq` (critical for parity).** `inv_freq` is non-persistent runtime state: the remote code rebuilds it from `base` and `dim` after loading and whenever the RoPE cache is rebuilt. The values are rounded through bf16 to match the checkpoint and the original vortex runtime, then held in fp32 so HF dtype casts leave them alone. Rotary multiplication is also performed in fp32 before casting back to the activation dtype, matching Flash Attention 2.
 - **`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, covered by `_keep_in_fp32_modules`, and explicitly restored to fp32 after `from_pretrained` completes.
 - **`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`, the same attention backend used by the default Arc Institute inference path; faster on long sequences; requires `flash-attn` installed.
+  - `"sdpa"`: calls `F.scaled_dot_product_attention`. Bit-exact with vortex's `use_flash_attn=False` path when both use the same SDPA backend (see Parity Verification).
+  - `"flash_attention_2"`: calls `flash_attn.flash_attn_qkvpacked_func`, the same attention backend used by the default Arc Institute inference path, and bit-exact with it; 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.
+- **`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. The trained FP8 scaling factors are stored as `fp8_*` buffers and handed to TE before the first FP8 call.
 - **Generation cache.** Unpadded prompts use recurrent MHA KV plus HCL/HCM/HCS FIR/IIR state matching the upstream Vortex cache path. Padded batches use full recomputation. KV tensors grow geometrically with the decoded length instead of allocating `max_seqlen` up front.
 - **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`. `transformer-engine[pytorch]` is required for this variant's FP8 input projections. `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"`).
+- **Dependencies.** `torch`, `transformers`, `numpy`, `safetensors`, `huggingface_hub`. `transformer-engine[pytorch]` is required for this variant's FP8 input projections. `accelerate` is optional but recommended if you want to load with `device_map="auto"` for multi-GPU sharding. `flash_attn` is optional: it's used for `attn_implementation="flash_attention_2"` and, when installed, for the RoPE kernel. Without it the model falls back to a pure-PyTorch RoPE with identical results. flash-attn 2.7 and 2.8 both work.
 
 ## Citation
 

configuration_evo2.py

@@ -78,6 +78,9 @@ class Evo2Config(PretrainedConfig):
         inference_mode: bool = True,
         # Projection precision (TE FP8 - required for 1b / 20b / 40b)
         use_fp8_input_projections: bool = False,
+        # TE scaling recipe for those projections: "delayed" matches the original
+        # Evo 2 (vortex); "current" also works when training across GPUs.
+        fp8_recipe: str = "delayed",
         # Backend toggles
         use_cache: bool = True,
         # Vortex used this flag at runtime. Our HF port ignores it entirely
@@ -161,6 +164,9 @@ class Evo2Config(PretrainedConfig):
         self.inference_mode = inference_mode
 
         self.use_fp8_input_projections = use_fp8_input_projections
+        if fp8_recipe not in ("delayed", "current"):
+            raise ValueError(f"fp8_recipe must be 'delayed' or 'current', got {fp8_recipe!r}")
+        self.fp8_recipe = fp8_recipe
 
         self.use_cache = use_cache
         self.use_flash_attn = use_flash_attn

hyena.py

@@ -412,6 +412,7 @@ class ParallelGatedConvBlock(nn.Module):
             bias=config.qkv_proj_bias,
             init_method=torch.nn.init.xavier_uniform_,
             use_fp8=config.get("use_fp8_input_projections", False),
+            fp8_recipe=config.get("fp8_recipe", "delayed"),
         )
 
         self.out_filter_dense = nn.Linear(

layers.py

@@ -143,8 +143,28 @@ if HAS_TE:
         telinear.get_workspace = get_workspace
         _TE_WORKSPACE_FIXUP_APPLIED = True
 
-    def set_format_recipe():
+    # Trained FP8 scaling state from the original checkpoint. TE's own
+    # _extra_state is a pickled blob, so we keep the fp32 values as int32 bit
+    # patterns instead: integer tensors survive dtype casts (bf16 loading,
+    # .to()) untouched. Shapes follow the HYBRID DelayedScaling recipe.
+    FP8_STATE_SHAPES = {
+        "fp8_scale_fwd": (3,),
+        "fp8_amax_history_fwd": (16, 3),
+        "fp8_scale_bwd": (2,),
+        "fp8_amax_history_bwd": (16, 2),
+    }
+
+    def set_format_recipe(kind: str = "delayed"):
+        """FP8 format and scaling recipe.
+
+        "delayed" is the recipe the original Evo 2 uses. Its scaling state is
+        shared across layers, so a gradient-enabled forward over layers on
+        different GPUs fails; "current" keeps no shared state and works there.
+        """
         fp8_format = Format.HYBRID
+        if kind == "current":
+            from transformer_engine.common.recipe import Float8CurrentScaling
+            return fp8_format, Float8CurrentScaling(fp8_format=fp8_format)
         fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max")
         return fp8_format, fp8_recipe
 
@@ -162,13 +182,14 @@ if HAS_TE:
             bias: bool = True,
             skip_bias_add: bool = False,
             use_fp8: bool = False,
+            fp8_recipe: str = "delayed",
             **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()
+                self.fp8_format, self.fp8_recipe = set_format_recipe(fp8_recipe)
 
             super().__init__(
                 in_features=input_size,
@@ -184,9 +205,43 @@ if HAS_TE:
                 return_bias=self.te_return_bias,
                 **kwargs,
             )
+            if use_fp8:
+                # Zeros mean "no trained state": TE then starts from scale 1.
+                for name, shape in FP8_STATE_SHAPES.items():
+                    self.register_buffer(name, torch.zeros(shape, dtype=torch.int32))
+                self._fp8_state_restored = False
+
+        def _restore_trained_fp8_state(self):
+            """Load the checkpoint's FP8 scales into TE, as vortex does.
+
+            With the delayed recipe, TE only updates scales when gradients are
+            on, so inference runs with these trained values, like the original.
+            """
+            self._fp8_state_restored = True
+            if not self.fp8_recipe.delayed() or not bool(self.fp8_scale_fwd.any()):
+                return
+            import pickle
+
+            def fp32(t):
+                return t.detach().cpu().view(torch.float32)
+
+            state = {
+                "scale_fwd": fp32(self.fp8_scale_fwd),
+                "amax_history_fwd": fp32(self.fp8_amax_history_fwd),
+                "scale_bwd": fp32(self.fp8_scale_bwd),
+                "amax_history_bwd": fp32(self.fp8_amax_history_bwd),
+                "recipe": self.fp8_recipe,
+                "extra_fp8_variables": {"fp8_checkpoint": True, "num_gemms": 1,
+                                        "fp8_max_fwd": 448, "fp8_max_bwd": 57344},
+            }
+            blob = torch.frombuffer(bytearray(pickle.dumps(state)), dtype=torch.uint8)
+            with torch.cuda.device(self.weight.device):
+                self.set_extra_state(blob)
 
         def forward(self, x):
             if self.use_fp8_input_projections:
+                if not self._fp8_state_restored:
+                    self._restore_trained_fp8_state()
                 with te.fp8_autocast(enabled=True, fp8_recipe=self.fp8_recipe):
                     out = super().forward(x)
             else:

model-fp8-state.safetensors

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

model.safetensors.index.json

@@ -6,261 +6,741 @@
   "weight_map": {
     "backbone.blocks.0.filter.h": "model-00001-of-00008.safetensors",
     "backbone.blocks.0.filter.short_filter_weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.0.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.0.mlp.l1.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.0.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.0.mlp.l2.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.0.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.0.mlp.l3.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.0.out_filter_dense.bias": "model-00001-of-00008.safetensors",
+    "backbone.blocks.0.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.0.out_filter_dense.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.0.post_norm.scale": "model-00001-of-00008.safetensors",
     "backbone.blocks.0.pre_norm.scale": "model-00001-of-00008.safetensors",
+    "backbone.blocks.0.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.0.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.0.projections.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.1.filter.D": "model-00001-of-00008.safetensors",
     "backbone.blocks.1.filter.h": "model-00001-of-00008.safetensors",
     "backbone.blocks.1.filter.short_filter_weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.1.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.1.mlp.l1.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.1.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.1.mlp.l2.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.1.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.1.mlp.l3.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.1.out_filter_dense.bias": "model-00001-of-00008.safetensors",
+    "backbone.blocks.1.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.1.out_filter_dense.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.1.post_norm.scale": "model-00001-of-00008.safetensors",
     "backbone.blocks.1.pre_norm.scale": "model-00001-of-00008.safetensors",
+    "backbone.blocks.1.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.1.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.1.projections.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.10.inner_mha_cls.Wqkv.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.inner_mha_cls.Wqkv.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.inner_mha_cls.Wqkv.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.inner_mha_cls.Wqkv.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.10.inner_mha_cls.Wqkv.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.10.inner_mha_cls.out_proj.bias": "model-00004-of-00008.safetensors",
     "backbone.blocks.10.inner_mha_cls.out_proj.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.10.inner_mha_cls.rotary_emb.inv_freq": "model-00004-of-00008.safetensors",
+    "backbone.blocks.10.mixer.dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mixer.dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mixer.dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mixer.dense.fp8_scale_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.10.mlp.l1.weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.10.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.10.mlp.l2.weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.10.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.10.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.10.mlp.l3.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.10.post_norm.scale": "model-00004-of-00008.safetensors",
     "backbone.blocks.10.pre_norm.scale": "model-00004-of-00008.safetensors",
     "backbone.blocks.11.filter.h": "model-00004-of-00008.safetensors",
     "backbone.blocks.11.filter.short_filter_weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.11.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.11.mlp.l1.weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.11.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.11.mlp.l2.weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.11.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.11.mlp.l3.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.11.out_filter_dense.bias": "model-00004-of-00008.safetensors",
+    "backbone.blocks.11.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.11.out_filter_dense.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.11.post_norm.scale": "model-00004-of-00008.safetensors",
     "backbone.blocks.11.pre_norm.scale": "model-00004-of-00008.safetensors",
+    "backbone.blocks.11.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.11.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.11.projections.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.12.filter.D": "model-00005-of-00008.safetensors",
     "backbone.blocks.12.filter.h": "model-00005-of-00008.safetensors",
     "backbone.blocks.12.filter.short_filter_weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.12.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.12.mlp.l1.weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.12.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.12.mlp.l2.weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.12.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.12.mlp.l3.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.12.out_filter_dense.bias": "model-00004-of-00008.safetensors",
+    "backbone.blocks.12.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.12.out_filter_dense.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.12.post_norm.scale": "model-00005-of-00008.safetensors",
     "backbone.blocks.12.pre_norm.scale": "model-00004-of-00008.safetensors",
+    "backbone.blocks.12.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.12.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.12.projections.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.filter.D": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.filter.log_poles": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.filter.residues": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.filter.short_filter_weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.13.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.13.mlp.l1.weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.13.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.13.mlp.l2.weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.13.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.13.mlp.l3.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.out_filter_dense.bias": "model-00005-of-00008.safetensors",
+    "backbone.blocks.13.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.13.out_filter_dense.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.post_norm.scale": "model-00005-of-00008.safetensors",
     "backbone.blocks.13.pre_norm.scale": "model-00005-of-00008.safetensors",
+    "backbone.blocks.13.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.13.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.13.projections.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.14.filter.h": "model-00005-of-00008.safetensors",
     "backbone.blocks.14.filter.short_filter_weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.14.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.14.mlp.l1.weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.14.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.14.mlp.l2.weight": "model-00005-of-00008.safetensors",
+    "backbone.blocks.14.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.14.mlp.l3.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.14.out_filter_dense.bias": "model-00005-of-00008.safetensors",
+    "backbone.blocks.14.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.14.out_filter_dense.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.14.post_norm.scale": "model-00005-of-00008.safetensors",
     "backbone.blocks.14.pre_norm.scale": "model-00005-of-00008.safetensors",
+    "backbone.blocks.14.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.14.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.14.projections.weight": "model-00005-of-00008.safetensors",
     "backbone.blocks.15.filter.D": "model-00006-of-00008.safetensors",
     "backbone.blocks.15.filter.h": "model-00006-of-00008.safetensors",
     "backbone.blocks.15.filter.short_filter_weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.15.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.15.mlp.l1.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.15.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.15.mlp.l2.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.15.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.15.mlp.l3.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.15.out_filter_dense.bias": "model-00005-of-00008.safetensors",
+    "backbone.blocks.15.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.15.out_filter_dense.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.15.post_norm.scale": "model-00006-of-00008.safetensors",
     "backbone.blocks.15.pre_norm.scale": "model-00005-of-00008.safetensors",
+    "backbone.blocks.15.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.15.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.15.projections.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.filter.D": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.filter.log_poles": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.filter.residues": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.filter.short_filter_weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.16.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.16.mlp.l1.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.16.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.16.mlp.l2.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.16.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.16.mlp.l3.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.out_filter_dense.bias": "model-00006-of-00008.safetensors",
+    "backbone.blocks.16.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.16.out_filter_dense.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.post_norm.scale": "model-00006-of-00008.safetensors",
     "backbone.blocks.16.pre_norm.scale": "model-00006-of-00008.safetensors",
+    "backbone.blocks.16.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.16.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.16.projections.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.17.inner_mha_cls.Wqkv.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.inner_mha_cls.Wqkv.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.inner_mha_cls.Wqkv.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.inner_mha_cls.Wqkv.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.17.inner_mha_cls.Wqkv.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.17.inner_mha_cls.out_proj.bias": "model-00006-of-00008.safetensors",
     "backbone.blocks.17.inner_mha_cls.out_proj.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.17.inner_mha_cls.rotary_emb.inv_freq": "model-00006-of-00008.safetensors",
+    "backbone.blocks.17.mixer.dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mixer.dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mixer.dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mixer.dense.fp8_scale_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.17.mlp.l1.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.17.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.17.mlp.l2.weight": "model-00006-of-00008.safetensors",
+    "backbone.blocks.17.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.17.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.17.mlp.l3.weight": "model-00006-of-00008.safetensors",
     "backbone.blocks.17.post_norm.scale": "model-00006-of-00008.safetensors",
     "backbone.blocks.17.pre_norm.scale": "model-00006-of-00008.safetensors",
     "backbone.blocks.18.filter.h": "model-00007-of-00008.safetensors",
     "backbone.blocks.18.filter.short_filter_weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.18.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.18.mlp.l1.weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.18.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.18.mlp.l2.weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.18.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.18.mlp.l3.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.18.out_filter_dense.bias": "model-00006-of-00008.safetensors",
+    "backbone.blocks.18.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.18.out_filter_dense.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.18.post_norm.scale": "model-00007-of-00008.safetensors",
     "backbone.blocks.18.pre_norm.scale": "model-00006-of-00008.safetensors",
+    "backbone.blocks.18.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.18.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.18.projections.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.19.filter.D": "model-00007-of-00008.safetensors",
     "backbone.blocks.19.filter.h": "model-00007-of-00008.safetensors",
     "backbone.blocks.19.filter.short_filter_weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.19.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.19.mlp.l1.weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.19.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.19.mlp.l2.weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.19.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.19.mlp.l3.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.19.out_filter_dense.bias": "model-00007-of-00008.safetensors",
+    "backbone.blocks.19.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.19.out_filter_dense.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.19.post_norm.scale": "model-00007-of-00008.safetensors",
     "backbone.blocks.19.pre_norm.scale": "model-00007-of-00008.safetensors",
+    "backbone.blocks.19.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.19.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.19.projections.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.2.filter.D": "model-00001-of-00008.safetensors",
     "backbone.blocks.2.filter.log_poles": "model-00001-of-00008.safetensors",
     "backbone.blocks.2.filter.residues": "model-00001-of-00008.safetensors",
     "backbone.blocks.2.filter.short_filter_weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.2.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.2.mlp.l1.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.2.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.2.mlp.l2.weight": "model-00001-of-00008.safetensors",
+    "backbone.blocks.2.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.2.mlp.l3.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.2.out_filter_dense.bias": "model-00001-of-00008.safetensors",
+    "backbone.blocks.2.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.2.out_filter_dense.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.2.post_norm.scale": "model-00001-of-00008.safetensors",
     "backbone.blocks.2.pre_norm.scale": "model-00001-of-00008.safetensors",
+    "backbone.blocks.2.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.2.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.2.projections.weight": "model-00001-of-00008.safetensors",
     "backbone.blocks.20.filter.D": "model-00007-of-00008.safetensors",
     "backbone.blocks.20.filter.log_poles": "model-00007-of-00008.safetensors",
     "backbone.blocks.20.filter.residues": "model-00007-of-00008.safetensors",
     "backbone.blocks.20.filter.short_filter_weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.20.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.20.mlp.l1.weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.20.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.20.mlp.l2.weight": "model-00007-of-00008.safetensors",
+    "backbone.blocks.20.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.20.mlp.l3.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.20.out_filter_dense.bias": "model-00007-of-00008.safetensors",
+    "backbone.blocks.20.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.20.out_filter_dense.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.20.post_norm.scale": "model-00007-of-00008.safetensors",
     "backbone.blocks.20.pre_norm.scale": "model-00007-of-00008.safetensors",
+    "backbone.blocks.20.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.20.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.20.projections.weight": "model-00007-of-00008.safetensors",
     "backbone.blocks.21.filter.h": "model-00008-of-00008.safetensors",
     "backbone.blocks.21.filter.short_filter_weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.21.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.21.mlp.l1.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.21.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.21.mlp.l2.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.21.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.21.mlp.l3.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.21.out_filter_dense.bias": "model-00007-of-00008.safetensors",
+    "backbone.blocks.21.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.21.out_filter_dense.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.21.post_norm.scale": "model-00008-of-00008.safetensors",
     "backbone.blocks.21.pre_norm.scale": "model-00007-of-00008.safetensors",
+    "backbone.blocks.21.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.21.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.21.projections.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.22.filter.D": "model-00008-of-00008.safetensors",
     "backbone.blocks.22.filter.h": "model-00008-of-00008.safetensors",
     "backbone.blocks.22.filter.short_filter_weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.22.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.22.mlp.l1.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.22.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.22.mlp.l2.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.22.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.22.mlp.l3.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.22.out_filter_dense.bias": "model-00008-of-00008.safetensors",
+    "backbone.blocks.22.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.22.out_filter_dense.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.22.post_norm.scale": "model-00008-of-00008.safetensors",
     "backbone.blocks.22.pre_norm.scale": "model-00008-of-00008.safetensors",
+    "backbone.blocks.22.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.22.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.22.projections.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.filter.D": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.filter.log_poles": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.filter.residues": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.filter.short_filter_weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.23.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.23.mlp.l1.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.23.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.23.mlp.l2.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.23.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.23.mlp.l3.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.out_filter_dense.bias": "model-00008-of-00008.safetensors",
+    "backbone.blocks.23.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.23.out_filter_dense.weight": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.post_norm.scale": "model-00008-of-00008.safetensors",
     "backbone.blocks.23.pre_norm.scale": "model-00008-of-00008.safetensors",
+    "backbone.blocks.23.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.23.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.23.projections.weight": "model-00008-of-00008.safetensors",
+    "backbone.blocks.3.inner_mha_cls.Wqkv.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.inner_mha_cls.Wqkv.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.inner_mha_cls.Wqkv.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.inner_mha_cls.Wqkv.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.3.inner_mha_cls.Wqkv.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.3.inner_mha_cls.out_proj.bias": "model-00001-of-00008.safetensors",
     "backbone.blocks.3.inner_mha_cls.out_proj.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.3.inner_mha_cls.rotary_emb.inv_freq": "model-00002-of-00008.safetensors",
+    "backbone.blocks.3.mixer.dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mixer.dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mixer.dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mixer.dense.fp8_scale_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.3.mlp.l1.weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.3.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.3.mlp.l2.weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.3.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.3.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.3.mlp.l3.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.3.post_norm.scale": "model-00002-of-00008.safetensors",
     "backbone.blocks.3.pre_norm.scale": "model-00001-of-00008.safetensors",
     "backbone.blocks.4.filter.h": "model-00002-of-00008.safetensors",
     "backbone.blocks.4.filter.short_filter_weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.4.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.4.mlp.l1.weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.4.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.4.mlp.l2.weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.4.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.4.mlp.l3.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.4.out_filter_dense.bias": "model-00002-of-00008.safetensors",
+    "backbone.blocks.4.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.4.out_filter_dense.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.4.post_norm.scale": "model-00002-of-00008.safetensors",
     "backbone.blocks.4.pre_norm.scale": "model-00002-of-00008.safetensors",
+    "backbone.blocks.4.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.4.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.4.projections.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.5.filter.D": "model-00002-of-00008.safetensors",
     "backbone.blocks.5.filter.h": "model-00002-of-00008.safetensors",
     "backbone.blocks.5.filter.short_filter_weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.5.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.5.mlp.l1.weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.5.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.5.mlp.l2.weight": "model-00002-of-00008.safetensors",
+    "backbone.blocks.5.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.5.mlp.l3.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.5.out_filter_dense.bias": "model-00002-of-00008.safetensors",
+    "backbone.blocks.5.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.5.out_filter_dense.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.5.post_norm.scale": "model-00002-of-00008.safetensors",
     "backbone.blocks.5.pre_norm.scale": "model-00002-of-00008.safetensors",
+    "backbone.blocks.5.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.5.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.5.projections.weight": "model-00002-of-00008.safetensors",
     "backbone.blocks.6.filter.D": "model-00003-of-00008.safetensors",
     "backbone.blocks.6.filter.log_poles": "model-00003-of-00008.safetensors",
     "backbone.blocks.6.filter.residues": "model-00003-of-00008.safetensors",
     "backbone.blocks.6.filter.short_filter_weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.6.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.6.mlp.l1.weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.6.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.6.mlp.l2.weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.6.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.6.mlp.l3.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.6.out_filter_dense.bias": "model-00002-of-00008.safetensors",
+    "backbone.blocks.6.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.6.out_filter_dense.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.6.post_norm.scale": "model-00003-of-00008.safetensors",
     "backbone.blocks.6.pre_norm.scale": "model-00002-of-00008.safetensors",
+    "backbone.blocks.6.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.6.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.6.projections.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.7.filter.h": "model-00003-of-00008.safetensors",
     "backbone.blocks.7.filter.short_filter_weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.7.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.7.mlp.l1.weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.7.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.7.mlp.l2.weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.7.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.7.mlp.l3.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.7.out_filter_dense.bias": "model-00003-of-00008.safetensors",
+    "backbone.blocks.7.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.7.out_filter_dense.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.7.post_norm.scale": "model-00003-of-00008.safetensors",
     "backbone.blocks.7.pre_norm.scale": "model-00003-of-00008.safetensors",
+    "backbone.blocks.7.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.7.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.7.projections.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.8.filter.D": "model-00003-of-00008.safetensors",
     "backbone.blocks.8.filter.h": "model-00003-of-00008.safetensors",
     "backbone.blocks.8.filter.short_filter_weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.8.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.8.mlp.l1.weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.8.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.8.mlp.l2.weight": "model-00003-of-00008.safetensors",
+    "backbone.blocks.8.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.8.mlp.l3.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.8.out_filter_dense.bias": "model-00003-of-00008.safetensors",
+    "backbone.blocks.8.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.8.out_filter_dense.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.8.post_norm.scale": "model-00003-of-00008.safetensors",
     "backbone.blocks.8.pre_norm.scale": "model-00003-of-00008.safetensors",
+    "backbone.blocks.8.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.8.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.8.projections.weight": "model-00003-of-00008.safetensors",
     "backbone.blocks.9.filter.D": "model-00004-of-00008.safetensors",
     "backbone.blocks.9.filter.log_poles": "model-00004-of-00008.safetensors",
     "backbone.blocks.9.filter.residues": "model-00004-of-00008.safetensors",
     "backbone.blocks.9.filter.short_filter_weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.9.mlp.l1.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l1.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l1.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l1.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.9.mlp.l1.weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.9.mlp.l2.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l2.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l2.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l2.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.9.mlp.l2.weight": "model-00004-of-00008.safetensors",
+    "backbone.blocks.9.mlp.l3.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l3.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l3.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.mlp.l3.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.9.mlp.l3.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.9.out_filter_dense.bias": "model-00003-of-00008.safetensors",
+    "backbone.blocks.9.out_filter_dense.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.out_filter_dense.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.out_filter_dense.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.out_filter_dense.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.9.out_filter_dense.weight": "model-00004-of-00008.safetensors",
     "backbone.blocks.9.post_norm.scale": "model-00004-of-00008.safetensors",
     "backbone.blocks.9.pre_norm.scale": "model-00003-of-00008.safetensors",
+    "backbone.blocks.9.projections.fp8_amax_history_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.projections.fp8_amax_history_fwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.projections.fp8_scale_bwd": "model-fp8-state.safetensors",
+    "backbone.blocks.9.projections.fp8_scale_fwd": "model-fp8-state.safetensors",
     "backbone.blocks.9.projections.weight": "model-00004-of-00008.safetensors",
     "backbone.embedding_layer.weight": "model-00001-of-00008.safetensors",
     "backbone.norm.scale": "model-00008-of-00008.safetensors"

modeling_evo2.py

@@ -314,8 +314,8 @@ class StripedHyena2(nn.Module):
             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.
+            ~4e-2 of noise on Q/K. We rebuild inv_freq here as vortex
+            runs it: the checkpoint's bf16 values held in fp32.
         """
         for k, p in self.named_parameters():
             if "log_poles" not in k and "residues" not in k:

rotary.py

@@ -4,6 +4,10 @@ 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.
+
+The original Evo 2 (vortex) loads ``inv_freq`` from the checkpoint, where it
+is stored in bf16, and runs RoPE on those rounded values held in fp32. We
+rebuild them the same way, so outputs match the original bit for bit.
 """
 
 from __future__ import annotations
@@ -20,9 +24,18 @@ except ImportError:
     _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 _checkpoint_inv_freq(inv_freq: torch.Tensor) -> torch.Tensor:
+    """Round exact fp32 frequencies through bf16, as stored in the checkpoint."""
+    return inv_freq.to(torch.bfloat16).to(torch.float32)
+
+
+def _rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
+    """[x1, x2] -> [-x2, x1] over halves, or over even/odd pairs if interleaved."""
+    if not interleaved:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    return torch.stack((-x2, x1), dim=-1).flatten(-2)
 
 
 @torch.no_grad()
@@ -32,14 +45,14 @@ def restore_rotary_inv_freq(
     device=None,
     invalidate_cache: bool = True,
 ) -> None:
-    """Rebuild ``inv_freq`` from RoPE metadata in fp32."""
+    """Rebuild ``inv_freq`` from RoPE metadata (bf16-rounded, held in fp32)."""
     if not hasattr(rotary_emb, "inv_freq"):
         return
     device = rotary_emb.inv_freq.device if device is None else device
     if hasattr(rotary_emb, "_compute_inv_freq"):
         inv_freq = rotary_emb._compute_inv_freq(device=device)
     else:
-        inv_freq = 1.0 / (
+        inv_freq = _checkpoint_inv_freq(1.0 / (
             float(rotary_emb.base)
             ** (
                 torch.arange(
@@ -51,7 +64,7 @@ def restore_rotary_inv_freq(
                 )
                 / int(rotary_emb.dim)
             )
-        )
+        ))
     rotary_emb.inv_freq = inv_freq.to(dtype=torch.float32)
     if invalidate_cache:
         rotary_emb._seq_len_cached = 0
@@ -65,7 +78,9 @@ def restore_rotary_inv_freq(
                 setattr(rotary_emb, name, None)
 
 
-def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+def _apply_rotary(
+    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
+) -> 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.
@@ -73,14 +88,38 @@ def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torc
     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()
+    if interleaved:
+        cos_full = torch.repeat_interleave(cos, 2, dim=-1).float()
+        sin_full = torch.repeat_interleave(sin, 2, dim=-1).float()
+    else:
+        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 = (x_rot_f * cos_full) + (_rotate_half(x_rot_f, interleaved) * sin_full)
     rotated = rotated.to(orig_dtype)
     return torch.cat((rotated, x_pass), dim=-1)
 
 
+def _set_cos_sin_tables(rope: nn.Module, freqs: torch.Tensor, seqlen: int, dtype) -> None:
+    """Fill a RoPE module's cos/sin caches from ``freqs``.
+
+    Same as flash_attn / vortex RotaryEmbedding: with xPos (``rope.scale``
+    set), queries use cos*scale and keys cos/scale, multiplied in fp32.
+    """
+    if rope.scale is None:
+        rope._cos_cached = torch.cos(freqs).to(dtype)
+        rope._sin_cached = torch.sin(freqs).to(dtype)
+        return
+    power = (
+        torch.arange(seqlen, dtype=rope.scale.dtype, device=rope.scale.device) - seqlen // 2
+    ) / rope.scale_base
+    scale = rope.scale.to(device=power.device) ** power[:, None]
+    rope._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+    rope._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+    rope._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+    rope._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+
 class _PureRotaryEmbedding(nn.Module):
     """Pure-PyTorch fallback RoPE. Mirrors flash_attn.layers.rotary.RotaryEmbedding's surface."""
 
@@ -94,10 +133,6 @@ class _PureRotaryEmbedding(nn.Module):
         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
@@ -106,17 +141,24 @@ class _PureRotaryEmbedding(nn.Module):
 
         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
+        scale = (
+            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
+            if scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale, persistent=False)
 
         self._seq_len_cached = 0
         self._cos_cached: torch.Tensor | None = None
         self._sin_cached: torch.Tensor | None = None
+        self._cos_k_cached: torch.Tensor | None = None
+        self._sin_k_cached: torch.Tensor | None = None
 
     def _compute_inv_freq(self, device=None) -> torch.Tensor:
-        return 1.0 / (
+        return _checkpoint_inv_freq(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 (
@@ -134,9 +176,7 @@ class _PureRotaryEmbedding(nn.Module):
                 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)
+            _set_cos_sin_tables(self, torch.outer(t, inv_freq), seqlen, dtype)
 
     def forward(
         self,
@@ -150,19 +190,44 @@ class _PureRotaryEmbedding(nn.Module):
         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]
+        span = slice(seqlen_offset, seqlen_offset + T)
+        cos = self._cos_cached[span][None, :, None, :]
+        sin = self._sin_cached[span][None, :, None, :]
+        if self.scale is None:
+            cos_k, sin_k = cos, sin
+        else:
+            cos_k = self._cos_k_cached[span][None, :, None, :]
+            sin_k = self._sin_k_cached[span][None, :, None, :]
         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)
+        q = _apply_rotary(q, cos, sin, self.interleaved)
+        k = _apply_rotary(k, cos_k, sin_k, self.interleaved)
         return torch.stack((q, k, v), dim=2)
 
 
 if _HAS_FLASH_ROTARY:
+    import inspect as _inspect
+
+    # flash-attn >= 2.8 dropped this argument and always uses fp32 positions.
+    _FLASH_TAKES_POS_IDX = (
+        "pos_idx_in_fp32" in _inspect.signature(_FlashRotaryEmbedding.__init__).parameters
+    )
+
     class RotaryEmbedding(_FlashRotaryEmbedding):
-        """Flash RoPE with exact fp32 frequency reconstruction."""
+        """Flash RoPE with the checkpoint's frequencies."""
+
+        def __init__(self, *args, pos_idx_in_fp32: bool = True, **kwargs):
+            if _FLASH_TAKES_POS_IDX:
+                kwargs["pos_idx_in_fp32"] = pos_idx_in_fp32
+            elif not pos_idx_in_fp32:
+                raise ValueError(
+                    "pos_idx_in_fp32=False needs flash-attn < 2.8, which always uses fp32 "
+                    "positions; uninstall flash-attn to use the PyTorch fallback instead."
+                )
+            super().__init__(*args, **kwargs)
+            self.pos_idx_in_fp32 = pos_idx_in_fp32
+
+        def _compute_inv_freq(self, device=None):
+            return _checkpoint_inv_freq(super()._compute_inv_freq(device=device))
 
         def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
             needs_update = (
@@ -218,21 +283,7 @@ class LinearlyScaledRotaryEmbedding(RotaryEmbedding):
             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)
+        _set_cos_sin_tables(self, torch.outer(t, inv_freq), seqlen, dtype)
 
 
 def swap_mha_rope(mha, new_rope=LinearlyScaledRotaryEmbedding, kwargs_new_rope=None):