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
from dataclasses import dataclass

@dataclass
class GroupedTokens:
    tokens: torch.Tensor
    tokens_per_expert_range: torch.tensor
    tokens_per_expert_range_cpu: torch.tensor
    indices: torch.Tensor
    token_gather_indices: torch.Tensor

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

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

configs = [
    # M, N, K, warps, stages
    # for AMD, num stages should be 1 or 2.
    (128, 128, 128, 4, 2),
    (128, 128, 64, 4, 2),
    (128, 128, 32, 4, 2),
    (256, 256, 64, 4, 2),
    (64, 128, 128, 4, 2),
    (128, 128, 128, 8, 2),
    (256, 256, 64, 8, 2),
    (64, 64, 128, 4, 2),
    (64, 64, 256, 4, 2),
    (128, 64, 256, 4, 2),
    (128, 128, 128, 4, 1),
    (128, 128, 64, 4, 1),
    (128, 128, 32, 4, 1),
    (256, 256, 64, 4, 1),
    (64, 128, 128, 4, 1),
    (128, 128, 128, 8, 1),
    (256, 256, 64, 8, 1),
    (64, 64, 128, 4, 1),
    (64, 64, 256, 4, 1),
    (128, 64, 256, 4, 1)
]

configs = [
    triton.Config({
        'BLOCK_SIZE_M': c[0],
        'BLOCK_SIZE_N': c[1],
        'BLOCK_SIZE_K': c[2],
        'num_warps': c[3],
        'num_stages': c[4],
        'NUM_SM': num_sms()
    }) for c in configs
]

@triton.autotune(
    configs=configs,
    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,
    ACTIVATION: tl.constexpr,
    # 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,
):
    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
            
            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])
                                
                a = tl.load(a_ptrs, mask=m_mask[:, None])
                b = tl.load(b_ptrs)
                accumulator = tl.dot(a, b, acc=accumulator)
                a_ptrs += BLOCK_SIZE_K
                b_ptrs += BLOCK_SIZE_K * ldb
                
            if ACTIVATION:
                accumulator = silu(accumulator)
                
            c = accumulator.to(tl.float16)

            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, :]

            tl.store(c_ptrs, c, mask=m_mask[:, None])

            # go to the next tile by advancing NUM_SM
            tile_idx += tl.num_programs(axis=0)

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

def get_a_addrs(group_A: GroupedTokens, group_size: int):
    a_addrs = []
    a_shapes = []
    for i in range(group_size):
        end_idx = group_A.tokens_per_expert_range_cpu[i]
        start_idx = 0 if i == 0 else group_A.tokens_per_expert_range_cpu[i - 1]
        a_addrs.append(group_A.tokens[start_idx:-1].data_ptr())
        M = end_idx - start_idx
        K = group_A.tokens.size(-1)
        a_shapes.append((M, K))
    return a_addrs, a_shapes
        
def get_b_addrs(group_B: List[torch.Tensor], group_size: int):
    b_addrs = []
    b_shapes = []
    for i in range(group_size):
        b_addrs.append(group_B[i].data_ptr())
        K = group_B[i].size(0)
        N = group_B[i].size(1)
        b_shapes.append((K, N))
    return b_addrs, b_shapes

def group_gemm_fn(group_A: GroupedTokens, group_B: List[torch.Tensor], activation):
    DEVICE=group_B[0].device
    
    group_size = len(group_B)
    
    A_addrs, A_shapes = get_a_addrs(group_A, group_size)
    B_addrs, B_shapes = get_b_addrs(group_B, group_size)
    
    C_addrs = []
    g_sizes = []
    g_lds = []
        
    c = torch.empty((group_A.tokens.size(0), group_B[0].size(-1)), device=DEVICE, dtype=group_A.tokens.dtype)
    
    c_offset = 0
    for i in range(group_size):
        assert A_shapes[i][1] == B_shapes[i][0]
        M, K = A_shapes[i]
        K, N = B_shapes[i]
        
        C_addrs.append(c[c_offset:].data_ptr()) # using [c_offset:] indexing to get view
        g_sizes += [M, N, K]
        g_lds += [K, N, c.stride(0)]
        c_offset += M
                
    # note these are device tensors
    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)

    # invoke like normal kernel???    
    #grid = lambda META: (triton.cdiv(c.size(0), META['BLOCK_SIZE_M']) * 
    #                     triton.cdiv(c.size(1), META['BLOCK_SIZE_N']), )
    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
    )
        
    return GroupedTokens(
        tokens=c,
        tokens_per_expert_range=group_A.tokens_per_expert_range,
        tokens_per_expert_range_cpu=group_A.tokens_per_expert_range_cpu,
        indices=group_A.indices,
        token_gather_indices=group_A.token_gather_indices,
    )

@torch.compile
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()
    tokens_per_expert_range = counts.cumsum(dim=0)
    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
) -> GroupedTokens:
    indices, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    token_dim = tokens.size(1)
    output = torch.gather(tokens, dim=0, index=token_indices.unsqueeze(1).expand(-1, token_dim))
    return GroupedTokens(
        tokens=output,
        tokens_per_expert_range=tokens_per_expert_range,
        tokens_per_expert_range_cpu=tokens_per_expert_range.detach().cpu(),
        indices=indices, 
        token_gather_indices=token_indices,
    )

@torch.compile
def expert_output_permute(
    grouped_tokens: GroupedTokens,
    expert_scores: torch.Tensor,
    num_experts_per_token: int,
    output: torch.Tensor
) -> torch.Tensor:
    grouped_tokens.tokens.mul_(expert_scores.view(-1, 1)[grouped_tokens.indices])
    output.scatter_reduce_(
        0,
        # index tensor doesn't broadcast, so repeat indices along hidden dim
        grouped_tokens.token_gather_indices.view(-1, 1).expand(-1, output.size(-1)),
        grouped_tokens.tokens,
        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_tokens: GroupedTokens = expert_input_permute(tokens, expert_indices, topk)

    gate: GroupedTokens = group_gemm_fn(group_tokens, gate_proj_weights, activation=True)
    
    # this fuses gate * (token @ W_up)
    up: GroupedTokens = group_gemm_fn(group_tokens, up_proj_weights, activation=False)
    
    up_gated = GroupedTokens(
        tokens=gate.tokens * up.tokens,
        tokens_per_expert_range=gate.tokens_per_expert_range,
        tokens_per_expert_range_cpu=gate.tokens_per_expert_range_cpu,
        indices=gate.indices,
        token_gather_indices=gate.token_gather_indices,
    )
    
    group_token: GroupedTokens = group_gemm_fn(up_gated, down_proj_weights, activation=False)

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

def custom_kernel(data):
    with torch.no_grad():
        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/tmplcrjdz7y
/usr/local/lib/python3.10/dist-packages/torch/_inductor/lowering.py:7086: 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(
W0606 00:13:52.211000 1559 torch/_inductor/utils.py:2019] [5/0] DeviceCopy in input program
W0606 00:13:57.003000 1559 torch/_inductor/utils.py:2019] [5/1] DeviceCopy in input program
