feat: upload RotorQuant-MLX-2bit (text-tower quantized; encoders BF16)

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,126 @@
+NVIDIA Open Model License Agreement
+This NVIDIA Open Model License Agreement (the “Agreement”) is a legal agreement between the Legal Entity You represent, or if no
+entity is identified, You and NVIDIA Corporation and its Affiliates (“NVIDIA”) and governs Your use of the Models that NVIDIA
+provides to You under this Agreement. NVIDIA and You are each a “party” and collectively the “parties.”
+NVIDIA models released under this Agreement are intended to be used permissively and enable the further development of AI
+technologies. Subject to the terms of this Agreement, NVIDIA confirms that:
+•
+
+Models are commercially useable.
+
+•
+
+You are free to create and distribute Derivative Models.
+
+•
+
+NVIDIA does not claim ownership to any outputs generated using the Models or Model Derivatives.
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of the Model or Derivative Model, or
+otherwise accepting the terms of this Agreement, you agree to be bound by this Agreement.
+1.
+
+Definitions. The following definitions apply to this Agreement:
+
+1.1.
+
+“Derivative Model” means all (a) modifications to the Model, (b) works based on the Model, and (c) any other derivative
+works of the Model. An output is not a Derivative Model.
+
+1.2.
+
+“Legal Entity” means 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 (a) the power, direct or indirect, to cause the
+direction or management of such entity, whether by contract or otherwise, or (b) ownership of fifty percent (50%) or more
+of the outstanding shares, or (c) beneficial ownership of such entity.
+
+1.3.
+
+“Model” means the machine learning model, software, checkpoints, learnt weights, algorithms, parameters, configuration
+files and documentation shared under this Agreement.
+
+1.4.
+
+“You” or “Your” means an individual or Legal Entity exercising permissions granted by this Agreement.
+
+2.
+
+Conditions for Use, License Grant, AI Ethics and IP Ownership.
+
+2.1.
+Conditions for Use. The Model and any Derivative Model are subject to additional terms as described in Section 2 and
+Section 3 of this Agreement and govern Your use. If You institute copyright or patent litigation against any entity (including a crossclaim or counterclaim in a lawsuit) alleging that the Model or a Derivative Model constitutes direct or contributory copyright or
+patent infringement, then any licenses granted to You under this Agreement for that Model or Derivative Model will terminate as of
+the date such litigation is filed. NVIDIA may update this Agreement to comply with legal and regulatory requirements at any time
+and You agree to either comply with any updated license or cease Your copying, use, and distribution of the Model and any
+Derivative Model.
+2.2.
+License Grant. The rights granted herein are explicitly conditioned on Your full compliance with the terms of this
+Agreement. Subject to the terms and conditions of this Agreement, NVIDIA hereby grants to You a perpetual, worldwide, nonexclusive, no-charge, royalty-free, revocable (as stated in Section 2.1) license to publicly perform, publicly display, reproduce, use,
+create derivative works of, make, have made, sell, offer for sale, distribute (through multiple tiers of distribution) and import the
+Model.
+2.3.
+AI Ethics. NVIDIA is committed to safety, trust and transparency in AI development. NVIDIA encourages You to (a) ensure
+that the product or service You develop, use, offer as a service or distributes meets the legal and ethical requirements of the
+relevant industry or use case, (b) take reasonable measures to address unintended bias and to mitigate harm to others, including
+underrepresented or vulnerable groups, and (c) inform users of the nature and limitations of the product or service. NVIDIA
+expressly prohibits the use of its products or services for any purpose in violation of applicable law or regulation, including but not
+limited to (a) illegal surveillance, (b) illegal collection or processing of biometric information without the consent of the subject
+where required under applicable law, or (c) illegal harassment, abuse, threatening or bullying of individuals or groups of individuals
+or intentionally misleading or deceiving others.
+2.4.
+NVIDIA owns the Model and any Model Derivatives created by NVIDIA. Subject to NVIDIA’s underlying ownership rights in
+the Model or its Model Derivatives, You are and will be the owner of Your Model Derivatives. NVIDIA claims no ownership rights in
+outputs. You are responsible for outputs and their subsequent uses. Except as expressly granted in this Agreement, (a) NVIDIA
+reserves all rights, interests and remedies in connection with the Model and (b) no other license or right is granted to you by
+implication, estoppel or otherwise.
+3.
+Redistribution. You may reproduce and distribute copies of the Model or Derivative Models thereof in any medium, with or
+without modifications, provided that You meet the following conditions:
+
+3.1.
+If you distribute the Model, You must give any other recipients of the Model a copy of this Agreement and include the
+following attribution notice within a “Notice” text file with such copies: “Licensed by NVIDIA Corporation under the NVIDIA Open
+Model License”; and
+3.2.
+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 Models as a whole, provided
+Your use, reproduction, and distribution of the Model otherwise complies with the conditions stated in this Agreement.
+4.
+Trademarks. This Agreement does not grant permission to use the trade names, trademarks, service marks, or product
+names of NVIDIA, except as required for reasonable and customary use in describing the origin of the Model and reproducing the
+content of the “Notice” text file.
+5.
+Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, NVIDIA provides the Model 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 Model, Derivative Models and outputs and
+assume any risks associated with Your exercise of permissions under this Agreement.
+6.
+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, will NVIDIA be
+liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a
+result of this Agreement or out of the use or inability to use the Model, Derivative Models or outputs (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 NVIDIA has been advised of the possibility of such damages.
+7.
+Indemnity. You will indemnify and hold harmless NVIDIA from and against any claim by any third party arising out of or
+related to your use or distribution of the Model, Model Derivatives or outputs.
+8.
+You.
+
+Feedback. NVIDIA appreciates your feedback, and You agree that NVIDIA may use it without restriction or compensation to
+
+9.
+Governing Law. This Agreement will be governed in all respects by the laws of the United States and the laws of the State
+of Delaware, without regard to conflict of laws principles or the United Nations Convention on Contracts for the International Sale of
+Goods. The state and federal courts residing in Santa Clara County, California will have exclusive jurisdiction over any dispute or
+claim arising out of or related to this Agreement, and the parties irrevocably consent to personal jurisdiction and venue in those
+courts; except that, either party may apply for injunctive remedies or an equivalent type of urgent legal relief in any jurisdiction.
+10.
+Trade and Compliance. You agree to comply with all applicable export, import, trade and economic sanctions laws and
+regulations, as amended, including without limitation U.S. Export Administration Regulations and Office of Foreign Assets Control
+regulations. These laws include restrictions on destinations, end-users and end-use.
+Version Release Date: June 14, 2024
+
+

README.md

@@ -0,0 +1,53 @@
+---
+license: other
+license_name: nvidia-open-model-license
+license_link: https://developer.download.nvidia.com/licenses/nvidia-open-model-license-agreement-june-2024.pdf
+base_model: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
+tags: [nemotron, multimodal, mamba2, moe, quantized, rotorquant, mlx]
+---
+
+# Nemotron-3-Nano-Omni-30B-A3B-Reasoning - RotorQuant MLX 2-bit
+
+MLX 2-bit quantization of the **text tower** of `Nemotron-3-Nano-Omni-30B-A3B-Reasoning` (`nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16`)
+with RotorQuant weight method. Apple Silicon native via `mlx-lm`.
+
+This variant covers the LLM backbone only. Vision (CRADIO v4-H) + audio (Parakeet-TDT-0.6B-v2)
+encoders are NOT included — MLX-VLM Nemotron-Omni model class is **pending upstream support**
+(no PR observed as of 2026-05-04). For multimodal inference, use the GGUF variants with
+`llama-mtmd-cli` instead.
+
+For the matched-KV stack — RotorQuant weights + RotorQuant KV-cache modifier —
+see [`majentik/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-RotorQuant-MLX-2bit-RQ-KV`](https://huggingface.co/majentik/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-RotorQuant-MLX-2bit-RQ-KV).
+For the runtime KV-cache modifier itself, see
+[`majentik/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-RotorQuant`](https://huggingface.co/majentik/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-RotorQuant).
+
+## Modality matrix
+
+| Modality | Encoder | Quantization in this variant |
+|---|---|---|
+| Text | LLM backbone (Mamba-2 + Transformer hybrid Sparse MoE) | per the variant suffix |
+| Image | CRADIO v4-H | **BF16** (kept full-precision in every non-GGUF variant; GGUF uses mmproj-F16 split file) |
+| Audio | Parakeet-TDT-0.6B-v2 | **BF16** (same rationale) |
+| Video | Parakeet-TDT-0.6B-v2 + frame sampler | **BF16** (≤ 2 min, 256 frames @ 2 FPS) |
+
+NVIDIA's official FP8 / NVFP4 recipe keeps both encoders + the cross-modal
+MLP projectors in BF16 to preserve multimodal accuracy. We follow that
+convention in every quantized variant we ship.
+
+## Runtime quirks
+
+### MLX-LM (text-only)
+
+This variant covers the LLM backbone only. Vision + audio encoders
+are NOT included — MLX-VLM Nemotron-Omni model class is
+**pending upstream support** (no PR observed as of 2026-05-04).
+
+Use the `mlx_lm.generate` API; `enable_thinking` is a runtime flag
+(see below).
+
+### Reasoning mode
+
+`enable_thinking` defaults to `True`. To disable extended reasoning
+(e.g., for latency-sensitive cases), pass `enable_thinking=False`
+to the chat template / generate call. No separate "no-think"
+variant card exists — this is a runtime flag, not a model variant.

audio_model.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  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.
+
+"""Sound/Audio model components for multimodal integration.
+
+This module provides the SoundEncoder (wrapping Parakeet from HuggingFace transformers)
+and SoundProjection (MLP to project audio embeddings to LLM hidden size).
+
+The Parakeet model in HuggingFace transformers is documented at:
+https://huggingface.co/docs/transformers/en/model_doc/parakeet
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from transformers import ParakeetEncoder, ParakeetEncoderConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class SquaredReLU(nn.Module):
+    """Squared ReLU activation function."""
+    def forward(self, x):
+        return torch.pow(torch.nn.functional.relu(x), 2)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = 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.eps)
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+
+class SoundProjection(nn.Module):
+    """MLP projection from sound encoder hidden size to LLM hidden size.
+    
+    Architecture: RMSNorm -> linear1 -> SquaredReLU -> linear2
+    
+    This matches the Megatron checkpoint conversion structure:
+    - sound_projection.norm.weight
+    - sound_projection.linear1.weight
+    - sound_projection.linear2.weight
+    - sound_projection.linear1.bias (optional)
+    - sound_projection.linear2.bias (optional)
+    """
+    
+    def __init__(
+        self,
+        sound_hidden_size: int,
+        projection_hidden_size: int,
+        llm_hidden_size: int,
+        bias: bool = True,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.norm = RMSNorm(sound_hidden_size, eps=eps)
+        self.linear1 = nn.Linear(sound_hidden_size, projection_hidden_size, bias=bias)
+        self.activation = SquaredReLU()
+        self.linear2 = nn.Linear(projection_hidden_size, llm_hidden_size, bias=bias)
+        
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Project sound embeddings to LLM embedding space.
+        
+        Args:
+            hidden_states: Sound encoder output [batch, seq_len, sound_hidden_size]
+        
+        Returns:
+            Projected embeddings [batch, seq_len, llm_hidden_size]
+        """
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.linear1(hidden_states)
+        hidden_states = self.activation(hidden_states)
+        hidden_states = self.linear2(hidden_states)
+        return hidden_states
+
+
+class SoundEncoder(nn.Module):
+    """Wrapper around the Parakeet encoder from HuggingFace transformers.
+    
+    The Parakeet model is an ASR model with a Fast Conformer encoder.
+    We use only the encoder portion to extract audio embeddings.
+    
+    Checkpoint structure:
+    - sound_encoder.encoder.feature_extractor.* -> Feature extraction (mel spectrogram)
+    - sound_encoder.encoder.pre_encode.* -> Pre-encoding convolutions
+    - sound_encoder.encoder.layers.* -> Conformer layers
+    
+    Reference: https://huggingface.co/docs/transformers/en/model_doc/parakeet
+    """
+    
+    def __init__(self, config=None):
+        super().__init__()
+        
+        if config is not None:
+            # Build from config - handle both dict and config object
+            if hasattr(config, '__dict__'):
+                # It's a config object, extract relevant params for ParakeetConfig
+                config_dict = {
+                    'attention_bias': getattr(config, 'attention_bias', False),
+                    'hidden_size': getattr(config, 'hidden_size', 1024),
+                    'num_attention_heads': getattr(config, 'num_attention_heads', 8),
+                    'num_hidden_layers': getattr(config, 'num_hidden_layers', 24),
+                    'intermediate_size': getattr(config, 'intermediate_size', 4096),
+                    'conv_kernel_size': getattr(config, 'conv_kernel_size', 31),
+                    'convolution_bias': getattr(config, 'convolution_bias', False),
+                    'feat_in': getattr(config, 'feat_in', 80),
+                    'subsampling_factor': getattr(config, 'subsampling_factor', 8),
+                    'subsampling_conv_channels': getattr(config, 'subsampling_conv_channels', 256),
+                    'subsampling_conv_kernel_size': getattr(config, 'subsampling_conv_kernel_size', 3),
+                    'subsampling_conv_stride': getattr(config, 'subsampling_conv_stride', 2),
+                    'num_mel_bins': getattr(config, 'num_mel_bins', 128),
+                    'scale_input': getattr(config, 'scale_input', False),
+                }
+            elif isinstance(config, dict):
+                config_dict = config
+            else:
+                config_dict = {}
+            
+            # Create ParakeetConfig with the extracted parameters
+            parakeet_config = ParakeetEncoderConfig(**config_dict)
+            self.config = parakeet_config
+            self.encoder = ParakeetEncoder(parakeet_config)
+        else:
+            raise ValueError(
+                "config must be provided, "
+                "and ParakeetEncoder must be available in transformers."
+            )
+    
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode audio features.
+        
+        Args:
+            input_features: Mel spectrogram features [batch, seq_len, feature_dim]
+            attention_mask: Optional attention mask [batch, seq_len]
+            
+        Returns:
+            Audio embeddings [batch, encoded_seq_len, hidden_size]
+        """
+        outputs = self.encoder(
+            input_features=input_features,
+            attention_mask=attention_mask,
+        )
+        # Return the last hidden state
+        return outputs.last_hidden_state
+    
+    @property
+    def hidden_size(self) -> int:
+        """Return the hidden size of the encoder."""
+        return self.config.hidden_size

chat_template.jinja

@@ -0,0 +1,273 @@
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{%- endmacro -%}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set reasoning_budget = reasoning_budget if reasoning_budget is defined else None %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{#- Scan messages for VLM thinking toggles to override enable_thinking -#}
+{%- set toggle = namespace(enable=enable_thinking) %}
+{%- for m in messages %}
+    {%- if m['role'] == 'user' or m['role'] == 'system' -%}
+        {%- if m['content'] is string -%}
+            {%- set c = m['content'] %}
+            {%- if '/think' in c.replace('</think>', '') -%}
+                {%- set toggle.enable = true -%}
+            {%- elif '/no_think' in c -%}
+                {%- set toggle.enable = false -%}
+            {%- endif -%}
+        {%- else -%}
+            {%- for part in m['content'] -%}
+                {%- if part['type'] == 'text' -%}
+                    {%- set c = part['text'] %}
+                    {%- if '/think' in c.replace('</think>', '') -%}
+                        {%- set toggle.enable = true -%}
+                    {%- elif '/no_think' in c -%}
+                        {%- set toggle.enable = false -%}
+                    {%- endif -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+{#- Prepare message iteration similar to LM template -#}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+    {%- if m["role"] == "user" %}
+        {%- set ns.last_user_idx = loop.index0 %}
+    {%- endif %}
+{%- endfor -%}
+
+{%- if messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{#- Recompute last_user_idx relative to loop_messages after handling system -#}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+    {%- if m["role"] == "user" %}
+        {%- set ns.last_user_idx = loop.index0 %}
+    {%- endif %}
+{%- endfor -%}
+{#- System preamble with LM formatting, sanitize thinking toggles -#}
+{%- if system_message is defined %}
+    {%- set sys_content = system_message | string %}
+    {%- set sys_content = sys_content.replace('</think>', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '</think>') %}
+    {{- "<|im_start|>system\n" + sys_content }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system\n" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "\n\n" }}
+    {%- endif %}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '\n<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '\n<parameter>' }}
+                {{- '\n<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '\n<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '\n</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {%- set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '\n<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '\n</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '\n</function>' }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+{%- endif -%}
+{%- if system_message is defined %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif -%}
+{#- Iterate conversation -#}
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {#- Use LM assistant handling -#}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>\n" ~ message.reasoning_content ~ "\n</think>\n" ~ (message.content | default('', true)) %}
+        {%- else %}
+            {%- set content = message.content | default('', true) %}
+            {%- if content is string -%}
+                {%- if '<think>' not in content and '</think>' not in content -%}
+                    {%- set content = "<think></think>" ~ content -%}
+                {%- endif -%}
+            {%- else -%}
+                {%- set content = content -%}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {{- '<|im_start|>assistant\n' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '\n' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c | trim %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '\n' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>\n<function=' ~ tool_call.name ~ '>\n' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>\n' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '\n</parameter>\n' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>\n</tool_call>\n' -}}
+                {%- endfor %}
+                {{- '<|im_end|>\n' }}
+        {%- else %}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant\n<|im_end|>\n' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {#- Build VLM multimodal content when content is a sequence -#}
+        {%- if message.content is string -%}
+            {%- set content = (message.content | string) %}
+        {%- else -%}
+            {%- set text_ns = namespace(val='') -%}
+            {%- set mm_content = '' -%}
+            {%- set counters = namespace(images=0, videos=0, audios=0) -%}
+            {%- for part in message['content'] -%}
+                {%- if part['type'] == 'image' or part['type'] == 'image_url' -%}
+                    {%- set counters.images = counters.images + 1 -%}
+                {%- elif part['type'] == 'video' or part['type'] == 'video_url' -%}
+                    {%- set counters.videos = counters.videos + 1 -%}
+                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' -%}
+                    {%- set counters.audios = counters.audios + 1 -%}
+                {%- elif part['type'] == 'text' -%}
+                    {%- set text_ns.val = text_ns.val + part['text'] -%}
+                {%- endif -%}
+            {%- endfor -%}
+            {%- if '<image>' in text_ns.val -%}
+                {%- set counters.images = 0 -%}
+            {%- endif -%}
+            {%- if '<video>' in text_ns.val -%}
+                {%- set counters.videos = 0 -%}
+            {%- endif -%}
+            {%- if '<so_embedding>' in text_ns.val -%}
+                {%- set counters.audios = 0 -%}
+            {%- endif -%}
+            {%- if counters.images > 1 -%}
+                {%- set image_tags = namespace(tags=[]) -%}
+                {%- for i in range(counters.images) -%}
+                    {%- set image_tags.tags = image_tags.tags + ['<image ' + (i + 1)|string + '><image>'] -%}
+                {%- endfor -%}
+                {%- set mm_content = ' '.join(image_tags.tags) + '\n' -%}
+            {%- elif counters.images == 1 -%}
+                {%- set mm_content = '<image>\n' -%}
+            {%- endif -%}
+            {%- set mm_content = mm_content + '<video>\n' * counters.videos -%}
+            {%- set mm_content = mm_content + '<so_embedding>\n' * counters.audios -%}
+            {%- set content = mm_content + text_ns.val.lstrip('\n') -%}
+        {%- endif -%}
+        {#- Sanitize thinking toggle directives from user/system content -#}
+        {%- set content = content.replace('</think>', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '</think>') -%}
+        {%- set content = content | trim -%}
+        {%- if message.role == "user" and loop.index0 == ns.last_user_idx and reasoning_budget is not none -%}
+            {{- content + '\n\n{thinking token budget: ' + (reasoning_budget | string) + '}' -}}
+        {%- else -%}
+            {{- content -}}
+        {%- endif -%}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>\n' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor -%}
+{#- Generation prompt using computed thinking toggle -#}
+{%- if add_generation_prompt %}
+    {%- if toggle.enable %}
+        {{- '<|im_start|>assistant\n<think>\n' }}
+    {%- else %}
+        {{- '<|im_start|>assistant\n<think></think>' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,371 @@
+{
+  "architectures": [
+    "NemotronH_Nano_Omni_Reasoning_V3"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration.NemotronH_Nano_Omni_Reasoning_V3_Config",
+    "AutoModel": "modeling.NemotronH_Nano_Omni_Reasoning_V3",
+    "AutoModelForCausalLM": "modeling.NemotronH_Nano_Omni_Reasoning_V3"
+  },
+  "max_sequence_length": 131072,
+  "downsample_ratio": 0.5,
+  "force_image_size": 512,
+  "patch_size": 16,
+  "use_thumbnail": true,
+  "eos_token_id": 11,
+  "model_type": "NemotronH_Nano_Omni_Reasoning_V3",
+  "ps_version": "v2",
+  "template": "n5h_5p5_nanov2",
+  "torch_dtype": "bfloat16",
+  "image_tag_type": "internvl",
+  "img_context_token_id": 18,
+  "video_context_token_id": 131081,
+  "img_context_token": "<image>",
+  "video_context_token": "<video>",
+  "img_start_token": "<img>",
+  "img_end_token": "</img>",
+  "vit_hidden_size": 1280,
+  "projector_hidden_size": 20480,
+  "norm_mean": [
+    0.48145466,
+    0.4578275,
+    0.40821073
+  ],
+  "norm_std": [
+    0.26862954,
+    0.26130258,
+    0.27577711
+  ],
+  "video_pruning_rate": 0.7,
+  "sound_context_token_id": 27,
+  "sound_context_token": "<so_embedding>",
+  "sound_config": {
+    "model_type": "parakeet",
+    "hidden_size": 1024,
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "intermediate_size": 4096,
+    "conv_kernel_size": 9,
+    "convolution_bias": false,
+    "subsampling_conv_channels": 256,
+    "subsampling_conv_kernel_size": 3,
+    "subsampling_conv_stride": 2,
+    "subsampling_factor": 8,
+    "num_mel_bins": 128,
+    "projection_hidden_size": 4096,
+    "projection_bias": false,
+    "sampling_rate": 16000
+  },
+  "llm_config": {
+    "architectures": [
+      "NemotronHForCausalLM"
+    ],
+    "auto_map": {
+      "AutoConfig": "configuration_nemotron_h.NemotronHConfig",
+      "AutoModelForCausalLM": "modeling_nemotron_h.NemotronHForCausalLM"
+    },
+    "model_type": "nemotron_h",
+    "bos_token_id": 1,
+    "chunk_size": 128,
+    "conv_kernel": 4,
+    "expand": 2,
+    "eos_token_id": 11,
+    "pad_token_id": 0,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.55.4",
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "head_dim": 128,
+    "hidden_dropout": 0.0,
+    "hidden_size": 2688,
+    "hybrid_override_pattern": "MEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEM*EMEMEMEM*EMEMEMEME",
+    "initializer_range": 0.02,
+    "intermediate_size": 1856,
+    "layer_norm_epsilon": 1e-05,
+    "mamba_head_dim": 64,
+    "mamba_hidden_act": "silu",
+    "mamba_num_heads": 64,
+    "mamba_proj_bias": false,
+    "max_position_embeddings": 262144,
+    "mlp_bias": false,
+    "mlp_hidden_act": "relu2",
+    "moe_intermediate_size": 1856,
+    "moe_shared_expert_intermediate_size": 3712,
+    "n_group": 1,
+    "n_groups": 8,
+    "n_routed_experts": 128,
+    "n_shared_experts": 1,
+    "norm_eps": 1e-05,
+    "norm_topk_prob": true,
+    "num_attention_heads": 32,
+    "num_experts_per_tok": 6,
+    "num_hidden_layers": 52,
+    "num_key_value_heads": 2,
+    "num_logits_to_keep": 1,
+    "partial_rotary_factor": 1.0,
+    "rescale_prenorm_residual": true,
+    "residual_in_fp32": false,
+    "rope_theta": 10000,
+    "routed_scaling_factor": 2.5,
+    "sliding_window": null,
+    "ssm_state_size": 128,
+    "tie_word_embeddings": false,
+    "time_step_floor": 0.0001,
+    "time_step_limit": [
+      0.0,
+      1e+30
+    ],
+    "time_step_max": 0.1,
+    "time_step_min": 0.001,
+    "topk_group": 1,
+    "use_bias": false,
+    "use_cache": true,
+    "use_conv_bias": true,
+    "use_mamba_kernels": true,
+    "vocab_size": 131072
+  },
+  "vision_config": {
+    "auto_map": {
+      "AutoConfig": "nvidia/C-RADIOv2-H--hf_model.RADIOConfig",
+      "AutoModel": "nvidia/C-RADIOv2-H--hf_model.RADIOModel"
+    },
+    "adaptor_configs": {},
+    "adaptor_names": null,
+    "architectures": [
+      "RADIOModel"
+    ],
+    "args": {
+      "aa": null,
+      "amp": true,
+      "amp_dtype": "bfloat16",
+      "amp_impl": "native",
+      "aug_repeats": 0,
+      "aug_splits": 0,
+      "bn_eps": null,
+      "bn_momentum": null,
+      "cache_dir": null,
+      "channels_last": false,
+      "checkpoint_hist": 10,
+      "chk_keep_forever": 100,
+      "class_map": "",
+      "clip_grad": null,
+      "clip_mode": "norm",
+      "cls_token_per_teacher": true,
+      "coco_annotations_file": "/datasets/coco2017-adlsa/annotations/captions_val2017.json",
+      "coco_image_dir": "/datasets/coco2017-adlsa/val2017",
+      "color_jitter": 0.4,
+      "cooldown_epochs": 0,
+      "cpe_max_size": 2048,
+      "crd_loss": false,
+      "crd_loss_weight": 0.8,
+      "crop_pct": null,
+      "cutmix": 0.0,
+      "cutmix_minmax": null,
+      "dataset_download": false,
+      "debug_full_knn": false,
+      "decay_epochs": 90,
+      "decay_milestones": [
+        90,
+        180,
+        270
+      ],
+      "decay_rate": 0.1,
+      "depchain": true,
+      "dist_bn": "reduce",
+      "dist_norm_weight": 0.0,
+      "distributed": true,
+      "drop": 0.0,
+      "drop_block": null,
+      "drop_connect": null,
+      "drop_path": null,
+      "dtype": "bfloat16",
+      "epoch_repeats": 0.0,
+      "eval": false,
+      "eval_metric": "knn_top1",
+      "eval_teacher": false,
+      "eval_teacher_only": false,
+      "eval_throughput": false,
+      "fast_norm": false,
+      "fd_loss_fn": "MSE",
+      "feature_normalization": "SHIP_NORM",
+      "feature_summarizer": "cls_token",
+      "feature_upscale_factor": null,
+      "force_new_wandb_id": false,
+      "force_spectral_reparam": true,
+      "freeze_bn": false,
+      "fsdp": false,
+      "fuser": "",
+      "gp": null,
+      "grad_accum_steps": 1,
+      "grad_checkpointing": false,
+      "head_init_bias": null,
+      "head_init_scale": null,
+      "head_warmup": 5,
+      "head_weight_decay": 0.001,
+      "hflip": 0.5,
+      "img_size": null,
+      "in_chans": null,
+      "initial_checkpoint": null,
+      "input_size": null,
+      "interpolation": "",
+      "layer_decay": null,
+      "local_rank": 0,
+      "log_interval": 50,
+      "log_mlflow": false,
+      "log_wandb": true,
+      "loss_auto_balance": false,
+      "lr_base": 0.1,
+      "lr_base_scale": "",
+      "lr_base_size": 256,
+      "lr_cycle_decay": 0.5,
+      "lr_cycle_limit": 1,
+      "lr_cycle_mul": 1.0,
+      "lr_k_decay": 1.0,
+      "lr_noise": null,
+      "lr_noise_pct": 0.67,
+      "lr_noise_std": 1.0,
+      "mean": null,
+      "mesa": false,
+      "min_lr": 0,
+      "mixup": 0.0,
+      "mixup_mode": "batch",
+      "mixup_off_epoch": 0,
+      "mixup_prob": 1.0,
+      "mixup_switch_prob": 0.5,
+      "mlp_hidden_size": 1520,
+      "mlp_num_inner": 3,
+      "mlp_version": "v2",
+      "model": "vit_huge_patch16_224",
+      "model_kwargs": {},
+      "model_norm": false,
+      "momentum": 0.9,
+      "no_aug": false,
+      "no_ddp_bb": true,
+      "no_prefetcher": false,
+      "no_resume_opt": false,
+      "num_classes": null,
+      "opt_betas": null,
+      "opt_eps": null,
+      "patience_epochs": 10,
+      "pin_mem": false,
+      "prefetcher": true,
+      "pretrained": false,
+      "rank": 0,
+      "ratio": [
+        0.75,
+        1.3333333333333333
+      ],
+      "recount": 1,
+      "recovery_interval": 0,
+      "register_multiple": 10,
+      "remode": "pixel",
+      "reprob": 0.0,
+      "reset_loss_state": false,
+      "resplit": false,
+      "save_images": false,
+      "scale": [
+        0.5,
+        1.0
+      ],
+      "sched": "cosine",
+      "seed": 42,
+      "smoothing": 0.1,
+      "spectral_heads": false,
+      "spectral_reparam": false,
+      "split_bn": false,
+      "start_epoch": null,
+      "std": null,
+      "stream_teachers": true,
+      "sync_bn": false,
+      "synchronize_step": false,
+      "teachers": [
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 378,
+          "model": "ViT-H-14-378-quickgelu",
+          "name": "clip",
+          "pretrained": "dfn5b",
+          "type": "open_clip",
+          "use_summary": true
+        },
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 378,
+          "model": "ViT-SO400M-14-SigLIP-384",
+          "name": "siglip",
+          "pretrained": "webli",
+          "type": "open_clip",
+          "use_summary": true
+        },
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 378,
+          "model": "dinov2_vitg14_reg",
+          "name": "dino_v2",
+          "type": "dino_v2",
+          "use_summary": true
+        },
+        {
+          "fd_normalize": false,
+          "feature_distillation": true,
+          "input_size": 1024,
+          "model": "vit-h",
+          "name": "sam",
+          "type": "sam",
+          "use_summary": false
+        }
+      ],
+      "torchcompile": null,
+      "torchscript": false,
+      "train_interpolation": "random",
+      "train_split": "train",
+      "tta": 0,
+      "use_coco": false,
+      "use_multi_epochs_loader": false,
+      "val_ema_only": false,
+      "val_split": "val",
+      "vflip": 0.0,
+      "vitdet_version": 1,
+      "wandb_entity": "",
+      "wandb_job_type": "",
+      "wandb_name": "",
+      "wandb_project": "",
+      "warmup_lr": 1e-05,
+      "warmup_prefix": false,
+      "worker_seeding": "all",
+      "workers": 8,
+      "world_size": 256,
+      "min_num_patches": 1024,
+      "max_num_patches": 13312
+    },
+    "feature_normalizer_config": null,
+    "inter_feature_normalizer_config": null,
+    "max_resolution": 2048,
+    "patch_size": 16,
+    "preferred_resolution": [
+      768,
+      768
+    ],
+    "torch_dtype": "bfloat16",
+    "version": "radio_v2.5-h",
+    "vitdet_window_size": null,
+    "min_num_patches": 1024,
+    "max_num_patches": 13312,
+    "video_target_num_patches": 1024,
+    "video_maintain_aspect_ratio": true,
+    "video_temporal_patch_size": 2,
+    "video_prompt_version": 2,
+    "separate_video_embedder": true
+  },
+  "quantization": {
+    "group_size": 64,
+    "bits": 2
+  },
+  "quantization_config": {
+    "group_size": 64,
+    "bits": 2
+  }
+}

configuration.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  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.
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from .configuration_nemotron_h import NemotronHConfig
+from .configuration_radio import RADIOConfig
+
+logger = logging.get_logger(__name__)
+
+
+class SoundConfig(PretrainedConfig):
+    """Configuration for the sound/audio model (Parakeet encoder + projection)."""
+    model_type = "parakeet"
+    
+    def __init__(
+        self,
+        # Parakeet encoder config
+        hidden_size: int = 1024,
+        num_attention_heads: int = 8,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 4096,
+        conv_kernel_size: int = 31,
+        feat_in: int = 80,  # Mel features
+        subsampling_factor: int = 8,
+        # Projection config  
+        projection_hidden_size: int = 20480,
+        projection_bias: bool = True,
+        # Audio processing
+        sampling_rate: int = 16000,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.conv_kernel_size = conv_kernel_size
+        self.feat_in = feat_in
+        self.subsampling_factor = subsampling_factor
+        self.projection_hidden_size = projection_hidden_size
+        self.projection_bias = projection_bias
+        self.sampling_rate = sampling_rate
+
+
+class NemotronH_Nano_Omni_Reasoning_V3_Config(PretrainedConfig):
+    model_type = 'NemotronH_Nano_Omni_Reasoning_V3'
+    is_composition = True
+
+    def __init__(
+        self,
+        vision_config=None,
+        llm_config=None,
+        sound_config=None,
+        force_image_size=None,
+        downsample_ratio=0.5,
+        template=None,
+        ps_version='v1',
+        image_tag_type="internvl",
+        projector_hidden_size=4096,
+        vit_hidden_size=1280,
+        attn_implementation="flash_attention_2",
+        video_pruning_rate: float = 0.0,
+        video_temporal_patch_size: int = 2,
+        # Sound/audio settings
+        sound_context_token_id: int = None,
+        sound_context_token: str = "<audio>",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        if vision_config is not None:
+            self.vision_config = RADIOConfig(**vision_config)
+        else:
+            self.vision_config = RADIOConfig()
+
+        # Handle both cases: when loading from JSON (llm_config is dict) and when called internally by transformers (llm_config is None)
+        if llm_config is not None:
+            self.llm_config = NemotronHConfig(**llm_config)
+        else:
+            self.llm_config = NemotronHConfig()
+
+        # Sound/audio model configuration
+        if sound_config is not None:
+            self.sound_config = SoundConfig(**sound_config)
+        else:
+            self.sound_config = None  # Sound model is optional
+
+        # Assign configuration values
+        self.force_image_size = force_image_size
+        self.downsample_ratio = downsample_ratio
+        self.template = template  # TODO move out of here and into the tokenizer
+        self.ps_version = ps_version  # Pixel shuffle version
+        self.image_tag_type = image_tag_type # TODO: into the tokenizer too?
+        self.projector_hidden_size = projector_hidden_size
+        self.vit_hidden_size = vit_hidden_size
+        self.video_pruning_rate = video_pruning_rate
+        self.video_temporal_patch_size = video_temporal_patch_size
+        
+        # Sound/audio token settings
+        self.sound_context_token_id = sound_context_token_id
+        self.sound_context_token = sound_context_token
+
+        self._attn_implementation = attn_implementation
+        self.vision_config.use_flash_attn = self._attn_implementation is not None and "flash_attention" in self._attn_implementation
+        self.llm_config._attn_implementation = self._attn_implementation
+
+    # vLLM's `NemotronH_Nano_VL_V2` implementation reads the language-model sub-config as
+    # `config.text_config`. Our HF config stores it as `config.llm_config`; expose an alias so the
+    # same config object loads under both loaders without having to duplicate the dict on disk.
+    @property
+    def text_config(self):
+        return self.llm_config

configuration_nemotron_h.py

@@ -0,0 +1,276 @@
+# coding=utf-8
+# Copyright 2024 AI21 Labs Ltd. and the HuggingFace Inc. team. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. 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.
+"""NemotronH model configuration"""
+
+import re
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class NemotronHConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`NemotronHModel`]. It is used to instantiate a
+    NemotronH 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 NemotronH-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 NemotronH model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`NemotronHModel`]
+        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 `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`):
+            The pattern of the hybrid model. The pattern is a string of characters where each character represents M: Mamba2, *: Attention, -: MLP
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        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.
+        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_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.
+    """
+
+    model_type = "nemotron_h"
+    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="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-",
+        num_attention_heads=32,
+        head_dim=128,
+        num_key_value_heads=8,  # nemo: num_query_groups
+        mlp_hidden_act="relu2",
+        attention_bias=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=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        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_chunk_size=128,
+        rescale_prenorm_residual=True,
+        n_routed_experts=8,
+        n_shared_experts=1,
+        moe_intermediate_size=7688,
+        moe_shared_expert_intermediate_size=7688,
+        num_experts_per_tok=2,
+        routed_scaling_factor=1.0,
+        n_group=1,
+        topk_group=1,
+        norm_topk_prob=True,
+        **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.head_dim = head_dim
+        self.sliding_window = sliding_window
+        self.max_position_embeddings = max_position_embeddings
+        self.attention_dropout = attention_dropout
+        self.hidden_dropout = hidden_dropout
+
+        # Validate hybrid_override_pattern
+        # M: Mamba2, *: Attention, -: MLP, E: MoE
+        assert len(self.hybrid_override_pattern) == self.num_hidden_layers, "hybrid_override_pattern must have the same length as num_hidden_layers"
+        assert re.match(r"^[*\-ME]+$", 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.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.chunk_size = mamba_chunk_size
+        self.rescale_prenorm_residual = rescale_prenorm_residual
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.moe_intermediate_size = moe_intermediate_size
+        self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.routed_scaling_factor = routed_scaling_factor
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.norm_topk_prob = norm_topk_prob
+
+        # Derived per-layer block type list. Transformers 5.6+ looks this up as `layer_types` on the
+        # config to pick the correct cache structure (linear attention vs full attention). MLP (stateless)
+        # layers are tagged as "moe" here only because the transformers cache validator rejects "mlp";
+        # from the cache's point of view, both MLP and MoE layers need no kv cache (they become
+        # LinearAttentionLayer with zero state).
+        self.layer_types = [
+            "mamba" if self.hybrid_override_pattern[i] == "M" else
+            "attention" if self.hybrid_override_pattern[i] == "*" else
+            "moe" if self.hybrid_override_pattern[i] == "-" else "moe"
+            for i in range(self.num_hidden_layers)
+        ]
+        # Per-layer semantic labels used by the modeling code (includes "mlp", which the
+        # transformers `layer_types` validator would reject — that's why `layer_types` above maps
+        # "-" → "moe" for cache purposes, while this attribute keeps the true label for the block
+        # dispatch in NemotronHBlock).
+        self.layers_block_type = [
+            "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)
+        ]
+
+        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,
+        )

configuration_radio.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from dataclasses import dataclass
+from typing import Optional, NamedTuple, Union, List, Dict
+
+from transformers import PretrainedConfig
+
+
+class Resolution(NamedTuple):
+    height: int
+    width: int
+
+
+@dataclass
+class RadioResource:
+    url: str
+    patch_size: int
+    max_resolution: int
+    preferred_resolution: Resolution
+    vitdet_num_windowed: Optional[int] = None
+    vitdet_num_global: Optional[int] = None
+
+
+RESOURCE_MAP = {
+    # RADIOv2.5
+    "radio_v2.5-b": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-b_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-l": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h-norm": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h-norm.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-g.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1792,
+        preferred_resolution=(896, 896),
+        vitdet_num_global=8,
+    ),
+    # RADIO
+    "radio_v2.1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.1_bf16.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v1.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1050,
+        preferred_resolution=Resolution(378, 378),
+    ),
+    # E-RADIO
+    "e-radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/eradio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+    # C-RADIO
+    "c-radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/C-RADIOv2-g/resolve/main/c-radio_v2-g_half.pth.tar",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=8,
+    ),
+    "c-radio_v3-l": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-L/resolve/main/c-radio-v3_l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+}
+
+DEFAULT_VERSION = "radio_v2.5-h"
+
+
+class RADIOConfig(PretrainedConfig):
+    """Pretrained Hugging Face configuration for RADIO models."""
+
+    def __init__(
+        self,
+        args: Optional[dict] = None,
+        version: Optional[str] = DEFAULT_VERSION,
+        patch_size: Optional[int] = None,
+        max_resolution: Optional[int] = None,
+        preferred_resolution: Optional[Resolution] = None,
+        adaptor_names: Union[str, List[str]] = None,
+        adaptor_configs: Dict[str, Dict[str, int]] = None,
+        vitdet_window_size: Optional[int] = None,
+        feature_normalizer_config: Optional[dict] = None,
+        inter_feature_normalizer_config: Optional[dict] = None,
+        **kwargs,
+    ):
+        self.args = args
+        for field in ["dtype", "amp_dtype"]:
+            if self.args is not None and field in self.args:
+                # Convert to a string in order to make it serializable.
+                # For example for torch.float32 we will store "float32",
+                # for "bfloat16" we will store "bfloat16".
+                self.args[field] = str(args[field]).split(".")[-1]
+        self.version = version
+        resource = RESOURCE_MAP[version]
+        self.patch_size = patch_size or resource.patch_size
+        self.max_resolution = max_resolution or resource.max_resolution
+        self.preferred_resolution = (
+            preferred_resolution or resource.preferred_resolution
+        )
+        self.adaptor_names = adaptor_names
+        self.adaptor_configs = adaptor_configs
+        self.vitdet_window_size = vitdet_window_size
+        self.feature_normalizer_config = feature_normalizer_config
+        self.inter_feature_normalizer_config = inter_feature_normalizer_config
+        super().__init__(**kwargs)

evs.py

@@ -0,0 +1,73 @@
+import torch
+from typing import Tuple
+
+class EfficientVideoSampling:
+    @staticmethod
+    def compute_retention_mask(
+        *,
+        video_embeds: torch.FloatTensor,
+        thw: torch.LongTensor,
+        spatial_merge_size: int,
+        q: float,
+    ):
+        """
+        Computes the retention mask for video embeddings based on the grid dimensions.
+
+        Args:
+            video_embeds (`torch.FloatTensor` of shape `(T * H * W, hidden_size)`):
+                The video embeddings to compute the retention mask for.
+            thw (`torch.LongTensor` of shape `(3)`):
+                The temporal, height and width of feature shape of each video in LLM.
+            spatial_merge_size (`int`): The spatial merge size of the video embeddings.
+                If embeddings will be downsampled *later*, this should be the downsampling factor.
+            q: (`float`): Pruning rate factor, indicating number of tokens to prune (remove)
+
+        Returns:
+            `torch.Tensor`: The retention mask for the video embeddings (T * H * W).
+                1 for tokens to keep, 0 for tokens to prune.
+        """
+        T, H, W = thw
+
+        # video_embeds = einops.rearrange(
+        #     video_embeds,
+        #     "(T H W) C -> T H W C",
+        #     T=T,
+        #     H=H // spatial_merge_size,
+        #     W=W // spatial_merge_size,
+        # )
+        # Use reshape instead of einops to avoid graph breaks
+        video_embeds = video_embeds.reshape(
+            T, H // spatial_merge_size, W // spatial_merge_size, video_embeds.size(-1)
+        )
+
+        # Core EVS
+        similarity = torch.nn.functional.cosine_similarity(
+            video_embeds[1:, ...], video_embeds[:-1, ...], dim=-1
+        )
+        dissimilarity = 1 - similarity
+
+        # Always ensure we include all tokens from the first frame
+        dissimilarity = torch.cat(
+            [255 * torch.ones_like(video_embeds[:1, :, :, 0]), dissimilarity], dim=0
+        )
+        dissimilarity_flat = dissimilarity.view(-1)
+
+        min_num_tokens = (H // spatial_merge_size) * (W // spatial_merge_size)  # a single frame
+        evs_num_tokens = int(T * min_num_tokens * (1 - q))
+        num_tokens_to_keep = max(min_num_tokens, evs_num_tokens)
+
+        order = torch.argsort(dissimilarity_flat,
+                              dim=-1,
+                              descending=True,
+                              stable=True)
+        topk_indices = order[:num_tokens_to_keep]
+
+        retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool)
+        retention_mask[topk_indices] = True
+        retention_mask = retention_mask.reshape(dissimilarity.size())
+
+        # print(
+        #     f"Computed retention mask of shape {retention_mask.shape=} with sparsity {retention_mask.float().mean().item():.4f} for {q=}",
+        # )
+        mask = retention_mask.view(-1)  # "T H W -> (T H W)"
+        return mask

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [2, 11],
+  "pad_token_id": 0,
+  "do_sample": true,
+  "temperature": 0.6,
+  "top_p": 0.95,
+  "max_new_tokens": 16384,
+  "reasoning_budget": 16384,
+  "reasoning_grace": 512,
+  "repetition_penalty": 1.0,
+  "transformers_version": "4.55.4"
+}

image_processing.py

@@ -0,0 +1,228 @@
+import math
+from typing import List, Optional, Union
+
+from PIL import Image
+import torch
+from transformers.image_processing_base import BatchFeature
+from transformers.image_processing_utils_fast import BaseImageProcessorFast
+from transformers.image_utils import make_list_of_images, get_image_type, ImageInput, ImageType
+from transformers.utils import TensorType
+import torchvision.transforms as T
+
+
+class NemotronH_Nano_Omni_Reasoning_V3ImageProcessor(BaseImageProcessorFast):
+    """
+    Dynamic-resolution image processor for the V3 omni model.
+
+    Each image is resized to a single tile whose patch-grid `(h_patches, w_patches)` is chosen to
+    land between `min_num_patches` and `max_num_patches` (on a 16×16-pixel grid), respecting
+    aspect ratio. This matches the algorithm in vLLM's `DynamicResolutionImageTiler`
+    (`vllm/model_executor/models/nano_nemotron_vl.py`) so HF and vLLM inference see identical pixel
+    inputs.
+    """
+
+    model_input_names = ["pixel_values"]
+
+    def __init__(
+        self,
+        norm_mean=None,
+        norm_std=None,
+        patch_size=16,
+        downsample_ratio=0.5,
+        min_num_patches=1024,
+        max_num_patches=13312,
+        max_model_len=16384,
+        video_target_num_patches=1024,
+        video_maintain_aspect_ratio=True,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.norm_mean = norm_mean
+        self.norm_std = norm_std
+        self.patch_size = patch_size
+        self.downsample_ratio = downsample_ratio
+        # Integer reduction factor for pixel_shuffle (downsample_ratio = 0.5 → factor 2).
+        self._downsample_factor = int(round(1.0 / downsample_ratio))
+        # Per-image patch-grid bounds (on the pre-pixel-shuffle 16×16 grid).
+        self.min_num_patches = min_num_patches
+        self.max_num_patches = max_num_patches
+        self.max_model_len = max_model_len
+        # Video frames use a separate (fixed) target-patch budget with aspect-ratio preserved.
+        # Matches vLLM's `_compute_aspect_preserving_size` in `nano_nemotron_vl.py`.
+        self.video_target_num_patches = video_target_num_patches
+        self.video_maintain_aspect_ratio = video_maintain_aspect_ratio
+
+    # Keep the PIL image through to `_preprocess` — we need PIL.resize (bicubic) to match vLLM's
+    # algorithm exactly; resizing a tensor via `torchvision.transforms.Resize` uses different
+    # kernels and breaks bit-exact agreement.
+    def _process_image(self, image: ImageInput, **kwargs):
+        if get_image_type(image) == ImageType.PIL:
+            if image.mode != "RGB":
+                image = image.convert("RGB")
+        return image
+
+    # transformers 5.6 renamed this hook from `_process_image` to `process_image`; alias both.
+    process_image = _process_image
+
+    # Toggled by `processing.py` around video calls (the strict `ImagesKwargs` validator won't let
+    # us thread a new kwarg down, so we use a flag on the instance instead).
+    _is_video_mode: bool = False
+
+    def _preprocess(
+        self,
+        images,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> BatchFeature:
+        """Port of vLLM's `DynamicResolutionImageTiler._images_to_pixel_values_lst`.
+
+        When `self._is_video_mode=True` (flipped by `processing.py` before the video call), each
+        input is resized using the **video** target-size rule (`video_target_num_patches`,
+        aspect-ratio preserved) instead of the image dynamic-res rule. This matches vLLM's split
+        between `video_to_pixel_values` (video path) and `DynamicResolutionImageTiler` (image
+        path).
+        """
+        is_video = self._is_video_mode
+        images = make_list_of_images(images)
+
+        target_sizes = []
+        if is_video:
+            for img in images:
+                target_w_patches, target_h_patches = self._compute_target_patches_video(img)
+                target_sizes.append((target_w_patches, target_h_patches))
+        else:
+            # Image path: per-image budget bounded by [min_num_patches, max_num_patches], with a
+            # global cap derived from `max_model_len` × pixel-shuffle factor².
+            num_tokens_available = self.max_model_len - 4  # match vLLM's reserve
+            budget = num_tokens_available * (self._downsample_factor ** 2)
+            budget = max(budget, self.min_num_patches * len(images))
+            max_budget = self.max_num_patches if (self.max_num_patches and self.max_num_patches > 0) else float("inf")
+            per_image_budget = [max(min(budget, max_budget), self.min_num_patches) for _ in images]
+            # Single-pass — vLLM has an iterative scale-down for the batch, but it rarely binds in
+            # single-image / small-batch inference.
+            for img, tokens_for_media in zip(images, per_image_budget):
+                target_w_patches, target_h_patches = self._compute_target_patches(img, tokens_for_media)
+                target_sizes.append((target_w_patches, target_h_patches))
+
+        import numpy as np
+        norm_mean = torch.tensor(self.norm_mean).view(1, 3, 1, 1)
+        norm_std = torch.tensor(self.norm_std).view(1, 3, 1, 1)
+
+        pixel_values_list = []
+        num_tokens_per_image = []
+        imgs_sizes = []
+        for img, (wp, hp) in zip(images, target_sizes):
+            target_w = wp * self.patch_size
+            target_h = hp * self.patch_size
+            # Use torch's antialiased bicubic interpolation to match vLLM's
+            # `_bicubic_resize_and_normalize` (`torch.nn.functional.interpolate`, `antialias=True`).
+            # PIL's bicubic uses a different kernel (and no antialiasing), producing visibly different
+            # pixel values that amplify through the 52-layer ViT / mamba stack and cause HF/vLLM
+            # outputs to diverge past the first few tokens.
+            arr = np.asarray(img, dtype=np.uint8)  # (H, W, 3)
+            t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0).to(dtype=torch.float32)  # (1, 3, H, W)
+            if t.shape[-2] != target_h or t.shape[-1] != target_w:
+                t = torch.nn.functional.interpolate(
+                    t, size=(target_h, target_w), mode="bicubic", align_corners=False, antialias=True
+                )
+            t = (t / 255.0 - norm_mean) / norm_std
+            pixel_values_list.append(t.squeeze(0))  # (3, H, W)
+            num_tokens_per_image.append((wp * hp) // (self._downsample_factor ** 2))
+            imgs_sizes.append((target_h, target_w))
+
+        # Stack if all images have the same target size (common for same-aspect-ratio batches);
+        # otherwise keep as a list of (3, H_i, W_i) tensors. The outer model's `extract_feature`
+        # handles both.
+        all_same_shape = all(t.shape == pixel_values_list[0].shape for t in pixel_values_list)
+        if all_same_shape:
+            pixel_values = torch.stack(pixel_values_list, dim=0)
+        else:
+            pixel_values = pixel_values_list
+
+        return BatchFeature(
+            data={
+                "pixel_values": pixel_values,
+                # One tile per image in dynamic mode — `num_tokens` is what the text-side
+                # placeholder expansion should use.
+                "num_patches": [1] * len(images),
+                "num_tokens": num_tokens_per_image,
+                "imgs_sizes": imgs_sizes,
+            },
+            tensor_type=(return_tensors if all_same_shape else None),
+        )
+
+    def _compute_target_patches(self, img: Image.Image, tokens_available: int):
+        """Port of `DynamicResolutionImageTiler.process_media` (image-only, no thumbnail)."""
+        orig_w, orig_h = img.width, img.height
+        # Ceil-ish: `round(x + 0.5)` == `floor(x) + 1` for non-integer x, `x` for integer.
+        closest_patch_h = round(orig_h / self.patch_size + 0.5)
+        closest_patch_w = round(orig_w / self.patch_size + 0.5)
+        patches = closest_patch_h * closest_patch_w
+
+        # Downscale to fit the token budget.
+        factor = min(math.sqrt(tokens_available / patches), 1.0)
+        target_h = math.floor(factor * closest_patch_h)
+        target_w = math.floor(factor * closest_patch_w)
+
+        # Scale up if below the per-image minimum.
+        if (
+            tokens_available > self.min_num_patches
+            and target_h * target_w < self.min_num_patches
+        ):
+            up = math.sqrt(self.min_num_patches / (target_h * target_w))
+            target_h = math.ceil(up * target_h)
+            target_w = math.ceil(up * target_w)
+
+        # Round each dim to a multiple of the pixel_shuffle factor so tokens divide evenly.
+        divisor = self._downsample_factor
+        rem_h = target_h % divisor
+        if rem_h:
+            inc_h = divisor - rem_h
+            if (target_h + inc_h) * target_w <= tokens_available:
+                target_h += inc_h
+            else:
+                target_h = max(divisor, target_h - rem_h)
+        rem_w = target_w % divisor
+        if rem_w:
+            inc_w = divisor - rem_w
+            if target_h * (target_w + inc_w) <= tokens_available:
+                target_w += inc_w
+            else:
+                target_w = max(divisor, target_w - rem_w)
+
+        return target_w, target_h
+
+    def _compute_target_patches_video(self, img: Image.Image):
+        """Port of vLLM's `_compute_aspect_preserving_size` for video frames.
+
+        Each frame is resized to roughly `video_target_num_patches` (default 1024) on the 16×16
+        grid, with aspect ratio preserved and dims snapped to a multiple of the pixel_shuffle
+        factor. For `maintain_aspect_ratio=False`, it falls back to a square of sqrt(target)
+        patches.
+        """
+        orig_w, orig_h = img.width, img.height
+        target = self.video_target_num_patches
+        divisor = self._downsample_factor  # 2 for pixel_shuffle
+        if self.video_maintain_aspect_ratio:
+            aspect_wh = orig_w / max(orig_h, 1)
+            ph = max(round(math.sqrt(target / aspect_wh)), 1)
+            pw = max(round(math.sqrt(target * aspect_wh)), 1)
+            if divisor > 1:
+                rem_h = ph % divisor
+                rem_w = pw % divisor
+                ph_up = ph + (divisor - rem_h if rem_h else 0)
+                ph_down = ph - rem_h
+                pw_up = pw + (divisor - rem_w if rem_w else 0)
+                pw_down = pw - rem_w
+                # Prefer rounding up when the up-rounded patch count still fits the target;
+                # otherwise round down (mirrors vLLM's logic exactly).
+                if ph_up * pw_up <= target:
+                    ph, pw = ph_up, pw_up
+                else:
+                    ph = max(divisor, ph_down)
+                    pw = max(divisor, pw_down)
+        else:
+            side = int(math.sqrt(target))
+            side = max(divisor, (side // divisor) * divisor)
+            ph = pw = side
+        return pw, ph

model-00001-of-00003.safetensors

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

model-00002-of-00003.safetensors

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

model-00003-of-00003.safetensors

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

modeling.py

@@ -0,0 +1,638 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  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.
+import os
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import transformers
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import AutoModel, AutoModelForCausalLM, GenerationConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration import NemotronH_Nano_Omni_Reasoning_V3_Config
+from .modeling_nemotron_h import NemotronHForCausalLM
+from .evs import EfficientVideoSampling
+from .audio_model import SoundEncoder, SoundProjection
+
+logger = logging.get_logger(__name__)
+
+
+"""
+The following code is adapted from the
+https://huggingface.co/OpenGVLab/InternVL2-Llama3-76B/blob/main/modeling_internvl_chat.py repository
+
+The chat function is adapted to handle NVLM 1-D tile-tagging design for dynamic high-resolution images.
+"""
+
+
+class SquaredReLU(nn.Module):
+    def forward(self, x):
+        return torch.pow(torch.nn.functional.relu(x), 2)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = 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.eps)
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+
+def version_cmp(v1, v2, op='eq'):
+    import operator
+
+    from packaging import version
+    op_func = getattr(operator, op)
+    return op_func(version.parse(v1), version.parse(v2))
+
+
+class NemotronH_Nano_Omni_Reasoning_V3(PreTrainedModel):
+    config_class = NemotronH_Nano_Omni_Reasoning_V3_Config
+    main_input_name = 'pixel_values'
+    _supports_flash_attn_2 = True
+    _supports_flash_attn = True
+    _no_split_modules = ['NemotronHBlock']
+
+    def __init__(self, config: NemotronH_Nano_Omni_Reasoning_V3_Config):
+        super().__init__(config)
+
+        assert version_cmp(transformers.__version__, '4.36.2', 'ge')
+        image_size = config.force_image_size
+        patch_size = config.patch_size
+        self.patch_size = patch_size
+        self.template = config.template
+        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
+        self.downsample_ratio = config.downsample_ratio
+        self.ps_version = config.ps_version
+        self.image_tag_type = config.image_tag_type
+        self.img_context_token_id = config.img_context_token_id
+        self.video_context_token_id = config.video_context_token_id
+
+        logger.info(f'num_image_token: {self.num_image_token}')
+        logger.info(f'ps_version: {self.ps_version}')
+
+        # Instantiate LM directly to avoid Hugging Face dynamic module lookup requiring a repo id.
+        self.language_model = NemotronHForCausalLM(config.llm_config)
+        self.vision_model = AutoModel.from_config(config.vision_config, trust_remote_code=True)
+        self.vision_model.model._initialize_weights = self.vision_model.model._init_weights  # WAR for transformers issue 38358
+        self.vision_model.radio_model.make_preprocessor_external()
+
+        # Attach a separate 3D patch projection for video frames. The RADIO ViT ships with only a 2D
+        # `embedder` (shape `[embed_dim, C·P²]`); this repo's checkpoint also carries a
+        # `video_embedder` (shape `[embed_dim, T·C·P²]`) used for temporally-packed video patches,
+        # so we construct the module here to make the weight bind. `T = video_temporal_patch_size`
+        # is the number of frames collapsed into each temporal patch.
+        self.video_temporal_patch_dim = config.video_temporal_patch_size
+        pg = self.vision_model.radio_model.model.patch_generator
+        pg.video_embedder = nn.Linear(
+            in_features=self.video_temporal_patch_dim * 3 * pg.patch_size * pg.patch_size,
+            out_features=pg.embed_dim,
+            bias=False,
+        )
+
+        # Align CPE position-embedding interpolation with Megatron training + vLLM inference.
+        # The `nvidia/C-RADIOv2-H` remote code uses `align_corners=True` in eval mode, but the V3
+        # checkpoint was trained against `align_corners=False` (see Megatron's `radio.py`). That
+        # single-flag mismatch shifts every pos_embed by a fraction of a cell, which compounds
+        # through 52 ViT layers and is the main cause of HF/vLLM divergence for video (where CPE
+        # mode is active — dynamic-res tubelets don't match the model's native 2048-sized grid).
+        self._patch_cpe_align_corners(pg)
+
+        self.vision_model = self.vision_model.to(self.language_model.config.torch_dtype)
+
+        self.drop_vision_class_token = True
+
+        # Construct the vision projection.
+        # Default
+        vit_hidden_size = config.vit_hidden_size
+        vision_projection_hidden_size = config.projector_hidden_size
+        llm_hidden_size = config.llm_config.hidden_size
+
+        self.video_pruning_rate = config.video_pruning_rate
+
+        self.mlp1 = nn.Sequential(
+            RMSNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, eps=1e-5),
+            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, vision_projection_hidden_size, bias=False),
+            SquaredReLU(),
+            nn.Linear(vision_projection_hidden_size, llm_hidden_size, bias=False)
+        )
+        self.mlp1 = self.mlp1.to(self.language_model.config.torch_dtype)
+
+        # Sound/audio model components (optional - only if sound_config is provided)
+        self.sound_context_token_id = getattr(config, 'sound_context_token_id', None)
+        if config.sound_config is not None:
+            sound_config = config.sound_config
+            sound_hidden_size = sound_config.hidden_size
+            sound_projection_hidden_size = sound_config.projection_hidden_size
+            
+            # Initialize sound feature extractor for converting raw audio to mel spectrograms
+            from transformers import ParakeetFeatureExtractor
+            sampling_rate = getattr(sound_config, 'sampling_rate', 16000)
+            feature_size = getattr(sound_config, 'num_mel_bins', 128)
+            self.sound_feature_extractor = ParakeetFeatureExtractor(
+                sampling_rate=sampling_rate,
+                feature_size=feature_size,
+            )
+            logger.info(f'Sound feature extractor initialized with sampling_rate={sampling_rate}, feature_size={feature_size}')
+            
+            # Initialize sound encoder - wraps Parakeet from transformers
+            self.sound_encoder = SoundEncoder(config=sound_config)
+            self.sound_encoder = self.sound_encoder.to(self.language_model.config.torch_dtype)
+            
+            # Initialize sound projection MLP
+            self.sound_projection = SoundProjection(
+                sound_hidden_size=sound_hidden_size,
+                projection_hidden_size=sound_projection_hidden_size,
+                llm_hidden_size=llm_hidden_size,
+                bias=sound_config.projection_bias,
+            )
+            self.sound_projection = self.sound_projection.to(self.language_model.config.torch_dtype)
+            
+            logger.info(f'Sound model initialized with hidden_size={sound_hidden_size}')
+        else:
+            self.sound_encoder = None
+            self.sound_projection = None
+            self.sound_feature_extractor = None
+
+        self.all_tied_weights_keys = {}
+
+    def forward(
+            self,
+            pixel_values: torch.FloatTensor,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_flags: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            inputs_embeds = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if inputs_embeds is None:
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        image_flags = image_flags.squeeze(-1)
+
+        B, N, C = inputs_embeds.shape
+        inputs_embeds = inputs_embeds.reshape(B * N, C)
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+
+        vit_batch_size = pixel_values.shape[0]
+        vit_embeds = self.extract_feature(pixel_values)
+
+        del pixel_values
+
+        if torch.distributed.get_rank() == 0:
+            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
+
+        vit_embeds = vit_embeds[image_flags == 1]
+        try:
+            inputs_embeds[selected] = inputs_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        except Exception as e:
+            vit_embeds = vit_embeds.reshape(-1, C)
+            print(f'warning: {e}, inputs_embeds[selected].shape={inputs_embeds[selected].shape}, '
+                  f'vit_embeds.shape={vit_embeds.shape}')
+            n_token = selected.sum()
+            inputs_embeds[selected] = inputs_embeds[selected] * 0.0 + vit_embeds[:n_token]
+
+        del vit_embeds
+
+        inputs_embeds = inputs_embeds.reshape(B, N, C)
+
+        outputs = self.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    @staticmethod
+    def _patch_cpe_align_corners(patch_generator) -> None:
+        """Monkey-patch `patch_generator._get_pos_embeddings` so the CPE-mode eval-path interpolation
+        uses `align_corners=False` (Megatron training + vLLM inference convention) instead of the
+        `align_corners=True` that the `nvidia/C-RADIOv2-H` remote code ships with.
+        """
+        import math
+        import torch.nn.functional as F
+
+        orig_method = patch_generator._get_pos_embeddings.__func__ if hasattr(
+            patch_generator._get_pos_embeddings, "__func__"
+        ) else patch_generator._get_pos_embeddings
+
+        def _get_pos_embeddings_aligned(self, batch_size, input_dims):
+            if (self.num_rows, self.num_cols) == input_dims:
+                return self.pos_embed
+            pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols, -1).permute(0, 3, 1, 2)
+
+            def window_select(pe):
+                if input_dims[0] < pe.shape[-2]:
+                    pe = pe[..., :input_dims[0], :]
+                if input_dims[1] < pe.shape[-1]:
+                    pe = pe[..., :, :input_dims[1]]
+                return pe
+
+            if self.cpe_mode:
+                if self.training:
+                    # Keep the original training-time jitter path (grid_sample + align_corners=True);
+                    # only patch the eval branch, which is what Megatron/vLLM use and where the bug is.
+                    return orig_method(self, batch_size, input_dims)
+                max_dim = max(input_dims)
+                pos_embed = F.interpolate(
+                    pos_embed.float(), size=(max_dim, max_dim), align_corners=False, mode="bilinear"
+                ).to(pos_embed.dtype)
+                pos_embed = window_select(pos_embed)
+            else:
+                pos_embed = window_select(pos_embed)
+
+            if pos_embed.shape[-2:] != input_dims:
+                pos_embed = F.interpolate(
+                    pos_embed.float(), size=input_dims, align_corners=False, mode="bilinear"
+                ).to(pos_embed.dtype)
+
+            pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
+            return pos_embed
+
+        import types
+        patch_generator._get_pos_embeddings = types.MethodType(_get_pos_embeddings_aligned, patch_generator)
+
+    def pixel_shuffle(self, x, scale_factor=0.5):
+        n, w, h, c = x.size()
+        # N, W, H, C --> N, W, H * scale, C // scale
+        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
+        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
+        x = x.permute(0, 2, 1, 3).contiguous()
+        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
+        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
+                   int(c / (scale_factor * scale_factor)))
+        if self.ps_version == 'v1':
+            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
+                          'which results in a transposed image.')
+        else:
+            x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+
+    def extract_feature(self, pixel_values):
+        """Run the ViT on a batch of image tiles.
+
+        Handles two layouts:
+        - A single 4D tensor `(B, 3, H, W)` with all tiles sharing the same spatial size (legacy
+          fixed-tile path **or** dynamic-resolution path when every image in the batch resizes to
+          the same target).
+        - A list of 4D tensors `[(1, 3, H_i, W_i), …]` when dynamic resolution picks different
+          target sizes per image. Each is run through the ViT independently and the output tokens
+          are concatenated along the sequence dim.
+
+        The patch grid `(h, w)` is computed from the actual input shape, not assumed square — this
+        is required for dynamic resolution where the tile aspect ratio matches the original image.
+        """
+        if isinstance(pixel_values, (list, tuple)):
+            outs = [self._extract_feature_single(pv) for pv in pixel_values]
+            return torch.cat(outs, dim=0)
+        return self._extract_feature_single(pixel_values)
+
+    def _extract_feature_single(self, pixel_values):
+        vit_embeds = self.vision_model(pixel_values).features
+        vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
+        # Compute patch grid from the input tile dims; pixel-shuffle needs the real (h, w).
+        patch_size = self.vision_model.radio_model.model.patch_generator.patch_size
+        B, _, H, W = pixel_values.shape
+        h = H // patch_size
+        w = W // patch_size
+        vit_embeds = vit_embeds.reshape(B, h, w, -1)
+        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
+        vit_embeds = vit_embeds.reshape(B, -1, vit_embeds.shape[-1])
+        vit_embeds = self.mlp1(vit_embeds)
+        return vit_embeds
+
+    def extract_video_feature(self, pixel_values_videos):
+        """
+        Extract features from video frames using the 3D `video_embedder`.
+
+        Consecutive `T = video_temporal_patch_dim` frames are packed into a single temporal patch
+        before the ViT, so the output has `N_frames // T` temporal units (each with the usual number
+        of spatial tokens) instead of one ViT output per frame.
+
+        Implementation trick: RADIO's patch_generator uses a channel-agnostic `Im2Patches` rearrange
+        followed by `self.embedder(patches)`. If we stack the T temporal frames into the channel
+        dim — `(N_frames, C, H, W)` → `(N_frames/T, T·C, H, W)` — the rearrange produces patches of
+        shape `(·, num_patches, T·C·P²)`, which is exactly what `video_embedder` expects. Temporarily
+        swapping `embedder ↔ video_embedder` lets us reuse the full ViT forward without duplicating
+        the transformer blocks, pos-embed handling, cls_token, etc.
+        """
+        pg = self.vision_model.radio_model.model.patch_generator
+        T = self.video_temporal_patch_dim
+        N, C, H, W = pixel_values_videos.shape
+
+        # Pad to a multiple of T by repeating the last frame so frame pairs align cleanly.
+        if N % T != 0:
+            pad = pixel_values_videos[-1:].expand(T - (N % T), -1, -1, -1)
+            pixel_values_videos = torch.cat([pixel_values_videos, pad], dim=0)
+            N = pixel_values_videos.shape[0]
+        num_groups = N // T
+
+        # Stack T frames into the channel dim. `.view` here preserves the (frame,channel) row-major
+        # layout → per-patch feature order is [t=0,c=0..C-1, t=1,c=0..C-1, ...], matching how the
+        # `video_embedder` weights are stored in the checkpoint.
+        x = pixel_values_videos.reshape(num_groups, T * C, H, W)
+
+        orig_embedder = pg.embedder
+        pg.embedder = pg.video_embedder
+        try:
+            vit_embeds = self.vision_model(x).features
+        finally:
+            pg.embedder = orig_embedder
+
+        # Same spatial post-processing as `extract_feature`. Compute `(h, w)` from the reshaped
+        # input so dynamic-res video frames (non-square patch grid) are handled correctly.
+        vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
+        patch_size = pg.patch_size
+        h = H // patch_size
+        w = W // patch_size
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        vit_embeds = self.mlp1(vit_embeds)
+        return vit_embeds
+
+    def extract_sound_feature(
+        self,
+        input_features: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Extract and project sound features from audio input.
+        
+        Args:
+            input_features: Mel spectrogram features [batch, seq_len, feature_dim]
+            attention_mask: Optional attention mask [batch, seq_len]
+            
+        Returns:
+            Sound embeddings projected to LLM hidden size [batch, encoded_seq_len, llm_hidden_size]
+        """
+        if self.sound_encoder is None:
+            raise RuntimeError("Sound encoder not initialized. Check if sound_config is provided.")
+        
+        # Encode audio features
+        sound_embeds = self.sound_encoder(input_features, attention_mask)
+        sound_embeds = sound_embeds.to(dtype=torch.bfloat16)
+        
+        # Project to LLM hidden size
+        sound_embeds = self.sound_projection(sound_embeds)
+        
+        return sound_embeds
+
+    @torch.no_grad()
+    def generate(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            pixel_values_videos: Optional[torch.FloatTensor] = None,
+            sound_clips: Optional[torch.FloatTensor] = None,
+            sound_length: Optional[torch.Tensor] = None,
+            input_ids: Optional[torch.FloatTensor] = None,
+            attention_mask: Optional[torch.LongTensor] = None,
+            generation_config: Optional[GenerationConfig] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            **generate_kwargs,
+    ) -> torch.LongTensor:
+        """Generate text given images, videos, and/or audio.
+        
+        Args:
+            pixel_values: Image pixel values [num_tiles, C, H, W]
+            pixel_values_videos: Video pixel values [num_frames, C, H, W]
+            sound_clips: Raw audio waveforms. Can be:
+                - A list of numpy arrays or torch tensors (one per audio clip)
+                - A single numpy array or torch tensor for a single audio clip
+                - Pre-extracted mel spectrogram features [batch, seq_len, num_mel_bins]
+            sound_length: Length of each audio clip in samples (optional, used for batched audio)
+            input_ids: Input token IDs [batch, seq_len]
+            attention_mask: Attention mask [batch, seq_len]
+            generation_config: Generation configuration
+            output_hidden_states: Whether to output hidden states
+            return_dict: Whether to return a dict
+            **generate_kwargs: Additional generation arguments
+            
+        Returns:
+            Generated token IDs
+        """
+        assert self.img_context_token_id is not None
+        
+        has_images = pixel_values is not None
+        has_videos = pixel_values_videos is not None
+        has_sound = sound_clips is not None and self.sound_encoder is not None
+        
+        if has_images or has_videos or has_sound:
+            image_vit_embeds, video_vit_embeds, sound_embeds = None, None, None
+            
+            # Process images
+            if has_images:
+                pixel_values = pixel_values.to(dtype=self.vision_model.config.torch_dtype)
+                image_vit_embeds = self.extract_feature(pixel_values)
+            
+            # Process videos
+            if has_videos:
+                pixel_values_videos = pixel_values_videos.to(dtype=self.vision_model.config.torch_dtype)
+                video_vit_embeds = self.extract_video_feature(pixel_values_videos)
+            
+            # Process sound/audio
+            if has_sound:
+                # Extract features from raw audio using the feature extractor
+                # Handle different input types:
+                # - list/tuple of waveforms
+                # - 1D tensor/array (single waveform)
+                # - 2D tensor [batch, samples] (batched raw waveforms)
+                # - 3D tensor [batch, seq_len, num_mel_bins] (pre-extracted features)
+                import numpy as np
+                
+                is_raw_waveform = False
+                if isinstance(sound_clips, (list, tuple)):
+                    # List of audio clips (waveforms)
+                    is_raw_waveform = True
+                    waveforms = sound_clips
+                elif isinstance(sound_clips, np.ndarray):
+                    # Numpy array - raw waveform
+                    is_raw_waveform = True
+                    waveforms = [sound_clips.squeeze()] if sound_clips.ndim > 1 else [sound_clips]
+                elif isinstance(sound_clips, torch.Tensor):
+                    if sound_clips.dim() == 1:
+                        # 1D tensor - single raw waveform
+                        is_raw_waveform = True
+                        waveforms = [sound_clips.cpu().numpy()]
+                    elif sound_clips.dim() == 2:
+                        # 2D tensor [batch, samples] - batched raw waveforms
+                        is_raw_waveform = True
+                        waveforms = [clip.cpu().numpy() for clip in sound_clips]
+                    else:
+                        # 3D tensor [batch, seq_len, num_mel_bins] - pre-extracted features
+                        is_raw_waveform = False
+                else:
+                    is_raw_waveform = False
+                
+                if is_raw_waveform:
+                    # Convert raw waveforms to mel spectrogram features
+                    audio_inputs = self.sound_feature_extractor(
+                        waveforms,
+                        sampling_rate=self.sound_feature_extractor.sampling_rate,
+                        return_tensors="pt",
+                    )
+                    sound_input_features = audio_inputs.input_features
+                    sound_attention_mask = audio_inputs.get("attention_mask", None)
+                else:
+                    # Already extracted features
+                    sound_input_features = sound_clips
+                    sound_attention_mask = None
+                
+                # Move to correct device and dtype
+                target_device = self.sound_encoder.encoder.subsampling.linear.weight.device
+                target_dtype = self.language_model.config.torch_dtype
+                
+                sound_input_features = sound_input_features.to(dtype=target_dtype, device=target_device)
+                if sound_attention_mask is not None:
+                    sound_attention_mask = sound_attention_mask.to(device=target_device)
+                
+                sound_embeds = self.extract_sound_feature(sound_input_features, sound_attention_mask)
+            
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+            B, N, C = inputs_embeds.shape
+            inputs_embeds = inputs_embeds.reshape(B * N, C)
+            input_ids_copy = input_ids.reshape(B * N)
+            
+            # Replace image tokens with image embeddings
+            if image_vit_embeds is not None:
+                image_mask = (input_ids_copy == self.img_context_token_id)
+                assert image_mask.sum() != 0, "No image tokens found in input_ids"
+                inputs_embeds[image_mask] = image_vit_embeds.reshape(-1, C).to(inputs_embeds.device, inputs_embeds.dtype)
+            
+            # Replace video tokens with video embeddings. The tokenizer has no distinct `<video>`
+            # token (`video_context_token_id` in config doesn't decode to any printable string), so
+            # the processor uses `<image>` (id = `img_context_token_id`) as the placeholder for
+            # video positions too. We rely on the caller passing `pixel_values_videos` (not
+            # `pixel_values`) to signal video vs. image — both share the same token id in the prompt.
+            if video_vit_embeds is not None:
+                if B > 1:
+                    raise NotImplementedError("Video is not supported for batch size > 1")
+                video_mask = (input_ids_copy == self.img_context_token_id)
+                assert video_mask.sum() != 0, "No video tokens found in input_ids"
+                inputs_embeds[video_mask] = video_vit_embeds.reshape(-1, C).to(inputs_embeds.device, inputs_embeds.dtype)
+            
+            # Replace sound tokens with sound embeddings.
+            # `sound_embeds` has shape `(B_sound, T_out_max, C)` where `T_out_max`
+            # is the encoder output length for the longest clip in the batch.
+            # When `B_sound > 1` the shorter clips have padding at the tail, so
+            # we must gather only the valid positions per row before scattering
+            # into `sound_mask`. The encoder's `_get_subsampling_output_length`
+            # converts each input mel-frame count (from the feature extractor's
+            # attention_mask) to its post-subsampling token count.
+            if sound_embeds is not None and self.sound_context_token_id is not None:
+                sound_mask = (input_ids_copy == self.sound_context_token_id)
+                assert sound_mask.sum() != 0, "No sound tokens found in input_ids"
+                if sound_embeds.dim() == 3 and sound_embeds.shape[0] > 1 and sound_attention_mask is not None:
+                    # `attention_mask.sum() = L_i // hop` per row, but
+                    # `ParakeetFeatureExtractor` pads each row to `1 + L_i // hop`
+                    # mel frames in single-call mode (the trailing frame comes
+                    # from STFT center padding) — and the existing batch=1 path
+                    # consumes that frame's embed too. Add 1 here to match.
+                    natural_input_lengths = sound_attention_mask.sum(-1) + 1
+                    output_lengths = self.sound_encoder.encoder._get_subsampling_output_length(natural_input_lengths)
+                    flat = torch.cat(
+                        [sound_embeds[i, : int(n)] for i, n in enumerate(output_lengths.tolist())],
+                        dim=0,
+                    )
+                else:
+                    flat = sound_embeds.reshape(-1, C)
+                assert sound_mask.sum().item() == flat.shape[0], (
+                    f"sound token count ({sound_mask.sum().item()}) != encoder output count ({flat.shape[0]})"
+                )
+                inputs_embeds[sound_mask] = flat.to(inputs_embeds.device, inputs_embeds.dtype)
+            
+            # Apply video pruning (EVS) if enabled
+            if video_vit_embeds is not None and self.video_pruning_rate > 0:  # EVS
+                h = w = int(video_vit_embeds.shape[1] ** 0.5)  # assumption here (and everywhere else) is that shape is square
+                evs_mask = EfficientVideoSampling.compute_retention_mask(
+                    video_embeds=video_vit_embeds,
+                    thw=(video_vit_embeds.shape[0], h, w),
+                    spatial_merge_size=1,  # we already work on vision embeddings, so no downsampling to follow
+                    q=self.video_pruning_rate,
+                )
+                print(f"pruning rate: {self.video_pruning_rate}, EVS mask: {evs_mask.sum().item()} tokens retained out of {evs_mask.numel()} total video tokens ({evs_mask.sum().item() / evs_mask.numel() * 100:.2f}%)")
+
+                retention_mask = torch.ones_like(input_ids_copy, dtype=torch.bool)
+                retention_mask[video_mask] = evs_mask.view(-1)
+                inputs_embeds = inputs_embeds[retention_mask].unsqueeze(0)  # adding batch=1
+                if attention_mask is not None:
+                    attention_mask = attention_mask[:, retention_mask].contiguous()
+                if input_ids is not None:
+                    input_ids = input_ids[:, retention_mask].contiguous()
+            else:
+                inputs_embeds = inputs_embeds.reshape(B, N, C)
+        else:
+            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+        
+        outputs = self.language_model.generate(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            use_cache=True,
+            **generate_kwargs,
+        )
+
+        return outputs

modeling_nemotron_h.py

@@ -0,0 +1,1368 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/nemotron_h/modular_nemotron_h.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_nemotron_h.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. 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.
+
+
+import contextlib
+import math
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+import copy
+
+from transformers import initialization as init
+from transformers.activations import ACT2FN
+from transformers.generation import GenerationMixin
+from transformers.integrations import (
+    lazy_load_kernel,
+    use_kernel_forward_from_hub,
+    use_kernel_func_from_hub,
+    use_kernelized_func,
+)
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.models.zamba2.modeling_zamba2 import Zamba2RMSNormGated
+from transformers.processing_utils import Unpack
+from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.import_utils import resolve_internal_import
+from transformers.utils.output_capturing import capture_outputs
+from .configuration_nemotron_h import NemotronHConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class NemotronHHybridDynamicCache:
+    """
+    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)`.
+    """
+
+    is_compileable = False
+
+    def __init__(
+        self, config: NemotronHConfig, batch_size: int, dtype: torch.dtype = torch.float16, device: str | None = None
+    ):
+        self.dtype = dtype
+        self.layers_block_type = config.layers_block_type
+        self.has_previous_state = False
+        self.intermediate_size = int(config.mamba_num_heads * config.mamba_head_dim)
+        self.ssm_state_size = config.ssm_state_size
+        self.conv_kernel_size = config.conv_kernel
+        self.n_mamba_heads = config.mamba_num_heads
+        self.transformer_layers = []
+        self._modules = {}
+        self._parameters = {}
+        self._buffers = {}
+        self.conv_states = {}
+        self.ssm_states = {}
+        for i in range(config.num_hidden_layers):
+            if self.layers_block_type[i] == "mamba":
+                # Only allocate mamba cache for mamba layers
+                self.conv_states[i] = torch.zeros(
+                    batch_size,
+                    self.intermediate_size + 2 * config.n_groups * self.ssm_state_size,
+                    self.conv_kernel_size,
+                    device=device,
+                    dtype=dtype,
+                )
+                self.ssm_states[i] = torch.zeros(
+                    batch_size,
+                    self.n_mamba_heads,
+                    config.mamba_head_dim,
+                    self.ssm_state_size,
+                    device=device,
+                    dtype=dtype,
+                )
+            else:
+                # For attention and moe layers, use empty tensors
+                self.conv_states[i] = torch.tensor([[]] * batch_size, device=device)
+                self.ssm_states[i] = torch.tensor([[]] * batch_size, device=device)
+
+            if self.layers_block_type[i] == "attention":
+                self.transformer_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 __len__(self):
+        return len(self.key_cache)
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: dict[str, Any] | None = 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."""
+        if self.get_seq_length() > 0:
+            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: int | None = 0) -> int:
+        """Returns the sequence length of the cached states. A layer index can be optionally passed."""
+        # take any layer that contains cache and not empty tensor
+        layer_idx = self.transformer_layers[0] if layer_idx not in self.transformer_layers else layer_idx
+        if len(self.key_cache) <= layer_idx or self.key_cache[layer_idx].numel() == 0:
+            return 0
+        return self.key_cache[layer_idx].shape[-2]
+
+    def get_mask_sizes(self, query_length, layer_idx: int) -> tuple[int, int]:
+        """Return the length and offset of the cache, used to generate the mask"""
+        # transformers >= 5.6 passes `query_length` as int; earlier versions passed a `cache_position` tensor.
+        if isinstance(query_length, torch.Tensor):
+            query_length = query_length.shape[0]
+        kv_offset = 0
+        kv_length = self.get_seq_length(layer_idx) + query_length
+        return kv_length, kv_offset
+
+    def update_conv_state(
+        self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor
+    ) -> torch.Tensor:
+        conv_state = self.conv_states[layer_idx]
+        cache_position = cache_position.clamp(0, self.conv_kernel_size - 1)
+
+        conv_state = conv_state.roll(shifts=-1, dims=-1)
+        conv_state[:, :, cache_position] = new_conv_state.to(conv_state.device)
+        self.conv_states[layer_idx].zero_()
+        self.conv_states[layer_idx] += conv_state
+        return self.conv_states[layer_idx]
+
+    def reset(self):
+        self.conv_states.zero_()
+        self.ssm_states.zero_()
+
+
+# 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
+
+
+class NemotronHMamba2Mixer(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: NemotronHConfig, layer_idx: int | None = None):
+        super().__init__()
+        self.config = config
+        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.use_mem_eff_path = True
+
+        self.n_groups = config.n_groups
+        self.head_dim = config.mamba_head_dim
+        self.num_heads = config.mamba_num_heads
+        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=self.conv_kernel_size,
+            groups=self.conv_dim,
+            padding=self.conv_kernel_size - 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 dependent
+
+        # 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.norm = Zamba2RMSNormGated(
+            self.intermediate_size, group_size=self.intermediate_size // self.n_groups, eps=config.layer_norm_epsilon
+        )
+        self.D = nn.Parameter(torch.ones(self.num_heads))
+
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias)
+
+        global causal_conv1d_update, causal_conv1d_fn
+        causal_conv1d = lazy_load_kernel("causal-conv1d")
+        causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
+        causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)
+
+        global selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
+        mamba_ssm = lazy_load_kernel("mamba-ssm")
+        selective_state_update = resolve_internal_import(
+            mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
+        )
+        mamba_chunk_scan_combined = resolve_internal_import(
+            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
+        )
+        mamba_split_conv1d_scan_combined = resolve_internal_import(
+            mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
+        )
+
+        global is_fast_path_available
+        is_fast_path_available = all(
+            (
+                selective_state_update,
+                mamba_chunk_scan_combined,
+                mamba_split_conv1d_scan_combined,
+                causal_conv1d_fn,
+                causal_conv1d_update,
+            )
+        )
+
+        if not is_fast_path_available:
+            logger.warning_once(
+                "The fast path is not available because one 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 cuda_kernels_forward(
+        self,
+        hidden_states: torch.Tensor,
+        cache_params: NemotronHHybridDynamicCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        # 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_to_remove = 2 * self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.num_heads
+
+        # getting projected states from cache if it exists
+        if cache_params is not None and cache_params.has_previous_state:
+            in_projected_states = self.in_proj(hidden_states.squeeze(1))  # (B 2D)
+            d_mlp = (in_projected_states.shape[-1] - d_to_remove) // 2
+            split_projection_dim = [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads]
+            _, _, gate, hidden_states_B_C, dt = torch.split(in_projected_states, split_projection_dim, dim=-1)
+
+            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,
+            )
+            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)
+            out = self.out_proj(hidden_states)[:, None, ...]
+        # if no cache is found, calling the kernel
+        else:
+            if attention_mask is not None and not torch.all(attention_mask == 1):
+                # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+                dtype = hidden_states.dtype
+                hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+            # 1. Gated MLP's linear projection
+            projected_states = self.in_proj(hidden_states)
+            A = -torch.exp(self.A_log.float())  # (num_heads) or (intermediate_size, state_size)
+            dt_limit_kwargs = {} if self.time_step_limit is None else {"dt_limit": self.time_step_limit}
+            if attention_mask is not None:
+                input_not_masked = torch.all(attention_mask == 1)
+            else:
+                input_not_masked = True
+
+            if self.use_mem_eff_path and self.training and cache_params is None and input_not_masked:
+                out, ssm_state = 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=None,
+                    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=True,
+                    **dt_limit_kwargs,
+                )
+
+            else:
+                gate, hidden_states_B_C, time_step = torch.split(
+                    projected_states,
+                    [self.intermediate_size, self.conv_dim, self.num_heads],
+                    dim=-1,
+                )
+
+                # 1D Convolution
+                if cache_params is not None:
+                    hidden_states_B_C_t = hidden_states_B_C.transpose(1, 2)
+                    conv_state = nn.functional.pad(
+                        hidden_states_B_C_t, (self.conv_kernel_size - hidden_states_B_C_t.shape[-1], 0)
+                    )
+                    cache_params.conv_states[self.layer_idx].copy_(conv_state)
+                if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:
+                    hidden_states_B_C = self.act(
+                        self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[:, :seq_len]
+                    )  # (B, L, self.d_inner + 2 * ngroups * d_state)
+                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)[:, :seq_len]
+                hidden_states, B, C = torch.split(
+                    hidden_states_B_C,
+                    [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+                    dim=-1,
+                )
+                if attention_mask is not None and not torch.all(attention_mask == 1):
+                    # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+                    dtype = hidden_states.dtype
+                    hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+                scan_output, ssm_state = mamba_chunk_scan_combined(
+                    hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+                    time_step,
+                    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=None,
+                    return_final_states=True,
+                    dt_bias=self.dt_bias,
+                    dt_softplus=True,
+                    **dt_limit_kwargs,
+                )
+                if ssm_state is not None and cache_params is not None:
+                    cache_params.ssm_states[self.layer_idx].copy_(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)
+                out = self.out_proj(scan_output)
+        return out
+
+    # fmt: off
+    def torch_forward(self, input_states, cache_params: NemotronHHybridDynamicCache | None=None, attention_mask: torch.Tensor | None=None):
+        batch_size, seq_len, _ = input_states.shape
+        dtype = input_states.dtype
+        # Gated MLP's linear projection
+        if cache_params is not None and cache_params.has_previous_state:
+            projected_states = self.in_proj(input_states.squeeze(1))
+        else:
+            if attention_mask is not None:
+                # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+                input_states = (input_states * attention_mask[:, :, None]).to(dtype)
+            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, dt = projected_states.split(
+                [d_mlp, d_mlp, self.intermediate_size,  self.conv_dim, self.num_heads], dim=-1
+        )
+
+        # Convolution sequence transformation
+        if cache_params is not None:
+            ssm_state = cache_params.ssm_states[self.layer_idx].clone()
+            ssm_state = ssm_state.to(hidden_states.device)
+            if cache_params.has_previous_state:
+                gate = gate.unsqueeze(1)
+                conv_state = cache_params.conv_states[self.layer_idx]                   # [batch, intermediate_size, conv_kernel_size]
+                conv_state = torch.roll(conv_state, shifts=-1, dims=-1)
+                # handle batched generation - states are copied through
+                conv_state[:, :, -1] = hidden_states[:, 0, :] if hidden_states.ndim == 3 else hidden_states
+                cache_params.conv_states[self.layer_idx].copy_(conv_state)
+                hidden_states = torch.sum(conv_state.to(projected_states.device) * self.conv1d.weight[:, 0, :], dim=-1)
+                if self.use_conv_bias:
+                    hidden_states += self.conv1d.bias
+                hidden_states = self.act(hidden_states).to(dtype)[:, None, ...]         # [batch, 1, intermediate_size] : decoding
+            else:
+                hidden_states = hidden_states.transpose(1,2)
+                conv_state = nn.functional.pad(
+                    hidden_states,
+                    (self.conv_kernel_size - hidden_states.shape[-1], 0)
+                )
+                cache_params.conv_states[self.layer_idx].copy_(conv_state)
+                hidden_states = self.act(self.conv1d(hidden_states).transpose(1,2))[:, :seq_len, :]     # [batch, intermediate_size, seq_len]
+                if attention_mask is not None:
+                    dtype = hidden_states.dtype
+                    # 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)
+        else:
+            ssm_state = torch.zeros(
+                (batch_size, self.num_heads, self.head_dim, self.ssm_state_size),
+                device=hidden_states.device, dtype=dtype
+            )
+            hidden_states = self.act(self.conv1d(hidden_states.transpose(1, 2))[..., :seq_len].transpose(1, 2))
+        hidden_states, B, C = torch.split(hidden_states, [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1)
+        A = -torch.exp(self.A_log.float())                            # [num_heads]
+        if cache_params is not None and cache_params.has_previous_state:
+            # Note: there is no need to pad parameter matrices here, as there is just one new token
+            # for batched generation
+            dt = dt[:, None, ...] if dt.ndim == 2 else 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_min) #, self.time_step_max)
+            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)
+
+            # 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]
+
+            # State calculation
+            cache_params.ssm_states[self.layer_idx].copy_(
+                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(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_min)
+            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))
+
+            # First, 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)
+
+
+            # Step 2: 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)
+
+            # Step 3: Compute Y_diag (apply to values)
+            Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(3)
+
+            # (right term of low-rank factorization of off-diagonal blocks; B terms)
+
+            decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
+            B_decay_contraction = B * decay_states.permute(0, 2, 3, 1)[..., None]
+            # permute back B * decay states
+            states = (B_decay_contraction.permute(0, 1, 3, 2, 4)[..., None]  * hidden_states.permute(0, 1, 3, 2, 4)[..., None, :]).sum(dim=3).permute(0, 1, 2, 4, 3)
+            if cache_params is not None and cache_params.has_previous_state:
+                previous_states = cache_params.ssm_states[self.layer_idx][:, None, ...]
+            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))))
+
+            states_permuted = states.permute(0, 2, 1, 3, 4)
+            result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2)
+            new_states = result.permute(0, 2, 1, 3, 4)
+            states, ssm_state = new_states[:, :-1], new_states[:, -1]
+
+            # 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)
+            # compute Yoff
+            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)
+            if ssm_state is not None and cache_params is not None:
+                cache_params.ssm_states[self.layer_idx].copy_(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: NemotronHHybridDynamicCache | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
+            return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
+
+        return self.torch_forward(hidden_states, cache_params, attention_mask)
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class NemotronHRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        NemotronHRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        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)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class NemotronHMLP(nn.Module):
+    def __init__(self, config, intermediate_size=None):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size or config.intermediate_size
+        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 forward(self, x):
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+
+class NemotronHMoE(nn.Module):
+    """
+    Mixture-of-Experts (MoE) module for NemotronH.
+
+    - Experts are stored as an nn.ModuleList of NemotronHMLP, to match the per-expert checkpoint
+      format (`experts.<idx>.up_proj.weight` / `experts.<idx>.down_proj.weight`).
+    - Optional latent projection wraps the experts.
+    """
+
+    def __init__(self, config, layer_idx: int | None = None):
+        super().__init__()
+        self.config = config
+        self.n_routed_experts = config.n_routed_experts
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.norm_topk_prob = config.norm_topk_prob
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.top_k = config.num_experts_per_tok
+
+        self.gate = NemotronHTopkRouter(config)
+
+        # Optional latent projection; `moe_latent_size` is absent on older configs.
+        moe_latent_size = getattr(config, "moe_latent_size", None)
+        if moe_latent_size is not None:
+            self.fc1_latent_proj = nn.Linear(config.hidden_size, moe_latent_size, bias=config.mlp_bias)
+            self.fc2_latent_proj = nn.Linear(moe_latent_size, config.hidden_size, bias=config.mlp_bias)
+            expert_input_dim = moe_latent_size
+        else:
+            self.fc1_latent_proj = nn.Identity()
+            self.fc2_latent_proj = nn.Identity()
+            expert_input_dim = config.hidden_size
+
+        # Shallow-copy config so each expert's NemotronHMLP runs in `expert_input_dim` dim without
+        # mutating the top-level config.
+        expert_config = copy.copy(config)
+        expert_config.hidden_size = expert_input_dim
+        self.experts = nn.ModuleList(
+            [NemotronHMLP(expert_config, intermediate_size=config.moe_intermediate_size) for _ in range(self.n_routed_experts)]
+        )
+        self.shared_experts = NemotronHMLP(config=config, intermediate_size=config.moe_shared_expert_intermediate_size)
+
+    def route_tokens_to_experts(self, router_logits):
+        router_logits = router_logits.sigmoid()
+        router_logits_for_choice = router_logits + self.gate.e_score_correction_bias
+        group_scores = (
+            router_logits_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group)
+            .topk(2, dim=-1)[0]
+            .sum(dim=-1)
+        )
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(-1, self.n_group, self.n_routed_experts // self.n_group)
+            .reshape(-1, self.n_routed_experts)
+        )
+        scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0)
+        topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+        topk_weights = router_logits.gather(1, topk_indices)
+        if self.norm_topk_prob:
+            denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
+            topk_weights /= denominator
+        topk_weights = topk_weights * self.routed_scaling_factor
+        return topk_indices, topk_weights
+
+    def forward(self, hidden_states):
+        residuals = hidden_states
+        orig_shape = hidden_states.shape
+        router_logits = self.gate(hidden_states)
+        topk_indices, topk_weights = self.route_tokens_to_experts(router_logits)
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+
+        # Route each token through its top-k experts (ModuleList version).
+        expert_inputs = self.fc1_latent_proj(hidden_states)
+        expert_outputs = torch.zeros_like(expert_inputs, dtype=topk_weights.dtype)
+
+        with torch.no_grad():
+            expert_mask = torch.nn.functional.one_hot(topk_indices, num_classes=self.n_routed_experts)
+            expert_mask = expert_mask.permute(2, 1, 0)  # (num_experts, top_k, num_tokens)
+            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero().squeeze(-1)
+
+        for expert_idx in expert_hit:
+            expert_idx = expert_idx.item()
+            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+            if token_idx.numel() == 0:
+                continue
+            current_state = expert_inputs[token_idx]
+            current_hidden_states = self.experts[expert_idx](current_state)
+            current_hidden_states = current_hidden_states * topk_weights[token_idx, top_k_pos, None]
+            expert_outputs.index_add_(0, token_idx, current_hidden_states.to(expert_outputs.dtype))
+
+        expert_outputs = expert_outputs.to(expert_inputs.dtype)
+        hidden_states = self.fc2_latent_proj(expert_outputs)
+
+        hidden_states = hidden_states.view(*orig_shape)
+        hidden_states = hidden_states + self.shared_experts(residuals)
+        return hidden_states
+
+
+class NemotronHTopkRouter(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.n_routed_experts = config.n_routed_experts
+
+        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size)))
+        self.register_buffer("e_score_correction_bias", torch.zeros(self.n_routed_experts))
+
+    def forward(self, hidden_states):
+        hidden_states = hidden_states.view(-1, self.config.hidden_size)
+        router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
+        return router_logits
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class NemotronHAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: NemotronHConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: NemotronHHybridDynamicCache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        if past_key_values is not None:
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, {"cache_position": cache_position}
+            )
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+MIXER_TYPES = {
+    "mamba": NemotronHMamba2Mixer,
+    "attention": NemotronHAttention,
+    "moe": NemotronHMoE,
+}
+
+
+class NemotronHBlock(GradientCheckpointingLayer):
+    """
+    A single transformer block in the NemotronH model.
+
+    This block can contain different types of mixers (Mamba, Attention, MLP, or MoE)
+    depending on the configuration. Each block applies pre-normalization followed by
+    the mixer, then adds a residual connection.
+
+    Args:
+        config (`NemotronHConfig`):
+            Model configuration specifying the block architecture.
+        layer_idx (`int`):
+            Index of this block in the model. Used to determine the block type from
+            `config.layers_block_type[layer_idx]`.
+    """
+
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.norm = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+
+        self.block_type = config.layers_block_type[layer_idx]
+        self.mixer = MIXER_TYPES[self.block_type](config, layer_idx=layer_idx)
+
+    def forward(
+        self,
+        hidden_states,
+        past_key_values: NemotronHHybridDynamicCache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        use_cache: bool | None = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        if hidden_states.device.type == "cuda":
+            # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue.
+            # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream.
+            # This leads to kernel reading uninitialized memory before the data transfer is complete.
+            stream_context = torch.cuda.stream(torch.cuda.default_stream(hidden_states.device))
+        else:
+            stream_context = contextlib.nullcontext()
+
+        with stream_context:
+            residual = hidden_states
+            hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype))
+
+            if self.block_type == "mamba":
+                hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask)
+            elif self.block_type == "attention":
+                hidden_states, _ = self.mixer(
+                    hidden_states=hidden_states,
+                    past_key_values=past_key_values,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    user_cache=use_cache,
+                    cache_position=cache_position,
+                    **kwargs,
+                )
+            else:
+                hidden_states = self.mixer(hidden_states)
+
+            hidden_states = residual + hidden_states
+            return hidden_states
+
+
+class NemotronHPreTrainedModel(PreTrainedModel):
+    config: NemotronHConfig
+    base_model_prefix = "backbone"
+    _no_split_modules = ["NemotronHBlock"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _is_stateful = True
+    _can_record_outputs = {
+        "hidden_states": NemotronHBlock,
+        "attentions": NemotronHAttention,
+    }
+    _keep_in_fp32_modules_strict = [
+        "e_score_correction_bias",
+    ]
+    _tied_weights_keys = {}
+    _keys_to_ignore_on_load_unexpected = [r"mtp.*"]
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights."""
+        super()._init_weights(module)
+        if isinstance(module, NemotronHMamba2Mixer):
+            # Initialize A_log and D parameters
+            A = torch.arange(1, self.config.mamba_num_heads + 1)
+            init.copy_(module.A_log, torch.log(A))
+            init.ones_(module.D)
+
+            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():
+                init.copy_(module.dt_bias, inv_dt)
+            module.dt_bias._no_reinit = True
+        elif isinstance(module, NemotronHTopkRouter):
+            init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            init.zeros_(module.e_score_correction_bias)
+        # Experts are now an nn.ModuleList of NemotronHMLP; their nn.Linear submodules are handled by
+        # the generic branch below, so no special init branch is needed.
+
+        if isinstance(module, nn.Linear):
+            if module.bias is not None:
+                if not getattr(module.bias, "_no_reinit", False):
+                    init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            init.normal_(module.weight, std=self.config.initializer_range)
+
+        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 == "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
+                    init.kaiming_uniform_(p, a=math.sqrt(5))
+                    with torch.no_grad():
+                        p_new = p / math.sqrt(self.config.num_hidden_layers)
+                        init.copy_(p, p_new)
+
+
+class NemotronHModel(NemotronHPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([NemotronHBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)])
+
+        self.norm_f = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon)
+
+        # Legacy checkpoints may use "embedding." (singular); rewrite to "embeddings." on load.
+        self._register_load_state_dict_pre_hook(self.load_hook)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def load_hook(self, state_dict, prefix, *args):
+        for k in list(state_dict.keys()):
+            if "embedding." in k:
+                state_dict[k.replace("embedding.", "embeddings.")] = state_dict.pop(k)
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embeddings = new_embeddings
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: NemotronHHybridDynamicCache | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPast:
+        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 use_cache and past_key_values is None:
+            past_key_values = NemotronHHybridDynamicCache(
+                config=self.config,
+                batch_size=inputs_embeds.shape[0],
+                dtype=inputs_embeds.dtype,
+                device=inputs_embeds.device,
+            )
+
+        hidden_states = inputs_embeds
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + hidden_states.shape[1], device=hidden_states.device
+            )
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+        mamba_mask = self._update_mamba_mask(attention_mask, cache_position)
+
+        # Map block types to their corresponding masks
+        block_type_to_mask = {
+            "mamba": mamba_mask,
+            "attention": causal_mask,
+            "moe": None,
+        }
+
+        for layer_idx, mixer_block in enumerate(self.layers):
+            layer_mask = block_type_to_mask[mixer_block.block_type]
+
+            hidden_states = mixer_block(
+                hidden_states,
+                attention_mask=layer_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+        hidden_states = self.norm_f(hidden_states)
+
+        if past_key_values is not None and not past_key_values.has_previous_state:
+            past_key_values.has_previous_state = True
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+    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 is not None and cache_position[0] > 0) or (
+            attention_mask is not None and torch.all(attention_mask == 1)
+        ):
+            mamba_mask = None
+        return mamba_mask
+
+
+# Adapted from transformers.models.jamba.modeling_jamba.JambaForCausalLM with Jamba->NemotronH, JAMBA->NEMOTRON_H
+class NemotronHForCausalLM(NemotronHPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {}
+
+    @classmethod
+    def _supports_default_dynamic_cache(cls) -> bool:
+        # This model supplies its own `NemotronHHybridDynamicCache` via `prepare_inputs_for_generation`.
+        # Preventing transformers from injecting a plain `DynamicCache` keeps our mixer code
+        # (which reads `cache_params.conv_states[layer_idx]` / `.has_previous_state`) happy.
+        return False
+
+    def __init__(self, config):
+        super().__init__(config)
+        # Attribute name matches the checkpoint key prefix ("backbone.").
+        self.backbone = NemotronHModel(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)
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: NemotronHHybridDynamicCache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs,
+    ) -> tuple | CausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, NemotronHForCausalLM
+
+        >>> model = NemotronHForCausalLM.from_pretrained("Zyphra/NemotronH-7B-v1")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Zyphra/NemotronH-7B-v1")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs = self.backbone(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        # Overwritten -- has a unique cache type, `NemotronHHybridDynamicCache`
+
+        if past_key_values is None:
+            past_key_values = NemotronHHybridDynamicCache(
+                self.config, input_ids.shape[0], dtype=self.dtype, device=self.device
+            )
+
+        kwargs["logits_to_keep"] = self.config.num_logits_to_keep
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            position_ids=position_ids,
+            use_cache=use_cache,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+        return model_inputs
+
+
+__all__ = ["NemotronHPreTrainedModel", "NemotronHModel", "NemotronHForCausalLM"]

preprocessor_config.json

@@ -0,0 +1,15 @@
+{
+  "image_processor_type": "NemotronH_Nano_Omni_Reasoning_V3ImageProcessor",
+  "auto_map": {
+    "AutoImageProcessor": "image_processing.NemotronH_Nano_Omni_Reasoning_V3ImageProcessor",
+    "AutoVideoProcessor": "video_processing.NemotronH_Nano_Omni_Reasoning_V3VideoProcessor",
+    "AutoProcessor": "processing.NemotronH_Nano_Omni_Reasoning_V3Processor"
+  },
+  "patch_size": 16,
+  "downsample_ratio": 0.5,
+  "norm_mean": [0.48145466, 0.4578275, 0.40821073],
+  "norm_std": [0.26862954, 0.26130258, 0.27577711],
+  "min_num_patches": 1024,
+  "max_num_patches": 13312,
+  "max_model_len": 16384
+}

processing.py

@@ -0,0 +1,509 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  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.
+from typing import Optional, Union, List
+import math
+import numpy as np
+import torch
+
+from transformers.feature_extraction_utils import BatchFeature
+from transformers.image_utils import ImageInput
+from transformers.processing_utils import ImagesKwargs, MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack, VideosKwargs
+from transformers.tokenization_utils_base import PreTokenizedInput, TextInput
+from transformers.video_utils import VideoInput
+
+# Audio input type - can be file paths, numpy arrays, or torch tensors
+AudioInput = Union[str, np.ndarray, torch.Tensor, List[str], List[np.ndarray], List[torch.Tensor]]
+
+
+class NemotronH_Nano_Omni_Reasoning_V3ImagesKwargs(ImagesKwargs):
+    min_pixels: Optional[int]
+    max_pixels: Optional[int]
+    patch_size: Optional[int]
+    temporal_patch_size: Optional[int]
+    merge_size: Optional[int]
+
+
+class NemotronH_Nano_Omni_Reasoning_V3AudioKwargs(ProcessingKwargs, total=False):
+    sampling_rate: Optional[int]
+
+
+class NemotronH_Nano_Omni_Reasoning_V3ProcessorKwargs(ProcessingKwargs, total=False):
+    images_kwargs: NemotronH_Nano_Omni_Reasoning_V3ImagesKwargs
+    videos_kwargs: VideosKwargs
+    audio_kwargs: NemotronH_Nano_Omni_Reasoning_V3AudioKwargs
+    _defaults = {
+        "text_kwargs": {
+            "padding": False,
+        },
+    }
+
+
+class NemotronH_Nano_Omni_Reasoning_V3Processor(ProcessorMixin):
+    r"""
+    Constructs a Nemotron-3-Nano-Omni-30B-A3B-Reasoning processor which wraps an image processor, audio feature extractor,
+    and a tokenizer into a single processor.
+    [`NemotronH_Nano_Omni_Reasoning_V3Processor`] offers all the functionalities of the image processor, audio processor,
+    and tokenizer. See the [`~NemotronH_Nano_Omni_Reasoning_V3Processor.__call__`] and [`~NemotronH_Nano_Omni_Reasoning_V3Processor.decode`]
+    for more information.
+    Args:
+        image_processor ([`AutoImageProcessor`], *optional*):
+            The image processor is a required input.
+        tokenizer ([`AutoTokenizer`], *optional*):
+            The tokenizer is a required input.
+        chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+            in a chat into a tokenizable string.
+        audio_sampling_rate (`int`, *optional*): Sampling rate for audio processing (default: 16000).
+        audio_subsampling_factor (`int`, *optional*): Subsampling factor for audio encoder (default: 8).
+        audio_hop_length (`int`, *optional*): Hop length in samples for feature extraction (default: 160).
+    """
+
+    attributes = ["image_processor", "tokenizer"]
+
+    image_processor_class = "AutoImageProcessor"
+    video_processor_class = "AutoVideoProcessor"
+    tokenizer_class = ("AutoTokenizer")
+
+    def __init__(
+        self,
+        image_processor=None,
+        tokenizer=None,
+        chat_template=None,
+        audio_sampling_rate: int = 16000,
+        audio_subsampling_factor: int = 8,
+        audio_hop_length: int = 160,
+        video_temporal_patch_dim: int = 2,
+        **kwargs
+    ):
+        # Number of frames collapsed into a single temporal patch by the model's `video_embedder`.
+        # The `<video>` expansion below issues one placeholder block per temporal patch.
+        self.video_temporal_patch_dim = video_temporal_patch_dim
+        self.image_token = "<image>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
+        self.video_token = "<video>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
+        self.audio_token = "<so_embedding>" if not hasattr(tokenizer, "audio_token") else tokenizer.audio_token
+        self.audio_start_token = "<so_start>"
+        self.audio_end_token = "<so_end>"
+        self.image_start_token = "<img>" if not hasattr(tokenizer, "image_start_token") else tokenizer.image_start_token
+        self.image_end_token = "</img>" if not hasattr(tokenizer, "image_end_token") else tokenizer.image_end_token
+        self.image_token_id = (
+            tokenizer.image_token_id
+            if getattr(tokenizer, "image_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.image_token)
+        )
+        self.video_token_id = (
+            tokenizer.video_token_id
+            if getattr(tokenizer, "video_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.video_token)
+        )
+        self.audio_token_id = (
+            tokenizer.audio_token_id
+            if getattr(tokenizer, "audio_token_id", None)
+            else tokenizer.convert_tokens_to_ids(self.audio_token)
+        )
+
+        # Audio processing parameters
+        self.audio_sampling_rate = audio_sampling_rate
+        self.audio_subsampling_factor = audio_subsampling_factor
+        self.audio_hop_length = audio_hop_length
+
+        super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+    def __call__(
+        self,
+        images: ImageInput = None,
+        text: Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]] = None,
+        videos: VideoInput = None,
+        audio: AudioInput = None,
+        **kwargs: Unpack[NemotronH_Nano_Omni_Reasoning_V3ProcessorKwargs],
+    ) -> BatchFeature:
+        """
+        Main method to prepare multimodal inputs (text, images, videos, audio) for the model. This method processes
+        text by replacing image/video/audio tokens with appropriate placeholder sequences, processes images and videos
+        through the image processor, and tokenizes the final text.
+
+        The method performs the following key operations:
+        1. Processes images using the image processor to get pixel values and patch counts
+        2. Processes videos using the image processor with max_num_tiles=1 to get video pixel values
+        3. Processes audio to compute the number of audio tokens based on duration
+        4. Replaces `<image>` tokens in text with `<img>` + image tokens + `</img>` sequences
+        5. Replaces `<video>` tokens in text with frame-by-frame descriptions including timestamps (if metadata provided)
+        6. Replaces `<audio>` tokens in text with repeated audio tokens based on duration
+        7. Tokenizes the processed text and combines all outputs
+
+        Args:
+            images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):
+                The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+                tensor. Both channels-first and channels-last formats are supported.
+            text (`str`, `List[str]`, *optional*):
+                The sequence or batch of sequences to be encoded. Each sequence should be a string. The text can contain
+                special tokens `<image>`, `<video>`, and `<audio>` that will be replaced with appropriate token sequences.
+            videos (`np.ndarray`, `torch.Tensor`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):
+                The video or batch of videos to be prepared. Each video should be a 4D NumPy array or PyTorch
+                tensor with shape (num_frames, channels, height, width). Both channels-first and channels-last formats
+                are supported. Note: Currently only supports batch size of 1 for videos.
+            audio (`str`, `np.ndarray`, `torch.Tensor`, `List[str]`, `List[np.ndarray]`, `List[torch.Tensor]`, *optional*):
+                The audio or batch of audio clips to be prepared. Can be file paths, numpy arrays (waveforms),
+                or torch tensors. Waveforms should be 1D arrays at the expected sampling rate.
+            images_kwargs (`Dict`, *optional*):
+                Additional keyword arguments for image processing, including:
+                - `min_pixels` (`int`, *optional*): Minimum number of pixels for image processing
+                - `max_pixels` (`int`, *optional*): Maximum number of pixels for image processing
+                - `patch_size` (`int`, *optional*): Size of patches for image processing
+                - `temporal_patch_size` (`int`, *optional*): Size of temporal patches
+                - `merge_size` (`int`, *optional*): Size for merging patches
+            videos_kwargs (`Dict`, *optional*):
+                Additional keyword arguments for video processing, including:
+                - `video_metadata` (`VideoMetadata`, *optional*): Metadata containing fps information for timestamp calculation
+            audio_kwargs (`Dict`, *optional*):
+                Additional keyword arguments for audio processing, including:
+                - `sampling_rate` (`int`, *optional*): Target sampling rate for audio
+            text_kwargs (`Dict`, *optional*):
+                Additional keyword arguments for text tokenization, including:
+                - `return_tensors` (`str` or [`~utils.TensorType`], *optional*): Framework for returned tensors ('tf', 'pt', 'np', 'jax')
+                - `padding` (`bool`, *optional*): Whether to pad sequences (defaults to False)
+
+        Returns:
+            [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+              `None`).
+            - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+            - **num_patches** -- Number of patches per image. Returned when `images` is not `None`.
+            - **pixel_values_videos** -- Pixel values of videos to be fed to a model. Returned when `videos` is not `None`.
+            - **sound_clips** -- Raw audio waveforms to be fed to a model. Returned when `audio` is not `None`.
+
+        Raises:
+            AssertionError: If videos are provided with batch size > 1 (not currently supported).
+
+        Note:
+            - Image tokens `<image>` in text are replaced with `<img>` + repeated image tokens + `</img>`
+            - Video tokens `<video>` in text are replaced with frame-by-frame descriptions
+            - Audio tokens `<audio>` in text are replaced with repeated audio placeholder tokens
+            - When video metadata with fps is provided, frame descriptions include timestamps
+            - Videos are processed with max_num_tiles=1 regardless of the images setting
+        """
+        output_kwargs = self._merge_kwargs(
+            NemotronH_Nano_Omni_Reasoning_V3ProcessorKwargs,
+            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+            **kwargs,
+        )
+        image_inputs, videos_inputs, audio_inputs = {}, {}, {}
+
+        if images is not None:
+            image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+            image_num_tokens = image_inputs["num_tokens"]
+
+        if videos is not None:
+            # One tile per frame, but sized with the **video** rule (`video_target_num_patches` +
+            # aspect ratio) — different from the image rule. We flip a flag on the image processor
+            # around the call rather than routing a new kwarg through `ImagesKwargs`, which is a
+            # strict-dataclass and rejects unknown fields.
+            self.image_processor._is_video_mode = True
+            try:
+                videos_inputs = self.image_processor(images=videos, **output_kwargs["images_kwargs"])
+            finally:
+                self.image_processor._is_video_mode = False
+            video_num_patches = [sum(videos_inputs["num_patches"])]
+            videos_inputs["pixel_values_videos"] = videos_inputs["pixel_values"]
+            del videos_inputs["pixel_values"]
+
+        # Process audio inputs
+        audio_num_tokens = []
+        if audio is not None:
+            audio_clips, audio_num_tokens = self._process_audio(audio, output_kwargs.get("audio_kwargs", {}))
+            # Keep as list of numpy arrays - don't let BatchFeature convert to tensor
+            # The model's generate function will handle conversion
+            audio_inputs["sound_clips"] = audio_clips
+
+        if not isinstance(text, list):
+            text = [text]
+
+        text = text.copy()  # below lines change text in-place
+        if images is not None:
+            index = 0
+            for i in range(len(text)):
+                while self.image_token in text[i]:
+                    n_tokens = image_num_tokens[index]
+                    text[i] = text[i].replace(
+                        self.image_token,
+                        self.image_start_token + "<|placeholder|>" * n_tokens + self.image_end_token,
+                        1,
+                    )
+                    index += 1
+                text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        if videos is not None:
+            assert len(text) == 1, "Video is not supported for batch size > 1"
+            video_metadata = output_kwargs.get("videos_kwargs", {}).get("video_metadata", None)
+            i = 0
+            index = 0
+            if self.video_token in text[i]:
+                # Matches vLLM's `get_video_repl` (`vllm/transformers_utils/processors/
+                # nano_nemotron_vl.py`): one `<img>…</img>` chunk per temporal patch (tubelet),
+                # labeled with the per-frame timestamps joined by " and " — capitalized "Frame"
+                # for the first frame in the tubelet, lowercase "frame" for the rest. No "This is
+                # a video:\n" prefix is emitted by the processor — it's expected to come from the
+                # client message, consistent with vLLM and training.
+                tokens_per_tubelet = videos_inputs["num_tokens"][0]
+                each_group = self.image_start_token + "<|placeholder|>" * tokens_per_tubelet + self.image_end_token
+                T = self.video_temporal_patch_dim
+                n_frames = video_num_patches[index]
+                n_groups = (n_frames + T - 1) // T
+
+                # vLLM formula: int(source_frame_idx) * int(1000 / source_fps) / 1000
+                # Requires source fps and source frame indices from video_metadata.
+                source_fps = video_metadata.fps if (video_metadata is not None and video_metadata.fps) else None
+                frames_indices = video_metadata.frames_indices if video_metadata is not None else None
+                if source_fps is not None:
+                    frame_duration_ms = int(1000.0 / source_fps)
+
+                frame_labels = []
+                for g in range(n_groups):
+                    parts = []
+                    for j in range(T):
+                        fi = g * T + j
+                        if fi >= n_frames:
+                            break  # last group may be short
+                        prefix = "Frame" if j == 0 else "frame"
+                        if source_fps is not None and frames_indices is not None and fi < len(frames_indices):
+                            ts = int(frames_indices[fi]) * frame_duration_ms / 1000.0
+                            parts.append(f"{prefix} {fi+1} sampled at {ts:.2f} seconds")
+                        elif source_fps is not None:
+                            ts = fi / source_fps
+                            parts.append(f"{prefix} {fi+1} sampled at {ts:.2f} seconds")
+                        else:
+                            parts.append(f"{prefix} {fi+1}")
+                    frame_labels.append(" and ".join(parts) + ": ")
+
+                video_prompt = ""
+                for g, label in enumerate(frame_labels):
+                    if g > 0:
+                        video_prompt += "\n"
+                    video_prompt += label + each_group
+
+                text[i] = text[i].replace(self.video_token, video_prompt, 1)
+            # The tokenizer has no real `<video>` token (the 131081 id in the config doesn't decode
+            # to any printable string), so we reuse `<image>` (id 18) as the placeholder. The outer
+            # model distinguishes image vs. video by which `pixel_values_*` arg was passed.
+            text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+        # Replace audio tokens with the correct number of placeholder tokens.
+        # The expansion loop is per-row, so batch size > 1 is supported as long
+        # as `audio_num_tokens` has one entry per `<so_embedding>` placeholder
+        # across the batch (in row-major order).
+        if audio is not None:
+            index = 0
+            for i in range(len(text)):
+                while self.audio_token in text[i]:
+                    num_tokens = audio_num_tokens[index] if index < len(audio_num_tokens) else 1
+                    # Replace <audio> with repeated audio tokens
+                    text[i] = text[i].replace(self.audio_token, self.audio_start_token + "<|audio_placeholder|>" * num_tokens + self.audio_end_token, 1)
+                    index += 1
+                text[i] = text[i].replace("<|audio_placeholder|>", self.audio_token)
+
+        return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+        text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+        # Build output - exclude audio from tensor conversion since it's raw waveforms
+        output_data = {**text_inputs, **image_inputs, **videos_inputs}
+        result = BatchFeature(data=output_data, tensor_type=return_tensors)
+
+        # Add audio clips separately (as list of numpy arrays, not tensors)
+        if audio_inputs:
+            result["sound_clips"] = audio_inputs["sound_clips"]
+
+        return result
+
+    def _process_audio(
+        self,
+        audio: AudioInput,
+        audio_kwargs: dict
+    ) -> tuple:
+        """Process audio inputs and compute the number of audio tokens.
+
+        Args:
+            audio: Audio input (file path, waveform array, or list thereof)
+            audio_kwargs: Additional audio processing arguments
+
+        Returns:
+            Tuple of (audio_clips, num_tokens_per_clip)
+        """
+        # Get sampling rate from kwargs or use default
+        sampling_rate = audio_kwargs.get("sampling_rate", self.audio_sampling_rate)
+
+        # Normalize audio to list
+        if not isinstance(audio, list):
+            audio = [audio]
+
+        audio_clips = []
+        num_tokens = []
+
+        for audio_item in audio:
+            # Load audio if it's a file path
+            if isinstance(audio_item, str):
+                waveform = self._load_audio(audio_item, sampling_rate)
+            elif isinstance(audio_item, torch.Tensor):
+                waveform = audio_item.numpy() if audio_item.dim() == 1 else audio_item.squeeze().numpy()
+            elif isinstance(audio_item, np.ndarray):
+                waveform = audio_item.squeeze() if audio_item.ndim > 1 else audio_item
+            else:
+                raise ValueError(f"Unsupported audio type: {type(audio_item)}")
+
+            audio_clips.append(waveform)
+
+            n_tokens = self._estimate_audio_num_embeddings(len(waveform))
+            num_tokens.append(max(1, n_tokens))  # At least 1 token
+
+        return audio_clips, num_tokens
+
+    def _estimate_audio_num_embeddings(self, audio_length_samples: int) -> int:
+        """Predict the exact number of `<so_embedding>` tokens that the sound
+        encoder will produce for an audio clip of `audio_length_samples` raw
+        samples. Replaces the previous heuristic, which under-counted by 1 for
+        certain lengths and tripped a shape mismatch in `modeling.py::generate`.
+
+        Mirrors `ParakeetFeatureExtractor` (center-padded STFT → ``1 + L // hop``
+        mel frames) followed by `ParakeetEncoder._get_subsampling_output_length`
+        (``log2(subsampling_factor)`` stages of stride-2 conv with kernel-size
+        ``subsampling_conv_kernel_size``, symmetric padding ``(kernel-1)//2`` on
+        each side).
+        """
+        # Mel frame count (center=True is HF Parakeet's default).
+        n_frames = 1 + audio_length_samples // self.audio_hop_length
+        # Conv-subsampling: ``log2(subsampling_factor)`` stages, stride 2,
+        # kernel ``subsampling_conv_kernel_size``, symmetric padding.
+        kernel_size = getattr(self, "audio_subsampling_conv_kernel_size", 3)
+        stride = getattr(self, "audio_subsampling_conv_stride", 2)
+        num_layers = int(math.log2(self.audio_subsampling_factor))
+        all_paddings = (kernel_size - 1) // 2 * 2
+        add_pad = all_paddings - kernel_size  # kernel=3, sym pad=1 → -1
+        L = n_frames
+        for _ in range(num_layers):
+            L = (L + add_pad) // stride + 1
+        return L
+
+    def _load_audio(self, audio_path: str, target_sr: int) -> np.ndarray:
+        """Load audio from file and resample if necessary.
+
+        Args:
+            audio_path: Path to audio file
+            target_sr: Target sampling rate
+
+        Returns:
+            Audio waveform as numpy array
+        """
+        try:
+            import librosa
+            waveform, sr = librosa.load(audio_path, sr=target_sr, mono=True)
+            return waveform
+        except ImportError:
+            pass
+
+        try:
+            import soundfile as sf
+            waveform, sr = sf.read(audio_path)
+            if waveform.ndim > 1:
+                waveform = waveform.mean(axis=1)  # Convert to mono
+            if sr != target_sr:
+                # Simple resampling using numpy
+                import scipy.signal
+                num_samples = int(len(waveform) * target_sr / sr)
+                waveform = scipy.signal.resample(waveform, num_samples)
+            return waveform.astype(np.float32)
+        except ImportError:
+            pass
+
+        raise ImportError(
+            "Audio loading requires either librosa or soundfile. "
+            "Install with: pip install librosa soundfile"
+        )
+
+    def _get_num_multimodal_tokens(self, image_sizes=None, video_sizes=None, **kwargs):
+        """
+        Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
+        Args:
+            image_sizes (`list[list[int]]`, *optional*):
+                The input sizes formatted as (height, width) per each image.
+            video_sizes (`list[list[int]]`, *optional*):
+                The input sizes formatted as (num_frames, height, width) per each video.
+        Returns:
+            `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
+            input modalities, along with other useful data.
+        """
+
+        vision_data = {}
+        if image_sizes is not None:
+            images_kwargs = NemotronH_Nano_Omni_Reasoning_V3ProcessorKwargs._defaults.get("images_kwargs", {})
+            images_kwargs.update(kwargs)
+            merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size
+
+            num_image_patches = [
+                self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
+                for image_size in image_sizes
+            ]
+            num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
+            vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
+        return MultiModalData(**vision_data)
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.batch_decode`]. Please
+        refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to the tokenizer's [`~PreTrainedTokenizer.decode`]. Please refer to
+        the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    def post_process_image_text_to_text(
+        self, generated_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False, **kwargs
+    ):
+        """
+        Post-process the output of the model to decode the text.
+
+        Args:
+            generated_outputs (`torch.Tensor` or `np.ndarray`):
+                The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+                or `(sequence_length,)`.
+            skip_special_tokens (`bool`, *optional*, defaults to `True`):
+                Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+            clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+                Whether or not to clean up the tokenization spaces. Argument passed to the tokenizer's `batch_decode` method.
+            **kwargs:
+                Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+        Returns:
+            `list[str]`: The decoded text.
+        """
+        return self.tokenizer.batch_decode(
+            generated_outputs,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs,
+        )
+
+    @property
+    def model_input_names(self):
+        tokenizer_input_names = self.tokenizer.model_input_names
+        image_processor_input_names = self.image_processor.model_input_names
+        names_from_processor = list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))
+        return names_from_processor + ["second_per_grid_ts"]
+
+
+__all__ = ["NemotronH_Nano_Omni_Reasoning_V3Processor"]

processing_utils.py

@@ -0,0 +1,83 @@
+from typing import List, Optional, Union, Any, Dict
+
+from PIL import Image
+import torch
+from transformers.image_processing_base import BatchFeature
+from transformers.image_processing_utils_fast import BaseImageProcessorFast, divide_to_patches
+from transformers.image_utils import (make_list_of_images, get_image_size,
+                                      get_image_type, ImageInput, ImageType, ChannelDimension)
+from transformers.utils import TensorType
+import torchvision.transforms as T
+
+
+def get_internvl_target_ratios(
+    min_num: int,
+    max_num: int,
+) -> list[tuple[int, int]]:
+    target_ratios = {(i, j)
+                     for n in range(min_num, max_num + 1)
+                     for i in range(1, n + 1)
+                     for j in range(1, n + 1) if min_num <= i * j <= max_num}
+    return sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_factor = float('-inf')
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        factor_based_on_area_n_ratio = min(
+            (ratio[0]*ratio[1]*image_size*image_size)/ area, 0.6
+            )* min(
+                target_aspect_ratio/aspect_ratio, aspect_ratio/target_aspect_ratio)
+        if factor_based_on_area_n_ratio > best_factor:
+            best_factor = factor_based_on_area_n_ratio
+            best_ratio = ratio
+    return best_ratio
+
+
+def calculate_targets(
+    orig_width: int,
+    orig_height: int,
+    target_ratios: list[tuple[int, int]],
+    image_size: int,
+) -> tuple[int, int, int]:
+    aspect_ratio = orig_width / orig_height
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio,
+        target_ratios,
+        width=orig_width,
+        height=orig_height,
+        image_size=image_size,
+    )
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    return blocks, target_width, target_height
+
+
+def dynamic_preprocess(image, image_size=512, max_num_tiles=12, use_thumbnail=True):
+    orig_height, orig_width = get_image_size(image, channel_dim=ChannelDimension.FIRST)
+    target_ratios = get_internvl_target_ratios(1, max_num_tiles)
+
+    blocks, target_width, target_height = calculate_targets(
+        orig_width,
+        orig_height,
+        target_ratios,
+        image_size
+    )
+    # resize the image
+    resized_img = T.Resize((target_width, target_height), interpolation=T.InterpolationMode.BICUBIC)(image)
+    patches = divide_to_patches(resized_img, image_size)
+    assert len(patches) == blocks
+    if use_thumbnail and len(patches) != 1:
+        thumbnail_img = T.Resize((image_size, image_size), interpolation=T.InterpolationMode.BICUBIC)(image)
+        patches.append(thumbnail_img)
+
+    return patches

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.json

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

video_io.py

@@ -0,0 +1,166 @@
+import os
+import base64
+import mimetypes
+from PIL import Image
+import io
+from transformers.video_utils import VideoMetadata
+
+
+def encode_pil_to_jpeg_data_url(pil_image):
+    from io import BytesIO
+    buf = BytesIO()
+    pil_image.save(buf, format="JPEG")
+    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:image/jpeg;base64,{b64}"
+
+
+def sample_video_frames_to_data_urls(video_path_local, fps=1, nframe=0, nframe_max=-1):
+    """
+    Sample frames from a video and return base64-encoded data URLs along with metadata.
+    
+    Args:
+        video_path_local: Path to the video file
+        fps: Target frames per second for sampling (if > 0, uses fps-based sampling)
+        nframe: Number of frames to sample (used if fps <= 0)
+        nframe_max: Maximum number of frames to sample
+    
+    Returns:
+        tuple: (frame_data_urls, metadata)
+        - frame_data_urls: List of base64-encoded frame images
+        - metadata: VideoMetadata dataclass containing info about the sampled frames:
+            - total_num_frames: Number of sampled frames
+            - fps: Effective frame rate of the sampled frames
+            - duration: Duration covered by the sampled frames (in seconds)
+            - video_backend: Backend used for video processing ('decord')
+    """
+    import numpy as np
+    from PIL import Image
+    import decord
+
+    vid = decord.VideoReader(video_path_local)
+    total_frames = len(vid)
+    video_fps = vid.get_avg_fps()
+    total_duration = total_frames / max(1e-6, video_fps)
+
+    if fps > 0:
+        required_frames = int(total_duration * fps)
+        desired_frames = max(1, required_frames)
+        if nframe_max > 0 and desired_frames > nframe_max:
+            desired_frames = nframe_max
+        if desired_frames >= total_frames:
+            indices = list(range(total_frames))
+        elif desired_frames == 1:
+            indices = [0]  # Always use first frame for single frame sampling
+        else:
+            # Generate evenly spaced indices and ensure uniqueness
+            raw_indices = np.linspace(0, total_frames - 1, desired_frames)
+            indices = list(np.unique(np.round(raw_indices).astype(int)))
+    else:
+        desired_frames = max(1, int(nframe) if nframe and nframe > 0 else 8)
+        if nframe_max > 0 and desired_frames > nframe_max:
+            desired_frames = nframe_max
+        if desired_frames >= total_frames:
+            indices = list(range(total_frames))
+        elif desired_frames == 1:
+            indices = [0]  # Always use first frame for single frame sampling
+        else:
+            # Generate evenly spaced indices and ensure uniqueness
+            raw_indices = np.linspace(0, total_frames - 1, desired_frames)
+            indices = list(np.unique(np.round(raw_indices).astype(int)))
+
+    images = [Image.fromarray(vid[i].asnumpy()) for i in indices]
+    frame_urls = [encode_pil_to_jpeg_data_url(im) for im in images]
+
+    sampled_num_frames = len(indices)
+
+    # Pass source fps and source frame indices so the processor can compute
+    # timestamps with vLLM's formula: int(source_frame_idx) * int(1000/source_fps) / 1000
+    metadata = VideoMetadata(
+        total_num_frames=sampled_num_frames,
+        fps=video_fps,
+        frames_indices=[int(i) for i in indices],
+        duration=total_duration,
+        video_backend=None,
+    )
+
+    return frame_urls, metadata
+
+
+def maybe_path_or_url_to_data_urls(path_or_url, fps=1, nframe=0, nframe_max=-1):
+    """
+    Convert a path or URL to data URLs, handling videos, images, and remote files.
+    
+    Args:
+        path_or_url: Path or URL to the media file
+        fps: Target frames per second for video sampling (if > 0, uses fps-based sampling)
+        nframe: Number of frames to sample from video (used if fps <= 0)
+        nframe_max: Maximum number of frames to sample
+    
+    Returns:
+        tuple: (data_urls, metadata)
+        - data_urls: List of base64-encoded data URLs
+        - metadata: VideoMetadata dataclass with video metadata or None for images
+    """
+    val = str(path_or_url or "")
+    low = val.lower()
+    
+    # Handle data URLs
+    if low.startswith("data:"):
+        if low.startswith("data:video/mp4"):
+            header, _, b64part = val.partition(",")
+            if not b64part:
+                return [val], None
+            import tempfile
+            tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
+            try:
+                tmp.write(base64.b64decode(b64part))
+                tmp.flush(); tmp.close()
+                return sample_video_frames_to_data_urls(tmp.name, fps=fps, nframe=nframe, nframe_max=nframe_max)
+            finally:
+                try:
+                    os.unlink(tmp.name)
+                except Exception:
+                    pass
+        return [val], None
+
+    # Remote URL
+    if low.startswith("http://") or low.startswith("https://"):
+        if low.endswith(".mp4"):
+            try:
+                import tempfile, urllib.request
+                with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmpf:
+                    urllib.request.urlretrieve(val, tmpf.name)
+                    local_path = tmpf.name
+                result = sample_video_frames_to_data_urls(local_path, fps=fps, nframe=nframe, nframe_max=nframe_max)
+                try:
+                    os.unlink(local_path)
+                except Exception:
+                    pass
+                return result
+            except Exception:
+                return [val], None
+        return [val], None
+
+    # Local path
+    if os.path.exists(val):
+        mime, _ = mimetypes.guess_type(val)
+        if mime and mime.startswith("image/"):
+            with open(val, "rb") as f:
+                b64 = base64.b64encode(f.read()).decode("utf-8")
+            return [f"data:{mime};base64,{b64}"], None
+        if mime == "video/mp4" or (mime is None and val.endswith(".mp4")):
+            return sample_video_frames_to_data_urls(val, fps=fps, nframe=nframe, nframe_max=nframe_max)
+        # Fallback: treat as binary image
+        with open(val, "rb") as f:
+            b64 = base64.b64encode(f.read()).decode("utf-8")
+        return [f"data:image/jpeg;base64,{b64}"], None
+
+    return [val], None
+
+
+def pil_image_from_base64(b64_str: str) -> Image.Image:
+    # Handle data URLs like "data:image/png;base64,...."
+    if b64_str.startswith('data:'):
+        b64_str = b64_str.split(',', 1)[1]
+    img_bytes = base64.b64decode(b64_str)
+    return Image.open(io.BytesIO(img_bytes))

video_processing.py

@@ -0,0 +1,166 @@
+# coding=utf-8
+# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+"""video processor class for Qwen2-VL."""
+
+import math
+from typing import Optional, Union
+
+from transformers.image_processing_utils import (
+    BatchFeature,
+)
+from transformers.image_utils import (
+    OPENAI_CLIP_MEAN,
+    OPENAI_CLIP_STD,
+    ChannelDimension,
+    SizeDict,
+    get_image_size,
+)
+from transformers.processing_utils import Unpack, VideosKwargs
+from transformers.utils import (
+    TensorType,
+    add_start_docstrings,
+    is_torch_available,
+    is_torchvision_available,
+    is_torchvision_v2_available,
+    is_vision_available,
+)
+from transformers.utils.import_utils import requires
+from transformers.video_processing_utils import (
+    BASE_VIDEO_PROCESSOR_DOCSTRING,
+    BaseVideoProcessor,
+)
+from transformers.video_utils import VideoMetadata, group_videos_by_shape, reorder_videos
+import torchvision.transforms as T
+
+from .processing_utils import get_internvl_target_ratios, calculate_targets
+
+
+if is_torchvision_available():
+    if is_torchvision_v2_available():
+        from torchvision.transforms.v2 import functional as F
+    else:
+        from torchvision.transforms import functional as F
+
+
+if is_torch_available():
+    import torch
+
+
+@requires(backends=("torchvision",))
+class NemotronH_Nano_Omni_Reasoning_V3VideoProcessor(BaseVideoProcessor):
+    model_input_names = ["pixel_values_videos", "video_grid_thw"]
+
+    def __init__(self, image_size=512, max_num_tiles=12, norm_mean=None, norm_std=None, **kwargs):
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        self.max_num_tiles = max_num_tiles
+        self.norm_mean = norm_mean
+        self.norm_std = norm_std
+
+    def _preprocess(
+        self,
+        videos: list["torch.Tensor"],
+        video_metadata: Union[list[VideoMetadata], list[dict]],
+        do_sample_frames: bool,
+        fps: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        device: Optional["torch.Tensor"] = None,
+        **kwargs,
+    ):
+        if do_sample_frames:
+            # Sample video frames
+            videos = [
+                self.sample_frames(
+                    video,
+                    metadata=metadata,
+                    num_frames=num_frames,
+                    fps=fps,
+                )
+                for video, metadata in zip(videos, video_metadata)
+            ]
+
+        # We need to sample frames first before moving to device, if `do_sample_frames=True`. Otherwise
+        # moving the whole video incurs high GPU mem usage for long videos
+        if device is not None:
+            videos = [video.to(device) for video in videos]
+
+        # Group videos by size for batched resizing
+        grouped_videos, grouped_videos_index = group_videos_by_shape(videos)
+        resized_videos_grouped = {}
+        processed_grids = {}
+        for shape, stacked_videos in grouped_videos.items():
+            height, width = get_image_size(stacked_videos[0], channel_dim=ChannelDimension.FIRST)
+            batch_size, grid_t, channel = stacked_videos.shape[:3]
+
+            target_ratios = get_internvl_target_ratios(1, self.max_num_tiles)
+            blocks, resize_width, resize_height = calculate_targets(
+                width,
+                height,
+                target_ratios,
+                self.image_size
+            )
+            stacked_videos = self.resize(
+                image=stacked_videos,
+                size=SizeDict(height=resize_height, width=resize_width),
+                interpolation=T.InterpolationMode.BICUBIC,
+            )
+            # stacked_videos = T.Resize((resize_width, resize_height), interpolation=T.InterpolationMode.BICUBIC)(stacked_videos)
+            norm_mean = torch.as_tensor(self.norm_mean, dtype=stacked_videos.dtype, device=stacked_videos.device).view(1, 1, 3, 1, 1)
+            norm_std  = torch.as_tensor(self.norm_std,  dtype=stacked_videos.dtype, device=stacked_videos.device).view(1, 1, 3, 1, 1)
+            stacked_videos = (stacked_videos - norm_mean) / norm_std
+            resized_videos_grouped[shape] = stacked_videos
+            grid_h, grid_w = resize_height // self.image_size, resize_width // self.image_size
+            processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+        resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index)
+        processed_grids = reorder_videos(processed_grids, grouped_videos_index)
+        pixel_values_videos = torch.cat(resized_videos, dim=0)
+        video_grid_thw = torch.tensor(processed_grids)
+
+        return BatchFeature(
+            data={"pixel_values_videos": pixel_values_videos, "video_grid_thw": video_grid_thw},
+            tensor_type=return_tensors,
+        )
+
+    def get_num_of_video_patches(self, num_frames: int, height: int, width: int):
+        """
+        A utility that returns number of video patches a given video size.
+
+        Args:
+            num_frames (`int`):
+                Number of frames in the input video.
+            height (`int`):
+                Height of the input video.
+            width (`int`):
+                Width of the input video.
+        Returns:
+            `Tuple(int, int)`: Number of placeholder tokens required and number of patches per image.
+        """
+        target_ratios = get_internvl_target_ratios(1, self.max_num_tiles)
+        blocks, _, _ = calculate_targets(
+            width,
+            height,
+            target_ratios,
+            self.image_size
+        )
+        return num_frames * blocks
+
+
+__all__ = ["NemotronH_Nano_Omni_Reasoning_V3VideoProcessor"]