[None][feat] Add Gemma 4 12B Unified (encoder-free multimodal) support

docs/source/models/supported-models.md

@@ -19,6 +19,7 @@ The following is a table of supported models for the PyTorch backend:
 | `Gemma3ForCausalLM`                  | Gemma 3                            | `google/gemma-3-1b-it`                       |
 | `Gemma3nForConditionalGeneration` [^7]| Gemma 3n                           | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it` |
 | `Gemma4ForConditionalGeneration`     | Gemma 4                            | `google/gemma-4-E2B-it`, `google/gemma-4-E4B-it`, `google/gemma-4-26B-A4B-it` [^6], `google/gemma-4-31B-it` [^6] |
+| `Gemma4UnifiedForConditionalGeneration` | Gemma 4 12B Unified (encoder-free) | `google/gemma-4-12B`, `google/gemma-4-12B-it`              |
 | `Glm4MoeForCausalLM`                 | GLM-4.5, GLM-4.6, GLM-4.7          | `THUDM/GLM-4-100B-A10B`                      |
 | `Glm4MoeLiteForCausalLM` [^5]        | GLM-4.7-Flash                      | `zai-org/GLM-4.7-Flash`                      |
 | `GlmMoeDsaForCausalLM`               | GLM-5                              | `zai-org/GLM-5`                              |
@@ -74,6 +75,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Glm4MoeLiteForCausalLM` [^5]    | Yes               | Yes        | Untested                   | Untested              | Yes             | No  | No               | No                | No     | Yes           | Untested         | Untested       | N/A                      | Untested              | Untested        |
 | `NemotronHForCausalLM`           | Yes               | Yes        | Yes                        | Yes                   | Yes             | Yes | No               | No                | No     | Yes           | Yes              | Yes            | N/A                      | Untested              | Untested        |
 | `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
+| `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | Yes | No               | No                | No     | Yes           | Untested         | Untested       | Yes                      | Untested              | Untested        |
 | `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | No               | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
 
@@ -97,6 +99,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Exaone4_5_ForConditionalGeneration` | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | No                        | L + I + V |
 | `Gemma3ForConditionalGeneration`     | Yes               | Yes        | N/A             | Yes           | Yes              | N/A            | Yes                   | No                        | L + I     |
 | `Gemma4ForConditionalGeneration`     | Untested          | Yes        | Yes             | Yes           | Untested         | No             | Untested              | No                        | L + I + V + A [^9] |
+| `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested        | Yes           | Untested         | No             | Untested              | No                        | L + I + A |
 | `HCXVisionForCausalLM`               | Yes               | Yes        | No              | Yes           | Yes              | Yes            | Yes                   | No                        | L + I     |
 | `LlavaLlamaModel (VILA)`             | Yes               | Yes        | No              | Yes           | Yes              | No             | Yes                   | No                        | L + I + V |
 | `LlavaNextForConditionalGeneration`  | Yes               | Yes        | Yes             | Yes           | Yes              | Yes            | Yes                   | Yes                       | L + I     |

tensorrt_llm/_torch/configs/__init__.py

@@ -1,21 +1,30 @@
 from tensorrt_llm._torch.configs.cosmos3 import Cosmos3Config
 from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config
 from tensorrt_llm._torch.configs.deepseekv4 import DeepseekV4Config
+from tensorrt_llm._torch.configs.gemma4_unified import (
+    Gemma4UnifiedAudioConfig,
+    Gemma4UnifiedConfig,
+    Gemma4UnifiedTextConfig,
+    Gemma4UnifiedVisionConfig,
+)
 from tensorrt_llm._torch.configs.laguna import LagunaConfig
 
 
 def _register_custom_configs_with_transformers() -> None:
     # Make AutoConfig.from_pretrained / AutoTokenizer.from_pretrained accept
     # model_types that TRT-LLM understands but upstream transformers does not
     # (DeepSeek-V3.2, Kimi K2, Laguna, and Cosmos3 omni ship config.json with
-    # these model_types and rely on TRT-LLM's local config workarounds).
+    # these model_types and rely on TRT-LLM's local config workarounds;
+    # likewise the gemma4_unified family).
     #
     # Without this, transformers 5.5.x falls back to a bare PreTrainedConfig
     # that lacks attributes like `max_position_embeddings`, and
     # AutoTokenizer.from_pretrained then raises AttributeError before any
     # tokenizer can be constructed. Bypass AutoConfig.register's model_type
     # consistency check for aliases (for example, DeepseekV3Config.model_type
     # is "deepseek_v3") by writing into the underlying mapping directly.
+    # Registration only fills gaps: when the installed transformers already
+    # ships a model_type, the native class wins.
     from transformers.models.auto.configuration_auto import CONFIG_MAPPING
     from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLVisionConfig
 
@@ -28,6 +37,10 @@ def _register_custom_configs_with_transformers() -> None:
         "kimi_k2": DeepseekV3Config,
         "deepseek_v4": DeepseekV4Config,
         "laguna": LagunaConfig,
+        "gemma4_unified": Gemma4UnifiedConfig,
+        "gemma4_unified_text": Gemma4UnifiedTextConfig,
+        "gemma4_unified_vision": Gemma4UnifiedVisionConfig,
+        "gemma4_unified_audio": Gemma4UnifiedAudioConfig,
     }
     # Cosmos3Config resolves vision sub-configs via ``qwen3_vl_vision``; that
     # alias is only present in newer transformers releases.
@@ -42,4 +55,13 @@ def _register_custom_configs_with_transformers() -> None:
 _register_custom_configs_with_transformers()
 del _register_custom_configs_with_transformers
 
-__all__ = ["Cosmos3Config", "DeepseekV3Config", "DeepseekV4Config", "LagunaConfig"]
+__all__ = [
+    "Cosmos3Config",
+    "DeepseekV3Config",
+    "DeepseekV4Config",
+    "Gemma4UnifiedAudioConfig",
+    "Gemma4UnifiedConfig",
+    "Gemma4UnifiedTextConfig",
+    "Gemma4UnifiedVisionConfig",
+    "LagunaConfig",
+]

tensorrt_llm/_torch/configs/gemma4_unified.py

@@ -0,0 +1,145 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+"""Config classes for Gemma 4 12B Unified (encoder-free multimodal).
+
+Registered with the transformers CONFIG_MAPPING (see `_torch/configs/__init__.py`)
+so `AutoConfig.from_pretrained` can parse a Gemma 4 12B checkpoint whenever the
+installed transformers does not ship the `gemma4_unified` model_types natively.
+
+All fields are read directly from the checkpoint's config.json, matching the
+attribute names used by `Gemma4UnifiedForConditionalGeneration`. The text
+backbone of the 12B is a standard dense Gemma 4 text model, so its sub-config
+reuses the native `Gemma4TextConfig`.
+"""
+
+from transformers import Gemma4TextConfig
+from transformers.configuration_utils import PretrainedConfig
+
+
+class Gemma4UnifiedTextConfig(Gemma4TextConfig):
+    """Text sub-config for Gemma 4 12B Unified.
+
+    The 12B text backbone is a standard dense Gemma 4 text model; only the
+    model_type string differs, so this is a pure alias of the native
+    `Gemma4TextConfig`.
+    """
+
+    model_type = "gemma4_unified_text"
+
+
+class Gemma4UnifiedVisionConfig(PretrainedConfig):
+    """Sub-config for the encoder-free vision projector."""
+
+    model_type = "gemma4_unified_vision"
+
+    def __init__(
+        self,
+        mm_embed_dim: int = 3840,
+        mm_posemb_size: int = 1120,
+        output_proj_dims: int = 3840,
+        patch_size: int = 16,
+        pooling_kernel_size: int = 3,
+        rms_norm_eps: float = 1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.mm_embed_dim = mm_embed_dim
+        self.mm_posemb_size = mm_posemb_size
+        self.output_proj_dims = output_proj_dims
+        self.patch_size = patch_size
+        self.pooling_kernel_size = pooling_kernel_size
+        self.rms_norm_eps = rms_norm_eps
+
+
+class Gemma4UnifiedAudioConfig(PretrainedConfig):
+    """Sub-config for the encoder-free audio projector.
+
+    `output_proj_dims` and `hidden_size` alias `audio_embed_dim` (the raw audio
+    frame width) when not given, matching the HF implementation; they are plain
+    attributes here so a checkpoint config.json that spells them out loads as-is.
+    """
+
+    model_type = "gemma4_unified_audio"
+
+    def __init__(
+        self,
+        audio_embed_dim: int = 640,
+        rms_norm_eps: float = 1e-6,
+        output_proj_dims: int | None = None,
+        hidden_size: int | None = None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.audio_embed_dim = audio_embed_dim
+        self.rms_norm_eps = rms_norm_eps
+        self.output_proj_dims = (
+            output_proj_dims if output_proj_dims is not None else audio_embed_dim
+        )
+        self.hidden_size = hidden_size if hidden_size is not None else audio_embed_dim
+
+
+class Gemma4UnifiedConfig(PretrainedConfig):
+    """Top-level config for Gemma 4 12B Unified (encoder-free multimodal).
+
+    Parses `config.json` fields required by
+    `Gemma4UnifiedForConditionalGeneration` without depending on any
+    natively shipped transformers class. The `text_config`, `vision_config`, and
+    `audio_config` sub-configs are reconstructed from nested dicts using the
+    shim classes above.
+    """
+
+    model_type = "gemma4_unified"
+
+    def __init__(
+        self,
+        text_config=None,
+        vision_config=None,
+        audio_config=None,
+        image_token_id: int = 258880,
+        audio_token_id: int = 258881,
+        video_token_id: int = 258884,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.image_token_id = image_token_id
+        self.audio_token_id = audio_token_id
+        self.video_token_id = video_token_id
+
+        # Sub-configs arrive as dicts from AutoConfig.from_pretrained; rebuild
+        # them with the classes above.
+        if text_config is not None:
+            if isinstance(text_config, dict):
+                self.text_config = Gemma4UnifiedTextConfig(**text_config)
+            else:
+                self.text_config = text_config
+        else:
+            self.text_config = None
+
+        if vision_config is not None:
+            if isinstance(vision_config, dict):
+                self.vision_config = Gemma4UnifiedVisionConfig(**vision_config)
+            else:
+                self.vision_config = vision_config
+        else:
+            self.vision_config = None
+
+        if audio_config is not None:
+            if isinstance(audio_config, dict):
+                self.audio_config = Gemma4UnifiedAudioConfig(**audio_config)
+            else:
+                self.audio_config = audio_config
+        else:
+            self.audio_config = None

tensorrt_llm/_torch/models/__init__.py

@@ -1,8 +1,9 @@
 import transformers
 
 # Importing _torch.configs triggers AutoConfig registration for TRT-LLM-only
-# model_types (deepseek_v32, kimi_k2) so AutoTokenizer.from_pretrained works
-# under transformers >= 5.5; see _torch/configs/__init__.py.
+# model_types (deepseek_v32, kimi_k2, gemma4_unified) so AutoConfig /
+# AutoTokenizer.from_pretrained work under transformers >= 5.5; see
+# _torch/configs/__init__.py.
 import tensorrt_llm._torch.configs  # noqa: F401
 
 from .modeling_afmoe import AfmoeForCausalLM
@@ -20,6 +21,9 @@
 from .modeling_exaone_moe import ExaoneMoeForCausalLM
 from .modeling_gemma3 import Gemma3ForCausalLM
 from .modeling_gemma3vl import Gemma3VLM
+from .modeling_gemma4 import Gemma4ForCausalLM
+from .modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration
+from .modeling_gemma4mm import Gemma4ForConditionalGeneration
 from .modeling_glm import Glm4MoeForCausalLM
 from .modeling_gpt_oss import GptOssForCausalLM
 from .modeling_hunyuan_dense import HunYuanDenseV1ForCausalLM
@@ -74,6 +78,9 @@
     "ExaoneMoeForCausalLM",
     "Gemma3ForCausalLM",
     "Gemma3VLM",
+    "Gemma4ForCausalLM",
+    "Gemma4ForConditionalGeneration",
+    "Gemma4UnifiedForConditionalGeneration",
     "HCXVisionForCausalLM",
     "LagunaForCausalLM",
     "HunYuanDenseV1ForCausalLM",
@@ -130,15 +137,3 @@
     print(
         f"Failed to import MllamaForConditionalGeneration as transformers.__version__ {transformers.__version__} < 4.45.1"
     )
-
-# Gemma4 requires transformers>=5.5.0 (native Gemma4 config/model classes).
-# Import silently on failure -- `get_model_architecture` in modeling_utils.py
-# raises a targeted "upgrade transformers" error only when the user actually
-# tries to load a Gemma4 model.
-try:
-    from .modeling_gemma4 import Gemma4ForCausalLM  # noqa
-    from .modeling_gemma4mm import Gemma4ForConditionalGeneration  # noqa
-
-    __all__.extend(["Gemma4ForCausalLM", "Gemma4ForConditionalGeneration"])
-except ImportError:
-    pass

tensorrt_llm/_torch/models/checkpoints/hf/gemma4_weight_mapper.py

@@ -27,6 +27,7 @@
 
 @register_mapper("HF", "Gemma4ForCausalLM")
 @register_mapper("HF", "Gemma4ForConditionalGeneration")
+@register_mapper("HF", "Gemma4UnifiedForConditionalGeneration")
 class Gemma4HfWeightMapper(HfWeightMapper):
     @property
     def _is_vlm(self) -> bool:

tensorrt_llm/_torch/models/modeling_gemma4.py

@@ -954,14 +954,17 @@ def get_model_defaults(cls, llm_args) -> dict:
     def _get_token_type_mask(self, mm_token_type_ids: torch.Tensor):
         """Build bidirectional attention mask from mm_token_type_ids.
 
-        mm_token_type_ids: 0=text, 1=image, 2=video (or any positive int for
-        a modality blob). Tokens within the same contiguous blob of the same
-        modality attend bidirectionally to each other.
+        mm_token_type_ids: 0=text, 1=image, 2=video, 3=audio. Only VISION
+        tokens (image/video) attend bidirectionally within their contiguous
+        blob; text and audio stay causal. Matches HF Gemma4, where
+        ``is_vision = (mm_token_type_ids == 1) | (mm_token_type_ids == 2)`` and
+        audio is left causal.
         """
         device = mm_token_type_ids.device
         token_type_ids = mm_token_type_ids.clone()
-        # We only care about non-zero (multimodal) tokens
-        is_mm = token_type_ids > 0
+        # Only vision tokens (image=1, video=2) get the bidirectional mask;
+        # audio (3) stays causal, matching HF Gemma4 (audio is not in is_vision).
+        is_mm = (token_type_ids == 1) | (token_type_ids == 2)
 
         # Detect blob boundaries: positions where type changes or goes 0->nonzero
         padded = torch.cat(