已合并
[sync] PR-40262: bugfix for profiler FLOPs when actual seq len can be zero #40379
ascend-robot创建于 7月7日
[sync] PR-40262: bugfix for profiler FLOPs when actual seq len can be zero #40379
已合并
共 1 个文件变更+1-1
| @@ -419,7 +419,7 @@ def _calculate_tnd_layout_flops( | |||
| 419 | _, _, v_d = v_shape | 419 | _, _, v_d = v_shape |
| 420 | q_lens = _parse_seq_len(actual_seq_qlen) | 420 | q_lens = _parse_seq_len(actual_seq_qlen) |
| 421 | kv_lens = _parse_seq_len(actual_seq_kvlen) | 421 | kv_lens = _parse_seq_len(actual_seq_kvlen) |
| 422 | - if len(q_lens) != len(kv_lens) or any(length <= 0 for length in q_lens + kv_lens): | 422 | + if len(q_lens) != len(kv_lens) or any(length < 0 for length in q_lens + kv_lens): |
| 423 | raise ValueError("actual_seq_qlen and actual_seq_kvlen must contain valid cumulative lengths") | 423 | raise ValueError("actual_seq_qlen and actual_seq_kvlen must contain valid cumulative lengths") |
| 424 | attention_scores = sum(q_len * kv_len for q_len, kv_len in zip(q_lens, kv_lens)) | 424 | attention_scores = sum(q_len * kv_len for q_len, kv_len in zip(q_lens, kv_lens)) |
| 425 | return int(2 * (q_heads or shape_q_heads) * (q_d + v_d) * attention_scores) | 425 | return int(2 * (q_heads or shape_q_heads) * (q_d + v_d) * attention_scores) |