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

Avoid PAGED topk metadata scans after MQA

Browse Source 
Prefill CP shared KV uses the PAGED fused topk path, but topk_transform still built RAGGED topk offsets before dispatching by method. That introduced cumsum/repeat_interleave work after MQA, showing up as DeviceScanInitKernel and host/device traffic in profiles even though PAGED topk only needs cu_seqlens_q. Move metadata construction into the selected branch and pass precomputed single-segment CP cu_seqlens overrides from NSA CP metadata.

Constraint: PAGED fused topk needs cu_seqlens_q but does not consume topk_indices_offset.\nRejected: Add a kernel to fuse repeat_interleave for PAGED | the offset is unused in the current path, so avoiding it is cheaper and safer.\nConfidence: high\nScope-risk: narrow\nDirective: Do not reintroduce unconditional topk_indices_offset construction in topk_transform; keep RAGGED-only metadata on the RAGGED branch.\nTested: python -m py_compile for modified files locally; g0034 container python3 -m py_compile for modified files; g0034 container python3 test/registered/unit/layers/test_nsa_cp_utils.py ran 23 tests OK.\nNot-tested: Full server profile after restart; full SGLang test suite.
This commit is contained in:
laoyao0822
2026-05-03 23:08:17 +08:00
parent 9eb9d82b51
commit a638d71d53
4 changed files with 131 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/layers/test_nsa_cp_utils.py
+87
View File
@@ -1,4 +1,5 @@
import unittest
import sys
from types import SimpleNamespace
from unittest.mock import patch
@@ -454,12 +455,14 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
shared_index_buffer=None,
shared_block_tables=None,
actual_seq_q_tensor=None,
actual_seq_q_cu_tensor=None,
):
topk_calls.append(
{
"kv_len": kv_len,
"actual_seq_q": actual_seq_q,
"actual_seq_q_tensor": actual_seq_q_tensor,
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
"shared_index_buffer": shared_index_buffer,
"shared_block_tables": shared_block_tables,
"current_index_kv": current_index_kv,
@@ -479,6 +482,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
kv_len_next=9,
actual_seq_q_prev=3,
actual_seq_q_next=2,
actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32),
actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
)
},
)()
@@ -505,6 +510,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertIsNone(topk_calls[0]["current_index_kv"])
self.assertEqual(topk_calls[0]["kv_len"], 5)
self.assertEqual(topk_calls[1]["kv_len"], 9)
self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3])
self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2])
self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
@@ -538,11 +545,13 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
shared_index_buffer=None,
shared_block_tables=None,
actual_seq_q_tensor=None,
actual_seq_q_cu_tensor=None,
):
topk_calls.append(
{
"current_index_kv": current_index_kv,
"actual_seq_q_tensor": actual_seq_q_tensor,
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
"shared_index_buffer": shared_index_buffer,
"shared_block_tables": shared_block_tables,
}
@@ -561,6 +570,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
kv_len_next=9,
actual_seq_q_prev=3,
actual_seq_q_next=2,
actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32),
actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
)
},
)()
@@ -581,8 +592,84 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
self.assertIs(topk_calls[1]["current_index_kv"], current_index_kv)
self.assertIsNone(topk_calls[0]["shared_index_buffer"])
self.assertIsNone(topk_calls[1]["shared_block_tables"])
self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3])
self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2])
self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
def test_paged_topk_transform_uses_cu_override_without_scan_metadata_ops(self):
import torch
from sglang.srt.layers.attention.nsa_backend import (
NSAMetadata,
NSAIndexerMetadata,
TopkTransformMethod,
)
cu_override = torch.tensor([0, 4], dtype=torch.int32)
attn_metadata = NSAMetadata(
page_size=64,
cache_seqlens_int32=torch.tensor([4], dtype=torch.int32),
max_seq_len_q=4,
max_seq_len_k=8,
cu_seqlens_q=torch.tensor([0, 4], dtype=torch.int32),
cu_seqlens_k=torch.tensor([0, 8], dtype=torch.int32),
page_table_1=torch.arange(8, dtype=torch.int32).view(1, 8),
real_page_table=torch.arange(8, dtype=torch.int32).view(1, 8),
nsa_cache_seqlens_int32=torch.tensor([4], dtype=torch.int32),
nsa_cu_seqlens_q=torch.arange(2, dtype=torch.int32),
nsa_cu_seqlens_k=torch.tensor([0, 4], dtype=torch.int32),
nsa_extend_seq_lens_list=[4],
nsa_seqlens_expanded=torch.arange(1, 5, dtype=torch.int32),
topk_indices_offset=torch.zeros(4, dtype=torch.int32),
)
metadata = NSAIndexerMetadata(
attn_metadata=attn_metadata,
topk_transform_method=TopkTransformMethod.PAGED,
)
logits = torch.zeros(4, 8)
lengths = torch.arange(1, 5, dtype=torch.int32)
expected = torch.full((4, 2), 7, dtype=torch.int32)
def fake_fused(**kwargs):
self.assertIs(kwargs["cu_seqlens_q"], cu_override)
self.assertIs(kwargs["lengths"], lengths)
return expected
fake_sgl_kernel = SimpleNamespace(
fast_topk_transform_fused=fake_fused,
fast_topk_transform_ragged_fused=lambda **_: (_ for _ in ()).throw(
AssertionError("ragged path should not run")
),
fast_topk_v2=lambda *_, **__: (_ for _ in ()).throw(
AssertionError("unfused path should not run")
),
)
with (
patch.dict(sys.modules, {"sgl_kernel": fake_sgl_kernel}),
patch(
"sglang.srt.layers.attention.nsa_backend.envs.SGLANG_NSA_FUSE_TOPK.get",
return_value=True,
),
patch(
"sglang.srt.layers.attention.nsa_backend.compute_cu_seqlens",
side_effect=AssertionError("paged override should skip cumsum"),
),
patch(
"torch.repeat_interleave",
side_effect=AssertionError("paged topk should not build ragged offsets"),
),
):
actual = metadata.topk_transform(
logits,
topk=2,
cu_seqlens_q=torch.tensor([4], dtype=torch.int32),
ke_offset=lengths,
cu_seqlens_q_topk_override=cu_override,
)
self.assertIs(actual, expected)
if __name__ == "__main__":
unittest.main()