Add config, code, tokenizer, index and model card

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or Derivative
+          Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

NOTICE

@@ -0,0 +1,25 @@
+CENO
+Copyright (c) 2025-2026 CENO Authors. All rights reserved.
+
+This product includes model code adapted from the Nemotron-H and Mamba2
+implementations distributed in Hugging Face Transformers
+(https://github.com/huggingface/transformers), with Hugging Face and NVIDIA
+copyright notices preserved. Hugging Face Transformers is licensed under the
+Apache License, Version 2.0.
+
+The CENO configuration and modeling modules under `ceno_model/ceno_hf/` are
+adapted from those implementations and are distributed under the same Apache
+License, Version 2.0, included in this distribution as the `LICENSE` file.
+
+The tokenizer module is derived from Arc Institute's Evo2 CharLevelTokenizer,
+also distributed under the Apache License, Version 2.0. Its upstream copyright
+notice is preserved in `ceno_model/ceno_hf/ceno_tokenizer.py`.
+
+Nemotron-H is © NVIDIA Corporation. CENO and CENO-P are derivative works that
+build on the Nemotron-H Mamba/Attention/MoE hybrid backbone and add
+multi-species alignment (MSA) post-training. All NVIDIA copyright notices and
+license terms are preserved as required by the Apache License, Version 2.0.
+
+The CENO normalization and CP merge helpers include adaptations of NVIDIA
+TransformerEngine (Apache-2.0). Its license is included in the corresponding
+third_party/TransformerEngine-LICENSE file.

README.md

@@ -0,0 +1,101 @@
+---
+license: apache-2.0
+library_name: transformers
+tags:
+  - dna
+  - genomics
+  - dna-language-model
+  - mamba
+  - moe
+pipeline_tag: text-generation
+---
+
+# CENO-9B-A3B-base
+
+**CENO-9B-A3B-base** is the base pretraining checkpoint of the 9B **CENO** DNA foundation model — a causal language model over genomic sequence built on a
+Nemotron-H **Mamba / Attention / Mixture-of-Experts** hybrid backbone (no MSA inputs).
+
+It is part of the **CENO** DNA foundation model family. Model code, the VEP pipeline, and a
+generation demo live in the companion [CENO code repository](https://github.com/CladeTeam/CENO).
+This checkpoint is standalone-loadable with `trust_remote_code=True` — the model code is
+bundled here.
+
+## Model details
+
+|  |  |
+|---|---|
+| Family | CENO (base) |
+| Training stage | Base pretraining (stage 2) |
+| Parameters | 8.9B (8928.6M) |
+| Precision | bfloat16 |
+| `model_type` | `ceno` |
+| Architecture class | `CENOForCausalLM` |
+| Auto-map (model) | `modeling_ceno.CENOForCausalLM` |
+| Auto-map (tokenizer) | `ceno_tokenizer.CENOCharLevelTokenizer` |
+
+## Architecture
+
+| Property | Value |
+|---|---|
+| Hidden layers | 42 |
+| Hidden size | 2560 |
+| Attention heads | 20 |
+| Intermediate size | 10240 |
+| Experts (MoE) | 8 (top-2 per token) |
+| Vocabulary | 512 (byte / character-level) |
+
+The backbone is a Mamba / Attention / Mixture-of-Experts hybrid (Nemotron-H
+architecture). The tokenizer is character-level, mapping DNA bases to their ASCII
+byte codes.
+
+## Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ckpt = "CladeTeam/CENO-9B-A3B-base"
+model = AutoModelForCausalLM.from_pretrained(ckpt, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(ckpt, trust_remote_code=True)
+
+ids = tokenizer.encode("ATCGATCG", return_tensors="pt")
+# out = model.generate(ids, max_new_tokens=128)   # needs a CUDA GPU (Mamba kernels)
+```
+
+> The Mamba layers require CUDA kernels, so forward passes and generation need a GPU.
+> Config, tokenizer, and weight loading are CPU-safe.
+
+## Intended use
+
+- **Base checkpoints (`CENO-*`)** — genomic-sequence generation and embedding extraction;
+  downstream adaptation (fine-tuning, probing) for genomics tasks.
+- **MSA checkpoints (`CENO-P-*`)** — variant effect prediction (VEP) by scoring wild-type
+  vs. variant sequences with delta log-likelihood. See the TraitGym VEP example in the
+  [CENO code repository](https://github.com/CladeTeam/CENO).
+
+## License
+
+Apache-2.0. The bundled model code is derived from NVIDIA's Nemotron-H Hugging Face
+implementation (Apache-2.0); the tokenizer is derived from the Arc Institute Evo2
+`CharLevelTokenizer` (Apache-2.0). See the `LICENSE` and `NOTICE` files in this repository
+for full attribution.
+
+
+<!-- CENO_INFERENCE_CURRENT -->
+## Current inference implementation
+
+Current source and usage instructions
+are maintained at [CladeTeam/CENO](https://github.com/CladeTeam/CENO).
+
+Loading
+uses the checkpoint dtype by default; explicit caller dtype overrides are honored.
+Use `torch_dtype="auto"`, `.cuda().eval()` and `torch.inference_mode()` for raw
+forward calls. `generate()` enters no-grad mode automatically.
+
+The HF precision path is single-GPU, unpadded, and permits at most 131072 tokens
+per prefill. Use the repository's vLLM entry for larger prefills.
+
+The tested HF stack is PyTorch 2.7.0a0 (NVIDIA 25.03), Transformers 4.48.3,
+mamba-ssm 2.2.4, causal-conv1d 1.2.2.post1 and FlashAttention 2.7.3. HF and vLLM
+use separate environments. Numerical correctness is not a guarantee of
+repetition-free or biologically functional generation.
+<!-- /CENO_INFERENCE_CURRENT -->

ceno_cp_math.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# NVIDIA TransformerEngine abbdd76 CP attention arithmetic, Triton adaptation.
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import libdevice
+
+
+@triton.jit
+def _lse_add(Total, Part, T: tl.constexpr, FULL: tl.constexpr,
+             OFFSET: tl.constexpr, H: tl.constexpr, BLOCK: tl.constexpr):
+    off=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
+    head=off//T;row=off%T;mask=off<H*T
+    ptr=Total+head*FULL+row+OFFSET
+    total=tl.load(ptr,mask=mask,other=0)
+    part=tl.load(Part+off,mask=mask,other=0).to(tl.float64)
+    hi=tl.maximum(total,part);lo=tl.minimum(total,part)
+    value=hi+libdevice.log1p(libdevice.exp(lo-hi))
+    tl.store(ptr,value,mask=mask)
+
+
+@triton.jit
+def _merge(Out,Part,Total,Lse,N:tl.constexpr,T:tl.constexpr,FULL:tl.constexpr,
+           OFFSET:tl.constexpr,ADD:tl.constexpr,H:tl.constexpr,D:tl.constexpr,
+           BLOCK:tl.constexpr):
+    off=tl.program_id(0)*BLOCK+tl.arange(0,BLOCK)
+    row=off//(H*D);head=(off//D)%H;mask=off<N
+    total=tl.load(Total+head*FULL+row+OFFSET,mask=mask,other=0)
+    lse=tl.load(Lse+head*T+row,mask=mask,other=0)
+    # TE's Torch 2.7 compiled kernel uses approximate exp and fused add/multiply.
+    # Current Torch's eager/compiled exp does not reproduce this BF16 rounding.
+    weight=tl.exp(lse-total)
+    value=tl.load(Part+off,mask=mask,other=0).to(tl.float32)
+    dest=Out+OFFSET*H*D+off
+    if ADD:
+        old=tl.load(dest,mask=mask,other=0).to(tl.float32)
+        result=tl.fma(value,weight,old)
+    else:result=value*weight
+    tl.store(dest,result,mask=mask)
+
+
+def combine(parts, offsets):
+    """Merge in ring order; every output update is stored back to BF16."""
+    first,lse=parts[0];length,h,d=first.shape
+    total=lse.double().clone()
+    for (_,part),offset in zip(parts[1:],offsets[1:]):
+        _lse_add[(triton.cdiv(part.numel(),256),)](
+            total,part,part.shape[1],length,offset,h,256,enable_fp_fusion=False)
+    total=total.float();out=torch.empty_like(first)
+    for i,((value,part),offset) in enumerate(zip(parts,offsets)):
+        _merge[(triton.cdiv(value.numel(),1024),)](
+            out,value,total,part,value.numel(),len(value),length,offset,i>0,h,d,1024,
+            enable_fp_fusion=False)
+    return out

ceno_norm.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# CENO Triton adaptation of TransformerEngine abbdd76 RMSNorm reduction order.
+"""Explicit, batch-independent TE RMSNorm arithmetic for BF16 CENO widths 512/1024.
+
+The training TE build chooses two reduction algorithms according to row count.
+Keep that choice explicit: ``te-general`` reproduces the original 5385-row
+reference; ``te-tuned`` reproduces the aligned training-shaped kernel. Never
+switch based on a serving batch's size. This is not a generic TE replacement.
+
+Source: NVIDIA/TransformerEngine@abbdd76, common/normalization/rmsnorm.
+See third_party/TransformerEngine-LICENSE for the upstream Apache-2.0 license.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _xor(value, offset: tl.constexpr):
+    # One warp owns one row. Changing num_warps/layout changes this reduction.
+    return tl.inline_asm_elementwise(
+        "shfl.sync.bfly.b32 $0, $1, $2, 31, -1;",
+        constraints="=f,f,r",
+        args=[value, offset],
+        dtype=tl.float32,
+        is_pure=True,
+        pack=1,
+    )
+
+
+@triton.jit
+def _te_norm(X, W, Y, EPS: tl.constexpr, TUNED: tl.constexpr, WIDTH: tl.constexpr):
+    row = tl.program_id(0)
+    lanes = tl.arange(0, 32)
+    mean = tl.full((32,), 0, tl.float32)
+    if TUNED:
+        for i in tl.static_range(WIDTH // 32):
+            cols = lanes * 8 + (i // 8) * 256 + i % 8
+            mean = mean + tl.load(X + row * WIDTH + cols).to(tl.float32)
+        for offset in tl.static_range(5):
+            mean = mean + _xor(mean, 1 << offset)
+        mean = mean * (1.0 / WIDTH)
+    variance = tl.full((32,), 0, tl.float32)
+    for i in tl.static_range(WIDTH // 32):
+        cols = lanes * 8 + (i // 8) * 256 + i % 8
+        x = tl.load(X + row * WIDTH + cols).to(tl.float32)
+        delta = x - mean if TUNED else x
+        variance = tl.fma(delta, delta, variance)
+    for offset in tl.static_range(5):
+        variance = variance + _xor(variance, 1 << offset)
+    moment = variance * (1.0 / WIDTH)
+    if TUNED:
+        moment = tl.fma(mean, mean, moment)
+    scale = tl.rsqrt(moment + EPS)
+    for i in tl.static_range(WIDTH // 32):
+        cols = lanes * 8 + (i // 8) * 256 + i % 8
+        x = tl.load(X + row * WIDTH + cols).to(tl.float32)
+        weight = tl.load(W + cols).to(tl.float32)
+        tl.store(Y + row * WIDTH + cols, (scale * x) * weight)
+
+
+def _te_rms_norm_impl(
+    x: torch.Tensor, weight: torch.Tensor, epsilon: float, mode: str
+) -> torch.Tensor:
+    if mode not in ("te-general", "te-tuned"):
+        raise ValueError(f"Unsupported TE RMSNorm mode: {mode}")
+    if (
+        not x.is_cuda
+        or x.dtype != torch.bfloat16
+        or x.shape[-1] not in (512, 1024)
+        or weight.dtype != x.dtype
+        or weight.shape != (x.shape[-1],)
+        or weight.device != x.device
+        or epsilon <= 0
+    ):
+        raise ValueError("TE RMSNorm requires CUDA BF16 width 512/1024 and a positive epsilon")
+    x, weight = x.contiguous(), weight.contiguous()
+    out = torch.empty_like(x)
+    rows = x.numel() // x.shape[-1]
+    if rows:
+        _te_norm[(rows,)](x, weight, out, epsilon, mode == "te-tuned", x.shape[-1], num_warps=1)
+    return out
+
+
+def _te_rms_norm_fake(x, weight, epsilon, mode):
+    return torch.empty_like(x, memory_format=torch.contiguous_format)
+
+
+# Multiple HF Hub repositories can import this file under different dynamic names.
+# Register the shared operator once per process.
+if hasattr(torch.ops.ceno, "te_rms_norm"):
+    te_rms_norm = torch.ops.ceno.te_rms_norm.default
+else:
+    te_rms_norm = torch.library.custom_op("ceno::te_rms_norm", mutates_args=())(
+        _te_rms_norm_impl
+    )
+    te_rms_norm.register_fake(_te_rms_norm_fake)

ceno_tokenizer.py

@@ -0,0 +1,571 @@
+# coding=utf-8
+# Copyright (c) 2025, Arc Institute. All rights reserved.
+# Copyright (c) 2026, CENO Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Byte-level (character-level) tokenizer for CENO."""
+
+import json
+import os
+from typing import List, Optional, Tuple, Union, Dict, Any
+import numpy as np
+import torch
+from transformers import PreTrainedTokenizer
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.json"}
+
+
+class CENOCharLevelTokenizer(PreTrainedTokenizer):
+    """
+    HuggingFace-style byte-level (character-level) tokenizer for CENO.
+    
+    This tokenizer converts text directly to byte values using numpy's fromstring,
+    which is perfect for DNA sequences and other character-level tasks.
+    
+    Args:
+        vocab_size (int): Size of the vocabulary (default: 512)
+        eos_token (str): End of sequence token
+        pad_token (str): Padding token  
+        unk_token (str): Unknown token
+        **kwargs: Additional arguments passed to PreTrainedTokenizer
+    """
+    
+    vocab_files_names = VOCAB_FILES_NAMES
+    
+    def __init__(
+        self,
+        vocab_size: int = 512,
+        eos_token: str = "<eos>",
+        pad_token: str = "<pad>",
+        unk_token: str = "<unk>",
+        **kwargs
+    ):
+        self._vocab_size = vocab_size
+        self.eod_id = 0
+        self.eos_id = 0
+        self.pad_id = 1
+        self.unk_id = 2
+        
+        # Build vocabulary - builds the CENO character mapping
+        self._vocab = self._build_vocab()
+        self._id_to_token = {v: k for k, v in self._vocab.items()}
+        
+        super().__init__(
+            eos_token=eos_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            **kwargs
+        )
+    
+    def _build_vocab(self) -> Dict[str, int]:
+        """Build vocabulary mapping characters to IDs"""
+        vocab = {}
+        
+        # Add special tokens
+        vocab["<unk>"] = 2
+        vocab["<pad>"] = 1
+        vocab["<eos>"] = 0
+        
+        # Add printable ASCII characters (32-126)
+        for i in range(32, min(127, self._vocab_size)):
+            vocab[chr(i)] = i
+        
+        # Add extended byte values as special tokens
+        for i in range(127, self._vocab_size):
+            vocab[f"<byte_{i}>"] = i
+            
+        return vocab
+    
+    def clamp(self, n: int) -> int:
+        """Clamp token ID to valid range, matching the CENO tokenizer implementation"""
+        return max(0, min(n, self._vocab_size - 1))
+    
+    @property
+    def vocab_size(self) -> int:
+        """Return vocabulary size"""
+        return self._vocab_size
+    
+    def get_vocab(self) -> Dict[str, int]:
+        """Return vocabulary dictionary"""
+        return self._vocab.copy()
+    
+    def _tokenize(self, text: str) -> List[int]:
+        """
+        Tokenize text using numpy's fromstring (byte-level tokenization).
+        Byte-level tokenization: text is converted directly to its ASCII byte IDs.
+        """
+        # Convert text to byte array using numpy (matches the CENO implementation)
+        token_ids = np.frombuffer(text.encode("utf-8"), dtype=np.uint8).tolist()
+        return token_ids
+    
+    def _convert_token_to_id(self, token: Union[str, int]) -> int:
+        """Convert token to ID"""
+        if isinstance(token, int):
+            return self.clamp(token)
+        
+        # Handle string tokens
+        if token in self._vocab:
+            return self._vocab[token]
+        
+        # Handle single characters
+        if len(token) == 1:
+            return self.clamp(ord(token))
+        
+        # Handle byte tokens
+        if token.startswith("<byte_") and token.endswith(">"):
+            try:
+                byte_val = int(token[6:-1])
+                return self.clamp(byte_val)
+            except ValueError:
+                pass
+        
+        # Return unknown token ID
+        return self._vocab.get(self.unk_token, 0)
+    
+    def _convert_id_to_token(self, index: int) -> str:
+        """Convert ID to token, CENO decode-token behavior"""
+        clamped_index = self.clamp(index)
+        
+        # Handle special cases before interpreting byte values.
+        if clamped_index == self.eos_id:
+            return self.eos_token
+        if clamped_index == self.pad_id:
+            return self.pad_token
+        if clamped_index == self.unk_id:
+            return self.unk_token
+        
+        # Convert to character if in printable range
+        if 32 <= clamped_index <= 126:
+            return chr(clamped_index)
+        
+        # Return byte token for extended range
+        return f"<byte_{clamped_index}>"
+    
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """Convert tokens back to string"""
+        result = []
+        for token in tokens:
+            if token in [self.pad_token, self.eos_token, self.unk_token]:
+                continue
+            elif token.startswith("<byte_") and token.endswith(">"):
+                try:
+                    byte_val = int(token[6:-1])
+                    result.append(chr(self.clamp(byte_val)))
+                except (ValueError, OverflowError):
+                    continue
+            else:
+                result.append(token)
+        return "".join(result)
+    
+    def tokenize(self, text: str, **kwargs) -> List[str]:
+        """
+        Tokenize text and return string tokens.
+        This wraps the numeric tokenization for HuggingFace compatibility.
+        """
+        # Get numeric tokens
+        numeric_tokens = self._tokenize(text)
+        
+        # Convert to string tokens
+        string_tokens = [self._convert_id_to_token(token_id) for token_id in numeric_tokens]
+        
+        return string_tokens
+    
+    def encode(
+        self,
+        text: str,
+        add_special_tokens: bool = True,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs
+    ) -> Union[List[int], torch.Tensor]:
+        """
+        Encode text to token IDs.
+        Core tokenization functionality of the CENO byte-level tokenizer.
+        """
+        # Tokenize to get numeric IDs directly
+        token_ids = self._tokenize(text)
+        
+        # Handle truncation
+        if truncation and max_length is not None:
+            token_ids = token_ids[:max_length]
+        
+        # Handle padding
+        if padding and max_length is not None:
+            if len(token_ids) < max_length:
+                token_ids.extend([self.pad_id] * (max_length - len(token_ids)))
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            return torch.tensor([token_ids], dtype=torch.long)
+        elif return_tensors == "np":
+            return np.array([token_ids], dtype=np.int64)
+        
+        return token_ids
+    
+    def decode(
+        self,
+        token_ids: Union[List[int], torch.Tensor, np.ndarray],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = True,
+        **kwargs
+    ) -> str:
+        """
+        Decode token IDs back to text.
+        CENO detokenization.
+        """
+        # Convert to list if tensor or numpy array
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+        elif isinstance(token_ids, np.ndarray):
+            token_ids = token_ids.tolist()
+        
+        # Convert IDs to tokens
+        tokens = [self._convert_id_to_token(token_id) for token_id in token_ids]
+        
+        # Filter special tokens if requested
+        if skip_special_tokens:
+            tokens = [
+                token for token in tokens 
+                if token not in [self.pad_token, self.eos_token, self.unk_token]
+            ]
+        
+        # Convert tokens to string
+        return self.convert_tokens_to_string(tokens)
+    
+    def batch_encode_plus(
+        self,
+        batch_text_or_text_pairs: Union[List[str], List[Tuple[str, str]]],
+        add_special_tokens: bool = True,
+        padding: bool = False,
+        truncation: bool = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs
+    ) -> BatchEncoding:
+        """Batch encode multiple texts"""
+        batch_outputs = []
+        
+        for text in batch_text_or_text_pairs:
+            if isinstance(text, tuple):
+                # Handle text pairs (not typically used for DNA sequences)
+                text = text[0]  # Just use first text for now
+            
+            encoded = self.encode(
+                text,
+                add_special_tokens=add_special_tokens,
+                padding=False,  # We'll handle padding after
+                truncation=truncation,
+                max_length=max_length,
+                return_tensors=None,
+            )
+            batch_outputs.append(encoded)
+        
+        # Handle batch padding
+        if padding and max_length is not None:
+            max_len = max_length
+        elif padding:
+            max_len = max(len(output) for output in batch_outputs)
+        else:
+            max_len = None
+        
+        if max_len is not None:
+            for i, output in enumerate(batch_outputs):
+                if len(output) < max_len:
+                    batch_outputs[i] = output + [self.pad_id] * (max_len - len(output))
+                elif len(output) > max_len:
+                    batch_outputs[i] = output[:max_len]
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            batch_outputs = torch.tensor(batch_outputs, dtype=torch.long)
+        elif return_tensors == "np":
+            batch_outputs = np.array(batch_outputs)
+        
+        return BatchEncoding({"input_ids": batch_outputs})
+    
+    def batch_decode(
+        self,
+        sequences: Union[List[List[int]], torch.Tensor, np.ndarray],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = True,
+        **kwargs
+    ) -> List[str]:
+        """Batch decode multiple sequences"""
+        # Convert to list format
+        if isinstance(sequences, torch.Tensor):
+            sequences = sequences.tolist()
+        elif isinstance(sequences, np.ndarray):
+            sequences = sequences.tolist()
+        
+        return [
+            self.decode(
+                sequence,
+                skip_special_tokens=skip_special_tokens,
+                clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+                **kwargs
+            )
+            for sequence in sequences
+        ]
+    
+    def save_pretrained(
+        self, 
+        save_directory: str, 
+        legacy_format: Optional[bool] = None,
+        filename_prefix: Optional[str] = None,
+        push_to_hub: bool = False,
+        **kwargs
+    ) -> Tuple[str]:
+        """
+        Save the tokenizer to a directory.
+        
+        Args:
+            save_directory (str): Directory to save the tokenizer
+            legacy_format (bool, optional): Whether to save in legacy format
+            filename_prefix (str, optional): Prefix for filenames
+            push_to_hub (bool): Whether to push to HuggingFace Hub
+            **kwargs: Additional arguments
+        
+        Returns:
+            Tuple[str]: Tuple of saved file paths
+        """
+        if not os.path.isdir(save_directory):
+            os.makedirs(save_directory, exist_ok=True)
+        
+        # Save vocabulary
+        vocab_file = os.path.join(
+            save_directory, 
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+        
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self._vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
+        
+        # Save tokenizer configuration
+        config_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + "tokenizer_config.json"
+        )
+        
+        tokenizer_config = {
+            "tokenizer_class": "CENOCharLevelTokenizer",
+            "vocab_size": self._vocab_size,
+            "eos_token": self.eos_token,
+            "pad_token": self.pad_token,
+            "unk_token": self.unk_token,
+            "eod_id": self.eod_id,
+            "eos_id": self.eos_id,
+            "pad_id": self.pad_id,
+            "model_max_length": getattr(self, 'model_max_length', 1000000),
+            "clean_up_tokenization_spaces": True,
+            "tokenize_chinese_chars": False,
+            "strip_accents": None,
+            "do_lower_case": False,
+            "do_basic_tokenize": False,
+            "never_split": None,
+            "tokenizer_type": "CharLevelTokenizer",
+            "name_or_path": save_directory,
+        }
+        
+        with open(config_file, "w", encoding="utf-8") as f:
+            json.dump(tokenizer_config, f, indent=2, ensure_ascii=False)
+        
+        # Save special tokens map
+        special_tokens_file = os.path.join(
+            save_directory,
+            (filename_prefix + "-" if filename_prefix else "") + "special_tokens_map.json"
+        )
+        
+        special_tokens_map = {
+            "eos_token": self.eos_token,
+            "pad_token": self.pad_token,
+            "unk_token": self.unk_token,
+        }
+        
+        with open(special_tokens_file, "w", encoding="utf-8") as f:
+            json.dump(special_tokens_map, f, indent=2, ensure_ascii=False)
+        
+        logger.info(f"Tokenizer saved to {save_directory}")
+        
+        return (vocab_file, config_file, special_tokens_file)
+    
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        cache_dir: Optional[str] = None,
+        force_download: bool = False,
+        local_files_only: bool = False,
+        token: Optional[str] = None,
+        revision: str = "main",
+        **kwargs
+    ):
+        """
+        Load a tokenizer from a pretrained model.
+        
+        Args:
+            pretrained_model_name_or_path (str): Path to directory containing tokenizer files
+                or name of a model on HuggingFace Hub
+            cache_dir (str, optional): Directory to cache downloaded files
+            force_download (bool): Whether to force download even if cached
+            local_files_only (bool): Whether to only use local files
+            token (str, optional): HuggingFace access token
+            revision (str): Model revision to use
+            **kwargs: Additional arguments
+        
+        Returns:
+            CENOCharLevelTokenizer: Loaded tokenizer instance
+        """
+        # Handle local directory
+        if os.path.isdir(pretrained_model_name_or_path):
+            model_path = pretrained_model_name_or_path
+        else:
+            # Try to download from HuggingFace Hub
+            try:
+                from huggingface_hub import snapshot_download
+                
+                model_path = snapshot_download(
+                    repo_id=pretrained_model_name_or_path,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    local_files_only=local_files_only,
+                    token=token,
+                    revision=revision,
+                )
+            except ImportError:
+                raise ImportError(
+                    "huggingface_hub is required to download models from the Hub. "
+                    "Install it with: pip install huggingface_hub"
+                )
+            except Exception as e:
+                logger.warning(f"Failed to download from HuggingFace Hub: {e}")
+                logger.warning("Falling back to local initialization...")
+                return cls(**kwargs)
+        
+        # Load tokenizer configuration
+        config_file = os.path.join(model_path, "tokenizer_config.json")
+        config = {}
+        
+        if os.path.exists(config_file):
+            with open(config_file, "r", encoding="utf-8") as f:
+                config = json.load(f)
+            logger.info(f"Loaded tokenizer config from {config_file}")
+        
+        # Load special tokens map
+        special_tokens_file = os.path.join(model_path, "special_tokens_map.json")
+        special_tokens = {}
+        
+        if os.path.exists(special_tokens_file):
+            with open(special_tokens_file, "r", encoding="utf-8") as f:
+                special_tokens = json.load(f)
+            logger.info(f"Loaded special tokens from {special_tokens_file}")
+        
+        # Load vocabulary
+        vocab_file = os.path.join(model_path, VOCAB_FILES_NAMES["vocab_file"])
+        vocab = None
+        
+        if os.path.exists(vocab_file):
+            with open(vocab_file, "r", encoding="utf-8") as f:
+                vocab = json.load(f)
+            logger.info(f"Loaded vocabulary from {vocab_file}")
+        
+        # Merge configurations (kwargs override file config)
+        init_kwargs = {
+            "vocab_size": config.get("vocab_size", 512),
+            "eos_token": special_tokens.get("eos_token", config.get("eos_token", "<eos>")),
+            "pad_token": special_tokens.get("pad_token", config.get("pad_token", "<pad>")),
+            "unk_token": special_tokens.get("unk_token", config.get("unk_token", "<unk>")),
+        }
+        
+        # Override with any provided kwargs
+        init_kwargs.update(kwargs)
+        
+        # Create tokenizer instance
+        tokenizer = cls(**init_kwargs)
+        
+        # Load custom vocabulary if available
+        if vocab is not None:
+            tokenizer._vocab = vocab
+            tokenizer._id_to_token = {v: k for k, v in vocab.items()}
+            logger.info("Loaded custom vocabulary")
+        
+        # Set additional attributes from config
+        if config:
+            tokenizer.eod_id = config.get("eod_id", 0)
+            tokenizer.eos_id = config.get("eos_id", 0)
+            tokenizer.pad_id = config.get("pad_id", 1)
+            if hasattr(tokenizer, 'model_max_length'):
+                tokenizer.model_max_length = config.get("model_max_length", 1000000)
+        
+        tokenizer.name_or_path = pretrained_model_name_or_path
+        logger.info(f"Successfully loaded tokenizer from {pretrained_model_name_or_path}")
+        
+        return tokenizer
+    
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """Save vocabulary to file (legacy method)"""
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        
+        vocab_file = os.path.join(
+            save_directory, 
+            (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+        
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self._vocab, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
+        
+        return (vocab_file,)
+    
+    @property
+    def unique_identifiers(self) -> Dict[str,Any]:
+        """
+        Megatron will call .unique_identifiers when it encounters
+        this object during its JSON‐dump of the dataset config.
+        Must be JSON-serializable.
+        """
+        return {
+            "tokenizer_class": self.__class__.__name__,
+            "name_or_path": getattr(self, "name_or_path", None),
+            "vocab_size": self.vocab_size,
+        }
+    # Compatibility methods for the CENO tokenizer interface
+    def tokenize_batch(self, text_batch: Union[List[str], str]) -> Union[List[List[int]], List[int]]:
+        """Batch tokenization matching the CENO tokenizer interface"""
+        if isinstance(text_batch, str):
+            return self._tokenize(text_batch)
+        return [self._tokenize(text) for text in text_batch]
+    
+    def detokenize(self, token_ids: Union[List[int], torch.Tensor]) -> str:
+        """Alias for decode method matching the CENO tokenizer interface"""
+        return self.decode(token_ids, skip_special_tokens=True)
+    
+    def detokenize_batch(self, token_ids_batch: Union[List[List[int]], torch.Tensor]) -> List[str]:
+        """Batch detokenization matching the CENO tokenizer interface"""
+        return self.batch_decode(token_ids_batch, skip_special_tokens=True)
+    
+    @property
+    def eod(self) -> int:
+        """End of document token ID"""
+        return self.eod_id
+    
+    @property
+    def eos(self) -> int:
+        """End of sequence token ID"""
+        return self.eos_id 

config.json

@@ -0,0 +1,84 @@
+{
+  "architectures": [
+    "CENOForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "attention_head_dim": null,
+  "auto_map": {
+    "AutoConfig": "configuration_ceno.CENOConfig",
+    "AutoModelForCausalLM": "modeling_ceno.CENOForCausalLM",
+    "AutoTokenizer": [
+      "ceno_tokenizer.CENOCharLevelTokenizer",
+      null
+    ]
+  },
+  "bos_token_id": 0,
+  "chunk_size": 128,
+  "conv_kernel": 4,
+  "eos_token_id": 0,
+  "expand": 2,
+  "head_dim": 128,
+  "hidden_dropout": 0.0,
+  "hidden_size": 2560,
+  "hybrid_override_pattern": "MEMEMEME*EMEMEMEM*MEMEMEME*EMEMEME*MEMEMEM",
+  "initializer_range": 0.02,
+  "intermediate_size": 10240,
+  "layer_norm_epsilon": 1e-05,
+  "mamba_head_dim": 64,
+  "mamba_hidden_act": "silu",
+  "mamba_in_proj_layernorm": true,
+  "mamba_num_heads": 80,
+  "mamba_proj_bias": false,
+  "mamba_ssm_cache_dtype": "bfloat16",
+  "max_position_embeddings": 1048576,
+  "mlp_bias": false,
+  "mlp_fc1_layernorm": false,
+  "mlp_hidden_act": "relu2",
+  "mlp_use_swiglu": false,
+  "model_type": "ceno",
+  "moe_intermediate_size": 10240,
+  "moe_router_dtype": "bfloat16",
+  "moe_router_pre_softmax": false,
+  "moe_shared_expert_intermediate_size": 10240,
+  "moe_top_k": 2,
+  "n_group": 1,
+  "n_groups": 8,
+  "n_routed_experts": 8,
+  "n_shared_experts": 0,
+  "norm_topk_prob": true,
+  "num_attention_heads": 20,
+  "num_experts": 8,
+  "num_experts_per_tok": 2,
+  "num_hidden_layers": 42,
+  "num_key_value_heads": 20,
+  "num_logits_to_keep": 1,
+  "pad_token_id": 1,
+  "qkv_layernorm": true,
+  "rescale_prenorm_residual": true,
+  "residual_in_fp32": false,
+  "routed_scaling_factor": 1.0,
+  "sliding_window": null,
+  "ssm_state_size": 128,
+  "tie_word_embeddings": false,
+  "time_step_floor": 0.0001,
+  "time_step_limit": [
+    0.0,
+    1e+30
+  ],
+  "time_step_max": 0.1,
+  "time_step_min": 0.001,
+  "topk_group": 1,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.48.3",
+  "use_bias": false,
+  "use_cache": true,
+  "use_conv_bias": true,
+  "use_mamba_kernels": true,
+  "vocab_size": 512,
+  "ceno_training_precision": false,
+  "ceno_mamba_gate_in_ssm": true,
+  "ceno_rmsnorm_mode": "torch",
+  "ceno_reference_cp_size": 1,
+  "ceno_reference_context_length": null
+}

configuration_ceno.py

@@ -0,0 +1,344 @@
+# coding=utf-8
+# Copyright 2024 AI21 Labs Ltd. and the HuggingFace Inc. team. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2026, CENO Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""CENO model configuration"""
+
+import re
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CENOConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`CENOModel`]. It is used to instantiate a
+    CENO model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the CENO-v0.1 model.
+
+    [todo](todo)
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 131072):
+            Vocabulary size of the CENO model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`CENOModel`]
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
+            model has a output word embedding layer.
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 21504):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 52):
+            Number of hidden layers in the Transformer encoder.
+        hybrid_override_pattern (`str`, *optional*, defaults to `"MMMM*MMMMM*MMMMM*MMMMM*MMMMMM"`):
+            The pattern of the hybrid model. The pattern is a string of characters where each character represents one layer: M=Mamba, *=Attention, -=MLP, E=MoE
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        attention_head_dim (`int`, *optional*, defaults to 128):
+            Dimension of each attention head.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used.
+        mlp_hidden_act (`str`, *optional*, defaults to "relu2"):
+            The non-linear activation function in the MLP layers.
+        mlp_use_swiglu (`bool`, *optional*, defaults to `False`):
+            Whether to use SwiGLU activation (gated linear unit) in MLP layers.
+        num_experts (`int`, *optional*, defaults to 8):
+            Number of experts for MoE layers.
+        top_k (`int`, *optional*, defaults to 2):
+            Number of experts to select for each token in MoE layers.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in attention layers.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in MLP layers.
+        use_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in the model.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
+            The epsilon used by the layer normalization layers.
+        residual_in_fp32 (`bool`, *optional*, defaults to `False`):
+            Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):
+            Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an
+            integer value, only last `num_logits_to_keep` logits will be calculated.
+        pad_token_id (`int`, *optional*, defaults to 1):
+            The id of the padding token.
+        bos_token_id (`int`, *optional*, defaults to 0):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 0):
+            The id of the "end-of-sequence" token.
+        sliding_window (`int`, *optional*, defaults to None):
+            Sliding window attention window size.
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
+            The maximum sequence length that this model might ever be used with.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        hidden_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the hidden states.
+        use_mamba_kernels (`bool`, *optional*, defaults to `True`):
+            Flag indicating whether or not to use the fast mamba kernels. These are available only if `mamba-ssm` and
+            `causal-conv1d` are installed, and the mamba modules are running on a CUDA device.
+        ssm_state_size (`int`, *optional*, defaults to 128):
+            The dimension of the mamba state space latents.
+        mamba_num_heads (`int`, *optional*, defaults to 128):
+            Number of heads in Mamba layers.
+        mamba_n_groups (`int`, *optional*, defaults to 8):
+            Number of groups in Mamba layers.
+        mamba_head_dim (`int`, *optional*, defaults to 64):
+            Dimension of each Mamba head.
+        mamba_d_conv (`int`, *optional*, defaults to 4):
+            The size of the mamba convolution kernel.
+        mamba_expand (`int`, *optional*, defaults to 2):
+            Expanding factor used to determine the mamba intermediate size.
+        mamba_hidden_act (`str`, *optional*, defaults to "silu"):
+            The non-linear activation function in the Mamba layers.
+        mamba_dt_min (`float`, *optional*, defaults to 0.001):
+            Minimum value for the time step in Mamba.
+        mamba_dt_max (`float`, *optional*, defaults to 0.1):
+            Maximum value for the time step in Mamba.
+        mamba_dt_limit (`tuple`, *optional*, defaults to (0.0, float("inf"))):
+            Limits for the time step in Mamba.
+        mamba_dt_init_floor (`float`, *optional*, defaults to 1e-4):
+            Floor value for time step initialization in Mamba.
+        mamba_conv_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the convolution layer of the mamba mixer block.
+        mamba_proj_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in the input and output projections of the mamba mixer block.
+        mamba_in_proj_layernorm (`bool`, *optional*, defaults to `False`):
+            Whether to apply an additional RMSNorm before the Mamba in-projection. This matches
+            Megatron's TELayerNormColumnParallelLinear fused Mamba in_proj norm when enabled.
+        mamba_chunk_size (`int`, *optional*, defaults to 256):
+            Size of chunks for Mamba processing.
+        rescale_prenorm_residual (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the pre-normalization residual connections.
+        num_experts (`int`, *optional*, defaults to 8):
+            Number of experts in MoE layers.
+        moe_top_k (`int`, *optional*, defaults to 2):
+            Number of experts to route to in MoE layers (renamed from top_k to avoid generation config conflicts).
+        qkv_layernorm (`bool`, *optional*, defaults to `False`):
+            Whether to apply an additional RMSNorm before QKV projections in attention layers. This matches
+            Megatron's TELayerNormColumnParallelLinear fused QKV norm when enabled.
+        intra_encoding_pattern (`str`, *optional*, defaults to `None`):
+            Per-layer switch (length must equal `num_hidden_layers`) to control which layers use intra-sequence
+            isolation. Characters in `{1, Y, y, S, s, T, t, +}` enable isolation; `{0, N, n, F, f, ., -}` disable.
+            If omitted, intra-sequence isolation is disabled (legacy behavior).
+    """
+
+    model_type = "ceno"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        tie_word_embeddings=False,
+        hidden_size=4096,
+        intermediate_size=21504,
+        num_hidden_layers=52,
+        hybrid_override_pattern="MMMM*MMMMM*MMMMM*MMMMM*MMMMMM",
+        num_attention_heads=32,
+        attention_head_dim=128,
+        num_key_value_heads=8,  # nemo: num_query_groups
+        mlp_hidden_act="relu2",
+        mlp_use_swiglu=False,
+        num_experts=8,
+        moe_top_k=2,
+        moe_router_pre_softmax=False,
+        attention_bias=False,
+        qkv_layernorm=False,
+        mlp_bias=False,
+        use_bias=False,
+        initializer_range=0.02, # nemo: init_method_std
+        layer_norm_epsilon=1e-5, # nemo: layernorm_epsilon
+        residual_in_fp32=False,  #  Megatron Core default value
+        use_cache=True,
+        num_logits_to_keep=1,
+        pad_token_id=1,
+        bos_token_id=0,
+        eos_token_id=0,
+        sliding_window=None,
+        max_position_embeddings=4096,
+        attention_dropout=0.0,
+        hidden_dropout=0.0, # * ADDED
+        use_mamba_kernels=True,
+        ssm_state_size=128, # mamba_state_size
+        mamba_num_heads=128,
+        mamba_n_groups=8,  # nemo: mamba_ssm_ngroups = num_heads
+        mamba_head_dim=64,
+        mamba_d_conv=4,
+        mamba_expand=2,
+        mamba_hidden_act="silu",
+        mamba_dt_min=0.001,
+        mamba_dt_max=0.1,
+        mamba_dt_limit=(0.0, float("inf")),
+        mamba_dt_init_floor=1e-4,
+        mamba_conv_bias=True,
+        mamba_proj_bias=False,
+        mamba_in_proj_layernorm=False,
+        mamba_chunk_size=256,
+        rescale_prenorm_residual=True,
+        ceno_training_precision=False,
+        ceno_mamba_gate_in_ssm=False,
+        ceno_rmsnorm_mode="torch",
+        ceno_reference_cp_size=1,
+        ceno_reference_context_length=None,
+        intra_encoding_pattern=None,
+        **kwargs,
+    ):
+        # Explicit inference reference; never inferred from a serving batch or prefix.
+        self.ceno_training_precision = ceno_training_precision
+        self.ceno_mamba_gate_in_ssm = ceno_mamba_gate_in_ssm
+        self.ceno_rmsnorm_mode = ceno_rmsnorm_mode
+        self.ceno_reference_cp_size = ceno_reference_cp_size
+        self.ceno_reference_context_length = ceno_reference_context_length
+        self.vocab_size = vocab_size
+        self.tie_word_embeddings = tie_word_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.hybrid_override_pattern = hybrid_override_pattern
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+        self.sliding_window = sliding_window
+        self.max_position_embeddings = max_position_embeddings
+        self.attention_dropout = attention_dropout
+        self.hidden_dropout = hidden_dropout
+        self.qkv_layernorm = qkv_layernorm
+
+        # Validate and process hybrid_override_pattern
+        # M: Mamba, *: Attention, -: MLP, E: MoE
+        if self.hybrid_override_pattern is None:
+            # Default pattern with even distribution of attention layers
+            self.hybrid_override_pattern = self._generate_default_pattern()
+        
+        # Pattern should be a string where each character represents one layer
+        # No processing needed since '-' represents MLP layers, not separators
+        
+        # For the default pattern, adjust it if num_hidden_layers doesn't match
+        if len(self.hybrid_override_pattern) != self.num_hidden_layers:
+            # If using default pattern and layers don't match, regenerate
+            if self.hybrid_override_pattern == "MMMM*MMMMM*MMMMM*MMMMM*MMMMMM":
+                self.hybrid_override_pattern = self._generate_default_pattern()
+            else:
+                # Only assert if it's a user-provided pattern
+                assert len(self.hybrid_override_pattern) == self.num_hidden_layers, f"hybrid_override_pattern length ({len(self.hybrid_override_pattern)}) must match num_hidden_layers ({self.num_hidden_layers})"
+        
+        assert re.match(r"^[M*\-E]+$", self.hybrid_override_pattern), "hybrid_override_pattern must only contain characters 'M', '*', '-', or 'E'"
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.mlp_hidden_act = mlp_hidden_act
+        self.mlp_use_swiglu = mlp_use_swiglu
+        self.num_experts = num_experts
+        self.moe_top_k = moe_top_k
+        # Megatron MoE routing: default is post-topk softmax (i.e., softmax after selecting top-k experts).
+        # This flag mirrors Megatron's `--moe-router-pre-softmax`.
+        self.moe_router_pre_softmax = moe_router_pre_softmax
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.use_bias = use_bias
+        self.initializer_range = initializer_range
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.residual_in_fp32 = residual_in_fp32
+
+        self.use_cache = use_cache
+        self.num_logits_to_keep = num_logits_to_keep
+
+        self.use_mamba_kernels = use_mamba_kernels
+        self.n_groups = mamba_n_groups
+        self.mamba_head_dim = mamba_head_dim
+        self.ssm_state_size = ssm_state_size
+        self.mamba_num_heads = mamba_num_heads
+        self.conv_kernel = mamba_d_conv
+        self.expand = mamba_expand
+        self.mamba_hidden_act = mamba_hidden_act
+        self.time_step_min = mamba_dt_min
+        self.time_step_max = mamba_dt_max
+        self.time_step_limit = mamba_dt_limit
+        self.time_step_floor = mamba_dt_init_floor
+        self.use_conv_bias = mamba_conv_bias
+        self.mamba_proj_bias = mamba_proj_bias
+        self.mamba_in_proj_layernorm = mamba_in_proj_layernorm
+        self.chunk_size = mamba_chunk_size
+        self.rescale_prenorm_residual = rescale_prenorm_residual
+        self.intra_encoding_pattern = intra_encoding_pattern
+        
+        # MoE parameters
+        self.num_experts = num_experts
+        self.moe_top_k = moe_top_k
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def _generate_default_pattern(self):
+        """Generate default hybrid pattern with even distribution of attention layers"""
+        # Default to mostly Mamba with some attention layers evenly distributed
+        pattern = ["M"] * self.num_hidden_layers
+        
+        # Place attention layers at roughly even intervals
+        if self.num_hidden_layers >= 4:
+            # For larger models, place attention every 4-5 layers
+            attention_interval = max(4, self.num_hidden_layers // 8)
+            for i in range(attention_interval - 1, self.num_hidden_layers, attention_interval):
+                pattern[i] = "*"
+        
+        return ''.join(pattern)
+
+    @property
+    def layers_block_type(self):
+        return [
+            "mamba" if self.hybrid_override_pattern[i] == "M" else
+            "attention" if self.hybrid_override_pattern[i] == "*" else
+            "mlp" if self.hybrid_override_pattern[i] == "-" else "moe"
+            for i in range(self.num_hidden_layers)]
+
+    @property
+    def intra_encoding_mask(self):
+        """
+        Per-layer boolean mask for intra-sequence isolation.
+        If `intra_encoding_pattern` is provided, it is used; otherwise isolation is disabled.
+        """
+        if getattr(self, "intra_encoding_pattern", None) is not None:
+            pat = self.intra_encoding_pattern
+            if len(pat) != self.num_hidden_layers:
+                raise ValueError(
+                    f"intra_encoding_pattern length ({len(pat)}) must match num_hidden_layers ({self.num_hidden_layers})"
+                )
+            enable_set = {"1", "Y", "y", "S", "s", "T", "t", "+"}
+            disable_set = {"0", "N", "n", "F", "f", ".", "-"}
+            return [c in enable_set if c in enable_set.union(disable_set) else True for c in pat]
+
+        # Default/legacy: no intra-sequence isolation
+        return [False for _ in range(self.num_hidden_layers)]

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 1,
+  "transformers_version": "4.48.3"
+}

model.safetensors.index.json

@@ -0,0 +1,570 @@
+{
+  "metadata": {
+    "total_size": 17857248160
+  },
+  "weight_map": {
+    "backbone.embeddings.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.A_log": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.D": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.conv1d.bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.conv1d.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.dt_bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.in_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.in_proj_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.0.mixer.out_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.0.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.0.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.1.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.1.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.2.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.2.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.3.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.3.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.4.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.4.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.5.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.5.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.6.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.6.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.7.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.experts.7.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.mixer.gate.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.1.pre_mlp_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.A_log": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.D": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.conv1d.bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.conv1d.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.dt_bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.in_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.in_proj_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.10.mixer.out_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.0.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.0.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.1.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.1.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.2.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.2.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.3.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.3.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.4.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.4.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.5.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.5.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.6.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.6.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.7.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.experts.7.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.11.mixer.gate.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.11.pre_mlp_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.A_log": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.D": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.conv1d.bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.conv1d.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.dt_bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.in_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.in_proj_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.12.mixer.out_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.0.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.0.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.1.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.1.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.2.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.2.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.3.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.3.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.4.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.4.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.5.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.5.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.6.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.6.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.7.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.experts.7.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.mixer.gate.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.13.pre_mlp_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.A_log": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.D": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.conv1d.bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.conv1d.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.dt_bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.in_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.in_proj_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.14.mixer.out_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.0.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.0.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.1.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.1.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.2.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.2.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.3.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.3.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.4.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.4.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.5.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.5.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.6.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.6.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.7.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.experts.7.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.mixer.gate.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.15.pre_mlp_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.A_log": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.D": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.conv1d.bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.conv1d.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.dt_bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.in_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.in_proj_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.16.mixer.out_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.mixer.k_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.mixer.o_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.mixer.q_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.mixer.qkv_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.mixer.v_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.17.pre_mlp_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.A_log": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.D": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.conv1d.bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.conv1d.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.dt_bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.in_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.in_proj_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.18.mixer.out_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.0.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.0.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.1.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.1.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.2.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.2.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.3.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.3.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.4.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.4.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.5.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.5.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.6.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.6.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.7.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.experts.7.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.mixer.gate.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.19.pre_mlp_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.2.mixer.A_log": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.D": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.conv1d.bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.conv1d.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.dt_bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.in_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.in_proj_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.2.mixer.out_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.20.mixer.A_log": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.D": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.conv1d.bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.conv1d.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.dt_bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.in_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.in_proj_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.20.mixer.out_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.0.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.0.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.1.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.1.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.2.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.2.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.3.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.3.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.4.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.4.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.5.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.5.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.6.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.6.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.7.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.experts.7.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.mixer.gate.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.21.pre_mlp_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.A_log": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.D": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.conv1d.bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.conv1d.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.dt_bias": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.in_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.in_proj_norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.22.mixer.out_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.0.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.0.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.1.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.1.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.2.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.2.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.3.down_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.3.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.4.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.4.up_proj.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.5.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.5.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.6.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.6.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.7.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.experts.7.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.23.mixer.gate.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.norm.weight": "model-00002-of-00004.safetensors",
+    "backbone.layers.23.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.A_log": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.D": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.conv1d.bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.conv1d.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.dt_bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.in_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.in_proj_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.24.mixer.out_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.0.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.0.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.1.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.1.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.2.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.2.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.3.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.3.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.4.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.4.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.5.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.5.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.6.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.6.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.7.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.experts.7.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.mixer.gate.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.25.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.mixer.k_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.mixer.o_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.mixer.q_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.mixer.qkv_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.mixer.v_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.26.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.0.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.0.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.1.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.1.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.2.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.2.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.3.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.3.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.4.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.4.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.5.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.5.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.6.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.6.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.7.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.experts.7.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.mixer.gate.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.27.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.A_log": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.D": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.conv1d.bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.conv1d.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.dt_bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.in_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.in_proj_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.28.mixer.out_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.0.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.0.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.1.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.1.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.2.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.2.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.3.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.3.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.4.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.4.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.5.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.5.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.6.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.6.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.7.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.experts.7.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.mixer.gate.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.29.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.0.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.0.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.1.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.1.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.2.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.2.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.3.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.3.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.4.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.4.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.5.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.5.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.6.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.6.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.7.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.experts.7.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.mixer.gate.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.3.pre_mlp_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.30.mixer.A_log": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.D": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.conv1d.bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.conv1d.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.dt_bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.in_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.in_proj_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.30.mixer.out_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.0.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.0.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.1.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.1.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.2.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.2.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.3.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.3.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.4.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.4.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.5.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.5.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.6.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.6.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.7.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.experts.7.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.mixer.gate.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.31.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.A_log": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.D": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.conv1d.bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.conv1d.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.dt_bias": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.in_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.in_proj_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.32.mixer.out_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.0.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.0.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.1.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.1.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.2.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.2.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.3.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.3.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.4.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.4.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.5.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.5.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.6.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.6.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.7.down_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.experts.7.up_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.mixer.gate.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.33.pre_mlp_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.34.mixer.k_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.34.mixer.o_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.34.mixer.q_proj.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.34.mixer.qkv_norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.34.mixer.v_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.34.norm.weight": "model-00003-of-00004.safetensors",
+    "backbone.layers.34.pre_mlp_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.A_log": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.D": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.conv1d.bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.conv1d.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.dt_bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.in_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.in_proj_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.35.mixer.out_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.0.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.0.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.1.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.1.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.2.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.2.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.3.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.3.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.4.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.4.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.5.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.5.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.6.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.6.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.7.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.experts.7.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.mixer.gate.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.36.pre_mlp_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.A_log": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.D": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.conv1d.bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.conv1d.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.dt_bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.in_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.in_proj_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.37.mixer.out_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.0.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.0.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.1.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.1.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.2.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.2.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.3.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.3.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.4.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.4.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.5.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.5.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.6.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.6.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.7.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.experts.7.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.mixer.gate.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.38.pre_mlp_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.A_log": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.D": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.conv1d.bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.conv1d.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.dt_bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.in_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.in_proj_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.39.mixer.out_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.4.mixer.A_log": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.D": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.conv1d.bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.conv1d.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.dt_bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.in_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.in_proj_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.4.mixer.out_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.0.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.0.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.1.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.1.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.2.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.2.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.3.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.3.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.4.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.4.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.5.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.5.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.6.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.6.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.7.down_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.experts.7.up_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.mixer.gate.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.40.pre_mlp_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.A_log": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.D": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.conv1d.bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.conv1d.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.dt_bias": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.in_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.in_proj_norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.norm.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.41.mixer.out_proj.weight": "model-00004-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.0.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.0.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.1.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.1.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.2.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.2.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.3.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.3.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.4.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.4.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.5.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.5.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.6.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.6.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.7.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.experts.7.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.mixer.gate.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.5.pre_mlp_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.A_log": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.D": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.conv1d.bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.conv1d.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.dt_bias": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.in_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.in_proj_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.6.mixer.out_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.0.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.0.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.1.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.1.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.2.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.2.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.3.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.3.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.4.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.4.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.5.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.5.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.6.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.6.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.7.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.experts.7.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.mixer.gate.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.7.pre_mlp_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.mixer.k_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.mixer.o_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.mixer.q_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.mixer.qkv_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.mixer.v_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.8.pre_mlp_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.0.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.0.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.1.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.1.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.2.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.2.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.3.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.3.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.4.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.4.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.5.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.5.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.6.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.6.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.7.down_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.experts.7.up_proj.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.mixer.gate.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.layers.9.pre_mlp_norm.weight": "model-00001-of-00004.safetensors",
+    "backbone.norm_f.weight": "model-00004-of-00004.safetensors",
+    "lm_head.weight": "model-00004-of-00004.safetensors"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "eos_token": "<eos>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

third_party/TransformerEngine-LICENSE

@@ -0,0 +1,175 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+

tokenizer_config.json

@@ -0,0 +1,26 @@
+{
+  "tokenizer_class": "CENOCharLevelTokenizer",
+  "vocab_size": 512,
+  "eos_token": "<eos>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>",
+  "eod_id": 0,
+  "eos_id": 0,
+  "pad_id": 1,
+  "unk_id": 2,
+  "model_max_length": 1000000000000000019884624838656,
+  "clean_up_tokenization_spaces": true,
+  "tokenize_chinese_chars": false,
+  "strip_accents": null,
+  "do_lower_case": false,
+  "do_basic_tokenize": false,
+  "never_split": null,
+  "tokenizer_type": "CharLevelTokenizer",
+  "name_or_path": "./ceno_tokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "ceno_tokenizer.CENOCharLevelTokenizer",
+      null
+    ]
+  }
+}

training_precision.py

@@ -0,0 +1,654 @@
+"""Opt-in HF inference against the same BF16 Megatron reference as vLLM.
+
+Use from the repository with the existing training image (Mamba 2.2.4 / FA2).
+Defaults in exported checkpoints are unchanged. This adapter is inference-only,
+single-GPU, BF16, unpadded uniform-length batches. Its prefill token budget is
+bounded to avoid the old Mamba kernels' large-stride indexing paths.
+"""
+import contextlib
+import contextvars
+import functools
+import os
+import types
+
+_CHUNK = 128
+
+
+def _local_module(name):
+    import importlib
+    return importlib.import_module(f"{__package__}.{name}" if __package__ else name)
+
+
+def configure_backend():
+    """Set the cuBLASLt workspace limit; call before the first CUDA GEMM.
+
+    The reference GEMMs run through cuBLASLt (see ``backend_scope``) with a
+    1 KiB workspace, which is what fixes their algorithm and hence their bits.
+    That is the only process-level setting. ``CUBLAS_WORKSPACE_CONFIG`` is not
+    set: it only sizes the legacy cuBLAS handle workspace, which these GEMMs do
+    not use (outputs are bit-identical with or without it), and under torch 2.7
+    it adds ~30-40 us of host time to every GEMM call, the main cost of B=1
+    decode. The torch backend flags are applied only inside precision-enabled
+    forwards, so the surrounding program (e.g. an FP32 head trained with TF32)
+    keeps its own settings.
+    """
+    os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
+
+
+@contextlib.contextmanager
+def backend_scope():
+    """Apply the reference GEMM flags inside the block and restore them on exit."""
+    import torch
+    matmul, cudnn = torch.backends.cuda.matmul, torch.backends.cudnn
+    has_cuda = torch.version.cuda is not None
+    saved = (matmul.allow_bf16_reduced_precision_reduction,
+             matmul.allow_fp16_reduced_precision_reduction,
+             matmul.allow_tf32, cudnn.allow_tf32,
+             torch.backends.cuda.preferred_blas_library() if has_cuda else None)
+    matmul.allow_bf16_reduced_precision_reduction = False
+    matmul.allow_fp16_reduced_precision_reduction = False
+    matmul.allow_tf32 = False
+    cudnn.allow_tf32 = False
+    if has_cuda:
+        torch.backends.cuda.preferred_blas_library("cublaslt")
+    try:
+        yield
+    finally:
+        (matmul.allow_bf16_reduced_precision_reduction,
+         matmul.allow_fp16_reduced_precision_reduction,
+         matmul.allow_tf32, cudnn.allow_tf32) = saved[:4]
+        if has_cuda:
+            torch.backends.cuda.preferred_blas_library(saved[4])
+
+
+def _pin_scan():
+    from mamba_ssm.ops.triton import ssd_bmm, ssd_chunk_scan, ssd_chunk_state, ssd_state_passing
+    choices = [
+        (ssd_chunk_state, "_chunk_cumsum_fwd_kernel", {"BLOCK_SIZE_H": 4}, 4, 2),
+        (ssd_chunk_state, "_chunk_state_fwd_kernel",
+         {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, 4, 4),
+        (ssd_chunk_scan, "_chunk_scan_fwd_kernel",
+         {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, 4, 4),
+        (ssd_bmm, "_bmm_chunk_fwd_kernel",
+         {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, 2, 4),
+        (ssd_state_passing, "_state_passing_fwd_kernel", {"BLOCK_SIZE": 256}, 4, 2),
+    ]
+    for module, name, wanted, warps, stages in choices:
+        kernel = getattr(module, name)
+        matches = [c for c in kernel.configs if c.kwargs == wanted
+                   and c.num_warps == warps and c.num_stages == stages]
+        if len(matches) != 1:
+            raise RuntimeError(f"Unsupported HF Mamba kernel profile: {name}")
+        kernel.configs[:] = matches
+        kernel.cache.clear()
+
+
+def _cache_buffers(cache):
+    if not hasattr(cache, "_ceno_chunks"):
+        cache._ceno_chunks = {}
+        original = cache.reorder_cache
+
+        def reorder(beam_idx):
+            result = original(beam_idx)
+            for state in cache._ceno_chunks.values():
+                for key in ("initial", "x", "dt", "B", "C", "z"):
+                    value = state[key]
+                    state[key] = value.index_select(0, beam_idx.to(value.device))
+                state["states_in"] = None  # derived from "initial"; rebuild lazily
+            return result
+
+        cache.reorder_cache = reorder
+    return cache._ceno_chunks
+
+
+def _install_scan_hooks(module):
+    """Namespace hooks are inactive unless a precision-enabled mixer sets context."""
+    import importlib
+    import torch
+    if hasattr(module, "_ceno_scan_context"):
+        return module._ceno_scan_context
+    context = contextvars.ContextVar("ceno_hf_mixer", default=None)
+    scan, update = module.mamba_chunk_scan_combined, module.selective_state_update
+
+    def scan_with_cache(x, dt, A, B, C, *args, **kwargs):
+        result = scan(x, dt, A, B, C, *args, **kwargs)
+        active = context.get()
+        if active is None or active[1] is None:
+            return result
+        mixer, cache = active
+        z = kwargs.get("z")
+        if z is None or result[1].dtype != torch.float32:
+            raise ValueError("HF precision needs the training gate and FP32 final state")
+        batch, length, heads, dim = x.shape
+        groups, state_dim = B.shape[2:]
+        count = length % _CHUNK
+        complete = length - count
+        if count == 0:
+            initial = result[1].clone()
+        elif complete:
+            prefix_kwargs = dict(kwargs, z=z[:, :complete].contiguous())
+            initial = scan(x[:, :complete].contiguous(), dt[:, :complete].contiguous(),
+                           A, B[:, :complete].contiguous(), C[:, :complete].contiguous(),
+                           *args, **prefix_kwargs)[1]
+        else:
+            initial = torch.zeros_like(result[1])
+        state = {"initial": initial, "pos": count, "states_in": None}
+        values = {"x": x, "dt": dt, "B": B, "C": C, "z": z}
+        for key, value in values.items():
+            buffer = torch.zeros((batch, _CHUNK, *value.shape[2:]),
+                                 device=x.device, dtype=value.dtype)
+            if count:
+                buffer[:, :count].copy_(value[:, -count:])
+            state[key] = buffer
+        _cache_buffers(cache)[mixer.layer_idx] = state
+        return result
+
+    # The combined scan's own building blocks. Off the chunk boundary only the
+    # output row is needed, so chunk_state/state_passing (which only produce the
+    # final state) are skipped; the remaining calls receive exactly the inputs the
+    # combined forward would pass them, so the row is bit-identical.
+    parts = importlib.import_module(scan.__module__)
+    fast = all(hasattr(parts, name) for name in
+               ("_chunk_cumsum_fwd", "_bmm_chunk_fwd", "_chunk_scan_fwd", "_state_passing_fwd"))
+
+    def chunk_row(mixer, state, pos):
+        if state.get("states_in") is None:
+            # Entering-chunk state exactly as _state_passing_fwd emits it for chunk 0
+            # (FP32 initial stored in C's dtype).
+            initial = state["initial"]
+            batch, heads, dim, dstate = initial.shape
+            flat = initial.reshape(batch, heads, dim * dstate)
+            out, _ = parts._state_passing_fwd(
+                torch.zeros((batch, 1, heads, dim * dstate), device=flat.device, dtype=torch.float32),
+                torch.zeros((batch, heads, 1), device=flat.device, dtype=torch.float32),
+                initial_states=flat, chunk_size=_CHUNK, out_dtype=state["C"].dtype)
+            state["states_in"] = out.view(batch, 1, heads, dim, dstate)
+        dA_cumsum, dt = parts._chunk_cumsum_fwd(
+            state["dt"], state["A"], _CHUNK, dt_bias=mixer.dt_bias, dt_softplus=True,
+            dt_limit=mixer.time_step_limit)
+        CB = parts._bmm_chunk_fwd(state["C"], state["B"], _CHUNK, seq_idx=None,
+                                  output_dtype=torch.float32)
+        D = mixer.D if mixer.D.stride(-1) == 1 else mixer.D.contiguous()
+        out, _ = parts._chunk_scan_fwd(CB, state["x"], dt, dA_cumsum, state["C"],
+                                       state["states_in"], D=D, z=state["z"], seq_idx=None)
+        return out[:, pos]
+
+    def update_from_chunk(state_tensor, x, dt, A, B, C, D, *args, **kwargs):
+        active = context.get()
+        if active is None:
+            return update(state_tensor, x, dt, A, B, C, D, *args, **kwargs)
+        mixer, cache = active
+        state = _cache_buffers(cache)[mixer.layer_idx]
+        state["decoded"] = True  # decode kernels have run eagerly once (graph warm-up)
+        pos = state["pos"]
+        # The recurrent update's output was never used and state_tensor is only
+        # read back at the chunk boundary below, so it is not run per token.
+        for key, value in (("x", x), ("dt", dt[:, :, 0]), ("B", B), ("C", C),
+                           ("z", kwargs["z"])):
+            if pos == 0:
+                state[key].zero_()
+            state[key][:, pos].copy_(value)
+        if "A" not in state:
+            state["A"] = -torch.exp(mixer.A_log.float())
+        if pos == _CHUNK - 1 or not fast:
+            y, final = scan(state["x"], state["dt"], state["A"],
+                            state["B"], state["C"], chunk_size=_CHUNK, D=mixer.D,
+                            z=state["z"], dt_bias=mixer.dt_bias, dt_softplus=True,
+                            dt_limit=mixer.time_step_limit, initial_states=state["initial"],
+                            return_final_states=True)
+            y = y[:, pos]
+        else:
+            y = chunk_row(mixer, state, pos)
+        if pos == _CHUNK - 1:
+            state["initial"].copy_(final)
+            state_tensor.copy_(final)
+            state["states_in"] = None
+        state["pos"] = (pos + 1) % _CHUNK
+        return y
+
+    module.mamba_chunk_scan_combined = scan_with_cache
+    module.selective_state_update = update_from_chunk
+    module._ceno_scan_context = context
+    module._ceno_scan_parts = parts if fast else None
+    return context
+
+
+class _DecodeGraph:
+    """Batch-1 cached decode replayed as CUDA graphs, one per span between attention layers.
+
+    Every kernel is the one the eager decode step launches, with the same shapes and
+    inputs, so replay is bit-identical; only the host work is removed. Three eager
+    pieces remain: the attention mixers (growing KV cache), the chunk-boundary step
+    (full scan, handled by the caller falling back to eager) and the chunk-buffer
+    reset at chunk position 0. Two data-dependent host decisions are moved onto the
+    device: the chunk position (an index tensor) and the MoE expert loop (every
+    expert runs on the single token; a non-selected expert's result is discarded by
+    torch.where, a selected one goes through the same index_add_ as eagerly).
+    """
+
+    _KEYS = ("x", "dt", "B", "C", "z")
+
+    def __init__(self, model, modeling, cache):
+        import torch
+        self.model, self.modeling, self.cache = model, modeling, cache
+        self.parts = modeling._ceno_scan_parts
+        self.layers = model.backbone.layers
+        self.chunks = cache._ceno_chunks
+        self.mamba = [i for i, b in enumerate(self.layers) if b.block_type == "mamba"]
+        self.attention = [i for i, b in enumerate(self.layers) if b.block_type == "attention"]
+        self.inference = torch.is_inference_mode_enabled()
+        self.key = self._key()
+        device = model.device
+        self.ids = torch.zeros((1, 1), device=device, dtype=torch.long)
+        self.pos = torch.zeros((1,), device=device, dtype=torch.long)
+        self.row0 = torch.zeros((1,), device=device, dtype=torch.long)
+        hidden = model.config.hidden_size
+        self.attn_out = [torch.zeros((1, 1, hidden), device=device, dtype=model.dtype)
+                         for _ in self.attention]
+        self.attn_in, self.attn_res, self.graphs = [None] * len(self.attention), [None] * len(self.attention), []
+        self.logits = None
+        self._capture()
+
+    def _key(self):
+        import torch
+        ptrs = [self.cache.conv_states[i].data_ptr() for i in self.mamba]
+        for i in self.mamba:
+            state = self.chunks[i]
+            ptrs += [state[k].data_ptr() for k in self._KEYS + ("initial", "A")]
+        return (torch.is_inference_mode_enabled(), tuple(ptrs))
+
+    def valid(self, cache):
+        return cache is self.cache and self._key() == self.key
+
+    # ----- graph bodies: same module calls / kernels as the eager decode path -----
+    def _mamba(self, i, hidden_states):
+        import torch
+        mixer, cache, state = self.layers[i].mixer, self.cache, self.chunks[i]
+        hidden_states = mixer._maybe_apply_in_proj_norm(hidden_states)
+        projected_states = mixer.in_proj(hidden_states)
+        batch_size = hidden_states.shape[0]
+        groups_time_state_size = mixer.n_groups * mixer.ssm_state_size
+        d_mlp = (projected_states.shape[-1] - 2 * mixer.intermediate_size
+                 - 2 * groups_time_state_size - mixer.num_heads) // 2
+        _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+            [d_mlp, d_mlp, mixer.intermediate_size, mixer.conv_dim, mixer.num_heads], dim=-1)
+        hidden_states_B_C = self.modeling.causal_conv1d_update(
+            hidden_states_B_C, cache.conv_states[i], mixer.conv1d.weight.squeeze(1),
+            mixer.conv1d.bias, mixer.activation)
+        x, B, C = torch.split(hidden_states_B_C,
+                              [mixer.intermediate_size, groups_time_state_size, groups_time_state_size], dim=-1)
+        values = (x.view(batch_size, mixer.num_heads, mixer.head_dim), dt,
+                  B.view(batch_size, mixer.n_groups, -1), C.view(batch_size, mixer.n_groups, -1),
+                  gate.view(batch_size, mixer.num_heads, mixer.head_dim))
+        for key, value in zip(self._KEYS, values):
+            state[key].index_copy_(1, self.pos, value.unsqueeze(1))
+        parts = self.parts
+        initial = state["initial"]
+        batch, heads, dim, dstate = initial.shape
+        flat = initial.reshape(batch, heads, dim * dstate)
+        states_in, _ = parts._state_passing_fwd(
+            torch.zeros((batch, 1, heads, dim * dstate), device=flat.device, dtype=torch.float32),
+            torch.zeros((batch, heads, 1), device=flat.device, dtype=torch.float32),
+            initial_states=flat, chunk_size=_CHUNK, out_dtype=state["C"].dtype)
+        states_in = states_in.view(batch, 1, heads, dim, dstate)
+        dA_cumsum, dt = parts._chunk_cumsum_fwd(
+            state["dt"], state["A"], _CHUNK, dt_bias=mixer.dt_bias, dt_softplus=True,
+            dt_limit=mixer.time_step_limit)
+        CB = parts._bmm_chunk_fwd(state["C"], state["B"], _CHUNK, seq_idx=None, output_dtype=torch.float32)
+        D = mixer.D if mixer.D.stride(-1) == 1 else mixer.D.contiguous()
+        out, _ = parts._chunk_scan_fwd(CB, state["x"], dt, dA_cumsum, state["C"], states_in,
+                                       D=D, z=state["z"], seq_idx=None)
+        y = out.index_select(1, self.pos).view(batch_size, mixer.num_heads * mixer.head_dim)
+        y = mixer.norm(y, None)
+        return mixer.out_proj(y)[:, None, ...]
+
+    def _moe(self, moe, hidden_states):
+        import torch
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        router_softmax = getattr(moe, "_ceno_router_softmax", torch.nn.functional.softmax)
+        flat = hidden_states.view(-1, hidden_dim)
+        if str(getattr(moe.config, "moe_router_dtype", "")).lower() == "fp32":
+            router_logits = torch.nn.functional.linear(flat.to(torch.float32), moe.gate.weight.to(torch.float32), bias=None)
+        else:
+            router_logits = moe.gate(flat).to(torch.float32)
+        topk_logits, selected = torch.topk(router_logits, moe.top_k, dim=-1)
+        weights = router_softmax(topk_logits, dim=-1)
+        final = torch.zeros_like(flat)
+        for expert_idx in range(moe.num_experts):
+            hit = selected[0] == expert_idx                       # (top_k,)
+            weight = torch.where(hit, weights[0], 0.).sum().view(1, 1).to(flat.dtype)
+            out = moe.experts[expert_idx](flat, per_token_scale=weight)
+            added = final.index_add(0, self.row0, out.to(hidden_states.dtype))
+            final = torch.where(hit.any(), added, final)
+        return final.view(batch_size, seq_len, hidden_dim)
+
+    def _span(self, start, end, hidden_states):
+        for i in range(start, end):
+            block = self.layers[i]
+            if block.block_type == "mamba":
+                residual = hidden_states
+                if block.norm is not None:
+                    hidden_states = block.norm(hidden_states.to(dtype=block.norm.weight.dtype))
+                if block.residual_in_fp32:
+                    residual = residual.to(self._fp32)
+                hidden_states = residual + self._mamba(i, hidden_states)
+            elif block.block_type in ("moe", "mlp"):
+                input_ln_out = block.norm(hidden_states.to(dtype=block.norm.weight.dtype))
+                residual = input_ln_out
+                if block.residual_in_fp32:
+                    residual = residual.to(self._fp32)
+                pre = block.pre_mlp_norm(input_ln_out.to(dtype=block.pre_mlp_norm.weight.dtype))
+                out = self._moe(block.mixer, pre) if block.block_type == "moe" else block.mixer(pre)
+                hidden_states = residual + out
+            else:
+                raise RuntimeError("attention layers are span boundaries")
+        return hidden_states
+
+    def _segment(self, k):
+        """Span k: after attention k-1 (or the embedding) up to attention k's input norm."""
+        bounds = [-1] + self.attention + [len(self.layers)]
+        start, end = bounds[k] + 1, bounds[k + 1]
+        if k == 0:
+            hidden_states = self.model.backbone.embeddings(self.ids)
+        else:
+            block = self.layers[bounds[k]]
+            hidden_states = self.attn_res[k - 1] + self.attn_out[k - 1]
+            hidden_states = block.pre_mlp_norm(hidden_states.to(dtype=block.pre_mlp_norm.weight.dtype))
+        hidden_states = self._span(start, end, hidden_states)
+        if k < len(self.attention):
+            block = self.layers[end]
+            residual = hidden_states
+            self.attn_in[k] = block.norm(hidden_states.to(dtype=block.norm.weight.dtype))
+            self.attn_res[k] = residual.to(self._fp32) if block.residual_in_fp32 else residual
+        else:
+            hidden_states = self.model.backbone.norm_f(hidden_states)
+            lm_head = self.model.lm_head
+            self.logits = lm_head(hidden_states.to(lm_head.weight.dtype)).float()
+
+    def _capture(self):
+        import torch
+        self._fp32 = torch.float32
+        segments = len(self.attention) + 1
+        # The warm-up really runs the bodies: it shifts the conv states and writes the
+        # chunk buffers. Snapshot everything it can write and restore it afterwards.
+        mutable = [self.cache.conv_states[i] for i in self.mamba]
+        mutable += [self.chunks[i][key] for i in self.mamba for key in self._KEYS]
+        snapshot = [t.clone() for t in mutable]
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            for k in range(segments):
+                self._segment(k)
+        torch.cuda.current_stream().wait_stream(stream)
+        for tensor, saved in zip(mutable, snapshot):
+            tensor.copy_(saved)
+        del snapshot
+        pool = torch.cuda.graph_pool_handle()
+        for k in range(segments):
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=pool, stream=stream):
+                self._segment(k)
+            self.graphs.append(graph)
+        torch.cuda.synchronize()
+
+    def step(self, input_ids, pos):
+        if pos == 0:
+            for i in self.mamba:
+                for key in self._KEYS:
+                    self.chunks[i][key].zero_()
+        self.ids.copy_(input_ids)
+        self.pos.fill_(pos)
+        for k, graph in enumerate(self.graphs):
+            graph.replay()
+            if k < len(self.attention):
+                mixer = self.layers[self.attention[k]].mixer
+                out = mixer(self.attn_in[k], attention_mask=None, past_key_value=self.cache,
+                            use_cache=True, cache_position=None)[0]
+                self.attn_out[k].copy_(out)
+        for i in self.mamba:
+            self.chunks[i]["pos"] = pos + 1
+        return self.logits.clone()
+
+
+def _decode_graph_step(model, modeling, kwargs):
+    """Return logits from the graph path, or None to run the eager forward."""
+    if os.environ.get("CENO_HF_DECODE_GRAPH", "1") == "0" or getattr(modeling, "_ceno_scan_parts", None) is None:
+        return None
+    if any(kwargs.get(k) is not None for k in ("inputs_embeds", "labels", "seq_idx")) or any(
+            kwargs.get(k) for k in ("output_attentions", "output_hidden_states")):
+        return None
+    if kwargs.get("use_cache") is False or (kwargs.get("use_cache") is None and not model.config.use_cache):
+        return None
+    ids = kwargs.get("input_ids")
+    cache = kwargs.get("cache_params") or kwargs.get("past_key_values")
+    position = kwargs.get("cache_position")
+    if (ids is None or tuple(ids.shape) != (1, 1) or cache is None or position is None
+            or position.numel() != 1 or model.config.output_hidden_states or model.config.output_attentions):
+        return None
+    chunks = getattr(cache, "_ceno_chunks", None)
+    layers = model.backbone.layers
+    mamba = [i for i, b in enumerate(layers) if b.block_type == "mamba"]
+    if not mamba or not chunks or any(i not in chunks or not chunks[i].get("decoded") for i in mamba):
+        return None
+    positions = {chunks[i]["pos"] for i in mamba}
+    if len(positions) != 1 or any(len(cache.key_cache[i].shape) != 4 for i in cache.attention_layers):
+        return None
+    pos = positions.pop()
+    if pos == _CHUNK - 1:
+        return None  # chunk boundary: full scan in the eager path
+    runner = getattr(cache, "_ceno_decode_graph", None)
+    with backend_scope():
+        if runner is None or not runner.valid(cache):
+            runner = cache._ceno_decode_graph = None
+            if pos == 0:  # the eager reset happens in step(); capture after it
+                for i in mamba:
+                    for key in _DecodeGraph._KEYS:
+                        chunks[i][key].zero_()
+            runner = cache._ceno_decode_graph = _DecodeGraph(model, modeling, cache)
+        return runner.step(ids, pos)
+
+
+def _attention(q, k, v, causal):
+    """Ordinary FA2 always; dense FA2 can change arithmetic with query length."""
+    import torch
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward
+    q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+    # Build [0, n] on the device: torch.tensor(..., device=cuda) copies from
+    # pageable host memory and synchronizes the stream on every call.
+    steps = torch.arange(2, device=q.device, dtype=torch.int32)
+    cuq, cuk = steps * len(q), steps * len(k)
+    result = _flash_attn_varlen_forward(
+        q, k, v, cuq, cuk, len(q), len(k), dropout_p=0., softmax_scale=q.shape[-1] ** -.5,
+        causal=causal, window_size_left=-1, window_size_right=0 if causal else -1,
+        softcap=0., alibi_slopes=None, return_softmax=False, block_table=None,
+    )
+    return result[0], result[1]
+
+
+def _cp_attention(q, k, v, reference_length):
+    import torch
+    combine = _local_module("ceno_cp_math").combine
+    cp = 8
+    block = reference_length // (2 * cp)
+    length, begin = len(k), len(k) - len(q)
+    out = torch.empty_like(q)
+
+    def segment(x, chunk):
+        return x[chunk * block:min((chunk + 1) * block, length)]
+
+    def part(x, rank, both):
+        first = segment(x, rank)
+        return torch.cat([first, segment(x, 2 * cp - rank - 1)]) if both else first
+
+    for rank in range(cp):
+        intervals = []
+        for chunk in (rank, 2 * cp - rank - 1):
+            start, end = max(begin, chunk * block), min(length, (chunk + 1) * block)
+            intervals.append((start - begin, end - begin) if end > start else (0, 0))
+        (a0, a1), (b0, b1) = intervals
+        first, second = a1 - a0, b1 - b0
+        if not first + second:
+            continue
+        query = torch.cat([q[a0:a1], q[b0:b1]])
+        pieces, offsets = [], []
+        for step in range(cp):
+            if step > rank and not second:
+                continue
+            src, both = (rank - step) % cp, step == 0 or step > rank
+            query_part = query if step <= rank else query[first:]
+            pieces.append(_attention(query_part, part(k, src, both), part(v, src, both), step == 0))
+            offsets.append(0 if step <= rank else first)
+        result = combine(pieces, offsets)
+        if first:
+            out[a0:a1].copy_(result[:first])
+        if second:
+            out[b0:b1].copy_(result[first:])
+    return out
+
+
+def enable_training_precision(model, *, norm_mode=None, reference_cp_size=None,
+                              reference_context_length=None):
+    """Attach the common inference implementation to an already loaded HF model."""
+    if hasattr(model, "ceno_precision_reference"):
+        return model
+    configure_backend()
+    import sys
+    import torch
+    norm_module = _local_module("ceno_norm")
+    te_rms_norm, te_rms_norm_eager = norm_module.te_rms_norm, norm_module._te_rms_norm_impl
+
+    config = model.config
+    modeling = sys.modules[type(model).__module__]
+    prefix = type(model).__name__.removesuffix("ForCausalLM")
+    norm_mode = norm_mode or getattr(config, "ceno_rmsnorm_mode", "te-tuned")
+    reference_cp_size = (getattr(config, "ceno_reference_cp_size", 1)
+                         if reference_cp_size is None else reference_cp_size)
+    reference_context_length = (getattr(config, "ceno_reference_context_length", None)
+                                if reference_context_length is None else reference_context_length)
+    expected = dict(hidden_size=(512, 1024), mamba_num_heads=(16, 32), mamba_head_dim=(64,),
+                    n_groups=(8,), ssm_state_size=(128,), chunk_size=(128,),
+                    num_attention_heads=(8, 16), num_experts=(8,), moe_top_k=(2,))
+    if (prefix not in ("NemotronH", "CENO")
+            or any(getattr(config, key, None) not in values for key, values in expected.items())
+            or config.num_key_value_heads != config.num_attention_heads):
+        raise ValueError("HF training precision supports ungrouped CENO 80M/300M/600M/1B")
+    if (getattr(model, "is_quantized", False)
+            or getattr(config, "moe_router_pre_softmax", False)
+            or getattr(config, "mlp_use_swiglu", False)
+            or getattr(config, "mlp_hidden_act", "relu2") not in ("relu2", "squared_relu")
+            or tuple(getattr(config, "time_step_limit", (0., float("inf")))) not in ((0., float("inf")), (0., 1e30))):
+        raise ValueError("HF precision requires unquantized post-topk ReLU2 and default Mamba dt limits")
+    if norm_mode not in ("te-general", "te-tuned") or reference_cp_size not in (1, 8):
+        raise ValueError("Choose an explicit norm and CP reference")
+    if reference_cp_size == 8 and (
+        norm_mode != "te-tuned" or not isinstance(reference_context_length, int)
+        or reference_context_length <= 0 or reference_context_length % 1024
+    ):
+        raise ValueError("CP8 requires te-tuned norm and a fixed reference window divisible by 1024")
+    if model.device.type != "cuda" or model.dtype != torch.bfloat16 or not modeling.is_fast_path_available:
+        raise ValueError("HF precision requires single-GPU BF16 and the CUDA Mamba/FA2 environment")
+    if any(p.device != model.device for p in model.parameters()):
+        raise ValueError("HF training precision requires one CUDA device")
+    mark = getattr(modeling, "_mark_backend_scope", None)
+    if mark is None:
+        raise ValueError("HF precision requires a modeling file with scoped backend flags")
+    mark(model)
+    config.ceno_mamba_gate_in_ssm = True
+    config.ceno_rmsnorm_mode = norm_mode
+    config.ceno_reference_cp_size = reference_cp_size
+    config.ceno_reference_context_length = reference_context_length
+    _pin_scan()
+    context = _install_scan_hooks(modeling)
+    def norm(self, x):
+        # Same Triton kernel; the registered op is only needed under torch.compile.
+        op = te_rms_norm if torch.compiler.is_compiling() else te_rms_norm_eager
+        return op(x, self.weight, self.variance_epsilon, norm_mode)
+
+    def softmax(x, dim=-1):
+        exp = torch.exp(x - torch.amax(x, dim=dim, keepdim=True))
+        return exp / torch.sum(exp, dim=dim, keepdim=True)
+
+    def attention(q, k, v, mask, query_length, **kwargs):
+        if kwargs.get("sliding_window") is not None or kwargs.get("dropout", 0) != 0:
+            raise ValueError("Only full causal inference attention is validated")
+        if reference_cp_size == 8 and k.shape[1] > reference_context_length:
+            raise ValueError("Sequence exceeds the fixed CP reference window")
+        results = []
+        for batch in range(q.shape[0]):
+            if reference_cp_size == 1 or k.shape[1] <= reference_context_length // 16:
+                results.append(_attention(q[batch], k[batch], v[batch], True)[0])
+            else:
+                results.append(_cp_attention(q[batch], k[batch], v[batch], reference_context_length))
+        return torch.stack(results)
+
+    for child in model.modules():
+        if isinstance(child, getattr(modeling, prefix + "RMSNorm")):
+            child.forward = types.MethodType(norm, child)
+        elif isinstance(child, getattr(modeling, prefix + "MoE")):
+            child._ceno_router_softmax = softmax
+        elif isinstance(child, getattr(modeling, prefix + "Attention")):
+            child.__class__ = getattr(modeling, prefix + "FlashAttention2")
+            child._flash_attn_uses_top_left_mask = False
+            child._ceno_attention = attention
+        elif isinstance(child, getattr(modeling, prefix + "Mamba2Mixer")):
+            original = child.cuda_kernels_forward
+
+            def mixer_forward(self, hidden_states, cache_params=None, *args,
+                              _original=original, **kwargs):
+                token = context.set((self, cache_params))
+                try:
+                    return _original(hidden_states, cache_params, *args, **kwargs)
+                finally:
+                    context.reset(token)
+
+            child.cuda_kernels_forward = types.MethodType(mixer_forward, child)
+
+    original_forward = model.forward
+
+    @functools.wraps(original_forward)
+    def checked_forward(*args, **kwargs):
+        if model.training or torch.is_grad_enabled():
+            raise ValueError("Training precision is inference-only; use eval and no_grad/inference_mode")
+        if kwargs.get("output_attentions") or model.config.output_attentions:
+            raise ValueError("Reload with attn_implementation='eager' to inspect attention maps")
+        ids = args[0] if args else kwargs.get("input_ids")
+        inputs = ids if ids is not None else kwargs.get("inputs_embeds")
+        if inputs is not None and inputs.shape[0] * inputs.shape[1] > 131072:
+            raise ValueError("HF precision prefill budget is 131072 tokens; use vLLM for larger prefill")
+        if kwargs.get("seq_idx") is not None:
+            raise ValueError("Base generation precision does not implement MSA segments")
+        mask = kwargs.get("attention_mask")
+        if mask is not None and not bool(torch.all(mask)):
+            raise ValueError("HF precision currently requires unpadded uniform-length batches")
+        if len(args) <= 1 and not (args and "input_ids" in kwargs):
+            graph_kwargs = dict(kwargs, input_ids=args[0]) if args else kwargs
+            logits = _decode_graph_step(model, modeling, graph_kwargs)
+            if logits is not None:
+                cache = kwargs.get("cache_params") or kwargs.get("past_key_values")
+                if kwargs.get("return_dict", model.config.use_return_dict) is False:
+                    return (logits, cache)
+                return modeling.CENOCausalLMOutput(logits=logits, cache_params=cache,
+                                                        past_key_values=cache)
+        return original_forward(*args, **kwargs)
+
+    model.forward = checked_forward
+    model.ceno_precision_reference = dict(dtype="bfloat16", norm=norm_mode,
+        cp_size=reference_cp_size, context_length=reference_context_length)
+    return model
+
+
+def load_model(model_dir, *, norm_mode="te-tuned", reference_cp_size=None,
+               reference_context_length=None, device="cuda"):
+    """Load the same implementation used by a configured AutoModel checkpoint."""
+    configure_backend()
+    import torch
+    CENOConfig = _local_module("configuration_ceno").CENOConfig
+    CENOForCausalLM = _local_module("modeling_ceno").CENOForCausalLM
+
+    config = CENOConfig.from_pretrained(model_dir)
+    model = CENOForCausalLM.from_pretrained(
+        model_dir, config=config, torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+    ).to(device).eval()
+    return enable_training_precision(model, norm_mode=norm_mode,
+        reference_cp_size=reference_cp_size, reference_context_length=reference_context_length)

vocab.json

@@ -0,0 +1,485 @@
+{
+  " ": 32,
+  "!": 33,
+  "\"": 34,
+  "#": 35,
+  "$": 36,
+  "%": 37,
+  "&": 38,
+  "'": 39,
+  "(": 40,
+  ")": 41,
+  "*": 42,
+  "+": 43,
+  ",": 44,
+  "-": 45,
+  ".": 46,
+  "/": 47,
+  "0": 48,
+  "1": 49,
+  "2": 50,
+  "3": 51,
+  "4": 52,
+  "5": 53,
+  "6": 54,
+  "7": 55,
+  "8": 56,
+  "9": 57,
+  ":": 58,
+  ";": 59,
+  "<": 60,
+  "<byte_127>": 127,
+  "<byte_128>": 128,
+  "<byte_129>": 129,
+  "<byte_130>": 130,
+  "<byte_131>": 131,
+  "<byte_132>": 132,
+  "<byte_133>": 133,
+  "<byte_134>": 134,
+  "<byte_135>": 135,
+  "<byte_136>": 136,
+  "<byte_137>": 137,
+  "<byte_138>": 138,
+  "<byte_139>": 139,
+  "<byte_140>": 140,
+  "<byte_141>": 141,
+  "<byte_142>": 142,
+  "<byte_143>": 143,
+  "<byte_144>": 144,
+  "<byte_145>": 145,
+  "<byte_146>": 146,
+  "<byte_147>": 147,
+  "<byte_148>": 148,
+  "<byte_149>": 149,
+  "<byte_150>": 150,
+  "<byte_151>": 151,
+  "<byte_152>": 152,
+  "<byte_153>": 153,
+  "<byte_154>": 154,
+  "<byte_155>": 155,
+  "<byte_156>": 156,
+  "<byte_157>": 157,
+  "<byte_158>": 158,
+  "<byte_159>": 159,
+  "<byte_160>": 160,
+  "<byte_161>": 161,
+  "<byte_162>": 162,
+  "<byte_163>": 163,
+  "<byte_164>": 164,
+  "<byte_165>": 165,
+  "<byte_166>": 166,
+  "<byte_167>": 167,
+  "<byte_168>": 168,
+  "<byte_169>": 169,
+  "<byte_170>": 170,
+  "<byte_171>": 171,
+  "<byte_172>": 172,
+  "<byte_173>": 173,
+  "<byte_174>": 174,
+  "<byte_175>": 175,
+  "<byte_176>": 176,
+  "<byte_177>": 177,
+  "<byte_178>": 178,
+  "<byte_179>": 179,
+  "<byte_180>": 180,
+  "<byte_181>": 181,
+  "<byte_182>": 182,
+  "<byte_183>": 183,
+  "<byte_184>": 184,
+  "<byte_185>": 185,
+  "<byte_186>": 186,
+  "<byte_187>": 187,
+  "<byte_188>": 188,
+  "<byte_189>": 189,
+  "<byte_190>": 190,
+  "<byte_191>": 191,
+  "<byte_192>": 192,
+  "<byte_193>": 193,
+  "<byte_194>": 194,
+  "<byte_195>": 195,
+  "<byte_196>": 196,
+  "<byte_197>": 197,
+  "<byte_198>": 198,
+  "<byte_199>": 199,
+  "<byte_200>": 200,
+  "<byte_201>": 201,
+  "<byte_202>": 202,
+  "<byte_203>": 203,
+  "<byte_204>": 204,
+  "<byte_205>": 205,
+  "<byte_206>": 206,
+  "<byte_207>": 207,
+  "<byte_208>": 208,
+  "<byte_209>": 209,
+  "<byte_210>": 210,
+  "<byte_211>": 211,
+  "<byte_212>": 212,
+  "<byte_213>": 213,
+  "<byte_214>": 214,
+  "<byte_215>": 215,
+  "<byte_216>": 216,
+  "<byte_217>": 217,
+  "<byte_218>": 218,
+  "<byte_219>": 219,
+  "<byte_220>": 220,
+  "<byte_221>": 221,
+  "<byte_222>": 222,
+  "<byte_223>": 223,
+  "<byte_224>": 224,
+  "<byte_225>": 225,
+  "<byte_226>": 226,
+  "<byte_227>": 227,
+  "<byte_228>": 228,
+  "<byte_229>": 229,
+  "<byte_230>": 230,
+  "<byte_231>": 231,
+  "<byte_232>": 232,
+  "<byte_233>": 233,
+  "<byte_234>": 234,
+  "<byte_235>": 235,
+  "<byte_236>": 236,
+  "<byte_237>": 237,
+  "<byte_238>": 238,
+  "<byte_239>": 239,
+  "<byte_240>": 240,
+  "<byte_241>": 241,
+  "<byte_242>": 242,
+  "<byte_243>": 243,
+  "<byte_244>": 244,
+  "<byte_245>": 245,
+  "<byte_246>": 246,
+  "<byte_247>": 247,
+  "<byte_248>": 248,
+  "<byte_249>": 249,
+  "<byte_250>": 250,
+  "<byte_251>": 251,
+  "<byte_252>": 252,
+  "<byte_253>": 253,
+  "<byte_254>": 254,
+  "<byte_255>": 255,
+  "<byte_256>": 256,
+  "<byte_257>": 257,
+  "<byte_258>": 258,
+  "<byte_259>": 259,
+  "<byte_260>": 260,
+  "<byte_261>": 261,
+  "<byte_262>": 262,
+  "<byte_263>": 263,
+  "<byte_264>": 264,
+  "<byte_265>": 265,
+  "<byte_266>": 266,
+  "<byte_267>": 267,
+  "<byte_268>": 268,
+  "<byte_269>": 269,
+  "<byte_270>": 270,
+  "<byte_271>": 271,
+  "<byte_272>": 272,
+  "<byte_273>": 273,
+  "<byte_274>": 274,
+  "<byte_275>": 275,
+  "<byte_276>": 276,
+  "<byte_277>": 277,
+  "<byte_278>": 278,
+  "<byte_279>": 279,
+  "<byte_280>": 280,
+  "<byte_281>": 281,
+  "<byte_282>": 282,
+  "<byte_283>": 283,
+  "<byte_284>": 284,
+  "<byte_285>": 285,
+  "<byte_286>": 286,
+  "<byte_287>": 287,
+  "<byte_288>": 288,
+  "<byte_289>": 289,
+  "<byte_290>": 290,
+  "<byte_291>": 291,
+  "<byte_292>": 292,
+  "<byte_293>": 293,
+  "<byte_294>": 294,
+  "<byte_295>": 295,
+  "<byte_296>": 296,
+  "<byte_297>": 297,
+  "<byte_298>": 298,
+  "<byte_299>": 299,
+  "<byte_300>": 300,
+  "<byte_301>": 301,
+  "<byte_302>": 302,
+  "<byte_303>": 303,
+  "<byte_304>": 304,
+  "<byte_305>": 305,
+  "<byte_306>": 306,
+  "<byte_307>": 307,
+  "<byte_308>": 308,
+  "<byte_309>": 309,
+  "<byte_310>": 310,
+  "<byte_311>": 311,
+  "<byte_312>": 312,
+  "<byte_313>": 313,
+  "<byte_314>": 314,
+  "<byte_315>": 315,
+  "<byte_316>": 316,
+  "<byte_317>": 317,
+  "<byte_318>": 318,
+  "<byte_319>": 319,
+  "<byte_320>": 320,
+  "<byte_321>": 321,
+  "<byte_322>": 322,
+  "<byte_323>": 323,
+  "<byte_324>": 324,
+  "<byte_325>": 325,
+  "<byte_326>": 326,
+  "<byte_327>": 327,
+  "<byte_328>": 328,
+  "<byte_329>": 329,
+  "<byte_330>": 330,
+  "<byte_331>": 331,
+  "<byte_332>": 332,
+  "<byte_333>": 333,
+  "<byte_334>": 334,
+  "<byte_335>": 335,
+  "<byte_336>": 336,
+  "<byte_337>": 337,
+  "<byte_338>": 338,
+  "<byte_339>": 339,
+  "<byte_340>": 340,
+  "<byte_341>": 341,
+  "<byte_342>": 342,
+  "<byte_343>": 343,
+  "<byte_344>": 344,
+  "<byte_345>": 345,
+  "<byte_346>": 346,
+  "<byte_347>": 347,
+  "<byte_348>": 348,
+  "<byte_349>": 349,
+  "<byte_350>": 350,
+  "<byte_351>": 351,
+  "<byte_352>": 352,
+  "<byte_353>": 353,
+  "<byte_354>": 354,
+  "<byte_355>": 355,
+  "<byte_356>": 356,
+  "<byte_357>": 357,
+  "<byte_358>": 358,
+  "<byte_359>": 359,
+  "<byte_360>": 360,
+  "<byte_361>": 361,
+  "<byte_362>": 362,
+  "<byte_363>": 363,
+  "<byte_364>": 364,
+  "<byte_365>": 365,
+  "<byte_366>": 366,
+  "<byte_367>": 367,
+  "<byte_368>": 368,
+  "<byte_369>": 369,
+  "<byte_370>": 370,
+  "<byte_371>": 371,
+  "<byte_372>": 372,
+  "<byte_373>": 373,
+  "<byte_374>": 374,
+  "<byte_375>": 375,
+  "<byte_376>": 376,
+  "<byte_377>": 377,
+  "<byte_378>": 378,
+  "<byte_379>": 379,
+  "<byte_380>": 380,
+  "<byte_381>": 381,
+  "<byte_382>": 382,
+  "<byte_383>": 383,
+  "<byte_384>": 384,
+  "<byte_385>": 385,
+  "<byte_386>": 386,
+  "<byte_387>": 387,
+  "<byte_388>": 388,
+  "<byte_389>": 389,
+  "<byte_390>": 390,
+  "<byte_391>": 391,
+  "<byte_392>": 392,
+  "<byte_393>": 393,
+  "<byte_394>": 394,
+  "<byte_395>": 395,
+  "<byte_396>": 396,
+  "<byte_397>": 397,
+  "<byte_398>": 398,
+  "<byte_399>": 399,
+  "<byte_400>": 400,
+  "<byte_401>": 401,
+  "<byte_402>": 402,
+  "<byte_403>": 403,
+  "<byte_404>": 404,
+  "<byte_405>": 405,
+  "<byte_406>": 406,
+  "<byte_407>": 407,
+  "<byte_408>": 408,
+  "<byte_409>": 409,
+  "<byte_410>": 410,
+  "<byte_411>": 411,
+  "<byte_412>": 412,
+  "<byte_413>": 413,
+  "<byte_414>": 414,
+  "<byte_415>": 415,
+  "<byte_416>": 416,
+  "<byte_417>": 417,
+  "<byte_418>": 418,
+  "<byte_419>": 419,
+  "<byte_420>": 420,
+  "<byte_421>": 421,
+  "<byte_422>": 422,
+  "<byte_423>": 423,
+  "<byte_424>": 424,
+  "<byte_425>": 425,
+  "<byte_426>": 426,
+  "<byte_427>": 427,
+  "<byte_428>": 428,
+  "<byte_429>": 429,
+  "<byte_430>": 430,
+  "<byte_431>": 431,
+  "<byte_432>": 432,
+  "<byte_433>": 433,
+  "<byte_434>": 434,
+  "<byte_435>": 435,
+  "<byte_436>": 436,
+  "<byte_437>": 437,
+  "<byte_438>": 438,
+  "<byte_439>": 439,
+  "<byte_440>": 440,
+  "<byte_441>": 441,
+  "<byte_442>": 442,
+  "<byte_443>": 443,
+  "<byte_444>": 444,
+  "<byte_445>": 445,
+  "<byte_446>": 446,
+  "<byte_447>": 447,
+  "<byte_448>": 448,
+  "<byte_449>": 449,
+  "<byte_450>": 450,
+  "<byte_451>": 451,
+  "<byte_452>": 452,
+  "<byte_453>": 453,
+  "<byte_454>": 454,
+  "<byte_455>": 455,
+  "<byte_456>": 456,
+  "<byte_457>": 457,
+  "<byte_458>": 458,
+  "<byte_459>": 459,
+  "<byte_460>": 460,
+  "<byte_461>": 461,
+  "<byte_462>": 462,
+  "<byte_463>": 463,
+  "<byte_464>": 464,
+  "<byte_465>": 465,
+  "<byte_466>": 466,
+  "<byte_467>": 467,
+  "<byte_468>": 468,
+  "<byte_469>": 469,
+  "<byte_470>": 470,
+  "<byte_471>": 471,
+  "<byte_472>": 472,
+  "<byte_473>": 473,
+  "<byte_474>": 474,
+  "<byte_475>": 475,
+  "<byte_476>": 476,
+  "<byte_477>": 477,
+  "<byte_478>": 478,
+  "<byte_479>": 479,
+  "<byte_480>": 480,
+  "<byte_481>": 481,
+  "<byte_482>": 482,
+  "<byte_483>": 483,
+  "<byte_484>": 484,
+  "<byte_485>": 485,
+  "<byte_486>": 486,
+  "<byte_487>": 487,
+  "<byte_488>": 488,
+  "<byte_489>": 489,
+  "<byte_490>": 490,
+  "<byte_491>": 491,
+  "<byte_492>": 492,
+  "<byte_493>": 493,
+  "<byte_494>": 494,
+  "<byte_495>": 495,
+  "<byte_496>": 496,
+  "<byte_497>": 497,
+  "<byte_498>": 498,
+  "<byte_499>": 499,
+  "<byte_500>": 500,
+  "<byte_501>": 501,
+  "<byte_502>": 502,
+  "<byte_503>": 503,
+  "<byte_504>": 504,
+  "<byte_505>": 505,
+  "<byte_506>": 506,
+  "<byte_507>": 507,
+  "<byte_508>": 508,
+  "<byte_509>": 509,
+  "<byte_510>": 510,
+  "<byte_511>": 511,
+  "<eos>": 0,
+  "<pad>": 1,
+  "<unk>": 2,
+  "=": 61,
+  ">": 62,
+  "?": 63,
+  "@": 64,
+  "A": 65,
+  "B": 66,
+  "C": 67,
+  "D": 68,
+  "E": 69,
+  "F": 70,
+  "G": 71,
+  "H": 72,
+  "I": 73,
+  "J": 74,
+  "K": 75,
+  "L": 76,
+  "M": 77,
+  "N": 78,
+  "O": 79,
+  "P": 80,
+  "Q": 81,
+  "R": 82,
+  "S": 83,
+  "T": 84,
+  "U": 85,
+  "V": 86,
+  "W": 87,
+  "X": 88,
+  "Y": 89,
+  "Z": 90,
+  "[": 91,
+  "\\": 92,
+  "]": 93,
+  "^": 94,
+  "_": 95,
+  "`": 96,
+  "a": 97,
+  "b": 98,
+  "c": 99,
+  "d": 100,
+  "e": 101,
+  "f": 102,
+  "g": 103,
+  "h": 104,
+  "i": 105,
+  "j": 106,
+  "k": 107,
+  "l": 108,
+  "m": 109,
+  "n": 110,
+  "o": 111,
+  "p": 112,
+  "q": 113,
+  "r": 114,
+  "s": 115,
+  "t": 116,
+  "u": 117,
+  "v": 118,
+  "w": 119,
+  "x": 120,
+  "y": 121,
+  "z": 122,
+  "{": 123,
+  "|": 124,
+  "}": 125,
+  "~": 126
+}