Emit kernel-aware NAS JSON artifacts

[Darkroom4364]

Summary

• add a machine-readable kernel-aware NAS report builder to scripts/nas_experiment.py
• support --json-output for deterministic comparison/ranking/kernel-cliff artifacts
• add tests for report schema, cross-hardware A100/T4/H100 ordering, and CLI artifact writes
• update the Kernel-aware architecture search: use kernel perf data to design models #81 research note with the artifact workflow

Refs #81

Tests

• python3 -m unittest tests.test_arch_cost_model
• python3 -m unittest tests.test_arch_cost_model tests.test_cli tests.test_ci_local_script
• python3 scripts/nas_experiment.py --hardware a100 --json-output /tmp/noeris-nas-a100.json
• python3 scripts/nas_experiment.py --hardware t4 --json-output /tmp/noeris-nas-t4.json
• python3 scripts/nas_experiment.py --hardware h100 --json-output /tmp/noeris-nas-h100.json
• python3 -m py_compile scripts/nas_experiment.py tests/test_arch_cost_model.py
• git diff --cached --check

Note: ruff is not installed locally.

Summary by CodeRabbit

• New Features

  • Added --json-output flag to generate machine-readable JSON reports containing latency predictions, performance rankings, timing breakdowns, and summary metrics for benchmark comparisons.

• Documentation

  • Updated instructions to document the JSON output feature and expected artifact contents for CI integration and benchmark tracking.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds JSON report output to the NAS experiment script. It introduces report-building functions to generate deterministic latency predictions and rankings across hardware targets, integrates them into the CLI with a --json-output flag, adds comprehensive test coverage, and updates the user-facing documentation.

Changes

NAS JSON Report Output

Layer / File(s)	Summary
Report building functions scripts/nas_experiment.py	Introduces _ranking_rows(), build_report() to orchestrate latency prediction and sweep computation, and write_report() for deterministic JSON serialization. Module docstring includes usage example with --hardware and --json-output.
CLI integration scripts/nas_experiment.py	Adds --json-output argument parsing and conditional report writing in main() with stdout confirmation message.
Test coverage tests/test_arch_cost_model.py	Adds JSON/subprocess/tempfile imports and three test cases: build_report schema validation, cross-hardware latency ordering, and CLI end-to-end artifact generation via subprocess.
Documentation updates docs/research/kernel-aware-nas.md	Updated A100 run command with --json-output, documented deterministic artifact contents for CI and benchmark packs, and added next steps item for multi-hardware comparison.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• 0sec-labs/noeris#105: Introduces the tile efficiency and tile penalty metadata (wasted_work_pct, tile_efficiency) in ArchitectureCostModel that the new JSON report generation depends on for sweep computation and summary fields.

Poem

🐰 A ranger hops through kernels bright,
JSON reports now gleam in light.
No more squinting at logs opaque—
Deterministic artifacts bake.
CI cheers as benchmarks flow,
Hardware sweeps steal the show! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.27% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'Emit kernel-aware NAS JSON artifacts' directly and concisely summarizes the primary change: adding functionality to emit JSON artifacts from the NAS experiment script.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feat/nas-json-artifacts

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

scripts/nas_experiment.py (1)

298-298: ⚡ Quick win

Consider building the report only when needed.

The report is built unconditionally at line 298 but only used when --json-output is provided. Moving the build_report() call inside the conditional block would avoid unnecessary computation when users don't request JSON output.

♻️ Proposed refactor

 model = ArchitectureCostModel(hardware=args.hardware)
-report = build_report(model)

 run_comparison(model)
 run_kernel_cliff_test(model)
 
 print(f"\n{'='*90}")
 print("  Key takeaways:")
 print("  - MLP (GeGLU) dominates for large models; attention dominates for long seq + many heads")
 print("  - Tile-unaligned dims (not multiple of 128) pay a measurable penalty")
 print("  - GQA (low num_kv_heads) saves attention time but QKV projection is still large")
 print("  - NAS should prefer hidden_dim, ffn_dim that are multiples of 128 (or 256)")
 print(f"{'='*90}\n")

 if args.json_output:
+    report = build_report(model)
     write_report(report, args.json_output)
     print(f"Wrote JSON artifact: {args.json_output}")

Also applies to: 311-313

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@scripts/nas_experiment.py` at line 298, The call to build_report(model) is
executed unconditionally and should be moved inside the branch that handles the
JSON output flag to avoid wasted work; modify the code so build_report(model) is
only invoked when the JSON output flag (e.g., args.json_output or the
'--json-output' branch) is true, create the report variable inside that
conditional, and update the blocks that currently reference report (lines
~311-313) to use that locally-built report instead of expecting a prebuilt
report.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@scripts/nas_experiment.py`:
- Line 298: The call to build_report(model) is executed unconditionally and
should be moved inside the branch that handles the JSON output flag to avoid
wasted work; modify the code so build_report(model) is only invoked when the
JSON output flag (e.g., args.json_output or the '--json-output' branch) is true,
create the report variable inside that conditional, and update the blocks that
currently reference report (lines ~311-313) to use that locally-built report
instead of expecting a prebuilt report.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 9d3febfb-4551-499c-9f98-5032eceacede

📥 Commits

Reviewing files that changed from the base of the PR and between 5946552 and fe2f1b5.

📒 Files selected for processing (3)

• docs/research/kernel-aware-nas.md
• scripts/nas_experiment.py
• tests/test_arch_cost_model.py