[Bugfix] Make Humming DeviceInfo lookups traceable under torch.compile (#61858)

Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

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

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)

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)