Support specifying tokamax gmm tile sizes in MaxText

[darisoy]

Description

This PR enables cleanly overriding the GMM tile sizes for both the forward and backward passes of Tokamax (ragged_dot) in MaxText using an elegant global heuristics monkey-patching approach.

Proposed Changes

• Heuristics Monkey-Patching (src/maxtext/layers/moe.py): Instead of introducing complex custom vjp function wrappers or polluting model layers with low-level compilation code, this PR dynamically monkey-patches PallasMosaicTpuRaggedDot._get_heuristics_config globally during RoutedMoE initialization.
• Dynamic Shapes & Pass Routing: The monkey-patch dynamically inspects the shapes of the operands (lhs/rhs) and the JAX dimension numbers. This allows it to detect the active pass (Forward, Backward DLHS, or Backward DRHS) and layer type (Weight Input wi or Weight Output wo) on the fly, applying your exact manual tile configurations:
  • wi_tile_fwd, wi_tile_dlhs, wi_tile_drhs configs.
  • wo_tile_fwd, wo_tile_dlhs, wo_tile_drhs configs.

FIXES: b/506157856

Tests

Verified compilation and trace collection via a 5-step synthetic pretraining run on a v6e-4 VM using the deepseek3-671b MoE model configuration with your manual tiling overrides:

JAX_PLATFORMS=tpu,cpu PYTHONPATH=src python3 -m maxtext.trainers.pre_train.train \
    run_name=darisoy-moe-train \
    base_output_directory=gs://maxtext-experiments-multipod \
    model_name=deepseek3-671b \
    dataset_type=synthetic \
    steps=5 \
    enable_checkpointing=False \
    enable_goodput_recording=False \
    attention_type=mla \
    sparse_matmul=True \
    megablox=True \
    per_device_batch_size=2 \
    base_num_decoder_layers=1 \
    first_num_dense_layers=0 \
    shared_experts=1 \
    base_mlp_dim=2048 \
    profiler=xplane \
    use_tokamax_gmm=True \
    override_model_config=True \
    max_target_length=4096 \
    ici_expert_parallelism=4 \
    use_ring_of_experts=True \
    tokamax_gmm_autotune=True \
    weight_dtype=bfloat16 \
    opt_type=sgd \
    wi_tile_fwd_batch_seq=128 \
    wi_tile_fwd_embed_dim=128 \
    wi_tile_fwd_mlp_dim=128 \
    wi_tile_dlhs_batch_seq=256 \
    wi_tile_dlhs_embed_dim=256 \
    wi_tile_dlhs_mlp_dim=256 \
    wi_tile_drhs_batch_seq=512 \
    wi_tile_drhs_embed_dim=512 \
    wi_tile_drhs_mlp_dim=512 \
    wo_tile_fwd_batch_seq: 128 \
    wo_tile_fwd_embed_dim: 128 \
    wo_tile_fwd_mlp_dim: 128 \
    wo_tile_dlhs_batch_seq: 256 \
    wo_tile_dlhs_embed_dim: 256 \
    wo_tile_dlhs_mlp_dim: 256 \
    wo_tile_drhs_batch_seq: 512 \
    wo_tile_drhs_embed_dim: 512 \
    wo_tile_drhs_mlp_dim: 512

• Execution Status: Completed successfully with no shape or compilation errors.
• Xprof Trace URL: https://xprof.corp.google.com/trace_viewer/darisoy-440103416108113419

See details in http://b/506157856

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• 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.

[codecov]

Codecov Report

❌ Patch coverage is 82.35294% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/moe.py	87.50%	2 Missing and 2 partials ⚠️
src/maxtext/models/deepseek_batchsplit_fp8.py	0.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

Lets try to find a more robust way to differentiate wi and wo - my idea is to monkey patch twice

gmm_up = monkey_patch_up(PallasMosaicTpuRaggedDot)
gmm_down = monkey_patch_down(PallasMosaicTpuRaggedDot) - see also comment

[gobbleturk]

can we not re-use the classes defined in moe.py? Is quantization different here?

cc @shuningjin

[gobbleturk]

Looks good, lets have the tokamax team take a look as well

[gobbleturk]

Note we currently don't support passing in input_buffer_count - perhaps this can be a separate add if we find significant lift. I think this defaults to 2 (with or without this change?)