From dc98a8953dff6d2f2787d122eb219e3089a7b7ad Mon Sep 17 00:00:00 2001 From: Suraj Kolla Date: Mon, 17 Aug 2026 11:57:02 -0700 Subject: [PATCH] Add support for Tokamax GMM v2 heuristic tiling function. # Description This change introduces a new configuration option, use_gmm_v2_heuristic_tiling, which allows MaxText to use the heuristic tiling functions from Tokamax GMM v2 and TGMM v2 instead of custom tile sizes. The option is integrated into the MoE layer, Megablox GMM operations, configuration validation, and documentation. An integration test has also been added to verify the heuristic tiling path. # Tests Expanded tokamax_test.py to cover model execution and XLA compilation with the GMM v2 heuristic tiling fallback. # Checklist Before submitting this PR, please make sure (put X in square brackets): - [X] I have performed a self-review of my code. - [X] I have necessary comments in my code, particularly in hard-to-understand areas. - [X] I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation. PiperOrigin-RevId: 966100631 --- .../core_concepts/moe_configuration.md | 4 +- src/maxtext/configs/base.yml | 3 ++ src/maxtext/configs/types.py | 8 ++++ src/maxtext/kernels/megablox/ops.py | 47 ++++++++++++++----- src/maxtext/layers/moe.py | 1 + tests/integration/tokamax_test.py | 20 ++++---- tests/unit/pyconfig_test.py | 8 ++++ 7 files changed, 69 insertions(+), 22 deletions(-) diff --git a/docs/reference/core_concepts/moe_configuration.md b/docs/reference/core_concepts/moe_configuration.md index 4f6945d425..489a5acaf2 100644 --- a/docs/reference/core_concepts/moe_configuration.md +++ b/docs/reference/core_concepts/moe_configuration.md @@ -93,6 +93,8 @@ MaxText implements an exact, paper-aligned version of DeepSeek V4's load balanci `use_gmm_v2`: If enabled, use the Tokamax GMM v2 kernel for grouped matrix multiplication. Requires `use_tokamax_gmm` to be True. +`use_gmm_v2_heuristic_tiling`: If enabled, use the heuristic tiling from Tokamax GMM v2. Recommended when not using custom tuned tile sizes. + `megablox`: If enabled, use Megablox for sparse matrix operations. Effective only when `use_tokamax_gmm` is False. `capacity_factor`: A scalar multiplier for expert capacity. Effective only when `sparse_matmul` is False. @@ -155,5 +157,5 @@ Implementation Support: - Tokamax Ragged Dot (Includes two implementations): - **GMM v1**: Uses Tokamax's native autotuner; does not accept manual tile sizes from MaxText. - - **GMM v2**: Supports all 18 manual tiling configurations. + - **GMM v2**: Supports all 18 manual tiling configurations. Optionally, use `use_gmm_v2_heuristic_tiling=True` for heuristic tiling. - Enabled for FP8 and BF16. diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index a902ec7892..d1156f4f55 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -274,6 +274,9 @@ merge_gating_gmm: false use_tokamax_gmm: false # Whether to use Tokamax GMM v2 for MoE kernel. Requires use_tokamax_gmm=true. use_gmm_v2: false +# Whether to use the heuristic tiling function from Tokamax GMM v2, when use_gmm_v2=true. +# When enabled, custom tile sizes (`wi_tile...`, `wo_tile...`) are ignored. +use_gmm_v2_heuristic_tiling: false norm_topk_prob: false # boolean to enable the top-k probability normalization. qwen3-specific normalization of router weights. diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index db9a68d1cd..fad06f105c 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -1041,6 +1041,11 @@ class MoEKernels(BaseModel): description="Whether to use Tokamax GMM v2 for MoE kernel.", ) + use_gmm_v2_heuristic_tiling: bool = Field( + False, + description="Whether to use the heuristic tiling from Tokamax GMM v2, when use_gmm_v2=true.", + ) + class DeepSeekMoE(BaseModel): """Configuration specific to DeepSeek-style MoE layers.""" @@ -4300,6 +4305,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de if self.use_batch_split_schedule: raise ValueError("GMM v2 is not supported with a batch split schedule.") + if self.use_gmm_v2_heuristic_tiling and not self.use_gmm_v2: + raise ValueError("`use_gmm_v2_heuristic_tiling=True` requires `use_gmm_v2=True`.") + for val in self.compress_ratios: if val != 0 and val < 4: raise ValueError(f"compress_ratio must be 0 (disabled) or >= 4, got {val}") diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index c717eda455..f88e0aa0ab 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -74,6 +74,7 @@ def gmm( qwix_rule: qwix.QtRule | None = None, use_manual_quantization: bool = False, # used in batchsplit use_gmm_v2: bool = False, + use_gmm_v2_heuristic_tiling: bool = False, partial_sum: jnp.ndarray | None = None, ): """Grouped matrix multiplication operation.""" @@ -105,7 +106,7 @@ def gmm( gmm_fwd_bwd = lambda *args: _gmm_fwd(*args)[0] # pylint: disable=C3001 gmm_fwd_bwd = jax.custom_vjp( gmm_fwd_bwd, - nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15), + nondiff_argnums=(3, 4, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16), ) gmm_fwd_bwd.defvjp(_gmm_fwd, functools.partial(_gmm_bwd, lhs.dtype, rhs.dtype)) return gmm_fwd_bwd( @@ -125,6 +126,7 @@ def gmm( lhs_vma_axes, rhs_vma_axes, use_gmm_v2, + use_gmm_v2_heuristic_tiling, partial_sum, ) @@ -161,6 +163,7 @@ def _gmm_fwd( lhs_vma_axes: tuple = tuple(), rhs_vma_axes: tuple = tuple(), use_gmm_v2: bool = False, + use_gmm_v2_heuristic_tiling: bool = False, partial_sum: jnp.ndarray | None = None, ) -> tuple[ jnp.ndarray, @@ -207,6 +210,7 @@ def _gmm_fwd( group_sizes, preferred_element_type, tiling, + use_gmm_v2_heuristic_tiling, group_offset, partial_sum, transpose_rhs, @@ -366,6 +370,7 @@ def _fwd_run_tokamax_v2( group_sizes: jnp.ndarray, preferred_element_type: jnp.dtype, tiling: tuple, + use_gmm_v2_heuristic_tiling: bool, group_offset: jnp.ndarray | None, partial_sum: jnp.ndarray | None, transpose_rhs: bool, @@ -382,18 +387,17 @@ def _fwd_run_tokamax_v2( rhs_operand = rhs_operand.qvalue rhs_scale = _fwd_prepare_rhs_scale(rhs, transpose_rhs=transpose_rhs) - custom_fwd_tiling = gmm_v2.TileSizes( - tile_m=tiling[0], - tile_k=tiling[1], - tile_n=tiling[2], - ) + if use_gmm_v2_heuristic_tiling: + fwd_tiling = gmm_v2.calculate_tiling + else: + fwd_tiling = gmm_v2.TileSizes(tile_m=tiling[0], tile_k=tiling[1], tile_n=tiling[2]) return gmm_v2.gmm_v2( lhs=lhs, # pyrefly: ignore[bad-argument-type] rhs=rhs_operand, # pyrefly: ignore[bad-argument-type] group_sizes=group_sizes, rhs_scale=rhs_scale, - tile_info=custom_fwd_tiling, + tile_info=fwd_tiling, preferred_element_type=preferred_element_type, partial_sum=partial_sum, group_offset=group_offset, @@ -449,6 +453,7 @@ def _gmm_bwd( lhs_vma_axes: tuple, rhs_vma_axes: tuple, use_gmm_v2: bool, + use_gmm_v2_heuristic_tiling: bool, residual: tuple[ jnp.ndarray | qpl.QArray, jnp.ndarray | qpl.QArray, @@ -495,6 +500,7 @@ def _gmm_bwd( use_manual_quantization, interpret, lhs_vma_axes, + use_gmm_v2_heuristic_tiling, ) # 4. DRHS Gradient Execution @@ -513,6 +519,7 @@ def _gmm_bwd( interpret, rhs_vma_axes, quantization_rule, + use_gmm_v2_heuristic_tiling, ) # 5. Output Formatting @@ -622,6 +629,7 @@ def _compute_dlhs( use_manual_quantization: bool, interpret: bool, lhs_vma_axes: tuple, + use_gmm_v2_heuristic_tiling: bool, ) -> jnp.ndarray: """Routes execution of DLHS based on backend choices.""" if use_tokamax_backend and not use_gmm_v2: @@ -634,7 +642,9 @@ def _compute_dlhs( use_manual_quantization, ) elif use_tokamax_backend and use_gmm_v2: - return _dlhs_run_tokamax_v2(dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, transpose_rhs) + return _dlhs_run_tokamax_v2( + dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, use_gmm_v2_heuristic_tiling, transpose_rhs + ) else: return _dlhs_run_megablox( dlhs_dout, rhs, group_sizes, group_offset, lhs_dtype, tiling, transpose_rhs, interpret, lhs_vma_axes @@ -707,6 +717,7 @@ def _dlhs_run_tokamax_v2( group_offset: jnp.ndarray | None, lhs_dtype: jax.typing.DTypeLike, tiling: tuple, + use_gmm_v2_heuristic_tiling: bool, transpose_rhs: bool, ) -> jnp.ndarray: """Executes Tokamax GMM V2 backend for DLHS = DLHS_dout @ RHS^T.""" @@ -714,7 +725,10 @@ def _dlhs_run_tokamax_v2( dlhs_rhs = rhs if transpose_rhs else rhs.swapaxes(1, 2) dlhs_lhs = dlhs_dout.qvalue if isinstance(dlhs_dout, qpl.QArray) else dlhs_dout - custom_dlhs_tiling = gmm_v2.TileSizes(tile_m=tiling[3], tile_k=tiling[4], tile_n=tiling[5]) + if use_gmm_v2_heuristic_tiling: + dlhs_tiling = gmm_v2.calculate_tiling + else: + dlhs_tiling = gmm_v2.TileSizes(tile_m=tiling[3], tile_k=tiling[4], tile_n=tiling[5]) dlhs = gmm_v2.gmm_v2( lhs=dlhs_lhs, @@ -722,7 +736,7 @@ def _dlhs_run_tokamax_v2( group_sizes=group_sizes, # rhs scale is already applied to dlhs_lhs rhs_scale=None, - tile_info=custom_dlhs_tiling, + tile_info=dlhs_tiling, preferred_element_type=lhs_dtype, # pyrefly: ignore[bad-argument-type] group_offset=group_offset, ) @@ -778,12 +792,15 @@ def _compute_drhs( interpret: bool, rhs_vma_axes: tuple, quantization_rule: qwix.QtRule | None, + use_gmm_v2_heuristic_tiling: bool, ) -> jnp.ndarray: """Routes execution of DRHS based on backend choices.""" if use_tokamax_backend and not use_gmm_v2: drhs = _drhs_run_tokamax_v1(drhs_dout, lhs, group_sizes, rhs_dtype, use_manual_quantization) elif use_tokamax_backend and use_gmm_v2: - drhs = _drhs_run_tokamax_v2(drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling) + drhs = _drhs_run_tokamax_v2( + drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, use_gmm_v2_heuristic_tiling + ) else: drhs = _drhs_run_megablox( drhs_dout, lhs, group_sizes, group_offset, num_actual_groups, rhs_dtype, tiling, interpret, rhs_vma_axes @@ -850,6 +867,7 @@ def _drhs_run_tokamax_v2( num_actual_groups: int, rhs_dtype: jax.typing.DTypeLike, tiling: tuple, + use_gmm_v2_heuristic_tiling: bool, ) -> jnp.ndarray: """Executes Tokamax TGMM V2 backend for DRHS = LHS^T @ DRHS_dout.""" drhs_rhs = drhs_dout.qvalue if isinstance(drhs_dout, qpl.QArray) else drhs_dout @@ -859,7 +877,10 @@ def _drhs_run_tokamax_v2( if isinstance(drhs_dout, qpl.QArray): rhs_scale = _drhs_prepare_bwd_scale(drhs_dout) - custom_drhs_tiling = gmm_v2.TileSizes(tile_m=tiling[6], tile_k=tiling[7], tile_n=tiling[8]) + if use_gmm_v2_heuristic_tiling: + drhs_tiling = tgmm_v2.calculate_tgmm_tiling + else: + drhs_tiling = gmm_v2.TileSizes(tile_m=tiling[6], tile_k=tiling[7], tile_n=tiling[8]) return tgmm_v2.tgmm_v2( lhs=drhs_lhs, @@ -870,7 +891,7 @@ def _drhs_run_tokamax_v2( precision=jax.lax.Precision.DEFAULT, preferred_element_type=rhs_dtype, # pyrefly: ignore[bad-argument-type] group_offset=group_offset, - tile_info=custom_drhs_tiling, + tile_info=drhs_tiling, ) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index da9e86e320..3eef3d4e17 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -1559,6 +1559,7 @@ def extract_vma(tensor): lhs_vma_axes=lhs_vma_axes, rhs_vma_axes=rhs_vma_axes, use_gmm_v2=self.config.use_gmm_v2, + use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling, partial_sum=partial_sum, interpret=megablox_interpret, ) diff --git a/tests/integration/tokamax_test.py b/tests/integration/tokamax_test.py index 784dc10607..3c6cab0529 100644 --- a/tests/integration/tokamax_test.py +++ b/tests/integration/tokamax_test.py @@ -37,16 +37,18 @@ class Train(parameterized.TestCase): "testcase_name": f"{base_name}_ep{ici_expert_parallelism}", "quantization": quantization, "use_gmm_v2": use_gmm_v2, + "use_gmm_v2_heuristic_tiling": use_gmm_v2_heuristic_tiling, "ici_expert_parallelism": ici_expert_parallelism, } - for base_name, quantization, use_gmm_v2, ici_expert_parallelism in [ - ("tokamax_v1_bf16", "", False, 1), - ("tokamax_v1_fp8", "fp8", False, 1), # not quantize gmm - ("tokamax_v1_fp8_full", "fp8_full", False, 1), # quantize gmm - ("tokamax_v2_bf16", "", True, 1), - ("tokamax_v2_fp8_full", "fp8_full", True, 1), - ("tokamax_v2_bf16", "", True, 2), - ("tokamax_v2_fp8_full", "fp8_full", True, 2), + for base_name, quantization, use_gmm_v2, use_gmm_v2_heuristic_tiling, ici_expert_parallelism in [ + ("tokamax_v1_bf16", "", False, False, 1), + ("tokamax_v1_fp8", "fp8", False, False, 1), # not quantize gmm + ("tokamax_v1_fp8_full", "fp8_full", False, False, 1), # quantize gmm + ("tokamax_v2_bf16", "", True, False, 1), + ("tokamax_v2_bf16_heuristic", "", True, True, 1), + ("tokamax_v2_fp8_full", "fp8_full", True, False, 1), + ("tokamax_v2_bf16", "", True, False, 2), + ("tokamax_v2_fp8_full", "fp8_full", True, False, 2), ] ) @pytest.mark.tpu_only @@ -54,6 +56,7 @@ def test_smoke_train( self, quantization: str, use_gmm_v2: bool, + use_gmm_v2_heuristic_tiling: bool, ici_expert_parallelism: int, ): """Smoke train with small config.""" @@ -84,6 +87,7 @@ def test_smoke_train( "megablox=False", "use_tokamax_gmm=True", f"use_gmm_v2={use_gmm_v2}", + f"use_gmm_v2_heuristic_tiling={use_gmm_v2_heuristic_tiling}", # tile sizes "wi_tile_fwd_batch_seq=128", "wi_tile_fwd_embed_dim=128", diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py index 2ded9c12d2..eecbc40d81 100644 --- a/tests/unit/pyconfig_test.py +++ b/tests/unit/pyconfig_test.py @@ -40,6 +40,14 @@ def test_empty_string_parse_as_empty_string(self): self.assertTrue(config.quantization is None or config.quantization == "") + def test_gmm_v2_heuristic_tiling_requires_gmm_v2(self): + with self.assertRaisesRegex(ValueError, "`use_gmm_v2_heuristic_tiling=True` requires `use_gmm_v2=True`."): + pyconfig.initialize( + [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()], + use_gmm_v2_heuristic_tiling=True, + use_gmm_v2=False, + ) + def test_managed_mldiagnostics_storage_path(self): # Test completely omitting the parameter (defaults to "" from base.yml) config_omitted = pyconfig.initialize(