Skip to content

Minor fix constant amax repr#1935

Open
mxinO wants to merge 2 commits into
mainfrom
mxin/fix-constant-amax-repr
Open

Minor fix constant amax repr#1935
mxinO wants to merge 2 commits into
mainfrom
mxin/fix-constant-amax-repr

Conversation

@mxinO

@mxinO mxinO commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: ?

The constant amax in kv cache is showing dynamic which is confusing, making it show the "value(const)".

Usage

# Add a code snippet demonstrating how to use this

Testing

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
  • Did you get Claude approval on this PR?: ✅ / ❌ / N/A

Additional Information

Summary by CodeRabbit

  • Bug Fixes

    • Updated quantized tensor summaries to show a constant maximum value when that mode is enabled.
    • Improved printed output so the quantizer representation now reflects the constant-amax setting more clearly.
  • Tests

    • Added coverage to verify the quantizer’s displayed representation includes the expected constant-amax indicator.

mxinO added 2 commits July 7, 2026 00:32
Constant-amax quantizers intentionally omit the _amax buffer, but the representation treated every missing buffer as runtime-dynamic. Format the same E4M3 maximum used by the quantization path so diagnostics describe the active behavior.

Constraint: Preserve true dynamic, static scalar, per-channel, and MX-format representations
Rejected: Register a synthetic _amax buffer | would change quantizer state and calibration semantics for a display-only fix
Confidence: high
Scope-risk: narrow
Reversibility: clean
Directive: Keep constant-amax reporting derived from the value used by _get_amax
Tested: tests/unit/torch/quantization/test_print.py (6 passed); focused Ruff, format, mypy, license, and Bandit hooks
Not-tested: Full ModelOpt unit suite
Signed-off-by: Meng Xin <mxin@nvidia.com>
A numeric value alone distinguishes constant mode from dynamic mode, but not from a stored calibrated scalar. Append a const marker so logs communicate both the effective value and its source without using per-channel bracket notation.

Constraint: Preserve the existing per-channel [min, max](count) representation
Rejected: Wrap the scalar in brackets | brackets already denote tensor ranges in quantizer diagnostics
Confidence: high
Scope-risk: narrow
Reversibility: clean
Tested: tests/unit/torch/quantization/test_print.py (6 passed); focused Ruff, format, mypy, license, and Bandit hooks
Not-tested: Full ModelOpt unit suite
Signed-off-by: Meng Xin <mxin@nvidia.com>
@copy-pr-bot

copy-pr-bot Bot commented Jul 7, 2026

Copy link
Copy Markdown

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

coderabbitai Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Adds a new conditional branch in TensorQuantizer._short_amax that returns a formatted constant string based on torch.finfo(torch.float8_e4m3fn).max when _use_constant_amax is enabled, along with a unit test verifying this behavior in the quantizer's repr() output.

Changes

Constant amax repr

Layer / File(s) Summary
Constant amax branch and test
modelopt/torch/quantization/nn/modules/tensor_quantizer.py, tests/unit/torch/quantization/test_print.py
_short_amax returns a formatted torch.finfo(torch.float8_e4m3fn).max string when _use_constant_amax is enabled; a new test verifies the quantizer's repr() includes amax=4.48e+02(const).

Estimated code review effort: 1 (Trivial) | ~3 minutes

Suggested reviewers: Quantization module maintainers familiar with TensorQuantizer repr formatting.

🚥 Pre-merge checks | ✅ 6
✅ Passed checks (6 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title is concise and accurately reflects the main change: a small fix to the constant amax representation.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns ✅ Passed No forbidden patterns found in the PR diff: no eval/exec, no nosec, no unsafe load flags, no trust_remote_code, and no dependency changes.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch mxin/fix-constant-amax-repr

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

@github-actions

github-actions Bot commented Jul 7, 2026

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1935/

Built to branch gh-pages at 2026-07-07 08:38 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@mxinO mxinO marked this pull request as ready for review July 7, 2026 08:42
@mxinO mxinO requested review from a team as code owners July 7, 2026 08:42
@mxinO mxinO requested a review from kinjalpatel27 July 7, 2026 08:42
@codecov

codecov Bot commented Jul 7, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.08%. Comparing base (bc5bc1a) to head (ffacac3).

Additional details and impacted files
@@            Coverage Diff             @@
##             main    #1935      +/-   ##
==========================================
- Coverage   77.69%   77.08%   -0.61%     
==========================================
  Files         517      517              
  Lines       57535    57537       +2     
==========================================
- Hits        44699    44350     -349     
- Misses      12836    13187     +351     
Flag Coverage Δ
examples 43.14% <50.00%> (-0.13%) ⬇️
gpu 57.95% <100.00%> (-0.68%) ⬇️
regression 14.95% <0.00%> (+0.11%) ⬆️
unit 55.16% <100.00%> (+<0.01%) ⬆️

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.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
modelopt/torch/quantization/nn/modules/tensor_quantizer.py (1)

1202-1213: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Update docstring to document the constant-amax case.

The docstring's return list omits the new (const) format now returned when _use_constant_amax is set.

📝 Suggested docstring update
     def _short_amax(self, fmt=".2e"):
         """Short description of amax.

         Returns:
             'dynamic': if _amax is not registered
             'amax': if _amax is per-tensor
             '[min, max](size)': if _amax is per-channel
+            'value(const)': if a constant amax is used (e.g. KV cache quantization)
         """
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@modelopt/torch/quantization/nn/modules/tensor_quantizer.py` around lines 1202
- 1213, The _short_amax method’s docstring is missing the constant-amax return
case now handled by the _use_constant_amax branch. Update the return description
in tensor_quantizer.TensorQuantizer._short_amax to document that it can return a
constant-amax string with the “(const)” suffix, alongside the existing
dynamic/per-tensor/per-channel cases, so the docstring matches the actual
behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@modelopt/torch/quantization/nn/modules/tensor_quantizer.py`:
- Around line 1202-1213: The _short_amax method’s docstring is missing the
constant-amax return case now handled by the _use_constant_amax branch. Update
the return description in tensor_quantizer.TensorQuantizer._short_amax to
document that it can return a constant-amax string with the “(const)” suffix,
alongside the existing dynamic/per-tensor/per-channel cases, so the docstring
matches the actual behavior.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8fbe8244-001c-42da-bf64-e143ab8c72e6

📥 Commits

Reviewing files that changed from the base of the PR and between bc5bc1a and ffacac3.

📒 Files selected for processing (2)
  • modelopt/torch/quantization/nn/modules/tensor_quantizer.py
  • tests/unit/torch/quantization/test_print.py

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant