Hugging Face's logo

Kernels:
kernels-community
/
aiter-kernels
Trusted publisher

Kernel card Files
xet
 Community
aiter-kernels / build /torch-rocm /gemm /basic /gemm_a16wfp4.py
kernels-bot's picture
kernels-bot
Uploaded using `kernel-builder`.
89181fc verified
Raw
History Blame
12.1 kB
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
from typing import Optional
import torch
import triton
from ...utils._triton import arch_info as arch_info
from ...utils.logger import AiterTritonLogger
from ...utils.common_utils import serialize_dict, deserialize_str
from ..._triton_kernels.gemm.basic.gemm_a16wfp4 import (
 _gemm_a16wfp4_kernel,
 _gemm_a16wfp4_preshuffle_kernel,
 _get_config,
)
from ..._triton_kernels.common.splitk_reduce import (
 _gemm_splitk_reduce_kernel,
)
from ...gemm.basic.gemm_afp4wfp4 import (
 get_splitk,
)
from ..._aiter_compat.torch_guard import torch_compile_guard
_LOGGER = AiterTritonLogger()
def gemm_a16wfp4_fake_tensor(
 x: torch.Tensor,
 w: torch.Tensor,
 w_scales: torch.Tensor,
 atomic_add: Optional[bool] = False,
 dtype: Optional[torch.dtype] = torch.bfloat16,
 y: Optional[torch.Tensor] = None,
 config: Optional[str] = None,
) -> torch.Tensor:
 if y is None:
 M, _ = x.shape
 N, _ = w.shape
 return torch.zeros((M, N), dtype=dtype, device=x.device)
 return y
@torch_compile_guard(gen_fake=gemm_a16wfp4_fake_tensor)
def gemm_a16wfp4_(
 x: torch.Tensor,
 w: torch.Tensor,
 w_scales: torch.Tensor,
 atomic_add: Optional[bool] = False,
 dtype: Optional[torch.dtype] = torch.bfloat16,
 y: Optional[torch.Tensor] = None,
 config: Optional[str] = None,
) -> torch.Tensor:
 """
 Computes matrix multiplication Y = X @ W^T with BF16 activations and FP4 weights.
 Key parameters:
 x (torch.Tensor): BF16/FP16 input matrix X with shape (M, K).
 Quantized to MXFP4 on-the-fly during GEMM.
 w (torch.Tensor): FP4 E2M1 weight matrix W with shape (N, K//2).
 w_scales (torch.Tensor): E8M0 per-group scale for w with shape (N, K//32).
 One scale per 32 elements in K dimension.
 atomic_add (Optional[bool]): use atomic_add for reduction
 dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16).
 y (Optional[torch.Tensor]): Pre-allocated output tensor with shape (M, N).
 config (Optional[str]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N,
 BLOCK_SIZE_K, GROUP_SIZE_M, NUM_KSPLIT, SPLITK_BLOCK_SIZE).
 Returns:
 y (torch.Tensor): Output with shape (M, N).
 """
_LOGGER.info(
 f"GEMM_A16WFP4: x={tuple(x.shape)} w={tuple(w.shape)} w_scale={tuple(w_scales.shape)} "
 )
assert arch_info.is_fp4_avail(), "MXFP4 is not available on your device"
M, K = x.shape
 N, K = w.shape
# inner kernel expects (K, N)
 w = w.T
if config is None:
 config, _ = _get_config(M, N, K)
 else:
 config = deserialize_str(config)
if y is None:
 if atomic_add:
 y = torch.zeros((M, N), dtype=dtype, device=x.device)
 else:
 y = torch.empty((M, N), dtype=dtype, device=x.device)
if config["NUM_KSPLIT"] > 1 and not atomic_add:
 SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
 K, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
 )
config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
 config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
 config["NUM_KSPLIT"] = NUM_KSPLIT
if config["BLOCK_SIZE_K"] >= 2 * K:
 config["BLOCK_SIZE_K"] = triton.next_power_of_2(2 * K)
 config["SPLITK_BLOCK_SIZE"] = 2 * K
 config["NUM_KSPLIT"] = 1
 config["BLOCK_SIZE_K"] = max(config["BLOCK_SIZE_K"], 64)
if config["NUM_KSPLIT"] > 1 and not atomic_add:
 y_pp = torch.empty(
 (config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=y.device
 )
 else:
 config["SPLITK_BLOCK_SIZE"] = 2 * K
 y_pp = None
grid = lambda META: ( # noqa: E731
 (
 META["NUM_KSPLIT"]
 * triton.cdiv(M, META["BLOCK_SIZE_M"])
 * triton.cdiv(N, META["BLOCK_SIZE_N"])
 ),
 )
 _gemm_a16wfp4_kernel[grid](
 x,
 w,
 y if y_pp is None else y_pp,
 w_scales,
 M,
 N,
 K,
 x.stride(0),
 x.stride(1),
 w.stride(0),
 w.stride(1),
 0 if y_pp is None else y_pp.stride(0),
 y.stride(0) if y_pp is None else y_pp.stride(1),
 y.stride(1) if y_pp is None else y_pp.stride(2),
 w_scales.stride(0),
 w_scales.stride(1),
 ATOMIC_ADD=atomic_add,
 **config,
 )
if config["NUM_KSPLIT"] > 1 and not atomic_add:
 REDUCE_BLOCK_SIZE_M = 16
 REDUCE_BLOCK_SIZE_N = 64
 # TODO: Need to debug - REDUCE_BLOCK_SIZE_N=128 with fp32 partials fails
 # NOTE: REDUCE_BLOCK_SIZE_N=16 gives best perf with fp32 partials and
 # REDUCE_BLOCK_SIZE_N=128 gives best perf with bf16 partials
 ACTUAL_KSPLIT = triton.cdiv(K, (config["SPLITK_BLOCK_SIZE"] // 2))
grid_reduce = (
 triton.cdiv(M, REDUCE_BLOCK_SIZE_M),
 triton.cdiv(N, REDUCE_BLOCK_SIZE_N),
 )
 _gemm_splitk_reduce_kernel[grid_reduce](
 y_pp,
 y,
 None,
 M,
 N,
 y_pp.stride(0),
 y_pp.stride(1),
 y_pp.stride(2),
 y.stride(0),
 y.stride(1),
 REDUCE_BLOCK_SIZE_M,
 REDUCE_BLOCK_SIZE_N,
 ACTUAL_KSPLIT,
 triton.next_power_of_2(config["NUM_KSPLIT"]),
 ADD_BIAS=False,
 activation="",
 use_activation=False,
 KERNEL_NAME="_gemm_afp4wfp4_reduce_kernel",
 )
return y
def gemm_a16wfp4(
 x: torch.Tensor,
 w: torch.Tensor,
 w_scales: torch.Tensor,
 atomic_add: Optional[bool] = False,
 dtype: Optional[torch.dtype] = torch.bfloat16,
 y: Optional[torch.Tensor] = None,
 config: Optional[dict] = None,
) -> torch.Tensor:
 config_hashable = serialize_dict(config) if config else None
 return gemm_a16wfp4_(x, w, w_scales, atomic_add, dtype, y, config_hashable)
def gemm_a16wfp4_preshuffle_fake_tensor(
 x: torch.Tensor,
 w: torch.Tensor,
 w_scales: torch.Tensor,
 dtype: Optional[torch.dtype] = torch.bfloat16,
 y: Optional[torch.Tensor] = None,
 config: Optional[str] = None,
 skip_reduce: Optional[bool] = False,
) -> torch.Tensor:
 M, K = x.shape
 N, _ = w.shape
config = deserialize_str(config)
num_ksplit = config["NUM_KSPLIT"]
 block_size_k = config["BLOCK_SIZE_K"]
if num_ksplit > 1:
 _, block_size_k, num_ksplit = get_splitk(K, block_size_k, num_ksplit)
if block_size_k >= 2 * K:
 num_ksplit = 1
if num_ksplit > 1 and skip_reduce:
 y_pp = torch.empty((num_ksplit, M, N), dtype=torch.float32, device=x.device)
 return y_pp
return torch.empty((M, N), dtype=dtype, device=x.device)
@torch_compile_guard(gen_fake=gemm_a16wfp4_preshuffle_fake_tensor)
def gemm_a16wfp4_preshuffle_(
 x: torch.Tensor,
 w: torch.Tensor,
 w_scales: torch.Tensor,
 prequant: Optional[bool] = True,
 dtype: Optional[torch.dtype] = torch.bfloat16,
 y: Optional[torch.Tensor] = None,
 config: Optional[str] = None,
 skip_reduce: Optional[bool] = False,
) -> torch.Tensor:
 """
 Computes matrix multiplication Y = X @ W^T with BF16 activations and FP4 weights.
 Key parameters:
 x (torch.Tensor): BF16/FP16 input matrix X with shape (M, K).
 Quantized to MXFP4 on-the-fly during GEMM.
 w (torch.Tensor): FP4 E2M1 weight matrix W with shape (N, K//2).
 w_scales (torch.Tensor): E8M0 per-group scale for w with shape (M//32, K).
 One scale per 32 elements in K dimension.
 dtype (Optional[torch.dtype]): Output datatype (BF16 or FP16).
 y (Optional[torch.Tensor]): Pre-allocated output tensor with shape (M, N).
 config (Optional[str]): Kernel tuning parameters (BLOCK_SIZE_M, BLOCK_SIZE_N,
 BLOCK_SIZE_K, GROUP_SIZE_M, NUM_KSPLIT, SPLITK_BLOCK_SIZE).
 skip_reduce (Optional[bool]): skip reduction, y becomes (SPK, M, N) where SPK is determined by config
 Returns:
 y (torch.Tensor): Output with shape (M, N).
 """
_LOGGER.info(
 f"GEMM_A16WFP4_PRESHUFFLE: x={tuple(x.shape)} w={tuple(w.shape)} w_scale={tuple(w_scales.shape)} "
 )
assert arch_info.is_fp4_avail(), "MXFP4 is not available on your device"
 assert prequant, "prequant == False is not supported yet"
M, K = x.shape
 N, K = w.shape
 N = N * 16
 K = K // 16
if config is None:
 config, _ = _get_config(M, N, K, True)
 else:
 config = deserialize_str(config)
if config["NUM_KSPLIT"] > 1:
 SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
 K, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
 )
config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
 config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
 config["NUM_KSPLIT"] = NUM_KSPLIT
if config["BLOCK_SIZE_K"] >= 2 * K:
 config["BLOCK_SIZE_K"] = triton.next_power_of_2(2 * K)
 config["SPLITK_BLOCK_SIZE"] = 2 * K
 config["NUM_KSPLIT"] = 1
 config["BLOCK_SIZE_N"] = max(config["BLOCK_SIZE_N"], 32)
return_y_pp = config["NUM_KSPLIT"] > 1 and skip_reduce
if config["NUM_KSPLIT"] > 1:
 y_pp = torch.empty(
 (config["NUM_KSPLIT"], M, N), dtype=torch.float32, device=x.device
 )
 else:
 config["SPLITK_BLOCK_SIZE"] = 2 * K
 y_pp = None
if y is None and not return_y_pp:
 y = torch.empty((M, N), dtype=dtype, device=x.device)
grid = lambda META: ( # noqa: E731
 (
 META["NUM_KSPLIT"]
 * triton.cdiv(M, META["BLOCK_SIZE_M"])
 * triton.cdiv(N, META["BLOCK_SIZE_N"])
 ),
 )
 _gemm_a16wfp4_preshuffle_kernel[grid](
 x,
 w,
 y if y_pp is None else y_pp,
 w_scales,
 M,
 N,
 K,
 x.stride(0),
 x.stride(1),
 w.stride(0),
 w.stride(1),
 0 if y_pp is None else y_pp.stride(0),
 y.stride(0) if y_pp is None else y_pp.stride(1),
 y.stride(1) if y_pp is None else y_pp.stride(2),
 w_scales.stride(0),
 w_scales.stride(1),
 PREQUANT=prequant,
 **config,
 )
if return_y_pp:
 return y_pp
 elif config["NUM_KSPLIT"] > 1:
 REDUCE_BLOCK_SIZE_M = 16
 REDUCE_BLOCK_SIZE_N = 64
 # TODO: Need to debug - REDUCE_BLOCK_SIZE_N=128 with fp32 partials fails
 # NOTE: REDUCE_BLOCK_SIZE_N=16 gives best perf with fp32 partials and
 # REDUCE_BLOCK_SIZE_N=128 gives best perf with bf16 partials
 ACTUAL_KSPLIT = triton.cdiv(K, (config["SPLITK_BLOCK_SIZE"] // 2))
grid_reduce = (
 triton.cdiv(M, REDUCE_BLOCK_SIZE_M),
 triton.cdiv(N, REDUCE_BLOCK_SIZE_N),
 )
 _gemm_splitk_reduce_kernel[grid_reduce](
 y_pp,
 y,
 None,
 M,
 N,
 y_pp.stride(0),
 y_pp.stride(1),
 y_pp.stride(2),
 y.stride(0),
 y.stride(1),
 REDUCE_BLOCK_SIZE_M,
 REDUCE_BLOCK_SIZE_N,
 ACTUAL_KSPLIT,
 triton.next_power_of_2(config["NUM_KSPLIT"]),
 ADD_BIAS=False,
 activation="",
 use_activation=False,
 KERNEL_NAME="_gemm_afp4wfp4_reduce_kernel",
 )
return y
def gemm_a16wfp4_preshuffle(
 x: torch.Tensor,
 w: torch.Tensor,
 w_scales: torch.Tensor,
 prequant: Optional[bool] = True,
 dtype: Optional[torch.dtype] = torch.bfloat16,
 y: Optional[torch.Tensor] = None,
 config: Optional[dict] = None,
 skip_reduce: Optional[bool] = False,
) -> torch.Tensor:
 if config is None:
 config_hashable = None
 M, _ = x.shape
 N, K = w.shape
 N = N * 16
 K = K // 16
 config, _ = _get_config(M, N, K, True)
 config_hashable = serialize_dict(config)
 return gemm_a16wfp4_preshuffle_(
 x, w, w_scales, prequant, dtype, y, config_hashable, skip_reduce
 )