Module psa 

PSA attention helper kernels and RuntimeOp wrappers.

Implements the data-rearrangement passes surrounding Flash Attention 2 for the PSABlock used in YOLO26 C2PSA layers.

Assumptions (derived from the ultralytics PSABlock / Attention module):

• head_dim = c / num_heads = 2 * key_dim
• QKV conv output has qkv_h = num_heads * 4 * KEY_DIM channels stored as NCHW [B, qkv_h, H, W].
• Per head h, channels [h*4*KEY_DIM : +KEY_DIM] are Q, [+KEY_DIM : +2*KEY_DIM] are K, [+2*KEY_DIM : +3*KEY_DIM] are V_lo, [+3*KEY_DIM : +4*KEY_DIM] are V_hi.

The V split trick lets us run FA2 with HEAD_DIM = key_dim twice (once for V_lo, once for V_hi) instead of needing HEAD_DIM = head_dim = 2*key_dim.

Structs

FlashAttn2PsaOp

CustomOp for FlashAttn2PsaRuntimeOp (V_lo or V_hi section).

FlashAttn2PsaRuntimeOp

RuntimeOp wrapping flash_attention2_forward for PSA attention.

PsaExtractVBackward

Backward of psa_extract_v_nchw: scatters d_v back to V_lo / V_hi channels of d_qkv.

PsaExtractVNchw

Extracts V channels from QKV NCHW into V NCHW.

PsaExtractVOp

CustomOp for PsaExtractVRuntimeOp.

PsaExtractVRuntimeOp

Runtime dispatch for the PSA V-extraction kernel (forward + backward).

PsaFa2Backward

Combined PSA Flash Attention 2 backward pass.

PsaMergeAttnBackward

Backward of psa_merge_attn_nchw: scatters d_merged back to d_lo and d_hi flat buffers.

PsaMergeAttnNchw

Merges two FA2 outputs (V_lo and V_hi attention results) into NCHW.

PsaMergeAttnOp

CustomOp for PsaMergeAttnRuntimeOp.

PsaMergeAttnRuntimeOp

Runtime dispatch for the PSA attention-merge kernel (forward + backward).

PsaPackQkv

Rearranges the NCHW QKV tensor into packed FA2 format.

PsaPackQkvBackward

Backward of psa_pack_qkv: scatters d_packed back to d_qkv.

PsaPackQkvOp

CustomOp for PsaPackQkvRuntimeOp.

PsaPackQkvRuntimeOp

Runtime dispatch for the PSA QKV-packing kernel (forward + backward).

Functions

psa_extract_v_backward

Backward of psa_extract_v_nchw: scatters d_v back to V_lo / V_hi channels of d_qkv.

psa_extract_v_nchw

Extracts V channels from QKV NCHW into V NCHW.

psa_fa2_backward

Combined PSA Flash Attention 2 backward pass.

psa_merge_attn_backward

Backward of psa_merge_attn_nchw: scatters d_merged back to d_lo and d_hi flat buffers.

psa_merge_attn_nchw

Merges two FA2 outputs (V_lo and V_hi attention results) into NCHW.

psa_pack_qkv

Rearranges the NCHW QKV tensor into packed FA2 format.

psa_pack_qkv_backward

Backward of psa_pack_qkv: scatters d_packed back to d_qkv.