Stop hand-porting kernels: migrate PyTorch and Triton kernels to AWS Trainium with an AWS Transform custom agent

Demand for generative AI has outpaced the supply of GPUs. Procurement delays and capacity blocks that sell out months ahead now gate how fast enterprises can ship AI workloads, and cost pressure compounds the problem at production scale. AWS Trainium gives customers a second path: Amazon EC2 Trn2 instances deliver 30 to 40 percent better price-performance than comparable GPU-based instances, and because AWS designs the silicon, customers gain access to capacity built and scaled by AWS directly. For many teams, moving inference and training to Trainium is now a business decision as much as an engineering one.

AWS Transform is an agentic service for migration and modernization that automates your move from VMware, Hyper-V, and bare metal physical servers to Amazon Elastic Compute Cloud (Amazon EC2), Amazon Elastic Block Store (Amazon EBS), and Amazon FSx. But migrating to Trainium also requires transforming the kernel, without which the migration can stall. This is hard because Triton and CUDA kernels are written for a GPU’s execution model: many lightweight threads, warps, shared memory, and implicit scheduling. Trainium works differently. It exposes explicit compute engines (tensor, vector, scalar) and an explicit on-chip memory hierarchy (SBUF, PSUM) that the kernel has to manage directly. Porting a kernel means re-expressing its data movement and tiling for Trainium’s hardware model which is hard for a generic migration tool because a numerically wrong version silently corrupts results and a slow one defeats the point of migrating at all. So the kernel becomes a ticket that reads “rewrite this for Trainium,” and the port stalls there. Custom kernels are a small fraction of a codebase but drive most of a model’s runtime performance, which makes this the highest-value gate in the migration.

This post introduces a custom AWS Transform agent that ports PyTorch and Triton kernels to Neuron Kernel Interface (NKI), the low-level programming interface for Trainium. Given a PyTorch nn.Module or a @triton.jit kernel from a repository, the agent generates an @nki.jit kernel, then compiles and runs it on a real Trainium device to verify that it produces the same numbers and profiles against the original to confirm it is fast. When either check fails, the agent feeds the compiler error or the profiling gap back in and iterates on the kernel. Only after a kernel passes both checks on silicon does it ship as a pull request with the evidence attached. Nothing is claimed that wasn’t measured on real hardware.

One result shaped the whole design. Asked to write an NKI kernel with no tools, Claude Opus 4.8 solves 6 percent of our 150-kernel benchmark. The same model, wrapped in a compile-verify-fix loop against a real Trainium device, compiles 83 percent of the benchmark; on the 57 weight-free tasks whose output can be checked from inputs alone, 88 percent match the PyTorch reference on-device (np.allclose, atol = rtol = 1e-3). We attribute this swing to the feedback loop.

By releasing this as a custom agent in AWS Transform, users can orchestrate kernel update as part of their migrations.

Design principles

Most agent systems are point solutions, one agent, one problem, one workflow. What we learned building this kernel agent is that the value compounds when the pieces are composable. The same Model Context Protocol (MCP) server that compiles and verifies kernels also serves the AWS Transform console, the AWS Transform (ATX) CLI, Claude Code, Kiro, and Codex, no fork required. The same reward endpoint that scored generations during reinforcement learning (RL) training validates candidates in production, so there is no train-to-production drift. The same model router that picks Claude Opus for novel ops falls back to a fine-tuned specialist for in-distribution shapes and a customer can pin either without touching the rest of the stack. Composability means each capability is independently deployable, independently valuable, and more powerful in combination than in isolation.

Solution overview

The plugin packages as a single Agent Plugin, a Skill (steering) plus an MCP server (tools), that registers as an agent in the AWS Transform agent registry, and is accessible from AWS Transform webapp, in Kiro via the AWS Transform Kiro Power, and also installs directly into Claude Code, Kiro, Codex, and any MCP-aware agent, as well as AWS Batch for fleet-scale runs. Under the hood it uses Amazon Bedrock AgentCore Runtime, Strands Agents, Claude on Amazon Bedrock, Bedrock Custom Model Import for a fine-tuned specialist, and a long-lived Trainium-backed verification service.

The agent runs the AWS Transform standard workflow (Discovery → Scope → Assessment → Requirements → Tasks → Execute → Diff/PR) with two human-in-the-loop gates rather than the usual one, because numerical equivalence on a kernel rewrite is a sign-off-worthy decision:

1. Discovery: scan repo for @triton.jit and hot nn.Module.forward
2. Scope (Gate 1): user picks kernels + model backend (auto / opus / specialist)
3. Assessment: skill lookup + dry-compile a stub per kernel
4. Requirements: draft accuracy + perf budget per kernel
5. Approval (Gate 2): user signs off the budget
6. Execute: per kernel: skill_lookup → generate → compile → verify → profile (on failure, feed the error back and retry)
7. Diff / PR: emit a unified diff or open a pull request

One MCP server, every surface

An MCP server-based interface contract makes it easy to plug this agent in the surface of your choice. AWS Transform is built on the same MCP-first pattern that allows importing of external agents, and Kiro also launched a Skill+MCP layout this year. So, we don’t need to fork the codebase when a surface changes. We mirror the layout of AWS Transform and Kiro exactly: MCP tool layer plus a SKILL.md steering layer, so the plugin installs from the same marketplace with the same /plugin install ergonomics.

A real Trainium device in the reward path

Compile, verify, and profile all go through a long-lived Flask service on a trn1.2xlarge. It is the only component that touches silicon, and it is the same endpoint our training pipeline used to score generations during reinforcement learning. The production reward signal is therefore identical to the training reward signal, no simulator-to-hardware drift. For a customer, that is the difference between “the agent says it works” and “we have a measured np.allclose pass and a 12 µs p50 latency on a real NeuronCore.”

The server uses subprocess-isolated handlers, compile, verify, profile, and baseline. Each validates the caller payload (256 KB source cap, identifier-style kernel name, a (shape, dtype) allowlist), runs the kernel in a fresh subprocess, and emits a single JSON line. A Neuron Executable File Format (NEFF) load failure or runtime fault crashes the subprocess, never the service. The baseline handler runs the user’s PyTorch reference under torch.compile(backend=“openxla”) on the same device, so the NKI-vs-PyTorch comparison is apples-to-apples on identical silicon.

Pluggable models with a router

The agent picks among three backends per invocation. Claude Opus on Bedrock is the default for novel ops and unusual shapes, no fine-tune required, and we inherit every new Claude release for free. A Bedrock Custom Model Import of Qwen3-Coder-30B with our supervised fine-tuning (SFT) specialist merged in handles in-distribution shapes at materially lower per-token cost. A small router in AgentCore picks per task: a high-confidence skill match (score ≥ 0.75) routes to the specialist; otherwise Opus takes it; two consecutive specialist failures escalate to Opus carrying the compiler error forward. Customers can pin a backend; the default is auto. The specialist routing is designed-in; the sample deployment ships Opus-only.

Prerequisites

• An AWS account with access to Claude on Amazon Bedrock (request it under Model access in the Bedrock console).
• An EC2 quota for at least one trn1.2xlarge.
• Permissions to create AgentCore Runtime endpoints, an IAM role for the AgentCore service, and a VPC security group scoped to the client CIDRs that need the reward port.
• The ATX CLI (AWS Transform) and agentcore CLI (Bedrock AgentCore) on your developer machine.
• Python 3.12+ with uv/uvx for the MCP server, and any MCP-aware agent (Claude Code, Kiro, or Codex are tested). Node 18+ and AWS CDK are also required for the CDK-based deployment.
• You do not need a GPU instance for the developer flow—all Trainium-touching work happens on the reward server.

Deploying the solution

The complete implementation is available in the sample-ATX-custom-NKI-Agent GitHub repository. Clone the repository and deploy the stack with a single command:

git clone https://github.com/aws-samples/sample-ATX-custom-NKI-Agent
cd sample-ATX-custom-NKI-Agent

# One-command deploy: provisions AgentCore and reward server CDK stacks
make deploy  # or: CICD_YES=1 ./cicd/deploy.sh --region us-east-1

# Install the MCP plugin into your IDE
/plugin marketplace add <your-marketplace>
/plugin install kernel-forge-aws-transform
# AWS Transform (atx) CLI surface — atx does not use /plugin; it runs a
# transformation definition and reads MCP servers from ~/.aws/atx/mcp.json
sed "s#REPLACE_WITH_ABSOLUTE_PATH#$(pwd)#" atx/mcp.json > ~/.aws/atx/mcp.json
atx mcp tools -s kernelforge-nki-mcp # confirm atx sees the 7 tools
atx custom def exec -p ./examples/cuda_rmsnorm_migration \
--configuration "additionalPlanContext=Follow atx/transformation-definition/transformation_definition.md"

Expected output:

• Step 1: gunicorn binds to port 5050 and boots one worker. The CDK stack’s health check confirms the server is reachable within the VPC. CDK deploys the reward server stack (Neuron DLAMI, systemd service, 4 gunicorn workers).
• Step 2: CDK deploys the AgentCore endpoint and auto-wires it to the reward server via repoint-agentcore.sh.
• Step 3: the plugin registers its kernel tools in the IDE, or atx mcp tools lists them for the CLI surface. What matters is that each step confirms the same three things: a reachable server, a live endpoint, and registered tools.

Then, in your agent: “Convert this Triton softmax kernel to NKI for Trainium.”

Composing with the AWS Transform migration agent

This agent is most useful as a post-migration step in a larger AWS Transform job, not as a standalone tool.

Consider a common case: a customer uses AWS Transform to migrate a GPU-based training workload to AWS Trainium for better price-performance. The model code, data pipelines, and infrastructure move cleanly, but the hand-tuned Triton or CUDA kernels were written for NVIDIA’s execution model and don’t run on Trainium. The migration succeeds at the infrastructure layer, but the custom kernels are now stranded: they have to be rewritten in NKI before the workload can serve traffic on the new hardware. This agent is built to close exactly that handoff where the infrastructure is migrated, but the kernel is stuck.

Every AWS Transform agent is a client of the same MCP contract, so chaining them is just passing the migrated repository from one job’s output to the next job’s input.

Table 1. Chaining a server migration into a kernel migration as a post-migration step.

In practice, the customer (or the delivery consultant driving the migration) runs the server migration first, reviews and merges that PR, then invokes this agent against the resulting repository. That invocation can happen from the AWS Transform console as a follow-up transformation, from the ATX CLI in the same pipeline, or by asking an IDE agent to “migrate the hot PyTorch kernels in this repo to Trainium.” The two jobs stay independent with different gates, different validation backends, but compose into one end-to-end GPU-to-Trainium path. Nothing about the kernel step depends on the server step having run first: this agent works equally well on a repository that AWS Transform never touched.

Results

All measurements below were taken on a single trn1.2xlarge in us-east-1, through the same reward path the agent uses at runtime.

A single kernel, end to end

We start with the simplest non-trivial NKI op: a row-wise softmax on a (128, 256) fp32 tensor, against a torch.softmax(x, dim=-1) reference.

Table 2. End-to-end reward path on a 128×256 fp32 softmax kernel.

Numerical agreement was max_abs_error = 2.57 × 10⁻⁷, max_rel_error = 9.85 × 10⁻⁶, 0 mismatched elements, well inside the atol = rtol = 1e-3 budget.

NKI vs. the PyTorch-on-Trainium baseline

A latency is a number; it is not a comparison until you put it next to what the customer would otherwise run on the same chip. The baseline handler runs the PyTorch reference under torch.compile(backend=“openxla”) on the same trn1.2xlarge, same dtype, same percentile schema.

Table 3. NKI kernel p50 vs. torch.compile(openxla) p50 (µs, fp32, same Trainium, 5 warmup + 50 iters). Each cell is kernel p50 / baseline p50 and the resulting speedup.

The curve is shape-dependent in a way that makes physical sense. At small shapes the baseline pays a fixed per-launch cost the hand-written kernel doesn’t, so the ratio is large; as per-call work grows, the baseline amortizes that overhead and the gap narrows. The kernel still wins at every shape we tested, 1.4 to 2.0× at the largest.

The headline: the loop matters more than the model

Three first-pass measurements are a confidence check, not a proof point. To stress the loop we ran the full 150-task subset of NKI Gen-Bench end to end on the same hardware, comparing Claude Opus 4.8 with and without the agent loop.

Table 4. Compile and numerical-verification rates on the 150-task subset of NKI Gen-Bench, measured on real Trainium (atol = rtol = 1e-3). Full results in Tang et al. (2026).

The last two rows are a fresh run of the released agent, reporting both compilation and on-device numerical verification. Numerical verification is reported on the 57 weight-free tasks whose reference can be reproduced from inputs alone; the remaining tasks carry learned weights not passed to the kernel and are scored on compilation. The 6 percent → 83 percent compile gap from identical weights is the clearest evidence we have that a tight feedback loop with ground-truth verification is the lever here.

What the agent can’t do (yet)

The hard tiers need the loop. Fused, multi-step kernels (Level 2) and full model components (Level 3) are near-zero in single-shot; the multi-turn loop is what makes them tractable at all, and pushing that further is the obvious next investment.

The reward server is single-Trn1 today. Concurrent requests are supported (4 workers with per-request UUID-based NEFF paths), but a multi-Trn fleet behind a load balancer is a deployment decision we describe but did not benchmark.

Kernels, not whole programs. Whole-program PyTorch to Neuron is torch-neuronx’s job. We do the kernel-level rewrites that block an otherwise-automatable migration, not a one-button “your CUDA app is now on Trainium.”

Clean up

Tear down both CDK stacks (reward server + AgentCore) in reverse order:

make destroy  # or: ./cicd/destroy.sh --region us-east-1

The local plugin install incurs no charges:

/plugin uninstall kernel-forge-aws-transform

Conclusion

We introduced an AWS Transform custom agent that converts PyTorch and Triton kernels to AWS Trainium NKI, compiling, verifying, and profiling every candidate on a real device behind two human approval gates. The reusable lessons are architectural: make the MCP server the contract so surfaces and backends can change underneath it, and don’t fake the device so the training-time and inference-time reward signals stay identical. And the measurement that should change how you think about agentic code generation: the same Claude Opus 4.8 that cold-solves 6% of our benchmark compiles 83% inside a verify-and-fix loop on real hardware, with 88% of the numerically-checkable kernels matching the reference on-device. The model is necessary, and the loop is what ships it.

References

Junjie Tang, Jun Huan, Hao Zhou, Yuhao Zhang, and Lin Wang. NKI-Agent: Domain-Specific Fine-Tuning and Agentic Tool Use for Neuron Kernel Generation. DL4C (Deep Learning for Code) Workshop at ICML 2026. https://arxiv.org/abs/2607.04395

About the authors