From dcfb92ddbcd2f5a6fe919f0810e4aa82315e462e Mon Sep 17 00:00:00 2001
From: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com>
Date: Mon, 5 Jan 2026 16:25:41 +0800
Subject: [PATCH] [diffusion] fix: fix using the sub-optimal fa kernel (#16382)

---
 .../runtime/layers/attention/backends/flash_attn.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
index 4b033e527..78297a380 100644
--- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
+++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py
@@ -201,18 +201,15 @@ class FlashAttentionImpl(AttentionImpl):
 
         if use_upstream:
             bsz, seqlen, nheads_q, d = q_shape
-            bsz_k, seqlen_k, nheads_k, d_k = k_shape
-            bsz_v, seqlen_v, nheads_v, d_v = v_shape
-            q_ = query.contiguous().reshape(bsz * seqlen, nheads_q, d)
-            k_ = key.contiguous().reshape(bsz * seqlen, nheads_k, d)
-            v_ = value.contiguous().reshape(bsz * seqlen, nheads_v, d)
-            cu = _get_cu_seqlens(q_.device.index, bsz, seqlen)
+            q_ = query.contiguous()
+            k_ = key.contiguous()
+            v_ = value.contiguous()
             out = flash_attn_varlen_func_upstream(
                 q_,
                 k_,
                 v_,
-                cu,
-                cu,
+                None,
+                None,
                 seqlen,
                 seqlen,
                 softmax_scale=self.softmax_scale,