Test Fused GDN Backward Pass Kernel w/ MaxText

src/maxtext/configs/base.yml

@@ -1344,6 +1344,8 @@ gdn_num_value_heads: 32
 gdn_chunk_size: 64
 # Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.
 use_qk_norm_in_gdn: true
+# Whether to use the fused analytical Pallas GDN kernel
+use_gdn_kernel: false
 # The ratio of dimension to apply ROPE on
 partial_rotary_factor: 1.0
 

src/maxtext/configs/models/qwen3.5-tiny.yml

@@ -0,0 +1,72 @@
+# Copyright 2026 Google LLC
+#
+# 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
+#
+#    https://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.
+
+# Tiny version of Qwen3.5 for fast testing and execution.
+
+decoder_block: "qwen3_5"
+
+# Core Architectural Parameters
+base_emb_dim: 256
+base_num_decoder_layers: 2
+base_num_query_heads: 4
+base_num_kv_heads: 2
+head_dim: 256
+vocab_size: 248320
+normalization_layer_epsilon: 1.0e-6
+
+# MoE Specific Parameters
+# Set base_mlp_dim to match base_moe_mlp_dim to pass validation for fully MoE models.
+base_mlp_dim: 256
+base_moe_mlp_dim: 256
+num_experts: 8
+shared_experts: 1
+num_experts_per_tok: 2
+norm_topk_prob: true
+
+# GatedDeltaNet Specific Parameters for Linear Attention (GDN)
+inhomogeneous_layer_cycle_interval: 2
+gdn_conv_kernel_dim: 4
+gdn_key_head_dim: 128
+gdn_value_head_dim: 128
+gdn_num_key_heads: 2
+gdn_num_value_heads: 4
+gdn_chunk_size: 64
+
+# RoPE Settings
+rope_max_timescale: 10000000
+partial_rotary_factor: 0.25
+
+# General Model Settings
+enable_dropout: false
+
+# Vision Encoder Configuration (need to set use_multimodal=true)
+vision_encoder_block: "qwen3_5"
+# Based on Qwen3.5 MoE Vision Model Config
+image_size_for_vit: 768
+hidden_size_for_vit: 1152
+intermediate_size_for_vit: 4304
+num_attention_heads_for_vit: 16
+num_hidden_layers_for_vit: 27
+num_channels_for_vit: 3
+patch_size_for_vit: 16
+temporal_patch_size_for_vit: 2
+spatial_merge_size_for_vit: 2
+out_hidden_size_for_vit: 256  # Projects to decoder emb_dim (256)
+num_position_embeddings_for_vit: 2304
+deepstack_visual_indexes_for_vit: []  # No deepstack for Qwen3.5 VL
+rope_theta_for_vit: 10000
+
+# MRoPE Settings (Multi-dimensional RoPE for multimodal)
+use_mrope: true
+mrope_section: [11, 11, 10]

src/maxtext/configs/types.py

@@ -282,6 +282,7 @@ class ProfilerType(str, Enum):
     "qwen3-custom-30b-a3b",
     "qwen3.5-35b-a3b",
     "qwen3.5-397b-a17b",
+    "qwen3.5-tiny",
     "gpt3-175b",
     "gpt3-22b",
     "gpt3-6b",
@@ -1144,6 +1145,10 @@ class Qwen3Next(BaseModel):
       True,
       description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.",
   )
+  use_gdn_kernel: bool = Field(
+      False,
+      description="Whether to use the fused analytical Pallas GDN kernel.",
+  )
   partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on")
 
 
@@ -4271,6 +4276,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "qwen3-vl-30b-a3b",
           "qwen3.5-35b-a3b",
           "qwen3.5-397b-a17b",
+          "qwen3.5-tiny",
           "maxtext-omni-gemma3-qwen3",
       )
       if self.model_name not in valid_mm_models and self.model_name != "default":

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py

@@ -1090,7 +1090,7 @@ def validate_inputs(
 
   assert group_offset.shape == (1,)
 
-  size_lhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype)
+  size_lhs_sublane = max(pltpu.get_tpu_info().get_sublane_tiling(lhs.dtype), 16)
   size_lhs_sublane = min(size_lhs_sublane, size_m)
   if fuse_act is not None:
     num_lanes = pltpu.get_tpu_info().num_lanes

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py

@@ -221,6 +221,9 @@ def make_tgmm_configs(
   size_lhs_sublane = min(size_lhs_sublane, size_m)
   size_rhs_sublane = pltpu.get_tpu_info().get_sublane_tiling(rhs.dtype)
   size_rhs_sublane = min(size_rhs_sublane, size_m)
+  common_sublane = min(size_lhs_sublane, size_rhs_sublane)
+  size_lhs_sublane = common_sublane
+  size_rhs_sublane = common_sublane
   assert size_lhs_sublane == size_rhs_sublane, (
       f"size_lhs_sublane should be the same as size_rhs_sublane {lhs.dtype=}," f" {rhs.dtype=}"
   )

src/maxtext/layers/decoders.py

@@ -724,6 +724,7 @@ def _apply_embedding(
             "qwen3-vl-30b-a3b",
             "qwen3.5-35b-a3b",
             "qwen3.5-397b-a17b",
+            "qwen3.5-tiny",
             "maxtext-omni-gemma3-qwen3",
         ]:
           y = mm_utils.merge_mm_embeddings(
@@ -744,6 +745,7 @@ def _apply_embedding(
             "qwen3-vl-30b-a3b",
             "qwen3.5-35b-a3b",
             "qwen3.5-397b-a17b",
+            "qwen3.5-tiny",
         ]:
           y = mm_utils.merge_mm_embeddings(
               text_embeddings=y,

src/maxtext/layers/nnx_decoders.py

@@ -1423,6 +1423,7 @@ def _apply_embedding(
             "qwen3-vl-30b-a3b",
             "qwen3.5-35b-a3b",
             "qwen3.5-397b-a17b",
+            "qwen3.5-tiny",
             "maxtext-omni-gemma3-qwen3",
         }:
           y = mm_utils.merge_mm_embeddings(
@@ -1442,6 +1443,7 @@ def _apply_embedding(
             "qwen3-vl-30b-a3b",
             "qwen3.5-35b-a3b",
             "qwen3.5-397b-a17b",
+            "qwen3.5-tiny",
         }:
           y = mm_utils.merge_mm_embeddings(
               text_embeddings=y,