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

Enable batched CP shared-KV current reuse correctness

Browse Source 
Batch-size > 1 cache-hit traffic must not lose the current/partial-current reuse fast path. This change extends the target MLA/index sync-correct path to validate batched current suffix rows, compose page-aligned prefix spans, and route batched index top-k through current-only or partial-current reuse without falling back to scalar guards.\n\nThe implementation keeps page as the minimum cache unit: prefix cache coverage is page-aligned, while current suffix rows are sliced by valid extend lengths so padded tail rows are not exposed to attention. The index top-k batch path now mirrors the single-request contract: current-only reuses current index KV directly, partial-current materializes once and shares the dense buffer across request segments.\n\nConstraint: CP shared-KV supports in-seq zigzag only; round-robin remains fail-fast for shared-KV.\nConstraint: No new collectives are introduced; this is a sync correctness path, not a new communication scheme.\nRejected: Keep bs>1 current_index_kv fail-fast | disables the cache-hit path that W4 is meant to restore.\nRejected: Pad requests to a common batch length | violates the page-granular contract and exposes padded tails to consumers.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce batch_size_not_one or batch_gt1 current-reuse guards without proving an equivalent batched fast path exists.\nTested: Local py_compile for touched runtime/test files.\nTested: Remote g0034 pytest test_nsa_cp_utils.py test_cp_shared_kv_runtime.py => 157 passed, 5 warnings, 2 subtests passed.\nTested: Remote g0034 pytest test_cp_shared_kv_layout.py => 27 passed, 3 warnings.\nNot-tested: Full ETE with live traffic; CUDA kernel-backed batched descriptors; L1 prefetch, HiCache load/backup, and draft/EAGLE bs>1 current reuse.
This commit is contained in:
laoyao0822
2026-06-03 07:06:21 +08:00
parent a7472c415f
commit 50d0008705
9 changed files with 1195 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan.md
+7 -6
View File
@@ -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
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md
+7 -6
View File
@@ -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 是单请求结构
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md
+56 -28
View File
@@ -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 保持
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+170 -35
View File
@@ -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,
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+73 -45
View File
@@ -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,
python/sglang/srt/layers/attention/nsa/utils.py
+224 -8
View File
@@ -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)
python/sglang/srt/layers/attention/nsa_backend.py
+24 -10
View File
@@ -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),
test/registered/unit/layers/test_nsa_cp_utils.py
+379 -12
View File
@@ -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
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+255 -2
View File
@@ -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