tvm_tirx.backend.cuda.lang.FlashAttentionLPTScheduler

tvm_tirx.backend.cuda.lang.FlashAttentionLPTScheduler#

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

Bases: BaseTileScheduler

LPT scheduler with L2 swizzle for causal flash attention.

Processes high-work Q blocks (with more KV blocks to attend to) first using Longest Processing Time (LPT) scheduling. Also applies L2 cache swizzle for better cache locality across batch*head dimensions.

The LPT aspect comes from reversing m_block order: lower Q blocks have more KV blocks to process due to causal masking, so processing them first balances load.

The scheduler is only applied to non-persistent kernels.

L2 Swizzle: Groups consecutive batch*head indices together for L2 locality.

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 (should equal total_tasks for causal)

  • l2_swizzle (int) – L2 swizzle factor for cache locality

update_current_m_n_idx(linear_idx)#

Convert linear index to (batch, head, m_block) with LPT + L2 swizzle.

init(cta_id)#

Initialize scheduler with CTA ID.

next_tile()#

Advance to the next tile.

Single-tile mode (num_ctas=None, the default): each CTA owns one task; terminate. Persistent mode (num_ctas=N): stride by N, like FlashAttentionLinearScheduler, while keeping the LPT + L2 swizzle index mapping.

valid()#

Check if there are more tiles to process.