Tpu Ops

• easydel.kernels.tpu_ops.__init__
  • pallas_paged_attention()
  • pallas_prefill_attention()
  • pallas_ragged_decode()
• Matmul Pallas
  • easydel.kernels.tpu_ops.matmul_pallas.__init__
  • easydel.kernels.tpu_ops.matmul_pallas._caller
• Paged Attention Pallas
  • easydel.kernels.tpu_ops.paged_attention_pallas.__init__
  • easydel.kernels.tpu_ops.paged_attention_pallas._forward_pallas
  • easydel.kernels.tpu_ops.paged_attention_pallas._paged_attention
• Ragged Attention Pallas
  • easydel.kernels.tpu_ops.ragged_attention_pallas.__init__
  • easydel.kernels.tpu_ops.ragged_attention_pallas._caller
  • easydel.kernels.tpu_ops.ragged_attention_pallas._forward_pallas
• Ring Attention Pallas
  • easydel.kernels.tpu_ops.ring_attention_pallas.__init__
  • easydel.kernels.tpu_ops.ring_attention_pallas._backward_pallas
  • easydel.kernels.tpu_ops.ring_attention_pallas._forward_pallas
  • easydel.kernels.tpu_ops.ring_attention_pallas._ring_attention
  • easydel.kernels.tpu_ops.ring_attention_pallas._utils