Revert the tail-chunk co-batching gate from 54c056af because allowing a continued chunk to share the next CP bs>1 batch reopens the mixed chunk/page-tail scheduler risks we are currently avoiding. Keep the independent real-prefix budget accounting so chunked requests still contribute their carried prefix to CP cached-token and buffer estimates.\n\nConstraint: Chunked-prefill requests must remain solo until the CP split/page-tail contract is revalidated for mixed batches.\nRejected: Full revert of 54c056af | it would also drop true-prefix budget accounting and under-estimate cache/buffer pressure for admitted chunks.\nConfidence: high\nScope-risk: moderate\nDirective: Do not reintroduce tail-chunk co-batching without tests covering page-tail split, CP buffer admission, and ETE chunked+cache-hit replay.\nTested: Local py_compile for environ.py, schedule_policy.py, cp_shared_kv_compose.py, test_prefill_adder.py, test_cp_shared_kv_compose_v2_8rank.py.\nTested: Remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_prefill_adder.py -> 25 passed.\nTested: Remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 157 passed, 2 subtests passed.\nNot-tested: Full mixed replay with chunked-prefill traffic after service restart.
专题 S4 (design: docs_internal/perf/prefill-compute-intensity-plan.md S4,
amended). Under FCFS a cold request joining a warm-led batch turns a
1-2s cache-hit forward into a 5-10s one, splitting the warm work into
the 新-cache-新 pattern. The policy prevents exactly that one thing:
- WARM candidates always admit (into a cold-led batch they are free
density — the cold extend dominates the forward anyway).
- COLD admits into an empty or cold-led batch (small colds co-batch
today; the FCFS head always starts a batch so the queue keeps moving).
- COLD into a WARM-led batch is skipped, bounded by a per-pass window
(W=16 skips), a head defer count (K=3 passes) and an age bound
(T=5s). On any bound the scan STOPS instead of force-admitting: the
cold waits for the same forward either way, but leads its own clean
batch next pass instead of polluting this one.
The skip is strictly post-match / pre-admit (after init_next_round_input,
before add_one_req): no lock, no allocation, no budget mutation to
unwind, and re-matching a skipped candidate next pass is exactly what
the scan already does after a cap rejection. Classification is the
in-scan match result (device prefix + host hit vs a 64-token floor) —
under FCFS+L2 no pre-scan signal exists, so this adds zero matching
work for inspected candidates. Disabled wholesale under priority
scheduling (the skip must not reorder across priority classes).
Three amendments vs the design draft, reasoned in the decision-table
docstring: cold+cold-led admits (STOP would regress today's small-cold
co-batching); starved heads STOP rather than force-admit (clean batch
boundaries at identical latency); priority interaction handled by
disabling rather than per-request comparison.
Decision logic is a pure function with table + bounds unit tests
(28/28 adder suite green). Default OFF.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
专题 S1 (design: docs_internal/perf/prefill-compute-intensity-plan.md
S1.0-S1.11). A batch containing a chunked prefill request has been
forced to bs=1 by the CP gate, so every chunk of a long prompt
monopolizes a forward while short-extend cache-hit continuations queue
— the direct cause of the replay TTFT tail (p90 19.3s / p99 49.2s at
91.8% cache hit). Yet mixed chunk batches already occur today (a
freshly-chunked request keeps earlier-admitted small requests), proving
the CP forward path is mixed-chunk-safe; only admission was asymmetric.
Three changes, the first flag-independent:
- add_chunked_req now seeds the budget with the chunk's TRUE prefix
(was 0), so the CP cached tally and the buffer estimator's mqa_logits
k_rows see the chunk's footprint before any later request is admitted
(landmine D1).
- New SGLANG_CP_PREFILL_MIX_CHUNKED (default OFF): with it on, the gate
admits requests after a chunked one and lets the existing CP caps
(extend / cached / buffer, now correctly seeded) bound the batch — a
FULL chunk still ends the scan by consuming the chunk-clamped extend
cap; only a tail chunk leaves headroom. A chunked prefix that is not
page-aligned (rare sub-page final-chunk tail) keeps its batch solo
(the CP page-aligned split would fail-fast otherwise).
- The symm staging capacity identity (admission extend cap + request
slack == staging pages) is asserted when the flag is on, locking the
coupling the design relies on (plan doc S1.4 I2).
Tests: 4 new adder units (budget seeding; tail chunk admits followers;
full chunk solo by budget; non-aligned prefix solo); the 8-rank
byte-exactness scenario gains a chunk-shaped request (2048-token
page-aligned carried prefix + 512 extend) — all four phases
(legacy/v2/symm/prefetch) byte-identical on g0033. Known pre-existing
cross-file pollution noted in problems.md P16.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
E2e caught one request wedged FOREVER in disagg_prefill_inflight_queue
(probes 21s apart with zero traffic both showed #inflight-req: 1, no
reap/timeout warnings ever logged). Mechanism, established by reading
the full state machine: prefill Success is set locally by the transfer
worker on the LAST chunk; if the decode peer is torn down between the
handshake and the prefill's final send(), add_transfer_request silently
drops the chunk (no transfer destinations) — Success becomes
unreachable. The only external rescue, the decode ABORT notification,
is best-effort (silently swallowed on send error, no-op if it races the
room registration), there is no prefill-side heartbeat of decode
sessions, and the sender's only timeout covers Bootstrapping — the
inflight queue itself has no liveness bound. The orphan pins the
request's KV pages and rides every poll collective.
Two fixes, both reaped through the existing Failed branch via the
CP/TP MIN-reduce poll consensus (Failed=0 wins, so one rank concluding
flips every rank together — rank-uniform by construction):
- add_transfer_request: a room with no transfer destinations that is
NOT already Success (the dummy-rank handshake marking) now concludes
Failed loudly instead of dropping the chunk silently.
- Inflight residency timeout: entries stuck in a non-terminal poll
state past SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT (default 300s,
matching the sibling BOOTSTRAP/WAITING timeouts) get sender.abort()
and reap on the next poll. Covers what the hardening cannot: lost
ABORT datagrams, decode crashes.
Known sibling gaps left for follow-up: the decode transfer queue has
no Transferring liveness bound, and an abort that matches no queue is
still a silent no-op (much narrower race than first thought — work
requests are ordered before control requests within a tick).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replaces the last remaining NCCL collective in the bs>1 compose layer loop
(the compact-current all-reduce) with a DeepEP-style counting barrier plus
a CUDA-IPC gather from peers' symmetric dense buffers, behind
SGLANG_CP_SHARED_KV_COMPOSE_SYMM (+ COMPOSE_ARENA, both default off).
- CpComposeArena.register_symm: fixes capacity at the pool-derived bound
(logical pages x dense page unit, overridable via
SGLANG_CP_SHARED_KV_SYMM_HEAP_MB), allocates slab + flags in CUDA-IPC
memory, exchanges handles once over the CP group; deterministic bump
carve means peer_base + my_offset addresses any peer's dense buffer with
no per-layer handshakes. Registration happens only in the token-KV
compose (uniform first-use point); growth after registration raises.
- Current pages are single-writer at page granularity under the
page-aligned in-seq split, so the exchange is the existing
gather_cuda_ipc_peer_pages with src==dst page ids and writer (compute
owner) ranks; writers are built per batch by
build_batch_current_page_writer_ranks and gated by
maybe_build_current_page_writer_ranks (page_aligned metadata required).
The barrier runs even with zero remote pages (counts must match).
- _agreed_tai_ipc_peer_ptrs: the per-rank IPC capability probe is now
agreed across the CP group (one-time MIN all-reduce per pool tensor) so
ranks can never split between gather and collective paths and deadlock
on mismatched NCCL shapes.
- ComposePlan cache re-anchored ON the slot_remap object (forward-batch
lifetime) instead of a module-level tensor-identity key, which could go
stale when a freed tensor's address is reused by the next batch.
Validation (g0034 syh-dev-new): tai-kernel cp_symm_barrier correctness
(200 adversarial iterations, rotating 10ms producer delays, phase-safety,
flags drained at quiescence) and perf (7.7us max-rank latency) both pass;
8-rank GPU test extends to the symm path - byte-identical to compose_v2
across 4 layers (parity halves exercised, slab registered); mem_cache
suite 464 passed with the only failures being a documented pre-existing
sys.modules stub pollution pair, reproduced identically with
SGLANG_CP_SHARED_KV_COMPOSE_V2=0.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The bs>1 partial-current materialize issued one sum-all-reduce per request
span per buffer per layer (48 collectives per F-layer at bs=12, 2880 per
batch, 419ms + 49ms launch gaps in the production trace), all inline on
the compute stream. The data is a partitioned gather, not a reduction:
every byte has exactly one producer.
compose_v2 (SGLANG_CP_SHARED_KV_COMPOSE_V2, default on) replaces this with:
- fast path: one tai-kernel CUDA-IPC slot-dense gather covering ALL prefix
spans (full-range descriptors, -1 sentinels zero-fill current slots and
replace the dense zero-fill) + ONE collective over the compact current
pages (uint8 byte view; exact because every byte is writer-exclusive).
- fallback (no peer IPC): local materialize of all prefix spans + ONE
whole-buffer sum-all-reduce (rows are still writer-exclusive pre-reduce).
IPC capability is decided once by the cached peer-pointer probe; after a
successful probe a failing gather raises (no per-call try/except).
cp_shared_kv_compose.py adds the per-batch ComposePlan descriptor cache
(layer-invariant, keyed on the slot_logical_pages identity) and the
CpComposeArena with tier-S carve discipline (deterministic bump, layer-
parity halves; default off) so the Step B symmetric-memory conversion is
a registration flip.
Microbenchmark (g0034 8xH200, traced 12-req batch, per batch): per-span
214ms -> fused AR 119ms -> IPC prefix + compact current 84ms; symm target
61ms. Validation: 143 unit tests incl. v2-contract twins, legacy siblings
and rank-merged simulations under both paths; mem_cache dir 432 passed;
8-rank GPU byte-exactness vs legacy with real NCCL + IPC (test/manual/
test_cp_shared_kv_compose_v2_8rank.py) passed with no fallback markers.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
CP shared-KV batching previously estimated MQA logits from full request
extend/context rows, which overstated memory because CP in-seq split only
computes each rank's two zigzag segments. Add CP-size aware row accounting
that mirrors the fused CP MQA materialization path and take the worst local
rank peak for scheduler admission.
Expose SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB as a more direct cap for one fp32
MQA logits chunk. Runtime and scheduler now both translate this GB cap into
chunk rows from the actual K rows, while keeping the old row cap as a
mutually-exclusive expert override.
Constraint: Scheduler admission must stay CUDA-sync-free and use static budget information only.
Rejected: Keep full-request q*k admission | it over-gates CP bs>1 batches because CP splits q rows per rank.
Rejected: Let rows and GB caps both apply | precedence would be ambiguous during tuning.
Confidence: medium
Scope-risk: moderate
Directive: Keep MQA logits admission tied to the fused CP MQA segment shape; do not revert to full request token counts.
Tested: Local py_compile for touched runtime, scheduler, estimator, and tests.
Tested: Local pytest test_cp_shared_kv_prefill_buffer_estimator.py: 8 passed.
Tested: Remote g0034 cjy-glm5-new py_compile and targeted estimator/runtime tests: 13 passed.
Not-tested: Full ETE high-cache-hit CP bs>1 load with SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_GB.
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP shared-KV bs>1 can build large fp32 MQA-logits temporaries from
DeepGEMM fp8_mqa_logits. The official SGLang path already chunks normal NSA
MQA logits by query rows behind a cached memory budget; carry the same budget
control into our NSA indexer and extend it to CP-ragged topk paths that use
row-wise topk_indices_offset_override.
This keeps the previous one-time cached memory-budget behavior rather than the
recent current-free-mem per-forward variant that regressed performance. A new
optional max-rows env provides an explicit hard cap for debugging or controlled
ETE runs without adding host syncs.
Constraint: DeepGEMM materializes fp32 [q, k] logits internally, so row chunking is the narrowest way to cap temporary memory
Rejected: Restore the reverted syh current-free-mem implementation | it changed hot-path heuristics and showed poor runtime performance
Rejected: Split by K/context dimension | would change topk semantics and require a different transform contract
Confidence: medium
Scope-risk: moderate
Directive: CP-ragged chunking relies on topk_indices_offset_override being row-addressed; do not route non-ragged CP paths through it without separate validation
Tested: Local py_compile for environ.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py
Tested: Remote g0034 cjy-glm5-new py_compile for environ.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py
Tested: Remote pytest TestCpSharedKVTaiMaterializeIntegration, 17 passed
Not-tested: CUDA ETE high-cache-hit bs>1 workload memory/performance after chunking
Co-authored-by: OmX <omx@oh-my-codex.dev>
Temporarily revert the syh MQA-logits row-chunking port because the observed
runtime performance is worse on the current bs>1 CP prefill workload. Keep the
history explicit so the optimization can be revisited after profiling identifies
where the extra overhead comes from.
This reverts commit 4e49751406 (Bound NSA MQA logits peak memory).
Constraint: Current priority is restoring the faster known path for remote ETE runs
Rejected: Keep chunking behind the force/debug env only | the production auto-chunk path still changes runtime heuristics and should not stay until profiled
Confidence: high
Scope-risk: narrow
Directive: Reintroduce logits chunking only with ETE performance evidence and forced-chunk equivalence coverage
Tested: Local py_compile for environ.py and nsa_indexer.py
Not-tested: Remote ETE performance after revert
Co-authored-by: OmX <omx@oh-my-codex.dev>
Paged and CP-ragged NSA indexer paths could materialize q x context fp32
MQA-logits buffers large enough to OOM high-cache-hit bs>1 prefill batches.
Port the syh branch chunking logic so paged and ragged paths split logits by
query rows when the estimated logits buffer exceeds the current free-memory
budget.
The free-memory query is cached on forward_batch so the OOM guard uses current
free memory without adding a torch.cuda.mem_get_info host sync on every layer.
The only new env kept from the syh commits is
SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS, which forces chunking for equivalence
validation.
Constraint: DeepGEMM fp8_mqa_logits still materializes fp32 logits internally, so limiting q rows is the least invasive way to cap peak memory
Rejected: Carry unrelated syh envs for page trace/source-fingerprint strictness | not part of the logits peak-memory fix
Rejected: Static mem_fraction-only budget | overestimates logits headroom shared with other forward activations
Confidence: medium
Scope-risk: moderate
Directive: Keep chunking row-split only; changing K/context partitioning needs topk_transform equivalence validation
Related: 40a0389a9c feat(nsa): chunk paged + CP-ragged MQA-logits by current-free-mem budget
Related: 108fa1f538 perf(nsa): cache MQA-logits free-mem budget per-forward
Tested: Local py_compile for environ.py and nsa_indexer.py
Tested: Remote g0034 cjy-glm5-new py_compile for environ.py and nsa_indexer.py
Not-tested: CUDA ETE run with forced SGLANG_NSA_MQA_LOGITS_CHUNK_FORCE_ROWS equivalence check
Not-tested: Full high-cache-hit bs>1 prefill OOM regression workload
Co-authored-by: OmX <omx@oh-my-codex.dev>
High cache-hit NSA sparse attention only gathers rows selected by topk_indices, but the FP8 RAGGED path dequantized the whole materialized K buffer first. Wire the syh in-place topk-only dequant path behind an explicit env gate so the sparse path can dequant referenced rows at their original row ids while keeping topk_indices unchanged. The old full-dequant path remains the default and the verification gate stays available for small correctness checks only.
Constraint: Production runs should enable only SGLANG_NSA_DEQUANT_ONLY_TOPK=1; VERIFY also runs the full-dequant reference and is too expensive for perf tests.
Rejected: Compact/remap topk buffer | data-dependent shape and remap add synchronization/aliasing risk; syh final in-place contract avoids both.
Confidence: medium
Scope-risk: moderate
Directive: Do not enable SGLANG_NSA_DEQUANT_ONLY_TOPK_VERIFY in throughput runs; use it only for small correctness probes.
Tested: python -m py_compile python/sglang/srt/environ.py python/sglang/srt/layers/attention/nsa_backend.py test/registered/unit/layers/test_nsa_dequant_only_topk.py
Tested: git diff --check
Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_dequant_only_topk.py -> 2 passed
Tested: remote g0034 CUDA smoke for tai_kernel.nsa_prefill.nsa_dequant_topk_inplace topk rows matched torch reference
Not-tested: Full GSM8K or replay ETE with SGLANG_NSA_DEQUANT_ONLY_TOPK=1
For large bs (production target: bs~10 x ~100k tokens), per-layer granularity does
10x79 = 790 submitTransfer calls + CUDA events + enqueues per forward on the forward
thread. Two overhead cuts:
- Group SGLANG_CP_SHARED_KV_PER_LAYER_GROUP (default 8) consecutive layers into ONE
RDMA submit: ~num_layers/K submits + events + enqueues instead of per-layer; same
bytes (page index lists are identical across layers). on_layer_end is O(1) at
non-boundary layers. The last partial group enqueues via the num_layers boundary;
any misses fall back to one batched sync submit.
- Scheduler hook skips reqs already registered (bs>1 batch-forming re-iterates the
same reqs ~9x -> was rebuilding the CP filter + context every time).
27 unit tests pass incl. grouping-boundary + batched-fallback.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add _transfer_layers_async: submit each layer's transfer non-blocking via the
async API (batch_transfer_async_submit), pipelining layers in the RDMA engine,
then wait for all once (wait_batch_transfers). Gated by the new
SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER env (default off); replaces the monolithic
all-layers batch_transfer_sync on that path. This is the transfer mechanism for
per-layer overlap (lever A) and removes the per-layer blocking-sync tax measured
in B1a; the forward-overlap hook (G2) builds on it next. Uses the safe async API,
never the OnCuda busy-wait/_exit path.
Unit-tested (test_per_layer_transfer.py, 5 cases): one submit per non-empty layer,
single wait-for-all, empty-layer skip, submit-failure drain + return -1, and
wait-status propagation.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The bs>1 path regressed from the beginning of a run, so log throughput alone is not enough to identify whether scheduler preparation, batch splits, index/top-k, sync compose, HiCache load planning, or valid-row selection is dominating. Add a separate env-gated timing channel with rate and slow-threshold controls so production-like runs can identify the hot operation without enabling verbose structural debug logs.\n\nInstrumentation is intentionally observational: it does not change batch planning, cache residency, transfer semantics, or fast-path eligibility. The timing logs are limited and can be filtered by SGLANG_CP_SHARED_KV_BS_GT1_TIMING_SLOW_MS.\n\nConstraint: Need runtime evidence without starting traffic from the agent side.\nConstraint: Existing bs=1 batch-plan behavior remains expected and unchanged.\nRejected: Reuse SGLANG_CP_SHARED_KV_BS_GT1_DEBUG | existing debug logs are structural and already too noisy for timing diagnosis.\nRejected: Add unconditional timing logs | would perturb the hot path and flood long ETE runs.\nConfidence: high\nScope-risk: narrow\nDirective: Keep this timing channel observational; do not use it to gate behavior or silently change batch policy.\nTested: g0034 docker py_compile for changed Python files.\nTested: g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py — 241 passed, 5 warnings, 2 subtests passed.\nNot-tested: Live ETE timing log interpretation under user-driven traffic.\nNot-tested: Nsight correlation with timing events.
The bs>1 prefill path has multiple coupled stages: scheduler admission, page-aligned batch planning, tensor splitting, direct cache writes, index top-k, MLA reuse, and disaggregated KV handoff. Add a default-off, rate-limited debug channel so production ETE runs can identify where batching or metadata semantics diverge without permanently increasing hot-path log volume.
Constraint: Logs must be default-off and rate-limited because these paths execute per-rank and often per-layer.
Rejected: Always-on INFO logs | would flood logs and add CPU overhead during normal prefill.
Rejected: Only scheduler-side logging | insufficient to distinguish planner, index, MLA, and transfer handoff failures.
Confidence: medium
Scope-risk: moderate
Directive: Keep bs>1 debug evidence env-gated; do not add unconditional per-layer or per-token logs in these paths.
Tested: Local py_compile for touched files
Tested: git diff --check
Tested: Remote py_compile and targeted NSA CP utility tests: 5 passed
Not-tested: Full ETE correctness with debug disabled
CP shared KV now keeps explicit L1 and host free-room targets so pressure is handled by planned eviction instead of repeated capacity-edge retries. The host allocator gains contiguous-preferred page reservation, L1 owner-lane allocation prefers contiguous physical pages, and CP HiCache metadata preserves pending backup safety for page-granular radix updates. Mooncake transfer stats and allocator microbenchmarks are included to make the remaining transfer bottlenecks measurable rather than inferred.
Constraint: CP shared KV uses decode CP size 1 with all prefill CP ranks participating in transfer, so L1/L2 cache residency must remain page-granular and avoid extra collectives.\nConstraint: Production HiCache can be hundreds of GB, so allocator metadata overhead must be visible before enabling aggressive contiguous allocation broadly.\nRejected: Evict only the exact deficit | this keeps the cache at the cliff and causes repeated evict/allocate pressure.\nRejected: Rely on allocator scans alone for contiguity | remote microbenchmarks show fragmented 220GB-equivalent host metadata can make contiguous-preferred scans multi-ms.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not increase L1/L2 free-room defaults or add new CP collectives without ETE evidence and transfer/allocator measurements.\nTested: python -m py_compile on touched runtime/test/benchmark files.\nTested: PYTHONPATH=. python -m pytest -q test/registered/unit/benchmark/test_cp_hicache_allocator_bench.py => 4 passed, 1 warning.\nTested: Remote g0034 log /mnt/beegfs/cjy/log/sglang_cp_hicache_20260601_233723.log shows active prefill process with L1/L2 free-room args, 702 HTTP 200 chat completions, 6272 prefill batches, and no fatal scheduler traceback in latest scan.\nTested: User-reported L1/L2 cache ETE validation passed on remote run.\nNot-tested: Full local pytest suite; local environment is missing several runtime dependencies.\nNot-tested: CUDA allocator microbenchmark during active production prefill process.\nNot-tested: Mooncake straggler fix; stats show transfer tail latency remains a separate bottleneck.
CP shared KV and HiCache now keep page-aligned physical ownership while preserving valid-token radix semantics. Repeated tiny EAGLE exact hits free duplicate tail pages instead of leaking one allocator page, owner-lane load-back uses page-vector admission/eviction, and single-DP idle schedulers avoid entering an unnecessary MLP-sync collective.
The commit also records the current page-aligned cache contract and adds gated decode-side EAGLE accept diagnostics so future accept-length collapses can be tied to draft KV/state transfer evidence instead of more prefill cache speculation.
Constraint: CP HiCache allocator ownership is page-granular while radix matching remains valid-token based.
Constraint: New diagnostics must be gated and must not alter normal EAGLE, transfer, or cache behavior.
Rejected: Padding short requests to cp_size or 2*cp_size pages | wastes KV capacity and still hides valid-tail lifecycle bugs.
Rejected: Adding more unconditional collectives to prove CP consistency | hot-path collectives previously caused severe performance risk.
Confidence: medium
Scope-risk: broad
Directive: Do not reintroduce silent fallback for CP shared KV/HiCache paths; warning-level fallback or fail-fast is intentional.
Tested: git diff --check
Tested: local py_compile for all modified Python files
Tested: remote g0034 container py_compile for modified Python/test files
Tested: remote g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/managers/test_scheduler_dp_attn_mixin.py => 114 passed, 5 warnings, 2 subtests passed
Not-tested: full ETE traffic rerun after this commit
Not-tested: CUDA/TAI kernel benchmark coverage for all production shapes
Async MLA/index prefetch is a scheduling optimization, not the correctness contract for target current reuse. Tiny cache-hit suffixes can skip async prefetcher creation while target partial-current reuse still composes page-slot prefix materialization with current KV rows synchronously. CP HiCache radix/device accounting now treats retained valid-tail pages as physical page spans so allocator state stays consistent when logical cache keys are shorter than the retained page.
Constraint: CP shared KV ownership and HiCache residency are page-granular while request-visible cache lengths remain valid-token lengths.
Constraint: Async prefetch can hang or regress on large-prefix tiny-extend traffic and must not be required for current reuse.
Rejected: Treat missing prefetcher as fail-fast for target partial-current reuse | disabled useful current reuse and broke tiny-prefix/tiny-suffix traffic.
Rejected: Keep async prefetcher object with synchronous consume mode | conflates prefetch object existence with current-layer correctness and hides fallback semantics.
Confidence: medium
Scope-risk: moderate
Directive: Do not make current-only or target partial-current reuse depend on MLA/index prefetcher creation; prefetcher objects mean async next-layer work exists.
Tested: Remote g0034 container py_compile for touched modules.
Tested: Remote g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 73 passed.
Tested: Remote g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -> 90 passed.
Not-tested: Latest full ETE traffic run with GLM-5.1 CP HiCache after this commit.
Not-tested: CUDA kernel-level performance impact of synchronous no-prefetch partial-current compose.
Prefix materialize prefetch has fixed launch and coordination overhead, so tiny cache-hit prefixes should stay on the simpler synchronous path. The default threshold is now expressed as a cached import-time token threshold while the existing page override remains available for workload-specific tuning.
Both MLA and index prefix prefetchers use the same page-size-aware threshold, so CP shared-KV prefix prefetch is enabled only when the prefix covers at least one page per CP lane and at least the configured token threshold.
Constraint: The default token threshold must be read once to avoid hot-path env lookups
Constraint: Existing page threshold override must continue to disable or force the gate
Rejected: Hard-code 1024 tokens | tuning needs to be possible without code changes
Rejected: Read token threshold on every maybe_create call | unnecessary hot-path env lookup
Confidence: medium
Scope-risk: narrow
Directive: Keep MLA and index prefetch threshold logic shared; do not let their gates diverge
Tested: Not rerun per user request; prior py_compile, diff check, and isolated env-cache check had passed before commit
Not-tested: full pytest in local environment due missing optional dependencies/kernel import constraints
Not-tested: GLM5 E2E throughput after this threshold change
CP shared KV and HiCache now use owner-lane metadata as the
authoritative capacity view for host write admission and GPU load-back
planning. This removes the debug scalar capacity env and keeps CP load-back
from relying on a rank-wide scalar collective when per-owner availability is
already known. The load-back planner also accounts for evicting child leaves
that unlock ancestor device residency, which fixes small lane deficits despite
large aggregate evictable capacity.
The commit also adds gated CPU timing logs for CP shared-KV MLA/index
prefetch and a CUDA microbenchmark for comparing dense all-reduce with
owner-packed all-gather layouts. The timing logs are intentionally behind the
existing MLA prefetch log env and should not be enabled for throughput
measurements.
Constraint: CP shared KV owner lanes require target/draft capacity decisions to preserve page_owners rather than total-token scalars
Constraint: CUDA collective benchmarks must run on target GPU hosts, not locally
Rejected: Keep SGLANG_CP_HICACHE_CAPACITY_DEBUG observer env | owner-lane admission now replaces that scalar debug path
Rejected: Add a silent scalar-allreduce fallback | unexpected owner-lane mismatch should fail fast or log loudly
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce CP capacity collectives on the scheduler hot path without proving the owner-lane metadata is insufficient
Directive: Disable SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH for end-to-end performance runs; it is diagnostic and high-volume
Tested: git diff --check
Tested: python -m py_compile on changed runtime/test/benchmark Python files
Tested: remote pytest -q test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py (81 passed, 5 warnings)
Not-tested: CUDA benchmark benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py
Not-tested: full GLM5 E2E throughput after this commit
CP HiCache write-through under shared KV was still using rank-wide collectives to decide host reservation eviction, and per-layer backup registration could be bypassed before the final forward boundary. This moves backup registration to the final run_batch pre-forward boundary, forwards it through SessionAwareCache, exposes fallback paths as explicit warnings, and introduces deterministic owner-lane capacity planning for CP host reservation.
Constraint: CP shared-KV ranks must keep target and draft host reservations owner-lane consistent without adding hot-path collective synchronization
Constraint: Remote CUDA validation must run in the g0034 container, not locally
Rejected: Keep reserve_slots_max all_reduce as the default admission path | observed reserve collectives reaching double-digit and occasional 100ms+ latency
Rejected: Silent post-forward catch-up backup | hides when per-layer forward-overlap backup is not actually active
Confidence: medium
Scope-risk: broad
Directive: Do not reintroduce CP HiCache hot-path collectives without a measured mismatch case and explicit fallback warning
Tested: py_compile for modified Python modules and CP HiCache metadata test file in remote g0034 container
Tested: python3 -m pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py -q in remote g0034 container (75 passed, 5 warnings)
Tested: git diff --check HEAD~1..HEAD
Not-tested: Local pytest blocked by missing pybase64 in the local environment
Not-tested: Full CP HiCache + MTP E2E after the no-collective reservation change
Co-authored-by: OmX <omx@oh-my-codex.dev>
The pure-Python convert_to_bigram_key list comprehension runs on every
cache_finished_req / cache_unfinished_req of every EAGLE radix variant
(radix_cache, hiradix_cache, swa_radix_cache), with token-list lengths
that scale with prompt + output. Scheduler profiles consistently flag it
as the largest mem_cache-side CPU hotspot.
This commit wires sglang.srt.mem_cache.utils.convert_to_bigram_key to
tai_kernel.radix.convert_to_bigram_key when the extension is importable,
falling back to the pure-Python implementation otherwise. The tai-kernel
path uses a pybind11 module that calls the CPython C API directly
(PyTuple_New + PyTuple_SET_ITEM with ref-stealing) rather than rebuilding
the list comprehension's bytecode-level tuple allocations. Measured 1.4x
at n=131k and up to 2.5x for n=1k on g0034 Python 3.11; allocator-bound
at large n because CPython's 2-tuple freelist already amortises the
construction. The int64-packing follow-up that bypasses tuple allocation
entirely is parked as a separate work item.
Runtime safety:
- The dispatcher catches any first-call JIT compile / runtime failure,
logs once, and falls through to the pure-Python path for the rest of
the process — JIT failures must degrade rather than crash a serving
loop.
- SGLANG_DISABLE_TAI_BIGRAM forces the Python path for bisecting.
Constraint: Output must be a real Python List[Tuple[int, int]] because
downstream radix dicts use the tuples as hashable keys.
Rejected: Pre-allocated tuple slab pool | CPython's per-interpreter 2-tuple
freelist already serves this case, and we cannot recycle tuples that
become radix-tree keys without changing the consumer.
Rejected: int64-packed keys this round | requires changes to RadixKey,
get_child_key_fn, key_match_fn, and EAGLE bigram detection; deserves its
own plan.
Confidence: high
Scope-risk: low
Directive: Keep _python_convert_to_bigram_key reachable; if the tai-kernel
path is ever removed, the EAGLE radix cache must continue to work
unchanged.
Tested: tai-kernel side validated on the cluster
(``python benchmark/radix/benchmark_convert_to_bigram_key.py --check``
prints byte-exact correctness then 1.4-2.5x speedup across sizes
128..131072 on g0034). The new
``test/registered/unit/mem_cache/test_convert_to_bigram_key.py``
exercises both dispatch paths via SGLANG_DISABLE_TAI_BIGRAM patching.
Not-tested: End-to-end EAGLE serving accuracy + scheduler-time delta on
a real workload.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Prefill CP only needs the local hidden shard for DeepSeek NextN draft extend. The change adds a draft shared-KV path that captures target hidden locally, feeds only the CP-local slice into the draft model, and keeps draft KV writes/transfers on the same shared logical-to-physical page mapping as target KV.\n\nDebug logs are gated behind SGLANG_CP_DRAFT_SHARED_KV_DEBUG and cover scheduler pool selection, KV manager buffer registration, local physical writes, prefill sender filtering, transfer pages, and decode commit metadata so ETE runs can prove draft KV is sharded rather than full-concatenated on a prefill rank.\n\nConstraint: Prefill runs CP while decode remains DP, so prefill must avoid full hidden/KV materialization but decode still receives full logical KV pages.\nRejected: Keep draft extend on full hidden state | preserves correctness but wastes prefill memory and defeats CP shared-KV intent.\nRejected: Transfer draft KV with a separate mapping | target and draft pools share req_to_token logical indices, so duplicating mapping adds risk without benefit.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not remove the debug logs until ETE evidence confirms draft MLA/index writes and transfer pages are CP-sharded on all ranks.\nTested: Remote compileall for changed CP draft, transfer, scheduler, NSA index, MLA write, and EAGLE files.\nNot-tested: Full GLM-5 EAGLE ETE with SGLANG_CP_DRAFT_SHARED_KV_DEBUG=1 after this logging addition; local pytest intentionally not run.
Phase8 prefetch was starting local materialization and the CP all-reduce from the same late hook, after rank-skewed MQA/topk/materialize work had already separated CP ranks. This commit separates local prefix materialization from reduce launch, starts next-layer prefetch earlier in MLA prepare, and explicitly launches pending MLA/index reduces at attention boundaries so the collectives can overlap the intended window instead of drifting behind skewed ranks.
Constraint: Prefetch remains optional and gated by the existing CP shared-KV prefetch environment controls.
Constraint: Profiling required per-source materialize all-reduce ranges, so NVTX labels are gated behind SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX.
Rejected: Launch all-reduce immediately in start_next_layer_prefix | this reproduced the old skew-sensitive timing and could put prefetch behind current-layer work.
Rejected: Remove later hooks entirely | they are still needed as fallbacks when the early MLA prepare hook is bypassed.
Confidence: medium
Scope-risk: moderate
Directive: Preserve the consume-time fallback launch; otherwise missed launch paths silently lose correctness or overlap.
Tested: Remote g0034 docker py_compile for changed SGLang files; python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (50 passed).
Not-tested: Full GLM5 multi-node prefill/decode profile after this exact commit.
The CP shared-KV path now has a gated tai-kernel replacement for NSA index
K/scale plus MQA range preparation, and Phase8 prefetch can skip tiny prefixes
that do not cover all CP lanes. The Phase9 plan documents the next scheduler
work for overlapping CP communication with peer-request attention windows.
Temporary diagnostic logs added while validating prefetch ownership and fused
index prepare routing were removed before committing so the runtime path does
not add log-only synchronization, log counters, or shape-reporting overhead.
Constraint: Production profiling showed small per-request CPU/GPU overhead from diagnostic logging and sync-prone debug counters.
Rejected: Keep fused-index prepare fallback/used logs behind a new env var | it leaves another runtime branch and logging surface for a path that should be benchmarked with profiler evidence instead.
Rejected: Keep owned page-count prefetch logs | they require sync-prone tensor reductions and were only useful for one-off diagnosis.
Confidence: medium
Scope-risk: moderate
Directive: Reintroduce CP shared-KV diagnostics only behind explicit debug paths, and avoid .item()/shape-heavy logging in hot prefill paths.
Tested: git diff --check for staged sglang-dev changes.
Tested: AST parse for environ.py, cp_shared_kv_prefetch.py, cp_shared_kv_runtime.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py.
Not-tested: Full unit test suite.
Not-tested: Multi-node GLM5 prefill/decode/router runtime after this exact commit.
The shared-KV prefill path now optionally calls tai_kernel.nsa_prefill.fused_store_mla_kv before falling back to logical_locs_to_physical plus set_mla_kv_buffer. The fast path supports packed FP8 and BF16/FP16 direct KV buffers, while debug mode and kernel failures still preserve the existing fallback behavior. Success logging was removed after path verification because per-layer/per-rank logs are too noisy in normal server runs.
Constraint: Runtime must remain safe when tai-kernel is absent or debug checks are enabled
Rejected: Keep success logs permanently | floods prefill logs once every rank/layer starts using the fast path
Confidence: high
Scope-risk: moderate
Directive: Keep fallback warnings; do not re-add per-layer success logs outside explicit debug instrumentation
Tested: g0034 container python -m py_compile python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
Tested: g0034 container PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (40 passed)
Not-tested: Full multi-node PD server throughput after log removal
CP shared KV now avoids the PyTorch sort/search remap for the single-request current-only path by deriving compact rows from page-level inverse mapping. The same change keeps sort NVTX attribution gated and splits high-frequency MoE sort markers behind a separate env var so profiling does not perturb normal runs.
Decode-side disaggregation prealloc also avoids rebuilding large token index tensors and records finer allocation timing, while compute-owner allocation/free tests cover the shared-KV page-lane behavior.
Constraint: The runtime tree used for validation is the remote /sgl-workspace/sglang-tai mount, which is not itself a Git repository, so these tracked files were synchronized into the local repo before commit.
Rejected: Keep torch.sort/searchsorted for current remap | it emits ATen/CCCL radixSortKVInPlace kernels in the attention hot path.
Rejected: Enable MoE sort NVTX under the generic sort env | the MoE preprocess sort is too frequent and can make profiling look like a hang.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce token-level torch.sort/searchsorted in CP shared-KV current remap without profiling the attention hot path under Nsight.
Tested: Remote container py_compile for modified runtime files; git diff --cached --check.
Not-tested: Full multi-node GLM5 PD throughput/profile rerun after the page-inverse current remap.
Phase8 only gains useful overlap when the next-layer MLA prefix prefetch is allowed to run until the next layer actually consumes the prefetched buffer. The old wait-after-attention switch let runtime configuration collapse the optimization back into current-layer tail latency, so the prefetch path now has one wait policy and the documentation records the implemented behavior.
Constraint: Phase8 should keep the production environment surface minimal while preserving the existing enable and debug-log knobs
Rejected: SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION | it reintroduced current-layer synchronous waiting and made profiling behavior depend on a nonessential policy knob
Confidence: medium
Scope-risk: narrow
Directive: Do not add another Phase8 wait policy knob without first proving the added policy improves end-to-end prefill latency under CP shared KV
Tested: Python AST parse for touched Python files
Tested: git diff --check
Not-tested: Full pytest and remote server integration were not run in this commit
Shared CP KV materialization remained on the critical path for cached
NSA/MLA prefill batches. This change introduces a one-layer-ahead
prefetcher that materializes the cached prefix for the next layer on a
separate CUDA stream and consumes it when that layer reaches attention.
The prefetch path keeps the existing dense page-table semantics, defers
waiting until the prefetched buffer is actually consumed, and uses the
TAI optimized materialize/remap helpers when enabled before falling back
to the torch implementation.
The implementation is intentionally gated by environment variables and
keeps layer-2-only probe logging for functional confirmation without
making normal profiling noisy.
Constraint: Prefill CP shared KV must preserve existing page-table and dense KV semantics for NSA paged topk attention
Constraint: The production performance path requires SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 and logging disabled
Rejected: Wait immediately after the producer layer attention | this truncated the overlap window and hid less work
Rejected: Torch-only prefetch materialize | it bypassed the optimized TAI materialize/remap path and could erase the expected win
Confidence: medium
Scope-risk: moderate
Directive: Do not evaluate Phase8 throughput with SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=1; use it only to confirm create/start/consume_hit behavior
Tested: Local AST parse for modified Python files
Tested: Local git diff --check
Tested: Remote g0034 container AST parse for modified files under /sgl-workspace/sglang-tai
Tested: Remote g0034 container pytest target covering Phase8 log env, TAI range materialize, optimized slot inverse/remap, and existing token TAI path
Not-tested: Full prefill/decode/router throughput after the TAI prefetch-path fix
Shared KV now relies on page-aligned CP metadata and compute-owner page allocation so persistent MLA KV and NSA index shards can be written by the rank that computed them. The compatibility read path keeps the dense full-view contract for existing topk and attention kernels, but removes duplicated prev/next index materialize, adds optional tai materialize integration, and tightens tests/docs around the fallback boundaries.
Constraint: Decode remains non-CP while prefill CP owns the shared-KV changes
Constraint: Existing attention/topk kernels still expect dense full-view KV/index inputs
Rejected: Change attention kernels to read owner-sharded KV directly | larger semantic change reserved for later phases
Rejected: Merge index K/scale storage with MLA KV storage | would couple topk and attention cache lifecycles before materialize overhead is isolated
Confidence: medium
Scope-risk: broad
Directive: Do not remove fallback logging or debug-gated assertions without reproducing long-context chunked/radix-hit paths
Tested: git diff --check --cached
Not-tested: Local pytest/runtime server verification not run in this commit step per current workflow constraints
Phase 2 shared KV keeps persistent KV/index sharded across CP ranks but
uses a full-view compatibility layer before NSA topk and MLA attention.
For current-only prefill chunks, the current KV/index tensors have already
been CP all-gathered and reranged before being written to the sharded
persistent pool. This change adds a guarded current-reuse path that remaps
logical current locs to compact tensor rows and skips the shared-KV
materialize path for current-only MLA and NSA indexer reads.
Constraint: Existing NSA/MLA kernels still consume full-view/compact page tables; history and mixed current/history batches must keep the Phase 2 fallback.
Rejected: Make all history attention shard-aware in this patch | that requires global topk merge and distributed sparse attention and belongs to a later phase.
Confidence: medium
Scope-risk: moderate
Reversibility: clean
Directive: Do not remove the Phase 2 fallback until mixed/history shared-KV paths have correctness and performance coverage.
Tested: python -m py_compile on modified Python files
Tested: git diff --check on staged modified files
Not-tested: local pytest collection is blocked by missing pybase64 in this environment.
Not-tested: full long-context chunked prefill/decode performance in this commit step.
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer.
Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary.
Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization
Constraint: Decode side remains non-CP and receives full KV through Mooncake
Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts
Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode
Confidence: medium
Scope-risk: broad
Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode
Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics
Tested: Remote py_compile for shared-KV touched Python files in g0034 container
Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected
Not-tested: Full GLM5 multi-node throughput/regression run after final doc update
Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
laoyao0822
leaveletandClaude Fable 5