Kernels

• easydel.kernels.__init__
• easydel.kernels.collective_matmul
  • MatrixMultiplyMethod
  • calculate_mesh_dimension_size()
  • compute_device_linear_index()
  • create_distributed_matmul()
  • perform_all_gather_matmul()
  • perform_reduce_scatter_matmul()
  • prepare_matrix_for_all_gather()
  • prepare_matrix_for_reduce_scatter()
  • run_all_tests()
  • test_all_gather_matmul()
  • test_reduce_scatter_matmul()
• Cpu Ops
  • easydel.kernels.cpu_ops.__init__
  • Flash Attention Jax
  • Ring Attention Jax
• easydel.kernels.flash_attention
  • AttentionConfig
  • Backend
  • FlashAttention
  • Platform
  • create_flash_attention()
  • free_gpu_in_process()
  • get_device_memory_usage()
• Gpu Ops
  • Deprecated
  • easydel.kernels.gpu_ops.__init__
  • easydel.kernels.gpu_ops._utils
  • Flash Attention Triton
  • Matmul Triton
• easydel.kernels.matmul
  • custom_dot_general_kernel()
  • matmul()
  • matmul_benchmark()
  • matmul_grad_test()
  • matmul_test()
  • replace_dot_general_with_matmul()
  • test_dot_general_replacer()
• easydel.kernels.ring_attention
  • ring_attention()
• easydel.kernels.rms_norm
  • basic_layer_norm()
  • rms_norm()
  • test_bwd_call()
  • test_fwd_call()
• Tpu Ops
  • easydel.kernels.tpu_ops.__init__
  • Matmul Pallas
  • Paged Attention Pallas
  • Ragged Attention Pallas
  • Ring Attention Pallas
• easydel.kernels.utils
  • get_stride()