All current NKI kernels are FP32. Trainium natively supports BF16 and FP16 with meaningful throughput advantages (~2× on Tensor Engine for BF16).
Deliverable: dtype-parametric versions of _complex_gemm_kernel, _complex_linear_kernel, _complex_mul_kernel, butterfly_stage_kernel. Tests with tolerance appropriate to the reduced precision.
Benchmark expected: 1.5-2× speedup on large GEMM and STFT paths.
All current NKI kernels are FP32. Trainium natively supports BF16 and FP16 with meaningful throughput advantages (~2× on Tensor Engine for BF16).
Deliverable: dtype-parametric versions of
_complex_gemm_kernel,_complex_linear_kernel,_complex_mul_kernel,butterfly_stage_kernel. Tests with tolerance appropriate to the reduced precision.Benchmark expected: 1.5-2× speedup on large GEMM and STFT paths.