[Common] MXFP8 kernel for grouped tensors

[Oleg-Goncharov]

Description

This PR adds a new kernel that supports MXFP8 quantization of grouped tensors.

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

• Added MXFP8 cast kernel for grouped tensors
• Added the 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

• Implements MXFP8 quantization kernel for grouped tensors, enabling efficient processing of multiple tensors with varying dimensions in a single operation
• Extends all major activation functions (GeLU, SiLU, ReLU, etc.) to support grouped tensor variants with comprehensive C API coverage
• Adds extensive test coverage with 777 test configurations validating correctness across different tensor shapes, scaling modes, and activation functions

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh	New CUDA kernel implementation for grouped tensor MXFP8 quantization with TMA operations and multi-stage processing
tests/cpp/operator/test_cast_mxfp8_grouped.cu	New comprehensive test suite with 777 test configurations validating grouped tensor quantization functionality
transformer_engine/common/cast/dispatch/quantize.cuh	Adds dispatch helper functions for grouped tensor quantization operations
transformer_engine/common/include/transformer_engine/cast.h	Extends C API with 7 new function declarations for grouped tensor quantization operations

Confidence score: 4/5

• This PR is generally safe to merge but requires careful review due to complex GPU kernel implementation and API extensions
• Score reflects sophisticated CUDA kernel code with architecture-specific optimizations and potential edge cases in TMA descriptor handling
• Pay close attention to the main CUDA kernel implementation and verify the API parameter type consistency for gradient functions

Sequence Diagram

sequenceDiagram
    participant User
    participant TestSuite as "GroupedFusedCastMXFP8TestSuite"
    participant Helper as "performTest<IType, OType>"
    participant KernelDispatch as "group_quantize_*_helper"
    participant Kernel as "group_quantize_mxfp8_kernel"
    participant GPU as "CUDA Device"

    User->>TestSuite: "Run test with parameters"
    TestSuite->>TestSuite: "Configure tensors and validation"
    TestSuite->>Helper: "performTest(processing_method, OP, ...)"
    Helper->>Helper: "Setup input/output tensors and reference data"
    Helper->>Helper: "Create NVTE grouped tensors"
    
    alt Processing Method: CAST_ONLY
        Helper->>KernelDispatch: "nvte_group_quantize(input, output)"
    else Processing Method: CAST_DBIAS
        Helper->>KernelDispatch: "nvte_group_quantize_dbias(grad, output, dbias)"
    else Processing Method: CAST_ACT
        Helper->>KernelDispatch: "nvte_group_gelu/silu/relu(input, output)"
    else Processing Method: CAST_DACT
        Helper->>KernelDispatch: "nvte_group_dgelu/dsilu/drelu(grad, input, output)"
    else Processing Method: CAST_DBIAS_DACT
        Helper->>KernelDispatch: "nvte_group_quantize_dbias_dgelu/dsilu/drelu(grad, input, output, dbias)"
    end
    
    KernelDispatch->>KernelDispatch: "group_quantize_fwd/bwd_helper"
    KernelDispatch->>Kernel: "update_tma_descriptors<<<num_tensors, 32>>>"
    GPU->>KernelDispatch: "TMA descriptors updated"
    KernelDispatch->>Kernel: "group_quantize_mxfp8_kernel<<<blocks, threads>>>"
    GPU->>Kernel: "Execute MXFP8 quantization with scaling"
    Kernel->>GPU: "Write quantized data and scaling factors"
    GPU->>KernelDispatch: "Kernel execution complete"
    
    alt IS_DBIAS enabled
        KernelDispatch->>GPU: "reduce_dbias(workspace, dbias)"
    end
    
    KernelDispatch->>Helper: "Quantized tensors ready"
    Helper->>Helper: "Compare GPU results with CPU reference"
    Helper->>TestSuite: "Validation complete"
    TestSuite->>User: "Test result"

Loading

[greptile-apps]

_{10 files reviewed, 6 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

logic: commented-out code in case statement creates ambiguity

The commented calculation rows_num = first_logical_dim / num_tensors; break; suggests this case should compute rows differently, but the fallthrough to VARYING_LAST_DIM (which sets rows_num = first_logical_dim) may not be the intended behavior. If fallthrough is intentional, remove the comment to clarify. If the division by num_tensors is needed for SAME_BOTH_DIMS, uncomment and add the break statement.

[greptile-apps]

_{10 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

style: missing newline after closing brace before constexpr declarations

This formatting issue breaks the visual separation between the conditional compilation block and the subsequent constant declarations. Add a newline for better readability.

Suggested change

	} constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X;
	}
	constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X;

_{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!}