Correct the tokamax.RaggedDotGroupSizes configs

The key generated for looking up the tokamax autotune cache in maxtext was a list, `"group_sizes":[1024, 1024, ...]`, while in the cache file, it is `"group_sizes":{"num_groups":256,"total_size":262144}`, causing the autotune cache miss.

Change it to `inputs.shape[0]` only will produce the key that has `"group_sizes":{"num_groups":256,"total_size":262144}`

PiperOrigin-RevId: 948385403

Author: Google-ML-Automation

src/maxtext/layers/moe.py

@@ -1401,10 +1401,9 @@ def get_tokamax_group_sizes(group_sizes, inputs, kernel):
       elif self.config.attention == "vllm_rpa":
         return group_sizes
       else:
-        num_groups = group_sizes.shape[0]
         return tokamax.RaggedDotGroupSizes(
             group_sizes,
-            (inputs.shape[0] // num_groups,) * num_groups,
+            inputs.shape[0],
         )
 
     def get_quantization_dtypes():