Update app.py

app.py

@@ -20,80 +20,119 @@ app.add_middleware(
 MODEL_REPO = "unsloth/Qwen3-4B-GGUF"
 MODEL_FILE = "Qwen3-4B-Q4_K_M.gguf"
 
-MAX_HISTORY = 4          # Giảm để tiết kiệm token xử lý
-MAX_CTX     = 4096       # Khung ngữ cảnh tối đa
-MAX_TOKENS  = 2048       # Giới hạn token sinh ra để giải phóng tài nguyên sớm
-THREADS     = 2          # Số luồng CPU, khớp với vCPU bạn có
+# ── Triết lý tối ưu ───────────────────────────────────────────────────────
+# RAM 18GB dư dả → nhét hết vào RAM, dùng prefix cache để CPU
+# không phải recompute system prompt mỗi request
+# n_batch = 4096 (sweet spot) — đủ để prefill nhanh mà không gây RAM spike
+# ─────────────────────────────────────────────────────────────────────────
+MAX_HISTORY = 6
+MAX_CTX     = 8192
+MAX_TOKENS  = 2048
+THREADS     = 2
+
+# System prompt cố định — sẽ được cache sẵn vào KV cache lúc startup
+# CPU chỉ tính 1 lần duy nhất, mọi request sau dùng lại cache này
+DEFAULT_SYSTEM = "Bạn là trợ lý AI, trả lời bằng tiếng Việt ngắn gọn."
 
 llm: Optional[Llama] = None
 
+
 @app.on_event("startup")
 async def startup_event():
     global llm
-    # 1. Tải model (xóa file cũ nếu lỗi)
+
     if os.path.exists(MODEL_FILE) and os.path.getsize(MODEL_FILE) < 1_000_000:
         os.remove(MODEL_FILE)
+
     if not os.path.exists(MODEL_FILE):
         print(f"Downloading {MODEL_FILE}...")
         hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE, local_dir=".")
         print("Download done!")
 
-    # 2. Khởi tạo model với các tham số tối ưu cho CPU và RAM
-    print("Loading model with RAM-optimized settings...")
+    print("Loading model — RAM-heavy, CPU-light mode...")
     llm = Llama(
-        model_path       = MODEL_FILE,
-        n_ctx            = MAX_CTX,
-        n_threads        = THREADS,
-        n_threads_batch  = THREADS,
-        n_batch          = MAX_CTX,      # Quan trọng: bằng với n_ctx để tối ưu cache
-        n_ubatch         = 512,          # Kích thước micro-batch, cân bằng với n_batch
-        n_gpu_layers     = 0,
-        verbose          = False,
-        use_mmap         = False,        # Tắt mmap, load toàn bộ model vào RAM
-        use_mlock        = True,         # Khóa model trong RAM, tránh swap ra ổ cứng
-        flash_attn       = True,         # Bật Flash Attention, giảm tải bộ nhớ
-        cache_type_k     = "q8_0",       # Lượng tử hóa KV cache (Key) xuống 8-bit
-        cache_type_v     = "q8_0",       # Lượng tử hóa KV cache (Value) xuống 8-bit
+        model_path      = MODEL_FILE,
+
+        # ── Context & batch ───────────────────────────────────────────────
+        n_ctx           = MAX_CTX,
+        n_batch         = 512 ,   # Nhỏ vừa tay CPU: 2 vCPU không bị nghẹt khi prefill
+        n_ubatch        = 512   ,    # Giữ nhỏ: ổn định hơn khi decode
+
+        # ── CPU ───────────────────────────────────────────────────────────
+        n_threads       = THREADS,
+        n_threads_batch = THREADS,
+        n_gpu_layers    = 0,
+
+        # ── RAM: load toàn bộ, khóa lại, không swap ──────────────────────
+        use_mmap        = False,
+        use_mlock       = True,
+
+        # ── KV Cache quantize — ăn RAM ít hơn, CPU vẫn nhẹ ───────────────
+        cache_type_k    = "q4_0",
+        cache_type_v    = "q4_0",
+
+        # ── Prefix cache: CPU tính system prompt 1 lần rồi thôi ──────────
+        last_n_tokens_size = 64,  # Cửa sổ detect prefix trùng
+
+        flash_attn      = True,
+        verbose         = False,
+    )
+
+    # ── Warm up prefix cache với system prompt ────────────────────────────
+    # Gọi 1 lần lúc startup để KV cache của system prompt được lưu sẵn
+    # Mọi request sau có cùng system prompt → CPU bỏ qua phần này hoàn toàn
+    print("Warming up prefix cache...")
+    warmup_msgs = [
+        {"role": "system", "content": DEFAULT_SYSTEM},
+        {"role": "user",   "content": "hi"},
+    ]
+    _ = llm.create_chat_completion(
+        messages   = warmup_msgs,
+        max_tokens = 1,
+        stream     = False,
     )
-    print("Model ready! All weights and caches are locked in RAM.")
+    print("Prefix cache warmed up! Model ready.")
+
 
 class Message(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
     prompt: str
     history: List[Message] = []
-    system_prompt: Optional[str] = None
+    system_prompt: Optional[str] = None  # Để None → tận dụng prefix cache
     max_tokens: int = MAX_TOKENS
     temperature: float = 0.7
     top_p: float = 0.9
 
+
 def build_messages(req: ChatRequest) -> list:
-    # System prompt ngắn gọn, tiết kiệm token
-    system = req.system_prompt or "Bạn là trợ lý AI, trả lời bằng tiếng Việt ngắn gọn."
+    # Dùng DEFAULT_SYSTEM nếu không truyền system_prompt
+    # → prefix cache luôn hit, CPU không recompute
+    system = req.system_prompt or DEFAULT_SYSTEM
     msgs = [{"role": "system", "content": system}]
-    
-    # Chỉ giữ lại lịch sử tối đa MAX_HISTORY tin nhắn (mỗi lượt gồm user + assistant)
-    recent_history = req.history[-(MAX_HISTORY * 2):]
-    for msg in recent_history:
+
+    recent = req.history[-(MAX_HISTORY * 2):]
+    for msg in recent:
         if msg.role in ("user", "assistant") and msg.content.strip():
             if msgs[-1]["role"] != msg.role:
                 msgs.append({"role": msg.role, "content": msg.content.strip()})
-    
-    # Tránh trùng lặp role user cuối
+
     if msgs[-1]["role"] == "user":
         msgs.pop()
     msgs.append({"role": "user", "content": req.prompt.strip()})
     return msgs
 
+
 @app.post("/chat")
 async def chat(req: ChatRequest):
     if llm is None:
         raise HTTPException(503, "Model chưa sẵn sàng, thử lại sau!")
     if not req.prompt.strip():
         raise HTTPException(400, "Prompt trống")
-    if len(req.prompt) > 4000:
+    if len(req.prompt) > 8000:
         raise HTTPException(400, "Prompt quá dài")
 
     messages = build_messages(req)
@@ -111,11 +150,11 @@ async def chat(req: ChatRequest):
                 delta = chunk["choices"][0]["delta"].get("content", "")
                 if delta:
                     full += delta
-                    yield f"data: {json.dumps({'delta': delta})}\n\n"
+                    yield f"data: {json.dumps({'delta': delta}, ensure_ascii=False)}\n\n"
         except Exception as e:
             yield f"data: {json.dumps({'delta': f'[Lỗi: {str(e)}]'})}\n\n"
         finally:
-            print(f">> Done: {full[:80]}")
+            print(f">> Done ({len(full)} chars): {full[:80]}")
             yield "data: [DONE]\n\n"
 
     return StreamingResponse(
@@ -124,17 +163,20 @@ async def chat(req: ChatRequest):
         headers    = {"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
     )
 
+
 @app.get("/")
 async def root():
     return {
-        "status"   : "ok" if llm else "loading",
-        "model"    : MODEL_FILE,
-        "message"  : "Model ready (RAM-optimized)!" if llm else "Model đang tải...",
+        "status"  : "ok" if llm else "loading",
+        "model"   : MODEL_FILE,
+        "message" : "Model ready (prefix cache active)!" if llm else "Model đang tải...",
     }
 
+
 @app.get("/health")
 async def health():
     return {"status": "healthy", "model_loaded": llm is not None}
 
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)