nki.compiler.enable_stack_allocator — AWS Neuron Documentation (original) (raw)

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