Refactor QKV Fusion Utilities to be LoRA-Aware by dg845 · Pull Request #14047 · huggingface/diffusers

dg845 · 2026-06-22T23:41:34Z

What does this PR do?

This PR refactors the attention QKV fusion utilities in AttentionMixin and AttentionModuleMixin to be more LoRA-aware. In particular, this PR adds guards when attempting to fuse/unfuse with a LoRA attached (because LoRAs cannot be easily transferred over when fusing/unfusing) and an inplace option to fuse without keeping copies of the split Q,K,V projections.

Changelist

Adds guards when attempting to fuse/unfuse Q,K,V with a LoRA attached (will raise an error).
Adds an inplace argument; if inplace=True, the module is modified to have only the fused QKV projection (e.g. to_qkv) with the split Q,K,V projections (e.g. to_q/to_k/to_v) removed. (inplace=False, the default, retains the current behavior).
Supports fusion in the case where add_k_proj and add_v_proj are present without add_q_proj also present (e.g. Wan models).
Adds get_qkv and get_added_qkv helper methods in AttentionModuleMixin which handles getting the Q, K, V (and added Q,K,V, for second stream projections in MM-DiT-style models like Flux) in both the fused and split case. This is intended to make it easier for attention processors to support both fused and split QKV.
Adds an experimental restore_checkpoint_fusion_state method to AttentionMixin to put models back in the fusion state of the original model checkpoint. A new _native_fused_projections attribute on AttentionModuleMixin is added to allow this state to be described. (The motivation is to make it easier to support PEFT adapters which target the original checkpoint structure.)

Partially addresses #14003.

Before submitting

Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
- Did you read the Coding with AI agents guide?
- Did you self-review the diff against .ai/review-rules.md?
Did you read the contributor guideline?
Did you read our philosophy doc? (important for complex PRs)
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?
Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul
@DN6

…e to original fusion state, module-level helpers to get Q,K,V in both fused and split cases

…duleMixin Tests are in tests/models/test_attention_mixins.py and cover the four minimal concrete AttentionModuleMixin fixtures (_MinimalSelfAttn, _MinimalCrossAttn, _MinimalAddedKVAttn, _MinimalAddedQKVAttn): TestAttentionModuleMixin (53 tests): - Idempotency of fuse_projections/unfuse_projections - Module attribute invariants for non-inplace and inplace paths - Weight/bias correctness: fused weight equals concatenation of split weights - Inplace round-trip weight preservation and storage-sharing (no copy on unfuse) - Cross-attention to_kv path, added-KV to_added_kv path (Wan-style), and added-QKV to_added_qkv path (Flux-style) - get_qkv and get_added_qkv numerical correctness in split and fused cases - LoRA guard: fuse_projections/unfuse_projections raise ValueError when PEFT-style lora_A/lora_B submodules are detected on split or fused projections TestAttentionMixin (6 tests): - fuse_qkv_projections/unfuse_qkv_projections propagate to all eligible blocks - restore_checkpoint_fusion_state respects _native_fused_projections=None/True/False per block, including mixed-state models Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

HuggingFaceDocBuilderDev · 2026-06-22T23:50:59Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

…n, Helios, ChronoEdit)

sayakpaul · 2026-06-24T07:04:46Z

+                delattr(self, "to_k")
+                delattr(self, "to_v")


Does it free the memory? If not, what is the purpose of deleting these attributes? Also, from what I understand this and some of the other refactors introduced in this PR aren't particularly for LoRA-awareness?

sayakpaul · 2026-06-24T07:05:03Z

+                delattr(self, "to_q")
+                delattr(self, "to_k")
+                delattr(self, "to_v")


Same as above.

sayakpaul · 2026-06-24T07:08:42Z

+    @staticmethod
+    def _has_active_lora(module: nn.Module) -> bool:
+        """Checks for the presence of PEFT-style LoRA modules without needing to import `peft`."""
+        return any("lora_A" in name or "lora_B" in name for name, _ in module.named_modules())


I think a better way to detect it is:

for name, mod in module.named_modules(): if isinstance(mod, BaseTunerLayer): ...

This way, we can cater towards any non-LoRA adapters in the future.

sayakpaul · 2026-06-24T07:10:54Z

            if isinstance(module, AttentionModuleMixin) and module._supports_qkv_fusion:
                module.unfuse_projections()

+    def restore_checkpoint_fusion_state(self, inplace: bool = False):


Could you provide an example of where / how this is used?

My mental model says:

Load the pipeline

Call the fuse_qkv_projections() method on pipe.transformer.

Then, before loading we call this method.

And then load the LoRA weights?

sayakpaul · 2026-06-24T07:13:13Z

@@ -0,0 +1,624 @@
+import pytest


Would it not make sense to make these a common mixin that can be added to models?

sayakpaul · 2026-06-24T07:15:19Z

+        return self.base(x) + self.lora_B(self.lora_A(x))
+
+
+class TestAttentionModuleMixin:


The tests are missing:

If outputs with and without fusion stay the same (they should).

If expected errors are raised when LoRA is attached.

Additionally, are we checking if the projection attributes like to_k, to_q, etc. are actually getting deleted for in-place fusions and are successfully restored when the user wants them?

sayakpaul · 2026-06-24T07:18:18Z

+        assert not hasattr(model.block2, "to_kv")
+
+    # -------------------------------------------------------------------------
+    # restore_checkpoint_fusion_state


Should the corresponding tests also check if we have the fused QKV projections rather than just checking the fused_projections attribute?

sayakpaul · 2026-06-24T07:18:37Z

+    def test_restore_checkpoint_noop_for_none(self, model):
+        # Default _native_fused_projections is None — state should be unchanged.
+        model.restore_checkpoint_fusion_state()
+        assert model.block1.fused_projections is False


Prefer

Suggested change

assert model.block1.fused_projections is False

assert not model.block1.fused_projections

dg845 and others added 5 commits June 21, 2026 07:46

Refactor QKV fusion: add inplace option, model-level helper to restor…

4e18884

…e to original fusion state, module-level helpers to get Q,K,V in both fused and split cases

Fix Claude review comments (typo, get_added_qkv docstring)

37c0f70

make style and make quality

b761f4a

Merge branch 'main' into refactor/lora-aware-qkv-fusion

a7512db

github-actions Bot added models tests size/L PR with diff > 200 LOC labels Jun 22, 2026

dg845 added 2 commits June 23, 2026 02:36

Add inplace argument to model-overriden methods for compatibility (Wa…

e41dee6

…n, Helios, ChronoEdit)

Use inplace=False in ChronoEdit for consistency

799b8cf

sayakpaul reviewed Jun 24, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor QKV Fusion Utilities to be LoRA-Aware#14047

Refactor QKV Fusion Utilities to be LoRA-Aware#14047
dg845 wants to merge 7 commits into
mainfrom
refactor/lora-aware-qkv-fusion

dg845 commented Jun 22, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Jun 22, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

sayakpaul Jun 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return self.base(x) + self.lora_B(self.lora_A(x))


		class TestAttentionModuleMixin:

	assert model.block1.fused_projections is False
	assert not model.block1.fused_projections

Conversation

dg845 commented Jun 22, 2026

What does this PR do?

Before submitting

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Jun 22, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants