[Stacked PR 4/5] Enable GDN backward pass kernel and hybrid precision in Qwen3 model

src/maxtext/configs/base.yml

@@ -123,6 +123,10 @@ grad_dtype: "float32"
 
 # activation dtypes.
 dtype: "bfloat16"
+
+# GDN precision configuration
+gdn_state_dtype: "float32"
+gdn_decay_dtype: "float32"
 # used to configure quantization in the transformer layers, defaults to null implying bf16.
 # possible alternative settings are as follows:
 # 'int8' for dynamic range quantization using 8-bits
@@ -1357,6 +1361,8 @@ gdn_chunk_size: 64
 use_qk_norm_in_gdn: true
 # The ratio of dimension to apply ROPE on
 partial_rotary_factor: 1.0
+# Whether to use GDN Pallas kernel
+use_gdn_kernel: false
 
 use_tokamax_splash: false
 # Setting this flag will use a non-pallas implementation.

src/maxtext/configs/pyconfig.py

@@ -337,6 +337,8 @@ def __init__(self, pydantic_config: types.MaxTextConfig):
     final_dict["dtype"] = jnp.dtype(final_dict["dtype"])
     final_dict["grad_dtype"] = jnp.dtype(final_dict["grad_dtype"])
     final_dict["weight_dtype"] = jnp.dtype(final_dict["weight_dtype"])
+    final_dict["gdn_state_dtype"] = jnp.dtype(final_dict.get("gdn_state_dtype", "float32"))
+    final_dict["gdn_decay_dtype"] = jnp.dtype(final_dict.get("gdn_decay_dtype", "float32"))
     final_dict["mu_dtype"] = (
         final_dict["weight_dtype"] if not final_dict["mu_dtype"] else jnp.dtype(final_dict["mu_dtype"])
     )

src/maxtext/configs/types.py

@@ -480,6 +480,11 @@ class DataTypes(BaseModel):
       description="If True, sets activations to float32 before the nonlinearity.",
   )
   dtype_mm: str = Field("float32", description="Data type for multimodal model's vision encoder")
+  gdn_state_dtype: DType = Field(DType.FLOAT32, description="The data type for GDN recurrent states.")
+  gdn_decay_dtype: DType = Field(
+      DType.FLOAT32,
+      description="The data type for GDN decay parameters (A_log, dt_bias).",
+  )
 
 
 class Quantization(BaseModel):
@@ -1154,6 +1159,10 @@ class Qwen3Next(BaseModel):
       description="Whether to apply L2 normalization to query and key tensors inside the Gated Delta Rule kernel.",
   )
   partial_rotary_factor: float = Field(1.0, description="The ratio of dimension to apply ROPE on")
+  use_gdn_kernel: bool = Field(
+      False,
+      description="Whether to use GDN Pallas kernel.",
+  )
 
 
 # ----------------------------------------------------------------------------