fix: preserve final RMS predictions during weight processing

[emerardd]

Description

Preserve predictions during compatibility-mode weight processing for supported mixed normalization models: LayerNorm blocks with a final RMSNorm (normalization_type="LN", final_rms=True).

Two centering operations incorrectly treated this topology as uniformly LayerNorm:

• Final norm folding centered each unembedding direction over d_model. This is equivalent after zero-mean LayerNorm output, but not after RMSNorm.
• Value-bias folding re-centered the attention output bias when center_writing_weights=True, even though the main writing-weight-centering path already excludes models with final_rms=True.

Both paths now respect final_rms. Final RMS scale is still folded into unembedding weights, and value biases are still folded into attention output biases; only the unsafe centering is skipped. Block LayerNorm folding, vocabulary-axis unembedding centering, and ordinary LayerNorm/RMSNorm behavior remain unchanged. No dependencies or public APIs change.

The cached-statistics analysis paths also respect the selected final norm: ordinary residual-stack normalization and the fused projected-neuron path share a target-aware centering check. Intermediate block LayerNorm centering is preserved. This keeps DLA and full residual decomposition consistent with the correctly preserved model logits. Writing-weight centering and post-value-fold output-bias cleanup likewise share one eligibility condition.

Reproduction and regression coverage

All reproductions use tiny, real Native models with random CPU weights and no checkpoint downloads. Extending the existing prediction-parity matrix to include mixed LN/final-RMS with nonidentity block norm parameters failed before the fix: maximum log-softmax difference was approximately 0.5501. Fixing only unembedding centering left approximately 0.0186 error, exposing the post-value-fold bias centering path.

Regression coverage now includes:

• Mixed LN/final RMS in the existing default compatibility-mode prediction-parity matrix, alongside the existing pure LN/RMS and other topology controls.
• Fold-only raw-logit parity with identity and nonidentity final RMS scale, asserting the unembedding weights absorb exactly that scale and the final norm scale becomes identity.
• Nonzero value-bias folding with and without norm folding, while writing-weight centering is requested, asserting prediction parity under final RMS.
• LN/LNPre blocks with final RMS under no-processing, fold-only and default compatibility modes; batched/batchless caches; DLA and expanded/unexpanded full residual decomposition reconciling with logits; vector/matrix projected-neuron parity against an uncentered oracle; and intermediate attention/MLP-input LN centering controls.

Validation

• Expanded Native, weight-processing, DLA-final-norm and activation-cache regression surface: 336 passed, 1 skipped, 53 deselected. The inherited cross-device skip requires CUDA or MPS; no skip was added. The 53 deselections are head-result tests outside this change's scope.
• The new ordinary cached-normalization regression failed six cases before the fix. Restoring old global centering in memory makes all four dedicated projected-neuron cases fail with assertion errors and zero setup/teardown errors.
• Mypy: no issues in 393 source files.
• Makefile-equivalent pycln/isort/Black checks on the review follow-up files and git diff --check passed.
• The complete unit-test and PR-review suites were not run. No GPU validation is claimed. The local run emitted existing SWIG dependency deprecation warnings.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

No documentation change is needed because this restores the existing prediction-preservation contract. The full-unit-suite checkbox remains unchecked because only the affected test surface was run.

[jlarson4]

@emerardd Thanks for fixing this bug. A couple comments related to some side-effects of the resolution.

[jlarson4]

Looks good to me! I will merge if/when CI is passing