# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""Python interface for GEMM extensions"""

from typing import Iterable, Optional, Tuple, Union, List
import os
import torch
import transformer_engine_torch as tex
import warnings
enable_lightop = os.getenv("NVTE_USE_LIGHTOP", "false").strip().lower() in ["true", "1"]
if enable_lightop:
    try:
        import lightop
        enable_lightop = True
    except ImportError:
        enable_lightop = False
from ..constants import TE_DType, TE_DType_To_Torch
from ..utils import get_sm_count, _empty_tensor
from transformer_engine.pytorch.triton.blockwise_int8_gemm_nt import w8a8_block_int8_matmul, w8a8_block_int8_matmul_batched
from transformer_engine.pytorch.triton.blockwise_int8_gemm_nt_wgrad import w8a8_block_int8_matmul_wgrad, w8a8_block_int8_matmul_wgrad_batched
from ..quantized_tensor import Quantizer
from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage
from ..tensor.utils import is_custom
from ..custom_recipes.gemm import custom_gemm
from ..tensor.storage.float8_tensor_storage import Float8TensorStorage
from ...debug.pytorch.debug_quantization import DebugQuantizer
from transformer_engine.pytorch.fp8 import blockwise_fp8_block_len
from transformer_engine.pytorch.triton.per_token_group_quant import (per_token_quant_fp8_to_int8, 
                                                                     per_token_quant_fp8_to_int8_v2, 
                                                                     per_token_quant_fp8_to_int8_opt, 
                                                                     channelwise_dequantize, 
                                                                     channelwise_dequantize_transA, 
                                                                     channelwise_dequantize_transA_float, 
                                                                     channelwise_dequantize_transB, 
                                                                     channelwise_dequantize_transA_add, 
                                                                     channelwise_dequantize_transA_float_add,
                                                                     tensorwise_dequantize,
                                                                     tensorwise_dequantize_add,
                                                                     tensorwise_dequantize_float,
                                                                     tensorwise_dequantize_float_add)
from transformer_engine.pytorch.utils import get_device_compute_capability

from transformer_engine.pytorch.fp8 import int8_simulation_fp8, int8_simulation_fp8_tensorwise
int8_simulation_fp8_tensorwise_batched = bool(int(os.getenv("NVTE_INT8_SIM_FP8_TENSORWISE_BATCHED", "0")))    
__all__ = [
    "general_gemm",
    "general_grouped_gemm",
    "batchgemm",
]

def w8a8_block_int8_matmul_wgrad_batched_native(A_list, B_list, As_list, Bs_list, C_list, accumulate, out_dtype=torch.float16):
    for i in range(len(C_list)):
        assert C_list[i] is not None
        if enable_lightop and get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128  and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
            C_list[i] = lightop.gemm_w8a8_wgrad_asm(
                A_list[i], B_list[i], As_list[i], Bs_list[i], C_list[i], accumulate, blockwise_fp8_block_len, out_dtype, "TN"
            )
        else:
            warnings.warn("Lightop is not available. Using default implementation for w8a8.")
            C_list[i], _ = w8a8_block_int8_matmul_wgrad(
                A_list[i], B_list[i], As_list[i], Bs_list[i], C_list[i], accumulate, blockwise_fp8_block_len,
                out_dtype,
                None
            )
    return C_list

def w8a8_int8_general_gemm(
    A: torch.Tensor,
    B: torch.Tensor,
    out_dtype: Optional[torch.dtype] = None,
    accumulate: bool = False,
    layout: str = "TN",
    out: Optional[torch.Tensor] = None) -> torch.Tensor:
    if layout == "TN":
        assert accumulate is False, "Accumulate not supported in w8a8_general_gemm with TN layout"
        assert out is None, "Output tensor not supported in w8a8_general_gemm with TN layout"
        qx_data = (B._rowwise_data.view(dtype=torch.int8))
        qw_data = (A._rowwise_data.view(dtype=torch.int8))
        ref_scales_x = B._rowwise_scale_inv
        ref_scales_w = A._rowwise_scale_inv
        if enable_lightop and get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128 and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
            y = lightop.gemm_w8a8_asm(qx_data, qw_data, ref_scales_x, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype, 'TN')
        else:
            warnings.warn("Lightop is not available. Using default implementation for w8a8.")
            y, _ = w8a8_block_int8_matmul(qx_data, qw_data, ref_scales_x, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype)
        return y
    elif layout == "NN":
        assert accumulate is False, "Accumulate not supported in w8a8_general_gemm with NN layout"
        assert out is None, "Output tensor not supported in w8a8_general_gemm with NN layout"
        if enable_lightop and get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128 and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
            qdout_data = (B._rowwise_data.view(dtype=torch.int8))
            qw_data = (A._rowwise_data.view(dtype=torch.int8))
            ref_scales_dout = B._rowwise_scale_inv
            ref_scales_w = A._rowwise_scale_inv
            y = lightop.gemm_w8a8_xgrad_asm(qdout_data, qw_data, ref_scales_dout, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype, 'NN')
        else:
            warnings.warn("Lightop is not available. Using default implementation for w8a8.")
            qdout_data = (B._rowwise_data.view(dtype=torch.int8))
            qw_data = (A._columnwise_data.view(dtype=torch.int8))
            ref_scales_dout = B._rowwise_scale_inv
            ref_scales_w = A._columnwise_scale_inv
            y, _ = w8a8_block_int8_matmul(qdout_data, qw_data, ref_scales_dout, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len],out_dtype)
        return y
    elif layout == "NT":
        qdout_data = (B._columnwise_data.view(dtype=torch.int8))
        qx_data = (A._columnwise_data.view(dtype=torch.int8))
        ref_scales_dout = B._columnwise_scale_inv
        ref_scales_x = A._columnwise_scale_inv
        if enable_lightop and get_device_compute_capability() >= (9, 3) and blockwise_fp8_block_len == 128 and ((out_dtype is torch.bfloat16) or (out_dtype is torch.float16)):
            out = lightop.gemm_w8a8_wgrad_asm(qdout_data, qx_data, ref_scales_dout, ref_scales_x, out, accumulate, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype, 'TN')
        else:
            warnings.warn("Lightop is not available. Using default implementation for w8a8.")
            out, _ = w8a8_block_int8_matmul_wgrad(qdout_data, qx_data, ref_scales_dout, ref_scales_x, out, accumulate, [blockwise_fp8_block_len, blockwise_fp8_block_len],out_dtype)
        return out
    else:
        raise ValueError(f"Unsupported layout {layout} in int8 simulation fp8")

def validate_gemm_scale(scale: Optional[float], required: bool) -> float:
    """Validate whether a GEMM scaling factor is consistent with its usage"""
    if required:
        return scale if scale is not None else 1.0
    if scale not in (0.0, None):
        raise ValueError("scale must be zero")
    return 0.0


def general_gemm(
    A: torch.Tensor,
    B: torch.Tensor,
    workspace: torch.Tensor,
    out_dtype: Optional[torch.dtype] = None,
    quantization_params: Optional[Quantizer] = None,
    gelu: bool = False,
    gelu_in: torch.Tensor = None,
    alpha: float = 1.0,
    beta: Optional[float] = None,
    accumulate: bool = False,
    layout: str = "TN",
    out: Optional[torch.Tensor] = None,
    bias: Optional[torch.Tensor] = None,
    use_split_accumulator: bool = False,
    grad: bool = False,
    ub: Union[tex.CommOverlap, tex.CommOverlapP2P] = None,
    ub_type: tex.CommOverlapType = None,
    extra_output: Optional[torch.Tensor] = None,
    bulk_overlap: bool = False,
) -> Iterable[Optional[torch.Tensor]]:
    """GEMM supporting fp8 inputs."""

    assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported."
    transa = layout[0] == "T"
    transb = layout[1] == "T"
    # assert quantization_params is None, "FP8 output not supported yet"

    alpha = validate_gemm_scale(alpha, True)
    beta = validate_gemm_scale(beta, accumulate)

    # if ub_type is not None:
    #     assert ub is not None, (
    #         f"{'AG+GEMM' if ub_type == tex.CommOverlapType.AG else 'GEMM+RS'} overlap requires"
    #         + "a valid `ub` communicator object."
    #     )

    if ub is not None:
        assert ub_type is not None, "Comm+GEMM overlap requires a valid `comm_type` argument."
        if ub_type == tex.CommOverlapType.RS:
            if not (bulk_overlap and not ub.is_fp8_ubuf()):
                assert extra_output is not None, "GEMM+RS overlap requires extra output tensor."

    if out is not None:
        if not out.is_contiguous():
            raise ValueError("Output tensor is not contiguous.")

    # If A or B are custom tensors -> dispatch to quantizers's qgemm implementation
    if is_custom(A) or is_custom(B):
        return custom_gemm(
            A,
            B,
            workspace,
            out_dtype,
            quantization_params,
            gelu,
            gelu_in,
            accumulate,
            layout,
            out,
            bias,
            use_split_accumulator,
            grad,
        )

    debug_quantizer = None
    if isinstance(quantization_params, DebugQuantizer):
        debug_quantizer = quantization_params
        quantization_params = quantization_params.parent_quantizer
        A = A.get_tensor(not transa)
        B = B.get_tensor(transb)

    # Use bfloat16 as default bias_dtype
    bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype]

    if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance(B, Float8BlockwiseQTensorStorage):
        # There is not use_split_accumulator == False
        # implementation for Float8BlockwiseQTensorStorage GEMM
        use_split_accumulator = True

        # Check that data format is supported
        if (
            A._data_format != tex.Float8BlockScaleTensorFormat.GEMM_READY
            or B._data_format != tex.Float8BlockScaleTensorFormat.GEMM_READY
        ):
            raise RuntimeError("GEMM with Float8BlockwiseQTensor requires GEMM_READY format")
    
    if int8_simulation_fp8 and (isinstance(A, Float8BlockwiseQTensorStorage) or isinstance(B, Float8BlockwiseQTensorStorage)):
        assert not gelu, "GELU not supported with int8 simulation"
        assert gelu_in is None, "GELU input not supported with int8 simulation"
        assert bias is None, "Bias not supported with int8 simulation"
        assert ub is None, "User buffer not supported with int8 simulation"
        assert ub_type is None, "User buffer type not supported with int8 simulation"
        assert extra_output is None, "Extra output not supported with int8 simulation"
        assert not bulk_overlap, "Bulk overlap not supported with int8 simulation"

        y = w8a8_int8_general_gemm(
            A,
            B,
            out_dtype,
            accumulate,
            layout,
            out
        )
        return y, None, None, None

    if int8_simulation_fp8 and (isinstance(A, Float8TensorStorage) or isinstance(B, Float8TensorStorage)) and int8_simulation_fp8_tensorwise:
        assert not gelu, "GELU not supported with int8 simulation"
        assert gelu_in is None, "GELU input not supported with int8 simulation"
        assert ub is None, "User buffer not supported with int8 simulation"
        assert ub_type is None, "User buffer type not supported with int8 simulation"
        assert extra_output is None, "Extra output not supported with int8 simulation"
        assert not bulk_overlap, "Bulk overlap not supported with int8 simulation"
        assert out_dtype is torch.bfloat16 or out_dtype is torch.float32, "Out_dtype must be bfloat16 or float32 for int8 simulation"
        
        args = (
            A,
            transa,  # transa
            B,
            transb,  # transb
            out,
            quantization_params,
            TE_DType[out_dtype] if out_dtype is not None else None,
            bias,
            bias_dtype,
            gelu,
            gelu_in,
            grad,  # grad
            workspace,
            workspace.shape[0],
            accumulate,
            True,
        )
        kwargs = {
            "comm_overlap": ub,
            "comm_type": ub_type,
            "extra_output": extra_output,
            "bulk_overlap": bulk_overlap,
        }
    
        out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs)
    
        if debug_quantizer is not None:
            out = debug_quantizer.process_gemm_output(out)
    
        return out, bias_grad, gelu_input, extra_output

    if int8_simulation_fp8 and (isinstance(A, Float8TensorStorage) or isinstance(B, Float8TensorStorage)):
        assert not gelu, "GELU not supported with int8 simulation"
        assert gelu_in is None, "GELU input not supported with int8 simulation"
        assert bias is None, "Bias not supported with int8 simulation"
        assert ub is None, "User buffer not supported with int8 simulation"
        assert ub_type is None, "User buffer type not supported with int8 simulation"
        assert extra_output is None, "Extra output not supported with int8 simulation"
        assert not bulk_overlap, "Bulk overlap not supported with int8 simulation"
        assert out_dtype is torch.bfloat16 or out_dtype is torch.float32, "Out_dtype must be bfloat16 or float32 for int8 simulation"

        if layout == "TN":
            assert out_dtype is torch.bfloat16
            out_shape = B._data.shape[:-1] + (A._data.shape[0], )
            x_int8, x_scales = per_token_quant_fp8_to_int8(B._data.view(dtype=TE_DType_To_Torch[B._fp8_dtype]), B._scale_inv, False)
            w_int8, w_scales = per_token_quant_fp8_to_int8(A._data.view(dtype=TE_DType_To_Torch[A._fp8_dtype]), A._scale_inv, False)
        
            y_int32 = tex.generic_gemm(
                w_int8,
                transa,
                x_int8,
                transb,
                None,
                None,
                TE_DType[torch.int32],
                bias,
                bias_dtype,
                gelu,
                gelu_in,
                grad,  # grad
                workspace,
                workspace.shape[0],
                False,
                use_split_accumulator,
            )[0]
            y = channelwise_dequantize_transB(x_scales, w_scales, y_int32)
            return y.view(out_shape), None, None, None
        
        elif layout == "NN":
            assert out_dtype is torch.bfloat16
            dx_shape = B._data.shape[:-1] + (A._data.shape[-1], )
            dy_int8, dy_scales = per_token_quant_fp8_to_int8(B._data.view(dtype=TE_DType_To_Torch[B._fp8_dtype]), B._scale_inv, False)
            w_int8, w_scales = per_token_quant_fp8_to_int8_opt(A._data.view(dtype=TE_DType_To_Torch[A._fp8_dtype]), A._scale_inv, False)
        
            dx_int32 = tex.generic_gemm(
                w_int8,
                transa,
                dy_int8,
                transb,
                None,
                None,
                TE_DType[torch.int32],
                bias,
                bias_dtype,
                gelu,
                gelu_in,
                grad,  # grad
                workspace,
                workspace.shape[0],
                False,
                use_split_accumulator,
            )[0]
            dx = channelwise_dequantize(dy_scales, w_scales, dx_int32)
            return dx.view(dx_shape), None, None, None

        elif layout == "NT":
            assert out_dtype is torch.bfloat16 or out_dtype is torch.float32
            dy_int8, dy_scales = per_token_quant_fp8_to_int8_opt(B._data.view(dtype=TE_DType_To_Torch[B._fp8_dtype]), B._scale_inv, False)
            x_int8, x_scales = per_token_quant_fp8_to_int8_opt(A._data.view(dtype=TE_DType_To_Torch[A._fp8_dtype]), A._scale_inv, False)
        
            dw_int32 = tex.generic_gemm(
                x_int8,
                transa,
                dy_int8,
                transb,
                None,
                None,
                TE_DType[torch.int32],
                bias,
                bias_dtype,
                gelu,
                gelu_in,
                grad,  # grad
                workspace,
                workspace.shape[0],
                False,  
                use_split_accumulator,
            )[0]
            if out_dtype is torch.bfloat16:
                if accumulate:
                    channelwise_dequantize_transA_add(dy_scales, x_scales, dw_int32, out)
                else:
                    out = channelwise_dequantize_transA(dy_scales, x_scales, dw_int32)
                    
            else:
                if accumulate:
                    channelwise_dequantize_transA_float_add(dy_scales, x_scales, dw_int32, out)     
                else:
                    out = channelwise_dequantize_transA_float(dy_scales, x_scales, dw_int32)     
            return out, None, None, None
        
        else:
            raise ValueError(f"Unsupported layout {layout} in int8 simulation fp8")

    args = (
        A,
        transa,  # transa
        B,
        transb,  # transb
        out,
        quantization_params,
        TE_DType[out_dtype] if out_dtype is not None else None,
        bias,
        bias_dtype,
        gelu,
        gelu_in,
        grad,  # grad
        workspace,
        workspace.shape[0],
        accumulate,
        use_split_accumulator,
    )
    kwargs = {
        "comm_overlap": ub,
        "comm_type": ub_type,
        "extra_output": extra_output,
        "bulk_overlap": bulk_overlap,
        "alpha": alpha,
        "beta": beta,
    }

    out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs)

    if debug_quantizer is not None:
        out = debug_quantizer.process_gemm_output(out)

    return out, bias_grad, gelu_input, extra_output


def general_grouped_gemm(
    A: List[torch.Tensor],
    B: List[torch.Tensor],
    out: List[torch.Tensor],
    out_dtype: torch.dtype,
    workspaces: List[torch.Tensor],
    layout: str = "TN",
    m_splits: Optional[List[int]] = None,
    gelu: bool = False,
    grad=False,
    accumulate: bool = False,
    bias: Optional[List[torch.Tensor]] = None,
    use_bias: bool = False,
    use_split_accumulator: bool = False,
    D_dtype: Optional[tex.DType] = None,
    single_output=False,
) -> Tuple[List[torch.Tensor], ...]:
    """
    TN layout Grouped GEMM with fp8 inputs.
    """
    num_gemms = len(A)

    transa = layout[0] == "T"
    transb = layout[1] == "T"

    empty_tensor = _empty_tensor()
    empty_tensors = [empty_tensor] * num_gemms

    # Use bfloat16 as default bias_dtype
    gelu_input = empty_tensors
    out_dtype = TE_DType[out[0].dtype] if D_dtype is None else D_dtype

    sm_count = get_sm_count()
    if grad and use_bias:
        grad_bias = [
            torch.empty(B[i].shape[1], dtype=out[0].dtype, device="cuda") for i in range(num_gemms)
        ]
    else:
        grad_bias = empty_tensors
    bias = bias if use_bias else empty_tensors
    if use_bias:
        bias_dtype = TE_DType[grad_bias[0].dtype] if grad else TE_DType[bias[0].dtype]
    else:
        bias_dtype = TE_DType[torch.bfloat16]

    if gelu:
        gelu_input = [
            torch.empty_like(o, dtype=bias_dtype, memory_format=torch.contiguous_format)
            for o in out
        ]  # this should differ with respect to single output

    if int8_simulation_fp8 and (isinstance(A[0], Float8BlockwiseQTensorStorage) or isinstance(B[0], Float8BlockwiseQTensorStorage)):
        assert len(set(m_splits)) == 1, "Int8 simulation groupgemm just surpport token pad as same as batchgemm for now."
        assert not gelu, "GELU not supported with int8 simulation groupgemm."
        assert not use_bias, "Bias not supported with int8 simulation groupgemm."

        if layout == "TN":
            qx_data = [
                b._rowwise_data.view(dtype=torch.int8) for b in B
            ]
            qw_data = [
                a._rowwise_data.view(dtype=torch.int8) for a in A
            ]
            ref_scales_x = [b._rowwise_scale_inv for b in B]
            ref_scales_w = [a._rowwise_scale_inv for a in A]
        
            num_gemms = len(A)
            seq_len = sum(m_splits) // num_gemms

            out[0] = w8a8_block_int8_matmul_batched(
                qx_data, qw_data, ref_scales_x, ref_scales_w, out[0].view(num_gemms, seq_len, out[0].size(-1)), [blockwise_fp8_block_len, blockwise_fp8_block_len],
                output_dtype=out_dtype
            )
            return out, bias, gelu_input
        
        elif layout == "NN":
            qdout_data = [
                b._rowwise_data.view(dtype=torch.int8) for b in B
            ]
            qw_data = [
                a._columnwise_data.view(dtype=torch.int8) for a in A
            ]
            ref_scales_dout = [b._rowwise_scale_inv for b in B]
            ref_scales_w = [a._columnwise_scale_inv for a in A]

            num_gemms = len(A)
            seq_len = sum(m_splits) // num_gemms
        
            out[0] = w8a8_block_int8_matmul_batched(
                qdout_data, qw_data, ref_scales_dout, ref_scales_w, out[0].view(num_gemms, seq_len, out[0].size(-1)), [blockwise_fp8_block_len, blockwise_fp8_block_len],
                output_dtype=out_dtype
            )
            return out, bias, gelu_input

        elif layout == "NT":
            qdout_data = [
                b._columnwise_data.view(dtype=torch.int8) for b in B
            ]
            qx_data = [
                a._columnwise_data.view(dtype=torch.int8) for a in A
            ]
            ref_scales_dout = [b._columnwise_scale_inv for b in B]
            ref_scales_x = [a._columnwise_scale_inv for a in A]
        
            out = w8a8_block_int8_matmul_wgrad_batched_native(
                qdout_data, qx_data, ref_scales_dout, ref_scales_x, out, accumulate, out_dtype
            )
            return out, bias, gelu_input
        
        else:
            raise ValueError(f"Unsupported layout {layout} in int8 simulation fp8")
    
    
    if int8_simulation_fp8 and (isinstance(A[0], Float8TensorStorage) or isinstance(B[0], Float8TensorStorage)) and int8_simulation_fp8_tensorwise and int8_simulation_fp8_tensorwise_batched:
        assert len(set(m_splits)) == 1, "Need token pad as same as batchgemm for NVTE_INT8_SIM_FP8_TENSORWISE_BATCHED."
        assert not gelu, "GELU not supported with int8 simulation groupgemm."
        assert not use_bias, "Bias not supported with int8 simulation groupgemm."
        assert TE_DType_To_Torch[out_dtype] is torch.bfloat16 or TE_DType_To_Torch[out_dtype] is torch.float32, "Out_dtype must be bfloat16 or float32 for int8 simulation"
        if layout == "TN":
            assert TE_DType_To_Torch[out_dtype] is torch.bfloat16
            qx_data_list, scales_x_list = [b._data.view(dtype=TE_DType_To_Torch[b._fp8_dtype]) for b in B], [b._scale_inv for b in B]
            w_data_list, scales_w_list = [a._data.view(dtype=TE_DType_To_Torch[a._fp8_dtype]) for a in A], [a._scale_inv for a in A]
            
            num_gemms = len(A)

            qx_data = torch.stack(qx_data_list).contiguous()
            w_data = torch.stack(w_data_list).contiguous()
            scales_x = torch.stack(scales_x_list).contiguous()
            scales_w = torch.stack(scales_w_list).contiguous()

            out[0] = tex.tensorwise_int8_batchgemm(
                w_data.view(-1, w_data.size(-1)),
                transa,
                qx_data.view(-1, qx_data.size(-1)),
                transb,
                scales_w,
                scales_x,
                out[0],
                num_gemms,
                None,
                out_dtype,
                None,
                bias_dtype,
                gelu,
                None,
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                accumulate,
                use_split_accumulator,
            )[0]
            return out, bias, gelu_input
        
        if layout == "NN":
            assert TE_DType_To_Torch[out_dtype] is torch.bfloat16
            qdout_data_list, scales_dout_list = [b._data.view(dtype=TE_DType_To_Torch[b._fp8_dtype]) for b in B], [b._scale_inv for b in B]
            w_data_list, scales_w_list = [a._data.view(dtype=TE_DType_To_Torch[a._fp8_dtype]) for a in A], [a._scale_inv for a in A]
            
            num_gemms = len(A)

            qdout_data = torch.stack(qdout_data_list).contiguous()
            w_data = torch.stack(w_data_list).contiguous()
            scales_dout = torch.stack(scales_dout_list).contiguous()
            scales_w = torch.stack(scales_w_list).contiguous()

            out[0] = tex.tensorwise_int8_batchgemm(
                w_data.view(-1, w_data.size(-1)),
                transa,
                qdout_data.view(-1, qdout_data.size(-1)),
                transb,
                scales_w,
                scales_dout,
                out[0],
                num_gemms,
                None,
                out_dtype,
                None,
                bias_dtype,
                gelu,
                None,
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                accumulate,
                use_split_accumulator,
            )[0]
            return out, bias, gelu_input
            
        elif layout == "NT":
            assert TE_DType_To_Torch[out_dtype] is torch.bfloat16 or TE_DType_To_Torch[out_dtype] is torch.float32
            qdout_data_list, scales_dout_list = [b._data.view(dtype=TE_DType_To_Torch[b._fp8_dtype]) for b in B], [b._scale_inv for b in B]
            qx_data_list, scales_x_list = [a._data.view(dtype=TE_DType_To_Torch[a._fp8_dtype]) for a in A], [a._scale_inv for a in A]
            
            num_gemms = len(A)

            qdout_data = torch.stack(qdout_data_list).contiguous()
            qx_data = torch.stack(qx_data_list).contiguous()
            scales_dout = torch.stack(scales_dout_list).contiguous()
            scales_x = torch.stack(scales_x_list).contiguous()
            dw = torch.stack(out).contiguous()

            dw = tex.tensorwise_int8_batchgemm(
                qx_data.view(-1, qx_data.size(-1)),
                transa,
                qdout_data.view(-1, qdout_data.size(-1)),
                transb,
                scales_x,
                scales_dout,
                dw.view(-1, dw.size(-1)),
                num_gemms,
                None,
                out_dtype,
                None,
                bias_dtype,
                gelu,
                None,
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                accumulate,
                use_split_accumulator,
            )[0]
            for i in range(num_gemms):
                out[i].copy_(dw[i])
            return out, bias, gelu_input


    if int8_simulation_fp8 and (isinstance(A[0], Float8TensorStorage) or isinstance(B[0], Float8TensorStorage)) and int8_simulation_fp8_tensorwise:
        assert not gelu, "GELU not supported with int8 simulation groupgemm."
        assert TE_DType_To_Torch[out_dtype] is torch.bfloat16 or TE_DType_To_Torch[out_dtype] is torch.float32, "Out_dtype must be bfloat16 or float32 for int8 simulation"
        bias = tex.te_general_grouped_gemm(
            A,
            transa,
            B,
            transb,
            out,
            out_dtype,
            m_splits,
            grad_bias if grad else bias,
            bias_dtype,
            single_output,
            gelu_input,  # this is pre_gelu_out
            grad,  # grad
            workspaces,
            workspaces[0].shape[0],
            accumulate,
            use_split_accumulator,
            sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))),
        )
    
        return out, bias, gelu_input
            

    if int8_simulation_fp8 and (isinstance(A[0], Float8TensorStorage) or isinstance(B[0], Float8TensorStorage)):
        assert len(set(m_splits)) == 1, "Int8 simulation groupgemm just surpport token pad as same as batchgemm for now."
        assert not gelu, "GELU not supported with int8 simulation groupgemm."
        assert not use_bias, "Bias not supported with int8 simulation groupgemm."
        assert TE_DType_To_Torch[out_dtype] is torch.bfloat16 or TE_DType_To_Torch[out_dtype] is torch.float32, "Out_dtype must be bfloat16 or float32 for int8 simulation"

        if layout == "TN":
            assert TE_DType_To_Torch[out_dtype] is torch.bfloat16
            qx_data_list = []
            w_data_list = []
            scales_x_list = []
            scales_w_list = []
            for b in B:
                x_int8, x_scales = per_token_quant_fp8_to_int8(b._data.view(dtype=TE_DType_To_Torch[b._fp8_dtype]), b._scale_inv, False)
                qx_data_list.append(x_int8)
                scales_x_list.append(x_scales)
            for a in A:
                w_int8, w_scales = per_token_quant_fp8_to_int8(a._data.view(dtype=TE_DType_To_Torch[a._fp8_dtype]), a._scale_inv, False)
                w_data_list.append(w_int8)
                scales_w_list.append(w_scales)
            
            num_gemms = len(A)
            seq_len = sum(m_splits) // num_gemms

            qx_data = torch.stack(qx_data_list).contiguous()
            w_data = torch.stack(w_data_list).contiguous()
            y_int32 = torch.empty((num_gemms, seq_len, out[0].size(-1)), dtype=torch.int32, device="cuda")

            y_int32 = tex.generic_batchgemm(
                w_data.view(-1, w_data.size(-1)),
                transa,
                qx_data.view(-1, qx_data.size(-1)),
                transb,
                y_int32.view(-1, y_int32.size(-1)),
                num_gemms,
                None,
                TE_DType[torch.int32],
                None,
                bias_dtype,
                gelu,
                None,
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                False,
                use_split_accumulator,
            )[0]
            
            out[0] = out[0].view(num_gemms, seq_len, out[0].size(-1))
            for i in range(num_gemms):
                out[0][i] = channelwise_dequantize_transB(scales_x_list[i], scales_w_list[i], y_int32[i])
            out[0] = out[0].view(-1, out[0].size(-1))
            return out, bias, gelu_input
        
        elif layout == "NN":
            assert TE_DType_To_Torch[out_dtype] is torch.bfloat16
            qdout_data_list = []
            w_data_list = []
            scales_dout_list = []
            scales_w_list = []
            for b in B:
                dy_int8, dy_scales = per_token_quant_fp8_to_int8(b._data.view(dtype=TE_DType_To_Torch[b._fp8_dtype]), b._scale_inv, False)
                qdout_data_list.append(dy_int8)
                scales_dout_list.append(dy_scales)
            for a in A:
                w_int8, w_scales = per_token_quant_fp8_to_int8_opt(a._data.view(dtype=TE_DType_To_Torch[a._fp8_dtype]), a._scale_inv, False)
                w_data_list.append(w_int8)
                scales_w_list.append(w_scales)

            num_gemms = len(A)
            seq_len = sum(m_splits) // num_gemms

            qdout_data = torch.stack(qdout_data_list).contiguous()
            w_data = torch.stack(w_data_list).contiguous()
            dx_int32 = torch.empty((num_gemms, seq_len, out[0].size(-1)), dtype=torch.int32, device="cuda")
        
            dx_int32 = tex.generic_batchgemm(
                w_data.view(-1, w_data.size(-1)),
                transa,
                qdout_data.view(-1, qdout_data.size(-1)),
                transb,
                dx_int32.view(-1, dx_int32.size(-1)),
                num_gemms,
                None,
                TE_DType[torch.int32],
                None,
                bias_dtype,
                gelu,
                None,
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                False,
                use_split_accumulator,
            )[0]

            out[0] = out[0].view(num_gemms, seq_len, out[0].size(-1))
            for i in range(num_gemms):
                out[0][i] = channelwise_dequantize(scales_dout_list[i], scales_w_list[i], dx_int32[i])
            out[0] = out[0].view(-1, out[0].size(-1))
            return out, bias, gelu_input

        elif layout == "NT":
            assert TE_DType_To_Torch[out_dtype] is torch.bfloat16 or TE_DType_To_Torch[out_dtype] is torch.float32
            qdout_data_list = []
            qx_data_list = []
            scales_dout_list = []
            scales_x_list = []
            for b in B:
                dy_int8, dy_scales = per_token_quant_fp8_to_int8_opt(b._data.view(dtype=TE_DType_To_Torch[b._fp8_dtype]), b._scale_inv, False)
                qdout_data_list.append(dy_int8)
                scales_dout_list.append(dy_scales)
            for a in A:
                x_int8, x_scales = per_token_quant_fp8_to_int8_opt(a._data.view(dtype=TE_DType_To_Torch[a._fp8_dtype]), a._scale_inv, False)
                qx_data_list.append(x_int8)
                scales_x_list.append(x_scales)

            num_gemms = len(A)

            qdout_data = torch.stack(qdout_data_list).contiguous()
            qx_data = torch.stack(qx_data_list).contiguous()
            dw_int32 = torch.empty((num_gemms, qdout_data.size(-1), qx_data.size(-1)), dtype=torch.int32, device="cuda")

            dw_int32 = tex.generic_batchgemm(
                qx_data.view(-1, qx_data.size(-1)),
                transa,
                qdout_data.view(-1, qdout_data.size(-1)),
                transb,
                dw_int32.view(-1, dw_int32.size(-1)),
                num_gemms,
                None,
                TE_DType[torch.int32],
                None,
                bias_dtype,
                gelu,
                None,
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                False,
                use_split_accumulator,
            )[0]        

            if TE_DType_To_Torch[out_dtype] is torch.bfloat16:
                if accumulate:
                    for i in num_gemms:
                        channelwise_dequantize_transA_add(scales_dout_list[i], scales_x_list[i], dw_int32[i], out[i])
                else:
                    for i in num_gemms:
                        out[i] = channelwise_dequantize_transA(scales_dout_list[i], scales_x_list[i], dw_int32[i])
            else:
                if accumulate:
                    for i in num_gemms:
                        channelwise_dequantize_transA_float_add(scales_dout_list[i], scales_x_list[i], dw_int32[i], out[i])     
                else:
                    for i in num_gemms:
                        out[i] = channelwise_dequantize_transA_float(scales_dout_list[i], scales_x_list[i], dw_int32[i])    
            return out, bias, gelu_input
        
        else:
            raise ValueError(f"Unsupported layout {layout} in int8 simulation fp8")

    bias = tex.te_general_grouped_gemm(
        A,
        transa,
        B,
        transb,
        out,
        out_dtype,
        m_splits,
        grad_bias if grad else bias,
        bias_dtype,
        single_output,
        gelu_input,  # this is pre_gelu_out
        grad,  # grad
        workspaces,
        workspaces[0].shape[0],
        accumulate,
        use_split_accumulator,
        sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))),
    )

    return out, bias, gelu_input

def batchgemm(
    A: List[torch.Tensor],
    B: List[torch.Tensor],
    out: List[torch.Tensor],
    dtype: torch.dtype,
    workspaces: List[torch.Tensor],
    gelu: bool = False,
    gelu_input: Optional[List[torch.Tensor]] = None,
    grad: bool = False,
    accumulate: bool = False,
    layout: str = "TN",
    bias: Optional[List[torch.Tensor]] = None,
    use_bias: bool = False,
) -> Tuple[Union[torch.Tensor, None], ...]:
    """Non FP8  batch GEMM."""

    assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported."
    transa = layout[0] == "T"
    transb = layout[1] == "T"
    num_gemms = len(A)
    empty_tensor = torch.Tensor()
    empty_tensors = [torch.Tensor()] * num_gemms

    if int8_simulation_fp8:
        assert 0, "If you want to use batchgemm in int8 simulation, please unset GROUPED_GEMM_BatchLinear and use moe groupgemm with pad token."
    
    if gelu and not grad:
        gelu_input = [
            torch.empty_like(o, dtype=dtype, memory_format=torch.contiguous_format) for o in out
        ]
    elif not gelu:
        gelu_input = empty_tensors

    if grad and use_bias:
        grad_bias = [
            torch.empty(B[i].shape[1], dtype=out[0].dtype, device="cuda") for i in range(num_gemms)
        ]
    else:
        grad_bias = empty_tensors

    bias = bias if use_bias else empty_tensors

    assert (
        A[0].dtype == dtype and B[0].dtype == dtype
    ), f"Expected dtype={dtype}, but found A.dtype={A[0].dtype} and B.dtype={B[0].dtype}"
    input_dtype = TE_DType[dtype]
    output_dtype = TE_DType[out[0].dtype]
    if use_bias:
        bias_dtype = TE_DType[grad_bias[0].dtype] if grad else TE_DType[bias[0].dtype]
    else:
        bias_dtype = output_dtype
    tex.te_batchgemm_ts(
        A,
        empty_tensor,
        0,  # A_offset
        input_dtype,
        transa,
        B,
        empty_tensor,
        0,  # B_offset
        input_dtype,
        transb,
        out,
        0,  # out_offset
        empty_tensor,  # out_scale
        output_dtype,
        empty_tensor,  # out_amax
        grad_bias if grad else bias,
        bias_dtype,
        gelu_input,  # gelu_input
        grad,
        workspaces,
        workspaces[0].shape[0],
        accumulate,
        False,  # use_split_accumulator
    )

    return out, grad_bias, gelu_input
