diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md index 30521c3fb..49cc421cd 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md @@ -84,17 +84,18 @@ Implication: Phase 2 must add batch-aware owner plans and allocation. Legacy all Implication: Phase 3 must split `out_cache_loc` per request, then concatenate this rank's local locs in the same local compute order. -### C7. Current/partial-current reuse is single-request +### C7. Current/partial-current reuse has been extended to target bs>1 - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:1764-1904` - - `_current_extend_kv_reuse_miss_reason()` returns `batch_size_not_one`. - - `current_extend_kv_rows_for_reuse()` requires one extend length. + - `_current_extend_kv_reuse_miss_reason()` no longer rejects `batch_size > 1`; it validates batch size, prefix/extend/seq lengths, and `out_cache_loc`. + - `current_extend_kv_rows_for_reuse()` returns flattened valid current rows from `sum(extend_seq_lens_cpu)`, so padded tails are not exposed to attention. - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:309-380` - - Index partial-current compose requires one positive page-aligned prefix. + - Index partial-current compose uses batch prefix slot spans for bs>1. + - Index top-k current-only bs>1 avoids materialize; partial-current bs>1 materializes once and reuses the dense buffer for all request segments. - `python/sglang/srt/layers/attention/nsa_backend.py:1940-2020` - - MLA partial-current compose has the same one-prefix fail-fast contract. + - MLA partial-current compose uses batch prefix slot spans for bs>1. -Implication: Phase 4/5 must first keep sync full materialize correct, then add batched current/partial-current reuse. +Implication: target MLA/index sync-correct current-only and partial-current reuse are covered. Remaining work is kernel-backed variable-length descriptors, L1 prefetch, HiCache load/backup, and draft/EAGLE. ### C8. L1 prefix prefetch is single-request 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 8dc58a9ce..263319b3d 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 @@ -187,19 +187,20 @@ target sync correctness 结论:Phase 3 必须以 request 为单位 split `out_cache_loc`,再拼接本 rank local loc。 -### C7. current / partial-current reuse 是单请求结构 +### C7. current / partial-current reuse 已从单请求结构扩展到 target bs>1 相关代码: - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py:1764-1904` - - `_current_extend_kv_reuse_miss_reason()` 返回 `batch_size_not_one`; - - `current_extend_kv_rows_for_reuse()` 只接受一个 extend length。 + - `_current_extend_kv_reuse_miss_reason()` 已移除 `batch_size_not_one`,改成 batch size、prefix/extend/seq len、`out_cache_loc` 的一致性校验; + - `current_extend_kv_rows_for_reuse()` 按 `sum(extend_seq_lens_cpu)` 返回 flattened valid current rows,避免 padded tail 暴露给 attention。 - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py:309-380` - - index partial-current compose 要求单个 positive page-aligned prefix。 + - index partial-current compose 对 bs>1 使用 batch prefix slot span; + - index top-k current-only bs>1 不 materialize,partial-current bs>1 只 materialize 一次并复用 dense buffer。 - `python/sglang/srt/layers/attention/nsa_backend.py:1940-2020` - - MLA partial-current compose 同样是单 prefix 合同。 + - MLA partial-current compose 对 bs>1 使用 batch prefix slot span。 -结论:Phase 5 需要做 batched current suffix slicing 和 remap。W4-1 可先只做 target index/top-k sync correctness。 +结论:target MLA/index 的同步正确路径已支持 bs>1 current-only 与 partial-current。剩余工作是把这套合同下沉到 variable-length TAI/SGL kernel、L1 prefetch、HiCache load/back up、draft/EAGLE。 ### C8. L1 shared-KV prefetch 是单请求结构 diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md index fa1b6d5c0..bcd549faa 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md @@ -3,9 +3,9 @@ > 日期:2026-06-03 > 分支:`cjy-cp-refactor` > 当前基线:`e4cf8d18b` -> 范围:W3 `local out_cache_loc + direct write`,W4-1 `target index/top-k sync correctness`。 +> 范围:W3 `local out_cache_loc + direct write`,W4-1 `target index/top-k sync correctness`,W4-2 `target current/partial-current reuse`。 -> **命名说明:** 本文的 **W4-1** 指顺序实现计划中的 **target index/top-k sync correctness**。并行派工文档/原定 W4 是 **current/partial-current reuse**;那部分在本文中仍视为后续阶段,不在本轮实现范围内。 +> **命名说明:** 本文的 **W4-1** 指顺序实现计划中的 **target index/top-k sync correctness**。并行派工文档/原定 W4 是 **current/partial-current reuse**;本轮已继续补齐 target MLA/index 的 bs>1 current-only 与 partial-current 同步正确性,记为 **W4-2**。 ## 0. 目标和非目标 @@ -14,12 +14,13 @@ 1. bs>1 的 CP shared-KV direct write 使用每个 request 独立的 page-aligned split,再按 request order 拼接本 rank local rows。 2. MLA KV direct write 和 index KV direct write 共享同一个 ForwardBatch 级 page/local-loc plan。 3. target index/top-k sync path 使用 per-request metadata,不能把 batch flatten 成一条长序列。 -4. W4-1 第一版只做 sync correctness,不启用 bs>1 current reuse、partial-current reuse、L1 prefetch、draft/EAGLE。 +4. W4-1 第一版只做 sync correctness;W4-2 打开 target MLA/index 的 bs>1 current reuse、partial-current reuse。 +5. L1 prefetch、draft/EAGLE、kernel-backed batched descriptor 仍不在本文实现范围内。 非目标: - 不实现 W2 owner-lane allocator;W3 单测使用 synthetic owner-valid `out_cache_loc`。 -- 不实现 bs>1 current/partial-current reuse。 +- 不实现 L1 prefetch、draft/EAGLE 的 bs>1 current/partial-current reuse。 - 不实现 batched MLA/index prefetch。 - 不实现 draft/EAGLE bs>1。 - 不新增 collective 来同步 batch plan。 @@ -218,26 +219,34 @@ for req_id: 后续性能优化可以再把多个 segment 合并成 batched top-k descriptor。 -### C10. 原定 W4 current/partial-current reuse 仍是单请求合同 +### C10. W4-2 已将 target current/partial-current reuse 扩展到 bs>1 `_maybe_materialize_shared_index_buffer()` 在 `nsa_indexer.py:309-517`。 -当 `current_index_kv is not None` 时,它要求: - -```python -len(prefix_lens_cpu) == 1 -positive page-aligned prefix -``` - -见 `nsa_indexer.py:341-355`。 - -因此 W4-1 第一版遇到 bs>1 + `current_index_kv is not None` 必须 fail-fast: +已修正的合同: ```text -[CP_SHARED_KV_FAIL_FAST][batch_gt1_index_current_reuse_unsupported] +current-only bs>1: + 不 materialize prefix,直接把 current_index_kv 传给每个 request segment 的 top-k。 + +partial-current bs>1: + 先用 batch prefix slot span 一次 materialize page-aligned prefix, + 再 splice flattened valid current rows。 ``` -不要在 W4-1 里临时拼 current reuse,否则会和 W5 的 partial/current reuse 工作混在一起。 +实现点: + +- `cp_shared_kv_runtime.current_extend_kv_rows_for_reuse()` 不再拒绝 `batch_size > 1`,按 `sum(extend_seq_lens_cpu)` 返回 valid current rows。 +- `cp_shared_kv_runtime.build_batch_prefix_slot_span()` 用 page-table row-major slot span 表达多个 request 的 prefix pages。 +- `materialize_prefix_and_reuse_current_kv_page_slots()` / `materialize_prefix_and_reuse_current_index_page_slots()` 支持 `prefix_slot_span`。 +- `nsa_backend.py` 的 MLA partial-current sync 对 bs>1 使用 batch prefix slot span。 +- `nsa_indexer.py` 的 index partial-current sync 同样使用 batch prefix slot span。 +- `_get_topk_in_seq_cp_pair_batch()` 不再对 `current_index_kv` fail-fast:current-only 直接透传,partial-current 先 materialize 后所有 segment 复用同一个 dense index buffer。 + +仍然不支持/未验证: + +- 这只是同步正确路径,不代表底层 TAI/SGL variable-length batched kernel 已完成。 +- L1 prefetch、draft/EAGLE 的 bs>1 current/partial-current reuse 尚未接入。 ### C11. `_build_batch_metadata_from_plan()` 的 batch cu tensor 不能直接传给单 segment top-k @@ -420,7 +429,7 @@ W4-1 依赖: - W1 batch metadata; - W3 local q/weights 已按 request boundary split; - `metadata.get_page_table_64()` 可返回 batch page table; -- 不启用 bs>1 current reuse / partial-current reuse。 +- W4-2 后允许 target bs>1 current reuse / partial-current reuse。 ### 4.2 输出合同 @@ -449,19 +458,26 @@ if getattr(metadata, "batch_size", 1) > 1: 保留原 scalar path。 -#### W4-1-S2:新增 `_get_topk_in_seq_cp_pair_batch()` +#### W4-1-S2/W4-2-S1:新增并扩展 `_get_topk_in_seq_cp_pair_batch()` 伪代码: ```python plan = get_cp_shared_kv_batch_plan(forward_batch) -if current_index_kv is not None: - fail-fast batch_gt1_index_current_reuse_unsupported - -shared_block_tables = metadata.get_page_table_64() -shared_index_buffer, shared_block_tables = self._maybe_materialize_shared_index_buffer( - forward_batch, layer_id, shared_block_tables -) +shared_index_buffer = None +shared_block_tables = None +current_index_kv_for_topk = current_index_kv +if current_index_kv is not None and not is_current_only_extend_batch(forward_batch): + current_index_kv_for_topk = None + shared_block_tables = metadata.get_page_table_64() + shared_index_buffer, shared_block_tables = self._maybe_materialize_shared_index_buffer( + forward_batch, layer_id, shared_block_tables, current_index_kv=current_index_kv + ) +elif current_index_kv is None: + shared_block_tables = metadata.get_page_table_64() + shared_index_buffer, shared_block_tables = self._maybe_materialize_shared_index_buffer( + forward_batch, layer_id, shared_block_tables + ) cursor = 0 outputs = [] @@ -551,12 +567,24 @@ mock `_maybe_materialize_shared_index_buffer()`,断言 bs>1 top-k 只调用一 构造某个 request 的 `actual_seq_q_next=0`,断言 `_get_topk_ragged_with_cp()` 不被调用,输出中该 segment 是 empty。 -#### T4:bs>1 current_index_kv fail-fast +#### T4:bs>1 partial-current index reuse 只 materialize 一次 传入 `current_index_kv`,断言: ```text -[CP_SHARED_KV_FAIL_FAST][batch_gt1_index_current_reuse_unsupported] +_maybe_materialize_shared_index_buffer(..., current_index_kv=...) 调用一次; +所有 request segment 的 top-k 复用同一个 shared_index_buffer/shared_block_tables; +segment 调用里 current_index_kv=None。 +``` + +#### T5:bs>1 current-only index reuse 不 materialize + +构造 `extend_prefix_lens_cpu=[0, 0]` 且 `seq_lens_cpu == extend_seq_lens_cpu`,传入 `current_index_kv`,断言: + +```text +不读取 page table; +不调用 _maybe_materialize_shared_index_buffer; +每个 request segment 直接收到 current_index_kv。 ``` #### T5:scalar tests 保持 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 881b1b283..4fa56647f 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 @@ -1783,11 +1783,8 @@ def _current_extend_kv_reuse_miss_reason(forward_batch) -> str | None: if not forward_mode.is_extend_without_speculative(): return "not_extend_without_speculative" - batch_size = int(getattr(forward_batch, "batch_size", 0)) - if batch_size != 1: - return "batch_size_not_one" - extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) out_cache_loc = getattr(forward_batch, "out_cache_loc", None) if extend_seq_lens_cpu is None: @@ -1796,22 +1793,42 @@ def _current_extend_kv_reuse_miss_reason(forward_batch) -> str | None: return "missing_seq_lens_cpu" if out_cache_loc is None: return "missing_out_cache_loc" - if len(extend_seq_lens_cpu) != 1: - return "extend_batch_not_one" - if int(seq_lens_cpu.numel()) != 1: - return "seq_lens_batch_not_one" - extend_len = int(extend_seq_lens_cpu[0]) - seq_len = int(seq_lens_cpu[0].item()) - if extend_len <= 0: - return "non_positive_extend_len" - if seq_len < extend_len: - return "seq_len_smaller_than_extend_len" + batch_size = int(getattr(forward_batch, "batch_size", 0) or 0) + extend_batch_size = len(extend_seq_lens_cpu) + seq_lens_batch_size = int(seq_lens_cpu.numel()) + if batch_size <= 0: + batch_size = extend_batch_size + if extend_batch_size != batch_size: + return "extend_batch_size_mismatch" + if seq_lens_batch_size != batch_size: + return "seq_lens_batch_size_mismatch" + if extend_prefix_lens_cpu is not None and len(extend_prefix_lens_cpu) != batch_size: + return "prefix_batch_size_mismatch" + + valid_current_rows = 0 + for req_id, (extend_len_raw, seq_len_raw) in enumerate( + zip(extend_seq_lens_cpu, seq_lens_cpu) + ): + extend_len = int(extend_len_raw) + seq_len = int(seq_len_raw.item()) + if extend_len <= 0: + return f"non_positive_extend_len_req_{req_id}" + if seq_len < extend_len: + return f"seq_len_smaller_than_extend_len_req_{req_id}" + if extend_prefix_lens_cpu is not None: + prefix_len = int(extend_prefix_lens_cpu[req_id]) + if prefix_len < 0: + return f"negative_prefix_len_req_{req_id}" + if prefix_len + extend_len != seq_len: + return f"prefix_extend_seq_len_mismatch_req_{req_id}" + valid_current_rows += extend_len + # ForwardBatch pads tensors such as out_cache_loc at the tail for CUDA graph # and CP alignment. The first extend_len rows still cover the valid current # suffix, so padded batches remain eligible for current reuse as long as the # valid suffix is fully present. - if int(out_cache_loc.numel()) < extend_len: + if int(out_cache_loc.numel()) < valid_current_rows: return "out_cache_loc_shorter_than_extend" return None @@ -1887,9 +1904,9 @@ def current_extend_kv_rows_for_reuse( return None extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) - if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) != 1: + if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) == 0: return None - valid_current_rows = int(extend_seq_lens_cpu[0]) + valid_current_rows = sum(int(x) for x in extend_seq_lens_cpu) if valid_current_rows <= 0: return None @@ -1903,6 +1920,84 @@ def current_extend_kv_rows_for_reuse( return valid_current_rows +def build_batch_prefix_slot_span( + *, + logical_pages: torch.Tensor, + prefix_lens_cpu, + page_size: int, +) -> tuple[int, int]: + """Return the flattened page-table slot span covering batched prefix pages. + + CP shared KV page-table slot layout is row-major by request. For bs>1 + partial-current reuse, each request has its own prefix length, so a scalar + ``prefix_pages`` is not sufficient. This helper returns the smallest + contiguous flattened slot span that covers all request prefix page ranges. + + The span can include current/suffix slots between two request prefixes. Those + slots are overwritten by current rows and masked before consumption; using a + single span keeps the synchronous fallback to one range collective instead of + one collective per request. + """ + + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + if prefix_lens_cpu is None: + raise ValueError("prefix_lens_cpu is required for batch prefix slot span") + batch_size = len(prefix_lens_cpu) + if batch_size == 0: + return (0, 0) + + if logical_pages.dim() == 1: + if batch_size != 1: + raise ValueError( + "1D logical_pages can only describe one request for prefix slot span: " + f"batch_size={batch_size} logical_pages_shape={tuple(logical_pages.shape)}" + ) + pages_per_request = int(logical_pages.numel()) + else: + if int(logical_pages.shape[0]) < batch_size: + raise ValueError( + "logical_pages has fewer rows than prefix_lens_cpu: " + f"rows={int(logical_pages.shape[0])} batch_size={batch_size}" + ) + pages_per_request = int(logical_pages.reshape(logical_pages.shape[0], -1).shape[1]) + + if pages_per_request < 0: + raise ValueError(f"pages_per_request must be non-negative: {pages_per_request}") + + start_slot: int | None = None + end_slot: int | None = None + for req_id, prefix_len_raw in enumerate(prefix_lens_cpu): + prefix_len = int(prefix_len_raw) + if prefix_len < 0: + raise ValueError( + f"prefix_lens_cpu contains negative prefix length: req_id={req_id} " + f"prefix_len={prefix_len}" + ) + if prefix_len % page_size != 0: + raise ValueError( + "CP shared KV batch prefix slot span requires page-aligned prefixes: " + f"req_id={req_id} prefix_len={prefix_len} page_size={page_size}" + ) + prefix_pages = prefix_len // page_size + if prefix_pages == 0: + continue + if prefix_pages > pages_per_request: + raise ValueError( + "prefix pages exceed per-request logical page-table width: " + f"req_id={req_id} prefix_pages={prefix_pages} " + f"pages_per_request={pages_per_request}" + ) + req_start = req_id * pages_per_request + req_end = req_start + prefix_pages + start_slot = req_start if start_slot is None else min(start_slot, req_start) + end_slot = req_end if end_slot is None else max(end_slot, req_end) + + if start_slot is None or end_slot is None: + return (0, 0) + return (start_slot, end_slot) + + def current_loc_remap_fast_path_args( forward_batch, ) -> tuple[int | None, int | None]: @@ -2892,6 +2987,7 @@ def materialize_prefix_and_reuse_current_kv_page_slots( layout: CpSharedKVLayout, page_size: int, prefix_pages: int, + prefix_slot_span: tuple[int, int] | None = None, layer_id: int | None = None, nvtx_source: str = "mla.partial_current_sync", ) -> tuple[torch.Tensor, torch.Tensor]: @@ -2905,11 +3001,28 @@ def materialize_prefix_and_reuse_current_kv_page_slots( """ total_slots = int(slot_remap.slot_logical_pages.numel()) - if prefix_pages < 0 or prefix_pages > total_slots: - raise ValueError( - "Invalid CP shared KV partial-current prefix range: " - f"prefix_pages={prefix_pages} total_slots={total_slots}" + if prefix_slot_span is None: + if prefix_pages < 0 or prefix_pages > total_slots: + raise ValueError( + "Invalid CP shared KV partial-current prefix range: " + f"prefix_pages={prefix_pages} total_slots={total_slots}" + ) + prefix_start_slot = 0 + prefix_end_slot = int(prefix_pages) + else: + prefix_start_slot, prefix_end_slot = ( + int(prefix_slot_span[0]), + int(prefix_slot_span[1]), ) + if ( + prefix_start_slot < 0 + or prefix_end_slot < prefix_start_slot + or prefix_end_slot > total_slots + ): + raise ValueError( + "Invalid CP shared KV partial-current prefix slot span: " + f"prefix_slot_span={prefix_slot_span} total_slots={total_slots}" + ) dense_kv_cache = kv_cache.new_zeros( (slot_remap.dense_num_pages * page_size, *kv_cache.shape[1:]) @@ -2920,8 +3033,8 @@ def materialize_prefix_and_reuse_current_kv_page_slots( slot_logical_pages=slot_remap.slot_logical_pages, layout=layout, page_size=page_size, - start_slot=0, - end_slot=prefix_pages, + start_slot=prefix_start_slot, + end_slot=prefix_end_slot, ) if not materialized_by_ipc: materialize_local_token_kv_page_slots_into( @@ -2930,11 +3043,15 @@ def materialize_prefix_and_reuse_current_kv_page_slots( slot_logical_pages=slot_remap.slot_logical_pages, layout=layout, page_size=page_size, - start_slot=0, - end_slot=prefix_pages, + start_slot=prefix_start_slot, + end_slot=prefix_end_slot, ) - prefix_rows = slot_range_to_token_slice(page_size, 0, prefix_pages) + prefix_rows = slot_range_to_token_slice( + page_size, + prefix_start_slot, + prefix_end_slot, + ) _all_reduce_materialized_buffer_range( dense_kv_cache, layout.cp_size, @@ -2979,17 +3096,35 @@ def materialize_prefix_and_reuse_current_index_page_slots( page_size: int, index_head_dim: int, prefix_pages: int, + prefix_slot_span: tuple[int, int] | None = None, layer_id: int | None = None, nvtx_source: str = "index.partial_current_sync", ) -> tuple[torch.Tensor, torch.Tensor]: """Synchronously compose prefix index materialization with current index rows.""" total_slots = int(slot_remap.slot_logical_pages.numel()) - if prefix_pages < 0 or prefix_pages > total_slots: - raise ValueError( - "Invalid CP shared KV index partial-current prefix range: " - f"prefix_pages={prefix_pages} total_slots={total_slots}" + if prefix_slot_span is None: + if prefix_pages < 0 or prefix_pages > total_slots: + raise ValueError( + "Invalid CP shared KV index partial-current prefix range: " + f"prefix_pages={prefix_pages} total_slots={total_slots}" + ) + prefix_start_slot = 0 + prefix_end_slot = int(prefix_pages) + else: + prefix_start_slot, prefix_end_slot = ( + int(prefix_slot_span[0]), + int(prefix_slot_span[1]), ) + if ( + prefix_start_slot < 0 + or prefix_end_slot < prefix_start_slot + or prefix_end_slot > total_slots + ): + raise ValueError( + "Invalid CP shared KV index partial-current prefix slot span: " + f"prefix_slot_span={prefix_slot_span} total_slots={total_slots}" + ) dense_page_buffer = page_buffer.new_zeros( (slot_remap.dense_num_pages, *page_buffer.shape[1:]) @@ -2999,8 +3134,8 @@ def materialize_prefix_and_reuse_current_index_page_slots( dense_page_buffer=dense_page_buffer, slot_logical_pages=slot_remap.slot_logical_pages, layout=layout, - start_slot=0, - end_slot=prefix_pages, + start_slot=prefix_start_slot, + end_slot=prefix_end_slot, ) if not materialized_by_ipc: materialize_local_paged_buffer_page_slots_into( @@ -3008,10 +3143,10 @@ def materialize_prefix_and_reuse_current_index_page_slots( dense_page_buffer=dense_page_buffer, slot_logical_pages=slot_remap.slot_logical_pages, layout=layout, - start_slot=0, - end_slot=prefix_pages, + start_slot=prefix_start_slot, + end_slot=prefix_end_slot, ) - prefix_rows = slot_range_to_page_slice(0, prefix_pages) + prefix_rows = slot_range_to_page_slice(prefix_start_slot, prefix_end_slot) _all_reduce_materialized_buffer_range( dense_page_buffer, layout.cp_size, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index acc3b4586..4f132da32 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -15,12 +15,14 @@ from sglang.jit_kernel.fused_store_index_cache import ( from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import index_buf_accessor from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + build_batch_prefix_slot_span, cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, + current_extend_kv_rows_for_reuse, filter_owned_logical_locs, get_or_build_shared_paged_buffer_slot_remap, is_current_only_extend_batch, @@ -338,36 +340,55 @@ class Indexer(MultiPlatformOp): if extend_lens_cpu is not None else None ) - if ( - prefix_lens_cpu is None - or len(prefix_lens_cpu) != 1 - or int(prefix_lens_cpu[0]) <= 0 - or int(prefix_lens_cpu[0]) % page_size != 0 - ): + prefix_lens_valid = prefix_lens_cpu is not None and len(prefix_lens_cpu) > 0 + if prefix_lens_valid: + prefix_lens_valid = all( + int(prefix_len) >= 0 and int(prefix_len) % page_size == 0 + for prefix_len in prefix_lens_cpu + ) and any(int(prefix_len) > 0 for prefix_len in prefix_lens_cpu) + if not prefix_lens_valid: raise RuntimeError( "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " - "CP shared KV index partial-current compose requires one " - "positive page-aligned prefix. " + "CP shared KV index partial-current compose requires " + "positive page-aligned prefix pages. " f"cp_rank={layout.cp_rank} layer_id={layer_id} " f"prefix_lens={prefix_lens} extend_lens={extend_lens} " f"logical_page_table_shape={tuple(logical_page_table.shape)} " f"page_size={page_size}" ) current_locs = forward_batch.out_cache_loc - if extend_lens_cpu is not None and len(extend_lens_cpu) == 1: - valid_current_rows = int(extend_lens_cpu[0]) - if ( - valid_current_rows > 0 - and valid_current_rows < int(current_locs.numel()) - and valid_current_rows <= int(current_index_kv[0].shape[0]) - and valid_current_rows <= int(current_index_kv[1].shape[0]) - ): - current_locs = current_locs[:valid_current_rows] - current_index_kv = ( - current_index_kv[0][:valid_current_rows], - current_index_kv[1][:valid_current_rows], - ) - prefix_pages = int(prefix_lens_cpu[0]) // page_size + valid_current_rows = current_extend_kv_rows_for_reuse( + forward_batch, + current_index_kv[0], + current_index_kv[1], + ) + if valid_current_rows is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " + "CP shared KV index partial-current compose received " + "current_index_kv that does not satisfy current reuse " + "metadata. " + f"cp_rank={layout.cp_rank} layer_id={layer_id} " + f"prefix_lens={prefix_lens} extend_lens={extend_lens} " + f"current_k_shape={tuple(current_index_kv[0].shape)} " + f"current_scale_shape={tuple(current_index_kv[1].shape)} " + f"out_cache_loc_shape={tuple(current_locs.shape)}" + ) + current_locs = current_locs[:valid_current_rows] + current_index_kv = ( + current_index_kv[0][:valid_current_rows], + current_index_kv[1][:valid_current_rows], + ) + prefix_slot_span = None + if len(prefix_lens_cpu) == 1: + prefix_pages = int(prefix_lens_cpu[0]) // page_size + else: + prefix_pages = 0 + prefix_slot_span = build_batch_prefix_slot_span( + logical_pages=logical_page_table, + prefix_lens_cpu=prefix_lens_cpu, + page_size=page_size, + ) if index_prefetcher is not None: prefetched = index_prefetcher.consume_prefix_with_current( layer_id=layer_id, @@ -422,6 +443,7 @@ class Indexer(MultiPlatformOp): page_size=page_size, index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, prefix_pages=prefix_pages, + prefix_slot_span=prefix_slot_span, layer_id=layer_id, ) ) @@ -432,12 +454,13 @@ class Indexer(MultiPlatformOp): cp_shared_kv_mla_prefetch_log( "index_partial_current_sync_compose cp_rank=%s layer=%s " "prefix_lens=%s extend_lens=%s prefix_pages=%s " - "current_rows=%s dense_pages=%s", + "prefix_slot_span=%s current_rows=%s dense_pages=%s", layout.cp_rank, layer_id, prefix_lens, extend_lens, prefix_pages, + prefix_slot_span, int(current_index_kv[0].shape[0]), int(materialized.shape[0]), ) @@ -1481,12 +1504,6 @@ class Indexer(MultiPlatformOp): cp_metadata = forward_batch.nsa_cp_metadata assert cp_metadata is not None batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1) - if current_index_kv is not None: - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][index_topk] " - "reason=batch_gt1_index_current_reuse_unsupported " - f"batch_size={batch_size} layer_id={layer_id}" - ) request_kv_len_prev = list(getattr(cp_metadata, "request_kv_len_prev", []) or []) request_kv_len_next = list(getattr(cp_metadata, "request_kv_len_next", []) or []) @@ -1510,14 +1527,31 @@ class Indexer(MultiPlatformOp): f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next}" ) - shared_block_tables = metadata.get_page_table_64() - shared_index_buffer, shared_block_tables = ( - self._maybe_materialize_shared_index_buffer( - forward_batch, - layer_id, - shared_block_tables, + shared_index_buffer = None + shared_block_tables = None + current_index_kv_for_topk = current_index_kv + if current_index_kv is not None and not is_current_only_extend_batch( + forward_batch + ): + current_index_kv_for_topk = None + shared_block_tables = metadata.get_page_table_64() + shared_index_buffer, shared_block_tables = ( + self._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id, + shared_block_tables, + current_index_kv=current_index_kv, + ) + ) + elif current_index_kv is None: + shared_block_tables = metadata.get_page_table_64() + shared_index_buffer, shared_block_tables = ( + self._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id, + shared_block_tables, + ) ) - ) outputs = [] cursor = 0 @@ -1558,7 +1592,7 @@ class Indexer(MultiPlatformOp): metadata, kv_len, segment_len, - current_index_kv=None, + 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, @@ -1943,14 +1977,8 @@ class Indexer(MultiPlatformOp): current_index_kv = None if self._can_reuse_current_index_kv(forward_batch): - extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) - valid_current_rows = int(forward_batch.out_cache_loc.numel()) - if extend_seq_lens_cpu is not None and len(extend_seq_lens_cpu) == 1: - valid_current_rows = min( - int(extend_seq_lens_cpu[0]), - valid_current_rows, - ) - if key.shape[0] >= valid_current_rows: + valid_current_rows = current_extend_kv_rows_for_reuse(forward_batch, key) + if valid_current_rows is not None and key.shape[0] >= valid_current_rows: current_k_fp8, current_k_scale = act_quant( key[:valid_current_rows].contiguous(), self.block_size, diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 6cae1459d..56a3dbb98 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -102,7 +102,38 @@ def is_nsa_prefill_cp_round_robin_split(): ) +def _is_cp_shared_kv_forward_batch(forward_batch: "ForwardBatch") -> bool: + return bool(getattr(forward_batch, "uses_cp_shared_kv", False)) + + +def _fail_if_cp_shared_kv_round_robin( + forward_batch: "ForwardBatch", + *, + op: str, +) -> None: + if forward_batch is None or not _is_cp_shared_kv_forward_batch(forward_batch): + return + try: + round_robin_split = is_nsa_prefill_cp_round_robin_split() + except ValueError: + round_robin_split = False + if not round_robin_split: + return + + error_msg = ( + "[CP_SHARED_KV_FAIL_FAST][round_robin_unsupported] " + "CP shared KV only supports in-seq zigzag split. " + f"op={op} nsa_prefill_cp_mode=round-robin-split" + ) + logger.error(error_msg) + raise RuntimeError(error_msg) + + def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"): + _fail_if_cp_shared_kv_round_robin( + forward_batch, + op="can_nsa_prefill_cp_round_robin_split", + ) if not forward_batch.forward_mode.is_context_parallel_extend(): return False cp_size = get_attention_cp_size() @@ -934,6 +965,7 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent( def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): + _fail_if_cp_shared_kv_round_robin(forward_batch, op="can_cp_split") if is_nsa_prefill_cp_round_robin_split(): cur_cp_seq_len = seq_len // cp_size assert seq_len % cp_size == 0, ( @@ -964,6 +996,7 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): + _fail_if_cp_shared_kv_round_robin(forward_batch, op="cp_split_and_rebuild_data") if is_nsa_prefill_cp_round_robin_split(): cp_size = get_attention_cp_size() assert input_.shape[0] % cp_size == 0, ( @@ -985,6 +1018,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor): + _fail_if_cp_shared_kv_round_robin(forward_batch, op="cp_split_and_rebuild_1d") try: round_robin_split = is_nsa_prefill_cp_round_robin_split() except ValueError: @@ -1188,6 +1222,7 @@ def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch") def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor): + _fail_if_cp_shared_kv_round_robin(forward_batch, op="cp_split_and_rebuild_position") if is_nsa_prefill_cp_round_robin_split(): cp_size = get_attention_cp_size() assert positions.shape[0] % cp_size == 0, ( @@ -1327,13 +1362,19 @@ def _cp_attn_tp_all_gather_padded_tensor( max_len = (total_len + attn_tp_size - 1) // attn_tp_size pad_size = max_len - input_.shape[0] if pad_size > 0: - input_ = F.pad(input_, (0, 0, 0, pad_size), mode="constant", value=0) + input_ = torch.cat( + ( + input_, + input_.new_zeros((pad_size, *input_.shape[1:])), + ), + dim=0, + ) with use_symmetric_memory( get_attention_cp_group(), disabled=not is_allocation_symmetric() ): input_tensor_all = torch.empty( max_len * attn_tp_size, - input_.shape[1], + *input_.shape[1:], device=input_.device, dtype=input_.dtype, ) @@ -1465,6 +1506,140 @@ def _torch_in_seq_all_gather_rerange( return output_tensor.view(-1, hidden_size) +def _raise_batch_rerange_error(reason: str, message: str, *args) -> None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][batch_gt1_full_rerange_unsupported] " + + f"reason={reason} " + + (message % args if args else message) + ) + + +def _torch_batch_in_seq_all_gather_rerange( + input_tensor_all: torch.Tensor, + forward_batch: "ForwardBatch", + *, + cp_size: int, +) -> torch.Tensor: + """Restore rank-major batched in-seq CP all-gather rows to request order. + + The row payload is intentionally opaque. MLA bf16 latent rows, packed fp8 + KV rows, and NSA index byte rows all share the same first-dimension token + order contract, so this helper only slices/copies rows and never interprets + the dtype or payload shape. + """ + + metadata = getattr(forward_batch, "nsa_cp_metadata", None) + batch_size = int(getattr(metadata, "batch_size", 1) or 1) + request_split_lists = getattr(metadata, "request_split_lists", None) + max_rank_len = getattr(metadata, "max_rank_len", None) + if metadata is None: + _raise_batch_rerange_error("missing_metadata", "nsa_cp_metadata is missing") + if batch_size <= 1: + _raise_batch_rerange_error( + "not_batch", + "batch-aware rerange requires batch_size > 1, got %s", + batch_size, + ) + if request_split_lists is None or len(request_split_lists) != batch_size: + _raise_batch_rerange_error( + "missing_request_split_lists", + "request_split_lists is missing or incomplete. batch_size=%s value=%s", + batch_size, + request_split_lists, + ) + if max_rank_len is None or len(max_rank_len) < cp_size: + _raise_batch_rerange_error( + "missing_max_rank_len", + "max_rank_len is missing or shorter than cp_size. cp_size=%s value=%s", + cp_size, + max_rank_len, + ) + if cp_size <= 0: + _raise_batch_rerange_error("bad_cp_size", "cp_size must be positive: %s", cp_size) + + split_lists: List[List[int]] = [] + for req_id, split_list in enumerate(request_split_lists): + if split_list is None or len(split_list) != cp_size * 2: + _raise_batch_rerange_error( + "bad_request_split", + "request split must have 2 * cp_size entries. req_id=%s " + "cp_size=%s split=%s", + req_id, + cp_size, + split_list, + ) + split_lists.append([int(x) for x in split_list]) + + max_rank_token = int(max_rank_len[0]) + if max_rank_token < 0: + _raise_batch_rerange_error( + "bad_max_rank_token", + "max_rank_len[0] must be non-negative, got %s", + max_rank_token, + ) + required_rows = max_rank_token * cp_size + if int(input_tensor_all.shape[0]) < required_rows: + _raise_batch_rerange_error( + "input_rows_short", + "input rows=%s required=%s max_rank_token=%s cp_size=%s", + int(input_tensor_all.shape[0]), + required_rows, + max_rank_token, + cp_size, + ) + + rank_request_offsets: List[List[int]] = [] + for source_rank in range(cp_size): + mirror = cp_size * 2 - source_rank - 1 + offsets: List[int] = [] + cursor = 0 + for split_list in split_lists: + offsets.append(cursor) + cursor += split_list[source_rank] + split_list[mirror] + if cursor > max_rank_token: + _raise_batch_rerange_error( + "rank_payload_exceeds_max", + "rank local rows exceed max_rank_token. rank=%s rows=%s " + "max_rank_token=%s", + source_rank, + cursor, + max_rank_token, + ) + rank_request_offsets.append(offsets) + + total_tokens = sum(sum(split_list) for split_list in split_lists) + output_tensor = input_tensor_all.new_empty( + (total_tokens, *input_tensor_all.shape[1:]) + ) + if total_tokens == 0: + return output_tensor + + output_request_base = 0 + for req_id, split_list in enumerate(split_lists): + segment_prefix = [0] + list(accumulate(split_list))[:-1] + for segment_id, segment_len in enumerate(split_list): + if segment_len <= 0: + continue + if segment_id < cp_size: + source_rank = segment_id + source_segment_offset = 0 + else: + source_rank = cp_size * 2 - segment_id - 1 + source_segment_offset = split_list[source_rank] + source_start = ( + source_rank * max_rank_token + + rank_request_offsets[source_rank][req_id] + + source_segment_offset + ) + output_start = output_request_base + segment_prefix[segment_id] + output_tensor[output_start : output_start + segment_len].copy_( + input_tensor_all[source_start : source_start + segment_len] + ) + output_request_base += sum(split_list) + + return output_tensor + + def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream): """ # for in-seq-split @@ -1492,14 +1667,47 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream): | token0, token1, token2, token3, token4, token5, token6, token7, ... | +-------------------------+ """ + _fail_if_cp_shared_kv_round_robin( + forward_batch, + op="cp_all_gather_rerange_output", + ) metadata = getattr(forward_batch, "nsa_cp_metadata", None) - if getattr(metadata, "batch_size", 1) > 1: - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][batch_gt1_full_rerange_unsupported] " - "CP shared-KV bs>1 must not use scalar full hidden/KV rerange. " - "Use batch-aware narrow output collection or add batch-aware full " - "rerange metadata/kernels for this consumer." + batch_size = int(getattr(metadata, "batch_size", 1) or 1) + if batch_size > 1: + if input_tensor.dim() < 2: + _raise_batch_rerange_error( + "bad_input_dim", + "input tensor must have token rows plus payload dims, got shape=%s", + tuple(input_tensor.shape), + ) + if getattr(metadata, "request_split_lists", None) is None: + _raise_batch_rerange_error( + "missing_request_split_lists", + "request_split_lists is required for batch-aware current rerange. " + "batch_size=%s", + batch_size, + ) + total_seq_lens = getattr(metadata, "total_seq_lens", None) + if total_seq_lens is None: + _raise_batch_rerange_error( + "missing_total_seq_lens", + "total_seq_lens is required for batch-aware current rerange. " + "batch_size=%s", + batch_size, + ) + input_tensor_all = _cp_attn_tp_all_gather_padded_tensor( + input_tensor, + total_seq_lens, + cp_size, + forward_batch, + stream, ) + return _torch_batch_in_seq_all_gather_rerange( + input_tensor_all, + forward_batch, + cp_size=cp_size, + ) + if is_nsa_prefill_cp_round_robin_split(): with use_symmetric_memory( get_attention_cp_group(), disabled=not is_allocation_symmetric() @@ -1599,6 +1807,10 @@ def prepare_input_dp_with_cp_dsa( forward_batch: "ForwardBatch" = None, page_size: int = None, ): + _fail_if_cp_shared_kv_round_robin( + forward_batch, + op="prepare_input_dp_with_cp_dsa", + ) if is_nsa_prefill_cp_round_robin_split(): return True """prepare_input_dp_with_cp_dsa-zigzag index @@ -1769,6 +1981,10 @@ def cp_collect_last_token_hidden( forward_batch: "ForwardBatch", cp_size: int, ) -> torch.Tensor: + _fail_if_cp_shared_kv_round_robin( + forward_batch, + op="cp_collect_last_token_hidden", + ) if is_nsa_prefill_cp_round_robin_split(): return _round_robin_collect_last_token(hidden_states, forward_batch, cp_size) return _in_seq_collect_last_token(hidden_states, forward_batch, cp_size) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 5cbb476e0..fac68f791 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -15,6 +15,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import ( CpSharedKVMlaPrefetcher, ) from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + build_batch_prefix_slot_span, build_current_loc_remap, cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, @@ -1987,16 +1988,18 @@ class NativeSparseAttnBackend( int(current_kv_cache.shape[0]), tuple(logical_page_table_1.shape), ) - if ( - prefix_lens_cpu is None - or len(prefix_lens_cpu) != 1 - or int(prefix_lens_cpu[0]) <= 0 - or int(prefix_lens_cpu[0]) % page_size != 0 - ): + prefix_lens_valid = prefix_lens_cpu is not None and len(prefix_lens_cpu) > 0 + if prefix_lens_valid: + prefix_lens_valid = all( + int(prefix_len) >= 0 + and int(prefix_len) % page_size == 0 + for prefix_len in prefix_lens_cpu + ) and any(int(prefix_len) > 0 for prefix_len in prefix_lens_cpu) + if not prefix_lens_valid: raise RuntimeError( "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync] " "CP shared KV MLA partial-current sync compose " - "requires one positive page-aligned prefix. " + "requires positive page-aligned prefix pages. " f"reason={reason} " f"cp_rank={forward_batch.cp_shared_kv_layout.cp_rank} " f"layer_id={layer.layer_id} " @@ -2007,7 +2010,16 @@ class NativeSparseAttnBackend( f"current_locs_shape={tuple(current_locs_for_reuse.shape)} " f"page_size={page_size}" ) - prefix_pages = int(prefix_lens_cpu[0]) // page_size + prefix_slot_span = None + if len(prefix_lens_cpu) == 1: + prefix_pages = int(prefix_lens_cpu[0]) // page_size + else: + prefix_pages = 0 + prefix_slot_span = build_batch_prefix_slot_span( + logical_pages=metadata.real_page_table, + prefix_lens_cpu=prefix_lens_cpu, + page_size=page_size, + ) slot_remap = get_or_build_shared_token_kv_slot_remap( forward_batch, kv_cache=kv_cache, @@ -2025,6 +2037,7 @@ class NativeSparseAttnBackend( layout=forward_batch.cp_shared_kv_layout, page_size=page_size, prefix_pages=prefix_pages, + prefix_slot_span=prefix_slot_span, layer_id=layer.layer_id, ) ) @@ -2037,14 +2050,15 @@ class NativeSparseAttnBackend( cp_shared_kv_mla_prefetch_log( "forward_partial_current_sync_compose cp_rank=%s " "layer=%s reason=%s prefix_lens=%s extend_lens=%s " - "prefix_pages=%s current_rows=%s kv_rows=%s " - "page_table_shape=%s", + "prefix_pages=%s prefix_slot_span=%s " + "current_rows=%s kv_rows=%s page_table_shape=%s", forward_batch.cp_shared_kv_layout.cp_rank, layer.layer_id, reason, prefix_lens, extend_lens, prefix_pages, + prefix_slot_span, int(current_kv_cache.shape[0]), int(kv_cache.shape[0]), tuple(page_table_1.shape), diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index db2496f73..5f7421770 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -19,11 +19,13 @@ from sglang.srt.layers.attention.nsa.utils import ( cp_collect_last_token_hidden, cp_split_and_rebuild_1d, cp_split_and_rebuild_data, + _torch_batch_in_seq_all_gather_rerange, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, get_cp_shared_kv_local_physical_out_cache_loc, get_cp_local_embedding_padded_token_count, pad_cp_local_input_ids_for_embedding, + prepare_input_dp_with_cp_dsa, split_tensor_by_cp_batch_plan, split_in_seq_cp_local_pair, ) @@ -543,6 +545,201 @@ class TestNSAInSeqCPUtils(unittest.TestCase): torch.zeros((8, 1)), 2, forward_batch, stream=None ) + def test_cp_shared_kv_prepare_rejects_round_robin_mode(self): + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + extend_seq_lens_cpu=[128], + extend_prefix_lens_cpu=[0], + token_to_kv_pool=SimpleNamespace(page_size=64), + forward_mode=Mode(), + ) + + with ( + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=True, + ), + self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[round_robin_unsupported\]", + ), + ): + prepare_input_dp_with_cp_dsa( + 128, + cp_rank=0, + cp_size=2, + seqs_len=[128], + forward_batch=forward_batch, + page_size=64, + ) + + def test_cp_shared_kv_all_gather_rejects_round_robin_mode(self): + import torch + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + nsa_cp_metadata=NSAContextParallelMetadata( + split_list=[2, 2, 2, 2], + zigzag_index=[0, 3], + reverse_split_len=[2, 2, 2, 2], + cp_reverse_index=[0, 2, 3, 1], + total_seq_lens=torch.tensor(8), + ), + ) + + with ( + patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=True, + ), + self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[round_robin_unsupported\]", + ), + ): + cp_all_gather_rerange_output( + torch.zeros((4, 1)), 2, forward_batch, stream=None + ) + + def test_batch_in_seq_all_gather_rerange_restores_request_order_bf16(self): + import torch + + cp_size = 2 + request_split_lists = [ + [2, 1, 3, 0], + [1, 2, 0, 1], + ] + input_tensor_all, expected = self._build_batch_rerange_case( + cp_size=cp_size, + request_split_lists=request_split_lists, + row_width=3, + dtype=torch.bfloat16, + ) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_split_lists=request_split_lists, + max_rank_len=[6, 6], + ) + ) + + actual = _torch_batch_in_seq_all_gather_rerange( + input_tensor_all, + forward_batch, + cp_size=cp_size, + ) + + self.assertEqual(actual.dtype, torch.bfloat16) + self.assertTrue(torch.equal(actual, expected)) + + def test_batch_in_seq_all_gather_rerange_treats_fp8_payload_as_opaque_rows(self): + import torch + + cp_size = 2 + request_split_lists = [ + [1, 2, 1, 0], + [2, 0, 1, 1], + ] + input_tensor_all, expected = self._build_batch_rerange_case( + cp_size=cp_size, + request_split_lists=request_split_lists, + row_width=5, + dtype=torch.uint8, + ) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_split_lists=request_split_lists, + max_rank_len=[4, 4], + ) + ) + + actual = _torch_batch_in_seq_all_gather_rerange( + input_tensor_all, + forward_batch, + cp_size=cp_size, + ) + + self.assertEqual(actual.dtype, torch.uint8) + self.assertTrue(torch.equal(actual, expected)) + + def _build_batch_rerange_case( + self, + *, + cp_size, + request_split_lists, + row_width, + dtype, + ): + import torch + + rank_tokens = [] + for rank in range(cp_size): + mirror = cp_size * 2 - rank - 1 + rank_tokens.append( + sum(split[rank] + split[mirror] for split in request_split_lists) + ) + max_rank_token = max(rank_tokens) + total_tokens = sum(sum(split) for split in request_split_lists) + input_tensor_all = torch.zeros( + (max_rank_token * cp_size, row_width), + dtype=dtype, + ) + expected = torch.empty((total_tokens, row_width), dtype=dtype) + + next_value = 1 + request_segments = [] + for split in request_split_lists: + segments = [] + for segment_len in split: + if dtype == torch.uint8: + rows = ( + torch.arange( + next_value, + next_value + segment_len * row_width, + dtype=torch.int64, + ) + .remainder(251) + .to(torch.uint8) + .view(segment_len, row_width) + ) + else: + rows = ( + torch.arange( + next_value, + next_value + segment_len * row_width, + dtype=torch.float32, + ) + .view(segment_len, row_width) + .to(dtype) + ) + next_value += segment_len * row_width + segments.append(rows) + request_segments.append(segments) + + output_cursor = 0 + for segments in request_segments: + for rows in segments: + expected[output_cursor : output_cursor + rows.shape[0]] = rows + output_cursor += rows.shape[0] + + for rank in range(cp_size): + mirror = cp_size * 2 - rank - 1 + rank_cursor = rank * max_rank_token + for segments in request_segments: + for segment_id in (rank, mirror): + rows = segments[segment_id] + input_tensor_all[ + rank_cursor : rank_cursor + rows.shape[0] + ] = rows + rank_cursor += rows.shape[0] + + return input_tensor_all, expected + def test_local_pair_split_uses_metadata_lengths_not_half_split(self): import torch @@ -1209,12 +1406,69 @@ class TestNSAInSeqCPUtils(unittest.TestCase): [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], ) - def test_indexer_in_seq_cp_pair_batch_rejects_current_index_reuse(self): + def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self): import torch from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer indexer = object.__new__(Indexer) + indexer.index_topk = 2 + current_index_kv = (torch.tensor([1]), torch.tensor([2])) + 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) + materialize_calls = [] + topk_calls = [] + + class Metadata: + def get_page_table_64(self): + return logical_pages + + def fake_materialize( + forward_batch, + layer_id, + logical_page_table, + current_index_kv=None, + ): + materialize_calls.append( + { + "layer_id": layer_id, + "logical_page_table": logical_page_table, + "current_index_kv": current_index_kv, + } + ) + return materialized_index, dense_pages + + def fake_get_topk( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + kv_len, + actual_seq_q, + cp_index=None, + current_index_kv=None, + shared_index_buffer=None, + shared_block_tables=None, + actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, + batch_idx=0, + ): + topk_calls.append( + { + "batch_idx": batch_idx, + "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) + + indexer._maybe_materialize_shared_index_buffer = fake_materialize + indexer._get_topk_ragged_with_cp = fake_get_topk + forward_batch = SimpleNamespace( batch_size=2, nsa_cp_metadata=NSAContextParallelMetadata( @@ -1225,20 +1479,133 @@ class TestNSAInSeqCPUtils(unittest.TestCase): request_actual_seq_q_next=[1, 3], ), ) + result = Indexer._get_topk_in_seq_cp_pair( + indexer, + forward_batch, + layer_id=7, + q_fp8=torch.empty(7, 1), + weights=torch.empty(7, 1), + metadata=Metadata(), + current_index_kv=current_index_kv, + ) - with self.assertRaisesRegex( - RuntimeError, - "CP_SHARED_KV_FAIL_FAST.*batch_gt1_index_current_reuse_unsupported", + 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.assertTrue( + all(call["current_index_kv"] is None for call in topk_calls) + ) + 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.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( + result.tolist(), + [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], + ) + + def test_indexer_in_seq_cp_pair_batch_reuses_current_index_without_materialize(self): + import torch + + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + current_index_kv = (torch.tensor([1]), torch.tensor([2])) + topk_calls = [] + + class Mode: + def is_extend_without_speculative(self): + return True + + class Metadata: + def get_page_table_64(self): + raise AssertionError("current-only index reuse should not read page table") + + def fake_materialize(*args, **kwargs): + raise AssertionError("current-only index reuse should not materialize") + + def fake_get_topk( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + kv_len, + actual_seq_q, + cp_index=None, + current_index_kv=None, + shared_index_buffer=None, + shared_block_tables=None, + actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, + batch_idx=0, ): - Indexer._get_topk_in_seq_cp_pair( - indexer, - forward_batch, - layer_id=7, - q_fp8=torch.empty(7, 1), - weights=torch.empty(7, 1), - metadata=SimpleNamespace(), - current_index_kv=(torch.empty(1), torch.empty(1)), + topk_calls.append( + { + "batch_idx": batch_idx, + "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) + + 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=Mode(), + extend_prefix_lens_cpu=[0, 0], + extend_seq_lens_cpu=[3, 4], + seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64), + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_kv_len_prev=[3, 4], + request_kv_len_next=[3, 4], + request_actual_seq_q_prev=[2, 1], + request_actual_seq_q_next=[1, 3], + ), + ) + result = Indexer._get_topk_in_seq_cp_pair( + indexer, + forward_batch, + layer_id=7, + q_fp8=torch.empty(7, 1), + weights=torch.empty(7, 1), + metadata=Metadata(), + current_index_kv=current_index_kv, + ) + + self.assertEqual(len(topk_calls), 4) + self.assertTrue( + all(call["current_index_kv"] is current_index_kv for call in topk_calls) + ) + self.assertTrue( + all(call["shared_index_buffer"] is None for call in topk_calls) + ) + 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( + [call["actual_seq_q_cu_tensor"].tolist() for call in topk_calls], + [[0, 2], [0, 1], [0, 1], [0, 3]], + ) + self.assertEqual( + result.tolist(), + [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], + ) 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 d14b47a61..5b5f02a6a 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 @@ -565,6 +565,62 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64) self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + def test_can_reuse_current_extend_kv_allows_batch_partial_cache_hits(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + can_reuse_current_extend_kv, + ) + + forward_batch = SimpleNamespace( + forward_mode=_FakeExtendForwardMode(), + batch_size=2, + extend_prefix_lens_cpu=[40384, 8192], + extend_seq_lens_cpu=[65, 128], + seq_lens_cpu=torch.tensor([40384 + 65, 8192 + 128], dtype=torch.int32), + out_cache_loc=torch.arange(256, dtype=torch.int64), + ) + + self.assertTrue(can_reuse_current_extend_kv(forward_batch)) + + forward_batch.out_cache_loc = torch.arange(192, dtype=torch.int64) + self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + + forward_batch.out_cache_loc = torch.arange(256, dtype=torch.int64) + forward_batch.seq_lens_cpu = torch.tensor([40384 + 64, 8192 + 128], dtype=torch.int32) + self.assertFalse(can_reuse_current_extend_kv(forward_batch)) + + def test_current_extend_kv_rows_for_reuse_accepts_batch_bf16_and_fp8_payloads(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + forward_batch = SimpleNamespace( + forward_mode=_FakeExtendForwardMode(), + batch_size=2, + extend_prefix_lens_cpu=[0, 1024], + extend_seq_lens_cpu=[65, 128], + seq_lens_cpu=torch.tensor([65, 1024 + 128], dtype=torch.int32), + out_cache_loc=torch.arange(256, dtype=torch.int64), + ) + bf16_k = torch.empty((193, 2, 4), dtype=torch.bfloat16) + fp8_like_k = torch.empty((193, 5), dtype=torch.uint8) + + with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True): + self.assertEqual( + runtime.current_extend_kv_rows_for_reuse( + forward_batch, + bf16_k, + fp8_like_k, + ), + 193, + ) + + self.assertIsNone( + runtime.current_extend_kv_rows_for_reuse( + forward_batch, + bf16_k[:192], + fp8_like_k, + ) + ) + def test_should_reuse_current_extend_kv_enables_draft_partial_cache_hit_suffix( self, ): @@ -692,6 +748,21 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertIn("prefix_pages=0", branch_source) self.assertNotIn("kv_cache = current_kv_cache", branch_source) + def test_mla_partial_current_sync_uses_batch_prefix_slot_span(self): + from pathlib import Path + + source = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa_backend.py" + ).read_text() + branch_start = source.index("eagle_draft_mla_branch = \"partial_current_sync\"") + branch_end = source.index("forward_partial_current_sync_compose", branch_start) + branch_source = source[branch_start:branch_end] + + self.assertIn("build_batch_prefix_slot_span", source) + self.assertIn("prefix_slot_span=", branch_source) + self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source) + def test_nsa_backend_topk_transform_uses_effective_forward_impl(self): from sglang.srt.layers.attention.nsa_backend import ( NativeSparseAttnBackend, @@ -942,7 +1013,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): batch_size=2, extend_prefix_lens_cpu=[64, 64], extend_seq_lens_cpu=[64, 64], - seq_lens_cpu=torch.tensor([128, 128], dtype=torch.int32), + seq_lens_cpu=torch.tensor([129, 128], dtype=torch.int32), out_cache_loc=torch.arange(128, dtype=torch.int64), ) @@ -952,7 +1023,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): joined = "\n".join(logs.output) self.assertIn("[CP_SHARED_KV_FALLBACK][current_reuse]", joined) - self.assertIn("batch_size_not_one", joined) + self.assertIn("prefix_extend_seq_len_mismatch_req_0", joined) def test_tai_index_mqa_prepare_fast_path_miss_logs_warning(self): from sglang.srt.environ import envs @@ -1210,6 +1281,188 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(mixed_kv[12:14], current_kv)) self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]]) + def test_batch_prefix_slot_span_covers_request_prefix_pages_only(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + logical_pages = torch.tensor( + [ + [1, 2, 5], + [9, 11, 0], + ], + dtype=torch.int64, + ) + + self.assertEqual( + runtime.build_batch_prefix_slot_span( + logical_pages=logical_pages, + prefix_lens_cpu=[8, 4], + page_size=4, + ), + (0, 4), + ) + self.assertEqual( + runtime.build_batch_prefix_slot_span( + logical_pages=logical_pages, + prefix_lens_cpu=[0, 4], + page_size=4, + ), + (3, 4), + ) + self.assertEqual( + runtime.build_batch_prefix_slot_span( + logical_pages=logical_pages, + prefix_lens_cpu=[0, 0], + page_size=4, + ), + (0, 0), + ) + + def test_materialize_batch_prefix_span_and_reuse_current_kv_page_slots(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + page_size = 4 + layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0) + kv_cache = torch.arange(0, 64, dtype=torch.float32).view(64, 1, 1) + logical_locs = torch.tensor( + [ + [4, 8, 20, 21, -1, -1], + [36, 44, 45, -1, -1, -1], + ], + dtype=torch.int64, + ) + current_locs = torch.tensor([20, 21, 44, 45], dtype=torch.int64) + current_kv = torch.arange(100, 104, dtype=torch.float32).view(4, 1, 1) + remap_logical_pages = torch.tensor( + [ + [1, 2, 5], + [9, 11, 0], + ], + dtype=torch.int64, + ) + slot_remap = runtime.build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=logical_locs, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=page_size, + ) + prefix_slot_span = runtime.build_batch_prefix_slot_span( + logical_pages=remap_logical_pages, + prefix_lens_cpu=[8, 4], + page_size=page_size, + ) + + with patch.object( + runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce + ): + mixed_kv, mixed_locs = ( + runtime.materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=logical_locs, + current_kv_cache=current_kv, + current_locs=current_locs, + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + prefix_pages=0, + prefix_slot_span=prefix_slot_span, + ) + ) + + self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1], [16, 20, 21, -1, -1, -1]]) + self.assertTrue(torch.equal(mixed_kv[4:8], kv_cache[4:8])) + self.assertTrue(torch.equal(mixed_kv[8:12], kv_cache[8:12])) + self.assertTrue(torch.equal(mixed_kv[16:20], kv_cache[36:40])) + self.assertTrue(torch.equal(mixed_kv[12:14], current_kv[:2])) + self.assertTrue(torch.equal(mixed_kv[20:22], current_kv[2:])) + + def test_materialize_batch_prefix_span_and_reuse_current_index_page_slots(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + page_size = 4 + index_head_dim = 4 + scale_bytes = 4 + page_bytes = page_size * index_head_dim + page_size * scale_bytes + layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0) + page_buffer = torch.zeros((16, page_bytes), dtype=torch.uint8) + page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8) + page_buffer[2] = torch.arange(40, 40 + page_bytes, dtype=torch.uint8) + page_buffer[9] = torch.arange(70, 70 + page_bytes, dtype=torch.uint8) + logical_pages = torch.tensor([[1, 2, 5], [9, 11, 0]], dtype=torch.int64) + slot_remap = runtime.build_shared_paged_buffer_slot_remap( + page_buffer, + logical_pages, + layout, + ) + prefix_slot_span = runtime.build_batch_prefix_slot_span( + logical_pages=logical_pages, + prefix_lens_cpu=[8, 4], + page_size=page_size, + ) + current_k = torch.tensor( + [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12], [13, 14, 15, 16]], + dtype=torch.uint8, + ) + current_scale = torch.tensor([[1.25], [2.5], [3.5], [4.5]], dtype=torch.float32) + + with patch.object( + runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce + ): + dense_page_buffer, dense_pages = ( + runtime.materialize_prefix_and_reuse_current_index_page_slots( + page_buffer=page_buffer, + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=torch.tensor([20, 21, 44, 45], dtype=torch.int64), + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + index_head_dim=index_head_dim, + prefix_pages=0, + prefix_slot_span=prefix_slot_span, + layer_id=2, + ) + ) + + scale_offset = page_size * index_head_dim + self.assertEqual(dense_pages.tolist(), [[1, 2, 3], [4, 5, 0]]) + self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1])) + self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[2])) + self.assertTrue(torch.equal(dense_page_buffer[4], page_buffer[9])) + self.assertTrue(torch.equal(dense_page_buffer[3, 0:4], current_k[0])) + self.assertTrue(torch.equal(dense_page_buffer[3, 4:8], current_k[1])) + self.assertTrue(torch.equal(dense_page_buffer[5, 0:4], current_k[2])) + self.assertTrue(torch.equal(dense_page_buffer[5, 4:8], current_k[3])) + expected_scale_bytes = ( + current_scale[:2].contiguous().view(torch.uint8).reshape(2, scale_bytes) + ) + self.assertTrue( + torch.equal( + dense_page_buffer[ + 3, + scale_offset : scale_offset + page_size * scale_bytes, + ].view(page_size, scale_bytes)[:2], + expected_scale_bytes, + ) + ) + + def test_index_partial_current_sync_uses_batch_prefix_slot_span(self): + from pathlib import Path + + source = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py" + ).read_text() + branch_start = source.index(" if current_index_kv is not None:") + branch_end = source.index(" return materialized, dense_pages", branch_start) + branch_source = source[branch_start:branch_end] + + self.assertIn("build_batch_prefix_slot_span", source) + self.assertIn("prefix_slot_span=", branch_source) + self.assertNotIn("or len(prefix_lens_cpu) != 1", branch_source) + def test_ipc_page_descriptor_builder_maps_slots_to_owner_physical_pages(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout