nki(matrix_transpose): NKI (Trainium) implementation

[bowencui123]

NKI (AWS Trainium) implementation of matrix_transpose, split out of the consolidated NKI branch cecilia/feature/nki-vector-add (nki-all-operators, #259) so each operator can be reviewed independently. Supersedes #183 (older per-operator branch: legacy neuronxcc.nki imports; this is the migrated import nki version).

Files: A benchmarks/operators/matrix_transpose/impl_nki.py

Status: verified on trn2 (default + autotune; NKI 0.350 ms @ 4096x1024 fp16)

• includes NkiAutotuner wiring (search space over the kernel's tile/loop parameter; defaults unchanged)

Implementation by @Cecilia123li. Timing/identity infrastructure: #261; Trainium peak/roofline infra: #262.

🤖 Generated with Claude Code

https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL

Autotune (already on this branch)

NKI tunables	Triton counterpart	note
tile (square nc_transpose tile <=128)	BLOCK_TILE	—

autotune=False keeps the previous constants. Validation on trn2, case 0:

# initial run
[matrix_transpose] default : verify=OK (2s) 
[matrix_transpose] autotune: verify=OK (0s) last_config={'tile': 128} trace_records=1 
STUB_EXIT=0
Params   |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=1024   | fp16     |     0.0681 |       0.3506 |        0.19

[bowencui123]

I think the current NKI implementation is algorithmically correct, but the copy path may be unnecessarily expensive compared with both the Torch/Neuron baseline and the Triton/cuTile implementations.

Triton and cuTile both effectively perform a direct contiguous load/store copy:

input -> output

The current NKI implementation instead explicitly stages every tile through SBUF:

HBM -> SBUF -> HBM

via:

a_tile = nl.ndarray((p_sz, f_sz), dtype=a_input.dtype, buffer=nl.sbuf)

nisa.dma_copy(
    dst=a_tile,
    src=a_input[p_start:p_end, f_start:f_end],
)

nisa.dma_copy(
    dst=hbm_result_tile[p_start:p_end, f_start:f_end],
    src=a_tile,
)

For a pure matrix copy, this extra SBUF staging does not appear necessary. nisa.dma_copy supports HBM-to-HBM copies, so it may be worth testing a direct path such as:

nisa.dma_copy(
    dst=hbm_result_tile[p_start:p_end, f_start:f_end],
    src=a_input[p_start:p_end, f_start:f_end],
)

The performance pattern also suggests this is a bandwidth/path-efficiency issue rather than an algorithmic-complexity issue. For example, at N=4096:

• fp16: torch_nki = 0.1143 ms, NKI = 0.1738 ms
• bf16: torch_nki = 0.1144 ms, NKI = 0.1739 ms
• fp32: torch_nki = 0.2168 ms, NKI = 0.3294 ms

So NKI is consistently around 1.5x slower, with a very similar ratio across dtypes.

There is also one baseline issue worth checking. The Torch implementation is:

B = A + 0

On Torch/XLA, A + 0 may be algebraically simplified by the compiler, potentially turning the Torch baseline into an optimized copy/output-materialization path rather than an actual elementwise-add kernel.

Could we please:

1. Inspect the generated Torch HLO to confirm whether A + 0 is eliminated.
2. Benchmark a direct HBM-to-HBM NKI dma_copy.
3. Compare that result against the current HBM -> SBUF -> HBM implementation.

I think this would clarify whether the current gap is mainly caused by unnecessary SBUF staging. From the TileBench symmetry perspective, the current implementation is algorithmically equivalent to Triton/cuTile, but the memory-movement strategy is not yet as closely aligned.