昇腾云服务实战：构建自动化大模型 Benchmark 评测体系，深度剖析 KV Cache 与长上下文吞吐极限

import time
import json
import re
import asyncio
import httpx
from dataclasses import dataclass, field
from typing import List, Optional

@dataclass
class PerfMetrics:
    ttft_ms: float = 0.0
    tpot_ms: float = 0.0
    total_latency_ms: float = 0.0
    output_tokens: int = 0
    generated_text: str = ""
    is_correct: Optional[bool] = None

class LLMBenchmarkRunner:
    def __init__(self, endpoint: str, api_key: str, model_name: str, concurrency: int = 32):
        self.endpoint = endpoint
        self.model_name = model_name
        self.semaphore = asyncio.Semaphore(concurrency)
        self.headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json"
        }

    async def evaluate_single(self, client: httpx.AsyncClient, prompt: str, ground_truth: str) -> PerfMetrics:
        payload = {
            "model": self.model_name,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.0,  # 评测模式锁定 Greedy Decoding 确保结果可复现
            "max_tokens": 1024,
            "stream": True
        }
        metrics = PerfMetrics()
        token_timestamps: List[float] = []
        chunks: List[str] = []

        async with self.semaphore:
            start_time = time.perf_counter()
            async with client.stream("POST", self.endpoint, headers=self.headers, json=payload, timeout=120.0) as resp:
                resp.raise_for_status()
                async for line in resp.aiter_lines():
                    if not line or not line.startswith("data: "):
                        continue
                    data_str = line[6:].strip()
                    if data_str == "[DONE]":
                        break
                    
                    now = time.perf_counter()
                    chunk_json = json.loads(data_str)
                    delta = chunk_json["choices"][0]["delta"].get("content", "")
                    if delta:
                        if not token_timestamps:
                            # 记录首字到达时间 (TTFT)
                            metrics.ttft_ms = (now - start_time) * 1000.0
                        token_timestamps.append(now)
                        chunks.append(delta)

            end_time = time.perf_counter()
            metrics.total_latency_ms = (end_time - start_time) * 1000.0
            metrics.output_tokens = len(token_timestamps)
            metrics.generated_text = "".join(chunks)

            if len(token_timestamps) >= 2:
                decode_duration_ms = (token_timestamps[-1] - token_timestamps[0]) * 1000.0
                metrics.tpot_ms = decode_duration_ms / (len(token_timestamps) - 1)

            # 针对 GSM8K / MATH 数据集的自动化答案提取与校验
            metrics.is_correct = self._verify_math_answer(metrics.generated_text, ground_truth)
            return metrics

    @staticmethod
    def _verify_math_answer(prediction: str, truth: str) -> bool:
        # 优先匹配 LaTeX \boxed{...} 或 GSM8K 的 #### 标记
        boxed_match = re.findall(r"\\boxed\{([^}]+)\}", prediction)
        if boxed_match:
            pred_val = boxed_match[-1].strip()
        else:
            hash_match = re.findall(r"####\s*(-?\d+\.?\d*)", prediction)
            pred_val = hash_match[-1].strip() if hash_match else ""
        return pred_val == truth.strip()

# 在华为云昇腾容器环境中启动推理服务，开启前缀缓存与 Chunked Prefill
python -m vllm.entrypoints.openai.api_server \
    --model /mnt/obs-models/Qwen2.5-72B-Instruct \
    --tensor-parallel-size 4 \
    --enable-prefix-caching \
    --enable-chunked-prefill \
    --max-num-batched-tokens 8192 \
    --gpu-memory-utilization 0.92