pub struct NadamStep {
pub name: &'static str,
pub id: String,
pub block_size: i32,
pub kernel_source: String,
pub entry_point_source: String,
pub source: String,
}Expand description
NAdam step (Nesterov-accelerated Adam).
All bias correction terms and Nesterov coefficients are precomputed on the host and passed as scalars:
exp_avg = beta1 * exp_avg + (1 - beta1) * g
exp_avg_sq = beta2 * exp_avg_sq + (1 - beta2) * g²
denom = sqrt(exp_avg_sq) / bias_corr2_sqrt + eps
p -= lr * (coeff_g * g + coeff_m * exp_avg) / denomPrecomputed on host:
bias_corr2_sqrt = sqrt(1 - beta2^t)coeff_g = (1 - mu_t) / (1 - mu_product)(grad contribution)coeff_m = mu_t1 / (1 - mu_product_next)(moment contribution)
Grid: [ceil(n_elements / BLOCK_SIZE), 1, 1].
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 “__”.
block_size: 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 NadamStep
impl Kernel for NadamStep
Source§type Args<'__a> = (*mut f32, *mut f32, *mut f32, *mut f32, i32, f32, f32, f32, f32, f32, f32, f32, f32)
type Args<'__a> = (*mut f32, *mut f32, *mut f32, *mut f32, i32, f32, f32, f32, f32, f32, f32, f32, f32)
This kernel’s launch-argument tuple type.
Source§fn kernel_source(&self) -> &str
fn kernel_source(&self) -> &str
This kernel’s body source, without the entry-point wrapper.
Source§fn entry_point_source(&self) -> &str
fn entry_point_source(&self) -> &str
Returns the Rust source of the generated C-ABI entry-point wrapper function.
Source§fn entry_point_name(&self) -> String
fn entry_point_name(&self) -> String
Returns the PTX symbol name for this kernel:
"{name}_entry_point".Auto Trait Implementations§
impl Freeze for NadamStep
impl RefUnwindSafe for NadamStep
impl Send for NadamStep
impl Sync for NadamStep
impl Unpin for NadamStep
impl UnsafeUnpin for NadamStep
impl UnwindSafe for NadamStep
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more