# SPDX-License-Identifier: MIT

import os
import json
import logging
import functools
from functools import partial
from typing import Any, Dict, List, Optional, Tuple
from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
import aiter.ops.triton.utils.arch_info as arch_info
from aiter import logger
from aiter.jit.utils.chip_info import get_cu_num
from aiter.ops.triton.utils.common_utils import save_kernel_path, has_kernel_cache, get_triton_cache_dir

import torch
import triton
import triton.language as tl

import numpy as np
import torch.distributed as dist

from triton.language.extra.hip import libdevice
from triton.language.extra import libshmem_device

from aiter.dist.parallel_state import GroupCoordinator

@functools.lru_cache
def get_w8a8_block_int8_configs(N: int, K: int, block_n: int,
                                block_k: int) -> Optional[dict[int, Any]]:
    """
    Return optimized configurations for the w8a8 block fp8 kernel.

    The return value will be a dictionary that maps an irregular grid of
    batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
    kernel on a given batch size bs, the closest batch size in the grid should
    be picked and the associated configuration chosen to invoke the kernel.
    """

    # First look up if an optimized configuration is available in the configs
    # directory
    # device_name = current_platform.get_device_name().replace(" ", "_")
    device_name = arch_info.get_device()
    device_name = "BW200" if device_name.lower().startswith("bw") else device_name

    # new config by arch and cu number
    arch = triton.runtime.driver.active.get_current_target().arch
    num_cu = get_cu_num()
    json_file_name = f"N={N},K={K},arch={arch},cu={num_cu},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json" # noqa: E501
    config_file_path = os.path.join(
        f"{AITER_TRITON_CONFIGS_PATH}", "gemm/block_w8a8", json_file_name
    )

    # Fallback to device config (to be removed)
    if not os.path.exists(config_file_path):
        json_file_name = f"N={N},K={K},device_name={device_name},dtype=int8_w8a8,block_shape=[{block_n}, {block_k}].json"  # noqa: E501

    config_file_path = os.path.join(
        f"{AITER_TRITON_CONFIGS_PATH}", "gemm/block_w8a8", json_file_name
    )

    if os.path.exists(config_file_path):
        with open(config_file_path) as f:
            #logger.info(
            #    "Using configuration from %s for W8A8 Block INT8 kernel.",
            #    config_file_path,
            #)
            # If a configuration has been found, return it
            return {int(key): val for key, val in json.load(f).items()}

    # If no optimized configuration is available, we will use the default
    # configuration
    logger.warning(
        ("Using default W8A8 Block INT8 kernel config. Performance might "
         "be sub-optimal! Config file not found at %s"),
        config_file_path,
    )
    return None

# for intra-node case
@triton.heuristics(
    values={
        "DIVISIBLE_M": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0,
        "DIVISIBLE_N": lambda args: args["N"] % args["BLOCK_SIZE_N"] == 0,
    }
)
@triton.jit
def _w8a8_block_matmul_allreduce_oneshot_kernel(
    # Pointers to inputs and output
    A,
    B,
    C,
    As,
    Bs,
    # Shape for matmul
    M,
    N,
    K,
    # Block size for block-wise quantization
    group_n,
    group_k,
    # Stride for inputs and output
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_As_m,
    stride_As_k,
    stride_Bs_k,
    stride_Bs_n,
    # shmem args
    shm_ctx,
    barrier_bufs_ptr,
    scatter_bufs_ptr,
    barrier_buf_ptr,
    scatter_buf_ptr,
    cur_rank: tl.constexpr,
    local_world_size: tl.constexpr,
    world_size: tl.constexpr,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    DIVISIBLE_M: tl.constexpr,
    DIVISIBLE_N: tl.constexpr,
    COMBINE_SCALE_LOAD: tl.constexpr = False
):
    """Triton-accelerated function used to perform linear operations (dot
    product) on input tensors `A` and `B` with block-wise quantization, and
    store the result in output tensor `C`.
    """

    # set shmem device ctx (rocSHMEM required when call libshmem.func)
    # libshmem_device.set_rocshmem_ctx(shm_ctx)

    node_id = cur_rank // local_world_size

    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    tl.assume(pid > 0)
    tl.assume(pid_m > 0)
    tl.assume(pid_n > 0)
    tl.assume(group_n > 0)
    tl.assume(group_k > 0)
    tl.assume(stride_am > 0)
    tl.assume(stride_ak > 0)
    tl.assume(stride_bk > 0)
    tl.assume(stride_bn > 0)
    tl.assume(stride_cm > 0)
    tl.assume(stride_cn > 0)
    tl.assume(stride_As_m > 0)
    tl.assume(stride_As_k > 0)
    tl.assume(stride_Bs_k > 0)
    tl.assume(stride_Bs_n > 0)

    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    if COMBINE_SCALE_LOAD:
        As_ptrs = As + offs_am[:, None] * stride_As_m
        offs_bsn = offs_bn // group_n
        Bs_ptrs = Bs + offs_bsn[:, None] * stride_Bs_n
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K), 2):
            a0 = tl.load(a_ptrs,
                        mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
                        other=0.0)
            b0 = tl.load(b_ptrs,
                        mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                        other=0.0)

            k_start = k * BLOCK_SIZE_K
            offs_ks = k_start // group_k + tl.arange(0, 2)
            a_s = tl.load(As_ptrs + offs_ks[None, :] * stride_As_k,
                          mask=offs_ks[None, :] <= (K - 1) // group_k,
                          other=0.0)
            b_s = tl.load(Bs_ptrs + offs_ks[None, :] * stride_Bs_k,
                          mask=offs_ks[None, :] <= (K - 1) // group_k,
                          other=0.0)
            a_s0, a_s1 = tl.split(a_s)
            b_s0, b_s1 = tl.split(b_s)

            accumulator += tl.dot(a0, b0).to(tl.float32) * a_s0[:, None] * b_s0[None, :]

            a0 = tl.load(a_ptrs + BLOCK_SIZE_K * stride_ak,
                        mask=offs_k[None, :] < K - (k + 1)* BLOCK_SIZE_K,
                        other=0.0)
            b0 = tl.load(b_ptrs + BLOCK_SIZE_K * stride_bk,
                        mask=offs_k[:, None] < K - (k + 1) * BLOCK_SIZE_K,
                        other=0.0)

            accumulator += tl.dot(a0, b0).to(tl.float32) * a_s1[:, None] * b_s1[None, :]

            a_ptrs += BLOCK_SIZE_K * stride_ak * 2
            b_ptrs += BLOCK_SIZE_K * stride_bk * 2
    else:
        As_ptrs = As + offs_am * stride_As_m
        offs_bsn = offs_bn // group_n
        Bs_ptrs = Bs + offs_bsn * stride_Bs_n
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
            a = tl.load(a_ptrs,
                        mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
                        other=0.0)
            b = tl.load(b_ptrs,
                        mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                        other=0.0)

            k_start = k * BLOCK_SIZE_K
            offs_ks = k_start // group_k
            a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
            b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

            accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :]

            a_ptrs += BLOCK_SIZE_K * stride_ak
            b_ptrs += BLOCK_SIZE_K * stride_bk


    if C.dtype.element_ty == tl.bfloat16:
        c = accumulator.to(tl.bfloat16)
    elif C.dtype.element_ty == tl.float16:
        c = accumulator.to(tl.float16)
    else:
        c = accumulator.to(tl.float32)

    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    offs_c = stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
    # c_ptrs = C + offs_c
    # scatter_buf_ptr = tl.load(scatter_bufs_ptr + cur_rank).to(tl.pointer_type(tl.float32))
    c_ptrs = scatter_buf_ptr + offs_c

    STORE_MASK_FREE: tl.constexpr = DIVISIBLE_M & DIVISIBLE_N
    if STORE_MASK_FREE:
        tl.store(c_ptrs, c)
        c_mask = None
    else:
        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)

    # Signal to other ranks that we have completed this tile
    for i in range(1, local_world_size):
        remote = (cur_rank + i) % local_world_size + node_id * local_world_size
        remote_base_ptr = tl.load(barrier_bufs_ptr + remote).to(tl.pointer_type(tl.int32))
        # remote_base_ptr = libshmem_device.remote_ptr(barrier_bufs_ptr, remote).to(tl.pointer_type(tl.int32))
        tl.atomic_add(remote_base_ptr + pid, 1, scope="sys", sem="release")

    # consumer
    # local_base_ptr = tl.load(barrier_bufs_ptr + cur_rank).to(tl.pointer_type(tl.int32))
    while tl.atomic_cas(barrier_buf_ptr + pid, world_size - 1, 0, scope="sys", sem="acquire") != (world_size - 1):
        pass

    # acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    acc = accumulator

    # intra-node reduce
    for i in range(1, local_world_size):
        rank_id = (cur_rank + i) % local_world_size + node_id * local_world_size
        scatter_buf_ptr = tl.load(scatter_bufs_ptr + rank_id).to(tl.pointer_type(tl.float32))
        # scatter_buf_ptr = libshmem_device.remote_ptr(scatter_bufs_ptr, rank_id).to(tl.pointer_type(tl.float32))
        acc += tl.load(scatter_buf_ptr + offs_c, mask=c_mask)
    acc_f16 = acc.to(tl.float16)

    if STORE_MASK_FREE:
        tl.store(C + offs_c, acc_f16)
    else:
        tl.store(C + offs_c, acc_f16, mask=c_mask, cache_modifier=".wt")


# for intra-node case
@triton.heuristics(
    values={
        "DIVISIBLE_M": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0,
        "DIVISIBLE_N": lambda args: args["N"] % args["BLOCK_SIZE_N"] == 0,
        'DIVISIBLE_K': lambda args: args['K'] % args['BLOCK_SIZE_K'] == 0,
    }
)
# @triton.jit
@triton.jit
def _w8a8_block_matmul_scatter_kernel(
    # Pointers to inputs and output
    A,
    B,
    C,
    As,
    Bs,
    # Shape for matmul
    M,
    N,
    K,
    # Block size for block-wise quantization
    group_n: tl.constexpr,
    group_k: tl.constexpr,
    # Stride for inputs and output
    stride_am,
    stride_ak,
    stride_bk,
    stride_bn,
    stride_cm,
    stride_cn,
    stride_As_m,
    stride_As_k,
    stride_Bs_k,
    stride_Bs_n,
    # shmem args
    shm_ctx,
    barrier_bufs_ptr,
    scatter_bufs_ptr,
    barrier_buf_ptr,
    scatter_buf_ptr,
    cur_rank: tl.constexpr,
    local_world_size: tl.constexpr,
    # Meta-parameters
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    DIVISIBLE_M: tl.constexpr,
    DIVISIBLE_N: tl.constexpr,
    DIVISIBLE_K: tl.constexpr,
    COMBINE_SCALE_LOAD: tl.constexpr = False,
    USE_MLS_LOAD: tl.constexpr = False
):
    """Triton-accelerated function used to perform linear operations (dot
    product) on input tensors `A` and `B` with block-wise quantization, and
    store the result in output tensor `C`.
    """

    # set shmem device ctx (rocSHMEM required when call libshmem.func)
    # libshmem_device.set_rocshmem_ctx(shm_ctx)
    # tid = libdevice.thread_idx(0) # noqa

    # node_id = cur_rank // local_world_size

    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    tiles_per_rank_n = tl.cdiv(num_pid_n, local_world_size)
    N_per_rank = tl.cdiv(N, local_world_size)
    rank_offset = pid_n // tiles_per_rank_n

    tl.assume(pid > 0)
    tl.assume(pid_m > 0)
    tl.assume(pid_n > 0)
    tl.assume(group_n > 0)
    tl.assume(group_k > 0)
    tl.assume(stride_am > 0)
    tl.assume(stride_ak > 0)
    tl.assume(stride_bk > 0)
    tl.assume(stride_bn > 0)
    tl.assume(stride_cm > 0)
    tl.assume(stride_cn > 0)
    tl.assume(stride_As_m > 0)
    tl.assume(stride_As_k > 0)
    tl.assume(stride_Bs_k > 0)
    tl.assume(stride_Bs_n > 0)

    if group_k > 0:
        tl.static_assert(BLOCK_SIZE_K <= group_k and group_k % BLOCK_SIZE_K == 0,
            "BLOCK_SIZE_K must be divisible by GROUP_SIZE_K")
    if COMBINE_SCALE_LOAD: # used for use_int8_w8a8
        tl.static_assert(stride_As_k == 1,
            "COMBINE_SCALE_LOAD implictly stride_As_k == 1!")
        tl.static_assert(DIVISIBLE_K == True and BLOCK_SIZE_K == group_k,
            "COMBINE_SCALE_LOAD only add and verify on block_k_diviable!")
    if USE_MLS_LOAD:
        tl.static_assert(DIVISIBLE_K == True and DIVISIBLE_N == True,
            "USE_MLS_LOAD must require block_k_diviable and block_n_diviable!")

    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

    mls_offs_k = 0
    if COMBINE_SCALE_LOAD:
        As_ptrs = As + offs_am[:, None] * stride_As_m
        offs_bsn = offs_bn // group_n
        Bs_ptrs = Bs + offs_bsn[:, None] * stride_Bs_n
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K), 2):
            a0 = tl.load(a_ptrs,
                        mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
                        other=0.0)
            if not USE_MLS_LOAD:
                b0 = tl.load(b_ptrs,
                            mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                            other=0.0)
            else:
                b0 = tl.matrix_load(
                            B,
                            shape=[K, N],
                            strides=[stride_bk, stride_bn],
                            block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N],
                            offsets=[mls_offs_k, (pid_n * BLOCK_SIZE_N) % N])

            k_start = k * BLOCK_SIZE_K
            offs_ks = k_start // group_k + tl.arange(0, 2)
            a_s = tl.load(As_ptrs + offs_ks[None, :] * stride_As_k,
                          mask=offs_ks[None, :] <= (K - 1) // group_k,
                          other=0.0)
            b_s = tl.load(Bs_ptrs + offs_ks[None, :] * stride_Bs_k,
                          mask=offs_ks[None, :] <= (K - 1) // group_k,
                          other=0.0)
            a_s0, a_s1 = tl.split(a_s)
            b_s0, b_s1 = tl.split(b_s)

            accumulator += tl.dot(a0, b0).to(tl.float32) * a_s0[:, None] * b_s0[None, :]

            a0 = tl.load(a_ptrs + BLOCK_SIZE_K * stride_ak,
                        mask=offs_k[None, :] < K - (k + 1)* BLOCK_SIZE_K,
                        other=0.0)
            if not USE_MLS_LOAD:
                b0 = tl.load(b_ptrs + BLOCK_SIZE_K * stride_bk,
                            mask=offs_k[:, None] < K - (k + 1) * BLOCK_SIZE_K,
                            other=0.0)
            else:
                b0 = tl.matrix_load(
                        B,
                        shape=[K, N],
                        strides=[stride_bk, stride_bn],
                        block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N],
                        offsets=[mls_offs_k + BLOCK_SIZE_K, (pid_n * BLOCK_SIZE_N) % N])

            accumulator += tl.dot(a0, b0).to(tl.float32) * a_s1[:, None] * b_s1[None, :]

            a_ptrs += BLOCK_SIZE_K * stride_ak * 2
            b_ptrs += BLOCK_SIZE_K * stride_bk * 2
            mls_offs_k += BLOCK_SIZE_K * 2
    else:
        As_ptrs = As + offs_am * stride_As_m
        offs_bsn = offs_bn // group_n
        Bs_ptrs = Bs + offs_bsn * stride_Bs_n
        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
            a = tl.load(a_ptrs,
                        mask=offs_k[None, :] < K - k * BLOCK_SIZE_K,
                        other=0.0)
            if not USE_MLS_LOAD:
                b = tl.load(b_ptrs,
                            mask=offs_k[:, None] < K - k * BLOCK_SIZE_K,
                            other=0.0)
            else:
                b = tl.matrix_load(
                            B,
                            shape=[K, N],
                            strides=[stride_bk, stride_bn],
                            block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N],
                            offsets=[mls_offs_k, (pid_n * BLOCK_SIZE_N) % N])

            k_start = k * BLOCK_SIZE_K
            offs_ks = k_start // group_k
            a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
            b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)

            accumulator += tl.dot(a, b).to(tl.float32) * a_s[:, None] * b_s[None, :]

            a_ptrs += BLOCK_SIZE_K * stride_ak
            b_ptrs += BLOCK_SIZE_K * stride_bk
            mls_offs_k += BLOCK_SIZE_K


    if C.dtype.element_ty == tl.bfloat16:
        c = accumulator.to(tl.bfloat16)
    elif C.dtype.element_ty == tl.float16:
        c = accumulator.to(tl.float16)
    else:
        c = accumulator.to(tl.float32)

    # map columns into the slice owned by cur_rank
    local_pid_n = pid_n % tiles_per_rank_n
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = (cur_rank * N_per_rank + local_pid_n * BLOCK_SIZE_N) + tl.arange(0, BLOCK_SIZE_N)
    offs_c = stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]

    scatter_ptr = tl.load(scatter_bufs_ptr + rank_offset).to(tl.pointer_type(tl.float32))
    c_ptrs = scatter_ptr + offs_c

    # # for debug
    # offs_cm1 = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    # offs_cn1 = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    # offs_c1 = stride_cm * offs_cm1[:, None] + stride_cn * offs_cn1[None, :]
    # c_ptrs1 = C + offs_c1
    # c1 = tl.load(c_ptrs1)

    STORE_MASK_FREE: tl.constexpr = DIVISIBLE_M & DIVISIBLE_N
    if STORE_MASK_FREE:
        tl.store(c_ptrs, c)
        c_mask = None
    else:
        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)

    '''
    # Signal to other ranks that we have completed this tile
    for i in range(1, local_world_size):
        remote = (cur_rank + i) % local_world_size + node_id * local_world_size
        remote_base_ptr = tl.load(barrier_bufs_ptr + remote).to(tl.pointer_type(tl.int32))
        # remote_base_ptr = libshmem_device.remote_ptr(barrier_bufs_ptr, remote).to(tl.pointer_type(tl.int32))
        tl.atomic_add(remote_base_ptr + pid, 1, scope="sys", sem="release")

    # consumer
    # local_base_ptr = tl.load(barrier_bufs_ptr + cur_rank).to(tl.pointer_type(tl.int32))
    while tl.atomic_cas(barrier_buf_ptr + pid, world_size - 1, 0, scope="sys", sem="acquire") != (world_size - 1):
        pass

    # acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    acc = accumulator

    # intra-node reduce
    for i in range(1, local_world_size):
        rank_id = (cur_rank + i) % local_world_size + node_id * local_world_size
        scatter_buf_ptr = tl.load(scatter_bufs_ptr + rank_id).to(tl.pointer_type(tl.float32))
        # scatter_buf_ptr = libshmem_device.remote_ptr(scatter_bufs_ptr, rank_id).to(tl.pointer_type(tl.float32))
        acc += tl.load(scatter_buf_ptr + offs_c, mask=c_mask)
    acc_f16 = acc.to(tl.float16)

    if STORE_MASK_FREE:
        tl.store(C + offs_c, acc_f16)
    else:
        tl.store(C + offs_c, acc_f16, mask=c_mask, cache_modifier=".wt")
    '''

def gemm_allreduce_w8a8(
    A: torch.Tensor,
    B: torch.Tensor,
    As: torch.Tensor,
    Bs: torch.Tensor,
    block_size: list[int],
    output_dtype: torch.dtype,
    barriers_buf: torch.tensor,
    scatters_buf: torch.tensor,
    barrier_buf: torch.tensor,
    scatter_buf: torch.tensor,
    max_tiles: int,
    tp_group: dist.ProcessGroup,
    local_world_size: int,
    shmem_ctx: Optional[np.intp] = None,
    configs: Optional[Dict] = None
) -> torch.Tensor:
    """This function performs matrix multiplication with block-wise
    quantization.

    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
    The output is returned in the specified `output_dtype`.

    Args:
        A: The input tensor, e.g., activation.
        B: The input tensor, e.g., weight.
        As: The per-token-group quantization scale for `A`.
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization. It should be
            2-dim, e.g., [128, 128].
        output_dytpe: The dtype of the returned tensor.

    Returns:
        torch.Tensor: The result of matmul.
    """

    # get from tp_group
    cur_rank = tp_group.rank()
    world_size = tp_group.size()

    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]

    assert A.shape[-1] == B.shape[-1]
    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
    M = A.numel() // A.shape[-1]

    assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
    N, K = B.shape
    assert triton.cdiv(N, block_n) == Bs.shape[0]
    assert triton.cdiv(K, block_k) == Bs.shape[1]

    C_shape = A.shape[:-1] + (N, )
    C = A.new_empty(C_shape, dtype=output_dtype)

    if configs is None:
        configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1])
    if configs:
        # If an optimal configuration map has been found, look up the
        # optimal config
        config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
    else:
        # Default config
        # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1]
        config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": block_size[0],
            "BLOCK_SIZE_K": block_size[1],
            "GROUP_SIZE_M": 2,
            "COMBINE_SCALE_LOAD": False,
            "num_warps": 4,
            "num_stages": 1,
        }

    # check
    cur_total_tiles = \
        triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"])

    assert cur_total_tiles <= max_tiles, (
        f"{cur_total_tiles=} should not be bypass {max_tiles=}, "
         "please check the min BLOCK_SIZE_M & BLOCK_SIZE_N in awq_gemm config and pass them to register_shmem")

    def grid(META):
        return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *
                triton.cdiv(N, META["BLOCK_SIZE_N"]), )

    _kernel = _w8a8_block_matmul_allreduce_oneshot_kernel[grid](
        A,
        B,
        C,
        As,
        Bs,
        M,
        N,
        K,
        block_n,
        block_k,
        A.stride(-2),
        A.stride(-1),
        B.stride(1),
        B.stride(0),
        C.stride(-2),
        C.stride(-1),
        As.stride(-2),
        As.stride(-1),
        Bs.stride(1),
        Bs.stride(0),
        # shmem args
        shmem_ctx,
        barriers_buf,
        scatters_buf,
        barrier_buf,
        scatter_buf,
        cur_rank,
        local_world_size,
        world_size,
        **config,
    )

    # # Create kernel path metadata, which mapping autotune's key to kernel path.
    # # input_dtype = str(input.dtype).split('.')[-1]
    # block_shape_n = B.shape[0] // Bs.shape[0]
    # block_shape_k = B.shape[1] // Bs.shape[1]
    # path_key = str((M, K, N, block_shape_n, block_shape_k, config['BLOCK_SIZE_M'], config['BLOCK_SIZE_N'], config['BLOCK_SIZE_K'], config['GROUP_SIZE_M'], config['COMBINE_SCALE_LOAD']))
    # arch = triton.runtime.driver.active.get_current_target().arch
    # save_kernel_path(f"{_kernel.name}-{arch}-cu{get_cu_num()}-blockwise-w8a8.json",
    #                path_key, os.path.basename(_kernel.perf_ir_path))

    return C

# support internode with 3 stages
# intra-scatter + inter-allreduce + intra-allgather
def gemm_allreduce_w8a8_v2(
    A: torch.Tensor,
    B: torch.Tensor,
    As: torch.Tensor,
    Bs: torch.Tensor,
    block_size: list[int],
    output_dtype: torch.dtype,
    barriers_buf: torch.tensor,
    scatters_buf: torch.tensor,
    reduces_buf: torch.tensor,
    barrier_buf: torch.tensor,
    scatter_buf: torch.tensor,
    reduce_buf: torch.tensor,
    max_seq_len: int,
    max_tiles: int,
    tp: GroupCoordinator,
    pp: GroupCoordinator,
    shmem_ctx: Optional[np.intp] = None,
    configs: Optional[Dict] = None
) -> torch.Tensor:
    """This function performs matrix multiplication with block-wise
    quantization.

    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
    The output is returned in the specified `output_dtype`.

    Args:
        A: The input tensor, e.g., activation.
        B: The input tensor, e.g., weight.
        As: The per-token-group quantization scale for `A`.
        Bs: The per-block quantization scale for `B`.
        block_size: The block size for per-block quantization. It should be
            2-dim, e.g., [128, 128].
        output_dytpe: The dtype of the returned tensor.

    Returns:
        torch.Tensor: The result of matmul.
    """

    # get from tp_group
    world_size = tp.world_size
    cur_rank = tp.rank_in_group
    # nnodes = world_size // local_world_size

    assert len(block_size) == 2
    block_n, block_k = block_size[0], block_size[1]

    assert A.shape[-1] == B.shape[-1]
    assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
    assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
    M = A.numel() // A.shape[-1]

    assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
    N, K = B.shape
    assert triton.cdiv(N, block_n) == Bs.shape[0]
    assert triton.cdiv(K, block_k) == Bs.shape[1]

    # C_shape = A.shape[:-1] + (N, )
    # C = A.new_empty(C_shape, dtype=output_dtype)
    # # for debug
    # C = A.new_zeros(C_shape, dtype=output_dtype)
    # rows = C.shape[0]
    # cols = C.shape[1]
    # increment = torch.arange(rows * cols).reshape(rows, cols).to(C.device)
    # C = C + increment

    if configs is None:
        configs = get_w8a8_block_int8_configs(N, K, block_size[0], block_size[1])
    if configs:
        # If an optimal configuration map has been found, look up the
        # optimal config
        config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
    else:
        # Default config
        # Block-wise quant: BLOCK_SIZE_K must be divisible by block_size[1]
        config = {
            "BLOCK_SIZE_M": 64,
            "BLOCK_SIZE_N": block_size[0],
            "BLOCK_SIZE_K": block_size[1],
            "GROUP_SIZE_M": 2,
            "COMBINE_SCALE_LOAD": False,
            "USE_MLS_LOAD": False,
            "num_warps": 4,
            "num_stages": 1,
        }

    # check
    cur_total_tiles = \
        triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"])

    assert cur_total_tiles <= max_tiles, (
        f"{cur_total_tiles=} should not be bypass {max_tiles=}, "
         "please check the min BLOCK_SIZE_M & BLOCK_SIZE_N in awq_gemm config and pass them to register_shmem")

    # test divisible case first
    assert triton.cdiv(N, config["BLOCK_SIZE_N"]) >= world_size, (
        f'{N=} is too small compared to {config["BLOCK_SIZE_N"]=}'
        'please make sure num_pid_n at least >= local_world_size for efficiency'
    )

    # Stage 1: intra node gemm + scatter
    # A * B -> scatter_buf
    def grid(META):
        return (triton.cdiv(M, META["BLOCK_SIZE_M"]) *
                triton.cdiv(N, META["BLOCK_SIZE_N"]), )
    _kernel_gemm = _w8a8_block_matmul_scatter_kernel[grid](
        A,
        B,
        reduce_buf, #C,
        As,
        Bs,
        M,
        N,
        K,
        block_n,
        block_k,
        A.stride(-2),
        A.stride(-1),
        B.stride(1),
        B.stride(0),
        reduce_buf.stride(-2),
        reduce_buf.stride(-1),
        As.stride(-2),
        As.stride(-1),
        Bs.stride(1),
        Bs.stride(0),
        # shmem args
        shmem_ctx,
        barriers_buf,
        scatters_buf,
        barrier_buf,
        scatter_buf,
        cur_rank,
        world_size,
        **config,
    )

    _kernel_barrier = barrier_all_ipc[(1, )](cur_rank, world_size, barriers_buf)

    N_per_rank = triton.cdiv(N, world_size)

    grid_reduce = lambda META: (
        triton.cdiv(M, META["BLOCK_SIZE_M"])
        * triton.cdiv(N_per_rank, META["BLOCK_SIZE_N"]),
    )
    _kernel_scatter_reduce = scatter_reduce_gather_kernel[grid_reduce](
        reduce_buf,
        reduces_buf,
        scatter_buf,
        M,
        N,
        N_per_rank,
        scatter_buf.stride(-2),
        scatter_buf.stride(-1),
        reduce_buf.stride(-2),
        reduce_buf.stride(-1),
        cur_rank=cur_rank,
        local_world_size=world_size,
        BLOCK_SIZE_M=32,
        BLOCK_SIZE_N=128,
        num_warps=16,
        num_stages=2,
    )

    # return C
    return reduce_buf[:M]
    # return reduce_buf # for debug


@triton.jit
def barrier_all_ipc(rank, num_ranks, comm_buf_base_ptrs):
    tid = libdevice.thread_idx(axis=0)  # noqa: F841
    for i in range(num_ranks):
        remote_base_ptr = tl.load(comm_buf_base_ptrs + i).to(tl.pointer_type(tl.int32))
        # remote_base_ptr = comm_buf_base_ptrs[i]
        while tl.atomic_cas(remote_base_ptr + rank, 0, 1, scope="sys", sem="release") != 0:
            pass

    for i in range(num_ranks):
        local_base_ptr = tl.load(comm_buf_base_ptrs + rank).to(tl.pointer_type(tl.int32))
        # local_base_ptr = comm_buf_base_ptrs[rank]
        while tl.atomic_cas(local_base_ptr + i, 1, 0, scope="sys", sem="acquire") != 1:
            pass

    tl.debug_barrier()

@triton.jit
def scatter_reduce_local_kernel(
    out_ptr,
    scatter_buf_ptr,
    M,
    N,
    N_per_rank,
    stride_scatter_m,
    stride_scatter_n,
    stride_out_m,
    stride_out_n,
    cur_rank: tl.constexpr,
    local_world_size: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    """Reduce all scatter buffers into a local output slice of shape (M, N_per_rank)."""

    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_per_rank, BLOCK_SIZE_N)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_local = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    # accum = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    # mask = (offs_m[:, None] < M) & (offs_n_local[None, :] < N)
    ptrs = scatter_buf_ptr + offs_m[:, None] * stride_scatter_m + offs_n_local[None, :] * stride_scatter_n
    accum = tl.load(ptrs)
    # for r in range(0, local_world_size):
    for r in range(1, local_world_size):
        offs_n_global = offs_n_local + r * N_per_rank
        mask = (offs_m[:, None] < M) & (offs_n_global[None, :] < N)
        ptrs = scatter_buf_ptr + offs_m[:, None] * stride_scatter_m + offs_n_global[None, :] * stride_scatter_n
        accum += tl.load(ptrs, mask=mask)

    # offs_n_local = offs_n_local + (cur_rank % local_world_size) * N_per_rank
    out_ptrs = out_ptr + offs_m[:, None] * stride_out_m + offs_n_local[None, :] * stride_out_n
    tl.store(out_ptrs, accum, mask=(offs_m[:, None] < M) & (offs_n_local[None, :] < N_per_rank))

@triton.heuristics(
    values={
        "DIVISIBLE_M": lambda args: args["M"] % args["BLOCK_SIZE_M"] == 0,
        "DIVISIBLE_N": lambda args: (
            args["N_per_rank"] % args["BLOCK_SIZE_N"] == 0
            and args["N"] % args["local_world_size"] == 0
        ),
    }
)
# @triton.jit
@triton.jit
def scatter_reduce_gather_kernel(
    out_ptr,
    outs_ptr,
    scatter_buf_ptr,
    M,
    N,
    N_per_rank,
    stride_scatter_m,
    stride_scatter_n,
    stride_out_m,
    stride_out_n,
    cur_rank: tl.constexpr,
    local_world_size: tl.constexpr,
    DIVISIBLE_M: tl.constexpr,
    DIVISIBLE_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    """Reduce all scatter buffers into a local output slice of shape (M, N_per_rank)."""

    # tid = libdevice.thread_idx(0)
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_per_rank, BLOCK_SIZE_N)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_local = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    tl.multiple_of(offs_m, BLOCK_SIZE_M)
    tl.multiple_of(offs_n_local, BLOCK_SIZE_N)
    tl.max_contiguous(offs_n_local, BLOCK_SIZE_N)

    STORE_MASK_FREE: tl.constexpr = DIVISIBLE_M & DIVISIBLE_N

    # accum = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    ptrs = scatter_buf_ptr + offs_m[:, None] * stride_scatter_m + offs_n_local[None, :] * stride_scatter_n
    if STORE_MASK_FREE:
        accum = tl.load(ptrs)
    else:
        mask = (offs_m[:, None] < M) & (offs_n_local[None, :] < N_per_rank)
        accum = tl.load(ptrs, mask)
    # scatter buffer add locally across N_per_rank
    for r in range(1, local_world_size):
        offs_n_global = offs_n_local + r * N_per_rank
        n_end_rank = (r + 1) * N_per_rank
        n_end = tl.minimum(N, n_end_rank)
        ptrs = scatter_buf_ptr + offs_m[:, None] * stride_scatter_m + offs_n_global[None, :] * stride_scatter_n
        if STORE_MASK_FREE:
            accum += tl.load(ptrs)
        else:
            mask = (offs_m[:, None] < M) & (offs_n_global[None, :] < n_end)
            accum += tl.load(ptrs, mask=mask)

    # scatter buffer push reduce results to local/peer rank
    offs_n_local = offs_n_local + cur_rank * N_per_rank
    n_end_rank = (cur_rank + 1) * N_per_rank
    n_end = tl.minimum(N, n_end_rank)
    for r in range(local_world_size):
        remote_out_ptr = tl.load(outs_ptr + r).to(tl.pointer_type(out_ptr.dtype.element_ty))
        out_ptrs = remote_out_ptr + offs_m[:, None] * stride_out_m + offs_n_local[None, :] * stride_out_n
        if STORE_MASK_FREE:
            tl.store(out_ptrs, accum)
        else:
            tl.store(out_ptrs, accum, mask=(offs_m[:, None] < M) & (offs_n_local[None, :] < n_end)) #, cache_modifier=".cs")

@triton.jit
def scatter_gather_local_kernel(
    out_ptr,
    scatter_buf_ptr,
    scatter_bufs_ptr,
    max_seq_len: tl.constexpr,
    M: tl.constexpr,
    N: tl.constexpr,
    N_per_rank: tl.constexpr,
    stride_scatter_m: tl.constexpr,
    stride_scatter_n: tl.constexpr,
    stride_out_m: tl.constexpr,
    stride_out_n: tl.constexpr,
    cur_rank: tl.constexpr,
    local_world_size: tl.constexpr,
    nnodes: tl.constexpr,
    DO_REDUCE: tl.constexpr, # need reduce locally if nnodes > 1
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    """Gather the peer ranks scatter_bufs M * N_per_rank and concat along N dim"""

    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N_per_rank, BLOCK_SIZE_N)
    pid_m = pid // num_pid_n
    pid_n = pid % num_pid_n

    # the cur_rank id inside the node (0 ~ local_world_size - 1)
    cur_node = cur_rank // local_world_size
    cur_rank_local = cur_rank % local_world_size

    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n_local = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)

    # first do local copy
    offs_n_cur_rank = offs_n_local + cur_rank_local * N_per_rank
    offs_in = offs_m[:, None] * stride_scatter_m + offs_n_local * stride_scatter_n
    ptrs = scatter_buf_ptr + offs_in
    mask = (offs_m[:, None] < M) & (offs_n_local[None, :] < N_per_rank)
    accum = tl.load(ptrs, mask=mask)

    # local reduce if nnodes > 1
    if DO_REDUCE:
        for i in range(1, nnodes):
            ptrs = scatter_buf_ptr + offs_in + i * max_seq_len * N_per_rank
            accum += tl.load(ptrs, mask=mask)

    offs_out = offs_m[:, None] * stride_out_m + offs_n_cur_rank[None, :] * stride_out_n
    out_ptrs = out_ptr + offs_out

    n_end_rank = (cur_rank_local + 1) * N_per_rank
    n_end = tl.minimum(N, n_end_rank)

    mask = (offs_m[:, None] < M) & (offs_n_local[None, :] < n_end)
    tl.store(out_ptrs, accum, mask)

    for r in range(1, local_world_size):
        peer_rank_local = (cur_rank_local + r) % local_world_size
        peer_rank = peer_rank_local + local_world_size * cur_node
        offs_n_global = offs_n_local + peer_rank_local * N_per_rank
        mask_in = (offs_m[:, None] < M) & (offs_n_local[None, :] < N_per_rank)
        # load remote ptr
        remote_ptr = tl.load(scatter_bufs_ptr + peer_rank).to(tl.pointer_type(tl.float32))
        offs_in = offs_m[:, None] * stride_scatter_m + offs_n_local[None, :] * stride_scatter_n
        ptrs = remote_ptr + offs_in
        a = tl.load(ptrs, mask=mask_in, other=0.0)
        offs_out = offs_m[:, None] * stride_out_m + offs_n_global[None, :] * stride_out_n
        n_end_rank = (peer_rank_local + 1) * N_per_rank
        n_end = tl.minimum(N, n_end_rank)
        mask_out = (offs_m[:, None] < M) & (offs_n_global[None, :] < n_end_rank)
        out_ptrs = out_ptr + offs_out
        tl.store(out_ptrs, a, mask_out)


@triton.jit
def putmem_kernel(
    shm_ctx,
    out_ptr, # inout
    max_seq_len: tl.constexpr,
    M: tl.constexpr,
    N_per_rank: tl.constexpr,
    cur_rank: tl.constexpr,
    local_world_size: tl.constexpr,
    nnodes: tl.constexpr,
    TOKEN_BYTES: tl.constexpr,
):
    # set rocshmem device ctx
    libshmem_device.set_rocshmem_ctx(shm_ctx)

    pid = tl.program_id(0)
    byte_cnt = M * N_per_rank * TOKEN_BYTES

    for i in range(1, nnodes):
        peer_node = (i + 1) % nnodes
        peer_rank = peer_node * local_world_size + cur_rank % local_world_size
        dst_ptr = out_ptr +  i * max_seq_len * N_per_rank
        libshmem_device.putmem_nbi_wg(
            dst_ptr,
            out_ptr,
            byte_cnt,
            peer_rank,
        )

    # libshmem_device.fence()
    return
