SRAM‑HBM 分块调度：大模型真正的瓶颈从来不是算力

马斯克在近期访谈中提到，AI 的最大约束是内存，而不是原始算力。很多文章还在比拼 TFLOPS 算力指标，但真正卡死长上下文、训练吞吐、推理延迟的，是 HBM 与片上 SRAM 之间的数据搬运开销CSDN博…。 很多人有一个误区：GPU 计算单元越强，大模型跑的就越快。现实恰恰相反，绝大多数大模型训练、长文本推理场景下，张量核心经常处于空闲等待状态。GPU 算的很快，但数据搬不过来，算力再高也发挥不出来，这就是内存墙。GPU 硬件天然分成两级存储。 HBM，高带宽显存，容量很大，几十 GB 级别，模型权重、KV 缓存、大矩阵全部存在这里。带宽很高，但和芯片内部比依然很慢，数据读写要跨芯片互联，访问延迟高。 SRAM，片上共享内存，就在 SM 计算单元旁边，带宽可以达到几十 TB/s，比 HBM 高出一个数量级，计算就在这里完成。致命短板是容量极小，单块 SM 的 SRAM 只有几百 KB，整块 GPU 加起来也就几十 MB，放不下完整的 Q、K、V 大矩阵。硬件矛盾就摆在这里：计算必须在 SRAM 内部完成，但是 SRAM 装不下完整矩阵；全部数据又只能保存在 HBM。原生实现注意力的时候，会把 Q、K 全部从 HBM 读取进来，在 SRAM 算出 QKᵀ，把巨大的注意力中间矩阵 S 写回 HBM，再读回来做 softmax，再乘 V，再次读写 HBM。 中间巨大的 S 矩阵反复读写 HBM，IO 次数爆炸。尤其是长上下文，token 数量 N 变大，S 矩阵是 N×N 规模，读写量呈平方上涨。这就是长文本越跑越慢的根本原因，不是计算变多，是内存搬运暴涨。分块调度，也就是 Tiling，核心思路很直白：既然 SRAM 装不下完整大矩阵，那就把大矩阵切出一小块一小块的 tile。每次只把一小部分 Q、K 从 HBM 加载到 SRAM，在片上完成矩阵运算，算完立刻把结果输出，不把完整的中间矩阵 S 存到 HBM 里面去CSDN博…。标准注意力的流程： 读取全部 QKV → HBM 生成完整 S 矩阵 → 写回 HBM → 读取 S 做 softmax → 再次读取乘 V 输出。大量中间张量反复落地 HBM。分块调度后的流程： 1、按照 SRAM 的实际可用容量，设定 tile 分块大小，不能超过片上内存上限。 2、循环分批，只把对应分块的 Q 切片、K 切片从 HBM 搬运进 SRAM。 3、在 SRAM 内部完成 QKᵀ矩阵乘法、online‑softmax，全程中间结果不写回 HBM。 4、直接在片上和 V 分块做加权求和，把最终输出 O 的小块写回 HBM。 5、循环迭代全部分块，拼接得到完整输出。关键点，这不是近似算法。FlashAttention 这套分块计算，数学结果和标准注意力完全等价，没有信息丢失，属于 Exact Attention 精准注意力，不靠稀疏、截断做妥协CSDN博…。反向传播还有一个关键手段：重计算 Recomputation。 原生反向传播，需要保存前向产生的全部中间矩阵用来求梯度，显存占用巨大。分块调度不保存庞大中间张量，反向阶段重新再跑一遍前向分块计算，用少量额外计算开销，换取巨量 HBM 读写与显存节省。用算力换 IO，在内存受限场景下收益极高。很多人分不清 MoE 瓶颈和注意力访存瓶颈，这里做一个区分。 注意力的主要矛盾：IO 瓶颈，开销大头来自 HBM‑SRAM 之间反复搬运张量，计算量反而不大，属于内存受限算子。优化核心就是 Tiling 分块，减少 HBM 访问次数。 MoE 混合专家的主要矛盾：稀疏动态调度 + 多卡通信。每个 token 路由到不同专家，访问地址不规则，负载不均衡，除了显存搬运，还引入跨卡专家权重通信开销。MoE 也会用到 SRAM‑HBM 分块，但是额外多一层路由调度、专家负载均衡的复杂度。这也是为什么 MoE 的 CUDA Kernel 优化难度远高于普通注意力。放到真实业务场景，4K 上下文感受不到差距，一旦拉到 128K、百万 token 长文本，分块调度的价值彻底拉开差距。N 变大之后，原生注意力的 HBM 读写量 N² 爆炸，GPU 算力跑不满，大部分时间在等待数据搬运。经过 Tiling 分块调度，把 HBM 访问复杂度从 O (N²) 压到 O (N)，同样硬件，长文本吞吐可以提升 2‑4 倍。但分块调度不是万能，工程上有大量现实约束。 tile 尺寸不能随便设置，tile 开太大，超过 SRAM 容量直接报错；tile 切的太小，加载‑计算的切换变多，调度开销上涨。要根据不同 GPU 的 SRAM 大小，HBM 带宽，头维度做动态调参。国产 NPU、不同架构 GPU，片上 SRAM 容量不一样，一套分块参数不能直接照搬，这也是很多算子移植踩坑的地方。算子融合也要配合分块调度，把加载、矩阵乘、softmax、输出写入融合进同一个 CUDA Kernel，减少 Kernel 启动开销。现在回看行业现状：很多评测只宣传 GPUTFLOPS 算力，很少去讲 SRAM 容量、HBM 带宽、分块调度实现。同样一张卡，是否开启 FlashAttention，长文本性能可以差几倍。算力再强，如果内存层次调度做不好，张量核心大部分时间空闲。大模型硬件竞争已经走到新阶段：比拼的不只是峰值算力，更看硬件内存层级，以及 CUDA 算子能不能用好 SRAM‑HBM 分块调度，尽可能减少昂贵的片外内存访问。长上下文能力，本质上一半来自架构创新，另一半来自对内存墙的工程对抗。 本篇锚点公式：

(O=\\text{FlashAttention}(Q,K,V)=\\text{Tiling}\\big(\\text{softmax}(\\frac{QK^\\mathrm{T}}{\\sqrt{d}})V\\big))

核心约束：tile_size ≤ SRAM 可用容量，尽量降低 HBM 读写次数。

下面是极简可跑的FlashAttention 分块调度伪代码，模拟 SRAM‑HBM Tiling 逻辑，不是真正 CUDA kernel，用来放在 CSDN 文章做演示，看懂分块迭代逻辑。 只演示前向，Online‑Softmax 分块计算，Exact Attention，无近似。

import torch

def flash_attention_tiling(Q, K, V, sram_tile_size:int):
\”\”\”
模拟 SRAM‑HBM 分块调度(Tiling)
Q,K,V: [seq_len, head_dim]
sram_tile_size: 单次能装进SRAM的token数量
真实CUDA：tile受SM的SRAM物理大小硬限制
\”\”\”
seq_len, d = Q.shape
O = torch.zeros_like(V) # 输出存HBM

# 沿着K/V维度分块，每次加载一小块进入SRAM
for j in range(0, seq_len, sram_tile_size):
j_end = min(j + sram_tile_size, seq_len)
# ———- HBM -> SRAM：只加载tile切片 ———-
K_j = K[j:j_end, :]
V_j = V[j:j_end, :]

# 沿着Q维度分块
for i in range(0, seq_len, sram_tile_size):
i_end = min(i + sram_tile_size, seq_len)
Q_i = Q[i:i_end, :]

# ———- 在SRAM内部做计算，不生成完整N×N S矩阵写回HBM ———-
S_ij = Q_i @ K_j.T / (d ** 0.5)
P_ij = torch.softmax(S_ij, dim=–1)
O_i_part = P_ij @ V_j

# 结果小块写回HBM，中间S_ij、P_ij随SRAM释放，不落地显存
O[i:i_end, :] += O_i_part
return O

上面是概念版 Python 模拟代码。真实 CUDA 算子不会在全局显存生成完整的 N×N 注意力矩阵，每一轮 tile 只把少量切片从 HBM 搬运到片上 SRAM，矩阵乘、softmax 全部在片上完成。中间大矩阵直接丢弃，不写回 HBM，以此把访存开销压下来。 关键现实约束： sram_tile_size