Merge pull request #3420 from AI-Hypercomputer:amandaliang

PiperOrigin-RevId: 884673900

Author: Google-ML-Automation

src/maxtext/models/deepseek_batchsplit.py

@@ -804,16 +804,11 @@ def gmm(
       input_buffer_count,
       combine_scopes,
   ):
-
-    tokamax_group_sizes = tokamax.RaggedDotGroupSizes(
-        group_sizes,
-        max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
-    )
     if config.use_qwix_quantization:
       output = megablox.gmm(
           lhs=inputs,
           rhs=kernel,
-          group_sizes=tokamax_group_sizes,
+          group_sizes=group_sizes,
           preferred_element_type=preferred_element_type,
           tiling=tiling,
           use_qwix_quantization=config.use_qwix_quantization,
@@ -827,7 +822,10 @@ def gmm(
       output = tokamax.ragged_dot(
           lhs=inputs,
           rhs=kernel,
-          group_sizes=tokamax_group_sizes,
+          group_sizes=tokamax.RaggedDotGroupSizes(
+              group_sizes,
+              max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
+          ),
           precision=jax.lax.Precision.DEFAULT,
           preferred_element_type=preferred_element_type,
           implementation="mosaic",