flashinfer.testing.attention_flops

flashinfer.testing.attention_flops(batch_size, qo_seqlen, kv_seqlen, head_dim_qk, head_dim_vo, num_qo_heads, causal)

计算给定注意力层的 FLOPs。假设批处理中所有序列长度相同

参数:

• batch_size (int) – 批大小。

• qo_seqlen (int) – 查询的序列长度。假设批处理中相同。

• kv_seqlen (int) – 键和值的序列长度。假设批处理中相同。

• head_dim_qk (int) – 查询和键的头部维度。

• head_dim_vo (int) – 值的头部维度。

• num_qo_heads (int) – 查询头数。

• causal (bool) – 是否使用因果掩码。对于因果掩码，FLOPs 会减半。

返回值:

该层的总 FLOPs。

返回值类型:

total_flops (int)