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:32 +08:00
parent 47bd2fdf1f
commit 91fa31bcac
6 changed files with 344 additions and 16 deletions
+6
View File
@@ -2204,6 +2204,12 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.cp_rank,
self.cp_size,
forward_batch.seq_lens_cpu.tolist(),
forward_batch=forward_batch,
page_size=getattr(
getattr(forward_batch, "token_to_kv_pool", None),
"page_size",
None,
),
)
with get_attn_tp_context().maybe_input_scattered(forward_batch):