Tejash Shah

Tejash Shah is a principal product manager within the AI Platform Software group at NVIDIA, responsible for managing JAX and MLX frameworks. Before NVIDIA, Tejash held software engineering roles at semiconductor companies. He holds five patents in a wide range of technological domains. He earned a master's degree in Computer Science from The University of Texas at Dallas and a bachelor’s degree in Information Technology from Gujarat University.
Avatar photo

Posts by Tejash Shah

Agentic AI / Generative AI Jul 10, 2026

Reducing High-Bandwidth Memory Bottlenecks in JAX-Based LLM Training with Host Offloading

Large language model (LLM) training workloads increasingly run into GPU memory limits before compute is fully used. Model weights, gradients, optimizer states,... 9 MIN READ
Decorative image.
Agentic AI / Generative AI Jun 08, 2026

Train Models Faster with JAX and MaxText Using NVFP4 on NVIDIA Blackwell

Pre-training frontier LLMs comes down to throughput. When training spans trillions of tokens across thousands of accelerators, every percentage point of step... 7 MIN READ
Agentic AI / Generative AI Feb 03, 2026

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
Developer Tools & Techniques Nov 13, 2025

Achieve CUTLASS C++ Performance with Python APIs Using CuTe DSL

CuTe, a core component of CUTLASS 3.x, provides a unified algebra for describing data layouts and thread mappings, and abstracts complex memory access patterns... 9 MIN READ
Data Center / Cloud Jul 18, 2025

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
Developer Tools & Techniques Jul 16, 2025

CUTLASS 3.x: Orthogonal, Reusable, and Composable Abstractions for GEMM Kernel Design

GEMM optimization on GPUs is a modular problem. Performant implementations need to specify hyperparameters such as tile shapes, math and copy instructions, and... 12 MIN READ