maxtext.kernels.attention package

Attention kernels.

Submodules

• maxtext.kernels.attention.jax_flash_attention module
  • flash_attention_block_masked()
  • mask_blocker()
• maxtext.kernels.attention.ragged_attention module
  • get_mha_cost_estimate()
  • reference_mqa()
  • reference_mha()
  • reference_gqa()
  • ragged_flash_attention_kernel()
  • ragged_mqa()
  • ragged_mha()
  • ragged_gqa()
• maxtext.kernels.attention.splash_attention_kernel module
  • SegmentIds
    • SegmentIds.q
    • SegmentIds.kv
  • get_kernel_name()
  • attention_reference()
  • attention_reference_custom()
  • make_attention_reference()
  • make_masked_mha_reference()
  • make_masked_mqa_reference()
  • QKVLayout
    • QKVLayout.HEAD_DIM_MINOR
    • QKVLayout.SEQ_MINOR
  • from_head_minor()
  • BlockSizes
    • BlockSizes.block_q
    • BlockSizes.block_kv
    • BlockSizes.block_kv_compute
    • BlockSizes.block_q_dkv
    • BlockSizes.block_kv_dkv
    • BlockSizes.block_kv_dkv_compute
    • BlockSizes.block_q_dq
    • BlockSizes.block_kv_dq
    • BlockSizes.use_fused_bwd_kernel
    • BlockSizes.q_layout
    • BlockSizes.k_layout
    • BlockSizes.v_layout
    • BlockSizes.has_backward_blocks
    • BlockSizes.get_default()
  • flash_attention_kernel()
  • SplashAttentionKernel
    • SplashAttentionKernel.manual_fwd()
    • SplashAttentionKernel.manual_bwd()
    • SplashAttentionKernel.manual_sharding_spec()
    • SplashAttentionKernel.tree_flatten()
    • SplashAttentionKernel.tree_unflatten()
  • make_splash_mha()
  • make_splash_mqa()
  • make_splash_mha_single_device()
  • make_splash_mqa_single_device()