This document is relevant for: Inf2, Trn1, Trn2

nki.compiler.enable_stack_allocator#

nki.compiler.enable_stack_allocator(func=None, log_level=50)[source]#

Use stack allocator to allocate the psum and sbuf tensors in the kernel.

Must use together with skip_middle_end_transformations.

from neuronxcc import nki

@nki.compiler.enable_stack_allocator
@nki.compiler.skip_middle_end_transformations
@nki.jit
def kernel(...):
  ...

This document is relevant for: Inf2, Trn1, Trn2