Sharding, parallelism, and JIT compilation on TPUs, with JAX
· 9 min read
A gentle background on Megatron-style tensor parallelism, DeepSpeed-Ulysses sequence parallelism, and how jax.jit really compiles a loop — illustrated with the "VAE decode hang" that turned out not to be a hang at all.
Engineering NotesInfrastructure
Understanding video diffusion architectures: DiT, MoT, and conditioning
· 10 min read
A background primer on diffusion transformers, text/image conditioning, and Mixture-of-Transformers — and a tour of what each model family in vidax actually uses.
Engineering NotesModeling
Weight offloading: trading bandwidth for memory
· 9 min read
What weight offloading is, why it exists (ZeRO-Offload), how vidax streams DiT weights from host RAM into HBM one chunk at a time, and what it actually costs — across five different model families, including one case where the same tool turned out to be the wrong one.
Engineering NotesInfrastructure