#Tensor Cores

rss · NVIDIA CUDA Technical Blog · 10月19日 17:00 「背景」 Mamba-2 兼顾 RNN 的固定大小隐状态与注意力式可并行训练，但前向既要沿序列递推状态，又要把主要计算落到 Tensor Core 上。NVIDIA cuTile 用 Python 的 tile 级 DSL 把线程与 warp 调度交给编译器，本文讲述 Jane Street 与 NVIDIA 如何用它实现并优化 Mamba-2 前向核。 「方案」 作者按块重写递推：在对数域用累积和 G 表示衰减，把状态跳转写成 h\\{t+T\}=exp\(G\last\)h\t+\(exp\(G\last−G\)·B\)X^T，输出写成 Y=exp\(G\)Ch\\{kT\}+\(exp\(G−G^T\)·\(CB^T\)·M\)X+DX，M 为因果掩码，从而把主体计算变成 Tensor Core 矩阵乘。实现分为三个核：chunk\cumsum 计算块内累积衰减并以 log2 尺度保存；fwd\pass\state 在每个网格点内顺序遍历块，用 ct.mma 完成状态跳转；compute\out 将 CB^T、携带状态、衰减掩码与 Dx 融合，避免把巨大的块间交互矩阵写回全局内存。作者还给 load/store 加延迟提示，利用 Blackwell Tensor Core 操作数不对称的特点，以转置恒等式把融合项 b\dt\d 放到 A 操作数来缓解共享内存争用；把 dt 与 decay\cumsum 改为按 token 连续以启用 TMA，并用自动调参搜索块大小与占用率。在单张 B300 上，该前向核比 Triton 版 mamba\ssm 快 1.14–1.82 倍，序列越长、状态越大收益越明显；但这些仅限内核级，端到端训练还取决于文中未覆盖的反向传播。 「启示」 作者表明，经过块化重写与硬件感知的布局、操作数调度，Mamba-2 前向可以高效映射到 Tensor Core；cuTile 也能在保持 Python 可读性的同时产出有竞争力的内核，而完整训练性能仍待反向核补齐。 标签: Mamba-2, CUDA Tile, GPU kernels, Tensor Cores, performance optimization

rss · NVIDIA CUDA Technical Blog · 10月29日 15:00 「背景」 机器学习原子间势（MLIP）用神经网络预测能量和力，在接近 DFT 精度的同时降低计算成本，但推理仍多采用 FP32。分子动力学对力误差高度敏感，误差会沿时间步累积，因此不能简单改用低精度 Tensor Core。 「方案」 NVIDIA 作者以 PET-MAD 为例，评测两条不改模型权重和架构的 Tensor Core FP32 路径：Blackwell 上的 BF16x9 把每个 FP32 操作数拆成三个 BF16 分量，计算九个乘积后以 FP32 累加；Ampere 上的 TF32x3 用前导和残差 TF32 分量计算三个交叉项并省略最小项，通过 CUTLASS 自定义线性层同时覆盖能量与梯度。在 26M 和 255M 参数、192 与 17,280 原子体系上，BF16x9 使大 GEMM 提速 2.2–2.5 倍，MD 步进在大体系提速 1.7 倍、小体系 1.2 倍；TF32x3 使 GEMM 提速约 1.7–1.9 倍，MD 步进在较大体系为 1.5 倍，255M 模型上为 1.4 倍。作者报告 BF16x9 的误差低于模型自身非确定性，TF32x3 的力误差约为 FP32 重编译差异的三倍（3.9×10⁻⁶ eV/Å）；CO/Pt\(111\)势能曲线与 NVE 能量守恒测试接近 FP32，而单次 TF32 出现波动和漂移，Cantor 合金弛豫中 TF32x3 用 218–222 步达到容差，单次 TF32 在 500 步内未收敛。作者也指出，收益取决于体系大小、GEMM 占比和内核调优，且 TF32x3 仍需按模型的精度与物理稳定性要求验证。 「启示」 作者的核心结论是，Tensor Core FP32 可作为无需重训练或改架构的独立性能层，在保持接近 FP32 物理行为的同时显著加速以 GEMM 为主的 MLIP 推理；其普适性仍需在更多模型和硬件上验证。 标签: Tensor Cores, mixed-precision, machine learning interatomic potentials, molecular dynamics, GPU benchmarking