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

[Perf] Skip page-table columns past kv length in DSA draft-extend metadata kernel (#31981)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-21 16:20:14 -07:00
committed by GitHub
parent 4a55fdba0b
commit 08cb081ed0
2 changed files with 29 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsa_metadata.py
+9
View File
@@ -524,6 +524,15 @@ def _fused_dsa_draft_extend_metadata_kernel(
mask=req_row < bs,
other=0,
).to(tl.int32)
# Skip column blocks past the request's kv length: no consumer reads there
# (attention and the indexer both stay within cache_seqlens).
kv_len = tl.load(
seq_lens + req_row * seq_lens_stride,
mask=req_row < bs,
other=0,
).to(tl.int32)
if col_block * BLOCK_N >= kv_len:
return
if STATIC_EXTEND_LEN:
prefix = req_row * qo_len
else:
test/registered/kernels/test_dsa_metadata.py
+20 -4
View File
@@ -304,10 +304,19 @@ class TestDSAMetadataKernels(CustomTestCase):
)
expected_dsa = _dsa_seqlens(expected_expanded, dsa_index_topk)
# Only the live prefix [:kv_len] per request is defined; the kernel
# leaves columns past kv_len untouched. All expanded rows of a request
# share its kv length.
row_kv_lens = torch.repeat_interleave(seq_lens.to(torch.int32), extend_seq_lens)
cols = torch.arange(max_seqlen_k, dtype=torch.int32, device=self.device)
live_mask = cols.view(1, -1) < row_kv_lens.view(-1, 1)
_assert_equal(cache_seqlens, expected_cache, "draft cache_seqlens")
_assert_equal(cu_seqlens_k, _cu_seqlens(expected_cache), "draft cu_seqlens_k")
_assert_equal(
page_table_1[:total_len], expected_page_table, "draft page_table_1"
page_table_1[:total_len][live_mask],
expected_page_table[live_mask],
"draft page_table_1 (live [:kv_len] prefix)",
)
_assert_equal(
seqlens_expanded[:total_len], expected_expanded, "draft seqlens_expanded"
@@ -321,10 +330,17 @@ class TestDSAMetadataKernels(CustomTestCase):
"draft dsa_cu_seqlens_k",
)
if real_page_size > 1:
# Real-page column real_col maps to source column real_col*real_page_size.
real_width = real_page_table.shape[1]
real_cols = torch.arange(real_width, dtype=torch.int32, device=self.device)
real_live_mask = (
real_cols.view(1, -1) * real_page_size
) < row_kv_lens.view(-1, 1)
expected_real = _real_page_table(expected_page_table, real_page_size)
_assert_equal(
real_page_table[:total_len],
_real_page_table(expected_page_table, real_page_size),
"draft real_page_table",
real_page_table[:total_len][real_live_mask],
expected_real[real_live_mask],
"draft real_page_table (live [:kv_len] prefix)",
)
def test_decode_matches_eager_reference(self):