Update app.py

app.py

@@ -13,7 +13,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from kernels import get_kernel
 
-vllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")
+#vllm_flash_attn3 = get_kernel("kernels-community/vllm-flash-attn3")
 
 #torch._dynamo.config.disable = True
 
@@ -28,7 +28,7 @@ def load_model():
         MODEL_ID,
         torch_dtype=torch.bfloat16,  # if device == "cuda" else torch.float32,
         device_map="auto",  # if device == "cuda" else None,
-        attn_implementation="flash_attention_3", # "flash_attention_2", # 
+        attn_implementation="kernels-community/vllm-flash-attn3", # "flash_attention_2", # 
     )  # .cuda()
     print(f"Selected device:", device)
     return model, tokenizer, device