GPU优化

Flash Attention解决了什么问题？ 标准注意力计算需要将整个N×N的注意力矩阵存储在GPU高带宽内存（HBM）中。对于长序列，这个矩阵非常大——128K序列长度的注意力矩阵需要约64GB HBM。GPU核心（SM）与HBM之间的数据搬运成为瓶颈。 Flash Attention的核心创新：不在HBM中实例化完整注意力矩阵，而是在SRAM中分块计算。 GPU内存层次 理解Flash Attention需要先理解GPU的内存层次： SRAM (片上共享内存) ├── 延迟: ~20 cycles ├── 带宽: ~19 TB/s (A100) └── 容量: ~192KB per SM HBM (高带宽内存) ├── 延迟: ~200+ cycles ├── 带宽: ~2 TB/s (A100) └── 容量: 80GB (A100) 标准注意力的问题：在HBM中读写O(n²)大小的矩阵，受限于2TB/s的HBM带宽。 Flash Attention的解决思路：将计算分块，每块在SRAM中完成，减少HBM访问次数。 算法原理 标准注意力计算 S = Q @ K^T / √d # [N, N] 注意力分数 P = softmax(S) # [N, N] 归一化 O = P @ V # [N, d] 输出 问题：S和P是N×N矩阵，需要完整存储在HBM中。 ...

注意力计算的性能瓶颈 标准注意力计算的核心瓶颈不在于计算量，而在于显存访问。注意力矩阵 QK^T 的形状为 [n, n]，对于 n=8192 的序列，仅这个矩阵就需要256MB显存（FP32）。GPU的计算速度远超显存带宽——A100的计算能力为312 TFLOPS，而HBM带宽仅2TB/s。 这就是经典的"内存墙"问题：注意力计算大部分时间不是在算，而是在等数据搬运。 Flash Attention V1：分块计算的突破 Tri Dao在2022年提出的Flash Attention V1核心思想是：通过分块计算避免实例化完整的注意力矩阵。 在线Softmax Flash Attention的关键创新是在线Softmax算法。标准Softmax需要先遍历一次数据求最大值，再遍历一次计算指数和归一化。在线Softmax通过维护一个运行最大值和运行和，在单次遍历中完成Softmax计算： def online_softmax(blocks): """在线Softmax：逐块更新""" m = float('-inf') # 运行最大值 l = 0.0 # 运行和 result = None for block in blocks: block_max = block.max() # 更新全局最大值 new_m = max(m, block_max) # 修正之前的累积值 alpha = math.exp(m - new_m) beta = math.exp(block_max - new_m) l = alpha * l + beta * block.sum() if result is None: result = beta * block else: result = alpha * result + beta * block m = new_m return result / l 分块策略 将Q、K、V分别切分为块，每次只加载一小块到SRAM（GPU片上高速缓存）中计算。外层循环遍历K/V块，内层循环遍历Q块： for each K_j, V_j block: load K_j, V_j to SRAM for each Q_i block: load Q_i to SRAM compute S = Q_i @ K_j^T update online softmax statistics compute partial O = P @ V_j write O_i back to HBM 这种策略使得HBM访问量从O(n²)降低到O(n²d/M)，其中M是SRAM大小。对于典型配置，这相当于减少了10-20倍的显存访问。 Flash Attention V2：效率的进一步提升 V2在V1基础上做了几个关键优化： 减少非矩阵乘法运算 GPU上矩阵乘法（GEMM）的效率远高于其他运算（如rescale、softmax）。V2重新组织了计算顺序，将rescale操作推迟到最后，减少了中间的rescale次数。 并行化改进 V1主要沿序列维度并行化，V2增加了批次和头维度的并行化，使得在长序列场景下能更好地利用GPU的并行能力。 前向和后向的分配优化 V2重新分配了前向和后向传播中的工作负载，减少了线程块之间的同步开销。 Flash Attention V3：FP8与异步流水线 2024年发布的Flash Attention V3针对H100 GPU的FP8张量核心和异步执行特性进行了深度优化。 ...

Flash Attention 3：让 GPU 跑满的注意力计算 Flash Attention 系列是近年来大模型工程领域最重要的优化之一。从 Flash Attention 1 到 3，每代都在逼近 GPU 硬件的理论极限。2026 年，Flash Attention 3 已经成为所有主流大模型推理和训练的标配。本文将深入解析其原理。 一、问题：标准注意力的内存瓶颈 1.1 GPU 内存层次 现代 GPU（如 H100）有复杂的内存层次： ┌─────────────────────────────────────────────┐ │ GPU 内存层次 │ ├─────────────────────────────────────────────┤ │ │ │ ┌─────────┐ 延迟: ~20 cycles │ │ │Register │ 带宽: ~30 TB/s │ │ │(SRAM) │ 容量: 256 KB/SM │ │ └─────────┘ │ │ ↑ │ │ ┌─────────┐ 延迟: ~200 cycles │ │ │L2 Cache │ 带宽: ~12 TB/s │ │ │ │ 容量: 50 MB │ │ └─────────┘ │ │ ↑ │ │ ┌─────────┐ 延迟: ~400+ cycles │ │ │ HBM │ 带宽: ~3.35 TB/s (H100) │ │ │(显存) │ 容量: 80 GB │ │ └─────────┘ │ │ │ └─────────────────────────────────────────────┘ 关键洞察：HBM 带宽只有 SRAM 的 1/9，但标准注意力几乎完全在 HBM 上操作。 ...

NVIDIA 的推理加速王牌 TensorRT-LLM 是 NVIDIA 在大模型推理领域的旗舰产品。它不是一个独立的推理引擎，而是基于 TensorRT 的 LLM 专用优化层——通过算子融合、精度优化、内存布局优化等技术，在 NVIDIA GPU 上榨取每一分性能。 2026 年，随着 Blackwell 架构 GPU 的普及，TensorRT-LLM 的优势进一步扩大——它对 Blackwell 的 Transformer Engine 和 FP4 精度提供了原生支持。 2026 核心特性 性能优势 特性 TensorRT-LLM 2026 vLLM 0.8 SGLang 0.3 峰值吞吐量 8,500 tok/s 4,200 tok/s 4,800 tok/s 首 Token 延迟 0.15s 0.5s 0.35s FP4 支持 ✅ (Blackwell) ❌ ❌ FP8 支持 ✅ (Hopper) ✅ ✅ 算子融合 深度 基础 基础 模型编译 AOT 编译 JIT JIT 多 GPU TP + PP + EP TP + PP TP + PP Blackwell 架构优化 # Blackwell B200 上的 FP4 推理 import tensorrt_llm as trtllm # 编译模型为 FP4 精度 builder = trtllm.Builder() config = builder.create_builder_config( precision="fp4", # FP4 量化 plugin_config=trtllm.PluginConfig( attention_plugin=True, nccl_plugin=True, gemm_plugin=True, rmsnorm_plugin=True, # Blackwell 专属 transformer_engine=True, moe_plugin=True, ), max_batch_size=256, max_input_len=32768, max_output_len=4096, max_num_tokens=8192, use_paged_context_fmha=True, # Paged Attention use_context_fmha=True, # Flash Attention multiple_profiles=True, # 多优化 Profile tensor_parallel=8, # 8 路张量并行 pipeline_parallel=1, ) # 编译（AOT，提前编译为优化引擎） engine = builder.build( model_dir="Qwen/Qwen3-72B-Instruct", config=config, output_dir="engines/qwen3-72b-fp4" ) # 编译后的引擎不可移植，绑定特定 GPU 架构 # 但性能比 JIT 方案高 30-60% 部署流程 1. 模型编译 # 步骤 1：从 HuggingFace 模型编译 TensorRT 引擎 import tensorrt_llm as trtllm from tensorrt_llm.models import QWenForCausalLM # 加载模型配置 model_config = trtllm.ModelConfig.from_pretrained( "Qwen/Qwen3-72B-Instruct" ) # 编译配置 build_config = trtllm.BuildConfig( max_input_len=32768, max_output_len=4096, max_batch_size=128, max_num_tokens=8192, opt_batch_size=32, opt_input_len=4096, # 精度配置 precision="fp8", # fp4/fp8/fp16 int8_kv_cache=True, # KV Cache 量化 # 插件配置 plugin_config=trtllm.PluginConfig( paged_kv_cache=True, attention_plugin=True, gemm_plugin=True, nccl_plugin=True, rmsnorm_plugin=True, rotary_plugin=True, remove_input_padding=True, # 移除 padding 优化 ), # 并行配置 tensor_parallel=4, pipeline_parallel=1, # 高级优化 use_fused_mlp=True, # MLP 算子融合 use_fused_qkv=True, # QKV 融合 use_dynamic_shape=True, # 动态形状 weight_sparsity=True, # 权重稀疏化 ) # 编译引擎 builder = trtllm.Builder() engine = builder.build_model( model_config=model_config, build_config=build_config, output_dir="engines/qwen3-72b-fp8-tp4" ) 2. 启动服务 # 使用 Triton Inference Server 部署 # 模型仓库结构 models/ └── qwen3-72b/ ├── config.pbtxt ├── 1/ │ └── model.py └── engines/ └── qwen3-72b-fp8-tp4/ ├── rank0.engine ├── rank1.engine ├── rank2.engine └── rank3.engine # config.pbtxt name: "qwen3-72b" backend: "tensorrtllm" max_batch_size: 128 input [ { name: "input_ids" data_type: TYPE_INT32 dims: [ -1 ] }, { name: "input_lengths" data_type: TYPE_INT32 dims: [ 1 ] } ] output [ { name: "output_ids" data_type: TYPE_INT32 dims: [ -1, -1 ] } ] dynamic_batching { preferred_batch_size: [ 4, 8, 16, 32 ] max_queue_delay_microseconds: 100000 } instance_group [ { count: 1 kind: KIND_GPU } ] parameters [ { key: "tensorrt_llm_model_dir" value: { string_value: "/models/qwen3-72b/engines/qwen3-72b-fp8-tp4" } }, { key: "max_output_len" value: { string_value: "4096" } }, { key: "temperature" value: { string_value: "0.7" } }, { key: "top_p" value: { string_value: "0.9" } } ] # 启动 Triton Server docker run --gpus all -p 8000:8000 -p 8001:8001 -p 8002:8002 \ -v /models:/models \ nvcr.io/nvidia/tritonserver:25.06-py3 \ tritonserver --model-repository=/models \ --backend-directory=/opt/tritonserver/backends \ --log-verbose=1 3. 客户端调用 import tritonclient.grpc as grpcclient import numpy as np class TensorRTLLMClient: def __init__(self, url="localhost:8001"): self.client = grpcclient.InferenceServerClient(url) def generate(self, prompt: str, max_tokens: int = 512, temperature: float = 0.7, stream: bool = True): # Tokenize input_ids = self.tokenizer.encode(prompt) inputs = [ grpcclient.InferInput("input_ids", [1, len(input_ids)], "INT32"), grpcclient.InferInput("input_lengths", [1, 1], "INT32"), ] inputs[0].set_data_from_numpy(np.array([input_ids], dtype=np.int32)) inputs[1].set_data_from_numpy(np.array([[len(input_ids)]], dtype=np.int32)) outputs = [grpcclient.InferRequestedOutput("output_ids")] # 推理 result = self.client.infer( model_name="qwen3-72b", inputs=inputs, outputs=outputs ) output_ids = result.as_numpy("output_ids") # Detokenize return self.tokenizer.decode(output_ids[0]) 性能优化 精度对比 精度 显存 (72B) 吞吐量 质量损失 推荐 GPU FP16 145 GB 3,200 tok/s 0% A100 80GB×2 FP8 75 GB 5,800 tok/s <1% H100/H200 INT4 AWQ 42 GB 4,500 tok/s ~3% 任意 FP4 38 GB 8,500 tok/s ~5% B200 算子融合效果 优化 吞吐量提升 延迟降低 基础（无融合） 基准 基准 QKV 融合 +15% -8% + MLP 融合 +25% -15% + RMSNorm 融合 +30% -20% + Rotary 融合 +35% -22% + 全部融合 +42% -28% 多 GPU 扩展 # 张量并行 + 流水线并行配置 config = trtllm.BuildConfig( tensor_parallel=4, # 4 路张量并行 pipeline_parallel=2, # 2 路流水线并行 # 总共 8 GPU # 专家并行（MoE 模型） moe_config=trtllm.MoEConfig( num_experts=256, expert_parallel_size=8, moe_plugin=True, ), ) 并行策略 GPU 数量 吞吐量 扩展效率 TP=1 1 3,200 100% TP=2 2 5,800 91% TP=4 4 9,500 74% TP=8 8 15,200 59% TP=4+PP=2 8 14,800 58% 与 vLLM 对比 维度 TensorRT-LLM vLLM 峰值性能 ⭐⭐⭐⭐⭐ ⭐⭐⭐⭐ 部署难度 高 低 模型支持 跟随 NVIDIA 跟随社区 编译时间 10-60 分钟 即时 灵活性 低（AOT） 高（JIT） 跨平台 仅 NVIDIA NVIDIA + AMD 生态 NVIDIA 生态 开源生态 成本 需要 NVIDIA GPU 任意 GPU 适用场景 最适合 极致性能需求：延迟和吞吐量最优 NVIDIA 纯净环境：充分利用 GPU 特性 固定模型部署：AOT 编译换取性能 大规模生产：Triton Server 集群部署 Blackwell 用户：FP4 独家支持 不太适合 快速迭代：每次模型变更需要重新编译 多 GPU 品牌：仅支持 NVIDIA 小团队：部署和调优门槛高 实验性模型：新模型架构支持滞后于 vLLM 成本敏感：需要 NVIDIA GPU 许可 总结 TensorRT-LLM 在 2026 年仍然是"NVIDIA GPU 上最快的推理引擎"。它的 AOT 编译、深度算子融合、FP4/FP8 支持，让它在峰值性能上领先 vLLM 60-100%。这个优势在 Blackwell 架构上更加明显。 ...

为什么需要 Flash Attention？ 标准 Attention 的瓶颈不在计算，而在显存 IO。考虑 $n = 8192$, $d = 128$ 的注意力计算： 操作 HBM 读写量 SRAM 计算 $S = QK^T$ 读 Q(8MB) + 读 K(8MB) + 写 S(512MB) ~16 GFLOPs $P = \text{softmax}(S)$ 读 S(512MB) + 写 P(512MB) ~8 MFLOPs $O = PV$ 读 P(512MB) + 读 V(8MB) + 写 O(8MB) ~16 GFLOPs 总计 ~1.6 GB ~32 GFLOPs A100 的 HBM 带宽 2 TB/s，SRAM 带宽 19 TB/s。计算只需 ~0.02ms，但 IO 需要 ~0.8ms——IO 是瓶颈的 40 倍。 ...

推理优化的三个维度 显存优化 ↑ | 计算优化 ←──→ 通信优化 LLM 推理的瓶颈不是计算，而是显存带宽。7B 模型生成 1 个 Token 需要读取全部 7B 参数，但只做 1 次矩阵乘法。 显存分析 7B 模型显存分布 模型权重 (FP16): 14 GB KV Cache (2K ctx): 1 GB KV Cache (8K ctx): 4 GB KV Cache (32K ctx): 16 GB ← 比模型还大！ KV Cache (128K ctx): 64 GB 总显存 (8K ctx): 18 GB → A100 40G 可以跑 总显存 (32K ctx): 30 GB → A100 40G 勉强 总显存 (128K ctx): 78 GB → 需要 A100 80G KV Cache 计算 def kv_cache_size( num_layers: int, # 层数 num_heads: int, # 注意力头数 head_dim: int, # 每头维度 seq_len: int, # 序列长度 batch_size: int, # 批大小 dtype_size: int = 2, # FP16 = 2 bytes ): """计算 KV Cache 显存""" # K 和 V 各一份 return 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size # Qwen3-7B 示例 size = kv_cache_size( num_layers=32, num_heads=32, head_dim=128, seq_len=8192, batch_size=1, ) print(f"{size / 1024**3:.2f} GB") # 4.00 GB 优化技术一：量化 INT8 量化 # 权重 INT8：14GB → 7GB # KV Cache INT8：显存减半 # 性能损失：1-3% from transformers import AutoModelForCausalLM import torch model = AutoModelForCausalLM.from_pretrained( "Qwen3-7B-Instruct", load_in_8bit=True, # 权重 INT8 device_map="auto", ) INT4 量化 # 权重 INT4：14GB → 3.5GB # 性能损失：3-8% # 适合消费级 GPU model = AutoModelForCausalLM.from_pretrained( "Qwen3-7B-Instruct", load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", ) AWQ 量化 # AWQ：比 INT4 更好的量化方案 # 保护重要权重，精度损失更小 # vLLM 原生支持 # 量化模型 python -m awq.quantize \ --model-path /models/Qwen3-7B \ --output-path /models/Qwen3-7B-AWQ \ --w-bits 4 # vLLM 加载 llm = LLM(model="/models/Qwen3-7B-AWQ", quantization="awq") 优化技术二：KV Cache 优化 PagedAttention（vLLM） # 传统：预分配最大长度的连续显存 # 浪费率：60%（大量预分配空间未使用） # PagedAttention：按需分配，像操作系统的分页机制 # 浪费率：<5% # vLLM 默认启用 llm = LLM( model="...", gpu_memory_utilization=0.9, # 95% 显存用于 KV Cache ) KV Cache 量化 # KV Cache 用 INT8 存储，显存减半 llm = LLM( model="...", kv_cache_dtype="int8", # 或 "fp8" ) KV Cache 共享 # 多个请求共享相同的 system prompt KV Cache # vLLM 的 prefix caching llm = LLM( model="...", enable_prefix_caching=True, ) # 效果： # Request 1: "You are a helpful assistant. [Q1]" → 生成 KV Cache # Request 2: "You are a helpful assistant. [Q2]" → 复用 KV Cache # 节省：system prompt 部分的计算和显存 优化技术三：计算优化 Flash Attention # Flash Attention：减少 GPU 显存读写 # 标准 Attention：O(n²) 显存 # Flash Attention：O(n) 显存，速度提升 2-4x # vLLM 默认启用 llm = LLM( model="...", enforce_eager=False, # 启用 CUDA Graph + Flash Attention ) CUDA Graph # CUDA Graph：预编译计算图，减少 kernel launch 开销 # 适合：固定形状的推理（批处理） llm = LLM( model="...", enforce_eager=False, # 启用 CUDA Graph ) # 延迟降低 10-20% Continuous Batching # 传统批处理：等所有请求完成才处理下一批 # Continuous Batching：每个 step 动态加入/移除请求 # vLLM 默认使用 continuous batching # 效果：吞吐量提升 3-10x # 原理： # Time 0: [A, B, C] 开始生成 # Time 5: A 完成，D 加入 → [B, C, D] # Time 8: B 完成，E 加入 → [C, D, E] # GPU 始终满载 优化技术四：Speculative Decoding # 原理：用小模型猜，大模型验 # 小模型生成 5 个 token → 大模型一次验证 5 个 # 如果全部正确：5x 加速 # 如果 3 个正确：3x 加速 llm = LLM( model="/models/Qwen3-7B-Instruct", speculative_model="/models/Qwen3-0.5B-Instruct", num_speculative_tokens=5, ) # 实测效果： # 简单任务（翻译、摘要）：2.0x 加速 # 复杂任务（推理、创作）：1.3x 加速 # 代码生成：1.8x 加速 优化技术五：模型架构优化 GQA（Grouped Query Attention） # MHA：每个 head 有独立的 K, V → KV Cache 大 # MQA：所有 head 共享 K, V → KV Cache 小，但质量降 # GQA：折中，g 个 head 共享 K, V # Qwen3 使用 GQA，KV Cache 减少为 MHA 的 1/4 # 显存节省 75%，质量损失 <1% 模型蒸馏 # 大模型蒸馏到小模型 # Teacher: Qwen3-72B → Student: Qwen3-7B # 保留 90% 能力，推理快 10x，显存省 10x # 适用场景：固定任务，不需要通用能力 基准测试 测试方法 import time import asyncio async def benchmark(llm, prompts, max_tokens=512): start = time.time() results = await asyncio.gather(*[ llm.generate(p, SamplingParams(max_tokens=max_tokens)) for p in prompts ]) elapsed = time.time() - start total_tokens = sum(len(r.outputs[0].token_ids) for r in results) return { "total_time": elapsed, "total_tokens": total_tokens, "throughput": total_tokens / elapsed, # tokens/s "avg_latency": elapsed / len(prompts), } 参考数据 配置 模型 吞吐量 延迟 A100 80G Qwen3-7B 5000 tok/s 50ms A100 80G Qwen3-72B 800 tok/s 200ms 4090 24G Qwen3-7B (AWQ) 3000 tok/s 80ms 4090 24G Qwen3-7B (FP16) OOM - 2x 4090 Qwen3-72B (AWQ) 400 tok/s 250ms 生产环境配置建议 # 配置决策树 def recommend_config(gpu, model_size, concurrency): if gpu == "A100 80G": if model_size <= 7: return {"tp": 1, "quant": "fp16", "max_seqs": 256} elif model_size <= 72: return {"tp": 1, "quant": "awq", "max_seqs": 64} elif gpu == "4090 24G": if model_size <= 7: return {"tp": 1, "quant": "awq", "max_seqs": 32} elif model_size <= 72: return {"tp": 2, "quant": "awq", "max_seqs": 16} return {"error": "configuration not supported"} 结论 LLM 推理优化的核心原则： ...