Merge pull request #3385 from AI-Hypercomputer:chengnuojin-remove-rep

Google-ML-Automation · Google-ML-Automation · commit 63f5ca9cb29b · 2026-03-11T17:16:14.000-07:00
PiperOrigin-RevId: 882275144
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
@@ -898,7 +898,7 @@ def gmm(
     ):
       tokamax_group_sizes = tokamax.RaggedDotGroupSizes(
           group_sizes,
-          representative_value=max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
+          max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
       )
       pad_length = self.config.wi_tile_fwd_batch_seq
       hs_shape = inputs.shape
diff --git a/src/maxtext/models/deepseek_batchsplit.py b/src/maxtext/models/deepseek_batchsplit.py
@@ -807,7 +807,7 @@ def gmm(
 
     tokamax_group_sizes = tokamax.RaggedDotGroupSizes(
         group_sizes,
-        representative_value=max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
+        max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),
     )
     if config.use_qwix_quantization:
       output = megablox.gmm(

Original file line number	Diff line number	Diff line change
`@@ -898,7 +898,7 @@ def gmm(`
`898`	`898`	`):`
`899`	`899`	`tokamax_group_sizes = tokamax.RaggedDotGroupSizes(`
`900`	`900`	`group_sizes,`
`901`		`- representative_value=max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),`
	`901`	`+ max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),`
`902`	`902`	`)`
`903`	`903`	`pad_length = self.config.wi_tile_fwd_batch_seq`
`904`	`904`	`hs_shape = inputs.shape`
Original file line number	Diff line number	Diff line change
`@@ -807,7 +807,7 @@ def gmm(`
`807`	`807`
`808`	`808`	`tokamax_group_sizes = tokamax.RaggedDotGroupSizes(`
`809`	`809`	`group_sizes,`
`810`		`- representative_value=max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),`
	`810`	`+ max_utils.generate_representative_group_sizes(inputs.shape[0], kernel.shape[0]),`
`811`	`811`	`)`
`812`	`812`	`if config.use_qwix_quantization:`
`813`	`813`	`output = megablox.gmm(`