chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix Illegal Memory Access when fa3 + spec + topk + page_size > 1 (#15469)

Browse Source 
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
This commit is contained in:
Yubo Wang
2025-12-24 00:13:57 +08:00
committed by GitHub
co-authored by Liangsheng Yin
parent dd620987d1
commit 762846531f
5 changed files with 133 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/flashattention_backend.py
+40 -63
View File
@@ -668,24 +668,28 @@ class FlashAttentionBackend(AttentionBackend):
self.forward_metadata_spec_decode_expand.cache_seqlens_int32 += (
expanded_last_page_lens
)
decode_length = self.speculative_step_id + 1
expand_page_table = cache_loc[:, :decode_length].clone()
strided_indices_expand = torch.arange(
0,
decode_length,
self.page_size,
# NOTE: the max decode length is speculative_num_steps - 1 (one token always generated by draft extend)
# and we leave one extra for last_page_len, which -> speculative_num_steps for the page table
expand_page_table = torch.zeros(
forward_batch.batch_size * self.topk,
self.speculative_num_steps,
dtype=torch.int32,
device=self.device,
)
last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand(
-1, expand_page_table.shape[1]
# shape: [bs, num_steps, topk] -> [bs x topk, num_steps]
cache_loc = forward_batch.out_cache_loc.view(
-1, self.speculative_num_steps
)
expand_page_table -= last_page_lens_broadcast
expand_page_table = (
expand_page_table[:, strided_indices_expand] // self.page_size
)
self.forward_metadata_spec_decode_expand.page_table = (
expand_page_table.to(torch.int32)
draft_decode_set_expand_metadata(
cache_seqlens_int32=self.forward_metadata_spec_decode_expand.cache_seqlens_int32,
page_table=expand_page_table,
last_page_lens=last_page_lens,
decode_length=decode_length,
cache_loc=cache_loc,
topk=self.topk,
page_size=self.page_size,
)
self.forward_metadata_spec_decode_expand.page_table = expand_page_table
self.forward_metadata = metadata
@@ -1404,23 +1408,12 @@ class FlashAttentionBackend(AttentionBackend):
),
"page_table": torch.zeros(
max_bs * self.topk,
decode_length,
decode_length + 1, # Additional page for last partial page
dtype=torch.int32,
device=self.device,
),
}
if self.page_size > 1:
# Used for indicing expand page_table
self.draft_decode_metadata_topk_expand["strided_indices_expand"] = (
torch.arange(
0,
self.speculative_num_steps,
self.page_size,
device=self.device,
)
)
if (
self.speculative_num_draft_tokens is not None
and self.speculative_num_draft_tokens > 0
@@ -1859,8 +1852,6 @@ class FlashAttentionBackend(AttentionBackend):
# First attention handles seq_lens - last_page_lens if page size > 1.
last_page_lens = seq_lens % self.page_size
seq_lens = seq_lens - last_page_lens
# last_page_lens_cpu = last_page_lens.max().item()
# seq_lens_cpu -= last_page_lens_cpu
metadata.cache_seqlens_int32.copy_(seq_lens)
# metadata.max_seq_len_q = self.topk, already set in capture
# metadata.cu_seqlens_q already set in capture
@@ -1886,25 +1877,18 @@ 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:
# Second attention handles last_page_len + decode part.
strided_indices_expand = (
self.draft_decode_metadata_topk_expand.get(
"strided_indices_expand"
)
)
update_draft_decode_set_expand_metadata_with_page_size(
metadata_expand.cache_seqlens_int32, # Modifies
metadata_expand.page_table, # Modifies
cache_loc,
last_page_lens,
strided_indices_expand,
decode_length,
bs,
self.topk,
self.page_size,
draft_decode_set_expand_metadata(
cache_seqlens_int32=metadata_expand.cache_seqlens_int32,
page_table=metadata_expand.page_table,
last_page_lens=last_page_lens,
decode_length=decode_length,
cache_loc=cache_loc,
topk=self.topk,
page_size=self.page_size,
)
else:
metadata_expand.page_table[: cache_loc.shape[0]].copy_(
num_seqs = cache_loc.shape[0]
metadata_expand.page_table[:num_seqs, :decode_length].copy_(
cache_loc[:, :decode_length]
)
# TODO: Handle local attention metadata for draft decode when llama4 eagle is supported
@@ -2571,30 +2555,23 @@ def normal_decode_set_metadata(
@torch.compile(dynamic=True, backend=get_compiler_backend())
def update_draft_decode_set_expand_metadata_with_page_size(
def draft_decode_set_expand_metadata(
cache_seqlens_int32: torch.Tensor, # Modifies
page_table: torch.Tensor, # Modifies
cache_loc: torch.Tensor,
last_page_lens: torch.Tensor,
strided_indices_expand: torch.Tensor,
decode_length: int,
bs: int,
cache_loc: torch.Tensor,
topk: int,
page_size: int,
):
expanded_last_page_lens = last_page_lens.repeat_interleave(topk)
cache_seqlens_int32.copy_(decode_length + expanded_last_page_lens)
expand_page_table = cache_loc[:, :decode_length].clone()
last_page_lens_broadcast = expanded_last_page_lens.unsqueeze(-1).expand(
-1, expand_page_table.shape[1]
)
num_seqs = expand_page_table.shape[0]
expand_page_table -= last_page_lens_broadcast[:num_seqs]
expand_page_table = (
expand_page_table[
:, strided_indices_expand[: (decode_length + page_size - 1) // page_size]
]
// page_size
)
max_seq_pages_expand = (decode_length + page_size - 1) // page_size
page_table[:num_seqs, :max_seq_pages_expand].copy_(expand_page_table)
cache_loc = (cache_loc // page_size).to(torch.int32)
if cache_loc.dim() == 1:
cache_loc = cache_loc.unsqueeze(0)
# 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]
positions = mask.cumsum(dim=1) - 1
num_seqs = cache_loc.shape[0]
page_table[:num_seqs, :].scatter_(1, positions, cache_loc)