opentau.policies.flash_attn_cuda

Custom CUDA FlashAttention with in-kernel block-causal masking.

This package provides a dependency-free (no flash-attn, no flex_attention) fused attention forward + backward whose block-causal mask is reconstructed inside the CUDA kernel from a compact per-token block-id representation. The dense (B, S, S) attention mask is therefore never materialized, which is both the memory win over the eager backend and a direct way to “handle the block causal attention masks in the code”.

Public API:

  • make_att_block_ids() – build the compact (q_blk, k_blk, q_valid, k_valid) representation from the same pad_masks / att_masks that make_att_2d_masks() consumes. Provably reproduces the dense mask: attend(i, j) == q_valid[i] & k_valid[j] & (k_blk[j] <= q_blk[i]).

  • flash_attn_blockmask() – autograd-aware attention given the compact representation.

  • is_available() – whether the CUDA kernel compiled (else fall back).

Masking convention matches make_att_2d_masks: True = attend. Cross- attention key columns get k_blk = INT32_MIN so they are always attended (subject to padding).

Functions

flash_attn_blockmask(q, k, v, q_blk, k_blk, ...)

Block-causal flash attention.

make_att_block_ids(pad_masks, att_masks[, ...])

Compact block-id representation equivalent to make_att_2d_masks.
opentau.policies.flash_attn_cuda.flash_attn_blockmask(q: Tensor, k: Tensor, v: Tensor, q_blk: Tensor, k_blk: Tensor, q_valid: Tensor, k_valid: Tensor, scale: float) → Tensor[source]

Block-causal flash attention.

Parameters:

  • q – (B, Sq, H, D) queries.

  • k – (B, Sk, Hkv, D) keys (GQA/MQA: H % Hkv == 0).

  • v – (B, Sk, Hkv, D) values.

  • q_blk – int32 (B, Sq) query block-ids.

  • k_blk – int32 (B, Sk) key block-ids (cross columns use INT_BLK_MIN).

  • q_valid – bool (B, Sq) query padding mask (True = real token).

  • k_valid – bool (B, Sk) key padding mask.

  • scale – softmax scale (typically head_dim ** -0.5).

Returns:

(B, Sq, H, D) attention output, same dtype as q.

opentau.policies.flash_attn_cuda.is_available() → bool[source]

True if the custom flash-attention CUDA kernel compiled and is usable.

opentau.policies.flash_attn_cuda.load_error() → str | None[source]

The last compilation error string, if any (for diagnostics/tests).

opentau.policies.flash_attn_cuda.make_att_block_ids(pad_masks: Tensor, att_masks: Tensor, n_cross_att_tokens: int | None = None, cross_att_pad_masks: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor][source]

Compact block-id representation equivalent to make_att_2d_masks.

Mirrors make_att_2d_masks exactly but returns four 1-D-per-token tensors instead of a dense 2-D mask. The kernel reconstructs the mask via attend(i, j) == q_valid[i] & k_valid[j] & (k_blk[j] <= q_blk[i]).

Parameters:

  • pad_masks – bool (B, N), True where the token is real (not padding).

  • att_masks – int/bool (B, N), 1 opens a new attention block, 0 continues the current block (the same convention as make_att_2d_masks).

  • n_cross_att_tokens – if set, prepend that many cross-attention key columns (always attended), matching the cross-attention branch of make_att_2d_masks.

  • cross_att_pad_masks – bool (B, n_cross_att_tokens) padding for the cross columns. Required iff n_cross_att_tokens is set.

Returns:

(q_blk, k_blk, q_valid, k_valid) where q_blk is int32 (B, N), k_blk is int32 (B, N) or (B, n_cross_att_tokens + N), q_valid is bool (B, N), k_valid matches k_blk’s length.