feat(ascend): add batch-invariant RMSNorm Ascend C operator

benchmarks/benchmark_rmsnorm.py

@@ -14,22 +14,30 @@
 except ImportError:
     HAS_CUDA_EXT = False
 
+try:
+    from rl_engine.kernels.ops.ascend.norm.rmsnorm import RMSNormAscendOp
+
+    HAS_ASCEND_EXT = True
+except (ImportError, OSError, RuntimeError):
+    HAS_ASCEND_EXT = False
+
 
 def bench(fn, x, w, dy, warmup=20, iters=100):
+    sync = torch.npu.synchronize if x.device.type == "npu" else torch.cuda.synchronize
     for _ in range(warmup):
         x.grad = None
         w.grad = None
         y = fn(x, w)
         y.backward(dy)
-    torch.cuda.synchronize()
+    sync()
 
     start = time.time()
     for _ in range(iters):
         x.grad = None
         w.grad = None
         y = fn(x, w)
         y.backward(dy)
-    torch.cuda.synchronize()
+    sync()
     return (time.time() - start) * 1000.0 / iters
 
 
@@ -41,7 +49,7 @@ def main():
     args = parser.parse_args()
 
     dtype = torch.float16 if args.dtype == "fp16" else torch.bfloat16
-    device = "cuda"
+    device = "npu" if torch.npu.is_available() else "cuda"
     T, H = args.T, args.H
 
     torch.manual_seed(0)
@@ -72,6 +80,15 @@ def make_inputs():
     else:
         print("cuda        : skipped, extension is not built")
 
+    if device == "npu":
+        ascend_op = RMSNormAscendOp() if HAS_ASCEND_EXT else None
+        if ascend_op is not None:
+            x, w = make_inputs()
+            t_asc = bench(lambda a, b: ascend_op(a, b), x, w, dy)
+            print(f"ascend      : {t_asc:.4f} ms | speedup vs ref: {t_ref / t_asc:.2f}x")
+        else:
+            print("ascend      : skipped, extension is not built")
+
 
 if __name__ == "__main__":
     main()

csrc/ascend/batch_invariant_logp_ascend.asc

@@ -307,3 +307,6 @@ std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(torch::Tensor log
     }
     return {logp, lse};
 }
+
+// The PYBIND11_MODULE for rl_engine._C_npu lives in npu_module.cpp so that
+// every Ascend op shares one compiled module.

csrc/ascend/npu_module.cpp

@@ -25,6 +25,10 @@ std::vector<torch::Tensor> deterministic_attention_ascend_forward(
 torch::Tensor prefix_shared_attention_ascend_forward(
     torch::Tensor q, torch::Tensor k, torch::Tensor v);
 
+torch::Tensor rmsnorm_ascend_forward(torch::Tensor x,
+                                     torch::Tensor weight,
+                                     torch::Tensor rstd);
+
 int64_t deterministic_collective_create(
     torch::Tensor staging, int64_t world_size, int64_t rank);
 void deterministic_collective_destroy(int64_t handle);
@@ -58,4 +62,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
     m.def("deterministic_collective_reduce",
           &deterministic_collective_reduce,
           "Fixed-tree ordered reduction over gathered rank tensors (Ascend)");
+    m.def("rmsnorm_ascend",
+          &rmsnorm_ascend_forward,
+          "Batch-invariant RMSNorm (Ascend C forward, rstd precomputed)");
 }