ejkernel.kernels._triton.native_sparse_attention._utilities

ejkernel.kernels._triton.native_sparse_attention._utilities.nsa_block_mask(block_indices: Array, block_counts: jax.jaxlib._jax.Array | int, cu_seqlens: Array, block_size: int)