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-23 08:13:57 -08:00
committed by GitHub
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

103
python/sglang/srt/layers/attention/flashattention_backend.py
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)

96
python/sglang/test/attention/test_flashattn_backend.py
View File

@@ -5,7 +5,7 @@ import torch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.flashattention_backend import (
FlashAttentionBackend,
update_draft_decode_set_expand_metadata_with_page_size,
draft_decode_set_expand_metadata,
)
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
from sglang.srt.layers.radix_attention import RadixAttention
@@ -350,8 +350,13 @@ class TestFlashAttentionBackend(CustomTestCase):
class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
def test_update_draft_decode_set_expand_metadata_with_page_size(self):
bs, topk, decode_length, page_size = 1, 2, 1, 4
"""
All the test cases examples have 1 additional cache location than the decode length.
This is to align with the current allocation logic. It does not affect the correctness.
"""
def test_draft_decode_set_expand_metadata(self):
bs, topk, page_size = 1, 2, 4
cases = [
(
@@ -364,56 +369,103 @@ class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
),
torch.tensor(
[
[5],
[7],
[5, 6],
[7, 8],
],
dtype=torch.int32,
),
1,
),
# Decode span multiple pages:
# duplicated kv cache: 24, 25, 26
# decode locations: 27, 28, 29, 30, 31, 32
# We need 3 pages in total.
(
torch.tensor(
[
[27, 28],
[35, 36],
[27, 28, 29, 30, 31, 32],
[35, 36, 37, 38, 39, 40],
],
dtype=torch.int32,
),
torch.tensor(
[
[6],
[8],
[6, 7, 8, 0, 0, 0],
[8, 9, 10, 0, 0, 0],
],
dtype=torch.int32,
),
5,
),
]
last_page_lens = torch.tensor([3], dtype=torch.int32)
strided_indices_expand = torch.arange(
0, decode_length, page_size, dtype=torch.long
)
for cache_loc, expected_page_table in cases:
for cache_loc, expected_page_table, decode_length in cases:
cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
page_table = torch.zeros(bs * topk, decode_length, dtype=torch.int32)
update_draft_decode_set_expand_metadata_with_page_size(
page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
draft_decode_set_expand_metadata(
cache_seqlens_int32=cache_seqlens_int32,
page_table=page_table,
cache_loc=cache_loc,
last_page_lens=last_page_lens,
strided_indices_expand=strided_indices_expand,
decode_length=decode_length,
bs=bs,
cache_loc=cache_loc,
topk=topk,
page_size=page_size,
)
expected_cache_seqlens = torch.tensor([4, 4], dtype=torch.int32)
expected_cache_seqlens = torch.tensor(
[decode_length + 3, decode_length + 3], dtype=torch.int32
)
self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
self.assertTrue(torch.equal(page_table, expected_page_table))
def test_update_draft_decode_set_expand_metadata_multi_batch(self):
"""
Ensure expand metadata works when batch size > 1 and last pages differ.
"""
bs, topk, decode_length, page_size = 3, 2, 3, 4
cache_loc = torch.tensor(
[
# First batch: last page duplicate is 1, consecutive pages
[1, 2, 3, 4],
[6, 7, 8, 9],
# Second batch: last page duplicate is 3, non-consecutive pages
[3, 8, 9, 10],
[14, 15, 16, 17],
# Third batch: last page duplicate is 0, consecutive pages
[0, 1, 2, 3],
[4, 5, 6, 7],
],
dtype=torch.int32,
)
cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
last_page_lens = torch.tensor([1, 3, 0], dtype=torch.int32)
page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
draft_decode_set_expand_metadata(
cache_seqlens_int32=cache_seqlens_int32,
page_table=page_table,
last_page_lens=last_page_lens,
decode_length=decode_length,
cache_loc=cache_loc,
topk=topk,
page_size=page_size,
)
expected_cache_seqlens = torch.tensor([4, 4, 6, 6, 3, 3], dtype=torch.int32)
expected_page_table = torch.tensor(
[
[0, 1, 0, 0],
[1, 2, 0, 0],
[0, 2, 0, 0],
[3, 4, 0, 0],
[0, 0, 0, 0],
[1, 0, 0, 0],
],
dtype=torch.int32,
)
self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
self.assertTrue(torch.equal(page_table, expected_page_table))
if __name__ == "__main__":
unittest.main()