[ROCm][Bugfix] Fix empty-tensor .max() crash in AITER FA (#47894)
Signed-off-by: Djordje Ramic <djoramic@amd.com>
djramic committed
Jul 8, 2026 at 23:58 UTC
49abadaedba2d88c6f35216ebef235edac8c1e82
1 file changed
+5
-1
vllm/v1/attention/backends/rocm_aiter_fa.py
+5
-1
@@ -602,7 +602,11 @@ class AiterFlashAttentionMetadataBuilder(
602
torch.cumsum(
603
chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32
604
)
605
- max_cum_tokens = cu_seq_lens_cpu[:, -1].max().item()
605
+ # Avoid .max() on an empty tensor when there is no context
606
+ # (num_chunks == 0, e.g. Whisper encoder's first pass).
607
+ max_cum_tokens = (
608
+ cu_seq_lens_cpu[:, -1].max().item() if num_chunks > 0 else 0
609
+ )
610
611
range_idx = torch.arange(max_cum_tokens, dtype=torch.int32)[None, None, :]
612
idx_to_batch_tensor = range_idx == cu_seq_lens_cpu[:, 1:][:, :, None]