[ROCm][Bugfix] Fix empty-tensor .max() crash in AITER FA (#47894)

Signed-off-by: Djordje Ramic <djoramic@amd.com>
2026-07-08 16:58:36 -05:00
parent 089e412878
commit 49abadaedb
1 changed files with 5 additions and 1 deletions
@@ -602,7 +602,11 @@ class AiterFlashAttentionMetadataBuilder(
            torch.cumsum(
                chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32
            )
-            max_cum_tokens = cu_seq_lens_cpu[:, -1].max().item()
+            # Avoid .max() on an empty tensor when there is no context
+            # (num_chunks == 0, e.g. Whisper encoder's first pass).
+            max_cum_tokens = (
+                cu_seq_lens_cpu[:, -1].max().item() if num_chunks > 0 else 0
+            )

            range_idx = torch.arange(max_cum_tokens, dtype=torch.int32)[None, None, :]
            idx_to_batch_tensor = range_idx == cu_seq_lens_cpu[:, 1:][:, :, None]