ejkernel.modules.operations package

• ejkernel.modules.operations.attention
  • Attention
  • attention()
• ejkernel.modules.operations.blocksparse_attention
  • BlockSparseAttention
  • blocksparse_attention()
• ejkernel.modules.operations.configs
  • AttentionConfig
  • BaseOperationConfig
  • BlockSparseAttentionConfig
  • FlashAttentionConfig
  • FlashMLAConfig
  • GLAttentionConfig
  • GroupedMatmulConfig
  • KernelDeltaAttentionConfig
  • LightningAttentionConfig
  • MeanPoolingConfig
  • NativeSparseAttentionConfig
  • PageAttentionConfig
  • PrefillPageAttentionConfig
  • RaggedDecodeAttentionConfig
  • RaggedPageAttentionv2Config
  • RaggedPageAttentionv3Config
  • RecurrentAttentionConfig
  • RingAttentionConfig
  • ScaledDotProductAttentionConfig
  • StateSpaceV1Config
  • StateSpaceV2Config
  • UnifiedAttentionConfig
  • get_safe_hash_int()
  • hash_fn()
• ejkernel.modules.operations.flash_attention
  • FlashAttention
  • flash_attention()
• ejkernel.modules.operations.gated_linear_attention
  • GLAttention
  • gla_attention()
• ejkernel.modules.operations.grouped_matmul
  • GroupedMatmul
  • grouped_matmul()
• ejkernel.modules.operations.kernel_delta_attention
  • KernelDeltaAttention
  • kda_attention()
  • kernel_delta_attention()
• ejkernel.modules.operations.lightning_attention
  • LightningAttention
  • lightning_attention()
• ejkernel.modules.operations.multi_head_latent_attention
  • FlashMLA
  • mla_attention()
• ejkernel.modules.operations.native_sparse_attention
  • NativeSparseAttention
  • native_sparse_attention()
• ejkernel.modules.operations.page_attention
  • PageAttention
  • page_attention()
• ejkernel.modules.operations.pooling
  • MeanPooling
  • mean_pooling()
• ejkernel.modules.operations.prefill_page_attention
  • PrefillPageAttention
  • prefill_page_attention()
• ejkernel.modules.operations.ragged_decode_attention
  • RaggedDecodeAttention
  • ragged_decode_attention()
• ejkernel.modules.operations.ragged_page_attention_v2
  • RaggedPageAttentionv2
  • ragged_page_attention_v2()
• ejkernel.modules.operations.ragged_page_attention_v3
  • RaggedPageAttentionv3
  • ragged_page_attention_v3()
• ejkernel.modules.operations.recurrent
  • RecurrentAttention
  • recurrent_attention()
• ejkernel.modules.operations.ring_attention
  • RingAttention
  • ring_attention()
• ejkernel.modules.operations.scaled_dot_product_attention
  • ScaledDotProductAttention
  • scaled_dot_product_attention()
• ejkernel.modules.operations.state_space_v1
  • StateSpaceV1
  • state_space_v1()
• ejkernel.modules.operations.state_space_v2
  • StateSpaceV2
  • state_space_v2()
• ejkernel.modules.operations.unified_attention
  • UnifiedAttention
  • unified_attention()