burn_dragon_kernel 0.5.0

Fused GPU kernel crate for burn_dragon execution paths
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

pub mod bench;
pub mod conv;
pub mod conv_runtime;
pub mod selective_scan_backward;
pub mod selective_scan_forward;

/// Mamba selective state-space kernels are planned to match the pinned upstream Mamba-1 reference
/// and its selective-scan implementation. The current forward path is tensorized and exact on
/// tested fixtures, but it is not yet a true fused selective-scan kernel.
pub const STATUS: &str = "default_tensorized_direct_path";
pub const FORWARD_ACCELERATION_AVAILABLE: bool =
    conv::AVAILABLE && selective_scan_forward::AVAILABLE;
pub const BACKWARD_ACCELERATION_AVAILABLE: bool = selective_scan_backward::AVAILABLE;
pub const UPSTREAM_REPO: &str = "https://github.com/state-spaces/mamba";
pub const UPSTREAM_COMMIT: &str = "c5afbdf";
pub const UPSTREAM_TARGET_KIND: &str = "mamba1_selective_scan";