# 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
try:
    import lightop
except ImportError:
    pass
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, w8a8_block_int8_matmul_wgrad_batched_native
from ..tensor.quantized_tensor import Quantizer
from ..tensor._internal.float8_blockwise_tensor_base import Float8BlockwiseQTensorBase
from ..tensor._internal.float8_tensor_base import Float8TensorBase
from ...debug.pytorch.debug_quantization import DebugQuantizer
from transformer_engine.pytorch.fp8 import blockwise_fp8_block_len
from transformer_engine.pytorch.utils import use_lightop_w8a8
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
tensorwise_int8_check = bool(int(os.getenv("NVTE_INT8_SIM_FP8_TENSORWISE_CHECK", "0")))    
__all__ = [
    "general_gemm",
    "general_grouped_gemm",
    "batchgemm",
]


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,
    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"

    # 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 int8_simulation_fp8 and (isinstance(A, Float8BlockwiseQTensorBase) or isinstance(B, Float8BlockwiseQTensorBase)):
        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"

        if layout == "TN":
            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 use_lightop_w8a8([blockwise_fp8_block_len, blockwise_fp8_block_len]):
                 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:
                y, _ = w8a8_block_int8_matmul(
                    qx_data, qw_data, ref_scales_x, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len],
                    output_dtype=out_dtype
                )
            return y, None, None, None
        
        elif layout == "NN":
            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
            if use_lightop_w8a8([blockwise_fp8_block_len, blockwise_fp8_block_len]):
                y = lightop.gemm_w8a8_asm(
                    qdout_data, qw_data, ref_scales_dout, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len], out_dtype, 'TN'
                )
            else:
                y, _ = w8a8_block_int8_matmul(
                    qdout_data, qw_data, ref_scales_dout, ref_scales_w, [blockwise_fp8_block_len, blockwise_fp8_block_len],
                    output_dtype=out_dtype
                )
            return y, None, None, None

        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 use_lightop_w8a8([blockwise_fp8_block_len, blockwise_fp8_block_len]):
                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:
                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],
                    output_dtype=out_dtype
                )
            return out, None, None, None
        
        else:
            raise ValueError(f"Unsupported layout {layout} in int8 simulation fp8")

    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.")

    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, Float8BlockwiseQTensorBase) or isinstance(B, Float8BlockwiseQTensorBase):
        # There is not use_split_accumulator == False
        # implementation for Float8BlockwiseQTensorBase 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, Float8TensorBase) or isinstance(B, Float8TensorBase)) 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 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"
        
        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, Float8TensorBase) or isinstance(B, Float8TensorBase)):
        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,
    }

    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], Float8BlockwiseQTensorBase) or isinstance(B[0], Float8BlockwiseQTensorBase)):
        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, [blockwise_fp8_block_len, blockwise_fp8_block_len],
                output_dtype=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, Float8TensorBase) or isinstance(B, Float8TensorBase)) and int8_simulation_fp8_tensorwise:
        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 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
            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,
                TE_DType[out_dtype],
                bias,
                bias_dtype,
                gelu,
                gelu_input[0],
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                accumulate,
                use_split_accumulator,
            )[0]
            return out, bias, gelu_input
        
        if layout == "NN":
            assert 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,
                TE_DType[out_dtype],
                bias,
                bias_dtype,
                gelu,
                gelu_input[0],
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                accumulate,
                use_split_accumulator,
            )[0]
            return out, bias, gelu_input
            
        elif layout == "NT":
            assert out_dtype is torch.bfloat16 or 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,
                TE_DType[out_dtype],
                bias,
                bias_dtype,
                gelu,
                gelu_input[0],
                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], Float8TensorBase) or isinstance(B[0], Float8TensorBase)):
        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 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
            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],
                bias[0],
                bias_dtype,
                gelu,
                gelu_input[0],
                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 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],
                bias[0],
                bias_dtype,
                gelu,
                gelu_input[0],
                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 out_dtype is torch.bfloat16 or 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],
                bias[0],
                bias_dtype,
                gelu,
                gelu_input[0],
                grad,  # grad
                workspaces[0],
                workspaces[0].shape[0],
                False,
                use_split_accumulator,
            )[0]        

            if 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
