Struct TiledAttention 

pub struct TiledAttention {
    pub m_state: DeviceBuffer,
    pub l_state: DeviceBuffer,
    pub o_state: DeviceBuffer,
    /* private fields */
}

Fields

m_state: DeviceBuffer

l_state: DeviceBuffer

o_state: DeviceBuffer

Implementations

impl TiledAttention

pub fn new(dims: TiledAttentionDims) -> Result<Self>

pub fn begin_step(&self, ctx: &CudaCtx, num_seqs: usize) -> Result<()>

Reset (m, l, o) for num_seqs sequences. Call once at the start of each decode step before the first step_tile.

pub fn begin_step_on_stream(&self, stream: u64, num_seqs: usize) -> Result<()>

pub fn paged_strides(&self) -> (i64, i64, i64)

Stride triple for the kernel-native paged K/V layout: [num_blocks, block_size, num_kv_heads, head_dim]. All in BF16 elements.

pub fn scratch_pool_strides(&self) -> (i64, i64, i64)

Stride triple for the scratch-pool slot layout [slot, K|V, kv_head, block_size, head_dim]. Reads of contiguous- per-(kv_head) groups stay efficient on disk; the kernel pays a single extra multiply per address calculation.

pub fn step_tile( &self, ctx: &CudaCtx, q: u64, k_pool: u64, v_pool: u64, tile_blocks: u64, tile_block_counts: u64, num_seqs: usize, blk_stride: i64, tok_stride: i64, kvh_stride: i64, last_block_valid_slots: i32, ) -> Result<()>

One tile of blocks across num_seqs sequences. Stride triple (blk_stride, tok_stride, kvh_stride) selects the K/V layout — see paged_strides and scratch_pool_strides.

pub fn step_tile_on_stream( &self, stream: u64, q: u64, k_pool: u64, v_pool: u64, tile_blocks: u64, tile_block_counts: u64, num_seqs: usize, blk_stride: i64, tok_stride: i64, kvh_stride: i64, last_block_valid_slots: i32, ) -> Result<()>

pub fn finalize( &self, ctx: &CudaCtx, output: u64, num_seqs: usize, ) -> Result<()>

Divide o_state by l_state and store as BF16 in output.

pub fn finalize_on_stream( &self, stream: u64, output: u64, num_seqs: usize, ) -> Result<()>

pub fn dims(&self) -> TiledAttentionDims