maxtext.kernels package

Subpackages

• maxtext.kernels.attention package
  • 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
      • get_kernel_name()
      • attention_reference()
      • attention_reference_custom()
      • make_attention_reference()
      • make_masked_mha_reference()
      • make_masked_mqa_reference()
      • QKVLayout
      • from_head_minor()
      • BlockSizes
      • flash_attention_kernel()
      • SplashAttentionKernel
      • make_splash_mha()
      • make_splash_mqa()
      • make_splash_mha_single_device()
      • make_splash_mqa_single_device()
• maxtext.kernels.megablox package
  • Submodules
    • maxtext.kernels.megablox.backend module
      • make_group_metadata()
      • gmm()
      • tgmm()
    • maxtext.kernels.megablox.common module
      • is_tpu()
      • tpu_kind()
      • tpu_generation()
      • supports_bfloat16_matmul()
      • assert_is_supported_dtype()
      • select_input_dtype()
    • maxtext.kernels.megablox.ops module
      • gmm()

Submodules

• maxtext.kernels.gather_reduce_pallas module
  • sc_gather_reduce()
• maxtext.kernels.gather_reduce_sc module
  • VectorTypeHelper
  • sc_gather_reduce()
• maxtext.kernels.sort_activations module