diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 66856e3d8..58b0af1bc 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index accd22966..835d8d544 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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() diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 433d6d41d..1c1128e76 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index ba1b4ee30..91339f571 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -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 ) diff --git a/test/manual/bench_cpu_gap_fixes.py b/test/manual/bench_cpu_gap_fixes.py new file mode 100644 index 000000000..6337dffd8 --- /dev/null +++ b/test/manual/bench_cpu_gap_fixes.py @@ -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()