[torch.compile] Reorganize vllm/compilation and tests/compile (0/N for vLLM IR) (#33731)

Signed-off-by: Luka Govedič <lgovedic@redhat.com>
Signed-off-by: ProExpertProg <luka.govedic@gmail.com>
Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com>

vllm/compilation/fix_functionalization.py → vllm/compilation/passes/utility/fix_functionalization.py

@@ -10,8 +10,8 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 
-from .fx_utils import is_func
-from .vllm_inductor_pass import VllmInductorPass
+from ..fx_utils import is_func
+from ..vllm_inductor_pass import VllmInductorPass
 
 logger = init_logger(__name__)
 

vllm/compilation/noop_elimination.py → vllm/compilation/passes/utility/noop_elimination.py

@@ -9,8 +9,8 @@ from torch.fx.experimental.symbolic_shapes import statically_known_true
 
 from vllm.logger import init_logger
 
-from .fx_utils import is_func
-from .vllm_inductor_pass import VllmInductorPass
+from ..fx_utils import is_func
+from ..vllm_inductor_pass import VllmInductorPass
 
 logger = init_logger(__name__)
 

vllm/compilation/post_cleanup.py → vllm/compilation/passes/utility/post_cleanup.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from torch import fx
 
-from vllm.compilation.vllm_inductor_pass import VllmInductorPass
+from ..vllm_inductor_pass import VllmInductorPass
 
 
 class PostCleanupPass(VllmInductorPass):

vllm/config/compilation.py

@@ -11,7 +11,7 @@ from typing import TYPE_CHECKING, Any, ClassVar, Literal
 from pydantic import Field, TypeAdapter, field_validator
 
 import vllm.envs as envs
-from vllm.compilation.inductor_pass import CallableInductorPass, InductorPass
+from vllm.compilation.passes.inductor_pass import CallableInductorPass, InductorPass
 from vllm.config.utils import (
     Range,
     config,
@@ -170,7 +170,9 @@ class PassConfig:
 
     @staticmethod
     def default_fi_allreduce_fusion_max_size_mb() -> dict[int, float]:
-        from vllm.compilation.collective_fusion import FI_ALLREDUCE_FUSION_MAX_SIZE_MB
+        from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
+            FI_ALLREDUCE_FUSION_MAX_SIZE_MB,
+        )
         from vllm.platforms import current_platform
 
         if not current_platform.is_cuda():

vllm/platforms/interface.py

@@ -191,7 +191,7 @@ class Platform:
         Get the pass manager class for this platform.
         It will be registered as a custom pass under the current_platform.pass_key.
         """
-        return "vllm.compilation.pass_manager.PostGradPassManager"
+        return "vllm.compilation.passes.pass_manager.PostGradPassManager"
 
     @classmethod
     def get_compile_backend(cls) -> str: