pub struct PsaExtractVNchw<D: Float> {
pub name: &'static str,
pub id: String,
pub key_dim: i32,
pub kernel_source: String,
pub entry_point_source: String,
pub source: String,
/* private fields */
}Expand description
Extracts V channels from QKV NCHW into V NCHW.
Input: qkv_ptr — [B, qkv_h, H, W] NCHW
Output: v_ptr — [B, c, H, W] NCHW (c = num_heads * 2 * KEY_DIM)
Grid: [BH * N, 1, 1]. Block: [KEY_DIM, 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 “__”.
key_dim: 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<D: Float> Kernel for PsaExtractVNchw<D>
impl<D: Float> Kernel for PsaExtractVNchw<D>
Source§type Args<'__a> = (*mut D, *mut D, i32, i32, i32, i32, i32)
type Args<'__a> = (*mut D, *mut D, i32, i32, i32, i32, i32)
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.
§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<D> Freeze for PsaExtractVNchw<D>
impl<D> RefUnwindSafe for PsaExtractVNchw<D>where
D: RefUnwindSafe,
impl<D> Send for PsaExtractVNchw<D>where
D: Send,
impl<D> Sync for PsaExtractVNchw<D>where
D: Sync,
impl<D> Unpin for PsaExtractVNchw<D>where
D: Unpin,
impl<D> UnsafeUnpin for PsaExtractVNchw<D>
impl<D> UnwindSafe for PsaExtractVNchw<D>where
D: UnwindSafe,
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