ejkernel.kernels._triton.native_sparse_attention package

• ejkernel.kernels._triton.native_sparse_attention._compression
  • Compressed Attention Process:
  • nsa_compression()
  • nsa_compression_bwd()
  • nsa_compression_fwd()
• ejkernel.kernels._triton.native_sparse_attention._interface
  • NSA Architecture:
  • Key Benefits:
  • Implementation Details:
  • apply_native_sparse_attention()
  • native_sparse_attention()
• ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd
  • bwd_triton_impl()
• ejkernel.kernels._triton.native_sparse_attention._triton_impl_fwd
  • fwd_triton_impl()
  • nsa_topk()
• ejkernel.kernels._triton.native_sparse_attention._utilities
  • nsa_block_mask()