diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 60a46b146..e06bcb449 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -616,11 +616,14 @@ class FlashAttentionBackend(AttentionBackend): # upper bound) so captured graphs keep a valid address; each replay # refills a [:num_tokens] view. if self.use_sliding_window_kv_pool: + assert forward_batch.out_cache_loc is not None m.swa_page_table = torch.zeros( (bs, self.max_num_pages), dtype=torch.int32, device=device ) self.full_cg_prefill_swa_out_cache_loc = torch.zeros( - (self.max_context_len,), dtype=torch.int64, device=device + (forward_batch.out_cache_loc.shape[0],), + dtype=torch.int64, + device=device, ) self.full_cg_prefill_metadata = m m = self.full_cg_prefill_metadata @@ -657,6 +660,11 @@ class FlashAttentionBackend(AttentionBackend): # SWA write targets for the new tokens (KVWriteLoc.swa_loc), refilled # into the pointer-stable buffer and bound as a [:num_tokens] view. num_out = forward_batch.out_cache_loc.shape[0] + assert_buffer_fits( + num_out, + self.full_cg_prefill_swa_out_cache_loc.shape[0], + "full-CG prefill SWA write-location buffer", + ) self.full_cg_prefill_swa_out_cache_loc[:num_out].copy_( self.token_to_kv_pool.translate_loc_from_full_to_swa( forward_batch.out_cache_loc @@ -686,7 +694,7 @@ class FlashAttentionBackend(AttentionBackend): seq_lens_cpu = forward_batch.seq_lens_cpu eager_max_k = ( seq_lens_cpu.max().item() - if seq_lens_cpu is not None + if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0 else self.max_context_len ) diff --git a/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py b/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py new file mode 100644 index 000000000..e44db9903 --- /dev/null +++ b/test/registered/unit/layers/attention/test_flashattention_graph_metadata.py @@ -0,0 +1,72 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=2, stage="base-b", runner_config="1-gpu-small") + + +class TestFlashAttentionGraphMetadata(CustomTestCase): + def test_full_prefill_swa_buffer_uses_capture_token_capacity(self): + backend = FlashAttentionBackend.__new__(FlashAttentionBackend) + backend.full_cg_prefill_metadata = None + backend.use_sliding_window_kv_pool = True + backend.max_num_pages = 4 + backend.max_context_len = 4 + backend.page_size = 1 + backend.req_to_token = torch.zeros((1, 4), dtype=torch.int32) + backend.token_to_kv_pool = SimpleNamespace( + translate_loc_from_full_to_swa=lambda locations: locations + ) + forward_batch = SimpleNamespace( + batch_size=1, + seq_lens=torch.zeros(1, dtype=torch.int64), + extend_seq_lens=torch.zeros(1, dtype=torch.int64), + seq_lens_cpu=torch.zeros(1, dtype=torch.int64), + req_pool_indices=torch.zeros(1, dtype=torch.int64), + out_cache_loc=torch.arange(8, dtype=torch.int64), + positions=torch.arange(8, dtype=torch.int64), + ) + + backend._init_full_cg_prefill_metadata(forward_batch, in_capture=True) + + self.assertEqual(backend.full_cg_prefill_swa_out_cache_loc.shape, (8,)) + forward_batch.out_cache_loc = torch.arange(9, dtype=torch.int64) + with self.assertRaisesRegex(AssertionError, "used 9 > capacity 8"): + backend._init_full_cg_prefill_metadata(forward_batch, in_capture=False) + + def test_empty_cpu_sequence_lengths_use_static_bound(self): + backend = FlashAttentionBackend.__new__(FlashAttentionBackend) + backend.max_context_len = 16 + backend.req_to_token_pool = SimpleNamespace( + req_to_token=torch.zeros((1, 16), dtype=torch.int32) + ) + backend.is_prefill_aware_swa = False + backend.has_swa = False + backend.use_sliding_window_kv_pool = False + backend._unified_dense = False + backend.page_size = 1 + backend._compute_scheduler_metadata = lambda *_: None + backend._maybe_init_local_attn_metadata = lambda *_: None + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.DECODE, + seq_lens=torch.empty(0, dtype=torch.int64), + batch_size=0, + seq_lens_cpu=torch.empty(0, dtype=torch.int64), + spec_info=None, + req_pool_indices=torch.empty(0, dtype=torch.int64), + encoder_lens=None, + ) + + backend.init_forward_metadata(forward_batch) + + self.assertEqual(backend.forward_metadata.max_seq_len_k, 16) + + +if __name__ == "__main__": + unittest.main()