[PyTorch][NVFP4][MOE] NVFP4 Grouped Quantize with Hadamard Transform

[zhongbozhu]

Description

• This PR depends on the merge of this PR: [PyTorch][NVFP4][MOE] NVFP4 Grouped Hadamard Amax Kernel #2351 which is a grouped kernel for generating NVFP4 global amax values.
• This PR is trying to build grouped-quantize kernels. There are two grouped quantize kernels needed for NVFP4: the regular quantize without Hadamard Transform, as well as a Hadamard Transform FP4 Cast fusion. This PR will prioritize the one with Hadamard and maybe support the other one (or in another future PR)

Status:

• Grouped Rowwise-only quantize without Hadamard (@Oleg-Goncharov for adding transpose mode to it in case RHT get disabled)
• Grouped Columnwise RHT-cast fusion: numerically correct
• Grouped Columnwise RHT-cast fusion: optimization of the kernel for better memory SOL (current memory SOL is low). It has been root caused to be related with running high precision math in the epilogue of the RHT gemm. Running fast math will boost the perf but comes with some minor numerical impact due to the difference between x / y != x * (1/y) even for FP32, but divides are slower. Use NVTE_RHT_CAST_FUSION_USE_FAST_MATH NVTE_USE_FAST_MATH env var to control this behavior.
• Further fuse Rowwise quantize & Colwise RHT transform & Colwise cast transpose into one kernel. This new kernel targeting SOL perf will further assume 128 padding in the token dimension. Current padding to NVFP4 is 64 multiple, but we might have to further increase it to 128 and then try other methods to remove the padding overhead to achieve overall SOL performance. We can eliminate padding overhead using the following tricks in the next section. Megatron PR to also bump up this padding.
• Make Row-quant and Col-quant use two different random numbers like this PR 2487 already did in unfused path, @negvet to review.

Notes:

• Why is zero padding bad?
  Zero padding means full tensor memory read and write, and often stems from the scaling factor layout of quantized recipes like mxfp8.

• How to remove the padding overhead?

1. Router padding: let the router outputs x multiple of tokens directly so that zero padding becomes an no-op. This method has been integrated to Megatron-core here.
2. Router padding to 128 multiple sounds scary and will raise concerns about if this will incur numerical issues. A better solution is to fuse the zero padding into the token permute kernel. This is still an on-going work, but Megatron-core with the HybridEP token dispatcher backend already supports this feature to pad to a multiple here.

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

[greptile-apps]

Greptile Summary

This PR implements grouped NVFP4 quantization kernels with Hadamard Transform (RHT) fusion for MOE workloads. The implementation provides three kernel variants:

Key Changes:

• Fully-fused kernel (group_row_cast_col_hadamard_transform_cast_fusion.cu): Combines rowwise quantization, columnwise RHT, and columnwise quantization into a single persistent kernel. Requires 128-multiple padding in token dimension for optimal memory SOL performance.
• Columnwise-only RHT fusion (group_hadamard_transform_cast_fusion.cu): Handles RHT + columnwise quantization with configurable fast math mode via NVTE_USE_FAST_MATH env var. Fast math trades minor numerical precision for better performance.
• Rowwise-only quantization (group_quantize_transpose_nvfp4.cuh): Grouped rowwise quantization without RHT, supporting multiple splits with proper tensor ID lookup.

RNG State Management:
The PR correctly addresses the critical issue of using different random seeds for rowwise vs columnwise quantization. When the fully-fused kernel cannot be used (non-128-aligned splits), separate RNG states are generated to ensure rowwise and colwise use independent random numbers.

Test Coverage:
Comprehensive tests validate correctness against reference implementation across multiple edge cases: zero tokens at various positions, uneven splits, RHT enabled/disabled, and random sign mask configurations.

Performance Notes:

• Hardcoded tile size selection based on specific (M, N) shapes in group_hadamard_transform_cast_fusion.cu:938-964
• Pre-RHT amax path disabled pending verification (line 866 in cast.cpp)
• 128-multiple padding requirement for full fusion enables better memory SOL but may add overhead that needs mitigation via router padding or fused token permute

Confidence Score: 4/5

• PR is safe to merge with minor considerations - well-tested core functionality with clear performance/accuracy tradeoffs documented
• Complex GPU kernel implementation with proper synchronization, memory management, and numerical validation. RNG state handling correctly addresses randomness independence. Comprehensive test coverage validates correctness. Score reduced from 5 due to: (1) hardcoded tile size tuning that doesn't scale, (2) disabled pre-RHT amax path pending verification, (3) fast math mode's numerical impact needs monitoring in production
• Pay close attention to group_hadamard_transform_cast_fusion.cu if adding new tensor shapes (requires manual tile size tuning). Monitor numerical stability when NVTE_USE_FAST_MATH=1

Important Files Changed

Filename	Overview
transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu	Fully-fused grouped kernel for rowwise+columnwise RHT quantization with persistent kernel scheduler. Complex but well-structured implementation with proper RNG handling.
transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu	Grouped columnwise RHT-cast fusion kernel using persistent tile processing. Implements fast math option for performance at slight numerical cost.
transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh	Grouped rowwise quantization kernel with TMA operations. Handles multi-tensor args and separate RNG states correctly.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Python bindings for split quantization with proper RNG state management. Correctly routes to fused vs unfused paths based on alignment and quantizer type.
tests/pytorch/nvfp4/test_nvfp4_group_quantize.py	Comprehensive test coverage for grouped quantization including edge cases (zero tokens, uneven splits, RHT on/off, random sign masks).

Sequence Diagram

sequenceDiagram
    participant User
    participant PyAPI as Python API<br/>(split_quantize)
    participant CSrc as C++ Bindings<br/>(cast.cpp)
    participant Alloc as Bulk Allocator
    participant RNG as RNG State Setup
    participant Kernel as CUDA Kernels
    
    User->>PyAPI: split_quantize(tensor, splits, quantizers)
    PyAPI->>CSrc: Convert tensors & quantizers
    
    alt NVFP4 Quantizers
        CSrc->>Alloc: bulk_allocate_nvfp4_tensors()
        Alloc-->>CSrc: Contiguous FP4 data + FP8 scales + amax
        
        alt 128-aligned splits + RHT enabled
            CSrc->>RNG: setup RNG states (shared for row+col)
            RNG-->>CSrc: Single RNG state tensor
            CSrc->>Kernel: nvte_group_hadamard_transform_cast_fusion<br/>(fully fused row+col)
            Kernel->>Kernel: Rowwise quantization
            Kernel->>Kernel: RHT + Colwise quantization
        else Unaligned or no RHT
            alt Need separate RNG
                CSrc->>RNG: setup separate row/col RNG states
                RNG-->>CSrc: Two RNG state tensors
            end
            CSrc->>Kernel: nvte_group_nvfp4_quantize_with_amax<br/>(rowwise)
            CSrc->>Kernel: nvte_group_hadamard_transform_cast_fusion_columnwise<br/>(colwise RHT)
        end
    else Other Quantizers
        CSrc->>Kernel: multi_tensor_quantize_impl<br/>(unfused path)
    end
    
    Kernel-->>CSrc: Quantized tensors
    CSrc-->>PyAPI: NVFP4Tensor objects
    PyAPI-->>User: List of quantized tensors

Loading

[greptile-apps]

_{16 files reviewed, 3 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

style: hardcoded tile sizes for specific tensor shapes

This switch statement hardcodes k_tile_size for specific (M, N) combinations. This approach:

• Doesn't scale to new shapes
• Requires manual tuning for each workload
• Falls back to generic logic that may be suboptimal

Consider implementing a heuristic-based selector that accounts for SM count, memory bandwidth, and tensor dimensions, or move this tuning to a configuration file.

[greptile-apps]

_{20 files reviewed, 6 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

style: hardcoded tile sizes don't scale to new shapes

This switch statement hardcodes k_tile_size for specific (M, N) combinations. Consider implementing a heuristic-based selector that accounts for SM count, memory bandwidth, and tensor dimensions, or move this tuning to a configuration file so it can be extended without code changes.

[greptile-apps]

style: math operations for recomputing scaling factors are expensive in persistent kernel

Consider caching pre-computed scaling factors for all possible tensor IDs at kernel start, since num_tensors is known and typically small (≤64). This would avoid recomputing ComputeGlobalEncodeScaleFP4 on every tensor boundary crossing.

[greptile-apps]

_{20 files reviewed, 8 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: Overflow check only validates >= 0 but doesn't catch actual int32_t overflow. Consider checking against INT_MAX or using size_t arithmetic. Should this check for actual overflow against INT_MAX rather than just non-negative values?

[greptile-apps]

Additional Comments (4)

1. transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu, line 79-92 (link)

   style: boundary guard prevents out-of-bounds but relies on implicit assumption

   The early return at line 85-87 prevents out-of-bounds access when offset >= boundary, but this assumes the host code correctly sets split_sections_range[num_tensors] as a sentinel value. While this works, the implicit dependency makes the code fragile to future changes.

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

2. transformer_engine/pytorch/csrc/extensions/cast.cpp, line 810-811 (link)

   logic: overflow check validates non-negative but doesn't prevent int32_t overflow

   NVTE_CHECK(kernel_args.split_sections_range[kernel_args.num_tensors + 1] >= 0, ...) only checks for negative values after overflow has occurred. With cumulative sums, overflow to negative is possible, but overflow within positive range (wrapping past INT_MAX) won't be caught. Should this validate against INT_MAX explicitly, or is the tensor size constrained elsewhere to make this impossible?

3. transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh, line 293-296 (link)

   style: tensor ID is fetched every stage even when unchanged

   GetTensorId is called for every stage to check if tensor boundary is crossed. While the check if (new_tensor_id != tensor_id) prevents unnecessary updates, the function call and boundary scan happen unconditionally. Could this be optimized by pre-computing which stages cross tensor boundaries, or would the complexity outweigh the benefit?

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

4. transformer_engine/pytorch/csrc/extensions/cast.cpp, line 813-815 (link)

   style: document the performance implications of 128-alignment

   The all_aligned_token_dim check determines whether the fully-fused high-performance kernel is used. Users should be aware when they're not getting optimal performance due to alignment. Should there be a runtime warning or logging when falling back to the unfused path due to alignment?

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

_{17 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (4)

1. transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu, line 75-86 (link)

   logic: no bounds check after loop termination

   The while loop exits when split_sections_range[tensor_id + 1] <= offset is false, but there's no verification that tensor_id < num_tensors before returning. If host sets incorrect split_sections_range or offset is corrupted, this could return an invalid tensor_id.

2. transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu, line 1001-1027 (link)

   style: hardcoded tile size tuning doesn't scale

   The switch-case hardcodes optimal k_tile_size for specific tensor shapes. This requires manual updates for new workloads and doesn't handle intermediate sizes well. The fallback logic at lines 1025-1027 only handles small tensors.

   Consider: Move tuning parameters to a config file, or implement auto-tuning based on SM count and tensor dimensions.

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

3. transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh, line 810-811 (link)

   logic: overflow check doesn't validate actual overflow

   The check only ensures non-negative, but split_sections_range[i+1] could legitimately overflow int32_t with large tensor sizes or many splits. Use size_t for accumulation or check against INT_MAX:

4. transformer_engine/pytorch/csrc/extensions/cast.cpp, line 712-760 (link)

   style: resource management for RNG states could leak on exceptions

   The StochasticRngStateResources struct holds tensors and vectors but lacks explicit cleanup. If an exception occurs after partial initialization, resources may not be properly released. Consider adding RAII wrappers or using smart pointers consistently.

_{20 files reviewed, 4 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Additional Comments (1)

1. transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu, line 1289-1292 (link)

   logic: potential memory leak if kernel launch fails

   If NVTE_CHECK throws/aborts at line 1290, the workspace allocated at line 1281 won't be freed. Move the free before the status check, or use RAII wrapper.

_{19 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[timmoon10]

/te-ci L1