Fix the bugs about operator registration about PyTorch Dispatcher

**Background:**

There are two principles about operator registration in PyTorch
- The same namespace can be only registered once by `TORCH_LIBRARY`
- The operator signatures can be only registered once by `def`

Therefore, for the first problem, we can use `TORCH_LIBRARY_FRAGMEN` to expand operators within the same NAMESPACE.

For the second problem, operator signatures are not required for registration, so we only need to explicitly define operators that are only defined in vllm-ascend.

Signed-off-by: fffrog <[email protected]>

Author: fffrog

csrc/torch_binding.cpp

@@ -38,7 +38,7 @@ AscendType get_dtype_from_torch(at::ScalarType scalarType)
     }
 }
 
-std::tuple<at::Tensor, at::Tensor> rotary_embedding(at::Tensor &positions, at::Tensor &query, at::Tensor &key,
+void rotary_embedding(at::Tensor &positions, at::Tensor &query, std::optional<at::Tensor> key,
     int64_t head_size, at::Tensor &cos_sin_cache,  bool is_neox)
 {
     int32_t deviceId = 0;
@@ -47,22 +47,23 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding(at::Tensor &positions, at::T
     TORCH_CHECK(
         positions_ndim == 1 || positions_ndim == 2,
         "positions must have shape [num_tokens] or [batch_size, seq_len]");
+    TORCH_CHECK(key.has_value(), "rotary_embedding: key must have a value");
     if (positions_ndim == 1) {
       TORCH_CHECK(
-          query.size(0) == positions.size(0) && key.size(0) == positions.size(0),
+          query.size(0) == positions.size(0) && key.value().size(0) == positions.size(0),
           "query, key and positions must have the same number of tokens");
     }
     if (positions_ndim == 2) {
       TORCH_CHECK(
           query.size(0) == positions.size(0) &&
-              key.size(0) == positions.size(0) &&
+              key.value().size(0) == positions.size(0) &&
               query.size(1) == positions.size(1) &&
-              key.size(1) == positions.size(1),
+              key.value().size(1) == positions.size(1),
           "query, key and positions must have the same batch_size and seq_len");
     }
     TORCH_CHECK(head_size % 32 == 0, "rotary_embedding: headSize should be divisible by 32");
     int query_hidden_size = query.numel() / num_tokens;
-    int key_hidden_size = key.numel() / num_tokens;
+    int key_hidden_size = key.value().numel() / num_tokens;
     TORCH_CHECK(query_hidden_size % head_size == 0);
     TORCH_CHECK(key_hidden_size % head_size == 0);
     TORCH_CHECK(is_neox == true, "rotary_embedding: neox=false is not supported as custom kernel in vllm-ascend");
@@ -72,18 +73,18 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding(at::Tensor &positions, at::T
     int num_kv_heads = key_hidden_size / head_size;
     TORCH_CHECK(num_heads % num_kv_heads == 0);
     at::Tensor query_dst = at::empty({num_tokens, num_heads, head_size}, query.options());
-    at::Tensor key_dst = at::empty({num_tokens, num_kv_heads, head_size}, key.options());
+    at::Tensor key_dst = at::empty({num_tokens, num_kv_heads, head_size}, key.value().options());
 
     int rot_dim = cos_sin_cache.size(1);
     int seq_dim_idx = positions_ndim - 1;
     int64_t *position_ids_ptr = positions.data_ptr<int64_t>();
     void *query_dst_ptr = query_dst.data_ptr();
     void *key_dst_ptr = key_dst.data_ptr();
     void *query_ptr = query.data_ptr();
-    void *key_ptr = key.data_ptr();
+    void *key_ptr = key.value().data_ptr();
     void *cos_sin_cache_ptr = cos_sin_cache.data_ptr();
     int64_t query_stride = query.stride(seq_dim_idx);
-    int64_t key_stride = key.stride(seq_dim_idx);
+    int64_t key_stride = key.value().stride(seq_dim_idx);
     int64_t dst_query_stride = query_dst.stride(0);
     int64_t dst_key_stride = key_dst.stride(0);
     at::ScalarType scalar_type = query.scalar_type();
@@ -104,7 +105,9 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding(at::Tensor &positions, at::T
         return 0;
     });
     cmd.Run();
-    return {query_dst, key_dst};
+
+    query.copy_(query_dst);
+    key.value().copy_(key_dst);
 }
 
 std::tuple<at::Tensor, at::Tensor> get_masked_input_and_mask(
@@ -385,43 +388,36 @@ at::Tensor sgmv_expand(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indic
 }
 } // namespace vllm_ascend
 
-TORCH_LIBRARY_EXPAND(_C, ops)
+TORCH_LIBRARY_FRAGMENT_EXPAND(_C, ops)
 {
-    // vLLM-Ascend custom ops
-    ops.def("weak_ref_tensor(Tensor input) -> Tensor");
-    ops.impl("weak_ref_tensor", torch::kPrivateUse1, &vllm_ascend::weak_ref_tensor);
-
-    // Rotary embedding
-    // Apply GPT-NeoX style rotary embedding to query and key.
-    ops.def(
-        "rotary_embedding(Tensor positions, Tensor! query,"
-        "                 Tensor! key, int head_size,"
-        "                 Tensor cos_sin_cache, bool is_neox) -> (Tensor query, Tensor key)");
-    ops.impl("rotary_embedding", torch::kPrivateUse1, &vllm_ascend::rotary_embedding);
-
     ops.def(
         "get_masked_input_and_mask(Tensor input, "
         "                         int org_vocab_start_index, "
         "                         int org_vocab_end_index, "
         "                         int num_org_vocab_padding, "
         "                         int added_vocab_start_index, "
         "                         int added_vocab_end_index) -> (Tensor masked_input, Tensor mask)");
-    ops.impl("get_masked_input_and_mask", torch::kPrivateUse1, &vllm_ascend::get_masked_input_and_mask);
-
     ops.def("bgmv_shrink(Tensor! x, Tensor! weight, Tensor! indices, Tensor! y, float scale) -> ()");
-    ops.impl("bgmv_shrink", torch::kPrivateUse1, &vllm_ascend::bgmv_shrink);
-
     ops.def(
         "bgmv_expand(Tensor! x, Tensor! weight, Tensor! indices, Tensor! y,"
         "            int slice_offset, int slice_size) -> Tensor");
-    ops.impl("bgmv_expand", torch::kPrivateUse1, &vllm_ascend::bgmv_expand);
-
     ops.def("sgmv_shrink(Tensor! x, Tensor! weight, Tensor! lora_indices, Tensor! seq_len, Tensor! y, float scale) -> ()");
-    ops.impl("sgmv_shrink", torch::kPrivateUse1, &vllm_ascend::sgmv_shrink);
-
     ops.def(
         "sgmv_expand(Tensor! x, Tensor! weight, Tensor! lora_indices, Tensor! seq_len, Tensor! y,"
         "            int slice_offset, int slice_size) -> Tensor");
+}
+
+TORCH_LIBRARY_IMPL_EXPAND(_C, PrivateUse1, ops)
+{
+    // vLLM-Ascend custom ops
+    ops.impl("weak_ref_tensor", torch::kPrivateUse1, &vllm_ascend::weak_ref_tensor);
+    // Rotary embedding
+    // Apply GPT-NeoX style rotary embedding to query and key.
+    ops.impl("rotary_embedding", torch::kPrivateUse1, &vllm_ascend::rotary_embedding);
+    ops.impl("get_masked_input_and_mask", torch::kPrivateUse1, &vllm_ascend::get_masked_input_and_mask);
+    ops.impl("bgmv_shrink", torch::kPrivateUse1, &vllm_ascend::bgmv_shrink);
+    ops.impl("bgmv_expand", torch::kPrivateUse1, &vllm_ascend::bgmv_expand);
+    ops.impl("sgmv_shrink", torch::kPrivateUse1, &vllm_ascend::sgmv_shrink);
     ops.impl("sgmv_expand", torch::kPrivateUse1, &vllm_ascend::sgmv_expand);
 }
 

csrc/torch_binding_meta.cpp

@@ -36,23 +36,25 @@
 namespace vllm_ascend {
 namespace meta {
 
-std::tuple<at::Tensor, at::Tensor> rotary_embedding_meta(
+void rotary_embedding_meta(
   at::Tensor &positions,
   at::Tensor &query,
-  at::Tensor &key,
+  std::optional<at::Tensor> key,
   int64_t head_size, 
   at::Tensor &cos_sin_cache,
   bool is_neox) {
+    TORCH_CHECK(key.has_value(), "rotary_embedding_meta: key must have a value");
     auto num_tokens = positions.sym_numel();
     auto query_hidden_size = query.sym_numel() / num_tokens;
-    auto key_hidden_size = key.sym_numel() / num_tokens;
+    auto key_hidden_size = key.value().sym_numel() / num_tokens;
 
     auto num_heads = query_hidden_size / head_size;
     auto num_kv_heads = key_hidden_size / head_size;
-    at::Tensor query_dst = at::empty_symint({num_tokens, num_heads, head_size}, query.options());
-    at::Tensor key_dst = at::empty_symint({num_tokens, num_kv_heads, head_size}, key.options());
 
-    return {query_dst, key_dst};
+    c10::SymIntArrayRef query_shape({num_tokens, num_heads, head_size});
+    c10::SymIntArrayRef key_shape({num_tokens, num_kv_heads, head_size});
+    query.resize__symint(query_shape);
+    key.value().resize__symint(key_shape);
 }
 
 std::tuple<at::Tensor, at::Tensor> get_masked_input_and_mask_meta(
@@ -99,4 +101,4 @@ namespace {
     ops.impl("sgmv_expand", &vllm_ascend::meta::sgmv_expand_meta);
 
 }
-}
+}

csrc/utils.h

@@ -13,6 +13,10 @@
 // A version of the TORCH_LIBRARY macro that expands the NAME, i.e. so NAME
 // could be a macro instead of a literal token.
 #define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
+//
+// A version of the TORCH_LIBRARY_FRAGMENT macro that expands the NAME, i.e. so NAME
+// could be a macro instead of a literal token.
+#define TORCH_LIBRARY_FRAGMENT_EXPAND(NAME, MODULE) TORCH_LIBRARY_FRAGMENT(NAME, MODULE)
 
 // A version of the TORCH_LIBRARY_IMPL macro that expands the NAME, i.e. so NAME
 // could be a macro instead of a literal token.

tests/e2e/singlecard/ops/test_rotary_embedding.py

@@ -182,7 +182,7 @@ def test_rotary_embedding_quant_with_leading_dim(
     )
 
     ref_query, ref_key = rope.forward_native(positions, query, key)
-    query, key = torch.ops._C.rotary_embedding(
+    torch.ops._C.rotary_embedding(
         positions,
         query,
         key,
@@ -239,16 +239,16 @@ def forward(
         # we simulated a simple attention layer to test if it can be seamlessly captured into aclgraph
         qkv = self.qkv_proj(hidden_states)
         q, k, v = qkv.chunk(3, dim=-1)
-        query, key = torch.ops._C.rotary_embedding(
+        q_shape = q.shape
+        torch.ops._C.rotary_embedding(
             positions,
             q,
             k,
             self.rope.head_size,
             self.rope.cos_sin_cache,
             self.rope.is_neox_style,
         )
-        query = query.view(q.shape)
-        key = key.view(k.shape)
+        query = q.view(q_shape)
         o = self.o_proj(query)
         return o
 

vllm_ascend/ops/rotary_embedding.py

@@ -51,7 +51,7 @@ def _rope_forward_oot(
     # adopt custom kernel path for rotary_embedding
     if _custom_rotary_embedding_enabled(query, neox_style,
                                         self.head_size) and not is_310p():
-        query, key = torch.ops._C.rotary_embedding(
+        torch.ops._C.rotary_embedding(
             positions,
             query,
             key,

vllm_ascend/torchair/ops/torchair_rotary_embedding.py

@@ -62,7 +62,7 @@ def rope_forward_oot(
     # adopt custom kernel path for rotary_embedding
     if custom_rotary_embedding_enabled(query, neox_style,
                                        self.head_size) and not is_310p():
-        query, key = torch.ops._C.rotary_embedding(
+        torch.ops._C.rotary_embedding(
             positions,
             query,
             key,