Fix KDA precision and MD checkpoint correctness

[andresnowak]

Summary

• keep KDA A_log and dt_bias in FP32 through mixed-precision wrapping and support mixed-dtype torch_dist optimizer state
  • Supposedly the idea is that this two values should be in FP32, FlashKDA library also wants them in fp32 (but one can just upcast the values and thats it, that doesn't matter)
  • *But again maybe htis part is okay, we can remove maybe this fix, Like I'm not sure if these two values do really really need fp32 precision on the operations they work on
• restore KEEL post-layer norms in hybrid decoder specs
• reject KDA configuration controls that are not implemented (has to do with GDN configs)
• reserve residual-output handling for actual out_proj weights
• persist MD preserve-init fixed update norms across ordinary and torch_dist checkpoint resumes
  • Old checkpoints that predate the MD fixed-norm state still load. They recompute the unavailable target from the resumed weight; new checkpoints restore it exactly.

Validation

• 18 focused regressions passed
• mixed BF16/native-FP32 optimizer torch_dist round trip passed
• pre-fix MD torch_dist checkpoint load passed
• rectangular gains resume matched with maximum parameter difference 0.0
• ruff check and git diff --check passed

This implementation and pull request were prepared by Codex.

[mac-mvak]

I've also thought of harmonizing bf16 vs fp32 with FlashKDA, I think that's a good idea