Feature(MInference): update the pycuda

app.py

@@ -5,18 +5,13 @@ subprocess.run(
     env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
     shell=True,
 )
-subprocess.run(
-    "pip install pycuda==2023.1",
-    shell=True,
-)
 
 import gradio as gr
 import os
 import spaces
-from transformers import GemmaTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
-from minference import MInference
 
 # Set an environment variable
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
@@ -63,8 +58,6 @@ h1 {
 model_name = "gradientai/Llama-3-8B-Instruct-262k"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")  # to("cuda:0") 
-minference_patch = MInference("minference", model_name)
-model = minference_patch(model)
 
 terminators = [
     tokenizer.eos_token_id,
@@ -87,6 +80,15 @@ def chat_llama3_8b(message: str,
     Returns:
         str: The generated response.
     """
+    if "has_patch" not in model.__dict__:
+        from minference import MInference
+        global model
+        subprocess.run(
+            "pip install pycuda==2023.1",
+            shell=True,
+        )
+        minference_patch = MInference("minference", model_name)
+        model = minference_patch(model)
     conversation = []
     for user, assistant in history:
         conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])

minference/modules/minference_forward.py

@@ -1,10 +1,16 @@
+# Copyright (c) 2024 Microsoft
+# Licensed under The MIT License [see LICENSE for details]
+
 import inspect
 import json
 import os
 from importlib import import_module
 
 from transformers.models.llama.modeling_llama import *
-from vllm.attention.backends.flash_attn import *
+from transformers.utils.import_utils import _is_package_available
+
+if _is_package_available("vllm"):
+    from vllm.attention.backends.flash_attn import *
 
 from ..ops.block_sparse_flash_attention import block_sparse_attention
 from ..ops.pit_sparse_flash_attention_v2 import vertical_slash_sparse_attention

minference/patch.py

@@ -780,6 +780,7 @@ def minference_patch(model, config):
         model.model, model.model.__class__
     )
     model.forward = forward_llama_for_causal_lm.__get__(model, model.__class__)
+    model.has_patch = True
 
     print("Patched model for minference..")
     return model