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<()>

2026-09-25: Set (m, l, o) of the first num_seqs sequences to (-inf, 0, 0) on ctx.stream, before the step’s first step_tile.

Errors

num_seqs > max_seqs, or a cuMemsetD32Async fails.

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

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

2026-09-25: (blk, tok, kvh) strides, in BF16 elements, of the paged K/V layout [num_blocks, block_size, num_kv_heads, head_dim].

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

2026-09-25: (blk, tok, kvh) strides, in BF16 elements, of the scratch pool slot layout [slot, K|V, kv_head, block_size, head_dim].

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<()>

2026-09-25: One tile of KV blocks for num_seqs sequences, folded into the running state. The strides select 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<()>

2026-09-25: Write o_state / l_state for num_seqs sequences to output as BF16, 0 where l is 0.

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

pub fn dims(&self) -> TiledAttentionDims