Upload app.py

app.py

@@ -14,6 +14,13 @@ from threading import Thread
 MODEL_NAME = os.getenv("MODEL_NAME", "dllm-hub/Qwen3-0.6B-diffusion-bd3lm-v0.1")
 IS_DIFFUSION = "diffusion" in MODEL_NAME.lower()
 
+# CPU optimizations: tune threads and disable denormals before any torch work
+if not torch.cuda.is_available():
+    torch.set_num_threads(os.cpu_count() or 4)
+    torch.set_num_interop_threads(1)
+    torch.set_flush_denormal(True)
+    print(f"CPU mode: using {torch.get_num_threads()} threads, interop=1, flush_denormal=True")
+
 # Dynamic initialization layer targeting Diffusion Language Models
 if IS_DIFFUSION:
     try:
@@ -79,13 +86,11 @@ def clone_past_key_values(pkv):
     """Clone KV-cache. Fast path for tuples and Cache objects; falls back to deepcopy."""
     if pkv is None:
         return None
-    # Fast path: legacy tuple format
     if isinstance(pkv, tuple):
         return tuple(
             (k.clone() if k is not None else None, v.clone() if v is not None else None)
             for k, v in pkv
         )
-    # Fast path: transformers Cache objects (DynamicCache, etc.)
     if hasattr(pkv, 'key_cache') and hasattr(pkv, 'value_cache'):
         try:
             new_cache = pkv.__class__()
@@ -97,7 +102,6 @@ def clone_past_key_values(pkv):
             return new_cache
         except Exception:
             pass
-    # Fallback
     return copy.deepcopy(pkv)
 
 
@@ -322,33 +326,39 @@ def generate_stream(model, tokenizer, prompt, steps=128, max_new_tokens=128, blo
 def load_model():
     global model, tokenizer, device
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    print(f"Initializing {MODEL_NAME} on {device}... (Diffusion Strategy Flag = {IS_DIFFUSION})")
+    # Use float32 on CPU (bfloat16 is very slow on CPU); keep bfloat16 on GPU
+    dtype = torch.bfloat16 if device == "cuda" else torch.float32
+
+    print(f"Initializing {MODEL_NAME} on {device} with dtype={dtype}... (Diffusion Strategy Flag = {IS_DIFFUSION})")
+
     if IS_DIFFUSION:
         model = AutoModelForMaskedLM.from_pretrained(
             MODEL_NAME,
-            torch_dtype=torch.bfloat16,
+            torch_dtype=dtype,
             trust_remote_code=True
         ).to(device).eval()
     else:
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            torch_dtype=torch.bfloat16,
+            torch_dtype=dtype,
             trust_remote_code=False
         ).to(device).eval()
-    # Compile model for faster inference — ONLY for standard causal models
-    if not IS_DIFFUSION:
+
+    # Compile model only on GPU causal models; skip on CPU and diffusion
+    if device == "cuda" and not IS_DIFFUSION:
         try:
             model = torch.compile(model, mode="reduce-overhead", fullgraph=False)
             print("Model compiled with torch.compile.")
         except Exception as e:
             print(f"torch.compile skipped: {e}")
     else:
-        print("Diffusion model loaded without torch.compile (custom FX code incompatible with Dynamo).")
+        print("torch.compile skipped (CPU or diffusion).")
+
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         trust_remote_code=IS_DIFFUSION
     )
-    print("Model compilation completed and loaded into memory workspace.")
+    print("Model loaded into memory workspace.")
 
 
 @app.route('/health', methods=['GET'])
@@ -371,7 +381,6 @@ def generate_text():
         {"role": "system", "content": system_prompt},
         {"role": "user", "content": prompt}
     ]
-    # enable_thinking=False for ALL routes to prevent Qwen3 from leaking internal monologue
     encoded = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
@@ -399,7 +408,8 @@ def generate_text():
             max_new_tokens=max_new_tokens,
             temperature=temperature,
             do_sample=True if temperature > 0 else False,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            use_cache=True,
         )
         generated_ids = output_ids[0, input_ids.shape[-1]:]
         generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
@@ -452,7 +462,8 @@ def generate_text_stream():
         output_ids = model.generate(
             input_ids, max_new_tokens=max_new_tokens, temperature=temperature,
             do_sample=True if temperature > 0 else False,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            use_cache=True,
         )
         generated_ids = output_ids[0, input_ids.shape[-1]:]
         generated_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
@@ -505,6 +516,7 @@ def generate_text_sse():
                 temperature=temperature,
                 do_sample=True if temperature > 0 else False,
                 pad_token_id=tokenizer.eos_token_id,
+                use_cache=True,
             )
 
             def _generate():
@@ -516,7 +528,7 @@ def generate_text_sse():
 
             accumulated = []
             for text in streamer:
-                if not text:  # skip empty chunks
+                if not text:
                     continue
                 accumulated.append(text)
                 current = "".join(accumulated)