[Code Health] Scene tests hold an NPU lock through their CPU phases, and recompile every kernel on every CI run

[ChaoWao]

Category

Technical Debt (cleanup, refactor)

Component

Tests

Description

Two related inefficiencies in how scene tests consume devices. Both were quantified while investigating why qwen3_14b_decode needed to be excluded from the a2a3 sweep (fixed in #1601); the numbers below are measurements, not estimates.

a) Kernel compilation has no cross-run cache.

_compile_cache is a module-level dict cleared at session end, so a callable is compiled once per pytest session and never reused across CI runs. Every run recompiles every kernel from source with ccec.

Measured for qwen3_14b_decode on this repo's a2a3 box — 36 incores plus one orchestration:

	Wall
ordinary incore, each	1.1–1.3 s
orchestration	6.6 s
vendor FAI kernel, AIV variant	8.4 s (1330 KiB)
total for this one case	59 s

The key machinery already exists — l3_compile_cache_key() composes a stable key from (qualname, name, platform, runtime, pto_isa_token). It is only ever used against an in-memory dict. Adding a source-content hash and persisting to build/cache/kernels/ would make this survive: build/ is gitignored and the self-hosted runners' workspaces persist across runs, so CI would pay compilation once per kernel change instead of once per PR.

b) The device lock is held through work that never touches the device.

ci.yml wraps the whole pytest invocation in task-submit ... --run "python -m pytest ...", so kernel compilation, fixture construction and golden computation all run while holding a card.

For qwen3_14b_decode after #1601:

Phase	Wall	Needs a device?
kernel compilation	59 s	no
generate_inputs (38 GiB fixture)	13 s	no
compute_golden (40 layers, thread-capped)	~43 s	no
host→device upload, device run, comparison	remainder	yes — the device is busy for tens of milliseconds

So ~115 s of card time per run, for ~36 ms of device work. This is not specific to qwen: test_hello_worker — Worker(level=2).init() then close(), no kernels at all — takes 4.1 s, and the eleven allreduce collectives take 17.9–20.6 s each despite exercising five different algorithms (onephase / twophase / ring / bidirectional_ring / ibing), which says the cost is fixed overhead rather than the collective.

The contention this creates is measurable on the same job (st-onboard-a2a3, ci run 30507320146, job 90761014912): 1187 s of 2297 s — 52% — was spent queued waiting for free dies, across 8 separate task-submit acquisitions. The first one waited 976 s for four cards.

compute_golden already runs before the device run (scene_test.py:1288, :1367), so it is hoistable in principle. Compilation is the harder half, and (a) is its prerequisite: a pre-compile pass outside the lock only helps if the artifacts survive into the locked run.

Location

• simpler_setup/scene_test.py:39 — _compile_cache is an in-memory dict; :66 clears it at session end; :985 / :1029 are its only read and write
• simpler_setup/scene_test.py:54 — l3_compile_cache_key(), the key machinery that exists but never reaches disk
• .github/workflows/ci.yml:504, :525, :604 — the task-submit ... --run "pytest ..." wrappers that put every CPU phase inside the device lock

Proposed Fix

Two steps, in order, because the second depends on the first.

a) Persist the compile cache. Extend l3_compile_cache_key() with a hash of the kernel sources and their include closure, and back it with build/cache/kernels/<key>. Cross-run reuse on the self-hosted runners; GitHub-hosted sim jobs would need actions/cache to benefit, which can come later or not at all.

b) Acquire the card after the CPU phases. With (a) in place, the shape becomes: compile (and ideally build fixtures and goldens) outside task-submit, then take the lock only for the device run. Options worth weighing — a warm-up pass before the locked invocation, versus the harness acquiring the lock lazily at first device touch. The first is far simpler and probably enough.

Worth noting one trade-off for (b): the current 8 separate acquisitions cost ~210 s of re-queueing across the five trailing smoke steps, so there is a competing argument for holding one lock for the whole job. That is the opposite direction and the two need deciding together — shorter total device-seconds for the fleet, versus shorter wall time for this job.

Priority

Medium (minor risk, should fix in next few releases)