# awq_inference.py
"""使用 AWQ 量化模型推理"""

from vllm import LLM, SamplingParams

# 直接加载 AWQ 量化模型（推荐从 HuggingFace 下载已量化的模型）
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct-AWQ",  # 已量化的 AWQ 模型
    quantization="awq",
    max_model_len=4096,
    gpu_memory_utilization=0.9,
)

sampling_params = SamplingParams(
    temperature=0.7,
    max_tokens=256,
)

outputs = llm.generate(["什么是量子计算？"], sampling_params)
print(outputs[0].outputs[0].text)

# 命令行启动 AWQ 模型
vllm serve Qwen/Qwen2.5-7B-Instruct-AWQ \
    --quantization awq \
    --served-model-name qwen-7b-awq \
    --max-model-len 4096

# awq_quantize.py
"""自行对模型进行 AWQ 量化"""

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "Qwen/Qwen2.5-7B-Instruct"
quant_path = "Qwen2.5-7B-Instruct-AWQ"

# 1. 加载模型
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 2. 准备校准数据（用于确定量化参数）
# 使用一小批代表性数据
calibration_data = [
    "人工智能是计算机科学的一个分支。",
    "机器学习使计算机能够从数据中学习。",
    "深度学习使用多层神经网络。",
    # ... 更多校准样本
]

# 3. 执行量化
model.quantize(
    tokenizer,
    quant_config={
        "zero_point": True,
        "q_group_size": 128,
        "w_bit": 4,
        "version": "GEMM",
    },
)

# 4. 保存量化模型
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

print(f"量化模型已保存到 {quant_path}")

# gptq_inference.py
"""使用 GPTQ 量化模型推理"""

from vllm import LLM, SamplingParams

llm = LLM(
    model="TheBloke/Llama-2-7B-Chat-GPTQ",
    quantization="gptq",
    max_model_len=4096,
    dtype="half",  # GPTQ 模型通常使用 FP16
)

sampling_params = SamplingParams(temperature=0.7, max_tokens=256)
outputs = llm.generate(["Hello, how are you?"], sampling_params)
print(outputs[0].outputs[0].text)

# 命令行启动
vllm serve TheBloke/Llama-2-7B-Chat-GPTQ \
    --quantization gptq \
    --dtype half \
    --served-model-name llama-7b-gptq

# gptq_quantize.py
"""GPTQ 量化示例"""

from transformers import AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

model_path = "Qwen/Qwen2.5-7B-Instruct"
quant_path = "Qwen2.5-7B-Instruct-GPTQ"

# 量化配置
quantize_config = BaseQuantizeConfig(
    bits=4,            # 量化位数：2, 3, 4, 8
    group_size=128,    # 分组大小
    damp_percent=0.01, # 阻尼系数
    desc_act=True,     # 按激活值降序排列
    sym=False,         # 非对称量化
)

# 加载模型
model = AutoGPTQForCausalLM.from_pretrained(model_path, quantize_config)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 准备校准数据
examples = [
    tokenizer("Hello world", return_tensors="pt"),
    # ... 更多校准样本
]

# 执行量化
model.quantize(examples)

# 保存
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

# fp8_inference.py
"""FP8 量化推理"""

from vllm import LLM, SamplingParams

# 方式一：加载已量化的 FP8 模型
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    quantization="fp8",
    max_model_len=4096,
    tensor_parallel_size=2,
)

# 方式二：在线 FP8 量化（vLLM 自动量化 FP16 模型）
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    quantization="fp8",
    # vLLM 会自动将 FP16 权重转换为 FP8
)

sampling_params = SamplingParams(temperature=0.7, max_tokens=256)
outputs = llm.generate(["What is FP8?"], sampling_params)
print(outputs[0].outputs[0].text)

# 命令行
vllm serve meta-llama/Llama-3.1-70B-Instruct \
    --quantization fp8 \
    --tensor-parallel-size 4

# fp8_kv_cache.py
"""FP8 KV Cache 减少 KV Cache 内存占用"""

llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    kv_cache_dtype="fp8",  # 使用 FP8 存储 KV Cache
    # 每个 block 可节省约 50% 的 KV Cache 内存
)

# 命令行启用 FP8 KV Cache
vllm serve model --kv-cache-dtype fp8

# int8_inference.py
"""INT8 量化推理"""

from vllm import LLM, SamplingParams

# vLLM 支持通过 BitsAndBytes 加载 INT8 模型
llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct",
    load_format="bitsandbytes",   # 使用 BitsAndBytes 加载
    quantization="bitsandbytes",  # INT8 量化
    max_model_len=4096,
)

# benchmark_quantization.py
"""量化性能基准测试"""

import time
from vllm import LLM, SamplingParams

def benchmark_model(model_name, quantization=None, prompts=None):
    """测试模型性能"""
    kwargs = {"model": model_name, "max_model_len": 2048}
    if quantization:
        kwargs["quantization"] = quantization
    
    llm = LLM(**kwargs)
    sampling_params = SamplingParams(temperature=0, max_tokens=128)
    
    # Warmup
    llm.generate(prompts[:2], sampling_params)
    
    # Benchmark
    start = time.time()
    outputs = llm.generate(prompts, sampling_params)
    elapsed = time.time() - start
    
    total_tokens = sum(len(o.outputs[0].token_ids) for o in outputs)
    
    return {
        "model": model_name,
        "quant": quantization or "none",
        "time": elapsed,
        "tokens": total_tokens,
        "throughput": total_tokens / elapsed,
    }

prompts = [f"Explain concept {i} in detail." for i in range(20)]

# 对比测试
results = [
    benchmark_model("Qwen/Qwen2.5-7B-Instruct", None, prompts),
    benchmark_model("Qwen/Qwen2.5-7B-Instruct-AWQ", "awq", prompts),
    benchmark_model("Qwen/Qwen2.5-7B-Instruct-GPTQ", "gptq", prompts),
]

for r in results:
    print(f"{r['quant']:8s} | {r['time']:.2f}s | {r['throughput']:.1f} tok/s")