fix lint

Signed-off-by: ganyi <pleaplusone.gy@gmail.com>

Author: ganyi1996ppo

vllm_ascend/ops/meta_registration.py renamed to vllm_ascend/meta_registration.py

@@ -16,7 +16,6 @@
 # Both approaches enable tracing, export, and shape inference in PyTorch and vLLM, which
 # is essential for supporting `torch.compile` and aclgraph.
 
-
 # How to add a new meta implementation in Python:
 # -------------------------------------
 # 1. Write a Python function that takes the same arguments as your operator, and returns
@@ -35,56 +34,53 @@
 # 4. Example meta implementations are provided below for rotary_embedding and get_masked_input_and_mask.
 #
 # 5. When developing new custom ops, always provide a meta implementation to enable tracing,
-#    export, and shape inference in PyTorch and vLLM to enable the capture of `torch.compile` 
+#    export, and shape inference in PyTorch and vLLM to enable the capture of `torch.compile`
 #    and aclgraph.
 #
 # For more details, see: https://pytorch.org/docs/stable/notes/extending.html#meta-tensors
 
 lib = Library("_C", "IMPL")
 
 
-def register_meta_if_necessary(ns:str, op_name: str, fn, overload: str = ""):
-  if overload != "":
-    op_name = op_name + "." + overload
-  schema_to_find = ns + "::" + op_name
-  meta_impl_list = torch._C._dispatch_get_registrations_for_dispatch_key("Meta")
-  if schema_to_find in meta_impl_list:
-    return 
-  lib.impl(op_name, fn, "Meta")
+def register_meta_if_necessary(ns: str, op_name: str, fn, overload: str = ""):
+    if overload != "":
+        op_name = op_name + "." + overload
+    schema_to_find = ns + "::" + op_name
+    meta_impl_list = torch._C._dispatch_get_registrations_for_dispatch_key(
+        "Meta")
+    if schema_to_find in meta_impl_list:
+        return
+    lib.impl(op_name, fn, "Meta")
 
-def rotary_embedding_meta(
-    positions: torch.Tensor,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    head_size: int,
-    cos_sin_cache: torch.Tensor,
-    is_neox: bool):
 
-  num_tokens = positions.numel()
-  query_hidden_size = query.numel() / num_tokens
-  key_hidden_size = key.numel() / num_tokens
-  num_heads = query_hidden_size / head_size
-  num_kv_heads = key_hidden_size / head_size
+def rotary_embedding_meta(positions: torch.Tensor, query: torch.Tensor,
+                          key: torch.Tensor, head_size: int,
+                          cos_sin_cache: torch.Tensor, is_neox: bool):
 
-  query_dst = torch.empty_like(query).view(num_tokens, num_heads, head_size)
-  key_dst = torch.empty_like(key).view(num_tokens, num_kv_heads, head_size)
-  return query_dst, key_dst
+    num_tokens = positions.numel()
+    query_hidden_size = query.numel() // num_tokens
+    key_hidden_size = key.numel() // num_tokens
+    num_heads = query_hidden_size // head_size
+    num_kv_heads = key_hidden_size // head_size
 
+    query_dst = torch.empty_like(query).view(num_tokens, num_heads, head_size)
+    key_dst = torch.empty_like(key).view(num_tokens, num_kv_heads, head_size)
+    return query_dst, key_dst
 
-def get_masked_input_and_mask_meta(
-    input: torch.Tensor,
-    org_vocab_start_index: int,
-    org_vocab_end_index: int,
-    num_org_vocab_padding: int,
-    added_vocab_start_index: int,
-    added_vocab_end_index: int):
 
-  masked_input = torch.empty_like(input)
-  mask = torch.empty_like(input).to(torch.bool)
+def get_masked_input_and_mask_meta(input: torch.Tensor,
+                                   org_vocab_start_index: int,
+                                   org_vocab_end_index: int,
+                                   num_org_vocab_padding: int,
+                                   added_vocab_start_index: int,
+                                   added_vocab_end_index: int):
 
-  return masked_input, mask
+    masked_input = torch.empty_like(input)
+    mask = torch.empty_like(input).to(torch.bool)
 
+    return masked_input, mask
 
 
 register_meta_if_necessary("_C", "rotary_embedding", rotary_embedding_meta)
-register_meta_if_necessary("_C", "get_masked_input_and_mask", get_masked_input_and_mask_meta)
+register_meta_if_necessary("_C", "get_masked_input_and_mask",
+                           get_masked_input_and_mask_meta)

vllm_ascend/utils.py

@@ -215,9 +215,9 @@ def enable_custom_op():
         return _CUSTOM_OP_ENABLED
     try:
         # register custom ops into torch_library here
-        import vllm_ascend.vllm_ascend_C  # type: ignore  # noqa: F401
         # register the meta implementation for custom kernel if necessary
-        import vllm_ascend.ops.meta_registration  # type: ignore  # noqa: F401
+        import vllm_ascend.meta_registration  # type: ignore  # noqa: F401
+        import vllm_ascend.vllm_ascend_C  # type: ignore  # noqa: F401
         _CUSTOM_OP_ENABLED = True
     except ImportError:
         _CUSTOM_OP_ENABLED = False