Fix Qwen flash attention CUDA gating

[johnstonconnor773-ops]

Qwen text generation can launch comfy-kitchen's flash attention decode kernel on CUDA 12 PyTorch builds even though the CUDA backend is disabled for those builds, causing a driver/runtime error on the first decode step.

This change requires CUDA 13+ before selecting the fixed KV-cache flash decode path, preserving the regular KV-cache fallback for older builds. It also adds focused regression coverage for both unsupported and supported CUDA versions.

Tests: python -m pytest tests-unit/comfy_test/test_llama_kv_cache.py -q

Fixes: #16449

[coderabbitai]

Understand this PR’s impact

Explore downstream dependencies and potential security impact with Blast Radius.

View blast radius →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: Comfy-Org/ComfyUI/.coderabbit.yaml

Review profile: ASSERTIVE

Plan: Advanced

Run ID: caff5366-063f-400e-9eb4-0f2655a3b4a7

📥 Commits

Reviewing files that changed from the base of the PR and between b0f4b7b and 5e4030c.

📒 Files selected for processing (2)

• comfy/text_encoders/llama.py
• tests-unit/comfy_test/test_llama_kv_cache.py

Included review availability: Your plan provides up to 10 included reviews per hour; 7 remain after this review.

📜 Recent review details

🧰 Additional context used

📓 Path-based instructions (3)

Core ML/diffusion engine.

⚙️ CodeRabbit configuration file

Files:

• comfy/text_encoders/llama.py

IMPORTANT: Only comment on issues directly introduced by this PR's code changes.

⚙️ CodeRabbit configuration file

Files:

• tests-unit/comfy_test/test_llama_kv_cache.py
• comfy/text_encoders/llama.py

Documentation and README edits should be concise, factual, and tied to the changed behavior.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

• tests-unit/comfy_test/test_llama_kv_cache.py
• comfy/text_encoders/llama.py

🔇 Additional comments (2)

comfy/text_encoders/llama.py (1)

882-883: LGTM!

tests-unit/comfy_test/test_llama_kv_cache.py (1)

24-40: LGTM!

---

📝 Walkthrough

Walkthrough

Llama2_.init_kv_cache now selects the flash-attention FixedKV cache only when fixed_kv is enabled, flash attention is available, and the CUDA runtime is version 13 or newer. Older or unavailable CUDA versions use the regular KV-cache layout. Unit tests cover CUDA 12.8 fallback behavior and CUDA 13.0 flash-cache behavior.

Priority: ➖ Normal

Severity of issue fixed: Medium

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 4 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the Qwen flash-attention CUDA gating fix, which is the main change.
Description check	✅ Passed	The description accurately explains the CUDA 13 requirement, fallback behavior, regression tests, and linked issue.
Linked Issues check	✅ Passed	The change addresses #16449. Llama2_.init_kv_cache now requires torch.version.cuda to be present and have a major version of at least 13 before it selects FixedKV flash decode. CUDA 12.8 therefo…
Out of Scope Changes check	✅ Passed	The pull request changes only Qwen/Llama KV-cache selection and adds regression tests for the linked CUDA-version behavior. These changes directly support #16449 and do not add unrelated functionality…

• Fix all pre-merge checks with AI

---

_{Comment @coderabbitai help to get the list of available commands.}