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`

Considering that all custom operators defined in the current repo are only used by Ascend, instead of defining a common operator pattern by vLLM, all accelerators then follow this operator and complete the implementation based on their respective hardware, which is conducive to module functional abstraction.

Therefore, we can rename the operator registration namespace to an Ascend-specific namespace.

Signed-off-by: FFFrog <ljw1101.vip@gmail.com>

Author: FFFrog

benchmarks/ops/ben_vocabparallelembedding.py

@@ -112,7 +112,7 @@ def ref_fn():
 
     # Define custom function
     def custom_fn():
-        return torch.ops._C.get_masked_input_and_mask(
+        return torch.ops._C_ascend.get_masked_input_and_mask(
             input_tensor,
             test_case["org_start"],
             test_case["org_end"],

csrc/torch_binding.cpp

@@ -141,7 +141,7 @@ std::tuple<at::Tensor, at::Tensor> get_masked_input_and_mask(
     TP2, rank 1:
                             |< -----------BASE----------- >|< -BASE PADDING- >|< -----------LORA PADDING----------- >|
     corresponding token_id: | 512 | 513 | 514 | ... | 1009 | -1  | ...  | -1  |  -1  | ... |  -1  | -1  | ... |   -1 |
-                     index: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 512  | ... | 519  | 520 | ... |  543 | 
+                     index: |  0  |  1  |  2  | ... | 497  | 498 | ...  | 511 | 512  | ... | 519  | 520 | ... |  543 |
     Parameters:
         org_vocab_start_index //base embeddings start
         org_vocab_end_index //base embeddings end
@@ -164,22 +164,22 @@ std::tuple<at::Tensor, at::Tensor> get_masked_input_and_mask(
     // Create output tensors
     at::Tensor masked_input = at::empty_like(input);
 	at::Tensor mask = at::empty_like(input).to(at::kBool);
-    
+
     // Get data pointers
     void *input_ptr = input.data_ptr();
     void *masked_input_ptr = masked_input.data_ptr();
     void *mask_ptr = mask.data_ptr();
-    
+
     // Get current stream
     aclrtStream stream = c10_npu::getCurrentNPUStream().stream();
-    
+
     // Get scalar type
     at::ScalarType scalar_type = input.scalar_type();
-    
+
     // Create and configure OpCommand
     at_npu::native::OpCommand cmd;
     cmd.Name("get_masked_input_and_mask");
-    cmd.SetCustomHandler([scalar_type, size, stream, 
+    cmd.SetCustomHandler([scalar_type, size, stream,
                          input_ptr, masked_input_ptr, mask_ptr,
                          org_vocab_start_index, org_vocab_end_index,
                          num_org_vocab_padding, added_vocab_start_index,
@@ -193,7 +193,7 @@ std::tuple<at::Tensor, at::Tensor> get_masked_input_and_mask(
         get_masked_input_and_mask_impl(
             stream,
             input_ptr,
-            masked_input_ptr, 
+            masked_input_ptr,
             mask_ptr,
             org_vocab_start_index,
             org_vocab_end_index,
@@ -203,7 +203,7 @@ std::tuple<at::Tensor, at::Tensor> get_masked_input_and_mask(
             size,
             loop_cnt,
             aiv_num);
-            
+
         return 0;
     });
     cmd.Run();
@@ -320,8 +320,8 @@ void sgmv_shrink(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indices, at
     aclrtStream stream = c10_npu::getCurrentNPUStream().stream();
     at_npu::native::OpCommand cmd;
     cmd.Name("sgmv_shrink");
-    cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, 
-                          seq_len_ptr, seq_len_size, y_ptr, 
+    cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size,
+                          seq_len_ptr, seq_len_size, y_ptr,
                           batch_size, input_hidden_token, lora_rank, scale_f]() -> int {
         auto dtype = get_dtype_from_torch(scalar_type);
         int device_id = 0;
@@ -330,7 +330,7 @@ void sgmv_shrink(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indices, at
         int num_tokens_per_core = (batch_size + aiv_num - 1) / aiv_num;
         TORCH_CHECK("num_tokens_per_core != 0", "num_tokens_per_core should not be 0");
         sgmv_shrink_impl(dtype, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size,
-                         y_ptr, batch_size, 
+                         y_ptr, batch_size,
                          num_tokens_per_core, input_hidden_token, lora_rank, scale_f);
         return 0;
     });
@@ -367,15 +367,15 @@ at::Tensor sgmv_expand(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indic
     aclrtStream stream = c10_npu::getCurrentNPUStream().stream();
     at_npu::native::OpCommand cmd;
     cmd.Name("sgmv_expand");
-    cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, y_ptr, y_out_ptr, 
+    cmd.SetCustomHandler([scalar_type, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, y_ptr, y_out_ptr,
                           batch_size, lora_rank, slice_offset, slice_size, output_full_dim]() -> int {
         auto dtype = get_dtype_from_torch(scalar_type);
         int device_id = 0;
         int64_t aiv_num = 0;
         TORCH_CHECK(aclGetDeviceCapability(device_id, ACL_DEVICE_INFO_VECTOR_CORE_NUM, &aiv_num) == ACL_SUCCESS);
         int num_tokens_per_core = (batch_size + aiv_num - 1) / aiv_num;
         TORCH_CHECK("num_tokens_per_core != 0", "num_tokens_per_core should not be 0");
-        sgmv_expand_impl(dtype, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, y_ptr, y_out_ptr, 
+        sgmv_expand_impl(dtype, stream, x_ptr, weight_ptr, lora_indices_ptr, lora_indices_size, seq_len_ptr, seq_len_size, y_ptr, y_out_ptr,
                          batch_size, num_tokens_per_core, lora_rank, slice_size, slice_offset, output_full_dim);
         return 0;
     });
@@ -384,7 +384,7 @@ at::Tensor sgmv_expand(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_indic
 }
 } // namespace vllm_ascend
 
-TORCH_LIBRARY_EXPAND(_C, ops)
+TORCH_LIBRARY_EXPAND(CONCAT(_C, _ascend), ops)
 {
     // vLLM-Ascend custom ops
     ops.def("weak_ref_tensor(Tensor input) -> Tensor");

csrc/torch_binding_meta.cpp

@@ -40,7 +40,7 @@ std::tuple<at::Tensor, at::Tensor> rotary_embedding_meta(
   at::Tensor &positions,
   at::Tensor &query,
   at::Tensor &key,
-  int64_t head_size, 
+  int64_t head_size,
   at::Tensor &cos_sin_cache,
   bool is_neox) {
     auto num_tokens = positions.sym_numel();
@@ -86,9 +86,9 @@ at::Tensor sgmv_expand_meta(at::Tensor &x, at::Tensor &weight, at::Tensor &lora_
 } // namespace vllm_ascend
 
 namespace {
-  // Register the meta implementations of the custom kernels for symbolic tracing, this will also 
+  // Register the meta implementations of the custom kernels for symbolic tracing, this will also
   // the custom kernel been captured into aclgraph
-  TORCH_LIBRARY_IMPL_EXPAND(_C, Meta, ops) {
+  TORCH_LIBRARY_IMPL_EXPAND(CONCAT(_C, _ascend), Meta, ops) {
     // Rotary embedding meta implementation
     ops.impl("rotary_embedding", &vllm_ascend::meta::rotary_embedding_meta);
     // Masked input and mask meta implementation
@@ -99,4 +99,4 @@ namespace {
     ops.impl("sgmv_expand", &vllm_ascend::meta::sgmv_expand_meta);
 
 }
-}
+}

tests/e2e/singlecard/ops/test_bgmv_expand.py

@@ -33,8 +33,8 @@ def test_bgmv_expand():
     y_npu = y.npu()
 
     y_out = bgmv_expand_cpu_impl(x, w, indices, y, 0, 128)
-    y_out_npu = torch.ops._C.bgmv_expand(x_npu, w_npu, indices_npu, y_npu, 0,
-                                         128)
+    y_out_npu = torch.ops._C_ascend.bgmv_expand(x_npu, w_npu, indices_npu,
+                                                y_npu, 0, 128)
 
     # Compare the results.
     torch.testing.assert_close(y_out_npu.cpu(),

tests/e2e/singlecard/ops/test_bgmv_shrink.py

@@ -33,7 +33,7 @@ def test_bgmv_shrink():
     y_npu = y.npu()
 
     y = bgmv_shrink_cpu_impl(x, w, indices, y, 0.5)
-    torch.ops._C.bgmv_shrink(x_npu, w_npu, indices_npu, y_npu, 0.5)
+    torch.ops._C_ascend.bgmv_shrink(x_npu, w_npu, indices_npu, y_npu, 0.5)
 
     # Compare the results.
     torch.testing.assert_close(y_npu.cpu(),

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(
+    query, key = torch.ops._C_ascend.rotary_embedding(
         positions,
         query,
         key,
@@ -239,7 +239,7 @@ 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(
+        query, key = torch.ops._C_ascend.rotary_embedding(
             positions,
             q,
             k,

tests/e2e/singlecard/ops/test_vocabparallelembedding.py

@@ -72,7 +72,7 @@ def test_get_masked_input_and_mask(
 
     # Get custom op result
     print("input_tensor:", input_tensor)
-    custom_masked_input, custom_mask = torch.ops._C.get_masked_input_and_mask(
+    custom_masked_input, custom_mask = torch.ops._C_ascend.get_masked_input_and_mask(
         input_tensor, test_case["org_start"], test_case["org_end"],
         test_case["padding"], test_case["added_start"], test_case["added_end"])
 

tests/ut/ops/test_rotary_embedding.py

@@ -94,7 +94,7 @@ def setUp(self):
         self.mock_self.cos_sin_cache = self.cos_sin_cache
         self.mock_self.is_neox_style = self.is_neox_style
 
-    @patch('torch.ops._C')
+    @patch('torch.ops._C_ascend')
     @patch('vllm_ascend.ops.rotary_embedding.is_310p', return_value=False)
     @patch('vllm_ascend.ops.rotary_embedding._custom_rotary_embedding_enabled',
            return_value=True)

tests/ut/torchair/ops/test_torchair_rotary_embedding.py

@@ -104,7 +104,7 @@ def test_rope_forward_oot_torchair_enabled_base(self,
         self.assertTrue(torch.equal(result_q, self.query))
         self.assertTrue(torch.equal(result_k, self.key))
 
-    @patch('torch.ops._C')
+    @patch('torch.ops._C_ascend')
     @patch(
         'vllm_ascend.torchair.ops.torchair_rotary_embedding.get_ascend_config')
     @patch('vllm_ascend.torchair.ops.torchair_rotary_embedding.is_310p',

vllm_ascend/compilation/acl_graph.py

@@ -15,7 +15,8 @@
 from vllm.forward_context import BatchDescriptor, get_forward_context
 from vllm.logger import logger
 from vllm.platforms import current_platform
-from vllm.utils import weak_ref_tensors
+
+from ..utils import weak_ref_tensors
 
 
 @dataclasses.dataclass
@@ -35,10 +36,10 @@ class ACLGraphWrapper:
 
     The workflow of this wrapper in the aclgraph dispatching is as follows:
     1. At initialization, a runtime mode is assigned to the wrapper (FULL or
-    PIECEWISE). 
-    2. At runtime, the wrapper receives a runtime_mode and a 
+    PIECEWISE).
+    2. At runtime, the wrapper receives a runtime_mode and a
     batch_descriptor(key) from the forward context and blindly trust them
-    for aclgraph dispatching. 
+    for aclgraph dispatching.
     3. If runtime_mode is NONE or runtime_mode does not match the mode of the
     wrapper, just call the runnable directly.
     4. Otherwise, i.e., the runtime_mode matches the mode of the wrapper,
@@ -47,9 +48,9 @@ class ACLGraphWrapper:
 
     Note: ACLGraphWrapper does not store persistent buffers or copy any
     runtime inputs into that buffers for replay. We assume implementing them
-    is done outside of the wrapper. That is because we do not make any 
+    is done outside of the wrapper. That is because we do not make any
     assumption on the dynamic shape (batch size) of the runtime inputs, as a
-    trade-off for staying orthogonal to compilation logic. Nevertheless, 
+    trade-off for staying orthogonal to compilation logic. Nevertheless,
     tracing and checking the input addresses to be consistent during replay is
     guaranteed when VLLM_LOGGING_LEVEL == "DEBUG".
     """