From 665cb02003575914fc5f49e7d47955b87cdb227e Mon Sep 17 00:00:00 2001 From: Shu Wang Date: Fri, 12 Dec 2025 21:11:09 -0600 Subject: [PATCH] Fix regression caused by fa3 block_table (#15009) --- .../srt/layers/attention/flashattention_backend.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index f3b44e895..98e0e9ecc 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -263,7 +263,14 @@ def make_local_attention_virtual_batches( np.arange(actual_batch_size, dtype=np.int32), local_blocks * pages_per_local_batch, ) - block_table_local = block_table[batch_indices, block_indices].view( + + # NOTE: https://github.com/pytorch/pytorch/pull/160256 causes performance + # regression when using numpy arrays (batch and block indices) to index into + # torch tensor (block_table). As a workaround, convert numpy arrays to torch + # tensor first, which recovers perf. + batch_indices_torch = torch.from_numpy(batch_indices) + block_indices_torch = torch.from_numpy(block_indices) + block_table_local = block_table[batch_indices_torch, block_indices_torch].view( virtual_batches, -1 )