Metadata-Version: 2.4
Name: moe-ops
Version: 0.1.0
Summary: Pure-Triton MoE operators: grouped GEMM, permute/unpermute, and fused kernels
Author-email: echogujy <echogujy@users.noreply.github.com>
License-Expression: MIT
Project-URL: Homepage, https://github.com/echogujy/moe_ops
Project-URL: Repository, https://github.com/echogujy/moe_ops
Project-URL: Bug Tracker, https://github.com/echogujy/moe_ops/issues
Keywords: moe,mixture-of-experts,triton,gpu,grouped-gemm,transformer
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3 :: Only
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: triton
Dynamic: license-file

# moe-ops

Pure-Triton implementations of Mixture-of-Experts (MoE) operator primitives — **grouped GEMM**, **token permute / unpermute**, and **fused top-k softmax** — designed for **training (forward + backward)**. Every operator is a differentiable `torch.autograd.Function` whose backward pass is also written in Triton, so gradients flow through the whole MoE layer end-to-end. No C++/CUDA extensions required — every kernel is a Triton `@triton.jit` kernel.

---

## Quick Install

Requires Python >= 3.10, PyTorch (CUDA build) + Triton, and a CUDA GPU (sm80+).

```bash
cd moe_ops
pip install -e .
```

---

## Quick Usage

```python
import torch
from moe_ops import fused_topk_softmax, permute, grouped_gemm, unpermute

# Fused router: top-k + softmax in one kernel
weights, indices = fused_topk_softmax(logits, K=6)

# Permute tokens by expert (counting sort)
permuted, row_id_map, offsets = permute(x, indices.to(torch.int32), E=64)

# Grouped GEMM (gate/up/down FFN) — all metadata stays on-GPU
gate = grouped_gemm(permuted, w_gate, offsets, trans_b=True)

# Recover original token order, weighted by routing prob
out = unpermute(down, row_id_map, weights, num_tokens=x.shape[0], num_topK=6)
```

---

## Why this library

MoE layers are built from a handful of primitive operators. PyTorch implements
each of them as *separate, unfused* library calls, and each one ships a
baseline that leaves measurable performance on the table. This library
re-implements the four MoE primitives in Triton with a specific optimization
per op:

| MoE primitive           | PyTorch baseline                                                                                                                                                                                                                                                        | Ours (Triton)                                                                                                                                                                                |
| ----------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| **top-k softmax** | `torch.topk` then `F.softmax` — **two separate passes**, intermediate top-k weights written to/read from global memory                                                                                                                                       | **Fused** single kernel: top-k selection + softmax done entirely in registers/SRAM, no global round-trip. Backward computes exact softmax gradients in one pass.                       |
| **permute**       | `torch.argsort` (sort) + index copy — **O(N·log N)** sort + a separate gather/copy                                                                                                                                                                            | **Counting sort + fused copy** — uses block-local `tl.histogram` + exclusive prefix sums, **O(N)**, and writes the permuted output in one pass. Lower complexity, one kernel. |
| **grouped GEMM**  | `torch.nn.functional.grouped_mm` — on **SM80 it drops to cuBLAS's grouped GEMM**, which needs the per-expert batch sizes on the **CPU (D2H sync)** and copies boundaries; the D2H dependency stalls the GPU pipeline. (SM90 dispatch is less clear-cut.) | Pure-Triton kernels keep all metadata**on-GPU** — `offsets` stays resident, **no CPU/GPU D2H round-trip** on the critical path.                                               |
| **unpermute**     | `index_select` + scatter-add — **atomic / index add on global memory**, one (or more) read-modify-write round-trips per token                                                                                                                                  | **Gather-add on SRAM**: rows are gathered into on-chip buffers, accumulated, and written back to global memory **once**.                                                         |

Net effect: the routing stages (top-k, permute, unpermute) — which are pure
memory movement and dominate the *non-GEMM* time in a PyTorch MoE — collapse
to a fraction of their original cost, and grouped GEMM no longer waits on CPU
metadata.

---

## Layout conventions

All ops share the stacked MoE layout:

| tensor      | shape                                           | meaning                                                                                       |
| ----------- | ----------------------------------------------- | --------------------------------------------------------------------------------------------- |
| `A`       | `[total_tokens, K]`                           | activations, all experts stacked along M                                                      |
| `B`       | `[E, N, K]` if `trans_b` else `[E, K, N]` | per-expert weights                                                                            |
| `offsets` | `[E + 1]` or `[E]`                          | cumulative group ends (`offsets[e]` = end row of expert `e`; starts with 0 for `E + 1`) |
| `C`       | `[total_tokens, N]`                           | stacked output                                                                                |

`bf16` and `fp16` are supported (bf16 is the MoE default).

---

## Operators

### Grouped GEMM (`grouped_gemm_ops.py`, `grouped_gemm_ops_sm90.py`)

- `grouped_gemm(A, B, offsets, trans_b=True) -> C` — Differentiable wrapper (`torch.autograd.Function`) supporting automatic architecture dispatch.
  - **SM90+ (Hopper)**: Dispatched to the TMA-accelerated grouped GEMM kernel using asynchronous bulk copy hardware (`tl.make_tensor_descriptor`) with built-in hardware zero-fill out-of-bounds boundary handling, avoiding CPU/GPU formatting overhead.
  - **SM80 (Ampere/Lovelace)**: Dispatched to the standard 3D Grid grouped GEMM kernel employing an optimized **3D Grid SM Layout** to avoid SM-CTA persistence bottlenecks.
  - Backward pass computes `grad_A` via parallel Triton GMM kernels (with TMA on Hopper) and `grad_B` via Triton GMM B-gradient kernel.

### Routing permute / unpermute (`permute_ops.py`, `unpermute_ops.py`)

- `permute(input, indices, num_out_tokens=0, E=None) -> (permuted, row_id_map, base)` — Differentiable wrapper employing a high-throughput **counting sort** algorithm (uses block-local hardware `tl.histogram` and exclusive prefix sums) to group token rows by expert.
- `unpermute(input, row_id_map, prob, num_tokens, num_topK, max_tokens=-1) -> out` — Recovers original token order using the inverted gather map and scales by routing probabilities.
- `permute_backward(grad_permuted, row_id_map, num_tokens, topK) -> grad_input` — Computes permute backward pass using scatter-add.

### Fused Top-K Softmax (`fused_topk_softmax_ops.py`)

- `fused_topk_softmax(logits, K, fp32_routing=False) -> (weights, indices)` — Fuses top-K selection and softmax normalization entirely in GPU registers/SRAM, avoiding intermediate global memory roundtrips. Backward pass computes exact softmax gradients and scatters them back to input gradients.
  - **`fp32_routing` (bool)**: If `True`, forces routing weight calculations and outputs to be done in high-precision `float32` regardless of activation dtype (e.g. `bfloat16`/`float16`). This prevents underflow/overflow training instability commonly observed in deep MoE models.

---

## Repository Layout

```
moe_ops/
├── moe_ops/                       # importable package
│   ├── __init__.py               # Unified exposed public APIs
│   ├── grouped_gemm_ops.py       # Standard 3D Grid Grouped GEMM
│   ├── grouped_gemm_ops_sm90.py  # Hopper TMA Grouped GEMM & Dispatcher
│   ├── permute_ops.py            # Primary Counting-Sort Permute
│   ├── unpermute_ops.py          # Primary Unpermute
│   └── fused_topk_softmax_ops.py # Fused Top-K Softmax
└── tests/
    ├── test_correctness.py       # All-op correctness harness (vs PyTorch refs)
    └── test_benchmark.py         # MoE breakdown perf vs torch.grouped_mm table
```

---

## Running Tests

Two consolidated scripts cover correctness and performance:

```bash
# 1. Correctness harness for all ops (vs PyTorch / f32 loop references)
#    --smoke uses small shapes for low-VRAM environments.
python tests/test_correctness.py [--smoke]

# 2. MoE forward-pass breakdown: moe_ops vs torch.grouped_mm, table output
#    --small uses E=16 topk=2 for low-VRAM environments.
python tests/test_benchmark.py [--small]
```
