build/torch-rocm/gemm/basic/gemm_afp4wfp4_pre_quant_atomic.py

# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

from typing import Optional
import torch
from ...utils.logger import AiterTritonLogger
from ...gemm.basic.gemm_a16wfp4 import (
    gemm_a16wfp4,
)

_LOGGER = AiterTritonLogger()


def gemm_afp4wfp4_pre_quant(
    x: torch.Tensor,
    w: torch.Tensor,
    w_scales: torch.Tensor,
    dtype: Optional[float] = torch.bfloat16,
    y: Optional[torch.Tensor] = None,
    config: Optional[dict] = None,
) -> torch.Tensor:
    _LOGGER.info(
        "gemm_afp4wfp4_pre_quant will be deprecated in future AITER release, please switch to gemm_a16wfp4"
    )
    return gemm_a16wfp4(x, w, w_scales, True, dtype, y, config)