fp8_blockscale_fused_moe

import torch
import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):

    def __init__(self, num_experts: int, intermediate_size: int, top_k: int, scale_block_n: int=128, scale_block_k: int=128) -> None:
        super().__init__()
        self.num_experts = num_experts
        self.intermediate_size = intermediate_size
        self.top_k = top_k
        self.scale_block_n = scale_block_n
        self.scale_block_k = scale_block_k

    def forward(self, hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, a_scale: torch.Tensor, fc1_scale: torch.Tensor, fc2_scale: torch.Tensor) -> torch.Tensor:
        (B, D) = hidden_states.shape
        top_k = topk_ids.shape[1]
        (expert, model_dim, inter_dim) = w2.shape
        blk_n = self.scale_block_n
        blk_k = self.scale_block_k
        h = hidden_states.float()
        h = h.view(B, -1, blk_k) * a_scale.unsqueeze(-1)
        h = h.view(B, -1)
        nblk_n = inter_dim // blk_n
        nblk_k = model_dim // blk_k
        fc1_s = fc1_scale.view(-1, 1).repeat(1, blk_n * blk_k).view(expert, -1, nblk_k, blk_n, blk_k)
        fc1_s = fc1_s.permute(0, 1, 3, 2, 4).reshape(expert, -1, model_dim)
        w1_dq = w1.float() * fc1_s
        fc2_s = fc2_scale.view(-1, 1).repeat(1, blk_k * blk_n).view(expert, nblk_k, nblk_n, blk_k, blk_n)
        fc2_s = fc2_s.permute(0, 1, 3, 2, 4).reshape(expert, model_dim, inter_dim)
        w2_dq = w2.float() * fc2_s
        h = h.view(B, 1, model_dim).repeat(1, top_k, 1)
        out = torch.zeros(B, top_k, D, dtype=torch.float32, device=hidden_states.device)
        for eid in range(expert):
            mask = topk_ids == eid
            if not mask.any():
                continue
            tokens = h[mask]
            act_input = tokens @ w1_dq[eid].T
            (gate, up) = act_input.split([inter_dim, inter_dim], dim=-1)
            act_out = F.silu(gate) * up
            out[mask] = act_out @ w2_dq[eid].T
        return (out * topk_weights.view(B, -1, 1)).sum(dim=1).to(hidden_states.dtype)

import torch

def _make_inputs(token_count: int, hidden_size: int, intermediate_size: int, num_experts: int, top_k: int, scale_block_n: int=128, scale_block_k: int=128) -> dict[str, torch.Tensor]:
    dt = torch.float8_e4m3fnuz
    fp8_max = torch.finfo(dt).max
    blk_n = scale_block_n
    blk_k = scale_block_k
    nblk_h_k = hidden_size // blk_k
    nblk_i_n = intermediate_size // blk_n
    nblk_2i_n = (intermediate_size * 2) // blk_n

    def _block_quant(x_fp32: torch.Tensor, group_dims: tuple) -> tuple:
        amax = x_fp32.abs().amax(dim=group_dims, keepdim=True).clamp(min=1e-12)
        scale = amax / fp8_max
        q = (x_fp32 / scale).clamp(-fp8_max, fp8_max).to(dt)
        return q, scale

    hidden_bf = torch.randn(token_count, hidden_size, dtype=torch.bfloat16, device='cuda') * 0.03
    h_blocks = hidden_bf.float().view(token_count, nblk_h_k, blk_k)
    h_q, h_scale = _block_quant(h_blocks, group_dims=(-1,))
    hidden_states = h_q.view(token_count, hidden_size)
    a_scale = h_scale.squeeze(-1).to(torch.float32)

    w1_bf = torch.randn(num_experts, intermediate_size * 2, hidden_size, dtype=torch.bfloat16, device='cuda')
    w1_blocks = w1_bf.float().view(num_experts, nblk_2i_n, blk_n, nblk_h_k, blk_k)
    w1_q, w1_scale = _block_quant(w1_blocks, group_dims=(2, 4))
    w1 = w1_q.view(num_experts, intermediate_size * 2, hidden_size)
    fc1_scale = w1_scale.squeeze(2).squeeze(-1).to(torch.float32)

    w2_bf = torch.randn(num_experts, hidden_size, intermediate_size, dtype=torch.bfloat16, device='cuda')
    w2_blocks = w2_bf.float().view(num_experts, nblk_h_k, blk_k, nblk_i_n, blk_n)
    w2_q, w2_scale = _block_quant(w2_blocks, group_dims=(2, 4))
    w2 = w2_q.view(num_experts, hidden_size, intermediate_size)
    fc2_scale = w2_scale.squeeze(2).squeeze(-1).to(torch.float32)

    topk_ids = torch.topk(torch.rand(token_count, num_experts, device='cuda'), top_k, dim=1).indices.to(torch.int32)
    topk_weights = torch.softmax(torch.randn(token_count, top_k, device='cuda'), dim=-1).to(torch.float32)

    return {'hidden_states': hidden_states, 'w1': w1, 'w2': w2, 'topk_weights': topk_weights, 'topk_ids': topk_ids, 'a_scale': a_scale, 'fc1_scale': fc1_scale, 'fc2_scale': fc2_scale}

def get_inputs() -> list[torch.Tensor]:
    return list(_make_inputs(token_count=32, hidden_size=2048, intermediate_size=768, num_experts=128, top_k=8).values())

def get_init_inputs() -> dict[str, object]:
    return {'num_experts': 128, 'intermediate_size': 768, 'top_k': 8}