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

Code before:


if is_torch_available():
    import torch
    import torch.nn as nn
    import triton
    import triton.language as tl
    from torch.nn import functional as F


logger = logging.get_logger(__name__)

# Global for the CUTLASS quantization kernel (lazily loaded)
_quantization_kernel = None


def _get_quantization_kernel():
    """Lazily load the CUTLASS quantization kernel from HuggingFace Hub."""
    global _quantization_kernel
    if _quantization_kernel is None:
        try:
            from .hub_kernels import get_kernel

            _quantization_kernel = get_kernel("RedHatAI/quantization")
        except Exception as e:
            logger.warning_once(f"Failed to load CUTLASS quantization kernel: {e}. Falling back to Triton.")
            _quantization_kernel = False  # Mark as unavailable
    return _quantization_kernel if _quantization_kernel else None


def _supports_cutlass_blockwise_fp8(block_size: list[int] | None) -> bool:
    """
    Check if CUTLASS blockwise FP8 matmul is supported for the given block size.

    CUTLASS blockwise kernels require:
    - SM90+ (Hopper or newer)
    - Block size [128, 128] for weights
    - Block size [1, 128] for activations (handled implicitly)
    """
    if not is_torch_available() or not torch.cuda.is_available():
        return False

    # Check block size compatibility - CUTLASS only supports [128, 128]
    if block_size is None:
        return False
    if len(block_size) != 2 or block_size[0] != 128 or block_size[1] != 128:
        return False

    # Check GPU capability (SM90+)
    try:
        capability = torch.cuda.get_device_capability()

Review comment (suggestion):
instead of a single kernel, can we try to create a dict where we store multiple kernels ? We can leave this for a follow-up PR but the idea would be to also move the triton kernels in kernels. 

Code after:


if is_torch_available():
    import torch
    import torch.nn as nn
    import triton
    import triton.language as tl
    from torch.nn import functional as F


logger = logging.get_logger(__name__)

# Global for the CUTLASS quantization kernel (lazily loaded)
_quantization_kernel = None


def _get_quantization_kernel():
    """Lazily load the CUTLASS quantization kernel from HuggingFace Hub."""
    global _quantization_kernel
    if _quantization_kernel is None:
        try:
            from .hub_kernels import get_kernel

            _quantization_kernel = get_kernel("RedHatAI/quantization")
        except Exception as e:
            logger.warning_once(f"Failed to load CUTLASS quantization kernel: {e}. Falling back to Triton.")
            _quantization_kernel = False  # Mark as unavailable
    return _quantization_kernel if _quantization_kernel else None


def _supports_cutlass(block_size: list[int] | None, output_dtype: torch.dtype) -> bool:
    """
    Check if CUTLASS blockwise FP8 matmul is supported for the given block size and output dtype.

    CUTLASS blockwise kernels require:
    - SM90+ (Hopper or newer)
    - Block size [128, 128] for weights
    - Block size [1, 128] for activations (handled implicitly)
    - Output dtype bfloat16 or float16
    """

    if not is_torch_available() or not torch.cuda.is_available() or not is_kernels_available():
        return False

    # CUTLASS only supports bfloat16/float16 output
    if output_dtype not in (torch.bfloat16, torch.float16):
        return False

    # Check block size compatibility - CUTLASS only supports [128, 128]
    if block_size is None:
        return False