Portability is a myth: why the best AI stack can't be hardware-agnostic

Only the math layer is portable. Operations like matmul, softmax, and loss are the same across TPU, GPU, CPU, and Trainium, but at hyperscale, getting only 30% MFU on two platforms with a generic kernel versus 90% MFU on one platform makes a difference of millions of dollars, the article points out.

• Layer 1: Math - The abstract operations in PyTorch/JAX are common.
• Layer 2: Compiler - XLA, torch.compile, and the Triton compiler lower a shared IR to each hardware backend, but the backends themselves are not shared.
• Layer 3: Hardware-native code - Kernels like Pallas, CuTile, and NKI determine MFU.

Pallas, CuTile, and NKI each directly expose the memory hierarchy and synchronization primitives of TPU, Blackwell, and Trainium respectively. TPU uses VMEM and DMA, MXU/VPU/XLU; Blackwell relies on TMEM, tcgen05, mbarrier, and elect.sync; and Trainium explicitly handles SBUF, PSUM, and NeuronCore partitioning. The author views programs at this level not as different implementations, but as different algorithms.

The evidence is concrete. The MoE grouped GEMM in MaxText MegaBlox is implemented on TPU with 29 Pallas GMM kernel calls and 282 lines of Pallas Python. In contrast, the MoE CUDA in flashinfer for Blackwell SM100 is written with 300 generated kernels totaling roughly 4 million lines, with bf16, fp8, TMA, warp specialization, and cooperative/pingpong strategies baked directly into file names and code.

The fact that NVIDIA has embraced CUDA C, CUTLASS/CuTe, Triton, and CuTile, and that AMD, even after adopting Triton, went on to build FlyDSL separately, supports the same conclusion. The final conclusion is simple: building a single general-purpose DSL would require encompassing the concepts of TPU, NVIDIA, and Trainium all at once, which only leads to more target-specific branching and escape hatches — while raising the level of abstraction only strips away performance. So the DSLs that succeed are thin Python wrappers around an ISA, and rather than waiting for a universal solution, hardware vendors should build their own Python DSLs tailored to their own chips, the article argues.