fix(diffusion): keep replay transition means at SDE-native fp32

[Yuanjia-Lee]

Summary

FlowDPPO's KL mask and its log-prob score different distributions for the same
transition. The SDE kernel computes the transition mean μ_θ in fp32 and scores
log-prob against that mean, but stage.replay() narrowed the returned
prev_sample_means to trajectory_dtype (typically bf16), and the mask
(μ_θ − μ_old)²/(2σ_t²) compares those narrowed means against
kl_mask_threshold: 1e-5 (sd3_flowdppo.yaml). One bf16 rounding of a
|μ|≈1.6 value has std 2.3e-3, the same size as the μ-delta the threshold is
sized for (RMS ≈ 2e-3 at the recipe's σ window): small legitimate updates get
masked out, sub-ULP deltas read as KL = 0, and the β>0 reference penalty on the
same means reads 0.

The fix removes that cast from every diffusion replay implementation: SD3,
Qwen-Image/Edit-Plus, WAN 2.1/2.2, HunyuanVideo 1.0/1.5, HunyuanImage 3,
Z-Image, FLUX.2-klein, LTX2, Boogu-Image, Bagel, and SenseNova U1, plus the
shared BatchedStepReplayMixin. It also preserves the rollout-owned old-mean
anchors in Bagel and SenseNova U1. MiniMax-H3 already returned un-narrowed
means. log_probs keeps its logprob_dtype cast.

Measured through the real replay→loss→optimizer code at the SD3 recipe's
hyperparameters: replay's log-probs recomputed from the returned means match
exactly after the fix (max|Δ| = 0.0) and deviate by ~2e-5 with the cast. A
small weight perturbation flipped the mask on 5/24 sample-steps (native KL
~1e-6, narrowed up to 1.25e-5) and reversed 61/256 gradient components (cosine
0.76); with no flips the two arms' gradients are identical (cosine 1.0), so the
reversal is the mask, not noise.

Related Issue

N/A

Test Plan

• SKIP=no-commit-to-branch pre-commit run --files <all 17 changed files> --show-diff-on-failure
  → all hooks pass locally; the pre-push hook also passes on the follow-up commit.
• Focused kernel check with bf16 sample/noise_pred verifies both the returned
  transition and prev_sample_mean remain torch.float32.
• Source sweep verifies no diffusion replay or rollout-owned old-mean anchor
  narrows prev_sample_means/sde_means to trajectory_dtype.
• Original SD3 harness (uncommitted, no tests/ tree by policy test: remove tests directory #99/test: remove tests directory #267): full
  SD3DiffusionStage.diffuse/replay (serial and batched, bit-identical) and
  the FlowDPPO loss + AdamW (lr 3e-4, wd 0) at eta 0.7 / threshold 1e-5 /
  10-step shift-3 schedule / B=8, with a small deterministic transformer
  standing in for the SD3.5 checkpoint; CPU and CUDA bf16-autocast on one 4090.
  Numbers are in the Summary.

Compatibility / Risk

• Means now stay at the SDE kernel's native fp32 across all replay stages.
  Consumers already align old/reference means to new_means.dtype or compute
  the KL in fp32; no dtype assertions exist. Numerics of mean-based KL paths
  change by design; log_probs are untouched.
• A retained mean anchor takes twice the memory when trajectory_precision is
  bf16. For the SD3 recipe's [8, 3, 16, 64, 64] micro-batch, the frozen CPU
  anchor is 6 MiB instead of 3 MiB. Bagel and SenseNova U1 also retain their
  rollout-owned old-mean anchors in fp32 rather than bf16.

Reviewer Notes

• The implementation now matches the repo-wide contract documented on
  ReplayResult.prev_sample_means and in unirl/sde/README.md; serial and
  batched replay no longer differ by dtype.
• Duplicate check: no open PR or issue covers mean precision. feat(diffusion): batch video replay timesteps #394 only adds
  per-step batching and carried the same cast in its new code; feat(opd): harden teachers for video KL, meta-init, and resume #432, perf(bagel): pack FlowGRPO replay steps #384, and
  refactor(diffusion): make role residency one choice per role, not per phase #428 do not touch this boundary.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.
  (Contract recorded in unirl/sde/README.md Gotchas and on
  ReplayResult.prev_sample_means.)

[CjhHa1]

LGTM