Tejash Shah

Tejash Shah is a principal product manager in the AI Platform Software group at NVIDIA, responsible for managing ML framework, JAX and ML compiler, XLA.
Avatar photo

Posts by Tejash Shah

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
Models / Libraries / Frameworks

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
Models / Libraries / Frameworks

CUTLASS: Principled Abstractions for Handling Multidimensional Data Through Tensors and Spatial Microkernels

In the era of generative AI, utilizing GPUs to their maximum potential is essential to training better models and serving users at scale. Often, these models... 12 MIN READ