Keshav Krishna

MS Computer Science · NYU Courant

I'm an MS Computer Science student at New York University (Courant Institute) with a 3.9 GPA. I work as a research assistant to Prof. Mohamed Zahran, focusing on ML-driven program phase prediction and workload forecasting for dynamic hardware resource optimization.

Before NYU, I spent two years as a Software Engineer at JioSaavn, where I built production-scale recommendation systems, content pipelines, and RAG applications serving over 100 million users. I hold a BTech in Computer Science from IIT Ropar.

My research interests lie at the intersection of machine learning and computer systems — specifically ML for hardware resource optimization, GPU kernel engineering, and efficient model inference. I have a single-author publication in Frontiers in AI on ML for cache management.

I'll be joining Moore Capital Management as a Risk Technology Intern this summer. I'm seeking Fall 2026 research or engineering opportunities in ML systems, risk technology, and high-performance AI infrastructure. Feel free to reach out.

Keshav Krishna

News

Jun 2026 Announced Summer 2026 role as Risk Technology Intern at Moore Capital Management, building tools around risk data and market risk workflows.
Mar 2026 FlashAttention-2 Triton kernel implementation completed — from-scratch forward pass with online softmax, backward pass with recomputation, benchmarked up to seq length 65K. Code
Jan 2026 Joined Prof. Mohamed Zahran's research group at NYU Courant, working on ML-driven program phase prediction for hardware resource optimization.
Jan 2026 Built a complete Transformer LM from scratch (BPE tokenizer, RoPE, SwiGLU, AdamW) — no nn.Linear, nn.Embedding, or torch.optim. Code
Dec 2025 Completed Pico-LLM — unified codebase for Small Language Model training with SFT, DPO/GRPO alignment on GSM8K, mechanistic interpretability tools, and KV-cache inference optimization. Code
Aug 2025 Started MS in Computer Science at NYU Courant.
Apr 2025 Single-author paper published in Frontiers in Artificial Intelligence: "Advancements in cache management: a review of machine learning innovations." 9+ citations.
Jul 2023 Started as Software Engineer (Data Science) at JioSaavn.
May 2023 Graduated with BTech in Computer Science from IIT Ropar.

Publications

Advancements in cache management: a review of machine learning innovations Single Author
Keshav Krishna
Frontiers in Artificial Intelligence, 2025.  9+ citations.
Paper DOI
Last Level Cache Intra-set Write Balancing for Non-Volatile Memory
Research with IIT Ropar, 2022–2023.
arXiv preprint, 2024.
arXiv
ML-Driven Program Phase Prediction for Dynamic Hardware Resource Optimization
Keshav Krishna, Mohamed Zahran
Ongoing research, NYU Courant, 2026–present.

Selected Projects

FlashAttention-2: Triton Kernel

From-scratch Triton GPU kernel with online softmax and tiled computation, reducing attention memory from O(N²) to O(N). Backward pass with activation recomputation and benchmarking up to 65K-token sequences.

Triton · PyTorch · CUDA · Nsight Systems

Transformer LM from Scratch

Complete Transformer LM without nn.Linear/nn.Embedding/torch.optim — BPE tokenizer, RoPE, RMSNorm, SwiGLU, AdamW, cosine LR. Ablations across pre/post-norm, RoPE vs NoPE, SwiGLU vs SiLU.

PyTorch · BPE · RoPE · SwiGLU · AdamW

Pico-LLM: SLM Training & Interpretability

Unified codebase for Small Language Model training with SFT, DPO/GRPO alignment on GSM8K. Mechanistic interpretability tools and KV-cache inference optimization.

PyTorch · DPO · RLHF · KV-Cache

Alignment & Reasoning RL

Reasoning-RL experiments focused on alignment workflows for mathematical and multi-step reasoning, complementing SFT/DPO/GRPO work from Pico-LLM.

PyTorch · RL · Alignment · Reasoning

Experience

Risk Technology Intern · Moore Capital Management, New York
Summer 2026

Summer intern building tools around risk data, market risk workflows, and collaboration between engineering and risk analysis teams.

Research Assistant · NYU Courant — Prof. Mohamed Zahran
Jan 2026 – Present
  • ML-driven program phase prediction and workload forecasting for dynamic hardware resource optimization in multi-core processors.
  • Time-series models on hardware performance counters (CPI, branch mispredictions, stall cycles) to predict execution phase transitions at runtime.
Software Engineer, Data Science · JioSaavn, Mumbai
Jul 2023 – Jul 2025
  • Built distributed content recommendation services for 12M daily users with Hive, Spark, Redis — driving 1M impressions, 800K clicks, and 500K streams per day.
  • Generated mood-based topic mixes for 7.6M daily users using TF-IDF, served to 100M app users with a 12.81% CTR.
  • Created custom BERT embeddings for songs and users, raising top-5 and top-10 recommendation accuracy to 61.82% and 67.21%.
  • Built a RAG system using LLaMA3-7B for intent classification with Stable Diffusion cover art generation.
  • Transitioned content moderation from SetFit to prompt engineering with Llama3-8B-instruct (85% accuracy).
  • Reduced infrastructure costs by migrating from Annoy to Qdrant and optimizing data lifecycle management.
Spark/HiveLlama 3BERTQdrantRedisFlask
Software Intern · GE Healthcare, Bangalore
May 2022 – Jul 2022
  • YOLO-based face de-identification pipeline for medical images — 90%+ accuracy on 20,000+ images.
  • People detection models in IR/RGB imagery (85%+ mAP) for hospital environment monitoring.
  • Docker-containerized modules with FastAPI endpoints for cross-project reuse.
YOLOComputer VisionFastAPIDocker

Education

New York University — Courant Institute Aug 2025 – May 2027
MS, Computer Science · GPA: 3.9
Machine Learning · Honors Analysis of Algorithms · Programming Languages
Indian Institute of Technology, Ropar Aug 2019 – Jun 2023
BTech, Computer Science & Engineering · GPA: 8.74 / 10
Merit Scholarship for top academic performance

Transcript

You can view my latest academic transcript here: Open Transcript (PDF).

Awards