From b96ef178aab2cd675e88637d77c62fbdb9bb0f28 Mon Sep 17 00:00:00 2001 From: wunianze666-netizen Date: Thu, 15 Oct 2026 13:45:22 +0800 Subject: [PATCH 1/3] fix(gguf): exclude host decode intermediates from device reservations --- .../backend/quantization/gguf/ggml_tensor.py | 11 ++-- .../quantization/gguf/test_ggml_tensor.py | 50 +++++++++++++++++++ 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/invokeai/backend/quantization/gguf/ggml_tensor.py b/invokeai/backend/quantization/gguf/ggml_tensor.py index 7bb5532194a..62dde18985a 100644 --- a/invokeai/backend/quantization/gguf/ggml_tensor.py +++ b/invokeai/backend/quantization/gguf/ggml_tensor.py @@ -206,7 +206,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 +# not listed get the worst measured value only when they have a device kernel. The numpy fallback +# decodes and casts on the host, so it needs no device intermediates beyond the result. The # host-side counterpart, for decoding a whole tensor at load, is `_PEAK_COPIES` in # `model_manager/load/quantized_embedding.py`. _GGMLQ = gguf.GGMLQuantizationType @@ -245,9 +246,11 @@ def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int: 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: + intermediates = 0.0 per_element = weight.compute_dtype.itemsize + intermediates peak = max(peak, math.ceil(weight.tensor_shape.numel() * per_element) + weight.quantized_data.nbytes) return peak diff --git a/tests/backend/quantization/gguf/test_ggml_tensor.py b/tests/backend/quantization/gguf/test_ggml_tensor.py index 32b1d537fba..db41b2c6fec 100644 --- a/tests/backend/quantization/gguf/test_ggml_tensor.py +++ b/tests/backend/quantization/gguf/test_ggml_tensor.py @@ -190,3 +190,53 @@ 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("qtype", [gguf.GGMLQuantizationType.IQ4_NL, gguf.GGMLQuantizationType.IQ4_XS]) +@pytest.mark.parametrize("compute_dtype", [torch.float16, torch.bfloat16, torch.float32]) +def test_host_dequant_reservation_counts_only_result_and_packed_weight(qtype, compute_dtype, monkeypatch): + """Host decoding must not reserve device intermediates; the shared node estimator must agree.""" + 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 + + # Fixed logical dimensions work for both block sizes, with a different packed row shape. + block_size, type_size = gguf.GGML_QUANT_SIZES[qtype] + packed = torch.zeros((512 // block_size, type_size), dtype=torch.uint8) + packed[:, 1] = 60 # The first two bytes are a nonzero float16 block scale (1.0). + logical_shape = torch.Size((2, 256)) + weight = GGMLTensor(packed, qtype, logical_shape, compute_dtype) + decode = gguf.quants.dequantize + expected = torch.from_numpy(decode(packed.numpy(), qtype)).reshape(logical_shape).to(compute_dtype) + calls = [] + + def decode_on_host(data, quantization_type): + calls.append(quantization_type) + return decode(data, quantization_type) + + monkeypatch.setattr(gguf.quants, "dequantize", decode_on_host) + decoded = weight.get_dequantized_tensor() + torch.testing.assert_close(decoded, expected, rtol=0, atol=0) + assert calls == [qtype] + assert decoded.dtype == compute_dtype + assert decoded.count_nonzero() > 0 + + linear = torch.nn.Linear(256, 2, bias=False) + linear.weight = torch.nn.Parameter(weight, requires_grad=False) + model = torch.nn.Sequential(linear) + expected_bytes = decoded.nbytes + packed.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 + + +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 From 421978cf7ca178cb2ae83a82c573ecdb92363ca1 Mon Sep 17 00:00:00 2001 From: wunianze666-netizen Date: Tue, 20 Oct 2026 12:51:07 +0800 Subject: [PATCH 2/3] fix(gguf): retain host decode reservation on shared memory Use the execution device budget for host fallback intermediates. Pin device-kernel byte regressions and trim redundant decoder checks. --- .../backend/quantization/gguf/ggml_tensor.py | 16 ++- .../quantization/gguf/test_ggml_tensor.py | 133 +++++++++++++----- 2 files changed, 112 insertions(+), 37 deletions(-) diff --git a/invokeai/backend/quantization/gguf/ggml_tensor.py b/invokeai/backend/quantization/gguf/ggml_tensor.py index 62dde18985a..795c5d65363 100644 --- a/invokeai/backend/quantization/gguf/ggml_tensor.py +++ b/invokeai/backend/quantization/gguf/ggml_tensor.py @@ -9,6 +9,8 @@ TORCH_COMPATIBLE_QTYPES, dequantize, ) +from invokeai.backend.util import level_zero +from invokeai.backend.util.devices import TorchDevice def dequantize_and_run(func, args, kwargs): @@ -206,8 +208,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 get the worst measured value only when they have a device kernel. The numpy fallback -# decodes and casts on the host, so it needs no device intermediates beyond the result. The +# not listed get the worst measured value when they have a device kernel. Host fallback intermediates +# consume no dedicated VRAM; shared-memory execution retains the prior conservative bound. The # host-side counterpart, for decoding a whole tensor at load, is `_PEAK_COPIES` in # `model_manager/load/quantized_embedding.py`. _GGMLQ = gguf.GGMLQuantizationType @@ -240,6 +242,7 @@ 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 @@ -250,7 +253,14 @@ def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int: if qtype in TORCH_COMPATIBLE_QTYPES or qtype in DEQUANTIZE_FUNCTIONS: intermediates = _DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT.get(qtype, _DEQUANT_INTERMEDIATE_BYTES_WORST) else: - intermediates = 0.0 + if host_intermediates is None: + # Streaming can leave weights on CPU; budget against the session's execution device. + device = TorchDevice.choose_torch_device() + has_dedicated_vram = device.type == "cuda" or ( + device.type == "xpu" and level_zero.xpu_device_is_integrated(device) is not True + ) + host_intermediates = 0.0 if has_dedicated_vram else _DEQUANT_INTERMEDIATE_BYTES_WORST + 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 diff --git a/tests/backend/quantization/gguf/test_ggml_tensor.py b/tests/backend/quantization/gguf/test_ggml_tensor.py index db41b2c6fec..74f4300995e 100644 --- a/tests/backend/quantization/gguf/test_ggml_tensor.py +++ b/tests/backend/quantization/gguf/test_ggml_tensor.py @@ -1,10 +1,14 @@ +from unittest.mock import Mock + import gguf import pytest import torch 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 level_zero 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,33 +158,42 @@ 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 -def test_the_shared_dequant_transient_counts_gguf_linears(): +def test_the_shared_dequant_transient_counts_gguf_linears(monkeypatch): """Every denoise and encoder node sizes its reservation with `peak_dequant_transient_bytes`, so a GGUF Linear -- dequantized into a full copy on every forward -- has to count there, as an int8 one does. It is the largest layer that counts, whichever scheme it uses, since the layers run one at a time.""" from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear + monkeypatch.setattr( + TorchDevice, "choose_torch_device", lambda: pytest.fail("No host fallback needs a device lookup") + ) + monkeypatch.setattr(level_zero, "xpu_device_is_integrated", lambda _: pytest.fail("No host fallback needs a probe")) 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) @@ -192,43 +205,95 @@ def test_the_shared_dequant_transient_counts_gguf_linears(): assert peak_dequant_transient_bytes(torch.nn.Sequential(torch.nn.Linear(64, 64)), torch.float32) == 0 -@pytest.mark.parametrize("qtype", [gguf.GGMLQuantizationType.IQ4_NL, gguf.GGMLQuantizationType.IQ4_XS]) -@pytest.mark.parametrize("compute_dtype", [torch.float16, torch.bfloat16, torch.float32]) -def test_host_dequant_reservation_counts_only_result_and_packed_weight(qtype, compute_dtype, monkeypatch): - """Host decoding must not reserve device intermediates; the shared node estimator must agree.""" +@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 - # Fixed logical dimensions work for both block sizes, with a different packed row shape. - block_size, type_size = gguf.GGML_QUANT_SIZES[qtype] - packed = torch.zeros((512 // block_size, type_size), dtype=torch.uint8) - packed[:, 1] = 60 # The first two bytes are a nonzero float16 block scale (1.0). - logical_shape = torch.Size((2, 256)) - weight = GGMLTensor(packed, qtype, logical_shape, compute_dtype) - decode = gguf.quants.dequantize - expected = torch.from_numpy(decode(packed.numpy(), qtype)).reshape(logical_shape).to(compute_dtype) - calls = [] - - def decode_on_host(data, quantization_type): - calls.append(quantization_type) - return decode(data, quantization_type) - - monkeypatch.setattr(gguf.quants, "dequantize", decode_on_host) + 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() - torch.testing.assert_close(decoded, expected, rtol=0, atol=0) - assert calls == [qtype] assert decoded.dtype == compute_dtype - assert decoded.count_nonzero() > 0 - linear = torch.nn.Linear(256, 2, bias=False) - linear.weight = torch.nn.Parameter(weight, requires_grad=False) model = torch.nn.Sequential(linear) - expected_bytes = decoded.nbytes + packed.nbytes + 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(level_zero, "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_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_resolves_execution_device_once_per_call(monkeypatch): + from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes + + choose_device = Mock(side_effect=[torch.device("xpu:2"), torch.device("cuda:1")]) + probe = Mock(return_value=True) + monkeypatch.setattr(TorchDevice, "choose_torch_device", choose_device) + monkeypatch.setattr(level_zero, "xpu_device_is_integrated", probe) + 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. + choose_device.assert_called_once_with() + assert peak_ggml_linear_dequant_transient_bytes(model) == 1312 # 512 * 2 + 288. + assert choose_device.call_count == 2 + probe.assert_called_once_with(torch.device("xpu:2")) + + 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 From d03748f3373919bda1f0830e7bfde17ab4f8ce35 Mon Sep 17 00:00:00 2001 From: wunianze666-netizen Date: Wed, 21 Oct 2026 00:34:35 +0800 Subject: [PATCH 3/3] fix(gguf): share dedicated-memory budget policy Keep host decode allowance independent of device-kernel coefficients. Use one device-budget predicate in GGUF estimates and model-cache decisions. --- .../load/model_cache/model_cache.py | 27 +++---------- .../backend/quantization/gguf/ggml_tensor.py | 17 +++++---- invokeai/backend/util/devices.py | 13 +++++++ .../test_model_cache_integrated_gpu.py | 14 +++++-- .../test_model_cache_offload_from_vram.py | 14 +++++-- .../quantization/gguf/test_ggml_tensor.py | 38 ++++++++++--------- tests/backend/util/test_devices.py | 25 ++++++++++++ 7 files changed, 94 insertions(+), 54 deletions(-) diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 3691f374425..226bdb5767c 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/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) diff --git a/invokeai/backend/quantization/gguf/ggml_tensor.py b/invokeai/backend/quantization/gguf/ggml_tensor.py index 795c5d65363..c7f0108575f 100644 --- a/invokeai/backend/quantization/gguf/ggml_tensor.py +++ b/invokeai/backend/quantization/gguf/ggml_tensor.py @@ -9,8 +9,7 @@ TORCH_COMPATIBLE_QTYPES, dequantize, ) -from invokeai.backend.util import level_zero -from invokeai.backend.util.devices import TorchDevice +from invokeai.backend.util.devices import TorchDevice, device_has_dedicated_vram def dequantize_and_run(func, args, kwargs): @@ -208,9 +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 get the worst measured value when they have a device kernel. Host fallback intermediates -# consume no dedicated VRAM; shared-memory execution retains the prior conservative bound. 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] = { @@ -230,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. @@ -256,10 +258,9 @@ def peak_ggml_linear_dequant_transient_bytes(model: torch.nn.Module) -> int: if host_intermediates is None: # Streaming can leave weights on CPU; budget against the session's execution device. device = TorchDevice.choose_torch_device() - has_dedicated_vram = device.type == "cuda" or ( - device.type == "xpu" and level_zero.xpu_device_is_integrated(device) is not True + host_intermediates = ( + 0.0 if device_has_dedicated_vram(device) else _HOST_DEQUANT_INTERMEDIATE_BYTES_PER_ELEMENT ) - host_intermediates = 0.0 if has_dedicated_vram else _DEQUANT_INTERMEDIATE_BYTES_WORST 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) diff --git a/invokeai/backend/util/devices.py b/invokeai/backend/util/devices.py index b648adc0c8f..a56524ef3b3 100644 --- a/invokeai/backend/util/devices.py +++ b/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. diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py b/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py index 2621aa0ca41..2e094b257bd 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_integrated_gpu.py +++ b/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]: diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py b/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py index 084f6a616b4..368e0620c13 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_offload_from_vram.py +++ b/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() diff --git a/tests/backend/quantization/gguf/test_ggml_tensor.py b/tests/backend/quantization/gguf/test_ggml_tensor.py index 74f4300995e..f3755e5213b 100644 --- a/tests/backend/quantization/gguf/test_ggml_tensor.py +++ b/tests/backend/quantization/gguf/test_ggml_tensor.py @@ -1,12 +1,10 @@ -from unittest.mock import Mock - import gguf import pytest import torch 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 level_zero +from invokeai.backend.util import devices from invokeai.backend.util.calc_tensor_size import calc_tensor_size from invokeai.backend.util.devices import TorchDevice @@ -173,17 +171,13 @@ def _gguf_linear( return linear -def test_the_shared_dequant_transient_counts_gguf_linears(monkeypatch): +def test_the_shared_dequant_transient_counts_gguf_linears(): """Every denoise and encoder node sizes its reservation with `peak_dequant_transient_bytes`, so a GGUF Linear -- dequantized into a full copy on every forward -- has to count there, as an int8 one does. It is the largest layer that counts, whichever scheme it uses, since the layers run one at a time.""" from invokeai.backend.quantization.dequantizing_linear import peak_dequant_transient_bytes from invokeai.backend.quantization.int8_convrot import Int8ConvrotLinear - monkeypatch.setattr( - TorchDevice, "choose_torch_device", lambda: pytest.fail("No host fallback needs a device lookup") - ) - monkeypatch.setattr(level_zero, "xpu_device_is_integrated", lambda _: pytest.fail("No host fallback needs a probe")) gguf_model = torch.nn.Sequential(_gguf_linear(64, 128), _gguf_linear(64, 32)) gguf_peak = peak_dequant_transient_bytes(gguf_model, torch.float32) @@ -248,7 +242,7 @@ def is_integrated(probed_device): assert probed_device == device return integrated - monkeypatch.setattr(level_zero, "xpu_device_is_integrated", is_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. @@ -256,6 +250,19 @@ def is_integrated(probed_device): 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 @@ -275,23 +282,20 @@ def test_dedicated_dequant_reservation_uses_largest_bytes_in_either_order(monkey assert peak_ggml_linear_dequant_transient_bytes(torch.nn.Sequential(*layers)) == 21696 -def test_host_dequant_reservation_resolves_execution_device_once_per_call(monkeypatch): +def test_host_dequant_reservation_follows_execution_device_changes(monkeypatch): from invokeai.backend.quantization.gguf.ggml_tensor import peak_ggml_linear_dequant_transient_bytes - choose_device = Mock(side_effect=[torch.device("xpu:2"), torch.device("cuda:1")]) - probe = Mock(return_value=True) - monkeypatch.setattr(TorchDevice, "choose_torch_device", choose_device) - monkeypatch.setattr(level_zero, "xpu_device_is_integrated", probe) + 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. - choose_device.assert_called_once_with() + device = torch.device("cuda:1") assert peak_ggml_linear_dequant_transient_bytes(model) == 1312 # 512 * 2 + 288. - assert choose_device.call_count == 2 - probe.assert_called_once_with(torch.device("xpu:2")) def test_new_device_dequant_kernel_keeps_conservative_intermediate_reservation(monkeypatch): diff --git a/tests/backend/util/test_devices.py b/tests/backend/util/test_devices.py index 5bca941e44b..9ee15b63462 100644 --- a/tests/backend/util/test_devices.py +++ b/tests/backend/util/test_devices.py @@ -15,6 +15,7 @@ TorchDevice, choose_precision, choose_torch_device, + device_has_dedicated_vram, device_is_integrated, torch_dtype, ) @@ -33,6 +34,30 @@ def test_device_choice(device_name): assert torch_device == torch.device(device_name) +@pytest.mark.parametrize( + ("device_name", "integrated", "expected"), + [ + pytest.param("cpu", None, False, id="cpu-shared"), + pytest.param("mps", None, False, id="mps-shared"), + pytest.param("cuda:1", True, True, id="rocm-apu-keeps-dedicated-budget"), + pytest.param("xpu:2", True, False, id="integrated-xpu-shared"), + pytest.param("xpu:2", False, True, id="discrete-xpu-dedicated"), + pytest.param("xpu:2", None, True, id="unknown-xpu-dedicated"), + ], +) +def test_device_memory_budget_classification(device_name, integrated, expected, monkeypatch): + device = torch.device(device_name) + + def probe(probed_device): + assert probed_device == device + return integrated + + monkeypatch.setattr("invokeai.backend.util.devices.xpu_device_is_integrated", probe) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _: SimpleNamespace(is_integrated=1)) + monkeypatch.setattr(torch.version, "hip", "7.15.26333") + assert device_has_dedicated_vram(device) is expected + + # ===== per-thread session device (multi-GPU worker pinning) ================