大模型推理从入门到实测

import torch
from torch import nn

torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = nn.Linear(4, 6).to(device).eval()
x = torch.randn(2, 3, 4, device=device)

with torch.inference_mode():
    y = model(x)

for name, tensor in [("input", x), ("weight", model.weight), ("output", y)]:
    print(name, tuple(tensor.shape), tensor.dtype, tensor.device)
assert y.shape == (2, 3, 6)

import torch
from torch.nn.functional import scaled_dot_product_attention
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
for text in ("大模型如何推理？", "How do language models infer?"):
    ids = tokenizer(text, add_special_tokens=False).input_ids
    print(text, ids, tokenizer.convert_ids_to_tokens(ids))

torch.manual_seed(0)
q, k, v = [torch.randn(1, 2, 4, 8) for _ in range(3)]
changed = v.clone()
changed[:, :, -1, :] += 10
base = scaled_dot_product_attention(q, k, v, is_causal=True)
new = scaled_dot_product_attention(q, k, changed, is_causal=True)
torch.testing.assert_close(base[:, :, :3, :], new[:, :, :3, :])
print("Q/K/V:", q.shape, "scores:", (1, 2, 4, 4), "output:", base.shape)

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

name = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name).eval()
prompt = tokenizer("The largest planet is", return_tensors="pt")
stop = model.generation_config.eos_token_id
stop_ids = set(stop if isinstance(stop, list) else [stop])

@torch.inference_mode()
def greedy(use_cache):
    ids = prompt.input_ids
    mask = prompt.attention_mask
    cache = None
    lengths = []
    for _ in range(8):
        current = ids if cache is None else ids[:, -1:]
        lengths.append(current.shape[-1])
        result = model(input_ids=current, attention_mask=mask,
                       past_key_values=cache, use_cache=use_cache)
        chosen = result.logits[:, -1, :].argmax(dim=-1, keepdim=True)
        ids = torch.cat([ids, chosen], dim=-1)
        mask = torch.cat([mask, torch.ones_like(chosen)], dim=-1)
        cache = result.past_key_values
        if chosen.item() in stop_ids:
            break
    return ids, lengths

with torch.inference_mode():
    reference = model.generate(**prompt, max_new_tokens=8, do_sample=False)
for use_cache in (False, True):
    ids, lengths = greedy(use_cache)
    print(use_cache, lengths, "matches generate:", torch.equal(ids, reference))
    print(tokenizer.decode(ids[0], skip_special_tokens=True))

layers, kv_heads, head_dim, bytes_per_value = 32, 8, 128, 2
bytes_per_token = 2 * layers * kv_heads * head_dim * bytes_per_value
print(f"{bytes_per_token:,} bytes / token")
for length in (4096, 32768):
    mib = bytes_per_token * length / 1024**2
    print(f"{length:,} tokens: {mib:,.0f} MiB")

# 终端一：启动服务，保持运行
sglang serve Qwen/Qwen2.5-0.5B-Instruct --host 127.0.0.1 --port 30000

# 终端二：安装客户端依赖，再运行 stream_client.py
python -m pip install openai

from time import perf_counter
from openai import OpenAI

client = OpenAI(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY",
                max_retries=0, timeout=60.0)
model = client.models.list().data[0].id
messages = [{"role": "user", "content": "用一句话解释 KV Cache。"}]
started = perf_counter()
first_content = None

with client.chat.completions.create(
    model=model, messages=messages, stream=True, max_tokens=64, temperature=0,
) as stream:
    for chunk in stream:
        piece = chunk.choices[0].delta.content if chunk.choices else None
        if piece:
            if first_content is None:
                first_content = perf_counter()
            print(piece, end="", flush=True)

ended = perf_counter()
print(f"\nE2E: {(ended - started) * 1000:.1f} ms")
if first_content is not None:
    print(f"First content: {(first_content - started) * 1000:.1f} ms")

request_id,status,start_ms,end_ms,input_tokens,output_tokens,ttft_ms
r01,ok,0,400,128,20,92
r02,ok,15,465,128,22,107
r03,timeout,20,520,128,,

import csv
from math import ceil

with open("requests.csv", newline="", encoding="utf-8") as file:
    rows = list(csv.DictReader(file))
if not rows:
    raise SystemExit("No requests recorded.")
window = (max(float(r["end_ms"]) for r in rows)
          - min(float(r["start_ms"]) for r in rows)) / 1000
if window <= 0:
    raise SystemExit("The measurement window must be positive.")
success = [r for r in rows if r["status"] == "ok"]
ttft = sorted(float(r["ttft_ms"]) for r in success if r["ttft_ms"])
tokens = sum(int(r["output_tokens"]) for r in success)

print(f"Success: {len(success)}/{len(rows)}")
print(f"Window: {window:.3f} s")
print(f"Successful output throughput: {tokens / window:.2f} token/s")
for p in (0.50, 0.95):
    if ttft:
        print(f"P{p * 100:.0f} TTFT: {ttft[ceil(p * len(ttft)) - 1]:.2f} ms"
              f" (nearest rank, n={len(ttft)})")