pub struct AttnV41 {
pub cfg: AttnV41Cfg,
/* private fields */
}Expand description
The attention runtime: kernels, RoPE tables, and workspaces for up to
max_tokens positions per call.
Fields§
§cfg: AttnV41CfgImplementations§
Source§impl AttnV41
impl AttnV41
Sourcepub fn decode_capturable(&self, w: &AttnV41LayerWeights) -> bool
pub fn decode_capturable(&self, w: &AttnV41LayerWeights) -> bool
Can this layer’s single-token step be captured into a CUDA graph? The kv and index sources compute the position’s compressor group and the index top-k on the host in the middle of the forward (position parity, a score download); every other layer is kernels and copies from end to end once its per-token inputs are read from the device.
Sourcepub fn decode_fixed_topk(&self, w: &AttnV41LayerWeights) -> usize
pub fn decode_fixed_topk(&self, w: &AttnV41LayerWeights) -> usize
The topk a captured step launches sparse_attn with: the widest
selection the layer can ever see, so the launch argument is constant
and the selection is padded with -1 (which the kernel skips: a padded
entry adds a zero to the denominator sum and nothing to the output,
so the result is the eager step’s bit for bit).
Sourcepub fn invalidate_decode_uploads(&mut self)
pub fn invalidate_decode_uploads(&mut self)
Forget what the captured decode step last uploaded; the next
Self::decode_prep re-uploads. Call after any eager forward (it
writes pos / head_pos / idx_dev for itself) and at a new step.
Sourcepub fn decode_prep(
&mut self,
w: &AttnV41LayerWeights,
shared: &SharedV41,
gpu: &dyn GpuBackend,
start_pos: usize,
stream: u64,
) -> Result<Option<DevicePtr>>
pub fn decode_prep( &mut self, w: &AttnV41LayerWeights, shared: &SharedV41, gpu: &dyn GpuBackend, start_pos: usize, stream: u64, ) -> Result<Option<DevicePtr>>
The HOST half of a capturable layer’s single-token step at start_pos:
the position into pos / head_pos and the padded selection (window
slots, then the shared index selection, then -1) into idx_dev, each
only when it differs from what is already there. Returns the compressed
rows the captured sparse_attn reads (None on ratio-0 layers) — a
pointer the graph bakes, for the caller to verify on every replay.
Sourcepub fn decode_prep_role(
&mut self,
ratio_positive: bool,
shared: &SharedV41,
gpu: &dyn GpuBackend,
start_pos: usize,
stream: u64,
) -> Result<Option<DevicePtr>>
pub fn decode_prep_role( &mut self, ratio_positive: bool, shared: &SharedV41, gpu: &dyn GpuBackend, start_pos: usize, stream: u64, ) -> Result<Option<DevicePtr>>
Self::decode_prep by selection class: ratio_positive layers read
the window plus the shared index selection from idx_dev, the others
the window alone from idx_dev_win; both classes can be prepared for
one multi-layer capture.
Sourcepub fn decode_body(
&self,
gpu: &dyn GpuBackend,
w: &AttnV41LayerWeights,
st: &AttnV41LayerState,
x: DevicePtr,
rows_b: Option<DevicePtr>,
stream: u64,
) -> Result<DevicePtr>
pub fn decode_body( &self, gpu: &dyn GpuBackend, w: &AttnV41LayerWeights, st: &AttnV41LayerState, x: DevicePtr, rows_b: Option<DevicePtr>, stream: u64, ) -> Result<DevicePtr>
The DEVICE half of a capturable layer’s single-token step: the
projections, the window ring write at pos % window (slot read on
the device), sparse attention over the padded selection and the output
projection into the returned buffer. No host work, no per-token
scalar arguments: capturable into a CUDA graph and replayed for any
position once decode_prep has updated pos / head_pos / idx_dev.
Source§impl AttnV41
impl AttnV41
Sourcepub fn forward(
&self,
gpu: &dyn GpuBackend,
w: &AttnV41LayerWeights,
st: &mut AttnV41LayerState,
shared: &mut SharedV41,
x: DevicePtr,
m: usize,
start_pos: usize,
stream: u64,
) -> Result<AttnV41Run>
pub fn forward( &self, gpu: &dyn GpuBackend, w: &AttnV41LayerWeights, st: &mut AttnV41LayerState, shared: &mut SharedV41, x: DevicePtr, m: usize, start_pos: usize, stream: u64, ) -> Result<AttnV41Run>
One layer’s attention for m tokens at start_pos: x is the normed
input [m, dim] bf16; the output is [m, dim] bf16 in run.out.