From 699fcdc936c49c2a24fed3a0702a7c40b71c24d0 Mon Sep 17 00:00:00 2001 From: Dmitrii Sergeev Date: Sat, 8 Aug 2026 00:06:48 +0300 Subject: [PATCH] Fix _pa_swa_prefill_lens off-by-one in FlashAttentionBackend (#33379) --- .../attention/flashattention_backend.py | 7 +- ...flashattention_pa_swa_prefill_lens_size.py | 110 ++++++++++++++++++ 2 files changed, 115 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 41ebbdea8..fddf9f99d 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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 diff --git a/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py new file mode 100644 index 000000000..f320700ec --- /dev/null +++ b/test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py @@ -0,0 +1,110 @@ +"""Regression test for PR #32208 / base-b-test-1-gpu-large CI failure. + +FlashAttentionBackend's prefill-aware-SWA scratch buffer +(``_pa_swa_prefill_lens``) is indexed directly by raw ``req_pool_idx`` values +(``self._pa_swa_prefill_lens[forward_batch.req_pool_indices[:batch_size]] = +...``), not by batch position. ``ReqToTokenPool`` reserves row 0 as a padding +slot and hands out real slots ``1..size``, so the valid index range is +``[0, size]`` -- the buffer must hold ``size + 1`` elements, not ``size``. + +Sizing it to ``size`` (the pre-fix code) is off-by-one: writing at +``req_pool_idx == size`` overflows the buffer, which crashes on CUDA with an +``index_put`` "index out of bounds" device assertion. This was latent under +the old head-first ``ReqToTokenPool.alloc()`` (index ``size`` was only +reachable once the pool was nearly saturated) and became immediately +reachable once ``alloc()`` switched to popping free slots from the tail +(``memory_pool.py``, "O(1) slot allocation in ReqToTokenPool.alloc()"): the +very first request into a fresh pool gets ``req_pool_idx == size``. +""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend +from sglang.srt.runtime_context import get_context +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small") + + +def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64): + """A minimal fake ModelRunner that reaches FlashAttentionBackend.__init__'s + is_prefill_aware_swa branch (mirrors how models like + python/sglang/srt/models/unlimited_ocr.py opt in).""" + device = "cuda" + req_to_token_pool = SimpleNamespace( + size=pool_size, + req_to_token=torch.zeros( + pool_size + 1, max_context_len, dtype=torch.int32, device=device + ), + ) + model_config = SimpleNamespace( + is_encoder_decoder=False, + context_len=max_context_len, + attention_arch=AttentionArch.MHA, + is_local_attention_model=False, + head_dim=8, + hf_text_config=SimpleNamespace( + num_attention_heads=2, attn_logit_softcapping=None + ), + get_num_kv_heads=lambda tp_size: 2, + ) + server_args = SimpleNamespace( + speculative_eagle_topk=None, + enable_deterministic_inference=False, + is_embedding=False, + chunked_prefill_size=8192, + disable_radix_cache=False, + enable_prefill_cp=False, + enable_dp_attention=False, + ) + return SimpleNamespace( + sliding_window_size=None, + model_config=model_config, + device=device, + req_to_token_pool=req_to_token_pool, + token_to_kv_pool=object(), # not a SWAKVPool instance -> use_sliding_window_kv_pool=False + # getattr(..., "full_v2p_page_table", None) is None -> unified_mla_hooks + # falls back to the static (disabled) hook set. + token_to_kv_pool_allocator=object(), + kv_cache_dtype=torch.float16, + kv_cache_dtype_str="auto", + page_size=1, + ps=SimpleNamespace(attn_cp_size=1, tp_size=1), + is_draft_worker=False, + server_args=server_args, + attention_chunk_size=None, + prefill_aware_swa=True, + ) + + +@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +class TestPrefillAwareSwaPrefillLensBound(CustomTestCase): + def test_buffer_covers_full_req_pool_idx_range(self): + pool_size = 8 + runner = _make_prefill_aware_swa_runner(pool_size=pool_size) + + # __init__ reads get_spec().speculative_num_draft_tokens, which comes + # from the published runtime-context config bag, not model_runner / + # server_args -- publish a real (dummy) ServerArgs rather than faking + # the accessor (see the sglang-runtime-context skill). + with get_context().override_server_args(): + backend = FlashAttentionBackend(runner) + + # req_pool_idx ranges over [0, pool_size] inclusive (row 0 is the + # reserved CUDA-graph padding slot; real requests use 1..pool_size). + self.assertEqual(backend._pa_swa_prefill_lens.shape[0], pool_size + 1) + + # This mirrors the exact write that crashed in CI: writing at the + # maximum valid req_pool_idx must stay in bounds. + max_req_pool_idx = torch.tensor([pool_size], device=runner.device) + backend._pa_swa_prefill_lens[max_req_pool_idx] = 7 + self.assertEqual(backend._pa_swa_prefill_lens[pool_size].item(), 7) + + +if __name__ == "__main__": + unittest.main()