이식성은 신화다: 최고의 AI 스택은 하드웨어 비의존적일 수 없는 이유

수학 레이어만은 이식 가능하다. matmul, softmax, loss 같은 연산은 TPU, GPU, CPU, Trainium에서 동일하지만, 범용 커널로 두 플랫폼에서 30% MFU만 내는 대신 한쪽에서 90% MFU를 내는 편이 초거대 스케일에서는 수백만 달러 차이를 만든다고 지적한다.

• Layer 1: Math - PyTorch/JAX의 추상 연산은 공통이다.
• Layer 2: Compiler - XLA, torch.compile, Triton compiler는 shared IR를 각 하드웨어 backend로 내리지만, backend는 공유되지 않는다.
• Layer 3: Hardware-native code - Pallas, CuTile, NKI 같은 커널이 MFU를 결정한다.

Pallas, CuTile, NKI는 각각 TPU, Blackwell, Trainium의 메모리 계층과 동기화 장치를 직접 드러낸다. TPU는 VMEM과 DMA, MXU/VPU/XLU를 쓰고, Blackwell은 TMEM, tcgen05, mbarrier, elect.sync에 의존하며, Trainium은 SBUF와 PSUM, NeuronCore partitioning을 명시적으로 다룬다. 저자는 이 레벨에서의 프로그램은 서로 다른 구현이 아니라 서로 다른 알고리즘이라고 본다.

증거도 구체적이다. MaxText MegaBlox의 MoE grouped GEMM은 TPU에서 29개의 Pallas GMM 커널 호출과 282줄의 Pallas Python으로 구현된다. 반대로 flashinfer의 MoE CUDA는 Blackwell SM100용 생성 커널 300개, 약 400만 줄 규모로 작성되며, bf16, fp8, TMA, warp specialization, cooperative/pingpong 전략이 파일명과 코드에 그대로 박혀 있다.

NVIDIA는 CUDA C, CUTLASS/CuTe, Triton, CuTile을, AMD는 Triton을 받아들인 뒤에도 FlyDSL을 따로 만들었다는 점도 같은 결론을 지지한다. 마지막 결론은 단순하다. 범용 DSL을 하나 만들면 TPU, NVIDIA, Trainium의 개념을 모두 담아야 해서 결국 대상별 분기와 escape hatch만 늘어나고, 반대로 추상화만 높이면 성능이 빠진다. 그래서 성공하는 DSL은 ISA를 Python으로 얇게 감싼 형태이며, 하드웨어 벤더는 범용성을 기다리기보다 자기 칩에 맞는 Python DSL을 직접 만들어야 한다고 주장한다.