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 3, 14:04 UTC
[Perf] Set IR Op Priority Once at Worker Init (#42631)
Signed-off-by: BadrBasowid <badr.basowid@gmail.com>
BadrBasowid committed
May 15, 2026 at 23:56 UTC
fb5bd03f51f986bf263c299abe3bc4e22407ab33
6 files changed
+129
-53
tests/ir/test_op.py
+45
@@ -272,6 +272,31 @@ class TestIrOpImplDispatch:
272
# Restored to empty
273
assert _custom_add.get_priority() == []
274
275
+ @pytest.mark.parametrize(
276
+ "default,override",
277
+ [
278
+ (["impl_even", "impl_b"], ["impl_a"]),
279
+ (["impl_a"], ["impl_even", "impl_b"]),
280
+ ],
281
+ )
282
+ def test_set_default_priority(
283
+ self, custom_add_op, default: list[str], override: list[str]
284
+ ):
285
+ _custom_add = custom_add_op
286
+ assert _custom_add.get_priority() == []
287
+
288
+ _custom_add.set_default(default)
289
+ assert _custom_add.get_priority() == default
290
+
291
+ # Priority doesn't change after exiting the set_priority context.
292
+ with _custom_add.set_priority(override):
293
+ assert _custom_add.get_priority() == override
294
+ assert _custom_add.get_priority() == default
295
+
296
+ # Should override the previous default.
297
+ _custom_add.set_default(override)
298
+ assert _custom_add.get_priority() == override
299
+
300
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
301
def test_dispatch_priority_order(self, custom_add_op, overload: str):
302
_custom_add = custom_add_op
@@ -381,6 +406,26 @@ class TestIrOpImplDispatch:
406
assert "priority not set" in message
407
408
409
+@pytest.mark.parametrize("default", [True, False])
410
+def test_set_default_torch_wrap(default: bool):
411
+ """set_default_torch_wrap permanently flips the global flag."""
412
+ original = vllm.ir.op._ENABLE_TORCH_WRAP
413
+ try:
414
+ vllm.ir.set_default_torch_wrap(default)
415
+ assert vllm.ir.op._ENABLE_TORCH_WRAP is default
416
+
417
+ # Flag doesn't change after exiting the enable_torch_wrap context.
418
+ with vllm.ir.enable_torch_wrap(not default):
419
+ assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default)
420
+ assert vllm.ir.op._ENABLE_TORCH_WRAP is default
421
+
422
+ # Should override the previous default.
423
+ vllm.ir.set_default_torch_wrap(not default)
424
+ assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default)
425
+ finally:
426
+ vllm.ir.op._ENABLE_TORCH_WRAP = original
427
+
428
+
429
@pytest.fixture
430
def custom_mm_op(fake_vllm_ir):
431
"""Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``)."""
vllm/config/kernel.py
+28
-17
@@ -20,8 +20,8 @@ logger = init_logger(__name__)
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.
23
+ forward pass. Each member is a list of strings, which will be installed
24
+ in worker init via vllm.ir.ops.<op_name>.set_default().
25
A single comma-separated string is accepted as well,
26
27
If specified manually, platform defaults will be appended to the lists.
@@ -67,30 +67,41 @@ class IrOpPriorityConfig:
67
assert all(isinstance(v, str) for v in value)
68
return value
69
70
- @contextlib.contextmanager
71
- def set_priority(self):
70
+ def _iter_op_priorities(self):
71
"""
73
- Context manager to set the IR op priority for all op members.
74
- It also imports IR kernel implementations for the current platform
75
- to ensure all implementations are made available.
72
+ Yield (IrOp, priority_list) for each field, after importing platform
73
+ kernels and validating each entry.
74
"""
75
from vllm.ir.op import IrOp
76
from vllm.platforms import current_platform
77
78
current_platform.import_ir_kernels()
79
80
+ for field in fields(self): # type: ignore[arg-type]
81
+ op_priority = getattr(self, field.name)
82
+ assert op_priority is not None, (
83
+ f"IR op priority for {field.name} must be set"
84
+ )
85
+ logger.debug("Setting IR op priority for %s to %s", field.name, op_priority)
86
+ yield IrOp.registry[field.name], op_priority
87
+
88
+ def set_default(self) -> None:
89
+ """
90
+ Permanently set the IR op priority for all op members.
91
+ """
92
+ for ir_op, op_priority in self._iter_op_priorities():
93
+ ir_op.set_default(op_priority)
94
+
95
+ @contextlib.contextmanager
96
+ def set_priority(self):
97
+ """
98
+ Context manager to set the IR op priority for all op members.
99
+ It also imports IR kernel implementations for the current platform
100
+ to ensure all implementations are made available.
101
+ """
102
with contextlib.ExitStack() as stack:
83
- for field in fields(self): # type: ignore[arg-type]
84
- op_priority = getattr(self, field.name)
85
- assert op_priority is not None, (
86
- f"IR op priority for {field.name} must be set"
87
- )
88
- logger.debug(
89
- "Setting IR op priority for %s to %s", field.name, op_priority
90
- )
91
- ir_op = IrOp.registry[field.name]
103
+ for ir_op, op_priority in self._iter_op_priorities():
104
stack.enter_context(ir_op.set_priority(op_priority))
93
-
105
yield
106
107
@classmethod
vllm/forward_context.py
+1
-8
@@ -10,7 +10,6 @@ from typing import Any
10
import torch
11
12
import vllm.envs as envs
13
-import vllm.ir
13
from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
@@ -320,13 +319,7 @@ def set_forward_context(
319
)
320
321
try:
323
- with (
324
- override_forward_context(forward_context),
325
- vllm_config.kernel_config.ir_op_priority.set_priority(),
326
- vllm.ir.enable_torch_wrap(
327
- vllm_config.compilation_config.ir_enable_torch_wrap
328
- ),
329
- ):
322
+ with override_forward_context(forward_context):
323
yield
324
finally:
325
global last_logging_time, batchsize_logging_interval
vllm/ir/__init__.py
+2
-2
@@ -1,6 +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
4
+from .op import enable_torch_wrap, register_op, set_default_torch_wrap
5
6
-__all__ = ["enable_torch_wrap", "register_op", "ops"]
6
+__all__ = ["enable_torch_wrap", "register_op", "set_default_torch_wrap", "ops"]
vllm/ir/op.py
+45
-26
@@ -51,6 +51,14 @@ _ENABLE_TORCH_WRAP: bool = True
51
"""Global override flag to control torch op layer wrapping."""
52
53
54
+def set_default_torch_wrap(enable: bool = True) -> None:
55
+ """
56
+ Permanently set the torch wrap flag.
57
+ """
58
+ global _ENABLE_TORCH_WRAP
59
+ _ENABLE_TORCH_WRAP = enable
60
+
61
+
62
@contextlib.contextmanager
63
def enable_torch_wrap(enable: bool = True):
64
"""
@@ -372,42 +380,53 @@ class IrOp:
380
"""Get the current dispatch priority for implementations for this op."""
381
return [p.provider for p in self._priority_impls]
382
375
- @contextlib.contextmanager
376
- def set_priority(self, priority: list[str]):
377
- """
378
- Context manager to set the dispatch priority for implementations for this op.
379
- """
383
+ def _filter_priority_impls(self, priority: list[str]) -> list["IrOpImpl"]:
384
assert all(p in self.impls for p in priority), (
385
"All providers in priority must be registered implementations."
386
)
387
+ filtered_impls: list[IrOpImpl] = []
388
+ for p in priority:
389
+ impl = self.impls[p]
390
+ if not impl.supported:
391
+ # Skip unsupported implementations
392
+ continue
393
384
- def filter_priority_impls(p_list: list[str]) -> list[IrOpImpl]:
385
- filtered_impls = []
386
- for p in p_list:
387
- impl = self.impls[p]
388
- if not impl.supported:
389
- # Skip unsupported implementations
390
- continue
394
+ filtered_impls.append(impl)
395
392
- filtered_impls.append(impl)
396
+ # If all args are supported, skip other implementations
397
+ if impl.supports_all_args:
398
+ return filtered_impls
399
394
- # If all args are supported, skip other implementations
395
- if impl.supports_all_args:
396
- return filtered_impls
400
+ logger.warning_once(
401
+ "Op %s: No implementation in priority list supports all args, "
402
+ "execution fallback to native is possible. To silence this warning, "
403
+ "explicitly add 'native' to the end of the priority list",
404
+ self.name,
405
+ )
406
+ filtered_impls.append(self.impls["native"])
407
+ return filtered_impls
408
398
- logger.warning_once(
399
- "Op %s: No implementation in priority list supports all args, "
400
- "execution fallback to native is possible. To silence this warning, "
401
- "explicitly add 'native' to the end of the priority list",
402
- self.name,
403
- )
404
- filtered_impls.append(self.impls["native"])
405
- return filtered_impls
409
+ def set_default(self, priority: list[str]) -> None:
410
+ """
411
+ Permanently set the dispatch priority for this op. Use this for
412
+ process-lifetime setup (e.g., worker startup). For scoped overrides,
413
+ use ``set_priority`` instead.
414
+ """
415
+ self._priority_impls = self._filter_priority_impls(priority)
416
+ logger.debug(
417
+ "Priority for vllm.ir.%s set to %s",
418
+ self.name,
419
+ lazy(lambda: [p.provider for p in self._priority_impls]),
420
+ )
421
407
- # Temporarily set priority
422
+ @contextlib.contextmanager
423
+ def set_priority(self, priority: list[str]):
424
+ """
425
+ Context manager to set the dispatch priority for implementations for this op.
426
+ """
427
old_priority_impls = self._priority_impls
428
try:
410
- self._priority_impls = filter_priority_impls(priority)
429
+ self._priority_impls = self._filter_priority_impls(priority)
430
logger.debug(
431
"Priority for vllm.ir.%s set to %s",
432
self.name,
vllm/v1/worker/worker_base.py
+8
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
7
import torch
8
import torch.nn as nn
9
10
+import vllm.ir
11
from vllm.config import VllmConfig, set_current_vllm_config
12
from vllm.logger import init_logger
13
from vllm.lora.request import LoRARequest
@@ -87,6 +88,13 @@ class WorkerBase:
88
self.device: torch.device | None = None
89
self.model_runner: nn.Module | None = None
90
91
+ # IR op priority and torch-wrap state are constant for the worker's
92
+ # lifetime.
93
+ vllm_config.kernel_config.ir_op_priority.set_default()
94
+ vllm.ir.set_default_torch_wrap(
95
+ vllm_config.compilation_config.ir_enable_torch_wrap
96
+ )
97
+
98
def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
99
"""Get specifications for KV cache implementation."""
100
raise NotImplementedError