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

@@ -308,9 +308,5 @@ std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(torch::Tensor log
     return {logp, lse};
 }
 
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
-{
-    m.def("batch_invariant_logp_ascend",
-          &batch_invariant_logp_ascend_forward,
-          "Batch-invariant selected-token log-probability (Ascend C forward)");
-}
+// 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

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// Pybind entry point for the rl_engine._C_npu extension. The Ascend C kernels
+// and their torch host wrappers live in the sibling *.asc files; this TU only
+// declares and binds them so every Ascend op shares one compiled module.
+//
+// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by
+// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu.
+
+#include <torch/extension.h>
+
+std::vector<torch::Tensor> batch_invariant_logp_ascend_forward(torch::Tensor logits,
+                                                               torch::Tensor target,
+                                                               int64_t ignore_index);
+
+std::vector<torch::Tensor> rmsnorm_ascend_forward(torch::Tensor x,
+                                                  torch::Tensor weight,
+                                                  double eps);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    m.def("batch_invariant_logp_ascend",
+          &batch_invariant_logp_ascend_forward,
+          "Batch-invariant selected-token log-probability (Ascend C forward)");
+    m.def("rmsnorm_ascend",
+          &rmsnorm_ascend_forward,
+          "Batch-invariant RMSNorm (Ascend C forward)");
+}