@wafer_ai: Anthropic is committing more than $100 billion over the next 10 years to AWS for 5GW of new capacity to train and serve…

Anthropic is committing more than $100 billion over the next 10 years to AWS for 5GW of new capacity to train and serve Claude on Trainium2 through Trainium4 chips. Here's How Kernels Work on Trainium: aws built their own CUDA. it's called NKI, and the programming model is completely different from anything NVIDIA. NKI (Neuron Kernel Interface) is a python DSL that gives you direct ISA access to NeuronCores on Trainium chips. it launched in sept 2024 and is currently in beta. if you know CUDA, NKI will feel alien to you. the biggest difference is that in NKI there are no threads. in CUDA you write one thread's perspective and launch thousands. in NKI there are no threads. no warps. no blocks. you write sequential code, top to bottom, and the compiler extracts parallelism by scheduling across 4 heterogeneous compute engines that run asynchronously in parallel. the Tensor Engine is a systolic array for matmul; it reads from SBUF and writes to PSUM. the Vector Engine handles elementwise ops, reductions, softmax. the Scalar Engine does activations and type casting. and the GpSimd Engine has 8 programmable 512-bit SIMD processors for anything custom. so instead of thousands of identical threads on identical SMs, you have one sequential instruction stream pipelined across specialized engines. the compiler figures out what can overlap. the memory model will also probably throw you off NeuronCores have zero hardware caches. all memory is software-managed. SBUF (State Buffer) gives you 24-32 MiB of on-chip SRAM per NeuronCore, organized into 128 physical partitions. PSUM is a 2 MiB accumulator buffer where matmul results land. and then HBM ranges from 32 GiB on Trainium1 to 144 GiB on Trainium3. every data movement is explicit. HBM to SBUF to compute to SBUF to HBM. no L1/L2 to save you. this is like CUDA shared memory except it's mandatory and there's no cache fallback. the 128 partition constraint defines NKI programming every tile in NKI is 2D: [partition_dim, free_dim]. the partition dimension is always exactly 128 elements — it maps to the 128 physical SBUF partitions. hard ISA constraint. for matmul, the contraction axis must map to the partition dimension. tile sizes are hardware constants: TILE_M=128, TILE_K=128, TILE_N=512. another thing that trips people up: most of NKI is compile-time. print() runs at compile time, not on the device. python loops over tensor.shape unroll at compile time. only nki.isa.* calls generate actual on-device instructions. python is the metaprogramming layer, not the runtime. deep dive 3/6 by @gpuemi