Files
sglang/python
laoyao0822 e4cf8d18b4 Preserve CP narrow output while planning real batches
Batch-size support needs request-first CP metadata; treating a batch as one long sequence breaks page ownership, top-k ranges, and phase1 compact output collection. This adds a batch CP plan that records per-request page-aligned splits, rank-local offsets, kv/actual-seq metadata, last-token owners, and flattened descriptors for downstream allocator/runtime workstreams.

The scalar full-rerange path now fail-fasts for batch metadata so bs>1 cannot silently discard the narrow-output optimization or restore hidden states with single-request assumptions.

Constraint: CP shared-KV cache state is page-owned and must preserve request boundaries under bs>1.

Rejected: Let bs>1 fall back to scalar full hidden rerange | it loses the phase1 communication reduction and uses wrong single-request metadata.

Rejected: Add a collective to confirm batch plans | all ranks can derive the same plan from CPU metadata and config.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove batch fail-fast guards until W2/W3 consumers use CPSharedKVBatchPlan end-to-end.

Tested: python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py

Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py -> 39 passed

Not-tested: full ETE bs>1 CP shared-KV runtime; W2/W3 allocator/direct-write consumers are not implemented yet
2026-06-03 01:21:43 +08:00
..