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:
@@ -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 逻辑。
|
||||
|
||||
---
|
||||
|
||||
|
||||
Reference in New Issue
Block a user