Optimize distillation perf for qwen3-30b (20%→26% MFU) and gpt-oss-20b (17%→19%)

[gagika]

Description

Performance optimization of the qwen3-30b-a3b and gpt-oss-20b distillation configs on
TPU v7x (tpu7x-4x4x4).

MFU on v7x:

• qwen3-30b-a3b: ~20% → ~26% for pdbs=6 and the pdbs=8 + activation-offload
  variant: ~22% → ~24%
• gpt-oss-20b: ~17% → ~19%

All knobs are profile-derived (xplane) and documented inline:

• context=device + custom remat — keep attention outputs on device so the
  backward pass skips the splash-forward re-runs (per-layer forward call count
  halves in the profile). The dominant win; the configs note the HBM frontiers
  (qwen fits to pdbs=6 with the teacher resident, gpt-oss is pdbs=1 at seq 32k).
• Megablox grouped-matmul tiles — qwen at the full dims (emb 2048 /
  moe-mlp 768), ~+10% together with the layout change; gpt-oss raises the
  batch-seq m-tile 512 → 1024 (its k/n dims are already full), ~+3%.
• Splash-attention — sa_q_layout: SEQ_MINOR on qwen; kv-compute sub-blocks
  2048 → 1024 on gpt-oss (~+2%; uniformly smaller blocks regress).
• Batch / mesh — qwen default moves pdbs 4 → 6 in the headroom freed by
  dropping host offload; gpt-oss mesh moves to dp2 × fsdp64.

Tests

30-step runs on tpu7x-4x4x4 (xpk), steady-state step times, each
delta measured against the baselines. Loss and perplexity unchanged from
baseline.

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

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @gagika, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

🤖 I'm sorry @gagika, but I was unable to process your request. Please see the logs for more details.

[vlad-karp]

what is the difference vs previous file config?

[gagika]

per device batch size is 8 here (yml name ends with pdbs8).
Previous one gets best MFU with pdbs=6, so we can default to that unless we need larger batch size for convergence.