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

[Perf] Skip blocks past per-request live length in full-width Triton kernels (#32109)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-22 22:00:25 -07:00
committed by GitHub
parent eb242b6c03
commit 1b63155efe
9 changed files with 187 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_normal_decode_set_metadata.py
+42 -3
View File
@@ -56,6 +56,22 @@ def reference_normal_decode_set_metadata(
swa_page_table[:, :max_seq_pages].copy_(swa_page_indices // page_size)
def page_table_live_mask(
seq_lens: torch.Tensor,
seq_len_delta: int,
page_size: int,
max_seq_pages: int,
width: int,
) -> torch.Tensor:
"""Per-row live region of the page table: the fused kernel self-guards on
seq_len and leaves columns past pages(seq_len + delta) untouched."""
live_pages = torch.clamp(
(seq_lens + seq_len_delta + page_size - 1) // page_size, max=max_seq_pages
)
cols = torch.arange(width, device=seq_lens.device)
return cols.view(1, -1) < live_pages.view(-1, 1)
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
class TestNormalDecodeSetMetadata(CustomTestCase):
"""Test fused Triton kernel in normal_decode_set_metadata."""
@@ -235,14 +251,26 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
f"cu_seqlens_k mismatch. Expected:\n{ref_data['cu_seqlens_k']}\nGot:\n{test_data['cu_seqlens_k']}",
)
live_mask = page_table_live_mask(
test_data["seq_lens"],
test_data["seq_len_delta"],
page_size,
test_data["max_seq_pages"],
test_data["page_table"].shape[1],
)
self.assertTrue(
torch.equal(test_data["page_table"], ref_data["page_table"]),
torch.equal(
test_data["page_table"][live_mask], ref_data["page_table"][live_mask]
),
f"page_table mismatch at bs={batch_size}, page_size={page_size}",
)
if has_swa:
self.assertTrue(
torch.equal(test_data["swa_page_table"], ref_data["swa_page_table"]),
torch.equal(
test_data["swa_page_table"][live_mask],
ref_data["swa_page_table"][live_mask],
),
f"swa_page_table mismatch at bs={batch_size}, page_size={page_size}",
)
@@ -412,7 +440,18 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
self.assertTrue(
torch.equal(test_data["cu_seqlens_k"], ref_data["cu_seqlens_k"])
)
self.assertTrue(torch.equal(test_data["page_table"], ref_data["page_table"]))
live_mask = page_table_live_mask(
test_data["seq_lens"],
0,
page_size,
test_data["max_seq_pages"],
test_data["page_table"].shape[1],
)
self.assertTrue(
torch.equal(
test_data["page_table"][live_mask], ref_data["page_table"][live_mask]
)
)
if __name__ == "__main__":
test/registered/attention/test_trtllm_mha_graph_metadata.py
+17 -7
View File
@@ -266,10 +266,8 @@ def _make_swa_mapping(pool_token_cap, seed):
@pytest.mark.parametrize("q_mode", [Q_MODE_NONE, Q_MODE_CUMSUM, Q_MODE_STRIDED])
@pytest.mark.parametrize("with_swa", [False, True])
# static_width=True exercises the production path: the backend passes the STATIC
# max_num_pages (full upper bound), not a per-batch dynamic width, so the kernel
# rewrites the whole page-table width every replay (tail pages beyond a request's
# seq are gathered but ignored by the attention kernel via cache_seqlens). The
# aten reference gathers the same full width, so this stays a bit-exact check.
# max_num_pages upper bound, not a per-batch dynamic width. The kernel self-guards
# on the device-side seqlen, so the checks compare only the live prefix per row.
@pytest.mark.parametrize("static_width", [False, True])
def test_metadata_correctness(bs, seqlen_offset, q_mode, with_swa, static_width):
if not torch.cuda.is_available():
@@ -372,9 +370,18 @@ def test_metadata_correctness(bs, seqlen_offset, q_mode, with_swa, static_width)
cu_k_ref[1:] = torch.cumsum(cache_seqlens_ref, dim=0, dtype=torch.int32)
torch.testing.assert_close(cu_seqlens_k, cu_k_ref, rtol=0, atol=0)
# ---- page_table ----
# ---- page_table (live [:pages(cache_seqlen)] prefix per row) ----
pt_ref, gathered = _ref_page_table(req_to_token, req_pool_indices, max_seq_pages)
torch.testing.assert_close(page_table[:, :max_seq_pages], pt_ref, rtol=0, atol=0)
live_pages = torch.clamp(
(cache_seqlens_ref.to(torch.int64) + PAGE_SIZE - 1) // PAGE_SIZE,
max=max_seq_pages,
)
live_mask = torch.arange(max_seq_pages, device=DEVICE).view(
1, -1
) < live_pages.view(-1, 1)
torch.testing.assert_close(
page_table[:, :max_seq_pages][live_mask], pt_ref[live_mask], rtol=0, atol=0
)
# ---- cu_seqlens_q ----
if q_mode == Q_MODE_CUMSUM:
@@ -392,7 +399,10 @@ def test_metadata_correctness(bs, seqlen_offset, q_mode, with_swa, static_width)
if with_swa:
swa_pt_ref = _ref_swa_page_table(gathered, swa_mapping)
torch.testing.assert_close(
swa_page_table[:, :max_seq_pages], swa_pt_ref, rtol=0, atol=0
swa_page_table[:, :max_seq_pages][live_mask],
swa_pt_ref[live_mask],
rtol=0,
atol=0,
)
# swa_out_cache_loc reference: translate real prefix, zero-fill tail.
test/registered/kernels/ops/attention/test_dsa_metadata.py
+47 -11
View File
@@ -102,18 +102,35 @@ class TestDSAMetadataKernels(CustomTestCase):
expected_page_table = req_to_token[req_pool_indices, :max_len].contiguous()
expected_dsa = _dsa_seqlens(expected_cache, dsa_index_topk)
# Compare only the live prefix [:seq_len]: whole blocks starting past
# the kv length are skipped (keep stale values), while the last
# partially live block still writes lanes past it -- the tail is
# unspecified either way, and consumers never read past cache_seqlens.
cols = torch.arange(max_len, dtype=torch.int32, device=self.device)
live_mask = cols.view(1, -1) < expected_cache.view(-1, 1)
_assert_equal(cache_seqlens, expected_cache, "decode cache_seqlens")
_assert_equal(cu_seqlens_k, _cu_seqlens(expected_cache), "decode cu_seqlens_k")
_assert_equal(page_table_1, expected_page_table, "decode page_table_1")
_assert_equal(
page_table_1[live_mask],
expected_page_table[live_mask],
"decode page_table_1 (live [:seq_len] prefix)",
)
_assert_equal(dsa_cache_seqlens, expected_dsa, "decode dsa_cache_seqlens")
_assert_equal(
dsa_cu_seqlens_k, _cu_seqlens(expected_dsa), "decode dsa_cu_seqlens_k"
)
if real_page_size > 1:
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
) < expected_cache.view(-1, 1)
expected_real = _real_page_table(expected_page_table, real_page_size)
_assert_equal(
real_page_table,
_real_page_table(expected_page_table, real_page_size),
"decode real_page_table",
real_page_table[real_live_mask],
expected_real[real_live_mask],
"decode real_page_table (live [:seq_len] prefix)",
)
def _check_target_verify(
@@ -194,19 +211,37 @@ class TestDSAMetadataKernels(CustomTestCase):
expected_expanded = expected_expanded.reshape(-1).contiguous()
expected_dsa = _dsa_seqlens(expected_expanded, dsa_index_topk)
# Compare only the live prefix [:seq_len + next_n] per expanded row:
# whole blocks starting past the kv length are skipped, the last
# partially live block may still write past it -- the tail is
# unspecified, and consumers never read past cache_seqlens.
row_kv_lens = torch.repeat_interleave(expected_cache, next_n)
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, "target cache_seqlens")
_assert_equal(cu_seqlens_k, _cu_seqlens(expected_cache), "target cu_seqlens_k")
_assert_equal(page_table_1, expected_page_table, "target page_table_1")
_assert_equal(
page_table_1[live_mask],
expected_page_table[live_mask],
"target page_table_1 (live [:seq_len + next_n] prefix)",
)
_assert_equal(seqlens_expanded, expected_expanded, "target seqlens_expanded")
_assert_equal(dsa_cache_seqlens, expected_dsa, "target dsa_cache_seqlens")
_assert_equal(
dsa_cu_seqlens_k, _cu_seqlens(expected_dsa), "target dsa_cu_seqlens_k"
)
if real_page_size > 1:
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,
_real_page_table(expected_page_table, real_page_size),
"target real_page_table",
real_page_table[real_live_mask],
expected_real[real_live_mask],
"target real_page_table (live [:seq_len + next_n] prefix)",
)
if fill_ctx_lens:
expected_ctx = expected_cache.view(bs, 1).expand(bs, next_n).contiguous()
@@ -304,9 +339,10 @@ 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.
# Compare only the live prefix [:kv_len]: whole blocks starting past
# kv_len are skipped, the last partially live block may still write
# past it -- the tail is unspecified, and consumers never read past
# cache_seqlens. 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)