flâneur

Xianbang Wang — research & writing

peppaking8.github.io · 2,187 words · saved by 2 readers

Xianbang Wang — MIT '29, deep learning · vision · RL. Personal site and technical blog.

A practical story about writing TPU Pallas kernels: from naive attention to native backward kernels, sparse attention, and a 1.57x nanoGPT training speedup. GPU users have CUDA and Triton tutorials everywhere; TPU users usually hear the simpler advice: write JAX and trust XLA. This post asks when that advice is still leaving performance on the table. We build a causal attention Pallas kernel, add the native backward pass needed for training, fuse the Q/K preprocessing and layout boundary, and integrate the result into a nanoGPT-style language model. The final v6e-16 training run improves from

saved by

related reading