perf(GMM-v2): fuse transpose_rhs into GMM-v2 kernel for DLHS path

[Insideyyy]

Description

This PR adds native transposed-RHS support to GMM v2, allowing the DLHS path to consume expert weights in their existing layout without materializing a full-tensor swapaxes.

It also enables transpose_rhs with fused gate/up activations by making weight slicing and concatenation aware of the RHS layout, while preserving the existing scale and bias semantics.

For sub-byte weights, the DLHS path continues to use the materialized non-transposed fallback because the packed representation only supports the standard [group, k, n] layout.

Tests

Added parameterized regression tests in tests/unit/pallas_mosaic_tpu_v2_kernel_test.py covering:

• Native transposed RHS versus the materialized-swapaxes baseline for BF16 and FP8.
• DLHS dispatch behavior for:
  • int4 using the materialized non-transposed fallback.
  • int8 using the native transposed-RHS path.
• Fused activation with a native transposed RHS for BF16 and FP8 using:
  • silu
  • gelu
  • swigluoai
• Fused activation with RHS scale and bias.

The TPU validation was run on TPU v7x. All tested native-transpose and fallback paths matched their corresponding materialized-layout references exactly.

The targeted tests can be reproduced with:

python -m pytest tests/unit/pallas_mosaic_tpu_v2_kernel_test.py \
  -k "gmm_native_rhs_transpose_matches_swapaxes or \
      dlhs_rhs_transpose_dispatch_by_dtype or \
      gmm_fused_activation_with_native_rhs_transpose"

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. Targeted GMM v2 regression tests were run on TPU instead.
• 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. No documentation changes are required for this internal kernel optimization.

[codecov]

Codecov Report

❌ Patch coverage is 57.44681% with 20 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...ernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py	59.52%	16 Missing and 1 partial ⚠️
src/maxtext/kernels/megablox/ops.py	40.00%	2 Missing and 1 partial ⚠️

📢 Thoughts on this report? Let us know!