Skip to content

[Common] Ensure quantization kernels handle noop properly - #3271

Open
kainzhong wants to merge 1 commit into
NVIDIA:mainfrom
kainzhong:noop_tensor_fix
Open

[Common] Ensure quantization kernels handle noop properly#3271
kainzhong wants to merge 1 commit into
NVIDIA:mainfrom
kainzhong:noop_tensor_fix

Conversation

@kainzhong

@kainzhong kainzhong commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Description

Fixes a few places where quantization treats the noop flag wrong.

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:

  • MXFP8 & Group MXFP8 now ignore noop with IS_DBIAS since the reduction kernel will run unconditionally, so if we skip quantization with IS_DBIAS the results would be wrong.
  • MXFP8's zero_scales_kernel now will ignore noop using the same condition as how the quantization ignores it (otherwise if we skip zeroing but not skip quantization the result could be wrong)
  • Blockwise FP8's group_block_scaled_1d_tma_kernel and group_block_scaled_2d_tma_kernel now ignore the noop flag when dbias workspace is nullptr (I suppose this means IS_DBIAS) for the same reason
  • Pass noop flags to FP8 quantization
  • Add tests to check if noop flag is honored

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

Comment thread transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh Outdated
@greptile-apps

greptile-apps Bot commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

The PR aligns noop handling across FP8, MXFP8, grouped MXFP8, and blockwise FP8 quantization.

  • Passes the noop flag into the aligned FP8 quantization kernel.
  • Keeps quantization, scale zeroing, and unconditional DBIAS reduction behavior consistent.
  • Adds CUDA-graph tests covering noop behavior across quantization formats, output layouts, and shapes.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the MXFP8 DBIAS path now passes a null noop pointer consistently to both quantization and padded-scale zeroing, so the previously reported stale-scale mismatch is fixed.

Important Files Changed

Filename Overview
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh Makes DBIAS and activation variants consistently ignore noop in both quantization and padded-scale zeroing, resolving the previously reported stale-scale failure.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh Prevents DBIAS-only grouped MXFP8 quantization from skipping workspace production required by unconditional reduction.
transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh Restricts noop early returns to launches without a DBIAS workspace so reduction inputs are always produced.
transformer_engine/common/cast/fp8/quantize_fp8.cuh Propagates noop into the aligned one-dimensional FP8 kernel and skips its output updates when requested.
tests/pytorch/test_quantized_tensor.py Adds CUDA-graph coverage confirming noop preserves quantized data and scale buffers across supported formats and layouts.

Reviews (3): Last reviewed commit: "[Common] Ensure quantization kernels han..." | Re-trigger Greptile

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
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