[DSV4 TOP 1/N] Make trainer query RMS batch invariant

[kaixih]

Summary

Makes the trainer-side DeepSeek-V4 512-wide query RMS batch invariant.

• Adds an explicit pairwise fixed-tree mean over the last dimension.
• Adds a shared unweighted DSV4 query-RMS helper.
• Uses the fixed tree after wq_b when Megatron batch_invariant_mode=True.
• Preserves the existing native .mean formula when the flag is false.
• Preserves output dtype and autograd.

This is the first split from #1788.

Motivation

DeepSeek-V4 true on-policy (TOP) requires exact prefill-recompute versus
training-forward log-probabilities. The query RMS reduction must keep the same
floating-point association when the same query row is evaluated under different
outer batch shapes.

Megatron already exposes --batch-invariant-mode through TransformerConfig.
This change makes the DSV4 trainer query RMS honor that existing opt-in mode.

Behavior

Default behavior is unchanged:

batch_invariant_mode=False -> native FP32 square + mean + rsqrt
batch_invariant_mode=True  -> FP32 square + fixed-tree mean + rsqrt

The fixed-tree path is not enabled by default because its unfused implementation
has a performance cost.

Scope

This PR contains the fixed-tree primitive and the trainer-side 512-wide query-RMS
wiring as one feature.

It intentionally does not include the matching SGLang rollout/prefill
query-RMS path. That backend change is required in a follow-up before this forms
a complete DSV4 TOP numerical contract or proves end-to-end prefill-versus-trainer
parity.

It also does not include deterministic TP collectives, FP8 alignment,
hyper-connections, MoE accumulation, log-softmax, or the final E2E.

Tests

• One focused FP32 behavior test uses the known GB200 counterexample shapes
  (3, 1, 4, 512) and (3, 2, 4, 512), requiring exact fixed-tree output
  equality. It intentionally does not assert native-path inequality because
  that depends on the device and selected kernel.
• test_dsv4_query_rms_is_batch_invariant: 1 passed on the current head.
• Pre-commit, including Black: passed.
• GB200 CUDA smoke on the current head (fed40325) with
  radixark/miles:latest and PyTorch 2.11.0+cu130: passed.
  • Focused test: 1 passed.
  • FP32 and BF16 query-RMS outputs were bitwise equal for the same rows
    evaluated alone and inside a larger batch.
  • Query-RMS microbenchmark: native 59.71 us, fixed tree 225.73 us
    (3.78x).
• git diff --check: clean

The unfused fixed-tree path has a measurable performance cost, so the production
wiring in this PR remains opt-in behind --batch-invariant-mode.

[gemini-code-assist]

Caution

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

[gemini-code-assist]

Caution

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

[kaixih]

@maocheng23 @yueming-yuan to review thx.