Merge deepseek_v4_core_primitives into dsv4-moe-routing-primitives

Author: parambole

src/maxtext/layers/linears.py

@@ -603,7 +603,7 @@ def __init__(
     # Grouped block-diagonal projection kernel parameters
     # Kernels are stored as a 3D tensor: [n_groups, in_features_per_group, out_features_per_group]
     kernel_shape = (n_groups, in_features_per_group, self.out_features_per_group)
-    self.weight = nnx.Param(
+    self.kernel = nnx.Param(
         kernel_init(
             rngs.params(),
             kernel_shape,
@@ -623,10 +623,10 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
       Projected tensor of shape [..., n_groups, out_features_per_group]
     """
     x = jnp.asarray(x, self.dtype)
-    weight = jnp.asarray(self.weight[...], self.dtype)
+    kernel = jnp.asarray(self.kernel[...], self.dtype)
 
     # Execute parallel group projection via optimized einsum broadcasting.
     # x: [..., g, i]
-    # weight: [g, i, o]
+    # kernel: [g, i, o]
     # output: [..., g, o]
-    return jnp.einsum("...gi,gio->...go", x, weight)
+    return jnp.einsum("...gi,gio->...go", x, kernel)

src/maxtext/layers/normalizations.py

@@ -243,7 +243,7 @@ def l2norm(x: Array, dim: int = -1, eps: float = 1e-6) -> Array:
 
 
 class DeepSeekV4RMSNorm(nnx.Module):
-  """RMS normalization for DeepSeek-V4 (equivalent to T5LayerNorm)."""
+  """RMS normalization for DeepSeek-V4."""
 
   def __init__(
       self,
@@ -257,8 +257,9 @@ def __init__(
     self.dtype = dtype
     self.weight_dtype = weight_dtype
 
-    # Initialize learnable scale weight to ones matching T5LayerNorm behavior
+    # Initialize learnable scale weight to ones
     self.weight = nnx.Param(jnp.ones((hidden_size,), dtype=weight_dtype))
+    self.scale = self.weight
 
   def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
     # [B, S, D] where D = hidden_size
@@ -296,6 +297,7 @@ def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
     # Calculate variance across features axis
     variance = jnp.mean(lax.square(x_f32), axis=-1, keepdims=True)  # [..., 1]
 
-    # Apply reciprocal square root and cast back to active precision
-    normalized = x_f32 * lax.rsqrt(variance + self.eps)  # [..., D]
-    return jnp.asarray(normalized, self.dtype)  # [..., D]
+    # Apply reciprocal square root, cast to active precision, and multiply
+    inv_norm = jnp.asarray(lax.rsqrt(variance + self.eps), self.dtype)  # [..., 1]
+    x_active = jnp.asarray(x, self.dtype)  # [..., D]
+    return x_active * inv_norm  # [..., D]