tvm_tirx.backend.cuda.lang.FlashAttentionLinearScheduler

class tvm_tirx.backend.cuda.lang.FlashAttentionLinearScheduler(*args, **kwargs)

Bases: BaseTileScheduler

Linear 3D scheduler for flash attention (batch, head, m_block).

Used for non-causal attention with simple linear decomposition. Maps linear_idx -> (batch_idx, head_idx, m_block_idx) using:

batch = linear_idx // (num_heads * num_m_blocks)
head = (linear_idx % (num_heads * num_m_blocks)) // num_m_blocks
m_block = linear_idx % num_m_blocks

Parameters:

• prefix (str) – Prefix for TIR variable names

• num_batches (int) – Number of batches

• num_heads (int) – Number of KV heads

• num_m_blocks (int) – Number of Q blocks (M dimension tiles)

• num_ctas (int) – Number of CTAs for persistent kernel stride

update_current_m_n_idx(linear_idx)

Convert linear index to (batch, head, m_block) coordinates.

init(cta_id)

Initialize scheduler with CTA ID.

next_tile()

Advance to next tile by striding by num_ctas.

valid()

Check if there are more tiles to process.