Enable compute-owner KV layout by page-aligning NSA CP split

Phase 5 needs each current KV page to have exactly one CP compute owner before local KV/index direct writes can be safe. This change teaches in-seq NSA prefill CP to produce page-aligned split metadata under shared-KV mode, threads page size into the metadata builders, and fixes local pair splitting so unequal page-aligned zigzag segments do not corrupt topk inputs.

Constraint: Phase 5 direct-write layout requires page ownership to be expressible at page granularity
Constraint: Short page-unit batches remain on the token-balanced fallback to avoid zero-page segment risk
Rejected: Split local q/weights by half | page-aligned zigzag segments can have unequal token counts
Confidence: medium
Scope-risk: moderate
Directive: Do not enable compute-owner direct writes unless nsa_cp_metadata.page_aligned is true and local loc ownership is verified
Tested: python3 -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py python/sglang/srt/models/deepseek_v2.py python/sglang/srt/models/deepseek_nextn.py test/registered/unit/layers/test_nsa_cp_utils.py
Not-tested: Local pytest collection is blocked in this environment by missing pybase64; container/runtime tests were not rerun during this commit step
This commit is contained in:
laoyao0822
2026-05-01 00:54:16 +08:00
parent 47bd2fdf1f
commit 91fa31bcac
6 changed files with 344 additions and 16 deletions

View File

@@ -234,7 +234,14 @@ Phase 4 MVP 建议:
fallback 到旧 token-average split
```
推荐先使用保守 gate,跑通后再放宽。
当前实现采用保守 gate
```text
如果 num_units < 2 * cp_size:
fallback 到旧 token-average split
```
原因是 CP 本身不适合短序列;真实长上下文场景下 page unit 数通常远大于 `2 * cp_size`,先保证每个 zigzag segment 至少拿到一个完整 page unit可以避免 zero-token segment 给通信、attention kernel 和后续 compute-owner layout 带来额外风险。后续如果要让短序列直接不走 CP应单独收紧 `can_cp_split(...)` 的启用阈值,而不是混入 Phase 4 的 page-aligned split 逻辑。
---