@wafer_ai: Anthropic 承诺在未来10年内向 AWS 投入超过1000亿美元，用于新增5GW容量以训练和提供服务…

Anthropic 承诺在未来10年内向AWS投入超过1000亿美元，用于新建5GW算力，通过Trainium2至Trainium4芯片训练和部署Claude。

Trainium上的Kernel是如何工作的：

AWS 自研了CUDA替代方案，名为NKI，其编程模型与NVIDIA完全不同。

NKI（Neuron Kernel Interface）是一个Python DSL，可直接访问Trainium芯片上NeuronCore的ISA指令集。它于2024年9月发布，目前处于测试阶段。如果你熟悉CUDA，会觉得NKI很陌生。

最大区别在于：NKI中没有线程概念。

CUDA中，你编写单个线程的视角，然后启动成千上万个线程。NKI中没有线程，没有warp，也没有block。你只需编写从上到下的顺序代码，编译器通过调度4个异步并行运行的异构计算引擎来提取并行性。

张量引擎（Tensor Engine）是一个用于矩阵乘法的脉动阵列，从SBUF读取数据并写入PSUM。向量引擎（Vector Engine）处理逐元素运算、规约、softmax。标量引擎（Scalar Engine）负责激活函数和类型转换。GpSimd引擎（GpSimd Engine）包含8个可编程的512位SIMD处理器，用于自定义操作。

因此，不再是相同SM上成千上万个相同的线程，而是通过专用引擎流水线化的一条顺序指令流。编译器决定哪些操作可以重叠。

内存模型同样令人困惑：

NeuronCore完全没有硬件缓存。所有内存均由软件管理。SBUF（状态缓冲区）为每个NeuronCore提供24-32 MiB的片上SRAM，分为128个物理分区。PSUM是一个2 MiB的累加器缓冲区，矩阵乘法的结果落在此处。HBM容量从Trainium1的32 GiB到Trainium3的144 GiB不等。

每次数据移动都是显式的：HBM → SBUF → 计算 → SBUF → HBM。没有L1/L2来帮你。这类似于CUDA的共享内存，但它是强制性的，且没有缓存回退。

128个分区的约束定义了NKI编程：

NKI中的每个tile都是2D的：[partition_dim, free_dim]。partition维度的大小固定为128个元素——对应128个物理SBUF分区。这是硬ISA约束。对于矩阵乘法，收缩轴必须映射到partition维度。tile大小是硬件常量：TILE_M=128, TILE_K=128, TILE_N=512。

另一个容易让人困惑的地方：NKI的大部分操作发生在编译时。print()在编译时运行，而非设备上。对tensor.shape的Python循环在编译时展开。只有nki.isa.*调用才会生成实际设备指令。Python是元编程层，而非运行时。

深度分析第3/6部分，作者 @gpuemi