fix: cast aq score to float64 to avoid overflow

[Fridah-nv]

What does this PR do?

Type of change: Bug fix

Mamba mixer layers produce gradients with magnitudes up to ~1e22. Squaring these in float32 caused overflow to inf.
Fix: Multiply grad * diff first in float32, then cast to float64 only for the square + sum step to minimize performance overhead.

Usage

# Add a code snippet demonstrating how to use this

Testing

python examples/llm_ptq/hf_ptq.py --pyt_ckpt_path nvidia/Nemotron-H-4B-Base-8K --qformat nvfp4_mse,fp8 --calib_size 64 --export_path ./output/nemotron-h-4b-fp8 --trust_remote_code --dataset cnn_dailymail --auto_quantize_bits 4.75

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

• Is this change backward compatible?: ✅ / ❌ / N/A
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ / ❌ / N/A
• Did you write any new necessary tests?: ✅ / ❌ / N/A
• Did you update Changelog?: ✅ / ❌ / N/A

Additional Information

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-03-28 19:54 UTC

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.16%. Comparing base (b1f9f01) to head (9407536).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1126   +/-   ##
=======================================
  Coverage   70.16%   70.16%           
=======================================
  Files         229      229           
  Lines       26008    26009    +1     
=======================================
+ Hits        18248    18249    +1     
  Misses       7760     7760           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.