Add MaxText mHC-lite Pallas/Mosaic TPU kernel and profiling benchmarks.

[copybara-service]

Add MaxText mHC-lite Pallas/Mosaic TPU kernel and profiling benchmarks.

This change implements the Manifold Constrained Hyper-Connections (mHC-lite) algorithm in MaxText using Pallas. Instead of a single monolithic kernel, the implementation is split into a pipeline of three forward kernels and three custom VJP backward kernels. This design allows the wrapped model branch (Attn + MoE) to run in between.

Algorithmic Step-by-Step

1. Forward Pass

• Kernel 1: Coefficient Computation (_coeff_fwd)
  • Inputs: Input activations x, folded projection weights phi (fuses RMSNorm scale and projection matrices), gate parameters, and permutation matrices.
  • Operations:
    1. Projects x using phi.
    2. Applies RMSNorm to the projected logits (fused projection and normalization avoids materializing normalized x in HBM).
    3. Computes sigmoid gates: h_pre (input gate) and h_post (output gate).
    4. Computes softmax weights for permutations and blends them to produce a mixed token-specific residual permutation matrix.
  • Outputs: h_pre (passed to Kernel 2), h_post and residual (saved in context for Kernel 3).
• Kernel 2: Pre-Branch Gating (_pre_apply_fwd)
  • Inputs: x, h_pre.
  • Operations: Applies h_pre to x and sums across the streams dimension, collapsing it.
  • Outputs: layer_input of shape (tokens, embedding) which is passed to the wrapped model branch.
• [Wrapped Model Branch] Runs standard JAX/XLA (e.g. Attention) on layer_input to produce layer_output.
• Kernel 3: Post-Branch Mixing (_post_apply_fwd)
  • Inputs: x, layer_output, h_post, residual.
  • Operations:
    1. Permutes and mixes the original input streams x using the blended residual matrix.
    2. Scales the wrapped branch layer_output by h_post.
    3. Combines the mixed residual and scaled branch output.
  • Outputs: Final mixed output streams of shape (tokens, streams, embedding).

2. Backward Pass (Custom VJP)

• Kernel 4: Post-Branch Backward (_post_apply_bwd)
  • Inputs: x, layer_output, h_post, residual, and the incoming gradient d_output.
  • Operations: Computes gradients for the inputs of the post-apply step.
  • Outputs: d_layer_output (sent to wrapped block backward) and cotangents d_x_post, d_h_post, d_residual (passed to pre-branch backward).
• [Wrapped Model Branch Backward] Runs standard JAX/XLA backward pass to compute d_layer_input.
• Kernel 5: Pre-Branch Gating Backward (_pre_apply_bwd)
  • Inputs: x, h_pre, d_layer_input, and d_x_post (accumulated gradient from post-branch).
  • Operations:
    1. Computes VJP of pre_apply to get d_x_pre and d_h_pre.
    2. Accumulates the gradients: d_x_acc = d_x_pre + d_x_post.
  • Outputs: d_h_pre and accumulated d_x_acc.
• Kernel 6: Coefficient Backward (_coeff_bwd)
  • Inputs: x, phi, parameters, permutations, cotangents (d_h_pre, d_h_post, d_residual), and d_x_acc.
  • Operations:
    1. Recomputes mhc_coeffs forward pass internally to regenerate intermediate activations (gates, logits, softmax weights), saving HBM read/write traffic.
    2. Computes VJP of mhc_coeffs to get gradients for x (d_x_coeff), phi, and parameters.
    3. Accumulates final input gradient: d_x = d_x_coeff + d_x_acc.
  • Outputs: Final d_x, d_phi, and parameter gradients.
• [Fold Norm Scale Backward] Maps d_phi to individual projection gradients using JAX.

Rematerialization (Remat) Behavior

The MHC-lite kernel implementation is fully compatible with JAX rematerialization (checkpointing).

• Remat Disabled:
  • Intermediate activations are retained in HBM for the backward pass.
  • The custom VJPs save the forward inputs and intermediate outputs. Specifically, pre saves x (shape [B, S, K, D]), phi, and h_pre; post saves x, layer_output, h_post, and residual.
  • This results in higher HBM usage.
• Remat Enabled (Block-level checkpointing):
  • Intermediate activations (including the MHC context containing h_post and residual, and the layer_output) are discarded during the forward pass.
  • During the backward pass, the forward pass of the block is recomputed: pre kernel runs again to recreate the context, the wrapped branch runs again to recreate layer_output, and post kernel runs again.
  • The custom VJPs then execute using these recomputed activations.
  • This significantly reduces HBM pressure at the cost of recomputation overhead.
  • Note: The MHC context is not explicitly named for selective checkpointing policies (like minimal_with_context), meaning it will be recomputed rather than saved when using such policies.

Numerical Verification against MaxText Baseline Layer

• Verified algorithmic and numerical equivalence between the 3-stage Pallas/Mosaic TPU kernel and MaxText's native baseline layer (ManifoldConstrainedHyperConnections with use_mhc_pallas_kernel=False).
• Compared all 13 forward output and backward gradient tensors (forward output out, input gradient grad_x, weight gradient grad_W, and 10 primal parameter gradients: grad_norm_scale, grad_pre_alpha, grad_pre_bias, grad_pre_scale, grad_post_alpha, grad_post_bias, grad_post_scale, grad_res_alpha, grad_res_bias, grad_res_scale).
• All 13 outputs and parameter gradients achieved > 99.9% cosine similarity (cos_sim >= 0.999) and bounded elementwise difference in bfloat16 across multiple test shapes ((b=2, s=4096, k=4, d=3072)).

Hardware Benchmarks & Optimal Block Configurations

1. TPU v6 (Trillium / Ghostlite)

• XProf Profile URL: https://xprof.corp.google.com/?session_id=dandragona-4394958857230269497
• Optimal Block Sizes: fwd_bs = 256, bwd_bs = 256
• Benchmark Results for Shape (batch=2, sequence=4096, streams=4, embedding=3072):
  • Native MaxText Baseline (JAX + remat): 8.866 ms per step
  • Pallas Kernel (fwd=256, bwd=256): 4.511 ms per step (1.97x speedup)

2. TPU v7 (Ironwood / Ghostfish)

• XProf & Percale / LLO Trace URL: https://xprof.corp.google.com/?session_id=dandragona-18083177279302912858
• Optimal Block Sizes: fwd_bs = 256, bwd_bs = 128
• Benchmark Results for Shape (batch=2, sequence=4096, streams=4, embedding=3072):
  • Native MaxText Baseline (JAX + remat): 6.481 ms per step | 968.58 MB HBM usage | 52.00 MB VMEM usage.
  • Pallas Kernel (fwd=256, bwd=128): 4.514 ms per step (1.44x speedup) | 381.01 MB HBM usage (2.54x HBM traffic reduction) | 63.94 MB VMEM usage.

[codecov]

Codecov Report

❌ Patch coverage is 92.64214% with 22 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/kernels/mhc/common.py	80.28%	7 Missing and 7 partials ⚠️
src/maxtext/kernels/mhc/api.py	72.41%	4 Missing and 4 partials ⚠️

📢 Thoughts on this report? Let us know!