perf(metax_v2): fuse HF RMSNorm/RoPE/softmax/SwiGLU launch-heavy chains by shallitbeso · Pull Request #1 · ZavierHung/PyTorch-Plugin-FL

shallitbeso · 2026-07-15T04:03:15Z

perf(metax_v2): fuse HF RMSNorm/RoPE/softmax/SwiGLU launch-heavy chains (Qwen3 decode +100%)

背景

MetaX C500 上 Qwen3-0.6B 的 eager decode 是 launch-overhead-bound：
745 launches/token，每 launch 3µs CPU，GPU 仅 50% busy。HF 的 RMSNorm/RoPE/
softmax/SwiGLU 都是手写的多算子链，每次调用发 5~14 个独立 aten launch，是主要 tps 成本。

本 PR 把这些链路各自融成单个 flagos::* 自定义 kernel（warpSize=64，fp32 内部累加，
fp16 in/out），通过 opt-in env 开关（默认关闭，行为不变）注入。全部与 stock HF token-exact。

改动

Commit	融合项	端到端 tps 效果
`bae4fdb`	mean kernel 正确性修复（warpSize=64 / shfl cross-slot / early-exit）	不影响 tps，投产前置修复
`c1de067`	融合 HF RMSNorm → `flagos::rms_norm`	+46.5%（58.51→85.72 tok/s，首次越过原生上界 68.64）
`ed823e2`	融合 RoPE + masked-softmax	RoPE +16.5%（→99.85，累计 +70.5%）；softmax +23.4%
`07afde0`	融合 SwiGLU + residual-add+RMSNorm	SwiGLU +3.3%；add_rms_norm +0.1%（噪声内）
`4a9b0d2`	RoPE 消 cos/sin broadcast + qkv/gate_up GEMM 合并	rope CPU −47%（147.9→78ms）；GEMM 削 4452 launch

累计：五融合全开相对 stock HF ≈ +100%，是 muxi 原生上界的约 2×。单 op 融合边际收益
序列 +46.5% → +16.5% → +23.4% → +3.3% → +0.1%，趋于收敛。

实现方式

每个融合：kernel .cuh + host .cu 拆分（host C++ 不拉入 .cuh，避开 mxcc）；
register.cc 注册 flagos:: custom op（HF 这些不是 aten op，无法经 dispatcher 拦截）；
cmake/MetaxKernels.cmake 把新 kernel 加入 v2 编译集。
Python 侧 torch_fl/optimizations.py：统一 import hook + monkey-patch HF 模块，
由 env 开关门控（FLAGOS_RMSNORM_FUSED / FLAGOS_ROPE_FUSED / FLAGOS_SOFTMAX_FUSED /
FLAGOS_SWIGLU_FUSED / FLAGOS_QKV_FUSED / FLAGOS_GATEUP_FUSED），默认全关。

数值验证

全部融合 fp16 rel err ≤ 8e-3（swiglu 9.4e-5；add_rms_norm new_residual bit-exact；
rope ≤1.6e-4；rmsnorm 3.9e-3），decode 输出与 stock HF token-exact。

测量协议

子进程隔离低噪声（8 warmup + 10 timed + discard-2 median，16-token greedy，cooldown 5s），
所有变体同一环境。热降频窗口下 softmax/swiglu 用同会话交错对照 ratio。
硬件：MetaX C500 单卡（warpSize=64）；模型 Qwen3-0.6B fp16 eager。

… bugs mean_kernel_v2.cuh had three coexisting bugs in MeanAlongDimBlockReduceBody and MeanAlongDimSplitKReduceBody that together made mean.dim return ~1/reduce instead of the true mean on MetaX C500 (wavefront=64): - early-exit guard `if (outer_idx>=outer || inner_idx>=inner) return;` killed most threads before the shuffle reduce when inner<blockDim, leaving shfl with no data. Replaced with `const bool valid=...`; all threads participate in the reduce (out-of-bounds contribute 0), only the write-back is guarded. - `inner_idx = blockIdx.y*blockDim.x + threadIdx.x` packed BLOCK inner slots per block, so __shfl_down_sync mixed sums from different slots (wrong even when inner>=BLOCK). Changed to `inner_idx = blockIdx.y` (one block per slot) with launch geometry (outer, inner, ...) instead of (outer, ceil(inner/BLOCK), ...). - __shfl_down_sync(0xffffffffu, ..., 16) / `lane=tid&31` / `WARPS=BLOCK/32` hardcoded warpSize=32, dropping lanes 32..63 on wavefront=64 hardware. Switched to `constexpr int WARP=64`, 64-bit masks, offset=WARP/2. Also fixed the SplitK cross-warp reduce `offset=(BLOCK/64)` typo -> WARPS/2. LaunchMeanAlongDim: added `inner <= kMaxGridV2` guard before BlockReduce, since the new per-slot geometry grows grid Y by up to 256x and large inner (e.g. mean(x, dim=0) on [4,100000]) would exceed the 65535 grid-Y limit and be silently truncated; falls back to FlatKernel in that case. cat_v2.cc: narrow the contiguous fast-path to require pre_dim==1 in addition to post_dim==1. The old condition (post_dim==1 only) mis-handled cat on the innermost dim with multiple rows (pre_dim>1): inputs are contiguous but row-interleaved in the output, so a single bulk memcpy wrote input 0's full N rows then input 1's N rows instead of interleaving per row. Verified on MetaX C500 with backends_metax_v2.conf: - mean(arange(1,129)) = 64.5 (was 0.0078125), dispatch -> metax_v2 - tests/manual/test_mean_v2_warpsize64_bug.py: 8 passed (was 7+1 xfail) - tests/integration/ops/test_mean_dispatch.py -m "anyplatform or metax": 5 passed - tests/integration/ops/test_cat_dispatch.py (non-ascend): all passed - Qwen3-0.6B decode no longer outputs all-'\!' (mean->RMSNorm chain restored) See docs/实验报告_mean_v2_warpsize64_bug_2026-07-06.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

HF *RMSNorm.forward is a hand-written pow/mean/add/rsqrt/mul chain that issues 5+ separate aten launches per call. On the launch-overhead-bound MetaX v2 backend (~3us CPU/launch, ~2.7k RMSNorm calls per 64-token Qwen3 decode) that launch count is the dominant tps cost. Add a fused single-launch kernel and wire it up end to end: - backends/metax_v2/rmsnorm_kernel_v2.cuh: two-pass row kernel, fp32 reduction (warpSize=64), input/output kept in model dtype. - rmsnorm.h / rmsnorm_v2.cu: host entry, split so register.cc (host C++) binds the op without pulling in the .cuh (mxcc). - register.cc: register flagos::rms_norm custom op (HF RMSNorm is not an aten op, so it cannot be intercepted through the dispatcher). - MetaxKernels.cmake: add rmsnorm_v2 to the v2 kernel set. - torch_fl/optimizations.py: opt-in monkey-patch of HF *RMSNorm.forward (import hook + explicit apply), gated by FLAGOS_RMSNORM_FUSED=1. - profile_infer_flagos.py: opt-in entry for benchmarking. Numerically equivalent to the HF fp32 path. Default behaviour unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Extend the RMSNorm fusion approach to the two remaining launch-heavy sub-chains in the Qwen3 eager decode path, both patched as module-level HF functions (not class methods): - RoPE: apply_rotary_pos_emb's unsqueeze + rotate_half + mul + add chain (12-14 launches/layer) -> single flagos::rope launch. - masked softmax: eager_attention_forward's scale + mask-add + softmax(fp32) + cast sub-chain (6-10 launches/layer) -> single flagos::masked_softmax launch (fp32 reduction, model-dtype output). Wire-up mirrors RMSNorm: kernel .cuh + host .cu split, flagos:: custom op registration in register.cc, v2 kernel set in MetaxKernels.cmake, and opt-in monkey-patch entries (FLAGOS_ROPE_FUSED / FLAGOS_SOFTMAX_FUSED). Collapse the three per-fusion import hooks into one _FusionImportHook: three cooperating meta path finders mutually re-entered on shared target modules and deadlocked modelling_qwen3 import when all env vars were set. The single finder targets the union of modules and applies every enabled patch in one exec_module wrap. Numerically equivalent to the HF fp32 path; default behaviour unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Add two fusions following the existing custom-op + Python-patch pattern: - flagos::swiglu — fuse HF gated-MLP's silu(gate)*up (a silu launch + an elementwise-mul launch) into one kernel. silu in fp32, product cast back to model dtype. Patches MLP.forward for silu-gated families (Qwen3/Qwen2/Llama/ Mistral); gelu/no-gate families intentionally excluded. - flagos::add_rms_norm — fuse the DecoderLayer post-attention residual + attn_out and post_attention_layernorm into one kernel returning (normed, new_residual). new_residual is rounded to fp16 before reuse, matching HF bit-exactly. Both cut launches on the launch-overhead-bound MetaX v2 backend. Measured on Qwen3-0.6B (interleaved ratio, cold GPU): SwiGLU +3.3% (real, cuts 2 kernels/ layer); add_rms_norm +0.1% (within noise — residual-add is a single ~3.8us elementwise, lower leverage). Both token-exact vs stock HF. Numerics (fp16 rel err <= 8e-3): swiglu 9.4e-5; add_rms_norm new_residual bit-exact, normed 1.1e-4. New check_swiglu_add_rmsnorm_numeric.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Two launch-cutting changes on the launch-overhead-bound MetaX v2 backend, both token-exact vs stock HF on Qwen3-0.6B. RoPE host-wrapper (rope_kernel_v2.cuh): ApplyRoPEOne previously expanded cos/sin to the full [batch,heads,seq,d] and materialized them via .contiguous() — 4 such allocations per layer (q/k x cos/sin). This was the sole top10 op with CPU >> CUDA (147.9ms CPU vs 11.7ms CUDA). Now the kernel broadcasts over the head dim internally (derive (batch,seq) from the q row index -> cos_row = batch*seq + seq_idx) and cos/sin stay compact [batch,1,seq,d]. Trace: flagos::rope CPU 147.9 -> 78.0ms (-47%), call count and CUDA time unchanged. Numerics: new check_rope_numeric.py covers decode/prefill, batch 1/4, GQA — fp16 rel err <= 1.6e-4. qkv/gate_up GEMM merge (optimizations.py, pure Python patch): q/k/v share the same input and gate/up are two halves, so their weights concat along out-dim into one GEMM + split (weight cached on the module, GQA split sizes handled; gate_up folds in the fused SwiGLU). Cuts 3+2 GEMM launches/layer to 1+1. Trace (5-fuse vs 7-fuse): aten::linear/mm/matmul each -4452 (=3x1484), linear+mm CPU -82.9ms. End-to-end marginal delta is within noise (each merged GEMM saves one dispatch; GEMM GPU time unchanged), matching the diminishing-returns pattern — kept behind FLAGOS_QKV_FUSED / FLAGOS_GATEUP_FUSED toggles. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Copilot

Pull request overview

This PR introduces an opt-in fusion layer for MetaX v2 to reduce HuggingFace Qwen3 (and related model families) decode launch overhead by replacing several launch-heavy Python/aten chains with single torch.ops.flagos::* custom kernels, installed via a unified import hook and gated by environment variables (default off).

Changes:

Add a unified Python monkey-patching/import-hook system (torch_fl/optimizations.py) to optionally replace HF RMSNorm, RoPE, masked softmax, SwiGLU, add+RMSNorm, and QKV / gate+up projection patterns with fused torch.ops.flagos::* calls.
Register new flagos custom ops in the C++ extension and add new MetaX v2 kernel implementations for RMSNorm, RoPE, masked softmax, SwiGLU, and add+RMSNorm.
Update MetaX v2 kernels/build lists and fix/adjust existing kernels (mean warpSize=64 correctness; cat fast-path gating).

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
torch_fl/optimizations.py	Adds env-gated import hook + monkey patches for fused RMSNorm/RoPE/softmax/SwiGLU/add+RMSNorm/QKV/gate_up paths.
torch_fl/init.py	Re-exports optimization patch APIs from `torch_fl.optimizations`.
csrc/aten/register.cc	Registers new `torch.ops.flagos.*` custom ops and binds PrivateUse1 implementations.
csrc/aten/rmsnorm.h	Declares host entry for fused RMSNorm op.
csrc/aten/rope.h	Declares host entry for fused RoPE op.
csrc/aten/masked_softmax.h	Declares host entry for fused masked-softmax op.
csrc/aten/swiglu.h	Declares host entry for fused SwiGLU op.
csrc/aten/add_rms_norm.h	Declares host entry for fused residual-add + RMSNorm op.
csrc/aten/backends/metax_v2/rmsnorm_v2.cu	mxcc wrapper for RMSNorm v2 kernel.
csrc/aten/backends/metax_v2/rmsnorm_kernel_v2.cuh	Implements RMSNorm v2 CUDA kernel.
csrc/aten/backends/metax_v2/rope_v2.cu	mxcc wrapper for RoPE v2 kernel.
csrc/aten/backends/metax_v2/rope_kernel_v2.cuh	Implements RoPE v2 CUDA kernel (compact cos/sin, no head broadcast materialization).
csrc/aten/backends/metax_v2/softmax_v2.cu	mxcc wrapper for masked softmax v2 kernel.
csrc/aten/backends/metax_v2/softmax_kernel_v2.cuh	Implements fused scale+mask+softmax+cast kernel.
csrc/aten/backends/metax_v2/swiglu_v2.cu	mxcc wrapper for SwiGLU v2 kernel.
csrc/aten/backends/metax_v2/swiglu_kernel_v2.cuh	Implements fused SwiGLU elementwise kernel.
csrc/aten/backends/metax_v2/add_rms_norm_v2.cu	mxcc wrapper for add+rmsnorm v2 kernel.
csrc/aten/backends/metax_v2/add_rms_norm_kernel_v2.cuh	Implements fused residual-add + RMSNorm kernel (returns normed + new residual).
csrc/aten/backends/metax_v2/mean_kernel_v2.cuh	Fixes mean reduction correctness for warpSize=64 and adjusts grid mapping.
csrc/aten/backends/metax_v2/cat_v2.cc	Fixes cat memcpy fast-path to avoid incorrect layout when pre_dim>1.
cmake/MetaxKernels.cmake	Adds new v2 kernels to the MetaX v2 build set.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+        if real_spec is None:
+            try:
+                real_spec = importlib.util.find_spec(fullname)
+            except (ImportError, ValueError):
+                real_spec = None


+    w = getattr(self, "_fused_qkv_weight", None)
+    if w is not None:
+        return w
+    q, k, v = self.q_proj, self.k_proj, self.v_proj
+    if q.bias is not None or k.bias is not None or v.bias is not None:
+        return None
+    with torch.no_grad():
+        w = torch.cat([q.weight, k.weight, v.weight], dim=0)
+    # Cache split sizes too (GQA: k/v out-dim < q out-dim).


+def _get_fused_gate_up_weight(self):
+    """Lazily build+cache the concatenated [gate; up] weight. None if biased."""
+    w = getattr(self, "_fused_gate_up_weight", None)
+    if w is not None:
+        return w
+    g, u = self.gate_proj, self.up_proj
+    if g.bias is not None or u.bias is not None:
+        return None
+    with torch.no_grad():
+        w = torch.cat([g.weight, u.weight], dim=0)
+    self._fused_gate_up_weight = w


+// mxcc-compiled definition of the fused masked-softmax host entry declared in
+// csrc/aten/softmax.h. Thin wrapper over the kernel in softmax_kernel_v2.cuh.


+  // Clamp grid to the hardware limit (decode rows are tiny; prefill can be
+  // large, e.g. batch=32, seq=2048 -> rows=65536).
+  const int blocks =
+      static_cast<int>(std::min<int64_t>(rows, metax::kMaxGridV2));
+
+  RmsNormRowKernel<scalar_t, acc_t><<<blocks, BLOCK, 0, stream>>>(
+      rows, hidden, eps, out_ptr, in_ptr, weight_ptr);


+  const int blocks =
+      static_cast<int>(std::min<int64_t>(rows, metax::kMaxGridV2));
+  MaskedSoftmaxKernel<scalar_t, acc_t, mask_t>
+      <<<blocks, BLOCK, 0, stream>>>(
+          rows, k_len, q_len, heads, scaling, out_ptr, in_ptr, mask_ptr,
+          ms0, ms1, ms2, ms3);


+  const int blocks =
+      static_cast<int>(std::min<int64_t>(rows, metax::kMaxGridV2));
+
+  AddRmsNormRowKernel<scalar_t, acc_t><<<blocks, BLOCK, 0, stream>>>(
+      rows, hidden, eps, normed_ptr, new_residual_ptr, res_ptr, hid_ptr,
+      weight_ptr);


+TORCH_LIBRARY(flagos, m) {
+  m.def("rms_norm(Tensor input, Tensor weight, float eps) -> Tensor");
+  m.def("rope(Tensor q, Tensor k, Tensor cos, Tensor sin, int unsqueeze_dim) -> (Tensor, Tensor)");
+  m.def("masked_softmax(Tensor attn_weights, Tensor? mask, float scaling) -> Tensor");
+  m.def("swiglu(Tensor gate, Tensor up) -> Tensor");
+  m.def("add_rms_norm(Tensor residual, Tensor hidden, Tensor weight, float eps) -> (Tensor, Tensor)");
+}
+
+TORCH_LIBRARY_IMPL(flagos, PrivateUse1, m) {
+  m.impl("rms_norm", WrapperRmsNorm);
+  m.impl("rope", WrapperRoPE);
+  m.impl("masked_softmax", WrapperMaskedSoftmax);
+  m.impl("swiglu", WrapperSwiGLU);
+  m.impl("add_rms_norm", WrapperAddRmsNorm);
+}


rms_norm / add_rms_norm / masked_softmax used row = blockIdx.x with an early-exit for row >= rows, but the launch clamps gridDim.x to kMaxGridV2 (65535). For rows > 65535 (e.g. prefill: batch*heads*q_len — a single 2048-token prompt with 32 heads is 65536) the tail rows were never computed, reading uninitialized device memory. Wrap each kernel body in a grid-stride loop over rows and add a trailing __syncthreads() so the next row can safely reuse the shared reduction scratch. Also fix the softmax_v2.cu header comment to point at masked_softmax.h (not softmax.h). Verified: new large-rows regression tests pass (rows 65600 / 67584), and Qwen3-0.6B decode is token-exact vs stock HF with all fusions enabled. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…caches _FusionImportHook.find_spec falls back to importlib.util.find_spec, which re-walks sys.meta_path and re-enters this finder on the same fullname — recursing until the stack blows. Add a per-thread re-entry flag so the nested call short-circuits to None and lets the real finders resolve. Also document that the fused QKV / gate_up weight caches are inference-only: they concat under torch.no_grad() and never invalidate, so they assume frozen weights (the decode-throughput use case) and must not be enabled for training/finetuning. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Cover torch.ops.flagos.{rms_norm,rope,masked_softmax,swiglu,add_rms_norm} against an fp32 CPU reference (fp16/bf16 I/O, fp32 reduction — matching the kernels). Each reduction op also has a *_large_rows case that pushes rows past kMaxGridV2 (65535) to guard the grid-stride tail. Marked metax so conftest skips them off the MetaX backend. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

zhangliheng and others added 5 commits July 15, 2026 11:15

Copilot AI review requested due to automatic review settings July 15, 2026 04:03

Copilot started reviewing on behalf of shallitbeso July 15, 2026 04:03 View session

Copilot AI reviewed Jul 15, 2026

View reviewed changes

zhangliheng and others added 3 commits July 15, 2026 14:27

Navigation Menu

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

perf(metax_v2): fuse HF RMSNorm/RoPE/softmax/SwiGLU launch-heavy chains#1
shallitbeso wants to merge 8 commits into
ZavierHung:metax-v2from
shallitbeso:metax-v2

shallitbeso commented Jul 15, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		// mxcc-compiled definition of the fused masked-softmax host entry declared in
		// csrc/aten/softmax.h. Thin wrapper over the kernel in softmax_kernel_v2.cuh.

Conversation

shallitbeso commented Jul 15, 2026

perf(metax_v2): fuse HF RMSNorm/RoPE/softmax/SwiGLU launch-heavy chains (Qwen3 decode +100%)

背景

改动

实现方式

数值验证

测量协议

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants