[Misc] Move lora patch file into lora module (#2797)

Cleanup useless file in patch module. Update the lora support list is OK
in vLLM Ascend, no need to patch vLLM


- vLLM version: v0.10.1.1
- vLLM main:
f4962a6d55

Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>

tests/ut/test_platform.py

@@ -559,9 +559,8 @@ class TestNPUPlatform(TestBase):
 
     def test_get_punica_wrapper(self):
         result = self.platform.get_punica_wrapper()
-        self.assertEqual(
-            result,
-            "vllm_ascend.lora.punica_wrapper.punica_npu.PunicaWrapperNPU")
+        self.assertEqual(result,
+                         "vllm_ascend.lora.punica_npu.PunicaWrapperNPU")
 
     @patch("torch.npu.reset_peak_memory_stats")
     @patch("torch.npu.max_memory_allocated")

vllm_ascend/lora/punica_npu.py

@@ -11,12 +11,14 @@ if is_310p():
                                          bgmv_shrink, sgmv_expand,
                                          sgmv_expand_slice, sgmv_shrink)
 else:
-    from vllm_ascend.lora.punica_wrapper.lora_ops import (
-        bgmv_expand, bgmv_expand_slice, bgmv_shrink, sgmv_expand,
-        sgmv_expand_slice, sgmv_shrink)
+    from vllm_ascend.lora.lora_ops import (bgmv_expand, bgmv_expand_slice,
+                                           bgmv_shrink, sgmv_expand,
+                                           sgmv_expand_slice, sgmv_shrink)
 
 from vllm.lora.punica_wrapper.punica_base import PunicaWrapperBase
 
+from vllm_ascend.lora.utils import refresh_all_lora_classes
+
 
 # The platforms that are compatible with the PyTorch-native implementation can
 # inherit this class
@@ -31,6 +33,7 @@ class PunicaWrapperNPU(PunicaWrapperBase):
                  device: Union[torch.device, str], **kwargs):
         PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches,
                                    device)
+        refresh_all_lora_classes()
 
     def _shrink_prefill(
         self,

vllm_ascend/lora/utils.py

@@ -1,28 +1,19 @@
 from typing import Optional
 
+import vllm
 from torch import nn
 from transformers import PretrainedConfig
 from vllm.config import LoRAConfig
 from vllm.lora.layers import (ColumnParallelLinearWithLoRA,
                               MergedColumnParallelLinearWithLoRA,
-                              RowParallelLinearWithLoRA)
+                              RowParallelLinearWithLoRA,
+                              VocabParallelEmbeddingWithLoRA)
 
 from vllm_ascend.ops.linear import (AscendColumnParallelLinear,
                                     AscendMergedColumnParallelLinear,
                                     AscendRowParallelLinear)
-
-
-class AscendRowParallelLinearWithLoRA(RowParallelLinearWithLoRA):
-
-    @classmethod
-    def can_replace_layer(
-        cls,
-        source_layer: nn.Module,
-        lora_config: LoRAConfig,
-        packed_modules_list: list,
-        model_config: Optional[PretrainedConfig],
-    ) -> bool:
-        return type(source_layer) is AscendRowParallelLinear
+from vllm_ascend.ops.vocab_parallel_embedding import \
+    AscendVocabParallelEmbedding
 
 
 class AscendColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
@@ -50,3 +41,37 @@ class AscendMergedColumnParallelLinearWithLoRA(
         model_config: Optional[PretrainedConfig],
     ) -> bool:
         return type(source_layer) is AscendMergedColumnParallelLinear
+
+
+class AscendRowParallelLinearWithLoRA(RowParallelLinearWithLoRA):
+
+    @classmethod
+    def can_replace_layer(
+        cls,
+        source_layer: nn.Module,
+        lora_config: LoRAConfig,
+        packed_modules_list: list,
+        model_config: Optional[PretrainedConfig],
+    ) -> bool:
+        return type(source_layer) is AscendRowParallelLinear
+
+
+class AscendVocabParallelEmbeddingWithLoRA(VocabParallelEmbeddingWithLoRA):
+
+    @classmethod
+    def can_replace_layer(
+        cls,
+        source_layer: nn.Module,
+        lora_config: LoRAConfig,
+        packed_modules_list: list,
+        model_config: Optional[PretrainedConfig],
+    ) -> bool:
+        return type(source_layer) is AscendVocabParallelEmbedding
+
+
+def refresh_all_lora_classes():
+    vllm.lora.utils._all_lora_classes.add(AscendColumnParallelLinearWithLoRA)
+    vllm.lora.utils._all_lora_classes.add(
+        AscendMergedColumnParallelLinearWithLoRA)
+    vllm.lora.utils._all_lora_classes.add(AscendRowParallelLinearWithLoRA)
+    vllm.lora.utils._all_lora_classes.add(AscendVocabParallelEmbeddingWithLoRA)

vllm_ascend/patch/__init__.py

@@ -46,6 +46,15 @@
 #       Need a PR to vllm to support get port from environment.
 #    Future Plan:
 #       Remove those patch when vllm merged them
+#   2. `torch.distributed.all_reduce`, `torch.distributed.broadcast`
+#    Why:
+#       tensor alignment for 310p
+#    How：
+#       rewrite all_reduce and broadcast in torch.distributed
+#    Related PR (if no, explain why):
+#       No, not ready yet.
+#    Future Plan:
+#       Find a better way to support tensor alignment for 310p without this patch.
 #
 # * Worker Patch:
 # ===============
@@ -86,19 +95,15 @@
 #       - https://github.com/vllm-project/vllm/pull/21591
 #    Future Plan:
 #       Revert it when vLLM merge #21591 and release new version
-# ** File: worker/patch_common/patch_linear.py **
+# ** File: worker/patch_common/patch_logits.py **
 # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-#   1. `vllm.model_executor.layers.linear.RowParallelLinear`
+#   1. `vllm._custom_ops.apply_repetition_penalties`
 #    Why:
-#       We need to fuse matmul and allreuce in `RowParallelLinear`
-#       to improve performance.
+#       apply_repetition_penalties in vLLM use tensor.is_cuda to check if tensor is on cuda. But the value is always True
+#       on ascend, thus we need to patch apply_repetition_penalties.
 #    How：
-#       Create a new class `AscendRowParallelLinear` that inherits from `RowParallelLinear`.
-#       In this class, we override the `forward` method to use
-#       torch_npu.npu_mm_all_reduce_base to replace matmul and allreduce.
+#       Remove the related cuda check in apply_repetition_penalties.
 #    Related PR (if no, explain why):
-#       - https://github.com/vllm-project/vllm-ascend/pull/1926
+#       - this is a bug by Ascend only. It can' be fixed in vLLM.
 #    Future Plan:
-#       Validate more models in all kinds of scenario,
-#       if performance is always improved, we can enable this patch by default and remove the env
-#       variable `VLLM_ASCEND_ENABLE_FUSE_MATMUL_ALLREDUCE` in the future.
+#       Fix this bug in torch-npu, bump torch-npu version and remove this patch.

vllm_ascend/patch/worker/patch_common/__init__.py

@@ -17,7 +17,4 @@
 
 import vllm_ascend.patch.worker.patch_common.patch_distributed  # noqa
 import vllm_ascend.patch.worker.patch_common.patch_logits  # noqa
-import vllm_ascend.patch.worker.patch_common.patch_lora  # noqa
-import vllm_ascend.patch.worker.patch_common.patch_lora_embedding  # noqa
-import vllm_ascend.patch.worker.patch_common.patch_lora_linear  # noqa
 import vllm_ascend.patch.worker.patch_common.patch_minicpm  # noqa

vllm_ascend/patch/worker/patch_common/patch_lora.py

@@ -1,15 +0,0 @@
-import vllm
-from vllm.lora.utils import _all_lora_classes
-
-from vllm_ascend.patch.worker.patch_common.patch_lora_embedding import \
-    AscendVocabParallelEmbeddingWithLoRA
-from vllm_ascend.patch.worker.patch_common.patch_lora_linear import (
-    AscendColumnParallelLinearWithLoRA,
-    AscendMergedColumnParallelLinearWithLoRA, AscendRowParallelLinearWithLoRA)
-
-_all_lora_classes.add(AscendRowParallelLinearWithLoRA)
-_all_lora_classes.add(AscendColumnParallelLinearWithLoRA)
-_all_lora_classes.add(AscendMergedColumnParallelLinearWithLoRA)
-_all_lora_classes.add(AscendVocabParallelEmbeddingWithLoRA)
-
-vllm.lora.utils._all_lora_classes = _all_lora_classes

vllm_ascend/patch/worker/patch_common/patch_lora_embedding.py

@@ -1,22 +0,0 @@
-from typing import Optional
-
-from torch import nn
-from transformers import PretrainedConfig
-from vllm.config import LoRAConfig
-from vllm.lora.layers import VocabParallelEmbeddingWithLoRA
-
-from vllm_ascend.ops.vocab_parallel_embedding import \
-    AscendVocabParallelEmbedding
-
-
-class AscendVocabParallelEmbeddingWithLoRA(VocabParallelEmbeddingWithLoRA):
-
-    @classmethod
-    def can_replace_layer(
-        cls,
-        source_layer: nn.Module,
-        lora_config: LoRAConfig,
-        packed_modules_list: list,
-        model_config: Optional[PretrainedConfig],
-    ) -> bool:
-        return type(source_layer) is AscendVocabParallelEmbedding

vllm_ascend/platform.py

@@ -270,7 +270,7 @@ class NPUPlatform(Platform):
 
     @classmethod
     def get_punica_wrapper(cls) -> str:
-        return "vllm_ascend.lora.punica_wrapper.punica_npu.PunicaWrapperNPU"
+        return "vllm_ascend.lora.punica_npu.PunicaWrapperNPU"
 
     @classmethod
     def get_current_memory_usage(cls,