Skip to content

03. GPU Architecture and Memory | GPU 物理架构与内存层级

难度: Hard | 环境: CPU-first | 标签: 硬件系统, GPU, 内存层级 | 目标人群: 需要理解 GPU 性能约束的学习者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行:

Open In ColabOpen In Studio (国内推荐:魔搭社区免费实例)


本节导读

本节把 GPU 计算单元、内存层级和数据搬运放在同一条分析链上,帮助判断一段模型计算主要受计算吞吐、显存带宽还是显存容量限制。这里的工作集(working set)指当前计算阶段需要频繁访问或暂时驻留的数据集合;学习时依次检查计算单元、数据来源、工作集大小和访问路径,再判断瓶颈来自哪类资源。 本节沿着“计算单元 → 内存层级 → 工作集 → 瓶颈判断”展开:先区分计算吞吐、带宽和容量,再用工作集分析 Attention 等计算阶段需要搬运和暂存哪些数据。完成后,你应能根据计算量、数据量和访问路径,初步判断一个算子更可能受哪类资源限制。

关键词: Tensor Core, SRAM, HBM

本节概念关系


前置阅读

导语: 先回顾 dtype、参数量和 FLOPs 的数量级关系,再从 GPU 执行层级和内存层级出发,判断计算、带宽与容量压力分别来自哪里。


Q1:为什么 GPU 计算能力提高后,Transformer 不一定同比变快?

点击展开查看解析

GPU 规格中的峰值算力只描述计算单元最快能完成多少运算,不表示数据已经及时送到计算单元。Transformer 由多个重复的计算模块组成,Attention 和 MLP 都会产生输入、权重、激活与中间结果;如果这些数据的读取、写回和复用跟不上,计算单元就会等待,峰值算力也无法转化为同等的程序加速。 因此,判断一个算子受什么限制,可以先问两个问题:它要完成多少计算,以及要从内存搬运多少数据。后面的 FLOPs、Bytes 和算术强度,就是把这两个问题量化。

不同 GPU 代际会在 dtype 支持、片上存储、异步搬运和互连能力上做不同取舍。这里先比较这些能力如何改变算子执行,再把具体设备规格作为数量级参考。

可以先用下面的三分法定位现象:

现象主要限制常见表现
数据装不下显存容量OOM,需要 offload、分片或缩短序列
数据搬得慢显存带宽或访存路径计算单元等待,带宽接近上限
计算做得慢计算吞吐CUDA Core 或 Tensor Core 接近饱和

Q1小验证:先区分计算路径和数据供给

先观察计算路径与内存路径各自承担什么工作,再用后面的层级表和算术强度判断具体瓶颈。

python
from typing import Dict

# 教学代理值:用于比较数量级,不代表具体 GPU 的实测规格。
TEACHING_BANDWIDTH = {'shared_memory': 19e12, 'l2_cache': 1.5e12, 'hbm': 1.5e12}
COMPUTE_PATHS = {'cuda_core': '线程级标量 / 向量计算', 'tensor_core': '矩阵乘加与低精度路径'}


def bytes_to_gb(bytes_val: float) -> float:
    """把字节转换为十进制 GB;只用于教学展示。"""
    if bytes_val < 0:
        raise ValueError('bytes_val 不能为负数')
    return bytes_val / 1e9


def analyze_memory_hierarchy() -> Dict[str, Dict[str, float]]:
    """计算内存层级的带宽数量级和 1KB 传输时间代理值。"""
    result = {}
    for mem_type, bandwidth in TEACHING_BANDWIDTH.items():
        result[mem_type] = {'bandwidth_tb_s': bandwidth / 1e12, 'time_for_1kb_ns': 1024 / bandwidth * 1e9}
    return result


mem_analysis = analyze_memory_hierarchy()
print('计算路径:', COMPUTE_PATHS)
for name, stats in mem_analysis.items():
    print(f"{name:15s}: {stats['bandwidth_tb_s']:6.1f} TB/s, 1KB transfer proxy: {stats['time_for_1kb_ns']:6.2f} ns")
def test_gpu_memory_practice():
    assert set(COMPUTE_PATHS) == {'cuda_core', 'tensor_core'}
    assert mem_analysis['shared_memory']['bandwidth_tb_s'] > mem_analysis['hbm']['bandwidth_tb_s']
    print('✅ 03 hardware path tests passed')


test_gpu_memory_practice()

数量级速览

这一张表用于比较硬件演进中的具体变化:计算路径、低精度支持、内存带宽和互连能力如何影响算子执行。具体带宽、算力、互连和 dtype 支持取决于 SKU、驱动和平台,应以对应硬件文档为准。

代际计算路径变化内存 / 互连变化对本节的意义
V100Tensor Core 开始用于低精度矩阵乘加GPU 内存系统继续服务大规模张量计算理解 CUDA Core 与 Tensor Core 的分工
A100BF16、TF32 等训练路径更成熟HBM 与片上缓存能力增强理解计算吞吐和数据供给的协同
H100FP8 与异步数据搬运能力增强HBM 和 GPU 互连继续提升理解低精度、带宽和搬运机制
Blackwell更低精度与更强矩阵计算路径互连与数据搬运能力继续增强观察硬件演进如何改变算子优化空间

Q2:Transformer 中的矩阵计算如何使用 Tensor Core?

点击展开查看解析

Transformer 的 Attention 和 MLP 都包含矩阵乘法。Tensor Core 是专门提高矩阵乘加吞吐的计算单元,更适合回答“矩阵乘法算得够不够快”,而不是直接回答“整个 Transformer 模块算得够不够快”。

CUDA Core 负责更通用的线程级计算;Tensor Core 主要服务于 FP16、BF16、FP8 等低精度数据类型下的矩阵乘加。具体支持的 dtype、数据块形状和吞吐取决于 GPU 架构,不能用一个固定的规格概括所有设备。

Tensor Core 只覆盖矩阵乘加这一段路径;数据准备、布局转换、Softmax 和结果写回仍然需要其他执行单元与内存层级配合。因此,矩阵乘法具备 Tensor Core 候选路径,不等于整个 Transformer 模块都会按同样比例加速。 本节小验证只比较不同矩阵形状和 dtype 对 FLOPs、输入输出规模以及 Tensor Core 候选路径的影响。Attention 工作集与序列长度的关系放到 Q4 单独观察。

可以把 GPU 的执行层级先简化为:

层级作用
SM调度线程块并组织片上计算资源
Warp以一组线程为单位执行指令
CUDA Core执行通用线程级计算
Tensor Core执行特定形状的矩阵乘加
### Q2小验证:矩阵形状与 Tensor Core 候选路径

先比较矩阵形状、dtype 和 FLOPs,再判断它是否具备低精度矩阵乘法的 Tensor Core 候选路径;这仍是机制模拟,不是硬件吞吐测试。

python
def estimate_matmul_workload(m, n, k, dtype='bf16'):
    """估算矩阵乘法工作量并标记 Tensor Core 候选路径。"""
    if min(m, n, k) <= 0:
        raise ValueError('矩阵维度必须为正数')
    if dtype not in {'fp16', 'bf16', 'fp8', 'fp32'}:
        raise ValueError('unsupported dtype')
    return {
        'flops': 2 * m * n * k,
        'input_elements': m * k + k * n,
        'output_elements': m * n,
        'dtype': dtype,
        'tensor_core_candidate': dtype in {'fp16', 'bf16', 'fp8'},
    }

for shape, dtype in [((1024, 1024, 4096), 'bf16'), ((1024, 1024, 4096), 'fp32'), ((2048, 4096, 4096), 'fp16')]:
    print(shape, dtype, estimate_matmul_workload(*shape, dtype))


def test_matmul_workload():
    low_precision = estimate_matmul_workload(1024, 1024, 4096, 'bf16')
    fp32 = estimate_matmul_workload(1024, 1024, 4096, 'fp32')
    assert low_precision['flops'] == fp32['flops']
    assert low_precision['tensor_core_candidate'] is True
    assert fp32['tensor_core_candidate'] is False
    print('✅ Tensor Core candidate tests passed')


test_matmul_workload()

Q3:GPU 的内存层级如何影响数据搬运?

点击展开查看解析

把内存层级先看成“离计算单元的距离不同”:越靠近计算单元,通常访问更快、容量更小;越远,容量更大,但搬运代价也更明显。优化的关键不是让所有数据都停在最快的层级,而是让会被重复使用的数据尽量靠近计算单元。

当一个算子每搬运很多字节只完成少量计算时,更容易受到全局显存带宽限制。可以用算术强度把这个判断量化:

text
Arithmetic Intensity = FLOPs / Bytes

算术强度较低,说明每搬一次数据完成的计算较少,更容易受带宽限制;算术强度较高,计算单元才更可能成为限制因素。同一 warp 的线程如果访问连续地址,更容易形成合并访问;访问跨度大或布局不匹配时,会增加全局显存事务。Triton 的 block 抽象可以帮助组织 tile 和数据搬运,但不会自动决定最优的 tile 大小或复用方式,具体选择仍需结合算子形状和 profiling 判断。资源占用过高还可能降低一个 SM 同时驻留的 warp 数量,即 occupancy。下面的表格把这些判断落到各级存储层次。

层级主要作用主要观察点典型风险
寄存器保存线程正在使用的少量数据临时变量和寄存器占用register spilling、occupancy 下降
Shared Memory(片上 SRAM)在线程块内复用数据tile 是否放得下、数据是否重复复用bank conflict、容量或占用过高
L2 Cache缓冲多个 SM 对全局数据的访问不同线程块能否复用相同数据命中率不足、反复访问显存
全局显存(通常由 HBM 提供)保存权重、激活和中间结果访问次数、带宽和容量带宽受限、容量不足

Q3小验证:用算术强度判断瓶颈

把一个算子的 FLOPs 与需要搬运的字节数放在一起,观察它更可能受计算吞吐还是显存带宽限制。

python
def estimate_arithmetic_intensity(flops: float, bytes_moved: float) -> float:
    """计算算术强度:每搬运 1 Byte 数据对应的 FLOPs。"""
    if flops < 0 or bytes_moved <= 0:
        raise ValueError('FLOPs 不能为负数,bytes_moved 必须为正数')
    return flops / bytes_moved


def roofline_estimate(flops: float, bytes_moved: float, peak_compute_tflops: float, memory_bandwidth_gb_s: float) -> Dict[str, float | str]:
    """用简化 Roofline 估算可达性能和瓶颈;不是 profiler 计数器结果。"""
    if flops < 0 or bytes_moved <= 0 or peak_compute_tflops <= 0 or memory_bandwidth_gb_s <= 0:
        raise ValueError('输入必须满足 FLOPs>=0,字节数、峰值算力和带宽均为正数')
    intensity = estimate_arithmetic_intensity(flops, bytes_moved)
    bandwidth_roof = intensity * memory_bandwidth_gb_s / 1e3
    return {'arithmetic_intensity': intensity, 'bandwidth_roof_tflops': bandwidth_roof, 'attainable_tflops': min(peak_compute_tflops, bandwidth_roof), 'bottleneck': 'compute' if peak_compute_tflops <= bandwidth_roof else 'memory'}


cases = {'elementwise_like': (2_000_000, 1_000_000), 'matmul_like': (2_000_000_000, 100_000_000)}
for name, (flops, moved) in cases.items():
    r = roofline_estimate(flops, moved, 300, 1500)
    print(f"{name}: intensity={r['arithmetic_intensity']:.1f} FLOPs/Byte, {r['bottleneck']} bound")


def test_roofline_estimate():
    assert roofline_estimate(2e6, 1e6, 300, 1500)['bottleneck'] == 'memory'
    result = roofline_estimate(2e12, 1e8, 300, 1500)
    assert result['bottleneck'] == 'compute' and result['attainable_tflops'] <= 300
    print('✅ roofline tests passed')


test_roofline_estimate()

Q4:Attention 工作集为什么会暴露显存瓶颈?

点击展开查看解析

模型运行时会产生权重、激活和中间张量,它们需要在计算单元与不同内存层级之间移动。Attention 是一个典型例子:序列变长时,完整的 score 矩阵按序列长度的平方增长,工作集和数据搬运都可能成为限制。

本节先用 Attention 工作集解释分块的必要性;需要继续学习分块工作集和 Online Softmax 时,阅读 Part 01 · 14 FlashAttention 显存模型Part 02 · 20 FlashAttention 模拟。 小验证先估算 Q/K/V、score 矩阵和输出的工作集,再比较完整 score 矩阵与单个 tile 的规模;它解释驻留方式的变化,不代表真实 kernel 的端到端性能。

Q4小验证:Attention 工作集与分块尺度

python
def calculate_attention_vram(seq_len: int, num_heads: int, head_dim: int, dtype_bytes: int = 2, batch_size: int = 1) -> Dict[str, float]:
    """估算标准 Attention 的工作集显存;不是完整模型的 CUDA peak memory。"""
    if min(seq_len, num_heads, head_dim, dtype_bytes, batch_size) <= 0:
        raise ValueError('所有形状和 dtype_bytes 都必须为正数')
    qkv = 3 * batch_size * seq_len * num_heads * head_dim * dtype_bytes
    matrix = batch_size * num_heads * seq_len * seq_len * dtype_bytes
    output = batch_size * seq_len * num_heads * head_dim * dtype_bytes
    flops = 4 * batch_size * num_heads * seq_len * seq_len * head_dim
    return {'qkv': bytes_to_gb(qkv), 'attention_matrix': bytes_to_gb(matrix), 'output': bytes_to_gb(output), 'total': bytes_to_gb(qkv + matrix + output), 'score_flops': flops, 'score_intensity': flops / matrix}


def attention_sensitivity(seq_lens, batch_size=1, num_heads=32, head_dim=128, dtype_bytes=2):
    """生成不同序列长度下的工作集表;不等同于完整模型峰值显存。"""
    rows = []
    for seq_len in seq_lens:
        r = calculate_attention_vram(seq_len, num_heads, head_dim, dtype_bytes, batch_size)
        rows.append({'seq_len': seq_len, 'qkv_gb': round(r['qkv'], 4), 'attention_matrix_gb': round(r['attention_matrix'], 4), 'total_working_set_gb': round(r['total'], 4), 'attention_share': round(r['attention_matrix'] / r['total'], 4), 'score_flops': r['score_flops'], 'score_intensity': round(r['score_intensity'], 2)})
    return rows


def attention_score_bytes(seq_len: int, dtype_bytes: int = 2) -> int:
    """估算完整 Attention score 矩阵的体积。"""
    if seq_len <= 0 or dtype_bytes <= 0:
        raise ValueError('seq_len 和 dtype_bytes 必须为正数')
    return seq_len * seq_len * dtype_bytes


def tile_bytes(block_size: int, dtype_bytes: int = 2) -> int:
    """估算一个方形 tile 的体积。"""
    if block_size <= 0 or dtype_bytes <= 0:
        raise ValueError('block_size 和 dtype_bytes 必须为正数')
    return block_size * block_size * dtype_bytes


def compare_attention_working_set(seq_len: int, block_size: int = 128, dtype_bytes: int = 2):
    """比较完整矩阵和单 tile 工作集;不模拟真实 FlashAttention kernel。"""
    score, tile = attention_score_bytes(seq_len, dtype_bytes), tile_bytes(block_size, dtype_bytes)
    return {'naive_score_mb': score / 1024 / 1024, 'tile_kb': tile / 1024, 'tile_count': (seq_len + block_size - 1) // block_size, 'working_set_reduction': score / tile}


for row in attention_sensitivity([512, 2048, 4096, 8192]):
    print(row)

for seq_len in [2048, 4096, 8192]:
    print(f"seq_len={seq_len}: {compare_attention_working_set(seq_len)}")


def test_attention_working_set():
    rows = attention_sensitivity([512, 1024, 2048])
    totals = [r['total_working_set_gb'] for r in rows]
    matrices = [r['attention_matrix_gb'] for r in rows]
    assert totals == sorted(totals) and matrices == sorted(matrices)
    assert rows[-1]['attention_share'] > rows[0]['attention_share']
    short, long = compare_attention_working_set(2048), compare_attention_working_set(4096)
    assert long['naive_score_mb'] > short['naive_score_mb']
    assert long['tile_kb'] == short['tile_kb'] and long['tile_count'] > short['tile_count']
    print('✅ attention working-set tests passed')


test_attention_working_set()

扩展 Q5:显存跨卡分摊后,通信代价从哪里产生?

点击展开查看解析

当单卡装不下模型或需要跨卡协作时,参数、激活或 KV Cache 可能需要在 GPU 之间交换。显存分摊并不会消除成本,而是把一部分成本转移到通信带宽、通信延迟和同步等待。

PCIe、NVLink 和 NVSwitch 可以提供不同的连接带宽与拓扑条件;具体差异取决于设备、链路数量和平台实现。本节扩展用于判断:如果分摊后需要频繁 All-Reduce、All-Gather 或状态同步,通信就可能成为新的瓶颈。下面的代码只估算 PCIe / NVLink 链路传输和重复同步的时间代理,不模拟集合通信的真实算法、拓扑或通信量;这些内容需要在分布式项目中验证。

通信方式典型用途主要代价
PCIeGPU 与 CPU、GPU 与 GPU 之间的数据传输带宽和延迟受主机拓扑影响
NVLink节点内 GPU 间的高速互连需要匹配的 GPU、链路和拓扑
NVSwitch在多 GPU 系统中提供更均匀的互连路径仍需考虑同步与通信量
All-Reduce / All-Gather聚合或交换并行计算中的参数、梯度或激活通信量、同步等待和关键路径阻塞
python
def pcie_vs_nvlink(payload_mb: float, pcie_gbps: float = 64, nvlink_gbps: float = 900, communication_rounds: int = 1, sync_overhead_ms: float = 0.0):
    """估算 PCIe / NVLink 链路时间和重复同步代价;不是实测。

    All-Reduce、All-Gather、NVSwitch 拓扑和真实 collective 通信量不在此代理模型中。"""
    if payload_mb <= 0 or pcie_gbps <= 0 or nvlink_gbps <= 0:
        raise ValueError('payload 和带宽必须为正数')
    if communication_rounds < 1 or sync_overhead_ms < 0:
        raise ValueError('communication_rounds 至少为 1,sync_overhead_ms 不能为负数')
    pcie_once = payload_mb * 8 / pcie_gbps
    nvlink_once = payload_mb * 8 / nvlink_gbps
    return {'pcie_once_ms': round(pcie_once, 3), 'nvlink_once_ms': round(nvlink_once, 3), 'pcie_total_ms': round((pcie_once + sync_overhead_ms) * communication_rounds, 3), 'nvlink_total_ms': round((nvlink_once + sync_overhead_ms) * communication_rounds, 3), 'communication_rounds': communication_rounds, 'evidence_level': 'link_transfer_proxy', 'collectives_modeled': False}


for payload in [64, 256, 1024]:
    print(payload, 'MB ->', pcie_vs_nvlink(payload, communication_rounds=4, sync_overhead_ms=0.05))
print('带宽差异只有在通信位于关键路径时才会影响端到端时间')


def test_interconnect_model():
    one, many = pcie_vs_nvlink(256, communication_rounds=1), pcie_vs_nvlink(256, communication_rounds=4)
    assert many['pcie_total_ms'] > one['pcie_total_ms'] and many['nvlink_total_ms'] > one['nvlink_total_ms']
    assert many['collectives_modeled'] is False and many['evidence_level'] == 'link_transfer_proxy'
    print('✅ interconnect model tests passed')


test_interconnect_model()

相关阅读

导语: 如果想把硬件层级判断继续接到访存优化、attention kernel 和推理验证上,可以按“官方文档 → 开源实现 → 论文”顺序阅读;这些入口都是可选扩展。

Released under the MIT License.