Skip to content

trnsparse 0.3.1 — NKI 0.3.0 namespace + simulator dispatch

Choose a tag to compare

View all tags
@scttfrdmn scttfrdmn released this 14 Apr 21:08
· 28 commits to main since this release
v0.3.1
570f62e

Clears `[Unreleased]` with the three-commit trnsparse#23 landing. No kernel-body changes; no public API changes; no numeric drift. Upgrade is safe for anyone already on v0.3.0.

Changed

  • Migrated NKI imports to the nki.* namespace (NKI 0.3.0 Stable, Neuron SDK 2.29, April 2026). Legacy neuronxcc.nki.* shim is no longer used. pyproject.toml [neuron] extra gains nki>=0.3.0 alongside neuronxcc>=2.24 and torch-neuronx>=2.9. Hosts without an nki wheel (macOS, non-Linux archs) still hit HAS_NKI=False and get the torch fallback.
  • test CI job now filters -m "not neuron and not nki_simulator" so each test runs in exactly one job.

Added

  • TRNSPARSE_USE_SIMULATOR=1 dispatch branch through nki.simulate(kernel)(np_args). Bypasses torch_xla + NEFF compile; kernels run on CPU for correctness iteration. Hardware still owns perf numbers.
  • nki-simulator CI job on ubuntu-latest — installs nki>=0.3.0 from the AWS pip index, runs the simulator suite on every push/PR. Kernel correctness gate without AWS cost. Catches Python-trace-level errors (bad kwargs, dropped ops, shape mismatches); MLIR verifier errors remain hardware-only.
  • tests/test_nki_sim.py — curated simulator suite (4 tests: CSR aligned + rectangular, BSR block-dense + block-diagonal). Skips cleanly off-hardware.
  • scripts/run_simulator_tests.sh — SSM runner for the simulator suite on trn1.
  • tests/conftest.py — registers the nki_simulator pytest marker.

Validation

  • 37 CPU tests pass unchanged.
  • 15 hardware tests on trn1.2xlarge — all green post-migration (60s).
  • 4 simulator tests green on both trn1 and ubuntu-latest.

Closes #23.

Assets 2
Loading