diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index f182e5147..2f7061fdb 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -791,7 +791,7 @@ class Llama2_(nn.Module): def init_kv_cache(self, batch, capacity, device, dtype): caches = [] - fixed_kv = self.fixed_kv and comfy_kitchen.flash_attention_decode_is_available(device) + fixed_kv = self.fixed_kv and comfy.model_management.is_device_cuda(device) and comfy_kitchen.flash_attention_decode_is_available(device) for _ in range(self.config.num_hidden_layers): if fixed_kv: key = torch.empty((batch, capacity, self.config.num_key_value_heads, self.config.head_dim), device=device, dtype=dtype) diff --git a/tests-unit/comfy_test/test_llama_init_kv_cache.py b/tests-unit/comfy_test/test_llama_init_kv_cache.py new file mode 100644 index 000000000..1c2941269 --- /dev/null +++ b/tests-unit/comfy_test/test_llama_init_kv_cache.py @@ -0,0 +1,39 @@ +"""Regression test for init_kv_cache crashing when the text encoder runs on CPU.""" + +from __future__ import annotations + +import torch + +from comfy.cli_args import args + +if not torch.cuda.is_available(): + args.cpu = True + +import comfy.ops as ops # noqa: E402 +import comfy.text_encoders.llama as llama # noqa: E402 + + +def _make_model(): + config = llama.Llama2Config( + vocab_size=16, hidden_size=8, intermediate_size=8, + num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2, + ) + config.head_dim = 4 + config.fixed_kv = True + return llama.Llama2_(config, device="cpu", dtype=torch.float32, ops=ops.manual_cast) + + +def test_init_kv_cache_on_cpu_does_not_probe_cuda_capability(monkeypatch): + def _raise_if_not_cuda(device): + if torch.device(device).type != "cuda": + raise ValueError(f"Expected a cuda device, but got: {device}") + return True + + monkeypatch.setattr( + llama.comfy_kitchen, "flash_attention_decode_is_available", _raise_if_not_cuda + ) + + model = _make_model() + past = model.init_kv_cache(1, 4, torch.device("cpu"), torch.float32) + + assert not isinstance(past[0], llama.FixedKV)