Update app.py

app.py

@@ -7,7 +7,7 @@ from threading import Thread
 import transformers.modeling_rope_utils as rope_utils
 from huggingface_hub import hf_hub_download
 
-# 1. Monkey-patch dynamic_rope_update if it's missing
+# 1. Monkey-patch dynamic_rope_update if missing
 if not hasattr(rope_utils, 'dynamic_rope_update'):
     def dynamic_rope_update(rope_forward):
         return rope_forward
@@ -18,7 +18,7 @@ model_id = "deepgrove/maple-preview"
 print("🔄 Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 
-# 2. Explicitly load and set the chat template
+# 2. Explicitly load chat template
 try:
     template_path = hf_hub_download(repo_id=model_id, filename="chat_template.jinja")
     with open(template_path, "r", encoding="utf-8") as f:
@@ -36,7 +36,7 @@ model = AutoModelForCausalLM.from_pretrained(
     attn_implementation="eager",
 )
 
-# 3. Patch the cu_seq_lens mismatch bug in fa3.py in memory
+# 3. Patch the cu_seq_lens mismatch bug in fa3.py
 def patch_fa3_module():
     for name, module in list(sys.modules.items()):
         if 'fa3' in name and 'maple' in name:
@@ -44,21 +44,40 @@ 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:
-                        # Flash Attention expects: (batch_size, q_len, num_heads, head_dim)
-                        total_tokens = query_states.shape[0] * query_states.shape[1]
+                    # FIX: Derive cu_seq_lens dynamically from query and key sequence lengths
+                    if cu_seq_lens_q is not None and query_states is not None and key_states is not None:
+                        # Flash Attention shape: (batch_size, seq_len, num_heads, head_dim)
+                        bsz = query_states.shape[0]
+                        q_len = query_states.shape[1]
+                        k_len = key_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)
-                        max_length_q = total_tokens
-                        max_length_k = total_tokens
+                        q_tokens = bsz * q_len
+                        k_tokens = bsz * k_len
+                        
+                        cu_seq_lens_q = torch.tensor([0, q_tokens], device=query_states.device, dtype=torch.int32)
+                        cu_seq_lens_k = torch.tensor([0, k_tokens], device=key_states.device, dtype=torch.int32)
+                        max_length_q = q_len
+                        max_length_k = k_len
                         
                     return orig_fa_forward(
                         query_states, key_states, value_states, attention_mask, query_length, is_causal,
@@ -115,4 +134,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    demo.queue().launch(ssr=False)