Struct FlashDecodingPlan 

pub struct FlashDecodingPlan<T: Element> { /* private fields */ }

FlashDecoding forward plan (Dao 2023).

Split-K parallel attention decode for seq_q = 1. Replaces both FlashSdpaPlan and FA2 at the decode regime — both of those tile the Q dimension and waste work when seq_q < 64.

When to use: autoregressive decoder inference token loop. After the prefill step (which uses [FlashSdpaPlan] with fa2 for the long initial context), each generated token calls this plan with seq_q = 1 and the full grown KV cache.

Dtypes: f16, bf16 (the only dtypes inference uses).

Shape limits: head_dim ≤ 128. Arbitrary B, H, K_len.

Workspace: non-zero. See Self::workspace_size.

Precision guarantee: f32 accumulators throughout the split AND combine kernels. Deterministic — each output cell is written by exactly one block; no atomicAdd.

Implementations

impl<T: Element> FlashDecodingPlan<T>

pub fn select( _stream: &Stream, desc: &FlashDecodingDescriptor, _pref: PlanPreference, ) -> Result<Self>

Pick a kernel for the supplied descriptor.

pub fn can_implement(&self, args: &FlashDecodingArgs<'_, T>) -> Result<()>

Validate args against the descriptor.

pub fn backend(&self) -> BackendKind

Backend selected by select.

pub fn sku(&self) -> &KernelSku

Kernel SKU descriptor.

pub fn workspace_size(&self) -> usize

Workspace requirement in bytes for the (split + combine) pipeline.

pub fn run( &self, stream: &Stream, workspace: Workspace<'_>, args: FlashDecodingArgs<'_, T>, ) -> Result<()>

Run the FlashDecoding pipeline.