Add optimized Pallas kernel for mHC

[dandragona]

Description

Adds mHC Pallas Kernel to MaxText: https://xprof.corp.google.com/trace_viewer/dandragona-2944179191919190860?view_start=59318.405&view_end=59445.703

Key changes:

• Adds the kernel implementation split into:
  • src/maxtext/kernels/residual/mhc_common.py (shared operations)
  • src/maxtext/kernels/residual/mhc_fwd_kernels.py (forward kernels & wrapper)
  • src/maxtext/kernels/residual/mhc_bwd_kernels.py (backward kernels & wrapper)
• Integrates the kernel into the MHC layer (src/maxtext/layers/mhc.py) controlled by the new enable_mhc_pallas_kernel config flag.
• Refactors files to comply with code quality requirements (pyink, pylint).

REQUIRES: TP = 1

mHC Pallas Kernels Architecture & Flow

The Manifold-Constrained Hyper Connections (mHC) algorithm introduces a structured residual pathway that projects activations before and after applying a layer (like Attention or MLP) to stabilize training in deep networks.

Because mHC computes dynamic projection coefficients on the fly based on the activations, a native JAX implementation would require materializing large intermediate tensors (such as sequence-length dimension matrices), causing high memory usage (OOM) and latency.

The custom TPU Pallas kernels optimize this by fusing operations along the sequence dimension and computing gradients manually via custom VJP.

1. Forward Pass Flow

The forward pass is split into three main steps:

1. Coefficient Computation (_coeff_fwd):
   • Takes the input activation X (shape [T, K, D]) where T is batch * seq, K is the number of heads/branches, and D is the feature dimension.
   • Projects X to compute pre-coefficients H_pre, post-coefficients H_post, and the residual projection matrix res_M.
   • Formulated as mhc_coeffs.
2. Pre-mapping Apply (_pre_apply_fwd):
   • Applies the pre-coefficients H_pre to the activations X to get the input to the layer: layer_in = sum(H_pre * X).
3. Post-mapping Apply (_post_apply_fwd):
   • Combines the original input X, the layer output layer_out, H_post, and the residual matrix res_M.
   • Computes the final block output: Y = res_M * X + H_post * layer_out.

                            Input X
                           /   │   \
                          /    │    \
                         /     │     \
                        v      │      v
                 [ _coeff_fwd ]│   [ _pre_apply_fwd ] <── H_pre
                  /    │   \   │      │
                 /     │    \  │      v
                v      v     v │   layer_in
            H_pre   H_post  res_M     │
              │        │     │ │      v
              │        │     │ │   [ Layer ]
              │        │     │ │      │
              │        │     │ │      v
              ▼        │     │ │   layer_out
          (to pre)     │     │ │      │
                       ▼     v v      ▼
                     [   _post_apply_fwd   ]
                            │
                            v
                         Output Y

2. Backward Pass Flow (VJP)

The backward pass is implemented via custom VJP (_mhc_pallas_vjp_bwd_impl) to avoid as autograd did not work natively:

1. Post-Apply Backward:
   • _post_apply_bwd computes the gradients for activations: dx (using res_M and d_out) and dlo (using H_post and d_out).
2. Pre-Apply Backward (_pre_apply_bwd):
   • Computes gradients for X and H_pre from the layer input gradient d_layer_in.
3. Coefficient Backward (_coeff_bwd):
   • Propagates gradients from dH_pre, dH_post, and dres_M back to the input X and all mHC projection parameters (phi) using jax.vjp on mhc_coeffs.

Tests

• Verified correctness using CPU-side tests in tests/unit/mhc_kernel_test.py.
• Verified that all code quality checks (pre-commit) pass locally.

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

❌ Patch coverage is 87.61905% with 26 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/mhc.py	40.00%	12 Missing and 3 partials ⚠️
src/maxtext/kernels/residual/mhc_common.py	88.23%	3 Missing and 3 partials ⚠️
src/maxtext/kernels/residual/mhc_fwd_kernels.py	92.18%	5 Missing ⚠️

📢 Thoughts on this report? Let us know!