fix: keep MoE routing and rollout versions consistent

[0z5a]

Summary

This fixes two training consistency failures:

• MoE routing previously cast the full softmax distribution to BF16 before topk. Close probabilities could collapse into ties and select a different expert set from the FP32 router decision.
• Online training mutated shared weights before publishing the new policy version. Generation could observe a partial update under the old version, and a rollout could become too stale while asynchronous reward scoring was in flight.

Changes

• Keep MoE router probabilities in FP32 through expert selection; cast only selected dispatch weights to the activation dtype.
• Validate routed/shared/activated expert counts and reject unsupported routing methods instead of silently applying greedy routing.
• Add an atomic apply_weight_update(version, callback) scheduler boundary shared by generation and optimizer updates.
• Capture the generation version explicitly, enforce monotonic versions and a configurable maximum policy lag, and revalidate after asynchronous reward scoring and before cache reuse.
• Make the final version check and rollout-cache publication one policy-snapshot operation, closing the update-between-check-and-insert TOCTOU window.
• Add rollout_max_policy_lag; the compatibility default is rollout_interval - 1.
• Add deterministic generation/update, snapshot/update, cache-publication/update, stale/future-version, optimizer-failure, and MoE near-tie regression tests.
• Fix the checked-in router benchmark so both tuple-returning routes are reproducible, then refresh its L20 evidence and digest bindings.

NVIDIA L20 benchmark

Environment: NVIDIA L20 (SM89), PyTorch 2.11.0+cu128, CUDA 12.8, BF16, seed 3407, 20 warmups, 100 timed iterations on GPU5.

Tokens	Experts	K	Old median	New median	Delta	Old mismatch, random	Old mismatch, near-tie	New mismatch
2,048	8	2	0.0481 ms	0.0481 ms	0.00%	0.0488%	95.9473%	0%
8,192	8	2	0.0512 ms	0.0492 ms	-4.00%	0.1465%	96.2280%	0%
8,192	64	8	0.0645 ms	0.0584 ms	-9.52%	0.1099%	100%	0%
32,768	64	8	0.1382 ms	0.1372 ms	-0.74%	0.1190%	100%	0%

Mismatch compares the selected expert set against an FP32-softmax/FP32-top-k reference. The maximum additional probability storage in this portfolio is 4 MiB. Reproduction code and raw output are checked in under benchmarks/training_consistency/ and benchmarks/results/.

Validation

• Existing full local suite before the final cache-race supplement: 640 passed, 103 skipped
• Final NVIDIA L20 GPU5 focused suite: 116 passed in 3.40s
• Ruff format check: all 166 Python files formatted
• Ruff import lint: passed
• InfraSWE comparison cell validation: passed
• InfraSWE Draft/system-path model tests after digest refresh: 41 passed

InfraSWE commit 811bc775ed5b3a6ec853219245f3469f78818020 validates the AstrAI system-path comparison cell. The refreshed candidate, acceptance, probe, and workload digests all match their checked-in source sets. The visible-evidence diagnostic ProjectFit remains 92.41/100 and BenchmarkTrust 97.40/100. Official scoring intentionally remains unresolved because the draft is unsealed and lacks five fresh-process replays, system-trace evidence, hidden probes, and a verified evidence manifest.

Scope limits

The benchmark is a single-process router microbenchmark and the async tests use deterministic thread interleavings. This PR does not claim a long-running multiprocess training soak or official InfraSWE certification.