Investigate async execution / mark_step pipelining

[scttfrdmn]

Sub-issue of #33. Current dispatch does matmul_kernel → implicit sync via .to(\"cpu\"). Unclear whether the 4 ms host→XLA transfer and 567 µs kernel actually pipeline or run serially.

If serial: explicit xm.mark_step() between transfer and kernel dispatch may let them overlap, shaving up to the kernel's worth of time off the critical path.

Action: instrument _nki_matmul with torch_xla.sync() calls between phases and compare to current wall-clock. Document findings; if pipelining is already happening, close with that note.

Referenced profile: #33 (comment)

[scttfrdmn]

The dispatch timing sweep from #33 answers this directly — closing with findings.

Short answer: kern and to_cpu already partially overlap; no additional pipelining is possible without fixing #39.

---

What the timing data shows

From docs/design/matmul_dispatch_findings.md (trn1.2xlarge, 1024² fp32):

Step	Time
to_xla (host→device)	1.502 ms
kern (XLA submission)	0.684 ms
to_cpu (device→host)	0.970 ms
Sum	3.156 ms
nki_tot (wall-clock)	2.935 ms

nki_tot (2.935 ms) < sum of steps (3.156 ms) by 0.22 ms. The gap means XLA graph submission (kern) and device→host DMA (to_cpu) already partially overlap — the Neuron Runtime begins streaming results back while the next graph is being submitted.

At 2048² pre-pinned: kern=0.973 + to_cpu=3.144 = 4.117 ms vs pinned_tot=4.116 ms. These are essentially equal — in the pre-pinned path, the steps are already maximally pipelined.

Why mark_step() can't be removed

_to_xla() calls xm.mark_step() when operands are already on XLA. This forces the lazy XLA graph to flush before the next kernel dispatch. Removing it would let the lazy evaluator fuse consecutive NKI kernel invocations into one graph — which is exactly the condition that triggers the trn2-only shared-memory codegen described in #39. The sync is a necessary correctness barrier on trn1, not a perf regression.

Conclusion

No improvement available at the mark_step level. The existing ~0.22 ms overlap at 1024² is as much async as the current architecture permits. The next lever for reducing per-dispatch overhead is #39 (upstream AWS fix) or Phase 3 batching (already shipped in v0.5.0 as multi_einsum homogeneous batching).

Pipelining note added to docs/design/matmul_kernel_profile.md.