neuron-aerial

Port of NVIDIA Aerial cuPHY L1 (5G PHY) kernels to AWS Neuron/Trainium via NKI.

Golden reference = Aerial testVectors/ (HDF5, git-lfs) + our NumPy reference layer. See CLAUDE.md for non-negotiables, PLAN.md for the execution backlog, and INTEGRATION_PLAN.md for the agentic tooling layer.

Quick start (no hardware)

pip install -e ".[dev]"      # runtime + pytest/ruff, no Neuron toolchain
pytest -m "not device"       # the CI gate — must pass without Trainium
ruff check .

The [neuron] extra (neuronx-cc, nki, torch-neuronx, jax) is only needed for simulator (nki.simulate_kernel) and device runs; it is installed from the Neuron pip index, not PyPI.

Layout

neuron-aerial/
├── conformance/          # TV loaders, tolerances.yaml, report generation
├── reference/            # NumPy golden implementations (one module per kernel)
├── kernels/              # NKI kernels (one dir per kernel: impl, test, bench, SPEC.md)
├── pipelines/            # fused PUSCH/PDSCH graphs
├── bench/                # throughput + cost-per-slot scripts, results/ (json)
├── third_party/aerial/   # submodule, read-only (Aerial test vectors, git-lfs)
└── NOTICES.md

Test vectors (P0.2)

Aerial test vectors live in third_party/aerial/testVectors/ (git-lfs submodule). A human must provision them once:

git submodule add <aerial-repo-url> third_party/aerial
git -C third_party/aerial lfs pull

Then generate the inventory:

python -m conformance.inventory third_party/aerial/testVectors \
    --json conformance/tv_inventory.json --markdown

This writes conformance/tv_inventory.json (every HDF5 file's datasets: name/shape/dtype, classified into TV families) and prints the family-count table for this README. Unreadable files are recorded with an error rather than skipped, so the scan is complete and auditable.

Testing model

• pytest -m "not device" is the CI gate — runs on plain hardware, no Trainium, no TV corpus.
• Device tests are marked @pytest.mark.device and run only on a trn2 dev box.
• Numerics debugging order (never skip a stage): NumPy-ref vs TV → simulator vs NumPy-ref → device vs simulator. Never debug device-vs-TV directly.