推理优化¶

import torch
import torch.nn as nn

class KVCache:
    """KV Cache实现"""

    def __init__(self, max_batch_size: int, max_seq_len: int, num_heads: int, head_dim: int):
        """
        初始化KV Cache

        Args:
            max_batch_size: 最大批次大小
            max_seq_len: 最大序列长度
            num_heads: 注意力头数
            head_dim: 每个头的维度
        """
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.num_heads = num_heads
        self.head_dim = head_dim

        # 初始化缓存
        self.cache_k = torch.zeros(
            max_batch_size,
            num_heads,
            max_seq_len,
            head_dim
        )
        self.cache_v = torch.zeros(
            max_batch_size,
            num_heads,
            max_seq_len,
            head_dim
        )

        # 当前位置
        self.current_pos = 0

    def update(self, k: torch.Tensor, v: torch.Tensor):
        """
        更新缓存

        Args:
            k: Key张量 [batch_size, num_heads, seq_len, head_dim]
            v: Value张量 [batch_size, num_heads, seq_len, head_dim]
        """
        batch_size = k.size(0)
        seq_len = k.size(2)

        # 更新缓存
        self.cache_k[:batch_size, :, self.current_pos:self.current_pos+seq_len, :] = k
        self.cache_v[:batch_size, :, self.current_pos:self.current_pos+seq_len, :] = v

        # 更新位置
        self.current_pos += seq_len

    def get(self) -> tuple:
        """
        获取缓存

        Returns:
            (k, v) 缓存的Key和Value
        """
        k = self.cache_k[:, :, :self.current_pos, :]
        v = self.cache_v[:, :, :self.current_pos, :]
        return k, v

    def clear(self):
        """清空缓存"""
        self.current_pos = 0
        self.cache_k.zero_()
        self.cache_v.zero_()

# 使用示例
kv_cache = KVCache(
    max_batch_size=1,
    max_seq_len=1024,
    num_heads=12,
    head_dim=64
)

# 生成过程
for i in range(10):
    # 模拟计算K和V
    k = torch.randn(1, 12, 1, 64)
    v = torch.randn(1, 12, 1, 64)

    # 更新缓存
    kv_cache.update(k, v)

    # 获取缓存用于注意力计算
    cached_k, cached_v = kv_cache.get()

    # 使用缓存的K和V进行注意力计算
    # ...

class PagedKVCache:
    """分页KV Cache"""

    def __init__(self, num_heads: int, head_dim: int, page_size: int = 128):
        """
        初始化分页KV Cache

        Args:
            num_heads: 注意力头数
            head_dim: 每个头的维度
            page_size: 页大小
        """
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.page_size = page_size

        # 页表
        self.page_table = {}

        # 页池
        self.page_pool_k = []
        self.page_pool_v = []

    def allocate_page(self) -> int:
        """分配页"""
        if not self.page_pool_k:
            # 创建新页
            page_k = torch.zeros(1, self.num_heads, self.page_size, self.head_dim)
            page_v = torch.zeros(1, self.num_heads, self.page_size, self.head_dim)
            self.page_pool_k.append(page_k)
            self.page_pool_v.append(page_v)
            return len(self.page_pool_k) - 1
        else:
            # 复用空闲页
            return self._find_free_page()

    def update(self, k: torch.Tensor, v: torch.Tensor, seq_idx: int):
        """
        更新缓存

        Args:
            k: Key张量
            v: Value张量
            seq_idx: 序列索引
        """
        # 计算页索引
        page_idx = seq_idx // self.page_size
        offset = seq_idx % self.page_size

        # 分配页
        if page_idx not in self.page_table:
            self.page_table[page_idx] = self.allocate_page()

        # 更新页
        actual_page_idx = self.page_table[page_idx]
        self.page_pool_k[actual_page_idx][:, :, offset:offset+1, :] = k
        self.page_pool_v[actual_page_idx][:, :, offset:offset+1, :] = v

    def get(self, seq_len: int) -> tuple:
        """
        获取缓存

        Args:
            seq_len: 序列长度

        Returns:
            (k, v) 缓存的Key和Value
        """
        # 计算需要的页数
        num_pages = (seq_len + self.page_size - 1) // self.page_size

        # 收集页
        k_pages = []
        v_pages = []

        for page_idx in range(num_pages):
            if page_idx in self.page_table:
                actual_page_idx = self.page_table[page_idx]
                k_pages.append(self.page_pool_k[actual_page_idx])
                v_pages.append(self.page_pool_v[actual_page_idx])

        # 拼接页
        k = torch.cat(k_pages, dim=2)[:, :, :seq_len, :]
        v = torch.cat(v_pages, dim=2)[:, :, :seq_len, :]

        return k, v

import torch
from collections import deque

class ContinuousBatchProcessor:
    """连续批处理器"""

    def __init__(self, max_batch_size: int, max_seq_len: int):
        """
        初始化连续批处理器

        Args:
            max_batch_size: 最大批次大小
            max_seq_len: 最大序列长度
        """
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len

        # 请求队列
        self.request_queue = deque()  # deque双端队列，两端操作O(1)

        # 活跃请求
        self.active_requests = {}

    def add_request(self, request_id: str, input_ids: torch.Tensor):
        """
        添加请求

        Args:
            request_id: 请求ID
            input_ids: 输入token IDs
        """
        self.request_queue.append({
            'id': request_id,
            'input_ids': input_ids,
            'seq_len': input_ids.size(0)
        })

    def get_batch(self) -> tuple[torch.Tensor, list[str]]:
        """
        获取批次

        Returns:
            (batch_input_ids, request_ids)
        """
        if not self.request_queue:
            return None, []

        # 选择请求组成批次
        batch_requests = []
        total_tokens = 0

        while self.request_queue and len(batch_requests) < self.max_batch_size:
            request = self.request_queue[0]

            # 检查是否超过最大序列长度
            if total_tokens + request['seq_len'] > self.max_seq_len:
                break

            # 添加到批次
            batch_requests.append(self.request_queue.popleft())
            total_tokens += request['seq_len']

        if not batch_requests:
            return None, []

        # 准备批次数据
        max_len = max(req['seq_len'] for req in batch_requests)
        batch_input_ids = torch.zeros(
            len(batch_requests),
            max_len,
            dtype=torch.long
        )

        request_ids = []
        for i, req in enumerate(batch_requests):  # enumerate同时获取索引和元素
            batch_input_ids[i, :req['seq_len']] = req['input_ids']
            request_ids.append(req['id'])

        return batch_input_ids, request_ids

    def complete_request(self, request_id: str):
        """
        完成请求

        Args:
            request_id: 请求ID
        """
        if request_id in self.active_requests:
            del self.active_requests[request_id]

# 使用示例
processor = ContinuousBatchProcessor(
    max_batch_size=4,
    max_seq_len=1024
)

# 添加请求
processor.add_request("req1", torch.randint(0, 1000, (50,)))
processor.add_request("req2", torch.randint(0, 1000, (100,)))
processor.add_request("req3", torch.randint(0, 1000, (75,)))

# 获取批次
batch_input_ids, request_ids = processor.get_batch()
print(f"批次大小: {len(request_ids)}")
print(f"批次形状: {batch_input_ids.shape}")

import time
import asyncio  # Python标准异步库

class DynamicBatchProcessor:
    """动态批处理器"""

    def __init__(self, max_batch_size: int, max_wait_time: float = 0.01):
        """
        初始化动态批处理器

        Args:
            max_batch_size: 最大批次大小
            max_wait_time: 最大等待时间(秒)
        """
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time

        # 请求队列
        self.request_queue = []
        self.last_batch_time = time.time()

    async def add_request(self, request_id: str, input_ids: torch.Tensor):  # async def定义协程函数
        """
        添加请求

        Args:
            request_id: 请求ID
            input_ids: 输入token IDs
        """
        self.request_queue.append({
            'id': request_id,
            'input_ids': input_ids,
            'seq_len': input_ids.size(0),
            'arrival_time': time.time()
        })

    async def get_batch(self) -> tuple[torch.Tensor, list[str]]:
        """
        获取批次

        Returns:
            (batch_input_ids, request_ids)
        """
        if not self.request_queue:
            return None, []

        current_time = time.time()
        time_since_last_batch = current_time - self.last_batch_time

        # 检查是否应该组批
        should_batch = (
            len(self.request_queue) >= self.max_batch_size or
            time_since_last_batch >= self.max_wait_time
        )

        if not should_batch:
            return None, []

        # 组批
        batch_requests = self.request_queue[:self.max_batch_size]
        self.request_queue = self.request_queue[len(batch_requests):]
        self.last_batch_time = current_time

        # 准备批次数据
        max_len = max(req['seq_len'] for req in batch_requests)
        batch_input_ids = torch.zeros(
            len(batch_requests),
            max_len,
            dtype=torch.long
        )

        request_ids = []
        for i, req in enumerate(batch_requests):
            batch_input_ids[i, :req['seq_len']] = req['input_ids']
            request_ids.append(req['id'])

        return batch_input_ids, request_ids

# 使用示例
async def main():
    processor = DynamicBatchProcessor(
        max_batch_size=4,
        max_wait_time=0.01
    )

    # 添加请求
    await processor.add_request("req1", torch.randint(0, 1000, (50,)))  # await等待异步操作完成
    await asyncio.sleep(0.005)
    await processor.add_request("req2", torch.randint(0, 1000, (100,)))

    # 获取批次
    batch_input_ids, request_ids = await processor.get_batch()
    print(f"批次大小: {len(request_ids)}")

# 运行
asyncio.run(main())  # 创建事件循环运行顶层协程

import torch
import torch.nn as nn

# 定义模型
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()  # super()调用父类方法
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# 创建模型
model = SimpleModel()

# 转换为TorchScript
scripted_model = torch.jit.script(model)

# 保存模型
scripted_model.save("model_scripted.pt")

# 加载并推理
loaded_model = torch.jit.load("model_scripted.pt")
input_data = torch.randn(1, 10)
output = loaded_model(input_data)

import torch
from torch2trt import torch2trt
import tensorrt as trt

# 创建模型
model = SimpleModel()
model.eval()

# 示例输入
example_input = torch.randn(1, 10).cuda()

# 转换为TensorRT
trt_model = torch2trt(
    model,
    [example_input],
    fp16_mode=True,
    max_workspace_size=1 << 25
)

# 保存模型
torch.save(trt_model.state_dict(), "model_trt.pth")

# 推理
output = trt_model(example_input)

import torch
import onnx
import onnxruntime as ort

# 导出模型为ONNX
model = SimpleModel()
model.eval()

dummy_input = torch.randn(1, 10)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}}
)

# 使用ONNX Runtime推理
ort_session = ort.InferenceSession("model.onnx")

input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# 推理
input_data = torch.randn(1, 10).numpy()
outputs = ort_session.run([output_name], {input_name: input_data})

# vLLM 快速启动服务
from vllm import LLM, SamplingParams

llm = LLM(model="Qwen/Qwen2.5-7B-Instruct", tensor_parallel_size=1)
params = SamplingParams(temperature=0.7, max_tokens=512)

prompts = ["解释什么是Transformer", "写一首关于AI的诗"]
outputs = llm.generate(prompts, params)

for output in outputs:
    print(output.outputs[0].text)

import sglang as sgl

# 方式1：离线批量推理
llm = sgl.Engine(model_path="Qwen/Qwen2.5-7B-Instruct")

prompts = ["什么是大模型推理优化？", "解释KV Cache的作用"]
outputs = llm.generate(prompts, sampling_params={"temperature": 0.7, "max_new_tokens": 256})
for out in outputs:
    print(out["text"])

llm.shutdown()

# 方式2：启动 OpenAI 兼容服务器
# 终端执行：
# python -m sglang.launch_server --model Qwen/Qwen2.5-7B-Instruct --port 8000

# 客户端调用（与 vLLM/OpenAI 完全兼容）
import openai
client = openai.OpenAI(base_url="http://localhost:8000/v1", api_key="none")

response = client.chat.completions.create(
    model="Qwen/Qwen2.5-7B-Instruct",
    messages=[{"role": "user", "content": "用JSON格式输出三个中国城市的信息"}],
    response_format={"type": "json_object"},  # 结构化输出
)
print(response.choices[0].message.content)

# 1. 启用 Torch.compile 加速
llm = sgl.Engine(
    model_path="meta-llama/Llama-3-8B-Instruct",
    enable_torch_compile=True,  # 推理速度提升 1.5 倍
)

# 2. 张量并行（多卡部署）
llm = sgl.Engine(
    model_path="meta-llama/Llama-3-70B-Instruct",
    tensor_parallel_size=4,  # 使用 4 卡并行
)

# 3. DeepSeek 模型优化（启用 MLA）
llm = sgl.Engine(
    model_path="deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    enable_mla=True,  # MLA 注意力加速 7 倍
)

# 4. FP8 量化部署
llm = sgl.Engine(
    model_path="meta-llama/Llama-3-8B-Instruct",
    quant_method="fp8",  # FP8 W8A8 量化
)

# 5. 结构化输出（xgrammar 加速）
json_schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "integer", "minimum": 0},
    },
    "required": ["name", "age"]
}
output = llm.generate(
    "生成一个AI工程师的简介",
    sampling_params={"max_new_tokens": 200, "json_schema": json_schema}
)

# SGLang 原生支持 JSON Schema 约束（Compressed FSM 加速）
from sglang import Engine

llm = Engine(model_path="Qwen/Qwen2.5-7B-Instruct")

json_schema = {
    "type": "object",
    "properties": {
        "name": {"type": "string"},
        "age": {"type": "integer", "minimum": 0},
        "skills": {"type": "array", "items": {"type": "string"}}
    },
    "required": ["name", "age", "skills"]
}

output = llm.generate(
    "生成一个AI工程师的简介",
    sampling_params={"max_new_tokens": 200, "json_schema": json_schema}
)
print(output[0]["text"])  # 保证输出符合 JSON Schema

import torch

def speculative_decode(target_model, draft_model, input_ids, gamma=5, temperature=1.0):
    """
    Speculative Decoding 简化实现
    target_model: 大模型（如 70B）
    draft_model: 小模型（如 7B，同系列）
    gamma: 每轮猜测的 token 数
    """
    generated = input_ids.clone()

    while len(generated[0]) < input_ids.shape[1] + 100:  # 简化：生成100个token
        # Step 1: Draft Model 自回归生成 γ 个候选 token
        draft_tokens = []
        draft_probs = []
        draft_input = generated.clone()

        for _ in range(gamma):
            with torch.no_grad():  # 禁用梯度计算，节省内存（推理时使用）
                logits = draft_model(draft_input).logits[:, -1, :] / temperature
                probs = torch.softmax(logits, dim=-1)
                token = torch.multinomial(probs, 1)
                draft_tokens.append(token)
                draft_probs.append(probs)
                draft_input = torch.cat([draft_input, token], dim=-1)

        # Step 2: Target Model 一次前向传播验证所有候选
        with torch.no_grad():
            all_input = torch.cat([generated] + draft_tokens, dim=-1)
            target_logits = target_model(all_input).logits / temperature

        # Step 3: 逐个验证，基于拒绝采样
        accepted = 0
        for i in range(gamma):
            target_prob = torch.softmax(target_logits[:, -(gamma - i + 1), :], dim=-1)
            draft_prob = draft_probs[i]
            token = draft_tokens[i]

            # 接受概率
            accept_prob = torch.min(
                torch.ones_like(target_prob),
                target_prob / (draft_prob + 1e-10)
            )

            if torch.rand(1).item() < accept_prob[0, token[0, 0]].item():
                generated = torch.cat([generated, token], dim=-1)
                accepted += 1
            else:
                # 从修正分布采样
                residual = torch.clamp(target_prob - draft_prob, min=0)
                residual = residual / residual.sum(dim=-1, keepdim=True)
                new_token = torch.multinomial(residual, 1)
                generated = torch.cat([generated, new_token], dim=-1)
                break

        if accepted == gamma:
            # 所有猜测都被接受，额外从 target 采样一个
            bonus_prob = torch.softmax(target_logits[:, -1, :], dim=-1)
            bonus_token = torch.multinomial(bonus_prob, 1)
            generated = torch.cat([generated, bonus_token], dim=-1)

    return generated

import torch
import torch.nn.functional as F

class EAGLEDrafting:
    """EAGLE 投机解码实现"""

    def __init__(self, target_model, draft_model, max_depth=3, branching=3):
        self.target_model = target_model
        self.draft_model = draft_model
        self.max_depth = max_depth
        self.branching = branching

    def build_draft_tree(self, input_ids, kv_cache=None):
        """
        构建 draft 树

        Args:
            input_ids: 输入 token IDs
            kv_cache: 现有的 KV Cache

        Returns:
            tree_tokens: 树中所有候选 token
            tree_probs: 对应概率分布
            path_to_accept: 每个分支的接受路径
        """
        tree_tokens = [input_ids]
        tree_probs = []

        current_ids = input_ids
        current_kv = kv_cache

        for depth in range(self.max_depth):
            # 批量预测所有叶子节点
            next_layer_tokens = []
            next_layer_probs = []

            for branch_idx, branch_ids in enumerate(tree_tokens):
                # 复用 KV Cache 加速
                logits = self.draft_model(branch_ids, past_key_values=current_kv)
                probs = F.softmax(logits[:, -1, :], dim=-1)

                # 选择 top-k 候选
                top_k_probs, top_k_ids = torch.topk(probs, self.branching, dim=-1)

                for i in range(self.branching):
                    next_token = top_k_ids[:, i:i+1]
                    next_layer_tokens.append(
                        torch.cat([branch_ids, next_token], dim=-1)
                    )
                    next_layer_probs.append(top_k_probs[:, i])

            tree_tokens = next_layer_tokens
            tree_probs = next_layer_probs

        return tree_tokens, tree_probs

    def verify_and_sample(self, tree_tokens, tree_probs):
        """
        并行验证树中所有候选，选择最优路径

        Returns:
            accepted_ids: 被接受的 token 序列
            accepted_count: 被接受的 token 数量
        """
        # 拼接所有候选序列
        max_len = max(t.shape[1] for t in tree_tokens)
        padded_tokens = [
            torch.cat([t, t.new_zeros(1, max_len - t.shape[1])], dim=-1)
            for t in tree_tokens
        ]
        batch_tokens = torch.cat(padded_tokens, dim=0)

        # Target Model 一次前向传播验证
        with torch.no_grad():
            target_logits = self.target_model(batch_tokens).logits

        # 计算每条路径的接受概率
        best_path_idx = 0
        best_path_score = 0.0

        for idx, (tokens, draft_prob) in enumerate(zip(tree_tokens, tree_probs)):
            # 计算这条路径的累积得分
            path_score = draft_prob.item()
            if path_score > best_path_score:
                best_path_score = path_score
                best_path_idx = idx

        return tree_tokens[best_path_idx], best_path_score

class MedusaPredictionHead(nn.Module):
    """Medusa 预测头"""

    def __init__(self, hidden_size, vocab_size, num_heads=3):
        super().__init__()
        self.num_heads = num_heads
        # 每个头有自己的输出层
        self.heads = nn.ModuleList([
            nn.Linear(hidden_size, vocab_size, bias=False)
            for _ in range(num_heads)
        ])

    def forward(self, hidden_states):
        """
        Args:
            hidden_states: [batch, seq_len, hidden]

        Returns:
            predictions: [batch, seq_len, num_heads, vocab_size]
        """
        return torch.stack([head(hidden_states) for head in self.heads], dim=2)


class MedusaModel(nn.Module):
    """带 Medusa 头的模型"""

    def __init__(self, base_model, num_medusa_heads=3):
        super().__init__()
        self.base_model = base_model
        self.hidden_size = base_model.config.hidden_size
        self.vocab_size = base_model.config.vocab_size

        # 添加 Medusa 预测头
        self.medusa_heads = MedusaPredictionHead(
            self.hidden_size,
            self.vocab_size,
            num_heads=num_medusa_heads
        )

    def forward(self, input_ids):
        # 主干网络前向传播
        outputs = self.base_model(input_ids, output_hidden_states=True)
        hidden_states = outputs.hidden_states[-1]  # 最后一层

        # Medusa 预测
        medusa_logits = self.medusa_heads(hidden_states)

        return {
            'main_logits': outputs.logits,
            'medusa_logits': medusa_logits
        }

    def medusa_decode(self, input_ids, temperature=1.0):
        """
        Medusa 解码流程

        1. 一次前向传播获取所有预测头的输出
        2. 构建候选序列
        3. 验证并接受
        """
        outputs = self.forward(input_ids)

        # 解析各预测头的输出
        medusa_logits = outputs['medusa_logits']  # [batch, seq, num_heads, vocab]

        candidates = []
        for head_idx in range(self.medusa_heads.num_heads):
            head_logits = medusa_logits[:, -1, head_idx, :] / temperature
            probs = F.softmax(head_logits, dim=-1)
            token = torch.multinomial(probs, 1)
            candidates.append(token)

        # 构建候选序列：[input, cand1, cand2, cand3]
        candidate_sequence = torch.cat([input_ids] + candidates, dim=-1)

        # Target 验证
        with torch.no_grad():
            verification_logits = self.base_model(candidate_sequence).logits

        # 逐个验证并接受
        accepted = []
        for i, cand_token in enumerate(candidates):
            target_prob = F.softmax(verification_logits[:, input_ids.shape[1] + i, :], dim=-1)
            cand_prob = F.softmax(medusa_logits[:, -1, i, :], dim=-1)

            accept_prob = torch.min(
                torch.ones_like(target_prob),
                target_prob / (cand_prob + 1e-10)
            )

            if torch.rand(1).item() < accept_prob[0, cand_token[0, 0]].item():
                accepted.append(cand_token)
            else:
                # 拒绝：从修正分布采样
                residual = torch.clamp(target_prob - cand_prob, min=0)
                residual = residual / residual.sum(dim=-1, keepdim=True)
                new_token = torch.multinomial(residual, 1)
                accepted.append(new_token)
                break

        return torch.cat([input_ids] + accepted, dim=-1)

class SelfSpeculativeDecoder:
    """自投机解码器"""

    def __init__(self, model, num_draft_layers, num_draft_tokens=5):
        """
        Args:
            model: 原始模型
            num_draft_layers: 用作 draft 的层数
            num_draft_tokens: 每轮猜测的 token 数
        """
        self.model = model
        self.num_draft_layers = num_draft_layers
        self.num_draft_tokens = num_draft_tokens
        self.total_layers = model.config.num_hidden_layers

    def draft_generation(self, input_ids, draft_kv_cache):
        """
        浅层 draft 生成
        """
        draft_tokens = []
        current_ids = input_ids
        current_kv = draft_kv_cache

        for _ in range(self.num_draft_tokens):
            # 只通过浅层
            hidden = self.model.model.embed_tokens(current_ids)
            for layer_idx in range(self.num_draft_layers):
                hidden = self.model.model.layers[layer_idx](hidden, current_kv)

            # 预测下一个 token
            logits = self.model.lm_head(hidden[:, -1, :])
            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, 1)

            draft_tokens.append(next_token)
            current_ids = torch.cat([current_ids, next_token], dim=-1)

        return draft_tokens

    def verify(self, input_ids, draft_tokens):
        """
        深层验证
        """
        # 拼接 draft tokens
        all_ids = torch.cat([input_ids] + draft_tokens, dim=-1)

        # 一次完整前向传播
        with torch.no_grad():
            outputs = self.model(all_ids, use_cache=True)

        # 逐个验证
        accepted_tokens = []
        for i, draft_token in enumerate(draft_tokens):
            target_logits = outputs.logits[:, input_ids.shape[1] + i, :]
            target_probs = F.softmax(target_logits, dim=-1)

            # 简化：直接贪心接受 top-1
            target_token = torch.argmax(target_probs, dim=-1, keepdim=True)

            if torch.all(target_token == draft_token):
                accepted_tokens.append(draft_token)
            else:
                # 拒绝：使用 target 的预测
                accepted_tokens.append(target_token)
                break

        return accepted_tokens

from vllm import LLM, SamplingParams

# 启用投机解码
llm = LLM(
    model="meta-llama/Llama-3-70B-Instruct",
    tensor_parallel_size=4,
    speculative_model="meta-llama/Llama-3-8B-Instruct",  # 小模型作为 draft
    num_speculative_tokens=5,  # 每轮猜测 5 个 token
)

# 推理（与普通调用完全相同）
params = SamplingParams(temperature=0.7, max_tokens=512)
output = llm.generate(["解释量子计算"], params)
print(output[0].outputs[0].text)

import sglang as sgl

# 启动带投机解码的 SGLang
# python -m sglang.launch_server \
#     --model-path meta-llama/Llama-3-70B-Instruct \
#     --tensor-parallel-size 4 \
#     --speculative-model meta-llama/Llama-3-8B-Instruct \
#     --speculative-num-tokens 5

# 或在 Engine 中配置
llm = sgl.Engine(
    model_path="meta-llama/Llama-3-70B-Instruct",
    tensor_parallel_size=4,
    speculative_model="meta-llama/Llama-3-8B-Instruct",
    speculative_num_tokens=5,
)

output = llm.generate("解释量子计算", sampling_params={"max_new_tokens": 512})

from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载带 Medusa 的模型
model = AutoModelForCausalLM.from_pretrained(
    "princeton-nlp/Llama-7B-Medusa",
    torch_dtype="auto",
    device_map="auto"
)

# Medusa 解码
tokenizer = AutoTokenizer.from_pretrained("princeton-nlp/Llama-7B-Medusa")
input_ids = tokenizer("解释量子计算", return_tensors="pt").input_ids.to(model.device)

# 使用 medusa_decode 而非标准 generate
outputs = model.medusa_decode(input_ids, medusa_temperature=0.7)
print(tokenizer.decode(outputs[0]))

class MoEInferenceArchitecture:
    """MoE推理架构设计"""

    def __init__(self, model_name: str):
        self.model_name = model_name
        self.expert_config = self.get_expert_config(model_name)

    def get_expert_config(self, model_name: str) -> dict:
        """获取各MoE模型的专家配置"""
        configs = {
            # DeepSeek-V3: 671B总参数，37B激活
            "DeepSeek-V3": {
                "total_params": "671B",
                "activated_params": "37B",
                "num_shared_experts": 1,
                "num_routed_experts": 256,
                "top_k": 8,
                "expert_parallel": "EP144",  # Decode阶段
                "部署单元": "18节点",
                "每卡专家数": "2路由 + 1共享",
            },

            # Mixtral 8x7B: 总46B，激活约12B
            "Mixtral-8x7B": {
                "total_params": "46.7B",
                "activated_params": "12B",
                "num_experts": 8,
                "top_k": 2,
                "expert_type": "SwiGLU",
                "部署": "单机8卡",
            },

            # DeepSeek-MoE-16B: 总16B，激活约2B
            "DeepSeek-MoE-16B": {
                "total_params": "16.4B",
                "activated_params": "2B",
                "num_shared_experts": 1,
                "num_routed_experts": 64,
                "top_k": 4,
                "特点": "细粒度专家，组合自由度更高",
            },
        }
        return configs.get(model_name, {})

    def design_ep_layout(self, num_gpus: int, model_name: str) -> dict:
        """
        设计专家并行布局

        Args:
            num_gpus: 可用GPU数量
            model_name: 模型名称

        Returns:
            专家并行配置
        """
        # DeepSeek-V3 推理部署配置
        if "DeepSeek-V3" in model_name:
            return {
                "prefill": {
                    "expert_parallel": 32,  # EP32
                    "dp": 32,
                    "部署单元": "4节点",
                    "冗余专家": 32,
                    "每卡路由专家": 9,
                    "每卡共享专家": 1,
                },
                "decode": {
                    "expert_parallel": 144,  # EP144
                    "dp": 144,
                    "部署单元": "18节点",
                    "冗余专家": 32,
                    "每卡路由专家": 2,
                    "每卡共享专家": 1,
                }
            }
        return {}

class ExpertLoadBalancer:
    """专家负载均衡器"""

    def __init__(self, num_experts: int, top_k: int):
        self.num_experts = num_experts
        self.top_k = top_k
        self.expert_counts = [0] * num_experts  # 统计各专家使用次数
        self.expert_bias = [0.0] * num_experts   # 专家偏置（用于负载均衡）

    def router(self, input_features: torch.Tensor, temperature: float = 1.0) -> tuple:
        """
        路由器：为每个token选择top_k个专家

        Args:
            input_features: 输入特征 [batch_size, seq_len, hidden_dim]
            temperature: 温度参数（控制分布平滑度）

        Returns:
            (expert_indices, weights): 选中的专家索引和权重
        """
        # 门控网络计算专家得分
        # 实际实现中这是一个线性层 W_g
        gates = self.compute_gates(input_features)  # [batch_size, seq_len, num_experts]

        # 添加偏置用于负载均衡（DeepSeek的无辅助损失策略）
        gates = gates + torch.tensor(self.expert_bias, device=gates.device)

        # 选择top_k专家
        top_k_gates, top_k_indices = torch.topk(gates, self.top_k, dim=-1)

        # Softmax归一化权重
        top_k_weights = F.softmax(top_k_gates / temperature, dim=-1)

        # 更新统计（用于监控和调整）
        self.update_statistics(top_k_indices)

        return top_k_indices, top_k_weights

    def update_statistics(self, expert_indices: torch.Tensor):
        """更新专家使用统计"""
        # 统计每个专家被选中的次数
        for idx in expert_indices.flatten().unique():
            self.expert_counts[idx.item()] += 1

    def adjust_bias(self, target_load: float = 1.0):
        """
        调整专家偏置以实现负载均衡

        DeepSeek的无辅助损失策略：
        - 根据实际负载动态调整专家偏置
        - 避免辅助损失对模型性能的干扰
        - 保持训练稳定性

        Args:
            target_load: 目标负载率
        """
        total_tokens = sum(self.expert_counts)
        if total_tokens == 0:
            return

        for i in range(self.num_experts):
            actual_load = self.expert_counts[i] / total_tokens
            # 如果实际负载低于目标，增加偏置使其更可能被选中
            # 如果实际负载高于目标，减少偏置使其更少被选中
            self.expert_bias[i] += (target_load - actual_load) * 0.1

        # 重置计数
        self.expert_counts = [0] * self.num_experts

    def get_load_balance_metrics(self) -> dict:
        """获取负载均衡指标"""
        total = sum(self.expert_counts)
        if total == 0:
            return {"status": "no_data"}

        loads = [c / total for c in self.expert_counts]
        return {
            "expert_loads": loads,
            "std": np.std(loads),  # 标准差越小越均衡
            "max_load": max(loads),
            "min_load": min(loads),
            "balance_ratio": min(loads) / max(loads) if max(loads) > 0 else 0,
        }