| 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) |