Index

A

Accelerator (class in xma.accelerator)

B

backward() (CustomOp static method)

bmm() (in module xma.functional.bmm)

C

ceil_divide() (in module xma.math) check_power_of_2() (in module xma.math) clamp() (in module xma.triton_utils.math) clip_gradients() (in module xma.torch_utils) continuous_count() (in module xma.functional.continuous_count) cpp_jit() (in module xma.jit) cpu (Accelerator attribute) cross_entropy() (in module xma.functional.cross_entropy)

ctx_needs_gradients() (in module xma.custom_op) ctx_save_for_backward() (in module xma.custom_op) cuda (Accelerator attribute) (KernelBackend attribute) CustomOp (class in xma.custom_op) cutotune() (in module xma.cutotune.tuner) CutoTuneConfig (class in xma.cutotune.config) CutoTuneParameter (class in xma.cutotune.parameter)

D

divide_if_divisible() (in module xma.math)

down_projection_triton_forward() (Experts method)

E

empty_like_contiguous() (in module xma.utils.tensor) enable_counters() (in module xma.counters) enable_kernels() (in module xma.inductor) ensure_contiguous() (in module xma.utils.contiguous)

Experts (class in xma.layers.moe) extra_repr() (Experts method) (GRU method) (RNN method) (XMAModule method)

F

forward() (CustomOp static method) (GRU method) (MoE method) (RNN method)

forward_backward_torch() (CustomOp static method) fused_linear_cross_entropy() (in module xma.functional.fused_linear_cross_entropy) fused_residual_add_rmsnorm() (in module xma.functional.fused_residual_add_rmsnorm) fused_residual_add_rmsnorm_backward_triton_kernel() (in module xma.functional.fused_residual_add_rmsnorm.triton_implementation.backward) fused_residual_add_rmsnorm_forward_triton_kernel() (in module xma.functional.fused_residual_add_rmsnorm.triton_implementation.forward)

G

get_accelerator() (Accelerator static method) get_alignment() (in module xma.utils.tensor) get_boolean_env_variable() (in module xma.utils.env) get_cartesian_product_cutotune_configs() (in module xma.cutotune.config) get_compatible_accelerator() (KernelBackend method) get_counter_value() (in module xma.counters) get_current_device() (Accelerator static method) get_cutotune_cache() (in module xma.cutotune.cache) get_fused_residual_add_rmsnorm_replacer() (in module xma.inductor) get_kernel_backend() (Accelerator static method)

get_key_values() (CutoTuneConfig method) get_max_seqlen_and_max_seqlen_tensor() (in module xma.utils.cu_seqlens) get_next_power_of_2() (in module xma.math) get_num_elements_and_hidden_size() (in module xma.utils.tensor) get_powers_of_2() (in module xma.math) get_ptx_from_triton_kernel() (in module xma.utils.ptx) get_rmsnorm_replacer() (in module xma.inductor) get_sm_count() (Accelerator static method) get_triton_num_warps() (in module xma.utils.settings) GRU (class in xma.layers.gru) gru() (in module xma.functional.gru)

I

increment_counter() (in module xma.counters) init_inductor() (in module xma.inductor) is_condition_valid() (CutoTuneConfig method) is_counter_enabled() (in module xma.counters)

is_cute_dsl_available() (in module xma.utils.packages) is_torch_neuronx_available() (in module xma.utils.packages) is_torch_xla_available() (in module xma.utils.packages) is_triton_available() (in module xma.utils.packages)

K

KernelBackend (class in xma.accelerator)

L

leaky_relu() (in module xma.triton_utils.math)

leaky_relu_backward() (in module xma.triton_utils.math)

M

matmul() (in module xma.triton_utils.matmul) module xma xma.accelerator xma.constants xma.counters xma.custom_op xma.cute_dsl_utils xma.cute_dsl_utils.math xma.cute_dsl_utils.utils xma.cutotune xma.cutotune.cache xma.cutotune.config xma.cutotune.parameter xma.cutotune.tuner xma.functional xma.functional.bmm xma.functional.bmm.triton_implementation xma.functional.continuous_count xma.functional.continuous_count.cuda_implementation xma.functional.cross_entropy xma.functional.cross_entropy.triton_implementation xma.functional.fused_linear_cross_entropy xma.functional.fused_residual_add_rmsnorm xma.functional.fused_residual_add_rmsnorm.triton_implementation xma.functional.fused_residual_add_rmsnorm.triton_implementation.backward xma.functional.fused_residual_add_rmsnorm.triton_implementation.forward xma.functional.gru xma.functional.gru.triton_implementation xma.functional.gru.triton_implementation.backward xma.functional.gru.triton_implementation.forward xma.functional.gru.utils xma.functional.rmsnorm xma.functional.rnn xma.functional.rnn.triton_implementation xma.functional.rnn.triton_implementation.backward xma.functional.rnn.triton_implementation.forward xma.functional.sequence_packing xma.functional.sequence_packing.cuda_implementation xma.functional.sequence_packing.triton_implementation xma.functional.softmax xma.functional.softmax.triton_implementation xma.functional.softmax.triton_implementation.backward xma.functional.softmax.triton_implementation.forward xma.functional.swiglu xma.functional.swiglu.cuda_implementation xma.functional.swiglu.cuda_implementation.backward xma.functional.swiglu.cuda_implementation.forward xma.functional.swiglu.nki_implementation xma.functional.swiglu.nki_implementation.backward xma.functional.swiglu.nki_implementation.forward xma.functional.swiglu.pallas_implementation xma.functional.swiglu.pallas_implementation.backward xma.functional.swiglu.pallas_implementation.forward xma.functional.swiglu.triton_implementation xma.functional.swiglu.triton_implementation.backward xma.functional.swiglu.triton_implementation.forward xma.functional.swiglu_packed xma.inductor xma.jit xma.layers xma.layers.gru xma.layers.moe xma.layers.moe.triton_implementation xma.layers.moe.triton_implementation.group_backward_kernel xma.layers.moe.triton_implementation.group_kernel xma.layers.moe.triton_implementation.scatter_kernel xma.layers.rnn xma.math xma.module xma.torch_utils xma.triton_utils xma.triton_utils.math xma.triton_utils.matmul xma.utils xma.utils.contiguous xma.utils.cu_seqlens xma.utils.debugging xma.utils.env xma.utils.packages xma.utils.ptx xma.utils.random xma.utils.settings xma.utils.tensor

MoE (class in xma.layers.moe)

N

nki (KernelBackend attribute)

P

pack_sequence() (in module xma.functional.sequence_packing) pack_unpack_sequence_triton_kernel() (in module xma.functional.sequence_packing.triton_implementation)

pallas (KernelBackend attribute) partialize_and_update_signature() (in module xma.inductor) print_gradient() (in module xma.utils.debugging)

R

reset_counters() (in module xma.counters) reset_parameters() (Experts method) (GRU method) (RNN method) rmsnorm() (in module xma.functional.rmsnorm)

RNN (class in xma.layers.rnn) rnn() (in module xma.functional.rnn) rocm (Accelerator attribute) (KernelBackend attribute) run() (CustomOp class method)

S

scattered_experts() (in module xma.layers.moe.triton_implementation) set_seed() (in module xma.utils.random) sigmoid() (in module xma.cute_dsl_utils.math) (in module xma.torch_utils) (in module xma.triton_utils.math) sigmoid_backward() (in module xma.triton_utils.math) silu() (in module xma.triton_utils.math) silu_backward() (in module xma.triton_utils.math) softmax() (in module xma.functional.softmax) softmax_backward_triton_kernel() (in module xma.functional.softmax.triton_implementation.backward) softmax_forward_triton_kernel() (in module xma.functional.softmax.triton_implementation.forward) swiglu() (in module xma.functional.swiglu)

swiglu_backward_cuda_jit() (in module xma.functional.swiglu.cuda_implementation.backward) swiglu_backward_cuda_kernel() (in module xma.functional.swiglu.cuda_implementation.backward) swiglu_backward_nki_kernel() (in module xma.functional.swiglu.nki_implementation.backward) swiglu_backward_pallas_jit() (in module xma.functional.swiglu.pallas_implementation.backward) swiglu_backward_pallas_kernel() (in module xma.functional.swiglu.pallas_implementation.backward) swiglu_forward_cuda_jit() (in module xma.functional.swiglu.cuda_implementation.forward) swiglu_forward_cuda_kernel() (in module xma.functional.swiglu.cuda_implementation.forward) swiglu_forward_nki_kernel() (in module xma.functional.swiglu.nki_implementation.forward) swiglu_forward_pallas_jit() (in module xma.functional.swiglu.pallas_implementation.forward) swiglu_forward_pallas_kernel() (in module xma.functional.swiglu.pallas_implementation.forward) swiglu_packed() (in module xma.functional.swiglu_packed) synchronize() (Accelerator static method)

T

tanh() (in module xma.cute_dsl_utils.math) (in module xma.torch_utils) (in module xma.triton_utils.math) tanh_backward() (in module xma.triton_utils.math) torch (KernelBackend attribute)

torch_forward() (Experts method) torch_tensor_to_cute_tensor() (in module xma.cute_dsl_utils.utils) tpu (Accelerator attribute) trainium (Accelerator attribute) triton (KernelBackend attribute)

U

unpack_sequence() (in module xma.functional.sequence_packing)

up_projection_triton_forward() (Experts method)

V

verify_accelerator() (KernelBackend method)

X

xma module xma.accelerator module xma.constants module xma.counters module xma.custom_op module xma.cute_dsl_utils module xma.cute_dsl_utils.math module xma.cute_dsl_utils.utils module xma.cutotune module xma.cutotune.cache module xma.cutotune.config module xma.cutotune.parameter module xma.cutotune.tuner module xma.functional module xma.functional.bmm module xma.functional.bmm.triton_implementation module xma.functional.continuous_count module xma.functional.continuous_count.cuda_implementation module xma.functional.cross_entropy module xma.functional.cross_entropy.triton_implementation module xma.functional.fused_linear_cross_entropy module xma.functional.fused_residual_add_rmsnorm module xma.functional.fused_residual_add_rmsnorm.triton_implementation module xma.functional.fused_residual_add_rmsnorm.triton_implementation.backward module xma.functional.fused_residual_add_rmsnorm.triton_implementation.forward module xma.functional.gru module xma.functional.gru.triton_implementation module xma.functional.gru.triton_implementation.backward module xma.functional.gru.triton_implementation.forward module xma.functional.gru.utils module xma.functional.rmsnorm module xma.functional.rnn module xma.functional.rnn.triton_implementation module xma.functional.rnn.triton_implementation.backward module xma.functional.rnn.triton_implementation.forward module xma.functional.sequence_packing module xma.functional.sequence_packing.cuda_implementation module xma.functional.sequence_packing.triton_implementation module xma.functional.softmax module xma.functional.softmax.triton_implementation module xma.functional.softmax.triton_implementation.backward module xma.functional.softmax.triton_implementation.forward module

xma.functional.swiglu module xma.functional.swiglu.cuda_implementation module xma.functional.swiglu.cuda_implementation.backward module xma.functional.swiglu.cuda_implementation.forward module xma.functional.swiglu.nki_implementation module xma.functional.swiglu.nki_implementation.backward module xma.functional.swiglu.nki_implementation.forward module xma.functional.swiglu.pallas_implementation module xma.functional.swiglu.pallas_implementation.backward module xma.functional.swiglu.pallas_implementation.forward module xma.functional.swiglu.triton_implementation module xma.functional.swiglu.triton_implementation.backward module xma.functional.swiglu.triton_implementation.forward module xma.functional.swiglu_packed module xma.inductor module xma.jit module xma.layers module xma.layers.gru module xma.layers.moe module xma.layers.moe.triton_implementation module xma.layers.moe.triton_implementation.group_backward_kernel module xma.layers.moe.triton_implementation.group_kernel module xma.layers.moe.triton_implementation.scatter_kernel module xma.layers.rnn module xma.math module xma.module module xma.torch_utils module xma.triton_utils module xma.triton_utils.math module xma.triton_utils.matmul module xma.utils module xma.utils.contiguous module xma.utils.cu_seqlens module xma.utils.debugging module xma.utils.env module xma.utils.packages module xma.utils.ptx module xma.utils.random module xma.utils.settings module xma.utils.tensor module xma_op() (in module xma.custom_op) XMAModule (class in xma.module)

Z

zeros_like_contiguous() (in module xma.utils.tensor)