Struct SparseKernel 

pub struct SparseKernel { /* private fields */ }

Manager for sparse attention kernel execution.

Handles the execution of attention computations with dynamic head sparsity, including index mapping caching for efficient repeated execution with the same sparsity pattern.

Performance Optimizations

• Pre-computes index mappings for gather/scatter operations
• Caches mappings by pattern hash for repeated use
• Supports flash attention for memory efficiency
• Handles different memory layouts (contiguous, channels-last, grouped)

Example

let mut kernel = SparseKernel::new(KernelConfig::default());
kernel.prepare(&mask, layer_idx)?;
let output = kernel.execute(&query, &key, &value, &mask, layer_idx)?;

Implementations

impl SparseKernel

pub fn new(config: KernelConfig) -> Self

Create a new kernel with config

pub fn prepare(&mut self, mask: &AttentionMask, layer: usize) -> Result<()>

Prepare kernel for a specific mask

pub fn execute( &self, mask: &AttentionMask, layer: usize, _q: &[f32], _k: &[f32], _v: &[f32], ) -> Result<Vec<f32>>

Execute sparse attention for a layer

This is a simulated implementation. In practice, this would:

1. Gather only active Q, K, V heads
2. Compute attention only for active heads
3. Scatter results back to full head positions

pub fn estimate_savings(&self, mask: &AttentionMask) -> ComputeEstimate

Estimate compute savings for a mask

pub fn config(&self) -> &KernelConfig

Get current configuration

pub fn clear_cache(&mut self)

Clear cached index mappings

Trait Implementations

impl Debug for SparseKernel

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.