xma

• xma.cute_dsl_utils
  • xma.cute_dsl_utils.math
    • sigmoid()
    • tanh()
  • xma.cute_dsl_utils.utils
    • torch_tensor_to_cute_tensor()
• xma.cutotune
  • xma.cutotune.cache
    • get_cutotune_cache()
  • xma.cutotune.config
    • CutoTuneConfig
      • CutoTuneConfig.get_key_values()
      • CutoTuneConfig.is_condition_valid()
    • get_cartesian_product_cutotune_configs()
  • xma.cutotune.parameter
    • CutoTuneParameter
  • xma.cutotune.tuner
    • cutotune()
• xma.functional
  • xma.functional.bmm
    • bmm()
  • xma.functional.continuous_count
    • continuous_count()
  • xma.functional.cross_entropy
    • cross_entropy()
  • xma.functional.fused_residual_add_rmsnorm
    • fused_residual_add_rmsnorm()
  • xma.functional.gru
    • gru()
  • xma.functional.rnn
    • rnn()
  • xma.functional.sequence_packing
    • pack_sequence()
    • unpack_sequence()
  • xma.functional.softmax
    • softmax()
  • xma.functional.swiglu
    • swiglu()
  • xma.functional.fused_linear_cross_entropy
    • fused_linear_cross_entropy()
  • xma.functional.rmsnorm
    • rmsnorm()
  • xma.functional.swiglu_packed
    • swiglu_packed()
• xma.layers
  • xma.layers.moe
    • Experts
      • Experts.down_projection_triton_forward()
      • Experts.extra_repr()
      • Experts.reset_parameters()
      • Experts.torch_forward()
      • Experts.up_projection_triton_forward()
    • MoE
      • MoE.forward()
  • xma.layers.gru
    • GRU
      • GRU.extra_repr()
      • GRU.forward()
      • GRU.reset_parameters()
  • xma.layers.rnn
    • RNN
      • RNN.extra_repr()
      • RNN.forward()
      • RNN.reset_parameters()
• xma.triton_utils
  • xma.triton_utils.math
    • clamp()
    • leaky_relu()
    • leaky_relu_backward()
    • sigmoid()
    • sigmoid_backward()
    • silu()
    • silu_backward()
    • tanh()
    • tanh_backward()
  • xma.triton_utils.matmul
    • matmul()
• xma.utils
  • xma.utils.contiguous
    • ensure_contiguous()
  • xma.utils.cu_seqlens
  • xma.utils.debugging
    • print_gradient()
  • xma.utils.env
    • get_boolean_env_variable()
  • xma.utils.packages
    • is_cute_dsl_available()
    • is_torch_neuronx_available()
    • is_torch_xla_available()
    • is_triton_available()
  • xma.utils.ptx
    • get_ptx_from_triton_kernel()
  • xma.utils.random
    • set_seed()
  • xma.utils.settings
    • get_triton_num_warps()
  • xma.utils.tensor
    • empty_like_contiguous()
    • get_alignment()
    • get_num_elements_and_hidden_size()
    • zeros_like_contiguous()

• xma.accelerator
  • Accelerator
    • Accelerator.cpu
    • Accelerator.cuda
    • Accelerator.get_accelerator()
    • Accelerator.get_current_device()
    • Accelerator.get_kernel_backend()
    • Accelerator.get_sm_count()
    • Accelerator.rocm
    • Accelerator.synchronize()
    • Accelerator.tpu
    • Accelerator.trainium
  • KernelBackend
    • KernelBackend.cuda
    • KernelBackend.get_compatible_accelerator()
    • KernelBackend.nki
    • KernelBackend.pallas
    • KernelBackend.rocm
    • KernelBackend.torch
    • KernelBackend.triton
    • KernelBackend.verify_accelerator()
• xma.constants
• xma.counters
  • enable_counters()
  • get_counter_value()
  • increment_counter()
  • is_counter_enabled()
  • reset_counters()
• xma.custom_op
  • CustomOp
    • CustomOp.backward()
    • CustomOp.forward()
    • CustomOp.forward_backward_torch()
    • CustomOp.run()
  • ctx_needs_gradients()
  • ctx_save_for_backward()
  • xma_op()
• xma.inductor
  • enable_kernels()
  • get_fused_residual_add_rmsnorm_replacer()
  • get_rmsnorm_replacer()
  • init_inductor()
  • partialize_and_update_signature()
• xma.jit
  • cpp_jit()
• xma.math
  • ceil_divide()
  • check_power_of_2()
  • divide_if_divisible()
  • get_next_power_of_2()
  • get_powers_of_2()
• xma.module
  • XMAModule
    • XMAModule.extra_repr()
• xma.torch_utils
  • clip_gradients()
  • sigmoid()
  • tanh()