@tduncan / Nvidia-Optimization-Stack-VLLM / commits / b8baa31a28

Hardware-agnostic model definition via HF transformer backend (1/N) (#49458)

Signed-off-by: Thomas Ortner <boh@zurich.ibm.com>

Thomas Ortner committed Aug 13, 2026 at 12:21 UTC b8baa31a28650b6f3189bc7fcb051e1360c9270c
10 files changed +691 -5
tests/models/transformers/test_layer_registry.py new
+177
@@ -0,0 +1,177 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the Transformers backend's hw-agnostic layer resolution.
4 +
5 +`layers._resolve` imports a layer symbol from
6 +`vllm.model_executor.hw_agnostic.layers.<module>` when `VLLM_USE_HW_AGNOSTIC`
7 +is set and the symbol exists, and otherwise falls back to
8 +`vllm.model_executor.layers.<module>`. These tests pin that contract and the
9 +logging that reports which source was used.
10 +"""
11 +
12 +import logging
13 +import sys
14 +import types
15 +
16 +import pytest
17 +import torch
18 +
19 +from vllm.model_executor.models.transformers import layers
20 +
21 +HW_MODULE = "vllm.model_executor.hw_agnostic.layers.layernorm"
22 +
23 +
24 +@pytest.fixture
25 +def fake_hw_layernorm(monkeypatch):
26 + """Inject a hw-agnostic `layernorm` module exposing a sentinel `RMSNorm`.
27 +
28 + A `SimpleNamespace` stands in for the module: `importlib.import_module`
29 + returns it from `sys.modules` and `getattr` resolves `RMSNorm`, while its
30 + attributes are set at construction (no `ModuleType` attribute-set that mypy
31 + rejects, no constant `setattr` that ruff rejects)."""
32 + module = types.SimpleNamespace(RMSNorm=type("HwRMSNorm", (), {}))
33 + monkeypatch.setitem(sys.modules, HW_MODULE, module)
34 + return module
35 +
36 +
37 +def test_falls_back_to_vllm_when_disabled(monkeypatch, fake_hw_layernorm):
38 + """Disabled: the vLLM class is used even if a hw-agnostic one exists."""
39 + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
40 + from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
41 +
42 + assert layers._resolve("layernorm", "RMSNorm") is VllmRMSNorm
43 +
44 +
45 +def test_uses_hw_agnostic_when_enabled(monkeypatch, fake_hw_layernorm, caplog):
46 + """Enabled and available: the hw-agnostic class is used and logged."""
47 + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
48 + with caplog.at_level(logging.INFO):
49 + resolved = layers._resolve("layernorm", "RMSNorm")
50 + assert resolved is fake_hw_layernorm.RMSNorm
51 + assert "Using hw-agnostic layer: RMSNorm" in caplog.text
52 +
53 +
54 +def test_falls_back_when_symbol_missing(monkeypatch, caplog):
55 + """Enabled but the symbol is not ported: fall back to vLLM and warn."""
56 + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
57 + # A hw-agnostic module without the requested attribute triggers fallback.
58 + empty = types.ModuleType(HW_MODULE)
59 + monkeypatch.setitem(sys.modules, HW_MODULE, empty)
60 + from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm
61 +
62 + with caplog.at_level(logging.WARNING):
63 + resolved = layers._resolve("layernorm", "RMSNorm")
64 + assert resolved is VllmRMSNorm
65 + assert "falling back to default" in caplog.text
66 +
67 +
68 +def test_act_and_mul_falls_back_for_unknown_activation(
69 + monkeypatch, default_vllm_config
70 +):
71 + """An activation with no hw-agnostic equivalent falls back to vLLM's.
72 +
73 + `default_vllm_config` supplies the config context the CustomOp needs.
74 + """
75 + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
76 + from vllm.model_executor.layers.activation import GeluAndMul
77 +
78 + assert isinstance(layers.get_act_and_mul_fn("gelu"), GeluAndMul)
79 +
80 +
81 +@pytest.fixture(scope="module")
82 +def tiny_llama_path(tmp_path_factory):
83 + """A randomly-initialized microscopic Llama saved to disk (with an ungated
84 + tokenizer) so vLLM can load it like any local checkpoint."""
85 + from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
86 +
87 + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
88 + config = LlamaConfig(
89 + vocab_size=tokenizer.vocab_size,
90 + hidden_size=128,
91 + intermediate_size=256,
92 + num_hidden_layers=2,
93 + num_attention_heads=4,
94 + rms_norm_eps=1e-6,
95 + hidden_act="silu",
96 + )
97 + torch.manual_seed(0)
98 + model = LlamaForCausalLM(config)
99 +
100 + path = tmp_path_factory.mktemp("tiny_llama")
101 + model.save_pretrained(path)
102 + tokenizer.save_pretrained(path)
103 + return str(path)
104 +
105 +
106 +# Registered names of the layers the backend can
107 +# currently route to hw-agnostic implementations.
108 +_COVERED_LAYERS = ("rms_norm", "silu_and_mul")
109 +
110 +
111 +def _layer_providers(model) -> dict[str, str]:
112 + """Map each covered layer type present in the model to the provider its
113 + implementation came from (``hw_agnostic`` or ``vllm``).
114 + """
115 +
116 + def provider_of(module) -> str | None:
117 + for cls in type(module).__mro__:
118 + if "hw_agnostic.layers" in cls.__module__:
119 + return "hw_agnostic"
120 + if ".model_executor.layers." in cls.__module__:
121 + return "vllm"
122 + return None
123 +
124 + providers: dict[str, str] = {}
125 + for module in model.modules():
126 + name = getattr(module, "name", None)
127 + if (
128 + name in _COVERED_LAYERS
129 + and name not in providers
130 + and (prov := provider_of(module)) is not None
131 + ):
132 + providers[name] = prov
133 + return providers
134 +
135 +
136 +def _serve(vllm_runner, model_path, prompts):
137 + """Serve the model through the backend; return (layer_providers, logprobs)."""
138 + with vllm_runner(
139 + model_path,
140 + model_impl="transformers",
141 + max_model_len=64,
142 + enforce_eager=True,
143 + gpu_memory_utilization=0.3,
144 + ) as runner:
145 + assert runner.llm.llm_engine.model_config.using_transformers_backend()
146 + providers = runner.apply_model(_layer_providers)[0]
147 + outputs = runner.generate_greedy_logprobs(
148 + prompts, max_tokens=32, num_logprobs=5
149 + )
150 + return providers, outputs
151 +
152 +
153 +def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llama_path):
154 + """Serving the tiny model with hw-agnostic layers matches the vLLM baseline."""
155 + # spawn: worker re-imports layers with the env set (see docstring).
156 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
157 + # apply_model pickles the introspection function.
158 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
159 + from ..utils import check_logprobs_close
160 +
161 + prompts = ["The capital of France is", "vLLM is"]
162 +
163 + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
164 + vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts)
165 + # Both replaceable layers present in a Llama block must be vLLM's here.
166 + assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"}
167 +
168 + monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
169 + hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts)
170 + assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"}
171 +
172 + check_logprobs_close(
173 + outputs_0_lst=vllm_outputs,
174 + outputs_1_lst=hw_outputs,
175 + name_0="vllm",
176 + name_1="hw_agnostic",
177 + )
vllm/envs.py
+5
@@ -125,6 +125,7 @@ if TYPE_CHECKING:
125 VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
126 VLLM_SKIP_P2P_CHECK: bool = False
127 VLLM_DISABLED_KERNELS: list[str] = []
128 + VLLM_USE_HW_AGNOSTIC: bool = False
129 VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE: bool = True
130 VLLM_DISABLE_PYNCCL: bool = False
131 VLLM_USE_OINK_OPS: bool = False
@@ -1200,6 +1201,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1201 if "VLLM_DISABLED_KERNELS" not in os.environ
1202 else os.environ["VLLM_DISABLED_KERNELS"].split(",")
1203 ),
1204 + # Selects hw-agnostic layers for HF transformer backend
1205 + "VLLM_USE_HW_AGNOSTIC": lambda: (
1206 + os.getenv("VLLM_USE_HW_AGNOSTIC", "False").lower() in ("true", "1")
1207 + ),
1208 "VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE": lambda: bool(
1209 int(os.getenv("VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE", "1"))
1210 ),
vllm/model_executor/hw_agnostic/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/hw_agnostic/custom_op.py new
+318
@@ -0,0 +1,318 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import functools
4 +import inspect
5 +
6 +import torch
7 +import torch.nn as nn
8 +
9 +from vllm.config import get_cached_compilation_config
10 +from vllm.logger import init_logger
11 +from vllm.model_executor.utils import maybe_disable_graph_partition
12 +from vllm.platforms import current_platform
13 +
14 +logger = init_logger(__name__)
15 +
16 +# Dictionary of all custom ops (classes, indexed by registered name).
17 +# To check if an op with a name is enabled, call .enabled() on the class.
18 +# Examples:
19 +# - MyOp.enabled()
20 +# - op_registry["my_op"].enabled()
21 +op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
22 +op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
23 +
24 +
25 +def maybe_get_oot_by_class(class_type: type) -> type:
26 + class_name = class_type.__name__
27 + if class_name in op_registry_oot:
28 + return op_registry_oot[class_name]
29 + return class_type
30 +
31 +
32 +class PluggableLayer(nn.Module):
33 + """
34 + Base class for pluggable layers.
35 +
36 + A PluggableLayer is a *module-composing* abstraction: it may instantiate other
37 + ``torch.nn.Module`` objects as sub-layers, and its functionality depends on
38 + these sub-layers following a generalized invocation sequence. Also, it is stateful
39 + and may hold parameters or buffers.
40 +
41 + Unlike :class:`CustomOp`, PluggableLayer does NOT provide per-platform
42 + ``forward_*`` dispatch. Instead, it supports out-of-tree (OOT) replacement
43 + of the entire layer class at instantiation time, allowing customized
44 + initialization and submodule composition.
45 + """
46 +
47 + def __new__(cls, *args, **kwargs):
48 + try:
49 + layer_class_name = cls.__name__
50 + except AttributeError:
51 + raise TypeError(
52 + f"Cannot instantiate '{cls.__name__}': its 'name' attribute "
53 + f"was not set, possibly because it was not decorated with "
54 + f"@PluggableLayer.register, or it's the PluggableLayer itself."
55 + ) from None
56 +
57 + if layer_class_name not in op_registry_oot:
58 + layer_cls_to_instantiate = cls
59 + else:
60 + layer_cls_to_instantiate = op_registry_oot[layer_class_name]
61 + logger.debug(
62 + "Instantiating pluggable layer: %s using %s",
63 + layer_class_name,
64 + str(layer_cls_to_instantiate),
65 + )
66 + return super().__new__(layer_cls_to_instantiate)
67 +
68 + # Decorator to register pluggable layers.
69 + @classmethod
70 + def register(cls, name: str):
71 + def decorator(op_cls):
72 + assert name not in op_registry, f"Duplicate op name: {name}"
73 + op_cls.name = name
74 + op_registry[name] = op_cls
75 + return op_cls
76 +
77 + return decorator
78 +
79 + # Decorator to register out-of-tree(oot) pluggable layers.
80 + # For OOT pluggable layers:
81 + # if in-tree layer class is registered with an oot_custom_layer,
82 + # the oot_custom_layer will be used instead.
83 + @classmethod
84 + def register_oot(cls, _decorated_layer_cls=None, name: str | None = None):
85 + def decorator(layer_cls):
86 + reg_name = name if name is not None else cls.__name__
87 + assert reg_name not in op_registry_oot, f"Duplicate layer name: {reg_name}"
88 + layer_cls.name = reg_name
89 + op_registry_oot[reg_name] = layer_cls
90 + return layer_cls
91 +
92 + if _decorated_layer_cls is None:
93 + # Called with parentheses: @PluggableLayer.register_oot()
94 + # or @PluggableLayer.register_oot(name="...")
95 + return decorator
96 + elif isinstance(_decorated_layer_cls, type): # Check if it's a class
97 + # Called without parentheses: @PluggableLayer.register_oot
98 + return decorator(_decorated_layer_cls)
99 + else:
100 + raise TypeError("Decorator can only be applied to classes.")
101 +
102 +
103 +class CustomOp(nn.Module):
104 + """
105 + Base class for custom ops.
106 + Dispatches the forward method to the appropriate backend.
107 + """
108 +
109 + def __new__(cls, *args, **kwargs):
110 + try:
111 + op_name = cls.__name__
112 + except AttributeError:
113 + raise TypeError(
114 + f"Cannot instantiate '{cls.__name__}': its 'name' attribute "
115 + f"was not set, possibly because it was not decorated with "
116 + f"@CustomOp.register, or it's the CustomOp base class itself."
117 + ) from None
118 +
119 + if op_name not in op_registry_oot:
120 + op_cls_to_instantiate = cls
121 + else:
122 + op_cls_to_instantiate = op_registry_oot[op_name]
123 + logger.debug(
124 + "Instantiating custom op: %s using %s",
125 + op_name,
126 + str(op_cls_to_instantiate),
127 + )
128 + return super().__new__(op_cls_to_instantiate)
129 +
130 + def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False):
131 + super().__init__()
132 + self._enforce_enable = enforce_enable
133 + self._forward_method = self.dispatch_forward(compile_native=compile_native)
134 +
135 + def forward(self, *args, **kwargs):
136 + return self._forward_method(*args, **kwargs)
137 +
138 + def forward_native(self, *args, **kwargs):
139 + """PyTorch-native implementation; OOT plugins override via
140 + :meth:`forward_oot`."""
141 + raise NotImplementedError
142 +
143 + def forward_oot(self, *args, **kwargs):
144 + # By default, OOT ops fall back to the PyTorch-native implementation.
145 + return self.forward_native(*args, **kwargs)
146 +
147 + def dispatch_forward(self, compile_native: bool):
148 + # The hw-agnostic CustomOp dispatches only between forward_native
149 + # (in-tree, used on every platform) and forward_oot (overridden by
150 + # an OOT plugin via register_oot).
151 + compilation_config = get_cached_compilation_config()
152 +
153 + # NOTE(shen-shanshan): CustomOp object can be enforce enabled, e.g.,
154 + # enable device-specific kernels in ViT models when enabling graph
155 + # mode. By default, it will follow the compilation_config to determine
156 + # whether enable itself.
157 + # This enforce_enable mechanism will be removed after we adding a
158 + # separate compilation_config for multi-modal part.
159 + enabled = self._enforce_enable or self.enabled()
160 + if enabled:
161 + compilation_config.enabled_custom_ops.update([self.__class__.name])
162 + else:
163 + compilation_config.disabled_custom_ops.update([self.__class__.name])
164 +
165 + if not enabled:
166 + # Compile forward_native to avoid eager torch ops if inside
167 + # opaque torch custom op (e.g. fused_moe, unified_attention, etc.)
168 + return self.maybe_compile(self.forward_native, enable=compile_native)
169 +
170 + if current_platform.is_out_of_tree():
171 + return self.forward_oot
172 + return self.forward_native
173 +
174 + def maybe_compile(self, fn, *, enable: bool = True):
175 + """
176 + Compile fn if compilation enabled.
177 + Useful for CustomOp instances called from within a torch custom op,
178 + meaning the forward call is hidden from the model-level torch.compile.
179 +
180 + NOTE: this does not enable fusion across ops, so opaque custom ops
181 + should still be unwrapped wherever possible.
182 + """
183 + from vllm.config.compilation import CompilationMode
184 +
185 + # Do not compile if compilation disabled
186 + if not enable:
187 + return fn
188 +
189 + # Do not compile if global compilation disabled
190 + compilation_config = get_cached_compilation_config()
191 + if compilation_config.mode == CompilationMode.NONE:
192 + return fn
193 +
194 + # If eager backend is used, do not compile either
195 + if compilation_config.backend == "eager":
196 + return fn
197 +
198 + compile_options = maybe_disable_graph_partition(
199 + current_platform.simple_compile_backend
200 + )
201 + backend = current_platform.simple_compile_backend
202 +
203 + dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None)
204 + if dynamic_arg_dims is not None:
205 + compiled_fn = torch.compile(
206 + fn,
207 + dynamic=False,
208 + backend=backend,
209 + options=compile_options,
210 + )
211 + sig = inspect.signature(fn)
212 +
213 + @functools.wraps(fn)
214 + def wrapper(*args, **kwargs):
215 + bound = sig.bind(*args, **kwargs)
216 + bound.apply_defaults()
217 + for name, dims in dynamic_arg_dims.items():
218 + arg = bound.arguments.get(name)
219 + if arg is not None and isinstance(arg, torch.Tensor):
220 + dims_list = [dims] if isinstance(dims, int) else dims
221 + for d in dims_list:
222 + real_d = arg.ndim + d if d < 0 else d
223 + torch._dynamo.mark_dynamic(arg, real_d)
224 + return compiled_fn(*args, **kwargs)
225 +
226 + return wrapper
227 +
228 + # dynamic=True to avoid recompilations
229 + return torch.compile(
230 + fn,
231 + dynamic=True,
232 + backend=backend,
233 + options=compile_options,
234 + )
235 +
236 + @classmethod
237 + def enabled(cls) -> bool:
238 + # if no name, then it was not registered
239 + compilation_config = get_cached_compilation_config()
240 + custom_ops = compilation_config.custom_ops
241 + if not hasattr(cls, "name"):
242 + logger.warning_once(
243 + "Custom op %s was not registered, which means it won't appear "
244 + "in the op registry. It will be enabled/disabled based on the "
245 + "global settings.",
246 + cls.__name__,
247 + )
248 + return CustomOp.default_on()
249 +
250 + enabled = f"+{cls.name}" in custom_ops
251 + disabled = f"-{cls.name}" in custom_ops
252 + assert not (enabled and disabled), f"Cannot enable and disable {cls.name}"
253 +
254 + return (CustomOp.default_on() or enabled) and not disabled
255 +
256 + @staticmethod
257 + def default_on() -> bool:
258 + """
259 + Behavior controlled by `CompilationConfig.custom_ops`: On by default if
260 + 'all', off by default if 'none'.
261 + When PyTorch Inductor is used, 'none' is the default value,
262 + otherwise 'all'.
263 + """
264 + compilation_config = get_cached_compilation_config()
265 + count_none = compilation_config.custom_ops.count("none")
266 + count_all = compilation_config.custom_ops.count("all")
267 + assert count_none + count_all == 1
268 +
269 + return not count_none > 0 or count_all > 0
270 +
271 + # Decorator to register custom ops.
272 + @classmethod
273 + def register(
274 + cls,
275 + name: str,
276 + dynamic_arg_dims: dict[str, int | list[int]] | None = None,
277 + ):
278 + def decorator(op_cls):
279 + assert name not in op_registry, f"Duplicate op name: {name}"
280 + op_cls.name = name
281 + op_cls._dynamic_arg_dims = dynamic_arg_dims
282 + op_registry[name] = op_cls
283 + return op_cls
284 +
285 + return decorator
286 +
287 + # Decorator to register out-of-tree(oot) custom ops.
288 + # For OOT custom ops:
289 + # if in-tree layer class is registered with an oot_custom_op layer,
290 + # the oot_custom_op layer will be used instead.
291 + # Example:
292 + # - @UnquantizedFusedMoEMethod.register_oot
293 + # class HPUUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod)
294 + # or
295 + # - @CustomOP.register_oot(name="UnquantizedFusedMoEMethod")
296 + @classmethod
297 + def register_oot(cls, _decorated_op_cls=None, name: str | None = None):
298 + def decorator(op_cls):
299 + reg_name = name if name is not None else cls.__name__
300 + assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}"
301 + op_cls.name = reg_name
302 + op_registry_oot[reg_name] = op_cls
303 + return op_cls
304 +
305 + if _decorated_op_cls is None:
306 + # Called with parentheses: @CustomOP.register_oot()
307 + # or @CustomOP.register_oot(name="...")
308 + # So, _decorated_op_cls is None.
309 + # We return the actual decorator function.
310 + return decorator
311 + elif isinstance(_decorated_op_cls, type): # Check if it's a class
312 + # Called without parentheses: @CustomOP.register_oot
313 + # The first argument is the class itself.
314 + # We call the 'decorator' function immediately with the class.
315 + return decorator(_decorated_op_cls)
316 + else:
317 + # Handle other unexpected cases if necessary
318 + raise TypeError("Decorator can only be applied to classes.")
vllm/model_executor/hw_agnostic/layers/__init__.py new
+2
@@ -0,0 +1,2 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/hw_agnostic/layers/activation.py new
+27
@@ -0,0 +1,27 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +import torch.nn.functional as F
5 +
6 +from vllm.model_executor.hw_agnostic.custom_op import CustomOp
7 +
8 +
9 +@CustomOp.register("silu_and_mul")
10 +class SiluAndMul(CustomOp):
11 + """SwiGLU: ``x -> silu(x[:d]) * x[d:]`` where ``d = x.shape[-1] // 2``."""
12 +
13 + def forward_native(self, x: torch.Tensor) -> torch.Tensor:
14 + d = x.shape[-1] // 2
15 + return F.silu(x[..., :d]) * x[..., d:]
16 +
17 +
18 +# Activation-and-mul ops keyed by HF activation name.
19 +_ACTIVATION_AND_MUL_REGISTRY = {
20 + "silu": SiluAndMul,
21 + "swish": SiluAndMul,
22 +}
23 +
24 +
25 +def get_act_and_mul_fn(act_fn_name: str) -> CustomOp:
26 + """Build the hw-agnostic activation-and-mul op named `act_fn_name`."""
27 + return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name.lower()]()
vllm/model_executor/hw_agnostic/layers/layernorm.py new
+90
@@ -0,0 +1,90 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +import torch.nn as nn
5 +from torch import Tensor
6 +
7 +from vllm.model_executor.hw_agnostic.custom_op import CustomOp
8 +
9 +
10 +@CustomOp.register("rms_norm")
11 +class RMSNorm(CustomOp):
12 + """``x -> w * x / sqrt(E[x^2] + eps)``. With ``residual``, fuses
13 + ``residual += x`` then RMSNorm and returns ``(normalized, residual)``."""
14 +
15 + def __init__(
16 + self,
17 + hidden_size: int,
18 + eps: float = 1e-6,
19 + var_hidden_size: int | None = None,
20 + has_weight: bool = True,
21 + dtype: torch.dtype | None = None,
22 + ) -> None:
23 + super().__init__()
24 + self.hidden_size = hidden_size
25 + self.variance_epsilon = eps
26 + self.variance_size_override = (
27 + None if var_hidden_size == hidden_size else var_hidden_size
28 + )
29 + weight_dtype = dtype or torch.get_default_dtype()
30 + self.has_weight = has_weight
31 + weight = torch.ones(hidden_size, dtype=weight_dtype)
32 + if has_weight:
33 + self.weight = nn.Parameter(weight)
34 + else:
35 + self.register_buffer("weight", weight, persistent=False)
36 +
37 + def _rms_norm(
38 + self,
39 + x: Tensor,
40 + weight: Tensor | None,
41 + epsilon: float,
42 + variance_size: int | None = None,
43 + ) -> Tensor:
44 + """Weighted root-mean-square layer normalization"""
45 + orig_dtype = x.dtype
46 + x = x.to(torch.float32)
47 + x_var = x if variance_size is None else x[..., :variance_size]
48 + variance = x_var.pow(2).mean(dim=-1, keepdim=True)
49 + x = x * torch.rsqrt(variance + epsilon)
50 + if weight is not None:
51 + x = x.to(weight.dtype) * weight
52 + return x.to(orig_dtype)
53 +
54 + def _fused_add_rms_norm(
55 + self,
56 + x: Tensor,
57 + x_residual: Tensor,
58 + weight: Tensor | None,
59 + epsilon: float,
60 + variance_size: int | None = None,
61 + ) -> tuple[Tensor, Tensor]:
62 + """Fused add and weighted root-mean-square layer normalization"""
63 + orig_dtype = x.dtype
64 + x = x.to(torch.float32)
65 + x = x + x_residual.to(torch.float32)
66 + x_residual = x.to(orig_dtype)
67 +
68 + x_var = x if variance_size is None else x[..., :variance_size]
69 + variance = x_var.pow(2).mean(dim=-1, keepdim=True)
70 + x = x * torch.rsqrt(variance + epsilon)
71 + if weight is not None:
72 + x = x.to(weight.dtype) * weight
73 + return x.to(orig_dtype), x_residual
74 +
75 + def forward_native(
76 + self,
77 + x: torch.Tensor,
78 + residual: torch.Tensor | None = None,
79 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
80 + weight = self.weight if self.has_weight else None
81 + epsilon = self.variance_epsilon
82 + variance_size = self.variance_size_override
83 +
84 + if residual is None:
85 + return self._rms_norm(x, weight, epsilon, variance_size)
86 + else:
87 + return self._fused_add_rms_norm(x, residual, weight, epsilon, variance_size)
88 +
89 + def extra_repr(self) -> str:
90 + return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}"
vllm/model_executor/models/transformers/fusers/glu.py
+2 -4
@@ -11,10 +11,7 @@ from torch import fx, nn
11 from transformers.activations import ACT2CLS
12
13 from vllm.logger import init_logger
14 -from vllm.model_executor.layers.activation import (
15 - _ACTIVATION_AND_MUL_REGISTRY,
16 - get_act_and_mul_fn,
17 -)
14 +from vllm.model_executor.layers.activation import _ACTIVATION_AND_MUL_REGISTRY
15 from vllm.model_executor.layers.linear import MergedColumnParallelLinear
16 from vllm.model_executor.models.transformers.fusers.base import StackedFuser
17 from vllm.model_executor.models.transformers.fx_utils import (
@@ -26,6 +23,7 @@ from vllm.model_executor.models.transformers.fx_utils import (
23 replace_expr,
24 single_self_call,
25 )
26 +from vllm.model_executor.models.transformers.layers import get_act_and_mul_fn
27 from vllm.model_executor.models.transformers.utils import (
28 log_replacement,
29 replace_linear_class,
vllm/model_executor/models/transformers/fusers/rms_norm.py
+4 -1
@@ -15,7 +15,6 @@ from vllm.distributed import (
15 )
16 from vllm.distributed.parallel_state import model_parallel_is_initialized
17 from vllm.distributed.utils import split_tensor_along_last_dim
18 -from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
18 from vllm.model_executor.models.transformers.fusers.base import BaseFuser
19 from vllm.model_executor.models.transformers.fx_utils import (
20 find_node,
@@ -25,6 +24,10 @@ from vllm.model_executor.models.transformers.fx_utils import (
24 peel,
25 trace,
26 )
27 +from vllm.model_executor.models.transformers.layers import (
28 + GemmaRMSNorm,
29 + RMSNorm,
30 +)
31
32 if TYPE_CHECKING:
33 from vllm.config import VllmConfig
vllm/model_executor/models/transformers/layers.py new
+64
@@ -0,0 +1,64 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Layer provider resolution for the Transformers modeling backend.
4 +
5 +When ``VLLM_USE_HW_AGNOSTIC`` is set, layer symbols are imported from
6 +``vllm.model_executor.hw_agnostic.layers.<module>``, falling back to
7 +``vllm.model_executor.layers.<module>`` for anything not yet ported. The
8 +resolved source of every symbol is logged so it is clear which layers run
9 +hw-agnostic and which fell back to vLLM.
10 +"""
11 +
12 +import importlib
13 +
14 +import vllm.envs as envs
15 +from vllm.logger import init_logger
16 +
17 +logger = init_logger(__name__)
18 +
19 +_HW_PKG = "vllm.model_executor.hw_agnostic.layers"
20 +_VLLM_PKG = "vllm.model_executor.layers"
21 +
22 +
23 +def _resolve(module: str, name: str):
24 + """Return `name` from the hw-agnostic `module` when enabled and available,
25 + else from vLLM. Logs which source was used."""
26 + if envs.VLLM_USE_HW_AGNOSTIC:
27 + try:
28 + obj = getattr(importlib.import_module(f"{_HW_PKG}.{module}"), name)
29 + logger.info("Using hw-agnostic layer: %s", name)
30 + return obj
31 + except (ImportError, AttributeError):
32 + logger.warning(
33 + "hw-agnostic layer %s is not available; falling back to default", name
34 + )
35 + return getattr(importlib.import_module(f"{_VLLM_PKG}.{module}"), name)
36 +
37 +
38 +RMSNorm = _resolve("layernorm", "RMSNorm")
39 +GemmaRMSNorm = _resolve("layernorm", "GemmaRMSNorm")
40 +
41 +
42 +def get_act_and_mul_fn(act_fn_name: str):
43 + """Fused activation-and-mul op for `act_fn_name`, preferring hw-agnostic.
44 +
45 + Resolved per call because the op is name-parameterized: an activation with
46 + no hw-agnostic equivalent falls back to vLLM individually.
47 + """
48 + if envs.VLLM_USE_HW_AGNOSTIC:
49 + try:
50 + from vllm.model_executor.hw_agnostic.layers.activation import (
51 + get_act_and_mul_fn as hw_fn,
52 + )
53 +
54 + fn = hw_fn(act_fn_name)
55 + logger.info_once("Using hw-agnostic activation: %s", act_fn_name)
56 + return fn
57 + except (ImportError, KeyError):
58 + logger.warning_once(
59 + "hw-agnostic activation %s is not available; falling back to vLLM",
60 + act_fn_name,
61 + )
62 + from vllm.model_executor.layers.activation import get_act_and_mul_fn as vllm_fn
63 +
64 + return vllm_fn(act_fn_name)