Skip to content

[Common] Group NVFP4 Quantize Kernels - #3458

Open
Oleg-Goncharov wants to merge 6 commits into
NVIDIA:mainfrom
Oleg-Goncharov:pr_nvfp4_group_kernels
Open

[Common] Group NVFP4 Quantize Kernels #3458
Oleg-Goncharov wants to merge 6 commits into
NVIDIA:mainfrom
Oleg-Goncharov:pr_nvfp4_group_kernels

Conversation

@Oleg-Goncharov

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

Oleg-Goncharov and others added 4 commits August 31, 2026 14:52
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
@greptile-apps

greptile-apps Bot commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR adds coordinated grouped NVFP4 1D quantization and transpose kernels, integrates them with grouped dispatch and layout infrastructure, and adds extensive native test coverage.

  • Adds grouped NVFP4 quantization, transpose, scaling, and TMA/layout support.
  • Restores the delayed-scaling, MXFP8, and FP8 block-scaling dispatch branches identified in the previous review.
  • Registers grouped NVFP4 transpose tests covering multiple group counts and shape representations.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported grouped scaling-mode dispatch branches are active at the current head.

Important Files Changed

Filename Overview
transformer_engine/common/cast/dispatch/quantize.cuh Integrates grouped NVFP4 dispatch and restores the previously disabled format-specific grouped quantization branches.
transformer_engine/common/cast/nvfp4/specialized/group_quantize_transpose_nvfp4_tuned_1D.cuh Implements the tuned grouped NVFP4 1D quantize-and-transpose launch path with shape and allocation validation.
transformer_engine/common/cast/core/grouped_layout.cuh Extends shared grouped-layout infrastructure used to map variable-sized groups to coordinated kernel work.
transformer_engine/common/cast/core/grouped_tma.cuh Extends grouped TMA support required by the new NVFP4 kernel path.
tests/cpp/operator/test_cast_nvfp4_transpose_grouped.cu Adds numerical coverage for grouped NVFP4 rowwise and columnwise output across group counts and shape representations.
tests/cpp/operator/CMakeLists.txt Registers the grouped NVFP4 transpose test in the native operator test executable.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  I[Grouped BF16 input] --> D[Grouped quantization dispatch]
  D --> K[Grouped NVFP4 quantize and transpose kernel]
  K --> R[Rowwise packed FP4 data and scales]
  K --> C[Columnwise packed FP4 data and scales]
  R --> G[Grouped GEMM consumers]
  C --> G
Loading

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

Comment thread transformer_engine/common/cast/dispatch/quantize.cuh Outdated
Signed-off-by: Oleg Goncharov <ogoncharov@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