flashinfer.testing.attention_tb_per_sec_with_actual_seq_lens

flashinfer.testing.attention_tb_per_sec_with_actual_seq_lens(actual_seq_lens_q, actual_seq_lens_kv, head_dim_qk, head_dim_vo, num_qo_heads, num_kv_heads, time, q_dtype=torch.bfloat16, kv_dtype=torch.bfloat16, o_dtype=torch.bfloat16)

计算给定注意力层在实际序列长度下的每秒 TB 性能。不假设批处理中所有序列长度都相同。

参数:

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

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

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

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

• num_qo_heads (int) – 查询头数。

• num_kv_heads (int) – 键和值的头数。

• time (float) – 毫秒级的执行时间。

• q_dtype (torch.dtype) – 查询的数据类型。

• kv_dtype (torch.dtype) – 键和值的数据类型。

• o_dtype (torch.dtype) – 输出的数据类型。

返回值:

该层的每秒 TB 数。

返回值类型:

tb_per_sec (float)