# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
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

import torch
import triton
import triton.language as tl

from triton.language.extra.hip import libdevice
from triton.language.extra import libshmem_device
import numpy as np
import torch.distributed as dist

AWQ_TRITON_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]

def reverse_awq_order(tensor: torch.Tensor) -> torch.Tensor:
    """Reverse the AWQ order of the given tensor.

    Args:
        tensor: Input tensor to reorder

    Returns:
        Reordered tensor with bits masked to 4 bits
    """
    bits = 4
    AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
    reverse_order_tensor = torch.arange(
        tensor.shape[-1],
        dtype=torch.int32,
        device=tensor.device,
    )
    reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
    reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
    reverse_order_tensor = reverse_order_tensor.view(-1)

    tensor = tensor[:, reverse_order_tensor] & 0xF
    return tensor

def awq_reorder_and_repack(
    qweight: torch.Tensor,
    qzeros: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Reorder and pack weights and zeros using AWQ order.
    
    This function unpacks the 4-bit quantized weights and zeros from int32,
    applies reverse_awq_order to reorder them, and then packs them.
    For weight, repack to [N, K//2]
    For zeros, repack to [K//G, N//2]
    Args:
        qweight: Quantized weight tensor of shape [K, N // 8] with dtype int32
        qzeros: Quantized zero points tensor of shape [K // G, N // 8] with dtype int32
        
    Returns:
        Tuple of (reordered_qweight, reordered_qzeros) both with dtype int8
    """
    bits = 4
    shifts = torch.arange(0, 32, bits, device=qweight.device)
    K = qweight.shape[0]
    N = qweight.shape[1] * 8
    G = K // qzeros.shape[0]
    
    # Unpack weights: [K, N//8] -> [K, N//8, 8] -> [K, N]
    iweights = torch.bitwise_right_shift(
        qweight[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    iweights = iweights.view(K, -1)
    
    # Unpack zeros: [K//G, N//8] -> [K//G, N//8, 8] -> [K//G, N]
    zeros = torch.bitwise_right_shift(
        qzeros[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    zeros = zeros.view(K//G, -1)
    
    # Apply reverse AWQ order to both tensors
    iweights = reverse_awq_order(iweights)
    zeros = reverse_awq_order(zeros)
    
    # Mask to 4 bits
    iweights = torch.bitwise_and(iweights, (2**bits) - 1)
    zeros = torch.bitwise_and(zeros, (2**bits) - 1)
    
    # Repack weight to int32 and pack along the K direction
    # [K, N] -> [N, K]
    iweights = iweights.transpose(1, 0).contiguous()
    # Reshape to [N, K//2, 2] for weights
    iweights_packed = iweights.view(N, -1, 2)

    # Repack zeros to int8 and pack along the N direction
    # Reshape to [K//G, N//2, 2] for zeros
    zeros_packed = zeros.view(K//G, -1, 2)
    
    # Pack 2 int4 values into int8 using bit shifts
    # Direct packing: pack in the order they appear after reordering
    packed_weights = torch.zeros([N, K//2], dtype=torch.int8, device=qweight.device)
    packed_zeros = torch.zeros([K//G, N//2], dtype=torch.int8, device=zeros.device)
    
    for i in range(2):
        packed_weights |= (iweights_packed[:, :, i].to(torch.int8) << (i * bits))
        packed_zeros |= (zeros_packed[:, :, i].to(torch.int8) << (i * bits))
    
    return packed_weights, packed_zeros

'''
@triton.autotune(
    configs=[
        triton.Config({
            "BLOCK_SIZE_N": BN,
            "BLOCK_SIZE_K": BK
        }, num_warps=num_warps, num_stages=num_stages)
        for BN in [16, 32, 64, 128, 256]
        for BK in [16, 32, 64, 128, 256]
        for num_warps in [1, 2, 4, 8, 16] for num_stages in [1, 2]
    ],
    key=["K", "N"],
    perf_debug=True,
)
'''
@triton.heuristics(values={
    "NUM_GROUPS": lambda args: triton.cdiv(args["BLOCK_SIZE_K"], args["group_size"]),
    "BLOCK_SIZE_K2": lambda args: args["BLOCK_SIZE_K"] // 2
})
@triton.jit
def awq_dequantize_kernel(
        qweight_ptr,  # quantized matrix
        scales_ptr,  # scales, per group
        zeros_ptr,  # zeros, per group
        result_ptr,  # Output matrix
        N,
        N2,
        K,
        K2,
        group_size: tl.constexpr,  # Should always be one of the supported group sizes
        NUM_GROUPS: tl.constexpr,
        BLOCK_SIZE_N: tl.constexpr,
        BLOCK_SIZE_K: tl.constexpr,
        BLOCK_SIZE_K2: tl.constexpr):

    # Setup the pids.
    pid_n = tl.program_id(axis=0)
    pid_k = tl.program_id(axis=1)

    tl.assume(pid_n >= 0)
    tl.assume(pid_k >= 0)
    tl.assume(N > 0)
    tl.assume(K > 0)
    tl.assume(N2 > 0)
    tl.assume(K2 > 0)
    tl.assume(BLOCK_SIZE_N > 0)
    tl.assume(BLOCK_SIZE_K > 0)
    tl.assume(BLOCK_SIZE_K2 > 0)
    tl.assume(group_size > 0)

    # Compute offsets and masks for qweight_ptr.
    offsets_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offsets_n = tl.max_contiguous(tl.multiple_of(offsets_n, BLOCK_SIZE_N), BLOCK_SIZE_N)
    offsets_k = pid_k * BLOCK_SIZE_K2 + tl.arange(0, BLOCK_SIZE_K2)
    offsets_k = tl.max_contiguous(tl.multiple_of(offsets_k, BLOCK_SIZE_K2), BLOCK_SIZE_K2)
    offsets = K2 * offsets_n[:, None] + offsets_k[None, :]

    masks_n = offsets_n < N
    masks_k = offsets_k < K2

    masks = masks_n[:, None] & masks_k[None, :]

    # Compute offsets and masks for result output ptr.
    result_offsets_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    result_offsets_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
    result_offsets = (N * result_offsets_k[:, None] + result_offsets_n[None, :]) # [K, N]

    result_masks_n = result_offsets_n < N
    result_masks_k = result_offsets_k < K
    result_masks = result_masks_k[:, None] & result_masks_n[None, :]

    # Load the weights.
    iweights = tl.load(qweight_ptr + offsets, masks, 0.0) #[BLOCK_SIZE_N, BLOCK_SIZE_K//2]
    iweights = tl.interleave(iweights, iweights) # [BLOCK_SIZE_N, BLOCK_SIZE_K]

    # Use this to compute a set of shifts that can be used to unpack and
    # reorder the values in iweights and zeros.
    shifts = tl.arange(0, 2) * 4
    bshifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_N * BLOCK_SIZE_K2, 2))
    bshifts = tl.reshape(bshifts, (BLOCK_SIZE_N, BLOCK_SIZE_K))

    # Unpack and reorder: shift out the correct 4-bit value and mask.
    iweights = (iweights >> bshifts) & 0xF

    # Compute zero offsets and masks.
    zero_offsets_k = pid_k * BLOCK_SIZE_K // group_size + tl.arange(0, NUM_GROUPS)
    zero_offsets_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    zero_offsets_n2 = zero_offsets_n // 2
    zero_offsets = N2 * zero_offsets_k[:, None] + zero_offsets_n2[None, :]

    zero_masks_k = zero_offsets_k < K//group_size
    zero_masks_n = zero_offsets_n < N
    zero_masks = zero_masks_k[:, None] & zero_masks_n[None, :]

    # Load the zeros.
    zeros = tl.load(zeros_ptr + zero_offsets, zero_masks, 0.0) # [NUM_GROUPS, BLOCK_SIZE_N]

    # Compute scale offsets and masks.
    scale_offsets_k = pid_k * BLOCK_SIZE_K // group_size + tl.arange(0, NUM_GROUPS)
    scale_offsets_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
    scale_offsets = N * scale_offsets_k[:, None] + scale_offsets_n[None, :]
    scale_masks_k = scale_offsets_k < K//group_size
    scale_masks_n = scale_offsets_n < N
    scale_masks = scale_masks_k[:, None] & scale_masks_n[None, :]

    # Load the scales.
    scales = tl.load(scales_ptr + scale_offsets, scale_masks, 0.0) # [NUM_GROUPS, BLOCK_SIZE_N]

    if NUM_GROUPS == 1:
        zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N)) # [BLOCK_SIZE_K, BLOCK_SIZE_N]
        scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N)) # [BLOCK_SIZE_K, BLOCK_SIZE_N]
    else:
        zeros = tl.broadcast_to(zeros[:, None, :], (NUM_GROUPS, group_size, BLOCK_SIZE_N))
        scales = tl.broadcast_to(scales[:, None, :], (NUM_GROUPS, group_size, BLOCK_SIZE_N))
        zeros = tl.reshape(zeros, [BLOCK_SIZE_K, BLOCK_SIZE_N])
        scales = tl.reshape(scales, [BLOCK_SIZE_K, BLOCK_SIZE_N])

    # Unpack and reorder: shift out the correct 4-bit value and mask.
    zshifts = (zero_offsets_n[None, :] % 2) * 4 # [1, BLOCK_SIZE_N]
    zeros = (zeros >> zshifts) & 0xF # [BLOCK_SIZE_K, BLOCK_SIZE_N]

    # Dequantize.
    iweights = (iweights.T - zeros) * scales
    iweights = iweights.to(result_ptr.type.element_ty)

    # Finally, store.
    tl.store(result_ptr + result_offsets, iweights, result_masks)

@triton.jit
def awq_gemm_kernel_inner(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, tile_idx, k_idx, iter_begin, iter_end,
                          M, N, N2, K, K2, not_reduce, GROUP_SIZE: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
                          BLOCK_SIZE_K: tl.constexpr,
                          NUM_GROUPS: tl.constexpr, USE_REDUCE_KERNEL: tl.constexpr = 0):

    if not USE_REDUCE_KERNEL:
        k_idx = 0

    tl.assume(tile_idx >= 0)
    tl.assume(k_idx >= 0)
    tl.assume(iter_begin >= 0)
    tl.assume(iter_end >= 0)
    tl.assume(M > 0)
    tl.assume(N > 0)
    tl.assume(K > 0)
    tl.assume(K2 > 0)
    tl.assume(N2 > 0)

    num_tile_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_tile_n = tl.cdiv(N, BLOCK_SIZE_N)

    tile_idx_m = tile_idx // num_tile_n
    tile_idx_n = tile_idx % num_tile_n

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

    # Create reverse AWQ order as tensor: [0, 4, 1, 5, 2, 6, 3, 7]
    # that will map given indices to the correct order.
    #reverse_awq_order_tensor = ((tl.arange(0, 2) * 4)[None, :] + tl.arange(0, 4)[:, None]).reshape(8)

    # Create the necessary shifts to use to unpack.
    #shifts = reverse_awq_order_tensor * 4
    shifts = tl.arange(0, 2) * 4

    #zshifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_K * (BLOCK_SIZE_N // 2), 2))
    #zshifts = tl.reshape(zshifts, (BLOCK_SIZE_K, BLOCK_SIZE_N)).T

    #bshifts = tl.broadcast_to(shifts[:, None], (8, (BLOCK_SIZE_K // 8) * BLOCK_SIZE_N))
    bshifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_N * (BLOCK_SIZE_K // 2), 2))
    bshifts = tl.reshape(bshifts, (BLOCK_SIZE_N, BLOCK_SIZE_K))

    # Offsets and masks.
    offsets_am = tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    masks_am = offsets_am < M

    #offsets_zn = tile_idx_n * (BLOCK_SIZE_N // 2) + tl.arange(0, BLOCK_SIZE_N // 2)
    offsets_zn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offsets_zn = offsets_zn // 2
    #offsets_zn = tl.max_contiguous(tl.multiple_of(offsets_zn, BLOCK_SIZE_N // 2), BLOCK_SIZE_N // 2)
    masks_zn = offsets_zn < N2

    offsets_bn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    #offsets_bn = tl.max_contiguous(tl.multiple_of(offsets_bn, BLOCK_SIZE_N), BLOCK_SIZE_N)
    masks_bn = offsets_bn < N

    offsets_sn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    #offsets_sn = tl.max_contiguous(tl.multiple_of(offsets_sn, BLOCK_SIZE_N), BLOCK_SIZE_N)
    masks_sn = offsets_sn < N

    offsets_ak = iter_begin * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
    #offsets_ak = tl.max_contiguous(tl.multiple_of(offsets_ak, BLOCK_SIZE_K), BLOCK_SIZE_K)
    offsets_a = K * offsets_am[:, None] + offsets_ak[None, :]

    offsets_bk = iter_begin * (BLOCK_SIZE_K // 2) + tl.arange(0, BLOCK_SIZE_K // 2)
    #offsets_bk = tl.max_contiguous(tl.multiple_of(offsets_bk, BLOCK_SIZE_K // 2), BLOCK_SIZE_K // 2)
    #offsets_b = offsets_bk[:, None] + K // 2 * offsets_bn[None, :]
    #offsets_b = K // 2 * offsets_bn[:, None] + offsets_bk[None, :]
    offsets_b = K2 * offsets_bn[:, None] + offsets_bk[None, :]
    zshifts = (offsets_bn[:, None] % 2) * 4 # [N, 1]
    zshifts = zshifts.T

    a_ptrs = a_ptr + offsets_a
    b_ptrs = b_ptr + offsets_b
    for k in range(iter_end - iter_begin):
        masks_ak = offsets_ak < K
        masks_bk = offsets_bk < K2
        masks_a = masks_am[:, None] & masks_ak[None, :]
        masks_b = masks_bn[:, None] & masks_bk[None, :]
        other_bzs = 0.0
        a = tl.load(a_ptrs, mask=masks_a, other=0.)
        b = tl.load(b_ptrs, masks_b, other_bzs) #[N, K//2]
        b = tl.interleave(b, b) # [N, K]

        # Dequantize b.
        offsets_szk = ((BLOCK_SIZE_K * k + iter_begin * BLOCK_SIZE_K) // GROUP_SIZE + tl.arange(0, NUM_GROUPS))
        masks_szk = offsets_szk < K // GROUP_SIZE
        masks_z = masks_szk[:, None] & masks_zn[None, :]
        masks_s = masks_szk[:, None] & masks_sn[None, :]
        #masks_z = masks_zn[:, None] & masks_szk[None, :]
        #masks_s = masks_sn[:, None] & masks_szk[None, :]

        offsets_z = N2 * offsets_szk[:, None] + offsets_zn[None, :]
        #offsets_z = K // GROUP_SIZE * offsets_zn[:, None] + offsets_szk[None, :]
        zeros_ptrs = zeros_ptr + offsets_z
        zeros = tl.load(zeros_ptrs, mask=masks_z, other=other_bzs) # [K//G, N]
        #zshifts = (offsets_bn[:, None] % 2) * 4 # [N, 1]
        #zeros = (zeros >> _zshifts) & 0xF # [N, K//G]

        '''
        zeros = zeros.T # [K//G, N//2]
        zeros = tl.interleave(zeros, zeros) # [K//G, N]
        zeros = zeros.T # [N, K//G]
        '''

        offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :]
        #offsets_s = K // GROUP_SIZE * offsets_sn[:, None] + offsets_szk[None, :]
        scales_ptrs = scales_ptr + offsets_s
        scales = tl.load(scales_ptrs, mask=masks_s, other=other_bzs) # [K//G, N]

        if NUM_GROUPS == 1:
            # Original efficient implementation for single group
            zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            #zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_N, BLOCK_SIZE_K))
            #scales = tl.broadcast_to(scales, (BLOCK_SIZE_N, BLOCK_SIZE_K))
        else:
            # Reshape to (NUM_GROUPS, 1, N) then broadcast to (NUM_GROUPS, group_size_in_block, N)
            # Reshape to (K//G, 1, N) then broadcast to (K//G, group_size_in_block, N)
            zeros = tl.broadcast_to(zeros[:, None, :], (NUM_GROUPS, GROUP_SIZE, BLOCK_SIZE_N))
            scales = tl.broadcast_to(scales[:, None, :], (NUM_GROUPS, GROUP_SIZE, BLOCK_SIZE_N))
            ## Reshape back to (BLOCK_SIZE_K, N)
            zeros = tl.reshape(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            scales = tl.reshape(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))

            # Reshape to (N, K//G, 1) then broadcast to (N, K//G, group_size_in_block)
            #zeros = tl.broadcast_to(zeros[:, :, None], (BLOCK_SIZE_N, NUM_GROUPS, GROUP_SIZE))
            #scales = tl.broadcast_to(scales[:, :, None], (BLOCK_SIZE_N, NUM_GROUPS, GROUP_SIZE))
            #zeros = tl.reshape(zeros, (BLOCK_SIZE_N, BLOCK_SIZE_K))
            #scales = tl.reshape(scales, (BLOCK_SIZE_N, BLOCK_SIZE_K))

        b = (b >> bshifts) & 0xF
        b = b.T
        zeros = (zeros >> zshifts) & 0xF
        b = (b - zeros) * scales
        b = b.to(a_ptr.type.element_ty)

        # Accumulate results.
        accumulator = tl.dot(a, b, accumulator, out_dtype=tl.float32)

        offsets_ak += BLOCK_SIZE_K
        offsets_bk += BLOCK_SIZE_K // 2
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K // 2

    c = accumulator.to(c_ptr.type.element_ty)
    offs_cm = tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    # compiler hints
    offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_SIZE_M), BLOCK_SIZE_M)
    offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_SIZE_N), BLOCK_SIZE_N)
    offs_c = M * N * k_idx + N * offs_cm[:, None] + offs_cn[None, :]
    c_ptrs = c_ptr + offs_c
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)

    if USE_REDUCE_KERNEL:
        tl.store(c_ptrs, c, mask=c_mask)
    else:
        if c_ptr.type.element_ty == tl.float16:
            tl.store(c_ptrs, c, mask=c_mask)
        elif not_reduce:
            tl.store(c_ptrs, c, mask=c_mask)
        else:
            tl.atomic_add(c_ptrs, c, mask=c_mask)

@triton.jit
def awq_gemm_scatter_kernel_inner(
    a_ptr, b_ptr,
    c_ptr,
    zeros_ptr, scales_ptr, tile_idx, k_idx, iter_begin, iter_end,
    M, N, N2, K, K2, not_reduce, 
    cur_rank,
    world_size,
    GROUP_SIZE: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    NUM_GROUPS: tl.constexpr, USE_REDUCE_KERNEL: tl.constexpr = 0):

    if not USE_REDUCE_KERNEL:
        k_idx = 0

    tl.assume(tile_idx >= 0)
    tl.assume(k_idx >= 0)
    tl.assume(iter_begin >= 0)
    tl.assume(iter_end >= 0)
    tl.assume(M > 0)
    tl.assume(N > 0)
    tl.assume(K > 0)
    tl.assume(K2 > 0)
    tl.assume(N2 > 0)

    num_tile_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_tile_n = tl.cdiv(N, BLOCK_SIZE_N)

    tile_idx_m = tile_idx // num_tile_n
    tile_idx_n = tile_idx % num_tile_n
    
    # NOTE: this may be used when M is large 
    #   currently M range is small (1~128), so does M_PER_RANK
    # M_PER_RANK = M // world_size
    # num_tile_m_per_rank = M_PER_RANK // BLOCK_SIZE_M
    # rank_offset = tile_idx_m // num_tile_m_per_rank

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    shifts = tl.arange(0, 2) * 4
    bshifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_N * (BLOCK_SIZE_K // 2), 2))
    bshifts = tl.reshape(bshifts, (BLOCK_SIZE_N, BLOCK_SIZE_K))

    # Offsets and masks.
    offsets_am = tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    masks_am = offsets_am < M

    offsets_zn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offsets_zn = offsets_zn // 2
    masks_zn = offsets_zn < N2

    offsets_bn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    masks_bn = offsets_bn < N

    offsets_sn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    masks_sn = offsets_sn < N

    offsets_ak = iter_begin * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
    offsets_a = K * offsets_am[:, None] + offsets_ak[None, :]

    offsets_bk = iter_begin * (BLOCK_SIZE_K // 2) + tl.arange(0, BLOCK_SIZE_K // 2)
    offsets_b = K2 * offsets_bn[:, None] + offsets_bk[None, :]
    zshifts = (offsets_bn[:, None] % 2) * 4 # [N, 1]
    zshifts = zshifts.T

    a_ptrs = a_ptr + offsets_a
    b_ptrs = b_ptr + offsets_b
    for k in range(iter_end - iter_begin):
        masks_ak = offsets_ak < K
        masks_bk = offsets_bk < K2
        masks_a = masks_am[:, None] & masks_ak[None, :]
        masks_b = masks_bn[:, None] & masks_bk[None, :]
        other_bzs = 0.0
        a = tl.load(a_ptrs, mask=masks_a, other=0.)
        b = tl.load(b_ptrs, masks_b, other_bzs) #[N, K//2]
        b = tl.interleave(b, b) # [N, K]

        # Dequantize b.
        offsets_szk = ((BLOCK_SIZE_K * k + iter_begin * BLOCK_SIZE_K) // GROUP_SIZE + tl.arange(0, NUM_GROUPS))
        masks_szk = offsets_szk < K // GROUP_SIZE
        masks_z = masks_szk[:, None] & masks_zn[None, :]
        masks_s = masks_szk[:, None] & masks_sn[None, :]

        offsets_z = N2 * offsets_szk[:, None] + offsets_zn[None, :]
        zeros_ptrs = zeros_ptr + offsets_z
        zeros = tl.load(zeros_ptrs, mask=masks_z, other=other_bzs) # [K//G, N]

        offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :]
        #offsets_s = K // GROUP_SIZE * offsets_sn[:, None] + offsets_szk[None, :]
        scales_ptrs = scales_ptr + offsets_s
        scales = tl.load(scales_ptrs, mask=masks_s, other=other_bzs) # [K//G, N]

        if NUM_GROUPS == 1:
            # Original efficient implementation for single group
            zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))
        else:
            # Reshape to (NUM_GROUPS, 1, N) then broadcast to (NUM_GROUPS, group_size_in_block, N)
            # Reshape to (K//G, 1, N) then broadcast to (K//G, group_size_in_block, N)
            zeros = tl.broadcast_to(zeros[:, None, :], (NUM_GROUPS, GROUP_SIZE, BLOCK_SIZE_N))
            scales = tl.broadcast_to(scales[:, None, :], (NUM_GROUPS, GROUP_SIZE, BLOCK_SIZE_N))
            ## Reshape back to (BLOCK_SIZE_K, N)
            zeros = tl.reshape(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            scales = tl.reshape(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))

        b = (b >> bshifts) & 0xF
        b = b.T
        zeros = (zeros >> zshifts) & 0xF
        b = (b - zeros) * scales
        b = b.to(a_ptr.type.element_ty)

        # Accumulate results.
        accumulator = tl.dot(a, b, accumulator, out_dtype=tl.float32)

        offsets_ak += BLOCK_SIZE_K
        offsets_bk += BLOCK_SIZE_K // 2
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K // 2

    c = accumulator.to(c_ptr.type.element_ty)
    
    # offs_cm = tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    # NOTE: plus cur_rank offset
    offs_cm = M * cur_rank + tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    
    # compiler hints
    offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_SIZE_M), BLOCK_SIZE_M)
    offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_SIZE_N), BLOCK_SIZE_N)
    
    # batch index
    offs_c = world_size * M * N * k_idx + N * offs_cm[:, None] + offs_cn[None, :]
    c_mask = (offs_cm[:, None] < M * world_size) & (offs_cn[None, :] < N)
    
    # scatter to other ranks
    for j in tl.range(world_size, flatten=True):
        peer_rank = (j + cur_rank + 1) % world_size # swizzle
        c_ptr_peer = libshmem_device.remote_ptr(c_ptr, peer_rank).to(tl.pointer_type(c_ptr.type.element_ty))
        c_ptrs = c_ptr_peer + offs_c
        tl.store(c_ptrs, c, mask=c_mask)

@triton.jit
def awq_gemm_kernel_streamk(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2,
                    GROUP_SIZE: tl.constexpr, NUM_CUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr, 
                    DP_TILES: tl.constexpr, DANGLING_TILES: tl.constexpr,
                    NUM_GROUPS: tl.constexpr, USE_REDUCE_KERNEL: tl.constexpr):

    pid = tl.program_id(axis=0)

    if pid < DP_TILES:
        iters_per_cta = tl.cdiv(K, BLOCK_SIZE_K)
        awq_gemm_kernel_inner(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, pid, 0, 0, iters_per_cta,
                              M, N, N2, K, K2, True, GROUP_SIZE, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, 
                              NUM_GROUPS, USE_REDUCE_KERNEL)
    else:
        iters_per_tile = tl.cdiv(K, BLOCK_SIZE_K)
        total_iters = iters_per_tile * DANGLING_TILES
        iters_per_cta = tl.cdiv(total_iters, NUM_CUS)

        iter_begin = (pid - DP_TILES) * iters_per_cta
        iter_end = tl.minimum(iter_begin + iters_per_cta, total_iters)

        while iter_begin < iter_end:
            tile_idx = iter_begin // iters_per_tile + DP_TILES
            tile_iter_begin = (tile_idx - DP_TILES) * iters_per_tile
            tile_iter_end = tile_iter_begin + iters_per_tile
            local_iter_begin = iter_begin - tile_iter_begin
            local_iter_end = tl.minimum(iter_end, tile_iter_end) - tile_iter_begin
            k_idx = tl.cdiv(local_iter_begin, iters_per_cta)
            awq_gemm_kernel_inner(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, tile_idx, k_idx, local_iter_begin, local_iter_end,
                                M, N, N2, K, K2, False, GROUP_SIZE, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, 
                                NUM_GROUPS, USE_REDUCE_KERNEL)
            iter_begin = tile_iter_end

@triton.jit
def awq_gemm_kernel_splitk(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2,
                    GROUP_SIZE: tl.constexpr, NUM_CUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
                    BLOCK_SIZE_K: tl.constexpr,
                    SPLITK: tl.constexpr, NUM_GROUPS: tl.constexpr, USE_REDUCE_KERNEL: tl.constexpr):

    pid = tl.program_id(axis=0)

    tiles_M = tl.cdiv(M, BLOCK_SIZE_M)
    tiles_N = tl.cdiv(N, BLOCK_SIZE_N)
    total_tiles = tiles_M * tiles_N
    tile_idx = pid % total_tiles

    iters_per_cta = tl.cdiv(K, BLOCK_SIZE_K * SPLITK)
    iter_begin = pid // total_tiles * iters_per_cta
    iter_end = iter_begin + iters_per_cta
    k_idx = tl.cdiv(iter_begin, iters_per_cta)

    awq_gemm_kernel_inner(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, tile_idx, k_idx, iter_begin, iter_end,
                          M, N, N2, K, K2, False, GROUP_SIZE, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
                          NUM_GROUPS, USE_REDUCE_KERNEL)
       
@triton.jit
def awq_gemm_kernel_splitk_fused(
    a_ptr, b_ptr, c_ptr,
    zeros_ptr, scales_ptr,
    out_ptr, barrier_ptr,
    M, N, N2, K, K2,
    GROUP_SIZE: tl.constexpr,
    NUM_CUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    SPLITK: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
    USE_REDUCE_KERNEL: tl.constexpr):

    pid = tl.program_id(axis=0)
    tiles_M = tl.cdiv(M, BLOCK_SIZE_M)
    tiles_N = tl.cdiv(N, BLOCK_SIZE_N)
    total_tiles = tiles_M * tiles_N

    if pid < total_tiles * SPLITK:
        tile_idx = pid % total_tiles
        iters_per_cta = tl.cdiv(K, BLOCK_SIZE_K * SPLITK)
        iter_begin = pid // total_tiles * iters_per_cta
        iter_end = iter_begin + iters_per_cta
        k_idx = tl.cdiv(iter_begin, iters_per_cta)
        awq_gemm_kernel_inner(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, tile_idx, k_idx, iter_begin, iter_end,
                            M, N, N2, K, K2, False, GROUP_SIZE, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
                            NUM_GROUPS, USE_REDUCE_KERNEL)
        # set barriers
        tile_idx_m = tile_idx // tiles_N
        tile_idx_n = tile_idx % tiles_N
        offset = total_tiles * k_idx + tiles_N * tile_idx_m + tile_idx_n
        tl.store(barrier_ptr + offset, 1, cache_modifier=".wt")
    else:
        pid = pid - total_tiles * SPLITK
        # reduce kernel
        pid_m = pid // tiles_N
        pid_n = pid % tiles_N

        offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        # compiler hints
        offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
        offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N)

        mask_m = offs_m < M
        mask_n = offs_n < N
        mask = mask_m[:, None] & mask_n[None, :]

        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        # reduce split k
        for batch_idx in range(SPLITK):
            batch_offset = batch_idx * M * N
            input_offsets = batch_offset + offs_m[:, None] * N + offs_n[None, :]
            # wait barrier
            offset = total_tiles * batch_idx + pid_m * tiles_N + pid_n
            while tl.load(barrier_ptr + offset, cache_modifier=".cv", volatile=True) != 1:
                pass
            input_data = tl.load(c_ptr + input_offsets, mask=mask, other=0.0)
            acc += input_data

        output_offsets = offs_m[:, None] * N + offs_n[None, :]
        acc_f16 = acc.to(tl.float16)
        tl.store(out_ptr + output_offsets, acc_f16, mask=mask)

@triton.jit
def awq_gemm_allreduce_kernel_splitk(
    a_ptr, b_ptr, c_ptr,
    zeros_ptr, scales_ptr,
    out_ptr, barrier_ptr,
    M, N, N2, K, K2,
    cur_rank, world_size,
    GROUP_SIZE: tl.constexpr,
    NUM_CUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    SPLITK: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
    USE_REDUCE_KERNEL: tl.constexpr):

    pid = tl.program_id(axis=0)
    tiles_M = tl.cdiv(M, BLOCK_SIZE_M)
    tiles_N = tl.cdiv(N, BLOCK_SIZE_N)
    total_tiles = tiles_M * tiles_N

    if pid < total_tiles * SPLITK:
        tile_idx = pid % total_tiles
        iters_per_cta = tl.cdiv(K, BLOCK_SIZE_K * SPLITK)
        iter_begin = pid // total_tiles * iters_per_cta
        iter_end = iter_begin + iters_per_cta
        k_idx = tl.cdiv(iter_begin, iters_per_cta)
        c_ptr_remote = tl.load(c_ptr + cur_rank, cache_modifier=".cg").to(tl.pointer_type(tl.float32)) # new added
        awq_gemm_kernel_inner(a_ptr, b_ptr, c_ptr_remote, zeros_ptr, scales_ptr, tile_idx, k_idx, iter_begin, iter_end,
                            M, N, N2, K, K2, False, GROUP_SIZE, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
                            NUM_GROUPS, USE_REDUCE_KERNEL)
        # set barriers
        tile_idx_m = tile_idx // tiles_N
        tile_idx_n = tile_idx % tiles_N
        offset = total_tiles * k_idx + tiles_N * tile_idx_m + tile_idx_n
        barrier_ptr_remote = tl.load(barrier_ptr + cur_rank).to(tl.pointer_type(tl.int32)) # new added
        tl.store(barrier_ptr_remote + offset, 1)
    else:
        pid = pid - total_tiles * SPLITK
        # reduce kernel
        pid_m = pid // tiles_N
        pid_n = pid % tiles_N

        offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
        offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
        # compiler hints
        offs_m = tl.max_contiguous(tl.multiple_of(offs_m, BLOCK_SIZE_M), BLOCK_SIZE_M)
        offs_n = tl.max_contiguous(tl.multiple_of(offs_n, BLOCK_SIZE_N), BLOCK_SIZE_N)

        mask_m = offs_m < M
        mask_n = offs_n < N
        mask = mask_m[:, None] & mask_n[None, :]

        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        # reduce split k
        # for batch_idx in range(SPLITK):
        #     batch_offset = batch_idx * M * N
        #     input_offsets = batch_offset + offs_m[:, None] * N + offs_n[None, :]
        #     # wait barrier
        #     offset = total_tiles * batch_idx + pid_m * tiles_N + pid_n
        #     while tl.load(barrier_ptr + offset, cache_modifier=".cv", volatile=True) != 1:
        #         pass
        #     input_data = tl.load(c_ptr + input_offsets, mask=mask, other=0.0)
        #     acc += input_data

        for i in tl.range(world_size, flatten=True):
            # swizzle (always compute cur_rank 1st)
            rank_id = (i + cur_rank) % world_size
            # remote_barrier_ptr = libshmem_device.remote_ptr(barrier_ptr, rank_id).to(tl.pointer_type(tl.int32))
            # remote_c_ptr = libshmem_device.remote_ptr(c_ptr, rank_id).to(tl.pointer_type(c_ptr.type.element_ty))
            remote_barrier_ptr = tl.load(barrier_ptr + rank_id).to(tl.pointer_type(tl.int32))
            remote_c_ptr = tl.load(c_ptr + rank_id).to(tl.pointer_type(tl.float32))
            for batch_idx in range(SPLITK):
                batch_offset = batch_idx * M * N
                input_offsets = batch_offset + offs_m[:, None] * N + offs_n[None, :]
                # wait barrier
                offset = total_tiles * batch_idx + pid_m * tiles_N + pid_n
                while tl.load(remote_barrier_ptr + offset, cache_modifier=".cv", volatile=True) != 1:
                    pass
                input_data = tl.load(remote_c_ptr + input_offsets, mask=mask, other=0.0, volatile=True)
                acc += input_data

        output_offsets = offs_m[:, None] * N + offs_n[None, :]
        acc_f16 = acc.to(tl.float16)
        tl.store(out_ptr + output_offsets, acc_f16, mask=mask)

@triton.jit
def awq_gemm_scatter_kernel_splitk(
    a_ptr, b_ptr,
    scatter_ptr,
    zeros_ptr, scales_ptr,
    M, N, N2, K, K2,
    cur_rank,
    world_size,
    GROUP_SIZE: tl.constexpr,
    NUM_CUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    SPLITK: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
    USE_REDUCE_KERNEL: tl.constexpr):

    pid = tl.program_id(axis=0)
    tiles_M = tl.cdiv(M, BLOCK_SIZE_M)
    tiles_N = tl.cdiv(N, BLOCK_SIZE_N)
    total_tiles = tiles_M * tiles_N
    iters_per_cta = tl.cdiv(K, BLOCK_SIZE_K * SPLITK)

    tile_idx = pid % total_tiles
    iter_begin = pid // total_tiles * iters_per_cta
    iter_end = iter_begin + iters_per_cta
    k_idx = tl.cdiv(iter_begin, iters_per_cta)
    
    # TODO: expand this func
    awq_gemm_scatter_kernel_inner(
        a_ptr, b_ptr, scatter_ptr,
        zeros_ptr, scales_ptr, tile_idx, k_idx, iter_begin, iter_end,
        M, N, N2, K, K2, False,
        cur_rank,
        world_size,
        GROUP_SIZE, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K,
        NUM_GROUPS, USE_REDUCE_KERNEL)

'''
@triton.autotune(
    configs=[
        triton.Config({
        }, num_warps=num_warps, num_stages=num_stages)
        for num_warps in [1, 2, 4, 8, 16] for num_stages in [1, 2]
    ],
    key=["M", "K", "N", "GROUP_SIZE", "BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K", "SCHEDULER", "SPLITK"],
    perf_debug=True,
    #enable=int(os.getenv("TRITON_DO_AUTOTUNING", 0)) == 1,
    #prune_configs_by={
    #    "early_config_prune": lambda configs, nargs, **kwargs: [
    #        config for config in configs
    #        # SCHEDULE=1 代表 STREAMK，不需要遍历那么多 SPLITK 的值
    #        if config.all_kwargs()["SCHEDULER"] == 0 or (config.all_kwargs()["SCHEDULER"] == 1 and config.all_kwargs()["SPLITK"] == 1)
    #    ]
    #} 
)
@triton.heuristics(values={
    "NUM_GROUPS": lambda args: triton.cdiv(args["BLOCK_SIZE_K"], args["GROUP_SIZE"])
})
'''
@triton.jit
def awq_gemm_kernel(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2,
                    GROUP_SIZE: tl.constexpr, NUM_CUS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr,
                    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
                    NUM_GROUPS: tl.constexpr,
                    DP_TILES: tl.constexpr, DANGLING_TILES: tl.constexpr,
                    SPLITK: tl.constexpr, SCHEDULER: tl.constexpr, USE_REDUCE_KERNEL: tl.constexpr):
        if SCHEDULER == 0:
            return awq_gemm_kernel_splitk(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2, GROUP_SIZE, NUM_CUS, BLOCK_SIZE_M,
                                          BLOCK_SIZE_N, BLOCK_SIZE_K,
                                          SPLITK, NUM_GROUPS, USE_REDUCE_KERNEL)
        else:
            return awq_gemm_kernel_streamk(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2, GROUP_SIZE, NUM_CUS, BLOCK_SIZE_M,
                                           BLOCK_SIZE_N, BLOCK_SIZE_K,
                                           DP_TILES, DANGLING_TILES, NUM_GROUPS, USE_REDUCE_KERNEL)

@triton.jit
def awq_gemm_kernel_fused(
    a_ptr, b_ptr, c_ptr,
    zeros_ptr,
    scales_ptr,
    out_ptr, barrier_ptr,
    M, N, N2, K, K2,
    cur_rank, world_size,
    GROUP_SIZE: tl.constexpr,
    NUM_CUS: tl.constexpr,
    # NUM_GEMM_CUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
    DP_TILES: tl.constexpr, DANGLING_TILES: tl.constexpr,
    SPLITK: tl.constexpr, SCHEDULER: tl.constexpr,
    USE_REDUCE_KERNEL: tl.constexpr):
        if SCHEDULER == 0:
            # return awq_gemm_kernel_splitk_fused(
            #     a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr,
            #     out_ptr, barrier_ptr,
            #     M, N, N2, K, K2,
            #     GROUP_SIZE,
            #     NUM_CUS,
            #     # NUM_GEMM_CUS,
            #     BLOCK_SIZE_M,
            #     BLOCK_SIZE_N, BLOCK_SIZE_K,
            #     SPLITK, NUM_GROUPS, USE_REDUCE_KERNEL)
            return awq_gemm_allreduce_kernel_splitk(
                a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr,
                out_ptr, barrier_ptr,
                M, N, N2, K, K2,
                cur_rank, world_size,
                GROUP_SIZE,
                NUM_CUS,
                # NUM_GEMM_CUS,
                BLOCK_SIZE_M,
                BLOCK_SIZE_N, BLOCK_SIZE_K,
                SPLITK, NUM_GROUPS, USE_REDUCE_KERNEL)
        else:
            # TODO: to be supported
            assert False
            return awq_gemm_kernel_streamk(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2, GROUP_SIZE, NUM_CUS, BLOCK_SIZE_M,
                                           BLOCK_SIZE_N, BLOCK_SIZE_K,
                                           DP_TILES, DANGLING_TILES, NUM_GROUPS, USE_REDUCE_KERNEL)

# refer from dist/04-tutorial
@triton.jit
def tile_id_to_index_range(
    tile_id,
    M,
    N,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    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 = tile_id // 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)

    tile_in_group = tile_id % num_pid_in_group
    pid_m = first_pid_m + (tile_in_group % group_size_m)
    pid_n = tile_in_group // group_size_m

    rm_start = pid_m * BLOCK_SIZE_M
    rn_start = pid_n * BLOCK_SIZE_N

    # clamp to the maximum valid index (M-1, N-1)
    max_m = M - 1
    max_n = N - 1

    # generate indices
    rm = rm_start + tl.arange(0, BLOCK_SIZE_M)
    rn = rn_start + tl.arange(0, BLOCK_SIZE_N)

    rm = tl.minimum(rm, max_m)
    rn = tl.minimum(rn, max_n)

    return rm, rn, rm_start, rn_start

@triton.jit
def offset_for_tile(local_tile_id, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M, M_local, N_local):
    rm, rn, rm_start, rn_start = tile_id_to_index_range(
        local_tile_id, M_local, N_local, BLOCK_SIZE_M, BLOCK_SIZE_N, GROUP_SIZE_M
    )
    c_mask = (rm[:, None] < M_local) & (rn[None, :] < N_local)
    return rm, rn, c_mask, rm_start, rn_start

@triton.jit
def extract_submask_and_offset(
    M_local,
    N_local,
    rm,
    rn,
    mask,
    rm_start,
    rn_start,
    start_row,
    start_col,
    SUB_BLOCK_SIZE_M: tl.constexpr,
    SUB_BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    stride_cm_local: tl.constexpr,
    stride_cn_local: tl.constexpr,
):
    # Create indices for the sub-block
    sub_rm = tl.arange(0, SUB_BLOCK_SIZE_M) + start_row
    sub_rn = tl.arange(0, SUB_BLOCK_SIZE_N) + start_col

    # clamp to the maximum valid index (M-1, N-1)
    sub_rm = tl.minimum(sub_rm, M_local - 1)
    sub_rn = tl.minimum(sub_rn, N_local - 1)

    # Create a 2D grid of indices for the sub-block
    sub_rm_2d = sub_rm[:, None]  # Shape: (SUB_BLOCK_SIZE_M, 1)
    sub_rn_2d = sub_rn[None, :]  # Shape: (1, SUB_BLOCK_SIZE_N)

    # Compute the sub-mask
    sub_mask = (sub_rm_2d < BLOCK_SIZE_M) & (sub_rn_2d < BLOCK_SIZE_N)

    # Compute the sub-offset relative to the start of the tile
    sub_offset = ((rm_start + sub_rm_2d) * stride_cm_local) + ((rn_start + sub_rn_2d) * stride_cn_local)

    return sub_mask, sub_offset

# GROUP_SIZE_M swizzle
@triton.jit
def _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M):
    group_id = tile_id // 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 + (tile_id % group_size_m)
    pid_m = first_pid_m + ((tile_id % num_pid_in_group) % group_size_m)
    pid_n = (tile_id % num_pid_in_group) // group_size_m
    return pid_m, pid_n

@triton.jit
def awq_gemm_allreduce_oneshot_kernel(
    shm_ctx,
    a_ptr, b_ptr, c_ptr,
    scatter_bufs_ptr, barrier_bufs_ptr,
    zeros_ptr, scales_ptr,
    M, N, N2, K, K2,
    cur_rank, world_size,
    GROUP_SIZE: tl.constexpr,
    NUM_CUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    SPLITK: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
    USE_REDUCE_KERNEL: tl.constexpr,
    DP_TILES: tl.constexpr, DANGLING_TILES: tl.constexpr, SCHEDULER: tl.constexpr,):

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

    start_pid = tl.program_id(axis=0)

    tiles_M = tl.cdiv(M, BLOCK_SIZE_M)
    tiles_N = tl.cdiv(N, BLOCK_SIZE_N)
    total_tiles = tiles_M * tiles_N
    # in case split K > 1
    # tile_idx = pid % total_tiles

    iters_per_cta = tl.cdiv(K, BLOCK_SIZE_K)
    iter_begin = start_pid // total_tiles * iters_per_cta
    # iter_end = iter_begin + iters_per_cta
    # k_idx = tl.cdiv(iter_begin, iters_per_cta)
    k_idx = 0
    loop_k = tl.cdiv(K, BLOCK_SIZE_K)

    # tl.assume(tile_idx >= 0)
    # tl.assume(k_idx >= 0)
    tl.assume(iter_begin >= 0)
    # tl.assume(iter_end >= 0)
    tl.assume(M > 0)
    tl.assume(N > 0)
    tl.assume(K > 0)
    tl.assume(K2 > 0)
    tl.assume(N2 > 0)

    num_tile_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_tile_n = tl.cdiv(N, BLOCK_SIZE_N)

    # group size M swizzle
    # num_pid_in_group = GROUP_SIZE_M * num_tile_n
    # tile_idx_m, tile_idx_n = _compute_pid(start_pid, num_pid_in_group, num_tile_m, GROUP_SIZE_M)

    # no swizzle impl
    tile_idx_m = start_pid // num_tile_n
    tile_idx_n = start_pid % num_tile_n

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

    shifts = tl.arange(0, 2) * 4

    bshifts = tl.broadcast_to(shifts[None, :], (BLOCK_SIZE_N * (BLOCK_SIZE_K // 2), 2))
    bshifts = tl.reshape(bshifts, (BLOCK_SIZE_N, BLOCK_SIZE_K))

    # Offsets and masks.
    offsets_am = tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    masks_am = offsets_am < M

    #offsets_zn = tile_idx_n * (BLOCK_SIZE_N // 2) + tl.arange(0, BLOCK_SIZE_N // 2)
    offsets_zn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offsets_zn = offsets_zn // 2
    masks_zn = offsets_zn < N2

    offsets_bn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    masks_bn = offsets_bn < N

    offsets_sn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    masks_sn = offsets_sn < N

    offsets_ak = iter_begin * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
    offsets_a = K * offsets_am[:, None] + offsets_ak[None, :]

    offsets_bk = iter_begin * (BLOCK_SIZE_K // 2) + tl.arange(0, BLOCK_SIZE_K // 2)
    offsets_b = K2 * offsets_bn[:, None] + offsets_bk[None, :]
    zshifts = (offsets_bn[:, None] % 2) * 4 # [N, 1]
    zshifts = zshifts.T

    a_ptrs = a_ptr + offsets_a
    b_ptrs = b_ptr + offsets_b
    # for k in range(iter_end - iter_begin):
    for k in range(0, loop_k):
        masks_ak = offsets_ak < K
        masks_bk = offsets_bk < K2
        masks_a = masks_am[:, None] & masks_ak[None, :]
        masks_b = masks_bn[:, None] & masks_bk[None, :]
        other_bzs = 0.0
        a = tl.load(a_ptrs, mask=masks_a, other=0.)
        b = tl.load(b_ptrs, masks_b, other_bzs) #[N, K//2]
        b = tl.interleave(b, b) # [N, K]

        # Dequantize b.
        offsets_szk = ((BLOCK_SIZE_K * k) // GROUP_SIZE + tl.arange(0, NUM_GROUPS))
        masks_szk = offsets_szk < K // GROUP_SIZE
        masks_z = masks_szk[:, None] & masks_zn[None, :]
        masks_s = masks_szk[:, None] & masks_sn[None, :]

        offsets_z = N2 * offsets_szk[:, None] + offsets_zn[None, :]
        zeros_ptrs = zeros_ptr + offsets_z
        zeros = tl.load(zeros_ptrs, mask=masks_z, other=other_bzs) # [K//G, N]

        offsets_s = N * offsets_szk[:, None] + offsets_sn[None, :]
        scales_ptrs = scales_ptr + offsets_s
        scales = tl.load(scales_ptrs, mask=masks_s, other=other_bzs) # [K//G, N]

        if NUM_GROUPS == 1:
            # Original efficient implementation for single group
            zeros = tl.broadcast_to(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            scales = tl.broadcast_to(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))
        else:
            zeros = tl.broadcast_to(zeros[:, None, :], (NUM_GROUPS, GROUP_SIZE, BLOCK_SIZE_N))
            scales = tl.broadcast_to(scales[:, None, :], (NUM_GROUPS, GROUP_SIZE, BLOCK_SIZE_N))
            ## Reshape back to (BLOCK_SIZE_K, N)
            zeros = tl.reshape(zeros, (BLOCK_SIZE_K, BLOCK_SIZE_N))
            scales = tl.reshape(scales, (BLOCK_SIZE_K, BLOCK_SIZE_N))

        b = (b >> bshifts) & 0xF
        b = b.T
        zeros = (zeros >> zshifts) & 0xF
        b = (b - zeros) * scales
        b = b.to(a_ptr.type.element_ty)

        # Accumulate results.
        accumulator = tl.dot(a, b, accumulator, out_dtype=tl.float32)

        offsets_ak += BLOCK_SIZE_K
        offsets_bk += BLOCK_SIZE_K // 2
        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K // 2

    # c = accumulator.to(c_ptr.type.element_ty)
    c = accumulator.to(tl.float32)
    
    offs_cm = tile_idx_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = tile_idx_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    # compiler hints
    offs_cm = tl.max_contiguous(tl.multiple_of(offs_cm, BLOCK_SIZE_M), BLOCK_SIZE_M)
    offs_cn = tl.max_contiguous(tl.multiple_of(offs_cn, BLOCK_SIZE_N), BLOCK_SIZE_N)
    offs_c = M * N * k_idx + N * offs_cm[:, None] + offs_cn[None, :]
    # c_ptrs = c_ptr + 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
    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, world_size):
        remote = (cur_rank + i) % 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 + start_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(local_base_ptr + start_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
    for i in range(1, world_size):
        rank_id = (cur_rank + i) % 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)
    tl.store(c_ptr + offs_c, acc_f16, mask=c_mask, cache_modifier=".wt")

# Fused-Sequential impl of gemm-scatter
# Each rank produce the partial gemm output and push to the peer rank
@triton.jit
def awq_gemm_scatter_kernel(
    a_ptr, b_ptr, scatter_ptr,
    zeros_ptr,
    scales_ptr,
    M, N, N2, K, K2,
    cur_rank,
    world_size,
    GROUP_SIZE: tl.constexpr,
    NUM_CUS: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
    NUM_GROUPS: tl.constexpr,
    DP_TILES: tl.constexpr, DANGLING_TILES: tl.constexpr,
    SPLITK: tl.constexpr, SCHEDULER: tl.constexpr,
    USE_REDUCE_KERNEL: tl.constexpr):
        if SCHEDULER == 0:
            return awq_gemm_scatter_kernel_splitk(
                a_ptr, b_ptr, scatter_ptr,
                zeros_ptr, scales_ptr,
                M, N, N2, K, K2,
                cur_rank,
                world_size,
                GROUP_SIZE,
                NUM_CUS,
                # NUM_GEMM_CUS,
                BLOCK_SIZE_M,
                BLOCK_SIZE_N, BLOCK_SIZE_K,
                SPLITK, NUM_GROUPS, USE_REDUCE_KERNEL)
        else:
            # TODO: to be supported
            assert False
            return awq_gemm_kernel_streamk(a_ptr, b_ptr, c_ptr, zeros_ptr, scales_ptr, M, N, N2, K, K2, GROUP_SIZE, NUM_CUS, BLOCK_SIZE_M,
                                           BLOCK_SIZE_N, BLOCK_SIZE_K,
                                           DP_TILES, DANGLING_TILES, NUM_GROUPS, USE_REDUCE_KERNEL)


# qweights - [N     , K // 2], int8
# scales   - [K // G, N     ], float16
# zeros    - [K // G, N // 2], int8
# result   - [K, N], float16
def awq_dequantize_triton(qweight: torch.Tensor,
                          scales: torch.Tensor,
                          zeros: torch.Tensor,
                          **kwargs) -> torch.Tensor:
    N = qweight.shape[0]
    K = qweight.shape[1] * 2
    group_size = K // scales.shape[0]

    assert K > 0 and N > 0
    assert scales.shape[0] == K // group_size and scales.shape[1] == N
    assert zeros.shape[0] == K // group_size and zeros.shape[1] == N // 2
    assert group_size <= K
    assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K

    configs = {
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 64, 
        "num_warps": 4, 
        "num_stages": 1
    }

    result = torch.empty(K, N, device=qweight.device, dtype=scales.dtype)

    grid = lambda META: (
        triton.cdiv(N, META['BLOCK_SIZE_N']),
        triton.cdiv(K, META['BLOCK_SIZE_K']),
    )
    awq_dequantize_kernel[grid](qweight,
                                scales,
                                zeros,
                                result,
                                N,
                                N//2,
                                K,
                                K//2,
                                group_size,
                                **configs
                                )

    return result

@functools.lru_cache
def get_w4a16_awq_gemm_config_filepath(N: int, K: int, GROUP_SIZE: int, **kwargs) -> str:
    device_name = arch_info.get_device()
    if device_name.lower().startswith("bw"):
        device_name = "BW200"
    if "k100" in device_name.lower():
        device_name = "K100_AI"
    json_file_name = f"awq_gemm_N={N},K={K},device_name={device_name},dtype=w4a16,group_size={GROUP_SIZE}.json"

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

@functools.lru_cache
def get_w4a16_awq_gemm_configs(
    N: int, K: int, GROUP_SIZE: 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.
    """
    config_file_path = get_w4a16_awq_gemm_config_filepath(N, K, GROUP_SIZE)
    if os.path.exists(config_file_path):
        # print(f"Teng config_file_path: {config_file_path}")
        with open(config_file_path) as f:
            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(
            f"\nUsing default W4A16 AWQ GEMM kernel config. Performance might "
            f"be sub-optimal! Config file not found at {config_file_path}")
    return None

# The inference function
# input   - [m, k]
# qweight - [n, k // 2]
# qzeros  - [k//g, n//2]
# scales  - [k//g, n]
def gemm_a16w4(input: torch.tensor,
               qweight: torch.tensor,
               scales: torch.tensor,
               qzeros: torch.tensor,
               not_used_placeholder: int = 0,
               configs: Optional[Dict] = None) -> torch.tensor:

    M, K = input.shape
    N = qweight.shape[0] # (N, K//2)
    group_size = K // qzeros.shape[0]

    default_config = {
        "BLOCK_SIZE_M": 16,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 32,
        "SCHEDULER": 0,
        "SPLITK": 1,
        "D_SHAPE": (M, N),
        "D_DTYPE": 16,
        "DP_TILES": 0,
        "DANGLING_TILES": 0,
        "NUM_CUS": 0,
        "NUM_CUS_STREAMK": 0,
        "NUM_GROUPS":(32 + group_size - 1) // group_size,
        "USE_REDUCE_KERNEL": False
    }
    if configs is None:
        configs = get_w4a16_awq_gemm_configs(N, K, group_size)
    config = configs[min(configs.keys(), key=lambda x: abs(x - M))] if configs else default_config
    # Make sure not getting this wrong from other configs
    d_shape = list(config["D_SHAPE"])
    d_shape[-2] = M
    config["D_SHAPE"] = d_shape
    # if config["SPLITK"] > 1 and config["USE_REDUCE_KERNEL"]:
    #     return awq_gemm_triton_fused_impl(
    #         input, qweight, scales, qzeros, config, config.copy(), awq_gemm_kernel_fused)

    return awq_gemm_triton_impl(input, qweight, scales, qzeros, config, config.copy(), awq_gemm_kernel)

def awq_gemm_triton_fused_impl(
    input: torch.tensor,
    qweight: torch.tensor,
    scales: torch.tensor,
    qzeros: torch.tensor,
    config: Dict,
    cfg4kernel: Dict,
    func) -> torch.tensor:
    M, K = input.shape
    N = qweight.shape[0] # (N, K//2)
    assert(qweight.is_contiguous())
    group_size = qweight.shape[1] * 2 // qzeros.shape[0]

    assert N > 0 and K > 0 and M > 0
    assert qweight.shape[1] == K // 2 and qweight.shape[0] == N
    assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 2
    assert scales.shape[0] == K // group_size and scales.shape[1] == N
    assert group_size <= K
    assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K

    num_cus = config["NUM_CUS"]
    d_shape = config["D_SHAPE"]
    d_dtype = config["D_DTYPE"]
    d_dtype = torch.float16 if d_dtype == 16 else torch.float32

    def grid(META):
        # tiles_M = (M + META["BLOCK_SIZE_M"] - 1) // META["BLOCK_SIZE_M"]
        # tiles_N = (N + META["BLOCK_SIZE_N"] - 1) // META["BLOCK_SIZE_N"]
        tiles_M = triton.cdiv(M, META["BLOCK_SIZE_M"])
        tiles_N = triton.cdiv(N, META["BLOCK_SIZE_N"])
        total_tiles = tiles_M * tiles_N
        if META["SCHEDULER"] == 0:
            # dp or splitk
            # add extra total_tiles for reduction
            return (total_tiles * META["SPLITK"] + total_tiles,)
        else:
            # TODO: not supported yet
            # streamk
            return (META["DP_TILES"] + config["NUM_CUS_STREAMK"],)

    result = torch.zeros(d_shape, dtype=d_dtype, device=input.device)

    cfg4kernel.pop("D_SHAPE", None)
    cfg4kernel.pop("D_DTYPE", None)
    cfg4kernel.pop("NUM_CUS_STREAMK", None)

    fn = func[grid]

    if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
        fn = partial(func.warmup, grid=grid)
    
    total_tiles_splitk = \
        triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv(N, config["BLOCK_SIZE_N"]) * config["SPLITK"]
    final_result = torch.zeros((M, N), dtype=torch.float16, device=input.device)
    barrier = torch.zeros((total_tiles_splitk, ), dtype=torch.int, device=input.device)

    fn(input,
       qweight,
       result,
       qzeros,
       scales,
       final_result, # new added
       barrier,      # new added
       M,
       N,
       N//2,
       K,
       K//2,
       group_size,
       **cfg4kernel)

    if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
        return

    return final_result

def awq_gemm_triton_impl(input: torch.tensor,
                        qweight: torch.tensor,
                        scales: torch.tensor,
                        qzeros: torch.tensor,
                        config: Dict,
                        cfg4kernel: Dict,
                        func) -> torch.tensor:
    M, K = input.shape
    N = qweight.shape[0] # (N, K//2)
    assert(qweight.is_contiguous())
    group_size = qweight.shape[1] * 2 // qzeros.shape[0]

    assert N > 0 and K > 0 and M > 0
    assert qweight.shape[1] == K // 2 and qweight.shape[0] == N
    assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 2
    assert scales.shape[0] == K // group_size and scales.shape[1] == N
    assert group_size <= K
    assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K
    
    # print("SPLITK:", config["SPLITK"])
    # print("REDUCE:", config["USE_REDUCE_KERNEL"])

    num_cus = config["NUM_CUS"]
    d_shape = config["D_SHAPE"]
    d_dtype = config["D_DTYPE"]
    d_dtype = torch.float16 if d_dtype == 16 else torch.float32

    #curr_num_cus = torch.cuda.get_device_properties("cuda").multi_processor_count
    #if num_cus > 0 and num_cus != curr_num_cus:
    #    print("AWQ_GEMM config tuned based on num_cus={num_cus}, but now running on num_cus={curr_num_cus}, may lead to bad performance!")

    def grid(META):
        tiles_M = (M + META["BLOCK_SIZE_M"] - 1) // META["BLOCK_SIZE_M"]
        tiles_N = (N + META["BLOCK_SIZE_N"] - 1) // META["BLOCK_SIZE_N"]
        total_tiles = tiles_M * tiles_N
        if META["SCHEDULER"] == 0:
            # dp or splitk
            return (total_tiles * META["SPLITK"],)
        else:
            # streamk
            return (META["DP_TILES"] + config["NUM_CUS_STREAMK"],)

    result = torch.zeros(d_shape, dtype=d_dtype, device=input.device)

    cfg4kernel.pop("D_SHAPE", None)
    cfg4kernel.pop("D_DTYPE", None)
    cfg4kernel.pop("NUM_CUS_STREAMK", None)

    fn = func[grid]

    if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
        fn = partial(func.warmup, grid=grid)

    fn(input,
       qweight,
       result,
       qzeros,
       scales,
       M,
       N,
       N//2,
       K,
       K//2,
       group_size,
       **cfg4kernel)

    if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
        return

    if result.ndim == 3:
        batch_size = result.shape[0]
        final_result = torch.zeros((M, N), dtype=torch.float16, device=input.device)
        awq_reduce_and_convert_triton(result, final_result, M, N, batch_size)
        return final_result
    else:
        result = result.to(torch.float16)
        return result

# awq gemm fused with allreduce
def gemm_allreduce_a16w4(
    input: torch.tensor,
    qweight: torch.tensor,
    scales: torch.tensor,
    qzeros: torch.tensor,
    barrier_buf: torch.tensor,
    scatter_buf: torch.tensor,
    max_tiles: int,
    tp_group: dist.ProcessGroup,
    shmem_ctx: Optional[np.intp] = None,
    configs: Optional[Dict] = None) -> torch.tensor:

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

    M, K = input.shape
    N = qweight.shape[0] # (N, K//2)
    group_size = K // qzeros.shape[0]

    default_config = {
        "BLOCK_SIZE_M": 16,
        "BLOCK_SIZE_N": 128,
        "BLOCK_SIZE_K": 32,
        "SCHEDULER": 0,
        "SPLITK": 1,
        "D_SHAPE": (M, N),
        "D_DTYPE": 16,
        "DP_TILES": 0,
        "DANGLING_TILES": 0,
        "NUM_CUS": 0,
        "NUM_CUS_STREAMK": 0,
        "NUM_GROUPS":(32 + group_size - 1) // group_size,
        "USE_REDUCE_KERNEL": False
    }
    if configs is None:
        configs = get_w4a16_awq_gemm_configs(N, K, group_size)
    config = configs[min(configs.keys(), key=lambda x: abs(x - M))] if configs else default_config

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

    output, compiled_kernel = awq_gemm_allreduce_triton_oneshot_impl(
        shmem_ctx,
        input, qweight, scales, qzeros,
        barrier_buf,
        scatter_buf,
        cur_rank, world_size,
        config)

    # NOTE: this seems confilct with hip stream with this error when using vLLM:
    #       [nvSHMEM/src/host/coll/barrier/barrier.cpp:21] cuda failed with operation
    #       not permitted when stream is capturing
    #       so do barrier outside in vLLM using vLLM api
    # if world_size > 1:
    #     dist.barrier()
        # pynvshmem.nvshmem_barrier_all()

    # if cur_rank == 0:
    #     logger.info(f"triton kernel regs: {compiled_kernel.n_regs}, spills: {compiled_kernel.n_spills}")
    return output

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

    for i in range(num_ranks):
        while tl.atomic_cas(comm_buf_ptr + i, 1, 0, scope="sys", sem="acquire") != 1:
            pass

    tl.debug_barrier()

def barrier_all_on_stream(
    rank,
    num_ranks,
    sync_buf_ptr,
    stream,
):
    with torch.cuda.stream(stream):
        barrier_all_ipc[(1, )](rank, num_ranks, sync_buf_ptr)


# scatter to other ranks
# scatter buffer
# [a0, # Rank 0 -> [a0 + b0]
#  a1,
#  b0, # Rank 1 -> [a1 + b1]
#  b1]
@triton.jit
def local_reduce_kernel(
    input_ptr,
    output_ptr,
    M, N,
    batch_size,
    cur_rank,
    world_size,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    tl.assume(M >= 0)
    tl.assume(N >= 0)
    tl.assume(batch_size >= 0)
    tl.assume(BLOCK_SIZE_M >= 0)
    tl.assume(BLOCK_SIZE_N >= 0)
    tl.assume(pid_m >= 0)
    tl.assume(pid_n >= 0)
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    
    mask_m = offs_m < (M * world_size)
    mask_n = offs_n < N
    mask = mask_m[:, None] & mask_n[None, :]
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    for batch_idx in range(batch_size):
        batch_offset = batch_idx * M * N * world_size
        input_offsets = batch_offset + offs_m[:, None] * N + offs_n[None, :]
        # reduce the data from each rank
        for rank_id in range(world_size):
            rank_offset = M * N * rank_id
            input_data = tl.load(input_ptr + input_offsets + rank_offset, mask=mask, other=0.0)
            acc += input_data
    
    # store mask
    mask_m = offs_m < M
    mask = mask_m[:, None] & mask_n[None, :]
    
    output_offsets = offs_m[:, None] * N + offs_n[None, :]
    acc_f16 = acc.to(tl.float16)
    tl.store(output_ptr + output_offsets, acc_f16, mask=mask)

def awq_gemm_allreduce_triton_oneshot_impl(
    shm_ctx: np.intp,
    input: torch.tensor,
    qweight: torch.tensor,
    scales: torch.tensor,
    qzeros: torch.tensor,
    barrier_buf: torch.tensor,
    scatter_buf: torch.tensor,
    cur_rank,
    world_size,
    cfg4kernel: Dict) -> torch.tensor:
    M, K = input.shape
    N = qweight.shape[0] # (N, K//2)
    assert(qweight.is_contiguous())
    group_size = qweight.shape[1] * 2 // qzeros.shape[0]

    assert N > 0 and K > 0 and M > 0
    assert qweight.shape[1] == K // 2 and qweight.shape[0] == N
    assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 2
    assert scales.shape[0] == K // group_size and scales.shape[1] == N
    assert group_size <= K
    assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K

    # num_cus = config["NUM_CUS"]
    # d_shape = config["D_SHAPE"]
    # d_dtype = config["D_DTYPE"]
    # d_dtype = torch.float16 if d_dtype == 16 else torch.float32

    cfg4kernel.pop("D_SHAPE", None)
    cfg4kernel.pop("D_DTYPE", None)
    cfg4kernel.pop("NUM_CUS_STREAMK", None)
    cfg4kernel["GROUP_SIZE_M"] = 1

    output = torch.empty((M, N), dtype=torch.float16, device=input.device)
    # grid = lambda META: (min(
    #     META["NUM_CUS"],
    #     triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
    # ), )
    grid = lambda META: (
        triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
    compiled = awq_gemm_allreduce_oneshot_kernel[grid](
        shm_ctx,
        input,
        qweight,
        output,
        scatter_buf,
        barrier_buf,
        qzeros,
        scales,
        M,
        N,
        N//2,
        K,
        K//2,
        cur_rank, world_size,
        group_size,
        **cfg4kernel)
    return output, compiled

def awq_gemm_allreduce_triton_fused_impl(
    input: torch.tensor,
    qweight: torch.tensor,
    scales: torch.tensor,
    qzeros: torch.tensor,
    sync_buf: torch.tensor,
    barrier_buf: torch.tensor,
    scatter_buf: torch.tensor,
    cur_rank,
    world_size,
    config: Dict,
    cfg4kernel: Dict,
    func) -> torch.tensor:
    M, K = input.shape
    N = qweight.shape[0] # (N, K//2)
    assert(qweight.is_contiguous())
    group_size = qweight.shape[1] * 2 // qzeros.shape[0]

    assert N > 0 and K > 0 and M > 0
    assert qweight.shape[1] == K // 2 and qweight.shape[0] == N
    assert qzeros.shape[0] == K // group_size and qzeros.shape[1] == N // 2
    assert scales.shape[0] == K // group_size and scales.shape[1] == N
    assert group_size <= K
    assert group_size in AWQ_TRITON_SUPPORTED_GROUP_SIZES or group_size == K

    num_cus = config["NUM_CUS"]
    d_shape = config["D_SHAPE"]
    d_dtype = config["D_DTYPE"]
    d_dtype = torch.float16 if d_dtype == 16 else torch.float32

    def grid(META):
        # tiles_M = (M + META["BLOCK_SIZE_M"] - 1) // META["BLOCK_SIZE_M"]
        # tiles_N = (N + META["BLOCK_SIZE_N"] - 1) // META["BLOCK_SIZE_N"]
        tiles_M = triton.cdiv(M, META["BLOCK_SIZE_M"])
        tiles_N = triton.cdiv(N, META["BLOCK_SIZE_N"])
        total_tiles = tiles_M * tiles_N
        if META["SCHEDULER"] == 0:
            # dp or splitk
            # add extra total_tiles for reduction
            return (total_tiles * META["SPLITK"] + total_tiles,)
        else:
            # TODO: not supported yet
            assert False
            # streamk
            return (META["DP_TILES"] + config["NUM_CUS_STREAMK"],)

    cfg4kernel.pop("D_SHAPE", None)
    cfg4kernel.pop("D_DTYPE", None)
    cfg4kernel.pop("NUM_CUS_STREAMK", None)

    fn = func[grid]

    if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
        fn = partial(func.warmup, grid=grid)

    # result = torch.zeros(d_shape, dtype=d_dtype, device=input.device)
    final_result = torch.zeros((M, N), dtype=d_dtype, device=input.device)

    # with torch.cuda.stream(current_stream):
    fn(input,
       qweight,
       scatter_buf, #result,
       qzeros,
       scales,
       final_result,
       barrier_buf,
       M,
       N,
       N//2,
       K,
       K//2,
       cur_rank,
       world_size,
       group_size,
       **cfg4kernel)

    # barrier_all_ipc[(1, )](cur_rank, world_size, sync_buf)
    # pynvshmem.nvshmem_barrier_all()

    # if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
    #     return
    
    # barrier_all_on_stream(cur_rank, world_size, sync_buf, current_stream)
    
    # # a seperate local reduce
    # grid = lambda META: (
    #     triton.cdiv(M, META['BLOCK_SIZE_M']),
    #     triton.cdiv(N, META['BLOCK_SIZE_N']),
    # )
    # BLOCK_SIZE_M = config["BLOCK_SIZE_M"] #32
    # BLOCK_SIZE_N = config["BLOCK_SIZE_N"] #128
    # num_warps = 16
    # with torch.cuda.stream(current_stream):
    #     local_reduce_kernel[grid](
    #         scatter_buf,
    #         final_result,
    #         M, N,
    #         config["SPLITK"],
    #         cur_rank,
    #         world_size,
    #         BLOCK_SIZE_M,
    #         BLOCK_SIZE_N,
    #         num_warps=num_warps
    #     )

    return final_result

# The tuning functions below
def prune_configs(configs, nargs, **kwargs):

    def _ceil_div(x, y):
        return (x + y - 1) // y

    def _prune(config):
        _config = config.all_kwargs()
        all_kwargs = {**_config, **kwargs, **nargs}

        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = all_kwargs["BLOCK_SIZE_M"], all_kwargs["BLOCK_SIZE_N"], all_kwargs["BLOCK_SIZE_K"]
        num_stages = all_kwargs["num_stages"]

        if num_stages > 1 and (BLOCK_SIZE_M * BLOCK_SIZE_K + BLOCK_SIZE_K * BLOCK_SIZE_N) * 2 > 16384:
            return True

    remained = [c for c in configs if not _prune(c)]
    return remained

def update_config(M, K, N, G, cfg):
    config = cfg.copy()
    # 根据基本的配置计算其他参数，一则用于 launch，一则避免 kernel 内重复计算
    config["NUM_GROUPS"] = (config["BLOCK_SIZE_K"] + G - 1) // G
    if config["SCHEDULER"] == 0 and config["SPLITK"] == 1:
        # dp
        config["DP_TILES"] = 0
        config["DANGLING_TILES"] = 0
        config["D_SHAPE"] = (M, N)
        config["D_DTYPE"] = 16
    elif config["SCHEDULER"] == 0 and config["SPLITK"] > 1:
        # splitk
        config["DP_TILES"] = 0
        config["DANGLING_TILES"] = 0
        config["D_DTYPE"] = 32
        config["D_SHAPE"] = (config["SPLITK"], M, N) if config["USE_REDUCE_KERNEL"] else (M, N)
    else:
        # streamk
        tiles_M = (M + config["BLOCK_SIZE_M"] - 1) // config["BLOCK_SIZE_M"]
        tiles_N = (N + config["BLOCK_SIZE_N"] - 1) // config["BLOCK_SIZE_N"]
        total_tiles = tiles_M * tiles_N
        dangling_tiles = max(0, total_tiles - config["NUM_CUS"]) % config["NUM_CUS"]
        dp_tiles = total_tiles - dangling_tiles
        if dangling_tiles == 0:
            # redirect to dp
            config["SCHEDULER"] = 0
            config["SPLITK"] = 1
            config["USE_REDUCE_KERNEL"] = 0
            config["DP_TILES"] = 0
            config["DANGLING_TILES"] = 0
            config["D_SHAPE"] = (M, N)
            config["D_DTYPE"] = 16
        else:
            # still streamk
            config["DP_TILES"] = dp_tiles
            config["DANGLING_TILES"] = dangling_tiles
            iters_per_tile = (K + config["BLOCK_SIZE_K"] - 1) // config["BLOCK_SIZE_K"]
            dangling_iters = iters_per_tile * config["DANGLING_TILES"]
            dangling_iters_per_cu = (dangling_iters + config["NUM_CUS"] - 1) // config["NUM_CUS"]
            num_cus_streamk = (dangling_iters + dangling_iters_per_cu - 1) // dangling_iters_per_cu
            config["NUM_CUS_STREAMK"] = num_cus_streamk
            num_cus_per_dangling_tile = (iters_per_tile + dangling_iters_per_cu - 1) // dangling_iters_per_cu + 1
            config["D_DTYPE"] = 32
            config["D_SHAPE"] = (num_cus_per_dangling_tile, M, N) if config["USE_REDUCE_KERNEL"] else (M, N)

    return config

'''
@triton.autotune(
    configs=[
        triton.Config({
            "BLOCK_SIZE_M": M,
            "BLOCK_SIZE_N": N,
        }, num_warps=num_warps, num_stages=num_stages)
        for M in [128, 64, 32,  16] for N in [512, 128, 64, 32, 16]\
        for num_warps in [1, 2, 4, 8, 16] for num_stages in [1, 2]
    ],
    key=["M", "N", "batch_size"],
    perf_debug=True
)
'''

@triton.jit
def awq_reduce_and_convert_kernel(
    input_ptr,
    output_ptr,
    M,
    N,
    batch_size,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    tl.assume(M >= 0)
    tl.assume(N >= 0)
    tl.assume(batch_size >= 0)
    tl.assume(BLOCK_SIZE_M >= 0)
    tl.assume(BLOCK_SIZE_N >= 0)
    tl.assume(pid_m >= 0)
    tl.assume(pid_n >= 0)
    
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    
    mask_m = offs_m < M
    mask_n = offs_n < N
    mask = mask_m[:, None] & mask_n[None, :]
    
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    for batch_idx in range(batch_size):
        batch_offset = batch_idx * M * N
        input_offsets = batch_offset + offs_m[:, None] * N + offs_n[None, :]
        input_data = tl.load(input_ptr + input_offsets, mask=mask, other=0.0)
        acc += input_data
    
    output_offsets = offs_m[:, None] * N + offs_n[None, :]
    acc_f16 = acc.to(tl.float16)
    tl.store(output_ptr + output_offsets, acc_f16, mask=mask)

def awq_reduce_and_convert_triton(
    input_tensor: torch.Tensor,
    output_tensor: torch.Tensor,
    M: int,
    N: int,
    batch_size: int = 1
) -> None:

    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']),
        triton.cdiv(N, META['BLOCK_SIZE_N']),
    )

    BLOCK_SIZE_M = 32
    BLOCK_SIZE_N = 128
    num_warps = 16

    awq_reduce_and_convert_kernel[grid](
        input_tensor,
        output_tensor,
        M,
        N,
        batch_size,
        BLOCK_SIZE_M,
        BLOCK_SIZE_N,
        num_warps=num_warps
    )
