[Common] Persistent Grouped NVFP4 quantization kernel

[Oleg-Goncharov]

Description

This PR adds a persistent grouped NVFP4 quantization + transpose kernel with static scheduling.
It is built on top of the PR#2738 [Common] Persistent Grouped MXFP8 quantization 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

• Added persistent grouped kernel
• Added test suite

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

[greptile-apps]

Greptile Summary

This PR introduces a new persistent grouped NVFP4 quantize+transpose kernel (group_quantize_transpose_nvfp4_tuned_1D.cuh) with static grid-stride scheduling, double-buffered TMA prefetch, and per-tensor scale-offset bookkeeping. It also promotes the dbias output parameter from NVTETensor to NVTEGroupedTensor across the grouped activation/quantization API.

The new kernel itself is well-designed. However, the PR contains several changes that appear to be development-time leftovers that would cause serious regressions if merged as-is:

• All pre-existing single-tensor quantization is silently broken: Both quantize_fwd_helper and quantize_bwd_helper in quantize.cuh have their entire bodies commented out, leaving empty no-op stubs for every scaling mode.
• MXFP8 grouped quantization is disabled: The NVTE_MXFP8_1D_SCALING cases in both group_quantize_fwd_helper and group_quantize_bwd_helper are commented out, meaning any grouped MXFP8 call hits NVTE_ERROR.
• Test CI coverage gutted: tests/cpp/operator/CMakeLists.txt comments out all 29 existing test sources, and tests/cpp/CMakeLists.txt restricts the fallback CUDA architecture to Blackwell-only (sm_100). These would prevent detection of any regressions introduced by the above changes.
• Debug utilities left in test code: dump_nvfp4_tensor_data and print_detailed_tensor_comparison in the new test file write files to disk and print verbose tables to stdout; these should not ship in production test code.
• Unused variable: truncated_pair is computed but never read in the reference quantize_nvfp4 helper.

Confidence Score: 1/5

• Not safe to merge — pre-existing quantization paths are silently disabled and the test suite is stripped down to a single new test file.
• The new kernel implementation looks sound in isolation, but the PR inadvertently (or as a development shortcut) comments out the entire bodies of quantize_fwd_helper and quantize_bwd_helper, disables the MXFP8 grouped dispatch, restricts test builds to Blackwell only, and comments out all pre-existing operator tests. These changes would silently break all non-grouped and all MXFP8 quantization in the library, and the disabled test suite would prevent CI from catching it.
• transformer_engine/common/cast/dispatch/quantize.cuh (empty fwd/bwd helpers, disabled MXFP8 case), tests/cpp/CMakeLists.txt (architecture restriction), tests/cpp/operator/CMakeLists.txt (all tests commented out).

Important Files Changed

Filename	Overview
tests/cpp/CMakeLists.txt	CUDA architecture narrowed to sm_100 only in the fallback path, silently dropping Volta/Ampere/Ada/Hopper coverage; old line left as a comment.
tests/cpp/operator/CMakeLists.txt	All 29 existing operator test sources commented out; only the new nvfp4 grouped test is compiled, completely suppressing regression coverage.
transformer_engine/common/cast/dispatch/quantize.cuh	Both quantize_fwd_helper and quantize_bwd_helper are now empty stubs (entire bodies commented out); MXFP8 grouped dispatch is also disabled — all pre-existing quantization paths are silently broken.
transformer_engine/common/cast/nvfp4/specialized/group_quantize_transpose_nvfp4_tuned_1D.cuh	New persistent grouped NVFP4 quantize+transpose kernel with static grid-stride scheduling, double-buffered TMA prefetch, and separate rowwise/colwise scaling helpers; well-structured with configurable TunableConfig.
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	Refactored to share ShapeRepresentation enum and JobDescriptor/BlockDescriptor structs from common; MXFP8 kernel now also has persistent-mode support, but is currently unreachable because the dispatch in quantize.cuh is commented out.
transformer_engine/common/cast/core/common.cuh	Adds ShapeRepresentation enum (moved from mxfp8) and a new group_reduce_dbias_kernel / grouped_reduce_dbias for multi-tensor dbias reduction; changes are self-contained and correct.
tests/cpp/operator/test_cast_nvfp4_transpose_grouped.cu	New grouped NVFP4 quantize+transpose test suite; logic is sound but includes unused truncated_pair variable and debug file-dump/verbose-print utilities that should not ship in production tests.
tests/cpp/operator/test_cast_mxfp8_grouped.cu	Updated to support per-tensor dbias comparison across all group sizes; the old single-tensor special-casing is replaced with a unified loop, which is cleaner and more correct.
transformer_engine/common/include/transformer_engine/cast.h	API signature for nvte_group_quantize_dbias updated: dbias parameter promoted from NVTETensor to NVTEGroupedTensor, consistent with the new per-tensor dbias layout.
transformer_engine/common/activation/gelu.cu	dbias parameters in nvte_group_quantize_dbias_dgelu / nvte_group_quantize_dbias_dqgelu and internal local nulls updated from NVTETensor to NVTEGroupedTensor; straightforward signature alignment.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["nvte_group_quantize()"] --> B["group_quantize_fwd_helper()"]
    B --> C{scaling_mode?}
    C -->|NVTE_NVFP4_1D_SCALING| D["nvfp4::group_quantize_transpose()"]
    C -->|NVTE_MXFP8_1D_SCALING| E["❌ NVTE_ERROR\n(case commented out)"]
    C -->|other| F["NVTE_ERROR"]

    D --> G{use_single_work_grid?}
    G -->|SAME_BOTH_DIMS or\nVARYING_FIRST_DIM| H["Single contiguous work grid\n(ctaid_Y * cols + ctaid_X)"]
    G -->|VARYING_LAST_DIM or\nVARYING_BOTH_DIMS| I["Per-element block grid\n(block_id * ELTS_PER_CHUNK)"]

    H --> J["Persistent kernel\n(static grid-stride loop)"]
    I --> J

    J --> K["decode_job()\nResolve tensor_id, rows, cols"]
    K --> L["is_job_valid()\nBounds check vs offsets"]
    L -->|invalid| M["Drain in-flight TMA\nthen exit"]
    L -->|valid| N["decode_block()\nblock_id_Y, block_id_X, offsets"]

    N --> O["Prime pipeline:\nPrefetch PREFETCH_STAGES stage(s)\nvia TMA global→shared"]
    O --> P["Main stage loop\n(STAGES = TILES_Y × TILES_X)"]

    P --> Q["process_colwise_stage()\nAmax + E4M3 scale + FP4 quantize"]
    P --> R["process_rowwise_stage()\nAmax + E4M3 scale + FP4 quantize"]
    Q --> S["store_output_stage()\nTMA shared→global commit"]
    R --> S

    S --> T{Last PREFETCH_STAGES?}
    T -->|yes| U["Prefetch stage-0\nof next job"]
    U --> P
    T -->|next job in stride| V["Advance ctaid\nstatic_next_block_id += stride"]
    V --> K
    T -->|no more jobs| W["cp_async_bulk_wait\nthen exit CTA"]

Loading

Comments Outside Diff (1)

1. transformer_engine/common/cast/dispatch/quantize.cuh, line 175-314 (link)

   quantize_bwd_helper body completely commented out — function is now a no-op

   Identical to quantize_fwd_helper: the full backward quantization dispatch (all scaling modes) has been commented out. Any backward quantization pass — including the critical IS_DBIAS and IS_DACT paths used during training — now silently does nothing. This must be restored before merging.

_{Last reviewed commit: 46fdb93}

[greptile-apps]

Development-only architecture restriction should not be merged

The old multi-architecture fallback (75 80 89 90) has been replaced with a Blackwell-only build (100). The original line is left behind as a commented-out breadcrumb. This means on any CI runner with CUDA < 12.8 the test binary will only compile for sm_100, silently skipping all Volta/Ampere/Ada/Hopper targets. This is clearly a local development shortcut and must be reverted before merging.

Suggested change

	# set(CMAKE_CUDA_ARCHITECTURES 75 80 89 90)
	set(CMAKE_CUDA_ARCHITECTURES 100)
	set(CMAKE_CUDA_ARCHITECTURES 75 80 89 90)

[greptile-apps]

All existing operator tests commented out

Every pre-existing test (test_cast.cu, test_cast_mxfp8.cu, test_cast_mxfp8_grouped.cu, test_normalization.cu, etc.) has been commented out, leaving only test_cast_nvfp4_transpose_grouped.cu in the build. This completely disables the existing regression suite and hides any breakage introduced by the changes in quantize.cuh and group_quantize_mxfp8.cuh. These comment-outs appear to be a local development convenience and must be reverted before merging.

The new test file should simply be added to the existing list, not substituted for it.

[greptile-apps]

quantize_fwd_helper body completely commented out — function is now a no-op

The entire implementation of quantize_fwd_helper (covering NVTE_DELAYED_TENSOR_SCALING, NVTE_MXFP8_1D_SCALING, NVTE_NVFP4_1D_SCALING, NVTE_BLOCK_SCALING_2D, and NVTE_BLOCK_SCALING_1D) has been replaced with block comments, leaving the function as an empty stub. Any call to single-tensor forward quantization (nvte_quantize, nvte_quantize_dbias, etc.) now silently does nothing.

This is a severe regression that would cause all non-grouped quantization code paths to silently produce uninitialized output. These commented-out lines must be restored.

[greptile-apps]

MXFP8 grouped quantization dispatch disabled

The NVTE_MXFP8_1D_SCALING case has been commented out from both group_quantize_fwd_helper (line ~410) and group_quantize_bwd_helper (line ~464). As a result, any call to grouped MXFP8 forward or backward quantization falls through to the default branch and calls NVTE_ERROR, breaking existing MXFP8 grouped workflows.

If this is intentional (e.g., MXFP8 is being refactored in parallel), the PR description and checklist should call this out explicitly as a breaking change, and the NVTE_MXFP8_1D_SCALING cases should at least be preserved in a functional state or guarded by a feature flag.

[greptile-apps]

Unused variable truncated_pair

truncated_pair is computed from cvt_fp4x2_to_double2(casted_to_e2m1_pair) but its value is never read. This will generate a compiler warning and the dead computation should be removed.

Suggested change

const double2 truncated_pair = cvt_fp4x2_to_double2(casted_to_e2m1_pair);

[greptile-apps]

Debug dump and verbose comparison utilities should be removed or hidden before merging

dump_nvfp4_tensor_data (line 234) writes test/ref/diff text files to disk during CI runs. print_detailed_tensor_comparison (line 271) prints element-by-element tables to stdout. Both functions are only conditionally exercised and print_detailed_tensor_comparison is already fully commented out at the call site (lines 323–324). These are development-time diagnostic aids that add noise, may fail in read-only CI environments, and were not part of any documented public API. They should be removed from the final PR, or at minimum wrapped in #ifdef NVTE_DEBUG_DUMP guards so they never run in CI.