## What does this PR do?
**Type of change:** ? <!-- Use one of the following: Bug fix, new
feature, new example, new tests, documentation. -->
**Overview:** ?
## Usage
<!-- You can potentially add a usage example below. -->
```python
# Add a code snippet demonstrating how to use this
```
## Testing
<!-- Mention how have you tested your change if applicable. -->
## Before your PR is "*Ready for review*"
<!-- If you haven't finished some of the above items you can still open
`Draft` PR. -->
- **Make sure you read and follow [Contributor
guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)**
and your commits are signed.
- **Is this change backward compatible?**: Yes/No <!--- If No, explain
why. -->
- **Did you write any new necessary tests?**: Yes/No
- **Did you add or update any necessary documentation?**: Yes/No
- **Did you update
[Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**:
Yes/No <!--- Only for new features, API changes, critical bug fixes or
bw breaking changes. -->
## Additional Information
<!-- E.g. related issue. -->
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
## Release Notes
* **New Features**
* Added NVFP4StaticQuantizer for improved 4-bit quantization with
enhanced precision control
* Introduced NVFP4MSECalibrator with flexible candidate generation for
calibration optimization
* **Improvements**
* Optimized GPU kernels for Hopper+ graphics cards with better
performance
* Extended Triton support to broader GPU compatibility
* Enhanced backward compatibility for restoring previously quantized
models
* **Tests**
* Added comprehensive test coverage for new quantizers and calibration
methods
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
---------
Signed-off-by: realAsma <akuriparambi@nvidia.com>
## 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>
realAsma