torch.utils.flop_counter.sdpa_flop_count#

torch.utils.flop_counter.sdpa_flop_count(query_shape, key_shape, value_shape)[source]#

Count flops for self-attention.

Supports GQA (grouped-query attention) where key/value have fewer heads than the query. The kernel broadcasts KV heads to match query heads.