This document is relevant for: Inf2, Trn1, Trn1n

NKI API Reference Manual

Summary of different NKI API sets:

• nki top-level module contains APIs to decorate and simulate NKI kernels as well as NKI object types.

• nki.language consists of high-level compute and data movement APIs designed for ease-of-use. nki.language allows NKI programmers to transition from NumPy/Triton implementation to NKI quickly without the need to fully understand underlying NeuronDevice architecture. Most language APIs invoke one or more nki.isa APIs (that is, NeuronDevice hardware instructions) under the hood.

• nki.isa consists of low-level APIs that highly resemble hardware instructions in NeuronDevice ISA (instruction set architecture) designed to provide fine control over the hardware. These APIs expose all the programmable input parameters of the corresponding hardware instructions and also enforce the same tile-size and layout requirements as specified in NeuronDevice ISA.

• Other documents:

  • NKI API Common Fields documents common NKI API input parameters such as data types and masks, as well as common API behavior such as type promotion.

  • NKI API Errors captures common error types that are thrown by the NKI kernel compilation frontend, including syntax, tile-size and layout violation errors.

• nki
  • Framework
  • Types
• nki.language
  • Memory operations
  • Creation operations
  • Math operations
  • Bitwise operations
  • Logical operations
  • Tensor manipulation operations
  • Sorting/Searching operations
  • Collective communication operations
  • Iterators
  • Memory Hierarchy
  • Others
  • Data Types
  • Constants
• nki.isa
  • NKI ISA
  • Engine Types
• NKI API Common Fields
  • Supported Data Types
  • Supported Math Operators
  • Supported Activation Functions
  • NKI API Masking
  • NKI Type Promotion
• NKI API Errors
  • err_1d_arange_not_supported
  • err_annotation_shape_mismatch
  • err_cannot_assign_to_index
  • err_control_flow_condition_depending_on_arange
  • err_dynamic_control_flow_not_supported
  • err_exceed_max_supported_dimension
  • err_failed_to_infer_tile_from_local_tensor
  • err_indirect_indices_free_dim
  • err_local_variable_used_out_of_scope
  • err_nested_kernel_with_spmd_grid
  • err_nki_api_outside_of_nki_kernel
  • err_num_partition_exceed_arch_limit
  • err_num_partition_mismatch
  • err_read_modify_write_on_kernel_parameter
  • err_size_of_dimension_exceed_arch_limit
  • err_store_dst_shape_smaller_than_other_shape
  • err_tensor_access_out_of_bound
  • err_tensor_output_not_written_to
  • err_unexpected_output_dependencies
  • err_unsupported_memory
  • err_unsupported_mixing_basic_advanced_tensor_indexing

This document is relevant for: Inf2, Trn1, Trn1n