TJ Xu

TJ Xu is a software engineer at NVIDIA working on horizontal scaling optimizations in XLA.
Avatar photo

Posts by TJ Xu

Agentic AI / Generative AI

Accelerating Long-Context Model Training in JAX and XLA

Large language models (LLMs) are rapidly expanding their context windows, with recent models supporting sequences of 128K tokens, 256K tokens, and beyond.... 9 MIN READ
Data Center / Cloud

Optimizing for Low-Latency Communication in Inference Workloads with JAX and XLA

Running inference with large language models (LLMs) in production requires meeting stringent latency constraints. A critical stage in the process is LLM decode,... 6 MIN READ