Skip to content

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

Open
kainzhong wants to merge 3 commits into
NVIDIA:mainfrom
kainzhong:noop_tensor_fix
Open

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

Conversation

@kainzhong

Copy link
Copy Markdown
Collaborator

Description

When quantization kernels are fused with act / dact / dbias, they should ignore the noop flag. Currently we mostly rely on the caller passing the right noop flag but we don't have guards in the kernel.

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:

  • Quantization kernels now ignore noop when fused
  • MXFP8 now also ignores noop with IS_DBIAS. Also zero_scales_kernel now will ignore noop when the quantization ignores it (otherwise if we skip zeroing but not skip quantization the result could be wrong)
  • FP8 kernels now correctly respect noop tensor when it's safe to skip execution
  • Add tests

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

This PR corrects noop handling across quantization kernels.

  • Fused activation and DBIAS paths now execute regardless of the noop flag.
  • Safe non-fused FP8 and MXFP8 paths continue to skip quantization when noop is set.
  • MXFP8 scale zeroing now follows the same effective noop behavior as quantization.
  • PyTorch CUDA-graph tests verify that skipped quantization preserves output buffers.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename Overview
transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh Aligns scale zeroing and quantization by clearing the effective noop pointer for fused activation and DBIAS variants.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh Prevents grouped DBIAS MXFP8 quantization from incorrectly returning early on noop.
transformer_engine/common/cast/fp8/quantize_fp8.cuh Propagates noop into safe non-fused FP8 kernels while suppressing it for activation-fused launches.
transformer_engine/common/cast/fp8/group_quantize_fp8.cuh Restricts grouped FP8 noop early returns to non-activation variants.
transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh Keeps DBIAS workspace generation active because the subsequent reduction launch is unconditional.
tests/pytorch/test_quantized_tensor.py Adds CUDA-graph coverage ensuring noop preserves quantized data and scale buffers across formats and layouts.

Reviews (2): Last reviewed commit: "Update transformer_engine/common/cast/mx..." | Re-trigger Greptile

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Kaining Zhong <44538064+kainzhong@users.noreply.github.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