Views
No views yet
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).torch.library ops (loadable via kernels). Reference op: transformers HYV3Experts.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).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 |
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 |
1import torch
2from kernels import get_kernel
3
4kernel = get_kernel("aymous/hy3-moe-experts", version=1, trust_remote_code=True)
5# routing (sigmoid top-8, scaled) is computed outside the op, as in HYV3Moe:
6# 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).