[kaix-nv]

What does this PR do?

Type of change: New feature.

Add rounding="nearest" | "stochastic" to QuantizerAttributeConfig and TensorQuantizer. Rounding remains independent of format and scale grouping, and nearest remains the default.

The stochastic policy is unbiased between adjacent representable values. It supports built-in integer and FP8 E4M3 fake quantization with unquantized scales, including blockwise grouping, plus dynamic block16 NVFP4. For NVFP4, E2M1 values are rounded stochastically while E4M3 block scales retain nearest rounding. Exact range endpoints remain exact.

Integer, FP8, and NVFP4 quantization share a rounding primitive after constructing the format grid. Stochastic rounding uses the PyTorch device RNG and preserves the configured straight-through gradient. The policy is saved and restored through existing ModelOpt quantizer metadata; training code remains responsible for saving RNG state.

Calibration and weight folding force nearest rounding so their results remain deterministic. Unsupported real/custom quantization, quantized-scale formats other than dynamic block16 NVFP4, learned scales, fused GDN/attention/FP8 GEMM paths, and inference export reject stochastic rounding rather than silently falling back to nearest. ReplaySSM and serving integration remain separate follow-ups.

Usage

from modelopt.torch.quantization.config import QuantizerAttributeConfig
from modelopt.torch.quantization.nn import TensorQuantizer

quantizer = TensorQuantizer(
    QuantizerAttributeConfig(
        num_bits=(4, 3),  # Use 8 for integer quantization.
        type="dynamic",
        block_sizes={-1: 128},
        rounding="stochastic",
    )
)

nvfp4 = TensorQuantizer(
    QuantizerAttributeConfig(
        num_bits=(2, 1),
        block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
        rounding="stochastic",
    )
)

Validation

Current head b35886db2b:

• 96 passed in the focused CPU TensorQuantizer suites.
• All checks in pre-commit run --from-ref origin/kaix/linear-attention-state-nvfp4 --to-ref HEAD passed, including Ruff, mypy, RST validation, and Bandit.
• git diff --check passed.
• No model-accuracy, serving-speed, or native fused-kernel claim is made by this PR.

Before your PR is "Ready for review"

• Backward compatible: Yes; nearest is unchanged for existing supported inputs and legacy checkpoint metadata.
• Copied code or new dependency: None.
• Necessary tests: Added or extended as listed above.
• Changelog: Updated under Quantization new features.
• Keep this PR as a draft until owner review is complete.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2829/
Built to branch gh-pages at 2026-10-19 22:26 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 91.11111% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 68.99%. Comparing base (fa97dd3) to head (e62473c).

Files with missing lines	Patch %	Lines
modelopt/torch/quantization/tensor_quant.py	93.65%	4 Missing ⚠️
...torch/quantization/backends/fp8_per_tensor_gemm.py	0.00%	2 Missing ⚠️
.../torch/quantization/nn/modules/tensor_quantizer.py	91.66%	1 Missing ⚠️
.../torch/sparsity/attention_sparsity/plugins/vllm.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@                          Coverage Diff                          @@
##           kaix/linear-attention-state-nvfp4    #2829      +/-   ##
=====================================================================
+ Coverage                              68.13%   68.99%   +0.86%     
=====================================================================
  Files                                    665      656       -9     
  Lines                                  75689    72186    -3503     
=====================================================================
- Hits                                   51567    49808    -1759     
+ Misses                                 24122    22378    -1744     

Flag	Coverage Δ
unit	59.65% <91.11%> (-0.05%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 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.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.