10+ years building inference systems, ML accelerator frameworks, and large-scale serving platforms. I work on the AWS Neuron SDK team, where my job is making large-model inference fast on Trainium and Inferentia. That means disaggregated prefill/decode serving in vLLM, KV cache transfer throughput, and decode latency on production deployments. I also work on making these chips run stock PyTorch, including a CPU-only compile path that moved compilation off the accelerator fleet.
Experience
AWS Neuron (Annapurna ML), Cupertino, CA
Disaggregated Inference (DI) and Inference Optimization (Technical Lead)
- Built disaggregated inference for vLLM on Neuron. Prefill and decode now run on separate Trainium nodes, with the KV cache moved device to device over NIXL to keep serving latency down.
- Designed the runtime batch operations API for KV cache transfers, which doubled transfer throughput. Tuned EFA utilization to 20 of 25 Gbps on Trn2 for inter-node tensor transfers.
- Led DI bring-up on next-generation Trainium. EFA-to-Neuron device mapping had changed, so multi-node inference needed runtime-level fixes first. Also profiled and tuned NIXL for small transfers.
- Got DI ready for product launches, including xPyD multi-node configurations and mixed GPU plus Trainium deployments, and validated it on the EKS production stack.
- Moved the NeuronX Distributed Inference (NxDI) backend from single-program, multiple-data execution to multiple-program, multiple-data once native PyTorch support landed. That opens up distributed topologies the old model could not express.
- Bringing up a large open-weight model on next-generation Trainium, including the framework and runtime integration work it needs.
- Evaluated vLLM sleep mode for serverless inference, where weights load and unload on demand to cut idle cost for serving and RLHF workloads.
- Won back per-step host time that asynchronous execution was supposed to hide. The gap came from redundant device-to-host copies plus DMA descriptor cost at submit, and I specified the I/O ring and copy-elision fixes.
- Found the real cause of a decode stall on a 64-way expert-parallel deployment. The team had put it down to prefill and load imbalance. It was a per-step collective in the data-parallel padding path, serializing request churn.
Trainium as a First-Class PyTorch Device (Technical Lead)
- Built CPU-only compilation for
torch.compile on Trainium. PyTorch has no such path, because a custom device backend has to initialize, allocate, and execute on real hardware before it can compile anything. Models now compile to a deployable NEFF artifact on an ordinary CPU host, so compilation, CI, and artifact builds stop taking accelerator capacity. It shipped as the public torch_neuronx.compiler API after cross-team review, and I caught a numerical divergence before release that a missing compiler flag would otherwise have hidden.
- Made Trainium run stock PyTorch. Models go through standard
torch.compile instead of a forked PyTorch/XLA stack, so customers keep the programming model they already know and the SDK maintains one compile path rather than two. Removing the legacy XLA framework dependency took sign-off from eight senior leaders across the compiler, runtime, container, and Inferentia teams. The two questions they pushed hardest on were what test coverage we would lose and how customers on the prior release stay supported.
- Unblocked a large mixture-of-experts model that would not fit in device memory on the native path, and brought accuracy back in line with the XLA baseline at default memory settings. The compiler was counting the key-value cache twice. The fix touched no model code, no kernels, and no allocator.
- Halved key-value cache input memory. The team had assumed PyTorch AOTAutograd forbids caches that share storage. It does not, and I validated the layout that works.
- Brought up native tensor-parallel vLLM serving end to end on Trainium3. That took fixing five worker integration gaps and showing the native backend can share a process with the lightweight runtime, with outputs matching to 1.5e-7.
- Own the release bar for this work. Split the program into three parallel workstreams and built the Grafana dashboard that holds the native path to XLA correctness and performance across model configurations.
Managing PyTorch for Neuron and Inference Optimization
- Owned PyTorch releases end to end across the Neuron stack, which means builds, dependency coordination, and compatibility testing across 10+ packages for every new PyTorch version.
- Ran the first native (non-XLA) PyTorch release for Neuron, coordinating the torch-neuronx, NxDI, eager mode, and torch.compile paths.
- Turned on parallel compilation for inference workloads to cut model compile time, and added huge page support to the runtime benchmarks.
- Chased down inference performance regressions, memory faults, and framework compatibility breaks across PyTorch and XLA releases to keep production stable.
PJRT / JAX Architecture, HLO Optimization Tooling, and vLLM
- Split the PJRT C-API into three modules of 6K to 9.5K lines each, which cut build times and made the pieces testable on their own. Upgraded the C-API from 0.68 to 0.73 for JAX v0.7.0 inference support.
- Built an HLO graph drift detection tool from scratch. It diffs HLO IR graphs, including NKI custom kernel ops, and catches compiler-level inference regressions across framework upgrades.
- Contributed to the vLLM Neuron plugin and NxDI: EFA device mapping fixes, NIXL cleanup, multi-node MPI communication for distributed inference, and FX graph optimization passes.
- Ran the NKI kernel compatibility analysis for the Trn2 to Trn3 move, so custom kernels carry over without rewrites.
Amazon Web Services, Bangalore, India
AWS Billing (CuBE and Monetization Authority)
- Built Monetization Authority (MTZA) from scratch, a Step Functions workflow engine that automates RFP generation and delivery for enterprise billing across thousands of payers. It drove $300M in free cash flow improvements.
- Built a payer classification system on BIDW Redshift queries and automated SPP payer mapping with ops-managed configuration, which recovered another $100M in free cash flow.
- Designed a transfer billing system that simplified billing for third-party sellers on AWS and improved the security posture of third-party accounts. Estimated to add $1.2B in annual revenue by 2027.
- Built CuBE payer mapping APIs (Smithy/Lambda), distributed map scaling for Step Functions, RFP deduplication, and credits beginning balance pipelines.
- Led the JDK 17 migration across 8+ packages, ran cross-functional teams of 6+ developers, and wrote the PRFAQs and architecture docs.
AWS OpenSearch and CloudSearch
- Containerized data plane components of AWS OpenSearch using S3, ECR, Lambda, and CodeBuild.
- Built an automated node diagnostic and self-healing framework that cut customer RCA tickets by 30%.
- Designed Bastion infrastructure on VPC, Elastic IPs, and NAT gateways. Built the Safe Data Plane Deployment Service.
- Root-caused high-impact customer cluster anomalies and drove LSE and outage scenarios to closure.
Grab, Bangalore, India
- Rebuilt the Settlement Platform from a worker-based design to an event-driven one on AWS Lambda, which made merchant settlement 60% faster.
- Designed the core post-payment processing unit for payment gateways and international merchants in Golang.
Yatra Online Pvt Ltd, Hyderabad, India
- Designed customized Header Service for corporate merchants using Node.js, Java, Spring, and HAProxy/Nginx routing.
Tata Consultancy Services, Hyderabad, India
- Built the TCS-Ericsson JCAT framework for testing Ericsson AXE IO embedded systems, plus a web console that recommended tests to run (Star Performer award).
- Customized the OpenStack console UI for the ATLAS project.
Technical Skills
ML & Inference
PyTorch, torch.compile / TorchDynamo / AOTAutograd / FX, PyTorch/XLA, JAX, vLLM, NxDI, NKI, PJRT C-API, HLO/StableHLO, NEFF, NIXL, EFA, Trainium, Inferentia, LLM Serving, KV Cache Optimization, Disaggregated Inference, Expert/Data/Tensor Parallelism, Graph Compilation, Accelerator Profiling
Languages
Python, C++, Java, Golang, JavaScript/TypeScript
AWS Services
Lambda, Step Functions, DynamoDB, S3, CloudFormation, CDK, ECR, EKS, SageMaker, Redshift, SNS/SQS, VPC
Infrastructure
Bazel, Docker, Kubernetes, GitHub Actions, CI/CD Pipelines, Distributed Systems, gRPC/Protobuf
Education
Jalpaiguri Government Engineering College, Jalpaiguri | CGPA: 8.5/10
Open Source and Community
- GitHub: github.com/rajkthakur. Contributor to the aws-neuron org and to a fork of pytorch/xla. Maintained the AWS Neuron sample notebooks.
- Amazon Bar Raiser: Certified interviewer holding the hiring bar across Amazon.
- AppSec Certified Engineer: Certified in application security. Ran the AppSec review for the LLM features in Amazon Quick Suite.
- Mentoring: Mentor across orgs, and interview for multiple teams and companies.