Sonic MoE Qwen 3.8 Max

"""Grouped GEMM + fused SwiGLU up-projection for MoE FFN (RTX PRO 6000 / SM120).

Custom persistent Triton kernel for the variable-length grouped GEMM

    out[start_e:end_e] = silu(x_e @ W_gate[e]) * (x_e @ W_up[e])

Design notes:
- The gate/up weights are repacked once (per weight version) into a single
  column-interleaved tensor W_comb[e] = [g_0, u_0, g_1, u_1, ...], so each
  output tile needs only ONE wide dot into a (BM, 2*BN) accumulator; the
  epilogue de-interleaves and applies SwiGLU. This halves the MMA issue
  stream versus a dual-accumulator kernel.
- Tiles are enumerated expert-major; a persistent grid walks the flattened
  tile space with an amortized O(E) cursor, so forward() never synchronizes
  with the host.
- SM120 (Blackwell workstation) has no tcgen05/TMEM; the kernel targets the
  warp-level mma.sync bf16 tensor-core path with cp.async pipelining.
- A fully masked generic kernel handles any shape where H % BK or I % BN != 0.
"""

from __future__ import annotations
import torch, torch.nn as nn, triton, triton.language as tl

NUM_SMS = torch.cuda.get_device_properties(0).multi_processor_count

# Per-shape tile configs (tuned on RTX PRO 6000; keyed by (H, I)).
CFG_TABLE = {
    (4096, 1536): dict(BM=128, BN=128, BK=32, warps=8, stages=3, grid_mult=1, maxnreg=None),
    (2048, 1024): dict(BM=128, BN=64, BK=64, warps=8, stages=3, grid_mult=1, maxnreg=232),
    (2048, 4096): dict(BM=128, BN=128, BK=32, warps=8, stages=3, grid_mult=1, maxnreg=None),
}
CFG_DEFAULT = dict(BM=128, BN=128, BK=32, warps=8, stages=3, grid_mult=1, maxnreg=None)


def _cfg_for(H, I):  # noqa: E741
    return CFG_TABLE.get((H, I), CFG_DEFAULT)


@triton.jit
def _kernel_v5(
    x_ptr, w_ptr, out_ptr, offs_ptr,
    stride_xm, stride_wn, stride_om,
    MAX_ITER,
    E: tl.constexpr, GRID,
    H: tl.constexpr, I: tl.constexpr,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    pid = tl.program_id(0)
    NB: tl.constexpr = I // BN
    total = 0
    for ee in range(E):
        n_e = tl.load(offs_ptr + ee + 1) - tl.load(offs_ptr + ee)
        total += ((n_e + BM - 1) // BM) * NB
    e_cur = 0
    base_cur = 0
    start_cur = tl.load(offs_ptr + 0)
    end_cur = tl.load(offs_ptr + 1)
    nm_cur = ((end_cur - start_cur) + BM - 1) // BM
    rn = tl.arange(0, 2 * BN)
    rm = tl.arange(0, BM)
    rk = tl.arange(0, BK)
    for i in range(0, MAX_ITER):
        t = pid + i * GRID
        if t < total:
            span_cur = nm_cur * NB
            while t >= base_cur + span_cur:
                base_cur += span_cur
                e_cur += 1
                start_cur = end_cur
                end_cur = tl.load(offs_ptr + e_cur + 1)
                nm_cur = ((end_cur - start_cur) + BM - 1) // BM
                span_cur = nm_cur * NB
            off_t = t - base_cur
            mb = off_t // NB
            nb = off_t % NB
            row0 = start_cur + mb * BM
            wrow0 = e_cur * H
            x_ptrs = x_ptr + (row0 + rm)[:, None] * stride_xm + rk[None, :]
            w_ptrs = w_ptr + wrow0.to(tl.int64) * (2 * I) + rk[:, None] * (2 * I) + (nb * 2 * BN + rn)[None, :]
            acc = tl.zeros((BM, 2 * BN), dtype=tl.float32)
            for _k in range(0, H // BK):
                x = tl.load(x_ptrs)
                w = tl.load(w_ptrs)
                acc = tl.dot(x, w, acc)
                x_ptrs += BK
                w_ptrs += BK * (2 * I)
            g, u = tl.split(tl.reshape(acc, (BM, BN, 2)))
            res = g * tl.sigmoid(g) * u
            out = res.to(tl.bfloat16)
            rno = tl.arange(0, BN)
            out_ptrs = out_ptr + (row0 + rm)[:, None] * stride_om + (nb * BN + rno)[None, :]
            if (row0 + BM) <= end_cur:
                tl.store(out_ptrs, out)
            else:
                row_mask = (row0 + rm) < end_cur
                tl.store(out_ptrs, out, mask=row_mask[:, None])


@triton.jit
def _kernel_generic(
    x_ptr, wg_ptr, wu_ptr, out_ptr, offs_ptr,
    stride_xm, stride_om,
    H, I, MAX_ITER,
    E: tl.constexpr, GRID,
    BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr,
):
    """Masked dual-accumulator variant for shapes with H % BK or I % BN != 0."""
    pid = tl.program_id(0)
    NB = tl.cdiv(I, BN)
    total = 0
    for ee in range(E):
        n_e = tl.load(offs_ptr + ee + 1) - tl.load(offs_ptr + ee)
        total += ((n_e + BM - 1) // BM) * NB
    e_cur = 0
    base_cur = 0
    start_cur = tl.load(offs_ptr + 0)
    end_cur = tl.load(offs_ptr + 1)
    nm_cur = ((end_cur - start_cur) + BM - 1) // BM
    rn = tl.arange(0, BN)
    rm = tl.arange(0, BM)
    rk = tl.arange(0, BK)
    for i in range(0, MAX_ITER):
        t = pid + i * GRID
        if t < total:
            span_cur = nm_cur * NB
            while t >= base_cur + span_cur:
                base_cur += span_cur
                e_cur += 1
                start_cur = end_cur
                end_cur = tl.load(offs_ptr + e_cur + 1)
                nm_cur = ((end_cur - start_cur) + BM - 1) // BM
                span_cur = nm_cur * NB
            off_t = t - base_cur
            mb = off_t // NB
            nb = off_t % NB
            row0 = start_cur + mb * BM
            wrow0 = e_cur.to(tl.int64) * H
            cols = nb * BN + rn
            x_ptrs = x_ptr + (row0 + rm)[:, None] * stride_xm + rk[None, :]
            w_ptrs_g = wg_ptr + wrow0 * I + rk[:, None] * I + cols[None, :]
            w_ptrs_u = wu_ptr + wrow0 * I + rk[:, None] * I + cols[None, :]
            acc_g = tl.zeros((BM, BN), dtype=tl.float32)
            acc_u = tl.zeros((BM, BN), dtype=tl.float32)
            col_mask = cols < I
            for kk in range(0, tl.cdiv(H, BK)):
                k_mask = (kk * BK + rk) < H
                x = tl.load(x_ptrs, mask=k_mask[None, :], other=0.0)
                wg = tl.load(w_ptrs_g, mask=k_mask[:, None] & col_mask[None, :], other=0.0)
                wu = tl.load(w_ptrs_u, mask=k_mask[:, None] & col_mask[None, :], other=0.0)
                acc_g = tl.dot(x, wg, acc_g)
                acc_u = tl.dot(x, wu, acc_u)
                x_ptrs += BK
                w_ptrs_g += BK * I
                w_ptrs_u += BK * I
            res = acc_g * tl.sigmoid(acc_g) * acc_u
            out = res.to(tl.bfloat16)
            out_ptrs = out_ptr + (row0 + rm)[:, None] * stride_om + cols[None, :]
            row_mask = (row0 + rm) < end_cur
            tl.store(out_ptrs, out, mask=row_mask[:, None] & col_mask[None, :])


class Model(nn.Module):
    def __init__(self, T_total, H, I, E, K):
        super().__init__()
        self.T_total, self.H, self.I, self.E, self.K = T_total, H, I, E, K
        self.W_gate = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
        self.W_up = nn.Parameter(torch.empty(E, H, I, dtype=torch.bfloat16))
        nn.init.normal_(self.W_gate, std=0.02)
        nn.init.normal_(self.W_up, std=0.02)
        self._w_comb = None
        self._w_key = None

    def _get_w_comb(self):
        key = (self.W_gate.data_ptr(), self.W_up.data_ptr(),
               self.W_gate._version, self.W_up._version)
        if self._w_key != key:
            E, H, I = self.E, self.H, self.I
            w = torch.empty(E, H, 2 * I, dtype=torch.bfloat16, device=self.W_gate.device)
            w[:, :, 0::2] = self.W_gate
            w[:, :, 1::2] = self.W_up
            self._w_comb = w
            self._w_key = key
        return self._w_comb

    def forward(self, hidden_states, expert_offsets):
        T_perm, H = hidden_states.shape
        I, E = self.I, self.E
        out = torch.empty(T_perm, I, dtype=torch.bfloat16, device=hidden_states.device)
        cfg = _cfg_for(H, self.I)
        BM, BN, BK = cfg["BM"], cfg["BN"], cfg["BK"]
        NB = triton.cdiv(I, BN)
        upper = (triton.cdiv(T_perm, BM) + E) * NB
        grid = (min(NUM_SMS * cfg["grid_mult"], upper),)
        max_iter = triton.cdiv(upper, grid[0])
        kw = dict(num_warps=cfg["warps"], num_stages=cfg["stages"])
        if cfg.get("maxnreg"):
            kw["maxnreg"] = cfg["maxnreg"]
        if H % BK == 0 and I % BN == 0:
            w_comb = self._get_w_comb()
            _kernel_v5[grid](
                hidden_states, w_comb, out, expert_offsets,
                hidden_states.stride(0), w_comb.stride(1), out.stride(0),
                max_iter, E=E, GRID=grid[0], H=H, I=I, BM=BM, BN=BN, BK=BK, **kw)
        else:
            _kernel_generic[grid](
                hidden_states, self.W_gate, self.W_up, out, expert_offsets,
                hidden_states.stride(0), out.stride(0),
                H, I, max_iter, E=E, GRID=grid[0], BM=BM, BN=BN, BK=BK, **kw)
        return out


T_total = 32768
H = 4096
I = 1536  # noqa: E741
E = 128
K = 8


def _build_routing(T_total: int, E: int, K: int, device: str = "cpu") -> torch.Tensor:
    T_perm = T_total * K
    base = T_perm // E
    rem = T_perm - base * E
    counts = torch.full((E,), base, dtype=torch.int32, device=device)
    counts[:rem] += 1
    offsets = torch.zeros(E + 1, dtype=torch.int32, device=device)
    offsets[1:] = torch.cumsum(counts, dim=0)
    return offsets


def get_inputs():
    T_perm = T_total * K
    hidden_states = torch.randn(T_perm, H, dtype=torch.bfloat16) * 0.1
    expert_offsets = _build_routing(T_total, E, K)
    return [hidden_states, expert_offsets]


def get_init_inputs():
    return [T_total, H, I, E, K]