pub struct MuonNsXtx {
pub name: &'static str,
pub id: String,
pub transpose: bool,
pub block_r: i32,
pub block_k: i32,
pub group_r: i32,
pub kernel_source: String,
pub entry_point_source: String,
pub source: String,
}Expand description
Gram matrix: T = X·Xᵀ (!TRANSPOSE) or T = Xᵀ·X (TRANSPOSE).
Both cases are expressed as A @ Aᵀ with different views of X:
!TRANSPOSE→ A = X[M, N], output T[M, M], contraction K = N.TRANSPOSE→ A = Xᵀ[N, M], output T[N, N], contraction K = M.
stride_xm is the row stride of X (typically = N for row-major [M, N]).
Grid: [ceil(R/BLOCK_R)² grouped by GROUP_R, 1, 1] where R = M or N.
Fields§
§name: &'static strThe kernel function’s name (e.g. "flash_attention2_forward").
id: StringUnique kernel identifier: fn_name + dtype(s) + const values joined by “__”.
transpose: boolCompile-time kernel constant, from the annotated fn’s const generics.
block_r: i32Compile-time kernel constant, from the annotated fn’s const generics.
block_k: i32Compile-time kernel constant, from the annotated fn’s const generics.
group_r: i32Compile-time kernel constant, from the annotated fn’s const generics.
kernel_source: StringThe original kernel function source.
entry_point_source: StringThe Rust source of the generated C-ABI entry-point wrapper function.
source: StringCombined source (kernel_source + "\n\n" + entry_point_source); used by the Kernel trait.
Implementations§
Trait Implementations§
Source§impl Kernel for MuonNsXtx
impl Kernel for MuonNsXtx
Source§type Args<'__a> = (*mut f32, *mut f32, i32, i32, i32)
type Args<'__a> = (*mut f32, *mut f32, i32, i32, i32)
Source§fn kernel_source(&self) -> &str
fn kernel_source(&self) -> &str
Source§fn entry_point_source(&self) -> &str
fn entry_point_source(&self) -> &str
Source§fn entry_point_name(&self) -> String
fn entry_point_name(&self) -> String
"{name}_entry_point".