feat(rotary_embedding): make query_out / key_out optional (inplace-default)

Align with vLLM's `RotaryEmbedding.forward(positions, query, key)`
signature by letting callers omit the output buffers — the kernel then
writes back in place on `query` / `key`.  This removes a signature
mismatch that forced vllm-infini to allocate and pass explicit out
tensors it doesn't need.

Base class signature:
  `query_out` / `key_out` → `std::optional<Tensor>` with `std::nullopt`
  default.  Shape / stride members fall back to `query` / `key` when the
  optional is empty.

All three Ascend impls resolve the optional to a concrete `Tensor` at
the top of `operator()` via `value_or(query)`:
  - impl=0 (aclnn V2):      skips the D2D memcpy in the inplace case
                            since `query.data() == q_out.data()`
  - impl=1 (ATB RopeParam): same short-circuit on the D2D copy
  - impl=2 (aclnnRopeWithSinCosCache): descriptors reuse `q_out` /
                            `k_out` pointers, so the kernel writes to
                            whichever tensor is resolved

Adds `test_rotary_embedding_inplace` covering both fp16 / bf16 on
impl=0 and impl=1.  Tolerance is atol=5e-3 — matches the V2 ~4 ULP
fp16 accumulator error documented in `kernel.h`.

Author: zhangyue

src/ascend/rotary_embedding/kernel.h

@@ -4,6 +4,7 @@
 #include <cassert>
 #include <cstddef>
 #include <cstring>
+#include <optional>
 #include <vector>
 
 #include "acl/acl.h"
@@ -38,11 +39,17 @@ class Operator<RotaryEmbedding, Device::Type::kAscend>
  public:
   Operator(const Tensor positions, const Tensor query, const Tensor key,
            const Tensor cos_sin_cache, int64_t head_size, int64_t rotary_dim,
-           bool is_neox_style, Tensor query_out, Tensor key_out)
+           bool is_neox_style,
+           std::optional<Tensor> query_out = std::nullopt,
+           std::optional<Tensor> key_out = std::nullopt)
       : RotaryEmbedding(positions, query, key, cos_sin_cache, head_size,
                         rotary_dim, is_neox_style, query_out, key_out),
         max_seq_len_{cos_sin_cache.size(0)},
         elem_sz_{cos_sin_cache.element_size()} {
+    // Resolve optional out buffers; when omitted, RoPE writes back in place
+    // on `query` / `key` — vLLM-style inplace semantics.
+    Tensor q_out = query_out.value_or(query);
+    Tensor k_out = key_out.value_or(key);
     assert(rotary_dim == head_size &&
            "Ascend `RotaryEmbedding` requires rotary_dim == head_size "
            "(partial rotation not implemented in this wrapper)");
@@ -85,9 +92,9 @@ class Operator<RotaryEmbedding, Device::Type::kAscend>
     cos_v2_cache_ = ascend::AclTensorCache({T, 1, D}, acl_dt, cos_dev_);
     sin_v2_cache_ = ascend::AclTensorCache({T, 1, D}, acl_dt, sin_dev_);
     q_cache_ = ascend::AclTensorCache({T, Nq, D}, acl_dt,
-                                      const_cast<void*>(query_out.data()));
+                                      const_cast<void*>(q_out.data()));
     k_cache_ = ascend::AclTensorCache({T, Nkv, D}, acl_dt,
-                                      const_cast<void*>(key_out.data()));
+                                      const_cast<void*>(k_out.data()));
   }
 
   ~Operator() {
@@ -112,10 +119,16 @@ class Operator<RotaryEmbedding, Device::Type::kAscend>
 
   void operator()(const Tensor positions, const Tensor query, const Tensor key,
                   const Tensor cos_sin_cache, int64_t head_size,
-                  int64_t rotary_dim, bool is_neox_style, Tensor query_out,
-                  Tensor key_out) const override {
+                  int64_t rotary_dim, bool is_neox_style,
+                  std::optional<Tensor> query_out,
+                  std::optional<Tensor> key_out) const override {
     auto stream = static_cast<aclrtStream>(stream_);
 
+    // Resolve optional out buffers (inplace on `query` / `key` when omitted).
+    // Non-const so `.data()` returns a writable `void*`.
+    Tensor q_out = query_out.value_or(query);
+    Tensor k_out = key_out.value_or(key);
+
     const int64_t T = query.size(0);
     const int64_t Nq = num_heads_;
     const int64_t Nkv = num_kv_heads_;
@@ -162,15 +175,15 @@ class Operator<RotaryEmbedding, Device::Type::kAscend>
     // Step 2: Copy q→q_out, k→k_out if not inplace (V2 operates inplace).
     size_t elem_sz = query.element_size();
 
-    if (query.data() != query_out.data()) {
-      aclrtMemcpyAsync(query_out.data(),
+    if (query.data() != q_out.data()) {
+      aclrtMemcpyAsync(q_out.data(),
                        static_cast<size_t>(T * Nq * D) * elem_sz, query.data(),
                        static_cast<size_t>(T * Nq * D) * elem_sz,
                        ACL_MEMCPY_DEVICE_TO_DEVICE, stream);
     }
 
-    if (key.data() != key_out.data()) {
-      aclrtMemcpyAsync(key_out.data(),
+    if (key.data() != k_out.data()) {
+      aclrtMemcpyAsync(k_out.data(),
                        static_cast<size_t>(T * Nkv * D) * elem_sz, key.data(),
                        static_cast<size_t>(T * Nkv * D) * elem_sz,
                        ACL_MEMCPY_DEVICE_TO_DEVICE, stream);
@@ -179,17 +192,17 @@ class Operator<RotaryEmbedding, Device::Type::kAscend>
     // Step 3: Apply V2 RoPE inplace on q_out and k_out.
     auto t_cos = cos_v2_cache_.get(cos_dev_);
     auto t_sin = sin_v2_cache_.get(sin_dev_);
-    auto t_q = q_cache_.get(query_out.data());
-    auto t_k = k_cache_.get(key_out.data());
+    auto t_q = q_cache_.get(q_out.data());
+    auto t_k = k_cache_.get(k_out.data());
 
     if (!v2_exec_) {
       aclnnApplyRotaryPosEmbV2GetWorkspaceSize(
           t_q, t_k, t_cos, t_sin, /*layout=*/4, const_cast<char*>("half"),
           &v2_ws_, &v2_exec_);
       aclSetAclOpExecutorRepeatable(v2_exec_);
     } else {
-      aclSetInputTensorAddr(v2_exec_, 0, t_q, query_out.data());
-      aclSetInputTensorAddr(v2_exec_, 1, t_k, key_out.data());
+      aclSetInputTensorAddr(v2_exec_, 0, t_q, q_out.data());
+      aclSetInputTensorAddr(v2_exec_, 1, t_k, k_out.data());
     }
 
     auto& arena = ascend::GetWorkspacePool().Ensure(stream, v2_ws_);

src/ascend/rotary_embedding/kernel_atb.h

@@ -7,6 +7,7 @@
 #include <cstddef>
 #include <cstdint>
 #include <cstring>
+#include <optional>
 #include <vector>
 
 #include "acl/acl.h"
@@ -58,7 +59,9 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
  public:
   Operator(const Tensor positions, const Tensor query, const Tensor key,
            const Tensor cos_sin_cache, int64_t head_size, int64_t rotary_dim,
-           bool is_neox_style, Tensor query_out, Tensor key_out)
+           bool is_neox_style,
+           std::optional<Tensor> query_out = std::nullopt,
+           std::optional<Tensor> key_out = std::nullopt)
       : RotaryEmbedding(positions, query, key, cos_sin_cache, head_size,
                         rotary_dim, is_neox_style, query_out, key_out),
         is_neox_style_{is_neox_style} {
@@ -149,10 +152,16 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
 
   void operator()(const Tensor positions, const Tensor query, const Tensor key,
                   const Tensor cos_sin_cache, int64_t head_size,
-                  int64_t rotary_dim, bool is_neox_style, Tensor query_out,
-                  Tensor key_out) const override {
+                  int64_t rotary_dim, bool is_neox_style,
+                  std::optional<Tensor> query_out,
+                  std::optional<Tensor> key_out) const override {
     auto stream = static_cast<aclrtStream>(stream_);
 
+    // Resolve optional out buffers (inplace on `query` / `key` when omitted).
+    // Non-const so `.data()` returns a writable `void*`.
+    Tensor q_out = query_out.value_or(query);
+    Tensor k_out = key_out.value_or(key);
+
     int64_t T = query.size(0);
     int64_t D = head_size;
 
@@ -202,15 +211,15 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
     // Step 2: Copy q->q_out, k->k_out if not in-place.
     size_t elem_sz = query.element_size();
 
-    if (query.data() != query_out.data()) {
-      aclrtMemcpyAsync(query_out.data(),
+    if (query.data() != q_out.data()) {
+      aclrtMemcpyAsync(q_out.data(),
                        static_cast<size_t>(T * hiddenQ) * elem_sz, query.data(),
                        static_cast<size_t>(T * hiddenQ) * elem_sz,
                        ACL_MEMCPY_DEVICE_TO_DEVICE, stream);
     }
 
-    if (key.data() != key_out.data()) {
-      aclrtMemcpyAsync(key_out.data(),
+    if (key.data() != k_out.data()) {
+      aclrtMemcpyAsync(k_out.data(),
                        static_cast<size_t>(T * hiddenK) * elem_sz, key.data(),
                        static_cast<size_t>(T * hiddenK) * elem_sz,
                        ACL_MEMCPY_DEVICE_TO_DEVICE, stream);
@@ -227,9 +236,9 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 1>
     uint64_t gathered_bytes = static_cast<uint64_t>(T * D) * elem_size_;
 
     atb::Tensor t_q =
-        ascend::toAtbTensor(q_2d_shape_, acl_dt_, query_out.data(), q_bytes);
+        ascend::toAtbTensor(q_2d_shape_, acl_dt_, q_out.data(), q_bytes);
     atb::Tensor t_k =
-        ascend::toAtbTensor(k_2d_shape_, acl_dt_, key_out.data(), k_bytes);
+        ascend::toAtbTensor(k_2d_shape_, acl_dt_, k_out.data(), k_bytes);
     atb::Tensor t_cos = ascend::toAtbTensor(cos_sin_gathered_shape_, acl_dt_,
                                             cos_dev_, gathered_bytes);
     atb::Tensor t_sin = ascend::toAtbTensor(cos_sin_gathered_shape_, acl_dt_,

src/ascend/rotary_embedding/kernel_sincos_cache.h

@@ -3,6 +3,7 @@
 
 #include <cassert>
 #include <cstdint>
+#include <optional>
 
 #include "acl/acl.h"
 #include "aclnn/aclnn_base.h"
@@ -41,10 +42,17 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 2>
  public:
   Operator(const Tensor positions, const Tensor query, const Tensor key,
            const Tensor cos_sin_cache, int64_t head_size, int64_t rotary_dim,
-           bool is_neox_style, Tensor query_out, Tensor key_out)
+           bool is_neox_style,
+           std::optional<Tensor> query_out = std::nullopt,
+           std::optional<Tensor> key_out = std::nullopt)
       : RotaryEmbedding(positions, query, key, cos_sin_cache, head_size,
                         rotary_dim, is_neox_style, query_out, key_out),
         max_seq_len_{cos_sin_cache.size(0)} {
+    // Resolve optional out buffers (inplace on `query` / `key` when omitted).
+    // Non-const so `.data()` returns a writable `void*`.
+    Tensor q_out = query_out.value_or(query);
+    Tensor k_out = key_out.value_or(key);
+
     const int64_t T = num_tokens_;
     const int64_t Nq = num_heads_;
     const int64_t Nkv = num_kv_heads_;
@@ -61,9 +69,9 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 2>
         {max_seq_len_, rotary_dim_}, acl_dt,
         const_cast<void*>(cos_sin_cache.data()));
     q_out_cache_ =
-        ascend::AclTensorCache({T, Nq * D}, acl_dt, query_out.data());
+        ascend::AclTensorCache({T, Nq * D}, acl_dt, q_out.data());
     k_out_cache_ =
-        ascend::AclTensorCache({T, Nkv * D}, acl_dt, key_out.data());
+        ascend::AclTensorCache({T, Nkv * D}, acl_dt, k_out.data());
   }
 
   ~Operator() {
@@ -83,10 +91,15 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 2>
 
   void operator()(const Tensor positions, const Tensor query, const Tensor key,
                   const Tensor cos_sin_cache, int64_t head_size,
-                  int64_t rotary_dim, bool is_neox_style, Tensor query_out,
-                  Tensor key_out) const override {
+                  int64_t rotary_dim, bool is_neox_style,
+                  std::optional<Tensor> query_out,
+                  std::optional<Tensor> key_out) const override {
     auto stream = static_cast<aclrtStream>(stream_);
 
+    // Resolve optional out buffers (inplace on `query` / `key` when omitted).
+    Tensor q_out = query_out.value_or(query);
+    Tensor k_out = key_out.value_or(key);
+
     // Refresh cached descriptors with the current-call data pointers —
     // `Operator::call()` cache matches on shape/stride/dtype, so one
     // instance may serve multiple calls with different underlying buffers.
@@ -95,8 +108,8 @@ class Operator<RotaryEmbedding, Device::Type::kAscend, 2>
     auto t_k = k_in_cache_.get(const_cast<void*>(key.data()));
     auto t_cache =
         cos_sin_cache_cache_.get(const_cast<void*>(cos_sin_cache.data()));
-    auto t_q_out = q_out_cache_.get(query_out.data());
-    auto t_k_out = k_out_cache_.get(key_out.data());
+    auto t_q_out = q_out_cache_.get(const_cast<void*>(q_out.data()));
+    auto t_k_out = k_out_cache_.get(const_cast<void*>(k_out.data()));
 
     uint64_t ws_size = 0;
     aclOpExecutor* executor = nullptr;

src/base/rotary_embedding.h

@@ -2,6 +2,7 @@
 #define INFINI_OPS_BASE_ROTARY_EMBEDDING_H_
 
 #include <cstddef>
+#include <optional>
 #include <vector>
 
 #include "operator.h"
@@ -13,10 +14,17 @@ class RotaryEmbedding : public Operator<RotaryEmbedding> {
   // Accepts 2D `[T, N*D]` (vLLM convention) or 3D `[T, N, D]`.
   // `num_heads_` and `num_kv_heads_` are derived from `numel / (T *
   // head_size)`.
+  //
+  // `query_out` / `key_out` are optional.  When omitted, the kernel writes
+  // back into `query` / `key` — matching vLLM's inplace
+  // `RotaryEmbedding.forward(positions, query, key)` signature.  Pass
+  // explicit out buffers only when the caller needs a separate
+  // destination.
   RotaryEmbedding(const Tensor positions, const Tensor query, const Tensor key,
                   const Tensor cos_sin_cache, int64_t head_size,
-                  int64_t rotary_dim, bool is_neox_style, Tensor query_out,
-                  Tensor key_out)
+                  int64_t rotary_dim, bool is_neox_style,
+                  std::optional<Tensor> query_out = std::nullopt,
+                  std::optional<Tensor> key_out = std::nullopt)
       : num_tokens_{query.size(0)},
         num_heads_{static_cast<int64_t>(query.numel()) /
                    (static_cast<int64_t>(query.size(0)) * head_size)},
@@ -28,12 +36,12 @@ class RotaryEmbedding : public Operator<RotaryEmbedding> {
         query_shape_{query.shape()},
         key_shape_{key.shape()},
         cos_sin_cache_shape_{cos_sin_cache.shape()},
-        query_out_shape_{query_out.shape()},
-        key_out_shape_{key_out.shape()},
+        query_out_shape_{query_out.value_or(query).shape()},
+        key_out_shape_{key_out.value_or(key).shape()},
         query_strides_{query.strides()},
         key_strides_{key.strides()},
-        query_out_strides_{query_out.strides()},
-        key_out_strides_{key_out.strides()} {
+        query_out_strides_{query_out.value_or(query).strides()},
+        key_out_strides_{key_out.value_or(key).strides()} {
     assert(
         (query.ndim() == 2 || query.ndim() == 3) &&
         "`RotaryEmbedding` requires query to be 2D [T, N*D] or 3D [T, N, D]");
@@ -47,8 +55,10 @@ class RotaryEmbedding : public Operator<RotaryEmbedding> {
   virtual void operator()(const Tensor positions, const Tensor query,
                           const Tensor key, const Tensor cos_sin_cache,
                           int64_t head_size, int64_t rotary_dim,
-                          bool is_neox_style, Tensor query_out,
-                          Tensor key_out) const = 0;
+                          bool is_neox_style,
+                          std::optional<Tensor> query_out = std::nullopt,
+                          std::optional<Tensor> key_out = std::nullopt)
+      const = 0;
 
  protected:
   Tensor::Size num_tokens_{0};

tests/test_rotary_embedding.py

@@ -562,3 +562,78 @@ def test_rotary_embedding_partial(
 
     _assert_close(q_out, ref_q, rtol, atol)
     _assert_close(k_out, ref_k, rtol, atol)
+
+
+@pytest.mark.parametrize("implementation_index", (0, 1))
+@pytest.mark.parametrize(
+    ("dtype", "rtol", "atol"),
+    (
+        # V2 accumulates ~4 ULP error in fp16 (kernel.h doc: max diff ~0.008);
+        # ATB `RopeParam` is similar.  Use atol=5e-3 for honest headroom.
+        (torch.float16, 1e-2, 5e-3),
+        (torch.bfloat16, 1e-2, 5e-3),
+    ),
+)
+@pytest.mark.parametrize("device", ("npu",))
+def test_rotary_embedding_inplace(implementation_index, dtype, rtol, atol, device):
+    """Verify the inplace path (`query_out` / `key_out` omitted).
+
+    Matches vLLM's `RotaryEmbedding.forward(positions, query, key)`
+    convention where the op mutates `query` / `key` directly.
+    """
+    if not (hasattr(torch, "npu") and torch.npu.is_available()):
+        pytest.skip("NPU not available")
+
+    active_indices = infini.ops.RotaryEmbedding.active_implementation_indices(device)
+
+    if implementation_index not in active_indices:
+        pytest.skip(
+            f"Implementation index={implementation_index} not active on this build"
+        )
+
+    num_tokens = 4
+    num_heads = 8
+    num_kv_heads = 8
+    head_size = 64
+    rotary_dim = head_size
+    max_seq_len = 32
+
+    positions = randint_strided(
+        0, max_seq_len, (num_tokens,), None, dtype=torch.int64, device=device
+    )
+    query = randn_strided(
+        (num_tokens, num_heads, head_size), None, dtype=dtype, device=device
+    )
+    key = randn_strided(
+        (num_tokens, num_kv_heads, head_size), None, dtype=dtype, device=device
+    )
+    cos_sin_cache = randn_strided(
+        (max_seq_len, rotary_dim), None, dtype=dtype, device=device
+    )
+
+    # Reference: apply RoPE to clones of the original inputs.
+    ref_q, ref_k = _ref_rotary_embedding(
+        positions,
+        query.clone(),
+        key.clone(),
+        cos_sin_cache,
+        head_size,
+        rotary_dim,
+        is_neox_style=True,
+    )
+
+    # Inplace call — no `query_out` / `key_out` supplied.
+    infini.ops.rotary_embedding(
+        positions,
+        query,
+        key,
+        cos_sin_cache,
+        head_size,
+        rotary_dim,
+        True,
+        implementation_index=implementation_index,
+        stream=get_npu_stream(query),
+    )
+
+    _assert_close(query, ref_q, rtol, atol)
+    _assert_close(key, ref_k, rtol, atol)