Skip to content

Spike: investigate SBUF residency for small-op dispatch overhead #47

Description

@scttfrdmn

Parent: #40. Research spike before implementation.

v0.7.0/v0.8.0 benchmarks show small ops (mask 64×32, GEMM 128, ComplexLinear 128→256) are 30-100× slower on NKI than CPU because of per-call XLA dispatch + tensor staging overhead.

Investigate:

  • Can intermediate tensors stay in SBUF across multiple kernel invocations?
  • What's the actual overhead breakdown (XLA trace time, HBM round-trip, kernel launch)?
  • Would batching short sequences of small ops into a single fused kernel help?

Deliverable: a decision memo (pick approach, or conclude the overhead is irreducible without a larger architectural change).

Metadata

Metadata

Assignees

No one assigned

    Labels

    nkiNKI kernel development (Trainium hardware)

    Type

    No type

    Projects

    No projects

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions