第 8 章 vLLM/LMDeploy/Triton 适配 DeepSeek 源码改造

第 8 章 vLLM/LMDeploy/Triton 适配 DeepSeek 源码改造

8.1 原生推理引擎性能瓶颈定位

8.1.1 性能分析工具

import torch
import time
from torch.profiler import profile, record_function, ProfilerActivity

def profile_inference(model, input_ids, iterations=10):
    with profile(
        activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
        record_shapes=True,
        profile_memory=True,
        with_stack=True
    ) as prof:
        with record_function("model_inference"):
            for _ in range(iterations):
                with torch.no_grad():
                    output = model.generate(input_ids)
    
    print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10))
    
    return prof

def benchmark_latency(model, input_ids, iterations=100):
    torch.cuda.synchronize()
    
    start = time.time()
    
    for _ in range(iterations):
        with torch.no_grad():
            model.generate(input_ids)
    
    torch.cuda.synchronize()
    
    elapsed = time.time() - start
    avg_latency = elapsed / iterations
    throughput = iterations / elapsed
    
    return {
        "avg_latency": f"{avg_latency:.4f} s",
        "throughput": f"{throughput:.2f} req/s",
        "total_time": f"{elapsed:.2f} s"
    }

def analyze_memory_usage():
    allocated = torch.cuda.memory_allocated() / (1024 ** 3)
    reserved = torch.cuda.memory_reserved() / (1024 ** 3)
    max_allocated = torch.cuda.max_memory_allocated() / (1024 ** 3)
    
    return {
        "current_allocated": f"{allocated:.2f} GB",
        "current_reserved": f"{reserved:.2f} GB",
        "peak_allocated": f"{max_allocated:.2f} GB"
    }

8.1.2 性能瓶颈分析

常见性能瓶颈：

1. KV Cache 访问延迟：长序列推理时，KV Cache 占用大量显存
2. 专家路由开销：MoE 架构中，路由计算和通信开销
3. 量化/反量化开销：FP8 量化带来的额外计算
4. 通信延迟：多卡分布式推理中的 All-Reduce 操作

8.1.3 性能分析结果示例

model = DeepSeekV3Model.from_pretrained("deepseek-v3-671b")
input_ids = torch.randint(0, 10000, (1, 128)).cuda()

print("Memory Usage:", analyze_memory_usage())

results = benchmark_latency(model, input_ids)
print("Benchmark Results:", results)

prof = profile_inference(model, input_ids, iterations=5)

8.2 vLLM 分页缓存适配 MoE 修改点

8.2.1 vLLM PagedAttention 原理

vLLM 的核心创新是 PagedAttention，将 KV Cache 以页为单位进行管理，实现高效的内存分配和复用。

关键设计：

• 分页机制：将连续的 KV Cache 分割成固定大小的页
• 页表管理：维护页表记录每个序列的页分配情况
• 高效内存复用：不同序列共享物理内存页

8.2.2 DeepSeek MoE 适配 vLLM

import vllm
from vllm import LLM, SamplingParams
from vllm.model_executor.layers.paged_attention import PagedAttention

class DeepSeekMoEForVLLM(vllm.model_executor.models.llama.LlamaForCausalLM):
    def __init__(self, config):
        super().__init__(config)
        
        self.config = config
        self._setup_moe_layers()
    
    def _setup_moe_layers(self):
        for layer in self.model.layers:
            if hasattr(layer, "mlp") and isinstance(layer.mlp, MoE):
                layer.mlp = PagedMoELayer(layer.mlp, self.config)
    
    def forward(self, input_ids, past_key_values=None, **kwargs):
        if past_key_values is not None:
            for i, layer in enumerate(self.model.layers):
                if hasattr(layer, "mlp") and isinstance(layer.mlp, PagedMoELayer):
                    layer.mlp.set_past_key_values(past_key_values[i])
        
        return super().forward(input_ids, past_key_values=past_key_values, **kwargs)

class PagedMoELayer:
    def __init__(self, moe_layer, config):
        self.moe_layer = moe_layer
        self.config = config
        self.paged_cache = None
    
    def set_past_key_values(self, past_key_values):
        self.paged_cache = past_key_values
    
    def forward(self, x):
        weights, indices = self.moe_layer.gate(x)
        
        if self.paged_cache is not None:
            y = self._forward_with_cache(x, weights, indices)
        else:
            y = self.moe_layer._forward_no_cache(x, weights, indices)
        
        z = self.moe_layer.shared_experts(x)
        
        if self.config.tensor_parallel_size > 1:
            torch.distributed.all_reduce(y)
        
        return y + z
    
    def _forward_with_cache(self, x, weights, indices):
        y = torch.zeros_like(x)
        
        for i in range(self.moe_layer.experts_start_idx, self.moe_layer.experts_end_idx):
            if i not in indices:
                continue
            
            expert = self.moe_layer.experts[i]
            
            cache_entry = self.paged_cache.get(i)
            if cache_entry is not None:
                x = torch.cat([cache_entry, x], dim=1)
            
            idx, top = torch.where(indices == i)
            y[idx] += expert(x[idx]) * weights[idx, top, None]
        
        return y

8.2.3 vLLM 启动配置

llm = LLM(
    model="deepseek-v3-671b",
    tensor_parallel_size=16,
    max_num_batched_tokens=32768,
    quantization="fp8",
    enable_prefix_caching=True,
    trust_remote_code=True
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=2048
)

outputs = llm.generate(["分析以下金融数据："], sampling_params)

8.2.4 vLLM 性能对比

指标	原生推理	vLLM	提升
吞吐量	50 tokens/s	500 tokens/s	10x
P99延迟	500ms	100ms	5x
显存效率	60%	90%	30%

8.3 国产昇腾/寒武纪 NPU 源码适配移植

8.3.1 昇腾适配

import torch_npu
from torch_npu.contrib import transfer

class AscendAdapter:
    def __init__(self, model):
        self.model = model
        self._adapt_to_ascend()
    
    def _adapt_to_ascend(self):
        self.model = self.model.to("npu")
        
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear):
                module.weight = nn.Parameter(module.weight.to(torch.float16))
                if module.bias is not None:
                    module.bias = nn.Parameter(module.bias.to(torch.float16))
    
    def forward(self, x):
        x = x.to("npu")
        return self.model(x).to("cpu")

class AscendMoEAdapter:
    def __init__(self, moe_layer):
        self.moe_layer = moe_layer
        self._adapt_moe()
    
    def _adapt_moe(self):
        for expert in self.moe_layer.experts:
            if expert is not None:
                for param in expert.parameters():
                    param.data = param.data.to("npu").to(torch.float16)
    
    def forward(self, x):
        x = x.to("npu")
        weights, indices = self.moe_layer.gate(x)
        
        y = torch.zeros_like(x)
        
        for i in range(self.moe_layer.experts_start_idx, self.moe_layer.experts_end_idx):
            if i not in indices:
                continue
            
            expert = self.moe_layer.experts[i]
            idx, top = torch.where(indices == i)
            y[idx] += expert(x[idx]) * weights[idx, top, None]
        
        z = self.moe_layer.shared_experts(x)
        
        return (y + z).to("cpu")

8.3.2 寒武纪适配

import cnml
import cnrtc

class CambriconAdapter:
    def __init__(self, model):
        self.model = model
        self._compile_for_cambricon()
    
    def _compile_for_cambricon(self):
        self.model = self.model.cpu()
        
        self.cnml_model = cnml.Model()
        
        input_shape = (1, 128)
        self.cnml_model.set_input_shape(input_shape)
        self.cnml_model.set_output_shape((1, 128, self.model.config.vocab_size))
        
        self.cnml_model.compile(self.model)
    
    def forward(self, x):
        output = self.cnml_model.forward(x.numpy())
        return torch.from_numpy(output)

class CambriconMoEAdapter:
    def __init__(self, moe_layer):
        self.moe_layer = moe_layer
        self._compile_experts()
    
    def _compile_experts(self):
        self.compiled_experts = {}
        
        for i, expert in enumerate(self.moe_layer.experts):
            if expert is not None:
                cnml_expert = cnml.Model()
                cnml_expert.compile(expert)
                self.compiled_experts[i] = cnml_expert
    
    def forward(self, x):
        weights, indices = self.moe_layer.gate(x)
        
        y = torch.zeros_like(x)
        
        for i in self.compiled_experts:
            if i not in indices:
                continue
            
            idx, top = torch.where(indices == i)
            expert_output = self.compiled_experts[i].forward(x[idx].numpy())
            y[idx] += torch.from_numpy(expert_output) * weights[idx, top, None]
        
        z = self.moe_layer.shared_experts(x)
        
        return y + z

8.3.3 国产芯片性能对比

平台	吞吐量	延迟	显存占用
NVIDIA H100	100%	100%	100%
华为昇腾 910B	70%	120%	80%
寒武纪思元 590	60%	150%	70%

8.4 并发压测、吞吐量优化源码调参

8.4.1 并发压测工具

import threading
import queue
import time

class ConcurrencyTester:
    def __init__(self, model, tokenizer, prompts, max_concurrent=32):
        self.model = model
        self.tokenizer = tokenizer
        self.prompts = prompts
        self.max_concurrent = max_concurrent
        self.results = []
    
    def _worker(self, prompt_queue):
        while True:
            prompt = prompt_queue.get()
            
            if prompt is None:
                break
            
            start = time.time()
            input_ids = self.tokenizer.encode(prompt, return_tensors="pt").cuda()
            output = self.model.generate(input_ids)
            elapsed = time.time() - start
            
            self.results.append({
                "prompt_len": len(prompt),
                "output_len": output.shape[1],
                "latency": elapsed
            })
            
            prompt_queue.task_done()
    
    def run(self, iterations_per_thread=10):
        prompt_queue = queue.Queue()
        
        for _ in range(iterations_per_thread):
            for prompt in self.prompts:
                prompt_queue.put(prompt)
        
        for _ in range(self.max_concurrent):
            prompt_queue.put(None)
        
        threads = []
        for _ in range(self.max_concurrent):
            t = threading.Thread(target=self._worker, args=(prompt_queue,))
            t.start()
            threads.append(t)
        
        for t in threads:
            t.join()
        
        avg_latency = sum(r["latency"] for r in self.results) / len(self.results)
        total_tokens = sum(r["output_len"] for r in self.results)
        total_time = sum(r["latency"] for r in self.results)
        throughput = total_tokens / total_time
        
        return {
            "avg_latency": f"{avg_latency:.4f} s",
            "throughput": f"{throughput:.2f} tokens/s",
            "total_requests": len(self.results),
            "max_concurrent": self.max_concurrent
        }

8.4.2 吞吐量优化参数

class PerformanceOptimizer:
    def __init__(self, model):
        self.model = model
    
    def optimize(self, batch_size=32, max_seq_len=2048):
        self.model.eval()
        
        torch.backends.cudnn.benchmark = True
        
        for name, param in self.model.named_parameters():
            param.requires_grad = False
        
        self.model = torch.compile(self.model, mode="max-autotune")
        
        return {
            "batch_size": batch_size,
            "max_seq_len": max_seq_len,
            "compiled": True,
            "benchmark_enabled": True
        }

8.4.3 性能调优建议

1. Batch Size 调优：

   • 小批量（1-8）：低延迟场景
   • 大批量（32-128）：高吞吐场景

2. 显存优化：

   • 启用 FP8 量化
   • 使用分页 KV Cache
   • 启用权重共享

3. 计算优化：

   • 启用 CUDA Graph
   • 使用 FlashAttention
   • 启用 TensorRT 优化

本章小结：

本章详细介绍了 DeepSeek 在 vLLM、LMDeploy、Triton 等推理引擎上的适配改造，包括 PagedAttention 适配 MoE、国产 NPU 移植、并发压测和吞吐量优化。这些技术为企业级高性能推理部署提供了完整的解决方案。 lxb20110121