feat(ascend): Qwen-Image qk_rmsnorm & multi_axis_rope kernels (issue #386)

[erfgss]

Contributes the qk_rmsnorm and multi_axis_rope rows of the Qwen-Image WS1 kernel table (issue #386, NPU/Ascend C + PyTorch reference + registry wiring).

What

Ascend C kernels (csrc/ascend/Qwen-Image/)

• qk_rmsnorm_ascend.asc — per-head parameter-free QK RMSNorm (head_dim=128). Following the established rmsnorm_ascend contract, the host computes rstd with the exact torch reference ops (fp32 mean of squares + torch.rsqrt); the kernel performs only order-free fp32 elementwise multiplies and one round-to-nearest-even cast, so the fused output is bitwise identical to the reference and batch-invariant (fixed tile, one row per block, chunk coalescing for narrow rows).
• multi_axis_rope_ascend.asc — multi-axis RoPE (axes [16, 56, 56], text on the grid diagonal). The axis split / per-axis frequencies / diagonal text placement live in the fp32 cos/sin tables built host-side; the kernel is the dtype-generic rotate-half primitive with forward (sin_sign=+1) and backward (transpose R^T, sin_sign=-1) entry points sharing one code path.
• npu_module.cpp: registers qk_rmsnorm_ascend, multi_axis_rope_ascend_forward, multi_axis_rope_ascend_backward.

Python ops & registry

• PyTorch references: NativeQkRmsNormOp, NativeMultiAxisRopeOp (+ qwen_image_positions diagonal/grid position builder and build_multi_axis_cos_sin table builder).
• Ascend autograd wrappers: QkRmsNormAscendOp, MultiAxisRopeAscendOp with clean fallback to the native references for non-NPU / unsupported inputs.
• registry.py: OpBackend entries ASCEND_QK_RMS_NORM / ASCEND_MULTI_AXIS_ROPE (+ native fallbacks); npu priority [ASCEND, PYTORCH_NATIVE] for op types qk_rmsnorm / multi_axis_rope; native-only candidates on cpu/cuda/rocm/musa (previously these op types would have resolved to the generic logp backend there).

Testing (tests/test_qwen_image_ops.py)

CPU (22 passed on the dev host):

• bitwise match vs independently written fp32 formulas (both ops, fp32/bf16/fp16)
• per-head independence, forward/backward roundtrip of the rotation
• batch invariance: rows are bitwise identical regardless of batch size/position
• explicit fp32 VJP (used by the Ascend backward) vs autograd
• issue shapes {1024², 1328², 1664×928} (patchified grids 64×64 / 83×83 / 104×58) end to end
• registry wiring (npu priority, non-NPU fallbacks)

NPU (15 tests, auto-activated with skipif when the kernels are compiled):

• forward bitwise equals the PyTorch reference at fp32/bf16/fp16
• on-device batch invariance
• backward matches the reference autograd under the "elementwise" gtest tolerance class

Run: python -m pytest tests/test_qwen_image_ops.py -v (NPU cases require KERNEL_ALIGN_FORCE_ASCEND=1 pip install --no-build-isolation -e .).

Issue acceptance checklist

• Forward matches the FP32 reference bitwise (CPU harness; on-NPU bitwise tests included, pending hardware run)
• Backward matches the reference (explicit fp32 VJP, tested vs autograd)
• Batch invariance tested (CPU + on-device cases)
• Shapes cover {1024², 1328², 1664×928}
• Trace records (reduction order, kernel fingerprint, backend) — to follow with the gtest harness integration
• CUDA / Triton implementations — not part of this NPU-scoped PR
• Benchmark — to follow

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: abe23889-fdc8-45f1-a436-0bfc125854f5

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}