Kernels
=======

.. toctree::
   :maxdepth: 2

   kernels.flash_attention_2
   kernels.rms_norm
   kernels.tpu_ops.pallas_gemm
   kernels.tpu_ops.pallas_ring_attention
   kernels.cpu_ops.jax_ring_attention
   kernels.cpu_ops.jax_mha_flash_attention_2
   kernels.utils
   kernels.gpu_ops.triton_gemm
   kernels.gpu_ops.pallas_gqa_flash_attention_2
   kernels.gpu_ops.triton_gqa_flash_attention_2
   kernels.gpu_ops.pallas_gemm
   kernels.gpu_ops.pallas_mha_flash_attention_2
   kernels.ring_attention
   kernels.gemm