From 763bb0a8f0053f8e68f13356871ad9c981f027a1 Mon Sep 17 00:00:00 2001
From: dualchimerra <sunya1331@gmail.com>
Date: Sun, 20 Sep 2026 14:35:02 +0200
Subject: [PATCH] perf(krea2): use cuDNN SDPA with native GQA

---
 docs/user-guide/training-tips.md |  2 +-
 modeling/krea2/krea2_modeling.py | 48 ++++++++++++++++++++++++++------
 tests/test_krea2_modeling.py     | 41 ++++++++++++++++++++++++++-
 3 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/docs/user-guide/training-tips.md b/docs/user-guide/training-tips.md
index 8f69085..b2c172f 100644
--- a/docs/user-guide/training-tips.md
+++ b/docs/user-guide/training-tips.md
@@ -512,7 +512,7 @@ grad_accum: 2
 resolution: 1024
 grad_checkpoint: true
 mixed_precision: "bf16"
-attention_backend: "none"  # 用 PyTorch SDPA（也可选 "xformers" / "flash_attn"）
+attention_backend: "none"  # PyTorch SDPA；Krea2 CUDA 路径使用 cuDNN + native GQA
 cache_latents: true
 ```
 
diff --git a/modeling/krea2/krea2_modeling.py b/modeling/krea2/krea2_modeling.py
index 3e9e55c..b9b0bbe 100644
--- a/modeling/krea2/krea2_modeling.py
+++ b/modeling/krea2/krea2_modeling.py
@@ -22,6 +22,7 @@
 import torch.nn.functional as F
 from einops import rearrange
 from torch import Tensor, nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 
 def _rope(pos: Tensor, dim: int, theta: float) -> Tensor:
@@ -44,6 +45,42 @@ def _apply_rope(q: Tensor, k: Tensor, freqs: Tensor) -> tuple[Tensor, Tensor]:
     return q_out.reshape_as(q).to(q.dtype), k_out.reshape_as(k).to(k.dtype)
 
 
+def _native_gqa_attention(
+    q: Tensor,
+    k: Tensor,
+    v: Tensor,
+    mask: Tensor | None,
+    *,
+    enable_gqa: bool,
+) -> Tensor:
+    """Run Krea2 attention without expanding its 12 K/V heads to 48.
+
+    The official Krea2 implementation pins CUDA to cuDNN SDPA and passes GQA
+    through natively.  Check support first so unusual shapes/devices retain
+    PyTorch's normal backend selection instead of failing at runtime.
+    """
+    kwargs = {
+        "attn_mask": mask,
+        "dropout_p": 0.0,
+        "is_causal": False,
+        "enable_gqa": enable_gqa,
+    }
+    if q.device.type == "cuda":
+        params = torch.backends.cuda.SDPAParams(
+            q,
+            k,
+            v,
+            mask,
+            0.0,
+            False,
+            enable_gqa,
+        )
+        if torch.backends.cuda.can_use_cudnn_attention(params):
+            with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+                return F.scaled_dot_product_attention(q, k, v, **kwargs)
+    return F.scaled_dot_product_attention(q, k, v, **kwargs)
+
+
 def _timestep_embedding(t: Tensor, dim: int, *, dtype: torch.dtype) -> Tensor:
     half = dim // 2
     freqs = torch.exp(
@@ -178,18 +215,13 @@ def forward(
         q, k = self.qknorm(q, k)
         if freqs is not None:
             q, k = _apply_rope(q, k, freqs)
-        if self.kvheads != self.heads:
-            repeat = self.heads // self.kvheads
-            k = k.repeat_interleave(repeat, dim=1)
-            v = v.repeat_interleave(repeat, dim=1)
 
-        out = F.scaled_dot_product_attention(
+        out = _native_gqa_attention(
             q,
             k,
             v,
-            attn_mask=mask,
-            dropout_p=0.0,
-            is_causal=False,
+            mask,
+            enable_gqa=self.kvheads != self.heads,
         )
         out = rearrange(out, "b h l d -> b l (h d)")
         return self.wo(out * torch.sigmoid(gate))
diff --git a/tests/test_krea2_modeling.py b/tests/test_krea2_modeling.py
index 98abd93..0a0c761 100644
--- a/tests/test_krea2_modeling.py
+++ b/tests/test_krea2_modeling.py
@@ -179,11 +179,50 @@ def test_forward_rejects_invalid_temporal_and_conditioning_shapes() -> None:
         )
 
 
-def test_attention_gqa_preserves_shape() -> None:
+def test_attention_gqa_preserves_shape_and_does_not_expand_kv(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
     attention = Attention(dim=64, heads=4, kvheads=2)
     x = torch.randn(2, 7, 64)
     mask = torch.ones(2, 1, 1, 7, dtype=torch.bool)
+
+    observed: dict[str, object] = {}
+
+    def _fake_sdpa(q, k, v, **kwargs):
+        observed["heads"] = (q.shape[1], k.shape[1], v.shape[1])
+        observed["enable_gqa"] = kwargs.get("enable_gqa")
+        return torch.zeros_like(q)
+
+    monkeypatch.setattr(krea2_modeling.F, "scaled_dot_product_attention", _fake_sdpa)
+
     assert attention(x, mask=mask).shape == x.shape
+    assert observed == {"heads": (4, 2, 2), "enable_gqa": True}
+
+
+def test_native_gqa_matches_explicit_kv_expansion_on_cpu() -> None:
+    torch.manual_seed(17)
+    q = torch.randn(2, 4, 7, 16)
+    k = torch.randn(2, 2, 7, 16)
+    v = torch.randn(2, 2, 7, 16)
+    mask = torch.ones(2, 1, 1, 7, dtype=torch.bool)
+
+    native = krea2_modeling._native_gqa_attention(
+        q,
+        k,
+        v,
+        mask,
+        enable_gqa=True,
+    )
+    expanded = torch.nn.functional.scaled_dot_product_attention(
+        q,
+        k.repeat_interleave(2, dim=1),
+        v.repeat_interleave(2, dim=1),
+        attn_mask=mask,
+        dropout_p=0.0,
+        is_causal=False,
+    )
+
+    torch.testing.assert_close(native, expanded)
 
 
 def test_modeling_layer_only_imports_torch_einops_and_stdlib() -> None: