Benchmarks show small-op dispatch overhead dominates: even a 64×32 mask op is 96× slower on NKI than CPU. Investigate keeping intermediate tensors in SBUF across multiple kernel invocations (fused ops, persistent buffers) so the host doesn't have to round-trip XLA tensors per call.
Benchmarks show small-op dispatch overhead dominates: even a 64×32 mask op is 96× slower on NKI than CPU. Investigate keeping intermediate tensors in SBUF across multiple kernel invocations (fused ops, persistent buffers) so the host doesn't have to round-trip XLA tensors per call.