fix(gguf): avoid device over-reservation for host decoding

invokeai/backend/model_manager/load/model_cache/model_cache.py

@@ -34,7 +34,7 @@
 )
 from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
 from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType
-from invokeai.backend.util.devices import TorchDevice
+from invokeai.backend.util.devices import TorchDevice, device_has_dedicated_vram
 from invokeai.backend.util.level_zero import xpu_device_is_integrated
 from invokeai.backend.util.logging import InvokeAILogger
 from invokeai.backend.util.prefix_logger_adapter import PrefixedLoggerAdapter
@@ -507,25 +507,10 @@ def __call__(
 _MIN_RELEASABLE_FOR_RESERVE = 256 * MB
 
 
-def _has_dedicated_vram(device: torch.device) -> bool:
-    """Whether a device has its own VRAM rather than sharing system memory with the CPU.
-
-    Intel integrated GPUs share memory with the CPU exactly as MPS does, so budgeting them as
-    if they had dedicated VRAM double-counts the same DRAM (once as "VRAM", once as RAM). When
-    the Level Zero probe cannot determine the answer, the previous discrete-GPU behaviour is
-    kept rather than guessing.
-    """
-    if device.type == "cuda":
-        return True
-    if device.type == "xpu":
-        return xpu_device_is_integrated(device) is not True
-    return False
-
-
 def _is_integrated_xpu(device: torch.device) -> bool:
     """Whether ``device`` is an Intel integrated GPU, i.e. one whose VRAM *is* system RAM.
 
-    Narrower than ``not _has_dedicated_vram()``: CPU and MPS also share memory, but their
+    Narrower than ``not device_has_dedicated_vram()``: CPU and MPS also share memory, but their
     keep-RAM-copy behaviour is long-standing and out of scope here.
     """
     return device.type == "xpu" and xpu_device_is_integrated(device) is True
@@ -1071,7 +1056,7 @@ def put(
         #   loading path there. partial_load_to_vram() respects vram_available, whereas the full-load path ignores it
         #   ("no choice but to try and fit it all"), and on Linux that overshoot is an uncatchable OOM-kill rather than
         #   a degraded load. Combined with keep_ram_copy=False below, streaming moves weights rather than copying them.
-        running_with_dedicated_vram = _has_dedicated_vram(effective_execution_device)
+        running_with_dedicated_vram = device_has_dedicated_vram(effective_execution_device)
         supports_partial_loading = running_with_dedicated_vram or _is_integrated_xpu(effective_execution_device)
 
         # On an integrated GPU the "VRAM" copy and the RAM copy are the same DRAM, so keeping a RAM
@@ -2368,7 +2353,7 @@ def _get_vram_available(self, working_mem_bytes: Optional[int], honor_cap: bool
             # for activations). Credit the reclaimable reserve, excluding intra-segment
             # fragmentation slack that a large contiguous allocation could not use.
             vram_available_to_process = vram_free + vram_allocated + self._get_reclaimable_allocator_bytes()
-        elif self._execution_device.type == "xpu" and _has_dedicated_vram(self._execution_device):
+        elif self._execution_device.type == "xpu" and device_has_dedicated_vram(self._execution_device):
             vram_allocated = torch.xpu.memory_allocated(self._execution_device)
             vram_free, _vram_total = TorchDevice.xpu_mem_get_info(self._execution_device)
             vram_available_to_process = vram_free + vram_allocated + self._get_reclaimable_allocator_bytes()
@@ -2495,7 +2480,7 @@ def _calc_ram_available_to_model_cache(self) -> int:
         total_vram_bytes: int | None = None
         if self._execution_device.type == "cuda":
             total_vram_bytes = torch.cuda.get_device_properties(self._execution_device).total_memory
-        elif self._execution_device.type == "xpu" and _has_dedicated_vram(self._execution_device):
+        elif self._execution_device.type == "xpu" and device_has_dedicated_vram(self._execution_device):
             # A total of 0 means "unknown" (the device couldn't report total_memory). Treat it like
             # the CPU case (None) so the RAM heuristic below doesn't compute a negative cache cap.
             #
@@ -3143,7 +3128,7 @@ def offload_models_from_vram_except(self, keep_model_keys: Collection[str]) -> i
 
         Returns the number of VRAM bytes freed.
         """
-        if not _has_dedicated_vram(self._execution_device):
+        if not device_has_dedicated_vram(self._execution_device):
             # CPU, MPS, integrated GPUs: their "VRAM" is system RAM, so moving models out frees nothing.
             return 0
         keep = set(keep_model_keys)

invokeai/backend/quantization/gguf/ggml_tensor.py

@@ -9,6 +9,7 @@
     TORCH_COMPATIBLE_QTYPES,
     dequantize,
 )
+from invokeai.backend.util.devices import TorchDevice, device_has_dedicated_vram
 
 
 def dequantize_and_run(func, args, kwargs):
@@ -206,8 +207,8 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 # What one dequantization allocates beyond its result, per weight element: the block kernels in
 # `utils.py` widen codes and scales through integer and float16 intermediates before the final cast.
 # Measured on CUDA (torch 2.13) for a 4608x53248 weight and rounded up; F32/F16 are a plain cast. Types
-# not listed (the IQ family decodes through the numpy fallback) get the worst measured value. The
-# host-side counterpart, for decoding a whole tensor at load, is `_PEAK_COPIES` in
+# not listed get the worst measured value when they have a device kernel. The host-side counterpart,
+# for decoding a whole tensor at load, is `_PEAK_COPIES` in
 # `model_manager/load/quantized_embedding.py`.
 _GGMLQ = gguf.GGMLQuantizationType
 _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT: dict[gguf.GGMLQuantizationType, float] = {
@@ -227,6 +228,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs):
 }
 _DEQUANT_INTERMEDIATE_BYTES_WORST = 8.0
 
+# NumPy decoding can retain float32 chunks alongside its full float32 output. Keep the existing
+# conservative shared-memory allowance separate from the device-kernel fallback.
+_HOST_DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT = 8.0
+
 
 def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int:
     """Peak bytes one forward transiently needs to dequantize this model's GGUF-packed Linear weights.
@@ -239,15 +244,24 @@ def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int:
     the device for the call. Zero when no Linear weight is packed.
     """
     peak = 0
+    host_intermediates: float | None = None
     for module in model.modules():
         if not isinstance(module, torch.nn.Linear):
             continue
         weight = module.weight
         if not isinstance(weight, GGMLTensor):
             continue
-        intermediates = _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT.get(
-            weight._ggml_quantization_type, _DEQUANT_INTERMEDIATE_BYTES_WORST
-        )
+        qtype = weight._ggml_quantization_type
+        if qtype in TORCH_COMPATIBLE_QTYPES or qtype in DEQUANTIZE_FUNCTIONS:
+            intermediates = _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT.get(qtype, _DEQUANT_INTERMEDIATE_BYTES_WORST)
+        else:
+            if host_intermediates is None:
+                # Streaming can leave weights on CPU; budget against the session's execution device.
+                device = TorchDevice.choose_torch_device()
+                host_intermediates = (
+                    0.0 if device_has_dedicated_vram(device) else _HOST_DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT
+                )
+            intermediates = host_intermediates
         per_element = weight.compute_dtype.itemsize + intermediates
         peak = max(peak, math.ceil(weight.tensor_shape.numel() * per_element) + weight.quantized_data.nbytes)
     return peak

invokeai/backend/util/devices.py

@@ -61,6 +61,19 @@ def device_is_integrated(device: torch.device) -> Optional[bool]:
     return None if flag is None else bool(flag)
 
 
+def device_has_dedicated_vram(device: torch.device) -> bool:
+    """Whether to budget device memory as dedicated VRAM rather than shared system RAM.
+
+    CUDA/ROCm retains its dedicated-memory budget. Integrated XPU shares system RAM; an unknown
+    XPU type keeps the discrete-device policy rather than guessing.
+    """
+    if device.type == "cuda":
+        return True
+    if device.type == "xpu":
+        return xpu_device_is_integrated(device) is not True
+    return False
+
+
 def _default_cuda_device() -> torch.device:
     """The CUDA device ``device: auto`` means: the current one, unless it is an integrated GPU and a discrete one exists.
 

tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py

@@ -7,6 +7,7 @@
 """
 
 import logging
+from contextlib import contextmanager
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch
 
@@ -56,12 +57,17 @@ def _make_cache(
         )
 
 
+@contextmanager
 def _patch_integrated(is_integrated: bool | None):
     """Stub the Level Zero probe. None models a device whose type cannot be determined."""
-    return patch(
-        "invokeai.backend.model_manager.load.model_cache.model_cache.xpu_device_is_integrated",
-        return_value=is_integrated,
-    )
+    with (
+        patch("invokeai.backend.util.devices.xpu_device_is_integrated", return_value=is_integrated),
+        patch(
+            "invokeai.backend.model_manager.load.model_cache.model_cache.xpu_device_is_integrated",
+            return_value=is_integrated,
+        ),
+    ):
+        yield
 
 
 def _logged_messages(logger: MagicMock) -> list[str]:

tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py

@@ -43,7 +43,7 @@ def move(cache_entry, vram_bytes_to_free, keep_required_weights_in_vram=None):
 
     with (
         patch.object(cache, "_move_model_to_ram", side_effect=move),
-        patch(f"{MODULE}._has_dedicated_vram", return_value=True),
+        patch(f"{MODULE}.device_has_dedicated_vram", return_value=True),
         patch(f"{MODULE}.TorchDevice.empty_cache") as empty_cache,
     ):
         yield keys, empty_cache
@@ -112,9 +112,15 @@ def test_offload_model_from_vram_still_moves_one_model_with_its_submodels(cache:
     assert sorted(keys) == ["flux", "flux:vae"]
 
 
-def test_a_device_whose_vram_is_system_ram_moves_nothing(cache: ModelCache):
+@pytest.mark.parametrize("device", ["cpu", "mps", "xpu:2"])
+def test_a_device_whose_vram_is_system_ram_moves_nothing(cache: ModelCache, device: str):
     """CPU, MPS and integrated GPUs: moving a model "to RAM" would copy it within the same memory."""
-    cache.put("flux:transformer", torch.nn.Linear(MB // 4, 1, bias=False))
-    with patch.object(cache, "_move_model_to_ram") as move:
+    _put(cache, "flux:transformer", 1)
+    with (
+        patch.object(cache, "_execution_device", torch.device(device)),
+        patch("invokeai.backend.util.devices.xpu_device_is_integrated", return_value=True),
+        patch.object(cache, "_move_model_to_ram", return_value=MB) as move,
+        patch(f"{MODULE}.TorchDevice.empty_cache"),
+    ):
         assert cache.offload_models_from_vram_except(set()) == 0
     move.assert_not_called()

tests/backend/quantization/gguf/test_ggml_tensor.py

@@ -4,7 +4,9 @@
 
 from invokeai.backend.quantization.gguf.ggml_tensor import GGMLTensor
 from invokeai.backend.quantization.gguf.utils import DEQUANTIZE_FUNCTIONS, TORCH_COMPATIBLE_QTYPES
+from invokeai.backend.util import devices
 from invokeai.backend.util.calc_tensor_size import calc_tensor_size
+from invokeai.backend.util.devices import TorchDevice
 
 
 def quantize_tensor(data: torch.Tensor, ggml_quantization_type: gguf.GGMLQuantizationType) -> GGMLTensor:
@@ -154,10 +156,17 @@ def test_ggml_tensor_numpy_fallback_dequantize_uses_logical_shape():
     assert t.get_dequantized_tensor().shape == logical_shape
 
 
-def _gguf_linear(in_features: int, out_features: int) -> torch.nn.Linear:
+def _gguf_linear(
+    in_features: int,
+    out_features: int,
+    qtype: gguf.GGMLQuantizationType = gguf.GGMLQuantizationType.Q8_0,
+    compute_dtype: torch.dtype = torch.float32,
+) -> torch.nn.Linear:
+    block_size, type_size = gguf.GGML_QUANT_SIZES[qtype]
+    packed = torch.zeros((out_features, in_features // block_size * type_size), dtype=torch.uint8)
     linear = torch.nn.Linear(in_features, out_features, bias=False)
     linear.weight = torch.nn.Parameter(
-        quantize_tensor(torch.zeros(out_features, in_features), gguf.GGMLQuantizationType.Q8_0), requires_grad=False
+        GGMLTensor(packed, qtype, torch.Size((out_features, in_features)), compute_dtype), requires_grad=False
     )
     return linear
 
@@ -170,17 +179,15 @@ def test_the_shared_dequant_transient_counts_gguf_linears():
     from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear
 
     gguf_model = torch.nn.Sequential(_gguf_linear(64, 128), _gguf_linear(64, 32))
-    largest = gguf_model[0].weight
     gguf_peak = peak_dequant_transient_bytes(gguf_model, torch.float32)
 
-    # At least the copy of the larger layer, in the weight's own compute dtype, plus its packed bytes;
-    # never the sum of both layers.
-    assert gguf_peak > largest.tensor_shape.numel() * largest.compute_dtype.itemsize + largest.quantized_data.nbytes
-    assert gguf_peak == peak_dequant_transient_bytes(torch.nn.Sequential(gguf_model[0]), torch.float32)
+    # 8192 elements * (4-byte result + 2-byte Q8 intermediates), plus 256 packed 34-byte blocks.
+    assert gguf_peak == 57856
+    assert peak_dequant_transient_bytes(torch.nn.Sequential(gguf_model[0]), torch.float32) == 57856
 
     # Whichever scheme holds the largest layer decides, in either order.
     larger_int8 = Int8ConvrotLinear(
-        weight=torch.zeros(4096, 4096, dtype=torch.int8), weight_scale=torch.ones(4096, 1), convrot=False
+        weight=torch.zeros(128, 128, dtype=torch.int8), weight_scale=torch.ones(128, 1), convrot=False
     )
     mixed = torch.nn.Sequential(gguf_model[0], larger_int8)
     assert peak_dequant_transient_bytes(mixed, torch.float32) == larger_int8.dequant_transient_bytes(torch.float32)
@@ -190,3 +197,115 @@ def test_the_shared_dequant_transient_counts_gguf_linears():
     assert smaller_int8.dequant_transient_bytes(torch.float32) < gguf_peak
     assert peak_dequant_transient_bytes(torch.nn.Sequential(smaller_int8, gguf_model[0]), torch.float32) == gguf_peak
     assert peak_dequant_transient_bytes(torch.nn.Sequential(torch.nn.Linear(64, 64)), torch.float32) == 0
+
+
+@pytest.mark.parametrize("compute_dtype", [torch.float16, torch.float32], ids=["fp16", "fp32"])
+def test_host_dequant_reservation_counts_only_result_and_packed_weight(compute_dtype, monkeypatch):
+    """Host decode intermediates do not consume dedicated VRAM, even with the weight still on CPU."""
+    from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes
+    from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes
+
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("cuda:1"))
+    linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, compute_dtype)
+    weight = linear.weight
+    assert weight.device.type == "cpu"
+    decoded = weight.get_dequantized_tensor()
+    assert decoded.dtype == compute_dtype
+
+    model = torch.nn.Sequential(linear)
+    expected_bytes = decoded.nbytes + weight.quantized_data.nbytes
+    assert peak_ggml_linear_dequant_transient_bytes(model) == expected_bytes
+    # GGUF carries its own dtype, even if a caller passes a different module-level dtype.
+    assert peak_dequant_transient_bytes(model, torch.float64) == expected_bytes
+
+
+@pytest.mark.parametrize(
+    ("execution_device", "integrated", "compute_dtype", "intermediate_bytes"),
+    [
+        pytest.param("cpu", None, torch.float32, 8, id="cpu-fp32"),
+        pytest.param("mps", None, torch.float16, 8, id="mps-fp16"),
+        pytest.param("xpu:2", True, torch.float32, 8, id="integrated-xpu-fp32"),
+        pytest.param("xpu:2", False, torch.float16, 0, id="discrete-xpu-fp16"),
+        pytest.param("xpu:2", None, torch.float16, 0, id="unknown-xpu-fp16"),
+    ],
+)
+def test_host_dequant_reservation_follows_execution_memory_budget(
+    execution_device, integrated, compute_dtype, intermediate_bytes, monkeypatch
+):
+    """Shared execution retains the prior conservative 8 bytes/element for host intermediates."""
+    from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes
+
+    device = torch.device(execution_device)
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: device)
+
+    def is_integrated(probed_device):
+        assert probed_device == device
+        return integrated
+
+    monkeypatch.setattr(devices, "xpu_device_is_integrated", is_integrated)
+    linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, compute_dtype)
+    assert linear.weight.device.type == "cpu"
+    # 512 logical elements and sixteen 18-byte IQ4_NL blocks; the weight owns the result dtype.
+    expected_bytes = 512 * (compute_dtype.itemsize + intermediate_bytes) + 288
+    assert peak_dequant_transient_bytes(linear, torch.float64) == expected_bytes
+
+
+def test_shared_host_dequant_reservation_is_independent_of_device_kernel_bound(monkeypatch):
+    """Lowering the device-kernel allowance must not shrink shared host working memory."""
+    from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes
+    from invokeai.backend.quantization.gguf import ggml_tensor
+
+    monkeypatch.setattr(ggml_tensor, "_DEQUANT_INTERMEDIATE_BYTES_WORST", 0.0)
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("mps"))
+    linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, torch.float16)
+    assert linear.weight.device.type == "cpu"
+    # 512 elements * (2-byte result + 8-byte host allowance), plus sixteen 18-byte packed blocks.
+    assert peak_dequant_transient_bytes(linear, torch.float64) == 5408
+
+
+def test_bf16_device_dequant_reservation_keeps_measured_intermediates():
+    from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes
+
+    linear = _gguf_linear(256, 2, gguf.GGMLQuantizationType.BF16, torch.float16)
+    # 512 elements * (2-byte result + 6-byte BF16 intermediates), plus 1024 packed bytes.
+    assert peak_dequant_transient_bytes(linear, torch.float64) == 5120
+
+
+def test_dedicated_dequant_reservation_uses_largest_bytes_in_either_order(monkeypatch):
+    from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes
+
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("cuda:1"))
+    larger_iq = _gguf_linear(256, 16, gguf.GGMLQuantizationType.IQ4_NL)
+    smaller_q8 = _gguf_linear(256, 12)
+    # IQ: 4096 * 4 + 2304 = 18688. Q8: 3072 * (4 + 2) + 3264 = 21696.
+    for layers in ((larger_iq, smaller_q8), (smaller_q8, larger_iq)):
+        assert peak_ggml_linear_dequant_transient_bytes(torch.nn.Sequential(*layers)) == 21696
+
+
+def test_host_dequant_reservation_follows_execution_device_changes(monkeypatch):
+    from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes
+
+    device = torch.device("xpu:2")
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: device)
+    monkeypatch.setattr(devices, "xpu_device_is_integrated", lambda _: True)
+    model = torch.nn.Sequential(
+        _gguf_linear(256, 2, gguf.GGMLQuantizationType.IQ4_NL, torch.float16),
+        _gguf_linear(256, 1, gguf.GGMLQuantizationType.IQ4_NL, torch.float16),
+    )
+
+    assert peak_ggml_linear_dequant_transient_bytes(model) == 5408  # 512 * (2 + 8) + 288.
+    device = torch.device("cuda:1")
+    assert peak_ggml_linear_dequant_transient_bytes(model) == 1312  # 512 * 2 + 288.
+
+
+def test_new_device_dequant_kernel_keeps_conservative_intermediate_reservation(monkeypatch):
+    from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes
+
+    qtype = gguf.GGMLQuantizationType.IQ4_NL
+    packed = torch.zeros((16, 18), dtype=torch.uint8)
+    weight = GGMLTensor(packed, qtype, torch.Size((2, 256)), torch.float16)
+    linear = torch.nn.Linear(256, 2, bias=False)
+    linear.weight = torch.nn.Parameter(weight, requires_grad=False)
+    # A newly registered device kernel has no measured table entry yet. Do not treat it as host-only.
+    monkeypatch.setitem(DEQUANTIZE_FUNCTIONS, qtype, DEQUANTIZE_FUNCTIONS[gguf.GGMLQuantizationType.Q8_0])
+    assert peak_ggml_linear_dequant_transient_bytes(linear) == 512 * (2 + 8) + packed.nbytes