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

csrc/ascend/Qwen-Image/multi_axis_rope_ascend.asc

@@ -0,0 +1,349 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+//
+// Qwen-Image MMDiT multi-axis RoPE, Ascend C (CANN) forward + backward kernels.
+//
+// Issue #386 kernel-table row `multi_axis_rope`: Qwen-Image splits the
+// head_dim=128 rotary axes into [16, 56, 56] (temporal, height, width) and
+// places text tokens on the grid diagonal. The axis split, the per-axis
+// frequency bases, and the diagonal text placement all live in the position
+// tables: the Python wrapper builds fp32 cos/sin caches of shape
+// [table_rows, D/2] with the exact reference formula
+// (per-axis inv_freq = theta^(-arange(0, dim, 2)/dim), concatenated over
+// axes, coordinates per token, cat(cos, cos) / cat(sin, sin) for the
+// rotate-half convention). This kernel is the dtype-generic rotate-half
+// apply primitive over those tables; for pair i in [0, D/2):
+//
+//   out[i]       = x[i]       * cos[i] - x[i+D/2] * sin[i] * sin_sign
+//   out[i+D/2]   = x[i+D/2]   * cos[i] + x[i]     * sin[i] * sin_sign
+//
+// sin_sign=+1 is the forward rotation; sin_sign=-1 is its transpose R^T,
+// which is exactly the backward rotation used for grad_x (R(theta)^T =
+// R(-theta)), exposed as multi_axis_rope_ascend_backward.
+//
+// Bit-exactness and batch invariance follow the same contract as
+// csrc/ascend/rope_ascend.asc: every row is processed end-to-end by exactly
+// one AI core block with a fixed tile order over D/2, so adding or moving
+// other batch rows cannot alter a row's instruction sequence, and the fp32
+// elementwise IEEE ops plus one round-to-nearest-even cast at the output
+// cannot introduce drift against the reference apply.
+//
+// Mirrors the CUDA-side contract:
+//   - input  : x [n_rows, D] contiguous, fp32 / bf16 / fp16;
+//              cos/sin [table_rows, D/2] fp32 (built by the Python wrapper)
+//   - output : y [n_rows, D] same dtype as x
+//
+// Build: see setup.py (AscendBuildExtension, bisheng -x asc), gated by
+// KERNEL_ALIGN_FORCE_ASCEND=1. Requires CANN toolkit + torch_npu.
+
+#include <algorithm>
+#include <type_traits>
+
+#include "kernel_operator.h"
+
+#include <torch/extension.h>
+
+#include "torch_npu/csrc/core/npu/NPUStream.h"
+
+namespace {
+
+constexpr uint32_t ROPE_TILE_HALF = 4096;
+constexpr int64_t ROPE_MAX_BLOCKS = 128;
+
+template <typename T>
+class KernelMultiAxisRope {
+public:
+    __aicore__ inline KernelMultiAxisRope(AscendC::TPipe* pipe) : pipe_(pipe) {}
+
+    __aicore__ inline void Init(GM_ADDR x,
+                                GM_ADDR cos,
+                                GM_ADDR sin,
+                                GM_ADDR out,
+                                int64_t numRows,
+                                int64_t tableRows,
+                                int64_t headDim,
+                                float sinSign)
+    {
+        numRows_ = numRows;
+        tableRows_ = tableRows;
+        headDim_ = headDim;
+        halfDim_ = headDim / 2;
+        sinSign_ = sinSign;
+        xGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(x));
+        cosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(cos));
+        sinGm_.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(sin));
+        outGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(out));
+
+        pipe_->InitBuffer(x1InBuf_, ROPE_TILE_HALF * sizeof(T));
+        pipe_->InitBuffer(x2InBuf_, ROPE_TILE_HALF * sizeof(T));
+        pipe_->InitBuffer(x1FpBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(x2FpBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(cosBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(sinBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(out1FpBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(out2FpBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(tmpFpBuf_, ROPE_TILE_HALF * sizeof(float));
+        pipe_->InitBuffer(out1Buf_, ROPE_TILE_HALF * sizeof(T));
+        pipe_->InitBuffer(out2Buf_, ROPE_TILE_HALF * sizeof(T));
+
+        // Mark the reusable input and output buffers as initially available.
+        AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(0);
+        AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(0);
+    }
+
+    __aicore__ inline void Process()
+    {
+        for (int64_t row = AscendC::GetBlockIdx(); row < numRows_; row += AscendC::GetBlockNum()) {
+            ProcessRow(row);
+        }
+        // Drain the final tile before the block exits and its UB is reclaimed.
+        AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(0);
+        AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(0);
+    }
+
+private:
+    __aicore__ inline void ProcessRow(int64_t row)
+    {
+        // row % tableRows_ selects the position cache: the wrapper flattens
+        // [B, H, S, D] so that the S index cycles over the table rows,
+        // covering the image (t, h, w) grid and the diagonal text positions
+        // alike — both are just rows of the cos/sin tables.
+        const int64_t tableRow = row % tableRows_;
+        for (int64_t start = 0; start < halfDim_; start += ROPE_TILE_HALF) {
+            const int64_t remaining = halfDim_ - start;
+            const uint32_t count = static_cast<uint32_t>(
+                remaining < ROPE_TILE_HALF ? remaining : ROPE_TILE_HALF);
+
+            // Previous vector reads are complete before MTE2 reuses input/cache buffers.
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(0);
+            CopyIn(row, tableRow, start, count);
+            AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(0);
+
+            // Previous MTE3 reads are complete before V reuses output buffers.
+            AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(0);
+            Compute(count);
+
+            // The next MTE2 tile may reuse its buffers after all vector reads finish.
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(0);
+            AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
+            AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);
+            CopyOut(row, start, count);
+            AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(0);
+        }
+    }
+
+    __aicore__ inline void CopyIn(int64_t row,
+                                  int64_t tableRow,
+                                  int64_t start,
+                                  uint32_t count)
+    {
+        AscendC::DataCopyExtParams xParams{
+            1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
+        AscendC::DataCopyExtParams fpParams{
+            1, static_cast<uint32_t>(count * sizeof(float)), 0, 0, 0};
+        AscendC::DataCopyPadExtParams<T> xPad{false, 0, 0, 0};
+        AscendC::DataCopyPadExtParams<float> fpPad{false, 0, 0, 0};
+
+        const int64_t xBase = row * headDim_ + start;
+        if constexpr (std::is_same_v<T, float>) {
+            AscendC::DataCopyPad(x1FpBuf_.Get<float>(), xGm_[xBase], xParams, xPad);
+            AscendC::DataCopyPad(
+                x2FpBuf_.Get<float>(), xGm_[xBase + halfDim_], xParams, xPad);
+        } else {
+            AscendC::DataCopyPad(x1InBuf_.Get<T>(), xGm_[xBase], xParams, xPad);
+            AscendC::DataCopyPad(
+                x2InBuf_.Get<T>(), xGm_[xBase + halfDim_], xParams, xPad);
+        }
+
+        const int64_t cacheBase = tableRow * halfDim_ + start;
+        AscendC::DataCopyPad(cosBuf_.Get<float>(), cosGm_[cacheBase], fpParams, fpPad);
+        AscendC::DataCopyPad(sinBuf_.Get<float>(), sinGm_[cacheBase], fpParams, fpPad);
+    }
+
+    __aicore__ inline void Compute(uint32_t count)
+    {
+        AscendC::LocalTensor<float> x1 = x1FpBuf_.Get<float>();
+        AscendC::LocalTensor<float> x2 = x2FpBuf_.Get<float>();
+        if constexpr (!std::is_same_v<T, float>) {
+            AscendC::Cast(x1, x1InBuf_.Get<T>(), AscendC::RoundMode::CAST_NONE, count);
+            AscendC::Cast(x2, x2InBuf_.Get<T>(), AscendC::RoundMode::CAST_NONE, count);
+        }
+
+        AscendC::LocalTensor<float> cos = cosBuf_.Get<float>();
+        AscendC::LocalTensor<float> sin = sinBuf_.Get<float>();
+        AscendC::LocalTensor<float> out1 = out1FpBuf_.Get<float>();
+        AscendC::LocalTensor<float> out2 = out2FpBuf_.Get<float>();
+        AscendC::LocalTensor<float> tmp = tmpFpBuf_.Get<float>();
+
+        AscendC::Muls(sin, sin, sinSign_, count);
+        AscendC::Mul(out1, x1, cos, count);
+        AscendC::Mul(tmp, x2, sin, count);
+        AscendC::Sub(out1, out1, tmp, count);
+        AscendC::Mul(out2, x2, cos, count);
+        AscendC::Mul(tmp, x1, sin, count);
+        AscendC::Add(out2, out2, tmp, count);
+
+        if constexpr (!std::is_same_v<T, float>) {
+            AscendC::Cast(
+                out1Buf_.Get<T>(), out1, AscendC::RoundMode::CAST_RINT, count);
+            AscendC::Cast(
+                out2Buf_.Get<T>(), out2, AscendC::RoundMode::CAST_RINT, count);
+        }
+    }
+
+    __aicore__ inline void CopyOut(int64_t row, int64_t start, uint32_t count)
+    {
+        AscendC::DataCopyExtParams outParams{
+            1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
+        const int64_t outBase = row * headDim_ + start;
+        if constexpr (std::is_same_v<T, float>) {
+            AscendC::DataCopyPad(outGm_[outBase], out1FpBuf_.Get<float>(), outParams);
+            AscendC::DataCopyPad(
+                outGm_[outBase + halfDim_], out2FpBuf_.Get<float>(), outParams);
+        } else {
+            AscendC::DataCopyPad(outGm_[outBase], out1Buf_.Get<T>(), outParams);
+            AscendC::DataCopyPad(
+                outGm_[outBase + halfDim_], out2Buf_.Get<T>(), outParams);
+        }
+    }
+
+    AscendC::TPipe* pipe_;
+    AscendC::GlobalTensor<T> xGm_;
+    AscendC::GlobalTensor<float> cosGm_;
+    AscendC::GlobalTensor<float> sinGm_;
+    AscendC::GlobalTensor<T> outGm_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> x1InBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> x2InBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> x1FpBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> x2FpBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> cosBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> sinBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> out1FpBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> out2FpBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpFpBuf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> out1Buf_;
+    AscendC::TBuf<AscendC::TPosition::VECCALC> out2Buf_;
+    int64_t numRows_;
+    int64_t tableRows_;
+    int64_t headDim_;
+    int64_t halfDim_;
+    float sinSign_;
+};
+
+}  // namespace
+
+extern "C" __global__ __vector__ void multi_axis_rope_ascend_kernel_fp32(
+    GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out,
+    int64_t numRows, int64_t tableRows, int64_t headDim, float sinSign)
+{
+    AscendC::TPipe pipe;
+    KernelMultiAxisRope<float> op(&pipe);
+    op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign);
+    op.Process();
+}
+
+extern "C" __global__ __vector__ void multi_axis_rope_ascend_kernel_fp16(
+    GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out,
+    int64_t numRows, int64_t tableRows, int64_t headDim, float sinSign)
+{
+    AscendC::TPipe pipe;
+    KernelMultiAxisRope<half> op(&pipe);
+    op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign);
+    op.Process();
+}
+
+extern "C" __global__ __vector__ void multi_axis_rope_ascend_kernel_bf16(
+    GM_ADDR x, GM_ADDR cos, GM_ADDR sin, GM_ADDR out,
+    int64_t numRows, int64_t tableRows, int64_t headDim, float sinSign)
+{
+    AscendC::TPipe pipe;
+    KernelMultiAxisRope<bfloat16_t> op(&pipe);
+    op.Init(x, cos, sin, out, numRows, tableRows, headDim, sinSign);
+    op.Process();
+}
+
+namespace {
+
+// Shared apply path for forward (sin_sign=+1) and backward (sin_sign=-1,
+// the transpose rotation). Keeps the TORCH_CHECK gate and launch geometry in
+// one place so both directions are byte-identical code paths.
+torch::Tensor MultiAxisRopeApply(torch::Tensor x,
+                                 torch::Tensor cos,
+                                 torch::Tensor sin,
+                                 float sinSign)
+{
+    TORCH_CHECK(x.is_privateuseone(), "multi_axis_rope: x must be on an NPU device");
+    TORCH_CHECK(x.dim() == 2, "multi_axis_rope: x must be 2-D [n_rows, D]");
+    TORCH_CHECK(x.is_contiguous(), "multi_axis_rope: x must be contiguous");
+    TORCH_CHECK(
+        x.scalar_type() == at::kHalf || x.scalar_type() == at::kBFloat16 ||
+            x.scalar_type() == at::kFloat,
+        "multi_axis_rope: x must be fp16, bf16, or fp32");
+    TORCH_CHECK(cos.is_privateuseone() && sin.is_privateuseone(),
+                "multi_axis_rope: cos/sin must be on an NPU device");
+    TORCH_CHECK(cos.device() == x.device() && sin.device() == x.device(),
+                "multi_axis_rope: x, cos, and sin must be on the same NPU device");
+    TORCH_CHECK(cos.scalar_type() == at::kFloat && sin.scalar_type() == at::kFloat,
+                "multi_axis_rope: cos/sin must be fp32");
+    TORCH_CHECK(cos.dim() == 2 && sin.dim() == 2,
+                "multi_axis_rope: cos/sin must be 2-D [table_rows, D/2]");
+    TORCH_CHECK(cos.is_contiguous() && sin.is_contiguous(),
+                "multi_axis_rope: cos/sin must be contiguous");
+    TORCH_CHECK(cos.sizes() == sin.sizes(),
+                "multi_axis_rope: cos/sin shapes must match");
+
+    const int64_t numRows = x.size(0);
+    const int64_t headDim = x.size(1);
+    TORCH_CHECK(headDim > 0 && headDim % 2 == 0,
+                "multi_axis_rope: head_dim must be a positive even number");
+    TORCH_CHECK(cos.size(1) == headDim / 2,
+                "multi_axis_rope: cos/sin last dimension must equal head_dim/2");
+
+    torch::Tensor out = at::empty_like(x);
+    if (numRows == 0) {
+        return out;
+    }
+
+    const int64_t tableRows = cos.size(0);
+    TORCH_CHECK(tableRows > 0,
+                "multi_axis_rope: cos/sin table must contain at least one row");
+    TORCH_CHECK(numRows % tableRows == 0,
+                "multi_axis_rope: n_rows must be divisible by the cos/sin table row count");
+    auto aclStream = c10_npu::getCurrentNPUStream().stream(true);
+    const uint32_t blockNum =
+        static_cast<uint32_t>(std::min(numRows, ROPE_MAX_BLOCKS));
+
+    auto* xPtr = reinterpret_cast<uint8_t*>(x.mutable_data_ptr());
+    auto* cosPtr = reinterpret_cast<uint8_t*>(cos.mutable_data_ptr());
+    auto* sinPtr = reinterpret_cast<uint8_t*>(sin.mutable_data_ptr());
+    auto* outPtr = reinterpret_cast<uint8_t*>(out.mutable_data_ptr());
+    if (x.scalar_type() == at::kFloat) {
+        multi_axis_rope_ascend_kernel_fp32<<<blockNum, nullptr, aclStream>>>(
+            xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sinSign);
+    } else if (x.scalar_type() == at::kHalf) {
+        multi_axis_rope_ascend_kernel_fp16<<<blockNum, nullptr, aclStream>>>(
+            xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sinSign);
+    } else {
+        multi_axis_rope_ascend_kernel_bf16<<<blockNum, nullptr, aclStream>>>(
+            xPtr, cosPtr, sinPtr, outPtr, numRows, tableRows, headDim, sinSign);
+    }
+    return out;
+}
+
+}  // namespace
+
+torch::Tensor multi_axis_rope_ascend_forward(torch::Tensor x,
+                                             torch::Tensor cos,
+                                             torch::Tensor sin)
+{
+    return MultiAxisRopeApply(std::move(x), std::move(cos), std::move(sin), 1.0f);
+}
+
+torch::Tensor multi_axis_rope_ascend_backward(torch::Tensor grad,
+                                              torch::Tensor cos,
+                                              torch::Tensor sin)
+{
+    // grad_x = R^T(grad) = rotate with the negated sin table.
+    return MultiAxisRopeApply(std::move(grad), std::move(cos), std::move(sin), -1.0f);
+}