Build files.

.gitignore

@@ -8,7 +8,6 @@ __pycache__/
 
 # Distribution / packaging
 .Python
-build/
 develop-eggs/
 dist/
 downloads/

build/torch-universal/scattermoe/__init__.py

@@ -0,0 +1,15 @@
+from .parallel_experts import flatten_sort_count, parallel_linear, ParallelExperts
+from . import parallel_experts
+from . import kernels
+from . import mlp
+from . import utils
+
+__all__ = [
+    "flatten_sort_count",
+    "parallel_linear",
+    "ParallelExperts",
+    "parallel_experts",
+    "kernels",
+    "mlp",
+    "utils"
+]

build/torch-universal/scattermoe/kernels/__init__.py

@@ -0,0 +1,3 @@
+from . import ops
+
+__all__ = ["ops"]

build/torch-universal/scattermoe/kernels/ops.py

@@ -0,0 +1,457 @@
+import torch
+import triton
+import triton.language as tl
+from typing import Optional
+
+BLOCK_M = 128
+ALLOW_TF32 = True
+
+
+
+@triton.jit
+def _compute_expert_block(
+    E_idx, E_mask,
+    M_in_idx,
+    N_block, N_mask,
+    X_ptr, stride_xm, stride_xk,
+    W_ptr, stride_we, stride_wk, stride_wn,
+    K,
+    acc,
+    no_k_mask,
+    BLOCK_K,
+    allow_tf32=True,
+):
+
+    K_block = tl.arange(0, BLOCK_K)
+    X_blk_ptrs = X_ptr + M_in_idx[:, None] * stride_xm + K_block[None, :] * stride_xk
+    W_blk_ptrs = W_ptr + K_block[:, None] * stride_wk + N_block[None, :] * stride_wn + E_idx * stride_we
+    iters = tl.cdiv(K, BLOCK_K)
+
+    for K_block_id in range(iters):
+        if no_k_mask:
+            x = tl.load(X_blk_ptrs, mask=E_mask[:, None])
+            w = tl.load(W_blk_ptrs, mask=N_mask[None, :])
+        else:
+            K_mask = (K_block_id * BLOCK_K + K_block) < K
+            x = tl.load(X_blk_ptrs, mask=E_mask[:, None] & K_mask[None, :])
+            w = tl.load(W_blk_ptrs, mask=K_mask[:, None] & N_mask[None, :])
+
+        X_blk_ptrs += BLOCK_K * stride_xk
+        W_blk_ptrs += BLOCK_K * stride_wk
+        acc = tl.dot(x, w, acc, allow_tf32=allow_tf32)
+    return acc
+
+
+def _scatter2scatter_configs():
+    return [
+        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
+    ]
+
+@triton.autotune(configs=_scatter2scatter_configs(), key=['M', 'N', 'K'], )
+@triton.heuristics({
+    "NO_K_MASK": lambda args: (args['K'] % args['BLOCK_K']) == 0,
+    "NO_N_MASK": lambda args: (args['N'] % args['BLOCK_N']) == 0,
+})
+@triton.jit
+def _scatter2scatter(
+    X_ptr, stride_xm: tl.constexpr, stride_xk: tl.constexpr,
+    W_ptr, stride_we, stride_wk: tl.constexpr, stride_wn: tl.constexpr,
+    Y_ptr, stride_ym: tl.constexpr, stride_yn: tl.constexpr,
+    B_ptr, stride_be: tl.constexpr, stride_bn: tl.constexpr,
+    grouped_idx_ptr, expert_idxs_ptr,
+    # block_start_idx_ptr,
+    FAN_OUT: tl.constexpr,
+    M, K: tl.constexpr, N: tl.constexpr, E: tl.constexpr,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    # OUT_M,
+    allow_tf32: tl.constexpr,
+    x_grouped: tl.constexpr, y_grouped: tl.constexpr,
+    NO_K_MASK: tl.constexpr, NO_N_MASK: tl.constexpr
+):
+    pid = tl.program_id(axis=0)
+
+    N_BLOCK_COUNT = tl.cdiv(N, BLOCK_N)
+    M_block_id = pid // N_BLOCK_COUNT
+    N_block_id = pid % N_BLOCK_COUNT
+
+    M_block = M_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    N_mask = N_block < N
+    M_boundary_mask = M_block < (FAN_OUT * M)
+    E_idxs = tl.load(expert_idxs_ptr + M_block, mask=M_boundary_mask, other=E)
+
+    no_k_mask = K % BLOCK_K == 0
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=ACC_TYPE)
+    E_first_idx = tl.min(E_idxs)
+    E_last_idx = tl.minimum(tl.max(E_idxs), E - 1)
+    M_idx = tl.load(grouped_idx_ptr + M_block, mask=M_boundary_mask).to(tl.int32)
+    for E_idx in range(E_first_idx, E_last_idx + 1):
+        E_mask = E_idxs == E_idx
+        E_M_idx = M_idx
+        if x_grouped:
+            M_in_idx = M_block
+        else:
+            M_in_idx = E_M_idx // FAN_OUT
+        acc = _compute_expert_block(
+            E_idx, E_mask,
+            M_in_idx, N_block, N_mask,
+            X_ptr, stride_xm, stride_xk,
+            W_ptr, stride_we, stride_wk, stride_wn,
+            K,
+            acc,
+            no_k_mask,
+            BLOCK_K,
+            allow_tf32=allow_tf32,
+        )
+
+    if B_ptr is not None:
+        B_blk_ptrs = B_ptr + E_idxs[:, None] * stride_be + N_block[None, :] * stride_bn
+        acc += tl.load(B_blk_ptrs, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+    if y_grouped:
+        M_out_idx = M_block
+    else:
+        M_out_idx = M_idx
+    Y_blk_ptrs = Y_ptr + (M_out_idx[:, None] * stride_ym + N_block[None, :] * stride_yn)
+    tl.store(Y_blk_ptrs, acc, mask=M_boundary_mask[:, None] & N_mask[None, :])
+
+def scatter2scatter(X, W, sorted_expert_idxs, sorted_scattered_idxs, k,
+                    b=None,
+                    x_grouped=False, y_grouped=False,
+                    out=None):
+    assert sorted_scattered_idxs.size(0) == sorted_expert_idxs.size(0)
+    assert sorted_scattered_idxs.size(0) == X.size(0) * k
+    # Pre-kernel setup
+    y_dim = W.size(-1)
+    L_scattered = sorted_expert_idxs.size(0)
+    if out is None:
+        output = torch.empty((L_scattered, y_dim), device=X.device, dtype=X.dtype)
+    else:
+        assert out.size(0) == L_scattered and out.size(1) == y_dim
+        output = out
+
+    scatter2scatter_compileable(output, W, X, k, sorted_expert_idxs, sorted_scattered_idxs,
+                                b, x_grouped, y_grouped)
+    return output
+
+
+@torch.library.custom_op("scattermoe::scatter2scatter", mutates_args={"output"})
+def scatter2scatter_compileable(
+        output: torch.Tensor,
+        W: torch.Tensor,
+        X: torch.Tensor,
+        k: int,
+        sorted_expert_idxs: torch.Tensor,
+        sorted_scattered_idxs: torch.Tensor,
+        b: Optional[torch.Tensor],
+        x_grouped: bool, y_grouped: bool) -> None:
+    def grid(META):
+        grid_num = (
+            triton.cdiv(sorted_expert_idxs.size(0), META["BLOCK_M"]) *
+            triton.cdiv(META['N'], META['BLOCK_N']),
+        )
+        return grid_num
+
+    if b is None:
+        b = None
+        stride_be = stride_bk = 0
+    else:
+        stride_be, stride_bk = b.stride()
+
+    _scatter2scatter[grid](
+        # X_ptr, stride_xm, stride_xk,
+        X, X.stride(0), X.stride(1),
+        # W_ptr, stride_we, stride_wk, stride_wn,
+        W, W.stride(0), W.stride(1), W.stride(2),
+        # Y_ptr, stride_ym, stride_yn,
+        output, output.stride(0), output.stride(1),
+        # B_ptr, stride_be, stride_bk
+        b, stride_be, stride_bk,
+        grouped_idx_ptr=sorted_scattered_idxs,
+        expert_idxs_ptr=sorted_expert_idxs,
+        # block_start_idx_ptr=padded_block_idxs,
+        FAN_OUT=k,
+        M=X.size(0),
+        K=X.size(1),
+        N=output.size(1), E=W.size(0),
+        BLOCK_M=BLOCK_M,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32,
+        x_grouped=x_grouped, y_grouped=y_grouped,
+    )
+
+
+def _config_XtY():
+    return [
+        triton.Config({'BLOCK_N': 128, 'BLOCK_K': 128, 'BLOCK_M': 32}, num_stages=4, num_warps=4),
+    ]
+
+def group_bwd_W(DY, X, expert_offsets, E, has_bias=False):
+    DWt = torch.zeros((E, DY.size(-1), X.size(-1)), device=DY.device, dtype=DY.dtype)
+    DW = DWt.permute(0, 2, 1)
+    if has_bias:
+        Db = torch.zeros((E, DY.size(-1)), device=DY.device, dtype=DY.dtype)
+    else:
+        Db = None
+    groupXtY_compileable(E, DW, Db, DY, X, expert_offsets)
+    return DW, Db
+
+
+@torch.library.custom_op("scattermoe::groupXtY", mutates_args={"DW"})
+def groupXtY_compileable(
+        E: int,
+        DW: torch.Tensor,
+        Db: Optional[torch.Tensor],
+        DY: torch.Tensor,
+        X: torch.Tensor,
+        expert_offsets: torch.Tensor) -> None:
+    def grid(META):
+        grid = (
+            E * triton.cdiv(META['K'], META['BLOCK_K']),
+            triton.cdiv(META['N'], META['BLOCK_N']),
+        )
+        return grid
+    
+    if Db is None:
+        stride_dbe = 0
+        stride_dbn = 0
+    else:
+        stride_dbe, stride_dbn = Db.stride()
+
+    _groupXtY[grid](
+        # DY_ptr, stride_dym, stride_dyk,
+        DY, DY.stride(0), DY.stride(1),
+        # X_ptr, stride_xm, stride_xn,
+        X, X.stride(0), X.stride(1),
+        # DW_ptr, stride_dwe, stride_dwk, stride_dwn,
+        DW, DW.stride(0), DW.stride(1), DW.stride(2),
+        # Db_ptr, stride_dwe, stride_dbn,
+        Db, stride_dbe, stride_dbn,
+        # expert_offsets_ptr,
+        expert_offsets,
+        # K: tl.constexpr, N: tl.constexpr,
+        M=DY.size(0), N=DY.size(-1), K=X.size(-1),
+        # ACC_TYPE: tl.constexpr,
+        ACC_TYPE=tl.float32,
+        allow_tf32=ALLOW_TF32
+    )
+
+
+@triton.autotune(configs=_config_XtY(), key=['M', 'N', 'K'], )
+@triton.heuristics({
+    "NO_K_MASK": lambda args: (args['K'] % args['BLOCK_K']) == 0,
+    "NO_N_MASK": lambda args: (args['N'] % args['BLOCK_N']) == 0,
+})
+@triton.jit
+def _groupXtY(
+    DY_ptr, stride_dym, stride_dyk,
+    X_ptr, stride_xm, stride_xn,
+    DW_ptr, stride_dwe, stride_dwk, stride_dwn,
+    Db_ptr, stride_dbe, stride_dbn,
+    expert_offsets_ptr,
+    M, K: tl.constexpr, N: tl.constexpr,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+    allow_tf32: tl.constexpr,
+    NO_K_MASK: tl.constexpr, NO_N_MASK: tl.constexpr
+):
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+    num0 = tl.num_programs(0)
+    num1 = tl.num_programs(1)
+    # pid1, pid0 = tl.swizzle2d(pid1, pid0, num1, num0, 128)
+    pid0, pid1 = tl.swizzle2d(pid0, pid1, num0, num1, 4)
+
+    K_BLOCK_COUNT = tl.cdiv(K, BLOCK_K)
+    E_idx = pid0 // K_BLOCK_COUNT
+    K_block_id = pid0 % K_BLOCK_COUNT
+    N_block_id = pid1
+
+    if E_idx == 0:
+        start_idx = 0
+    else:
+        start_idx = tl.load(expert_offsets_ptr + E_idx - 1).to(tl.int32)
+    end_idx = tl.load(expert_offsets_ptr + E_idx).to(tl.int32)
+
+
+    if end_idx > start_idx:
+        M_block = tl.max_contiguous(start_idx + tl.arange(0, BLOCK_M), BLOCK_M)
+
+        K_block = K_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+        K_mask = K_block < K
+        K_block = tl.max_contiguous(tl.multiple_of(K_block % K, BLOCK_K), BLOCK_K)
+
+        N_block = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+        N_mask = N_block < N
+        N_block = tl.max_contiguous(tl.multiple_of(N_block % N, BLOCK_N), BLOCK_N)
+
+        M_idxs = M_block
+        xt_blk_ptrs = X_ptr + K_block[:, None] * stride_xn + M_idxs[None, :] * stride_xm
+        dy_blk_ptrs = DY_ptr + M_idxs[:, None] * stride_dym + N_block[None, :] * stride_dyk
+        if (Db_ptr is not None) and (K_block_id == 0):
+            _xty_and_bias(
+                E_idx, start_idx, end_idx,
+                M_block,
+                K_block, K_mask, N_block, N_mask, 
+                dy_blk_ptrs, stride_dym,
+                xt_blk_ptrs, stride_xm,
+                DW_ptr, stride_dwe, stride_dwk, stride_dwn,
+                Db_ptr, stride_dbe, stride_dbn,
+                BLOCK_M, BLOCK_N, BLOCK_K, ACC_TYPE,
+                allow_tf32, NO_K_MASK, NO_N_MASK,
+                compute_bias=True
+            )
+        else:
+            _xty_and_bias(
+                E_idx, start_idx, end_idx,
+                M_block,
+                K_block, K_mask, N_block, N_mask, 
+                dy_blk_ptrs, stride_dym,
+                xt_blk_ptrs, stride_xm,
+                DW_ptr, stride_dwe, stride_dwk, stride_dwn,
+                Db_ptr, stride_dbe, stride_dbn,
+                BLOCK_M, BLOCK_N, BLOCK_K, ACC_TYPE,
+                allow_tf32, NO_K_MASK, NO_N_MASK,
+                compute_bias=False
+            )
+
+
+@triton.jit
+def _xty_and_bias(
+        E_idx, start_idx, end_idx,
+        M_block,
+        K_block, K_mask, N_block, N_mask, 
+        dy_blk_ptrs, stride_dym,
+        xt_blk_ptrs, stride_xm,
+        DW_ptr, stride_dwe, stride_dwk, stride_dwn,
+        Db_ptr, stride_dbe, stride_dbn,
+        BLOCK_M, BLOCK_N, BLOCK_K, ACC_TYPE,
+        allow_tf32, NO_K_MASK, NO_N_MASK,
+        compute_bias: tl.constexpr
+    ):
+
+    if compute_bias:
+        db_acc = tl.zeros((BLOCK_N,), dtype=ACC_TYPE)
+    else:
+        db_acc = None
+
+    acc = tl.zeros((BLOCK_K, BLOCK_N), dtype=ACC_TYPE)
+    iters = tl.cdiv(end_idx - start_idx, BLOCK_M)
+    for i in range(0, iters):
+        M_mask = (i * BLOCK_M + M_block) < end_idx
+        if NO_K_MASK:
+            xt = tl.load(xt_blk_ptrs, mask=M_mask[None, :])
+        else:
+            xt = tl.load(xt_blk_ptrs, mask=K_mask[:, None] & M_mask[None, :])
+        if NO_N_MASK:
+            dy = tl.load(dy_blk_ptrs, mask=M_mask[:, None])
+        else:
+            dy = tl.load(dy_blk_ptrs, mask=M_mask[:, None] & N_mask[None, :])
+            
+        acc += tl.dot(xt, dy, out_dtype=ACC_TYPE, allow_tf32=allow_tf32)
+
+        xt_blk_ptrs += BLOCK_M * stride_xm
+        dy_blk_ptrs += BLOCK_M * stride_dym
+
+        if compute_bias:
+            db_acc += tl.sum(dy, axis=0)
+
+    DW_blk_ptrs = DW_ptr + E_idx * stride_dwe + K_block[:, None] * stride_dwk + N_block[None, :] * stride_dwn
+    acc = acc.to(DW_blk_ptrs.dtype.element_ty)
+    tl.store(DW_blk_ptrs, acc, mask=K_mask[:, None] & N_mask[None, :])
+    if compute_bias:
+        Db_blk_ptrs =  Db_ptr + E_idx * stride_dbe + N_block * stride_dbn
+        tl.store(Db_blk_ptrs, db_acc, mask=N_mask)
+
+
+
+def _config_grouping():
+    return [
+        triton.Config({'BLOCK_N': 256, 'BLOCK_K': 128}, num_stages=4, num_warps=4),
+        # triton.Config({'BLOCK_N': 128, 'BLOCK_K': 64}, num_stages=4, num_warps=4),
+        # triton.Config({'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=4, num_warps=4),
+    ]
+
+def group(A, sorted_expert_idxs, coeff=None, fan_out=1, out=None):
+    N = sorted_expert_idxs.size(0)
+    K = A.size(1)
+    assert A.size(0) * fan_out == N
+    if out is not None:
+        Y = out
+    else:
+        Y = torch.empty((N, K), dtype=A.dtype, device=A.device)
+    group_compileable(A, K, N, Y, coeff, coeff is not None, fan_out, sorted_expert_idxs)
+    return Y
+
+
+@torch.library.custom_op("scattermoe::group", mutates_args={"Y"})
+def group_compileable(
+        A: torch.Tensor,
+        K: int,
+        N: int,
+        Y: torch.Tensor,
+        coeff: torch.Tensor, has_coeff: bool,
+        fan_out: int,
+        sorted_expert_idxs: torch.Tensor) -> None:
+    def grid(META):
+        grid_num = (triton.cdiv(META['N'], META['BLOCK_N']),)
+        return grid_num
+    _group[grid](
+        # A_ptr, stride_an, stride_ai,
+        A, A.stride(0), A.stride(1), has_coeff, coeff, fan_out,
+        # Y_ptr, stride_yn, stride_yk,
+        Y, Y.stride(0), Y.stride(1),
+        # grouped_idx_ptr,
+        sorted_expert_idxs,
+        # N: tl.constexpr, K: tl.constexpr,
+        N, K
+    )
+
+
+@triton.autotune(configs=_config_grouping(), key=['K'])
+@triton.heuristics({
+    "NO_K_MASK": lambda args: (args['K'] % args['BLOCK_K']) == 0
+})
+@triton.jit
+def _group(
+    src_ptr, stride_sn, stride_sk, has_coeff: tl.constexpr, coeff_ptr, FAN_OUT: tl.constexpr,
+    tgt_ptr, stride_tn, stride_ti,
+    grouped_idx_ptr,
+    N, K: tl.constexpr,
+    BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    NO_K_MASK: tl.constexpr
+):
+    pid = tl.program_id(axis=0)
+
+    N_block_id = pid
+    N_blk = N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)
+    N_mask = N_blk < N
+    N_blk = tl.max_contiguous(tl.multiple_of(N_blk % N, BLOCK_N), BLOCK_N)
+    N_idx = tl.load(grouped_idx_ptr + N_blk, mask=N_mask, other=0)
+
+    K_blk = tl.arange(0, BLOCK_K)
+    src_blk_ptrs = src_ptr + (N_idx // FAN_OUT)[:, None] * stride_sn + K_blk[None, :] * stride_sk
+    tgt_blk_ptrs = tgt_ptr + N_blk[:, None] * stride_tn + K_blk[None, :] * stride_ti
+
+    if has_coeff:
+        c = tl.load(coeff_ptr + N_idx, mask=N_mask)[:, None]
+
+    iters = tl.cdiv(K, BLOCK_K)
+    for i in range(0, iters):
+        if NO_K_MASK or i < iters - 1:
+            block = tl.load(src_blk_ptrs, mask=N_mask[:, None])
+            if has_coeff:
+                block *= c
+            tl.store(tgt_blk_ptrs, block, mask=N_mask[:, None])
+
+        else:
+            K_mask = (i * BLOCK_K + K_blk) < K
+            mask = N_mask[:, None] & K_mask[None, :]
+            block = tl.load(src_blk_ptrs, mask=mask)
+            if has_coeff:
+                block *= c
+            tl.store(tgt_blk_ptrs, block, mask=mask)
+        src_blk_ptrs += BLOCK_K * stride_sk
+        tgt_blk_ptrs += BLOCK_K * stride_ti

build/torch-universal/scattermoe/kernels/single.py

@@ -0,0 +1,59 @@
+import torch
+import triton
+import triton.language as tl
+
+@triton.jit
+def _single2scatter(
+    X_ptr, stride_xm, stride_xk,
+    W_ptr, stride_we, stride_wk, stride_wn,
+    Y_ptr, stride_ym, stride_yn,
+    expert_idxs_ptr,
+    FAN_OUT: tl.constexpr,
+    K: tl.constexpr, N: tl.constexpr, E: tl.constexpr,
+    BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
+    ACC_TYPE: tl.constexpr,
+):
+    pid0 = tl.program_id(axis=0)
+    pid1 = tl.program_id(axis=1)
+
+    N_block_id = pid0
+    if FAN_OUT == 1:
+        in_idx = pid1
+    else:
+        in_idx = 0
+    out_idx = pid1
+
+    K_block = tl.arange(0, BLOCK_K)
+    N_block = tl.max_contiguous(tl.multiple_of((N_block_id * BLOCK_N + tl.arange(0, BLOCK_N)) % N, BLOCK_N), BLOCK_N)
+    E_idx = tl.load(expert_idxs_ptr + pid1)
+    X_blk_ptrs = X_ptr + in_idx * stride_xm + K_block[:, None] * stride_xk
+    W_blk_ptrs = W_ptr + E_idx * stride_we + K_block[:, None] * stride_wk + N_block[None, :] * stride_wn
+    acc = tl.zeros((1, BLOCK_N), dtype=ACC_TYPE)
+    for K_block_id in range(0, tl.cdiv(K, BLOCK_K)):
+        x = tl.load(X_blk_ptrs)
+        w = tl.load(W_blk_ptrs)
+        acc += tl.sum(x * w, axis=0)[None, :]
+        X_blk_ptrs += BLOCK_K * stride_xk
+        W_blk_ptrs += BLOCK_K * stride_wk
+    Y_blk_ptrs = Y_ptr + out_idx * stride_ym + N_block[None, :] * stride_yn
+    tl.store(Y_blk_ptrs, acc)
+
+def single2scatter(X, W, expert_idxs):
+    E, xdim, ydim = W.size()
+    k = expert_idxs.size(1)
+    assert X.size(0) == k or X.size(0) == 1
+    Y = torch.empty((k, ydim), device=X.device, dtype=X.dtype)
+    BLOCK_N = 128
+    BLOCK_K = 128
+    grid = ydim // BLOCK_N, k
+    _single2scatter[grid](
+        X, X.stride(0), X.stride(1),
+        W, W.stride(0), W.stride(1), W.stride(2),
+        Y, Y.stride(0), Y.stride(1),
+        expert_idxs,
+        FAN_OUT=Y.size(0) // X.size(0),
+        K=xdim, N=ydim, E=E,
+        BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
+        ACC_TYPE=tl.float32
+    )
+    return Y

build/torch-universal/scattermoe/mlp.py

@@ -0,0 +1,96 @@
+import torch
+from torch import nn
+
+from .parallel_experts import ParallelExperts, flatten_sort_count
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        input_size,
+        hidden_size,
+        num_experts,
+        top_k,
+        bias=False,
+        activation=None,
+    ):
+        super(MLP, self).__init__()
+
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.experts = ParallelExperts(num_experts, input_size, hidden_size, bias=bias)
+        self.output_experts = ParallelExperts(num_experts, hidden_size, input_size, bias=bias)
+        self.top_k = min(top_k, self.num_experts)
+        self.activation = activation
+
+    def extra_repr(self):
+        return 'k={}'.format(self.top_k)
+
+    def forward(self, x: torch.Tensor, expert_p: torch.Tensor, expert_idxs: torch.Tensor):
+        x_shape = x.size()
+        x = x.view(-1, x_shape[-1])
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = \
+            flatten_sort_count(expert_idxs, num_experts=self.num_experts)
+
+        h = self.experts(
+            x, self.top_k,
+            sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            grouped_out=True
+        )
+        h = self.activation(h)
+        y = self.output_experts(
+            h, 1, sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            gates=expert_p,
+        )
+        y = y.view(*x_shape[:-1], y.size(-1))
+        return y
+
+class GLUMLP(nn.Module):
+    def __init__(
+        self, 
+        input_size, 
+        hidden_size,
+        num_experts,
+        top_k,
+        bias=False,
+        activation=nn.SiLU(),
+    ):
+        super(GLUMLP, self).__init__()
+
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.experts = ParallelExperts(num_experts, input_size, 2 * hidden_size, bias=bias)
+        self.output_experts = ParallelExperts(num_experts, hidden_size, input_size, bias=bias)
+        self.top_k = min(top_k, self.num_experts)
+        self.activation = activation
+
+    def extra_repr(self):
+        return 'k={}'.format(self.top_k)
+
+    def forward(self, x: torch.Tensor, expert_p: torch.Tensor, expert_idxs: torch.Tensor):
+        x_shape = x.size()
+        x = x.view(-1, x_shape[-1])
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = \
+            flatten_sort_count(expert_idxs, num_experts=self.num_experts)
+
+
+        h, gates  = self.experts(
+            x, self.top_k,
+            sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            grouped_out=True
+        ).chunk(2, dim=-1)
+        h = self.activation(gates) * h
+        y = self.output_experts(
+            h, 1, sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True,
+            gates=expert_p,
+        )
+        y = y.view(*x_shape[:-1], y.size(-1))
+        return y
+

build/torch-universal/scattermoe/parallel_experts.py

@@ -0,0 +1,182 @@
+import torch
+import torch.nn as nn
+from . import kernels
+from typing import Optional
+
+@torch.library.custom_op("scattermoe::bincount", mutates_args={})
+def compileable_bincount(x: torch.Tensor, minlength: int) -> torch.Tensor:
+        return x.bincount(minlength=minlength)
+
+@compileable_bincount.register_fake
+def _(x: torch.Tensor, minlength: int) -> torch.Tensor:
+    return torch.empty(minlength, dtype=torch.long, device=x.device)
+
+@torch.compile
+def flatten_sort_count(expert_idxs: torch.Tensor, num_experts: int):
+    with torch.no_grad():
+        flattened_expert_idxs = expert_idxs.flatten()
+        sorted_expert_idxs, sorted_scattered_idxs = torch.sort(flattened_expert_idxs)
+        expert_counts = compileable_bincount(flattened_expert_idxs, minlength=num_experts)
+        expert_offsets = expert_counts.cumsum(-1)
+        return sorted_expert_idxs, sorted_scattered_idxs, expert_offsets
+
+
+
+class ParallelLinear(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, 
+        x: torch.Tensor, expert_weights: torch.Tensor, k: int,
+        sorted_expert_idxs: torch.Tensor, sorted_scattered_idxs: torch.Tensor,
+        expert_offsets: torch.Tensor,
+        expert_biases: Optional[torch.Tensor]=None,
+        gates: Optional[torch.Tensor]=None,
+        grouped_in: bool =False, grouped_out: bool=False,
+    ):
+        with torch.device(x.device):
+            output = kernels.ops.scatter2scatter(
+                X=x, W=expert_weights,
+                b=expert_biases, k=k,
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                x_grouped=grouped_in, y_grouped=grouped_out
+            )
+            if gates is not None:
+                output_expanded = output.view(gates.size(0), gates.size(1), output.size(-1))
+                output = (gates.unsqueeze(1) @ output_expanded).squeeze(1)
+            else:
+                output_expanded = None
+
+            ctx.save_for_backward(
+                x, expert_weights,
+                expert_biases,
+                sorted_expert_idxs,
+                sorted_scattered_idxs,
+                expert_offsets,
+                gates,
+                output_expanded
+            )
+            ctx.grouped_in = grouped_in
+            ctx.grouped_out = grouped_out
+            ctx.k = k
+        return output
+    @staticmethod
+    def backward(ctx, grad_out: torch.Tensor):
+        with torch.device(grad_out.device):
+            (x, expert_weights, expert_biases,
+             sorted_expert_idxs,
+             sorted_scattered_idxs,
+             expert_offsets,
+             gates, output_expanded) = ctx.saved_tensors
+            k = ctx.k
+            grouped_in = ctx.grouped_in
+            grouped_out = ctx.grouped_out
+            # print("backward")
+
+            if gates is not None:
+                # calculate gates gradient
+                # d_gates = torch.bmm(output_expanded, grad_out[:, :, None]).squeeze(-1)
+                d_gates = (output_expanded @ grad_out.unsqueeze(-1)).squeeze(-1)
+                gates_flat = gates.flatten()
+                gate_fan = gates.size(1)
+                grouped_grad_out = output_expanded.flatten(0, 1) # reuse expanded buffer later
+            else:
+                d_gates = None
+                gates_flat = None
+                gate_fan = 1
+                grouped_grad_out = None
+
+            if grouped_out:
+                grouped_grad_out = grad_out
+            else:
+                grouped_grad_out = kernels.ops.group(grad_out, sorted_scattered_idxs,
+                                                     fan_out=gate_fan, coeff=gates_flat,
+                                                     out=grouped_grad_out)
+            if grouped_in:
+                grouped_x = x
+                d_expanded_input = None
+            else:
+                grouped_x = kernels.ops.group(x, sorted_scattered_idxs, fan_out=k)
+                d_expanded_input = grouped_x
+
+            d_weights, d_biases = kernels.ops.group_bwd_W(
+                DY=grouped_grad_out, X=grouped_x,
+                expert_offsets=expert_offsets,
+                E=expert_weights.size(0),
+                has_bias=expert_biases is not None
+            )
+
+
+            d_expanded_input = kernels.ops.scatter2scatter(
+                X=grouped_grad_out, x_grouped=True,
+                W=expert_weights.permute(0, 2, 1),
+                sorted_expert_idxs=sorted_expert_idxs,
+                sorted_scattered_idxs=sorted_scattered_idxs,
+                k=1,
+                y_grouped=grouped_in,
+                out=d_expanded_input # Reuse grouped_x buffer
+            )
+
+            if k == 1:
+                d_input = d_expanded_input
+            else:
+                d_input = d_expanded_input.view(x.size(0), k, d_expanded_input.size(-1)).sum(-2)
+        # print("backward end.")
+        return (
+            # x, expert_weights,
+            d_input, d_weights,
+            # k, sorted_expert_idxs, sorted_scattered_idxs, expert_offsets,
+            None, None, None, None, 
+            # bias, gates
+            d_biases, d_gates,
+            # grouped_in, grouped_out,
+            None, None
+        )
+
+def parallel_linear(inputs, expert_weights, k,
+                    sorted_expert_idxs, sorted_scattered_idxs,
+                    expert_offsets,
+                    expert_biases=None,
+                    gates=None, grouped_in=False, grouped_out=False):
+    results = ParallelLinear.apply(inputs, expert_weights, k,
+                                   sorted_expert_idxs, sorted_scattered_idxs,
+                                   expert_offsets,
+                                   expert_biases,
+                                   gates, grouped_in, grouped_out)
+    return results
+
+class ParallelExperts(nn.Module):
+    def __init__(self, num_experts, input_size, output_size, bias=False) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size))
+
+        if bias:
+            self.bias = nn.Parameter(torch.empty(num_experts, output_size))
+        else:
+            self.bias = None
+
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.output_size = output_size
+        self.reset_parameters()
+
+    def extra_repr(self):
+        return 'num_experts={}, input_size={}, output_size={}'.format(
+            self.num_experts, self.input_size, self.output_size)
+
+    def reset_parameters(self) -> None:
+        nn.init.normal_(self.weight, std=0.02)
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(self, inputs, k, sorted_expert_idxs, sorted_scattered_idxs,
+                expert_offsets,
+                gates=None, grouped_in=False, grouped_out=False):
+
+        results = parallel_linear(
+            inputs, self.weight.permute(0, 2, 1), k,
+            sorted_expert_idxs, sorted_scattered_idxs, expert_offsets,
+            expert_biases=self.bias,
+            gates=gates, grouped_in=grouped_in, grouped_out=grouped_out
+        )
+        return results

build/torch-universal/scattermoe/utils/__init__.py

@@ -0,0 +1,3 @@
+from . import replace_moe
+
+__all__ = ['replace_moe']

build/torch-universal/scattermoe/utils/replace_moe.py

@@ -0,0 +1,209 @@
+import torch
+from .. import parallel_linear, flatten_sort_count
+from torch.nn import functional as F
+from torch import nn
+
+import logging
+
+def replace_function(cls, fun_name):
+    def decorator(fun):
+        def _fun(*args, **kwargs):
+            filename = fun.__code__.co_filename
+            name = fun.__name__
+            logging.info(f"Replacing `{cls.__name__}.{fun_name}` with {filename}:{name}")
+            setattr(cls, fun_name, fun)
+            return fun(*args, **kwargs)
+        setattr(cls, fun_name, _fun)
+    return decorator
+
+try:
+    from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts
+    @replace_function(cls=GptOssExperts, fun_name='forward')
+    def gpt_oss_forward(self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None) -> torch.Tensor:
+        batch_size = hidden_states.shape[0]
+        hidden_states = hidden_states.reshape(-1, self.hidden_size)  # (num_tokens, hidden_size)
+        k = router_indices.shape[1]
+        selected_weights = torch.gather(routing_weights, dim=1, index=router_indices)
+        router_indices = router_indices.flatten()
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = \
+            flatten_sort_count(router_indices, num_experts=self.num_experts)
+
+        gate_up = parallel_linear(
+            hidden_states, self.gate_up_proj, k,
+            sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            expert_biases=self.gate_up_proj_bias,
+            grouped_in=False, grouped_out=True, 
+        )
+
+        gate, up = gate_up[..., ::2], gate_up[..., 1::2]
+        gate = gate.clamp(min=None, max=self.limit)
+        up = up.clamp(min=-self.limit, max=self.limit)
+        glu = gate * torch.sigmoid(gate * self.alpha)
+        gated_output_ = (up + 1) * glu
+
+        out_scattered = parallel_linear(
+            gated_output_, self.down_proj, 1,
+            sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            expert_biases=self.down_proj_bias,
+            grouped_in=True, grouped_out=False,
+            gates=selected_weights,
+        )
+
+        next_states = out_scattered.view(batch_size, -1, self.hidden_size)
+        return next_states
+except Exception:
+    logging.info("Failed to replace GptOssExperts")
+
+
+try: 
+    from transformers.models.granitemoehybrid.modeling_granitemoehybrid import GraniteMoeHybridMoE
+    @replace_function(cls=GraniteMoeHybridMoE, fun_name='forward')
+    def granite_moe_forward(self, layer_input):
+        bsz, length, emb_size = layer_input.size()
+        layer_input = layer_input.reshape(-1, emb_size)
+        # compute the top_k routing decision
+        router_logits = self.router.layer(layer_input)
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.router.top_k, dim=-1)
+        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(layer_input.dtype)
+        sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = \
+            flatten_sort_count(selected_experts, num_experts=self.router.num_experts)
+
+        # compute experts
+        gates, h = parallel_linear(
+            layer_input, self.input_linear.weight.transpose(2, 1),
+            self.router.top_k,
+            sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=False, grouped_out=True,
+        ).chunk(2, dim=-1)
+        h = self.activation(gates) * h
+        layer_output = parallel_linear(
+            h, self.output_linear.weight.transpose(2, 1),
+            1,
+            sorted_expert_idxs, sorted_scattered_idxs,
+            expert_offsets,
+            grouped_in=True, grouped_out=False,
+            gates=routing_weights
+        )
+        layer_output = layer_output.view(bsz, length, emb_size)
+        return layer_output, router_logits
+except Exception:
+    logging.info("Failed to replace GraniteMoeHybridMoE")
+
+# TODO consolidating params into tensor. OOMs.
+# try: 
+#     from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeSparseMoeBlock
+#     fun = Qwen3MoeSparseMoeBlock.__init__
+#     def qwen3moe__init__(self: Qwen3MoeSparseMoeBlock, config):
+#         fun(self, config)
+
+#         weight_refs = {}
+#         def assemble_weights(self):
+#             down_proj_weights = [weight_refs[i, "down_proj"] for i in range(len(self.experts))]
+#             up_proj_weights = [weight_refs[i, "up_proj"] for i in range(len(self.experts))]
+#             gate_proj_weights = [weight_refs[i, "gate_proj"] for i in range(len(self.experts))]
+#             if (all(w is not None for w in down_proj_weights) and
+#                 all(w is not None for w in up_proj_weights) and
+#                 all(w is not None for w in gate_proj_weights)):
+#                 self.act_fn = self.experts[0].act_fn
+#                 mid_size = up_proj_weights[0].size(0)
+#                 self.down_proj = nn.Parameter(torch.stack(down_proj_weights))
+#                 gate_up_proj = torch.empty(
+#                     self.num_experts, 2 * mid_size, config.hidden_size,
+#                     dtype=up_proj_weights[0].dtype
+#                 )
+#                 for i in range(self.num_experts - 1, -1, -1):
+#                     gate_up_proj[i, :mid_size] = gate_proj_weights[i]
+#                     gate_up_proj[i, mid_size:] = up_proj_weights[i]
+#                     del self.experts[i]
+#                 self.gate_up_proj = nn.Parameter(gate_up_proj)
+#                 del self.experts
+        
+#         for i, e in enumerate(self.experts):
+#             def weight_tracker(expert_id, expert_weight_name):
+#                 def fun(module, err_msgs):
+#                     id_tup = (expert_id, expert_weight_name)
+#                     assert id_tup in weight_refs
+#                     weight_refs[id_tup] = module.weight
+#                     assemble_weights(self)
+#                 return fun
+#             e.gate_proj.register_load_state_dict_post_hook(weight_tracker(i, "gate_proj"))
+#             weight_refs[i, "gate_proj"] = None
+#             e.up_proj.register_load_state_dict_post_hook(weight_tracker(i, "up_proj"))
+#             weight_refs[i, "up_proj"] = None
+#             e.down_proj.register_load_state_dict_post_hook(weight_tracker(i, "down_proj"))
+#             weight_refs[i, "down_proj"] = None
+
+#         # def hook(module, state_dict, prefix, local_metadata):
+#         #     local_keys = [k for k in state_dict.keys() if k.startswith(prefix)]
+#         #     expert_keys = [k for k in local_keys if 'mlp.' in k]
+#         #     if len(expert_keys) == 0: # deleted
+#         #         assert prefix + "down_proj" in state_dict
+#         #         assert prefix + "gate_up_proj" in state_dict
+#         # self.register_state_dict_post_hook(hook)
+    
+#     state_dict_fun = Qwen3MoeSparseMoeBlock.state_dict
+#     def qwen3moe_state_dict(self: Qwen3MoeSparseMoeBlock, *args, destination=None, prefix="", keep_vars=False):
+#         destination = state_dict_fun(self, *args, destination, prefix, keep_vars)
+#         local_keys = [k for k in destination.keys() if k.startswith(prefix)]
+#         if ((prefix + "down_proj" in local_keys) and (prefix + "gate_up_proj" in local_keys)):
+#             # need to break down
+#             for i in range(self.num_experts):
+#                 down_proj_name = prefix + f"experts.{i}.down_proj.weight"
+#                 up_proj_name = prefix + f"experts.{i}.up_proj.weight"
+#                 gate_proj_name = prefix + f"experts.{i}.gate_proj.weight"
+#                 destination[down_proj_name] = destination[prefix + "down_proj"][i]
+#                 gate_proj, up_proj = destination[prefix + "gate_up_proj"][i].chunk(2, dim=1)
+#                 destination[gate_proj_name] = gate_proj.contiguous()
+#                 destination[up_proj_name] = up_proj.contiguous()
+#             del destination[prefix + "down_proj"]
+#             del destination[prefix + "gate_up_proj"]
+
+#         return destination
+#     Qwen3MoeSparseMoeBlock.state_dict = qwen3moe_state_dict
+
+
+#     Qwen3MoeSparseMoeBlock.__init__ = qwen3moe__init__
+
+#     def qwen3moe_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+#         """ """
+#         batch_size, sequence_length, hidden_dim = hidden_states.shape
+#         hidden_states = hidden_states.view(-1, hidden_dim)
+#         # router_logits: (batch * sequence_length, n_experts)
+#         router_logits = self.gate(hidden_states)
+
+#         routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+#         routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+#         if self.norm_topk_prob:  # only diff with mixtral sparse moe block!
+#             routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+#         # we cast back to the input dtype
+#         routing_weights = routing_weights.to(hidden_states.dtype)
+#         sorted_expert_idxs, sorted_scattered_idxs, expert_offsets = flatten_sort_count(selected_experts, num_experts=self.num_experts)
+
+#         gate_up = parallel_linear(
+#             hidden_states, self.gate_up_proj.transpose(1, 2), self.top_k,
+#             sorted_expert_idxs, sorted_scattered_idxs,
+#             expert_offsets,
+#             grouped_in=False, grouped_out=True, 
+#         )
+
+#         _gate, up = gate_up.chunk(2, dim=-1)
+#         intermediate = self.act_fn(_gate) * up
+
+#         final_hidden_states = parallel_linear(
+#             intermediate, self.down_proj.transpose(1, 2), 1,
+#             sorted_expert_idxs, sorted_scattered_idxs,
+#             expert_offsets,
+#             grouped_in=True, grouped_out=False,
+#             gates=routing_weights,
+#         )
+#         final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+#         return final_hidden_states, router_logits
+#     Qwen3MoeSparseMoeBlock.forward = qwen3moe_forward
+# except Exception as e:
+#     pass
+