CP shared-KV: global-position-rotated zigzag owner-lane balancing (default on)
In NSA in-seq-split prefill-CP shared-KV the owner builder restarted every
chunk/request at owner 0, starving physical lane 0 -> load-back livelock (g0057:
~64k cp_load_back_owner_lane_capacity_failed + ~4k owner_lane_exhausted). Rotate
the owner/compute assignment by phase=((extend_prefix_len//page_size)+req_index)
%cp_size so chunks/requests/draft pages spread across lanes, keeping owner==compute
(the local_loc_owner_mismatch invariant) by rotating the zigzag compute split by the
same phase. Cache-safe: reload replays recorded page_owners, never recomputes phase.
- owner builder + helpers (cp_owner_lane_phase two-term prefix+req_index,
cp_in_seq_reverse_map, cp_rotated_per_rank_actual_token) in
cp_shared_kv_compute_owner.py
- zigzag compute split, gather-back reverse map, bs>1 torch rerange, last-token
owner, request_phases on NSAContextParallelMetadata, scalar-path + last-token
fail-loud guards in nsa/utils.py
- mqa-logits prefill buffer estimator made rotation-aware (per-request phase) so it
no longer under-sizes the admission budget
- flag SGLANG_CP_SHARED_KV_OWNER_ROTATION (default True; =0 is the legacy escape
hatch -- phase forced to 0 = byte-for-byte legacy, needs no new kernel)
- startup require_tai_kernel_version("0.0.2") gate (utils/common.py, model_runner)
- unit tests (11): tiny-chunk pathology [40,0x7]->[5x8], synchronized-batch spread
[0x8]->[0..7], gather-back reassembly, per-rank-token
Requires tai-kernel >= 0.0.2 (phase-aware in_seq_all_gather_rerange). Reviewed by 3
adversarial agents (owner==compute consistency, gather-back equivalence, integration
safety): SHIP-WITH-NITS, all findings fixed.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
leavelet
014948c5d7