server.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import uvicorn
import threading
import time
import uuid
import json
from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse
from fastapi.responses import StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field, ConfigDict
from typing import List, Optional, Union, Dict, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
from transformers import BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from torch.utils.checkpoint import checkpoint
from transformers.modeling_outputs import CausalLMOutputWithPast
import asyncio
# ==========================================
# 1. 初始化 FastAPI App 與 跨域(CORS) 設定
# ==========================================
app = FastAPI(title="OpenAI-Compatible FrankenQwen Server", version="1.0")

# 開啟 CORS，讓前端 UI (如 Open WebUI) 可以順利連線
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
MAX_CONCURRENT_REQUESTS = 4
gpu_semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
def get_vram_status():
    import torch
    if not torch.cuda.is_available():
        return {"status": "CUDA 不可用"}
    
    # 轉換為 GiB (除以 1024^3)
    allocated = torch.cuda.memory_allocated() / (1024 ** 3)
    reserved = torch.cuda.memory_reserved() / (1024 ** 3)
    peak = torch.cuda.max_memory_allocated() / (1024 ** 3)
    
    return {
        "current_tensor_used_gib": round(allocated, 2),
        "pytorch_reserved_gib": round(reserved, 2),
        "peak_vram_used_gib": round(peak, 2)
    }
# ==========================================
# 2. 原始模型組件與邏輯 (保持不變)
# ==========================================
class HiddenStateCompressor(nn.Module):
    def __init__(self, dim, ratio=4):
        super().__init__()
        self.ratio = ratio
        self.wgate = nn.Linear(dim, 1, bias=False)
        self.ape = nn.Parameter(torch.empty(ratio, 1))
        nn.init.normal_(self.ape, std=0.02)
        
    def forward(self, x):
        B, L, D = x.shape
        assert L % self.ratio == 0, "輸入長度必須是壓縮倍率的整數倍"
        orig_dtype = x.dtype
        target_dtype = self.wgate.weight.dtype
        x_reshaped = x.view(B, L // self.ratio, self.ratio, D).to(target_dtype)
        scores = self.wgate(x_reshaped) + self.ape.view(1, 1, self.ratio, 1)
        scores_fp32 = scores.to(torch.float32)
        scores_probs = F.softmax(scores_fp32, dim=2).to(target_dtype)
        compressed_x = (x_reshaped * scores_probs).sum(dim=2)
        return compressed_x.to(orig_dtype)

def get_custom_qwen_forward(compressor_module, window_size=128, ratio=4):
    from collections.abc import Callable
    from transformers.models.qwen3_5.modeling_qwen3_5 import ALL_ATTENTION_FUNCTIONS, apply_rotary_pos_emb, eager_attention_forward, FlashAttentionKwargs
    from transformers.processing_utils import Unpack
    from transformers.cache_utils import Cache

    def forward2(
        self, hidden_states: torch.Tensor, position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None, past_key_values: Cache | None = None, **kwargs: Unpack[FlashAttentionKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        cos, sin = position_embeddings
        bsz, q_len, _ = hidden_states.size()
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        hidden_states = hidden_states.contiguous()
        query_states, gate = torch.chunk(
            self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2), 2, dim=-1
        )
        gate = gate.contiguous().reshape(*input_shape, -1)
        query_states = self.q_norm(query_states.contiguous().view(hidden_shape)).transpose(1, 2)
        key_states, value_states = None, None
        do_compress = getattr(self, "compress_enabled", True)

        if do_compress:
            if q_len > 1:  # Prefill
                window = min(window_size, q_len)
                global_len = q_len - window
                cutoff = global_len - (global_len % ratio)
                compressed = compressor_module(hidden_states[:, :cutoff, :]) if cutoff > 0 else hidden_states[:, :0, :]
                residual = hidden_states[:, cutoff:global_len, :] if global_len > cutoff else hidden_states[:, :0, :]
                tail = hidden_states[:, global_len:, :]
                kv_hidden = torch.cat([compressed, residual, tail], dim=1).contiguous()

                compressed_pos = torch.arange(ratio - 1, cutoff, ratio, device=hidden_states.device, dtype=torch.long) if cutoff > 0 else torch.arange(0, device=hidden_states.device, dtype=torch.long)
                residual_pos = torch.arange(cutoff, global_len, device=hidden_states.device, dtype=torch.long)
                tail_pos = torch.arange(global_len, q_len, device=hidden_states.device, dtype=torch.long)
                kv_positions = torch.cat([compressed_pos, residual_pos, tail_pos], dim=0)
                kv_seq_len = kv_positions.size(0)

                q_pos = torch.arange(q_len, device=hidden_states.device).view(q_len, 1)
                k_pos = kv_positions.view(1, kv_seq_len)
                custom_causal_mask = (q_pos >= k_pos)
                additive_mask = torch.zeros((q_len, kv_seq_len), dtype=hidden_states.dtype, device=hidden_states.device)
                additive_mask.masked_fill_(~custom_causal_mask, torch.finfo(hidden_states.dtype).min)
                additive_mask = additive_mask.view(1, 1, q_len, kv_seq_len).expand(bsz, 1, q_len, kv_seq_len)

                attention_mask = (attention_mask[..., kv_positions] + additive_mask) if attention_mask is not None else additive_mask
                kv_hidden_shape = (bsz, kv_seq_len, -1, self.head_dim)
                key_states = self.k_norm(self.k_proj(kv_hidden).view(kv_hidden_shape)).transpose(1, 2).contiguous()
                value_states = self.v_proj(kv_hidden).view(kv_hidden_shape).transpose(1, 2).contiguous()

                query_states, _ = apply_rotary_pos_emb(query_states, query_states, cos, sin)
                k_cos = cos[:, kv_positions, :] if cos.dim() == 3 else cos[kv_positions, :]
                k_sin = sin[:, kv_positions, :] if sin.dim() == 3 else sin[kv_positions, :]
                _, key_states = apply_rotary_pos_emb(key_states, key_states, k_cos, k_sin)
                if past_key_values is not None:
                    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
            else:  # Decode
                new_key = self.k_norm(self.k_proj(hidden_states).view(bsz, 1, -1, self.head_dim)).transpose(1, 2)
                new_value = self.v_proj(hidden_states).view(bsz, 1, -1, self.head_dim).transpose(1, 2)
                query_states, _ = apply_rotary_pos_emb(query_states, query_states, cos, sin)
                _, new_key = apply_rotary_pos_emb(new_key, new_key, cos, sin)

                target_dtype = hidden_states.dtype
                new_key = new_key.to(target_dtype).contiguous()
                new_value = new_value.to(target_dtype).contiguous()

                if past_key_values is not None:
                    if not hasattr(past_key_values, "uncompressed_buffer_k"):
                        past_key_values.uncompressed_buf_h = {}
                        past_key_values.uncompressed_buffer_k = {}
                        past_key_values.uncompressed_buffer_v = {}

                    layer_idx = self.layer_idx
                    buf_k = past_key_values.uncompressed_buffer_k.get(layer_idx, torch.empty((bsz, self.config.num_key_value_heads, 0, self.head_dim), device=new_key.device, dtype=target_dtype))
                    buf_v = past_key_values.uncompressed_buffer_v.get(layer_idx, torch.empty((bsz, self.config.num_key_value_heads, 0, self.head_dim), device=new_value.device, dtype=target_dtype))
                    buf_h = past_key_values.uncompressed_buf_h.get(layer_idx, torch.empty((bsz, 0, hidden_states.size(-1)), device=hidden_states.device, dtype=hidden_states.dtype))
                    
                    buf_h = torch.cat([buf_h, hidden_states], dim=1)
                    buf_k = torch.cat([buf_k, new_key], dim=2)
                    buf_v = torch.cat([buf_v, new_value], dim=2)

                    buffer_len = buf_k.shape[2]
                    if getattr(self, "compress_enabled", True) and buffer_len >= (window_size + ratio):
                        to_compress_h = buf_h[:, :ratio, :]
                        compressed_h = compressor_module(to_compress_h)
                        compressed_k = self.k_norm(self.k_proj(compressed_h).view(bsz, 1, -1, self.head_dim)).transpose(1, 2)
                        compressed_v = self.v_proj(compressed_h).view(bsz, 1, -1, self.head_dim).transpose(1, 2)

                        past_key_values.update(compressed_k, compressed_v, layer_idx)
                        buf_h = buf_h[:, ratio:, :]
                        buf_k = buf_k[:, :, ratio:, :]
                        buf_v = buf_v[:, :, ratio:, :]

                    past_key_values.uncompressed_buffer_k[layer_idx] = buf_k
                    past_key_values.uncompressed_buffer_v[layer_idx] = buf_v
                    past_key_values.uncompressed_buf_h[layer_idx] = buf_h

                    layer_cache = past_key_values.layers[layer_idx] if hasattr(past_key_values, "layers") else None
                    history_k = layer_cache.keys if layer_cache else past_key_values.key_cache[layer_idx]
                    history_v = layer_cache.values if layer_cache else past_key_values.value_cache[layer_idx]

                    key_states = torch.cat([history_k, buf_k], dim=2) if history_k is not None and history_k.shape[2] > 0 else buf_k
                    value_states = torch.cat([history_v, buf_v], dim=2) if history_v is not None and history_v.shape[2] > 0 else buf_v
                    kv_seq_len = past_key_values.get_seq_length(self.layer_idx) + buf_k.shape[2]
                else:
                    key_states, value_states = new_key, new_value
        else:
            kv_hidden_shape = (bsz, q_len, -1, self.head_dim)
            key_states = self.k_norm(self.k_proj(hidden_states).view(kv_hidden_shape)).transpose(1, 2).contiguous()
            value_states = self.v_proj(hidden_states).view(kv_hidden_shape).transpose(1, 2).contiguous()
            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
            if past_key_values is not None:
                key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
            kv_seq_len = key_states.shape[2]

        target_dtype = hidden_states.dtype
        query_states, key_states, value_states = query_states.to(target_dtype).contiguous(), key_states.to(target_dtype).contiguous(), value_states.to(target_dtype).contiguous()

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(self.config._attn_implementation, eager_attention_forward)
        attn_output, attn_weights = attention_interface(
            self, query_states, key_states, value_states, attention_mask,
            dropout=0.0 if not self.training else self.attention_dropout, scaling=self.scaling, **kwargs,
        )
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = attn_output * torch.sigmoid(gate)
        return self.o_proj(attn_output), attn_weights

    return forward2

def get_chunked_causal_forward(chunk_size=512):
    def custom_causal_forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None, return_dict=None, **kwargs):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, **kwargs)
        hidden_states = outputs[0]
        loss, logits = None, None

        if labels is not None:
            shift_hidden_states = hidden_states[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            seq_len = shift_hidden_states.shape[1]
            total_loss = 0.0
            active_tokens = (shift_labels != -100).sum().clamp(min=1)

            def compute_chunk_loss(h_chunk, lbl_chunk):
                c_logits = self.lm_head(h_chunk).float()
                return F.cross_entropy(c_logits.reshape(-1, c_logits.size(-1)), lbl_chunk.reshape(-1), reduction='sum')

            for i in range(0, seq_len, chunk_size):
                end_i = min(i + chunk_size, seq_len)
                c_loss = checkpoint(compute_chunk_loss, shift_hidden_states[:, i:end_i, :], shift_labels[:, i:end_i], use_reentrant=False)
                total_loss += c_loss
            loss = total_loss / active_tokens
        else:
            logits = self.lm_head(hidden_states)

        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
    return custom_causal_forward

def load_custom_checkpoint(model, filepath="/vault/franken_qwen2B_compressor_checkpoint.pt"):
    if not os.path.exists(filepath):
        print(f"⚠️ 找不到權重檔 {filepath}！將使用原始權重。")
        return model
    print(f"\n📂 正在載入客製化權重: {filepath}")
    trainable_state_dict = torch.load(filepath, map_location="cpu", weights_only=True)
    clean_state_dict = { (k[7:] if k.startswith('module.') else k): v for k, v in trainable_state_dict.items() }
    model.load_state_dict(clean_state_dict, strict=False)
    print("✅ 客製化權重完美掛載！")
    return model

def build_franken_qwen(qwen_id="Qwen/Qwen3.6-27B", device="cuda"):
    quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)
    qwen_model = AutoModelForCausalLM.from_pretrained(qwen_id, dtype=torch.bfloat16, device_map="auto", quantization_config=quantization_config)
    #qwen_model = prepare_model_for_kbit_training(qwen_model)
    
    qwen_model.compressors = nn.ModuleList()
    for i, layer in enumerate(qwen_model.model.layers):
        if i % 4 == 3:
            compressor = HiddenStateCompressor(dim=qwen_model.config.hidden_size, ratio=4).to(device, dtype=torch.bfloat16)
            qwen_model.compressors.append(compressor)
            custom_fw = get_custom_qwen_forward(compressor, window_size=128, ratio=4)
            layer.self_attn.forward = custom_fw.__get__(layer.self_attn, type(layer.self_attn))
            
    qwen_model.forward = get_chunked_causal_forward(chunk_size=512).__get__(qwen_model, type(qwen_model))
    
    for param in qwen_model.parameters(): param.requires_grad = False
    for compressor in qwen_model.compressors:
        for param in compressor.parameters(): param.requires_grad = True
    if qwen_id=="Qwen/Qwen3.6-27B":
        lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM")
        qwen_model = get_peft_model(qwen_model, lora_config)
        return load_custom_checkpoint(qwen_model, filepath="./Downloads/franken_qwen27B_compressor_lora.pt")
    elif qwen_id=="Qwen/Qwen3.5-2B":
        return load_custom_checkpoint(qwen_model, filepath="./Downloads/franken_qwen2B_compressor_checkpoint.pt")
    qwen_model.enable_input_require_grads()
    return load_custom_checkpoint(qwen_model, filepath="/vault/franken_qwen27B_compressor_lora.pt")

# 全域模型初始化
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
MODEL_ID = "Qwen/Qwen3.5-2B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = build_franken_qwen(qwen_id=MODEL_ID, device="cuda")
model=model.get_base_model() if hasattr(model, "get_base_model") else model
model=torch.compile(model)
model.eval()
print(get_vram_status())
# ==========================================
# 3. OpenAI 官方相容的 Pydantic 資料模型定義
# ==========================================
class ChatMessage(BaseModel):
    role: str
    content: str

class ChatCompletionRequest(BaseModel):
    # 🌟 關鍵優化：extra="allow" 允許 lm-eval 傳入額外的 OpenAI 參數 (如 n, logprobs) 而不報錯
    model_config = ConfigDict(extra="allow") 
    model: str
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    max_tokens: Optional[int] = Field(default=512, alias="max_tokens")
    stream: Optional[bool] = False
    stop: Optional[Union[str, List[str]]] = None # 🌟 接收 lm-eval 的停止符號
class CompletionRequest(BaseModel):
    # 🌟 針對標準文字補全的 Request Schema
    model_config = ConfigDict(extra="allow")
    model: str
    prompt: Union[str, List[str]]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.9
    max_tokens: Optional[int] = Field(default=512, alias="max_tokens")
    stream: Optional[bool] = False
    stop: Optional[Union[str, List[str]]] = None
# ==========================================
# 4. OpenAI 相容核心端點 (/v1/chat/completions)
# ==========================================
@app.exception_handler(Exception)
async def universal_exception_handler(request: Request, exc: Exception):
    return JSONResponse(
        status_code=500,
        content={"error": {"message": f"Internal Server Error: {str(exc)}", "type": "server_error"}}
    )
def prepare_gen_config(request, inputs):
    """公用生成參數建立函式"""
    gen_config = {
        **inputs,
        "max_new_tokens": request.max_tokens,
        "do_sample": True if request.temperature > 0 else False,
        "temperature": request.temperature if request.temperature > 0 else None,
        "top_p": request.top_p,
        "repetition_penalty": 1.05,
        "use_cache": True,
        "pad_token_id": tokenizer.eos_token_id,
        "eos_token_id": [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids("<|im_end|>")]
    }
    # 🌟 處理 lm-eval 傳過來的 stop 截斷字串
    if request.stop:
        gen_config["stop_strings"] = [request.stop] if isinstance(request.stop, str) else request.stop
        gen_config["tokenizer"] = tokenizer
    return gen_config
@app.post("/v1/chat/completions")
async def chat_completions_endpoint(request: ChatCompletionRequest):
    async with gpu_semaphore: # 佇列鎖管控
        formatted_messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        input_text = tokenizer.apply_chat_template(formatted_messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(input_text, return_tensors="pt", add_special_tokens=False).to(model.device)
        gen_config = prepare_gen_config(request, inputs)

        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        gen_config["streamer"] = streamer
        threading.Thread(target=lambda: torch.no_grad()(model.generate)(**gen_config)).start()

        if not request.stream:
            # 🌟 優化：將同步阻塞的 "".join 丟進線程池，解放 Event Loop
            full_text = await asyncio.to_thread(lambda: "".join([new_text for new_text in streamer]))
            return {
                "id": f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", "created": int(time.time()), "model": request.model,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": full_text}, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": inputs["input_ids"].shape[1], "completion_tokens": len(tokenizer.encode(full_text)), "total_tokens": inputs["input_ids"].shape[1] + len(tokenizer.encode(full_text))}
            }
        else:
            async def openai_chat_stream():
                chunk_id, created_time = f"chatcmpl-{uuid.uuid4()}", int(time.time())
                yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
                
                # 🌟 優化：建立安全不阻塞事件循環的取 Token 函式
                def get_next_token():
                    try:
                        return next(streamer), False
                    except StopIteration:
                        return "", True

                while True:
                    # 🌟 關鍵：每次拿 Token 都讓出主執行緒，讓監控端點可以插隊進來執行
                    new_text, is_done = await asyncio.to_thread(get_next_token)
                    if is_done:
                        break
                    if new_text:
                        yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'index': 0, 'delta': {'content': new_text}, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
                        
                yield f"data: {json.dumps({'id': chunk_id, 'object': 'chat.completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'index': 0, 'delta': {}, 'finish_reason': 'stop'}]}, ensure_ascii=False)}\n\n"
                yield "data: [DONE]\n\n"
            return StreamingResponse(openai_chat_stream(), media_type="text/event-stream")

# ----- 端點 B：文字接龍 (/v1/completions) -----
@app.post("/v1/completions")
async def completions_endpoint(request: CompletionRequest):
    async with gpu_semaphore: # 佇列鎖管控

        input_text = request.prompt[0] if isinstance(request.prompt, list) else request.prompt
        inputs = tokenizer(input_text, return_tensors="pt", add_special_tokens=False).to(model.device)
        gen_config = prepare_gen_config(request, inputs)

        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        gen_config["streamer"] = streamer
        threading.Thread(target=lambda: torch.no_grad()(model.generate)(**gen_config)).start()

        if not request.stream:
            full_text = "".join([new_text for new_text in streamer])
            return {
                "id": f"cmpl-{uuid.uuid4()}", "object": "text_completion", "created": int(time.time()), "model": request.model,
                "choices": [{"text": full_text, "index": 0, "logprobs": None, "finish_reason": "stop"}],
                "usage": {"prompt_tokens": inputs["input_ids"].shape[1], "completion_tokens": len(tokenizer.encode(full_text)), "total_tokens": inputs["input_ids"].shape[1] + len(tokenizer.encode(full_text))}
            }
        else:
            async def openai_text_stream():
                chunk_id, created_time = f"cmpl-{uuid.uuid4()}", int(time.time())
                for new_text in streamer:
                    if new_text:
                        yield f"data: {json.dumps({'id': chunk_id, 'object': 'text_completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'text': new_text, 'index': 0, 'logprobs': None, 'finish_reason': None}]}, ensure_ascii=False)}\n\n"
                yield f"data: {json.dumps({'id': chunk_id, 'object': 'text_completion.chunk', 'created': created_time, 'model': request.model, 'choices': [{'text': '', 'index': 0, 'logprobs': None, 'finish_reason': 'stop'}]}, ensure_ascii=False)}\n\n"
                yield "data: [DONE]\n\n"
            return StreamingResponse(openai_text_stream(), media_type="text/event-stream")

# ==========================================
# 5. 補齊 OpenAI 的 /v1/models 查詢端點 (許多 UI 啟動時會先呼叫這個)
# ==========================================
@app.get("/v1/models")
async def list_models():
    return {
        "object": "list",
        "data": [
            {
                "id": MODEL_ID,
                "object": "model",
                "created": int(time.time()),
                "owned_by": "custom"
            }
        ]
    }
@app.get("/v1/debug/vram")
async def vram_monitor():
    """隨時動態查詢當前 GPU 記憶體狀態"""
    return {
        "object": "vram_status",
        "data": get_vram_status()
    }
if __name__ == "__main__":
    uvicorn.run(app, host="127.0.0.1", port=8001)