This repository contains AI agents and skills for developing on AWS Neuron (Trainium/Inferentia) hardware, including NKI kernel development, profiling and debugging. For an overview of Neuron Agentic Development and the tools it offers for agent-enabled workflows with Neuron, see the overview of Neuron Agentic Development in the the public Neuron docs.
# from neuron pypi repository
pip install --upgrade neuron-agentic-development \
--extra-index-url https://pip.repos.neuron.amazonaws.com
# from wheel
pip install neuron_agentic_development-1.0-py3-none-any.whl
# from local github clone
git clone https://github.com/aws-neuron/neuron-agentic-development.git
cd neuron-agentic-development
pip install .
# Then deploy to your preferred tool:
deploy-neuron-agentic-development-to-kiro
# or
deploy-neuron-agentic-development-to-claude| Agent | Description |
|---|---|
| neuron-nki-agent | Unified NKI kernel development agent. Full lifecycle: writing kernels from PyTorch/NumPy/natural language, debugging compilation errors, profiling performance, optimizing bottlenecks, migrating between API versions, analyzing Perfetto traces, and NKI documentation lookup. |
| neuron-nki-writer-agent | NKI kernel authoring and modification. Translates from PyTorch/NumPy/natural language, adds shape/dtype support, refactors tiling strategies, and implements new features following Beta 3 API patterns. |
| neuron-nki-debugger-agent | Autonomous NKI kernel compilation error debugging. Analyzes compiler errors, searches documentation and code examples for fixes, applies corrections following simplicity over performance, and validates fixes. |
| neuron-nki-profile-analysis-agent | Profile and analyze NKI kernels on Neuron hardware. Captures execution traces, computes performance bounds, identifies bottleneck engines, and runs investigations to localize inefficiencies to NKI source lines. |
| neuron-framework-autoport-agent | A model porting agent to port GPU-compatible models to functionally accurate implementation on Neuron. Executes the full porting workflow including architecture analysis, implementation, compilation, inference testing, and validation. |
| neuron-framework-profiling-vllm-neuron-agent | Profiles vLLM-Neuron models on Trainium. Runs benchmarks with profiling enabled, captures device-level and host-level traces, and identifies output artifacts for downstream analysis. |
| neuron-framework-profile-analysis-vllm-neuron-agent | Ingests vLLM-Neuron profile artifacts and analyzes them to localize the performance bottleneck (compute, collective, or host). |
| neuron-parallelism-planner-agent | Plans tensor/pipeline/context-parallelism and data-parallel replica configurations for transformer LLM inference on Neuron, reconciling latency and throughput targets into a recommended recipe. |
| neuron-hardware-entitlement-agent | Computes the per-operator hardware-entitlement floor for a chosen inference parallelism configuration on Neuron and identifies which operator binds, to direct optimization effort. |
| Skill | Description |
|---|---|
| neuron-nki-writing | Write and modify NKI kernels. Covers new kernel creation from PyTorch/NumPy/natural language, editing existing kernels, adding shape/dtype support, refactoring tiling strategies, and implementing new features. |
| neuron-nki-debugging | Debug NKI compilation errors on Neuron hardware. |
| neuron-nki-docs | Research NKI documentation for API lookups, tutorials, error codes, and architecture details. |
| neuron-nki-profiling | Profile NKI kernels to analyze performance on Neuron hardware. |
| neuron-nki-profile-querying | Query and analyze NKI kernel profile data from neuron-explorer parquet files via SQL and Python. |
| neuron-framework-autoport | Port a GPU compatible model to NeuronX Distributed Inference for AWS Trainium/Inferentia. Handles the full workflow including architecture analysis, NeuronX implementation, compilation, inference testing, and accuracy validation. |
| neuron-framework-equivalence | Verifies functional equivalence between two implementations of the same model using a hierarchical algorithm. |
| neuron-framework-profiling-vllm-neuron | Profile vLLM-Neuron models on Trainium by capturing device and system traces from a live server or an offline benchmark run. |
| neuron-framework-profile-analysis-vllm-neuron | Ingest vLLM-Neuron execution artifacts with neuron-explorer and analyze them: identify prefill/decode graphs, compute per-engine bounds and DMA efficiency, and produce structured workload reports. |
| neuron-parallelism-planner | Plan tensor/pipeline/context-parallelism and context-bucket configurations for transformer LLM inference on Neuron using an analytic hardware-entitlement model. Estimates latency, throughput, and per-device memory before running on hardware. |
| neuron-hardware-entitlement | Compute per-operator compute and memory work for a single inference parallelism configuration and attribute a deployed graph's actual memory usage back to those floors, to localize which operator binds. |
We are evaluating the external contribution process. All capabilities undergo internal verification to ensure technical accuracy, security, and architectural alignment. In the interim, we welcome feedback and feature requests via Issues
Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.