Repository: huggingface/transformers
File: src/transformers/integrations/finegrained_fp8.py (Python)

Code before:

def w8a8_block_fp8_matmul(
    A: torch.Tensor,
    B: torch.Tensor,
    As: torch.Tensor,
    Bs: torch.Tensor,
    block_size: list[int],
    output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """
    Dispatch to CUTLASS or Triton for block-wise FP8 matmul.

    Uses CUTLASS when:
    - Block size is [128, 128] (the only size CUTLASS supports)
    - Running on SM90+ (Hopper or newer)
    - The CUTLASS kernel is available
    - Output dtype is bfloat16 or float16 (CUTLASS requirement)
    - Tensor dimensions are compatible (divisible by 16)

    Otherwise falls back to Triton.
    """
    # Check if we can use CUTLASS
    # CUTLASS only supports bfloat16/float16 output
    cutlass_compatible_dtype = output_dtype in (torch.bfloat16, torch.float16)

    if cutlass_compatible_dtype and _supports_cutlass_blockwise_fp8(block_size):
        kernel = _get_quantization_kernel()
        if kernel is not None:
            try:
                # CUTLASS expects:
                # - A: [M, K] row-major, float8_e4m3fn
                # - B: [K, N] column-major, float8_e4m3fn
                # - As: [M, K//128] M-major (activation scales)
                # - Bs: [K//128, N//128] K-major (weight scales)

                # Reshape A to 2D if needed
                original_shape = A.shape
                M = A.numel() // A.shape[-1]
                K = A.shape[-1]
                N = B.shape[0]

                # CUTLASS requires dimensions divisible by 16
                if K % 16 != 0 or N % 16 != 0:
                    raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16")

                A_2d = A.view(M, K).contiguous()
                # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1
                # Our B is [N, K] row-major. Make it contiguous first, then transpose.
                # B.contiguous() gives [N, K] with stride=(K,1)
                # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major
                # Do NOT call .contiguous() after .t() as it would make it row-major!

Review comment (suggestion):
let's put the dtype check in `_supports_cutlass_blockwise_fp8`

Code after:

    return C


def w8a8_block_fp8_matmul(
    A: torch.Tensor,
    B: torch.Tensor,
    As: torch.Tensor,
    Bs: torch.Tensor,
    block_size: list[int],
    output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """
    Dispatch to CUTLASS or Triton for block-wise FP8 matmul.

    Uses CUTLASS when:
    - Block size is [128, 128] (the only size CUTLASS supports)
    - Running on SM90+ (Hopper or newer)
    - The CUTLASS kernel is available
    - Output dtype is bfloat16 or float16 (CUTLASS requirement)
    - Tensor dimensions are compatible (divisible by 16)

    Otherwise falls back to Triton.
    """

    if _supports_cutlass(block_size, output_dtype):
        kernel = _get_quantization_kernel()
        if kernel is not None:
            try:
                # CUTLASS expects:
                # - A: [M, K] row-major, float8_e4m3fn
                # - B: [K, N] column-major, float8_e4m3fn
                # - As: [M, K//128] M-major (activation scales)
                # - Bs: [K//128, N//128] K-major (weight scales)

                # Reshape A to 2D if needed
                original_shape = A.shape
                M = A.numel() // A.shape[-1]
                K = A.shape[-1]
                N = B.shape[0]

                # CUTLASS requires dimensions divisible by 16
                if K % 16 != 0 or N % 16 != 0:
                    raise ValueError(f"CUTLASS requires K ({K}) and N ({N}) divisible by 16")

                A_2d = A.view(M, K).contiguous()
                # B needs to be column-major for CUTLASS: [K, N] with stride(0)==1
                # Our B is [N, K] row-major. Make it contiguous first, then transpose.
                # B.contiguous() gives [N, K] with stride=(K,1)
                # B.contiguous().t() gives [K, N] with stride=(1,K) which is column-major
                # Do NOT call .contiguous() after .t() as it would make it row-major!