From a3dc260c2acd194eb661a3430631d521db6a8c43 Mon Sep 17 00:00:00 2001
From: chelsealong <chelsealong@126.com>
Date: Sat, 15 Aug 2026 03:24:18 +0000
Subject: [PATCH] Fix crash when Llama-based text encoder KV cache init runs on
 CPU

init_kv_cache probed comfy_kitchen.flash_attention_decode_is_available()
unconditionally, which calls torch.cuda.get_device_capability() and raises
ValueError when the execution device is CPU (e.g. MiniMax Music3 text
encoder offloaded to CPU on low-VRAM GPUs). Guard the probe with
comfy.model_management.is_device_cuda(), matching the pattern already
used elsewhere in this file (ar.py's cuda_device check).

Fixes #15607
---
 comfy/text_encoders/llama.py                  |  2 +-
 .../comfy_test/test_llama_init_kv_cache.py    | 39 +++++++++++++++++++
 2 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 tests-unit/comfy_test/test_llama_init_kv_cache.py

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)