## What does this PR do?
**Type of change:** Fix MCore KV Cache Quantization: Amax Device
Placement Bug; Code clean up; Distributed Sync of KVCache Quantizer
params; unittest expansion to hybrid models
**Overview:** Fixes bugs preventing MCore KV Cache quantization from
working during checkpoint restore.
### Bug Chain
**Bug 1:** `is_enabled = self.weight_quantizer.is_enabled if
hasattr(self, "weight_quantizer") else False`
No `weight_quantizer` for KV-cache-only quant → `is_enabled=False` →
metadata not saved → `modelopt_post_restore()` never called. *(Thanks to
@jenchen13 )*
**Bug 2:** After fixing Bug 1, `_amax` restored on CPU (via
`_reset_pytorch_state_from_metadata`). Fallback
`_calibrate_quantizers()` never called because `_amax` exists.
**Bug 3:** Even if called, `_calibrate_quantizers()` fails —
`core_attention` has no parameters → can't determine device/dtype.
### The Fix
1. Remove `is_enabled` check entirely — disabled modules may still need
metadata restore. Explicitly skip `output_layer` from extra state
callbacks (never quantized)
2. Set `dtype`/`device` on `core_attention` from parent Attention
module, `modelopt_post_restore()` calls `self.to(device, dtype)`
3. Remove dead `_calibrate_quantizers()` code (will bring back similar
logic for KV cache affine quantization)
### Previous Unit Test Was Wrong
`model_test` was `mtq.quantize()`'d, not `mto.restore()`'d. Never tested
actual restore path.
### Additional Fixes
- Amax sync across DP/TP for KV cache quantizers
- `flash_decode` auto-disabled
### Code Cleanup
Removed ~100 lines of dead code.
## Testing
1. MCore KV Cache QAD with Nano V3 + Context Parallel works
2. Unit tests: hybrid models, KV+GEMM configs, correct restore workflow,
backward pass validation
## Before your PR is "*Ready for review*"
- **Is this change backward compatible?**: Yes
- **Did you write any new necessary tests?**: Yes
- **Did you add or update any necessary documentation?**: No
- **Did you update Changelog?**: Yes
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
Co-authored-by: Asma Thekkumpate <akuriparambi@cw-dfw-cs-001-vscode-02.cm.cluster>
