Machine Learning Engineer·Pune, India
Shrirang Mahajan*
I train language models from scratch, fine-tune them for the real world, and ship the production ML systems that put them to work — from an LLM that runs in your browser to models serving users at scale.
Currently Machine Learning Engineer II at Skylark Labs.
I'd rather train the model
than just prompt one.
Three things I've come to believe about building ML and LLM systems — the kind of engineer I'm trying to be.
Provider, not consumer.
Most "AI engineering" today means orchestrating APIs that someone else built — wrapping endpoints, chaining prompts, paying for inference. I'd rather understand the model end to end: pretrain it, fine-tune it, and serve it. TinyGPT is a 95M-parameter LLM I pretrained from scratch on 8 GB of VRAM — and when I wanted to know what really happens under
.cuda(), I wrote Tensorax, a tensor library from scratch. Depth, not just glue.Proof of work over pedigree.
A degree describes potential. Code describes what you've actually shipped. The bar I hold myself to is whether someone can
pip installmy work — or talk to a model I trained — and have it just run. What I can build, train, and ship matters more than where I studied it.The right hard path.
I use AI to learn faster. I do not use it to write code I can't read. When I pretrained TinyGPT on a single 8 GB GPU, I learned attention internals, mixed precision, and gradient accumulation by hand — not by autocomplete. Going one layer deeper than I strictly need to is the whole point. A shortcut that skips the part where you become better isn't a shortcut.
A few things I've built recently.
From hand-written CUDA kernels and an LLM inference server to hybrid search and an LLM trained on a single 8 GB GPU. I build things to understand how they work underneath.
- 01
LLM Serving Engine
- Designed an LLM inference server with continuous batching, KV caching and admission control, serving Qwen2.5-0.5B over FastAPI.
- Load-tested it against vLLM on TTFT, TPOT and throughput, sustaining 788 tok/s at 64 concurrent requests on one 8 GB GPU, and shedding load with HTTP 429 past capacity.
PyTorchFastAPIvLLMContinuous BatchingLLM Inference - 02›
Tensorax
- Built a standalone tensor computation library from scratch with custom C++/CUDA kernels and Python bindings, without relying on NumPy, PyTorch, or cuBLAS.
- Implemented complete automatic differentiation (autograd) with computational graph and PyTorch-like API including neural network layers, optimizers, and loss functions.
- Hand-wrote a Tensor Core flash attention kernel in inline PTX (mma.sync, ldmatrix, cp.async) that reaches 24.7 TFLOPS, 0.79× cuDNN at long sequences, without calling any external math library.
- Published on PyPI with a complete neural network training pipeline including SGD, Adam optimizers, and multiple activation functions.
- Available on GitHub with comprehensive documentation and examples.
C++CUDAPythonPyBind11GPU Computing - 03
Kiosk
- Building a news search and recommendation product over 90K articles. Search runs hybrid BM25 + dense retrieval fused with RRF, adds query expansion only when results look weak, and reranks the top results with a cross-encoder.
- The home feed retrieves separately for each of the reader's interests and merges the results, then ranks for relevance, freshness and MMR diversity. Serves the feed at 99 ms P95.
- Detects breaking news without any click data: a story counts as breaking when many different outlets cover it in a short window. Groups articles into stories that keep the same ID across updates, so readers can follow them.
- Adds RAG chat with cited answers (vLLM, Groq or Gemini) and a benchmark suite that tracks nDCG, MRR and latency stage by stage.
QdrantSentence TransformersFastAPIReactHybrid Search - 04›
TinyGPT
- Pretrained TinyGPT, a 95M-parameter GPT (RoPE, GQA, RMSNorm), from scratch in PyTorch on ~2.1B tokens on one RTX 3070 Ti (8GB VRAM), using mixed precision, gradient accumulation and Liger fused kernels.
- Instruction-tuned it on Stanford Alpaca (val perplexity 7.3), exported it to ONNX, and run it in the browser via WebGPU on this site.
- Deployed an interactive Streamlit UI on a free-tier CPU with dynamic model switching. Published the model on HuggingFace and shared on GitHub and Streamlit Cloud.
PyTorchStreamlitHuggingFaceMixed-Precision Training
My 95M-parameter LLM, running on your device.
The model I pretrained from scratch on ~2.1B tokens of OpenWebText using a single RTX 3070 Ti, then instruction-tuned on Stanford Alpaca. Click to download the ONNX weights (~536 MB) — inference then runs locally via WebGPU (or WASM fallback). Fluent English, often wrong facts — it's small.
A note — at 95M parameters, TinyGPT speaks English fluently but isn't reliable on facts. It will write a coherent sentence about anything, and confidently make things up. Treat it as a coherence demo, not a knowledge model.
Click to download the model and run it locally.
Essays, with illustrations.
Long-form posts on attention, embeddings, and pre-training LLMs on consumer hardware — with the occasional dive under the hood.
- Mar2026
Pre-training a 95M Parameter LLM on a Consumer GPU
15 min readLLM · PyTorch · Pre-training - Jan2026
Writing CUDA Kernels from Scratch: A Beginner's Guide
10 min readCUDA · GPU Programming · Matrix Multiplication - Apr2025
Llama 2 Explained
8 min readLlama · Meta AI · Llama2 - Apr2025
Llama 4: The Natively Multimodal AI
5 min readLlama · Meta AI · Llama4
A short résumé.
Ask about this site.
A hosted LLM grounded in this site's content — every project, role, blog post, and skill on this page is in its context. It can also fetch fresh GitHub data and send you to the right page when you ask for it. Much better informed about me than TinyGPT above — though, like any language model, still capable of getting things wrong.
I know this entire site. Ask me about Shrirang's work, or pick a prompt below.
Answers are generated by a third-party language model and can be wrong. What you type is sent to that provider for processing — please don't enter anything confidential or personal. How this data is handled.
Thanks for scrolling. If something here resonated — a project, an essay, a problem you're trying to solve — I'd love to hear from you.
— Shrirang