llm-algo-8

深入探讨的 Triton Flash Attention、PagedAttention、W8A16 Quantization GEMM 以及 Multi-LoRA 四大核心算子，你已经触及了现代大模型推理引擎（如 vLLM, SGLang, TensorRT-LLM）最底层的"四根支柱"。

---

第一阶段：夯实认知地基（高频复习与内化）

目标：将零散的算子知识串联成网，形成条件反射级的直觉。不要急于写新代码，先确保现有知识无死角。

建立"算子-瓶颈"映射表（高频查阅）

不要孤立记忆算子，要记忆它解决的物理瓶颈。建议整理并反复回顾以下矩阵：

算子	解决的核心瓶颈	关键数学/算法思想	硬件依赖	典型应用场景
Flash Attention	HBM 带宽 + O(N2) 显存	Tiling + Online Safe Softmax	SRAM 大小, Tensor Core	Prefill, 长文本训练
PagedAttention	KV Cache 显存碎片	虚拟内存分页 + 间接寻址	L2 Cache, 内存控制器	Decode, 高并发 Serving
W8A16 GEMM	权重读取带宽	On-the-fly Dequant + 融合	INT8/FP16 混合精度单元	Memory Bound 线性层
Multi-LoRA	多租户切换开销	Token-wise 指针路由 + 低秩分解	寄存器文件, 分支预测	SaaS 多模型并发

手推核心公式（深度研究）

代码会过时，但数学不会。必须能够脱离代码白板推导以下内容：

• Online Softmax 修正系数 α ：为什么 emold−mnew 能严格等价于全局 Softmax？尝试用归纳法证明。
• PagedAttention 地址翻译 ：给定逻辑 Token ID，写出从 block_table 到物理 HBM 地址的完整计算公式。
• 量化误差传播：推导 Per-channel vs Per-tensor 在矩阵乘法中的误差上界差异。

绘制数据流图（可视化内化）

对每个算子，画出 HBM ↔ SRAM ↔ Register 三级存储的数据流动图。重点标注：

• 哪些数据只进 SRAM 不回 HBM（如 FlashAttn 的 P 矩阵、W8A16 的反量化中间态）
• 哪些访存是间接的（如 PagedAttn 的 KV load、Multi-LoRA 的权重 pool 索引）
• 累加器的精度转换点（FP32 acc → FP16 store）

---

第二阶段：结构化实践体系（从复现到变体）

目标：通过"破坏-修复-扩展"循环，把教程代码变成自己的肌肉记忆。

基础复现与数值对齐（必做）

• 盲写测试：不看参考代码，仅凭数学公式和接口签名，独立实现四个算子的 Triton 版本。
• 多维度验证：不仅验证正确性，还要验证 边界条件：
  • SeqLen 非 BLOCK_SIZE 整数倍
  • HeadDim = 96 / 192 等非 2 的幂
  • Batch 内 LoRA 索引全相同 / 全不同
  • Context Length = 0 或 1 的极端情况

"破坏性"实验（深度理解的关键）

故意引入错误，观察现象，建立 debugging 直觉：

故意引入的错误	预期现象	学到的教训
FlashAttn m_i 初始化为 0	长序列/负值 qk 时结果偏移	Online Softmax 对初始值的敏感性
PagedAttn KV load other=-inf	NaN 或 Inf	Mask 填充值语义区分（数据 vs 分数）
W8A16 acc 用 FP16	K>1024 时误差指数增长	累加器精度的必要性
Multi-LoRA 去掉 mask	尾块垃圾值	分块计算的边界保护
所有算子去掉 sm_scale	Softmax 梯度消失/溢出	数值稳定性工程细节

渐进式扩展项目（实战能力提升）

按难度递增完成以下改造，每个都需附带 Benchmark 对比：

1. Causal Flash Attention ：在内层循环添加 mask = offs_m[:,None] >= offs_n[None,:]，验证自回归正确性。
2. Group-wise W8A16：将 Per-channel scale 改为每 128 列共享，修改 scale 索引逻辑，对比精度与性能权衡。
3. GQA/MQA PagedAttention：支持 KV Head 数 < Q Head 数，添加 head 分组归约逻辑。
4. Fused Base+LoRA：将 Base Model GEMM 与 LoRA GEMM 融合到一个 Kernel，用流水线隐藏 LoRA 访存延迟。
5. W8A16 + Multi-LoRA 融合：LoRA 权重也做 INT8 量化，在 SRAM 内同时完成反量化和路由。

---

第三阶段：深度研究与性能分析

目标：从"能跑"到"跑得快"，建立科学的性能分析方法论。

掌握 Profiling 工具链

• Nsight Compute：学会看 Roofline Model、Memory Workload Analysis、Source-level 热点。
• Triton Benchmark ：熟练使用 triton.testing.do_bench，理解 warmup、quantiles 的含义。
• 理论峰值计算：能根据 GPU 型号（如 A100/H100）手算 HBM 带宽上限和 Tensor Core TFLOPS 上限，判断算子是 Memory Bound 还是 Compute Bound。

建立性能分析模板

对每个算子，回答以下问题并形成文档：

• 当前实现的 Arithmetic Intensity (AI) 是多少？处于 Roofline 的哪个区域？
• HBM 读写总量是否与理论最优一致？有无冗余访存？
• SRAM 占用是否接近上限？能否增大 Tile Size 提升复用？
• Tensor Core 利用率如何？是否存在 padding/warp divergence？
• 与 cuBLAS / FlashAttention-2 / vLLM 官方实现的差距在哪？原因是什么？

前沿论文追踪清单

保持对以下方向的敏感度：

• FlashAttention-3 / Hopper 架构优化：WGMMA, TMA, Asynchronous Pipeline
• KV Cache 压缩：KV Quantization, Token Pruning, StreamingLLM
• 高效 LoRA：DoRA, LoRA+, rsLoRA, Merge-aware Serving
• Speculative Decoding：Draft Model 与 Target Model 的协同调度

---

第四阶段：系统级整合思维

目标：跳出单算子视角，理解推理引擎的全局设计。

理解算子间的交互

• PagedAttention 的 Block Size 如何影响 Flash Attention 的 Tiling 效率？
• W8A16 量化后的权重布局如何与 Multi-LoRA 的内存池兼容？
• Continuous Batching 调度器如何决定何时触发 Prefill Kernel vs Decode Kernel？

阅读工业级源码

带着问题读代码，而非泛读：

• vLLM ：paged_attention_kernel.cu/triton → 看工业级间接寻址和 GQA 实现
• SGLang：RadixAttention → 看 Prefix Caching 如何与 PagedAttn 结合
• FlashAttention-2/3：官方 Triton/CUDA 实现 → 看极致优化的 Warp Specialization
• Punica / S-LoRA：看 Multi-LoRA 的生产级调度与内存管理

构建个人知识库

建议使用 Notion/Obsidian 维护以下结构：

LLM Inference Optimization/
├── 01_Math_Foundations/     # Online Softmax, 量化理论, 低秩分解
├── 02_Hardware_Arch/        # GPU 存储层次, Tensor Core, Nsight 笔记
├── 03_Kernel_Implementations/
│   ├── FlashAttention/      # 代码 + 踩坑记录 + Benchmark
│   ├── PagedAttention/
│   ├── QuantizedGEMM/
│   └── MultiLoRA/
├── 04_Performance_Analysis/ # Roofline 分析, Profiling 截图, 调优日志
├── 05_System_Integration/   # vLLM/SGLang 源码笔记, 调度策略
└── 06_Paper_Notes/          # 前沿论文摘要与复现计划

---

最终建议 ：这四个算子构成了现代 LLM Serving 的技术底座。不要追求一次性全部精通，而是以"解决问题"为导向螺旋上升。当你遇到实际的性能瓶颈或部署需求时，回到对应的模块深入钻研，这种"按需深挖"比线性学习效率更高。保持动手写代码、保持 Profiling、保持对硬件的敬畏，这三点是通往系统级 AI 工程师的必经之路。

Triton 进阶：Fused Softmax 与数值稳定性实战

**为什么标准 Softmax 慢？**标准的 Softmax 实现是典型的 Memory Bound（访存受限） 算子。对于矩阵的一行数据，朴素实现需要访问 HBM（显存）三次：

1. Pass 1 : 遍历整行寻找 x m a x x_{max} xmax
2. Pass 2 : 遍历整行计算 ∑ e x i − x m a x ∑e^{x_i−x_{max}} ∑exi−xmax
3. Pass 3 : 遍历整行计算 e x i − x m a x / s u m e^{x_i−x_{max}}/sum exi−xmax/sum

优化核心 ：利用 GPU 片上 SRAM（Shared Memory）。将整行数据一次性 Load 到 SRAM，在 SRAM 内部完成 Max、Exp、Sum、Div 所有计算，最终只写回 HBM 一次。访存次数从 3R+1W 降低为 1R+1W。

数值稳定性问题 (Safe Softmax) : 直接计算 e x i e^{xi} exi 极易导致浮点数溢出（Overflow）。例如当 xi=50 时， e 50 ≈ 5.18 × 10 21 e^{50}≈5.18×10^{21} e50≈5.18×1021 ，在 FP16/BF16 下直接变为 Inf 或 NaN。

Safe Softmax 数学等价性推导 , 需要证明减去最大值不改变 Softmax 结果： s o f t m a x ( x ) i = e x i ∑ j e x j = e x i − c ⋅ e c ∑ j ( e x j − c ⋅ e c ) = e x i − c ∑ j e x j − c softmax(x)_i=\frac{e^{x_i}}{∑_je^{x_j}}=\frac{e^{x_i−c}⋅e^c}{∑_j(e^{x_j−c}⋅e^c)}=\frac{e^{x_i−c}}{∑_je^{x_j−c}} softmax(x)i=∑jexjexi=∑j(exj−c⋅ec)exi−c⋅ec=∑jexj−cexi−c 。 其中 c=max⁡(x) 。该变换将指数运算的输入平移到 (−∞,0] 区间，彻底消除上溢风险。

---

算法流程与并行模型

Triton 行级并行策略 ： Triton 采用 1D Grid + Row-wise Parallelism：

• Grid : (M,)，每个 Program Instance 处理矩阵的一行
• Block: 每个 Block 负责一行数据的完整归约
• SRAM: 整行数据驻留片上，归约操作由编译器自动映射到 Warp-level Shuffle 指令

执行流程图
#mermaid-svg-gVYbVvxaLxiN11jp{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-gVYbVvxaLxiN11jp .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-gVYbVvxaLxiN11jp .error-icon{fill:#552222;}#mermaid-svg-gVYbVvxaLxiN11jp .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-gVYbVvxaLxiN11jp .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-gVYbVvxaLxiN11jp .marker{fill:#333333;stroke:#333333;}#mermaid-svg-gVYbVvxaLxiN11jp .marker.cross{stroke:#333333;}#mermaid-svg-gVYbVvxaLxiN11jp svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-gVYbVvxaLxiN11jp p{margin:0;}#mermaid-svg-gVYbVvxaLxiN11jp .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-gVYbVvxaLxiN11jp .cluster-label text{fill:#333;}#mermaid-svg-gVYbVvxaLxiN11jp .cluster-label span{color:#333;}#mermaid-svg-gVYbVvxaLxiN11jp .cluster-label span p{background-color:transparent;}#mermaid-svg-gVYbVvxaLxiN11jp .label text,#mermaid-svg-gVYbVvxaLxiN11jp span{fill:#333;color:#333;}#mermaid-svg-gVYbVvxaLxiN11jp .node rect,#mermaid-svg-gVYbVvxaLxiN11jp .node circle,#mermaid-svg-gVYbVvxaLxiN11jp .node ellipse,#mermaid-svg-gVYbVvxaLxiN11jp .node polygon,#mermaid-svg-gVYbVvxaLxiN11jp .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-gVYbVvxaLxiN11jp .rough-node .label text,#mermaid-svg-gVYbVvxaLxiN11jp .node .label text,#mermaid-svg-gVYbVvxaLxiN11jp .image-shape .label,#mermaid-svg-gVYbVvxaLxiN11jp .icon-shape .label{text-anchor:middle;}#mermaid-svg-gVYbVvxaLxiN11jp .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-gVYbVvxaLxiN11jp .rough-node .label,#mermaid-svg-gVYbVvxaLxiN11jp .node .label,#mermaid-svg-gVYbVvxaLxiN11jp .image-shape .label,#mermaid-svg-gVYbVvxaLxiN11jp .icon-shape .label{text-align:center;}#mermaid-svg-gVYbVvxaLxiN11jp .node.clickable{cursor:pointer;}#mermaid-svg-gVYbVvxaLxiN11jp .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-gVYbVvxaLxiN11jp .arrowheadPath{fill:#333333;}#mermaid-svg-gVYbVvxaLxiN11jp .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-gVYbVvxaLxiN11jp .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-gVYbVvxaLxiN11jp .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-gVYbVvxaLxiN11jp .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-gVYbVvxaLxiN11jp .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-gVYbVvxaLxiN11jp .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-gVYbVvxaLxiN11jp .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-gVYbVvxaLxiN11jp .cluster text{fill:#333;}#mermaid-svg-gVYbVvxaLxiN11jp .cluster span{color:#333;}#mermaid-svg-gVYbVvxaLxiN11jp div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-gVYbVvxaLxiN11jp .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-gVYbVvxaLxiN11jp rect.text{fill:none;stroke-width:0;}#mermaid-svg-gVYbVvxaLxiN11jp .icon-shape,#mermaid-svg-gVYbVvxaLxiN11jp .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-gVYbVvxaLxiN11jp .icon-shape p,#mermaid-svg-gVYbVvxaLxiN11jp .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-gVYbVvxaLxiN11jp .icon-shape .label rect,#mermaid-svg-gVYbVvxaLxiN11jp .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-gVYbVvxaLxiN11jp .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-gVYbVvxaLxiN11jp .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-gVYbVvxaLxiN11jp :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} tl.load + mask
tl.max axis=0
Subtract x_max
tl.exp
tl.sum axis=0
Divide by G
tl.store + mask
HBM: Input Matrix Row
SRAM: Local Buffer
Safe Softmax Compute
x_max Scalar
Safe Row
Numerator
Denominator Scalar
Softmax Output
HBM: Output Matrix Row

---

完整代码实现与深度解析

Kernel 实现

import torch
import triton
import triton.language as tl

@triton.jit
def fused_softmax_kernel(
    output_ptr, input_ptr, 
    input_row_stride, output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr,
):
    # ==================== 1. 行定位与指针计算 ====================
    row_idx = tl.program_id(0)
    row_start_ptr = input_ptr + row_idx * input_row_stride
    
    # ==================== 2. 构造索引与边界掩码 ====================
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = row_start_ptr + col_offsets
    mask = col_offsets < n_cols
    
    # 关键技巧: other=-float('inf') 
    # 越界位置填充 -inf，保证 max 时被忽略，exp(-inf)=0 不影响 sum
    row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
    
    # ==================== 3. Safe Softmax 四步曲 ====================
    # Step 1: 行内归约求最大值 (编译为 warp shuffle)
    row_max = tl.max(row, axis=0)
    
    # Step 2: 平移防溢出
    safe_row = row - row_max
    
    # Step 3: 计算分子 (指数)
    numerator = tl.exp(safe_row)
    
    # Step 4: 归约求分母 & 最终除法
    denominator = tl.sum(numerator, axis=0)
    softmax_output = numerator / denominator
    
    # ==================== 4. 写回 HBM ====================
    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    output_ptrs = output_row_start_ptr + col_offsets
    tl.store(output_ptrs, softmax_output, mask=mask)

Host 端启动配置

def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    M, N = x.shape
    y = torch.empty_like(x)
    
    # BLOCK_SIZE 必须是 2 的幂，且 >= N
    BLOCK_SIZE = triton.next_power_of_2(N)
    
    # 动态调整 num_warps 以匹配计算密度
    num_warps = 4
    if BLOCK_SIZE >= 2048:
        num_warps = 8
    if BLOCK_SIZE >= 4096:
        num_warps = 16
        
    grid = (M,)
    
    fused_softmax_kernel[grid](
        y, x,
        x.stride(0), y.stride(0),
        N,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps
    )
    return y

---

关键技术点对照表

技术点	错误/朴素做法	正确/推荐做法	原因说明
越界填充	other=0.0	other=-float('inf')	0 会影响 max 结果；-inf 在 max 中被忽略，exp(-inf)=0 不影响 sum
归约轴	axis=1 或不指定	axis=0	Load 后是一维向量，只有一个轴
BLOCK_SIZE	任意整数	next_power_of_2(N)	GPU Warp 大小要求对齐，避免硬件浪费
数值安全	直接 tl.exp(row)	先减 row_max 再 exp	防止 FP16/BF16 指数溢出
num_warps	固定值	根据 BLOCK_SIZE 动态调整	大块需更多线程隐藏延迟；小块过多 warp 反而增加调度开销
内存布局	列优先访问	行优先连续访问	合并访存(Coalesced Access)最大化带宽利用率

---

常见陷阱

1. other=0 导致的静默错误
   • 现象：当 N < BLOCK_SIZE 时，Softmax 结果偏小
   • 原因：填充的 0 参与了 max 运算，若真实最大值 < 0，则 row_max=0，平移失效
   • 修复：始终使用 other=-float('inf')
2. BLOCK_SIZE 未对齐
   • 现象：Kernel 启动失败或性能骤降
   • 原因：Triton 要求 BLOCK_SIZE 为 2 的幂
   • 修复：triton.next_power_of_2(N)
3. 大 N 场景 OOM / 寄存器溢出
   • 现象：N > 8192 时编译失败或占用过高
   • 原因：单行全量加载超出 SRAM/寄存器容量
   • 修复：改用 Online Softmax / Tiled Softmax（分块循环归约），本节为教学简化版
4. 精度验证不充分
   • 现象：与 PyTorch 对比误差 > 1e-3
   • 原因：FP16 下累积误差
   • 修复：测试时使用 FP32 基准；生产环境接受 FP16 合理误差范围 (~1e-3)

自查 Checklist

• mask 是否同时用于 tl.load 和 tl.store？
• other 参数是否为 -float('inf')？
• tl.max 和 tl.sum 的 axis 是否为 0？
• BLOCK_SIZE 是否通过 next_power_of_2 计算？
• num_warps 是否随 BLOCK_SIZE 动态调整？
• 是否在 GPU 上与 torch.softmax 做了数值对齐验证？
• 是否进行了 Benchmark 确认加速比 > 1x？

---

本节实现假设 N <= BLOCK_SIZE，适用于 Attention Score 等中等长度场景。对于长序列（如 N=65536），需升级为：
#mermaid-svg-xhFtiYp3YbXQCRtN{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-xhFtiYp3YbXQCRtN .error-icon{fill:#552222;}#mermaid-svg-xhFtiYp3YbXQCRtN .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-xhFtiYp3YbXQCRtN .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-xhFtiYp3YbXQCRtN .marker{fill:#333333;stroke:#333333;}#mermaid-svg-xhFtiYp3YbXQCRtN .marker.cross{stroke:#333333;}#mermaid-svg-xhFtiYp3YbXQCRtN svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-xhFtiYp3YbXQCRtN p{margin:0;}#mermaid-svg-xhFtiYp3YbXQCRtN .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN .cluster-label text{fill:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN .cluster-label span{color:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN .cluster-label span p{background-color:transparent;}#mermaid-svg-xhFtiYp3YbXQCRtN .label text,#mermaid-svg-xhFtiYp3YbXQCRtN span{fill:#333;color:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN .node rect,#mermaid-svg-xhFtiYp3YbXQCRtN .node circle,#mermaid-svg-xhFtiYp3YbXQCRtN .node ellipse,#mermaid-svg-xhFtiYp3YbXQCRtN .node polygon,#mermaid-svg-xhFtiYp3YbXQCRtN .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-xhFtiYp3YbXQCRtN .rough-node .label text,#mermaid-svg-xhFtiYp3YbXQCRtN .node .label text,#mermaid-svg-xhFtiYp3YbXQCRtN .image-shape .label,#mermaid-svg-xhFtiYp3YbXQCRtN .icon-shape .label{text-anchor:middle;}#mermaid-svg-xhFtiYp3YbXQCRtN .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-xhFtiYp3YbXQCRtN .rough-node .label,#mermaid-svg-xhFtiYp3YbXQCRtN .node .label,#mermaid-svg-xhFtiYp3YbXQCRtN .image-shape .label,#mermaid-svg-xhFtiYp3YbXQCRtN .icon-shape .label{text-align:center;}#mermaid-svg-xhFtiYp3YbXQCRtN .node.clickable{cursor:pointer;}#mermaid-svg-xhFtiYp3YbXQCRtN .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-xhFtiYp3YbXQCRtN .arrowheadPath{fill:#333333;}#mermaid-svg-xhFtiYp3YbXQCRtN .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-xhFtiYp3YbXQCRtN .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-xhFtiYp3YbXQCRtN .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xhFtiYp3YbXQCRtN .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-xhFtiYp3YbXQCRtN .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xhFtiYp3YbXQCRtN .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-xhFtiYp3YbXQCRtN .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-xhFtiYp3YbXQCRtN .cluster text{fill:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN .cluster span{color:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-xhFtiYp3YbXQCRtN .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-xhFtiYp3YbXQCRtN rect.text{fill:none;stroke-width:0;}#mermaid-svg-xhFtiYp3YbXQCRtN .icon-shape,#mermaid-svg-xhFtiYp3YbXQCRtN .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-xhFtiYp3YbXQCRtN .icon-shape p,#mermaid-svg-xhFtiYp3YbXQCRtN .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-xhFtiYp3YbXQCRtN .icon-shape .label rect,#mermaid-svg-xhFtiYp3YbXQCRtN .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-xhFtiYp3YbXQCRtN .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-xhFtiYp3YbXQCRtN .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-xhFtiYp3YbXQCRtN :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Online Softmax
分块加载 Tile
局部 max + sum
跨块修正因子
全局归一化
单次写回

Online Softmax 核心思想 ：维护运行中的 (max_so_far, sum_so_far)，每加载一个新 Tile 就更新这两个状态并修正已有累加值，最终只需遍历一次即可完成归约，且 SRAM 占用恒定为 BLOCK_SIZE。

---

学习建议 ：先跑通本节代码并通过测试，再手动修改 other=0 观察错误输出，加深对 Mask 机制的理解。最后尝试将 num_warps 设为固定值进行 Benchmark 对比，体会动态调优的意义。

Triton 设计模式：从算子实现到系统思维

学习目标：跳出"写 Kernel"的战术层面，建立"看数据流、定融合边界、选分块策略"的战略设计能力。本页是连接基础算子与工业级 LLM 优化的认知桥梁。

四大核心设计模式复盘 ： 把算子抽象为四种基本范式。在实际工程中，复杂算子往往是这些范式的组合。

模式	代表算子	核心特征	关键 Triton 原语	典型瓶颈与对策
Element-wise	VecAdd, ReLU	元素独立，无数据依赖	tl.load, tl.store	Memory Bound → 融合相邻算子
Reduction	Softmax, LayerNorm	跨线程/跨块数据聚合	tl.max, tl.sum, mask	归约开销 → SRAM 内归约 + Safe Math
GEMM	MatMul, Linear	二维分块 + K 维累加	tl.dot, Block Ptr	Compute Bound → Tiling + Pipeline
Fusion	Fused Softmax, SwiGLU	中间结果驻留 SRAM	组合上述原语	HBM 往返 → 扩大融合边界

模式识别心法 ： 拿到一个新算子时，不要急着写代码。先问自己：

• 它的数据依赖图是什么样的？（独立 / 行内归约 / 块间累加）
• 它的算术强度（FLOPs/Byte）是高还是低？
• 它在计算图中处于什么位置？前后是否有可融合的 Element-wise 操作？

---

三条黄金设计原则

原则一：先看数据流，再看算子

• 错误思维："我要实现一个 SwiGLU，去找 SwiGLU 的公式。"
• 正确思维："SwiGLU 的数据流是 Read X/W → GEMM → Element-wise Mul → Write。中间张量是否可以不写回 HBM？"
• 本质：算子是数学概念，数据流才是硬件执行的物理现实。

原则二：先决定融合边界，再决定分块

• 融合边界决定了哪些计算共享同一份 SRAM 数据。
• 分块大小必须在融合边界确定后才能优化。如果融合边界选错了，分块再精细也只是在优化一个次优架构。
• 决策顺序 ：分析计算图 → 划定融合范围 → 估算 SRAM 需求 → 确定 Tile Size

原则三：先保证正确性，再谈 Autotune

• 踩坑警告：在功能未验证前做 Autotune 是浪费时间。错误的 Kernel 调参只会得到"快速的错误结果"。
• 自查流程：
  1. 小尺寸 CPU/GPU 数值对齐（FP32 基准）
  2. 边界条件测试（非对齐尺寸、空 Tensor）
  3. 性能 Profiling 确认瓶颈类型
  4. 最后才启动 Autotune

---

设计实战：简化版 SwiGLU

output = (x @ W) * sigmoid(x @ W)

注：这是去掉门控分支的 SwiGLU 简化版，真实 SwiGLU 为 (x @ W_gate) * sigmoid(x @ W_up) @ W_down

设计思路解析（无需完整代码）

Q1: 属于哪类模式？ GEMM + Element-wise Fusion。核心是一个矩阵乘法后紧跟逐元素激活和乘法。

Q2: 融合边界在哪里？

• 融合边界 ：GEMM 的输出 (x @ W) 绝不写回 HBM ，直接在 SRAM 中完成 sigmoid 和 multiply。
• 收益：省去一次 O(M×N) 的 HBM 读写，对于 Memory Bound 的激活函数部分，加速比可达 2-3x。

Q3: Grid 和 Block 设计

• Grid : (M // BLOCK_M, N // BLOCK_N) --- 按输出矩阵的二维分块调度
• Block : (BLOCK_M, BLOCK_N) --- 每个 Block 负责输出矩阵的一个 Tile
• K 维 : 在 Block 内部循环累加，for k in range(0, K, BLOCK_K)

Q4: 大输入维度的分块策略考量

维度	过大时的挑战	应对策略
K 很大	单次 GEMM 累加链过长，寄存器压力	K 维分块循环 + Software Pipelining
N 很大	输出 Tile 占满 SRAM，无法容纳融合算子	缩小 BLOCK_N，或拆分融合阶段
M 很大	Grid 足够大，但单 Block 利用率可能不足	增大 BLOCK_M 提升计算密度
SRAM 紧张	GEMM Tile + 融合算子临时变量超限	分阶段融合：先写回 SRAM 再加载做 Element-wise

关键洞察 ：融合不是免费的。每多融合一个算子，SRAM 中就多一份中间状态。融合的极限由 SRAM 容量决定，而非数学上的可能性。

---

知识图谱：从模式到工业场景

#mermaid-svg-mroVMopRChujnmtS{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-mroVMopRChujnmtS .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-mroVMopRChujnmtS .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-mroVMopRChujnmtS .error-icon{fill:#552222;}#mermaid-svg-mroVMopRChujnmtS .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-mroVMopRChujnmtS .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-mroVMopRChujnmtS .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-mroVMopRChujnmtS .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-mroVMopRChujnmtS .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-mroVMopRChujnmtS .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-mroVMopRChujnmtS .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-mroVMopRChujnmtS .marker{fill:#333333;stroke:#333333;}#mermaid-svg-mroVMopRChujnmtS .marker.cross{stroke:#333333;}#mermaid-svg-mroVMopRChujnmtS svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-mroVMopRChujnmtS p{margin:0;}#mermaid-svg-mroVMopRChujnmtS .edge{stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .section--1 rect,#mermaid-svg-mroVMopRChujnmtS .section--1 path,#mermaid-svg-mroVMopRChujnmtS .section--1 circle,#mermaid-svg-mroVMopRChujnmtS .section--1 polygon,#mermaid-svg-mroVMopRChujnmtS .section--1 path{fill:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section--1 text{fill:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .node-icon--1{font-size:40px;color:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .section-edge--1{stroke:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth--1{stroke-width:17;}#mermaid-svg-mroVMopRChujnmtS .section--1 line{stroke:hsl(60, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-0 rect,#mermaid-svg-mroVMopRChujnmtS .section-0 path,#mermaid-svg-mroVMopRChujnmtS .section-0 circle,#mermaid-svg-mroVMopRChujnmtS .section-0 polygon,#mermaid-svg-mroVMopRChujnmtS .section-0 path{fill:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-mroVMopRChujnmtS .section-0 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-0{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-0{stroke:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-0{stroke-width:14;}#mermaid-svg-mroVMopRChujnmtS .section-0 line{stroke:hsl(240, 100%, 83.5294117647%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-1 rect,#mermaid-svg-mroVMopRChujnmtS .section-1 path,#mermaid-svg-mroVMopRChujnmtS .section-1 circle,#mermaid-svg-mroVMopRChujnmtS .section-1 polygon,#mermaid-svg-mroVMopRChujnmtS .section-1 path{fill:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-1 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-1{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-1{stroke:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-1{stroke-width:11;}#mermaid-svg-mroVMopRChujnmtS .section-1 line{stroke:hsl(260, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-2 rect,#mermaid-svg-mroVMopRChujnmtS .section-2 path,#mermaid-svg-mroVMopRChujnmtS .section-2 circle,#mermaid-svg-mroVMopRChujnmtS .section-2 polygon,#mermaid-svg-mroVMopRChujnmtS .section-2 path{fill:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-2 text{fill:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .node-icon-2{font-size:40px;color:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .section-edge-2{stroke:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-2{stroke-width:8;}#mermaid-svg-mroVMopRChujnmtS .section-2 line{stroke:hsl(90, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-3 rect,#mermaid-svg-mroVMopRChujnmtS .section-3 path,#mermaid-svg-mroVMopRChujnmtS .section-3 circle,#mermaid-svg-mroVMopRChujnmtS .section-3 polygon,#mermaid-svg-mroVMopRChujnmtS .section-3 path{fill:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-3 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-3{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-3{stroke:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-3{stroke-width:5;}#mermaid-svg-mroVMopRChujnmtS .section-3 line{stroke:hsl(120, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-4 rect,#mermaid-svg-mroVMopRChujnmtS .section-4 path,#mermaid-svg-mroVMopRChujnmtS .section-4 circle,#mermaid-svg-mroVMopRChujnmtS .section-4 polygon,#mermaid-svg-mroVMopRChujnmtS .section-4 path{fill:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-4 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-4{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-4{stroke:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-4{stroke-width:2;}#mermaid-svg-mroVMopRChujnmtS .section-4 line{stroke:hsl(150, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-5 rect,#mermaid-svg-mroVMopRChujnmtS .section-5 path,#mermaid-svg-mroVMopRChujnmtS .section-5 circle,#mermaid-svg-mroVMopRChujnmtS .section-5 polygon,#mermaid-svg-mroVMopRChujnmtS .section-5 path{fill:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-5 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-5{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-5{stroke:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-5{stroke-width:-1;}#mermaid-svg-mroVMopRChujnmtS .section-5 line{stroke:hsl(180, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-6 rect,#mermaid-svg-mroVMopRChujnmtS .section-6 path,#mermaid-svg-mroVMopRChujnmtS .section-6 circle,#mermaid-svg-mroVMopRChujnmtS .section-6 polygon,#mermaid-svg-mroVMopRChujnmtS .section-6 path{fill:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-6 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-6{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-6{stroke:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-6{stroke-width:-4;}#mermaid-svg-mroVMopRChujnmtS .section-6 line{stroke:hsl(210, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-7 rect,#mermaid-svg-mroVMopRChujnmtS .section-7 path,#mermaid-svg-mroVMopRChujnmtS .section-7 circle,#mermaid-svg-mroVMopRChujnmtS .section-7 polygon,#mermaid-svg-mroVMopRChujnmtS .section-7 path{fill:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-7 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-7{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-7{stroke:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-7{stroke-width:-7;}#mermaid-svg-mroVMopRChujnmtS .section-7 line{stroke:hsl(270, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-8 rect,#mermaid-svg-mroVMopRChujnmtS .section-8 path,#mermaid-svg-mroVMopRChujnmtS .section-8 circle,#mermaid-svg-mroVMopRChujnmtS .section-8 polygon,#mermaid-svg-mroVMopRChujnmtS .section-8 path{fill:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-8 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-8{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-8{stroke:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-8{stroke-width:-10;}#mermaid-svg-mroVMopRChujnmtS .section-8 line{stroke:hsl(330, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-9 rect,#mermaid-svg-mroVMopRChujnmtS .section-9 path,#mermaid-svg-mroVMopRChujnmtS .section-9 circle,#mermaid-svg-mroVMopRChujnmtS .section-9 polygon,#mermaid-svg-mroVMopRChujnmtS .section-9 path{fill:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-9 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-9{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-9{stroke:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-9{stroke-width:-13;}#mermaid-svg-mroVMopRChujnmtS .section-9 line{stroke:hsl(0, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-10 rect,#mermaid-svg-mroVMopRChujnmtS .section-10 path,#mermaid-svg-mroVMopRChujnmtS .section-10 circle,#mermaid-svg-mroVMopRChujnmtS .section-10 polygon,#mermaid-svg-mroVMopRChujnmtS .section-10 path{fill:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-10 text{fill:black;}#mermaid-svg-mroVMopRChujnmtS .node-icon-10{font-size:40px;color:black;}#mermaid-svg-mroVMopRChujnmtS .section-edge-10{stroke:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .edge-depth-10{stroke-width:-16;}#mermaid-svg-mroVMopRChujnmtS .section-10 line{stroke:hsl(30, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-mroVMopRChujnmtS .disabled,#mermaid-svg-mroVMopRChujnmtS .disabled circle,#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:lightgray;}#mermaid-svg-mroVMopRChujnmtS .disabled text{fill:#efefef;}#mermaid-svg-mroVMopRChujnmtS .section-root rect,#mermaid-svg-mroVMopRChujnmtS .section-root path,#mermaid-svg-mroVMopRChujnmtS .section-root circle,#mermaid-svg-mroVMopRChujnmtS .section-root polygon{fill:hsl(240, 100%, 46.2745098039%);}#mermaid-svg-mroVMopRChujnmtS .section-root text{fill:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .section-root span{color:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .section-2 span{color:#ffffff;}#mermaid-svg-mroVMopRChujnmtS .icon-container{height:100%;display:flex;justify-content:center;align-items:center;}#mermaid-svg-mroVMopRChujnmtS .edge{fill:none;}#mermaid-svg-mroVMopRChujnmtS .mindmap-node-label{dy:1em;alignment-baseline:middle;text-anchor:middle;dominant-baseline:middle;text-align:center;}#mermaid-svg-mroVMopRChujnmtS :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Triton 设计体系
基础模式 01-06
Element-wise
Reduction
GEMM
Fusion
设计原则
数据流优先
融合先于分块
正确性先于调优
工业应用 07-11
Attention 优化
FlashAttention
PagedAttention
推理优化
KV Cache
Speculative Decoding
Quantization

阶段	核心问题	前置模式
进阶 A	如何让单次 Attention 计算更快？	GEMM + Reduction + Fusion
进阶 B	如何让推理服务更省显存、更灵活？	Element-wise + 内存管理

阅读建议：在进入 Attention 优化之前，确保你能不看笔记回答以下三个问题：

1. 给定任意算子表达式，能否画出其数据流图并标出融合边界？
2. Reduction 算子中 mask 的 other 值应如何根据运算语义选择？
3. 当 SRAM 不足以容纳完整融合时，有哪些降级策略？

如果任一问题不确定，请回顾对应章节后再继续。设计模式的价值不在于记住，而在于遇到新问题时能自动激活。

Triton Fused RoPE：融合旋转位置编码实战

为什么 PyTorch 原生 RoPE 慢？ 标准 PyTorch 实现虽然语义清晰，但在推理时存在严重的 Memory Bound 问题：

# PyTorch 原生实现的隐藏开销
x_evens = x[..., 0::2]      #  非连续切片 → 触发 copy 或 stride 访问
x_odds = x[..., 1::2]       #  同上
out_evens = x_evens * cos - x_odds * sin  #  分配临时 Tensor
out_odds = x_evens * sin + x_odds * cos   #  再次分配
y[..., 0::2] = out_evens    #  非连续写入 → 带宽浪费
y[..., 1::2] = out_odds     #  同上

优化核心 ：RoPE 是典型的 Element-wise + 配对依赖 算子。通过 Triton 融合，我们可以在 SRAM 中完成"加载偶奇对 → 旋转计算 → 原地写回"全流程，消除所有中间张量分配和非连续访存。

In-place 修改的战略价值 在大模型推理（尤其是 Prefill 阶段）中，激活值张量可能占据数 GB 显存。In-place RoPE 相比 Out-of-place 版本：

• 节省 50% 激活显存（无需分配输出张量）
• 减少 1 次 HBM 写操作（直接覆盖原地址）
• 这正是 vLLM / TensorRT-LLM 等引擎的标配做法

---

数学原理与内存布局

复数旋转的几何直觉 RoPE 的本质是将相邻特征对 (x2i,x2i+1) 视为复数 z=x2i+j⋅x2i+1 ，乘以旋转因子 ejθ： z ′ = z ⋅ e j θ = ( x 2 i + j ⋅ x 2 i + 1 ) ( c o s ⁡ θ + j s i n ⁡ θ ) z′=z⋅e^{j_θ}=(x_{2i}+j⋅x_{2i+1})(cos⁡θ+jsin⁡θ) z′=z⋅ejθ=(x2i+j⋅x2i+1)(cos⁡θ+jsin⁡θ)

展开实部与虚部即得旋转公式： x 2 i ′ = x 2 i c o s ⁡ θ − x 2 i + 1 s i n ⁡ θ , x 2 i + 1 ′ = x 2 i s i n ⁡ θ + x 2 i + 1 c o s ⁡ θ x_{2i}′=x_{2i}cos⁡θ−x_{2i+1}sin⁡θ,x_{2i+1}′=x_{2i}sin⁡θ+x_{2i+1}cos⁡θ x2i′=x2icos⁡θ−x2i+1sin⁡θ,x2i+1′=x2isin⁡θ+x2i+1cos⁡θ
#mermaid-svg-42AEms3K5SzS0I2z{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-42AEms3K5SzS0I2z .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-42AEms3K5SzS0I2z .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-42AEms3K5SzS0I2z .error-icon{fill:#552222;}#mermaid-svg-42AEms3K5SzS0I2z .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-42AEms3K5SzS0I2z .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-42AEms3K5SzS0I2z .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-42AEms3K5SzS0I2z .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-42AEms3K5SzS0I2z .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-42AEms3K5SzS0I2z .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-42AEms3K5SzS0I2z .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-42AEms3K5SzS0I2z .marker{fill:#333333;stroke:#333333;}#mermaid-svg-42AEms3K5SzS0I2z .marker.cross{stroke:#333333;}#mermaid-svg-42AEms3K5SzS0I2z svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-42AEms3K5SzS0I2z p{margin:0;}#mermaid-svg-42AEms3K5SzS0I2z .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-42AEms3K5SzS0I2z .cluster-label text{fill:#333;}#mermaid-svg-42AEms3K5SzS0I2z .cluster-label span{color:#333;}#mermaid-svg-42AEms3K5SzS0I2z .cluster-label span p{background-color:transparent;}#mermaid-svg-42AEms3K5SzS0I2z .label text,#mermaid-svg-42AEms3K5SzS0I2z span{fill:#333;color:#333;}#mermaid-svg-42AEms3K5SzS0I2z .node rect,#mermaid-svg-42AEms3K5SzS0I2z .node circle,#mermaid-svg-42AEms3K5SzS0I2z .node ellipse,#mermaid-svg-42AEms3K5SzS0I2z .node polygon,#mermaid-svg-42AEms3K5SzS0I2z .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-42AEms3K5SzS0I2z .rough-node .label text,#mermaid-svg-42AEms3K5SzS0I2z .node .label text,#mermaid-svg-42AEms3K5SzS0I2z .image-shape .label,#mermaid-svg-42AEms3K5SzS0I2z .icon-shape .label{text-anchor:middle;}#mermaid-svg-42AEms3K5SzS0I2z .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-42AEms3K5SzS0I2z .rough-node .label,#mermaid-svg-42AEms3K5SzS0I2z .node .label,#mermaid-svg-42AEms3K5SzS0I2z .image-shape .label,#mermaid-svg-42AEms3K5SzS0I2z .icon-shape .label{text-align:center;}#mermaid-svg-42AEms3K5SzS0I2z .node.clickable{cursor:pointer;}#mermaid-svg-42AEms3K5SzS0I2z .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-42AEms3K5SzS0I2z .arrowheadPath{fill:#333333;}#mermaid-svg-42AEms3K5SzS0I2z .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-42AEms3K5SzS0I2z .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-42AEms3K5SzS0I2z .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-42AEms3K5SzS0I2z .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-42AEms3K5SzS0I2z .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-42AEms3K5SzS0I2z .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-42AEms3K5SzS0I2z .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-42AEms3K5SzS0I2z .cluster text{fill:#333;}#mermaid-svg-42AEms3K5SzS0I2z .cluster span{color:#333;}#mermaid-svg-42AEms3K5SzS0I2z div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-42AEms3K5SzS0I2z .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-42AEms3K5SzS0I2z rect.text{fill:none;stroke-width:0;}#mermaid-svg-42AEms3K5SzS0I2z .icon-shape,#mermaid-svg-42AEms3K5SzS0I2z .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-42AEms3K5SzS0I2z .icon-shape p,#mermaid-svg-42AEms3K5SzS0I2z .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-42AEms3K5SzS0I2z .icon-shape .label rect,#mermaid-svg-42AEms3K5SzS0I2z .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-42AEms3K5SzS0I2z .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-42AEms3K5SzS0I2z .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-42AEms3K5SzS0I2z :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 视为复数
× e^jθ
展开
(x_even, x_odd)
z = x_even + j·x_odd
z' = z·(cosθ + j·sinθ)
(x'_even, x'_odd)

关键约束 ：偶数位和奇数位必须成对处理，不能独立看待。这决定了我们的内存访问模式必须是交错的。

Triton 并行策略选择

策略	Grid 维度	每 Program 处理	适用场景	本教程选择
Token-Head 级	(seq_len * n_heads,)	一个 Head 向量 (head_dim)	head_dim ≤ BLOCK_SIZE
Token 级	(seq_len,)	所有 Head	head_dim 很小，Head 数多	并行度不足
Element 级	(total_elements,)	单个元素	无配对依赖	无法成对计算

选择理由 ：LLaMA 系列 head_dim=128，恰好可一次性装入 SRAM（BLOCK_SIZE=128），每个 Program 独立完成一个 Head 的完整旋转，无需跨 Block 通信。

---

Kernel 实现

import torch
import triton
import triton.language as tl

@triton.jit
def fused_rope_kernel(
    t_ptr, cos_ptr, sin_ptr,
    seq_len, n_heads, head_dim,
    BLOCK_SIZE: tl.constexpr,
):
    # ==================== 1. 定位当前 Token-Head ====================
    pid = tl.program_id(0)
    t_offset = pid * head_dim              # 特征起始偏移
    token_idx = pid // n_heads             # 用于索引 cos/sin
    
    # ==================== 2. 构造交错索引 ====================
    half_dim = head_dim // 2
    # 核心技巧：偶奇交错寻址
    evens = tl.arange(0, BLOCK_SIZE // 2) * 2      # [0, 2, 4, ...]
    odds  = evens + 1                               # [1, 3, 5, ...]
    mask  = evens < head_dim                        # 防越界
    
    # ==================== 3. 加载数据到 SRAM ====================
    x_evens = tl.load(t_ptr + t_offset + evens, mask=mask)
    x_odds  = tl.load(t_ptr + t_offset + odds,  mask=mask)
    
    freq_offset = token_idx * half_dim + tl.arange(0, BLOCK_SIZE // 2)
    freq_mask   = tl.arange(0, BLOCK_SIZE // 2) < half_dim
    cos_vals = tl.load(cos_ptr + freq_offset, mask=freq_mask)
    sin_vals = tl.load(sin_ptr + freq_offset, mask=freq_mask)
    
    # ==================== 4. 复数旋转计算 ====================
    # TODO 1: 执行旋转公式
    out_evens = x_evens * cos_vals - x_odds * sin_vals
    out_odds  = x_evens * sin_vals + x_odds * cos_vals
    
    # ==================== 5. In-place 写回 ====================
    # TODO 2: 原地修改，不分配新张量
    tl.store(t_ptr + t_offset + evens, out_evens, mask=mask)
    tl.store(t_ptr + t_offset + odds,  out_odds,  mask=mask)

Host 端封装

def triton_apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    """
    x:   (seq_len, n_heads, head_dim) --- 会被 in-place 修改
    cos: (seq_len, head_dim // 2)
    sin: (seq_len, head_dim // 2)
    """
    seq_len, n_heads, head_dim = x.shape
    
    # 必须保证内存连续，否则 stride 计算失效
    assert x.is_contiguous(), "Input tensor must be contiguous!"
    assert cos.is_contiguous() and sin.is_contiguous()
    
    BLOCK_SIZE = triton.next_power_of_2(head_dim)
    grid = (seq_len * n_heads,)
    
    fused_rope_kernel[grid](
        x, cos, sin,
        seq_len, n_heads, head_dim,
        BLOCK_SIZE=BLOCK_SIZE,
    )
    return x  # 返回同一张量引用，强调 in-place 语义

---

技术点	错误/朴素做法	正确/推荐做法	原因说明
索引构造	分别 arange(0, half) 再手动 ×2	evens = arange * 2; odds = evens + 1	一条指令生成两组交错索引，减少 ALU 开销
cos/sin 形状	(seq_len, head_dim) 全量存储	(seq_len, head_dim//2) 半量存储	偶奇共享同一频率，节省 50% 常量显存
内存连续性	直接使用可能非连续的 view	assert x.is_contiguous()	Kernel 假设连续内存，非连续会导致静默错误
写回方式	分配新 Tensor 再赋值	tl.store 到原指针	In-place 省显存 + 少一次 HBM 写
Mask 粒度	对 evens 和 odds 分别构造 mask	仅对 evens 构造，odds 复用	evens < head_dim ⟺ odds < head_dim+1，当 head_dim 为偶数时等价
Grid 设计	二维 (seq_len, n_heads)	一维 (seq_len * n_heads,)	简化索引计算，避免多维 Grid 调度开销

---

常见陷阱

1. 非连续输入导致静默错误
   • 现象：结果数值完全错乱，但无报错
   • 原因：x.transpose() 或 x[:, ::2] 产生的 view 不连续，Kernel 按连续偏移寻址读到错误数据
   • 修复：Host 端强制 .contiguous() 或加 assert
2. cos/sin 频率索引错位
   • 现象：部分位置编码正确，部分偏移
   • 原因：freq_offset 计算时用了 pid 而非 token_idx = pid // n_heads
   • 修复：牢记 cos/sin 是按 token 位置索引的，与 head 无关
3. head_dim 非 2 的幂
   • 现象：编译失败或越界
   • 原因：BLOCK_SIZE = next_power_of_2(head_dim) 后，half_dim 可能不等于 BLOCK_SIZE // 2
   • 修复：mask 始终基于实际 head_dim，而非 BLOCK_SIZE
4. In-place 与 Autograd 冲突
   • 现象：训练时 backward 报错 "tensor modified inplace"
   • 原因：PyTorch 自动微分不允许对需要梯度的张量做 in-place 修改
   • 修复：训练时用 out-of-place 版本；推理时才用 in-place

自查 Checklist

• 输入张量是否已确认 contiguous？
• token_idx 是否正确通过 pid // n_heads 计算？
• evens 和 odds 的 mask 是否覆盖了所有有效位置？
• cos/sin 的形状是否为 (seq_len, head_dim // 2)？
• 旋转公式的符号是否正确？（even: cos - sin, odd: sin + cos）
• 是否与 PyTorch 参考实现在 FP32 下做了数值对齐？
• Benchmark 是否体现了 in-place 的显存优势？

---

#mermaid-svg-zSLC7EhKhq1XUUWR{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-zSLC7EhKhq1XUUWR .error-icon{fill:#552222;}#mermaid-svg-zSLC7EhKhq1XUUWR .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-zSLC7EhKhq1XUUWR .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-zSLC7EhKhq1XUUWR .marker{fill:#333333;stroke:#333333;}#mermaid-svg-zSLC7EhKhq1XUUWR .marker.cross{stroke:#333333;}#mermaid-svg-zSLC7EhKhq1XUUWR svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-zSLC7EhKhq1XUUWR p{margin:0;}#mermaid-svg-zSLC7EhKhq1XUUWR .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR .cluster-label text{fill:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR .cluster-label span{color:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR .cluster-label span p{background-color:transparent;}#mermaid-svg-zSLC7EhKhq1XUUWR .label text,#mermaid-svg-zSLC7EhKhq1XUUWR span{fill:#333;color:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR .node rect,#mermaid-svg-zSLC7EhKhq1XUUWR .node circle,#mermaid-svg-zSLC7EhKhq1XUUWR .node ellipse,#mermaid-svg-zSLC7EhKhq1XUUWR .node polygon,#mermaid-svg-zSLC7EhKhq1XUUWR .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-zSLC7EhKhq1XUUWR .rough-node .label text,#mermaid-svg-zSLC7EhKhq1XUUWR .node .label text,#mermaid-svg-zSLC7EhKhq1XUUWR .image-shape .label,#mermaid-svg-zSLC7EhKhq1XUUWR .icon-shape .label{text-anchor:middle;}#mermaid-svg-zSLC7EhKhq1XUUWR .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-zSLC7EhKhq1XUUWR .rough-node .label,#mermaid-svg-zSLC7EhKhq1XUUWR .node .label,#mermaid-svg-zSLC7EhKhq1XUUWR .image-shape .label,#mermaid-svg-zSLC7EhKhq1XUUWR .icon-shape .label{text-align:center;}#mermaid-svg-zSLC7EhKhq1XUUWR .node.clickable{cursor:pointer;}#mermaid-svg-zSLC7EhKhq1XUUWR .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-zSLC7EhKhq1XUUWR .arrowheadPath{fill:#333333;}#mermaid-svg-zSLC7EhKhq1XUUWR .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-zSLC7EhKhq1XUUWR .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-zSLC7EhKhq1XUUWR .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-zSLC7EhKhq1XUUWR .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-zSLC7EhKhq1XUUWR .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-zSLC7EhKhq1XUUWR .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-zSLC7EhKhq1XUUWR .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-zSLC7EhKhq1XUUWR .cluster text{fill:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR .cluster span{color:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-zSLC7EhKhq1XUUWR .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-zSLC7EhKhq1XUUWR rect.text{fill:none;stroke-width:0;}#mermaid-svg-zSLC7EhKhq1XUUWR .icon-shape,#mermaid-svg-zSLC7EhKhq1XUUWR .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-zSLC7EhKhq1XUUWR .icon-shape p,#mermaid-svg-zSLC7EhKhq1XUUWR .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-zSLC7EhKhq1XUUWR .icon-shape .label rect,#mermaid-svg-zSLC7EhKhq1XUUWR .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-zSLC7EhKhq1XUUWR .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-zSLC7EhKhq1XUUWR .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-zSLC7EhKhq1XUUWR :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Batch 维度
GQA/MQA
变长序列
FP8 推理
Prefill+Decode 混合
基础 Fused RoPE
生产环境需求
展平为 total_tokens 统一调度
cos/sin 按 KV Head 数索引
传入 cu_seqlens 替代固定 seq_len
旋转计算保持 FP16/BF16 精度
动态 Grid 适配不同 batch 组成

理论加速比分析

指标	PyTorch 原生	Triton Fused	提升来源
HBM 读次数	3 (x + cos + sin)	3 (相同)	---
HBM 写次数	2 (非连续)	1 (连续 in-place)	50% 写带宽
临时分配	4 个中间 Tensor	0	显存 + GC 开销
Kernel Launch	5+ 次	1 次	启动开销
预期加速	1.0x	2-3x	综合以上因素

1. 先跑通本节代码，验证数值正确性
2. 故意传入非连续张量，观察错误行为，加深理解
3. 将 head_dim 改为 64/256，测试 BLOCK_SIZE 自适应能力
4. 对比 in-place vs out-of-place 版本的显存占用（torch.cuda.memory_allocated()）
5. 思考：如果要将 RoPE 进一步融合进 Attention Kernel，数据流应如何调整？

Triton Flash Attention：SRAM 内的在线注意力计算

标准 Attention 的致命瓶颈 标准 Attention 需要物化 N×N 的注意力矩阵 P=softmax(QK^T) ：

• 显存复杂度: O(N2)，序列长度翻倍，显存消耗翻四倍
• 访存瓶颈: P 矩阵必须写回 HBM 再读回做 PV 乘法，带宽浪费严重

Flash Attention 的核心洞察

关键思想 ：我们不需要完整的 P 矩阵，只需要最终的 O=PV。通过 Tiling（分块） + Online Softmax（在线归约） ，可以在 SRAM 中增量式地计算输出，永远不将 N×N 矩阵写入 HBM。

指标	标准 Attention	Flash Attention	提升来源
显存复杂度	O(N2)	O(N)	不物化 P 矩阵
HBM 访问	O(N2)	O(N2d/M)	Tiling + SRAM 复用
计算复杂度	O(N2d)	O(N2d)	相同，但常数更小
实际加速	1.0x	2-4x	减少 HBM 往返

---

Online Softmax 数学推导（核心难点）

Flash Attention 的灵魂是 Online Safe Softmax。当新块到来导致全局最大值更新时，如何修正历史累加值？

状态变量定义

对于 Q 的每一行，维护三个状态量：

符号	含义	初始值	更新规则
mi	当前已知的全局最大值	−∞	max⁡(mold,mj)
li	修正后的指数和（分母）	0	l o l d ⋅ α + ∑ e s j − m n e w l_{old}⋅α+∑e^{s_j−m_{new}} lold⋅α+∑esj−mnew
acc	修正后的加权输出累加器	0	a c c o l d ⋅ α + P j V j acc_{old}⋅α+P_jV_j accold⋅α+PjVj

修正系数 α的推导 当处理第 j 个 KV 块时，若 mnew>mold ，之前按 mold 计算的指数被高估了。修正系数为：α=emold−mnew

为什么这个修正有效？ 举例说明：假设已处理两块数据，真实最大值为 m3：

• 第1块按 m1 计算：贡献为 es1−m1
• 第2块到来， m2=max⁡(m1,mblock2)，第1块修正为 es1−m1⋅em1−m2=es1−m2
• 第3块到来， m3>m2，前两块再乘 em2−m3，总修正为 e s 1 − m 1 ⋅ e m 1 − m 2 ⋅ e m 2 − m 3 = e s 1 − m 3 e^{s1−m1}⋅e^{m1−m2}⋅e^{m2−m3}=e^{s1−m3} es1−m1⋅em1−m2⋅em2−m3=es1−m3

归纳结论：无论最大值更新多少次，累积修正系数的连乘始终等于 emfirst_used−mfinal，数学上严格等价于一次性全局 Softmax。

完整状态更新流程
#mermaid-svg-jDrlc4EY8VfilFAj{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-jDrlc4EY8VfilFAj .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-jDrlc4EY8VfilFAj .error-icon{fill:#552222;}#mermaid-svg-jDrlc4EY8VfilFAj .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-jDrlc4EY8VfilFAj .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-jDrlc4EY8VfilFAj .marker{fill:#333333;stroke:#333333;}#mermaid-svg-jDrlc4EY8VfilFAj .marker.cross{stroke:#333333;}#mermaid-svg-jDrlc4EY8VfilFAj svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-jDrlc4EY8VfilFAj p{margin:0;}#mermaid-svg-jDrlc4EY8VfilFAj .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-jDrlc4EY8VfilFAj .cluster-label text{fill:#333;}#mermaid-svg-jDrlc4EY8VfilFAj .cluster-label span{color:#333;}#mermaid-svg-jDrlc4EY8VfilFAj .cluster-label span p{background-color:transparent;}#mermaid-svg-jDrlc4EY8VfilFAj .label text,#mermaid-svg-jDrlc4EY8VfilFAj span{fill:#333;color:#333;}#mermaid-svg-jDrlc4EY8VfilFAj .node rect,#mermaid-svg-jDrlc4EY8VfilFAj .node circle,#mermaid-svg-jDrlc4EY8VfilFAj .node ellipse,#mermaid-svg-jDrlc4EY8VfilFAj .node polygon,#mermaid-svg-jDrlc4EY8VfilFAj .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-jDrlc4EY8VfilFAj .rough-node .label text,#mermaid-svg-jDrlc4EY8VfilFAj .node .label text,#mermaid-svg-jDrlc4EY8VfilFAj .image-shape .label,#mermaid-svg-jDrlc4EY8VfilFAj .icon-shape .label{text-anchor:middle;}#mermaid-svg-jDrlc4EY8VfilFAj .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-jDrlc4EY8VfilFAj .rough-node .label,#mermaid-svg-jDrlc4EY8VfilFAj .node .label,#mermaid-svg-jDrlc4EY8VfilFAj .image-shape .label,#mermaid-svg-jDrlc4EY8VfilFAj .icon-shape .label{text-align:center;}#mermaid-svg-jDrlc4EY8VfilFAj .node.clickable{cursor:pointer;}#mermaid-svg-jDrlc4EY8VfilFAj .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-jDrlc4EY8VfilFAj .arrowheadPath{fill:#333333;}#mermaid-svg-jDrlc4EY8VfilFAj .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-jDrlc4EY8VfilFAj .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-jDrlc4EY8VfilFAj .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-jDrlc4EY8VfilFAj .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-jDrlc4EY8VfilFAj .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-jDrlc4EY8VfilFAj .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-jDrlc4EY8VfilFAj .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-jDrlc4EY8VfilFAj .cluster text{fill:#333;}#mermaid-svg-jDrlc4EY8VfilFAj .cluster span{color:#333;}#mermaid-svg-jDrlc4EY8VfilFAj div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-jDrlc4EY8VfilFAj .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-jDrlc4EY8VfilFAj rect.text{fill:none;stroke-width:0;}#mermaid-svg-jDrlc4EY8VfilFAj .icon-shape,#mermaid-svg-jDrlc4EY8VfilFAj .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-jDrlc4EY8VfilFAj .icon-shape p,#mermaid-svg-jDrlc4EY8VfilFAj .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-jDrlc4EY8VfilFAj .icon-shape .label rect,#mermaid-svg-jDrlc4EY8VfilFAj .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-jDrlc4EY8VfilFAj .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-jDrlc4EY8VfilFAj .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-jDrlc4EY8VfilFAj :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 是
否
加载 KV Block j
S_j = Q_block × K_j^T × scale
m_j = rowmax(S_j)
m_new = max(m_i, m_j)
α = exp(m_i - m_new)
P_j = exp(S_j - m_new)
l_new = l_i × α + rowsum(P_j)
acc = acc × α + P_j × V_j
更新 m_i=l_i=acc
还有 KV 块?
O = acc / l_i

---

Kernel 实现

import torch
import triton
import triton.language as tl

@triton.jit
def flash_attn_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, sm_scale,
    Out_ptr,
    seqlen_q, seqlen_k, head_dim,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_DMODEL: tl.constexpr,
):
    # ==================== 1. Q Block 定位与加载 ====================
    start_m = tl.program_id(0)
    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d = tl.arange(0, BLOCK_DMODEL)
    
    q_ptrs = Q_ptr + (offs_m[:, None] * BLOCK_DMODEL + offs_d[None, :])
    q = tl.load(q_ptrs)  # (BLOCK_M, BLOCK_DMODEL)
    
    # ==================== 2. 初始化 Online Softmax 状态 ====================
    acc = tl.zeros((BLOCK_M, BLOCK_DMODEL), dtype=tl.float32)
    m_i = tl.zeros((BLOCK_M,), dtype=tl.float32) - float('inf')  # ⚠️ 必须 -inf
    l_i = tl.zeros((BLOCK_M,), dtype=tl.float32)
    
    # ==================== 3. KV Block 循环 ====================
    num_n_blocks = tl.cdiv(seqlen_k, BLOCK_N)
    
    for start_n in range(0, num_n_blocks):
        offs_n = start_n * BLOCK_N + tl.arange(0, BLOCK_N)
        
        k_ptrs = K_ptr + (offs_n[:, None] * BLOCK_DMODEL + offs_d[None, :])
        v_ptrs = V_ptr + (offs_n[:, None] * BLOCK_DMODEL + offs_d[None, :])
        k = tl.load(k_ptrs)  # (BLOCK_N, BLOCK_DMODEL)
        v = tl.load(v_ptrs)  # (BLOCK_N, BLOCK_DMODEL)
        
        # TODO 1: 注意力分数 S = Q @ K^T * scale
        qk = tl.dot(q, tl.trans(k))  # (BLOCK_M, BLOCK_N)
        qk *= sm_scale
        
        # TODO 2: Online Softmax 状态更新
        m_block = tl.max(qk, axis=1)           # 当前块行最大值
        m_new = tl.maximum(m_i, m_block)       # 全局最大值更新
        
        # TODO 3: 修正系数 + 指数 + 分母更新
        p = tl.exp(qk - m_new[:, None])        # 安全的 softmax 分子
        alpha = tl.exp(m_i - m_new)            # 历史修正系数
        l_new = l_i * alpha + tl.sum(p, axis=1) # 更新分母
        
        # TODO 4: 修正历史输出 + 累加新贡献
        acc = acc * alpha[:, None]             # 修正历史累加器
        acc += tl.dot(p.to(v.dtype), v)        # FP32→FP16 平衡精度与带宽
        
        # 状态传递到下一轮
        m_i = m_new
        l_i = l_new
    
    # ==================== 4. 最终归一化并写回 ====================
    acc = acc / l_i[:, None]
    
    out_ptrs = Out_ptr + (offs_m[:, None] * BLOCK_DMODEL + offs_d[None, :])
    tl.store(out_ptrs, acc.to(Out_ptr.dtype.element_ty))

Host 端封装

def triton_flash_attention(q, k, v, sm_scale):
    """
    2D per-head 简化版。4D 输入请先在 wrapper 中 flatten batch+head 维度。
    q/k/v: (seqlen, head_dim), dtype=float16/bfloat16
    """
    seqlen_q, head_dim = q.shape
    seqlen_k, _ = k.shape
    out = torch.empty_like(q)
    
    BLOCK_M = 64
    BLOCK_N = 64
    BLOCK_DMODEL = triton.next_power_of_2(head_dim)
    
    grid = (triton.cdiv(seqlen_q, BLOCK_M),)
    
    flash_attn_fwd_kernel[grid](
        q, k, v, sm_scale, out,
        seqlen_q, seqlen_k, head_dim,
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_DMODEL=BLOCK_DMODEL,
    )
    return out

---

关键技术点对照表

技术点	错误/朴素做法	正确/推荐做法	原因说明
m_i 初始值	0.0	-float('inf')	若所有 qk < 0，max(0, qk)=0 导致修正错误
alpha 广播	acc * alpha	acc * alpha[:, None]	alpha 是 (BLOCK_M,)，acc 是 (BLOCK_M, D)，需显式扩展
P×V 精度	全程 FP32	p.to(v.dtype) 后 dot	FP32 softmax → FP16 matmul，兼顾数值稳定与 Tensor Core 吞吐
sm_scale 时机	softmax 后再乘	dot 后立即乘	防止 qk 值过大导致 exp 溢出，scale 应在 exp 之前
BLOCK_DMODEL	直接用 head_dim	next_power_of_2(head_dim)	tl.dot 要求维度对齐；多余位置用 mask 或零填充
l_i 除零保护	直接 / l_i	确认 l_i > 0	理论上 l_i ≥ 1（至少一个 exp(0)=1），但极端 NaN 输入可能为 0

---

常见陷阱

1. m_i 初始化为 0 导致的静默错误
   • 现象：短序列结果正确，长序列或特定分布下误差飙升
   • 原因：当某行所有 qk 值均为负数时，max(0, negative) = 0，修正系数 α=e0−mnewα =e 0−mne**w 不等于正确的 e−∞−mnew=0e −∞−mne**w=0
   • 修复：始终初始化为 -float('inf')
2. alpha 广播维度错误
   • 现象：编译报错或结果全错
   • 原因：alpha shape=(BLOCK_M,)，acc shape=(BLOCK_M, D)，Triton 不会自动广播
   • 修复：alpha[:, None] 显式扩展为列向量
3. P 矩阵未转换精度直接 matmul
   • 现象：性能下降 30-50%
   • 原因：FP32 × FP16 的 dot 无法充分利用 Tensor Core
   • 修复：tl.dot(p.to(v.dtype), v)，softmax 保持 FP32 精度，matmul 降级到 V 的精度
4. 非对齐 head_dim 导致越界
   • 现象：CUDA error 或 NaN
   • 原因：head_dim=96 时 BLOCK_DMODEL=128，超出部分读到垃圾数据
   • 修复：生产版本需在 load/store 处添加 mask；本节教学版假设可整除
5. KV 循环顺序错误
   • 现象：Causal Mask 场景下结果错误
   • 原因：Causal 模式下 KV 块的遍历范围依赖 Q 块位置
   • 修复：Causal 版本内层循环应为 range(0, min(num_n_blocks, cdiv((start_m+1)*BLOCK_M, BLOCK_N)))

自查 Checklist

• m_i 是否初始化为 -float('inf')？
• alpha 是否使用了 [:, None] 广播？
• sm_scale 是否在 tl.exp 之前应用？
• p.to(v.dtype) 是否在 tl.dot 之前执行？
• BLOCK_DMODEL 是否为 2 的幂？
• 最终 acc / l_i[:, None] 是否有广播？
• 是否与 PyTorch 参考实现在 FP32 基准下做了数值对齐（误差 < 1e-3 for FP16）？
• Benchmark 是否体现了相对标准实现的加速？

---

本节实现是理解 Flash Attention 的最小完整版本。工业级实现还需解决：
#mermaid-svg-FtIg2i9gy7RexILQ{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-FtIg2i9gy7RexILQ .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-FtIg2i9gy7RexILQ .error-icon{fill:#552222;}#mermaid-svg-FtIg2i9gy7RexILQ .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-FtIg2i9gy7RexILQ .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-FtIg2i9gy7RexILQ .marker{fill:#333333;stroke:#333333;}#mermaid-svg-FtIg2i9gy7RexILQ .marker.cross{stroke:#333333;}#mermaid-svg-FtIg2i9gy7RexILQ svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-FtIg2i9gy7RexILQ p{margin:0;}#mermaid-svg-FtIg2i9gy7RexILQ .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-FtIg2i9gy7RexILQ .cluster-label text{fill:#333;}#mermaid-svg-FtIg2i9gy7RexILQ .cluster-label span{color:#333;}#mermaid-svg-FtIg2i9gy7RexILQ .cluster-label span p{background-color:transparent;}#mermaid-svg-FtIg2i9gy7RexILQ .label text,#mermaid-svg-FtIg2i9gy7RexILQ span{fill:#333;color:#333;}#mermaid-svg-FtIg2i9gy7RexILQ .node rect,#mermaid-svg-FtIg2i9gy7RexILQ .node circle,#mermaid-svg-FtIg2i9gy7RexILQ .node ellipse,#mermaid-svg-FtIg2i9gy7RexILQ .node polygon,#mermaid-svg-FtIg2i9gy7RexILQ .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-FtIg2i9gy7RexILQ .rough-node .label text,#mermaid-svg-FtIg2i9gy7RexILQ .node .label text,#mermaid-svg-FtIg2i9gy7RexILQ .image-shape .label,#mermaid-svg-FtIg2i9gy7RexILQ .icon-shape .label{text-anchor:middle;}#mermaid-svg-FtIg2i9gy7RexILQ .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-FtIg2i9gy7RexILQ .rough-node .label,#mermaid-svg-FtIg2i9gy7RexILQ .node .label,#mermaid-svg-FtIg2i9gy7RexILQ .image-shape .label,#mermaid-svg-FtIg2i9gy7RexILQ .icon-shape .label{text-align:center;}#mermaid-svg-FtIg2i9gy7RexILQ .node.clickable{cursor:pointer;}#mermaid-svg-FtIg2i9gy7RexILQ .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-FtIg2i9gy7RexILQ .arrowheadPath{fill:#333333;}#mermaid-svg-FtIg2i9gy7RexILQ .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-FtIg2i9gy7RexILQ .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-FtIg2i9gy7RexILQ .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-FtIg2i9gy7RexILQ .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-FtIg2i9gy7RexILQ .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-FtIg2i9gy7RexILQ .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-FtIg2i9gy7RexILQ .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-FtIg2i9gy7RexILQ .cluster text{fill:#333;}#mermaid-svg-FtIg2i9gy7RexILQ .cluster span{color:#333;}#mermaid-svg-FtIg2i9gy7RexILQ div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-FtIg2i9gy7RexILQ .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-FtIg2i9gy7RexILQ rect.text{fill:none;stroke-width:0;}#mermaid-svg-FtIg2i9gy7RexILQ .icon-shape,#mermaid-svg-FtIg2i9gy7RexILQ .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-FtIg2i9gy7RexILQ .icon-shape p,#mermaid-svg-FtIg2i9gy7RexILQ .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-FtIg2i9gy7RexILQ .icon-shape .label rect,#mermaid-svg-FtIg2i9gy7RexILQ .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-FtIg2i9gy7RexILQ .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-FtIg2i9gy7RexILQ .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-FtIg2i9gy7RexILQ :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Batch + Multi-Head
Causal Mask
变长序列
GQA/MQA
反向传播
性能极致
FP8 推理
教学版 2D Per-Head
生产需求
4D Wrapper: flatten (B,H) → 2D Grid
内层循环上界裁剪 + mask 填充 -inf
cu_seqlens 索引 + 动态 Block 范围
KV Head 映射: kv_head_idx = head_idx // group_size
保存 LSE (log-sum-exp) 用于 dQ/dK/dV 重计算
Warp Specialization + Software Pipelining
Dequant on-the-fly + FP32 Accumulator

1. 先跑通本节代码，验证数值正确性
2. 手动追踪一个小例子（BLOCK_M=2, BLOCK_N=2, 3个KV块），手算 m/l/acc 的更新过程，与代码输出对照
3. 将 m_i 改为 0.0，观察哪些 case 会出错，加深对 Online Softmax 的理解
4. 尝试添加 Causal Mask：在内层循环中对 qk 应用 tl.where(mask, qk, -inf)
5. 思考：为什么 Flash Attention 的反向传播可以不存储 N×NN ×N 的 P 矩阵？（提示：LSE + 重计算）

Triton PagedAttention：KV Cache 间接寻址与解码优化

为什么需要 PagedAttention？ 在传统 LLM 推理中，KV Cache 必须预分配连续显存。这导致两个致命问题：

• 内部碎片：为最大可能长度预分配，实际使用率往往 < 40%
• 外部碎片：不同长度的请求释放后留下无法复用的空洞
• 并发瓶颈：显存浪费直接限制了 Batch Size，吞吐量上不去

PagedAttention 的核心洞察

操作系统虚拟内存思想迁移 ：将 KV Cache 切分为固定大小的物理块（Physical Block） ，通过块映射表（Block Table） 实现逻辑序列到物理内存的间接寻址。逻辑上连续的 Token 在物理上可以分散存储，彻底消除碎片。

指标	传统连续 KV Cache	PagedAttention	提升来源
显存利用率	40-60%	> 95%	消除内/外部碎片
并发承载	受限于最大序列长度	受限于总 Token 数	按需分配物理块
内存分配	每请求一次大分配	全局池 + 小块分配	减少 malloc 开销
计算内核	连续访存 FlashAttn	间接寻址 + Online Softmax	适配碎片化布局

Decoding vs Prefill 的计算差异

本节聚焦 Decoding 阶段，其计算特性与 Prefill 截然不同：

特性	Prefill (Flash Attention)	Decoding (PagedAttention)
Q 形状	(BLOCK_M, head_dim) 矩阵	(head_dim,) 单向量
KV 遍历	Tiling 分块循环	按物理块逐块循环
内存访问	连续或规则分块	间接寻址（非连续）
计算瓶颈	Compute Bound (GEMM)	Memory Bound (带宽)
Grid 设计	(num_q_blocks,) 1D	(batch, num_heads) 2D

---

物理存储模型与间接寻址

双层地址空间
#mermaid-svg-PXjt6IIOuCewV8UB{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-PXjt6IIOuCewV8UB .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-PXjt6IIOuCewV8UB .error-icon{fill:#552222;}#mermaid-svg-PXjt6IIOuCewV8UB .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-PXjt6IIOuCewV8UB .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-PXjt6IIOuCewV8UB .marker{fill:#333333;stroke:#333333;}#mermaid-svg-PXjt6IIOuCewV8UB .marker.cross{stroke:#333333;}#mermaid-svg-PXjt6IIOuCewV8UB svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-PXjt6IIOuCewV8UB p{margin:0;}#mermaid-svg-PXjt6IIOuCewV8UB .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-PXjt6IIOuCewV8UB .cluster-label text{fill:#333;}#mermaid-svg-PXjt6IIOuCewV8UB .cluster-label span{color:#333;}#mermaid-svg-PXjt6IIOuCewV8UB .cluster-label span p{background-color:transparent;}#mermaid-svg-PXjt6IIOuCewV8UB .label text,#mermaid-svg-PXjt6IIOuCewV8UB span{fill:#333;color:#333;}#mermaid-svg-PXjt6IIOuCewV8UB .node rect,#mermaid-svg-PXjt6IIOuCewV8UB .node circle,#mermaid-svg-PXjt6IIOuCewV8UB .node ellipse,#mermaid-svg-PXjt6IIOuCewV8UB .node polygon,#mermaid-svg-PXjt6IIOuCewV8UB .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-PXjt6IIOuCewV8UB .rough-node .label text,#mermaid-svg-PXjt6IIOuCewV8UB .node .label text,#mermaid-svg-PXjt6IIOuCewV8UB .image-shape .label,#mermaid-svg-PXjt6IIOuCewV8UB .icon-shape .label{text-anchor:middle;}#mermaid-svg-PXjt6IIOuCewV8UB .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-PXjt6IIOuCewV8UB .rough-node .label,#mermaid-svg-PXjt6IIOuCewV8UB .node .label,#mermaid-svg-PXjt6IIOuCewV8UB .image-shape .label,#mermaid-svg-PXjt6IIOuCewV8UB .icon-shape .label{text-align:center;}#mermaid-svg-PXjt6IIOuCewV8UB .node.clickable{cursor:pointer;}#mermaid-svg-PXjt6IIOuCewV8UB .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-PXjt6IIOuCewV8UB .arrowheadPath{fill:#333333;}#mermaid-svg-PXjt6IIOuCewV8UB .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-PXjt6IIOuCewV8UB .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-PXjt6IIOuCewV8UB .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-PXjt6IIOuCewV8UB .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-PXjt6IIOuCewV8UB .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-PXjt6IIOuCewV8UB .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-PXjt6IIOuCewV8UB .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-PXjt6IIOuCewV8UB .cluster text{fill:#333;}#mermaid-svg-PXjt6IIOuCewV8UB .cluster span{color:#333;}#mermaid-svg-PXjt6IIOuCewV8UB div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-PXjt6IIOuCewV8UB .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-PXjt6IIOuCewV8UB rect.text{fill:none;stroke-width:0;}#mermaid-svg-PXjt6IIOuCewV8UB .icon-shape,#mermaid-svg-PXjt6IIOuCewV8UB .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-PXjt6IIOuCewV8UB .icon-shape p,#mermaid-svg-PXjt6IIOuCewV8UB .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-PXjt6IIOuCewV8UB .icon-shape .label rect,#mermaid-svg-PXjt6IIOuCewV8UB .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-PXjt6IIOuCewV8UB .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-PXjt6IIOuCewV8UB .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-PXjt6IIOuCewV8UB :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 物理内存池 (Global)
Block Table
逻辑空间 (Per Sequence)
Token 0-15

Logical Block 0
Token 16-31

Logical Block 1
Token 32-47

Logical Block 2
btb,0 = 5
btb,1 = 12
btb,2 = 3
Block 3
Block 5
Block 12
...其他块...

物理池内存布局

KV Cache 物理池是一个四维连续张量：

k_cache / v_cache: [num_blocks, block_size, num_heads, head_dim]
                      ↑            ↑           ↑          ↑
                  stride(0)    stride(1)   stride(2)  stride(3)

给定 physical_block_idx、token_offset、head_idx、dim_idx，物理地址为：

addr=base+p⋅s0+t⋅s1+h⋅s2+d⋅s3addr=base+p ⋅s 0+t ⋅s 1+h ⋅s 2+d ⋅s3

关键区别 ：与 Flash Attention 直接用指针偏移不同，PagedAttention 的 pp 来自 运行时查表，编译器无法静态优化，这就是"间接寻址"的含义。

---

Kernel 实现

import torch
import triton
import triton.language as tl
import math

@triton.jit
def paged_attention_decoding_kernel(
    out_ptr, q_ptr, k_cache_ptr, v_cache_ptr,
    block_tables_ptr, context_lens_ptr,
    sm_scale,
    stride_k_block, stride_k_seq, stride_k_head, stride_k_dim,
    stride_bt_batch, stride_bt_block,
    BLOCK_SIZE: tl.constexpr, HEAD_DIM: tl.constexpr,
):
    # ==================== 1. 定位当前 Batch + Head ====================
    batch_idx = tl.program_id(0)
    head_idx = tl.program_id(1)
    
    # ==================== 2. 获取上下文长度 & 加载 Query ====================
    context_len = tl.load(context_lens_ptr + batch_idx)
    num_logical_blocks = tl.cdiv(context_len, BLOCK_SIZE)
    
    # Decoding: Q 只有 1 个 token，shape = (HEAD_DIM,)
    q_offset = (batch_idx * tl.num_programs(1) + head_idx) * HEAD_DIM \
               + tl.arange(0, HEAD_DIM)
    q = tl.load(q_ptr + q_offset)  # (HEAD_DIM,)
    
    # ==================== 3. 初始化 Online Softmax 状态 ====================
    m_i = -float('inf')       # ⚠️ 必须 -inf，不能是 0
    l_i = 0.0
    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
    
    # ==================== 4. 逐物理块循环 ====================
    for logical_block_idx in range(num_logical_blocks):
        
        # TODO 1: 查表获取物理块索引 (Indirect Memory Access)
        bt_offset = batch_idx * stride_bt_batch \
                    + logical_block_idx * stride_bt_block
        physical_block_idx = tl.load(block_tables_ptr + bt_offset)
        
        # 构造当前块内的 token mask
        start_token_idx = logical_block_idx * BLOCK_SIZE
        token_offsets = tl.arange(0, BLOCK_SIZE)
        physical_token_idx = start_token_idx + token_offsets
        mask = physical_token_idx < context_len
        
        # TODO 2: 计算 KV 物理地址并加载
        # k_cache shape: [num_blocks, block_size, num_heads, head_dim]
        k_offset = (physical_block_idx * stride_k_block
                    + token_offsets[:, None] * stride_k_seq
                    + head_idx * stride_k_head
                    + tl.arange(0, HEAD_DIM)[None, :])
        
        k = tl.load(k_cache_ptr + k_offset, mask=mask[:, None], other=0.0)
        v = tl.load(v_cache_ptr + k_offset, mask=mask[:, None], other=0.0)
        
        # TODO 3: 注意力分数 + Mask
        # Decoding: q 是向量，用 element-wise mul + sum 代替 matmul
        qk = tl.sum(q[None, :] * k, axis=1) * sm_scale  # (BLOCK_SIZE,)
        qk = tl.where(mask, qk, -float('inf'))           #无效位置设 -inf
        
        # TODO 4: Online Softmax 状态更新
        m_block = tl.max(qk, axis=0)
        m_new = tl.maximum(m_i, m_block)
        alpha = tl.exp(m_i - m_new)                      # 历史修正系数
        p = tl.exp(qk - m_new)                           # 当前块 softmax 分子
        l_new = l_i * alpha + tl.sum(p, axis=0)          # 更新分母
        
        # TODO 5: 修正历史输出 + 累加新贡献
        acc = acc * alpha + tl.sum(p[:, None] * v, axis=0)
        
        # 状态传递
        m_i = m_new
        l_i = l_new
    
    # ==================== 5. 最终归一化并写回 ====================
    acc = acc / l_i
    tl.store(out_ptr + q_offset, acc.to(out_ptr.dtype.element_ty))

Host 端封装

def triton_paged_attention_decode(q, k_cache, v_cache, 
                                   block_tables, context_lens, block_size):
    """
    q:            (batch, num_heads, head_dim)
    k_cache:      (num_blocks, block_size, num_heads, head_dim)
    v_cache:      同上
    block_tables: (batch, max_logical_blocks), int32
    context_lens: (batch,), int32
    """
    batch_size, num_heads, head_dim = q.shape
    out = torch.empty_like(q)
    sm_scale = 1.0 / math.sqrt(head_dim)
    
    grid = (batch_size, num_heads)
    
    paged_attention_decoding_kernel[grid](
        out, q, k_cache, v_cache, block_tables, context_lens, sm_scale,
        k_cache.stride(0), k_cache.stride(1), 
        k_cache.stride(2), k_cache.stride(3),
        block_tables.stride(0), block_tables.stride(1),
        BLOCK_SIZE=block_size, HEAD_DIM=head_dim,
    )
    return out

---

技术点	错误/朴素做法	正确/推荐做法	原因说明
Q-K 点积	tl.dot(q, k.T)	tl.sum(q[None,:]*k, axis=1)	Decoding Q 是向量，dot 要求至少 2D；element-wise + reduce 更高效
Mask 填充值	other=0.0 for qk	tl.where(mask, qk, -inf)	0.0 参与 softmax 会窃取概率质量；-inf → exp(-inf)=0 才是正确的
KV load 填充	other=-inf	other=0.0	KV 是数据不是分数，填充 0 不影响 dot 结果；填 -inf 会导致 NaN
m_i 初始值	0.0	-float('inf')	同 Flash Attention，防止全负 qk 时修正错误
物理块索引类型	float / int64	int32	block_tables 通常是 int32；类型不匹配导致 tl.load 失败
stride 传递	硬编码	从 .stride() 动态获取	物理池可能被 slice/view，stride 不一定是理论值
Grid 设计	1D flatten	2D (batch, heads)	Decoding 每程序处理 1 个 Q 向量，2D Grid 自然映射

---

常见陷阱

1. Mask 填充值混淆（最高频错误）
   • 现象：结果数值偏小或 NaN
   • 原因：tl.load KV 时用 other=-inf，导致 q*k 产生 NaN；或 qk 未 mask 就直接 softmax，无效位置的 0.0 参与了归一化
   • 修复：KV load 用 other=0.0；qk 计算后用 tl.where(mask, qk, -inf)
2. Decoding 误用 tl.dot
   • 现象：编译报错 "invalid dot operands"
   • 原因：tl.dot 要求两个操作数都是 ≥2D 且满足矩阵乘维度约束；Decoding 的 Q 是 1D 向量
   • 修复：改用 tl.sum(q[None,:] * k, axis=1)
3. block_tables stride 硬编码
   • 现象：batch > 1 时第二个序列结果全错
   • 原因：假设 stride_bt_batch == max_logical_blocks，但实际 tensor 可能有 padding
   • 修复：始终通过 block_tables.stride(0) 传入
4. context_len 边界 off-by-one
   • 现象：最后一个有效 token 被 mask 掉，或多算一个无效 token
   • 原因：physical_token_idx < context_len 写成了 <= 或起始索引计算错误
   • 修复：手算验证 logical_block_idx=0 和最后一个块的 token 范围
5. acc 精度丢失
   • 现象：长序列下误差累积
   • 原因：acc 使用 FP16 累加
   • 修复：acc 和 m_i/l_i 始终用 FP32，仅在 store 时转回目标 dtype

自查 Checklist

• KV tl.load 的 other 是否为 0.0？
• qk 是否用 tl.where(mask, qk, -inf) 做了有效屏蔽？
• Q-K 点积是否用了向量友好的 sum(mul) 而非 dot？
• m_i 是否初始化为 -float('inf')？
• acc, m_i, l_i 是否都是 FP32？
• physical_block_idx 是否通过 tl.load 从 block_tables 读取？
• stride 参数是否全部从 host 端动态传入？
• 是否与 PyTorch 参考实现做了数值对齐（FP16 误差 < 2e-3）？

---

#mermaid-svg-pBJK4Y65Nu21OXxN{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-pBJK4Y65Nu21OXxN .error-icon{fill:#552222;}#mermaid-svg-pBJK4Y65Nu21OXxN .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-pBJK4Y65Nu21OXxN .marker{fill:#333333;stroke:#333333;}#mermaid-svg-pBJK4Y65Nu21OXxN .marker.cross{stroke:#333333;}#mermaid-svg-pBJK4Y65Nu21OXxN svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-pBJK4Y65Nu21OXxN p{margin:0;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge{stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 path{fill:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 text{fill:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon--1{font-size:40px;color:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge--1{stroke:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth--1{stroke-width:17;}#mermaid-svg-pBJK4Y65Nu21OXxN .section--1 line{stroke:hsl(60, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 path{fill:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-0{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-0{stroke:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-0{stroke-width:14;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-0 line{stroke:hsl(240, 100%, 83.5294117647%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 path{fill:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-1{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-1{stroke:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-1{stroke-width:11;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-1 line{stroke:hsl(260, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 path{fill:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 text{fill:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-2{font-size:40px;color:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-2{stroke:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-2{stroke-width:8;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 line{stroke:hsl(90, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 path{fill:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-3{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-3{stroke:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-3{stroke-width:5;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-3 line{stroke:hsl(120, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 path{fill:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-4{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-4{stroke:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-4{stroke-width:2;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-4 line{stroke:hsl(150, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 path{fill:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-5{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-5{stroke:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-5{stroke-width:-1;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-5 line{stroke:hsl(180, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 path{fill:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-6{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-6{stroke:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-6{stroke-width:-4;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-6 line{stroke:hsl(210, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 path{fill:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-7{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-7{stroke:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-7{stroke-width:-7;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-7 line{stroke:hsl(270, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 path{fill:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-8{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-8{stroke:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-8{stroke-width:-10;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-8 line{stroke:hsl(330, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 path{fill:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-9{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-9{stroke:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-9{stroke-width:-13;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-9 line{stroke:hsl(0, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 polygon,#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 path{fill:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 text{fill:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .node-icon-10{font-size:40px;color:black;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-edge-10{stroke:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .edge-depth-10{stroke-width:-16;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-10 line{stroke:hsl(30, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled circle,#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:lightgray;}#mermaid-svg-pBJK4Y65Nu21OXxN .disabled text{fill:#efefef;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-root rect,#mermaid-svg-pBJK4Y65Nu21OXxN .section-root path,#mermaid-svg-pBJK4Y65Nu21OXxN .section-root circle,#mermaid-svg-pBJK4Y65Nu21OXxN .section-root polygon{fill:hsl(240, 100%, 46.2745098039%);}#mermaid-svg-pBJK4Y65Nu21OXxN .section-root text{fill:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-root span{color:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .section-2 span{color:#ffffff;}#mermaid-svg-pBJK4Y65Nu21OXxN .icon-container{height:100%;display:flex;justify-content:center;align-items:center;}#mermaid-svg-pBJK4Y65Nu21OXxN .edge{fill:none;}#mermaid-svg-pBJK4Y65Nu21OXxN .mindmap-node-label{dy:1em;alignment-baseline:middle;text-anchor:middle;dominant-baseline:middle;text-align:center;}#mermaid-svg-pBJK4Y65Nu21OXxN :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} LLM 推理显存优化
KV Cache 管理
PagedAttention ← 本节
消除碎片
间接寻址
Decoding 优化
Prefix Caching
跨请求共享 KV 块
Radix Tree 索引
模型压缩
量化 Quantization ← 下一节
INT8 / FP8 权重
KV Cache 量化
剪枝 / 蒸馏
计算优化
Flash Attention ← 上一节
Prefill 加速
Speculative Decoding
推测解码
Continuous Batching
动态组批

PagedAttention 与量化的关系

关键认知 ：PagedAttention 解决的是 "怎么更好地用显存" ，Quantization 解决的是 "怎么让模型本身更省显存"。两者正交互补：

• PagedAttention 让 KV Cache 利用率从 ~50% → >95%
• INT8 量化让权重显存减半
• 组合效果：70B 模型在单卡 A100-80GB 上可服务更多并发请求

---

特性	教学版	vLLM 工业版
Block Table 查找	每次循环 tl.load	预加载到 SRAM / 寄存器
KV Cache 布局	[B,S,H,D] 分离 K/V	[B,S,H,KV_D] K/V 交错存储
GQA/MQA 支持	仅 MHA	KV Head 映射 + 分组归约
多步 Decoding	单 Token	Speculative Decode 多 Token
Autotune	固定 BLOCK_SIZE	按 head_dim/context_len 自动选择
Warp 调度	默认	Warp Specialization 隐藏访存延迟
KV 量化	FP16	FP8/INT8 on-the-fly dequant

1. 先跑通本节代码，重点理解间接寻址的实现方式
2. 故意交换 KV load 的 other 值和 qk mask 的填充值，观察两种错误的不同表现
3. 将 block_tables 改为非连续分配（如随机打乱物理块索引），验证间接寻址的正确性
4. 思考：如果要将此 Kernel 扩展支持 GQA（Grouped Query Attention），哪些地方需要修改？（提示：KV head 索引 ≠ Q head 索引）
5. 对比 Flash Attention（08）和本节代码，列出所有因 Decoding 单 Token 特性而产生的简化点

Triton W8A16 Quantization GEMM：即时反量化融合矩阵乘

**为什么 Weight-Only 量化是推理标配？**LLM 推理（尤其是 Decoding）是典型的 Memory Bound 任务。瓶颈不在算力，而在 HBM 带宽。

模型	FP16 权重	INT8 权重	A100-80GB 单卡	核心收益
7B	14 GB	7 GB	充裕	更大 KV Cache / Batch
13B	26 GB	13 GB	可跑	从"勉强"变"从容"
70B	140 GB	70 GB	可单卡！	跨代际部署能力

关键洞察 ：W8A16 的核心价值不是"算得更快"，而是 "读得更少"。INT8 权重体积减半 → HBM 读取量减半 → Memory Bound 场景下吞吐提升。

传统方案 vs On-the-fly 融合
#mermaid-svg-OuK47TqJKHeyJij9{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-OuK47TqJKHeyJij9 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-OuK47TqJKHeyJij9 .error-icon{fill:#552222;}#mermaid-svg-OuK47TqJKHeyJij9 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-OuK47TqJKHeyJij9 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-OuK47TqJKHeyJij9 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-OuK47TqJKHeyJij9 .marker.cross{stroke:#333333;}#mermaid-svg-OuK47TqJKHeyJij9 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-OuK47TqJKHeyJij9 p{margin:0;}#mermaid-svg-OuK47TqJKHeyJij9 .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-OuK47TqJKHeyJij9 .cluster-label text{fill:#333;}#mermaid-svg-OuK47TqJKHeyJij9 .cluster-label span{color:#333;}#mermaid-svg-OuK47TqJKHeyJij9 .cluster-label span p{background-color:transparent;}#mermaid-svg-OuK47TqJKHeyJij9 .label text,#mermaid-svg-OuK47TqJKHeyJij9 span{fill:#333;color:#333;}#mermaid-svg-OuK47TqJKHeyJij9 .node rect,#mermaid-svg-OuK47TqJKHeyJij9 .node circle,#mermaid-svg-OuK47TqJKHeyJij9 .node ellipse,#mermaid-svg-OuK47TqJKHeyJij9 .node polygon,#mermaid-svg-OuK47TqJKHeyJij9 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-OuK47TqJKHeyJij9 .rough-node .label text,#mermaid-svg-OuK47TqJKHeyJij9 .node .label text,#mermaid-svg-OuK47TqJKHeyJij9 .image-shape .label,#mermaid-svg-OuK47TqJKHeyJij9 .icon-shape .label{text-anchor:middle;}#mermaid-svg-OuK47TqJKHeyJij9 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-OuK47TqJKHeyJij9 .rough-node .label,#mermaid-svg-OuK47TqJKHeyJij9 .node .label,#mermaid-svg-OuK47TqJKHeyJij9 .image-shape .label,#mermaid-svg-OuK47TqJKHeyJij9 .icon-shape .label{text-align:center;}#mermaid-svg-OuK47TqJKHeyJij9 .node.clickable{cursor:pointer;}#mermaid-svg-OuK47TqJKHeyJij9 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-OuK47TqJKHeyJij9 .arrowheadPath{fill:#333333;}#mermaid-svg-OuK47TqJKHeyJij9 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-OuK47TqJKHeyJij9 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-OuK47TqJKHeyJij9 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-OuK47TqJKHeyJij9 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-OuK47TqJKHeyJij9 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-OuK47TqJKHeyJij9 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-OuK47TqJKHeyJij9 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-OuK47TqJKHeyJij9 .cluster text{fill:#333;}#mermaid-svg-OuK47TqJKHeyJij9 .cluster span{color:#333;}#mermaid-svg-OuK47TqJKHeyJij9 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-OuK47TqJKHeyJij9 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-OuK47TqJKHeyJij9 rect.text{fill:none;stroke-width:0;}#mermaid-svg-OuK47TqJKHeyJij9 .icon-shape,#mermaid-svg-OuK47TqJKHeyJij9 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-OuK47TqJKHeyJij9 .icon-shape p,#mermaid-svg-OuK47TqJKHeyJij9 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-OuK47TqJKHeyJij9 .icon-shape .label rect,#mermaid-svg-OuK47TqJKHeyJij9 .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-OuK47TqJKHeyJij9 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-OuK47TqJKHeyJij9 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-OuK47TqJKHeyJij9 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} ✅ On-the-fly 融合
半带宽读取
INT8 W

(HBM)
Fused Kernel

SRAM内反量化+GEMM
FP16 Scales
FP16 X
❌ 传统两步法
全量反量化
再次读取
INT8 W

(HBM)
FP16 W'

(HBM)
GEMM Kernel
FP16 X

指标	传统两步法	On-the-fly 融合	差异
HBM 读 W	2× (INT8 + FP16)	1× (仅 INT8)	省一半带宽
额外显存	O(K×N) FP16 临时张量	0	省 ~14GB@70B
Kernel Launch	2 次	1 次	减少启动开销
计算开销	无额外	SRAM 内 type cast + mul	几乎免费

---

量化公式与数据布局

Per-channel Symmetric Quantization 本节采用最基础的逐列对称量化 ：Y=X⋅Wfp16≈X⋅(int8(W)⊙s)其中 s∈R1×N 是每列的缩放因子， ⊙ 表示逐列广播乘法。内存布局

张量	Shape	Dtype	说明
X	(M, K)	FP16	激活值，保持高精度
W_int8	(K, N)	INT8	量化权重，体积减半
scales	(N,)	FP16	每列一个缩放因子
Y	(M, N)	FP16	输出，FP32 累加后转回

Scale 粒度选择：Per-channel（每列一个 scale）比 Per-tensor（全局一个）精度高很多，且 scale 加载成本极低（每个 BLOCK_N 只读一次），是工程实践的首选基线。

---

Kernel 实现

import torch
import triton
import triton.language as tl

@triton.jit
def w8a16_gemm_kernel(
    x_ptr, w_int8_ptr, scales_ptr, y_ptr,
    M, N, K,
    stride_xm, stride_xk,
    stride_wk, stride_wn,
    stride_ym, stride_yn,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
):
    # ==================== 1. Grid 定位 ====================
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    
    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    
    # Scale 只依赖 N 维度，提前计算指针，避免在 K 循环中重复计算
    scale_ptrs = scales_ptr + offs_n
    
    # FP32 累加器，防止 FP16 精度丢失
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    
    # ==================== 2. K 维分块循环 ====================
    for k in range(0, tl.cdiv(K, BLOCK_K)):
        offs_k = k * BLOCK_K + tl.arange(0, BLOCK_K)
        
        # X: (BLOCK_M, BLOCK_K)
        x_ptrs = x_ptr + (offs_m[:, None] * stride_xm 
                          + offs_k[None, :] * stride_xk)
        # W: (BLOCK_K, BLOCK_N)
        w_ptrs = w_int8_ptr + (offs_k[:, None] * stride_wk 
                               + offs_n[None, :] * stride_wn)
        
        x = tl.load(x_ptrs)           # FP16
        w_int8 = tl.load(w_ptrs)      # INT8 ← 半带宽读取
        
        # TODO 1: SRAM 内即时反量化
        w_fp16 = w_int8.to(x.dtype)          # INT8 → FP16 (寄存器级转换)
        scales = tl.load(scale_ptrs)          # (BLOCK_N,)
        w_fp16 = w_fp16 * scales[None, :]     # 逐列广播乘 scale
        
        # TODO 2: 融合矩阵乘累加
        acc += tl.dot(x, w_fp16)              # FP16 dot → FP32 accumulate
    
    # ==================== 3. 写回结果 ====================
    y_ptrs = y_ptr + (offs_m[:, None] * stride_ym 
                      + offs_n[None, :] * stride_yn)
    tl.store(y_ptrs, acc.to(tl.float16))

Host 端封装

def triton_w8a16_gemm(x: torch.Tensor, w_int8: torch.Tensor, 
                       scales: torch.Tensor):
    """
    x:       (M, K), FP16
    w_int8:  (K, N), INT8
    scales:  (N,),   FP16
    """
    M, K = x.shape
    _, N = w_int8.shape
    y = torch.empty((M, N), device=x.device, dtype=torch.float16)
    
    BLOCK_M = 16
    BLOCK_N = 64
    BLOCK_K = 64
    
    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
    
    w8a16_gemm_kernel[grid](
        x, w_int8, scales, y,
        M, N, K,
        x.stride(0), x.stride(1),
        w_int8.stride(0), w_int8.stride(1),
        y.stride(0), y.stride(1),
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K,
    )
    return y

---

SRAM 内执行流详解

理解这个 Kernel 的关键在于看清每一轮 K 循环内 SRAM 中发生了什么：
#mermaid-svg-YpOXFfNGks02SJr7{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-YpOXFfNGks02SJr7 .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-YpOXFfNGks02SJr7 .error-icon{fill:#552222;}#mermaid-svg-YpOXFfNGks02SJr7 .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-YpOXFfNGks02SJr7 .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-YpOXFfNGks02SJr7 .marker{fill:#333333;stroke:#333333;}#mermaid-svg-YpOXFfNGks02SJr7 .marker.cross{stroke:#333333;}#mermaid-svg-YpOXFfNGks02SJr7 svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-YpOXFfNGks02SJr7 p{margin:0;}#mermaid-svg-YpOXFfNGks02SJr7 .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-YpOXFfNGks02SJr7 .cluster-label text{fill:#333;}#mermaid-svg-YpOXFfNGks02SJr7 .cluster-label span{color:#333;}#mermaid-svg-YpOXFfNGks02SJr7 .cluster-label span p{background-color:transparent;}#mermaid-svg-YpOXFfNGks02SJr7 .label text,#mermaid-svg-YpOXFfNGks02SJr7 span{fill:#333;color:#333;}#mermaid-svg-YpOXFfNGks02SJr7 .node rect,#mermaid-svg-YpOXFfNGks02SJr7 .node circle,#mermaid-svg-YpOXFfNGks02SJr7 .node ellipse,#mermaid-svg-YpOXFfNGks02SJr7 .node polygon,#mermaid-svg-YpOXFfNGks02SJr7 .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-YpOXFfNGks02SJr7 .rough-node .label text,#mermaid-svg-YpOXFfNGks02SJr7 .node .label text,#mermaid-svg-YpOXFfNGks02SJr7 .image-shape .label,#mermaid-svg-YpOXFfNGks02SJr7 .icon-shape .label{text-anchor:middle;}#mermaid-svg-YpOXFfNGks02SJr7 .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-YpOXFfNGks02SJr7 .rough-node .label,#mermaid-svg-YpOXFfNGks02SJr7 .node .label,#mermaid-svg-YpOXFfNGks02SJr7 .image-shape .label,#mermaid-svg-YpOXFfNGks02SJr7 .icon-shape .label{text-align:center;}#mermaid-svg-YpOXFfNGks02SJr7 .node.clickable{cursor:pointer;}#mermaid-svg-YpOXFfNGks02SJr7 .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-YpOXFfNGks02SJr7 .arrowheadPath{fill:#333333;}#mermaid-svg-YpOXFfNGks02SJr7 .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-YpOXFfNGks02SJr7 .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-YpOXFfNGks02SJr7 .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-YpOXFfNGks02SJr7 .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-YpOXFfNGks02SJr7 .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-YpOXFfNGks02SJr7 .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-YpOXFfNGks02SJr7 .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-YpOXFfNGks02SJr7 .cluster text{fill:#333;}#mermaid-svg-YpOXFfNGks02SJr7 .cluster span{color:#333;}#mermaid-svg-YpOXFfNGks02SJr7 div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-YpOXFfNGks02SJr7 .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-YpOXFfNGks02SJr7 rect.text{fill:none;stroke-width:0;}#mermaid-svg-YpOXFfNGks02SJr7 .icon-shape,#mermaid-svg-YpOXFfNGks02SJr7 .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-YpOXFfNGks02SJr7 .icon-shape p,#mermaid-svg-YpOXFfNGks02SJr7 .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-YpOXFfNGks02SJr7 .icon-shape .label rect,#mermaid-svg-YpOXFfNGks02SJr7 .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-YpOXFfNGks02SJr7 .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-YpOXFfNGks02SJr7 .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-YpOXFfNGks02SJr7 :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 否
是
tl.load INT8 W block

(半带宽 HBM→SRAM)
w_int8.to(fp16)

(寄存器类型转换)
tl.load scales

(极小开销)
w_fp16 * scalesNone,:

(SRAM 内广播乘)
tl.load FP16 X block
tl.dot(X, W_fp16)

(Tensor Core FP16→FP32)
acc += result

(FP32 累加器)
K 循环结束?
acc.to(fp16) → tl.store

核心要点 ：反量化的 type cast + multiply 发生在寄存器/SRAM 中，从未产生过完整的 FP16 权重矩阵。这就是"On-the-fly"的含义------数据以 INT8 形态穿越 HBM，以 FP16 形态参与计算，中间态只在片上存在。

---

技术点	错误/朴素做法	正确/推荐做法	原因说明
累加器精度	FP16 acc	FP32 acc	INT8 反量化 + 大量累加，FP16 会严重截断
Scale 加载位置	K 循环内每次加载	K 循环外预计算指针	Scale 只依赖 N，与 K 无关，避免冗余访存
类型转换时机	load 后立即转	dot 前转即可	保持 INT8 形态更久，理论上利于编译器优化
Scale 广播	scales * w_fp16	w_fp16 * scales[None, :]	显式 [None, :] 确保广播语义正确
dot 输入类型	FP32 × FP32	FP16 × FP16	Tensor Core FP16 指令吞吐远高于 FP32
Grid 划分	1D flatten	2D (M//BM, N//BN)	标准 GEMM tiling，利于 L2 Cache 复用
W 存储顺序	Row-major (K,N)	Column-major or (K,N) contiguous	确保 N 维度连续，配合 BLOCK_N 向量化读取

---

常见陷阱

1. FP16 累加器导致精度崩溃
   • 现象：小矩阵正确，大矩阵（K≥1024）误差指数增长
   • 原因：FP16 有效位仅 10 bit，数千次累加后低位完全丢失
   • 修复：acc 始终声明为 dtype=tl.float32
2. Scale 在 K 循环内重复加载
   • 现象：性能下降 5-10%
   • 原因：Scale 只与 N 有关，每轮 K 迭代都重新 load 浪费带宽和指令
   • 修复：将 scale_ptrs 计算移到 K 循环外
3. INT8 符号扩展问题
   • 现象：部分值异常偏大
   • 原因：某些硬件/编译器对 INT8→FP16 转换不做符号扩展
   • 修复：Triton 的 .to(tl.float16) 默认正确处理有符号 INT8；若用自定义 cast 需验证
4. Scale shape 不匹配
   • 现象：编译报错或广播错误
   • 原因：scales 是 (N,) 但忘了 [None, :] 扩展为 (1, BLOCK_N)
   • 修复：scales[None, :] 显式广播
5. Benchmark 误判
   • 现象：Triton W8A16 比 PyTorch FP16 GEMM 还慢
   • 原因：小矩阵是 Compute Bound，反量化开销 > 带宽节省；PyTorch FP16 GEMM 调用高度优化的 cuBLAS
   • 修复：W8A16 的收益只在 Memory Bound 场景（大 K/N、小 M）体现；Benchmark 应包含典型 LLM Linear 尺寸

自查 Checklist

• acc 是否为 FP32？
• scale_ptrs 是否在 K 循环外计算？
• w_int8.to(x.dtype) 是否在 tl.dot 之前完成？
• scales 广播是否使用了 [None, :]？
• tl.dot 的两个操作数是否都是 FP16？
• Store 时是否做了 acc.to(tl.float16)？
• 是否与 PyTorch 参考实现做了数值对齐（误差 < 1e-3）？
• Benchmark 是否覆盖了 Memory Bound 的典型尺寸？

---

#mermaid-svg-lIHWIv213S6K5aSj{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-lIHWIv213S6K5aSj .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-lIHWIv213S6K5aSj .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-lIHWIv213S6K5aSj .error-icon{fill:#552222;}#mermaid-svg-lIHWIv213S6K5aSj .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-lIHWIv213S6K5aSj .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-lIHWIv213S6K5aSj .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-lIHWIv213S6K5aSj .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-lIHWIv213S6K5aSj .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-lIHWIv213S6K5aSj .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-lIHWIv213S6K5aSj .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-lIHWIv213S6K5aSj .marker{fill:#333333;stroke:#333333;}#mermaid-svg-lIHWIv213S6K5aSj .marker.cross{stroke:#333333;}#mermaid-svg-lIHWIv213S6K5aSj svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-lIHWIv213S6K5aSj p{margin:0;}#mermaid-svg-lIHWIv213S6K5aSj .edge{stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .section--1 rect,#mermaid-svg-lIHWIv213S6K5aSj .section--1 path,#mermaid-svg-lIHWIv213S6K5aSj .section--1 circle,#mermaid-svg-lIHWIv213S6K5aSj .section--1 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section--1 path{fill:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section--1 text{fill:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon--1{font-size:40px;color:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge--1{stroke:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth--1{stroke-width:17;}#mermaid-svg-lIHWIv213S6K5aSj .section--1 line{stroke:hsl(60, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-0 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-0 path,#mermaid-svg-lIHWIv213S6K5aSj .section-0 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-0 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-0 path{fill:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-lIHWIv213S6K5aSj .section-0 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-0{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-0{stroke:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-0{stroke-width:14;}#mermaid-svg-lIHWIv213S6K5aSj .section-0 line{stroke:hsl(240, 100%, 83.5294117647%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-1 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-1 path,#mermaid-svg-lIHWIv213S6K5aSj .section-1 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-1 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-1 path{fill:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-1 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-1{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-1{stroke:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-1{stroke-width:11;}#mermaid-svg-lIHWIv213S6K5aSj .section-1 line{stroke:hsl(260, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-2 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-2 path,#mermaid-svg-lIHWIv213S6K5aSj .section-2 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-2 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-2 path{fill:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-2 text{fill:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-2{font-size:40px;color:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-2{stroke:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-2{stroke-width:8;}#mermaid-svg-lIHWIv213S6K5aSj .section-2 line{stroke:hsl(90, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-3 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-3 path,#mermaid-svg-lIHWIv213S6K5aSj .section-3 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-3 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-3 path{fill:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-3 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-3{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-3{stroke:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-3{stroke-width:5;}#mermaid-svg-lIHWIv213S6K5aSj .section-3 line{stroke:hsl(120, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-4 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-4 path,#mermaid-svg-lIHWIv213S6K5aSj .section-4 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-4 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-4 path{fill:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-4 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-4{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-4{stroke:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-4{stroke-width:2;}#mermaid-svg-lIHWIv213S6K5aSj .section-4 line{stroke:hsl(150, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-5 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-5 path,#mermaid-svg-lIHWIv213S6K5aSj .section-5 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-5 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-5 path{fill:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-5 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-5{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-5{stroke:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-5{stroke-width:-1;}#mermaid-svg-lIHWIv213S6K5aSj .section-5 line{stroke:hsl(180, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-6 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-6 path,#mermaid-svg-lIHWIv213S6K5aSj .section-6 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-6 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-6 path{fill:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-6 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-6{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-6{stroke:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-6{stroke-width:-4;}#mermaid-svg-lIHWIv213S6K5aSj .section-6 line{stroke:hsl(210, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-7 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-7 path,#mermaid-svg-lIHWIv213S6K5aSj .section-7 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-7 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-7 path{fill:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-7 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-7{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-7{stroke:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-7{stroke-width:-7;}#mermaid-svg-lIHWIv213S6K5aSj .section-7 line{stroke:hsl(270, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-8 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-8 path,#mermaid-svg-lIHWIv213S6K5aSj .section-8 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-8 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-8 path{fill:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-8 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-8{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-8{stroke:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-8{stroke-width:-10;}#mermaid-svg-lIHWIv213S6K5aSj .section-8 line{stroke:hsl(330, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-9 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-9 path,#mermaid-svg-lIHWIv213S6K5aSj .section-9 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-9 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-9 path{fill:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-9 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-9{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-9{stroke:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-9{stroke-width:-13;}#mermaid-svg-lIHWIv213S6K5aSj .section-9 line{stroke:hsl(0, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-10 rect,#mermaid-svg-lIHWIv213S6K5aSj .section-10 path,#mermaid-svg-lIHWIv213S6K5aSj .section-10 circle,#mermaid-svg-lIHWIv213S6K5aSj .section-10 polygon,#mermaid-svg-lIHWIv213S6K5aSj .section-10 path{fill:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-10 text{fill:black;}#mermaid-svg-lIHWIv213S6K5aSj .node-icon-10{font-size:40px;color:black;}#mermaid-svg-lIHWIv213S6K5aSj .section-edge-10{stroke:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .edge-depth-10{stroke-width:-16;}#mermaid-svg-lIHWIv213S6K5aSj .section-10 line{stroke:hsl(30, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-lIHWIv213S6K5aSj .disabled,#mermaid-svg-lIHWIv213S6K5aSj .disabled circle,#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:lightgray;}#mermaid-svg-lIHWIv213S6K5aSj .disabled text{fill:#efefef;}#mermaid-svg-lIHWIv213S6K5aSj .section-root rect,#mermaid-svg-lIHWIv213S6K5aSj .section-root path,#mermaid-svg-lIHWIv213S6K5aSj .section-root circle,#mermaid-svg-lIHWIv213S6K5aSj .section-root polygon{fill:hsl(240, 100%, 46.2745098039%);}#mermaid-svg-lIHWIv213S6K5aSj .section-root text{fill:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .section-root span{color:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .section-2 span{color:#ffffff;}#mermaid-svg-lIHWIv213S6K5aSj .icon-container{height:100%;display:flex;justify-content:center;align-items:center;}#mermaid-svg-lIHWIv213S6K5aSj .edge{fill:none;}#mermaid-svg-lIHWIv213S6K5aSj .mindmap-node-label{dy:1em;alignment-baseline:middle;text-anchor:middle;dominant-baseline:middle;text-align:center;}#mermaid-svg-lIHWIv213S6K5aSj :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} LLM 推理优化
算得更快
Prefill
Decode
Operator Fusion
省得更多
W8A16 量化 ← 本节
On-the-fly Dequant
Per-channel Scale
W4A16 / GPTQ / AWQ
KV Cache 量化
服务更灵活
Continuous Batching
Speculative Decoding
Prefix Caching

量化技术的演进路线

级别	方案	Scale 粒度	精度	速度	复杂度
基线	W8A16 Per-channel	每列	★★★★	★★★	低 ← 本节
进阶	W8A16 Group-wise	每 G 列共享	★★★★☆	★★★	中
高级	W4A16 GPTQ/AWQ	每 G 列 + zero_point	★★★★	★★★★	高
前沿	W8A8 / FP8	Per-token × Per-channel	★★★★★	★★★★★	很高

1. 先跑通本节代码，验证数值正确性
2. 将 acc 改为 FP16，观察不同 K 值下的误差变化曲线，直观感受精度损失
3. 对比 M=1 (Decode) 和 M=4096 (Prefill) 两种场景的加速比，理解 Memory Bound vs Compute Bound
4. 思考：如果要扩展到 Group-wise 量化（每 128 列共享一个 scale），Kernel 需要改哪些地方？（提示：scale 索引从 offs_n 变为 offs_n // group_size）
5. 思考：W8A16 与 PagedAttention 能否组合？如果可以，KV Cache 也用 INT8 存储，解码时的 Kernel 需要怎样改造？

Triton Multi-LoRA：Token 级动态路由与融合推理

多租户 LoRA 推理的困境 在 SaaS 场景中，成百上千个用户可能同时请求不同的 LoRA 适配器（代码生成、翻译、摘要等）。传统做法面临两难：

方案	做法	问题
串行执行	按 LoRA 分组，逐组调用 linear()	GPU 利用率低，小 Batch 无法打满算力
Padding 对齐	所有请求 pad 到相同长度再 batch	浪费大量无效计算和显存
权重切换	每次请求前加载对应 LoRA 权重	HBM 带宽被权重搬运占满，计算停滞

Multi-LoRA 的核心洞察
关键思想 ：将所有 LoRA 权重预加载到统一的**显存池（Weight Pool）**中，通过 Token 级路由索引 在一次 Kernel 调用中完成异构 LoRA 的并行计算。Batch 内每个 Token 可以属于不同的 LoRA，但共享同一个 Kernel Launch。
#mermaid-svg-VkMkmaJbuzAaKwCZ{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-VkMkmaJbuzAaKwCZ .error-icon{fill:#552222;}#mermaid-svg-VkMkmaJbuzAaKwCZ .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-VkMkmaJbuzAaKwCZ .marker{fill:#333333;stroke:#333333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .marker.cross{stroke:#333333;}#mermaid-svg-VkMkmaJbuzAaKwCZ svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-VkMkmaJbuzAaKwCZ p{margin:0;}#mermaid-svg-VkMkmaJbuzAaKwCZ .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .cluster-label text{fill:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .cluster-label span{color:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .cluster-label span p{background-color:transparent;}#mermaid-svg-VkMkmaJbuzAaKwCZ .label text,#mermaid-svg-VkMkmaJbuzAaKwCZ span{fill:#333;color:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .node rect,#mermaid-svg-VkMkmaJbuzAaKwCZ .node circle,#mermaid-svg-VkMkmaJbuzAaKwCZ .node ellipse,#mermaid-svg-VkMkmaJbuzAaKwCZ .node polygon,#mermaid-svg-VkMkmaJbuzAaKwCZ .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .rough-node .label text,#mermaid-svg-VkMkmaJbuzAaKwCZ .node .label text,#mermaid-svg-VkMkmaJbuzAaKwCZ .image-shape .label,#mermaid-svg-VkMkmaJbuzAaKwCZ .icon-shape .label{text-anchor:middle;}#mermaid-svg-VkMkmaJbuzAaKwCZ .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .rough-node .label,#mermaid-svg-VkMkmaJbuzAaKwCZ .node .label,#mermaid-svg-VkMkmaJbuzAaKwCZ .image-shape .label,#mermaid-svg-VkMkmaJbuzAaKwCZ .icon-shape .label{text-align:center;}#mermaid-svg-VkMkmaJbuzAaKwCZ .node.clickable{cursor:pointer;}#mermaid-svg-VkMkmaJbuzAaKwCZ .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .arrowheadPath{fill:#333333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-VkMkmaJbuzAaKwCZ .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-VkMkmaJbuzAaKwCZ .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-VkMkmaJbuzAaKwCZ .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-VkMkmaJbuzAaKwCZ .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .cluster text{fill:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ .cluster span{color:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-VkMkmaJbuzAaKwCZ .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-VkMkmaJbuzAaKwCZ rect.text{fill:none;stroke-width:0;}#mermaid-svg-VkMkmaJbuzAaKwCZ .icon-shape,#mermaid-svg-VkMkmaJbuzAaKwCZ .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-VkMkmaJbuzAaKwCZ .icon-shape p,#mermaid-svg-VkMkmaJbuzAaKwCZ .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-VkMkmaJbuzAaKwCZ .icon-shape .label rect,#mermaid-svg-VkMkmaJbuzAaKwCZ .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-VkMkmaJbuzAaKwCZ .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-VkMkmaJbuzAaKwCZ .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-VkMkmaJbuzAaKwCZ :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} Single Triton Kernel
统一 LoRA 权重池
并发请求
Token 0

LoRA #2 (代码)
Token 1

LoRA #0 (翻译)
Token 2

LoRA #1 (摘要)
Token 3

LoRA #2 (代码)
LoRA #0
LoRA #1
LoRA #2
Program(0) → Pool2
Program(1) → Pool0
Program(2) → Pool1
Program(3) → Pool2

推理优化主线收束 本节是推理优化系列的收束点，将前面所有技术串联为完整的多租户服务栈：

章节	解决的问题	在本节中的体现
07 RoPE	位置编码融合	LoRA 作用于 Linear 层，RoPE 在 Attention 层独立处理
08 FlashAttn	Attention 加速	LoRA 通常挂载在 QKV/FFN 投影层，与 Attn 解耦
09 PagedAttn	KV Cache 碎片管理	不同 LoRA 请求共享同一 KV Cache 池
10 Quantization	权重体积压缩	LoRA 池本身也可量化（W8A16 + LoRA）
11 Multi-LoRA	多租户动态路由	本节：让以上优化在多租户场景下高效复用

---

内存池模型与指针路由

三维权重池布局

lora_a_pool: (num_loras, R, IN_DIM)     ← 降维矩阵
lora_b_pool: (num_loras, OUT_DIM, R)    ← 升维矩阵

给定 lora_idx，对应 LoRA 的基地址为：

baseA=poolA+lora_idx×stride_a_pool

为什么用 stride 而非直接索引？ 因为 PyTorch 张量可能被 slice/view，实际 stride 不一定等于理论值。始终通过 .stride() 传入，保证正确性。

Token 级路由的执行流 每个 Triton Program 负责 (1 个 Token, 1 个输出块)：
#mermaid-svg-RMBGQWvf8jV5ccDd{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-RMBGQWvf8jV5ccDd .error-icon{fill:#552222;}#mermaid-svg-RMBGQWvf8jV5ccDd .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-RMBGQWvf8jV5ccDd .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-RMBGQWvf8jV5ccDd .marker{fill:#333333;stroke:#333333;}#mermaid-svg-RMBGQWvf8jV5ccDd .marker.cross{stroke:#333333;}#mermaid-svg-RMBGQWvf8jV5ccDd svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-RMBGQWvf8jV5ccDd p{margin:0;}#mermaid-svg-RMBGQWvf8jV5ccDd .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd .cluster-label text{fill:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd .cluster-label span{color:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd .cluster-label span p{background-color:transparent;}#mermaid-svg-RMBGQWvf8jV5ccDd .label text,#mermaid-svg-RMBGQWvf8jV5ccDd span{fill:#333;color:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd .node rect,#mermaid-svg-RMBGQWvf8jV5ccDd .node circle,#mermaid-svg-RMBGQWvf8jV5ccDd .node ellipse,#mermaid-svg-RMBGQWvf8jV5ccDd .node polygon,#mermaid-svg-RMBGQWvf8jV5ccDd .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-RMBGQWvf8jV5ccDd .rough-node .label text,#mermaid-svg-RMBGQWvf8jV5ccDd .node .label text,#mermaid-svg-RMBGQWvf8jV5ccDd .image-shape .label,#mermaid-svg-RMBGQWvf8jV5ccDd .icon-shape .label{text-anchor:middle;}#mermaid-svg-RMBGQWvf8jV5ccDd .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-RMBGQWvf8jV5ccDd .rough-node .label,#mermaid-svg-RMBGQWvf8jV5ccDd .node .label,#mermaid-svg-RMBGQWvf8jV5ccDd .image-shape .label,#mermaid-svg-RMBGQWvf8jV5ccDd .icon-shape .label{text-align:center;}#mermaid-svg-RMBGQWvf8jV5ccDd .node.clickable{cursor:pointer;}#mermaid-svg-RMBGQWvf8jV5ccDd .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-RMBGQWvf8jV5ccDd .arrowheadPath{fill:#333333;}#mermaid-svg-RMBGQWvf8jV5ccDd .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-RMBGQWvf8jV5ccDd .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-RMBGQWvf8jV5ccDd .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-RMBGQWvf8jV5ccDd .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-RMBGQWvf8jV5ccDd .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-RMBGQWvf8jV5ccDd .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-RMBGQWvf8jV5ccDd .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-RMBGQWvf8jV5ccDd .cluster text{fill:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd .cluster span{color:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-RMBGQWvf8jV5ccDd .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-RMBGQWvf8jV5ccDd rect.text{fill:none;stroke-width:0;}#mermaid-svg-RMBGQWvf8jV5ccDd .icon-shape,#mermaid-svg-RMBGQWvf8jV5ccDd .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-RMBGQWvf8jV5ccDd .icon-shape p,#mermaid-svg-RMBGQWvf8jV5ccDd .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-RMBGQWvf8jV5ccDd .icon-shape .label rect,#mermaid-svg-RMBGQWvf8jV5ccDd .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-RMBGQWvf8jV5ccDd .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-RMBGQWvf8jV5ccDd .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-RMBGQWvf8jV5ccDd :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} pid_m = Token ID

pid_n = Output Block ID
lora_idx = load(indices + pid_m)
base_A = pool_A + idx × stride_A

base_B = pool_B + idx × stride_B
K循环: x @ A → h_r

(IN_DIM 分块累加)
h_r @ B → acc

(OUT_DIM 分块)
store(out, acc)

---

Kernel 实现

import torch
import triton
import triton.language as tl

@triton.jit
def fused_multi_lora_kernel(
    x_ptr, out_ptr,
    lora_a_pool_ptr, lora_b_pool_ptr,
    lora_indices_ptr,
    M, IN_DIM, OUT_DIM, R: tl.constexpr,
    stride_x_m, stride_x_in,
    stride_out_m, stride_out_dim,
    stride_a_pool, stride_a_r, stride_a_in,
    stride_b_pool, stride_b_out, stride_b_r,
    BLOCK_IN: tl.constexpr, BLOCK_OUT: tl.constexpr,
):
    pid_m = tl.program_id(0)   # Token 维度
    pid_n = tl.program_id(1)   # Output 分块维度

    offs_n = pid_n * BLOCK_OUT + tl.arange(0, BLOCK_OUT)

    # ==================== TODO 1: 读取当前 Token 的 LoRA 索引 ====================
    lora_idx = tl.load(lora_indices_ptr + pid_m)

    # ==================== TODO 2: 计算内存池中该 LoRA 的基地址 ====================
    a_pool_base = lora_a_pool_ptr + lora_idx * stride_a_pool
    b_pool_base = lora_b_pool_ptr + lora_idx * stride_b_pool

    # ==================== TODO 3: x @ A → h_r (低秩中间激活) ====================
    h_r = tl.zeros((R,), dtype=tl.float32)
    num_k_blocks = tl.cdiv(IN_DIM, BLOCK_IN)

    for k in range(num_k_blocks):
        offs_in = k * BLOCK_IN + tl.arange(0, BLOCK_IN)
        in_mask = offs_in < IN_DIM

        # 加载 x 的一行片段: (BLOCK_IN,)
        x_ptrs = x_ptr + pid_m * stride_x_m + offs_in * stride_x_in
        x_val = tl.load(x_ptrs, mask=in_mask, other=0.0)

        # 加载 A 的一个块: (R, BLOCK_IN)
        offs_r = tl.arange(0, R)
        a_ptrs = (a_pool_base
                  + offs_r[:, None] * stride_a_r
                  + offs_in[None, :] * stride_a_in)
        a_val = tl.load(a_ptrs, mask=in_mask[None, :], other=0.0)

        #  向量×矩阵归约: (BLOCK_IN,) × (R, BLOCK_IN) → (R,)
        h_r += tl.sum(x_val[None, :] * a_val, axis=1)

    # ==================== TODO 4: h_r @ B → acc (投影回输出空间) ====================
    offs_r = tl.arange(0, R)
    out_mask = offs_n < OUT_DIM

    # 加载 B 的一个块: (BLOCK_OUT, R)
    b_ptrs = (b_pool_base
              + offs_n[:, None] * stride_b_out
              + offs_r[None, :] * stride_b_r)
    b_val = tl.load(b_ptrs, mask=out_mask[:, None], other=0.0)

    #  向量×矩阵归约: (R,) × (BLOCK_OUT, R) → (BLOCK_OUT,)
    acc = tl.sum(h_r[None, :] * b_val, axis=1)

    # ==================== TODO 5: 写回结果 ====================
    out_ptrs = out_ptr + pid_m * stride_out_m + offs_n * stride_out_dim
    tl.store(out_ptrs, acc.to(out_ptr.dtype.element_ty), mask=out_mask)

Host 端封装

def triton_multi_lora_forward(
    x: torch.Tensor,           # (M, IN_DIM)
    lora_a_pool: torch.Tensor, # (num_loras, R, IN_DIM)
    lora_b_pool: torch.Tensor, # (num_loras, OUT_DIM, R)
    lora_indices: torch.Tensor # (M,), int32
):
    M, IN_DIM = x.shape
    _, OUT_DIM, _ = lora_b_pool.shape
    _, R, _ = lora_a_pool.shape

    out = torch.empty((M, OUT_DIM), device=x.device, dtype=x.dtype)

    BLOCK_IN = 64
    BLOCK_OUT = 64

    grid = (M, triton.cdiv(OUT_DIM, BLOCK_OUT))

    fused_multi_lora_kernel[grid](
        x, out,
        lora_a_pool, lora_b_pool,
        lora_indices,
        M, IN_DIM, OUT_DIM, R,
        x.stride(0), x.stride(1),
        out.stride(0), out.stride(1),
        lora_a_pool.stride(0), lora_a_pool.stride(1), lora_a_pool.stride(2),
        lora_b_pool.stride(0), lora_b_pool.stride(1), lora_b_pool.stride(2),
        BLOCK_IN=BLOCK_IN, BLOCK_OUT=BLOCK_OUT,
    )
    return out

---

关键技术点对照表

技术点	错误/朴素做法	正确/推荐做法	原因说明
路由粒度	Batch 级统一 LoRA	Token 级 lora_indices[pid_m]	支持同一 Batch 内混合多种 LoRA
权重寻址	Python 层切片后传指针	Kernel 内 idx * stride 偏移	避免 Host 端拷贝，一次 Launch 搞定
x@A 计算	tl.dot(x, A.T)	tl.sum(x[None,:]*A, axis=1)	R 很小(8~64)，dot 开销大于收益；element-wise+reduce 更灵活
h_r 精度	FP16 累加	FP32 累加	两次低秩乘加的误差会叠加，FP32 保精度
Mask 方向	统一 mask	x/A 用 in_mask，B/out 用 out_mask	IN_DIM 和 OUT_DIM 边界独立，不可混用
Grid 设计	(M*N,) 1D	(M, cdiv(N,BLOCK_OUT)) 2D	自然映射 Token×OutputBlock，便于 pid 解码
R 的处理	运行时变量	tl.constexpr	R 通常固定且小，编译期展开可消除循环开销

---

常见陷阱

1. lora_idx 越界无保护
   • 现象：随机 CUDA Error / NaN
   • 原因：lora_indices 中的值 ≥ num_loras，指针偏移到非法地址
   • 修复：Host 端校验 assert (lora_indices < num_loras).all()；Kernel 内不做分支判断以保持性能
2. h_r 使用 FP16 累加
   • 现象：小 R 正确，R≥32 时误差飙升
   • 原因：x@A 的结果已经是缩放过的小值，FP16 累加丢失低位
   • 修复：h_r 和 acc 都声明为 FP32
3. A/B 的 stride 顺序搞反
   • 现象：结果全错但不报错
   • 原因：lora_a_pool shape=(L,R,IN)，stride(1)=IN, stride(2)=1；若误用 stride(1) 作为 IN 维步长则完全错位
   • 修复：严格按 stride_a_r 和 stride_a_in 分别索引两个维度
4. OUT_DIM 尾块未 Mask
   • 现象：输出末尾出现垃圾值
   • 原因：BLOCK_OUT 不整除 OUT_DIM 时，超出部分未屏蔽
   • 修复：tl.store 和 B 的 tl.load 都必须带 mask=offs_n < OUT_DIM
5. 误用 tl.dot 做低秩乘法
   • 现象：编译失败或性能反而下降
   • 原因：tl.dot 要求最小维度 ≥ 16 且对齐；R=8 时不满足约束
   • 修复：R < 16 时用 tl.sum(mul, axis)；R ≥ 16 且对齐时可考虑 tl.dot

自查 Checklist

• lora_idx 是否通过 tl.load(lora_indices_ptr + pid_m) 获取？
• A/B 基地址是否用 idx * stride_pool 计算？
• h_r 和 acc 是否为 FP32？
• x@A 的 K 循环是否有正确的 in_mask？
• h_r@B 是否有正确的 out_mask？
• tl.store 是否做了类型转换 .to(out_ptr.dtype.element_ty)？
• Host 端是否校验了 lora_indices 的范围？
• 是否与 PyTorch for-loop 参考实现做了数值对齐？

---

验证 Multi-LoRA 的价值需要合理的对比口径：

# 推荐的 Benchmark 三件套
# 1. 串行基线：按 LoRA 分组，逐组 torch.nn.functional.linear
# 2. Padding 基线：pad 到最大 batch 后统一计算
# 3. Multi-LoRA 融合：本节实现

# 关键观测指标
# - 端到端延迟 (ms)
# - GPU SM 利用率 (%)
# - HBM 吞吐量 (GB/s)
# - 随 Batch Size 增长的吞吐曲线

• 小 Batch（≤8）：Multi-LoRA ≈ 串行，优势不明显
• 中 Batch（32-128）：Multi-LoRA 2-5x 加速，SM 利用率显著提升
• 大 Batch（≥256）：加速比趋于稳定，瓶颈转为 Memory Bound

---

特性	教学版	S-LoRA / Punica 工业版
路由粒度	Token 级	Token 级 + Sequence 级缓存
权重池	单一 FP16 池	分层池（热/冷 LoRA）+ 按需换入
量化支持	无	W8A16 / W4A16 LoRA 权重
Rank 适配	固定 R	混合 Rank 动态 padding / 分组
Base Model 融合	分离计算	Base GEMM + LoRA GEMM 流水线重叠
Autotune	固定 BLOCK	按 R/IN/OUT 自动选择最优 Tile
调度策略	静态 Grid	Continuous Batching 感知路由

推理优化全景总结
#mermaid-svg-9B1HZ8tJ5jv17Sdn{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .error-icon{fill:#552222;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .marker{fill:#333333;stroke:#333333;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .marker.cross{stroke:#333333;}#mermaid-svg-9B1HZ8tJ5jv17Sdn svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-9B1HZ8tJ5jv17Sdn p{margin:0;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge{stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 path{fill:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 text{fill:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon--1{font-size:40px;color:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge--1{stroke:hsl(240, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth--1{stroke-width:17;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section--1 line{stroke:hsl(60, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 path{fill:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-0{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-0{stroke:hsl(60, 100%, 73.5294117647%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-0{stroke-width:14;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-0 line{stroke:hsl(240, 100%, 83.5294117647%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 path{fill:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-1{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-1{stroke:hsl(80, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-1{stroke-width:11;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-1 line{stroke:hsl(260, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 path{fill:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 text{fill:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-2{font-size:40px;color:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-2{stroke:hsl(270, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-2{stroke-width:8;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 line{stroke:hsl(90, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 path{fill:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-3{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-3{stroke:hsl(300, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-3{stroke-width:5;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-3 line{stroke:hsl(120, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 path{fill:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-4{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-4{stroke:hsl(330, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-4{stroke-width:2;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-4 line{stroke:hsl(150, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 path{fill:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-5{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-5{stroke:hsl(0, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-5{stroke-width:-1;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-5 line{stroke:hsl(180, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 path{fill:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-6{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-6{stroke:hsl(30, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-6{stroke-width:-4;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-6 line{stroke:hsl(210, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 path{fill:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-7{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-7{stroke:hsl(90, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-7{stroke-width:-7;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-7 line{stroke:hsl(270, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 path{fill:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-8{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-8{stroke:hsl(150, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-8{stroke-width:-10;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-8 line{stroke:hsl(330, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 path{fill:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-9{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-9{stroke:hsl(180, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-9{stroke-width:-13;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-9 line{stroke:hsl(0, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 polygon,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 path{fill:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 text{fill:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .node-icon-10{font-size:40px;color:black;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-edge-10{stroke:hsl(210, 100%, 76.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge-depth-10{stroke-width:-16;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-10 line{stroke:hsl(30, 100%, 86.2745098039%);stroke-width:3;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:lightgray;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .disabled text{fill:#efefef;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-root rect,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-root path,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-root circle,#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-root polygon{fill:hsl(240, 100%, 46.2745098039%);}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-root text{fill:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-root span{color:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .section-2 span{color:#ffffff;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .icon-container{height:100%;display:flex;justify-content:center;align-items:center;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .edge{fill:none;}#mermaid-svg-9B1HZ8tJ5jv17Sdn .mindmap-node-label{dy:1em;alignment-baseline:middle;text-anchor:middle;dominant-baseline:middle;text-align:center;}#mermaid-svg-9B1HZ8tJ5jv17Sdn :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} LLM Serving 优化栈
计算加速
Prefill
Fused RoPE
Operator Fusion
显存管理
KV Cache
Weights
多租户复用
Multi-LoRA ← 本节收束
Token-wise Routing
Weight Pool
Batched Heterogeneous Inference
系统调度
Continuous Batching
Prefix Caching
Speculative Decoding