Support gqa on all attention backends, drop support for pytorch 2.4#14772
Support gqa on all attention backends, drop support for pytorch 2.4#14772comfyanonymous wants to merge 3 commits into
Conversation
comfyanonymous
requested review from
Kosinkadink,
alexisrolland,
guill,
kijai and
rattus128
as code owners
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (6)
📜 Recent review details⏰ Context from checks skipped due to timeout. (9)
🧰 Additional context used📓 Path-based instructions (3)**/*.py📄 CodeRabbit inference engine (AGENTS.md)
Files:
**⚙️ CodeRabbit configuration file
Files:
⚙️ CodeRabbit configuration file
Files:
comfy/**⚙️ CodeRabbit configuration file
Files:
🧠 Learnings (2)📚 Learning: 2026-02-21T14:01:41.482ZApplied to files:
📚 Learning: 2026-05-13T12:31:45.069ZApplied to files:
🔇 Additional comments (6)
📝 WalkthroughWalkthroughChangesThis PR updates attention handling to use shared GQA-aware reshaping and backend kwargs across core attention paths, SageAttention, FlashAttention, and several downstream callers that previously repeated K/V heads manually. It also removes the Related issues: None identified from the provided information. Related PRs: None identified from the provided information. Suggested labels: attention, refactor, gqa, torch-compat Suggested reviewers: Reviewers familiar with Poem: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 inconclusive)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy/ldm/modules/attention.py`:
- Around line 646-647: The GQA path in attention.py expands K/V too early in the
attention flow, so the fallback branch that routes through PyTorch for large
head dims or small sequence lengths still pays the extra VRAM cost. Move the
`_repeat_kv_for_gqa` handling out of the early `enable_gqa` block and apply it
only on the non-fallback path, keeping `k` and `v` unexpanded before the
`sdpa`/PyTorch routing decision in the attention function.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: ccce84a8-1fec-46a5-9497-abc29e3cbb36
📒 Files selected for processing (3)
README.mdcomfy/ldm/modules/attention.pycomfy/text_encoders/gpt_oss.py
📜 Review details
⏰ Context from checks skipped due to timeout. (7)
- GitHub Check: test (windows-2022)
- GitHub Check: test (ubuntu-latest)
- GitHub Check: test (ubuntu-latest)
- GitHub Check: test (macos-latest)
- GitHub Check: test (windows-latest)
- GitHub Check: Run Pylint
- GitHub Check: Run Pylint
🧰 Additional context used
📓 Path-based instructions (3)
**
⚙️ CodeRabbit configuration file
**: ## Engineering Style
- Keep changes small and direct. Most fixes should touch the narrowest code path
that explains the bug, performance issue, dtype issue, model-format issue, or
user-facing behavior.- Change the least amount of files possible. A change that touches many files is
more likely to be a bad change than a good one unless the broader scope is
directly required.- Prefer practical fixes over broad architecture work. Add abstractions only
when they remove real repeated logic or match an existing ComfyUI pattern.- Prefer fewer dependencies. Do not add new dependencies to ComfyUI unless they
are absolutely necessary.- Delete obsolete code aggressively when newer infrastructure makes it useless.
Remove dead fallbacks, migration paths, unused options, debug prints, and
compatibility branches that are no longer needed. Do not leave dead branches,
unreachable code, or functions that are never called. If code is not
necessary for the current behavior, remove it.- Revert or disable problematic behavior quickly when it breaks users. It is
better to remove a broken feature path than keep a complicated partial fix.- Preserve existing APIs, node names, model-loading behavior, file layout, and
workflow compatibility unless the change is explicitly about replacing them.- Code must look hand-written for this repository. Changes that read like
generic AI-generated code will be rejected automatically: unnecessary helper
layers, vague names, boilerplate comments, defensive branches without a real
failure mode, broad rewrites, or code that ignores the local style.Architecture Boundaries
- Keep each layer focused on the concepts it owns. Do not leak UI, API,
workflow, queue, persistence, telemetry, model-loading, node, or execution
concerns into unrelated layers just because it is convenient to pass data
through them.- Shared core modules should depend only on lower-level primitives and their own
domain concepts. Highe...
Files:
README.mdcomfy/text_encoders/gpt_oss.pycomfy/ldm/modules/attention.py
⚙️ CodeRabbit configuration file
**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing awith:block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.
Files:
README.mdcomfy/text_encoders/gpt_oss.pycomfy/ldm/modules/attention.py
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Keep imports at module scope; avoid inline imports unless they are already part of an established optional-backend probe or are needed to avoid an import cycle.
Do not add unnecessarytry/exceptblocks; use them only for optional dependency, platform, or backend capability detection when the program has a useful fallback, and prefer specific exception types when changing new code.
Remove workarounds for PyTorch versions that ComfyUI no longer officially supports; if a workaround does not have a comment naming the exact supported PyTorch version(s), remove it.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently producing lower quality output.
Match the existing local style in the file you edit; this codebase tolerates long lines, simple helper functions, module-level state, and direct tensor operations when they make the code easier to follow.
Keep comments sparse and useful; strip useless comments that restate the code or describe obvious behavior, while allowing short TODOs that name the concrete missing follow-up.
Treat dtype, device placement, VRAM usage, and offloading behavior as core correctness concerns; check CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low VRAM implications when touching shared execution or loading code.
Prefer native ComfyUI formats and existing quantization/offload helpers over adding parallel code paths; usecomfy.quant_ops,comfy.model_management,comfy.memory_management,comfy.pinned_memory,comfy_aimdo, andcomfy-kitchenhelpers where they already solve the problem.
Use optimized comfy-kitchen ops where they improve performance without changing expected dtype, device, memory, or interface behavior.
All models should use the optimized attention function selected by ComfyUI; higher-level code must not inspect function identity, names, modules, or implementation details to decide behavior.
Apply the same opacity rule to similar patterns beyond...
Files:
comfy/text_encoders/gpt_oss.pycomfy/ldm/modules/attention.py
comfy/**
⚙️ CodeRabbit configuration file
comfy/**: Core ML/diffusion engine. Focus on:
- Backward compatibility (breaking changes affect all custom nodes)
- Memory management and GPU resource handling
- Performance implications in hot paths
- Thread safety for concurrent execution
Files:
comfy/text_encoders/gpt_oss.pycomfy/ldm/modules/attention.py
🧠 Learnings (2)
📚 Learning: 2026-02-21T14:01:41.482Z
Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.
Applied to files:
comfy/text_encoders/gpt_oss.pycomfy/ldm/modules/attention.py
📚 Learning: 2026-05-13T12:31:45.069Z
Learnt from: rattus128
Repo: Comfy-Org/ComfyUI PR: 13802
File: comfy/pinned_memory.py:19-30
Timestamp: 2026-05-13T12:31:45.069Z
Learning: When reviewing code that uses comfy/pinned_memory.py’s `HostBuffer.extend(size=..., reallocate=...)`: by default (`reallocate` is not True / False), `extend(size=...)` is a *relative increment* that grows the buffer by `size` bytes—so slicing like `[offset:offset+size]` after `hostbuf.extend(size=size)` is correct and the argument should not be rewritten to `offset + size`. Only in the single-segment reallocation mode (`reallocate=True`, e.g., as used by `resize_pin_buffer()` in `comfy/model_management.py`) should `size` be treated as an *absolute target* and the call/arguments should be checked accordingly.
Applied to files:
comfy/text_encoders/gpt_oss.pycomfy/ldm/modules/attention.py
🔇 Additional comments (4)
README.md (1)
232-232: LGTM!comfy/ldm/modules/attention.py (1)
3-3: LGTM!Also applies to: 22-27, 93-130, 194-206, 264-273, 340-348, 470-505, 529-540, 568-596, 675-692, 711-724, 734-763
comfy/text_encoders/gpt_oss.py (2)
15-15: LGTM!
112-128: LGTM!
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@comfy/ldm/modules/attention.py`:
- Around line 671-679: The fallback path in attention_pytorch is still retaining
expanded q_s/k_s/v_s tensors from the skip_reshape + enable_gqa branch,
increasing peak VRAM before the retry. Remove the conditional guard and free
q_s, k_s, and v_s unconditionally right before the fallback so the materialized
K/V copies are released even when skip_reshape is true; keep the change
localized to the attention_pytorch fallback handling and the q_s/k_s/v_s
assignments.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: aa1384aa-e35d-4f51-b04f-7335bd9ae921
📒 Files selected for processing (1)
comfy/ldm/modules/attention.py
📜 Review details
⏰ Context from checks skipped due to timeout. (9)
- GitHub Check: Run Pylint
- GitHub Check: test (ubuntu-latest)
- GitHub Check: test
- GitHub Check: test (windows-2022)
- GitHub Check: test (macos-latest)
- GitHub Check: test (macos-latest)
- GitHub Check: test (ubuntu-latest)
- GitHub Check: test (windows-latest)
- GitHub Check: Run Pylint
🧰 Additional context used
📓 Path-based instructions (3)
**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
**/*.py: Keep imports at module scope; avoid inline imports unless they are already part of an established optional-backend probe or are needed to avoid an import cycle.
Do not add unnecessarytry/exceptblocks; use them only for optional dependency, platform, or backend capability detection when the program has a useful fallback, and prefer specific exception types when changing new code.
Remove workarounds for PyTorch versions that ComfyUI no longer officially supports; if a workaround does not have a comment naming the exact supported PyTorch version(s), remove it.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently producing lower quality output.
Match the existing local style in the file you edit; this codebase tolerates long lines, simple helper functions, module-level state, and direct tensor operations when they make the code easier to follow.
Keep comments sparse and useful; strip useless comments that restate the code or describe obvious behavior, while allowing short TODOs that name the concrete missing follow-up.
Treat dtype, device placement, VRAM usage, and offloading behavior as core correctness concerns; check CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low VRAM implications when touching shared execution or loading code.
Prefer native ComfyUI formats and existing quantization/offload helpers over adding parallel code paths; usecomfy.quant_ops,comfy.model_management,comfy.memory_management,comfy.pinned_memory,comfy_aimdo, andcomfy-kitchenhelpers where they already solve the problem.
Use optimized comfy-kitchen ops where they improve performance without changing expected dtype, device, memory, or interface behavior.
All models should use the optimized attention function selected by ComfyUI; higher-level code must not inspect function identity, names, modules, or implementation details to decide behavior.
Apply the same opacity rule to similar patterns beyond...
Files:
comfy/ldm/modules/attention.py
**
⚙️ CodeRabbit configuration file
**: ## Engineering Style
- Keep changes small and direct. Most fixes should touch the narrowest code path
that explains the bug, performance issue, dtype issue, model-format issue, or
user-facing behavior.- Change the least amount of files possible. A change that touches many files is
more likely to be a bad change than a good one unless the broader scope is
directly required.- Prefer practical fixes over broad architecture work. Add abstractions only
when they remove real repeated logic or match an existing ComfyUI pattern.- Prefer fewer dependencies. Do not add new dependencies to ComfyUI unless they
are absolutely necessary.- Delete obsolete code aggressively when newer infrastructure makes it useless.
Remove dead fallbacks, migration paths, unused options, debug prints, and
compatibility branches that are no longer needed. Do not leave dead branches,
unreachable code, or functions that are never called. If code is not
necessary for the current behavior, remove it.- Revert or disable problematic behavior quickly when it breaks users. It is
better to remove a broken feature path than keep a complicated partial fix.- Preserve existing APIs, node names, model-loading behavior, file layout, and
workflow compatibility unless the change is explicitly about replacing them.- Code must look hand-written for this repository. Changes that read like
generic AI-generated code will be rejected automatically: unnecessary helper
layers, vague names, boilerplate comments, defensive branches without a real
failure mode, broad rewrites, or code that ignores the local style.Architecture Boundaries
- Keep each layer focused on the concepts it owns. Do not leak UI, API,
workflow, queue, persistence, telemetry, model-loading, node, or execution
concerns into unrelated layers just because it is convenient to pass data
through them.- Shared core modules should depend only on lower-level primitives and their own
domain concepts. Highe...
Files:
comfy/ldm/modules/attention.py
⚙️ CodeRabbit configuration file
**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing awith:block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.
Files:
comfy/ldm/modules/attention.py
comfy/**
⚙️ CodeRabbit configuration file
comfy/**: Core ML/diffusion engine. Focus on:
- Backward compatibility (breaking changes affect all custom nodes)
- Memory management and GPU resource handling
- Performance implications in hot paths
- Thread safety for concurrent execution
Files:
comfy/ldm/modules/attention.py
🧠 Learnings (2)
📚 Learning: 2026-02-21T14:01:41.482Z
Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.
Applied to files:
comfy/ldm/modules/attention.py
📚 Learning: 2026-05-13T12:31:45.069Z
Learnt from: rattus128
Repo: Comfy-Org/ComfyUI PR: 13802
File: comfy/pinned_memory.py:19-30
Timestamp: 2026-05-13T12:31:45.069Z
Learning: When reviewing code that uses comfy/pinned_memory.py’s `HostBuffer.extend(size=..., reallocate=...)`: by default (`reallocate` is not True / False), `extend(size=...)` is a *relative increment* that grows the buffer by `size` bytes—so slicing like `[offset:offset+size]` after `hostbuf.extend(size=size)` is correct and the argument should not be rewritten to `offset + size`. Only in the single-segment reallocation mode (`reallocate=True`, e.g., as used by `resize_pin_buffer()` in `comfy/model_management.py`) should `size` be treated as an *absolute target* and the call/arguments should be checked accordingly.
Applied to files:
comfy/ldm/modules/attention.py
🔇 Additional comments (1)
comfy/ldm/modules/attention.py (1)
93-129: LGTM!
| if skip_reshape: | ||
| q_s = q | ||
| if kwargs.get("enable_gqa", False): | ||
| k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3) | ||
| else: | ||
| k_s, v_s = k, v | ||
| else: | ||
| q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False)) | ||
| q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s)) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the attention implementation around the reported lines.
sed -n '630,720p' comfy/ldm/modules/attention.py
# Show the surrounding function boundaries for context.
python3 - <<'PY'
from pathlib import Path
path = Path("comfy/ldm/modules/attention.py")
lines = path.read_text().splitlines()
for i, line in enumerate(lines, 1):
if "def " in line and 620 <= i <= 720:
print(f"{i}: {line}")
PYRepository: Comfy-Org/ComfyUI
Length of output: 3557
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
path = Path("comfy/ldm/modules/attention.py")
text = path.read_text().splitlines()
targets = [
"_repeat_kv_for_gqa",
"def attention_pytorch",
]
for target in targets:
for i, line in enumerate(text, 1):
if line.startswith(f"def {target}") or line.startswith(f"def {target}("):
start = max(1, i - 20)
end = min(len(text), i + 120)
print(f"\n### {target} @ {i}\n")
for j in range(start, end + 1):
print(f"{j}: {text[j-1]}")
break
PYRepository: Comfy-Org/ComfyUI
Length of output: 6086
Free q_s/k_s/v_s unconditionally before the fallback. The skip_reshape + enable_gqa branch materializes expanded K/V copies, and attention_pytorch re-expands them again on fallback. Dropping the if not skip_reshape guard is safe and trims peak VRAM on the failure path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@comfy/ldm/modules/attention.py` around lines 671 - 679, The fallback path in
attention_pytorch is still retaining expanded q_s/k_s/v_s tensors from the
skip_reshape + enable_gqa branch, increasing peak VRAM before the retry. Remove
the conditional guard and free q_s, k_s, and v_s unconditionally right before
the fallback so the materialized K/V copies are released even when skip_reshape
is true; keep the change localized to the attention_pytorch fallback handling
and the q_s/k_s/v_s assignments.
Sources: Coding guidelines, Path instructions
1 participant
No description provided.