ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd
ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd
-
ejkernel.kernels._triton.native_sparse_attention._triton_impl_bwd.bwd_triton_impl(q: Array, k: Array, v: Array, o: Array, lse: Array, do: Array, block_indices: Array, block_counts: jax.jaxlib._jax.Array | int, block_size: int = 64, softmax_scale: float | None = None, cu_seqlens: jax.jaxlib._jax.Array | None = None, token_indices: jax.jaxlib._jax.Array | None = None)[source]