diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md index 263319b3d..29112caec 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md @@ -651,6 +651,15 @@ Phase 0 -> Phase 2 5. **workspace pressure：** bs>1 materialize buffer 更大，必须测显存和临时 buffer，不只看 latency。 6. **collective 顺序：** 所有 CP rank 必须按完全相同的 layer/request 顺序执行 collective。 7. **fallback 可见性：** 开发阶段优先 fail-fast；生产 fallback 必须 warning 且限频。 +8. **W2 allocator CPU descriptor overhead：** 当前 W2 correctness path 仍在 + allocation 热路径中用 Python list/tensor 临时对象构造 batch owner descriptor： + `extend_lens` / `extend_prefix_lens` list、`flat_page_compute_owners` + list、`positions_by_owner` list，以及每个 owner lane 的 `position_tensor`。 + 这仍复用了 owner free/release bucket 本地记账，比旧全量 scan/sort 路径轻， + 但在线上大量 200-2000 token 短 extend 叠 batch 时可能成为 CPU hot path。 + 后续应让 W1 batch plan 直接携带 page-owner descriptor，W2 消费 + tensor/固定 buffer，并用 kernel 或固定 workspace 完成 selected-pages + gather/scatter，避免每次 allocation 重建 Python descriptor。 --- diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index a5111a268..5d9417f7d 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md @@ -253,6 +253,44 @@ out_cache_loc: torch.Tensor # flattened request order - CP shared-KV bs>1 进入 allocator 时不再出现 `multi_batch` fallback。 - W3 direct write 可以信任 `out_cache_loc` 的 owner-lane 合同。 +### 当前实现状态与性能风险 + +PR #10 已把 W2 正确性路径接入： + +- bs>1 会按 request 独立调用 in-seq page owner planner； +- owner list 按 request order flatten； +- `alloc_extend_compute_owner()` 删除 bs=1 hard guard； +- allocator 复用已有 `alloc_extend_naive()` 生成 flattened `out_cache_loc`； +- supported CP shared-KV case 不再走 `multi_batch` legacy fallback。 + +但当前实现仍有一段 Python/control-path descriptor overhead，主要包括： + +1. `alloc_paged_token_slots_extend()` 每次 allocation 从 CPU tensor `.item()` 构造 + `extend_lens` / `extend_prefix_lens` Python lists。 +2. `build_batch_in_seq_page_compute_owners()` 用 Python loop 生成 + `flat_page_compute_owners: List[int]`。 +3. `_select_compute_owner_pages()` 再用 Python loop 构造 + `required_by_owner` / `positions_by_owner`。 +4. 每个 owner lane 会临时构造 `position_tensor`，再把 lane pages scatter 回 + `selected_pages`。 + +这比旧的全量 free-page scan / `torch.isin` / sort-merge 路径轻，因为 allocator +仍使用 `_owner_free_pages` / `_owner_release_pages` 本地 bucket 记账；但它还不是最终 +descriptor 化 fast path。在线上大量 200-2000 token 短 extend 叠 batch 时，allocation +频率高，这段 Python list/tensor descriptor 构造可能重新成为 CPU hot path。 + +后续优化方向： + +```text +W1 batch plan 直接携带 page-owner descriptor + -> W2 allocator 消费 tensor/descriptor 而不是重新 Python 推导 + -> selected_pages gather/scatter 用固定 buffer 或 kernel 完成 + -> 避免每 tick/request 反复创建 Python lists 和临时 position tensors +``` + +这个风险不影响当前 W2 correctness，但需要在 bs>1 ETE/perf 阶段用 micro benchmark +和线上 trace 单独验证。 + --- ## 6. W3：CP split/rebuild + direct write fast path @@ -662,3 +700,48 @@ runtime / kernel 都消费 CPSharedKVBatchPlan descriptors - 设计 variable-length descriptor benchmark，覆盖 bf16/fp8、page_first_direct、random/owner-lane pages。 第一批完成后，再开始 W3/W4/W6/W7 draft 接入。 + +--- + +## 15. 2026-06-03 W5/W4 kernel 接入状态更新 + +### 已接入并验证 + +1. **NSA index top-k 的 batch cp_index 调用已从 per request/segment MQA top-k 收敛为单次调用。** + - SGLang 路径：`python/sglang/srt/layers/attention/nsa/nsa_indexer.py::_get_topk_in_seq_cp_pair_batch`。 + - 语义：多个 request 的 prev/next CP segment 先 compact 成一个 `cp_index` descriptor，再调用一次 `_get_topk_ragged_with_cp`。 + - 已有测试：`test_indexer_in_seq_cp_pair_batch_*` 与 `test_indexer_ragged_cp_index_current_batch_does_not_materialize`。 + +2. **index partial/current compose 的 current-slot fill 已接入 TAI kernel。** + - SGLang 路径：`cp_shared_kv_runtime.fill_current_index_page_slots`。 + - TAI kernel：`tai_kernel.nsa_prefill.cp_shared_kv_materialize.fill_current_index_page_slots`。 + - 输入是 flatten 后的 current rows，可一次处理 bs>1 的 current suffix；不再在 CUDA 上默认走 PyTorch advanced indexing。 + - fallback 仍为显式 warning：`[CP_SHARED_KV_FALLBACK][tai_materialize] reason=index_current_fill_*`。 + - 远端 CUDA smoke：通过 64/page slot compose 的 K/scale 字节布局检查。 + +3. **shared-index batch cp_index 的 K/S + range prepare 已接入 TAI batch descriptor kernel。** + - SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_index_batch`。 + - TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_kv_and_range_batch`。 + - descriptor：`batch_indices, kv_lens, q_starts, q_lens, k_bases, q_bases, total_kv_len, total_q_count, max_kv_len, max_q_len`。 + - 关键约束：`max_kv_len/max_q_len` 由 Python planner 传入，避免在 wrapper 内对 CUDA tensor 做 `.max().item()` 同步。 + - 输出合同：最终仍是按 segment concat 的 dense K buffer / scale buffer；`ks` 为 segment 的 K base；`ke_offset` 为 segment-local causal end offset，调用端计算 `ke = ks + ke_offset`。 + - 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。 + +### 仍需继续处理的 kernel/runtime 缺口 + +1. **current-only cp_index 的 K/S compact copy 已接入 TAI batch kernel。** + - SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_current_index_batch`。 + - TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_current_kv_and_range_batch`。 + - 语义：从 flatten current rows 按 segment `current_base/current_len` 拷到 concat K/S buffer，同时生成 `ks/ke_offset`。 + - 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。 + +2. **HiCache per-layer D2H backup 仍按 reservation/node 提交。** + - `scheduler._prepare_hicache_write_backups_before_forward()` 仍 `for req in batch.reqs: prepare_write_backup_for_req(req)`。 + - `cache_controller.submit_write_cp_layer()` 每个 reservation 在每层调用 host pool backup。 + - 底层 TAI direct transfer kernel 接收 flatten indices，本身可以吃 batched descriptor；但 runtime 还没有把多个 request/reservation 合并成一个 layer-level descriptor。 + +3. **L2->L1 load/prefetch 需要继续确认实际 queue 合并粒度。** + - `cache_controller.load_cp()` / `start_loading()` 已比 backup 更接近 batched queue，但仍需用 ETE/NVTX 确认 layer-level descriptor 是否足够大，是否存在 per-node launch。 + +4. **MLA current/partial-current 和 direct store 已有 flatten kernel 基础，但还需用 bs>1 ETE 验证没有 fallback hot path。** + - 特别关注 FP8/BF16 两种 dtype、draft/EAGLE 路径、prefetch 关闭时 full/current reuse 是否仍启用。 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 4fa56647f..2d3532c21 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 @@ -577,6 +577,26 @@ def _load_tai_index_mqa_prepare_kernel(): return None +@lru_cache(maxsize=1) +def _load_tai_index_mqa_prepare_batch_kernel(): + try: + from tai_kernel.nsa_prefill import prepare_cp_mqa_kv_and_range_batch + + return prepare_cp_mqa_kv_and_range_batch + except Exception: + return None + + +@lru_cache(maxsize=1) +def _load_tai_index_mqa_current_prepare_batch_kernel(): + try: + from tai_kernel.nsa_prefill import prepare_cp_mqa_current_kv_and_range_batch + + return prepare_cp_mqa_current_kv_and_range_batch + except Exception: + return None + + @lru_cache(maxsize=1) def _load_tai_index_mqa_range_kernel(): try: @@ -800,6 +820,214 @@ def try_tai_prepare_cp_mqa_index( return None +def try_tai_prepare_cp_mqa_current_index_batch( + *, + current_index_k: torch.Tensor, + current_index_scale: torch.Tensor, + current_bases: torch.Tensor, + kv_lens: torch.Tensor, + q_starts: torch.Tensor, + q_lens: torch.Tensor, + k_bases: torch.Tensor, + q_bases: torch.Tensor, + total_kv_len: int, + total_q_count: int, + max_kv_len: int, + max_q_len: int, + index_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Try the TAI batched current-index compact + MQA range prepare path.""" + + if not cp_shared_kv_tai_index_mqa_prepare_enabled(): + _log_tai_index_mqa_prepare_fallback( + "current_batch_env_disabled", + "CP shared KV TAI batched current-index prepare fast path is disabled; " + "falling back to Python slice/cat. segments=%s total_kv_len=%s " + "total_q_count=%s index_head_dim=%s", + int(current_bases.numel()), + total_kv_len, + total_q_count, + index_head_dim, + limit=1, + ) + return None + + kernel = _load_tai_index_mqa_current_prepare_batch_kernel() + if kernel is None: + _log_tai_index_mqa_prepare_fallback( + "current_batch_kernel_missing", + "CP shared KV TAI batched current-index prepare kernel is unavailable; " + "falling back to Python slice/cat. segments=%s total_kv_len=%s " + "total_q_count=%s index_head_dim=%s", + int(current_bases.numel()), + total_kv_len, + total_q_count, + index_head_dim, + limit=1, + ) + return None + + if index_head_dim != 128: + _log_tai_index_mqa_prepare_fallback( + "current_batch_unsupported_layout", + "CP shared KV TAI batched current-index prepare supports " + "index_head_dim=128 only; falling back to Python slice/cat. " + "index_head_dim=%s", + index_head_dim, + limit=4, + ) + return None + + try: + return kernel( + _contiguous_for_tai(current_index_k), + _contiguous_for_tai(current_index_scale), + _contiguous_for_tai(current_bases), + _contiguous_for_tai(kv_lens), + _contiguous_for_tai(q_starts), + _contiguous_for_tai(q_lens), + _contiguous_for_tai(k_bases), + _contiguous_for_tai(q_bases), + total_kv_len=int(total_kv_len), + total_q_count=int(total_q_count), + max_kv_len=int(max_kv_len), + max_q_len=int(max_q_len), + index_head_dim=int(index_head_dim), + ) + except Exception as exc: + _log_tai_index_mqa_prepare_fallback( + "current_batch_kernel_failed", + "CP shared KV TAI batched current-index prepare failed; falling back " + "to Python slice/cat. error=%s segments=%s total_kv_len=%s " + "total_q_count=%s", + exc, + int(current_bases.numel()), + total_kv_len, + total_q_count, + ) + return None + + +def try_tai_prepare_cp_mqa_index_batch( + *, + index_buffer: torch.Tensor, + block_tables: torch.Tensor, + batch_indices: torch.Tensor, + kv_lens: torch.Tensor, + q_starts: torch.Tensor, + q_lens: torch.Tensor, + k_bases: torch.Tensor, + q_bases: torch.Tensor, + total_kv_len: int, + total_q_count: int, + max_kv_len: int, + max_q_len: int, + page_size: int, + index_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Try the TAI batched GetK/GetS + MQA range prepare path. + + This is the bs>1 / multi-CP-segment counterpart of + :func:`try_tai_prepare_cp_mqa_index`: all segment descriptors are flattened + and one kernel family prepares the concatenated K/S buffer and ks/ke ranges. + """ + + if not cp_shared_kv_tai_index_mqa_prepare_enabled(): + _log_tai_index_mqa_prepare_fallback( + "batch_env_disabled", + "CP shared KV TAI batched index MQA prepare fast path is disabled; " + "falling back to per-segment GetK/GetS. segments=%s total_kv_len=%s " + "total_q_count=%s page_size=%s index_head_dim=%s", + int(batch_indices.numel()), + total_kv_len, + total_q_count, + page_size, + index_head_dim, + limit=1, + ) + return None + + kernel = _load_tai_index_mqa_prepare_batch_kernel() + if kernel is None: + _log_tai_index_mqa_prepare_fallback( + "batch_kernel_missing", + "CP shared KV TAI batched index MQA prepare kernel is unavailable; " + "falling back to per-segment GetK/GetS. segments=%s total_kv_len=%s " + "total_q_count=%s page_size=%s index_head_dim=%s", + int(batch_indices.numel()), + total_kv_len, + total_q_count, + page_size, + index_head_dim, + limit=1, + ) + return None + + if index_buffer.dtype != torch.uint8: + _log_tai_index_mqa_prepare_fallback( + "batch_unsupported_dtype", + "CP shared KV TAI batched index MQA prepare requires uint8 page " + "buffer; falling back to per-segment GetK/GetS. dtype=%s", + index_buffer.dtype, + limit=4, + ) + return None + if not index_buffer.is_contiguous() or not block_tables.is_contiguous(): + _log_tai_index_mqa_prepare_fallback( + "batch_non_contiguous", + "CP shared KV TAI batched index MQA prepare requires contiguous " + "index buffer and block tables; falling back to per-segment GetK/GetS. " + "index_shape=%s index_stride=%s table_shape=%s table_stride=%s", + tuple(index_buffer.shape), + tuple(index_buffer.stride()), + tuple(block_tables.shape), + tuple(block_tables.stride()), + limit=4, + ) + return None + if index_head_dim != 128 or page_size != 64: + _log_tai_index_mqa_prepare_fallback( + "batch_unsupported_layout", + "CP shared KV TAI batched index MQA prepare supports page_size=64 " + "and index_head_dim=128 only; falling back to per-segment GetK/GetS. " + "page_size=%s index_head_dim=%s", + page_size, + index_head_dim, + limit=4, + ) + return None + + try: + return kernel( + index_buffer, + _contiguous_for_tai(block_tables), + _contiguous_for_tai(batch_indices), + _contiguous_for_tai(kv_lens), + _contiguous_for_tai(q_starts), + _contiguous_for_tai(q_lens), + _contiguous_for_tai(k_bases), + _contiguous_for_tai(q_bases), + total_kv_len=int(total_kv_len), + total_q_count=int(total_q_count), + max_kv_len=int(max_kv_len), + max_q_len=int(max_q_len), + page_size=int(page_size), + index_head_dim=int(index_head_dim), + ) + except Exception as exc: + _log_tai_index_mqa_prepare_fallback( + "batch_kernel_failed", + "CP shared KV TAI batched index MQA prepare failed; falling back to " + "per-segment GetK/GetS. error=%s segments=%s total_kv_len=%s " + "total_q_count=%s", + exc, + int(batch_indices.numel()), + total_kv_len, + total_q_count, + ) + return None + + def try_tai_prepare_cp_mqa_range( *, valid_q_count: int, @@ -1222,6 +1450,64 @@ def fill_current_kv_page_slots_and_remap_locs( return dense_kv_cache, mixed_locs, current_mask +def _try_tai_fill_current_index_page_slots( + *, + dense_page_buffer: torch.Tensor, + current_index_k: torch.Tensor, + current_index_scale: torch.Tensor, + current_locs: torch.Tensor, + page_inverse: torch.Tensor, + page_size: int, + index_head_dim: int, +) -> torch.Tensor | None: + if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("fill_current_index_page_slots") + return None + + kernels = _load_tai_materialize_kernels() + if kernels is None: + return None + fill_kernel = getattr(kernels, "fill_current_index_page_slots", None) + if fill_kernel is None: + _log_tai_materialize_fallback( + "index_fill_current_missing", + "CP shared KV tai index current-slot fill kernel is unavailable; " + "falling back to torch reference. Upgrade tai-kernel to keep this " + "hot path off PyTorch. page_size=%s index_head_dim=%s current_rows=%s " + "dense_pages=%s", + page_size, + index_head_dim, + int(current_locs.numel()), + int(dense_page_buffer.shape[0]), + limit=1, + ) + return None + + try: + return fill_kernel( + _contiguous_for_tai(dense_page_buffer), + _contiguous_for_tai(current_index_k), + _contiguous_for_tai(current_index_scale), + _contiguous_for_tai(current_locs.reshape(-1)), + _contiguous_for_tai(page_inverse), + page_size=int(page_size), + index_head_dim=int(index_head_dim), + ) + except Exception as exc: + _log_tai_materialize_fallback( + "index_fill_current_failed", + "CP shared KV tai index current-slot fill failed; falling back to " + "torch reference. error=%s page_size=%s index_head_dim=%s " + "current_rows=%s dense_pages=%s", + exc, + page_size, + index_head_dim, + int(current_locs.numel()), + int(dense_page_buffer.shape[0]), + ) + return None + + def fill_current_index_page_slots( *, dense_page_buffer: torch.Tensor, @@ -1244,6 +1530,17 @@ def fill_current_index_page_slots( current_rows = int(current_locs.numel()) if current_rows == 0: return dense_page_buffer + tai_result = _try_tai_fill_current_index_page_slots( + dense_page_buffer=dense_page_buffer, + current_index_k=current_index_k, + current_index_scale=current_index_scale, + current_locs=current_locs, + page_inverse=page_inverse, + page_size=page_size, + index_head_dim=index_head_dim, + ) + if tai_result is not None: + return tai_result if dense_page_buffer.is_cuda: _log_tai_materialize_fallback( "index_current_fill_torch_reference_cuda", diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 4f132da32..248a4b791 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -32,7 +32,9 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( should_reuse_current_extend_kv, tensor_debug_checksum, tensor_debug_summary, + try_tai_prepare_cp_mqa_current_index_batch, try_tai_prepare_cp_mqa_index, + try_tai_prepare_cp_mqa_index_batch, try_tai_prepare_cp_mqa_range, ) from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor @@ -1139,8 +1141,15 @@ class Indexer(MultiPlatformOp): actual_seq_q_list = [] batch_idx_list = [] - if current_index_kv is not None and cp_index is not None: - current_index_kv = None + if ( + current_index_kv is not None + and cp_index is not None + and (shared_index_buffer is not None or shared_block_tables is not None) + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_cp_index_ambiguous_current_and_shared_buffer" + ) if current_index_kv is None: if shared_index_buffer is not None or shared_block_tables is not None: if shared_index_buffer is None or shared_block_tables is None: @@ -1177,32 +1186,234 @@ class Indexer(MultiPlatformOp): and forward_batch.extend_seq_lens_cpu is not None ) if cp_index is not None: - # TODO Multi-batch support has accuracy issues - for batch_idx, start_seq_position, end_seq_position in cp_index: + current_req_offsets: Optional[List[int]] = None + if current_index_kv is not None: + current_req_offsets = [] + current_cursor = 0 + for extend_len in forward_batch.extend_seq_lens_cpu: + current_req_offsets.append(current_cursor) + current_cursor += int(extend_len) + + segment_records: List[Tuple[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 = [] + 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 offset == 0 and batch_idx != 0: - offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1] - k_fp8 = index_buf_accessor.GetK.execute( - forward_batch.token_to_kv_pool, - index_buffer, - seq_len=end_seq_position, - page_indices=block_tables[batch_idx], - ) - k_scale = index_buf_accessor.GetS.execute( - forward_batch.token_to_kv_pool, - index_buffer, - seq_len=end_seq_position, - page_indices=block_tables[batch_idx], + 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, + ) ) + 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) + k_cursor += kv_len_i + q_cursor += extend_seq_len + + 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.int64, 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), + page_size=page_size, + index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, + ) + if tai_batch_prepared is not None: + 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 + with self._with_real_sm_count(): + logits = deep_gemm.fp8_mqa_logits( + q_fp8, + kv_fp8, + weights, + ks, + ke, + clean_logits=False, + ) + topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + batch_idx_list=batch_idx_list, + ) + 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 + ] + current_index_head_dim = getattr( + forward_batch.token_to_kv_pool, + "index_head_dim", + int(current_index_kv[0].reshape(current_index_kv[0].shape[0], -1).shape[1]), + ) + tai_current_prepared = try_tai_prepare_cp_mqa_current_index_batch( + current_index_k=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), + 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 + with self._with_real_sm_count(): + logits = deep_gemm.fp8_mqa_logits( + q_fp8, + kv_fp8, + weights, + ks, + ke, + clean_logits=False, + ) + topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q, + ke_offset=ke_offset, + batch_idx_list=batch_idx_list, + ) + return topk_result + + for ( + batch_idx, + start_seq_position, + end_seq_position, + extend_seq_len, + _kv_len_i, + segment_k_base, + _segment_q_base, + ) in segment_records: + if current_index_kv is None: + k_fp8 = index_buf_accessor.GetK.execute( + forward_batch.token_to_kv_pool, + index_buffer, + seq_len=end_seq_position, + page_indices=block_tables[batch_idx], + ) + k_scale = index_buf_accessor.GetS.execute( + forward_batch.token_to_kv_pool, + index_buffer, + seq_len=end_seq_position, + page_indices=block_tables[batch_idx], + ) + else: + if pre_chunk_offset != 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_cp_index_current_with_prefix " + f"batch_idx={batch_idx} prefix={pre_chunk_offset}" + ) + assert current_req_offsets is not None + req_base = current_req_offsets[batch_idx] + req_end = req_base + int(end_seq_position) + if ( + req_end > int(current_index_kv[0].shape[0]) + or req_end > int(current_index_kv[1].shape[0]) + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_cp_index_current_rows_short " + f"batch_idx={batch_idx} req_end={req_end} " + f"k_rows={int(current_index_kv[0].shape[0])} " + f"scale_rows={int(current_index_kv[1].shape[0])}" + ) + k_fp8 = current_index_kv[0][req_base:req_end].contiguous() + k_scale = current_index_kv[1][req_base:req_end].contiguous() - extend_seq_len = end_seq_position - start_seq_position ks = torch.full( - (extend_seq_len,), offset, dtype=torch.int32, device="cuda" + (extend_seq_len,), + segment_k_base, + dtype=torch.int32, + device=q_fp8.device, ) k_fp8_list.append(k_fp8) k_scale_list.append(k_scale) @@ -1211,14 +1422,13 @@ class Indexer(MultiPlatformOp): start_seq_position + 1, end_seq_position + 1, dtype=torch.int32, - device="cuda", + device=q_fp8.device, ) ke_offset_list.append(ke_offset) actual_seq_q = torch.tensor( - [extend_seq_len], dtype=torch.int32, device="cuda" + [extend_seq_len], dtype=torch.int32, device=q_fp8.device ) actual_seq_q_list.append(actual_seq_q) - batch_idx_list.append(batch_idx) k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn) k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1) @@ -1553,68 +1763,83 @@ class Indexer(MultiPlatformOp): ) ) - outputs = [] cursor = 0 + output_cursor = 0 + cp_index: List[Tuple[int, int, int]] = [] + compact_q_chunks = [] + compact_weight_chunks = [] + compact_output_spans: List[Tuple[int, int]] = [] + current_only = ( + current_index_kv_for_topk is not None + and is_current_only_extend_batch(forward_batch) + ) + page_table_1 = None if current_only else metadata.get_page_table_1() - def call_segment( + def collect_segment( *, req_id: int, segment_len: int, kv_len: int, - ) -> torch.Tensor: - nonlocal cursor + ) -> None: + nonlocal cursor, output_cursor segment_len = int(segment_len) kv_len = int(kv_len) q_segment = q_fp8[cursor : cursor + segment_len] weights_segment = weights[cursor : cursor + segment_len] cursor += segment_len if segment_len == 0: - return torch.empty( - (0, self.index_topk), - dtype=torch.int32, - device=q_fp8.device, + return + + seq_len = int(forward_batch.seq_lens_cpu[req_id].item()) + extend_seq_len = int(forward_batch.extend_seq_lens_cpu[req_id]) + cp_kv_end = seq_len - extend_seq_len + kv_len + if current_only: + logical_kv_limit = seq_len + else: + if page_table_1 is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_gt1_missing_page_table_1" + ) + if page_table_1.dim() == 1: + logical_kv_limit = min(seq_len, int(page_table_1.numel())) + else: + logical_kv_limit = min(seq_len, int(page_table_1.shape[1])) + valid_q_count = _compute_contiguous_valid_cp_query_count( + cp_kv_end=cp_kv_end, + actual_seq_q=segment_len, + logical_kv_limit=logical_kv_limit, + ) + if valid_q_count <= 0: + return + + start_abs = cp_kv_end - segment_len + end_abs = start_abs + valid_q_count + pre_chunk_offset = seq_len - extend_seq_len + cp_index.append( + ( + req_id, + start_abs - pre_chunk_offset, + end_abs - pre_chunk_offset, ) - actual_seq_q_tensor = torch.tensor( - [segment_len], - dtype=torch.int32, - device=q_fp8.device, - ) - actual_seq_q_cu_tensor = torch.tensor( - [0, segment_len], - dtype=torch.int32, - device=q_fp8.device, - ) - return self._get_topk_ragged_with_cp( - forward_batch, - layer_id, - q_segment, - weights_segment, - metadata, - kv_len, - segment_len, - current_index_kv=current_index_kv_for_topk, - shared_index_buffer=shared_index_buffer, - shared_block_tables=shared_block_tables, - actual_seq_q_tensor=actual_seq_q_tensor, - actual_seq_q_cu_tensor=actual_seq_q_cu_tensor, - batch_idx=req_id, ) + compact_q_chunks.append(q_segment[:valid_q_count]) + compact_weight_chunks.append(weights_segment[:valid_q_count]) + compact_output_spans.append((output_cursor, valid_q_count)) for req_id in range(batch_size): - outputs.append( - call_segment( - req_id=req_id, - segment_len=request_actual_seq_q_prev[req_id], - kv_len=request_kv_len_prev[req_id], - ) + collect_segment( + req_id=req_id, + segment_len=request_actual_seq_q_prev[req_id], + kv_len=request_kv_len_prev[req_id], ) - outputs.append( - call_segment( - req_id=req_id, - segment_len=request_actual_seq_q_next[req_id], - kv_len=request_kv_len_next[req_id], - ) + output_cursor += int(request_actual_seq_q_prev[req_id]) + collect_segment( + req_id=req_id, + segment_len=request_actual_seq_q_next[req_id], + kv_len=request_kv_len_next[req_id], ) + output_cursor += int(request_actual_seq_q_next[req_id]) if cursor != int(q_fp8.shape[0]) or cursor != int(weights.shape[0]): raise RuntimeError( @@ -1624,9 +1849,46 @@ class Indexer(MultiPlatformOp): f"q_tokens={int(q_fp8.shape[0])} weights_tokens={int(weights.shape[0])}" ) - if not outputs: - return torch.empty((0, self.index_topk), dtype=torch.int32, device=q_fp8.device) - return torch.cat(outputs, dim=0) + result = torch.full( + (output_cursor, self.index_topk), + -1, + dtype=torch.int32, + device=q_fp8.device, + ) + if not compact_q_chunks: + return result + + compact_q = torch.cat(compact_q_chunks, dim=0) + compact_weights = torch.cat(compact_weight_chunks, dim=0) + compact_rows = int(compact_q.shape[0]) + compact_topk = self._get_topk_ragged_with_cp( + forward_batch, + layer_id, + compact_q, + compact_weights, + metadata, + 0, + compact_rows, + cp_index=cp_index, + current_index_kv=current_index_kv_for_topk, + shared_index_buffer=shared_index_buffer, + shared_block_tables=shared_block_tables, + actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, + ) + if int(compact_topk.shape[0]) != compact_rows: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_gt1_compact_topk_rows_mismatch " + f"expected={compact_rows} got={int(compact_topk.shape[0])}" + ) + compact_cursor = 0 + for output_start, valid_q_count in compact_output_spans: + result[output_start : output_start + valid_q_count] = compact_topk[ + compact_cursor : compact_cursor + valid_q_count + ] + compact_cursor += valid_q_count + return result def forward_indexer( self, diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 5f7421770..0dddcb75c 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -1203,6 +1203,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase): def get_page_table_64(self): return logical_pages + def get_page_table_1(self): + return torch.empty((2, 1000), dtype=torch.int32) + + def get_page_table_1(self): + return torch.empty((2, 1000), dtype=torch.int32) + + def get_page_table_1(self): + return torch.empty((2, 1000), dtype=torch.int32) + def fake_materialize(forward_batch, layer_id, logical_page_table): materialize_calls.append((layer_id, logical_page_table)) return materialized_index, dense_pages @@ -1304,6 +1313,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase): def get_page_table_64(self): return logical_pages + def get_page_table_1(self): + return torch.empty((2, 1000), dtype=torch.int32) + def fake_materialize(forward_batch, layer_id, logical_page_table): materialize_calls.append((layer_id, logical_page_table)) return materialized_index, dense_pages @@ -1329,6 +1341,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): "batch_idx": batch_idx, "kv_len": kv_len, "actual_seq_q": actual_seq_q, + "cp_index": cp_index, "q": q_fp8.flatten().tolist(), "weights": weights.flatten().tolist(), "actual_seq_q_tensor": actual_seq_q_tensor, @@ -1338,17 +1351,20 @@ class TestNSAInSeqCPUtils(unittest.TestCase): "current_index_kv": current_index_kv, } ) - return torch.full( - (actual_seq_q, 2), - len(topk_calls), - dtype=torch.int32, - ) + rows = int(q_fp8.shape[0]) + return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk forward_batch = SimpleNamespace( batch_size=2, + forward_mode=SimpleNamespace( + is_extend_without_speculative=lambda: True, + ), + extend_prefix_lens_cpu=[0, 0], + extend_seq_lens_cpu=[1000, 1000], + seq_lens_cpu=torch.tensor([1000, 1000], dtype=torch.int64), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, kv_len_prev=100, @@ -1357,8 +1373,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase): actual_seq_q_next=1, actual_seq_q_prev_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), actual_seq_q_next_cu_tensor=torch.tensor([0, 1], dtype=torch.int32), - request_kv_len_prev=[100, 300], - request_kv_len_next=[200, 400], + request_kv_len_prev=[2, 1], + request_kv_len_next=[3, 4], request_actual_seq_q_prev=[2, 1], request_actual_seq_q_next=[1, 3], ), @@ -1378,32 +1394,26 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0][1], logical_pages) + self.assertEqual(len(topk_calls), 1) + self.assertEqual(topk_calls[0]["batch_idx"], 0) + self.assertEqual(topk_calls[0]["actual_seq_q"], 7) self.assertEqual( - [ - ( - call["batch_idx"], - call["kv_len"], - call["actual_seq_q"], - call["q"], - call["weights"], - call["actual_seq_q_tensor"].tolist(), - call["actual_seq_q_cu_tensor"].tolist(), - ) - for call in topk_calls - ], - [ - (0, 100, 2, [0.0, 1.0], [100.0, 101.0], [2], [0, 2]), - (0, 200, 1, [2.0], [102.0], [1], [0, 1]), - (1, 300, 1, [3.0], [103.0], [1], [0, 1]), - (1, 400, 3, [4.0, 5.0, 6.0], [104.0, 105.0, 106.0], [3], [0, 3]), - ], + topk_calls[0]["cp_index"], + [(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)], ) + self.assertEqual(topk_calls[0]["q"], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + self.assertEqual( + topk_calls[0]["weights"], + [100.0, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0], + ) + self.assertIsNone(topk_calls[0]["actual_seq_q_tensor"]) + self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"]) self.assertTrue(all(call["shared_index_buffer"] is materialized_index for call in topk_calls)) self.assertTrue(all(call["shared_block_tables"] is dense_pages for call in topk_calls)) self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls)) self.assertEqual( result.tolist(), - [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], + [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self): @@ -1413,7 +1423,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): indexer = object.__new__(Indexer) indexer.index_topk = 2 - current_index_kv = (torch.tensor([1]), torch.tensor([2])) + current_index_kv = (torch.arange(7), torch.arange(7)) logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) materialized_index = torch.tensor([11], dtype=torch.int32) dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) @@ -1424,6 +1434,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase): def get_page_table_64(self): return logical_pages + def get_page_table_1(self): + return torch.empty((2, 512), dtype=torch.int32) + def fake_materialize( forward_batch, layer_id, @@ -1458,23 +1471,33 @@ class TestNSAInSeqCPUtils(unittest.TestCase): topk_calls.append( { "batch_idx": batch_idx, + "actual_seq_q": actual_seq_q, + "cp_index": cp_index, + "q_rows": int(q_fp8.shape[0]), "current_index_kv": current_index_kv, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, } ) - return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32) + rows = int(q_fp8.shape[0]) + return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk forward_batch = SimpleNamespace( batch_size=2, + forward_mode=SimpleNamespace( + is_extend_without_speculative=lambda: True, + ), + extend_prefix_lens_cpu=[64, 64], + extend_seq_lens_cpu=[3, 4], + seq_lens_cpu=torch.tensor([67, 68], dtype=torch.int64), nsa_cp_metadata=NSAContextParallelMetadata( batch_size=2, - request_kv_len_prev=[100, 300], - request_kv_len_next=[200, 400], + request_kv_len_prev=[2, 1], + request_kv_len_next=[3, 4], request_actual_seq_q_prev=[2, 1], request_actual_seq_q_next=[1, 3], ), @@ -1492,7 +1515,13 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(len(materialize_calls), 1) self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages) self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv) - self.assertEqual(len(topk_calls), 4) + self.assertEqual(len(topk_calls), 1) + self.assertEqual(topk_calls[0]["actual_seq_q"], 7) + self.assertEqual( + topk_calls[0]["cp_index"], + [(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)], + ) + self.assertEqual(topk_calls[0]["q_rows"], 7) self.assertTrue( all(call["current_index_kv"] is None for call in topk_calls) ) @@ -1502,14 +1531,11 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertTrue( all(call["shared_block_tables"] is dense_pages for call in topk_calls) ) - self.assertEqual([call["batch_idx"] for call in topk_calls], [0, 0, 1, 1]) - self.assertEqual( - [call["actual_seq_q_cu_tensor"].tolist() for call in topk_calls], - [[0, 2], [0, 1], [0, 1], [0, 3]], - ) + self.assertEqual([call["batch_idx"] for call in topk_calls], [0]) + self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"]) self.assertEqual( result.tolist(), - [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], + [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) def test_indexer_in_seq_cp_pair_batch_reuses_current_index_without_materialize(self): @@ -1519,7 +1545,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): indexer = object.__new__(Indexer) indexer.index_topk = 2 - current_index_kv = (torch.tensor([1]), torch.tensor([2])) + current_index_kv = (torch.arange(7), torch.arange(7)) topk_calls = [] class Mode: @@ -1552,13 +1578,17 @@ class TestNSAInSeqCPUtils(unittest.TestCase): topk_calls.append( { "batch_idx": batch_idx, + "actual_seq_q": actual_seq_q, + "cp_index": cp_index, + "q_rows": int(q_fp8.shape[0]), "current_index_kv": current_index_kv, "shared_index_buffer": shared_index_buffer, "shared_block_tables": shared_block_tables, "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, } ) - return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32) + rows = int(q_fp8.shape[0]) + return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2) indexer._maybe_materialize_shared_index_buffer = fake_materialize indexer._get_topk_ragged_with_cp = fake_get_topk @@ -1587,7 +1617,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): current_index_kv=current_index_kv, ) - self.assertEqual(len(topk_calls), 4) + self.assertEqual(len(topk_calls), 1) self.assertTrue( all(call["current_index_kv"] is current_index_kv for call in topk_calls) ) @@ -1597,16 +1627,305 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertTrue( all(call["shared_block_tables"] is None for call in topk_calls) ) - self.assertEqual([call["batch_idx"] for call in topk_calls], [0, 0, 1, 1]) + self.assertEqual(topk_calls[0]["batch_idx"], 0) + self.assertEqual(topk_calls[0]["actual_seq_q"], 7) self.assertEqual( - [call["actual_seq_q_cu_tensor"].tolist() for call in topk_calls], - [[0, 2], [0, 1], [0, 1], [0, 3]], + topk_calls[0]["cp_index"], + [(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], ) + self.assertEqual(topk_calls[0]["q_rows"], 7) + self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"]) self.assertEqual( result.tolist(), - [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], + [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) + def test_indexer_ragged_cp_index_current_batch_does_not_materialize(self): + import contextlib + import torch + from types import SimpleNamespace + from unittest.mock import patch + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + indexer._with_real_sm_count = lambda: contextlib.nullcontext() + + deep_gemm_calls = [] + + def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False): + deep_gemm_calls.append( + { + "q_rows": int(q_fp8.shape[0]), + "kv_rows": int(kv_fp8[0].shape[0]), + "weights_rows": int(weights.shape[0]), + "ks": ks.tolist(), + "ke": ke.tolist(), + } + ) + return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32) + + class Metadata: + def get_page_table_64(self): + raise AssertionError("current cp_index path must not materialize index pages") + + def topk_transform(self, logits, topk, **kwargs): + return ( + torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) + .view(-1, 1) + .repeat(1, topk) + ) + + forward_batch = SimpleNamespace( + token_to_kv_pool=SimpleNamespace(page_size=64), + seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), + extend_seq_lens_cpu=[3, 4], + ) + q_fp8 = torch.empty((7, 1), dtype=torch.float32) + weights = torch.empty((7, 1, 1), dtype=torch.float32) + current_index_kv = ( + torch.arange(7, dtype=torch.uint8).view(7, 1), + torch.arange(7, dtype=torch.float32).view(7, 1), + ) + + with patch.object( + nsa_indexer, + "deep_gemm", + SimpleNamespace(fp8_mqa_logits=fake_logits), + ): + result = Indexer._get_topk_ragged_with_cp( + indexer, + forward_batch, + layer_id=7, + q_fp8=q_fp8, + weights=weights, + metadata=Metadata(), + kv_len=0, + actual_seq_q=7, + cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], + current_index_kv=current_index_kv, + ) + + self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) + self.assertEqual(len(deep_gemm_calls), 1) + self.assertEqual(deep_gemm_calls[0]["q_rows"], 7) + self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7) + self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14) + self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10]) + self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14]) + + def test_indexer_ragged_cp_index_shared_batch_uses_tai_prepare_once(self): + import contextlib + import torch + from types import SimpleNamespace + from unittest.mock import patch + + from sglang.srt.layers.attention.nsa import index_buf_accessor, nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + indexer._with_real_sm_count = lambda: contextlib.nullcontext() + + prepare_calls = [] + deep_gemm_calls = [] + + def fake_prepare(**kwargs): + prepare_calls.append(kwargs) + total_kv_len = int(kwargs["total_kv_len"]) + total_q_count = int(kwargs["total_q_count"]) + return ( + torch.zeros((total_kv_len, 1), dtype=torch.uint8), + torch.zeros((total_kv_len,), dtype=torch.float32), + torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32), + torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32), + ) + + def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False): + deep_gemm_calls.append( + { + "q_rows": int(q_fp8.shape[0]), + "kv_rows": int(kv_fp8[0].shape[0]), + "weights_rows": int(weights.shape[0]), + "ks": ks.tolist(), + "ke": ke.tolist(), + } + ) + return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32) + + class Metadata: + def topk_transform(self, logits, topk, **kwargs): + return ( + torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) + .view(-1, 1) + .repeat(1, topk) + ) + + forward_batch = SimpleNamespace( + token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1), + seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), + extend_seq_lens_cpu=[3, 4], + ) + q_fp8 = torch.empty((7, 1), dtype=torch.float32) + weights = torch.empty((7, 1, 1), dtype=torch.float32) + shared_index_buffer = torch.zeros((8, 264), dtype=torch.uint8) + shared_block_tables = torch.arange(8, dtype=torch.int64).view(2, 4) + + with patch.object( + nsa_indexer, + "try_tai_prepare_cp_mqa_index_batch", + side_effect=fake_prepare, + create=True, + ), patch.object( + index_buf_accessor.GetK, + "execute", + side_effect=AssertionError("batched path must not call per-segment GetK"), + ), patch.object( + index_buf_accessor.GetS, + "execute", + side_effect=AssertionError("batched path must not call per-segment GetS"), + ), patch.object( + nsa_indexer, + "deep_gemm", + SimpleNamespace(fp8_mqa_logits=fake_logits), + ): + result = Indexer._get_topk_ragged_with_cp( + indexer, + forward_batch, + layer_id=7, + q_fp8=q_fp8, + weights=weights, + metadata=Metadata(), + kv_len=0, + actual_seq_q=7, + cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], + shared_index_buffer=shared_index_buffer, + shared_block_tables=shared_block_tables, + ) + + self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) + self.assertEqual(len(prepare_calls), 1) + call = prepare_calls[0] + self.assertIs(call["index_buffer"], shared_index_buffer) + self.assertIs(call["block_tables"], shared_block_tables) + self.assertEqual(call["batch_indices"].tolist(), [0, 0, 1, 1]) + self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4]) + self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1]) + self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3]) + self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10]) + self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4]) + self.assertEqual(call["total_kv_len"], 14) + self.assertEqual(call["total_q_count"], 7) + self.assertEqual(call["max_kv_len"], 4) + self.assertEqual(call["max_q_len"], 3) + self.assertEqual(len(deep_gemm_calls), 1) + self.assertEqual(deep_gemm_calls[0]["q_rows"], 7) + self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7) + self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14) + self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10]) + self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14]) + + def test_indexer_ragged_cp_index_current_batch_uses_tai_compact_once(self): + import contextlib + import torch + from types import SimpleNamespace + from unittest.mock import patch + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + indexer._with_real_sm_count = lambda: contextlib.nullcontext() + + prepare_calls = [] + deep_gemm_calls = [] + + def fake_prepare(**kwargs): + prepare_calls.append(kwargs) + total_kv_len = int(kwargs["total_kv_len"]) + return ( + torch.zeros((total_kv_len, 1), dtype=torch.uint8), + torch.zeros((total_kv_len,), dtype=torch.float32), + torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32), + torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32), + ) + + def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False): + deep_gemm_calls.append( + { + "q_rows": int(q_fp8.shape[0]), + "kv_rows": int(kv_fp8[0].shape[0]), + "ks": ks.tolist(), + "ke": ke.tolist(), + } + ) + return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32) + + class Metadata: + def get_page_table_64(self): + raise AssertionError("current cp_index path must not materialize index pages") + + def topk_transform(self, logits, topk, **kwargs): + return ( + torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32) + .view(-1, 1) + .repeat(1, topk) + ) + + forward_batch = SimpleNamespace( + token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1), + seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), + extend_seq_lens_cpu=[3, 4], + ) + current_index_kv = ( + torch.arange(7, dtype=torch.uint8).view(7, 1), + torch.arange(7, dtype=torch.float32).view(7, 1), + ) + + with patch.object( + nsa_indexer, + "try_tai_prepare_cp_mqa_current_index_batch", + side_effect=fake_prepare, + create=True, + ), patch.object( + nsa_indexer, + "deep_gemm", + SimpleNamespace(fp8_mqa_logits=fake_logits), + ): + result = Indexer._get_topk_ragged_with_cp( + indexer, + forward_batch, + layer_id=7, + q_fp8=torch.empty((7, 1), dtype=torch.float32), + weights=torch.empty((7, 1, 1), dtype=torch.float32), + metadata=Metadata(), + kv_len=0, + actual_seq_q=7, + cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)], + current_index_kv=current_index_kv, + ) + + self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]]) + self.assertEqual(len(prepare_calls), 1) + call = prepare_calls[0] + self.assertIs(call["current_index_k"], current_index_kv[0]) + self.assertIs(call["current_index_scale"], current_index_kv[1]) + self.assertEqual(call["current_bases"].tolist(), [0, 0, 3, 3]) + self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4]) + self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1]) + self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3]) + self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10]) + self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4]) + self.assertEqual(call["total_kv_len"], 14) + self.assertEqual(call["total_q_count"], 7) + self.assertEqual(len(deep_gemm_calls), 1) + self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14) + self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10]) + self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14]) + def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self): import torch diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 5b5f02a6a..4a4f60f18 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -1191,6 +1191,55 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]]) self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]]) + def test_fill_current_index_page_slots_uses_tai_kernel_when_available(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + class FakeKernels: + calls = [] + + @staticmethod + def fill_current_index_page_slots(*args, **kwargs): + FakeKernels.calls.append((args, kwargs)) + dense_page_buffer = args[0] + dense_page_buffer[2, 0] = 99 + return dense_page_buffer + + dense_page_buffer = torch.zeros((3, 32), dtype=torch.uint8) + current_k = torch.ones((2, 4), dtype=torch.uint8) + current_scale = torch.ones((2, 1), dtype=torch.float32) + current_locs = torch.tensor([8, 9], dtype=torch.int64) + page_inverse = torch.tensor([0, -1, 2], dtype=torch.int64) + + with patch.object( + runtime, + "_tai_materialize_runtime_enabled", + return_value=True, + ), patch.object( + runtime, + "_load_tai_materialize_kernels", + return_value=FakeKernels, + ): + result = runtime.fill_current_index_page_slots( + dense_page_buffer=dense_page_buffer, + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=current_locs, + page_inverse=page_inverse, + page_size=4, + index_head_dim=4, + ) + + self.assertIs(result, dense_page_buffer) + self.assertEqual(int(result[2, 0]), 99) + self.assertEqual(len(FakeKernels.calls), 1) + args, kwargs = FakeKernels.calls[0] + self.assertIs(args[0], dense_page_buffer) + self.assertTrue(torch.equal(args[1], current_k)) + self.assertTrue(torch.equal(args[2], current_scale)) + self.assertTrue(torch.equal(args[3], current_locs)) + self.assertTrue(torch.equal(args[4], page_inverse)) + self.assertEqual(kwargs, {"page_size": 4, "index_head_dim": 4}) + def test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails(self): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime