aymous/hy3-moe-experts

Hy3 MoE experts

Fused MoE routed experts for tencent/Hy3 (hy_v3, 299B-A21B): 192 routed experts, top-8, plain SiLU SwiGLU, router-weight scale + scatter-add. The router (HYV3TopKRouter: linear + sigmoid + top-8, normalized by sum, scaled by router_scaling_factor=2.826) and the shared expert stay outside the op — already-optimal cuBLAS/linear paths. Two entry points, both consuming (hidden_states, top_k_index, top_k_weights): hyv3_moe_experts (grouped GEMM, prefill) and hyv3_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 kernels). Reference op: transformers HYV3Experts.

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

Mean latency (torch.utils.benchmark blocked_autorange); op_eager/op_compile are the torch reference, custom is this kernel. Hidden 4096, per-expert intermediate 1536, 192 experts top-8. Verified vs the reference (atol=rtol=3e-2: the fp32 kernel accumulation vs the bf16 reference index_add_, amplified by the 2.826 router scaling — the kernel is the more accurate side).

Prefill — n_tokens = 2048, hyv3_moe_experts: 47.66 ms, 1.39× vs op_compile, verified.

workload	mean (ms)
op_eager (reference)	65.48
op_compile	66.47
custom	47.66

Decode — n_tokens = 1, hyv3_moe_experts_decode: 1.315 ms, 1.90× vs op_compile, verified.

workload	mean (ms)
op_eager (reference)	2.427
op_compile	2.505
custom	1.315

Usage

import torch
from kernels import get_kernel

kernel = get_kernel("aymous/hy3-moe-experts", version=1, trust_remote_code=True)
# routing (sigmoid top-8, scaled) is computed outside the op, as in HYV3Moe:
#   out = kernel.hyv3_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 × 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).