r"""
Uses pytorch nested tensor, which should just be a faster
version of the group gemm idea.
"""

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.nested._NestedTensor
    indices: torch.Tensor
    token_gather_indices: torch.Tensor

@torch.compile
def get_token_indices(expert_indices, num_experts_per_token):
    flat_expert_indices = expert_indices.view(-1)
    indices = flat_expert_indices.argsort()
    token_indices = indices // num_experts_per_token
    counts = flat_expert_indices.bincount()
    # torch.NestedTensor init offsets expects zero prefix. This stays on device
    tokens_per_expert_range = torch.empty(counts.size(0) + 1, device=counts.device, dtype=counts.dtype)
    torch.cumsum(counts, dim=0, out=tokens_per_expert_range[1:])
    tokens_per_expert_range[0] = 0
    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
) -> torch.nested._NestedTensor:
    indices, tokens_per_expert_range, token_indices = get_token_indices(expert_indices, num_experts_per_token)
    gather_tokens = torch.gather(tokens, dim=0, index=token_indices.unsqueeze(1).expand(-1, tokens.size(1)))
    routed_tokens = torch.nested.nested_tensor_from_jagged(values=gather_tokens, offsets=tokens_per_expert_range)
    return GroupedTokens(
        tokens=routed_tokens,
        indices=indices,
        token_gather_indices=token_indices
    )

@torch.compile
def expert_output_permute(
    grouped_tokens: GroupedTokens,
    expert_scores: torch.Tensor,
    output: torch.Tensor,
) -> torch.Tensor:
    values = grouped_tokens.tokens.values()
    values.mul_(expert_scores.view(-1, 1)[grouped_tokens.indices])
    output.scatter_reduce_(
        0,
        # `scatter_reduce_` doesn't broadcast indices, so repeat indices along hidden dim. `.expand` returns a view
        grouped_tokens.token_gather_indices.view(-1, 1).expand(-1, output.size(-1)),
        values,
        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)]
    
    # NOTE: ideally, I would like to avoid copying the weights...
    # I'm not sure how to do that... I think 
    gate_proj_tensor = torch.stack(gate_proj_weights, dim=0)
    up_proj_tensor = torch.stack(up_proj_weights, dim=0)
    down_proj_tensor = torch.stack(down_proj_weights, dim=0)

    # 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: torch.nested._NestedTensor = F.silu(torch.bmm(group_tokens.tokens, gate_proj_tensor))
    up: torch.nested._NestedTensor = torch.bmm(group_tokens.tokens, up_proj_tensor)

    # this requires gate and up to have the SAME offsets tensor    
    up_gated: torch.nested._NestedTesnor = up * gate
    
    down: torch.nested._NestedTensor = torch.bmm(up_gated, down_proj_tensor)
    
    group_tokens = GroupedTokens(
        tokens=down,
        indices=group_tokens.indices,
        token_gather_indices=group_tokens.token_gather_indices
    )
    
    expert_output_permute(
        group_tokens, 
        expert_scores, 
        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/tmpmfamk0mt
/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(
