Merge pull request #2733 from AI-Hypercomputer:hengtaoguo-kvcache

PiperOrigin-RevId: 836754833

Author: Google-ML-Automation

src/MaxText/layers/gemma3.py

@@ -382,7 +382,7 @@ def __init__(
         matmul_precision=self.config.matmul_precision,
         rngs=self.rngs,
     )
-    self.Dropout_0 = nnx.Dropout(rate=self.config.dropout_rate)
+    self.Dropout_0 = Dropout(rate=self.config.dropout_rate, rngs=self.rngs)
     self.Dense_1 = DenseGeneral(
         in_features_shape=self.config.intermediate_size_for_vit,
         out_features_shape=self.config.hidden_size_for_vit,
@@ -452,12 +452,12 @@ def __init__(
         config=self.config,
         rngs=self.rngs,
     )
-    self.Dropout_0 = nnx.Dropout(self.config.dropout_rate, rngs=self.rngs)
+    self.Dropout_0 = Dropout(rate=self.config.dropout_rate, rngs=self.rngs)
 
   def __call__(self, x: jax.Array, deterministic: bool = False) -> jax.Array:
     y = self.LayerNorm_0(x)
 
-    y = self.MultiHeadDotProductAttention_0(inputs_q=y, inputs_kv=y, deterministic=deterministic)
+    y, _ = self.MultiHeadDotProductAttention_0(inputs_q=y, inputs_kv=y, deterministic=deterministic)
     y = self.Dropout_0(y, deterministic=deterministic)
     x = x + y
 
@@ -634,7 +634,7 @@ def __init__(
         width=self.config.hidden_size_for_vit,
         dtype=self.config.dtype_mm,
     )
-    self.Dropout_0 = nnx.Dropout(self.config.dropout_rate, rngs=self.rngs)
+    self.Dropout_0 = Dropout(rate=self.config.dropout_rate, rngs=self.rngs)
     self.Transformer = Encoder(
         config=self.config,
         mesh=self.mesh,