[megablox] Remove redundant expert-weight transpose copy in unquantized gmm_v2 DLHS

[copybara-service]

[megablox] Remove redundant expert-weight transpose copy in unquantized gmm_v2 DLHS

The tokamax gmm_v2 kernel has no native transpose_rhs support, so the V2 DLHS
gradient path (_dlhs_run_tokamax_v2) manually materializes rhs.swapaxes(1, 2)
on every backward step. For the MoE expert weight this is a full-tensor HBM copy
of a large [num_experts, k, n] array (e.g. bf16[16, 7168, 2048]), emitted by
XLA as a standalone copy with op_name="rhs".

The V1 tokamax path already computes the identical DLHS with no transpose: it
calls tokamax.ragged_dot_general with DLHS_RAGGED_DOT_DIM_NUMS, contracting
over the n dimension directly. This change routes the unquantized gmm_v2 DLHS
(quantization_rule is None) through that transpose-free V1 kernel. Quantized
paths are unchanged and keep the V2 kernel, because their scale handling
(_bwd_prepare_inputs / _dlhs_scale_grad_by_rhs_scale) is gmm_v2-specific.

Verified on TPU with a DeepSeek-scale shape (lhs[m=16384, k=7168],
weight[g=16, k=7168, n=2048], bf16): the full expert-weight transpose copy is
removed (1 -> 0), gradients are bit-identical (max relative error 0.0), and the
full gmm fwd+bwd is ~3% faster. The output is bit-identical, so this is a pure
efficiency change with no numerical or convergence impact.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!