We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
block_mask
AttentionEngine(block_size=BLOCK_SIZE)( q, # [batch, seqlenq, head, dimqk] k, # [batch, seqlenkv, head, dimqk] v, # [batch, seqlenkv, head, dimv] block_mask, # [batch, seqlenq//BLOCK_SIZE, seqlenkv//BLOCK_SIZE] )
block_indices
AttentionEngine(block_size=BLOCK_SIZE)( q, # [batch, seqlenq, head, dimqk] k, # [batch, seqlenkv, head, dimqk] v, # [batch, seqlenkv, head, dimv] block_indices, # [batch, seqlenq//BLOCK_SIZE, head, MAX_BLOCKS] selected_block_num, # [batch, head] )
block_maskblock_indices