Update app.py

app.py

@@ -44,30 +44,16 @@ def patch_fa3_module():
                 orig_fa_forward = module._flash_attention_forward
                 
                 def patched_fa_forward(
-                    query_states,
-                    key_states,
-                    value_states,
-                    attention_mask,
-                    query_length,
-                    is_causal,
-                    dropout=0.0,
-                    position_ids=None,
-                    softmax_scale=None,
-                    sliding_window=None,
-                    use_top_left_mask=False,
-                    softcap=None,
-                    deterministic=None,
-                    cu_seq_lens_q=None,
-                    cu_seq_lens_k=None,
-                    max_length_q=None,
-                    max_length_k=None,
-                    target_dtype=None,
-                    **kwargs,
+                    query_states, key_states, value_states, attention_mask, query_length, is_causal,
+                    dropout=0.0, position_ids=None, softmax_scale=None, sliding_window=None,
+                    use_top_left_mask=False, softcap=None, deterministic=None,
+                    cu_seq_lens_q=None, cu_seq_lens_k=None, max_length_q=None,
+                    max_length_k=None, target_dtype=None, **kwargs,
                 ):
                     # FIX: Override cu_seq_lens to perfectly match the actual number of tokens in query_states
                     if cu_seq_lens_q is not None and query_states is not None:
-                        # query_states shape at this point is (batch_size, num_heads, q_len, head_dim)
-                        total_tokens = query_states.shape[0] * query_states.shape[2]
+                        # Flash Attention expects: (batch_size, q_len, num_heads, head_dim)
+                        total_tokens = query_states.shape[0] * query_states.shape[1]
                         
                         cu_seq_lens_q = torch.tensor([0, total_tokens], device=query_states.device, dtype=torch.int32)
                         cu_seq_lens_k = torch.tensor([0, total_tokens], device=key_states.device, dtype=torch.int32)