Kunjan Patel

I worked on model bring up and optimization of mainly OSS diffusion and some MOE models. My team and my main responsibility is implementing the forward pass, training pipeline in JAX and optimizing its performance on TPUs and GPUs.

• Worked in partnership Nvidia's MLE's to optimize SDXL and Flux models for H100 and H200 GPUs
• This involved performing roofline analysis, calculating MFU looking at perfetto and xprof traces to find bottlenecks.
• For SDXL I focused on optimizing U-net, this involved experimenting with custom pallas kernel for convolutions.
• Other improvements involved fixing gaps between train steps, fixing MFU calculation and long jax jit compile times due to constant folding.
• For flux this involved analyzing the xprof to find and fix the shardings to removed unexpected activation all gathers for FSDP parallelism.

Worked on Optimizing WAN 2.1 diffusion model training and inference on latest.

• Read wan training technical report and began experiments for supporting long context length. Ran experiments with multiple sharding strategies ( tensor parallelism, context parallelism, sequence parallelism, hybrid parallelism).
• Experimented with oss ring attention in lax and jax, identified performance issues and implemented fused kernel of ring and flash attention in Pallas.

Previously worked on Distributed Inference Team

• Experimented and developed solution for disaggregated inference
• Contributor to vllm
• Worked on smart router for lora and did perfomance analysis of dynamic lora switching, roofline analysis of GMMV and BGMV punica kernels in VLLM.

Worked on Tanium Cloud team , wrote bazel libraries to support tanium cloud product

Joined as new grad and implemented a Kafka pipeline to ingest AWS logs and using a logical model create firewall policies.