LLM 推理优化 2026：从 Prefill-Decode 分离到投机解码的全栈技术指南

"""
Prefill-Decode 分离架构部署配置
使用 vLLM 0.8 的 PD 分离模式
"""
from vllm import LLM, SamplingParams
from vllm.distributed import PDSeparationConfig

# ── 1. Prefill 节点配置 ──
prefill_config = {
    "model": "meta-llama/Llama-4-70B-Instruct",
    "tensor_parallel_size": 4,        # 4 卡张量并行
    "gpu_memory_utilization": 0.95,
    "max_model_len": 32768,
    "dtype": "bfloat16",
    "kv_transfer_config": {
        "role": "prefill",             # Prefill 角色
        "connector": "nccl",           # 使用 NCCL 传输 KV Cache
        "target_nodes": ["decode-0", "decode-1", "decode-2"],
    },
    # Prefill 专用优化
    "enable_chunked_prefill": True,    # 分块 Prefill（减少内存峰值）
    "chunk_size": 2048,                # 每块 2048 tokens
    "scheduler": "prefill_optimized",  # Prefill 优化的调度策略
}

# ── 2. Decode 节点配置 ──
decode_config = {
    "model": "meta-llama/Llama-4-70B-Instruct",
    "tensor_parallel_size": 2,        # Decode 只需 2 卡
    "gpu_memory_utilization": 0.90,
    "max_model_len": 32768,
    "dtype": "bfloat16",
    "kv_transfer_config": {
        "role": "decode",              # Decode 角色
        "connector": "nccl",
        "source_nodes": ["prefill-0"],
    },
    # Decode 专用优化
    "speculative_config": {
        "model": "meta-llama/Llama-4-8B-Instruct",  # 8B 草稿模型
        "num_speculative_tokens": 6,   # 每次推测 6 个 token
        "acceptance_threshold": 0.7,   # 接受率阈值
    },
    "scheduler": "decode_optimized",   # Decode 优化的调度策略
    "max_batch_size": 256,             # 更大的批处理大小
}

# ── 3. 启动 PD 分离服务 ──
def start_pd_separation_service():
    """启动 PD 分离推理服务"""
    
    # 启动 Prefill 节点
    prefill_engine = LLM(
        **prefill_config,
        worker_cls="vllm.worker.PrefillWorker",
    )
    
    # 启动 Decode 节点（可多实例）
    decode_engines = []
    for i in range(3):
        engine = LLM(
            **decode_config,
            worker_cls="vllm.worker.DecodeWorker",
        )
        decode_engines.append(engine)
    
    # 启动负载均衡器
    from vllm.serving import PDRouter
    
    router = PDRouter(
        prefill_engines=[prefill_engine],
        decode_engines=decode_engines,
        routing_strategy="min_decode_queue",  # 路由到队列最短的 Decode 节点
        health_check_interval=5,
    )
    
    # 启动 API 服务
    router.serve(
        host="0.0.0.0",
        port=8000,
        api_type="openai",  # 兼容 OpenAI API
    )

# ── 4. 性能监控 ──
class PDMetrics:
    """PD 分离架构的性能指标"""
    
    def __init__(self):
        self.metrics = {
            "prefill_latency_ms": [],      # Prefill 延迟
            "decode_latency_per_token": [], # 每 token 解码延迟
            "ttft_ms": [],                  # 首 token 延迟
            "throughput_tokens_per_sec": [], # 吞吐量
            "kv_transfer_ms": [],           # KV Cache 传输延迟
            "spec_acceptance_rate": [],     # 投机解码接受率
        }
    
    def report(self):
        """生成性能报告"""
        import numpy as np
        
        report = {}
        for key, values in self.metrics.items():
            if values:
                report[key] = {
                    "mean": np.mean(values),
                    "p50": np.percentile(values, 50),
                    "p99": np.percentile(values, 99),
                }
        
        print("═══ PD 分离架构性能报告 ═══")
        for metric, stats in report.items():
            print(f"{metric}:")
            print(f"  均值: {stats['mean']:.2f}")
            print(f"  P50:  {stats['p50']:.2f}")
            print(f"  P99:  {stats['p99']:.2f}")
        
        return report

"""
KV Cache 优化：PagedAttention v2 配置
"""
from vllm import LLM, SamplingParams

# ── 1. PagedAttention v2 配置 ──
llm = LLM(
    model="meta-llama/Llama-4-70B-Instruct",
    tensor_parallel_size=4,
    
    # PagedAttention v2 配置
    block_manager_config={
        "version": "v2",                  # 使用 v2 分页管理
        "block_size": "dynamic",          # 动态分页粒度
        "min_block_size": 8,              # 最小页大小
        "max_block_size": 64,             # 最大页大小
        "enable_prefix_sharing": True,    # 跨请求前缀共享
        "prefix_sharing_strategy": "radix",  # 基于 Radix Tree 的共享
    },
    
    # KV Cache 压缩
    kv_cache_config={
        "dtype": "auto",                  # 原始精度（跟随模型）
        "compression": {
            "enabled": True,
            "key_dtype": "int8",          # Key 用 INT8
            "value_dtype": "int4",        # Value 用 INT4
            "calibration_method": "minmax",  # 量化校准方法
            "group_size": 128,            # 分组量化大小
        },
        # 内存预算
        "memory_budget_gb": 40,           # KV Cache 总预算 40GB
        "per_request_budget_gb": 2,       # 每个请求预算 2GB
        "eviction_policy": "lru",         # 淘汰策略
    },
    
    # 注意力稀疏化（当 KV Cache 预算耗尽时）
    attention_config={
        "sparse_attention": {
            "enabled": True,
            "strategy": "topk_heads",     # 保留最重要的注意力头
            "keep_ratio": 0.7,            # 保留 70% 的头
            "trigger_threshold": 0.9,     # 内存使用 >90% 时触发
        },
    },
    
    gpu_memory_utilization=0.95,
    max_model_len=32768,
)

# ── 2. 前缀共享效果演示 ──
def demo_prefix_sharing():
    """演示跨请求前缀共享的内存节省"""
    
    # 共享的 System Prompt（约 2000 tokens）
    system_prompt = "你是一个专业的 AI 助手..."
    
    # 多个用户请求（共享同一个 system prompt）
    requests = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": "解释量子计算"},
    ]
    
    requests_2 = [
        {"role": "system", "content": system_prompt},  # 相同！
        {"role": "user", "content": "解释相对论"},
    ]
    
    # 不使用前缀共享：每个请求独立计算 system prompt 的 KV Cache
    # 内存占用：2 × KV(system_prompt) + KV(user_1) + KV(user_2)
    
    # 使用前缀共享：system prompt 的 KV Cache 只计算一次
    # 内存占用：1 × KV(system_prompt) + KV(user_1) + KV(user_2)
    # 节省：约 2000 tokens 的 KV Cache
    
    outputs = llm.chat(requests, SamplingParams(temperature=0.7))
    outputs_2 = llm.chat(requests_2, SamplingParams(temperature=0.7))
    
    # 查看内存使用情况
    cache_stats = llm.get_kv_cache_stats()
    print(f"总页数: {cache_stats['total_pages']}")
    print(f"共享页数: {cache_stats['shared_pages']}")
    print(f"共享率: {cache_stats['shared_pages'] / cache_stats['total_pages']:.1%}")
    print(f"节省内存: {cache_stats['saved_memory_gb']:.1f} GB")

# ── 3. KV Cache 压缩精度测试 ──
def test_kv_compression_accuracy():
    """测试 KV Cache 压缩对输出质量的影响"""
    
    test_prompts = [
        "解释 Transformer 的自注意力机制",
        "写一首关于春天的诗",
        "分析 2026 年 AI 芯片市场趋势",
    ]
    
    # 无压缩基线
    llm_fp16 = LLM(
        model="meta-llama/Llama-4-70B-Instruct",
        kv_cache_config={"compression": {"enabled": False}},
    )
    
    # INT8 Key + INT4 Value 压缩
    llm_compressed = LLM(
        model="meta-llama/Llama-4-70B-Instruct",
        kv_cache_config={"compression": {
            "enabled": True,
            "key_dtype": "int8",
            "value_dtype": "int4",
        }},
    )
    
    for prompt in test_prompts:
        out_fp16 = llm_fp16.generate(prompt)
        out_comp = llm_compressed.generate(prompt)
        
        # 计算输出相似度
        similarity = compute_semantic_similarity(out_fp16, out_comp)
        print(f"Prompt: {prompt[:30]}...")
        print(f"  语义相似度: {similarity:.3f}")
        print(f"  内存节省: {get_memory_savings():.1f}x")

"""
生产级 LLM 推理服务架构
支持百万级日活、自动扩缩容、请求优先级
"""
import asyncio
from dataclasses import dataclass
from enum import Enum
from typing import Optional
import time

# ── 1. 请求模型 ──
@dataclass
class InferenceRequest:
    request_id: str
    prompt: str
    max_tokens: int
    temperature: float
    priority: int  # 0=VIP, 1=Normal, 2=Background
    created_at: float = time.time()
    user_id: Optional[str] = None
    
    @property
    def estimated_prefill_tokens(self) -> int:
        """估算 Prefill token 数"""
        return len(self.prompt) // 4  # 粗略估算：4 字符 ≈ 1 token

# ── 2. 动态批处理器 ──
class DynamicBatcher:
    """动态批处理器：合并多个请求以提高 GPU 利用率"""
    
    def __init__(
        self,
        max_batch_size: int = 256,
        max_wait_ms: int = 10,        # 最大等待时间 10ms
        max_total_tokens: int = 128000,  # 批次最大 token 数
    ):
        self.max_batch_size = max_batch_size
        self.max_wait_ms = max_wait_ms
        self.max_total_tokens = max_total_tokens
        
        self.pending_queue: list[InferenceRequest] = []
        self.batch_ready_event = asyncio.Event()
    
    async def submit(self, request: InferenceRequest):
        """提交请求到批处理队列"""
        self.pending_queue.append(request)
        
        # 按优先级排序（VIP 优先）
        self.pending_queue.sort(key=lambda r: r.priority)
        
        # 检查是否满足立即批处理条件
        if self._should_flush():
            self.batch_ready_event.set()
    
    def _should_flush(self) -> bool:
        """判断是否应该立即发送批次"""
        if not self.pending_queue:
            return False
        
        # 条件 1：队列已满
        if len(self.pending_queue) >= self.max_batch_size:
            return True
        
        # 条件 2：token 总量达到上限
        total_tokens = sum(r.estimated_prefill_tokens for r in self.pending_queue)
        if total_tokens >= self.max_total_tokens:
            return True
        
        # 条件 3：等待时间超过阈值
        oldest_request = self.pending_queue[0]
        wait_time_ms = (time.time() - oldest_request.created_at) * 1000
        if wait_time_ms >= self.max_wait_ms:
            return True
        
        return False
    
    async def get_next_batch(self) -> list[InferenceRequest]:
        """获取下一个批次"""
        while not self._should_flush():
            self.batch_ready_event.clear()
            await self.batch_ready_event.wait()
        
        batch = self.pending_queue[:self.max_batch_size]
        self.pending_queue = self.pending_queue[self.max_batch_size:]
        return batch

# ── 3. 推理调度器 ──
class InferenceScheduler:
    """推理调度器：管理 PD 分离集群"""
    
    def __init__(self):
        self.batcher = DynamicBatcher()
        self.prefill_nodes: list[str] = []  # Prefill 节点列表
        self.decode_nodes: list[str] = []   # Decode 节点列表
        self.node_load: dict[str, int] = {} # 节点负载
    
    async def schedule_request(self, request: InferenceRequest):
        """调度单个请求"""
        # 提交到批处理器
        await self.batcher.submit(request)
    
    async def run_serving_loop(self):
        """主服务循环"""
        while True:
            # 获取下一个批次
            batch = await self.batcher.get_next_batch()
            
            if not batch:
                continue
            
            # 选择 Prefill 节点（负载最低）
            prefill_node = self._select_least_loaded_node("prefill")
            
            # 执行 Prefill
            kv_cache = await self._execute_prefill(prefill_node, batch)
            
            # 选择 Decode 节点（队列最短）
            decode_node = self._select_least_loaded_node("decode")
            
            # 传输 KV Cache 并执行 Decode
            results = await self._execute_decode(decode_node, batch, kv_cache)
            
            # 返回结果
            for request, result in zip(batch, results):
                await self._send_response(request, result)
    
    def _select_least_loaded_node(self, node_type: str) -> str:
        """选择负载最低的节点"""
        nodes = self.prefill_nodes if node_type == "prefill" else self.decode_nodes
        return min(nodes, key=lambda n: self.node_load.get(n, 0))
    
    async def _execute_prefill(self, node, batch):
        """执行 Prefill 阶段"""
        # 实际实现中调用 vLLM/TensorRT-LLM API
        pass
    
    async def _execute_decode(self, node, batch, kv_cache):
        """执行 Decode 阶段"""
        pass
    
    async def _send_response(self, request, result):
        """发送响应给用户"""
        pass

# ── 4. 自动扩缩容 ──
class AutoScaler:
    """根据负载自动扩缩容 GPU 节点"""
    
    def __init__(
        self,
        scheduler: InferenceScheduler,
        min_decode_nodes: int = 2,
        max_decode_nodes: int = 32,
        target_gpu_utilization: float = 0.7,
        scale_up_threshold: float = 0.85,  # GPU 利用率 >85% 扩容
        scale_down_threshold: float = 0.3,  # GPU 利用率 <30% 缩容
        cooldown_seconds: int = 300,        # 扩缩容冷却时间 5 分钟
    ):
        self.scheduler = scheduler
        self.min_nodes = min_decode_nodes
        self.max_nodes = max_decode_nodes
        self.target_util = target_gpu_utilization
        self.scale_up_threshold = scale_up_threshold
        self.scale_down_threshold = scale_down_threshold
        self.cooldown = cooldown_seconds
        self.last_scale_time = 0
    
    async def check_and_scale(self):
        """检查并执行扩缩容"""
        current_time = time.time()
        
        # 冷却期内不操作
        if current_time - self.last_scale_time < self.cooldown:
            return
        
        # 计算平均 GPU 利用率
        avg_util = self._get_avg_gpu_utilization()
        queue_length = len(self.scheduler.batcher.pending_queue)
        current_nodes = len(self.scheduler.decode_nodes)
        
        if avg_util > self.scale_up_threshold and current_nodes < self.max_nodes:
            # 扩容
            new_nodes = self._calculate_scale_up_count(avg_util, queue_length)
            await self._add_decode_nodes(new_nodes)
            self.last_scale_time = current_time
            print(f"🔼 扩容: +{new_nodes} 节点 (GPU util: {avg_util:.1%})")
        
        elif avg_util < self.scale_down_threshold and current_nodes > self.min_nodes:
            # 缩容
            remove_count = self._calculate_scale_down_count(avg_util, current_nodes)
            await self._remove_decode_nodes(remove_count)
            self.last_scale_time = current_time
            print(f"🔽 缩容: -{remove_count} 节点 (GPU util: {avg_util:.1%})")
    
    def _get_avg_gpu_utilization(self) -> float:
        """获取平均 GPU 利用率"""
        # 实际实现中通过 pynvml 或 DCGM 获取
        return 0.0
    
    def _calculate_scale_up_count(self, util, queue_len) -> int:
        """计算需要扩容的节点数"""
        # 目标：将利用率降到 target_util
        current_capacity = len(self.scheduler.decode_nodes)
        needed_capacity = int(current_capacity * util / self.target_util) + 1
        return min(needed_capacity - current_capacity, 4)  # 每次最多加 4 个
    
    def _calculate_scale_down_count(self, util, current_count) -> int:
        """计算需要缩容的节点数"""
        target_capacity = max(int(current_count * util / self.target_util) + 1, self.min_nodes)
        return min(current_count - target_capacity, 2)  # 每次最多减 2 个
    
    async def _add_decode_nodes(self, count: int):
        """添加 Decode 节点"""
        pass
    
    async def _remove_decode_nodes(self, count: int):
        """移除 Decode 节点（优雅关闭，等待当前请求完成）"""
        pass

# LCLM (Latent Context Models) 推理示例
# pip install transformers latent-context

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from latent_context import LatentContextCompressor

# 加载基础模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B",
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B")

# 初始化 LCLM 压缩器
compressor = LatentContextCompressor(
    compression_ratio=16,  # 16x 压缩
    model_path="latent-context/lclm-llama3-8b"
)

# 长上下文场景：100K token 的 RAG 文档
long_document = load_large_document("./rag_context.txt")  # ~100K tokens
document_tokens = tokenizer(long_document, return_tensors="pt",
                            truncation=True, max_length=100_000)

print(f"原始 token 数: {document_tokens['input_ids'].shape[1]}")
# 原始 token 数: 100000

# 压缩为潜在表示
compressed = compressor.compress(document_tokens)
print(f"压缩后 latent 维度: {compressed.shape}")
# 压缩后 latent 维度: torch.Size([1, 6250, 4096])
# 100K tokens → 6250 latent vectors (16x 压缩)

# 使用压缩后的上下文进行推理
query = "这份文档中关于推理优化的关键结论是什么？"
inputs = tokenizer(query, return_tensors="pt").to("cuda")

output = model.generate(
    **inputs,
    latent_context=compressed,  # 传入压缩后的上下文
    max_new_tokens=512,
    use_cache=True
)

answer = tokenizer.decode(output[0], skip_special_tokens=True)
print(f"回答: {answer}")

# === 性能对比 ===
# | 方案 | 显存占用 | 延迟 | 质量 |
# |------|---------|------|------|
# | 完整 KV Cache | OOM (>80GB) | N/A | N/A |
# | LCLM 16x | ~5GB | 62% | 97% NDCG |
# | 分块 RAG | ~12GB | 85% | 89% NDCG |

# DFlash + FP4 量化推理配置
# 模拟 MiMo UltraSpeed 的推理加速方案

from dataclasses import dataclass

@dataclass
class DFlashConfig:
    """DFlash 块级投机解码配置"""
    block_size: int = 64          # 每次生成 64 个 token 的块
    draft_model: str = "mico-8b"  # 轻量草稿模型
    fp4_expert_quant: bool = True # FP4 专家层量化
    acceptance_threshold: float = 0.65
    max_speculative_tokens: int = 128

class DFlashInference:
    """DFlash 块级投机解码引擎"""
    
    def __init__(self, model_path: str, config: DFlashConfig):
        self.config = config
        # 加载主模型（FP4 量化专家层）
        self.target_model = self._load_quantized(model_path)
        # 加载草稿模型
        self.draft_model = self._load_draft(config.draft_model)
    
    def _load_quantized(self, path: str):
        """加载 FP4 量化模型"""
        from tile_rt import QuantizedModel
        return QuantizedModel.from_pretrained(
            path,
            expert_quantization="fp4",    # 仅量化 MoE 专家层
            attention_dtype="fp16",        # 注意力层保持 FP16
            embedding_dtype="int8",        # 嵌入层 INT8
        )
    
    def generate_block(self, prompt: str) -> str:
        """块级生成：一次生成 N 个 token"""
        # 第一步：草稿模型快速生成候选块
        candidates = self.draft_model.generate(
            prompt,
            max_tokens=self.config.block_size,
            temperature=0.8
        )
        
        # 第二步：目标模型一次性验证整个块
        verification = self.target_model.verify_block(
            prompt=prompt,
            candidates=candidates,
            threshold=self.config.acceptance_threshold
        )
        
        # 第三步：返回接受的 token + 从拒绝位置重新采样
        accepted = verification.accepted_tokens
        if verification.rejected_at >= 0:
            # 从拒绝位置用目标模型重新采样
            resampled = self.target_model.generate(
                prompt + accepted,
                max_tokens=1
            )
            accepted += resampled
        
        return accepted
    
    def benchmark(self, prompts: list, n_runs: int = 10):
        """性能基准测试"""
        import time
        
        results = {"tokens": [], "latency_ms": [], "throughput": []}
        
        for _ in range(n_runs):
            for prompt in prompts:
                start = time.perf_counter()
                output = self.generate_block(prompt)
                elapsed = (time.perf_counter() - start) * 1000
                
                results["tokens"].append(len(output.split()))
                results["latency_ms"].append(elapsed)
                results["throughput"].append(
                    len(output.split()) / (elapsed / 1000)
                )
        
        import statistics
        print(f"=== DFlash + FP4 推理基准 ===")
        print(f"平均吞吐量: {statistics.mean(results['throughput']):.1f} tok/s")
        print(f"P50 延迟: {statistics.median(results['latency_ms']):.1f} ms")
        print(f"P99 延迟: {sorted(results['latency_ms'])[int(len(results['latency_ms'])*0.99)]:.1f} ms")
        
        return results

# 使用示例
config = DFlashConfig(
    block_size=64,
    draft_model="mico-8b",
    fp4_expert_quant=True,
    acceptance_threshold=0.65
)
engine = DFlashInference("xiaomi/mimo-v2.5-pro", config)
# 预期结果：~1000 tok/s（8x GPU 节点）