Comfy-Org / Comfy-Org/comfy-kitchen

Flash Attention Decode availability detection raises ValueError for CPU devices instead of returning a boolean

Open Beginner friendly
#131 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
220
Forks
91
Avg merge
1d 7h
Merged PRs (30d)
12

Description

### Issue
The issue occurs in the following function:
https://github.com/Comfy-Org/comfy-kitchen/blob/7d86acf60c88fd6c3c733c0e54db22ef74b8d77f/comfy_kitchen/flash_attention.py#L17

When the `device` argument is CPU and torch is built with cuda 13.0 or later, the following error is raised at the return statement:
```ValueError: Expected a cuda device, but got: cpu```

### Example
When using the ACE-Step 1.5 workflow template in ComfyUI, executing the `TextEncodeAceStepAudio1.5` node results in the error.
This occurs because the text encoders are running on the CPU.
- Log
```python
[INFO] pytorch version: 2.13.0+cu130
[INFO] Device: cuda:0 NVIDIA GeForce GTX 1660 Ti : cudaMallocAsync

[INFO] got prompt
[INFO] VAE load device: cuda:0, offload device: cpu, dtype: torch.float32
[INFO] Requested to load ACE15TEModel_
[INFO] loaded completely; 4673.04 MB loaded, full load: True
[INFO] CLIP/text encoder model load device: cpu, offload device: cpu, current: cpu, dtype: torch.float16
[ERROR] !!! Exception during processing !!! Expected a cuda device, but got: cpu
[ERROR] Traceback (most recent call last):
File "Path\ComfyUI_windows_portable\ComfyUI\execution.py", line 545, in execute
output_data, output_ui, has_subgraph, has_pending_tasks = await get_output_data(prompt_id, unique_id, obj, input_data_all, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "Path\ComfyUI_windows_portable\ComfyUI\execution.py", line 344, in get_output_data
return_values = await _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, obj.FUNCTION, allow_interrupt=True, execution_block_cb=execution_block_cb, pre_execute_cb=pre_execute_cb, v3_data=v3_data)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "Path\ComfyUI_windows_portable\ComfyUI\execution.py", line 318, in _async_map_node_over_list
await process_inputs(input_dict, i)
File "Path\ComfyUI_windows_portable\ComfyUI\execution.py", line 306, in process_inputs
result = f(**inputs)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy_api\internal\__init__.py", line 149, in wrapped_func
return method(locked_class, **inputs)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy_api\latest\_io.py", line 1990, in EXECUTE_NORMALIZED
to_return = cls.execute(*args, **kwargs)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy_extras\nodes_ace.py", line 60, in execute
conditioning = clip.encode_from_tokens_scheduled(tokens)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy\sd.py", line 341, in encode_from_tokens_scheduled
pooled_dict = self.encode_from_tokens(tokens, return_pooled=return_pooled, return_dict=True)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy\sd.py", line 410, in encode_from_tokens
o = self.cond_stage_model.encode_token_weights(tokens)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy\text_encoders\ace15.py", line 335, in encode_token_weights
audio_codes = generate_audio_codes(getattr(self, self.lm_model, self.qwen3_06b), token_weight_pairs["lm_prompt"], token_weight_pairs["lm_prompt_negative"], min_tokens=lm_metadata["min_tokens"], max_tokens=lm_metadata["min_tokens"], seed=lm_metadata["seed"], cfg_scale=lm_metadata["cfg_scale"], temperature=lm_metadata["temperature"], top_p=lm_metadata["top_p"], top_k=lm_metadata["top_k"], min_p=lm_metadata["min_p"])
File "Path\ComfyUI_windows_portable\ComfyUI\comfy\text_encoders\ace15.py", line 165, in generate_audio_codes
return sample_manual_loop_no_classes(model, ids, cfg_scale=cfg_scale, temperature=temperature, top_p=top_p, top_k=top_k, min_p=min_p, seed=seed, min_tokens=min_tokens, max_new_tokens=max_tokens)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy\text_encoders\ace15.py", line 64, in sample_manual_loop_no_classes
past_key_values = model.transformer.model.init_kv_cache(embeds_batch, embeds.shape[1] + max_new_tokens, device, execution_dtype)
File "Path\ComfyUI_windows_portable\ComfyUI\comfy\text_encoders\llama.py", line 802, in init_kv_cache
fixed_kv = self.fixed_kv and comfy_kitchen.flash_attention_decode_is_available(device)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^
File "Path\ComfyUI_windows_portable\python_embeded\Lib\site-packages\comfy_kitchen\flash_attention.py", line 22, in is_available
return torch.cuda.get_device_capability(device) >= _MINIMUM_CAPABILITY
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^
File "Path\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\cuda\__init__.py", line 720, in get_device_capability
prop = get_device_properties(device)
File "Path\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\cuda\__init__.py", line 738, in get_device_properties
device = _get_device_index(device, optional=True)
File "Path\ComfyUI_windows_portable\python_embeded\Lib\site-packages\torch\cuda\_utils.py", line 582, in _get_device_index
raise ValueError(f"Expected a cuda device, but got: {device}")
ValueError: Expected a cuda device, but got: cpu
```

Contributor guide

Open the contributing guide

Research direction

Start in comfy_kitchen/flash_attention.py at flash_attention_decode_is_available, referenced around line 17, and inspect how the device argument reaches the CUDA capability check. Verify the CPU case returns a boolean without raising ValueError, including the ACE-Step text-encoder path described in the issue.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, performance
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
80/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.