From 26b9053dcc5c0022944c0afd9a02040e642cef35 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 6 Jun 2026 03:24:34 -0400 Subject: [PATCH] [Spec] Fix fa3 EAGLE draft-decode expand page_table scatter OOB for topk>1 + page_size>1 (#27360) --- python/sglang/srt/debug_utils/pr_fix_toggle.py | 11 +++++++++++ .../layers/attention/flashattention_backend.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/python/sglang/srt/debug_utils/pr_fix_toggle.py b/python/sglang/srt/debug_utils/pr_fix_toggle.py index 23ce0f9a9..3eee2837d 100644 --- a/python/sglang/srt/debug_utils/pr_fix_toggle.py +++ b/python/sglang/srt/debug_utils/pr_fix_toggle.py @@ -66,9 +66,20 @@ patches: """ +_PR_REVERT_YAML_27360 = """ +patches: + - target: sglang.srt.layers.attention.flashattention_backend.FlashAttentionBackend._apply_cuda_graph_metadata + edits: + - match: | + cache_loc = cache_loc[:, :decode_length] + replacement: "" +""" + + _PR_FIX_REVERT_YAML: Dict[int, str] = { 25015: _PR_REVERT_YAML_25015, 26329: _PR_REVERT_YAML_26329, + 27360: _PR_REVERT_YAML_27360, } diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 2a9b1e854..55891709b 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -2102,6 +2102,20 @@ class FlashAttentionBackend(AttentionBackend): # shape: [bs, num_steps, topk] -> [bs x topk, num_steps] cache_loc = out_cache_loc.view(-1, self.speculative_num_steps) if self.page_size > 1: + # Only the draft tokens produced up to this step are live; + # cache_loc arrives num_steps-wide. Slice so the scatter fills at + # most decode_length of the (decode_length + 1) expand page_table + # columns -- without this the extra distinct pages overflow the row. + cache_loc = cache_loc[:, :decode_length] + assert ( + cache_loc.shape[1] <= metadata_expand.page_table.shape[1] + ), ( + f"draft expand page_table too narrow: cache_loc width " + f"{cache_loc.shape[1]} > " + f"{metadata_expand.page_table.shape[1]} columns " + f"(decode_length + 1); page_size={self.page_size}, " + f"topk={self.topk}, num_steps={self.speculative_num_steps}" + ) draft_decode_set_expand_metadata( cache_seqlens_int32=metadata_expand.cache_seqlens_int32, page_table=metadata_expand.page_table, @@ -2731,6 +2745,8 @@ def draft_decode_set_expand_metadata( cache_loc = (cache_loc // page_size).to(torch.int32) if cache_loc.dim() == 1: cache_loc = cache_loc.unsqueeze(0) + # cache_loc is pre-sliced to decode_length by the caller, so the scatter fills at + # most decode_length of the (decode_length + 1) page_table columns. # Vectorized torch.unique_consecutive: track value change points then scatter mask = torch.ones_like(cache_loc, dtype=torch.bool) mask[:, 1:] = cache_loc[:, 1:] != cache_loc[:, :-1]