@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 40bb175027
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 5, 17:15 UTC

[vLLM IR] 1/N Implement IR skeleton and rms_norm op (#33825)

Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com> Signed-off-by: chzhang <chaojun.zhang@intel.com> Signed-off-by: Luka Govedic <luka.govedic@gmail.com> Co-authored-by: Xinyu Chen <xinyu1.chen@intel.com> Co-authored-by: Chaojun Zhang <chaojun.zhang@intel.com> Co-authored-by: Luka Govedič <ProExpertProg@h100-01.nemg-001.lab.rdu2.dc.redhat.com>

Luka Govedič committed Mar 31, 2026 at 22:15 UTC 40bb17502736a9b6127cefc535303285b13e0630
49 files changed +2177 -265
.buildkite/test_areas/kernels.yaml
+10
@@ -2,6 +2,16 @@ group: Kernels
2 depends_on:
3 - image-build
4 steps:
5 +- label: vLLM IR Tests
6 + timeout_in_minutes: 10
7 + working_dir: "/vllm-workspace/"
8 + source_file_dependencies:
9 + - vllm/ir
10 + - vllm/kernels
11 + commands:
12 + - pytest -v -s tests/ir
13 + - pytest -v -s tests/kernels/ir
14 +
15 - label: Kernels Core Operation Test
16 timeout_in_minutes: 75
17 source_file_dependencies:
.github/CODEOWNERS
+4
@@ -13,6 +13,9 @@
13 /vllm/model_executor/layers/rotary_embedding.py @vadiklyutiy
14 /vllm/model_executor/model_loader @22quinn
15 /vllm/model_executor/layers/batch_invariant.py @yewentao256
16 +/vllm/ir @ProExpertProg
17 +/vllm/kernels/ @ProExpertProg @tjtanaa
18 +/vllm/kernels/helion @ProExpertProg @zou3519
19 /vllm/multimodal @DarkLight1337 @ywang96 @NickLucche @tjtanaa
20 /vllm/vllm_flash_attn @LucasWilkinson @MatthewBonanni
21 CMakeLists.txt @tlrmchlsmth @LucasWilkinson
@@ -74,6 +77,7 @@ CMakeLists.txt @tlrmchlsmth @LucasWilkinson
77 /tests/entrypoints @DarkLight1337 @robertgshaw2-redhat @aarnphm @NickLucche
78 /tests/evals @mgoin @vadiklyutiy
79 /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256
80 +/tests/kernels/ir @ProExpertProg @tjtanaa
81 /tests/models @DarkLight1337 @ywang96
82 /tests/multimodal @DarkLight1337 @ywang96 @NickLucche
83 /tests/quantization @mgoin @robertgshaw2-redhat @yewentao256 @pavanimajety
tests/compile/backend.py
+12 -4
@@ -8,7 +8,7 @@ from copy import deepcopy
8
9 import depyf
10 from torch import fx
11 -from torch._ops import OpOverload
11 +from torch._ops import OpOverload, OpOverloadPacket
12 from torch.fx._utils import lazy_format_graph_code
13
14 from vllm.compilation.passes.fx_utils import find_op_nodes
@@ -90,7 +90,9 @@ class TestBackend:
90 # assign by reference, will reflect the final state of the graph
91 self.final_graph = graph
92
93 - def check_before_ops(self, ops: Sequence[OpOverload], fully_replaced=True):
93 + def check_before_ops(
94 + self, ops: Sequence[OpOverload | OpOverloadPacket], fully_replaced=True
95 + ):
96 for op in ops:
97 num_pre = len(list(find_op_nodes(op, self.graph_pre_pass)))
98 num_post = len(list(find_op_nodes(op, self.graph_post_pass)))
@@ -99,13 +101,19 @@ class TestBackend:
101 if fully_replaced:
102 assert num_post == 0, f"Unexpected op {op.name()} in post-pass graph"
103
102 - def check_after_ops(self, ops: Sequence[OpOverload]):
104 + def check_after_ops(self, ops: Sequence[OpOverload | OpOverloadPacket]):
105 for op in ops:
106 num_pre = len(list(find_op_nodes(op, self.graph_pre_pass)))
107 num_post = len(list(find_op_nodes(op, self.graph_post_pass)))
108 assert num_pre == 0, f"Unexpected op {op.name()} in pre-pass graph"
109 assert num_post > 0, f"Op {op.name()} not found in post-pass graph"
110
109 - def op_count(self, op: OpOverload, before=False) -> int:
111 + def op_count(self, op: OpOverload | OpOverloadPacket, before=False) -> int:
112 graph = self.graph_pre_pass if before else self.graph_post_pass
113 return len(list(find_op_nodes(op, graph)))
114 +
115 + def print_graphs(self):
116 + print("=== Graph before custom passes ===")
117 + print(self.graph_pre_pass.python_code(root_module="self", verbose=True).src)
118 + print("=== Graph after custom passes ===")
119 + print(self.graph_post_pass.python_code(root_module="self", verbose=True).src)
tests/compile/fusions_e2e/test_tp1_quant.py
+3
@@ -99,6 +99,8 @@ def test_tp1_fp8_fusions(
99 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
100 model_kwargs["load_format"] = "dummy"
101 model_kwargs["max_model_len"] = 1024
102 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
103 +
104 compilation_config = dict(
105 use_inductor_graph_partition=inductor_graph_partition,
106 custom_ops=custom_ops.split(","),
@@ -166,6 +168,7 @@ def test_tp1_fp4_fusions(
168 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
169 model_kwargs["load_format"] = "dummy"
170 model_kwargs["max_model_len"] = 1024
171 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
172
173 compilation_config = dict(
174 use_inductor_graph_partition=inductor_graph_partition,
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+3
@@ -68,6 +68,7 @@ def test_tp2_ar_rms_fp8_fusions(
68 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
69 model_kwargs["load_format"] = "dummy"
70 model_kwargs["max_model_len"] = 1024
71 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
72
73 compilation_config = dict(
74 use_inductor_graph_partition=inductor_graph_partition,
@@ -128,6 +129,7 @@ def test_tp2_ar_rms_fp4_fusions(
129 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
130 model_kwargs["load_format"] = "dummy"
131 model_kwargs["max_model_len"] = 1024
132 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
133
134 compilation_config = dict(
135 use_inductor_graph_partition=inductor_graph_partition,
@@ -182,6 +184,7 @@ def test_tp2_ar_rms_fusions(
184 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
185 model_kwargs["load_format"] = "dummy"
186 model_kwargs["max_model_len"] = 1024
187 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
188
189 compilation_config = dict(
190 use_inductor_graph_partition=inductor_graph_partition,
tests/compile/fusions_e2e/test_tp2_async_tp.py
+2
@@ -58,6 +58,7 @@ def test_tp2_async_tp_fp8_fusions(
58 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
59 model_kwargs["load_format"] = "dummy"
60 model_kwargs["max_model_len"] = 1024
61 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
62
63 compilation_config = dict(
64 use_inductor_graph_partition=inductor_graph_partition,
@@ -121,6 +122,7 @@ def test_tp2_async_tp_fusions(
122 model_kwargs["hf_overrides"] = hf_overrides(n_layers)
123 model_kwargs["load_format"] = "dummy"
124 model_kwargs["max_model_len"] = 1024
125 + model_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}
126
127 compilation_config = dict(
128 use_inductor_graph_partition=inductor_graph_partition,
tests/compile/passes/distributed/test_sequence_parallelism.py
+7 -7
@@ -9,7 +9,6 @@ from tests.compile.backend import TestBackend
9 from tests.utils import TestFP8Layer, multi_gpu_test
10 from vllm.compilation.passes.fusion.rms_quant_fusion import RMSNormQuantFusionPass
11 from vllm.compilation.passes.fusion.sequence_parallelism import SequenceParallelismPass
12 -from vllm.compilation.passes.fx_utils import find_auto_fn
12 from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
13 from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
14 from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass
@@ -86,13 +85,14 @@ class TestAllReduceRMSNormModel(torch.nn.Module):
85 ]
86
87 def ops_in_model(self):
89 - if RMSNorm.enabled():
90 - return [
91 - torch.ops._C.rms_norm.default,
88 + return (
89 + [torch.ops.vllm_ir.rms_norm]
90 + + [
91 torch.ops._C.fused_add_rms_norm.default,
92 ]
94 - else:
95 - return []
93 + if RMSNorm.enabled()
94 + else []
95 + )
96
97
98 class TestAllReduceRMSNormStaticQuantFP8Model(torch.nn.Module):
@@ -321,4 +321,4 @@ def sequence_parallelism_pass_on_test_model(
321 assert backend.op_count(op, before=False) == 4
322
323 for op in model.ops_in_model():
324 - find_auto_fn(backend.graph_post_pass.nodes, op)
324 + assert backend.op_count(op, before=False) > 0
tests/compile/passes/ir/__init__.py
tests/compile/passes/ir/test_lowering.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +from torch import nn
6 +
7 +import vllm.kernels # noqa: F401 to register kernels
8 +from vllm import ir
9 +from vllm.compilation.passes.ir.lowering_pass import (
10 + VllmIRLoweringPass,
11 +)
12 +from vllm.config import get_current_vllm_config
13 +from vllm.ir import ops
14 +from vllm.platforms import current_platform
15 +
16 +from ...backend import TestBackend
17 +
18 +
19 +class Model(nn.Module):
20 + def __init__(self, hidden_size=16, *args, **kwargs):
21 + super().__init__(*args, **kwargs)
22 + self.hidden_size = hidden_size
23 + self.weight = torch.ones(hidden_size, dtype=torch.bfloat16)
24 +
25 + def forward(self, x):
26 + x1 = x + 4.0
27 + x2 = ops.rms_norm(x1, self.weight, 1e-5)
28 + x3 = x2 * 5.0
29 + # no weight
30 + x4 = ops.rms_norm(x3, None, 1e-5)
31 + x5 = x4 / 2.0
32 + # dispatch to native due to variance_size parameter
33 + x6 = ops.rms_norm(x5, self.weight, 1e-5, self.hidden_size // 2)
34 + return x6 + 3.0
35 +
36 +
37 +@pytest.mark.parametrize("rms_provider", ops.rms_norm.supported_providers())
38 +def test_lowering_rms_norm(rms_provider, default_vllm_config):
39 + torch.set_default_device(current_platform.device_type)
40 +
41 + lowering_pass = VllmIRLoweringPass(get_current_vllm_config())
42 + backend = TestBackend(lowering_pass)
43 + backend_unlowered = TestBackend()
44 +
45 + model = Model()
46 + x = torch.randn(8, 16, dtype=torch.bfloat16)
47 + with (
48 + ops.rms_norm.set_priority([rms_provider, "native"]),
49 + ir.enable_torch_wrap(True),
50 + ):
51 + compiled_model = torch.compile(model, backend=backend, fullgraph=True)
52 + compiled_unlowered_model = torch.compile(
53 + model, backend=backend_unlowered, fullgraph=True
54 + )
55 + output = compiled_model(x)
56 + output_unlowered = compiled_unlowered_model(x)
57 +
58 + selected = lowering_pass.selected_impls["rms_norm"]
59 + assert len(selected) == 3
60 + assert selected["rms_norm"] == rms_provider
61 + assert selected["rms_norm_1"] == rms_provider
62 + assert selected["rms_norm_2"] == "native"
63 +
64 + # Compiled function guards on global value, avoid recompilation
65 + with ir.enable_torch_wrap(True):
66 + output2 = compiled_model(x)
67 +
68 + torch.testing.assert_close(output_unlowered, output)
69 + torch.testing.assert_close(output_unlowered, output2)
tests/compile/passes/test_fusion.py
+10 -8
@@ -6,6 +6,7 @@ import pytest
6 import torch
7
8 import vllm.config
9 +import vllm.ir.ops
10 import vllm.plugins
11 from tests.compile.backend import TestBackend
12 from tests.utils import TestBlockFP8Layer, TestFP8Layer
@@ -51,7 +52,6 @@ from vllm.utils.deep_gemm import (
52
53 FP8_DTYPE = current_platform.fp8_dtype()
54
54 -RMS_OP = torch.ops._C.rms_norm.default
55 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
56
57 # Kernel and group_shape combinations: (kernel, group_shape)
@@ -246,10 +246,8 @@ class TestModel(torch.nn.Module):
246 ]
247
248 def ops_in_model_before_partial(self):
249 - return (
250 - [RMS_OP, RMS_ADD_OP]
251 - if self.enable_rms_norm_custom_op
252 - else [torch.ops.aten.rsqrt]
249 + return [torch.ops.vllm_ir.rms_norm] + (
250 + [RMS_ADD_OP] if self.enable_rms_norm_custom_op else [torch.ops.aten.rsqrt]
251 )
252
253
@@ -340,7 +338,10 @@ def test_fusion_rmsnorm_quant(
338 ),
339 )
340
343 - with vllm.config.set_current_vllm_config(vllm_config):
341 + with (
342 + vllm.config.set_current_vllm_config(vllm_config),
343 + vllm_config.kernel_config.ir_op_priority.set_priority(),
344 + ):
345 # Setup device before model creation
346 torch.set_default_device("cuda")
347 torch.set_default_dtype(dtype)
@@ -370,8 +371,9 @@ def test_fusion_rmsnorm_quant(
371 # Hence, we check only 2 add nodes are left (final fused rmsnorm add).
372 if not enable_rms_norm_custom_op:
373 n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))
373 - # 7 = 1 (RMS) + 3x2 (3xRMS_ADD, 2 each)
374 - assert n_add_nodes(backend.graph_pre_pass) == 7
374 + # rms_norm is IR, not included
375 + # 6 = 3x2 (3xRMS_ADD, 2 each)
376 + assert n_add_nodes(backend.graph_pre_pass) == 6
377 assert n_add_nodes(backend.graph_post_pass) == 2
378
379
tests/compile/passes/test_qk_norm_rope_fusion.py
+8 -10
@@ -3,11 +3,11 @@
3
4 import pytest
5 import torch
6 +from torch._ops import OpOverload, OpOverloadPacket
7
8 from tests.compile.backend import TestBackend
9 from vllm.compilation.passes.fusion.matcher_utils import (
10 FLASHINFER_ROTARY_OP,
10 - RMS_OP,
11 ROTARY_OP,
12 )
13 from vllm.compilation.passes.fusion.qk_norm_rope_fusion import (
@@ -100,13 +100,8 @@ class QKNormRoPETestModel(torch.nn.Module):
100 q, k = self.rotary_emb(positions, q, k)
101 return q, k, v
102
103 - def ops_in_model_before(self) -> list[torch._ops.OpOverload]:
104 - ops = []
105 - if self.enable_rms_norm_custom_op:
106 - ops.append(RMS_OP)
107 - else:
108 - ops.append(RSQRT_OP)
109 -
103 + def ops_in_model_before(self) -> list[OpOverload | OpOverloadPacket]:
104 + ops: list[OpOverload | OpOverloadPacket] = [torch.ops.vllm_ir.rms_norm]
105 if self.enable_rope_custom_op:
106 if self.rotary_emb.use_flashinfer:
107 ops.append(FLASHINFER_ROTARY_OP)
@@ -116,7 +111,7 @@ class QKNormRoPETestModel(torch.nn.Module):
111 ops.append(INDEX_SELECT_OP)
112 return ops
113
119 - def ops_in_model_after(self) -> list[torch._ops.OpOverload]:
114 + def ops_in_model_after(self) -> list[OpOverload | OpOverloadPacket]:
115 return [FUSED_QK_ROPE_OP]
116
117
@@ -166,7 +161,10 @@ def test_qk_norm_rope_fusion(
161 num_heads, num_kv_heads, head_dim = 16, 4, 128
162 T = 5
163
169 - with set_current_vllm_config(vllm_config):
164 + with (
165 + set_current_vllm_config(vllm_config),
166 + vllm_config.kernel_config.ir_op_priority.set_priority(),
167 + ):
168 model = QKNormRoPETestModel(
169 num_heads=num_heads,
170 num_kv_heads=num_kv_heads,
tests/conftest.py
+23
@@ -1622,3 +1622,26 @@ def fresh_vllm_cache(monkeypatch, use_fresh_inductor_cache):
1622 def enable_pickle(monkeypatch):
1623 """`LLM.apply_model` requires pickling a function."""
1624 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
1625 +
1626 +
1627 +@pytest.fixture(scope="function")
1628 +def disable_log_dedup(monkeypatch):
1629 + """
1630 + Disable log deduplication such that warning_once and info_once always print.
1631 + """
1632 +
1633 + # Patch logger._print_warning_once to remove the lru_cache decorator
1634 + from vllm import logger
1635 +
1636 + original_print_warning_once = logger._print_warning_once
1637 + original_print_info_once = logger._print_info_once
1638 + original_print_debug_once = logger._print_debug_once
1639 +
1640 + logger._print_warning_once = original_print_warning_once.__wrapped__
1641 + logger._print_info_once = original_print_info_once.__wrapped__
1642 + logger._print_debug_once = original_print_debug_once.__wrapped__
1643 +
1644 + yield
1645 + logger._print_warning_once = original_print_warning_once
1646 + logger._print_info_once = original_print_info_once
1647 + logger._print_debug_once = original_print_debug_once
tests/engine/test_arg_utils.py
+17
@@ -523,3 +523,20 @@ def test_human_readable_model_len():
523 for invalid in ["1a", "pwd", "10.24", "1.23M", "1.22T"]:
524 with pytest.raises(ArgumentError):
525 parser.parse_args(["--max-model-len", invalid])
526 +
527 +
528 +def test_ir_op_priority():
529 + from vllm.config.kernel import IrOpPriorityConfig, KernelConfig
530 +
531 + ir_op_priority = IrOpPriorityConfig(rms_norm=["vllm_c"])
532 + cfg1 = EngineArgs(ir_op_priority=ir_op_priority).create_engine_config()
533 + cfg2 = EngineArgs(
534 + kernel_config=KernelConfig(ir_op_priority=ir_op_priority)
535 + ).create_engine_config()
536 + assert cfg1.kernel_config.ir_op_priority == cfg2.kernel_config.ir_op_priority
537 +
538 + with pytest.raises(ValueError, match="rms_norm"):
539 + _ = EngineArgs(
540 + ir_op_priority=ir_op_priority,
541 + kernel_config=KernelConfig(ir_op_priority=ir_op_priority),
542 + ).create_engine_config()
tests/ir/test_op.py new
+497
@@ -0,0 +1,497 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import importlib.util
4 +import logging
5 +from pathlib import Path
6 +from typing import Any
7 +
8 +import pytest
9 +import torch
10 +from torch import fx
11 +from torch.fx.experimental.proxy_tensor import make_fx
12 +
13 +import vllm.ir.op
14 +from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl
15 +
16 +# This should not exist
17 +assert "_custom_add" not in IrOp.registry
18 +
19 +
20 +class CustomError(Exception):
21 + pass
22 +
23 +
24 +@vllm.ir.register_op
25 +def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
26 + return x + y
27 +
28 +
29 +def test_registration_overloads():
30 + assert all(
31 + n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"]
32 + )
33 +
34 + # Calling with decorator
35 + @vllm.ir.register_op()
36 + def _custom_sub(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
37 + return x - y
38 +
39 + assert _custom_sub.name == "_custom_sub"
40 + assert _custom_sub is IrOp.registry["_custom_sub"]
41 +
42 + # Custom name
43 + @vllm.ir.register_op(name="_custom_mul")
44 + def custom_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
45 + return x * y
46 +
47 + assert custom_mul.name == "_custom_mul"
48 + assert custom_mul is IrOp.registry["_custom_mul"]
49 +
50 + # Direct construction does not register directly
51 + def _custom_div(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
52 + return x / y
53 +
54 + custom_div = IrOp("_custom_div", _custom_div)
55 + assert custom_div.name == "_custom_div"
56 + assert "_custom_div" not in IrOp.registry
57 +
58 + # Duplicate op registration not allowed
59 + with pytest.raises(AssertionError):
60 +
61 + @vllm.ir.register_op
62 + def _custom_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
63 + return x * y - 100
64 +
65 +
66 +def test_no_kw_only_args():
67 + # kw-only args not supported
68 + with pytest.raises(ValueError, match="keyword-only arguments"):
69 +
70 + @vllm.ir.register_op
71 + def _custom_kwarg_op(
72 + x: torch.Tensor, y: torch.Tensor, *, kwarg: int = 0
73 + ) -> torch.Tensor:
74 + return x + y + kwarg
75 +
76 + assert "_custom_kwarg_op" not in IrOp.registry
77 +
78 +
79 +class TestIrOpCustomAdd:
80 + # Registration invariants
81 + def test_decorated_object(self):
82 + """Make sure that referring directly to an op is correct"""
83 + assert isinstance(_custom_add, IrOp)
84 + assert "_custom_add" in IrOp.registry
85 + assert _custom_add is IrOp.registry["_custom_add"]
86 +
87 + def test_torch_op_is_registered(self):
88 + assert hasattr(torch.ops.vllm_ir, "_custom_add")
89 + assert callable(torch.ops.vllm_ir._custom_add.default)
90 +
91 + # Semantic correctness
92 + def test_semantics_match_native(self):
93 + x = torch.randn(4, 5)
94 + y = torch.randn(4, 5)
95 +
96 + # Calls native by default
97 + out = _custom_add(x, y)
98 + ref = x + y
99 +
100 + torch.testing.assert_close(out, ref)
101 +
102 + # -------------------------
103 + # Implementation registration
104 + # -------------------------
105 +
106 + def test_register_impl_is_non_intrusive(self):
107 + @_custom_add.register_impl("dummy_provider")
108 + def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
109 + return x + y + 123
110 +
111 + assert "dummy_provider" in _custom_add.impls
112 + assert isinstance(_custom_add.impls["dummy_provider"], IrOpImpl)
113 +
114 + x = torch.ones(2, 2)
115 + y = torch.ones(2, 2)
116 +
117 + # Native semantics must still hold
118 + torch.testing.assert_close(_custom_add(x, y), x + y)
119 +
120 + def test_schema_contains_tensor_signature(self):
121 + schema = _custom_add._schema_str
122 +
123 + assert "Tensor" in schema
124 + assert "-> Tensor" in schema
125 +
126 + # -------------------------
127 + # FX visibility
128 + # -------------------------
129 +
130 + @pytest.mark.parametrize("enable_torch_wrap", [True, False])
131 + @pytest.mark.parametrize("symbolic_trace", [True, False])
132 + def test_trace_sees_single_custom_op(
133 + self, symbolic_trace: bool, enable_torch_wrap: bool
134 + ):
135 + def fn(x, y):
136 + return _custom_add(x, y)
137 +
138 + def find_fn(target: Any, gm: fx.GraphModule):
139 + return gm.graph.find_nodes(op="call_function", target=target)
140 +
141 + with pytest.raises(CustomError), vllm.ir.enable_torch_wrap(enable_torch_wrap):
142 + if symbolic_trace:
143 + gm = torch.fx.symbolic_trace(fn)
144 + else:
145 + gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2))
146 +
147 + x1, y1 = torch.rand(5, 4), torch.rand(5, 4)
148 + out_fx = gm(x1, y1)
149 + out_eager = fn(x1, y1)
150 +
151 + # raise error to check enable_torch_wrap context restored correctly
152 + raise CustomError
153 +
154 + # check behavior matches eager in all cases
155 + torch.testing.assert_close(out_fx, out_eager)
156 +
157 + # check that IR nodes only appear if enable_torch_wrap=True
158 + ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm)
159 + if enable_torch_wrap:
160 + assert len(ir_nodes) == 1, gm.code
161 + else:
162 + assert len(ir_nodes) == 0, gm.code
163 +
164 + # with torch wrapping enabled (default), IR nodes appear
165 + if symbolic_trace:
166 + gm = torch.fx.symbolic_trace(fn)
167 + else:
168 + gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2))
169 +
170 + ir_nodes = find_fn(torch.ops.vllm_ir._custom_add.default, gm)
171 + assert len(ir_nodes) == 1, gm.code
172 +
173 +
174 +@_custom_add.register_impl("impl_a")
175 +def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
176 + return x + y + 10
177 +
178 +
179 +@_custom_add.register_impl("impl_b")
180 +def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
181 + return x + y + 20
182 +
183 +
184 +@_custom_add.register_impl("impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0)
185 +def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
186 + return x + y + 50
187 +
188 +
189 +class TestIrOpImplDispatch:
190 + def test_register_impl(self):
191 + assert "impl_a" in _custom_add.impls
192 + impl = _custom_add.impls["impl_a"]
193 +
194 + assert impl is impl_a
195 + assert impl.op is _custom_add
196 + assert impl.provider == "impl_a"
197 + assert callable(impl.impl_fn)
198 +
199 + # Test duplicate registration rejected
200 + with pytest.raises(AssertionError):
201 +
202 + @_custom_add.register_impl("impl_a")
203 + def impl_a_dup(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
204 + return x + y + 30
205 +
206 + # Check the original impl is still intact
207 + assert _custom_add.impls["impl_a"] is impl_a
208 +
209 + # Check support all args
210 + assert impl_a.supports_all_args
211 + assert impl_b.supports_all_args
212 + assert not impl_even.supports_all_args
213 +
214 + def test_reserved_provider_rejected(self):
215 + for provider in RESERVED_PROVIDERS:
216 + with pytest.raises(AssertionError):
217 +
218 + @_custom_add.register_impl(provider)
219 + def bad_impl(x, y):
220 + return x + y
221 +
222 + def test_set_priority_scoped(self):
223 + assert _custom_add.get_priority() == []
224 +
225 + with _custom_add.set_priority(["impl_even", "impl_b"]):
226 + assert _custom_add.get_priority() == ["impl_even", "impl_b"]
227 +
228 + # Check nesting
229 + with _custom_add.set_priority(["impl_b"]):
230 + assert _custom_add.get_priority() == ["impl_b"]
231 +
232 + # Restored
233 + assert _custom_add.get_priority() == ["impl_even", "impl_b"]
234 +
235 + # Check that exception restores priority
236 + with pytest.raises(CustomError), _custom_add.set_priority(["impl_a"]):
237 + assert _custom_add.get_priority() == ["impl_a"]
238 + raise CustomError
239 +
240 + # Restored again
241 + assert _custom_add.get_priority() == ["impl_even", "impl_b"]
242 +
243 + # Restored to empty
244 + assert _custom_add.get_priority() == []
245 +
246 + def test_dispatch_priority_order(self):
247 + x = torch.tensor(1, dtype=torch.int32)
248 + y = torch.tensor(2, dtype=torch.int32)
249 +
250 + with _custom_add.set_priority(["impl_b", "impl_a"]):
251 + assert _custom_add.dispatch(x, y) is impl_b
252 + out1 = _custom_add(x, y)
253 + out2 = torch.ops.vllm_ir._custom_add(x, y)
254 +
255 + with _custom_add.set_priority(["impl_a"]):
256 + assert _custom_add.dispatch(x, y) is impl_a
257 + out3 = _custom_add(x, y)
258 + out4 = torch.ops.vllm_ir._custom_add(x, y)
259 +
260 + # impl_b
261 + assert out1.item() == 1 + 2 + 20
262 + assert out2.item() == 1 + 2 + 20
263 + # impl_a
264 + assert out3.item() == 1 + 2 + 10
265 + assert out4.item() == 1 + 2 + 10
266 +
267 + def test_unsupported_impl_filtered(self):
268 + @_custom_add.register_impl("unsupported", supported=False)
269 + def impl_bad(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
270 + return x + y + 999
271 +
272 + x = torch.tensor(1, dtype=torch.int32)
273 + y = torch.tensor(2, dtype=torch.int32)
274 +
275 + with _custom_add.set_priority(["unsupported", "impl_a"]):
276 + assert _custom_add.get_priority() == ["impl_a"]
277 + out = _custom_add(x, y)
278 +
279 + # impl_bad skipped → impl_a
280 + assert out.item() == 1 + 2 + 10
281 +
282 + def test_supports_args_runtime_dispatch_and_warning(
283 + self, caplog_vllm: pytest.LogCaptureFixture
284 + ):
285 + x1 = torch.ones((2, 2), dtype=torch.int32)
286 + y1 = torch.full((2, 2), 2, dtype=torch.int32)
287 +
288 + x2 = torch.ones((2, 3), dtype=torch.int32)
289 + y2 = torch.full((2, 3), 2, dtype=torch.int32)
290 +
291 + with (
292 + caplog_vllm.at_level(logging.WARNING),
293 + _custom_add.set_priority(["impl_even"]),
294 + ):
295 + # Test the warning about native fallback is logged (before even dispatching)
296 + assert len(caplog_vllm.records) == 1
297 + message = caplog_vllm.records[0].message
298 + assert "_custom_add" in message
299 + assert "fallback to native" in message
300 + assert "priority" in message
301 +
302 + # Check dispatching
303 + assert _custom_add.get_priority() == ["impl_even", "native"]
304 + assert _custom_add.dispatch(x1, y1) is impl_even
305 + assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"]
306 +
307 + out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even
308 + out2 = _custom_add(x2, y2) # size(1) == 3 → native fallback
309 +
310 + # no other warnings
311 + assert len(caplog_vllm.records) == 1
312 + assert torch.all(out1 == 1 + 2 + 50)
313 + assert torch.all(out2 == 1 + 2)
314 +
315 + def test_default_priority(
316 + self, caplog_vllm: pytest.LogCaptureFixture, disable_log_dedup
317 + ):
318 + # Make sure logs are not deduplicated to properly test the warning
319 + x = torch.tensor([3], dtype=torch.int32)
320 + y = torch.tensor([4], dtype=torch.int32)
321 +
322 + # No priority set → falls back to native
323 + assert _custom_add.get_priority() == []
324 + with caplog_vllm.at_level(logging.WARNING):
325 + # Native by default
326 + assert _custom_add.dispatch(x, y) is _custom_add.impls["native"]
327 + out = _custom_add(x, y)
328 +
329 + # Check dispatching to native by default
330 + assert out.item() == 3 + 4
331 +
332 + # Check warning
333 + assert len(caplog_vllm.records) == 2
334 + message = caplog_vllm.records[0].message.lower()
335 + assert "_custom_add" in message
336 + assert "priority not set" in message
337 +
338 +
339 +@vllm.ir.register_op
340 +def _custom_mm(
341 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
342 +) -> torch.Tensor:
343 + tmp = x @ y
344 + return tmp if bias is None else tmp + bias
345 +
346 +
347 +def test_default_args():
348 + # Test that default args are properly applied when dispatching and calling
349 + @_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)
350 + def impl_mm(
351 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
352 + ) -> torch.Tensor:
353 + tmp = x @ y
354 + return tmp + 50 if bias is None else tmp + bias + 100
355 +
356 + x1 = torch.tensor([1, 2], dtype=torch.int32)
357 + x2 = torch.tensor([3, 4], dtype=torch.int32)
358 +
359 + # Test that supports_args receives the defaulted args
360 + assert impl_mm.supports_args(x1, x2)
361 + with _custom_mm.set_priority(["impl_mm", "native"]):
362 + assert _custom_mm.dispatch(x1, x2) is impl_mm
363 +
364 +
365 +def test_bad_impl_registrations():
366 + # Check bad schema
367 + with pytest.raises(ValueError, match="does not match native schema"):
368 +
369 + @_custom_mm.register_impl("impl_mm_bad_schema")
370 + def impl_mm_bad_schema(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
371 + return x @ y - 1
372 +
373 + with pytest.raises(ValueError, match="does not match native schema"):
374 +
375 + @_custom_mm.register_impl("impl_mm_bad_schema_2")
376 + def impl_mm_bad_schema_2(
377 + x: torch.Tensor, y: torch.Tensor, b: torch.Tensor | None = None
378 + ) -> torch.Tensor:
379 + return x @ y + b - 2
380 +
381 + with pytest.raises(ValueError, match="does not match native schema"):
382 +
383 + @_custom_mm.register_impl("impl_mm_bad_schema_3")
384 + def impl_mm_bad_schema_3(
385 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor
386 + ) -> torch.Tensor:
387 + return x @ y + bias - 5
388 +
389 + # check supports_args with incorrect params
390 + with pytest.raises(ValueError, match="supports_args must be a callable"):
391 +
392 + @_custom_mm.register_impl("impl_mm_bad_supports_args", supports_args=True)
393 + def impl_mm_bad_supports_args(
394 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
395 + ) -> torch.Tensor:
396 + return x @ y + 10
397 +
398 + with pytest.raises(ValueError, match="number of parameters"):
399 +
400 + @_custom_mm.register_impl(
401 + "impl_mm_bad_supports_args_2", supports_args=lambda x, y: True
402 + )
403 + def impl_mm_bad_supports_args(
404 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
405 + ) -> torch.Tensor:
406 + return x @ y + 10
407 +
408 + with pytest.raises(ValueError, match="keyword-only parameters"):
409 +
410 + @_custom_mm.register_impl(
411 + "impl_mm_bad_supports_args_3", supports_args=lambda x, y, *, b: True
412 + )
413 + def impl_mm_bad_supports_args_2(
414 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
415 + ) -> torch.Tensor:
416 + return x @ y + 20
417 +
418 + with pytest.raises(ValueError, match="does not match native parameter"):
419 +
420 + @_custom_mm.register_impl(
421 + "impl_mm_bad_supports_args_4", supports_args=lambda x, y, b: True
422 + )
423 + def impl_mm_bad_supports_args_4(
424 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
425 + ) -> torch.Tensor:
426 + return x @ y + 30
427 +
428 + with pytest.raises(ValueError, match="does not match native default"):
429 +
430 + @_custom_mm.register_impl(
431 + "impl_mm_bad_supports_args_5", supports_args=lambda x, y, bias=1: True
432 + )
433 + def impl_mm_bad_supports_args_5(
434 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
435 + ) -> torch.Tensor:
436 + return x @ y + 40
437 +
438 + assert set(_custom_mm.impls.keys()) == {"impl_mm", "native"}
439 +
440 +
441 +IMPL_OOT_SRC = """
442 +import torch
443 +
444 +@_custom_mm.register_impl("impl_mm_oot")
445 +def impl_mm_oot(
446 + x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
447 +) -> torch.Tensor:
448 + return x @ y - 99
449 +"""
450 +
451 +
452 +def load_custom_mm_module(file_path: Path):
453 + spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path)
454 + assert spec is not None
455 + module = importlib.util.module_from_spec(spec)
456 +
457 + # Inject the variable into the module's global namespace
458 + # This allows the @_custom_mm.register_impl decorator to work
459 + module._custom_mm = _custom_mm # type: ignore[attr-defined]
460 +
461 + # Execute the file; this triggers the decorator
462 + assert spec.loader is not None
463 + spec.loader.exec_module(module)
464 + return module
465 +
466 +
467 +def test_uuid_and_oot(tmp_path: Path):
468 + file_path = tmp_path / "_custom_mm_oot.py"
469 + file_path.write_text(IMPL_OOT_SRC)
470 +
471 + assert "impl_mm_oot" not in _custom_mm.impls
472 + _ = load_custom_mm_module(file_path)
473 + assert "impl_mm_oot" in _custom_mm.impls
474 +
475 + uuid = _custom_mm.impls["impl_mm_oot"].uuid()
476 + del _custom_mm.impls["impl_mm_oot"]
477 +
478 + # Replace file source
479 + file_path.write_text(IMPL_OOT_SRC + " # added file source")
480 + assert "impl_mm_oot" not in _custom_mm.impls
481 + _ = load_custom_mm_module(file_path)
482 + assert "impl_mm_oot" in _custom_mm.impls
483 +
484 + uuid1 = _custom_mm.impls["impl_mm_oot"].uuid()
485 + assert uuid1 != uuid
486 + del _custom_mm.impls["impl_mm_oot"]
487 +
488 + # Back to original
489 + file_path.write_text(IMPL_OOT_SRC)
490 + assert "impl_mm_oot" not in _custom_mm.impls
491 + _ = load_custom_mm_module(file_path)
492 + assert "impl_mm_oot" in _custom_mm.impls
493 +
494 + uuid2 = _custom_mm.impls["impl_mm_oot"].uuid()
495 + assert uuid2 == uuid
496 + assert uuid2 != uuid1
497 + del _custom_mm.impls["impl_mm_oot"]
tests/kernels/ir/test_layernorm.py new
+129
@@ -0,0 +1,129 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +import torch
5 +
6 +# This registers op implementations
7 +import vllm.kernels # noqa: F401
8 +from tests.kernels.allclose_default import get_default_rtol
9 +from vllm import ir
10 +from vllm.platforms import current_platform
11 +
12 +
13 +def rms_norm_inputs(n_tokens: int, hidden_size: int, dtype: torch.dtype):
14 + x = torch.randn(n_tokens, hidden_size, dtype=dtype)
15 + weight = torch.rand(hidden_size, dtype=dtype)
16 + return x, weight
17 +
18 +
19 +rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn
20 +
21 +
22 +@pytest.mark.skipif(
23 + not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
24 + reason="Currently only kernels on CUDA, ROCm and XPU",
25 +)
26 +def test_rms_norm_registration():
27 + expected = {
28 + "native": True,
29 + "vllm_c": current_platform.is_cuda_alike(),
30 + "aiter": current_platform.is_rocm(),
31 + "oink": False,
32 + "xpu_kernels": current_platform.is_xpu(),
33 + }
34 +
35 + actual = {
36 + provider: impl.supported for provider, impl in ir.ops.rms_norm.impls.items()
37 + }
38 +
39 + assert actual == expected
40 +
41 +
42 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
43 +@pytest.mark.parametrize("n_tokens", [1, 8, 17])
44 +@pytest.mark.parametrize("hidden_size", [16, 4096, 8192])
45 +@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
46 +@pytest.mark.skipif(
47 + not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
48 + reason="Currently only kernels on CUDA, ROCm and XPU",
49 +)
50 +class TestRMSNorm:
51 + @classmethod
52 + def setup_class(cls, **kwargs):
53 + torch.set_default_device(current_platform.device_type)
54 +
55 + def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
56 + x, weight = rms_norm_inputs(4, 8, dtype)
57 + out = rms_norm_native(x, weight, epsilon=epsilon)
58 +
59 + # Check shape, dtype, device
60 + assert out.shape == x.shape
61 + assert out.dtype == x.dtype
62 + assert out.device == x.device
63 +
64 + # Check the scaling property of rms norm
65 + out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon)
66 + torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3)
67 +
68 + # Check behavior with and without weight
69 + weight1 = torch.ones_like(weight)
70 + out3 = rms_norm_native(x, weight1, epsilon=epsilon)
71 + out4 = rms_norm_native(x, None, epsilon=epsilon)
72 + torch.testing.assert_close(out3, out4)
73 +
74 + @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels"])
75 + def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
76 + impl = ir.ops.rms_norm.impls[provider]
77 + if not impl.supported:
78 + pytest.skip(f"{provider} impl not supported on this platform")
79 +
80 + x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
81 + args = (x, weight, epsilon, None)
82 +
83 + assert impl.supported
84 +
85 + if provider == "aiter" and dtype not in [torch.float16, torch.bfloat16]:
86 + assert not impl.supports_args(*args)
87 + return
88 +
89 + assert impl.supports_args(*args)
90 +
91 + out_impl = impl.impl_fn(*args)
92 + out_native = rms_norm_native(*args)
93 +
94 + torch.testing.assert_close(
95 + out_impl, out_native, rtol=get_default_rtol(out_impl), atol=1e-3
96 + )
97 +
98 + # check that dispatched call matches direct call
99 + with ir.ops.rms_norm.set_priority([provider, "native"]):
100 + out_impl2 = ir.ops.rms_norm(*args)
101 +
102 + # exact match
103 + torch.testing.assert_close(out_impl2, out_impl, rtol=0.0, atol=0.0)
104 +
105 + # none of these support variance_size override
106 + assert not impl.supports_args(x, weight, epsilon, 4)
107 + assert not impl.supports_args(x, weight, epsilon, variance_size=4)
108 +
109 + # test weight=None behavior
110 + out_impl_no_weight = impl.impl_fn(x, None, epsilon)
111 + out_impl_unit_weight = impl.impl_fn(x, torch.ones_like(weight), epsilon)
112 + torch.testing.assert_close(
113 + out_impl_no_weight,
114 + out_impl_unit_weight,
115 + rtol=get_default_rtol(out_impl_no_weight),
116 + atol=2e-4,
117 + )
118 +
119 + @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"])
120 + def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
121 + if not ir.ops.rms_norm.impls[provider].supported:
122 + pytest.skip(f"{provider} impl not supported on this platform")
123 +
124 + x, weight = rms_norm_inputs(n_tokens, hidden_size, dtype)
125 + args = (x, weight, epsilon, None)
126 +
127 + # When checking the torch op, we have to set priority and use dispatch
128 + with ir.ops.rms_norm.set_priority([provider, "native"]):
129 + torch.library.opcheck(torch.ops.vllm_ir.rms_norm, args)
tests/model_executor/test_enabled_custom_ops.py
+4 -9
@@ -27,7 +27,6 @@ from vllm.model_executor.layers.layernorm import (
27 RMSNorm,
28 dispatch_rocm_rmsnorm_func,
29 fused_add_rms_norm,
30 - rms_norm,
30 )
31 from vllm.platforms import current_platform
32
@@ -156,7 +155,7 @@ def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
155 assert topk_func == vllm_topk_sigmoid
156
157
159 -@pytest.mark.parametrize("add_residual", [True, False])
158 +@pytest.mark.parametrize("add_residual", [False])
159 @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
160 @pytest.mark.parametrize("use_rocm_aiter", [True, False])
161 @pytest.mark.skipif(
@@ -165,7 +164,7 @@ def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
164 def test_rms_norm_dispatch(
165 add_residual: bool, dtype: torch.dtype, use_rocm_aiter: bool
166 ):
168 - rms_norm_func = dispatch_rocm_rmsnorm_func(add_residual, dtype, use_rocm_aiter)
167 + rms_norm_func = dispatch_rocm_rmsnorm_func(dtype, use_rocm_aiter)
168
169 should_use_rocm_aiter = (
170 current_platform.is_rocm()
@@ -173,11 +172,7 @@ def test_rms_norm_dispatch(
172 and dtype in RMS_NORM_SUPPORTED_DTYPES
173 )
174
176 - if add_residual and should_use_rocm_aiter:
175 + if should_use_rocm_aiter:
176 assert rms_norm_func == rocm_aiter_ops.rms_norm2d_with_add
178 - elif should_use_rocm_aiter:
179 - assert rms_norm_func == rocm_aiter_ops.rms_norm
180 - elif add_residual:
181 - assert rms_norm_func == fused_add_rms_norm
177 else:
183 - assert rms_norm_func == rms_norm
178 + assert rms_norm_func == fused_add_rms_norm
tests/test_config.py
+68
@@ -6,12 +6,14 @@ import os
6 from dataclasses import MISSING, Field, asdict, dataclass, field
7 from unittest.mock import patch
8
9 +import pydantic
10 import pytest
11 from pydantic import ValidationError
12
13 from vllm.compilation.backends import VllmBackend
14 from vllm.config import (
15 CompilationConfig,
16 + KernelConfig,
17 ModelConfig,
18 ParallelConfig,
19 PoolerConfig,
@@ -21,6 +23,7 @@ from vllm.config import (
23 update_config,
24 )
25 from vllm.config.compilation import CompilationMode, CUDAGraphMode
26 +from vllm.config.kernel import IrOpPriorityConfig
27 from vllm.config.load import LoadConfig
28 from vllm.config.utils import get_field
29 from vllm.config.vllm import (
@@ -1077,6 +1080,39 @@ def test_vllm_config_explicit_overrides():
1080 assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
1081
1082
1083 +def test_fusion_pass_op_priority():
1084 + """This test checks that custom op enablement & IR op priority
1085 + correctly control default fusions"""
1086 +
1087 + # Default config, O2, rms_norm+quant fusion disabled
1088 + cfg1 = VllmConfig()
1089 + assert not cfg1.compilation_config.pass_config.fuse_norm_quant
1090 +
1091 + # rms_norm manually enabled, O1, rms_norm+quant fusion enabled
1092 + cfg2 = VllmConfig(
1093 + optimization_level=OptimizationLevel.O1,
1094 + compilation_config=CompilationConfig(
1095 + custom_ops=["+rms_norm"],
1096 + ),
1097 + )
1098 + assert cfg2.compilation_config.pass_config.fuse_norm_quant
1099 +
1100 + # using custom kernel for RMSNorm via IR:
1101 + # Note that vLLM IR only supports the non-residual rms_norm for now;
1102 + # soon this will be resolved.
1103 + cfg3 = VllmConfig(
1104 + kernel_config=KernelConfig(
1105 + ir_op_priority=IrOpPriorityConfig(rms_norm=["vllm_c"])
1106 + )
1107 + )
1108 + assert cfg3.compilation_config.pass_config.fuse_norm_quant
1109 +
1110 + # block-fp8 model should enable quant_fp8 automatically
1111 + cfg4 = VllmConfig(model_config=ModelConfig("Qwen/Qwen3-4B-FP8"))
1112 + assert "+quant_fp8" in cfg4.compilation_config.custom_ops
1113 + assert cfg4.compilation_config.pass_config.fuse_norm_quant
1114 +
1115 +
1116 def test_scheduler_config_init():
1117 with pytest.raises(ValidationError):
1118 # Positional InitVars missing
@@ -1171,3 +1207,35 @@ def test_eagle_draft_model_config():
1207 assert draft_model_config.hf_text_config.model_type == "eagle"
1208 assert draft_model_config.architectures == ["EagleLlamaForCausalLM"]
1209 assert draft_model_config.architecture == "EagleLlamaForCausalLM"
1210 +
1211 +
1212 +def test_ir_op_priority_default():
1213 + """Test that IR op priority defaults are set correctly."""
1214 + from vllm.config.kernel import IrOpPriorityConfig
1215 +
1216 + # Assert default is applied to ops
1217 + priority_config = IrOpPriorityConfig.with_default(["vllm_c", "native"])
1218 + assert priority_config.rms_norm == ["vllm_c", "native"]
1219 +
1220 + # Assert single ops override the default
1221 + assert IrOpPriorityConfig.with_default(
1222 + ["vllm_c", "native"], rms_norm=["oink", "native"]
1223 + ) == IrOpPriorityConfig(rms_norm=["oink", "native"])
1224 +
1225 +
1226 +def test_ir_op_priority_str():
1227 + """Test that passing a comma-delimited string works"""
1228 + from vllm.config.kernel import IrOpPriorityConfig
1229 +
1230 + priority_config = IrOpPriorityConfig(rms_norm="vllm_c")
1231 + assert priority_config.rms_norm == ["vllm_c"]
1232 +
1233 + priority_config = IrOpPriorityConfig(rms_norm="vllm_c,native")
1234 + assert priority_config.rms_norm == ["vllm_c", "native"]
1235 +
1236 + priority_config = IrOpPriorityConfig(rms_norm=" native, vllm_c ")
1237 + assert priority_config.rms_norm == ["native", "vllm_c"]
1238 +
1239 + with pytest.raises(pydantic.ValidationError):
1240 + # must be list of only strings
1241 + priority_config = IrOpPriorityConfig(rms_norm=["vllm_c", 4, "native"])
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+16 -14
@@ -3,6 +3,7 @@
3 import contextlib
4 from importlib.util import find_spec
5 from types import ModuleType
6 +from typing import Any
7
8 import torch
9 import torch._inductor.pattern_matcher as pm
@@ -10,6 +11,7 @@ import torch.fx as fx
11 from torch._higher_order_ops.auto_functionalize import auto_functionalized
12 from torch._inductor.pattern_matcher import PatternMatcherPass
13
14 +import vllm.ir.ops
15 from vllm.config import VllmConfig
16 from vllm.config.utils import Range
17 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
@@ -28,7 +30,7 @@ from vllm.utils.torch_utils import (
30
31 from ..inductor_pass import enable_fake_mode
32 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
31 -from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
33 +from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8
34
35 FP8_DTYPE = current_platform.fp8_dtype()
36
@@ -258,6 +260,12 @@ class BasePattern:
260 self.tp = get_tp_group()
261 self.tp_size = get_tensor_model_parallel_world_size()
262
263 + def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
264 + return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs)
265 +
266 + def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor:
267 + return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs)
268 +
269
270 class AllReduceRMSNormPattern(BasePattern):
271 """
@@ -276,20 +284,17 @@ class AllReduceRMSNormPattern(BasePattern):
284 super().__init__(dtype, device)
285 self.epsilon = epsilon
286 self.allreduce_params = allreduce_params
279 - self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
287
288 def get_inputs(self) -> list[torch.Tensor]:
282 - input, weight = self.rmsnorm_matcher.inputs()
283 -
284 - # input goes through allreduce first, always 16-bit
285 - return [input.to(self.dtype), weight]
289 + # input, weight
290 + return [self.empty(5, 16), self.empty(16)]
291
292 def register(self, pm_pass: PatternMatcherPass) -> None:
293 def pattern(
294 input: torch.Tensor, weight: torch.Tensor
295 ) -> tuple[torch.Tensor, torch.Tensor]:
296 allreduce_output = tensor_model_parallel_all_reduce(input)
292 - rms = self.rmsnorm_matcher(allreduce_output, weight)
297 + rms = vllm.ir.ops.rms_norm(allreduce_output, weight, self.epsilon)
298
299 return rms, allreduce_output
300
@@ -407,15 +412,13 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
412 self.epsilon = epsilon
413 self.allreduce_params = allreduce_params
414 self.quant_dtype = torch.float8_e4m3fn
410 - self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
415 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
416
417 def get_inputs(self) -> list[torch.Tensor]:
414 - input, weight = self.rmsnorm_matcher.inputs()
418 _, scale = self.quant_matcher.inputs()
419
417 - # input goes through allreduce first, always 16-bit
418 - return [input.to(self.dtype), weight, scale]
420 + # input, weight
421 + return [self.empty(5, 16), self.empty(16), scale]
422
423 def register(self, pm_pass: PatternMatcherPass) -> None:
424 def pattern(
@@ -424,7 +427,7 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):
427 scale: torch.Tensor,
428 ) -> tuple[torch.Tensor, torch.Tensor]:
429 all_reduce = tensor_model_parallel_all_reduce(input)
427 - rms = self.rmsnorm_matcher(all_reduce, weight)
430 + rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon)
431 quant, _ = self.quant_matcher(rms, scale)
432 return quant, all_reduce
433
@@ -553,7 +556,6 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
556 super().__init__(dtype, device)
557 self.epsilon = epsilon
558 self.allreduce_params = allreduce_params
556 - self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
559
560 def get_inputs(self) -> list[torch.Tensor]:
561 input = torch.empty([1, 16, 16], device=self.device, dtype=self.dtype)
@@ -575,7 +577,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
577 output_scale: torch.Tensor,
578 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
579 all_reduce = tensor_model_parallel_all_reduce(input)
578 - rms = self.rmsnorm_matcher(all_reduce, weight)
580 + rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon)
581 quant_out_tuple = auto_functionalized(
582 STATIC_FP4_QUANT_OP,
583 input=rms,
vllm/compilation/passes/fusion/matcher_utils.py
-64
@@ -26,7 +26,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
26 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
27 from vllm.platforms import current_platform
28
29 -RMS_OP = torch.ops._C.rms_norm.default
29 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
30 ROTARY_OP = torch.ops._C.rotary_embedding.default
31 FLASHINFER_ROTARY_OP = torch.ops.vllm.flashinfer_rotary_embedding.default
@@ -160,69 +159,6 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
159 return result
160
161
163 -class MatcherRMSNorm(MatcherCustomOp):
164 - def __init__(
165 - self,
166 - epsilon: float,
167 - enabled: bool | None = None,
168 - match_rocm_aiter: bool = False,
169 - ) -> None:
170 - if enabled is None:
171 - enabled = RMSNorm.enabled()
172 -
173 - super().__init__(enabled)
174 - self.epsilon = epsilon
175 - self._rmsnorm_op = RMS_OP
176 - self.match_rocm_aiter = match_rocm_aiter
177 -
178 - if match_rocm_aiter:
179 - self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_op()
180 -
181 - def inputs(self) -> list[torch.Tensor]:
182 - input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
183 - weight = self.empty(16)
184 - return [input, weight]
185 -
186 - def forward_rocm_aiter(
187 - self,
188 - input: torch.Tensor,
189 - weight: torch.Tensor,
190 - ) -> torch.Tensor:
191 - return self._rmsnorm_op(
192 - x=input,
193 - weight=weight,
194 - variance_epsilon=self.epsilon,
195 - )
196 -
197 - def forward_custom(
198 - self,
199 - input: torch.Tensor,
200 - weight: torch.Tensor,
201 - ) -> torch.Tensor:
202 - if self.match_rocm_aiter:
203 - return self.forward_rocm_aiter(input, weight)
204 -
205 - result = torch.empty_like(input)
206 - _, result = auto_functionalized(
207 - self._rmsnorm_op,
208 - result=result,
209 - input=input,
210 - weight=weight,
211 - epsilon=self.epsilon,
212 - )
213 -
214 - return result
215 -
216 - def forward_native(
217 - self,
218 - input: torch.Tensor,
219 - weight: torch.Tensor,
220 - ) -> torch.Tensor:
221 - return RMSNorm.forward_static(
222 - input, self.epsilon, input.size(-1), self.model_dtype, weight
223 - )
224 -
225 -
162 class MatcherFusedAddRMSNorm(MatcherCustomOp):
163 def __init__(
164 self,
vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
+4 -4
@@ -10,6 +10,7 @@ from torch import fx
10 from torch._higher_order_ops.auto_functionalize import auto_functionalized
11 from torch._inductor.pattern_matcher import PatternMatcherPass
12
13 +import vllm.ir.ops
14 from vllm.config import VllmConfig, get_layers_from_vllm_config
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.attention import Attention
@@ -17,7 +18,7 @@ from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
18
19 from ..inductor_pass import enable_fake_mode
20 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
20 -from .matcher_utils import MatcherRMSNorm, MatcherRotaryEmbedding
21 +from .matcher_utils import MatcherRotaryEmbedding
22 from .rms_quant_fusion import empty_bf16, empty_fp32, empty_i64
23
24 logger = init_logger(__name__)
@@ -64,7 +65,6 @@ class QkNormRopePattern:
65 self.q_size = self.num_heads * self.head_dim
66 self.kv_size = self.num_kv_heads * self.head_dim
67 self.eps = eps
67 - self.rmsnorm_matcher = MatcherRMSNorm(eps)
68 self.is_neox = is_neox
69 self.rope_flashinfer = rope_flashinfer
70 self.rope_matcher = MatcherRotaryEmbedding(
@@ -129,14 +129,14 @@ class QkNormRopePattern:
129 q_by_head = q.view(
130 *q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim
131 )
132 - q_normed_by_head = self.rmsnorm_matcher(q_by_head, q_weight)
132 + q_normed_by_head = vllm.ir.ops.rms_norm(q_by_head, q_weight, self.eps)
133 q_flat = q_normed_by_head.view(q.shape)
134
135 # K path: view -> RMS -> view back to k.shape
136 k_by_head = k.view(
137 *k.shape[:-1], k.shape[-1] // self.head_dim, self.head_dim
138 )
139 - k_normed_by_head = self.rmsnorm_matcher(k_by_head, k_weight)
139 + k_normed_by_head = vllm.ir.ops.rms_norm(k_by_head, k_weight, self.eps)
140 k_flat = k_normed_by_head.view(k.shape)
141
142 # RoPE: apply to flattened q/k
vllm/compilation/passes/fusion/rms_quant_fusion.py
+18 -20
@@ -9,6 +9,7 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized
9 from torch._inductor.pattern_matcher import PatternMatcherPass
10 from torch._ops import OpOverload
11
12 +import vllm.ir.ops
13 from vllm.config import VllmConfig, get_current_vllm_config
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -30,7 +31,6 @@ from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
31 from .matcher_utils import (
32 MatcherFusedAddRMSNorm,
33 MatcherQuantFP8,
33 - MatcherRMSNorm,
34 )
35
36 logger = init_logger(__name__)
@@ -54,7 +54,6 @@ def empty_i64(*args: Any, **kwargs: Any) -> torch.Tensor:
54 return torch.empty(*args, **kwargs, dtype=torch.int64, device="cuda")
55
56
57 -RMS_OP = torch.ops._C.rms_norm.default
57 RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
58
59 QUANT_OPS: dict[QuantKey, OpOverload] = {
@@ -131,11 +130,9 @@ class RMSNormQuantPattern:
130 assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"
131 self.FUSED_OP = FUSED_OPS[key]
132
134 - self.rmsnorm_matcher = (
135 - MatcherRMSNorm(epsilon)
136 - if not key.fused_add
137 - else MatcherFusedAddRMSNorm(epsilon)
138 - )
133 + if key.fused_add:
134 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)
135 +
136 self.quant_matcher = MatcherQuantFP8(
137 key.quant,
138 has_col_major_scales=has_col_major_scales,
@@ -161,16 +158,12 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
158 def pattern(
159 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
160 ) -> torch.Tensor:
164 - result_rms = self.rmsnorm_matcher(input, weight)
161 + result_rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
162 return self.quant_matcher(result_rms, scale)[0]
163
164 def replacement(
165 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
166 ) -> torch.Tensor:
170 - # In case we're matching native rms-norm, conversions might be
171 - # optimized out. We convert here just to be safe.
172 - input = input.to(dtype=self.model_dtype)
173 -
167 result = torch.empty(
168 input.shape, device=input.device, dtype=self.quant_dtype
169 )
@@ -187,8 +180,8 @@ class RMSNormStaticQuantPattern(RMSNormQuantPattern):
180 return at[1]
181
182 inputs = [
190 - # input, weight
191 - *self.rmsnorm_matcher.inputs(),
183 + empty_bf16(5, 16), # input
184 + empty_bf16(16), # weight
185 self.quant_matcher.inputs()[1], # scale
186 ]
187 pattern(*inputs)
@@ -391,7 +384,7 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
384 def pattern(
385 input: torch.Tensor, weight: torch.Tensor, scale: torch.Tensor
386 ) -> tuple[torch.Tensor, torch.Tensor]:
394 - result_rms = self.rmsnorm_matcher(input, weight)
387 + result_rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
388 result = torch.empty(
389 result_rms.shape,
390 device=result_rms.device,
@@ -442,12 +435,14 @@ class RMSNormGroupQuantPattern(RMSNormQuantPattern):
435 # result, scale
436 return at[1], at[2]
437
445 - scale = self.quant_matcher.empty_f32(1, 1)
446 -
438 pm.register_replacement(
439 pattern,
440 replacement,
450 - self.rmsnorm_matcher.inputs() + [scale],
441 + [
442 + empty_bf16(5, 16), # input
443 + empty_bf16(16), # weight
444 + self.quant_matcher.empty_f32(1, 1), # scale
445 + ],
446 pm.fwd_only,
447 pm_pass,
448 )
@@ -472,7 +467,7 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
467 def pattern(
468 input: torch.Tensor, weight: torch.Tensor
469 ) -> tuple[torch.Tensor, torch.Tensor]:
475 - result_rms = self.rmsnorm_matcher(input, weight)
470 + result_rms = vllm.ir.ops.rms_norm(input, weight, self.epsilon)
471 # result, scale
472 return self.quant_matcher(result_rms) # type: ignore[no-any-return]
473
@@ -502,7 +497,10 @@ class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
497 pm.register_replacement(
498 pattern,
499 replacement,
505 - self.rmsnorm_matcher.inputs(),
500 + [
501 + empty_bf16(5, 16), # input
502 + empty_bf16(16), # weight
503 + ],
504 pm.fwd_only,
505 pm_pass,
506 )
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+22 -10
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4
5 import torch
6 import torch._inductor.pattern_matcher as pm
@@ -24,7 +25,6 @@ from .act_quant_fusion import ActivationQuantPattern
25 from .matcher_utils import (
26 MatcherFusedAddRMSNorm,
27 MatcherQuantFP8,
27 - MatcherRMSNorm,
28 MatcherSiluAndMul,
29 )
30 from .rms_quant_fusion import (
@@ -41,17 +41,23 @@ class AiterRMSNormQuantPattern:
41 ):
42 self.epsilon = epsilon
43 self.quant_dtype = key.quant.dtype
44 + self.device = torch.device("cuda")
45
45 - self.rmsnorm_matcher = (
46 - MatcherRMSNorm(epsilon, match_rocm_aiter=True)
47 - if not key.fused_add
48 - else MatcherFusedAddRMSNorm(epsilon, match_rocm_aiter=True)
49 - )
46 + if key.fused_add:
47 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(
48 + epsilon, match_rocm_aiter=True
49 + )
50 self.quant_matcher = MatcherQuantFP8(
51 key.quant,
52 match_rocm_aiter=match_aiter_quant,
53 )
54
55 + def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
56 + return torch.empty(*args, dtype=torch.bfloat16, device=self.device, **kwargs)
57 +
58 + def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor:
59 + return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs)
60 +
61
62 class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
63 """AITER RMSNorm + Dynamic Quantization pattern."""
@@ -79,7 +85,7 @@ class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
85 input: torch.Tensor,
86 weight: torch.Tensor,
87 ) -> tuple[torch.Tensor, torch.Tensor]:
82 - result_rms = self.rmsnorm_matcher(input, weight)
88 + result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
89 result, scale = self.quant_matcher(result_rms)
90 return result, scale
91
@@ -99,7 +105,8 @@ class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
105 pm.register_replacement(
106 pattern,
107 replacement,
102 - self.rmsnorm_matcher.inputs(),
108 + # input, weight
109 + [self.empty(5, 16), self.empty(16)],
110 pm.fwd_only,
111 pm_pass,
112 )
@@ -188,7 +195,7 @@ class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
195 input: torch.Tensor,
196 weight: torch.Tensor,
197 ) -> tuple[torch.Tensor, torch.Tensor]:
191 - result_rms = self.rmsnorm_matcher(input, weight)
198 + result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
199 result, scale = self.quant_matcher(result_rms)
200 return result, scale
201
@@ -206,7 +213,12 @@ class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
213 return at[0], at[1]
214
215 pm.register_replacement(
209 - pattern, replacement, self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass
216 + pattern,
217 + replacement,
218 + # input, weight
219 + [self.empty(5, 16), self.empty(16)],
220 + pm.fwd_only,
221 + pm_pass,
222 )
223
224
vllm/compilation/passes/fusion/sequence_parallelism.py
+18 -17
@@ -10,6 +10,7 @@ import torch._inductor.pattern_matcher as pm
10 import torch.fx as fx
11 from torch._inductor.pattern_matcher import PatternMatcherPass
12
13 +import vllm.ir.ops
14 from vllm.config import VllmConfig
15 from vllm.config.utils import Range
16 from vllm.distributed import get_tp_group, tensor_model_parallel_all_reduce
@@ -22,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 from ..inductor_pass import enable_fake_mode
24 from ..utility.noop_elimination import NoOpEliminationPass
25 from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
25 -from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
26 +from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8
27
28 logger = init_logger(__name__)
29
@@ -122,35 +123,38 @@ class _SequenceParallelPatternHelper:
123 x, dim=0, world_size=self.tp_size, group_name=self.tp_group.unique_name
124 )
125
126 + def empty(self, *args: Any, **kwargs: Any) -> torch.Tensor:
127 + return torch.empty(*args, dtype=self.dtype, device=self.device, **kwargs)
128 +
129 + def empty_f32(self, *args: Any, **kwargs: Any) -> torch.Tensor:
130 + return torch.empty(*args, dtype=torch.float32, device=self.device, **kwargs)
131 +
132
133 class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):
134 def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:
135 super().__init__(epsilon, dtype, device)
129 - self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
136
137 def get_inputs(self) -> list[torch.Tensor]:
132 - input = torch.empty([1, 8, 4], device=self.device, dtype=self.dtype)
133 - arg3_1 = torch.empty([4], device=self.device, dtype=self.dtype)
134 -
135 - return [input, arg3_1]
138 + # input, weight
139 + return [self.empty([1, 8, 4]), self.empty([4])]
140
141 def register(self, pm_pass: PatternMatcherPass) -> None:
142 def pattern(
143 input: torch.Tensor,
140 - arg3_1: torch.Tensor,
144 + weight: torch.Tensor,
145 ) -> tuple[torch.Tensor, torch.Tensor]:
146 all_reduce = self._all_reduce(input)
143 - rmsnorm = self.rmsnorm_matcher(all_reduce, arg3_1)
147 + rmsnorm = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon)
148
149 return rmsnorm, all_reduce
150
151 def replacement(
152 input: torch.Tensor,
149 - arg3_1: torch.Tensor,
153 + weight: torch.Tensor,
154 ) -> tuple[torch.Tensor, torch.Tensor]:
155 reduce_scatter = self._reduce_scatter(input)
156
153 - rmsnorm = self.rmsnorm_matcher(reduce_scatter, arg3_1)
157 + rmsnorm = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon)
158 all_gather = self._all_gather(rmsnorm)
159 return all_gather, reduce_scatter
160
@@ -222,14 +226,11 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
226 device: str | None,
227 ) -> None:
228 super().__init__(epsilon, dtype, device)
225 - self.rmsnorm_matcher = MatcherRMSNorm(epsilon)
229 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)
230
231 def get_inputs(self) -> list[torch.Tensor]:
229 - input = torch.zeros([1, 8, 4], device=self.device, dtype=self.dtype)
230 - weight = torch.empty([4], device=self.device, dtype=self.dtype)
231 - scale = torch.tensor(1.0, device=self.device, dtype=torch.float32)
232 - return [input, weight, scale]
232 + # input, weight, scale
233 + return [self.empty([1, 8, 4]), self.empty([4]), self.empty_f32([1, 1])]
234
235 def register(self, pm_pass: PatternMatcherPass) -> None:
236 def pattern(
@@ -238,7 +239,7 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
239 scale: torch.Tensor,
240 ) -> tuple[torch.Tensor, torch.Tensor]:
241 all_reduce = self._all_reduce(input)
241 - rms = self.rmsnorm_matcher(all_reduce, weight)
242 + rms = vllm.ir.ops.rms_norm(all_reduce, weight, self.epsilon)
243 quant, _ = self.quant_matcher(rms, scale)
244 return quant, all_reduce
245
@@ -248,7 +249,7 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):
249 scale: torch.Tensor,
250 ) -> tuple[torch.Tensor, torch.Tensor]:
251 reduce_scatter = self._reduce_scatter(input)
251 - rms = self.rmsnorm_matcher(reduce_scatter, weight)
252 + rms = vllm.ir.ops.rms_norm(reduce_scatter, weight, self.epsilon)
253 quant, _ = self.quant_matcher(rms, scale)
254 all_gather = self._all_gather(quant)
255
vllm/compilation/passes/ir/__init__.py
vllm/compilation/passes/ir/lowering_pass.py new
+158
@@ -0,0 +1,158 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections import defaultdict
4 +from collections.abc import Iterable
5 +
6 +from torch import fx
7 +from torch._inductor.pattern_matcher import (
8 + CallFunctionVarArgs,
9 + Match,
10 + PatternMatcherPass,
11 + register_graph_pattern,
12 +)
13 +from torch._ops import OpOverload, OpOverloadPacket
14 +
15 +from vllm.config import VllmConfig
16 +from vllm.ir.op import IrOp
17 +from vllm.logger import init_logger
18 +from vllm.logging_utils import lazy
19 +
20 +from ..vllm_inductor_pass import VllmInductorPass
21 +
22 +logger = init_logger(__name__)
23 +
24 +
25 +def get_default_overload(op: OpOverload | OpOverloadPacket) -> OpOverload:
26 + if isinstance(op, OpOverloadPacket):
27 + return op.default
28 + assert isinstance(op, OpOverload), "Expected an OpOverload or OpOverloadPacket"
29 + return op
30 +
31 +
32 +def get_ir_op(node: fx.Node) -> IrOp | None:
33 + if node.op != "call_function":
34 + return None
35 +
36 + if not isinstance(node.target, (OpOverload, OpOverloadPacket)):
37 + return None
38 +
39 + op_overload = get_default_overload(node.target)
40 + if op_overload.namespace != "vllm_ir":
41 + return None
42 +
43 + op_name = op_overload._opname
44 + if op_name not in IrOp.registry:
45 + logger.warning(
46 + "Unknown vLLM IR op %s, there's likely an issue with torch registration, "
47 + "or a torch custom op was registered in the vllm_ir namespace by mistake.",
48 + op_name,
49 + )
50 + return None
51 +
52 + ir_op = IrOp.registry[op_name]
53 + return ir_op
54 +
55 +
56 +class VllmIRLoweringPass(VllmInductorPass):
57 + """
58 + This pass lowers vLLM IR ops to their implementations the priority list.
59 + """
60 +
61 + def __init__(self, vllm_config: VllmConfig) -> None:
62 + super().__init__(vllm_config)
63 + self.patterns = PatternMatcherPass(self.pass_name)
64 + self.selected_impls: dict[str, dict[str, str]] = defaultdict(lambda: {})
65 + self.ops = [ir_op.torch_op for ir_op in IrOp.registry.values()]
66 +
67 + # Look for any call_function node where the target is a vLLM IR op.
68 + # Then, lower_matched_op will select, trace, and insert the implementation.
69 + register_graph_pattern(
70 + CallFunctionVarArgs(self.ops),
71 + pass_dict=self.patterns,
72 + )(self.lower_matched_op)
73 +
74 + def lower_matched_op(self, match: Match, *args, **kwargs):
75 + # TODO(luka) I think args and kwargs are for the match, but just use the node?
76 +
77 + assert len(match.nodes) == 1, "Expected single node match"
78 + node = match.nodes[0]
79 + ir_op = get_ir_op(node)
80 + assert ir_op is not None, "Expected vLLM IR op"
81 + assert not node.kwargs # I think there should never be kwargs here
82 +
83 + # Select and record the implementation, using fake args
84 + fake_args = fx.map_arg(node.args, lambda arg: arg.meta["val"])
85 + ir_op_impl = ir_op.dispatch(*fake_args)
86 + self.selected_impls[ir_op.name][node.name] = ir_op_impl.provider
87 +
88 + # replace_by_example wants node args, not the fake tensors
89 + # TODO(luka): Use aot_export_module to get functionalized graph
90 + # TODO(luka): Cache the fx_replacement to avoid re-tracing the same impl
91 +
92 + # Defaults not present on node.args but required for replacement tracing
93 + bound_args = ir_op._py_signature.bind(*node.args)
94 + bound_args.apply_defaults()
95 + match.replace_by_example(ir_op_impl.impl_fn, bound_args.args)
96 +
97 + @VllmInductorPass.time_and_log
98 + def __call__(self, graph: fx.Graph) -> None:
99 + # clear at the beginning instead of end, so that tests can inspect
100 + self.selected_impls.clear()
101 +
102 + count = self.patterns.apply(graph)
103 + logger.debug("VllmIRLoweringPass lowered %d vLLM IR nodes", count)
104 +
105 + # TODO write self.selected_impls to depyf/tlparse dir
106 + def count_items(impls: Iterable[str]) -> dict[str, int]:
107 + counts: dict[str, int] = defaultdict(lambda: 0)
108 + for impl in impls:
109 + counts[impl] += 1
110 + return counts
111 +
112 + def print_count(counts: dict[str, int]) -> str:
113 + # e.g., "impl1*3,impl2"
114 + impl_count = lambda i, c: f"{i}" if c == 1 else f"{i}*{c}"
115 + return ",".join(impl_count(impl, count) for impl, count in counts.items())
116 +
117 + logger.debug(
118 + "Selected implementations: %s",
119 + lazy(
120 + lambda: ", ".join(
121 + f"{op}={print_count(count_items(impls_by_node.values()))}"
122 + for op, impls_by_node in self.selected_impls.items()
123 + )
124 + ),
125 + )
126 +
127 + failed_nodes: list[fx.Node] = []
128 + failed_ops: set[str] = set()
129 + # Check no vllm_ir nodes were left in the graph
130 + for node in graph.nodes:
131 + if (ir_op := get_ir_op(node)) is None:
132 + continue
133 +
134 + failed_nodes.append(node)
135 + failed_ops.add(ir_op.name)
136 +
137 + if failed_nodes or failed_ops:
138 + logger.warning("Failed to lower vLLM IR ops: %s", ",".join(failed_ops))
139 + logger.warning("Full node list: %s", failed_nodes)
140 +
141 + def uuid(self) -> str:
142 + """
143 + IR op priority & impl sources affect lowering pass output,
144 + so we include them in the cache key.
145 + """
146 + priorities = {name: op.get_priority() for name, op in IrOp.registry.items()}
147 + priorities_str = ";".join(
148 + f"{name}={','.join(p)}" for name, p in priorities.items()
149 + )
150 +
151 + impl_uuids_str = ";".join(
152 + f"{name}={
153 + ','.join(IrOp.registry[name].impls[provider].uuid() for provider in p)
154 + }"
155 + for name, p in priorities.items()
156 + )
157 +
158 + return f"{super().uuid()}|{priorities_str}|{impl_uuids_str}"
vllm/compilation/passes/pass_manager.py
+17 -3
@@ -14,6 +14,7 @@ from vllm.logger import init_logger
14 from vllm.platforms import current_platform
15 from vllm.utils.system_utils import set_env_var
16
17 +from .ir.lowering_pass import VllmIRLoweringPass
18 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
19
20 if rocm_aiter_ops.is_enabled():
@@ -99,8 +100,17 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
100 else:
101 logger.debug("Skipping %s with compile range %s", pass_, compile_range)
102
102 - # post-cleanup goes before fix_functionalization
103 - # because it requires a functional graph
103 + # perform the first post-cleanup before IR lowering to clean up fusion artifacts
104 + # and make sure no dead IR ops are lowered.
105 + self.post_cleanup(graph)
106 + VllmInductorPass.dump_prefix += 1
107 +
108 + # lowering before cleanup so DCE can clean up lowered ops.
109 + # DCE handles mutating ops correctly as well.
110 + self.ir_lowering(graph)
111 + VllmInductorPass.dump_prefix += 1
112 +
113 + # clean up after lowering again
114 self.post_cleanup(graph)
115 VllmInductorPass.dump_prefix += 1
116
@@ -152,7 +162,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
162 self.passes += [SplitCoalescingPass(config)]
163 self.passes += [QKNormRoPEFusionPass(config)]
164
155 - # needs a functional graph
165 + self.ir_lowering = VllmIRLoweringPass(config)
166 self.post_cleanup = PostCleanupPass(config)
167 self.fix_functionalization = FixFunctionalizationPass(config)
168
@@ -171,6 +181,10 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
181 state: dict[str, Any] = {"pass_config": self.pass_config.compute_hash()}
182 for pass_ in self.passes:
183 passes.append(pass_.uuid())
184 +
185 + passes.append(self.post_cleanup.uuid())
186 + passes.append(self.ir_lowering.uuid())
187 + passes.append(self.post_cleanup.uuid())
188 passes.append(self.fix_functionalization.uuid())
189
190 # Include the compile range in the uuid to ensure that inductor
vllm/compilation/passes/vllm_inductor_pass.py
+1
@@ -152,6 +152,7 @@ class VllmPatternMatcherPass(VllmInductorPass):
152 f"auto_functionalized as auto_functionalized\n"
153 f"from torch._inductor.pattern_matcher import *\n"
154 f"vllm = torch.ops.vllm",
155 + "vllm_ir = torch.ops.vllm_ir",
156 file=f,
157 )
158
vllm/config/compilation.py
+10
@@ -466,6 +466,15 @@ class CompilationConfig:
466 disabled when running with Inductor: mode>CompilationMode.NONE and
467 backend="inductor".
468 Inductor generates (fused) Triton kernels for disabled custom ops."""
469 +
470 + ir_enable_torch_wrap: bool = None # type: ignore[assignment]
471 + """If True, enable vllm_ir torch custom op wrapping during the forward pass.
472 + When False, torch custom op wrapping is disabled, allowing Dynamo to trace the
473 + selected implementation directly or avoiding torch custom op overhead in eager mode.
474 + Defaults to True when using Inductor with vllm-compile
475 + (backend=="inductor" and mode == VLLM_COMPILE), False otherwise.
476 + """
477 +
478 splitting_ops: list[str] | None = None
479 """A list of ops to exclude from cudagraphs, used in piecewise compilation.
480
@@ -830,6 +839,7 @@ class CompilationConfig:
839 "cudagraph_mode",
840 "max_cudagraph_capture_size",
841 "use_inductor_graph_partition",
842 + "ir_enable_torch_wrap",
843 mode="wrap",
844 )
845 @classmethod
vllm/config/kernel.py
+142 -19
@@ -1,13 +1,106 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
3 +import contextlib
4 from collections.abc import Callable
5 -from typing import Any, Literal
5 +from dataclasses import asdict, fields
6 +from typing import TYPE_CHECKING, Any, Literal
7 +
8 +from pydantic import Field, field_validator
9 +
10 +from vllm.config.utils import config, get_hash_factors, hash_factors
11 +from vllm.logger import init_logger
12 +
13 +if TYPE_CHECKING:
14 + from vllm.config import VllmConfig
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +@config
20 +class IrOpPriorityConfig:
21 + """
22 + Configuration for vLLM IR op priority for dispatching/lowering during the
23 + forward pass. Each member is a list of strings, which will be passed to
24 + vllm.ir.ops.<op_name>.set_priority() for the duration of the forward pass.
25 + A single comma-separated string is accepted as well,
26 +
27 + If specified manually, platform defaults will be appended to the lists.
28 + See KernelConfig.set_platform_defaults().
29 + """
30 +
31 + rms_norm: list[str] = Field(default_factory=list)
32 + """Priority list for vllm.ir.ops.rms_norm"""
33 +
34 + def compute_hash(self) -> str:
35 + """
36 + Produces a hash unique to the pass configuration.
37 + Any new fields that affect compilation should be added to the hash.
38 + Any future fields that don't affect compilation should be excluded.
39 +
40 + Also, manually add IR op impl UUIDs to make sure they affect the compile cache.
41 + """
42 + factors = get_hash_factors(self, set())
43 +
44 + # Implementations are hidden from Dynamo,
45 + # so they don't show up in the traced files list.
46 + from vllm.ir.op import IrOp
47 +
48 + assert "_impls" not in factors
49 + factors["_impls"] = {
50 + name: {
51 + provider: IrOp.registry[name].impls[provider].uuid() for provider in p
52 + }
53 + for name, p in asdict(self).items()
54 + }
55 +
56 + return hash_factors(factors)
57
7 -from pydantic import field_validator
58 + @field_validator("*", mode="before")
59 + @classmethod
60 + def _to_list_str(cls, value: str | list[str]):
61 + if isinstance(value, str):
62 + value = value.replace(" ", "").split(",")
63 +
64 + assert all(isinstance(v, str) for v in value)
65 + return value
66 +
67 + @contextlib.contextmanager
68 + def set_priority(self):
69 + """
70 + Context manager to set the IR op priority for all op members.
71 + It also imports vllm.kernels to ensure all implementations are made available.
72 + """
73 + import vllm.kernels # noqa: F401, registers IR op implementations
74 + from vllm.ir.op import IrOp
75 +
76 + with contextlib.ExitStack() as stack:
77 + for field in fields(self):
78 + op_priority = getattr(self, field.name)
79 + assert op_priority is not None, (
80 + f"IR op priority for {field.name} must be set"
81 + )
82 + logger.debug(
83 + "Setting IR op priority for %s to %s", field.name, op_priority
84 + )
85 + ir_op = IrOp.registry[field.name]
86 + stack.enter_context(ir_op.set_priority(op_priority))
87 +
88 + yield
89 +
90 + @classmethod
91 + def with_default(
92 + cls, default: list[str], /, **kwargs: list[str]
93 + ) -> "IrOpPriorityConfig":
94 + """
95 + A helper to create an IrOpPriorityConfig where fields not specified in kwargs
96 + use the given default list.
97 + """
98 + for field in fields(cls):
99 + if field.name not in kwargs:
100 + kwargs[field.name] = list(default)
101 +
102 + return cls(**kwargs)
103
9 -from vllm.config.utils import config
10 -from vllm.utils.hashing import safe_hash
104
105 MoEBackend = Literal[
106 "auto",
@@ -26,6 +119,12 @@ MoEBackend = Literal[
119 class KernelConfig:
120 """Configuration for kernel selection and warmup behavior."""
121
122 + ir_op_priority: IrOpPriorityConfig = Field(default_factory=IrOpPriorityConfig)
123 + """
124 + vLLM IR op priority for dispatching/lowering during the forward pass.
125 + Platform defaults appended automatically during VllmConfig.__post_init__.
126 + """
127 +
128 enable_flashinfer_autotune: bool = None # type: ignore[assignment]
129 """If True, run FlashInfer autotuning during kernel warmup."""
130
@@ -51,21 +150,17 @@ class KernelConfig:
150
151 def compute_hash(self) -> str:
152 """
54 - WARNING: Whenever a new field is added to this config,
55 - ensure that it is included in the factors list if
56 - it affects the computation graph.
57 -
58 - Provide a hash that uniquely identifies all the configs
59 - that affect the structure of the computation
60 - graph from input ids/embeddings to the final hidden states,
61 - excluding anything before input ids/embeddings and after
62 - the final hidden states.
153 + Produces a hash unique to the pass configuration.
154 + Any new fields that affect compilation should be added to the hash.
155 + Any future fields that don't affect compilation should be excluded.
156 """
64 - # no factors to consider.
65 - # this config will not affect the computation graph.
66 - factors: list[Any] = []
67 - hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
68 - return hash_str
157 + ignored_factors = {
158 + "enable_flashinfer_autotune",
159 + "ir_op_priority", # handled separately below
160 + }
161 + factors = get_hash_factors(self, ignored_factors)
162 + factors["ir_op_priority"] = self.ir_op_priority.compute_hash()
163 + return hash_factors(factors)
164
165 @field_validator("enable_flashinfer_autotune", mode="wrap")
166 @classmethod
@@ -74,3 +169,31 @@ class KernelConfig:
169 if value is None:
170 return value
171 return handler(value)
172 +
173 + def set_platform_defaults(self, vllm_config: "VllmConfig") -> None:
174 + """Set platform-specific defaults for the kernel config."""
175 + from vllm.platforms import current_platform
176 +
177 + platform_op_priority = current_platform.get_default_ir_op_priority(vllm_config)
178 + logger.debug(
179 + "Setting platform-specific IR op priority defaults: %s, user-defined: %s",
180 + platform_op_priority,
181 + self.ir_op_priority,
182 + )
183 + for op_name, op_priority in asdict(platform_op_priority).items():
184 + current_op_priority: list[str] = getattr(self.ir_op_priority, op_name)
185 + if current_op_priority is None:
186 + setattr(self.ir_op_priority, op_name, op_priority)
187 + else:
188 + # Append platform-specific priorities
189 + # Must be idempotent because vllm_config.set_platform_defaults() may be
190 + # called multiple times (due to VllmConfig.__post_init__ manual call).
191 + unique_op_priority = [
192 + op for op in op_priority if op not in current_op_priority
193 + ]
194 + current_op_priority.extend(unique_op_priority)
195 +
196 + logger.info(
197 + "Final IR op priority after setting platform defaults: %s",
198 + self.ir_op_priority,
199 + )
vllm/config/vllm.py
+23 -4
@@ -95,9 +95,11 @@ def enable_norm_fusion(cfg: "VllmConfig") -> bool:
95 """Enable if either RMS norm or quant FP8 custom op is active;
96 otherwise Inductor handles fusion."""
97
98 - return cfg.compilation_config.is_custom_op_enabled(
99 - "rms_norm"
100 - ) or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
98 + return (
99 + cfg.compilation_config.is_custom_op_enabled("rms_norm")
100 + or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
101 + or cfg.kernel_config.ir_op_priority.rms_norm[0] != "native"
102 + )
103
104
105 def enable_act_fusion(cfg: "VllmConfig") -> bool:
@@ -417,6 +419,10 @@ class VllmConfig:
419 vllm_factors.append(self.compilation_config.compute_hash())
420 else:
421 vllm_factors.append("None")
422 + if self.kernel_config:
423 + vllm_factors.append(self.kernel_config.compute_hash())
424 + else:
425 + vllm_factors.append(None)
426 if self.kv_transfer_config:
427 vllm_factors.append(self.kv_transfer_config.compute_hash())
428 else:
@@ -890,6 +896,13 @@ class VllmConfig:
896 else:
897 self.compilation_config.mode = CompilationMode.NONE
898
899 + # By default, enable torch wrapping only when using custom Inductor lowering
900 + if self.compilation_config.ir_enable_torch_wrap is None:
901 + self.compilation_config.ir_enable_torch_wrap = (
902 + self.compilation_config.mode == CompilationMode.VLLM_COMPILE
903 + and self.compilation_config.backend == "inductor"
904 + )
905 +
906 if all(s not in self.compilation_config.custom_ops for s in ("all", "none")):
907 if (
908 self.compilation_config.backend == "inductor"
@@ -899,6 +912,11 @@ class VllmConfig:
912 else:
913 self.compilation_config.custom_ops.append("all")
914
915 + # This populates IR op priorities,
916 + # must happen after compilation mode and backend are decided,
917 + # but before fusion defaults are applied as those may depend on op priority.
918 + self.kernel_config.set_platform_defaults(self)
919 +
920 default_config = OPTIMIZATION_LEVEL_TO_CONFIG[self.optimization_level]
921 self._apply_optimization_level_defaults(default_config)
922 if self.kernel_config.enable_flashinfer_autotune is None:
@@ -1706,7 +1724,8 @@ class VllmConfig:
1724 f"enable_prefix_caching={self.cache_config.enable_prefix_caching}, "
1725 f"enable_chunked_prefill={self.scheduler_config.enable_chunked_prefill}, " # noqa
1726 f"pooler_config={self.model_config.pooler_config!r}, "
1709 - f"compilation_config={self.compilation_config!r}"
1727 + f"compilation_config={self.compilation_config!r}, "
1728 + f"kernel_config={self.kernel_config!r}"
1729 )
1730
1731 def validate_block_size(self) -> None:
vllm/engine/arg_utils.py
+23 -2
@@ -8,7 +8,7 @@ import functools
8 import json
9 import sys
10 from collections.abc import Callable
11 -from dataclasses import MISSING, dataclass, fields, is_dataclass
11 +from dataclasses import MISSING, asdict, dataclass, fields, is_dataclass
12 from itertools import permutations
13 from types import UnionType
14 from typing import (
@@ -70,7 +70,7 @@ from vllm.config.cache import (
70 PrefixCachingHashAlgo,
71 )
72 from vllm.config.device import Device
73 -from vllm.config.kernel import MoEBackend
73 +from vllm.config.kernel import IrOpPriorityConfig, MoEBackend
74 from vllm.config.lora import MaxLoRARanks
75 from vllm.config.model import (
76 ConvertOption,
@@ -401,6 +401,7 @@ class EngineArgs:
401 max_cudagraph_capture_size: int | None = get_field(
402 CompilationConfig, "max_cudagraph_capture_size"
403 )
404 + ir_op_priority: IrOpPriorityConfig = get_field(KernelConfig, "ir_op_priority")
405 # Note: Specifying a custom executor backend by passing a class
406 # is intended for expert use only. The API may change without
407 # notice.
@@ -657,6 +658,9 @@ class EngineArgs:
658 self.weight_transfer_config = WeightTransferConfig(
659 **self.weight_transfer_config
660 )
661 + if isinstance(self.ir_op_priority, dict):
662 + self.ir_op_priority = IrOpPriorityConfig(**self.ir_op_priority)
663 +
664 # Setup plugins
665 from vllm.plugins import load_general_plugins
666
@@ -1293,6 +1297,7 @@ class EngineArgs:
1297 title="KernelConfig",
1298 description=KernelConfig.__doc__,
1299 )
1300 + kernel_group.add_argument("--ir-op-priority", **kernel_kwargs["ir_op_priority"])
1301 kernel_group.add_argument(
1302 "--enable-flashinfer-autotune",
1303 **kernel_kwargs["enable_flashinfer_autotune"],
@@ -1917,6 +1922,22 @@ class EngineArgs:
1922 if self.moe_backend != "auto":
1923 kernel_config.moe_backend = self.moe_backend
1924
1925 + # Transfer top-level ir_op_priority into KernelConfig.ir_op_priority
1926 + for op_name, op_priority in asdict(self.ir_op_priority).items():
1927 + # Empty means unset
1928 + if not op_priority:
1929 + continue
1930 +
1931 + # Priority cannot be set 2x for the same op
1932 + if getattr(kernel_config.ir_op_priority, op_name):
1933 + raise ValueError(
1934 + f"Op priority for {op_name} specified via both ir_op_priority "
1935 + f"and KernelConfig.ir_op_priority, only one allowed at a time."
1936 + )
1937 +
1938 + # Set the attribute
1939 + setattr(kernel_config.ir_op_priority, op_name, op_priority)
1940 +
1941 load_config = self.create_load_config()
1942
1943 # Pass reasoning_parser into StructuredOutputsConfig
vllm/forward_context.py
+8 -1
@@ -10,6 +10,7 @@ from typing import Any
10 import torch
11
12 import vllm.envs as envs
13 +import vllm.ir
14 from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
15 from vllm.logger import init_logger
16 from vllm.platforms import current_platform
@@ -378,7 +379,13 @@ def set_forward_context(
379 )
380
381 try:
381 - with override_forward_context(forward_context):
382 + with (
383 + override_forward_context(forward_context),
384 + vllm_config.kernel_config.ir_op_priority.set_priority(),
385 + vllm.ir.enable_torch_wrap(
386 + vllm_config.compilation_config.ir_enable_torch_wrap
387 + ),
388 + ):
389 yield
390 finally:
391 global last_logging_time, batchsize_logging_interval
vllm/ir/__init__.py new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from . import ops
4 +from .op import enable_torch_wrap, register_op
5 +
6 +__all__ = ["enable_torch_wrap", "register_op", "ops"]
vllm/ir/op.py new
+414
@@ -0,0 +1,414 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import contextlib
4 +import inspect
5 +from collections.abc import Callable
6 +from pathlib import Path
7 +from typing import Any, ClassVar, overload
8 +
9 +import torch
10 +from torch.library import Library, infer_schema
11 +
12 +from vllm.ir.util import hash_source, weak_cache
13 +from vllm.logger import init_logger
14 +from vllm.logging_utils import lazy, tensors_str_no_data
15 +
16 +vllm_ir_lib = Library("vllm_ir", "FRAGMENT")
17 +
18 +logger = init_logger(__name__)
19 +
20 +RESERVED_PROVIDERS = ["native", "unfused"]
21 +"""Providers that are reserved and cannot be used for custom implementations."""
22 +
23 +_ENABLE_TORCH_WRAP: bool = True
24 +"""Global override flag to control torch op layer wrapping."""
25 +
26 +
27 +@contextlib.contextmanager
28 +def enable_torch_wrap(enable: bool = True):
29 + """
30 + Context manager to enable/disable torch custom op wrapping for vLLM IR ops.
31 + When torch wrapping is disabled, the torch custom op layer is skipped
32 + and IR ops dispatch directly to the implementation.
33 + Helpful for avoiding torch dispatch overhead in eager mode
34 + and avoiding the need for lowering for platforms not using Inductor.
35 + """
36 +
37 + global _ENABLE_TORCH_WRAP
38 + old = _ENABLE_TORCH_WRAP
39 + try:
40 + _ENABLE_TORCH_WRAP = enable
41 + yield
42 + finally:
43 + _ENABLE_TORCH_WRAP = old
44 +
45 +
46 +# 0-param decorator overload
47 +@overload
48 +def register_op(f: Callable[..., Any]) -> "IrOp": ...
49 +
50 +
51 +# parametrized decorator overload
52 +@overload
53 +def register_op(
54 + *,
55 + name: str | None = None,
56 +) -> Callable[[Callable[..., Any]], "IrOp"]: ...
57 +
58 +
59 +def register_op(
60 + f: Callable | None = None,
61 + *,
62 + name: str | None = None,
63 +) -> "IrOp | Callable[[Callable], IrOp]":
64 + """
65 + Register a new vLLM IR op.
66 +
67 + :param f: the native implementation of the op
68 + :param name: the name of the op, defaults to the function name
69 + :return: the IrOp object if f is provided, otherwise a decorator
70 +
71 + Example usage:
72 + ```python
73 + @vllm.ir.register_op
74 + def my_op(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
75 + return x + y
76 +
77 +
78 + @vllm.ir.register_op(name="custom_mul")
79 + def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
80 + return x * y"""
81 +
82 + def decorator(_f: Callable):
83 + op_name: str = _f.__name__ if name is None else name
84 + assert op_name not in IrOp.registry
85 + op = IrOp(op_name, _f)
86 + IrOp.registry[op_name] = op
87 + return op
88 +
89 + if f is not None:
90 + return decorator(f)
91 +
92 + return decorator
93 +
94 +
95 +class IrOp:
96 + registry: ClassVar[dict[str, "IrOp"]] = {}
97 +
98 + name: str
99 + impls: dict[str, "IrOpImpl"]
100 +
101 + def __init__(self, name: str, native_impl: Callable):
102 + self._py_signature = inspect.signature(native_impl)
103 + if any(
104 + p.kind == inspect.Parameter.KEYWORD_ONLY
105 + for p in self._py_signature.parameters.values()
106 + ):
107 + raise ValueError(
108 + f"Op {name} has keyword-only arguments which are not currently "
109 + f"supported. That's because kwargs are not allowed during lowering."
110 + )
111 +
112 + self.name = name
113 + self.impls: dict[str, IrOpImpl] = {}
114 + self._priority_impls: list[IrOpImpl] = []
115 + self._schema_str = infer_schema(native_impl, mutates_args=[])
116 +
117 + # native implementation
118 + self.impls["native"] = IrOpImpl(
119 + self, "native", native_impl, supported=True, supports_args=None
120 + )
121 +
122 + # By default, fake routes directly to native,
123 + # can be overridden by register_fake
124 + self._fake_fn = native_impl
125 +
126 + # torch registration
127 + vllm_ir_lib.define(self.name + self._schema_str)
128 + # CompositeExplicitAutograd is not decomposed
129 + # by ATen IR normalization in AOTAutograd
130 + vllm_ir_lib.impl(
131 + self.name, self._inner_call, dispatch_key="CompositeExplicitAutograd"
132 + )
133 + vllm_ir_lib._register_fake(self.name, self._fake_call)
134 + assert hasattr(torch.ops.vllm_ir, name)
135 + self.torch_op: torch._ops.OpOverload = getattr(torch.ops.vllm_ir, name).default
136 +
137 + def register_fake(self, fn: Callable) -> Callable:
138 + """
139 + Register a fake impl for the torch custom op. If this method is not called,
140 + the native implementation is used directly for the fake implementation.
141 + """
142 + self._fake_fn = fn
143 + return fn
144 +
145 + def _fake_call(self, *args, **kwargs) -> Any:
146 + """
147 + Call to the fake implementation of the op. We use indirection because we want
148 + users to be able to register fake later but also want it to fall back to native
149 + directly by default, instead of going through the dispatching mechanism.
150 + """
151 + return self._fake_fn(*args, **kwargs)
152 +
153 + def register_impl(
154 + self,
155 + provider: str,
156 + *,
157 + supported: bool = True,
158 + supports_args: Callable[..., bool] | None = None,
159 + ):
160 + """
161 + Register an implementation for this custom op.
162 + :param provider: The name of the provider, must be unique.
163 + :param supported: Static support check, use this to check platform support.
164 + :param supports_args: Dynamic arg support check, used for types and shapes.
165 + :return: A decorator that registers the implementation.
166 +
167 + The decorated function must have the same semantics and signature as
168 + the native implementation.
169 +
170 + The provider name must be unique and not one of the RESERVED_PROVIDERS.
171 + The supported and supports_args parameters should not be used to implement
172 + custom enablement logic based on global state (e.g. environment variables).
173 + Instead, supported param should only be used to check for platform support
174 + (e.g. whether a specific hardware or library is available).
175 + supports_args should be used to check whether the provided arguments are
176 + compatible with the implementation.
177 + For custom enablement logic, set op impl priority.
178 +
179 + Example:
180 + ```python
181 + @my_op.register_impl("my_provider", supported=torch.cuda.is_available())
182 + def my_provider_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: ...
183 + ```
184 +
185 + """
186 + assert provider not in RESERVED_PROVIDERS, (
187 + f"Provider name {provider} is reserved."
188 + )
189 +
190 + def _register_impl(f: Callable):
191 + impl = IrOpImpl(self, provider, f, supported, supports_args)
192 + self.impls[provider] = impl
193 +
194 + if self.get_priority():
195 + logger.warning(
196 + "Warning: registering new impl %s for op %s while priority is set.",
197 + provider,
198 + self.name,
199 + )
200 +
201 + return impl
202 +
203 + return _register_impl
204 +
205 + def _inner_call(self, *args, **kwargs) -> Any:
206 + """
207 + Eager call to torch op lands here. When torch wrapping is disabled,
208 + __call__ routes straight here instead of going through torch op dispatching.
209 + """
210 + impl = self.dispatch(*args, **kwargs)
211 + return impl.impl_fn(*args, **kwargs)
212 +
213 + def apply_arg_defaults(self, args) -> tuple:
214 + """
215 + Return args with default values applied.
216 + Defaults are taken from the native implementation signature.
217 +
218 + SHOULD NOT BE USED IN THE DISPATCH PATH (SLOW).
219 + Only for Inductor lowering.
220 + """
221 + bound_args = self._py_signature.bind(*args)
222 + bound_args.apply_defaults()
223 + return bound_args.args
224 +
225 + def dispatch(self, *args, **kwargs) -> "IrOpImpl":
226 + """
227 + Dispatch to the appropriate implementation based on current priority
228 + and argument support checks. Returns the selected IrOpImpl.
229 +
230 + THIS FUNCTION IS ON THE HOT PATH (OP DISPATCH), MUST BE FAST.
231 + """
232 + if not self._priority_impls:
233 + if not torch.compiler.is_compiling():
234 + # Logging not compatible with Dynamo tracing
235 + # (this code is exposed when torch wrapping is disabled)
236 + logger.warning_once(
237 + "Priority not set for op %s, using native implementation.",
238 + self.name,
239 + )
240 + return self.impls["native"]
241 +
242 + for impl in self._priority_impls:
243 + if not impl.supported:
244 + raise ValueError(
245 + f"Implementation {impl.provider} for op {self.name} not supported. "
246 + f"All implementations in priority list must be supported."
247 + )
248 + if impl.supports_args(*args, **kwargs):
249 + return impl
250 +
251 + if not torch.compiler.is_compiling():
252 + logger.debug(
253 + "Skipping provider %s because it does not support "
254 + "%s with args=%s kwargs=%s",
255 + impl.provider,
256 + self.name,
257 + lazy(lambda: tensors_str_no_data(args)),
258 + lazy(lambda: tensors_str_no_data(kwargs)),
259 + )
260 +
261 + raise RuntimeError(
262 + "Priority set incorrectly: the last implementation must "
263 + "support all args (can be native). This is likely an internal bug"
264 + )
265 +
266 + def __call__(self, *args, **kwargs) -> Any:
267 + if not _ENABLE_TORCH_WRAP:
268 + return self._inner_call(*args, **kwargs)
269 +
270 + return self.torch_op(*args, **kwargs)
271 +
272 + def get_priority(self) -> list[str]:
273 + """Get the current dispatch priority for implementations for this op."""
274 + return [p.provider for p in self._priority_impls]
275 +
276 + @contextlib.contextmanager
277 + def set_priority(self, priority: list[str]):
278 + """
279 + Context manager to set the dispatch priority for implementations for this op.
280 + """
281 + assert all(p in self.impls for p in priority), (
282 + "All providers in priority must be registered implementations."
283 + )
284 +
285 + def filter_priority_impls(p_list: list[str]) -> list[IrOpImpl]:
286 + filtered_impls = []
287 + for p in p_list:
288 + impl = self.impls[p]
289 + if not impl.supported:
290 + # Skip unsupported implementations
291 + continue
292 +
293 + filtered_impls.append(impl)
294 +
295 + # If all args are supported, skip other implementations
296 + if impl.supports_all_args:
297 + return filtered_impls
298 +
299 + logger.warning_once(
300 + "Op %s: No implementation in priority list supports all args, "
301 + "execution fallback to native is possible. To silence this warning, "
302 + "explicitly add 'native' to the end of the priority list",
303 + self.name,
304 + )
305 + filtered_impls.append(self.impls["native"])
306 + return filtered_impls
307 +
308 + # Temporarily set priority
309 + old_priority_impls = self._priority_impls
310 + try:
311 + self._priority_impls = filter_priority_impls(priority)
312 + yield
313 + finally:
314 + self._priority_impls = old_priority_impls
315 +
316 + def supported_providers(self) -> list[str]:
317 + return [p.provider for p in self.impls.values() if p.supported]
318 +
319 +
320 +class IrOpImpl:
321 + def __init__(
322 + self,
323 + op: IrOp,
324 + provider: str,
325 + impl_fn: Callable,
326 + supported: bool,
327 + supports_args: Callable[..., bool] | None,
328 + ):
329 + assert provider not in op.impls, (
330 + f"Implementation for provider {provider} already registered."
331 + )
332 + # Native also uses this path, so we allow it here.
333 + assert provider == "native" or provider not in RESERVED_PROVIDERS
334 +
335 + # Enforce the exact same schema as the native implementation.
336 + # This takes care of names, types, and defaults.
337 + schema = infer_schema(impl_fn, mutates_args=[])
338 + if schema != op._schema_str:
339 + raise ValueError(
340 + f"Implementation for provider {provider} has schema '{schema}' which "
341 + f"does not match native schema '{op._schema_str}' for op {op.name}."
342 + )
343 +
344 + if supports_args is not None:
345 + if not callable(supports_args):
346 + raise ValueError(
347 + f"supports_args for provider {provider} must be a callable"
348 + )
349 +
350 + # We also manually validate the supports_args signature.
351 + # Matching signatures allow faster dispatch on the hotpath.
352 +
353 + # Check that supports_args does not have keyword-only parameters
354 + supports_args_signature = inspect.signature(supports_args)
355 + params = supports_args_signature.parameters
356 + if any(p.kind == inspect.Parameter.KEYWORD_ONLY for p in params.values()):
357 + raise ValueError(
358 + f"supports_args for provider {provider} "
359 + f"cannot have keyword-only parameters"
360 + )
361 +
362 + # Check that supports_args has the same total number of parameters
363 + op_params = op._py_signature.parameters
364 + if len(params) != len(op_params):
365 + raise ValueError(
366 + f"supports_args for provider {provider} must have the same number "
367 + f"of parameters ({len(params)}) as the native implementation "
368 + f"({len(op_params)})"
369 + )
370 +
371 + # Check that names and defaults match for supports_args
372 + for p, op_p in zip(params.values(), op_params.values()):
373 + if p.name != op_p.name:
374 + raise ValueError(
375 + f"supports_args for provider {provider} has parameter "
376 + f"'{p.name}' which does not match native parameter "
377 + f"'{op_p.name}'"
378 + )
379 + if p.default != op_p.default:
380 + raise ValueError(
381 + f"supports_args for provider {provider} has parameter "
382 + f"'{p.name}' with default {p.default} which does not match "
383 + f"native default {op_p.default}'"
384 + )
385 +
386 + self.op = op
387 + self.provider = provider
388 + self.impl_fn = impl_fn
389 + self.supported = supported
390 + self._supports_args = supports_args
391 +
392 + @property
393 + def supports_all_args(self) -> bool:
394 + """Check if this implementation supports all args unconditionally."""
395 + return self._supports_args is None
396 +
397 + def supports_args(self, *args, **kwargs) -> bool:
398 + if self._supports_args is None:
399 + return True
400 +
401 + return self._supports_args(*args, **kwargs)
402 +
403 + @weak_cache
404 + def uuid(self):
405 + """
406 + Compile-time hash to uniquely determine whether the implementation has changed.
407 + Used by vllm-compile hash mechanism and torch.compile lowering pass uuid to
408 + control the vLLM compile cache and AOTAutograd/Inductor caches respectively.
409 +
410 + Source file contents do not change so we cache uuid.
411 + TODO(luka): Cache the file hash as multiple impls are likely in the same file.
412 + """
413 + sources = [Path(inspect.getfile(self.impl_fn))]
414 + return hash_source(*sources)
vllm/ir/ops/__init__.py new
+5
@@ -0,0 +1,5 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from .layernorm import rms_norm
4 +
5 +__all__ = ["rms_norm"]
vllm/ir/ops/layernorm.py new
+22
@@ -0,0 +1,22 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +from torch import Tensor
5 +
6 +from ..op import register_op
7 +
8 +
9 +@register_op
10 +def rms_norm(
11 + x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
12 +) -> Tensor:
13 + """Weighted root-mean-square layer normalization"""
14 + orig_dtype = x.dtype
15 + x = x.to(torch.float32)
16 + x_var = x if variance_size is None else x[..., :variance_size]
17 + variance = x_var.pow(2).mean(dim=-1, keepdim=True)
18 + x = x * torch.rsqrt(variance + epsilon)
19 + x = x.to(orig_dtype)
20 + if weight is not None:
21 + x = x * weight
22 + return x
vllm/ir/util.py new
+61
@@ -0,0 +1,61 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 +import hashlib
5 +import inspect
6 +import types
7 +import weakref
8 +from pathlib import Path
9 +from typing import Any
10 +
11 +
12 +def hash_source(*srcs: str | Any) -> str:
13 + """
14 + Utility method to hash the sources of functions or objects.
15 + :param srcs: strings or objects to add to the hash.
16 + Objects and functions have their source inspected.
17 + :return:
18 + """
19 + hasher = hashlib.sha256()
20 + for src in srcs:
21 + if src is None:
22 + src_str = "None"
23 + elif isinstance(src, str):
24 + src_str = src
25 + elif isinstance(src, Path):
26 + src_str = src.read_text()
27 + elif isinstance(src, (types.FunctionType, type)):
28 + src_str = inspect.getsource(src)
29 + else:
30 + # object instance
31 + src_str = inspect.getsource(src.__class__)
32 + hasher.update(src_str.encode("utf-8"))
33 + return hasher.hexdigest()
34 +
35 +
36 +def weak_lru_cache(maxsize: int | None = 128, typed: bool = False):
37 + """
38 + LRU Cache decorator that keeps a weak reference to 'self'.
39 + This avoids memory leakage, which happens when functools.lru_cache
40 + stores a reference to self in the global cache.
41 +
42 + Taken from: https://stackoverflow.com/a/68052994/5082708
43 + """
44 +
45 + def wrapper(func):
46 + @functools.lru_cache(maxsize, typed)
47 + def _func(_self, *args, **kwargs):
48 + return func(_self(), *args, **kwargs)
49 +
50 + @functools.wraps(func)
51 + def inner(self, *args, **kwargs):
52 + return _func(weakref.ref(self), *args, **kwargs)
53 +
54 + return inner
55 +
56 + return wrapper
57 +
58 +
59 +def weak_cache(user_function, /):
60 + """Simple weak equivalent to functools.cache"""
61 + return weak_lru_cache(maxsize=None)(user_function)
vllm/kernels/__init__.py
+4
@@ -1,3 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Kernel implementations for vLLM."""
4 +
5 +from . import aiter_ops, oink_ops, vllm_c, xpu_ops
6 +
7 +__all__ = ["vllm_c", "aiter_ops", "oink_ops", "xpu_ops"]
vllm/kernels/aiter_ops.py new
+79
@@ -0,0 +1,79 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 +
5 +import torch
6 +from torch import Tensor
7 +from torch.library import Library
8 +
9 +from vllm import ir
10 +from vllm.platforms import current_platform
11 +from vllm.utils.torch_utils import direct_register_custom_op
12 +
13 +current_platform.import_kernels()
14 +
15 +
16 +def is_aiter_found() -> bool:
17 + from importlib.util import find_spec
18 +
19 + return find_spec("aiter") is not None
20 +
21 +
22 +aiter_lib = Library("vllm_aiter", "FRAGMENT")
23 +"""
24 +This library holds torch custom ops for wrapped AITER ops.
25 +Many AITER ops want to remain invisible to torch.compile even after lowering.
26 +They are thus wrapped into torch custom ops inside the IR op implementations.
27 +"""
28 +
29 +direct_register_aiter_op = functools.partial(
30 + direct_register_custom_op, target_lib=aiter_lib
31 +)
32 +"""Syntactic sugar for registering AITER custom ops."""
33 +
34 +AITER_SUPPORTED = is_aiter_found()
35 +"""Most kernels in this file are supported if AITER is installed."""
36 +
37 +rms_no_var_16bit_only = (
38 + lambda x, weight, epsilon, variance_size=None: variance_size is None
39 + and x.dtype
40 + in (
41 + torch.float16,
42 + torch.bfloat16,
43 + )
44 +)
45 +"""AITER rms_norm only supports float16 and bfloat16 acts and no var_size override."""
46 +
47 +
48 +@ir.ops.rms_norm.register_impl(
49 + "aiter", supports_args=rms_no_var_16bit_only, supported=AITER_SUPPORTED
50 +)
51 +def rms_norm(
52 + x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
53 +) -> Tensor:
54 + assert variance_size is None
55 + assert x.dtype in (torch.float16, torch.bfloat16)
56 + if weight is None:
57 + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
58 + return torch.ops.vllm_aiter.rms_norm(x, weight, epsilon)
59 +
60 +
61 +def _rms_norm_impl(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor:
62 + from aiter import rms_norm
63 +
64 + if x.dim() > 2:
65 + x_original_shape = x.shape
66 + x = x.reshape(-1, x_original_shape[-1])
67 + x = rms_norm(x, weight, variance_epsilon)
68 + return x.reshape(x_original_shape)
69 +
70 + return rms_norm(x, weight, variance_epsilon)
71 +
72 +
73 +def _rms_norm_fake(x: Tensor, weight: Tensor, variance_epsilon: float) -> Tensor:
74 + return torch.empty_like(x)
75 +
76 +
77 +direct_register_aiter_op(
78 + op_name="rms_norm", op_func=_rms_norm_impl, fake_impl=_rms_norm_fake
79 +)
vllm/kernels/oink_ops.py new
+77
@@ -0,0 +1,77 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm import ir
6 +from vllm.platforms import current_platform
7 +
8 +OINK_AVAILABLE = current_platform.has_device_capability(100) and hasattr(
9 + torch.ops, "oink"
10 +)
11 +
12 +
13 +def has_oink_op(name: str) -> bool:
14 + """Check if a specific oink op is registered."""
15 + return OINK_AVAILABLE and hasattr(torch.ops.oink, name)
16 +
17 +
18 +def _can_view_as_2d(x: torch.Tensor) -> bool:
19 + """Return True if x.view(-1, x.shape[-1]) is viewable (no copy)."""
20 + if x.dim() < 2:
21 + return False
22 + if x.dim() == 2:
23 + return True
24 + # For a view(-1, N) to be valid, all leading dims must be contiguous with
25 + # respect to each other (size-1 dims are ignored).
26 + for dim in range(x.dim() - 1):
27 + # Strides for size-1 dims are irrelevant and can be arbitrary.
28 + if x.size(dim + 1) != 1 and x.stride(dim) != x.stride(dim + 1) * x.size(
29 + dim + 1
30 + ):
31 + return False
32 + return True
33 +
34 +
35 +def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
36 + """Return True if x_2d meets Oink's pointer-path stride constraints."""
37 + if x_2d.dim() != 2:
38 + return False
39 + if x_2d.stride(1) != 1:
40 + return False
41 + # Match Oink's vectorization constraint: stride(0) divisible by 256b.
42 + if x_2d.dtype in (torch.float16, torch.bfloat16):
43 + divby = 16
44 + elif x_2d.dtype == torch.float32:
45 + divby = 8
46 + else:
47 + return False
48 + return (x_2d.stride(0) % divby) == 0
49 +
50 +
51 +oink_rms_supported = (
52 + lambda x, weight, epsilon, variance_size=None: variance_size is None
53 + and weight is not None
54 + and x.dim() >= 2
55 + and x.dtype == weight.dtype
56 + and weight.is_contiguous()
57 + and _can_view_as_2d(x)
58 + and _is_oink_stride_compatible_2d(x.view(-1, x.shape[-1]))
59 +)
60 +"""
61 +Oink rms only supports 2d-like inputs with contiguous weight
62 +and no variance_size override.
63 +"""
64 +
65 +
66 +@ir.ops.rms_norm.register_impl(
67 + "oink", supports_args=oink_rms_supported, supported=has_oink_op("rmsnorm")
68 +)
69 +def rms_norm(
70 + x: torch.Tensor,
71 + weight: torch.Tensor | None,
72 + epsilon: float,
73 + variance_size: int | None = None,
74 +) -> torch.Tensor:
75 + assert variance_size is None
76 + x_2d = x.view(-1, x.shape[-1])
77 + return torch.ops.oink.rmsnorm(x_2d, weight, epsilon).view_as(x)
vllm/kernels/vllm_c.py new
+30
@@ -0,0 +1,30 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +from torch import Tensor
5 +
6 +from vllm import ir
7 +from vllm.platforms import current_platform
8 +
9 +current_platform.import_kernels()
10 +
11 +CUDA_ALIKE = current_platform.is_cuda_alike()
12 +"""Most kernels in this file are supported on all CUDA-alike platforms."""
13 +
14 +rms_no_var_size = lambda x, weight, epsilon, variance_size=None: variance_size is None
15 +"""vLLM kernel does not support variance_size parameter."""
16 +
17 +
18 +@ir.ops.rms_norm.register_impl(
19 + "vllm_c", supports_args=rms_no_var_size, supported=CUDA_ALIKE
20 +)
21 +def rms_norm(
22 + x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
23 +) -> Tensor:
24 + if weight is None:
25 + # Kernel requires weight tensor, pass ones
26 + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
27 + assert variance_size is None
28 + output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
29 + torch.ops._C.rms_norm(output, x, weight, epsilon)
30 + return output
vllm/kernels/xpu_ops.py new
+36
@@ -0,0 +1,36 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +from torch import Tensor
5 +
6 +from vllm import ir
7 +from vllm.platforms import current_platform
8 +
9 +current_platform.import_kernels()
10 +
11 +
12 +def is_xpu_kernels_found() -> bool:
13 + from importlib.util import find_spec
14 +
15 + return find_spec("vllm_xpu_kernels") is not None
16 +
17 +
18 +XPU_KERNELS_SUPPORTED = is_xpu_kernels_found()
19 +"""Kernels in this file are supported if vLLM XPU kernels are installed."""
20 +
21 +rms_no_var = lambda x, weight, epsilon, variance_size=None: variance_size is None
22 +
23 +
24 +@ir.ops.rms_norm.register_impl(
25 + "xpu_kernels", supports_args=rms_no_var, supported=XPU_KERNELS_SUPPORTED
26 +)
27 +def rms_norm(
28 + x: Tensor, weight: Tensor | None, epsilon: float, variance_size: int | None = None
29 +) -> Tensor:
30 + if weight is None:
31 + # Kernel requires weight tensor, pass ones
32 + weight = torch.ones(x.shape[-1], device=x.device, dtype=x.dtype)
33 + assert variance_size is None
34 + output = torch.empty(x.shape, device=x.device, dtype=x.dtype)
35 + torch.ops._C.rms_norm(output, x, weight, epsilon)
36 + return output
vllm/logging_utils/__init__.py
+2
@@ -8,6 +8,7 @@ from vllm.logging_utils.access_log_filter import (
8 from vllm.logging_utils.formatter import ColoredFormatter, NewLineFormatter
9 from vllm.logging_utils.lazy import lazy
10 from vllm.logging_utils.log_time import logtime
11 +from vllm.logging_utils.torch_tensor import tensors_str_no_data
12
13 __all__ = [
14 "NewLineFormatter",
@@ -16,4 +17,5 @@ __all__ = [
17 "create_uvicorn_log_config",
18 "lazy",
19 "logtime",
20 + "tensors_str_no_data",
21 ]
vllm/logging_utils/torch_tensor.py new
+10
@@ -0,0 +1,10 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +
6 +def tensors_str_no_data(arg: Any):
7 + from torch._tensor_str import printoptions
8 +
9 + with printoptions(threshold=1, edgeitems=0):
10 + return str(arg)
vllm/model_executor/layers/layernorm.py
+27 -69
@@ -6,7 +6,9 @@ import torch
6 import torch.nn as nn
7 import torch.nn.functional as F
8
9 -from vllm import _oink_ops, envs
9 +# Import kernels
10 +import vllm.kernels # noqa: F401
11 +from vllm import _oink_ops, envs, ir
12 from vllm._aiter_ops import rocm_aiter_ops
13 from vllm.logger import init_logger
14 from vllm.model_executor.custom_op import CustomOp
@@ -51,23 +53,6 @@ def _is_oink_stride_compatible_2d(x_2d: torch.Tensor) -> bool:
53 return (x_2d.stride(0) % divby) == 0
54
55
54 -def rms_norm(
55 - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
56 -) -> torch.Tensor:
57 - from vllm import _custom_ops as ops
58 -
59 - if envs.VLLM_BATCH_INVARIANT:
60 - return rms_norm_batch_invariant(x, weight, variance_epsilon)
61 - out = torch.empty_like(x)
62 - ops.rms_norm(
63 - out,
64 - x,
65 - weight,
66 - variance_epsilon,
67 - )
68 - return out
69 -
70 -
56 def fused_add_rms_norm(
57 x: torch.Tensor,
58 residual: torch.Tensor,
@@ -105,23 +90,16 @@ def poly_norm(
90 return out
91
92
108 -def dispatch_rocm_rmsnorm_func(
109 - with_fused_add: bool, dtype: torch.dtype, use_aiter: bool = False
110 -):
93 +def dispatch_rocm_rmsnorm_func(dtype: torch.dtype, use_aiter: bool = False):
94 use_aiter = use_aiter and dtype in [
95 torch.float16,
96 torch.bfloat16,
97 ]
98
116 - if use_aiter and with_fused_add:
117 - return rocm_aiter_ops.rms_norm2d_with_add
99 if use_aiter:
119 - return rocm_aiter_ops.rms_norm
120 -
121 - # fall back to CUDA implementation
122 - if with_fused_add:
100 + return rocm_aiter_ops.rms_norm2d_with_add
101 + else:
102 return fused_add_rms_norm
124 - return rms_norm
103
104
105 # --8<-- [start:rms_norm]
@@ -158,20 +136,14 @@ class RMSNorm(CustomOp):
136
137 if current_platform.is_rocm():
138 aiter_rmsnorm_enabled = rocm_aiter_ops.is_rmsnorm_enabled()
161 - self.rocm_norm_func = dispatch_rocm_rmsnorm_func(
162 - with_fused_add=False,
163 - dtype=weight_dtype,
164 - use_aiter=aiter_rmsnorm_enabled,
165 - )
139 self.rocm_norm_func_with_add = dispatch_rocm_rmsnorm_func(
167 - with_fused_add=True, dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled
140 + dtype=weight_dtype, use_aiter=aiter_rmsnorm_enabled
141 )
142
143 # Optional: enable Oink Blackwell RMSNorm custom-op fast path on
144 # compatible CUDA devices (e.g., SM100) when the external Oink
145 # package is available. This is detected once at construction time
146 # to avoid per-call device queries in the hot path.
174 - self._use_oink_rmsnorm = False
147 self._use_oink_fused_add_rmsnorm = False
148 if (
149 not current_platform.is_rocm()
@@ -203,7 +175,6 @@ class RMSNorm(CustomOp):
175 try:
176 device_index = torch.accelerator.current_device_index()
177 if _oink_ops.is_oink_available_for_device(device_index):
206 - self._use_oink_rmsnorm = True
178 self._use_oink_fused_add_rmsnorm = (
179 _oink_ops.has_fused_add_rms_norm()
180 )
@@ -215,7 +186,6 @@ class RMSNorm(CustomOp):
186 "RMSNorm; falling back to vLLM RMSNorm. Error: %s",
187 e,
188 )
218 - self._use_oink_rmsnorm = False
189 self._use_oink_fused_add_rmsnorm = False
190
191 @staticmethod
@@ -270,6 +240,10 @@ class RMSNorm(CustomOp):
240 residual: torch.Tensor | None = None,
241 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
242 """PyTorch-native implementation equivalent to forward()."""
243 + if residual is None:
244 + return ir.ops.rms_norm(
245 + x, self.weight.data, self.variance_epsilon, self.variance_size_override
246 + )
247
248 return self.forward_static(
249 x,
@@ -286,35 +260,14 @@ class RMSNorm(CustomOp):
260 x: torch.Tensor,
261 residual: torch.Tensor | None = None,
262 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
263 + if residual is None and not envs.VLLM_BATCH_INVARIANT:
264 + return ir.ops.rms_norm(
265 + x, self.weight.data, self.variance_epsilon, self.variance_size_override
266 + )
267 +
268 if self.variance_size_override is not None:
269 return self.forward_native(x, residual)
270
292 - # Optional Oink SM100 fast path (no residual). This path is
293 - # torch.compile-friendly via torch.ops.oink.rmsnorm and preserves
294 - # 2D layouts (including padded rows) when using the Oink
295 - # pointer-based kernel.
296 - if (
297 - residual is None
298 - and getattr(self, "_use_oink_rmsnorm", False)
299 - and x.is_cuda
300 - and x.dim() >= 2
301 - and self.has_weight
302 - and not envs.VLLM_BATCH_INVARIANT
303 - and self.weight.data.dtype == x.dtype
304 - and self.weight.data.is_contiguous()
305 - ):
306 - orig_shape = x.shape
307 - hidden_size = orig_shape[-1]
308 - if _can_view_as_2d(x):
309 - x_2d = x.view(-1, hidden_size)
310 - if _is_oink_stride_compatible_2d(x_2d):
311 - y_2d = _oink_ops.rmsnorm(
312 - x_2d,
313 - self.weight.data,
314 - self.variance_epsilon,
315 - )
316 - return y_2d.view(orig_shape)
317 -
271 # Optional Oink SM100 fast path (fused residual-add + RMSNorm, in-place).
272 # This mirrors vLLM's fused_add_rms_norm semantics by mutating both
273 # `x` (normalized output) and `residual` (residual-out buffer).
@@ -356,29 +309,34 @@ class RMSNorm(CustomOp):
309 )
310 return x, residual
311
359 - add_residual = residual is not None
360 - if add_residual:
312 + if residual is not None:
313 return fused_add_rms_norm(
314 x, residual, self.weight.data, self.variance_epsilon
315 )
316 else:
365 - return rms_norm(x, self.weight.data, self.variance_epsilon)
317 + assert envs.VLLM_BATCH_INVARIANT
318 + return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)
319
320 def forward_hip(
321 self,
322 x: torch.Tensor,
323 residual: torch.Tensor | None = None,
324 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
325 + if residual is None and not envs.VLLM_BATCH_INVARIANT:
326 + return ir.ops.rms_norm(
327 + x, self.weight.data, self.variance_epsilon, self.variance_size_override
328 + )
329 +
330 if self.variance_size_override is not None:
331 return self.forward_native(x, residual)
332
375 - add_residual = residual is not None
376 - if add_residual:
333 + if residual is not None:
334 return self.rocm_norm_func_with_add(
335 x, residual, self.weight.data, self.variance_epsilon
336 )
337 else:
381 - return self.rocm_norm_func(x, self.weight.data, self.variance_epsilon)
338 + assert envs.VLLM_BATCH_INVARIANT
339 + return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)
340
341 def forward_xpu(
342 self,
vllm/platforms/cuda.py
+21
@@ -30,6 +30,7 @@ from .interface import DeviceCapability, Platform, PlatformEnum
30 if TYPE_CHECKING:
31 from vllm.config import VllmConfig
32 from vllm.config.cache import CacheDType
33 + from vllm.config.kernel import IrOpPriorityConfig
34 from vllm.v1.attention.selector import AttentionSelectorConfig
35 else:
36 VllmConfig = None
@@ -550,6 +551,26 @@ class CudaPlatformBase(Platform):
551 def use_custom_op_collectives(cls) -> bool:
552 return True
553
554 + @classmethod
555 + def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig:
556 + from vllm.config.compilation import CompilationMode
557 + from vllm.config.kernel import IrOpPriorityConfig
558 +
559 + # Native used by default when compiling,
560 + # use vllm_c kernels where available when no codegen
561 + cc = vllm_config.compilation_config
562 + using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
563 + default = ["native"] if using_inductor else ["vllm_c", "native"]
564 +
565 + # Use oink if enabled for rms_norm
566 + # TODO(Laurawly/luka): remove this env var,
567 + # users can just use IR op priority directly
568 + rms_norm = default
569 + if envs.VLLM_USE_OINK_OPS:
570 + rms_norm = ["oink"] + default
571 +
572 + return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
573 +
574
575 # NVML utils
576 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/interface.py
+11
@@ -17,6 +17,7 @@ if TYPE_CHECKING:
17 from torch.distributed import PrefixStore, ProcessGroup
18
19 from vllm.config import VllmConfig
20 + from vllm.config.kernel import IrOpPriorityConfig
21 from vllm.inputs import EngineInput
22 from vllm.pooling_params import PoolingParams
23 from vllm.sampling_params import SamplingParams
@@ -931,6 +932,16 @@ class Platform:
932 "num_compute_units is not implemented for the current platform."
933 )
934
935 + @classmethod
936 + def get_default_ir_op_priority(
937 + cls, vllm_config: "VllmConfig"
938 + ) -> "IrOpPriorityConfig":
939 + """Get the default IR op priority for the current platform."""
940 + from vllm.config.kernel import IrOpPriorityConfig
941 +
942 + # Native always used by default. Platforms can override this behavior.
943 + return IrOpPriorityConfig.with_default(["native"])
944 +
945
946 class UnspecifiedPlatform(Platform):
947 _enum = PlatformEnum.UNSPECIFIED
vllm/platforms/rocm.py
+30
@@ -19,6 +19,7 @@ from .interface import DeviceCapability, Platform, PlatformEnum
19
20 if TYPE_CHECKING:
21 from vllm.config import VllmConfig
22 + from vllm.config.kernel import IrOpPriorityConfig
23 from vllm.v1.attention.selector import AttentionSelectorConfig
24
25 logger = init_logger(__name__)
@@ -903,3 +904,32 @@ class RocmPlatform(Platform):
904 @classmethod
905 def use_custom_op_collectives(cls) -> bool:
906 return True
907 +
908 + @classmethod
909 + def get_default_ir_op_priority(
910 + cls, vllm_config: "VllmConfig"
911 + ) -> "IrOpPriorityConfig":
912 + from vllm.config.compilation import CompilationMode
913 + from vllm.config.kernel import IrOpPriorityConfig
914 +
915 + # Native used by default when compiling,
916 + # use vllm_c kernels where available when no codegen
917 + # TODO(luka/TJ) use aiter, vllm_c, native by default on ROCm
918 + cc = vllm_config.compilation_config
919 + using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
920 + default = ["native"] if using_inductor else ["vllm_c", "native"]
921 +
922 + # This (mostly) preserves previous CustomOp behavior
923 + # Necessary on ROCm because it's common that users
924 + # enable rms_norm to use the aiter kernel.
925 + # TODO(luka/TJ) remove env vars completely
926 + if (
927 + cc.is_custom_op_enabled("rms_norm")
928 + and envs.VLLM_ROCM_USE_AITER
929 + and envs.VLLM_ROCM_USE_AITER_RMSNORM
930 + ):
931 + rms_norm = ["aiter"] + default
932 + else:
933 + rms_norm = default
934 +
935 + return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
vllm/platforms/xpu.py
+16
@@ -21,6 +21,7 @@ from .interface import DeviceCapability, Platform, PlatformEnum
21
22 if TYPE_CHECKING:
23 from vllm.config import VllmConfig
24 + from vllm.config.kernel import IrOpPriorityConfig
25 from vllm.v1.attention.selector import AttentionSelectorConfig
26 else:
27 VllmConfig = None
@@ -257,6 +258,21 @@ class XPUPlatform(Platform):
258 )
259 return "vllm.distributed.device_communicators.xpu_communicator.XpuCommunicator" # noqa
260
261 + @classmethod
262 + def get_default_ir_op_priority(
263 + cls, vllm_config: "VllmConfig"
264 + ) -> "IrOpPriorityConfig":
265 + from vllm.config.compilation import CompilationMode
266 + from vllm.config.kernel import IrOpPriorityConfig
267 +
268 + # Native used by default when compiling,
269 + # use fused kernels where available when no codegen
270 + cc = vllm_config.compilation_config
271 + using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE
272 + default = ["native"] if using_inductor else ["xpu_kernels", "native"]
273 +
274 + return IrOpPriorityConfig.with_default(default)
275 +
276 @classmethod
277 def device_count(cls) -> int:
278 return torch.xpu.device_count()