The substantial overhead of NKI kernel construction in the Python call stack.

[dinghongsong]

Describe the bug

As previously discussed, the current optimized NKI matrix multiplication kernel is still slower than the matrix multiplication using torch.matmul or torch.einsum measured in the benchmark_sampling of the NxDI for two reasons:

1. In torch.matmul, the tensor data on the state buffer is directly loaded into the TensorMatrix engine for matrix multiplication. However, in the NKI kernel, NKI enforces that all inputs and outputs to a kernel must be within HBM. This means that the tensor data must be taken from the state buffer and written to HBM outside the kernel scope. Then, within the kernel scope, the tensor data is loaded back from HBM into the state buffer.

2. The NKI kernel construction in the Python call stack causes significant overhead.

Expected Behavior

Under Python's time.time() measurement, the NKI matrix multiplication kernel should be faster than torch.matmul. The NKI kernel construction overhead in the Python call stack should be minimal.

Current Behavior

Under Python's time.time() measurement, the NKI matrix multiplication kernel is much slower than torch.matmul.

Reproduction Steps

torch_matmul.py:

from neuronxcc import nki
import neuronxcc.nki.language as nl
import neuronxcc.nki.isa as ni
import numpy as np
import torch 
import time
import neuronxcc.nki.isa as nisa
from torch_xla.core import xla_model as xm
import os
os.environ["NEURON_FRAMEWORK_DEBUG"] = "1"
os.environ["NEURON_CC_FLAGS"]= " --disable-dge "


if __name__ == "__main__":
  K, M, N = 128, 128, 512
  device = xm.xla_device()
  cpu = torch.device('cpu')

  A = torch.rand((M, K), dtype=torch.bfloat16, device=device)
  B = torch.rand((K, N), dtype=torch.bfloat16, device=device)

  for _ in range(100):
    start_time = time.time()
    output_torch = torch.matmul(A, B)
    xm.mark_step()
    xm.wait_device_ops()
    end_time = time.time()
    print(f"output_torch={output_torch}")
    print("torch matmul time (ms): ", (end_time - start_time ) * 1000)

nki_matmul.py:

from neuronxcc import nki
import neuronxcc.nki.language as nl
import neuronxcc.nki.isa as ni
import numpy as np
import torch 
import time
import neuronxcc.nki.isa as nisa
from torch_xla.core import xla_model as xm
import os
os.environ["NEURON_FRAMEWORK_DEBUG"] = "1"
os.environ["NEURON_CC_FLAGS"]= " --disable-dge "

@nki.jit
def nki_matmul_basic_(lhsT, rhs):
  """NKI kernel to compute a 128x128x512 matrix multiplication operation

  Args:
      lhsT: an input tensor of shape [128,128], a left hand side argument of the
        matrix multiplication, delivered transposed for optimal performance
      rhs: an input tensor of shape [128,512], a right hand side argument of the
        matrix multiplication
  Returns:
      result: the resulting output tensor of shape [128,512]
  """
  result = nl.ndarray((128, 512), dtype=lhsT.dtype, buffer=nl.shared_hbm)

  i_lhsT_p, i_lhsT_f = nl.mgrid[0:128, 0:128]
  i_rhs_p, i_rhs_f = nl.mgrid[0:128, 0:512]
  i_out_p, i_out_f = nl.mgrid[0:128, 0:512]

 
  lhs_tile = nl.load(lhsT[i_lhsT_p, i_lhsT_f])
  rhs_tile = nl.load(rhs[i_rhs_p, i_rhs_f])

  
  result_psum = nisa.nc_matmul(lhs_tile, rhs_tile)

 
  result_sbuf = nl.copy(result_psum, dtype=result.dtype)
  nl.store(result[i_out_p, i_out_f], value=result_sbuf)

  return result


if __name__ == "__main__":
  K, M, N = 128, 128, 512
  device = xm.xla_device()
  cpu = torch.device('cpu')

  A = torch.rand((M, K), dtype=torch.bfloat16, device=device)
  B = torch.rand((K, N), dtype=torch.bfloat16, device=device)
  
  for _ in range(100):
    start_time = time.time()
    output_nki = nki_matmul_basic_(A.T, B)
    xm.mark_step()
    xm.wait_device_ops()
    end_time = time.time()
    print(f"output_nki={output_nki}")
    print("nki matmul time (ms): ", (end_time - start_time ) * 1000)

Regression Issue

• Select this option if this issue appears to be a regression.

Possible Solution

No response

Additional Information/Context

No response

neuronx-cc version used

NeuronX Compiler version 2.16.372.0+4a9b2326

Framework(s) and their versions used (JAX, PyTorch, etc..)

No response

[EmilyWebber]

Hi Dinghong! What's happening here is that the Torch invocation is using the highest levels of the NeuronX compiler, where we have already pushed many optimizations for common operations like matmul.

Because NKI enters the compiler at a lower level, these same optimizations are not immediately present.

However when you use the nki.benchmark decorator, which bypasses any overhead introduced by the framework while also passing multiple invocations to easily catch the compiler cache, then you can clearly see the speed up.

We will add documentation for both of these aspects.

[dinghongsong]

Hi Emily,

Thank you for your reply！

1. Yes, by adding the nki.benchmark decorator, we can obtain the execution time of the NKI kernel, and the measured time is consistent with the time reported by the Neuron Profiler.

   For example, under the nki.benchmark decorator, the measured time for nki_matmul_fully_optimized_ is 6.98 ms (P99). This matches the result I got when using torch.tensor as input instead of nt.tensor and running it with the Neuron Profiler — also 6.98 ms.

However, when using torch.tensor as input, the profiler also shows an additional overhead of around 4 ms, which corresponds to Reason 1.

2. Most of the time, we need to integrate NKI kernels into models written with PyTorch and measure the end-to-end runtime. In this case, we are not measuring the time of a single kernel, but the end-to-end execution time of the entire model. To do this, we need to use benchmark_sampling from the NxDI, which is based on Python's time.time() for measurement. However, after integrating the NKI kernel, we observed that the measured runtime is significantly slower compared to using the original torch.matmul or torch.einsum, due to Reason 2.

   This phenomenon has not only been observed by us, but also by others.

Can the NKI team help address this issue, so that we can achieve greater speedups when integrating NKI kernels into models? This might better showcase the advantages of NKI and Trainium.

[mrkcath-aws]

Hi @dinghongsong the NKI team is looking into this. Thank you for your patience!