Merge pull request #4542 from AI-Hypercomputer:mdformat_fix

PiperOrigin-RevId: 951070936

Author: Google-ML-Automation

.pre-commit-config.yaml

@@ -57,7 +57,17 @@ repos:
         args: ['--number']
         additional_dependencies: [mdformat-myst, mdformat-ruff]
         files: (docs/.)
-        exclude: docs/guides/checkpointing_solutions.md|docs/guides.md|docs/run_maxtext.md
+        # Exclude all files that use complex MyST directives that mdformat breaks
+        exclude: >
+          (?x)^(
+            docs/guides/checkpointing_solutions\.md|
+            docs/guides\.md|
+            docs/run_maxtext\.md|
+            docs/guides/monitoring_and_debugging\.md|
+            docs/reference/core_concepts.md|
+            docs/reference/architecture.md|
+            docs/reference/models.md
+          )$
 
   - repo: https://github.com/adrienverge/yamllint
     rev: v1.35.0

docs/reference/core_concepts/moe_configuration.md

@@ -99,7 +99,7 @@ Dropping:
 
 `mlp_bias`: If enabled, add learnable bias terms for MLP matmul. Originally implemented to support the GPT-OSS model architecture.
 
-`prefuse_moe_weights`: If enabled alongside `sparse_matmul=True`, fuses the two FFN1 grouped GEMMs (wi\_0 and wi\_1) into a single grouped GEMM call. Expert weights are stored in a concatenated `(num_experts, embed_dim, 2 * mlp_dim)` shape, so input activations are loaded from HBM once per forward pass instead of twice. Backend-agnostic (works with Megablox, JAX Ragged Dot, and Tokamax). When used with `attention=vllm_rpa`, the fused weight tensor is passed directly to the vLLM-TPU serving kernel without splitting.
+`prefuse_moe_weights`: If enabled alongside `sparse_matmul=True`, fuses the two FFN1 grouped GEMMs (wi_0 and wi_1) into a single grouped GEMM call. Expert weights are stored in a concatenated `(num_experts, embed_dim, 2 * mlp_dim)` shape, so input activations are loaded from HBM once per forward pass instead of twice. Backend-agnostic (works with Megablox, JAX Ragged Dot, and Tokamax). When used with `attention=vllm_rpa`, the fused weight tensor is passed directly to the vLLM-TPU serving kernel without splitting.
 
 `use_batch_split_schedule` (experimental): If enabled, split batch into micro-batches to hide communications that yields performance benefits.