Update app.py

app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import torch
 import spaces
 import sys
-import re
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 import transformers.modeling_rope_utils as rope_utils
@@ -37,32 +36,55 @@ model = AutoModelForCausalLM.from_pretrained(
     attn_implementation="eager",
 )
 
-# 3. Patch the seq_len bug in fa3.py in memory
+# 3. Patch the cu_seq_lens mismatch bug in fa3.py in memory
 def patch_fa3_module():
-    try:
-        fa3_path = hf_hub_download(repo_id=model_id, filename="fa3.py")
-        with open(fa3_path, "r", encoding="utf-8") as f:
-            fa3_source = f.read()
-
-        # The model authors made a typo: query is (bsz, num_heads, q_len, head_dim)
-        # so q_len is at index 2, not 1!
-        if "seq_len = query.shape[1]" in fa3_source:
-            fa3_source = fa3_source.replace("seq_len = query.shape[1]", "seq_len = query.shape[2]")
-            print("Fixed seq_len bug in fa3 source string.")
-        else:
-            fa3_source = re.sub(r'seq_len\s*=\s*query\.shape\[1\]', 'seq_len = query.shape[2]', fa3_source)
-            print("Attempted regex fix on fa3 source string.")
-
-        # Inject the patched source code directly into the loaded module in sys.modules
-        for name, module in list(sys.modules.items()):
-            if 'fa3' in name and 'maple' in name:
-                print(f"Injecting patched fa3 into {name}")
-                if not hasattr(module, '__package__') or not module.__package__:
-                    module.__package__ = name.rsplit('.', 1)[0]
-                exec(compile(fa3_source, fa3_path, 'exec'), vars(module))
-                print("✅ Injection successful.")
-    except Exception as e:
-        print(f"Failed to patch fa3 module: {e}")
+    for name, module in list(sys.modules.items()):
+        if 'fa3' in name and 'maple' in name:
+            if hasattr(module, '_flash_attention_forward'):
+                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,
+                ):
+                    # 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]
+                        
+                        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)
+                        max_length_q = total_tokens
+                        max_length_k = total_tokens
+                        
+                    return orig_fa_forward(
+                        query_states, key_states, value_states, attention_mask, query_length, is_causal,
+                        dropout, position_ids, softmax_scale, sliding_window, use_top_left_mask, softcap,
+                        deterministic, cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k, target_dtype, **kwargs
+                    )
+                    
+                module._flash_attention_forward = patched_fa_forward
+                print("✅ Patched _flash_attention_forward to fix cu_seq_lens mismatch.")
+                return True
+    print("⚠️ Could not find fa3 module to patch.")
+    return False
 
 patch_fa3_module()