Skip to content

[Pytorch][Common] Option to Disable 2nd level Scale in NVFP4 - #3297

Open
vthumbe1503 wants to merge 3 commits into
NVIDIA:mainfrom
vthumbe1503:disable_2d_scaling
Open

[Pytorch][Common] Option to Disable 2nd level Scale in NVFP4#3297
vthumbe1503 wants to merge 3 commits into
NVIDIA:mainfrom
vthumbe1503:disable_2d_scaling

Conversation

@vthumbe1503

Copy link
Copy Markdown
Collaborator

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

Please list the changes introduced in this PR:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
@greptile-apps

greptile-apps Bot commented Jul 31, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Adds an option to omit NVFP4 second-level scaling and represent unit global scaling with absent amax buffers.

  • Propagates optional amax metadata through PyTorch quantizers, storage, grouped allocation, and conversion paths.
  • Updates quantization, dequantization, Hadamard-fusion, recipe, transpose, and grouped-GEMM kernels to handle null amax pointers.
  • Adds exactness tests for quantization, grouped quantization, output reuse, and GEMM behavior.

Confidence Score: 4/5

The restricted-max grouped-GEMM scaling defect should be fixed before merging because supported 4over6 NVFP4 configurations can produce incorrectly scaled results.

Grouped GEMM accepts operands quantized with an E4M3 maximum of 256, but the changed alpha path always normalizes their amax values using the unrestricted maximum of 448, breaking the encode/decode scaling contract.

Files Needing Attention: transformer_engine/common/gemm/cublaslt_grouped_gemm.cu

Important Files Changed

Filename Overview
transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Adds null-amax grouped-GEMM support, but its fixed 448-based alpha normalization mis-scales supported 256-max NVFP4 operands.
transformer_engine/pytorch/csrc/quantizer.cpp Propagates the new quantizer option through allocation, conversion, amax computation, and reuse paths.
transformer_engine/pytorch/tensor/nvfp4_tensor.py Exposes and copies the new option while disabling the incompatible row-scaled mode.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py Makes amax metadata optional and preserves missing-amax semantics across usage and transpose handling.
transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh Interprets null amax as the unit-global-scale sentinel while retaining row-scaled validation.
transformer_engine/pytorch/csrc/extensions/cast.cpp Omits amax allocation and computation across grouped and split quantization when second-level scaling is disabled.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Q["NVFP4Quantizer"] -->|"disable_second_level_scale=false"| A["Allocate and compute amax"]
  Q -->|"disable_second_level_scale=true"| N["Leave amax absent"]
  A --> K["NVFP4 quantize/dequantize kernels"]
  N --> K
  K -->|"present amax"| S["Two-level global + block scaling"]
  K -->|"null amax"| U["Unit global scale + block scaling"]
  S --> G["GEMM alpha correction"]
  U --> G
Loading

Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile

Comment on lines +1404 to +1406
constexpr float kUnitGlobalScaleAmax =
transformer_engine::nvfp4::unit_global_scale_amax<transformer_engine::fp8e4m3,
transformer_engine::fp4e2m1>();

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.

P1 Restricted-max GEMM scaling mismatch

When grouped GEMM receives 4over6 NVFP4 operands configured with nvfp4_e4m3_max=256, this path normalizes their amax values using the unrestricted 448*6 maximum instead of the 256*6 value used during quantization, causing numerically mis-scaled GEMM output.

Knowledge Base Used: Common core (transformer_engine/common)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

We dont have have the 4over6 support today for the graph safe path today(group_quantize + grouped_gemm). Do we need that support now @timmoon10 @ptrendx ? I think it would just increase the scope of this PR so it shouldnt be added now.

@vthumbe1503
vthumbe1503 requested review from ptrendx and tdophung July 31, 2026 22:01
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