flashinfer.testing.attention_flops_with_actual_seq_lens

flashinfer.testing.attention_flops_with_actual_seq_lens(actual_seq_lens_q, actual_seq_lens_kv, head_dim_qk, head_dim_vo, num_qo_heads, causal)

计算给定注意力层级的 FLOPs,其中实际序列长度作为 1D 张量提供。

参数:
  • actual_seq_lens_q (torch.Tensor) – 查询的实际序列长度数组。

  • actual_seq_lens_kv (torch.Tensor) – 键和值的实际序列长度数组。

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

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

  • num_qo_heads (int) – 查询头数。

  • causal (bool) – 是否使用因果掩码。

  • 注意 – 对于解码,Causal 必须为 False,因为此函数假定 qo_seqlen == kv_seqlen。

返回值:

该层级的总 FLOPs。

返回值类型:

total_flops (int)