aymous/unlimited-ocr-moe-experts

DeepSeek-V2 MoE experts

Fused MoE routed experts for the DeepSeek-V2 architecture — the language backbone of baidu/Unlimited-OCR: 64 routed experts, top-6, plain SiLU SwiGLU, router-weight scale + scatter-add. The router (softmax top-6, DeepseekV2Moe.route_tokens_to_experts, norm_topk_prob=False) and the 2 shared experts stay outside the op — a single cuBLAS MLP and a tiny linear, already optimal. Both entry points consume (hidden_states, top_k_index, top_k_weights): deepseekv2_moe_experts (grouped GEMM, prefill shapes) and deepseekv2_moe_experts_decode (fused gather-GEMV for n_tokens ~ 1–4 — two persistent-shape kernels, no per-expert cuBLAS loop, no D2H sync).

Built with kernel-builder as native torch.library ops (loadable via the kernels library). Reference op: transformers DeepseekV2Experts.

Benchmark (NVIDIA GB10 / sm_121, aarch64, CUDA 13.0, torch 2.12, bf16)

Mean latency (torch.utils.benchmark blocked_autorange) — op_eager / op_compile are the torch reference (eager and torch.compile); custom is this kernel. Hidden 1280, per-expert intermediate 896, 64 experts top-6. Correctness verified against the reference (atol=rtol=2e-2).

Prefill shape — n_tokens = 2048, deepseekv2_moe_experts: 4.191 ms, 2.26× vs op_compile, verified.

workload	mean (ms)
op_eager (reference)	9.119
op_compile	9.490
custom	4.191

Decode shape — n_tokens = 1, deepseekv2_moe_experts_decode: 0.183 ms, 4.40× vs op_compile, verified.

workload	mean (ms)
op_eager (reference)	0.770
op_compile	0.805
custom	0.183

Usage

import torch
from kernels import get_kernel

kernel = get_kernel("aymous/unlimited-ocr-moe-experts", version=1)
# routing (softmax top-6) is computed outside the op, as in DeepseekV2Moe:
#   out = kernel.deepseekv2_moe_experts_decode(x, gate_up_proj, down_proj, top_k_index, top_k_weights, top_k)

Built for aarch64 across torch 2.11/2.12 x CUDA 12.6/12.8/13.0/13.2, GPU archs 7.5-9.0 + 10.0/11.0/12.0/12.1+PTX (runs Turing through Blackwell / NVIDIA GB10). Benchmarked on GB10 (sm_121).