vllm - 💡(How to fix) Fix [RFC]: Refactor PassManager infrastructure [1 comments, 1 participants]

# Form 1: Python class
from vllm.config import CompilationConfig
from vllm.compilation.passes.fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass

compilation_config = CompilationConfig(
    pass_pipeline=[
        QKNormRoPEFusionPass,
    ],
)

---

from vllm.config import CompilationConfig

compilation_config = CompilationConfig(
    pass_pipeline=[
        "vllm.compilation.passes.fusion.qk_norm_rope_fusion.QKNormRoPEFusionPass",
    ],
)

---

class PassManager(CustomGraphPass):
    def __init__(
        self, passes: list[tuple[str, InductorPass]] | None = None
    ) -> None:
        self.passes: list[tuple[str, InductorPass]] = list(passes or [])
        self.hook: str = ""
        self._config_hash: str = ""

    def bind_to_hook(self, hook: str) -> "PassManager":
        self.hook = hook
        return self

    @classmethod
    def from_config(
        cls, config: "VllmConfig", hook: str | None = None
    ) -> "PassManager":
        """Build a configured pass manager from ``VllmConfig`` for a hook."""
        pass_pipeline = config.compilation_config.pass_pipeline
        manager = cls()
        if hook is not None:
            manager.bind_to_hook(hook)

        with set_current_vllm_config(config, check_compile=False):
            if pass_pipeline is None:
                pass_pipeline = cls.default_pass_pipeline(config)
            manager._config_hash = cls.compute_config_hash(config, pass_pipeline)
            for item in pass_pipeline:
                manager.passes.append(_resolve_pass(item))

        return manager

    @classmethod
    def compute_config_hash(
        cls, config: "VllmConfig", pass_pipeline: list[InductorPass] | None
    ) -> str:
        if config.compilation_config.pass_pipeline is None:
            return config.compilation_config.pass_config.compute_hash()
        return ""

    @classmethod
    def default_pass_pipeline(cls, config: "VllmConfig") -> list[InductorPass]:
        return []

    def describe_pipeline(self) -> list[str]:
        return [pass_id for pass_id, _ in self.passes]

    def __call__(self, graph: fx.Graph) -> None:
        compile_range = get_pass_context().compile_range
        for _, pass_ in self.passes:
            if pass_.is_applicable_for_range(compile_range):
                pass_(graph)

    def uuid(self) -> str:
        state: dict[str, Any] = {
            "manager": f"{type(self).__module__}.{type(self).__qualname__}",
            "hook": self.hook,
            "config_hash": self._config_hash,
            "compile_range": str(get_pass_context().compile_range),
            "passes": [
                {
                    "pass_id": pass_id,
                    "uuid": pass_.uuid(),
                }
                for pass_id, pass_ in self.passes
            ],
        }
        return InductorPass.hash_dict(state)

---

class PostGradPassManager(PassManager):

    @classmethod
    def compute_config_hash(
        cls, config: VllmConfig, pass_pipeline: list[InductorPass] | None
    ) -> str:
        return config.compilation_config.pass_config.compute_hash()

    @classmethod
    def default_pass_pipeline(cls, config: VllmConfig) -> list[InductorPass]:
        pass_config = config.compilation_config.pass_config
        pass_pipeline: list[InductorPass] = []

        # ... platform-specific fusion passes ...

        if pass_config.fuse_allreduce_rms and current_platform.is_cuda():
            pass_pipeline.append(AllReduceFusionPass)

        if pass_config.fuse_minimax_qk_norm and current_platform.is_cuda():
            pass_pipeline.append(MiniMaxQKNormPass)

        # ... other passes ...

        pass_pipeline.append(PostCleanupPass)
        pass_pipeline.append(VllmIRLoweringPass)
        pass_pipeline.append(PostCleanupPass)
        pass_pipeline.append(FixFunctionalizationPass)
        return pass_pipeline

---

@config
class PassConfig:
    def compute_hash(self) -> str:
        return hash_factors(get_hash_factors(self, set()))

@config
class PostGradPassConfig(PassConfig):
    fuse_norm_quant: bool = None
    fuse_act_quant: bool = None
    fuse_allreduce_rms: bool = None
    # ...

---

@config
class MyOOTPassConfig(PassConfig):
    enable_my_pass: bool = True

---

class MyPlatformPassManager(PassManager):

    @classmethod
    def compute_config_hash(
        cls, config: VllmConfig, pass_pipeline: list[InductorPass] | None
    ) -> str:
        pass_config = config.compilation_config.pass_config
        assert isinstance(pass_config, MyOOTPassConfig)
        return pass_config.compute_hash()

    @classmethod
    def default_pass_pipeline(cls, config: VllmConfig) -> list[InductorPass]:
        pass_config = config.compilation_config.pass_config
        assert isinstance(pass_config, MyOOTPassConfig)
        pipeline = []
        if pass_config.enable_my_pass:
            pipeline.append(MyPass)
        return pipeline