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

Cut per-layer CPU on the prefill launch path: validators, plans, spans

Browse Source 
From the nsys CPU-gap attribution (launch thread, one 78-layer forward:
374ms API time; 642 cudaStreamSynchronize blocking 89.5ms and overlapping
122ms of the 505ms GPU idle; ~44ms pure-Python before concat_mla_absorb_q):

- memory_pool_host: skip validate_page_aligned_token_indices on CUDA
  tensors in _get_indexer_page_indices and _prepare_load_page_indices —
  torch.any/torch.equal there cost a queue-deep cudaStreamSynchronize per
  layer-group submit (~0.42ms each, ~12.7ms/forward measured). Same
  construction-based-invariant guard the CacheController pair check
  already documents; CPU/test tensors stay validated.
- nsa_indexer: per-batch _CpRaggedIndexPlan replaces the per-F-layer
  rebuild of the O(total-q-tokens) topk offset list and the 6-7 int32
  ragged descriptor tensors (segment records, kv_lens/q_starts/q_lens/
  k_bases/q_bases/current_bases). All inputs are batch metadata; the plan
  is anchored on the forward batch with a content key over cp_index.
- nsa_indexer forward_indexer: read seq_lens_cpu instead of a device
  seq_lens[i].item() per request per layer (one stream sync each).
- cp_shared_kv_runtime: get_or_build_batch_slot_spans caches the
  layer-invariant prefix/current slot spans per batch (the builders read
  logical_pages only for its shape); nsa_backend x3 + nsa_indexer call
  sites switched.

Microbenchmark (idle H200, traced batch shape bs=12 / 44.6K q tokens,
test/manual/bench_cpu_gap_fixes.py, equality-checked): validator path
197.1us -> 59.2us per submit under a busy queue (x3.3); ragged plan
3238.6us -> 36.2us per layer (x90, ~128ms launch-thread time per forward
at 40 F-layers); slot spans 20.1us -> 0.5us (x41). Layer suites A/B vs
HEAD: identical failure set (5 pre-existing CPU-tensor indexer tests),
no regressions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-11 21:01:00 +00:00
parent b356773d2f
commit a24111a5f4
5 changed files with 519 additions and 159 deletions
Download Patch File Download Diff File Expand all files Collapse all files

45
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -2787,6 +2787,51 @@ def build_batch_prefix_slot_span(
return (start_slot, end_slot)
def get_or_build_batch_slot_spans(
forward_batch,
*,
logical_pages: torch.Tensor,
prefix_lens_cpu,
extend_lens_cpu,
page_size: int,
want_prefix: bool,
) -> tuple[list[tuple[int, int]] | None, list[tuple[int, int]]]:
"""Per-batch cache for the layer-invariant slot-span builders.
The builders read ``logical_pages`` only for its SHAPE; together with the
batch-scoped ``prefix/extend`` lens that makes the spans identical for
every layer of a forward — rebuilding the per-request Python loops per
layer was part of the measured pre-attention CPU gap.
"""
key = (tuple(logical_pages.shape), int(page_size), bool(want_prefix))
cache = getattr(forward_batch, "_cp_batch_slot_spans_cache", None)
if cache is None:
cache = {}
forward_batch._cp_batch_slot_spans_cache = cache
hit = cache.get(key)
if hit is not None:
return hit
prefix_spans = (
build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens_cpu,
page_size=page_size,
)
if want_prefix
else None
)
current_spans = build_batch_current_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
result = (prefix_spans, current_spans)
cache[key] = result
return result
def build_batch_prefix_slot_spans(
*,
logical_pages: torch.Tensor,

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

@@ -26,6 +26,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_should_log_layer,
current_extend_kv_rows_for_reuse,
filter_owned_logical_locs,
get_or_build_batch_slot_spans,
get_or_build_shared_paged_buffer_slot_remap,
is_current_only_extend_batch,
log_cp_draft_shared_kv_debug,
@@ -199,6 +200,162 @@ def _build_current_index_request_bases(forward_batch: ForwardBatch) -> List[int]
return current_req_offsets
class _CpRaggedIndexPlan:
"""Layer-invariant ragged CP index descriptors, built once per batch.
Everything here is a pure function of batch metadata (``cp_index``,
``seq_lens_cpu``, ``extend_seq_lens_cpu``, the owner-lane request bases) —
the per-layer indexer used to rebuild the O(total-q-tokens)
``topk_indices_offset`` list and 6-7 int32 descriptor tensors from Python
lists on every F-layer (measured as a large share of the ~44 ms/forward
pre-attention Python gap).
"""
__slots__ = (
"segment_records",
"topk_indices_offset_override",
"batch_indices",
"kv_lens",
"q_starts",
"q_lens",
"k_bases",
"q_bases",
"current_bases",
"actual_seq_q",
"total_kv_len",
"total_q_count",
"max_kv_len",
"max_q_len",
)
def _build_cp_ragged_index_plan(
forward_batch: ForwardBatch,
cp_index,
device: torch.device,
current_req_offsets: Optional[List[int]],
) -> _CpRaggedIndexPlan:
seq_lens_cpu_list = forward_batch.seq_lens_cpu.tolist()
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
segment_records: List[Tuple[int, int, int, int, int, int, int, int]] = []
batch_idx_list: List[int] = []
kv_lens_list: List[int] = []
q_starts_list: List[int] = []
q_lens_list: List[int] = []
k_bases_list: List[int] = []
q_bases_list: List[int] = []
topk_offset_list: List[int] = []
request_kv_bases: List[int] = []
request_kv_base = 0
for seq_len in seq_lens_cpu_list:
request_kv_bases.append(int(request_kv_base))
request_kv_base += int(seq_len)
k_cursor = 0
q_cursor = 0
for raw_batch_idx, start_seq_position, end_seq_position in cp_index:
batch_idx = int(raw_batch_idx)
if batch_idx < 0 or batch_idx >= len(request_kv_bases):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_cp_index_bad_batch_idx "
f"batch_idx={batch_idx} seq_lens={seq_lens_cpu_list}"
)
pre_chunk_offset = int(seq_lens_cpu_list[batch_idx]) - int(
extend_seq_lens_cpu[batch_idx]
)
start_seq_position += pre_chunk_offset
end_seq_position += pre_chunk_offset
if end_seq_position < start_seq_position:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_cp_index_bad_segment "
f"batch_idx={batch_idx} start={start_seq_position} "
f"end={end_seq_position}"
)
extend_seq_len = int(end_seq_position - start_seq_position)
kv_len_i = int(end_seq_position)
segment_records.append(
(
batch_idx,
int(start_seq_position),
int(end_seq_position),
extend_seq_len,
kv_len_i,
k_cursor,
q_cursor,
int(pre_chunk_offset),
)
)
batch_idx_list.append(batch_idx)
kv_lens_list.append(kv_len_i)
q_starts_list.append(int(start_seq_position))
q_lens_list.append(extend_seq_len)
k_bases_list.append(k_cursor)
q_bases_list.append(q_cursor)
topk_offset_list.extend([request_kv_bases[batch_idx]] * extend_seq_len)
k_cursor += kv_len_i
q_cursor += extend_seq_len
plan = _CpRaggedIndexPlan()
plan.segment_records = segment_records
plan.topk_indices_offset_override = torch.tensor(
topk_offset_list, dtype=torch.int32, device=device
)
def _i32(values: List[int]) -> torch.Tensor:
return torch.tensor(values, dtype=torch.int32, device=device)
plan.batch_indices = _i32(batch_idx_list)
plan.kv_lens = _i32(kv_lens_list)
plan.q_starts = _i32(q_starts_list)
plan.q_lens = _i32(q_lens_list)
plan.k_bases = _i32(k_bases_list)
plan.q_bases = _i32(q_bases_list)
plan.actual_seq_q = plan.q_lens
plan.current_bases = (
_i32([int(current_req_offsets[b]) for b in batch_idx_list])
if current_req_offsets is not None
else None
)
plan.total_kv_len = k_cursor
plan.total_q_count = q_cursor
plan.max_kv_len = max(kv_lens_list, default=0)
plan.max_q_len = max(q_lens_list, default=0)
return plan
def _get_or_build_cp_ragged_index_plan(
forward_batch: ForwardBatch,
cp_index,
device: torch.device,
current_req_offsets: Optional[List[int]],
) -> _CpRaggedIndexPlan:
"""Per-batch cache of the ragged index plan, anchored on the batch.
The key is content-based (``cp_index`` may be rebuilt per layer); at
bs<=segments it is a handful of small tuples — negligible vs. rebuilding
the descriptor tensors.
"""
key = (
tuple((int(b), int(s), int(e)) for b, s, e in cp_index),
str(device),
current_req_offsets is not None,
)
plans = getattr(forward_batch, "_cp_ragged_index_plans", None)
if plans is None:
plans = {}
forward_batch._cp_ragged_index_plans = plans
plan = plans.get(key)
if plan is None:
plan = _build_cp_ragged_index_plan(
forward_batch, cp_index, device, current_req_offsets
)
plans[key] = plan
return plan
def _select_batch_topk_query_lengths(
*,
cp_metadata,
@@ -615,22 +772,18 @@ class Indexer(MultiPlatformOp):
current_req_id = torch.zeros_like(current_locs, dtype=torch.long)
else:
current_req_id = current_req_id[: int(current_locs.shape[0])]
prefix_slot_spans = None
current_slot_spans = build_batch_current_slot_spans(
logical_pages=logical_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
else:
prefix_pages = 0
prefix_slot_spans = build_batch_prefix_slot_spans(
logical_pages=logical_page_table,
prefix_lens_cpu=prefix_lens_cpu,
page_size=page_size,
)
prefix_slot_spans, current_slot_spans = get_or_build_batch_slot_spans(
forward_batch,
logical_pages=logical_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=len(prefix_lens_cpu) > 1,
)
if index_prefetcher is not None:
prefetched = index_prefetcher.consume_prefix_with_current(
layer_id=layer_id,
@@ -1597,101 +1750,31 @@ class Indexer(MultiPlatformOp):
forward_batch
)
segment_records: List[Tuple[int, int, int, int, int, int, int, int]] = []
batch_idx_list = []
kv_lens_list = []
q_starts_list = []
q_lens_list = []
k_bases_list = []
q_bases_list = []
topk_offset_list = []
request_kv_bases: List[int] = []
request_kv_base = 0
for seq_len in forward_batch.seq_lens_cpu.tolist():
request_kv_bases.append(int(request_kv_base))
request_kv_base += int(seq_len)
k_cursor = 0
q_cursor = 0
for raw_batch_idx, start_seq_position, end_seq_position in cp_index:
batch_idx = int(raw_batch_idx)
pre_chunk_offset = (
forward_batch.seq_lens_cpu[batch_idx].item()
- forward_batch.extend_seq_lens_cpu[batch_idx]
)
start_seq_position += pre_chunk_offset
end_seq_position += pre_chunk_offset
if end_seq_position < start_seq_position:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_cp_index_bad_segment "
f"batch_idx={batch_idx} start={start_seq_position} "
f"end={end_seq_position}"
)
extend_seq_len = int(end_seq_position - start_seq_position)
kv_len_i = int(end_seq_position)
segment_records.append(
(
batch_idx,
int(start_seq_position),
int(end_seq_position),
extend_seq_len,
kv_len_i,
k_cursor,
q_cursor,
int(pre_chunk_offset),
)
)
batch_idx_list.append(batch_idx)
kv_lens_list.append(kv_len_i)
q_starts_list.append(int(start_seq_position))
q_lens_list.append(extend_seq_len)
k_bases_list.append(k_cursor)
q_bases_list.append(q_cursor)
if batch_idx < 0 or batch_idx >= len(request_kv_bases):
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][index_topk] "
"reason=batch_cp_index_bad_batch_idx "
f"batch_idx={batch_idx} seq_lens={forward_batch.seq_lens_cpu.tolist()}"
)
topk_offset_list.extend(
[request_kv_bases[batch_idx]] * extend_seq_len
)
k_cursor += kv_len_i
q_cursor += extend_seq_len
topk_indices_offset_override = torch.tensor(
topk_offset_list, dtype=torch.int32, device=q_fp8.device
plan = _get_or_build_cp_ragged_index_plan(
forward_batch,
cp_index,
q_fp8.device,
current_req_offsets,
)
segment_records = plan.segment_records
topk_indices_offset_override = plan.topk_indices_offset_override
if current_index_kv is None:
assert index_buffer is not None
assert block_tables is not None
descriptor_device = q_fp8.device
tai_batch_prepared = try_tai_prepare_cp_mqa_index_batch(
index_buffer=index_buffer,
block_tables=block_tables,
batch_indices=torch.tensor(
batch_idx_list, dtype=torch.int32, device=descriptor_device
),
kv_lens=torch.tensor(
kv_lens_list, dtype=torch.int32, device=descriptor_device
),
q_starts=torch.tensor(
q_starts_list, dtype=torch.int32, device=descriptor_device
),
q_lens=torch.tensor(
q_lens_list, dtype=torch.int32, device=descriptor_device
),
k_bases=torch.tensor(
k_bases_list, dtype=torch.int32, device=descriptor_device
),
q_bases=torch.tensor(
q_bases_list, dtype=torch.int32, device=descriptor_device
),
total_kv_len=k_cursor,
total_q_count=q_cursor,
max_kv_len=max(kv_lens_list, default=0),
max_q_len=max(q_lens_list, default=0),
batch_indices=plan.batch_indices,
kv_lens=plan.kv_lens,
q_starts=plan.q_starts,
q_lens=plan.q_lens,
k_bases=plan.k_bases,
q_bases=plan.q_bases,
total_kv_len=plan.total_kv_len,
total_q_count=plan.total_q_count,
max_kv_len=plan.max_kv_len,
max_q_len=plan.max_q_len,
page_size=page_size,
index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
)
@@ -1699,9 +1782,6 @@ class Indexer(MultiPlatformOp):
k_fp8_u8, k_scale, ks, ke_offset = tai_batch_prepared
k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
kv_fp8 = (k_fp8, k_scale)
actual_seq_q = torch.tensor(
q_lens_list, dtype=torch.int32, device=q_fp8.device
)
ke = ks + ke_offset
topk_result = self._mqa_logits_topk_ragged_chunked(
metadata,
@@ -1717,10 +1797,7 @@ class Indexer(MultiPlatformOp):
return topk_result
else:
assert current_req_offsets is not None
descriptor_device = q_fp8.device
current_bases_list = [
int(current_req_offsets[batch_idx]) for batch_idx in batch_idx_list
]
assert plan.current_bases is not None
current_index_head_dim = getattr(
forward_batch.token_to_kv_pool,
"index_head_dim",
@@ -1729,37 +1806,22 @@ class Indexer(MultiPlatformOp):
tai_current_prepared = try_tai_prepare_cp_mqa_current_index_batch(
current_index_k=_current_index_k_for_tai(current_index_kv[0]),
current_index_scale=current_index_kv[1],
current_bases=torch.tensor(
current_bases_list, dtype=torch.int32, device=descriptor_device
),
kv_lens=torch.tensor(
kv_lens_list, dtype=torch.int32, device=descriptor_device
),
q_starts=torch.tensor(
q_starts_list, dtype=torch.int32, device=descriptor_device
),
q_lens=torch.tensor(
q_lens_list, dtype=torch.int32, device=descriptor_device
),
k_bases=torch.tensor(
k_bases_list, dtype=torch.int32, device=descriptor_device
),
q_bases=torch.tensor(
q_bases_list, dtype=torch.int32, device=descriptor_device
),
total_kv_len=k_cursor,
total_q_count=q_cursor,
max_kv_len=max(kv_lens_list, default=0),
max_q_len=max(q_lens_list, default=0),
current_bases=plan.current_bases,
kv_lens=plan.kv_lens,
q_starts=plan.q_starts,
q_lens=plan.q_lens,
k_bases=plan.k_bases,
q_bases=plan.q_bases,
total_kv_len=plan.total_kv_len,
total_q_count=plan.total_q_count,
max_kv_len=plan.max_kv_len,
max_q_len=plan.max_q_len,
index_head_dim=current_index_head_dim,
)
if tai_current_prepared is not None:
k_fp8_u8, k_scale, ks, ke_offset = tai_current_prepared
k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
kv_fp8 = (k_fp8, k_scale)
actual_seq_q = torch.tensor(
q_lens_list, dtype=torch.int32, device=q_fp8.device
)
ke = ks + ke_offset
topk_result = self._mqa_logits_topk_ragged_chunked(
metadata,
@@ -2492,8 +2554,15 @@ class Indexer(MultiPlatformOp):
q_len_start = 0
seq_lens_cpu = forward_batch.seq_lens_cpu
for i in range(forward_batch.batch_size):
seq_len = forward_batch.seq_lens[i].item()
# seq_lens is a device tensor; indexing .item() there would cost a
# cudaStreamSynchronize per request per layer.
seq_len = (
int(seq_lens_cpu[i])
if seq_lens_cpu is not None
else int(forward_batch.seq_lens[i].item())
)
q_len = (
forward_batch.extend_seq_lens_cpu[i]
if forward_batch.forward_mode.is_extend()

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

@@ -30,6 +30,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
current_loc_remap_fast_path_args,
filter_owned_logical_locs,
get_cp_shared_kv_token_loc_req_id,
get_or_build_batch_slot_spans,
get_or_build_shared_token_kv_slot_remap,
is_current_only_extend_batch,
is_packed_fp8_mla_kv_cache,
@@ -2092,7 +2093,8 @@ class NativeSparseAttnBackend(
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
)
current_slot_spans = build_batch_current_slot_spans(
_, current_slot_spans = get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=getattr(
forward_batch, "extend_prefix_lens_cpu", None
@@ -2101,6 +2103,7 @@ class NativeSparseAttnBackend(
forward_batch, "extend_seq_lens_cpu", None
),
page_size=page_size,
want_prefix=False,
)
kv_cache, page_table_1 = (
materialize_prefix_and_reuse_current_kv_page_slots(
@@ -2230,21 +2233,29 @@ class NativeSparseAttnBackend(
f"current_locs_shape={tuple(current_locs_for_reuse.shape)} "
f"page_size={page_size}"
)
prefix_slot_spans = None
current_slot_spans = build_batch_current_slot_spans(
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=False,
)
)
else:
prefix_pages = 0
prefix_slot_spans = build_batch_prefix_slot_spans(
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
page_size=page_size,
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=True,
)
)
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
@@ -2606,19 +2617,17 @@ class NativeSparseAttnBackend(
)
if len(prefix_lens_cpu) == 1:
prefix_pages = int(prefix_lens_cpu[0]) // page_size
prefix_slot_spans = None
else:
prefix_pages = 0
prefix_slot_spans = build_batch_prefix_slot_spans(
prefix_slot_spans, current_slot_spans = (
get_or_build_batch_slot_spans(
forward_batch,
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
want_prefix=len(prefix_lens_cpu) > 1,
)
current_slot_spans = build_batch_current_slot_spans(
logical_pages=metadata.real_page_table,
prefix_lens_cpu=prefix_lens_cpu,
extend_lens_cpu=extend_lens_cpu,
page_size=page_size,
)
logical_locs_row_ids = build_flattened_request_row_ids(
metadata.indexer_seq_lens_cpu,

33
python/sglang/srt/mem_cache/memory_pool_host.py
View File

@@ -448,10 +448,18 @@ class HostKVCache(abc.ABC):
) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
if self.layout not in ("page_first_direct", "layer_page_first"):
return None, None
validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
validate_page_aligned_token_indices(
device_indices, self.page_size, "device_indices"
)
# Page alignment is construction-based on the hot path (see
# CacheController._validate_page_aligned_pair): the generic validator
# uses Tensor truth values and would cudaStreamSynchronize per call on
# CUDA tensors, so validate CPU/test tensors only.
if not host_indices.is_cuda:
validate_page_aligned_token_indices(
host_indices, self.page_size, "host_indices"
)
if not device_indices.is_cuda:
validate_page_aligned_token_indices(
device_indices, self.page_size, "device_indices"
)
host_page_indices = (
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
)
@@ -2205,10 +2213,19 @@ class NSATokenToKVPoolHost(MLATokenToKVPoolHost):
def _get_indexer_page_indices(self, host_indices, device_indices):
if host_indices.numel() == 0:
return host_indices, device_indices
validate_page_aligned_token_indices(host_indices, self.page_size, "host_indices")
validate_page_aligned_token_indices(
device_indices, self.page_size, "device_indices"
)
# Same construction-based invariant as _prepare_load_page_indices:
# this runs per layer(-group) on the write-through hot path, and the
# generic validator costs ~0.4 ms of cudaStreamSynchronize per call on
# CUDA tensors (measured: ~12.7 ms/forward). Validate CPU/test
# tensors only.
if not host_indices.is_cuda:
validate_page_aligned_token_indices(
host_indices, self.page_size, "host_indices"
)
if not device_indices.is_cuda:
validate_page_aligned_token_indices(
device_indices, self.page_size, "device_indices"
)
host_page_indices = (
host_indices.reshape(-1, self.page_size)[:, 0] // self.page_size
)

220
test/manual/bench_cpu_gap_fixes.py Normal file
View File

@@ -0,0 +1,220 @@
#!/usr/bin/env python3
"""Micro-benchmark for the pre-attention CPU-gap fixes (task #14).
Scenario mirrors the traced production batch: bs=12, prefix 640..26304 tok,
extends ~3.7K tok, cp8 in-seq-split (2*cp segments/request), page 64.
1. page-aligned validator skip on CUDA tensors
(memory_pool_host._get_indexer_page_indices hot path) — measured with a
busy GPU queue, because torch.any/.equal sync for the whole queue.
2. ragged index descriptor plan: rebuild-per-layer (old) vs per-batch cache.
3. slot-span builders: rebuild-per-layer (old) vs per-batch cache.
Run (single GPU is enough):
PYTHONPATH=python python test/manual/bench_cpu_gap_fixes.py
"""
from __future__ import annotations
import time
from types import SimpleNamespace
import torch
from sglang.srt.layers.attention.nsa import nsa_indexer
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
build_batch_current_slot_spans,
build_batch_prefix_slot_spans,
get_or_build_batch_slot_spans,
)
from sglang.srt.mem_cache.page_index_utils import (
validate_page_aligned_token_indices,
)
DEV = torch.device("cuda", 0)
PAGE = 64
BS = 12
CP = 8
PREFIX_LENS = [19200, 256] + [26304] * 10
EXTEND_LENS = [3776, 7360] + [3347] * 10
REPS = 200
def timed(fn, reps=REPS, warmup=20):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(reps):
fn()
torch.cuda.synchronize()
return (time.perf_counter() - t0) / reps * 1e6 # us
def make_fake_batch():
fb = SimpleNamespace()
fb.seq_lens_cpu = torch.tensor(
[p + e for p, e in zip(PREFIX_LENS, EXTEND_LENS)], dtype=torch.int64
)
fb.extend_seq_lens_cpu = list(EXTEND_LENS)
fb.extend_prefix_lens_cpu = list(PREFIX_LENS)
return fb
def make_cp_index():
# 2*CP zigzag segments per request over the extend, page-aligned-ish.
cp_index = []
for req, extend in enumerate(EXTEND_LENS):
seg = max(PAGE, (extend // (2 * CP)) // PAGE * PAGE)
pos = 0
while pos < extend:
end = min(pos + seg, extend)
cp_index.append((req, pos, end))
pos = end
return cp_index
def bench_validator():
n_pages = 600 # ~one layer-group submit worth of pages
starts = torch.arange(n_pages, device=DEV, dtype=torch.int64) * PAGE
indices = (
starts[:, None] + torch.arange(PAGE, device=DEV, dtype=torch.int64)
).reshape(-1)
# Busy queue: enqueue ~0.5ms of GEMM before each validator call, the way
# the real submit lands behind a layer's compute.
a = torch.randn(2048, 2048, device=DEV, dtype=torch.bfloat16)
b = torch.randn(2048, 2048, device=DEV, dtype=torch.bfloat16)
def old_path():
for _ in range(4):
a @ b
validate_page_aligned_token_indices(indices, PAGE, "bench")
starts2 = indices.reshape(-1, PAGE)[:, 0] // PAGE
return starts2
def new_path():
for _ in range(4):
a @ b
if not indices.is_cuda:
validate_page_aligned_token_indices(indices, PAGE, "bench")
starts2 = indices.reshape(-1, PAGE)[:, 0] // PAGE
return starts2
# measure WALL time per call without trailing torch.cuda.synchronize in
# the loop (the sync inside the validator is exactly what we measure).
def wall(fn, reps=60, warmup=10):
for _ in range(warmup):
fn()
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(reps):
fn()
t1 = time.perf_counter() # NO sync: CPU-side blocking is the metric
torch.cuda.synchronize()
return (t1 - t0) / reps * 1e6
old = wall(old_path)
new = wall(new_path)
print(
f"1. validator (busy queue, {n_pages} pages): old={old:8.1f}us "
f"new={new:8.1f}us speedup x{old/new:.1f}"
)
def bench_ragged_plan():
fb = make_fake_batch()
cp_index = make_cp_index()
def old_build():
# the pre-fix per-layer work: full python loop + 7 tensor H2Ds
return nsa_indexer._build_cp_ragged_index_plan(fb, cp_index, DEV, None)
fb2 = make_fake_batch()
def cached():
return nsa_indexer._get_or_build_cp_ragged_index_plan(
fb2, cp_index, DEV, None
)
old = timed(old_build)
new = timed(cached)
n_tokens = sum(EXTEND_LENS)
print(
f"2. ragged index plan (bs={BS}, {len(cp_index)} segs, {n_tokens} q tok): "
f"per-layer rebuild={old:8.1f}us cached={new:8.1f}us speedup x{old/new:.0f}"
)
def bench_spans():
fb = make_fake_batch()
pages_per_req = max(
(p + e + PAGE - 1) // PAGE for p, e in zip(PREFIX_LENS, EXTEND_LENS)
)
logical_pages = torch.zeros((BS, pages_per_req), dtype=torch.int64)
def old_build():
prefix = build_batch_prefix_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=PREFIX_LENS,
page_size=PAGE,
)
current = build_batch_current_slot_spans(
logical_pages=logical_pages,
prefix_lens_cpu=PREFIX_LENS,
extend_lens_cpu=EXTEND_LENS,
page_size=PAGE,
)
return prefix, current
def cached():
return get_or_build_batch_slot_spans(
fb,
logical_pages=logical_pages,
prefix_lens_cpu=PREFIX_LENS,
extend_lens_cpu=EXTEND_LENS,
page_size=PAGE,
want_prefix=True,
)
old = timed(old_build, reps=2000)
new = timed(cached, reps=2000)
print(
f"3. slot spans (bs={BS}): per-layer rebuild={old:8.1f}us "
f"cached={new:8.1f}us speedup x{old/new:.0f}"
)
def main():
torch.cuda.init()
print(f"device: {torch.cuda.get_device_name(0)}")
bench_validator()
bench_ragged_plan()
bench_spans()
# equality check: cached plan tensors match a fresh build
fb = make_fake_batch()
cp_index = make_cp_index()
p1 = nsa_indexer._build_cp_ragged_index_plan(fb, cp_index, DEV, None)
p2 = nsa_indexer._get_or_build_cp_ragged_index_plan(fb, cp_index, DEV, None)
assert torch.equal(p1.topk_indices_offset_override, p2.topk_indices_offset_override)
assert torch.equal(p1.kv_lens, p2.kv_lens) and torch.equal(p1.q_bases, p2.q_bases)
assert p1.segment_records == p2.segment_records
s1 = build_batch_current_slot_spans(
logical_pages=torch.zeros((BS, 512), dtype=torch.int64),
prefix_lens_cpu=PREFIX_LENS,
extend_lens_cpu=EXTEND_LENS,
page_size=PAGE,
)
_, s2 = get_or_build_batch_slot_spans(
SimpleNamespace(),
logical_pages=torch.zeros((BS, 512), dtype=torch.int64),
prefix_lens_cpu=PREFIX_LENS,
extend_lens_cpu=EXTEND_LENS,
page_size=PAGE,
want_prefix=False,
)
assert s1 == s2
print("EQUALITY CHECKS PASS")
if __name__ == "__main__":
main()