ejkernel.kernels._pallas.tpu.blocksparse_attention package

• ejkernel.kernels._pallas.tpu.blocksparse_attention._info
  • MaskInfo
  • process_dynamic_mask()
  • process_dynamic_mask_dkv()
  • process_mask()
  • process_mask_dkv()
• ejkernel.kernels._pallas.tpu.blocksparse_attention._kernel
  • BlockSizes
  • QKVLayout
  • SegmentIds
  • SplashAttentionKernel
  • attention_reference()
  • attention_reference_custom()
  • blocksparse_attention()
  • flash_attention_kernel()
  • from_head_minor()
  • get_kernel_name()
  • make_attention_reference()
  • make_masked_mha_reference()
  • make_masked_mqa_reference()
  • make_splash_mha()
  • make_splash_mha_single_device()
  • make_splash_mqa()
  • make_splash_mqa_single_device()
• ejkernel.kernels._pallas.tpu.blocksparse_attention._masks
  • CausalMask
  • ChunkedCausalMask
  • FullMask
  • LocalMask
  • LogicalAnd
  • LogicalOr
  • Mask
  • MultiHeadMask
  • NumpyMask
  • make_causal_mask()
  • make_chunk_attention_mask()
  • make_local_attention_mask()
  • make_random_mask()