Initial release: CENO-1B-base (Apache-2.0)

CENO base, 1B params (stage2), bfloat16. Bundles trust_remote_code model code + tokenizer. Derived from NVIDIA Nemotron-H (Apache-2.0); tokenizer derived from Arc Institute Evo2 CharLevelTokenizer (Apache-2.0).

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,18 @@
+CENO Checkpoint — CENO-1B-base
+Copyright (c) 2025-2026 CENO Authors. All rights reserved.
+
+This directory bundles a trained model checkpoint (CENO base model) together with a
+copy of the CENO model code (configuration, modeling, and tokenizer modules).
+
+The bundled model code is derived from NVIDIA's Nemotron-H HuggingFace
+implementation, which is licensed under the Apache License, Version 2.0.
+Nemotron-H is © NVIDIA Corporation. The CENO model code is a derivative work
+distributed under the same Apache License, Version 2.0, included in this
+directory as the `LICENSE` file.
+
+The tokenizer module (`ceno_tokenizer.py`) is derived from Arc Institute's
+Evo2 CharLevelTokenizer, which is licensed under the Apache License, Version 2.0.
+
+All upstream copyright notices and license terms (NVIDIA; Arc Institute) are
+preserved in the corresponding source files as required by the Apache License,
+Version 2.0.

README.md

@@ -0,0 +1,76 @@
+---
+license: apache-2.0
+library_name: transformers
+tags: ["dna", "genomics", "dna-language-model", "mamba", "moe"]
+pipeline_tag: text-generation
+language:
+  - dna
+---
+
+# CENO-1B-base
+
+**CENO-1B-base** is a checkpoint of the CENO base DNA foundation model (base (stage 2)). It is a plain causal language model over genomic sequence on a Nemotron-H Mamba/Attention/MoE hybrid backbone, with no MSA inputs.
+
+This checkpoint is part of the **CENO** DNA foundation model family. The model
+code, VEP pipeline, and generation demo live in the sibling [CENO code repository](https://github.com/CladeTeam/CENO); this
+directory is standalone-loadable via `trust_remote_code=True` (the model code is
+bundled here).
+
+## Model details
+
+| | |
+|---|---|
+| Family | CENO (base) |
+| Stage | base (stage 2) |
+| Parameters | 1302.4M |
+| Precision | bfloat16 |
+| Weights | `model.safetensors` |
+| `model_type` | `ceno` |
+| `architectures` | `CENOForCausalLM` |
+| `auto_map` → model | `modeling_ceno.CENOForCausalLM` |
+| `auto_map` → tokenizer | `ceno_tokenizer.CENOCharLevelTokenizer` |
+
+### Architecture
+
+| Hidden layers | 38 |
+| Context length | 1048576 |
+| Vocab size | 512 |
+| Attention heads | 16 |
+| Intermediate size | 4096 |
+| Num experts (MoE) | 8 |
+| Experts per token | 2 |
+
+The backbone is a Mamba / Attention / Mixture-of-Experts hybrid (Nemotron-H
+architecture). The tokenizer is byte-level (character-level), mapping DNA
+characters to their ASCII byte codes (vocab size 512).
+
+## Loading
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+ckpt = "CENO-1B-base"   # path to this directory
+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 GPU (Mamba kernels)
+```
+
+> The Mamba layers require CUDA kernels, so forward / generation needs 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) on genomics tasks.
+- **MSA checkpoints (CENO-P-*)**: variant effect prediction (VEP) by scoring
+  wild-type vs. variant sequences with delta log-likelihood, using the MSA
+  scoring path. 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
+HuggingFace implementation (Apache-2.0); the tokenizer is derived from Arc
+Institute's Evo2 CharLevelTokenizer (Apache-2.0). See the `LICENSE` and `NOTICE`
+files in this directory for full attribution.

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,79 @@
+{
+  "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": 64,
+  "hidden_dropout": 0.0,
+  "hidden_size": 1024,
+  "hybrid_override_pattern": "MEMEM*EMEMEM*EMEMEMEMEM*EMEMEM*EMEMEME",
+  "initializer_range": 0.02,
+  "intermediate_size": 4096,
+  "layer_norm_epsilon": 1e-05,
+  "mamba_head_dim": 64,
+  "mamba_hidden_act": "silu",
+  "mamba_in_proj_layernorm": true,
+  "mamba_num_heads": 32,
+  "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": 4096,
+  "moe_router_dtype": "bfloat16",
+  "moe_router_pre_softmax": false,
+  "moe_shared_expert_intermediate_size": 4096,
+  "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": 16,
+  "num_experts": 8,
+  "num_experts_per_tok": 2,
+  "num_hidden_layers": 38,
+  "num_key_value_heads": 16,
+  "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,
+    1e30
+  ],
+  "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
+}

configuration_ceno.py

@@ -0,0 +1,333 @@
+# 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 0):
+            The id of the padding token.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            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,
+        intra_encoding_pattern=None,
+        **kwargs,
+    ):
+        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

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

modeling_ceno.py

@@ -0,0 +1,2002 @@
+# coding=utf-8
+# Copyright 2024 HuggingFace Inc. team.
+# 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.
+"""PyTorch CENO model."""
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import DynamicCache  # we need __iter__ and __len__ of pkv
+
+# Define squared_relu / relu2 activation if not already in ACT2FN
+def squared_relu(x: torch.Tensor) -> torch.Tensor:
+    """Squared ReLU activation: (max(0, x))^2"""
+    return torch.pow(torch.relu(x), 2)
+
+# Register relu2/squared_relu in ACT2FN if not present
+if "relu2" not in ACT2FN:
+    ACT2FN["relu2"] = squared_relu
+if "squared_relu" not in ACT2FN:
+    ACT2FN["squared_relu"] = squared_relu
+from transformers.generation.utils import GenerationMixin
+from transformers.modeling_attn_mask_utils import (
+    AttentionMaskConverter,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (
+    ModelOutput,
+    add_code_sample_docstrings,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+)
+from transformers.utils.import_utils import (
+    is_causal_conv1d_available,
+    is_flash_attn_2_available,
+    is_flash_attn_greater_or_equal_2_10,
+    is_mamba_2_ssm_available,    
+)
+try:
+    # Preferred import when using this repo as a package.
+    from ceno_hf.configuration_ceno import CENOConfig
+except Exception:
+    try:
+        # Fallback for relative import contexts.
+        from .configuration_ceno import CENOConfig
+    except Exception:
+        # Last resort for running from within the directory.
+        from configuration_ceno import CENOConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.mamba.modeling_mamba2.modeling_mamba2.py with MAMBA2->CENO,Mamba2->CENO
+# For Mamba2 components Mamba2->CENOMamba2
+if is_mamba_2_ssm_available():
+    from mamba_ssm.ops.triton.selective_state_update import selective_state_update
+    from mamba_ssm.ops.triton.ssd_combined import mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
+else:
+    mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined, selective_state_update = None, None, None
+
+try:
+    #from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated
+    from mamba_ssm.ops.triton.layernorm_gated import rmsnorm_fn
+except Exception:
+    # mamba_ssm is optional at import time (it may fail to initialize on a
+    # GPU-less machine, e.g. triton driver init). The kernel is only needed
+    # inside MambaRMSNormGated.forward, so defer the hard error to call time.
+    rmsnorm_fn = None
+
+if is_causal_conv1d_available():
+    from causal_conv1d import causal_conv1d_fn, causal_conv1d_update
+else:
+    causal_conv1d_update, causal_conv1d_fn = None, None
+
+if is_flash_attn_2_available():
+    from transformers.modeling_flash_attention_utils import _flash_attention_forward
+
+is_fast_path_available = all(
+    (
+        selective_state_update,
+        mamba_chunk_scan_combined,
+        mamba_split_conv1d_scan_combined,
+        causal_conv1d_fn,
+        causal_conv1d_update,
+    )
+)
+
+
+def _supports_cuda_fast_path(device: torch.device) -> bool:
+    if not is_fast_path_available or device.type != "cuda" or not torch.cuda.is_available():
+        return False
+    try:
+        major, minor = torch.cuda.get_device_capability(device)
+    except Exception:
+        return False
+    return major >= 8
+
+
+_CHECKPOINT_FOR_DOC = "nvidia/Nemotron-H-56B-Base-8K"
+_CONFIG_FOR_DOC = "CENOConfig"
+
+
+# Helper methods for segment sum computation
+
+
+def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int):
+    """
+    Padding x tensor with `pad_size` on the seq_len dim (dim=1)
+
+    Assumes that we only have tensors of either size 4 or 3
+    """
+    pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0)
+
+    return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0)
+
+
+def reshape_into_chunks(input_tensor, pad_size, chunk_size):
+    """
+    Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and
+    simultaneously splitting it into chunk sequences.
+
+    Assumes that we only have tensors of either size 4 or 3
+    """
+    # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...]
+    input_tensor = pad_tensor_by_size(input_tensor, pad_size)
+
+    if len(input_tensor.shape) == 3:
+        # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads]
+        return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2])
+    else:
+        # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> [bsz, -1, chunk_size, num_heads, head_dim or state_size]
+        return input_tensor.reshape(
+            input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3]
+        )
+
+
+def segment_sum(input_tensor):
+    """
+    More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
+    """
+    chunk_size = input_tensor.size(-1)
+    # 1. expand input tensor to have an additional dimension and repeat along that dimension
+    # [..., chunk_size] -> [..., chunk_size, chunk_size]
+    input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size)
+    # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag
+    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1)
+    input_tensor = input_tensor.masked_fill(~mask, 0)
+    # 3. compute actual cumsum
+    tensor_segsum = torch.cumsum(input_tensor, dim=-2)
+
+    # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time)
+    mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0)
+    tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf)
+    return tensor_segsum
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+    """
+    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+    """
+    if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+        dtype = hidden_states.dtype
+        hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+    return hidden_states
+
+# Copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py
+class HybridMambaAttentionDynamicCache(DynamicCache):
+    """
+    A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache
+    (which has a constant shape regardless of seq_len).
+
+    This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`
+    and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor
+    For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,
+    while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).
+    For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),
+    while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,
+    and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`.
+    """
+
+    def __init__(self, config, batch_size, dtype=torch.float16, device=None):
+        super().__init__()
+        self.dtype = dtype
+        self.hybrid_override_pattern = config.hybrid_override_pattern
+        self.has_previous_state = False  # only used by mamba
+        intermediate_size = config.expand * config.hidden_size
+        ssm_state_size = config.ssm_state_size
+        conv_kernel_size = config.conv_kernel
+        self.conv_kernel_size = conv_kernel_size
+        self.conv_states = []
+        self.ssm_states = []
+        self.transformer_layers = []
+        self.attention_layers = []
+        for i in range(config.num_hidden_layers):
+            if self.hybrid_override_pattern[i] == "M":
+                # Mamba layer
+                self.conv_states += [
+                    torch.zeros(batch_size, intermediate_size, conv_kernel_size, device=device, dtype=dtype)
+                ]
+                self.ssm_states += [
+                    torch.zeros(batch_size, intermediate_size, ssm_state_size, device=device, dtype=dtype)
+                ]
+            else:
+                # Attention or MLP layer
+                self.conv_states += [torch.tensor([[]] * batch_size, device=device)]
+                self.ssm_states += [torch.tensor([[]] * batch_size, device=device)]
+                self.transformer_layers.append(i)
+                if self.hybrid_override_pattern[i] == "*":
+                    self.attention_layers.append(i)
+
+        self.key_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]
+        self.value_cache = [torch.tensor([[]] * batch_size, device=device) for _ in range(config.num_hidden_layers)]
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Update the cache
+        if self.key_cache[layer_idx].shape[-1] == 0:
+            self.key_cache[layer_idx] = key_states
+            self.value_cache[layer_idx] = value_states
+        else:
+            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=2)
+            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=2)
+
+        return self.key_cache[layer_idx], self.value_cache[layer_idx]
+
+    def reorder_cache(self, beam_idx: torch.LongTensor):
+        """Reorders the cache for beam search, given the selected beam indices."""
+        for layer_idx in range(len(self.key_cache)):
+            device = self.key_cache[layer_idx].device
+            self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.value_cache[layer_idx].device
+            self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device))
+
+            device = self.conv_states[layer_idx].device
+            self.conv_states[layer_idx] = self.conv_states[layer_idx].index_select(0, beam_idx.to(device))
+            device = self.ssm_states[layer_idx].device
+            self.ssm_states[layer_idx] = self.ssm_states[layer_idx].index_select(0, beam_idx.to(device))
+
+    def get_seq_length(self, layer_idx: Optional[int] = 0) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        if self.attention_layers:
+            # Prefer an attention layer with a populated cache.
+            if layer_idx in self.attention_layers and self.key_cache[layer_idx].numel() > 0:
+                return self.key_cache[layer_idx].shape[-2]
+            for attn_idx in self.attention_layers:
+                if self.key_cache[attn_idx].numel() > 0:
+                    return self.key_cache[attn_idx].shape[-2]
+            return 0
+        # Fallback for models without attention layers.
+        if layer_idx in self.transformer_layers:
+            return self.key_cache[layer_idx].shape[-2] if self.key_cache[layer_idx].numel() > 0 else 0
+        if self.transformer_layers:
+            idx = self.transformer_layers[0]
+            return self.key_cache[idx].shape[-2] if self.key_cache[idx].numel() > 0 else 0
+        return 0
+
+    def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor]]:
+        raise NotImplementedError("HybridMambaAttentionDynamicCache does not have a legacy cache equivalent.")
+
+    @classmethod
+    def from_legacy_cache(cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None) -> "DynamicCache":
+        raise NotImplementedError("HybridMambaAttentionDynamicCache does not have a legacy cache equivalent.")
+
+    # Copied from modeling_mamba2.py
+    def update_conv_state(
+        self, layer_idx: int, new_conv_state: torch.Tensor, cache_init: bool = False
+    ) -> torch.Tensor:
+        if cache_init:
+            target_device = self.conv_states[layer_idx].device
+            self.conv_states[layer_idx] = new_conv_state.to(target_device)
+        else:
+            self.conv_states[layer_idx] = self.conv_states[layer_idx].roll(shifts=-1, dims=-1)
+            target_device = self.conv_states[layer_idx].device
+            self.conv_states[layer_idx][:, :, -1] = new_conv_state[:, 0, :].to(target_device)
+        return self.conv_states[layer_idx]
+
+    def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor):
+        target_device = self.ssm_states[layer_idx].device
+        self.ssm_states[layer_idx] = new_ssm_state.to(target_device)
+        return self.ssm_states[layer_idx]
+
+    def reset(self):
+        self.conv_states.zero_()
+        self.ssm_states.zero_()
+
+class MambaRMSNormGated(torch.nn.Module):
+    def __init__(self, hidden_size, group_size, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+        self.group_size = group_size
+
+    # jan28b version
+    def forward(self, hidden_states, gate=None):
+        if rmsnorm_fn is None:
+            input_dtype = hidden_states.dtype
+            x = hidden_states.float()
+            if gate is not None:
+                x = x * torch.nn.functional.silu(gate.float())
+            if self.group_size and x.shape[-1] % self.group_size == 0:
+                x = x.reshape(*x.shape[:-1], -1, self.group_size)
+                x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.variance_epsilon)
+                x = x.reshape(*x.shape[:-2], -1)
+            else:
+                x = x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.variance_epsilon)
+            return (x * self.weight.float()).to(input_dtype)
+        return rmsnorm_fn(x=hidden_states,
+                          weight=self.weight,
+                          bias=None, # No bias
+                          z=gate,
+                          eps=self.variance_epsilon,
+                          group_size=self.group_size,
+                          norm_before_gate=False
+        )
+
+class CENOMamba2Mixer(nn.Module):
+    """
+    Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
+    A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
+    ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
+    and is why Mamba is called **selective** state spaces)
+    """
+
+    def __init__(self, config: CENOConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.num_heads = config.mamba_num_heads
+        self.hidden_size = config.hidden_size
+        self.ssm_state_size = config.ssm_state_size
+        self.conv_kernel_size = config.conv_kernel
+        self.intermediate_size = config.mamba_num_heads * config.mamba_head_dim
+        self.layer_idx = layer_idx
+        self.use_conv_bias = config.use_conv_bias
+        self.activation = config.mamba_hidden_act
+        self.act = ACT2FN[config.mamba_hidden_act]
+
+        self.layer_norm_epsilon = config.layer_norm_epsilon
+        self.in_proj_norm = None
+        if getattr(config, "mamba_in_proj_layernorm", False):
+            self.in_proj_norm = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+
+        self.n_groups = config.n_groups
+        self.head_dim = config.mamba_head_dim
+        self.chunk_size = config.chunk_size
+
+        self.time_step_limit = config.time_step_limit
+        self.time_step_min = config.time_step_min
+        self.time_step_max = config.time_step_max
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            bias=config.use_conv_bias,
+            kernel_size=config.conv_kernel,
+            groups=self.conv_dim,
+            padding=config.conv_kernel - 1,
+        )
+
+        # projection of the input hidden states
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size,
+            projection_size,
+            bias=config.use_bias,
+        )
+        # selective projection used to make dt, B and C input dependant
+
+        # time step projection (discretization)
+        # instantiate once and copy inv_dt in init_weights of PretrainedModel
+        self.dt_bias = nn.Parameter(torch.ones(self.num_heads))
+
+        # S4D real initialization. These are not discretized!
+        # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+        A = torch.arange(1, self.num_heads + 1)
+        self.A_log = nn.Parameter(torch.log(A))
+        self.A_log._no_weight_decay = True
+        self.norm = MambaRMSNormGated(self.intermediate_size, eps=self.layer_norm_epsilon, group_size=self.intermediate_size // self.n_groups)
+        self.D = nn.Parameter(torch.ones(self.num_heads))
+        self.D._no_weight_decay = True
+
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
+        self.use_bias = config.use_bias
+
+        if not is_fast_path_available:
+            logger.warning_once(
+                "The fast path is not available because on of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+                " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
+                " https://github.com/Dao-AILab/causal-conv1d"
+            )
+
+    def _maybe_apply_in_proj_norm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.in_proj_norm is None:
+            return hidden_states
+        return self.in_proj_norm(hidden_states.to(dtype=self.in_proj_norm.weight.dtype))
+
+    def cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params: Optional[HybridMambaAttentionDynamicCache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+    ):
+        # 1. Gated MLP's linear projection
+        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+        hidden_states = self._maybe_apply_in_proj_norm(hidden_states)
+        projected_states = self.in_proj(hidden_states)
+
+        # Set up dimensions for reshapes later
+        batch_size, seq_len, _ = hidden_states.shape
+        groups_time_state_size = self.n_groups * self.ssm_state_size
+        d_mlp = (
+            projected_states.shape[-1]
+            - 2 * self.intermediate_size
+            - 2 * self.n_groups * self.ssm_state_size
+            - self.num_heads
+        ) // 2
+
+        # Single step calculations via cache
+        if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+            _, _, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+                [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
+            )
+
+            # 2. Convolution sequence transformation
+            hidden_states_B_C = causal_conv1d_update(
+                hidden_states_B_C,
+                cache_params.conv_states[self.layer_idx],
+                self.conv1d.weight.squeeze(1),
+                self.conv1d.bias,
+                self.activation,
+            )
+
+            hidden_states, B, C = torch.split(
+                hidden_states_B_C,
+                [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+                dim=-1,
+            )
+
+            # 3. SSM transformation
+            A = -torch.exp(self.A_log.float())  # (nheads,)
+            A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+            dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+            D = self.D[:, None, ...].expand(-1, self.head_dim)
+            B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+            C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+            hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
+            hidden_states = selective_state_update(
+                cache_params.ssm_states[self.layer_idx],
+                hidden_states_reshaped,
+                dt,
+                A,
+                B,
+                C,
+                D,
+                z=None,
+                dt_bias=dt_bias,
+                dt_softplus=True,
+            )
+            hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)
+            hidden_states = self.norm(hidden_states, gate)
+
+            # 4. Final linear projection
+            out = self.out_proj(hidden_states)[:, None, ...]
+
+        # Fused calculations or step by step if no initialized cache is found
+        else:
+            A = -torch.exp(self.A_log.float())  # (num_heads) or (intermediate_size, state_size)
+            dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+
+            # 2-4. Fused kernel for conv1d, SSM, and the final projection
+            if self.training and cache_params is None:
+                out = mamba_split_conv1d_scan_combined(
+                    projected_states,
+                    self.conv1d.weight.squeeze(1),
+                    self.conv1d.bias,
+                    self.dt_bias,
+                    A,
+                    D=self.D,
+                    chunk_size=self.chunk_size,
+                    seq_idx=seq_idx,
+                    activation=self.activation,
+                    rmsnorm_weight=self.norm.weight,
+                    rmsnorm_eps=self.norm.variance_epsilon,
+                    outproj_weight=self.out_proj.weight,
+                    outproj_bias=self.out_proj.bias,
+                    headdim=self.head_dim,
+                    ngroups=self.n_groups,
+                    norm_before_gate=False,
+                    return_final_states=False,
+                    **dt_limit_kwargs,
+                )
+
+            else:
+                _, _, gate, hidden_states_B_C, dt = projected_states.split(
+                    [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
+                )
+
+                # 2. Convolution sequence transformation
+                # Init cache
+                if cache_params is not None:
+                    hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
+                    conv_states = nn.functional.pad(
+                        hidden_states_B_C_transposed,
+                        (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0),
+                    )
+                    cache_params.update_conv_state(
+                        layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True
+                    )
+
+                if self.activation not in ["silu", "swish"]:
+                    hidden_states_B_C = self.act(
+                        self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2)
+                    )
+                else:
+                    hidden_states_B_C = causal_conv1d_fn(
+                        x=hidden_states_B_C.transpose(1, 2),
+                        weight=self.conv1d.weight.squeeze(1),
+                        bias=self.conv1d.bias,
+                        activation=self.activation,
+                    ).transpose(1, 2)
+                hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+                hidden_states, B, C = torch.split(
+                    hidden_states_B_C,
+                    [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+                    dim=-1,
+                )
+
+                # 3. SSM transformation
+                scan_output, ssm_state = mamba_chunk_scan_combined(
+                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                    dt,
+                    A,
+                    B.view(batch_size, seq_len, self.n_groups, -1),
+                    C.view(batch_size, seq_len, self.n_groups, -1),
+                    chunk_size=self.chunk_size,
+                    D=self.D,
+                    z=None,
+                    seq_idx=seq_idx,
+                    return_final_states=True,
+                    dt_bias=self.dt_bias,
+                    dt_softplus=True,
+                    **dt_limit_kwargs,
+                )
+
+                # Init cache
+                if ssm_state is not None and cache_params is not None:
+                    cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state)
+
+                scan_output = scan_output.view(batch_size, seq_len, -1)
+
+                # Multiply "gate" branch and apply extra normalization layer
+                scan_output = self.norm(scan_output, gate)
+
+                # 4. Final linear projection
+                out = self.out_proj(scan_output)
+        return out
+
+    # fmt: off
+    def torch_forward(self, input_states, cache_params: Optional[HybridMambaAttentionDynamicCache]=None, cache_position:Optional[torch.LongTensor]=None, attention_mask: Optional[torch.Tensor]=None):
+        batch_size, seq_len, _ = input_states.shape
+        dtype = input_states.dtype
+
+        # 1. Gated MLP's linear projection
+        input_states = apply_mask_to_padding_states(input_states, attention_mask)
+        input_states = self._maybe_apply_in_proj_norm(input_states)
+        projected_states = self.in_proj(input_states)
+        d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size-self.num_heads) // 2
+        _, _, gate, hidden_states_B_C, dt = projected_states.split(
+                [d_mlp, d_mlp, self.intermediate_size,  self.conv_dim, self.num_heads], dim=-1
+        )
+
+        # 2. Convolution sequence transformation
+        if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+            cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=hidden_states_B_C, cache_init=False)
+
+            # We need to guarantee that anything regarding the cache is on the same device
+            conv_states = cache_params.conv_states[self.layer_idx].to(device=self.conv1d.weight.device)
+
+            hidden_states_B_C = torch.sum(
+                conv_states * self.conv1d.weight.squeeze(1), dim=-1
+            )
+            if self.use_conv_bias:
+                hidden_states_B_C = hidden_states_B_C + self.conv1d.bias
+            hidden_states_B_C = self.act(hidden_states_B_C)
+        else:
+            # Init cache
+            if cache_params is not None:
+                hidden_states_B_C_transposed = hidden_states_B_C.transpose(1, 2)
+                conv_states = nn.functional.pad(
+                    hidden_states_B_C_transposed, (cache_params.conv_kernel_size - hidden_states_B_C_transposed.shape[-1], 0)
+                )
+                cache_params.update_conv_state(layer_idx=self.layer_idx, new_conv_state=conv_states, cache_init=True)
+
+            hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C.transpose(1, 2))[..., :seq_len].transpose(1, 2))
+
+        hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+        hidden_states, B, C = torch.split(
+            hidden_states_B_C,
+            [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
+            dim=-1
+        )
+
+        # 3. SSM transformation
+        A = -torch.exp(self.A_log.float())                            # [num_heads]
+        if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+            # We need to guarantee that anything regarding the cache is on the same device
+            cache_device = cache_params.ssm_states[self.layer_idx].device
+
+            # Note: there is no need to pad parameter matrices here, as there is just one new token
+            # for batched generation
+            dt = dt[:, 0, :][:, None, ...]
+            dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)
+            # [num_heads] -> [num_heads, head_dim]
+            dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)
+
+            dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))
+            dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+            A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+            # [bsz, num_heads, head_dim, state_size]
+            dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)
+
+            # Discretize B
+            # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->
+            # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]
+            B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]
+            B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
+            B = B.reshape(batch_size, -1, B.shape[-1])
+            # [bsz, num_heads, head_dim, state_size]
+            dB = dt[..., None] * B[..., None, :]
+
+            # Discretize x into dB
+            # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
+            hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
+            dBx = (dB * hidden_states[..., None]).to(device=cache_device)
+
+            # State calculation
+            cache_params.update_ssm_state(
+                layer_idx=self.layer_idx,
+                new_ssm_state=cache_params.ssm_states[self.layer_idx] * dA + dBx
+            )
+
+            # Subsequent output
+            # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
+            C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]
+            C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()
+            C = C.reshape(batch_size, -1, C.shape[-1])
+            # [bsz, num_heads, head_dim]
+
+            ssm_states = cache_params.ssm_states[self.layer_idx].to(device=C.device, dtype=C.dtype)  # Shape: [b, h, d, n]
+            # Reshape ssm_states to merge the first two dimensions
+            ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size)  # Shape: [b*h, d, n]
+            C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1)  # Shape: [b*h, n, 1]
+            y = torch.bmm(ssm_states_reshaped, C_reshaped)
+            y = y.view(batch_size, self.num_heads, self.head_dim)
+
+            # D skip connection
+            # [num_heads] -> [num_heads, head_dim]
+            D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
+            y = (y + hidden_states * D).to(y.dtype)
+
+            # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]
+            y = y.reshape(batch_size, -1)[:, None, ...]
+        else:
+            # begin ssd naive implementation without einsums
+            dt = nn.functional.softplus(dt + self.dt_bias)
+            dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+            hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
+            B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+            C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+            B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
+            C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
+            pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size
+
+            D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+            # Discretize x and A
+            hidden_states = hidden_states * dt[..., None]
+            A = A.to(hidden_states.dtype) * dt
+
+            # Rearrange into blocks/chunks
+            hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]
+
+            # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
+            A = A.permute(0, 3, 1, 2)
+            A_cumsum = torch.cumsum(A, dim=-1)
+
+            # 1. Compute the output for each intra-chunk (diagonal blocks)
+            # This is the analog of a causal mask
+            L = torch.exp(segment_sum(A))
+
+            # Contraction of C and B to get G (attention-weights like)
+            G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :]  # shape: (b, c, l, s, h, n)
+            G = G_intermediate.sum(dim=-1)  # shape: (b, c, l, s, h)
+
+            # Compute M, equivalent to applying attention mask to weights
+            M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
+            M = M_intermediate.sum(dim=-1)
+
+            # Compute Y_diag (apply to values)
+            Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3)
+
+            # 2. Compute the state for each intra-chunk
+            # (right term of low-rank factorization of off-diagonal blocks; B terms)
+            decay_states = torch.exp((A_cumsum[:, :, :, -1:] - A_cumsum))
+            B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]
+            states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)
+
+            # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+            # (middle term of factorization of off-diag blocks; A terms)
+            if cache_params is not None and cache_position is not None and cache_position[0] > 0:
+                previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...].to(device=states.device)
+            else:
+                previous_states = torch.zeros_like(states[:, :1])
+            states = torch.cat([previous_states, states], dim=1)
+            decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))
+            decay_chunk = decay_chunk.transpose(1, 3)
+            new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+            states, ssm_state = new_states[:, :-1], new_states[:, -1]
+
+            # 4. Compute state -> output conversion per chunk
+            # (left term of low-rank factorization of off-diagonal blocks; C terms)
+            state_decay_out = torch.exp(A_cumsum)
+            C_times_states = (C[..., None, :] * states[:, :, None, ...])
+            state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
+            Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])
+
+            # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+            y = Y_diag + Y_off
+            # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]
+            y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)
+
+            y = y + D_residual
+            # Cutting off padded chunks
+            if pad_size > 0:
+                y = y[:, :seq_len, :, :]
+            y = y.reshape(batch_size, seq_len, -1)
+
+            # Init cache
+            if ssm_state is not None and cache_params is not None:
+                cache_params.update_ssm_state(layer_idx=self.layer_idx, new_ssm_state=ssm_state)
+
+        scan_output = self.norm(y, gate)
+
+        # end ssd naive
+
+        # 4. Final linear projection
+        contextualized_states = self.out_proj(scan_output.to(dtype))  # [batch, seq_len, hidden_size]
+        return contextualized_states
+    # fmt: on
+
+    def forward(
+        self,
+        hidden_states,
+        cache_params: Optional[HybridMambaAttentionDynamicCache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+    ):
+        if self.config.use_mamba_kernels and _supports_cuda_fast_path(self.in_proj.weight.device):
+            return self.cuda_kernels_forward(hidden_states, cache_params, cache_position, attention_mask, seq_idx=seq_idx)
+        dtype = hidden_states.dtype
+        if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+            # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+            hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+        return self.torch_forward(hidden_states, cache_params, cache_position, attention_mask)
+
+
+class CENORMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        CENORMSNorm is equivalent to T5LayerNorm and LlamaRMSNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        # Weights are in float32
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+class CENOBlock(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.residual_in_fp32 = config.residual_in_fp32
+        # Optional second norm to mirror Megatron pre_mlp_layernorm for non-Mamba blocks
+        self.pre_mlp_norm = None
+
+        # M: Mamba, *: Attention, -: MLP, E: MoE
+        self.block_type = config.layers_block_type[layer_idx]
+        if self.block_type == "mamba" and getattr(config, "mamba_in_proj_layernorm", False):
+            self.norm = None
+        else:
+            self.norm = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        if self.block_type == "mamba":
+            self.mixer = CENOMamba2Mixer(config, layer_idx=layer_idx)
+        elif self.block_type == "attention":
+            self.mixer = CENO_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx=layer_idx)
+            self.pre_mlp_norm = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        elif self.block_type == "mlp":
+            self.mixer = CENOMLP(
+                config,
+                layer_idx=layer_idx,
+                use_fc1_norm=True,
+            )
+            self.pre_mlp_norm = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        elif self.block_type == "moe":
+            self.mixer = CENOMoE(config, layer_idx=layer_idx)
+            self.pre_mlp_norm = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        else:
+            raise ValueError(f"Invalid layer pattern {config.hybrid_override_pattern[layer_idx]}")
+
+    def forward(
+        self,
+        hidden_states,
+        cache_params: Optional[HybridMambaAttentionDynamicCache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+    ):
+        with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)):
+            # * Use torch.cuda.stream() to avoid NaN issues when using multiple GPUs
+            
+            if self.block_type == "mamba":
+                # Mamba block: residual = x, output = x + mamba(norm(x))
+                residual = hidden_states
+                if self.norm is not None:
+                    hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+                if self.residual_in_fp32:
+                    residual = residual.to(torch.float32)
+                hidden_states = self.mixer(
+                    hidden_states,
+                    cache_params=cache_params,
+                    cache_position=cache_position,
+                    attention_mask=attention_mask,
+                    seq_idx=seq_idx,
+                )
+                hidden_states = residual + hidden_states
+                
+            elif self.block_type == "attention":
+                # Attention block in Megatron (TransformerLayer with mlp=IdentityOp):
+                # 1. input_ln_out = input_layernorm(x)
+                # 2. attn_out = self_attention(input_ln_out)
+                # 3. hidden = residual + attn_out (self_attn_bda applies residual)
+                # 4. residual = hidden
+                # 5. pre_mlp_out = pre_mlp_layernorm(hidden)
+                # 6. mlp_out = IdentityOp(pre_mlp_out) = pre_mlp_out
+                # 7. final = IdentityFuncOp(...)(mlp_out, residual, ...) = mlp_out (no residual!)
+                # So final output = pre_mlp_layernorm(x + attention(input_layernorm(x)))
+                #
+                # self.norm = input_layernorm
+                # self.pre_mlp_norm = pre_mlp_layernorm
+                residual = hidden_states
+                input_ln_out = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+                if self.residual_in_fp32:
+                    residual = residual.to(torch.float32)
+                attn_out, _, _ = self.mixer(
+                    input_ln_out,
+                    attention_mask=attention_mask,
+                    past_key_value=cache_params,
+                    use_cache=cache_params is not None,
+                    cache_position=cache_position,
+                )
+                hidden_after_attn = residual + attn_out
+                # Apply pre_mlp_layernorm and that's the final output (no second residual)
+                hidden_states = self.pre_mlp_norm(hidden_after_attn.to(dtype=self.pre_mlp_norm.weight.dtype))
+                
+            elif self.block_type == "mlp":
+                # MLP block in Megatron: residual = input_ln_out, output = input_ln_out + mlp(pre_mlp_ln(input_ln_out))
+                input_ln_out = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+                residual = input_ln_out
+                if self.residual_in_fp32:
+                    residual = residual.to(torch.float32)
+                pre_mlp_out = self.pre_mlp_norm(input_ln_out.to(dtype=self.pre_mlp_norm.weight.dtype))
+                mlp_out = self.mixer(pre_mlp_out)
+                hidden_states = residual + mlp_out
+                
+            elif self.block_type == "moe":
+                # MoE block in Megatron: uses TransformerLayer structure
+                # When self_attention=IdentityOp:
+                #   1. input_ln_out = input_layernorm(x)
+                #   2. hidden = input_ln_out (no residual add)
+                #   3. residual = hidden = input_ln_out
+                #   4. pre_mlp_out = pre_mlp_layernorm(hidden)
+                #   5. moe_out = MoE(pre_mlp_out)
+                #   6. output = residual + moe_out = input_ln(x) + MoE(pre_mlp_ln(input_ln(x)))
+                #
+                # self.norm = input_layernorm
+                # self.pre_mlp_norm = pre_mlp_layernorm
+                input_ln_out = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+                # In Megatron, residual is set to input_ln_out (not original x)
+                # because IdentityFuncOp for self_attn_bda skips residual addition
+                residual = input_ln_out
+                if self.residual_in_fp32:
+                    residual = residual.to(torch.float32)
+                # Apply pre_mlp_layernorm before MoE
+                pre_mlp_out = self.pre_mlp_norm(input_ln_out.to(dtype=self.pre_mlp_norm.weight.dtype))
+                moe_out = self.mixer(pre_mlp_out)
+                hidden_states = residual + moe_out
+                
+            else:
+                raise ValueError(f"Invalid block_type: {self.block_type}")
+
+            return hidden_states
+
+
+# Updated to support SwiGLU activation and per-token scaling (for MoE routing weights)
+class CENOMLP(nn.Module):
+    def __init__(self, config, layer_idx: Optional[int] = None, use_fc1_norm: bool = True):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.fc1_norm = None
+        if use_fc1_norm:
+            self.fc1_norm = CENORMSNorm(self.hidden_size, eps=config.layer_norm_epsilon)
+        
+        # Check if using SwiGLU (gated activation)
+        self.use_swiglu = getattr(config, 'mlp_use_swiglu', False)
+        
+        if self.use_swiglu:
+            # For SwiGLU, we need both gate and up projections
+            # Note: In Megatron, these are combined in fc1 with 2*intermediate_size output
+            # Here we split them into separate projections for clarity
+            self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+            self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+            self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+            self.act_fn = ACT2FN[config.mlp_hidden_act]
+        else:
+            # Standard MLP
+            self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+            self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+            self.act_fn = ACT2FN[config.mlp_hidden_act]
+
+    def _maybe_apply_fc1_norm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.fc1_norm is None:
+            return hidden_states
+        return self.fc1_norm(hidden_states.to(dtype=self.fc1_norm.weight.dtype))
+
+    def forward(self, x, per_token_scale=None):
+        """
+        Forward pass with optional per-token scaling (for MoE routing weights).
+        
+        In Megatron MoE, routing weights are applied AFTER activation but BEFORE fc2/down_proj.
+        This is different from applying after the full MLP, and is critical for numerical equivalence.
+        
+        Args:
+            x: input hidden states
+            per_token_scale: optional [num_tokens, 1] tensor of routing weights
+        """
+        x = self._maybe_apply_fc1_norm(x)
+
+        if self.use_swiglu:
+            # SwiGLU: act_fn(gate_proj(x)) * up_proj(x)
+            intermediate = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
+        else:
+            # Standard MLP: act_fn(up_proj(x))
+            up_out = self.up_proj(x)
+            intermediate = self.act_fn(up_out)
+        
+        # Apply per-token scaling AFTER activation, BEFORE down_proj (Megatron behavior)
+        if per_token_scale is not None:
+            intermediate = intermediate * per_token_scale
+        
+        output = self.down_proj(intermediate)
+        return output
+
+
+# Standard HuggingFace-style MoE implementation for CENO
+class CENOMoE(nn.Module):
+    def __init__(self, config, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.num_experts = getattr(config, 'num_experts', 8)
+        self.top_k = getattr(config, 'moe_top_k', 2)
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        
+        # Router/gate
+        self.gate = nn.Linear(self.hidden_size, self.num_experts, bias=False)
+        
+        # Expert networks
+        self.experts = nn.ModuleList([
+            CENOMLP(config, layer_idx, use_fc1_norm=False) for _ in range(self.num_experts)
+        ])
+        
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+        
+        # Flatten for easier processing
+        hidden_states_flat = hidden_states.view(-1, hidden_dim)
+        
+        # Router logits
+        # Megatron default behavior (when --moe-router-pre-softmax is NOT set):
+        #   - select top-k on logits
+        #   - compute softmax *within* the selected top-k logits (post-topk softmax)
+        # If moe_router_pre_softmax=True, we instead do global softmax then take top-k (legacy HF behavior).
+        #
+        # NOTE: For strict HF-vs-Megatron validation, it's sometimes necessary to compute router logits
+        # in fp32 (matmul in fp32) to avoid bf16 GEMM rounding drift that can be amplified by routing
+        # and the squared-relu expert MLP. This is controlled by config.moe_router_dtype.
+        router_dtype = str(getattr(self.config, "moe_router_dtype", "")).lower()
+        if router_dtype == "fp32":
+            # bias=False; use explicit fp32 matmul for reproducibility across backends.
+            router_logits = torch.nn.functional.linear(
+                hidden_states_flat.to(torch.float32),
+                self.gate.weight.to(torch.float32),
+                bias=None,
+            )  # [T, num_experts] fp32
+        else:
+            # default: run the module in the model dtype, then cast logits to fp32 for softmax/topk
+            router_logits = self.gate(hidden_states_flat).to(torch.float32)  # [T, num_experts]
+        pre_softmax = bool(getattr(self.config, "moe_router_pre_softmax", False))
+
+        if pre_softmax:
+            # Pre-softmax routing: softmax over all experts, then take top-k.
+            probs = torch.nn.functional.softmax(router_logits, dim=-1)
+            routing_weights, selected_experts = torch.topk(probs, self.top_k, dim=-1)
+            routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20)
+        else:
+            # Post-topk softmax routing: take top-k logits first, then softmax over top-k.
+            topk_logits, selected_experts = torch.topk(router_logits, self.top_k, dim=-1)
+            routing_weights = torch.nn.functional.softmax(topk_logits, dim=-1)
+        
+        # Initialize output
+        final_hidden_states = torch.zeros_like(hidden_states_flat)
+        
+        # Route tokens to experts
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+        
+        # Process each expert
+        for expert_idx in range(self.num_experts):
+            expert_layer = self.experts[expert_idx]
+            
+            # Get tokens for this expert from all top-k positions
+            idx, top_x = torch.where(expert_mask[expert_idx])
+            
+            if top_x.shape[0] == 0:
+                continue
+                
+            # Index the correct hidden states
+            current_state = hidden_states_flat[top_x]  # [num_selected, hidden_dim]
+            
+            # Get the routing weights for these tokens
+            # Shape: [num_selected, 1] for per-token scaling inside the MLP
+            current_weights = routing_weights[top_x, idx].unsqueeze(-1).to(current_state.dtype)
+            
+            # Pass routing weights to expert MLP - they are applied AFTER activation, BEFORE down_proj
+            # This matches Megatron's MoE behavior for numerical equivalence
+            current_hidden_states = expert_layer(
+                current_state, 
+                per_token_scale=current_weights
+            )
+            
+            # Accumulate outputs
+            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+            
+        # Reshape back
+        return final_hidden_states.view(batch_size, seq_len, hidden_dim)
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class CENOAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: CENOConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+                "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        if config.attention_head_dim is not None:
+            self.head_dim = config.attention_head_dim
+        else:
+            self.head_dim = config.hidden_size // config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.is_causal = True
+
+        self.qkv_norm = None
+        if getattr(config, "qkv_layernorm", False):
+            self.qkv_norm = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
+        self.o_proj = nn.Linear(self.head_dim * self.num_heads, self.hidden_size, bias=config.attention_bias)
+
+    def _maybe_apply_qkv_norm(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self.qkv_norm is None:
+            return hidden_states
+        return self.qkv_norm(hidden_states.to(dtype=self.qkv_norm.weight.dtype))
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        # position_embeddings: Tuple[torch.Tensor, torch.Tensor], #TODO
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[HybridMambaAttentionDynamicCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        bsz, q_len, _ = hidden_states.size()
+
+        hidden_states = self._maybe_apply_qkv_norm(hidden_states)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if past_key_value is not None:
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        causal_mask = attention_mask
+        if attention_mask is not None:  # no matter the length, we just slice it
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        is_causal = True if causal_mask is None and q_len > 1 else False
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=causal_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        #attn_output = attn_output.view(bsz, q_len, self.hidden_size)
+        attn_output = attn_output.view(bsz, q_len, self.num_heads * self.head_dim)
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+# Adapted from transformers.models.mistral.modeling_mistral.MistralFlashAttention2 with Mistral->Jamba
+#class JambaFlashAttention2(JambaAttention):
+class CENOFlashAttention2(CENOAttention):
+    """
+    Jamba flash attention module. This module inherits from `JambaAttention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+        # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+        # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+        self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[HybridMambaAttentionDynamicCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        bsz, q_len, _ = hidden_states.size()
+
+        hidden_states = self._maybe_apply_qkv_norm(hidden_states)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        # Flash attention requires the input to have the shape
+        # batch_size x seq_length x head_dim x hidden_dim
+        # therefore we just need to keep the original shape
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if past_key_value is not None:
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx)
+
+        # repeat k/v heads if n_kv_heads < n_heads
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+        dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
+        # therefore the input hidden states gets silently casted in float32. Hence, we need
+        # cast them back in float16 just to be sure everything works as expected.
+        input_dtype = query_states.dtype
+        if input_dtype == torch.float32:
+            if torch.is_autocast_enabled():
+                target_dtype = torch.get_autocast_gpu_dtype()
+            # Handle the case where the model is quantized
+            elif hasattr(self.config, "_pre_quantization_dtype"):
+                target_dtype = self.config._pre_quantization_dtype
+            else:
+                target_dtype = self.q_proj.weight.dtype
+
+            logger.warning_once(
+                f"The input hidden states seems to be silently casted in float32, this might be related to"
+                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+                f" {target_dtype}."
+            )
+
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            value_states = value_states.to(target_dtype)
+
+        # Reashape to the expected shape for Flash Attention
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        attn_output = _flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            q_len,
+            dropout=dropout_rate,
+            sliding_window=getattr(self.config, "sliding_window", None),
+            is_causal=self.is_causal,
+            use_top_left_mask=self._flash_attn_uses_top_left_mask,
+        )
+
+        #attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.head_dim).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+# Adapted from transformers.models.mistral.modeling_mistral.MistralSdpaAttention with Mistral->Jamba
+#class JambaSdpaAttention(JambaAttention):
+class CENOSdpaAttention(CENOAttention):
+    """
+    Jamba attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
+    `JambaAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
+    SDPA API.
+    """
+
+    # Adapted from CENOAttention.forward
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[HybridMambaAttentionDynamicCache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if output_attentions:
+            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
+            logger.warning_once(
+                "CENOModel is using CENOSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
+                'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
+            )
+            return super().forward(
+                hidden_states=hidden_states,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_value=past_key_value,
+                output_attentions=output_attentions,
+                use_cache=use_cache,
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        hidden_states = self._maybe_apply_qkv_norm(hidden_states)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+        if past_key_value is not None:
+            key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx)
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        causal_mask = attention_mask
+        if attention_mask is not None:
+            causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
+
+        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
+        # Reference: https://github.com/pytorch/pytorch/issues/112577.
+        if query_states.device.type == "cuda" and attention_mask is not None:
+            query_states = query_states.contiguous()
+            key_states = key_states.contiguous()
+            value_states = value_states.contiguous()
+
+        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an inline conditional assignment
+        # in SDPA to support both torch.compile's dynamic shapes and full graph options. An inline conditional prevents dynamic shapes from compiling.
+        # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
+        is_causal = True if self.is_causal and causal_mask is None and q_len > 1 else False
+
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=causal_mask,
+            dropout_p=self.attention_dropout if self.training else 0.0,
+            is_causal=is_causal,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(bsz, q_len, self.hidden_size)
+
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, None, past_key_value
+
+
+CENO_ATTENTION_CLASSES = {
+    "eager": CENOAttention,
+    "flash_attention_2": CENOFlashAttention2,
+    "sdpa": CENOSdpaAttention,
+}
+
+# Copied from transformers.models.mamba.modeling_mamba2.Mamba2PreTrainedModel
+class CENOPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = CENOConfig
+    base_model_prefix = "backbone"
+    _no_split_modules = ["CENOBlock"]
+    supports_gradient_checkpointing = True
+    _is_stateful = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        if isinstance(module, CENOMamba2Mixer):
+            module.A_log._no_weight_decay = True
+            module.D._no_weight_decay = True
+
+            dt = torch.exp(
+                torch.rand(self.config.mamba_num_heads)
+                * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min))
+                + math.log(self.config.time_step_min)
+            ).clamp(min=self.config.time_step_floor)
+
+            # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
+            inv_dt = dt + torch.log(-torch.expm1(-dt))
+            with torch.no_grad():
+                module.dt_bias.copy_(inv_dt)
+            module.dt_bias._no_reinit = True
+
+        if isinstance(module, nn.Linear):
+            if module.bias is not None:
+                if not getattr(module.bias, "_no_reinit", False):
+                    nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=self.config.initializer_range)
+
+        # TODO: Check
+        if self.config.rescale_prenorm_residual:
+            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+            #
+            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+            for name, p in module.named_parameters():
+                if name in ["out_proj.weight"]:
+                    # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                    # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                    # We need to reinit p since this code could be called multiple times
+                    # Having just p *= scale would repeatedly scale it down
+                    nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                    with torch.no_grad():
+                        p /= math.sqrt(self.config.num_hidden_layers)
+
+
+@dataclass
+# Copied from transformers.models.mamba.modeling_mamba2.Mamba2Output with MAMBA2->CENO,Mamba2->CENO
+class CENOOutput(ModelOutput):
+    """
+    Class for the CENO model outputs.
+
+    Args:
+        last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        cache_params (`HybridMambaAttentionDynamicCache`):
+            The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
+            avoid providing the old `input_ids`.
+
+            Includes both the State space model state matrices after the selective scan, and the Convolutional states
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+    """
+
+    last_hidden_state: Optional[torch.FloatTensor] = None
+    cache_params: Optional[HybridMambaAttentionDynamicCache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+@dataclass
+# Copied from transformers.models.mamba2.modeling_mamba2.MambaCausalLMOutput with Mamba2->CENO
+class CENOCausalLMOutput(ModelOutput):
+    """
+    Base class for causal language model (or autoregressive) outputs.
+
+    Args:
+        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+            Language modeling loss (for next-token prediction).
+        logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        cache_params (`HybridMambaAttentionDynamicCache`):
+            The state of the model at the last time step. Can be used in a forward method with the next `input_ids` to
+            avoid providing the old `input_ids`.
+
+            Includes both the State space model state matrices after the selective scan, and the Convolutional states
+        hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+    """
+
+    loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    cache_params: Optional[HybridMambaAttentionDynamicCache] = None
+    past_key_values: Optional[HybridMambaAttentionDynamicCache] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
+CENO_START_DOCSTRING = r"""
+
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`CENOConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+CENO_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`, *optional*):
+            Indices of input sequence tokens in the vocabulary.
+
+            If `cache_params.seqlen_offset>0`, only `input_ids` that do not have their past calculated should be passed as
+            `input_ids`.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        position_ids (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings.
+        cache_params (`HybridMambaAttentionDynamicCache`, *optional*):
+            If passed along, the model uses the previous state in all the blocks (which will give the output for the
+            `input_ids` provided as if the model add `state_input_ids + input_ids` as context).
+        use_cache (`bool`, *optional*):
+            If set to `True`, the `cache_params` is returned and can be used to quickly generate the next logits.
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        cache_position (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            The position of the current input in the cache. This is used to ensure that the cache is correctly updated.
+            If `cache_params` is passed, `cache_position` should also be passed.
+        attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+        seq_idx (`torch.IntTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Optional per-token sequence ids used to isolate different sub-sequences. Tokens sharing the same id are
+            allowed to attend to each other; different ids are blocked. If omitted, no intra-sequence isolation is used.
+"""
+
+
+@add_start_docstrings(
+    "The bare CENO Model transformer outputting raw hidden-states without any specific head on top.",
+    CENO_START_DOCSTRING,
+)
+class CENOModel(CENOPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([CENOBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)])
+
+        self.gradient_checkpointing = False
+        self.norm_f = CENORMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+        # Initialize weights and apply final processing
+        self._register_load_state_dict_pre_hook(self.load_hook)
+        self.post_init()
+
+    def load_hook(self, state_dict, prefix, *args):
+        for k in state_dict:
+            if "embedding." in k:
+                state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k)
+                break
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embeddings = new_embeddings
+
+    @add_start_docstrings_to_model_forward(CENO_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=CENOOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        cache_params: Optional[HybridMambaAttentionDynamicCache] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CENOOutput]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        # use_cache = use_cache if use_cache is not None else self.config.use_cache
+        use_cache = use_cache if use_cache is not None else (self.config.use_cache if not self.training else False)
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):  # ^ is python for xor
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embeddings(input_ids)
+
+        if self.gradient_checkpointing and self.training and use_cache:
+            logger.warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
+            )
+            use_cache = False
+
+        # From zamba_modeling.py
+        if use_cache and cache_params is None:
+            logger.warning_once(
+                "CENO requires an initialized `CENOHybridDynamicCache` to return a cache. None was "
+                "provided, so no cache will be returned."
+            )
+
+        hidden_states = inputs_embeds
+
+        if cache_position is None:
+            cache_position = torch.arange(hidden_states.shape[1], device=hidden_states.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = self._update_causal_mask(attention_mask, inputs_embeds, cache_position)
+        mamba_mask = self._update_mamba_mask(attention_mask, cache_position)
+        if seq_idx is not None:
+            seq_idx = seq_idx.to(device=hidden_states.device, dtype=torch.int32)
+            if seq_idx.shape[1] > hidden_states.shape[1]:
+                # seq_idx may be longer (extended for generate decode steps);
+                # truncate to current hidden_states length (prefill or decode step)
+                seq_idx = seq_idx[:, :hidden_states.shape[1]]
+            if seq_idx.shape[1] != hidden_states.shape[1]:
+                raise ValueError(
+                    f"seq_idx shape {seq_idx.shape[:2]} must match hidden_states shape {hidden_states.shape[:2]}"
+                )
+        intra_encoding_mask = getattr(self.config, "intra_encoding_mask", None)
+        if intra_encoding_mask is None:
+            intra_encoding_mask = [False for _ in range(len(self.layers))]
+
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        # Until HERE
+
+        for layer_idx, mixer_block in enumerate(self.layers):
+            use_seq_idx = False
+            if seq_idx is not None:
+                use_seq_idx = bool(intra_encoding_mask[layer_idx])
+
+            current_layer_seq_idx = seq_idx if use_seq_idx else None
+            # Depending on the layer type we opt for 2D base attention mask (Mamba) or 4D causal mask (Attention)
+            if mixer_block.block_type == "mamba":
+                layer_mask = mamba_mask
+            elif mixer_block.block_type == "attention":
+                if current_layer_seq_idx is not None:
+                    layer_mask = self._merge_seq_idx_attention_mask(
+                        current_layer_seq_idx, causal_mask, hidden_states.shape[1], hidden_states.dtype
+                    )
+                else:
+                    layer_mask = causal_mask
+            elif mixer_block.block_type == "mlp":
+                layer_mask = None
+            elif mixer_block.block_type == "moe":
+                layer_mask = None
+            else:
+                raise ValueError(f"Invalid block_type: {mixer_block.block_type}")
+
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states = self._gradient_checkpointing_func(
+                    mixer_block.__call__, hidden_states, cache_params, cache_position, layer_mask, current_layer_seq_idx
+                )
+            else:
+                hidden_states = mixer_block(
+                    hidden_states,
+                    cache_params=cache_params,
+                    cache_position=cache_position,
+                    attention_mask=layer_mask,
+                    seq_idx=current_layer_seq_idx,
+                )
+
+            # TODO: Store attentions
+            # if output_attentions:
+            #     if layer_outputs[1] is not None:
+            #         # append attentions only of attention layers. Mamba layers return `None` as the attention weights
+            #         all_self_attns += (layer_outputs[1],)
+
+            # TODO (Check): should it happen before the forward pass?
+            # if output_hidden_states:
+            #     all_hidden_states = all_hidden_states + (hidden_states,)
+
+        hidden_states = self.norm_f(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, cache_params, all_hidden_states] if v is not None)
+
+        return CENOOutput(
+            last_hidden_state=hidden_states,
+            cache_params=cache_params if use_cache else None,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+    # Copied from transformers.models.jamba.modeling_jamba.JambaModel._update_causal_mask
+    def _update_causal_mask(self, attention_mask, input_tensor, cache_position):
+        if self.config._attn_implementation == "flash_attention_2":
+            if attention_mask is not None and 0.0 in attention_mask:
+                return attention_mask
+            return None
+
+        dtype, device = input_tensor.dtype, input_tensor.device
+        min_dtype = torch.finfo(dtype).min
+        sequence_length = input_tensor.shape[1]
+        target_length = cache_position[-1] + 1
+
+        causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
+        if sequence_length != 1:
+            causal_mask = torch.triu(causal_mask, diagonal=1)
+        causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
+        causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
+        if attention_mask is not None:
+            causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit
+            if attention_mask.dim() == 2:
+                mask_length = attention_mask.shape[-1]
+                padding_mask = causal_mask[..., :mask_length].eq(0.0) * attention_mask[:, None, None, :].eq(0.0)
+                causal_mask[..., :mask_length] = causal_mask[..., :mask_length].masked_fill(padding_mask, min_dtype)
+
+        if (
+            self.config._attn_implementation == "sdpa"
+            and attention_mask is not None
+            and attention_mask.device.type == "cuda"
+        ):
+            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
+            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
+            # Details: https://github.com/pytorch/pytorch/issues/110213
+            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
+
+        return causal_mask
+
+    def _update_mamba_mask(self, attention_mask, cache_position):
+        """
+        No need for zeroing states when
+            1. Cached forward
+            2. Attending to all inputs
+        """
+        mamba_mask = attention_mask
+        if cache_position[0] > 0 or (attention_mask is not None and torch.all(attention_mask == 1)):
+            mamba_mask = None
+        return mamba_mask
+
+    def _merge_seq_idx_attention_mask(
+        self,
+        seq_idx: torch.Tensor,
+        base_mask: Optional[torch.Tensor],
+        target_len: int,
+        dtype: torch.dtype,
+    ):
+        """
+        Build an intra-sequence visibility mask from seq_idx and merge with an existing causal/padding mask.
+        """
+        seq_idx = seq_idx[:, :target_len]
+        same_segment = seq_idx[:, None, :, None] == seq_idx[:, None, None, :]
+        seg_mask = torch.where(
+            same_segment,
+            torch.zeros_like(same_segment, dtype=dtype),
+            torch.full_like(same_segment, torch.finfo(dtype).min, dtype=dtype),
+        )
+        if base_mask is None:
+            return seg_mask
+        return base_mask[..., :target_len] + seg_mask
+
+
+@add_start_docstrings(
+    """
+    The CENO Model transformer with a language modeling head on top (linear layer with weights not tied to the input
+    embeddings).
+    """,
+    CENO_START_DOCSTRING,
+)
+class CENOForCausalLM(CENOPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.backbone = CENOModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.backbone.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        return self.backbone.set_input_embeddings(new_embeddings)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def get_decoder(self):
+        return self.backbone
+
+    def set_decoder(self, decoder):
+        self.backbone = decoder
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        seq_idx=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        # Copy from https://github.com/huggingface/transformers/blob/main/src/transformers/models/jamba/modeling_jamba.py
+        # Overwitten -- uses `cache_params` as opposed to `past_key_values`
+        empty_past_kv = past_key_values is None
+
+        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
+        # Exception 1: when passing input_embeds, input_ids may be missing entries
+        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
+        # Exception 3: with synced GPUs cache_position may go out of bounds, but we only want dummy token in that case.
+        #              (we can't check exception 3 while compiling)
+        if not empty_past_kv:
+            if (
+                inputs_embeds is not None  # Exception 1
+                or cache_position[-1] >= input_ids.shape[1]  # Exception 3
+            ):
+                input_ids = input_ids[:, -cache_position.shape[0] :]
+                if seq_idx is not None:
+                    seq_idx = seq_idx[:, -cache_position.shape[0] :]
+            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
+                input_ids = input_ids[:, cache_position]
+                if seq_idx is not None:
+                    seq_idx = seq_idx[:, cache_position]
+        else:
+            past_key_values = HybridMambaAttentionDynamicCache(
+                self.config, input_ids.shape[0], self.dtype, device=self.device
+            )
+
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if not empty_past_kv:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and empty_past_kv:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids.contiguous()}  # `contiguous()` needed for compilation use cases
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+                "seq_idx": seq_idx,
+                "logits_to_keep": self.config.num_logits_to_keep,
+                "cache_position": cache_position,
+            }
+        )
+        return model_inputs
+
+    @add_start_docstrings_to_model_forward(CENO_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=CENOCausalLMOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        cache_params: Optional[HybridMambaAttentionDynamicCache] = None,
+        past_key_values: Optional[HybridMambaAttentionDynamicCache] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        seq_idx: Optional[torch.Tensor] = None,
+        **kwargs,  # for now we need this for generation
+    ) -> Union[Tuple, CENOCausalLMOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if cache_params is None and past_key_values is not None:
+            cache_params = past_key_values
+
+        ceno_outputs = self.backbone(
+            input_ids,
+            cache_params=cache_params,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            attention_mask=attention_mask,
+            seq_idx=seq_idx,
+        )
+        hidden_states = ceno_outputs[0]
+
+        # TODO: Check zamba_modeling.py: https://github.com/huggingface/transformers/blob/d7188ba600e36d3fd191b12e19f1b3bb81a8404f/src/transformers/models/zamba/modeling_zamba.py#L1284C1-L1286C2
+        #logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float()
+        logits = self.lm_head(hidden_states.to(self.lm_head.weight.dtype)).float()
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + ceno_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return CENOCausalLMOutput(
+            loss=loss,
+            logits=logits,
+            cache_params=ceno_outputs.cache_params,
+            past_key_values=ceno_outputs.cache_params,
+            hidden_states=ceno_outputs.hidden_states,
+            attentions=ceno_outputs.attentions,
+        )

special_tokens_map.json

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

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
+    ]
+  }
+}

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
+}