perf(nki): switch matvec_kernel to Tensor Engine via nisa.nc_matmul (… #48
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: CI | |
| on: | |
| push: | |
| branches: [main] | |
| pull_request: | |
| jobs: | |
| lint: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v6 | |
| - uses: actions/setup-python@v6 | |
| with: | |
| python-version: "3.12" | |
| - uses: astral-sh/setup-uv@v5 | |
| with: | |
| enable-cache: true | |
| - run: uv pip install --system "ruff>=0.6" | |
| - run: ruff check . | |
| - run: ruff format --check . | |
| test: | |
| runs-on: ubuntu-latest | |
| strategy: | |
| matrix: | |
| python-version: ["3.10", "3.11", "3.12"] | |
| steps: | |
| - uses: actions/checkout@v6 | |
| - uses: actions/setup-python@v6 | |
| with: | |
| python-version: ${{ matrix.python-version }} | |
| - uses: astral-sh/setup-uv@v5 | |
| with: | |
| enable-cache: true | |
| - run: uv pip install --system -e ".[dev]" | |
| - run: pytest tests/ -v -x --tb=short -m "not neuron and not nki_simulator" --cov=trnsolver --cov-report=xml | |
| - name: Upload coverage reports to Codecov | |
| uses: codecov/codecov-action@v5 | |
| with: | |
| token: ${{ secrets.CODECOV_TOKEN }} | |
| slug: trnsci/trnsolver | |
| # Runs NKI kernels through nki.simulate(kernel)(numpy_args) on CPU. | |
| # Catches Python-trace-level errors (bad kwargs, dropped ops, shape | |
| # mismatches) pre-merge without AWS round-trips. MLIR verifier errors | |
| # remain hardware-only (simulator explicitly skips compile). NKI 0.3.0 | |
| # Stable is the canonical namespace — package is `nki`, not `neuronxcc.nki`. | |
| nki-simulator: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v6 | |
| - uses: actions/setup-python@v6 | |
| with: | |
| python-version: "3.12" | |
| - uses: astral-sh/setup-uv@v5 | |
| with: | |
| enable-cache: true | |
| - name: Install trnsolver + NKI simulator deps | |
| # --index-strategy unsafe-best-match: uv normally locks a package to | |
| # the first index it's found on; nki's transitive deps (e.g. islpy) | |
| # exist on both the Neuron repo and PyPI. Match pip's multi-index | |
| # behavior so resolution picks the newest available wheel per package. | |
| run: | | |
| uv pip install --system -e ".[dev]" | |
| uv pip install --system --index-strategy unsafe-best-match "nki>=0.3.0" --extra-index-url https://pip.repos.neuron.amazonaws.com | |
| - name: Run simulator-backed kernel tests | |
| env: | |
| TRNSOLVER_USE_SIMULATOR: "1" | |
| run: pytest tests/ -v -m nki_simulator --tb=short --cov=trnsolver --cov-report=xml --cov-append | |
| - name: Upload coverage reports to Codecov | |
| uses: codecov/codecov-action@v5 | |
| with: | |
| token: ${{ secrets.CODECOV_TOKEN }} | |
| slug: trnsci/trnsolver | |
| flags: simulator | |
| # pytest-benchmark bit-rot gate. Runs the benchmark suite with minimum | |
| # rounds on ubuntu-latest to catch import / fixture / API drift before | |
| # it hits the real perf measurement path. No perf gate — just execution. | |
| # Real perf numbers come from the hardware (`run_neuron_tests.sh`) and | |
| # g5 (`run_cuda_tests.sh`) runners. | |
| bench-smoke: | |
| runs-on: ubuntu-latest | |
| steps: | |
| - uses: actions/checkout@v6 | |
| - uses: actions/setup-python@v6 | |
| with: | |
| python-version: "3.12" | |
| - uses: astral-sh/setup-uv@v5 | |
| with: | |
| enable-cache: true | |
| - run: uv pip install --system -e ".[dev]" | |
| - name: Run benchmark smoke (exclude neuron + cuda) | |
| # --override-ini: bench_*.py is outside the test_*.py default pattern; | |
| # override rather than pollute the pyproject test config. | |
| run: | | |
| pytest benchmarks/ -v --benchmark-only -m "not neuron and not cuda" \ | |
| --override-ini 'python_files=bench_*.py' \ | |
| --benchmark-min-rounds=2 \ | |
| --benchmark-json=bench.json \ | |
| --tb=short | |
| - name: Upload benchmark JSON | |
| if: always() | |
| uses: actions/upload-artifact@v4 | |
| with: | |
| name: bench-smoke-results | |
| path: bench.json | |
| retention-days: 14 |