..Celebrating the release of TorchNeuron
Original Source of the blog post: Celebrating the release of TorchNeuron
TorchNeuron is here, and we’re excited about what we can build with it. At Loka, we’re already bringing its native PyTorch workflow into our research and client work. Our journey on AWS Silicon started with real customer workloads and took us through custom kernels, applied research, and conversations with Annapurna Labs about the tools we wanted to use. A native PyTorch backend was high on that list, and now it’s finally here.
Getting started with AWS Silicon
We got started at the end of 2024 with two client projects on AWS Silicon. Sebastian Gonzalez Aseretto and Nicolás Roldán Fajardo worked on one. On the other project, I worked with Zafir Stojanovski and Henrique Ribeiro Delgado da Silva. Real customer workloads gave us a practical way to learn what these chips could do and where we needed to dig deeper.
Zafir, Henrique, and I focused on migrating BEVFusion to AWS Neuron. The Neuron Kernel Interface (NKI) had just become available, so we experimented with NKI kernels and wrote custom C++ operations too. I enjoyed working through those details and getting closer to the hardware.
We put what we’d learned into a report and presentation for Annapurna Labs, covering what worked well and where the developer experience could improve. A native PyTorch backend was high on our list: we wanted to bring our models and training code over using the tools we already knew. We sent the deck over, and it started making its way around the team.
One moment at AWS Summit New York in 2025 caught me completely off guard. We were talking with the Annapurna Labs team at their booth about our work when they recognized us as the Loka team behind the presentation. Then one of them pulled out his phone and brought up our deck. I was stunned. It was the same presentation Zafir, Henrique, and I had put together, now being shown back to us in the middle of the conversation. I hadn’t expected them to recognize our work, let alone pull up the deck right there. We even have a photo from that moment.
The native PyTorch workflow we wanted
We stayed in touch after New York, and Annapurna Labs gave us early access to the first TorchNeuron beta. And trust me, the appeal was immediate. With TorchNeuron’s native backend installed, we can move a supported model and its input tensors onto the Neuron device with .to("neuron") and start experimenting in eager mode. For compiled workloads, we keep that device placement and use torch.compile(model, backend="neuron"). We’re working through our training and inference code with familiar PyTorch tooling, without an XLA-specific execution loop. After our conversations in 2024 and 2025, getting to run our own code this way felt especially satisfying.
That extends to distributed training. TorchNeuron supports the PyTorch distributed APIs, including DDP, FSDP, and DTensor, with sharding configured for the NeuronCores we use.
Putting TorchNeuron to work
We’ve been putting that workflow to use in our applied research. One experiment I especially enjoyed was post-training Tiny ESMC in pure PyTorch, which worked like a charm. We’ve also been experimenting with SAM3, GENA-LM, and Boltz 2.2.1.
ESMC-300M gave us a concrete example of that workflow. For zero-shot protein-variant scoring, we kept the model in PyTorch, switched off Flash Attention to use the standard attention path, and compiled the transformer blocks with torch.compile. No model rewrite was needed. A single Trainium2 instance reached 490.8 variants per second, enough to score roughly 42.4 million a day at that rate.
At batch size 16 and sequence length 512, that was 18.6% higher throughput than our H100 baseline, with nearly identical ROC-AUC. The cost per variant was about 60% lower at the Capacity Block reservation rates we used in August 2026. For the code, setup, and full benchmarks, take a look at our ESMC-on-Trainium2 post. João Correia, Telmo Felgueira, Tiago Gonçalves, and I wrote it with Jim Burtoft and Louise Ping, our friends at AWS’s Annapurna Labs.
Tiago took that ESMC work to Sword AI Summit 2026 in Portugal. In his talk, “Foundation Models on AWS Trainium,” he shared our experience running the model on Trainium2.
What’s next
And we’re still cooking. Alongside our applied research, we’re working with two major clients, one in food delivery and one in streaming entertainment. We’re stress-testing ranking models in native PyTorch with one and preparing to work on LLMs with the other. The early ranking-model results have us excited. Our next steps are to keep testing the inference code, compare serving configurations, and get the LLM work underway.
Tiago and I are now competing in Phase 2 of the AWS Trainium Frontier Competition with The Optimal Batch @ Loka. As of October 19, 2026, we’re third out of ten teams on the public leaderboard. It’s exciting to see how far we’ve come, and we’re still pushing to improve our results.
What excites me now is how much more we can explore with the PyTorch tools we already know, whether we’re testing a research idea or building for a client. There’s still hardware-specific work to do, from checking compatibility to tuning performance, and we’re eager to get into it. I’m grateful to everyone at Loka and Annapurna Labs who’s been part of this journey since 2024. We’ll keep sharing what we learn as the client work and the competition progress.
Give TorchNeuron a try on AWS Silicon. And if you’d like a hand bringing your own models and workloads over, reach out to us at Loka. We’d love to help you get them running. Let’s put TorchNeuron to work.