perf(metax_v2): fuse HF RMSNorm/RoPE/softmax/SwiGLU launch-heavy chains - #1

Open

shallitbeso wants to merge 8 commits into

ZavierHung:metax-v2from

shallitbeso:metax-v2

Commits on Jul 15, 2026

fix(metax_v2): mean kernel early-exit + shfl cross-slot + warpSize=32 bugs

zhangliheng
and
claude
committed
bae4fdb
feat(metax_v2): fuse HF RMSNorm into single flagos::rms_norm kernel

zhangliheng
and
claude
committed
5d71415
feat(metax_v2): fuse HF RoPE and masked-softmax into single kernels

zhangliheng
and
claude
committed
16e7d42
feat(metax_v2): fuse SwiGLU and residual-add+RMSNorm into single kernels

zhangliheng
and
claude
committed
31f0cfa
perf(metax_v2): drop RoPE cos/sin broadcast; fuse qkv/gate_up GEMMs

zhangliheng
and
claude
committed
53da4a3
fix(metax_v2): grid-stride reduction kernels over rows

zhangliheng
and
claude
committed
d43bafc
fix(optimizations): guard import-hook recursion; note inference-only caches

zhangliheng
and
claude
committed
8bf4b21
test(metax_v2): add correctness tests for the five fused ops

zhangliheng
and
claude
committed
cc69422