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

D

divide_if_divisible() (in module xma.math)

down_projection_triton_forward() (Experts method)

E

enable_counters() (in module xma.counters) enable_kernels() (in module xma.inductor) Experts (class in xma.layers.moe)

extra_repr() (Experts method) (GRU method) (RNN method)

F

forward() (CustomOp static method) (GRU method) (LinearAttention method) (M2RNN 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)

G

get_accelerator() (Accelerator static method) get_cartesian_product_xtune_configs() (in module xma.xtuner.config) get_compatible_accelerator() (KernelBackend method) get_counter_value() (in module xma.counters) get_current_device() (Accelerator static method) get_fused_residual_add_rmsnorm_replacer() (in module xma.inductor) get_kernel_backend() (Accelerator static method)

get_key_values() (XTuneConfig method) get_next_power_of_2() (in module xma.math) get_powers_of_2() (in module xma.math) get_rmsnorm_replacer() (in module xma.inductor) get_sm_count() (Accelerator static method) get_xtune_cache() (in module xma.xtuner.cache) GRU (class in xma.layers.gru.module) gru() (in module xma.layers.gru.op)

I

increment_counter() (in module xma.counters)

is_condition_valid() (XTuneConfig method) is_counter_enabled() (in module xma.counters)

K

KernelBackend (class in xma.accelerator)

L

linear_attention() (in module xma.layers.linear_attention.op)

LinearAttention (class in xma.layers.linear_attention.module)

M

M2RNN (class in xma.layers.m2rnn.layer) m2rnn() (in module xma.layers.m2rnn.op) module xma xma.accelerator xma.counters xma.custom_op xma.functional xma.functional.bmm xma.functional.continuous_count xma.functional.cross_entropy xma.functional.fused_linear_cross_entropy xma.functional.fused_residual_add_rmsnorm xma.functional.p_norm xma.functional.rmsnorm xma.functional.sequence_packing xma.functional.softmax xma.functional.swiglu xma.functional.swiglu_packed xma.inductor xma.jit xma.layers xma.layers.gru xma.layers.gru.module xma.layers.gru.op xma.layers.linear_attention xma.layers.linear_attention.module xma.layers.linear_attention.op xma.layers.m2rnn xma.layers.m2rnn.layer xma.layers.m2rnn.op xma.layers.moe xma.layers.rnn xma.layers.rnn.module xma.layers.rnn.op xma.math xma.optimizers xma.optimizers.sgd xma.optimizers.sgd.module xma.optimizers.sgd.op xma.xtuner xma.xtuner.cache xma.xtuner.config xma.xtuner.parameter xma.xtuner.tuner

MoE (class in xma.layers.moe) mps (Accelerator attribute) (KernelBackend attribute)

N

nki (KernelBackend attribute)

P

p_norm() (in module xma.functional.p_norm) pack_sequence() (in module xma.functional.sequence_packing)

pallas (KernelBackend attribute) partialize_and_update_signature() (in module xma.inductor)

R

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

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

S

SGD (class in xma.optimizers.sgd.module) sgd() (in module xma.optimizers.sgd.op) softmax() (in module xma.functional.softmax)

step() (SGD method) swiglu() (in module xma.functional.swiglu) swiglu_packed() (in module xma.functional.swiglu_packed) synchronize() (Accelerator static method)

T

torch (KernelBackend attribute) torch_forward() (Experts method)

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.counters module xma.custom_op module xma.functional module xma.functional.bmm module xma.functional.continuous_count module xma.functional.cross_entropy module xma.functional.fused_linear_cross_entropy module xma.functional.fused_residual_add_rmsnorm module xma.functional.p_norm module xma.functional.rmsnorm module xma.functional.sequence_packing module xma.functional.softmax module xma.functional.swiglu module xma.functional.swiglu_packed module xma.inductor module xma.jit module xma.layers module xma.layers.gru module xma.layers.gru.module module xma.layers.gru.op module

xma.layers.linear_attention module xma.layers.linear_attention.module module xma.layers.linear_attention.op module xma.layers.m2rnn module xma.layers.m2rnn.layer module xma.layers.m2rnn.op module xma.layers.moe module xma.layers.rnn module xma.layers.rnn.module module xma.layers.rnn.op module xma.math module xma.optimizers module xma.optimizers.sgd module xma.optimizers.sgd.module module xma.optimizers.sgd.op module xma.xtuner module xma.xtuner.cache module xma.xtuner.config module xma.xtuner.parameter module xma.xtuner.tuner module xma_op() (in module xma.custom_op) xtune() (in module xma.xtuner.tuner) XTuneConfig (class in xma.xtuner.config) XTunedFunction (class in xma.xtuner.tuner) XTuneParameter (class in xma.xtuner.parameter)