r"""
Uses a group gemm, instead of a loop over experts.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, Tuple, List, Optional
import math
import einops
import triton
import triton.language as tl

def num_sms():
    # MI300 has 304 CUs
    return 304

@triton.jit
def silu(tile):
    # tl.exp needs float32 or float64.
    return tile * (1.0 / (1.0 + tl.exp(-tile)))

@triton.autotune(
    configs=[
        triton.Config({
            'BLOCK_SIZE_M': 128,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 84,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 128,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 128,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 64,
            'BLOCK_SIZE_N': 64,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 84,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 64,
            'BLOCK_SIZE_N': 64,
            'BLOCK_SIZE_K': 32,
            'NUM_SM': 128,
        }),
        triton.Config({
            'BLOCK_SIZE_M': 128,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 64,
            'NUM_SM': num_sms(),
        }),
        triton.Config({
            'BLOCK_SIZE_M': 64,
            'BLOCK_SIZE_N': 128,
            'BLOCK_SIZE_K': 64,
            'NUM_SM': num_sms(),
        }),
    ],
    key=['group_size'],
)
@triton.jit
def grouped_matmul_kernel(
    # device tensor of matrices pointers
    group_a_ptrs,
    group_b_ptrs,
    group_c_ptrs,
    # device tensor of gemm sizes. its shape is [group_size, 3]
    # dim 0 is group_size, dim 1 is the values of <M, N, K> of each gemm
    group_gemm_sizes,
    # device tensor of leading dimension sizes. its shape is [group_size, 3]
    # dim 0 is group_size, dim 1 is the values of <lda, ldb, ldc> of each gemm
    g_lds,
    # number of gemms
    group_size,
    # number of virtual SM
    NUM_SM: tl.constexpr,
    # tile sizes
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    ACTIVATION: tl.constexpr,
    MULT_C: tl.constexpr
):
    r"""
    This is taken and modified from the triton tutorial. 
      - adds masking to get working with different shapes.
      - fuses the SiLU activation function. 
    """
    tile_idx = tl.program_id(0)
    last_problem_end = 0
    for g in range(group_size):
        # get the gemm size of the current problem
        gm = tl.load(group_gemm_sizes + g * 3)
        gn = tl.load(group_gemm_sizes + g * 3 + 1)
        gk = tl.load(group_gemm_sizes + g * 3 + 2)
        num_m_tiles = tl.cdiv(gm, BLOCK_SIZE_M)
        num_n_tiles = tl.cdiv(gn, BLOCK_SIZE_N)
        num_tiles = num_m_tiles * num_n_tiles
        # iterate through the tiles in the current gemm problem
        while (tile_idx >= last_problem_end and tile_idx < last_problem_end + num_tiles):
            # pick up a tile from the current gemm problem
            k = gk
            lda = tl.load(g_lds + g * 3)
            ldb = tl.load(g_lds + g * 3 + 1)
            ldc = tl.load(g_lds + g * 3 + 2)
            a_ptr = tl.load(group_a_ptrs + g).to(tl.pointer_type(tl.float16))
            b_ptr = tl.load(group_b_ptrs + g).to(tl.pointer_type(tl.float16))
            c_ptr = tl.load(group_c_ptrs + g).to(tl.pointer_type(tl.float16))
            # figure out tile coordinates
            tile_idx_in_gemm = tile_idx - last_problem_end
            tile_m_idx = tile_idx_in_gemm // num_n_tiles
            tile_n_idx = tile_idx_in_gemm % num_n_tiles

            # do regular gemm here
            offs_am = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
            offs_bn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
            offs_k = tl.arange(0, BLOCK_SIZE_K)
            a_ptrs = a_ptr + offs_am[:, None] * lda + offs_k[None, :]
            b_ptrs = b_ptr + offs_k[:, None] * ldb + offs_bn[None, :]
            
            m_mask = offs_am < gm
            n_mask = offs_bn < gn
            
            accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
            for kk in range(0, tl.cdiv(k, BLOCK_SIZE_K)):
                # hint to Triton compiler to do proper loop pipelining
                tl.multiple_of(a_ptrs, [16, 16])
                tl.multiple_of(b_ptrs, [16, 16])
                
                k_mask = kk + offs_k < gk
                a_mask = m_mask[:, None] and k_mask[None, :]
                b_mask = k_mask[:, None] and n_mask[None, :]
                
                a = tl.load(a_ptrs, mask=a_mask)
                b = tl.load(b_ptrs, mask=b_mask)
                accumulator += tl.dot(a, b)
                a_ptrs += BLOCK_SIZE_K
                b_ptrs += BLOCK_SIZE_K * ldb
                
            if ACTIVATION:
                accumulator = silu(accumulator)

            offs_cm = tile_m_idx * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
            offs_cn = tile_n_idx * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
            c_ptrs = c_ptr + ldc * offs_cm[:, None] + offs_cn[None, :]
            c_mask = m_mask[:, None] and n_mask[None, :]

            if MULT_C:
                # maybe this is bad since it's an extra load...
                c = tl.load(c_ptrs, mask=c_mask)
                c *= accumulator.to(tl.float16)
            else:
                c = accumulator.to(tl.float16)
            
            tl.store(c_ptrs, c, mask=c_mask)

            # go to the next tile by advancing NUM_SM
            tile_idx += NUM_SM

        # get ready to go to the next gemm problem
        last_problem_end = last_problem_end + num_tiles

def group_gemm_fn(
        group_A: List[torch.Tensor],
        group_B: List[torch.Tensor],
        activation: bool,
        group_C_init: Optional[torch.Tensor] = None
    ):
    assert len(group_A) == len(group_B)

    DEVICE = 'cuda'
    group_size = len(group_A)

    A_addrs = []
    B_addrs = []
    C_addrs = []
    g_sizes = []
    g_lds = []
    group_C = []
    for i in range(group_size):        
        A = group_A[i]
        B = group_B[i]
        assert A.shape[1] == B.shape[0]
        M, K = A.shape
        K, N = B.shape
        
        A_addrs.append(A.data_ptr())
        B_addrs.append(B.data_ptr())
        
        if group_C_init is None:
            C = torch.empty((M, N), device=DEVICE, dtype=A.dtype)
        else:
            C = group_C_init[i]
            
        group_C.append(C)
        C_addrs.append(C.data_ptr())
        
        g_sizes += [M, N, K]
        g_lds += [A.stride(0), B.stride(0), C.stride(0)]
        
    # note these are device tensors,
    # TODO: does this involve device transfers???
    d_a_ptrs = torch.tensor(A_addrs, device=DEVICE)
    d_b_ptrs = torch.tensor(B_addrs, device=DEVICE)
    d_c_ptrs = torch.tensor(C_addrs, device=DEVICE)
    d_g_sizes = torch.tensor(g_sizes, dtype=torch.int32, device=DEVICE)
    d_g_lds = torch.tensor(g_lds, dtype=torch.int32, device=DEVICE)
    # we use a fixed number of CTA, and it's auto-tunable
    grid = lambda META: (META['NUM_SM'], )
    grouped_matmul_kernel[grid](
        d_a_ptrs,
        d_b_ptrs,
        d_c_ptrs,
        d_g_sizes,
        d_g_lds,
        group_size,
        ACTIVATION=activation,
        MULT_C=group_C_init is not None
    )
    
    return group_C

def get_token_indices(expert_indices, num_experts_per_token):
    flat_expert_indices = expert_indices.view(-1)
    indices = flat_expert_indices.argsort()
    counts = flat_expert_indices.bincount().cpu().numpy() # TODO: why move to cpu?
    tokens_per_expert_range = counts.cumsum()
    token_indices = indices // num_experts_per_token
    return indices, tokens_per_expert_range, token_indices

@torch.compile
def expert_input_permute(
    tokens: torch.Tensor, 
    expert_indices: torch.Tensor, 
    num_experts_per_token: int
) -> List[torch.Tensor]:
    _, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    
    expert_data = []
    for expert_id, end_idx in enumerate(tokens_per_expert_range):
        start_idx = 0 if expert_id == 0 else tokens_per_expert_range[expert_id - 1]
        if start_idx == end_idx:
            # triton group gemm needs len(group_A) == len(group_B),
            # adding an empty list is just a hack to ensure the group sizes are the
            # same. This could/should also be done by not passing in weights if the
            # corresponding expert gets no tokens.
            expert_data.append(torch.empty((0, tokens.size(-1))))
            continue
        
        exp_token_idxs = token_indices[start_idx:end_idx]
        expert_tokens = tokens[exp_token_idxs]
        expert_data.append(expert_tokens)
        
    return expert_data, tokens_per_expert_range

@torch.compile
def expert_output_permute(
    expert_data: List[torch.Tensor], 
    expert_indices: torch.Tensor, 
    expert_weights: torch.Tensor, 
    num_experts_per_token: int,
    output: torch.Tensor,
) -> torch.Tensor:
    hidden_dim = expert_data[0].size(-1)
    flat_expert_weights = expert_weights.view(-1, 1)

    indices, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    
    for expert_id, end_idx in enumerate(tokens_per_expert_range):
        start_idx = 0 if expert_id == 0 else tokens_per_expert_range[expert_id - 1]
        if start_idx == end_idx:
            continue
        
        # multiply token by by expert weights and reduce
        expert_data[expert_id].mul_(flat_expert_weights[indices[start_idx:end_idx]])
        exp_token_idxs = token_indices[start_idx:end_idx]
        
        # directly accumulate onto output tensor
        output.scatter_reduce_(
            0,
            exp_token_idxs.view(-1, 1).repeat(1, hidden_dim),
            expert_data[expert_id],
            reduce='sum'
        )
    
    return output

@torch.compile
def compute_shared_expert(weights, tokens):
    w_gate = weights['shared_experts.0.weight']
    w_up = weights['shared_experts.1.weight']
    w_down = weights['shared_experts.2.weight']

    gate = F.silu(tokens @ w_gate)
    up = tokens @ w_up
    down = (gate * up) @ w_down
    return down

@torch.compile
def compute_routing(weights, tokens, topk):
    # For some reason, router weights are already transposed?? So transpose here...
    w_g = weights['router.weight'].t()
    logits = tokens @ w_g
    scores = logits.softmax(dim=-1)
    topk_scores, topk_indices = torch.topk(scores, k=topk, dim=-1, sorted=False)
    return topk_scores, topk_indices

def compute_routed_experts(
        weights, 
        tokens, 
        topk, 
        num_experts,
        expert_indices,
        expert_scores,
        output
    ):
    gate_proj_weights = [weights[f'experts.{i}.0.weight'] for i in range(num_experts)]
    up_proj_weights = [weights[f'experts.{i}.1.weight'] for i in range(num_experts)]
    down_proj_weights = [weights[f'experts.{i}.2.weight'] for i in range(num_experts)]
    
    # get tokens for each expert, but compute expert transforms in parallel with group gemms.
    group_token, _ = expert_input_permute(tokens, expert_indices, topk)

    gate = group_gemm_fn(group_token, gate_proj_weights, activation=True)
    
    # this fuses gate * (token @ W_up)
    up_gated = group_gemm_fn(group_token, up_proj_weights, activation=False, group_C_init=gate)
    
    # TODO: this should be another group kernel, or possibly combined with the down project kernel
    #lhs = [t_gate * t_up for (t_gate, t_up) in zip(gate, up)]
    group_token = group_gemm_fn(up_gated, down_proj_weights, activation=False)

    # routed experts get accumulated onto output
    expert_output_permute(
        group_token, 
        expert_indices, 
        expert_scores, 
        topk, 
        output)

def custom_kernel(data):
    input_tensor, weights, config = data
    num_experts = config["n_routed_experts"]
    topk = config["n_experts_per_token"]

    hidden_size = input_tensor.size(2)
    original_shape = input_tensor.size()
    
    flat_tokens = input_tensor.view(-1, hidden_size)

    shared_output = compute_shared_expert(weights, flat_tokens)
    scores, indices = compute_routing(weights, flat_tokens, topk)

    compute_routed_experts(
        weights,
        flat_tokens,
        topk,
        num_experts,
        indices,
        scores,
        shared_output
    )

    return shared_output.view(original_shape)

# Input generation for the reference code

def generate_input(
    dhidden: int,
    dexpert: int,
    nroutedexperts: int,
    nsharedexperts: int,
    nexpertspertoken: int,
    bs: int,
    seqlen: int,
    seed: int
):

    # Really dumb but for now _ isn't parsing correctly.
    d_hidden = dhidden
    d_expert = dexpert
    n_routed_experts = nroutedexperts
    n_shared_experts = nsharedexperts
    n_experts_per_token = nexpertspertoken
    batch_size = bs
    seq_len = seqlen

    config = {
        "d_hidden": d_hidden,
        "d_expert": d_expert,
        "n_routed_experts": n_routed_experts,
        "n_shared_experts": n_shared_experts,
        "n_experts_per_token": n_experts_per_token,
        "batch_size": batch_size,
        "seq_len": seq_len,
    }

    gen = torch.Generator(device='cuda')
    gen.manual_seed(seed)

    num_experts = n_routed_experts
    expert_dim = d_expert
    weights = {}

    input_tensor = torch.randn(
        (batch_size, seq_len, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ).contiguous()

    # Initialize router weights
    weights['router.weight'] = torch.randn(
        (num_experts, d_hidden),
        device="cuda",
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    for i in range(num_experts):
        weights[f'experts.{i}.0.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.1.weight'] = torch.randn(
            (d_hidden, expert_dim),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(expert_dim)

        weights[f'experts.{i}.2.weight'] = torch.randn(
            (expert_dim, d_hidden),
            device='cuda',
            dtype=torch.float16,
            generator=gen
        ) / math.sqrt(d_hidden)
    
    weights['shared_experts.0.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.1.weight'] = torch.randn(
        (d_hidden, expert_dim * n_shared_experts),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(expert_dim * n_shared_experts)
    weights['shared_experts.2.weight'] = torch.randn(
        (expert_dim * n_shared_experts, d_hidden),
        device='cuda',
        dtype=torch.float16,
        generator=gen
    ) / math.sqrt(d_hidden)

    return (input_tensor, weights, config)

if __name__ == '__main__':
    input = generate_input(128, 128, 4, 1, 2, 8, 8, 0)
    output = custom_kernel(input)
python eval.py test /tmp/tmpa6_s47g2
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7078: UserWarning: 
Online softmax is disabled on the fly since Inductor decides to
split the reduction. Cut an issue to PyTorch if this is an
important use case and you want to speed it up with online
softmax.

  warnings.warn(
W0529 07:08:13.448000 970 torch/_inductor/utils.py:2018] [4/0] DeviceCopy in input program
W0529 07:08:18.794000 970 torch/_inductor/utils.py:2018] [4/1] DeviceCopy in input program
