Fix _pa_swa_prefill_lens off-by-one in FlashAttentionBackend (#33379)

2026-08-07 14:06:48 -07:00
parent 62a28197c0
commit 699fcdc936
2 changed files with 115 additions and 2 deletions
@@ -242,9 +242,12 @@ class FlashAttentionBackend(AttentionBackend):
                "Prefill-aware SWA requires page_size=1, "
                f"got page_size={self.page_size}"
            )
-            max_bs = model_runner.req_to_token_pool.size
+            # Indexed by raw req_pool_idx values (see the write below and
+            # _build_pa_page_table), which range over [0, size] (row 0 is the
+            # reserved padding slot) -- so this needs size+1, not size.
+            max_req_pool_idx = model_runner.req_to_token_pool.size
            self._pa_swa_prefill_lens = torch.zeros(
-                max_bs, dtype=torch.int32, device=model_runner.device
+                max_req_pool_idx + 1, dtype=torch.int32, device=model_runner.device
            )
            self._pa_swa_max_prefill_len = 0