The nisa.nc_matmul in NKI is not as well-optimized as torch.matmul.

[dinghongsong]

Describe the bug

As the core computation kernels of LLMs, the matrix multiplication in NKI is significantly slower than torch.matmul, as measured by Python's time.time()

Relevant discussions can be found here (1), (2), (3).

I performed matrix multiplication with shapes [128,128] and [128,512] using both torch.matmul and nisa.nc_matmul, and obtained the respective web UI outputs with Neuron Profile as follows.

torch_matmul:

nki_matmul:

Their differences mainly appear after the x-axis reaches 20,000 (the area highlighted in the red box). The reason for nisa.nc_matmul being slower may be that it is not as well-optimized as torch.matmul.

Expected Behavior

Under Python's time.time() measurement, nki_matmul should be faster than torch.matmul.

Current Behavior

The matrix multiplication in NKI is significantly slower than torch.matmul, as measured by Python's time.time()

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)
    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..)

Python version 3.10.12
HWM version 2.16.0.372+4a9b2326
NumPy version 1.25.2
Running on trn1

[jluntamazon]

The NKI matmult API is a low level primitive that provides direct control and may not always produce optimal results. In a torch-native (non-NKI) compilation, many factors may change the performance characteristics of the matrix multiplication since the compiler uses heuristics to select the optimal execution. For NKI matrix multiplication help, see the following tutorial: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/tutorials/matrix_multiplication.html

In this example, the difference in the profiles has less to do with the matmult behavior and more to do with the DMA activity surrounding the matmult:

• In the torch-native profile, the tensor data on state buffer is directly loaded into the TensorMatrix engine for the matrix multiplication.
• In the NKI profile, you will observe that there are additional memory transfers that are not present in the torch-native profile. The API guarantee that NKI enforces is that all inputs and outputs to a kernel must be within HBM. This means that when you compile the NEFF it must take the tensor data which is on state buffer and write it to HBM outside of the kernel scope. Then within the kernel scope, the tensor data is loaded from HBM back into state buffer.

This behavior is expected and is currently enforced by the compiler. The reason why this performs badly in this scenario is because the computation size is small. In general, when the computation is large, this interfacing behavior should be relatively minor compared to the compute. For kernels that already use HBM-allocated tensors, there is no additional overhead.

Note that NKI is not an automatic optimization system, but rather a low-level programming interface that gives direct control over hardware resources. It's similar to CUDA in this respect - it provides the capabilities to write highly optimized code, but the optimization responsibility lies with the kernel developer.

The performance you achieve with NKI will directly reflect how you structure and optimize your kernel code. Some key optimization strategies you might need to consider:

• Tile size selection
• Memory access patterns
• Pipeline scheduling
• Local memory allocation

If you're looking for optimized matrix multiplication out of the box, torch.matmul is a good choice as it uses pre-optimized implementations. However, if you need custom behavior or want to optimize for specific matrix sizes or patterns, NKI gives you the fine-grained control to potentially exceed standard library performance.

[dinghongsong]

Hi @jluntamazon ,

Thanks for your reply!

You mentioned that the reason why this performs badly in this scenario is because the computation size is small. However, when performing matrix multiplication with larger dimensions, nki_matmul is still much slower than torch.matmul measured by Python's time.time(), even with the most optimized version from the tutorial. You can check the previous discussion here.

[jluntamazon]

In the previous discussion on #45, the timing code is incorrect. By default, torch-xla uses lazy execution to trace the graph and the NEFF is not executed until the output tensor requires some form of materialization.

The line below is lazily tracking the matmult operation and does not compile/execute on device:

output_torch = torch.matmul(A, B)

One way that you can ensure that you have compiled and executed the model is by using mark_step (this will be extremely slow on the first execution due to compilation):

output_torch = torch.matmul(A, B)
xm.mark_step()
# Alternative: torch_xla._XLAC._xla_sync_multi([output_torch], ['xla'], True, True)

This form of microbenchmarking is still not accurate since a cached graph will not be populated in the runtime before the first call. With a careful benchmarking setup we will still see a >2x difference in the timings. I measure 8.66ms on torch and 36.39ms on NKI. When I inspect the call stack, the NKI overhead is caused by the kernel construction in the python call stack (~29ms). When inspecting the NEFF, we can find that the NKI code has more performant device execution than the default torch heuristic in this case (torch=8.08ms, nki=6.95ms).

Torch Profile:

NKI Profile:

For accurate microbenchmarking, the recommendation is to be precise in defining the exact set of operations that is being measured. You can confirm from the profiler "Layer Summary" if you have correctly captured the desired operations. Secondly, benchmarking single operations is usually not useful since you are can easily measure overheads unrelated to the device execution.

[dinghongsong]

Thank you very much for your reply.

So for models with operators written using NKI, such as Llama, how do you think the time should be measured?
How should the speedup be measured before and after using the NKI kernel?

[dinghongsong]

In addition, if you don't mind, could you help me with another issue?

[jluntamazon]

So for models with operators written using NKI, such as Llama, how do you think the time should be measured?

For inference, we use tracing to avoid python overheads associated with building the HLO/NEFF. If you are using NxDI to build a model, using the NeuronBaseForCausalLM should allow you to trace the entire model into a single fixed graph and then never re-execute the python code. See: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-inference/developer_guides/onboarding-models.html

For training, we use normal python execution (with caching) which can incur some of the overheads you see in microbenchmarks. However, in the training context most of the the python time can be overlapped with device execution for different batches. The recommendation here is to measure the execution time over multiple batch iterations.

How should the speedup be measured before and after using the NKI kernel?

The least error-prone way we measure in a development setting is by looking at the profile results of the block which you are interested in. Once a model is fully developed, check end-to-end time. For example, see the NxDI benchmarking API: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/libraries/nxd-inference/developer_guides/onboarding-models.html#benchmarking

[dinghongsong]

I have reviewed the source code of benchmark_sampling in NxDI benchmarking API here, and it actually uses Python's time.time() to measure end-to-end execution time.

However, as previously mentioned, under Python's time.time() measurement, torch.matmul is significantly faster than nki_matmul. The NKI overhead is caused by kernel construction in the Python call stack.

Is it possible to reduce the overhead of kernel construction so that nki_matmul can be faster than torch.matmul under Python's time.time() measurement?

[dinghongsong]

Moreover, I found that one major drawback of developing with neuronx_distributed is the difficulty of setting breakpoints and stepping through the code in VSCode, as is possible with regular PyTorch code. This significantly reduces development efficiency, largely due to its lazy execution.

Do you have any suggestions to help with debugging in this case?

[chongmni-aws]

Hi @dinghongsong ,
The original questions on latency measurement have been addressed. Regarding debugging and putting breakpoints in the source code of NxD: this is required by the current design of NxD to generate a performant graph to run on Neuron. You can also explore eager mode execution: https://awsdocs-neuron.readthedocs-hosted.com/en/latest/frameworks/torch/torch-neuronx/programming-guide/training/pytorch-neuron-debug.html#using-eager-debug-mode