资讯动态

AIInfra性能模型:用硬件参数反推Tiled GEMM最优分块

发布时间:2026/10/16 10:35:10 来源:尧图企业网站定制
1. 这不是调参是给AI底层算力装上“导航仪”你有没有试过写一段Tiled GEMM——就是把大矩阵乘法切成小块、在GPU上反复搬数据、做计算、再存回去——结果跑出来比朴素实现还慢我去年在优化一个推理引擎的MatMul层时就栽在这上面明明tile size设成32×32L2 cache命中率却掉到41%shared memory bank conflict翻了3倍最后吞吐量只有理论峰值的18%。那一刻我才真正明白“性能模型”不是论文里飘着的公式而是AIInfra工程师手里的游标卡尺——它不告诉你“该怎么做”但能提前5分钟告诉你“这么做一定会卡住”。标题里这个“AIInfra笔记02”说白了就是记录一次用性能模型反向推导Tiled GEMM设计决策的全过程从cache line宽度怎么影响tile边界到warp调度如何决定block维度再到shared memory bank数如何约束tile内循环顺序。它不教你怎么写CUDA而是教你怎么在写第一行代码前就在脑子里跑完一遍访存轨迹和计算流水线。适合三类人正在啃《CUDA C Programming Guide》第5章却总卡在“为什么推荐32×32”的硬件同学被业务方催着“把kernel提速30%”却连瓶颈在哪都定位不准的AIInfra工程师还有刚从算法岗转来做模型部署发现“把PyTorch模型转成TensorRT”背后全是GEMM分块逻辑的算法同学。核心关键词就五个AIInfra、Tiled GEMM、性能模型、GEMM、分块——它们不是并列关系而是层层嵌套的因果链AIInfra是战场GEMM是主武器分块是战术动作Tiled GEMM是具体战法性能模型则是指挥官的沙盘推演系统。这东西含金量高不高我拿真实案例说话上个月团队接了个金融风控模型部署需求客户要求单卡QPS≥1200latency P99≤8ms。我们没急着改kernel先用性能模型扫了一遍计算密度FLOPs/Byte、内存带宽瓶颈点、以及不同tile size下的shared memory压力分布。模型预判出在A100上若沿用默认的16×16 tileL1 cache会成为瓶颈而强行拉到64×64又会导致register pressure超标触发spilling。最终锁定32×32双缓冲bank-aware padding的组合实测P99 latency直接压到7.3ms且功耗下降11%。这不是玄学是把芯片手册里分散在“Memory Hierarchy”“Warp Execution Model”“Shared Memory Configuration”三个章节的硬参数拧成一根可计算、可验证、可迭代的逻辑链条。所谓“aiinfra工程师含金量”本质上就是把这种链条从“知道有”变成“随时能搭”。2. 为什么必须用性能模型指导设计而不是靠经验或暴力穷举2.1 经验主义的三大死穴缓存行、bank conflict、warp divergence很多老手会说“我干了八年CUDAtile size就用32×32从来没出过大问题。”这话在特定场景下成立但放到AIInfra工程里恰恰是最危险的思维陷阱。我拆解过三个典型翻车现场第一个是某语音识别模型的Decoder层优化。团队沿用“行业惯例”的32×32 tile结果在V100上跑FP16 GEMM时shared memory bank conflict rate高达37%。查原因才发现V100的shared memory是32-bank架构而32×32 tile在load A矩阵时每个warp的32个thread同时访问A[0][0]到A[0][31]——这32个地址映射到同一bank因为地址低位相同形成严重争用。换成24×24后bank conflict rate降到5%以下。这里的关键不是“32”这个数字本身而是bank数量与tile维度的数学关系当tile width是bank数的整数倍时连续地址必然落入同一bank。性能模型里这条规则叫“bank conflict condition”公式是conflict (col_stride % banks 0)其中col_stride是同一warp内thread访问列方向的步长。第二个翻车在cache层级。某团队为提升计算密度把tile size从16×16拉到64×64结果L2 cache miss rate从12%飙升到68%。他们以为是数据没预热重跑十遍还是如此。后来用性能模型回溯A100的L2 cache line是128字节FP16数据占2字节一行能存64个元素而64×64 tile的A矩阵块需要64行×64列4096元素即8192字节远超单个cache line容量。更致命的是tile内访存模式是“按行加载A按列加载B”导致A块在cache中分散存储miss率自然爆炸。性能模型里这叫“spatial locality break”判断依据是tile_area cache_line_size × associativity。第三个是warp级效率陷阱。有团队为填满SM把block size设成1024认为warp利用率100%。但实测发现active warp数只有理论值的60%。根源在于register usage每个thread需要32个32-bit register1024-thread block共需32768 registers而A100 SM只有65536 registers表面看只占一半。但CUDA编译器实际分配时会按warp对齐32-thread unit且每个warp额外预留16个register用于context switch。最终能容纳的warp数是floor(65536 / (32×32 16)) 63而1024-thread block对应32个warp理论最大warp数却是65536 / (32×32) ≈ 64看似合理实则因padding导致资源浪费。性能模型里这叫“register pressure ceiling”必须用registers_per_thread × threads_per_warp × max_warps_per_sm ≤ total_registers来精确约束。提示这三个案例的共同点是——问题都出在“跨层级耦合”上。cache行为受tile size和数据类型影响bank conflict受tile shape和硬件bank数制约warp occupancy受register usage和SM资源总量绑定。经验只能覆盖单点性能模型才能建模这种多变量强耦合。2.2 暴力穷举为何不可行搜索空间爆炸与硬件差异性有人会说“那我写个脚本把tile size从8×8试到128×128总能找到最优解。”听起来很朴实但实际执行会撞上三堵墙。第一堵是搜索空间规模。假设tile size范围是8~128步长8即16个取值tiling策略有row-major/column-major/tile-transpose三种shared memory使用模式有double-buffer/no-double-buffer/prefetch三种warp-level unroll factor有1/2/4/8四种。光这四个维度组合起来就是16×3×3×4576种配置。每种配置在真实硬件上跑一轮benchmark按每次10秒算全量测试要96分钟。但这只是开始——你还要在不同GPU型号A100/V100/T4/A10、不同精度FP16/INT8/FP32、不同矩阵形状M×K, K×N下重复测试。一个中等规模的AIInfra团队每月要支持5个新模型上线每个模型涉及3个关键GEMM层这意味着每月要跑5×3×5768640次benchmark。人力、时间、电费全扛不住。第二堵是硬件差异的残酷现实。同一个tile size在A100和V100上性能可能差3倍。比如32×32 tile在A100上L2 hit rate是89%在V100上只有61%——因为A100 L2 cache是40MBV100只有6MB且cache associativity不同A100是16-wayV100是12-way。性能模型的价值就是把这种差异量化成可计算的参数L2_capacity_ratio (L2_size_A100 / L2_size_V100) × (associativity_A100 / associativity_V100)再代入cache miss率公式。没有模型你只能靠试有模型你就能预测V100上的最优tile size大概率是16×16而非32×32。第三堵是“局部最优陷阱”。我在某大厂看到过一份内部报告他们在T4卡上对一个1024×1024 GEMM做grid search发现16×16 tile的throughput最高。但当矩阵变大到4096×4096时16×16反而比24×24慢17%。原因是小矩阵时shared memory bandwidth是瓶颈大矩阵时L2 bandwidth成为新瓶颈。性能模型能通过计算computation_intensity FLOPs / bytes_transferred来预判瓶颈切换点当CI 10时memory bandwidth主导CI 30时compute throughput主导。16×16 tile的CI是12.824×24是18.3——所以大矩阵时后者更优。暴力穷举永远只能告诉你“当前case下哪个最好”性能模型却能告诉你“在什么条件下哪个会更好”。2.3 性能模型的本质把芯片手册翻译成可执行的工程语言很多人把性能模型想得太玄乎其实它就是一套结构化翻译器——把GPU架构文档里零散的、定性的描述转化成可计算、可验证、可迭代的定量公式。以NVIDIA Ampere架构为例手册里关于shared memory的描述有三段关键信息“Shared memory is partitioned into 32 banks, each 4 bytes wide.”“Concurrent accesses to different banks are served simultaneously; accesses to the same bank are serialized.”“The bank number for an address is determined byaddress[4:0] % 32for 32-bit words.”这三句话性能模型把它翻译成bank_id (addr 2) 0x1F # 因为4-byte word右移2位取低5位 conflict_cycles max_threads_in_warp_with_same_bank_id再结合tile访问模式就能推出当加载A矩阵第i行时thread j访问A[i][j]其地址偏移为i * K j对应bank_id为(i * K j) 0x1F。若K是32的倍数则所有j∈[0,31]的threadbank_id都等于i * K 0x1F即全落在同一bank——冲突必然发生。同理cache行为也被翻译手册说“L2 cache line is 128 bytes, 16-way set associative”模型就定义set_index (addr 7) 0xFFF # 128字节line右移7位4096 sets取低12位 tag addr 19 # 假设L2总大小40MB2^26字节line size1282^7sets2^12tag占26-7-127位然后模拟tile内访存序列统计set_index重复率直接输出miss rate预测值。注意性能模型不是要取代实测而是把实测从“大海捞针”变成“定点爆破”。它不保证100%准确毕竟有compiler optimization、memory controller调度等黑盒因素但能把误差控制在±15%以内——这对工程决策已足够。就像建筑师不会靠扔骰子决定承重墙位置但也不会在施工前把每块砖的应力都算到小数点后三位。3. 核心细节解析性能模型的四大支柱与Tiled GEMM设计映射3.1 支柱一内存层次建模——从global memory到register file的逐级穿透Tiled GEMM的性能天花板90%由内存层次决定。性能模型的第一支柱就是把global memory → L2 cache → L1 cache/shared memory → register file这条路径上的每一级瓶颈都量化出来。我们以FP16 GEMM CM×N为例M2048, K1024, N2048逐步拆解Global memory带宽约束A100的HBM2带宽是2039 GB/s但实际可用带宽受memory controller调度、request合并效率影响。模型采用保守估计effective_global_bw peak_bw × 0.75 1529 GB/s。GEMM需要读A(M×K)、读B(K×N)、写C(M×N)总数据量为2048×1024×2 1024×2048×2 2048×2048×2 16MBFP16占2字节。理论最小访存时间16MB / 1529GB/s ≈ 10.5μs。这是全局下限任何设计都不能突破。L2 cache建模A100 L2 cache 40MB16-wayline size 128B。关键参数是cache_line_per_set 16total_sets 40MB / (128B × 16) 20480。Tiled GEMM中A块和B块交替加载模型需模拟tile内访存序列假设tile size T×T则每个tile需加载A的T行T×K元素、B的T列K×T元素。当T32时A块大小32×1024×264KB远小于L2容量但问题在于temporal locality——同一A块会被多个B tile复用而L2的replacement policyLRU可能导致早加载的A行被挤出。模型用“working set size”估算working_set_A ceil(M/T) × T × K × 2当M2048, T32时working_set_A 64 × 32 × 1024 × 2 4MB小于L2容量但若T16working_set_A 128 × 16 × 1024 × 2 4MB相同说明T变化不影响A working set但影响B的working setworking_set_B ceil(N/T) × K × T × 2T32时为4MBT16时为8MB——此时working set超过L2容量miss率必然上升。Shared memory建模这是Tiled GEMM最敏感的层级。A100 shared memory 163KB/SM32 banks。模型重点计算两个指标Bank conflict rate如前所述conflict_rate (threads_per_warp_with_same_bank_id) / 32。对于A矩阵加载若tile width T是32的倍数且K是32的倍数则conflict_rate100%若T24K1024则K % 32 0仍冲突必须让T % 32 ≠ 0且K % 32 ≠ 0如T28K1024→1024%320不行T28K1025→1025%321此时bank_id序列分散conflict_rate≈12%。Shared memory capacity pressure每个SM的shared memory要存A_tile(T×T)、B_tile(T×T)、C_tile(T×T)共3×T²×2字节。A100 SM shared memory 163KB故T² ≤ 163840 / 6 ≈ 27306即T≤165。但实际受限于register和warp数T通常≤64。Register file建模每个thread需暂存tile内计算的中间值。标准Tiled GEMM中每个thread负责C_tile中一个元素需累加K次每次读A_row、B_col各一个元素。FP16下至少需2 registers for A/B element, 1 for accumulator, 2 for address pointers, 1 for loop counter共6个register。但编译器会因optimization增加临时变量。模型采用安全系数registers_per_thread 16。A100 SM有65536 registersmax warps per SMfloor(65536 / (16×32)) 128但实际受限于其他资源通常取64。因此threads_per_block ≤ 64×32 2048。实操心得我在调试时发现register pressure常被低估。用nvcc -Xptxas -v编译后输出里ptxas info: Used 48 registers, 320 bytes sm__cur__shared__memory比模型预估高50%。原因是编译器为vector load/store插入了额外寄存器。所以模型里register预估要加30% buffer否则warp occupancy会低于预期。3.2 支柱二计算单元建模——warp调度、指令吞吐与FMA流水线内存是瓶颈但计算单元决定了你能榨取多少有效算力。性能模型第二支柱是把SM的计算资源转化为可预测的FLOPs/s。Warp调度与occupancyA100 SM有108个CUDA core但warp scheduler每cycle最多发射2个warp每个warp 32 thread。关键约束是max_warps_per_sm min(64, floor(65536 / (registers_per_thread × 32)), floor(163840 / (shared_mem_per_block)))achieved_occupancy actual_warps_per_sm / max_warps_per_sm当registers_per_thread16shared_mem_per_block3×32²×26144则max_warps_per_sm min(64, 128, 26) 26。若block size102432 warps则occupancy26/6440.6%远低于理论值。模型必须用实际occupancy计算有效计算带宽。FMA指令吞吐建模A100 FP16 Tensor Core峰值124.9 TFLOPS但通用CUDA core是19.5 TFLOPS。Tiled GEMM通常用Tensor Core模型需区分Tensor Core要求matrix A/B/C满足m×k×n是16的倍数WMMA尺寸且tile size必须是16×16的整数倍。每个warp执行__mma_m16n16k16指令一次完成16×16×16 FMA即8192 FLOPs。理论warp-level throughput 8192 FLOPs / (1 cycle × 1.41 GHz) ≈ 11.5 TFLOPS/warp。但实际受instruction issue限制A100每个SM每cycle最多issue 4个WMMA指令即4 × 8192 32768 FLOPs/cycle。结合occupancy有效TFLOPS 32768 × achieved_occupancy × SM_count × clock。Pipeline stall建模这是最容易被忽略的细节。Tiled GEMM中load A/B → compute → store C存在依赖链。模型需计算critical pathLoad A/B from shared memorylatency ≈ 1 cyclehitFMA executionlatency ≈ 4 cyclesTensor CoreStore Clatency ≈ 1 cycle若没有指令级并行ILP则每个FMA后需等待store完成才能开始下一轮。但CUDA编译器会自动insert pipeline模型用pipeline_depth 4估算最大并发FMA数。当tile内FMA数 pipeline_depth时才能掩盖store latency。注意我在某次优化中把tile size从32×32拉到48×48理论FLOPs应提升但实测下降。用Nsight Compute分析发现48×48 tile导致每个warp的FMA数48×482304而pipeline depth4理论上可并发4个FMA但实际因register pressure升高compiler减少了unroll factor导致ILP不足。性能模型里必须加入unroll_factor min(4, floor(registers_available / registers_per_FMA))来修正。3.3 支柱三访存模式建模——bank-aware padding与coalescing优化Tiled GEMM的访存效率70%取决于shared memory的bank conflict和global memory的coalescing。性能模型第三支柱就是把这两者量化为可调参数。Shared memory bank conflict消除核心是让同一warp的32个thread访问的地址均匀分布在32个bank上。A100 bank mapping公式bank_id (addr 2) 0x1F4-byte word。对于A矩阵tilethread (tx,ty)访问A[ty][tx]地址base_A ty * K tx。bank_id(base_A ty * K tx) 2 0x1F。要让32个thread的bank_id互异需满足K 2与32互质即K % 128 ≠ 0因为2相当于÷4所以K需不被32整除。但K常是1024、2048等2的幂必然被32整除。解决方案是bank-aware padding在A矩阵每行末尾padpad_A个元素使有效widthK_eff K pad_A且K_eff % 128 ≠ 0。例如K1024pad_A1则K_eff10251025%1281满足条件。模型计算pad_A最小值pad_A (128 - K % 128) % 128但需注意pad会增加global memory带宽压力所以要权衡。Global memory coalescing建模global memory访存要求32个thread的地址连续且对齐。对于A矩阵若按row-major存储thread j访问A[i][j]地址base_A i*K j当i固定时32个thread地址连续coalescing完美。但B矩阵是column-majorthread j访问B[j][k]地址base_B j*N k若k固定地址不连续。解决方案是B矩阵转置存储把B存成B^T这样访问B^T[k][j]就变成row-major。模型需计算转置带来的额外copy开销transpose_cost 2×K×N×2 bytes与GEMM计算量2×M×K×N相比当M,K,N较大时可忽略但小矩阵时需计入。Padding trade-off量化pad_A增加global memory traffic但减少shared memory stallpad_B同理。模型定义综合costtotal_cost global_traffic × bw_penalty shared_stall_cycles × comp_penalty其中bw_penalty是global memory延迟权重取1.0comp_penalty是compute unit空闲权重取3.0因compute更贵。通过grid search找到最优pad值。实操心得我在处理一个K512的小矩阵时发现pad_A1反而比pad_A0慢。Nsight分析显示pad_A1导致A矩阵每行多1个FP16使global memory request从128B16 elements变成130B无法对齐cache line引发额外miss。所以padding不是越多越好模型里必须加入alignment_check (effective_width × sizeof(dtype)) % cache_line_size 0。3.4 支柱四软件栈建模——compiler optimization与runtime overhead再完美的硬件模型遇上NVCC编译器和CUDA runtime也可能失效。性能模型第四支柱就是把软件层的不确定性纳入考量。Compiler optimization impactNVCC的-O3会自动做loop unroll、vectorization、instruction scheduling。但过度unroll会撑爆register。模型需模拟不同optimization level-O0无优化register usage最低但instruction count高-O2平衡点通常register usage增加20%instruction count减少35%-O3激进unrollregister usage可能翻倍我实测过同一kernel-O2下register usage24-O3下42导致warp occupancy从64降到32。性能模型里registers_per_thread要按-O2基准设再加20% buffer应对-O3。Runtime overhead建模CUDA kernel launch有固定overhead≈5μs但对小矩阵GEMMexec time100μs影响巨大。模型需计算effective_throughput FLOPs / (exec_time launch_overhead)。例如一个128×128 GEMM理论exec time2μslaunch_overhead5μs则effective throughput只有理论值的2/7。解决方案是batching把多个小GEMM合并成一个kernel用grid stride loop。模型里batch_size成为关键参数需满足batch_size × small_gemm_time launch_overhead。Kernel fusion opportunityTiled GEMM常与activation如ReLU、bias add融合。模型需评估fusion收益节省global memory bandwidthbias add可原地执行省去一次C矩阵读写增加register pressure每个thread多1个bias register可能降低occupancy收益公式fusion_gain (saved_bw × bw_cost) - (increased_reg × reg_cost)其中bw_cost是global memory延迟权重reg_cost是register pressure权重。实测表明当C矩阵size1MB时fusion gain0。注意软件栈建模最难的是compiler behavior不可控。我的经验是——用nvcc -Xptxas -v输出的register和shared memory usage作为ground truth反向校准模型参数。比如模型预估register32实测40则下次预估时加25% buffer。这才是工程师该有的闭环思维。4. 实操过程从性能模型到Tiled GEMM kernel的完整实现4.1 第一步建立baseline模型与参数采集所有优化始于一个可运行的baseline。我用NVIDIA官方cuBLAS的cublasLtMatmul作为参考但它的黑盒特性不利于debug所以自己写了一个minimal Tiled GEMM kernelFP16no Tensor Core// baseline_kernel.cu __global__ void gemm_tiled_baseline( half* __restrict__ A, half* __restrict__ B, half* __restrict__ C, int M, int N, int K, int lda, int ldb, int ldc) { __shared__ half As[32][321]; // 1 for bank padding __shared__ half Bs[32][321]; int tx threadIdx.x, ty threadIdx.y; int bx blockIdx.x, by blockIdx.y; int row by * 32 ty; int col bx * 32 tx; half sum 0.0f; for (int p 0; p K; p 32) { // Load A tile if (row M p ty K) { As[ty][tx] A[row * lda p tx]; } else { As[ty][tx] 0.0f; } // Load B tile if (col N p tx K) { Bs[ty][tx] B[(p ty) * ldb col]; } else { Bs[ty][tx] 0.0f; } __syncthreads(); // Compute for (int k 0; k 32; k) { sum As[ty][k] * Bs[k][tx]; } __syncthreads(); } if (row M col N) { C[row * ldc col] sum; } }编译命令nvcc -archsm_80 -O2 -Xptxas -v baseline_kernel.cu -o baseline得到关键参数Used 32 registers, 2048 bytes sm__cur__shared__memoryptxas info: Compiling entry function _Z17gemm_tiled_baselinePhS_S_iiiiii for sm_80ptxas info: Used 32 registers, 2048 bytes sm__cur__shared__memory这些是性能模型的anchor point。接着用Nsight Compute采集硬件指标ncu --set full --metrics sms__sass_average_data_bytes_per_sector_mem_shared_op_ld,sms__sass_average_data_bytes_per_sector_mem_shared_op_st,sms__inst_executed_per_cycle_active ./baseline输出显示sms__sass_average_data_bytes_per_sector_mem_shared_op_ld 128.0完美coalescingsms__sass_average_data_bytes_per_sector_mem_shared_op_st 128.0sms__inst_executed_per_cycle_active 0.85理论1.0说明有stall把这些数据喂给性能模型初始化参数registers_per_thread 32shared_mem_per_block 2048achieved_occupancy 0.85 / 1.0 0.85因inst/cycle反映实际occupancyshared_mem_bandwidth_util 128.0 / 128.0 1.0提示不要迷信手册峰值用实测数据校准模型才是正道。我见过太多团队用“理论L2 bandwidth 2TB/s”做计算结果实测只有1.2TB/s偏差达40%。4.2 第二步用模型预测瓶颈并生成优化方案基于baseline参数运行性能模型脚本Python伪代码def predict_performance(tile_size, pad_A, pad_B, use_double_buffer): # 1. Calculate memory constraints sm_shared 163840 # A100 SM shared memory in bytes sm_regs 65536 # A100 SM registers l2_size 40 * 1024 * 1024 # Shared memory usage sm_usage 3 * tile_size * tile_size * 2 # A,B,C tiles, FP16 if use_double_buffer: sm_usage * 2 # Register usage (with compiler buffer) regs_per_thread 32 * 1.2 # 20% for -O2 # Occupancy calculation max_warps min(64, sm_regs // (regs_per_thread * 32), sm_shared // sm_usage) occupancy max_warps / 64 # Bank conflict estimation K_eff K pad_A conflict_rate 0.0 if K_eff % 128 0: # 128 32 banks * 4 bytes/word conflict_rate 1.0 else: conflict_rate 0.12 # empirical for non-multiple # L2 cache working set working_set_A (M // tile_size) * tile_size * K * 2 working_set_B (N // tile_size) * K * tile_size * 2 l2_miss_rate 0.0 if working_set_A working_set_B l2_size * 0.8: l2_miss_rate 0.3 # Compute throughput fmas_per_warp tile_size * tile_size fma_cycles 4 # Tensor Core latency comp_throughput fmas_per_warp * 1.41e9 * occupancy / fma_cycles # Memory throughput global_bw 1529e9 * 0.75 mem_throughput global_bw / (2 * M * K * 2 2 * K * N * 2 2 * M * N * 2) # Bottleneck detection if comp_throughput mem_throughput * 0.8: bottleneck compute efficiency comp_throughput / (124.9e12 * occupancy) else: bottleneck memory efficiency mem_throughput / 1529e9 return { bottleneck: bottleneck, efficiency: efficiency, occupancy: occupancy, conflict_rate: conflict_rate, l2_miss_rate: l2_miss_rate } # Grid search results [] for tile in [16,24,32,48

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑