diff --git a/.buildkite/test_areas/quantization.yaml b/.buildkite/test_areas/quantization.yaml index cdf8e7bbe7c7..490e20534f0f 100644 --- a/.buildkite/test_areas/quantization.yaml +++ b/.buildkite/test_areas/quantization.yaml @@ -11,6 +11,7 @@ steps: - csrc/ - vllm/model_executor/layers/quantization - vllm/models/deepseek_v41/quant_config.py + - vllm/utils/humming.py - tests/quantization commands: # Pin torchao to the version compatible with the CI PyTorch and CUDA stack. diff --git a/tests/quantization/test_humming_compile.py b/tests/quantization/test_humming_compile.py new file mode 100644 index 000000000000..d69973b68af2 --- /dev/null +++ b/tests/quantization/test_humming_compile.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.utils.import_utils import has_humming + +pytestmark = pytest.mark.skipif( + not (has_humming() and current_platform.is_cuda()), + reason="requires humming and CUDA", +) + + +def test_humming_input_quant_compiles_fullgraph(): + """Humming's forward helpers look up ``DeviceInfo`` on every call; vLLM + traces them with ``fullgraph=True`` (vllm-project/humming#147).""" + import vllm.utils.humming as hm + + weight_schema = hm.HummingWeightSchema(b_dtype="int4", weight_scale_group_size=128) + input_schema = hm.HummingInputSchema(a_dtype="int8") + from humming.transform import prepare_layer_config + + config = prepare_layer_config( + shape_n=256, + shape_k=256, + weight_schema=weight_schema, + input_schema=input_schema, + torch_dtype=torch.bfloat16, + device=torch.device("cuda"), + ) + x = torch.randn(16, 256, dtype=torch.bfloat16, device="cuda") + + def quant(t: torch.Tensor): + return hm.may_process_input(config, t) + + expected = quant(x) + actual = torch.compile(quant, fullgraph=True)(x) + for a, e in zip(actual, expected): + if e is None: + assert a is None + else: + torch.testing.assert_close(a, e) diff --git a/vllm/utils/humming.py b/vllm/utils/humming.py index 8004812ea488..0cd80a9be559 100644 --- a/vllm/utils/humming.py +++ b/vllm/utils/humming.py @@ -109,10 +109,70 @@ def create_humming_locks(device: torch.device | str) -> torch.Tensor: ) +_patched = False + + +def _patch_humming() -> None: + """Make humming's per-call ``DeviceInfo(device)`` lookups traceable. + + The forward path (``LayerConfig.check_device``, PDL heuristics) builds a + ``DeviceInfo`` on every call, which breaks vLLM's fullgraph compile: + Dynamo recurses forever in ``DeviceInfo.__getattr__`` + (vllm-project/humming#147) and, with that fixed (#148), still cannot + trace the extension's ``_DeviceInfo`` constructor. Hand out one cached + instance per device instead; ``prepare_layer_config`` creates it eagerly + during weight processing. + + TODO: remove after bumping humming to a release with + vllm-project/humming#149. + """ + global _patched + if _patched: + return + _patched = True + + import humming.config.config + import humming.transform + import humming.tune + import humming.tune.process_input + import humming.utils.smem + from humming.device import DeviceInfo, get_device_index + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError( + f"{type(self).__name__!r} object has no attribute {name!r}" + ) + return getattr(self._get_value(), name) + + DeviceInfo.__getattr__ = __getattr__ + + device_infos: dict[int, DeviceInfo] = {} + + def cached_device_info(index: int | torch.device | None = None) -> DeviceInfo: + if isinstance(index, torch.device) and index.type != "cuda": + return DeviceInfo(index) + device_index = get_device_index(index) + if device_index not in device_infos: + device_infos[device_index] = DeviceInfo(device_index) + return device_infos[device_index] + + for module in ( + humming.config.config, + humming.transform, + humming.tune, + humming.tune.process_input, + humming.utils.smem, + ): + if getattr(module, "DeviceInfo", None) is DeviceInfo: + module.DeviceInfo = cached_device_info + + def __getattr__(name: str) -> Any: spec = _EXPORTS.get(name) if spec is None: raise AttributeError(f"module 'vllm.utils.humming' has no attribute {name!r}") + _patch_humming() if ":" in spec: mod_path, attr = spec.split(":", 1) obj = getattr(importlib.import_module(mod_path), attr)