chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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

10
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -873,6 +873,7 @@ class Indexer(MultiPlatformOp):
shared_index_buffer: Optional[torch.Tensor] = None,
shared_block_tables: Optional[torch.Tensor] = None,
actual_seq_q_tensor: Optional[torch.Tensor] = None,
actual_seq_q_cu_tensor: Optional[torch.Tensor] = None,
) -> torch.Tensor:
if TYPE_CHECKING:
assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
@@ -1064,6 +1065,7 @@ class Indexer(MultiPlatformOp):
clean_logits=False,
)
topk_indices_offset_override = None
cu_seqlens_q_topk_override = None
if (
getattr(getattr(metadata, "topk_transform_method", None), "name", None)
== "RAGGED"
@@ -1074,8 +1076,13 @@ class Indexer(MultiPlatformOp):
# post-MQA metadata kernels entirely.
topk_indices_offset_override = ks
actual_seq_q_tensor = None
elif valid_q_count == actual_seq_q and actual_seq_q_cu_tensor is not None:
cu_seqlens_q_topk_override = actual_seq_q_cu_tensor
elif actual_seq_q_tensor is None or valid_q_count != actual_seq_q:
actual_seq_q_tensor = ke_offset.new_full((1,), valid_q_count)
cu_seqlens_q_topk_override = ke_offset.new_empty((2,))
cu_seqlens_q_topk_override[0] = 0
cu_seqlens_q_topk_override[1] = actual_seq_q_tensor.reshape(-1)[0]
elif actual_seq_q_tensor.ndim == 0:
actual_seq_q_tensor = actual_seq_q_tensor.reshape(1)
valid_topk_result = metadata.topk_transform(
@@ -1085,6 +1092,7 @@ class Indexer(MultiPlatformOp):
cu_seqlens_q=actual_seq_q_tensor,
ke_offset=ke_offset,
topk_indices_offset_override=topk_indices_offset_override,
cu_seqlens_q_topk_override=cu_seqlens_q_topk_override,
)
if valid_q_count == actual_seq_q:
topk_result = valid_topk_result
@@ -1154,6 +1162,7 @@ class Indexer(MultiPlatformOp):
shared_index_buffer=shared_index_buffer,
shared_block_tables=shared_block_tables,
actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor,
actual_seq_q_cu_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_cu_tensor,
)
topk_result_next = self._get_topk_ragged_with_cp(
@@ -1168,6 +1177,7 @@ class Indexer(MultiPlatformOp):
shared_index_buffer=shared_index_buffer,
shared_block_tables=shared_block_tables,
actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor,
actual_seq_q_cu_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_cu_tensor,
)
return torch.cat([topk_result_prev, topk_result_next], dim=0)

10
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -180,6 +180,8 @@ class NSAContextParallelMetadata:
kv_len_next_tensor: torch.Tensor = None
actual_seq_q_prev_tensor: torch.Tensor = None
actual_seq_q_next_tensor: torch.Tensor = None
actual_seq_q_prev_cu_tensor: torch.Tensor = None
actual_seq_q_next_cu_tensor: torch.Tensor = None
total_seq_lens: torch.Tensor = None
page_aligned: bool = False
page_size: int = 1
@@ -921,6 +923,12 @@ def prepare_input_dp_with_cp_dsa(
actual_seq_q_next_tensor = torch.tensor(actual_seq_q_next).to(
device="cuda", dtype=torch.int32
)
actual_seq_q_prev_cu_tensor = torch.tensor(
[0, actual_seq_q_prev], device="cuda", dtype=torch.int32
)
actual_seq_q_next_cu_tensor = torch.tensor(
[0, actual_seq_q_next], device="cuda", dtype=torch.int32
)
nsa_cp_metadata = NSAContextParallelMetadata(
split_list=split_list,
@@ -937,6 +945,8 @@ def prepare_input_dp_with_cp_dsa(
kv_len_next_tensor=kv_len_next_tensor,
actual_seq_q_prev_tensor=actual_seq_q_prev_tensor,
actual_seq_q_next_tensor=actual_seq_q_next_tensor,
actual_seq_q_prev_cu_tensor=actual_seq_q_prev_cu_tensor,
actual_seq_q_next_cu_tensor=actual_seq_q_next_cu_tensor,
total_seq_lens=comm_total_seq_lens,
page_aligned=page_split_info.page_aligned,
page_size=page_split_info.page_size,

37
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -357,6 +357,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
ke_offset: torch.Tensor = None,
batch_idx_list: List[int] = None,
topk_indices_offset_override: Optional[torch.Tensor] = None,
cu_seqlens_q_topk_override: Optional[torch.Tensor] = None,
) -> torch.Tensor:
from sgl_kernel import (
fast_topk_transform_fused,
@@ -364,19 +365,6 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
fast_topk_v2,
)
if topk_indices_offset_override is not None:
cu_topk_indices_offset = topk_indices_offset_override
cu_seqlens_q_topk = None
elif cu_seqlens_q is not None:
cu_seqlens_q = cu_seqlens_q.to(torch.int32)
cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q)
cu_topk_indices_offset = torch.repeat_interleave(
cu_seqlens_q_topk[:-1],
cu_seqlens_q,
)
else:
cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q
cu_topk_indices_offset = self.attn_metadata.topk_indices_offset
if ke_offset is not None:
seq_lens_topk = ke_offset
else:
@@ -389,6 +377,14 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
if not envs.SGLANG_NSA_FUSE_TOPK.get() or self.force_unfused_topk:
return fast_topk_v2(logits, seq_lens_topk, topk, row_starts=ks)
elif self.topk_transform_method == TopkTransformMethod.PAGED:
if cu_seqlens_q_topk_override is not None:
cu_seqlens_q_topk = cu_seqlens_q_topk_override
elif cu_seqlens_q is not None:
if cu_seqlens_q.dtype != torch.int32:
cu_seqlens_q = cu_seqlens_q.to(torch.int32)
cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q)
else:
cu_seqlens_q_topk = self.attn_metadata.cu_seqlens_q
# NOTE(dark): if fused, we return a transformed page table directly
validate_paged_topk = (
self.validate_paged_topk
@@ -415,6 +411,21 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
_validate_paged_topk_transform_output(topk_result, page_table_size_1)
return topk_result
elif self.topk_transform_method == TopkTransformMethod.RAGGED:
if topk_indices_offset_override is not None:
cu_topk_indices_offset = topk_indices_offset_override
elif cu_seqlens_q is not None:
if cu_seqlens_q.dtype != torch.int32:
cu_seqlens_q = cu_seqlens_q.to(torch.int32)
if cu_seqlens_q_topk_override is not None:
cu_seqlens_q_topk = cu_seqlens_q_topk_override
else:
cu_seqlens_q_topk = compute_cu_seqlens(cu_seqlens_q)
cu_topk_indices_offset = torch.repeat_interleave(
cu_seqlens_q_topk[:-1],
cu_seqlens_q,
)
else:
cu_topk_indices_offset = self.attn_metadata.topk_indices_offset
return fast_topk_transform_ragged_fused(
score=logits,
lengths=seq_lens_topk,

87
test/registered/unit/layers/test_nsa_cp_utils.py
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()