torch.utils.flop_counter.sdpa_flop_count

torch.utils.flop_counter.sdpa_flop_count(query_shape, key_shape, value_shape) [исходный код]

Подсчитывает количество операций с плавающей запятой для механизма внимания к себе.

Поддерживает GQA (внимание с группированными запросами), при котором у ключей и значений меньше голов, чем у запроса. Ядро распространяет головы KV, чтобы их количество соответствовало количеству голов запроса.