Greptile Overview

Greptile Summary

This PR implements MXFP8 quantization for grouped tensors, adding a new GPU kernel that uses TMA (Tensor Memory Accelerator) descriptors for efficient data transfer and O(log N) binary search for tensor identification in grouped tensor scenarios.

Key Changes:

• Added group_quantize_mxfp8.cuh with the core MXFP8 grouped quantization kernel supporting rowwise/columnwise scaling
• Implemented TMA descriptor update mechanism with O(log N) complexity for varying tensor shapes
• Extended C API in cast.h with 7 new grouped tensor quantization functions
• Added grouped variants for activation functions (GeLU, ReLU, SiLU, QGeLU, SReLU) with dbias support
• Comprehensive test suite with reference implementation covering multiple shape representations

Issues from Previous Comments:
Previous review threads identified several concerns that appear to remain unresolved:

• Binary search underflow risk when current_offset < offsets_ptr[0] (line 85)
• Uninitialized shape_rep variable if none of the four shape conditions match (line 755-764)
• Commented-out code creating ambiguity in switch statement (line 104)
• Typos: "gropued" instead of "grouped" in multiple locations in cast.h
• Missing newline after conditional compilation block

Architecture:
The implementation uses a two-kernel approach: first updating TMA descriptors per tensor, then launching the main quantization kernel that uses binary search to identify which tensor each block processes. The kernel supports both single-tensor (constant last dimension) and multi-tensor cases with different optimization paths.

Confidence Score: 3/5

• This PR requires addressing several logic issues before merging, particularly around variable initialization and edge case handling
• Score reflects substantial new functionality with comprehensive tests, but multiple unresolved concerns from previous review including potential binary search underflow, uninitialized variables, and typos that need to be addressed
• Primary attention needed on transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh for initialization and edge case issues, and transformer_engine/common/include/transformer_engine/cast.h for typo corrections

Important Files Changed

Sequence Diagram

sequenceDiagram
    participant User
    participant API as C API Layer<br/>(cast.cu)
    participant Dispatcher as Dispatch Layer<br/>(quantize.cuh)
    participant Kernel as MXFP8 Kernel<br/>(group_quantize_mxfp8.cuh)
    participant GPU as GPU Device

    User->>API: nvte_group_quantize(input, output, stream)
    API->>Dispatcher: group_quantize_fwd_helper()
    Dispatcher->>Dispatcher: Convert NVTEGroupedTensor to GroupedTensor*
    Dispatcher->>Dispatcher: Check scaling_mode == NVTE_MXFP8_1D_SCALING
    
    alt Multi-tensor case (VARYING_LAST_DIM or VARYING_BOTH_DIMS)
        Dispatcher->>Kernel: update_tma_descriptors<<<num_tensors, 32>>>()
        Kernel->>GPU: Launch descriptor update kernel
        loop For each tensor in group
            GPU->>GPU: modify_base_tensor_map()<br/>Update tensor map for each tensor's data pointer
        end
        GPU-->>Kernel: TMA descriptors updated
    end
    
    Dispatcher->>Kernel: group_quantize_mxfp8_kernel<<<blocks, 128>>>()
    Kernel->>GPU: Launch main quantization kernel
    
    loop For each block
        GPU->>GPU: get_current_tensor_id()<br/>Binary search to find tensor ID
        GPU->>GPU: Acquire TMA fence for tensor map
        GPU->>GPU: TMA load input data to shared memory
        
        alt COLWISE_SCALING
            GPU->>GPU: Compute column-wise AMAX
            GPU->>GPU: Generate E8M0 scale factor
            GPU->>GPU: Quantize to MXFP8 with column-wise scale
            GPU->>GPU: TMA store to global memory
        end
        
        alt ROWWISE_SCALING
            GPU->>GPU: Compute row-wise AMAX
            GPU->>GPU: Generate E8M0 scale factor
            GPU->>GPU: Quantize to MXFP8 with row-wise scale
            GPU->>GPU: TMA store to global memory
        end
    end
    
    GPU-->>Dispatcher: Quantization complete
    Dispatcher-->>API: Return
    API-->>User: Return

/te-ci L1 pytorch

@Oleg-Goncharov, I have tested grouped_quantize from a Pytorch binding created for nvte_grouped_quantize and it works fine for all four cases of (first_dims, last_dims). And the changes in the PR look ok to me based on what I could understand. Could we merge this @Oleg-Goncharov @ptrendx ?

/te-ci pytorch

/te-ci pytorch

Thank you for checking it, @vthumbe1503. I’m working on extending dbias support to output a split grouped tensor, since the kernel currently accumulates dbias into a single tensor. Let’s merge this once that’s in.

@Oleg-Goncharov Please open a new PR with the proper dbias support - let's try to minimize the review effort.