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