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