feat(ascend): support non-neox rotaryMode via ATB RopeParam rotaryCoeff

RotaryEmbedding impl=1 (ATB Rope) now plumbs both rotary styles:
- is_neox_style=true  -> rotaryCoeff=2          (half split + cat)
- is_neox_style=false -> rotaryCoeff=head_size  (interleave)

The cos/sin expand path also branches: neox layout duplicates the
half values front/back, while interleave layout repeats each value
pair-wise.  Test skip is narrowed to impl=0 only, which still uses
aclnnApplyRotaryPosEmbV2 (declares "interleave" but only implements
"half").

G (partial rotary) skip message updated to reflect that neither
aclnn nor ATB fused APIs support rotary_dim < head_size.

Author: zhangyue

src/ascend/rotary_embedding/kernel_atb.h

@@ -44,11 +44,14 @@ namespace infini::ops {
 // gathered `[T, D]` tensors to ATB Rope.  The `seqlen` input is a single
 // int32 element equal to T (all tokens treated as one batch).
 //
-// Restrictions (implementation choices, not ATB API limits):
+// Restrictions:
 //   - `rotary_dim` must equal `head_size` (full rotation only).  ATB
 //     RopeParam supports `rotaryCoeff=2/4/head_size/head_size_2` per the
-//     CANN 8.5 ATB docs; this wrapper plumbs only `rotaryCoeff=2`.
-//   - `is_neox_style` must be true.
+//     CANN 8.5 ATB docs.  This wrapper plumbs:
+//       * `rotaryCoeff=2` when `is_neox_style=true`  (half split + cat)
+//       * `rotaryCoeff=head_size` when `is_neox_style=false` (interleave)
+//     Partial rotary (`rotary_dim < head_size`) is not supported by either
+//     the aclnn or ATB fused APIs; callers must pad to `head_size` upstream.
 template <>
 class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
     : public RotaryEmbedding {
@@ -57,11 +60,10 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
            const Tensor cos_sin_cache, int64_t head_size, int64_t rotary_dim,
            bool is_neox_style, Tensor query_out, Tensor key_out)
       : RotaryEmbedding(positions, query, key, cos_sin_cache, head_size,
-                        rotary_dim, is_neox_style, query_out, key_out) {
+                        rotary_dim, is_neox_style, query_out, key_out),
+        is_neox_style_{is_neox_style} {
     assert(rotary_dim == head_size &&
            "ATB `RotaryEmbedding` requires rotary_dim == head_size");
-    assert(is_neox_style &&
-           "ATB `RotaryEmbedding` requires neox style (rotaryCoeff=2)");
 
     const int64_t D = head_size_;
     const size_t elem_sz = cos_sin_cache.element_size();
@@ -110,10 +112,14 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
     cos_out_cache_ = ascend::AclTensorCache({T, D}, acl_dt_, cos_dev_);
     sin_out_cache_ = ascend::AclTensorCache({T, D}, acl_dt_, sin_dev_);
 
-    // Create the ATB Rope operation.
+    // Create the ATB Rope operation.  `rotaryCoeff` selects the rotation
+    // pattern: 2 for neox (split-then-rotate halves), `head_size` for
+    // interleave (pair-wise rotate adjacent elements).
     atb::infer::RopeParam param;
-    param.rotaryCoeff = 2;  // Neox half-rotation.
-    param.cosFormat = 0;    // Inference mode.
+    param.rotaryCoeff = is_neox_style
+                            ? 2
+                            : static_cast<int32_t>(D);
+    param.cosFormat = 0;  // Inference mode.
     atb::Status s = atb::CreateOperation(param, &op_);
 
     assert(s == atb::NO_ERROR && "atb::CreateOperation(Rope) failed");
@@ -254,8 +260,16 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
   }
 
  private:
-  // D2H copy cos_sin_cache, split into cos/sin, neox-expand, and upload to
-  // device.  Called once at construction.
+  // D2H copy cos_sin_cache, split into cos/sin, expand to `[max_seq_len, D]`
+  // in the layout that ATB Rope expects for the chosen `rotaryCoeff`, and
+  // upload to device.  Called once at construction.
+  //
+  // For `rotaryCoeff=2` (neox): cos tensor holds the same `half_D` values
+  // duplicated front/back — `[c0 .. c_{half-1}, c0 .. c_{half-1}]`.
+  //
+  // For `rotaryCoeff=head_size` (interleave): cos tensor holds each of the
+  // `half_D` values repeated pair-wise —
+  // `[c0, c0, c1, c1, .., c_{half-1}, c_{half-1}]`.
   void uploadCosSinCache(const Tensor cos_sin_cache) const {
     const int64_t D = head_size_;
     const int64_t half_D = D / 2;
@@ -277,16 +291,35 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
         const auto* s_src = cache_host.data() +
                             static_cast<size_t>(p * D + half_D + j) * elem_sz;
 
-        std::memcpy(cos_host.data() + static_cast<size_t>(p * D + j) * elem_sz,
-                    c_src, elem_sz);
-        std::memcpy(
-            cos_host.data() + static_cast<size_t>(p * D + half_D + j) * elem_sz,
-            c_src, elem_sz);
-        std::memcpy(sin_host.data() + static_cast<size_t>(p * D + j) * elem_sz,
-                    s_src, elem_sz);
-        std::memcpy(
-            sin_host.data() + static_cast<size_t>(p * D + half_D + j) * elem_sz,
-            s_src, elem_sz);
+        if (is_neox_style_) {
+          // Neox layout: [c_j ... , c_j ...] front/back duplication.
+          std::memcpy(
+              cos_host.data() + static_cast<size_t>(p * D + j) * elem_sz,
+              c_src, elem_sz);
+          std::memcpy(
+              cos_host.data() + static_cast<size_t>(p * D + half_D + j) * elem_sz,
+              c_src, elem_sz);
+          std::memcpy(
+              sin_host.data() + static_cast<size_t>(p * D + j) * elem_sz,
+              s_src, elem_sz);
+          std::memcpy(
+              sin_host.data() + static_cast<size_t>(p * D + half_D + j) * elem_sz,
+              s_src, elem_sz);
+        } else {
+          // Interleave layout: each value repeated pair-wise.
+          std::memcpy(
+              cos_host.data() + static_cast<size_t>(p * D + 2 * j) * elem_sz,
+              c_src, elem_sz);
+          std::memcpy(
+              cos_host.data() + static_cast<size_t>(p * D + 2 * j + 1) * elem_sz,
+              c_src, elem_sz);
+          std::memcpy(
+              sin_host.data() + static_cast<size_t>(p * D + 2 * j) * elem_sz,
+              s_src, elem_sz);
+          std::memcpy(
+              sin_host.data() + static_cast<size_t>(p * D + 2 * j + 1) * elem_sz,
+              s_src, elem_sz);
+        }
       }
     }
 
@@ -296,6 +329,8 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
                 ACL_MEMCPY_HOST_TO_DEVICE);
   }
 
+  bool is_neox_style_;
+
   atb::Operation* op_ = nullptr;
 
   // Neox-expanded cos/sin tables on device: [max_seq_len, D].

tests/test_rotary_embedding.py

@@ -164,10 +164,12 @@ def test_rotary_embedding_full(
                 f"Implementation index={implementation_index} not active on this build"
             )
 
-    if device == "npu" and not is_neox_style:
+    # Only implementation 0 (`aclnnApplyRotaryPosEmbV2`) is still limited to
+    # `rotaryMode="half"`; implementation 1 (ATB `RopeParam`) plumbs
+    # `rotaryCoeff=head_size` for the non-neox (interleave) case.
+    if device == "npu" and not is_neox_style and implementation_index == 0:
         pytest.skip(
-            'Ascend `RotaryEmbedding` wrappers only plumb `rotaryMode="half"` '
-            "through the underlying V2/ATB APIs."
+            'Ascend `aclnnApplyRotaryPosEmbV2` only supports `rotaryMode="half"`'
         )
 
     # `aclnnApplyRotaryPosEmbV2` accumulates with ~4 ULP error for float16.
@@ -486,7 +488,13 @@ def test_rotary_embedding_partial(
         pytest.skip("NPU not available")
 
     if device == "npu":
-        pytest.skip("Ascend aclnnApplyRotaryPosEmbV2 requires rotary_dim == head_size")
+        pytest.skip(
+            "Partial rotary (`rotary_dim < head_size`) is not supported by "
+            "any Ascend fused API: `aclnnApplyRotaryPosEmbV2`, "
+            "`aclnnRotaryPositionEmbedding`, and ATB `RopeParam` all require "
+            "`cos.D == sin.D == x.D`.  A decomposed implementation is "
+            "forbidden by project policy."
+        )
 
     num_tokens = 16
     max_seq_len = 64