Fixes for Triton implementations

Signed-off-by: Luka Govedič <lgovedic@redhat.com>

tests/compile/passes/ir/test_lowering.py

@@ -3,6 +3,7 @@
 import pytest
 import torch
 from torch import nn
+from torch._library.triton import set_wrap_triton_enabled
 
 import vllm.kernels  # noqa: F401 to register kernels
 from vllm import ir
@@ -47,6 +48,7 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config):
     with (
         ops.rms_norm.set_priority([rms_provider, "native"]),
         ir.enable_torch_wrap(True),
+        set_wrap_triton_enabled(False),  # set by default in forward context
     ):
         compiled_model = torch.compile(model, backend=backend, fullgraph=True)
         compiled_unlowered_model = torch.compile(

vllm/compilation/passes/ir/lowering_pass.py

@@ -10,6 +10,7 @@ from torch._inductor.pattern_matcher import (
     PatternMatcherPass,
     register_graph_pattern,
 )
+from torch._library.triton import set_wrap_triton_enabled
 from torch._ops import OpOverload, OpOverloadPacket
 
 from vllm.config import VllmConfig
@@ -92,14 +93,19 @@ class VllmIRLoweringPass(VllmInductorPass):
         # Defaults not present on node.args but required for replacement tracing
         bound_args = ir_op._py_signature.bind(*node.args)
         bound_args.apply_defaults()
-        match.replace_by_example(ir_op_impl.impl_fn, bound_args.args)
+        match.replace_by_example(
+            ir_op_impl.impl_fn, bound_args.args, run_functional_passes=False
+        )
 
     @VllmInductorPass.time_and_log
     def __call__(self, graph: fx.Graph) -> None:
         # clear at the beginning instead of end, so that tests can inspect
         self.selected_impls.clear()
 
-        count = self.patterns.apply(graph)
+        # Triton wrap is disabled in the forward context, enable it during lowering.
+        # This way make_fx replacement tracing handles the Triton kernel correctly.
+        with set_wrap_triton_enabled(True):
+            count = self.patterns.apply(graph)
         logger.debug("VllmIRLoweringPass lowered %d vLLM IR nodes", count)
 
         # TODO write self.selected_impls to depyf/tlparse dir

vllm/compilation/passes/pass_manager.py

@@ -113,8 +113,8 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
         VllmInductorPass.dump_prefix += 1
 
         # clean up after lowering again
-        self.post_cleanup(graph)
-        VllmInductorPass.dump_prefix += 1
+        # self.post_cleanup(graph)
+        # VllmInductorPass.dump_prefix += 1
 
         # always run fix_functionalization last
         self.fix_functionalization(graph)
@@ -190,7 +190,7 @@ class PostGradPassManager(CustomGraphPass):  # type: ignore[misc]
 
         passes.append(self.post_cleanup.uuid())
         passes.append(self.ir_lowering.uuid())
-        passes.append(self.post_cleanup.uuid())
+        # passes.append(self.post_cleanup.uuid())
         passes.append(self.fix_functionalization.uuid())
 
         # Include the compile range in the uuid to ensure that inductor

vllm/forward_context.py

@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
 from typing import Any
 
 import torch
+from torch._library.triton import set_wrap_triton_enabled
 
 import vllm.envs as envs
 import vllm.ir
@@ -326,6 +327,7 @@ def set_forward_context(
             vllm.ir.enable_torch_wrap(
                 vllm_config.compilation_config.ir_enable_torch_wrap
             ),
+            set_wrap_triton_enabled(False),
         ):
             yield
     finally: