Port the DeepGEMM "deprecate-from-sgl-kernel, separate sgl-deep-gemm wheel"
migration (upstream PR #24268 and follow-ons) so our code runs against the
torch-2.11 dev-cu13 image, which ships DeepGEMM as the separate sgl-deep-gemm
0.1.2 wheel (import name still `deep_gemm`). Verified against upstream/main HEAD,
NOT the introducing PR — the wheel API drifted between 0.0.1 and 0.1.2.
Ports (all verified against HEAD = wheel 0.1.2):
- nsa_indexer/nsa_backend: paged-MQA context_lens must be (N_total, 1). The
0.0.1 form (batch_size, next_n) DEADLOCKS fp8_paged_mqa_logits on next_n>=2
(our EAGLE deploy uses next_n=4 on SM90/H200, which does not take the SM100-
only DG-native broadcast path). Matches HEAD's _to_2d_context_lens.
- compile_utils warmup: hasattr-guard the dropped get/set_compile_mode API;
pass m_indices positionally to m_grouped_fp8_gemm_nt_contiguous.
- fp8_utils.transform_scale_ue8m0: restore TMA-aligned stride when the DLPack
round-trip collapses a size-1 trailing dim.
- moe_runner/deep_gemm: guard the SBO masked-gemm return unpack when overlap is
inactive (#26839) — reachable via our --enable-single-batch-overlap.
- entrypoint: enable DeepGEMM PDL by default, hasattr-guarded (#23979).
- engine: require sglang-kernel >= 0.4.3 (first sgl-deep-gemm-era kernel); the
migrated (N_total,1) paths would misbehave on the old bundled DeepGEMM.
Verified-unchanged symbols left as-is: fp8_mqa_logits, fp8_gemm_nt, bf16_gemm_*,
get_mk_alignment_for_contiguous_layout, transform_sf_into_required_layout,
get/set_num_sms, the masked-gemm signature. Runtime validation pending the
torch-2.11 image (tai-kernel rebuild + harness).
Refs: WI-2026-06-07-001
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 32818784d332b5c48faf8027247cd4cd9a0f48cd)
The GSM8K/cache-hit debugging pass added request-correlation logs across scheduler, radix, HiCache, Mooncake, and prefill handoff. The root cause work has moved to request-slot remap semantics, so those high-cardinality traces are no longer needed in the runtime diff.
This removes the temporary tracing helpers and call sites while leaving existing fail-fast checks, fallback warnings, and explicit debug/timing infrastructure that is still part of normal CP shared-KV diagnostics.
Constraint: Production CP hot paths should not carry investigation-only request signatures or transfer summaries
Rejected: Keep all debug logs gated by env | even gated logs increase maintenance surface and encourage stale diagnosis paths
Confidence: high
Scope-risk: moderate
Directive: Reintroduce request-correlation logs only as a narrow opt-in probe with a planned removal point
Tested: Local py_compile for touched runtime files; remote py_compile; remote pytest test_cp_shared_kv_runtime.py, test_nsa_cp_utils.py, test_cp_shared_kv_layout.py => 263 passed, 5 warnings, 2 subtests passed
Not-tested: Long-running ETE log-volume comparison after removal
Warm-cache bs>1 requests can carry duplicate logical page ids across request rows. A batch-global page inverse collapses those request slots and can alias current/prefix KV across requests.
The runtime now carries compact row-scoped sorted page descriptors and remaps flattened logical locs with request row ids where needed, while retaining the legacy global inverse for no-row-context paths.
Constraint: Avoid the rejected dense [batch, logical_page_capacity] inverse because bs up to 10 and long contexts make that memory cost unacceptable
Rejected: Keep global page_inverse for bs>1 duplicate pages | it is lossy and matches the GSM8K warm-cache corruption shape
Rejected: Allocate page_inverse_by_row | correctness-safe but too much GPU memory for production
Confidence: medium
Scope-risk: moderate
Directive: Any future TAI materialize fast path for bs>1 duplicate pages must consume row-scoped descriptors or an equivalent request-row key
Tested: Local py_compile for touched runtime files; remote py_compile; remote pytest test_cp_shared_kv_runtime.py, test_nsa_cp_utils.py, test_cp_shared_kv_layout.py => 263 passed, 5 warnings, 2 subtests passed
Not-tested: Full GSM8K ETE after this cleanup pass
Warm-cache GSM8K failures needed request-to-log correlation across scheduler, prefill transfer, and Mooncake CP filtering. The added diagnostics are gated by existing CP shared-KV debug envs and include bounded token signatures plus transfer page summaries so future debugging can identify whether a failed request hit L1, L2, or transfer truncation paths.\n\nThe findings document records the ruled-out hypotheses and the RAGGED current-row contract failure, preventing repeated log archaeology after context compaction.\n\nConstraint: Production hot paths must not emit these logs unless SGLANG_CP_SHARED_KV_BS_GT1_DEBUG or existing CP shared-KV debug gates are enabled.\nRejected: Add a new debug env | reuse of existing bs>1 debug gate avoids more runtime switches.\nRejected: Store full token ids in logs | bounded signatures are enough for correlation without huge logs.\nConfidence: medium\nScope-risk: narrow\nDirective: Keep these diagnostics gated and bounded; do not convert them to unconditional INFO logs.\nTested: Remote container py_compile as part of the RAGGED fix validation.\nNot-tested: Long-running production log-volume impact with debug enabled.
FP8 flashmla_sparse uses flattened RAGGED page tables that include both cached prefix and the just-computed current suffix. The old cache-hit path materialized the whole flattened range from persistent KV, which could read current rows through the wrong contract under CP shared-KV and compute padding.\n\nThis change makes the RAGGED path use the page-slot partial-current compose contract: prefix pages are materialized from cache slots while current rows are sourced from fresh k/k_rope and packed for FP8 when needed. A new helper accepts the actual current-row contracts seen by attention code: already-local valid rows, CP-local compute-padded rows, or unsplit global valid rows.\n\nConstraint: CP shared-KV stores and consumes cache at page granularity, while attention current rows may be valid-token tensors rather than cache-write local compute rows.\nRejected: Full materialize prefix+current from persistent KV | it can read current suffix through stale or unordered persistent cache state.\nRejected: Reuse select_cp_local_valid_rows_for_cache_write directly | it only accepts CP-local compute-padded rows and killed prefill when RAGGED supplied global valid rows.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not route RAGGED FP8 cache-hit current suffix through full materialize without proving current write ordering and row-contract compatibility.\nTested: Local py_compile for touched runtime/test files.\nTested: Remote container pytest for RAGGED current compose and compute-padding global-current row selection.\nNot-tested: Full GSM8K warm-cache ETE after restart.
The cache-hit accuracy drop is not isolated to one obvious kernel boundary, so the scheduler/radix/HiCache path needs request-correlated evidence. Add default-off, rate-limited CP shared-KV debug logs for prefix matching, valid-tail flooring, pending backup attach/commit, load-back planning, and unfinished-request cache insertion. The investigation notes capture current GSM8K pass1/pass2 evidence and rejected interpretations to avoid repeated log archaeology.
Constraint: Logs must be default-off and bounded because these paths are scheduler/control hot paths.
Rejected: Always-on INFO tracing | would distort the exact latency and cache-hit behavior under investigation.
Confidence: medium
Scope-risk: moderate
Directive: Remove or keep env-gated only after the GSM8K cache-hit root cause is closed; do not leave unbounded hot-path logs.
Tested: python -m py_compile on changed runtime files.
Not-tested: Local pytest blocked before collection by missing orjson dependency; fresh restarted GSM8K verification still pending.
(cherry picked from commit bfc6d1473dc2a5d72bc3a8d6fca1e2429537be0e)
CP load-back already handles exact owner-lane deficits before admitting L2->L1 reloads. Treating free-room target misses as synchronous evict work adds avoidable latency to cache-hit requests even when the current load can fit. The load-back path now records the free-room miss as advisory and leaves proactive replenishment to non-critical eviction paths.
Constraint: GSM8K/cache-hit runs are sensitive to CPU/control-path latency during L2->L1 reload.
Rejected: Evict to free-room on every load-back | preserves room but puts synchronous eviction on the latency-sensitive cache-hit path.
Confidence: medium
Scope-risk: moderate
Directive: Exact owner-lane deficits remain mandatory; free-room targets must not become blocking in load-back without ETE latency evidence.
Tested: python -m py_compile on changed runtime files.
Not-tested: ETE latency after this isolated change not rerun in this commit step.
(cherry picked from commit 2f42acba0d748d2d747f0531963267d86c483cde)
Batch in-seq CP rerange has two lengths under compute padding: source payload length includes synthetic rows used only to keep CP compute well-shaped, while output length must expose only valid request rows. The torch fallback now computes rank-major source offsets from compute splits and output placement from valid splits.
Constraint: Tiny extend batching can add compute-padding rows that must not appear in restored valid token order.
Rejected: Use valid splits for source offsets | following requests on the same rank are shifted when a previous request has padded mirror rows.
Confidence: medium
Scope-risk: narrow
Directive: Batch rerange implementations must distinguish source compute splits from valid output splits.
Tested: python -m py_compile on changed runtime files.
Not-tested: Local pytest blocked before collection by missing orjson dependency.
(cherry picked from commit 31e741477503caa52f3a23acdc1286f46079043c)
Decode queue compaction receives req_to_token rows after the prefill side has already populated cached prefix slots. Cache-hit requests therefore need the extend/suffix slice, not the leading prefix slice, when building the prebuilt transfer chunk.
Constraint: Prefill/decode disaggregation shares req_to_token rows across cached prefix and new suffix positions.
Rejected: Keep slicing from zero | cache-hit requests would copy prefix KV locs into the prebuilt suffix chunk.
Confidence: medium
Scope-risk: narrow
Directive: Do not change prepare_for_prebuilt slicing without testing cache-hit req_to_token layouts.
Tested: python -m py_compile on changed runtime files.
Not-tested: Local pytest blocked before collection by missing orjson dependency.
(cherry picked from commit 416112b617fabe71e8cff7484794af73f3e84440)
The transfer worker iterates every non-dummy decode info for a room and calls
per_layer_mgr.finish() once per info, but register_per_layer_transfer registers
exactly one context per room/chunk (built for one info's dst_kv_indices). This is
only sound when there is exactly one non-dummy info (required_dst_info_num == 1).
With decode attn_tp < prefill attn_tp a single prefill rank holds >1 non-dummy
infos; finishing once-per-info would over-pop chunk contexts and under-deliver KV
to the other infos. Make the assumption explicit: register only when there is one
non-dummy info, otherwise fall back to the monolithic post-forward transfer (which
fans out to all infos correctly). Found by an independent first-principles audit.
Adds TestRegisterGuardSingleInfo (2-info fallback, 1-info register, all-dummy
fallback) exercising the real MooncakeKVManager.register_per_layer_transfer.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The per-layer KV transfer registration hardcoded chunk_page_start=0 when
filtering CP-shared-KV owned pages. The CP filter's second return (`positions`)
are absolute full-sequence page positions built from chunk_page_start, and the
transfer indexes the FULL-request dst_kv_indices by those absolute positions
(mirroring the monolithic send(), which passes chunk_page_start=index_slice.start
— the cumulative page offset). With start=0, chunk N>0's positions were
chunk-local, so its KV was written onto chunk 0's decode pages, corrupting the
decode output. Non-chunked requests (single chunk, start=0) were unaffected,
matching the observed symptom (non-chunked byte-identical, chunked garbage).
Fix: chunk_page_start = chunk_key // page_size, where chunk_key is the chunk's
start_send_idx (page-aligned), making it exactly the monolithic index_slice.start.
Verified: opus first-principles code audit; empirical mapping-invariant on the
deployed modules (per-chunk == whole-request for all 8 CP ranks; old start=0
sends chunk1 to chunk0's dst); 2 new regression tests (TestChunkedDstMapping).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Per code review (HiCache load is layer-by-layer & correctly ordered into the compute
stream before the backup hook; current-reuse is a within-forward read that doesn't
rewrite pool pages): unify registration to the EXACT range this forward's
send_kv_chunk transmits — req_to_token[start_send_idx:end_idx], page-floored for a
non-last chunk. Non-chunked = one full range; chunked = one range per chunk. Drop
the is_chunked/start_send_idx skip.
To avoid the review's collision risk (chunk N still finishing when chunk N+1
registers), the manager keys contexts per (room, start_send_idx): _active[room] is a
FIFO deque of (chunk_key, ctx); register dedups the same chunk but appends a new one;
finish(room) pops the FRONT (chunks finish in send order — no chunk key needed in the
transfer_worker); drop drains all the room's chunks; on_layer_end enqueues for all
active chunk contexts (per-ctx note_enqueued dedup keeps each chunk's own events).
28 unit tests pass incl. chunked FIFO + per-chunk dedup.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
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>
E2E at high cache-hit + concurrency exposed a vicious cycle: a context stays active
until finish, but the notifier re-fires every layer on every subsequent forward and
note_enqueued counted each, so _enqueued (the finish target) grew by the layer count
each forward (target=3042-5538 observed) faster than 4 workers can drain -> finish
never reaches it -> 30s timeout -> context stays active -> repeat. TTFT p99 = 116s.
Fix: note_enqueued(layer_id) dedups per layer (target caps at the layer count); the
first fire for a layer is from the request's own forward so its event is correct.
Also guard register() against overwriting an active room (was leaking + re-registering,
registered=917 for ~100 reqs).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The premise holds at realistic context: at ~8k tokens + high cache-hit the KV
transfer is ~30% of TTFT (712ms), sequential after the forward. Lever A was scoped
to prefix==0 so it never activated there. Broaden it to cover the FULL sequence
(cached prefix + new tokens): the per-layer backup hook fires after each layer's
full processing, so the HiCache-loaded prefix KV and forward-written new KV are both
final — transferring all of layer L's pages then overlaps the dominant prefix
transfer with the forward. Add a sonnet-bench verify mode for output-equality.
Gated by the flag; correctness verified next.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Log per-request finish breakdown (wait_workers / fallback / wait_rdma) to prove
WHY the per-layer transfer stage (~193ms) doesn't shrink vs the monolithic (~172ms)
— i.e. whether the workers fall behind the forward (submits land late, no overlap)
or the RDMA itself is the cost. INFO-gated.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
E2E diagnostic was precise: "finish TIMEOUT processed=78/79", submit_failed=0 —
the per-layer notifier fires 78x but kv_data_ptrs has 79 layers (the 79th is the
MTP/nextn EAGLE buffer: present in kv_data_ptrs so the monolithic send moves it,
but it doesn't fire the per-layer hook). The old completion required all num_layers,
so it both hung to the timeout AND would silently miss that layer's KV (corruption).
Redesign: gate completion on the ACTUAL enqueued count (note_enqueued), and in
finish() SYNCHRONOUSLY transfer any layers the notifier didn't fire for (KV is fully
written post-forward, no event needed). Net: the per-layer set == kv_data_ptrs,
byte-identical to the monolithic send; robust to any model firing fewer hooks than
KV buffers. The fired layers stay overlapped with the forward.
Unit tests updated (25 pass): fallback transfers the missed layers; submit failures
still report -1.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
E2E (clean run) showed finished ret=-1 at exactly the 30s timeout for every
request: finish() hung because _worker_step swallowed submit_layer exceptions
without counting the layer toward completion (so _processed never reached
num_layers). Fixes:
- submit_layer: try/finally that ALWAYS counts the layer (completion can never
hang on a per-layer error) and LOGS the actual exception.
- PerLayerTransferManager worker_init: torch.cuda.set_device on each worker thread
(likely cause — event.synchronize() needs the device set on these fresh threads,
unlike the transfer_worker where A1's engine call worked).
- finish() logs processed/num_layers on timeout to separate exception-failure from
notifier-undercount.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
E2E caught it: registered=8 (per-layer activated correctly) then the prefill
crashed with "Boolean value of Tensor with no values is ambiguous" — req.prefix_indices
is a tensor and `prefix or []` evaluated its truthiness. Use a None+len check
(safe for None/list/tensor), and wrap each request's registration in try/except so
a setup error degrades to the monolithic transfer instead of crashing the scheduler.
Also guard start_send_idx with getattr.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Complete the lever-A hot-path integration behind SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER:
- PerLayerTransferContext: add num_layers + a completion event so finish() waits
until ALL layers are processed before wait_batch_transfers (never races ahead of
the worker threads and silently drops in-flight layers); times out to FAILURE.
- PerLayerTransferManager: has_room (for the swap) + drop (abort/failure drain so
outstanding RDMA finishes before pages are reclaimed).
- MooncakeKVManager.register_per_layer_transfer: build + register a context before
the forward, reusing send()'s exact CP filter (no re-derivation).
- transfer_worker: when a room is per-layer-active, wait those transfers (finish)
instead of the monolithic send_kvcache -- no double-send; aux/state/completion
unchanged. The skip path drops the context on abort/failure.
- prefill scheduler: _register_per_layer_transfers(batch) before run_batch, scoped
(first impl) to single-forward, no-cached-prefix requests (the notifier transfers
forward-written pages; chunked/cached-prefix are HiCache-loaded -> lever B).
Unit-tested (25 cases). e2e output-equality + TTFT verification next.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add MooncakeKVManager.build_per_layer_context: assembles a PerLayerTransferContext
from the SAME CP-filtered (prefill_kv_indices, dst_kv_indices) the post-forward
transfer uses — so the bytes moved are byte-identical to the monolithic path, and
the CP owner mapping is NOT re-derived (eliminating the #1 correctness risk). It
mirrors the MLA branch of _send_kvcache_generic exactly (get_mla_kv_ptrs_with_pp +
group_concurrent_contiguous + build_layer_blocks, verified set_transfer_blocks-
identical). Returns None for MHA / unregistered decode / empty owned set.
Unit-tested (4 cases): per-layer address correctness + the None guards. Remaining
A3-step3: call this in the send/scheduler flow (register before forward, finish
after, skip the main-KV monolithic send), then output-equality + TTFT verification.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add build_layer_blocks: the pure per-layer transfer-address computation (src/dst
addrs + lengths for layer L's owned page blocks), the core of the context's
get_blocks closure. Mirrors the mooncake set_transfer_blocks math; the page index
lists are identical across layers, so only the per-layer base ptr + item_len
change. Unit-tested (3 cases incl. the cross-layer invariant). 27 per-layer/async
tests green total.
The remaining A3 step assembles get_blocks from the scheduler's per-request data
(transfer_infos dst indices + decode_kv_args_table dst ptrs + out_cache_loc src +
CP owner filter) before run_batch, and reconciles finish() with send_kv_chunk —
the hot-path integration, to be verified by the bitwise-equivalence harness.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Wire PerLayerTransferManager into the prefill bootstrap queue: when
SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER is on, create the manager (real
torch.cuda.Event / current_stream) and register its on_layer_end on the KV pool's
layer_backup_notifiers, exposing it as kv_manager.per_layer_transfer_manager. The
forward's per-layer end hook now reaches the manager. Additive and a no-op until a
request registers a transfer context (A3-step2): on_layer_end returns early when no
contexts are active (verified). Remaining A3 (get_blocks + setup-before-run_batch +
finish reconciled with send_kv_chunk) builds on this.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add PerLayerTransferManager: owns a worker thread-pool + the active per-request
contexts for the current forward batch, registered as a KV-pool notifier.
on_layer_end (forward thread) records ONE CUDA event on the compute stream and
enqueues (ctx, layer, event) per active context; workers do the event-wait +
async submit OFF the forward thread. finish(room) waits the request's transfers.
event_factory/current_stream injected for unit-testability (no CUDA needed).
Unit-tested (5 manager cases, 11 total in test_cp_per_layer_transfer.py):
per-active-ctx enqueue with the event recorded on the stream, worker-step submit +
mark-failed-on-exception, finish pop + idempotency, no-op when idle. The A3
scheduler wiring (notifier registration + setup-before-forward + finish-after +
no-double-send reconcile) is the remaining hot-path step; plan locked in
lever-a-implementation-plan.md.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add PerLayerTransferContext: the per-request coordinator for overlapped per-layer
KV transfer. submit_layer(layer_id, event) waits the layer's CUDA write event (on
a background thread, never the compute stream) before async-submitting that
layer's RDMA via the G1 path, so the transfer never reads a layer before its
write kernel finished — the core correctness invariant for the forward overlap.
finish() waits all accumulated batch_ids; idempotent per layer; fails closed.
Unit-tested (test_cp_per_layer_transfer.py, 6 cases): event-wait-before-submit
ordering, idempotency, empty-layer skip, finish-waits-all, submit-failure stop,
wait-failure propagation. The scheduler/notifier wiring (A2-wiring + A3) builds
on this.
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>
Add batch_transfer_async_submit() (non-blocking submit -> batch_id) and
wait_batch_transfers() (block-until-all, graceful) to MooncakeTransferEngine,
exposing the mooncake async API for the upcoming per-layer overlapped transfer.
Deliberately avoids batch_transfer_*_on_cuda: its CUDA host callback busy-waits
on the stream and calls _exit(1) on transfer failure (verified in the mooncake
source transfer_engine_py.cpp), which is unsafe for production (a decode-side
abort mid-transfer would crash the prefill). The async submit + wait pair lets
per-layer submits pipeline in the RDMA engine, then waits once on a background
thread with graceful failure.
Both bindings raise a clear upgrade error if the engine predates the API and
degrade to -1 on transient errors. Mock-engine unit tested (8 cases).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port the decode->prefill half of upstream sgl-project/sglang #27372 (the worker
skip-Failed guard half landed in cb1a03f0a). In multi-prefill/multi-decode PD, a
decode-initiated abort frees the decode's KV pages back to the allocator, but the
prefill never learns of it (request_status is per-process), so the prefill keeps
RDMA-writing the remaining chunks into pages the decode may have reallocated to a
different live request -> KV corruption. The already-ported worker guard is a no-op
here because nothing sets the prefill room to Failed.
Now the decode receiver sends a 4-field b"ABORT" notification to every prefill peer
(on abort() and on the poll() WaitingForInput timeout, at most once); the prefill
marks the room Failed (so the worker guard fires) and replies b"ABORT_ACK". The
ABORT branch is handled before the unconditional waiting_req_bytes[3] decode in
bootstrap_thread (a 4-field ABORT would otherwise crash the else branch), and
ABORT_ACK before the 3-tuple unpack in the decode thread.
De-entangled from the staging/tracing infra this branch lacks. Component-tested
(test_mooncake_abort_protocol.py: format, send-once guard, abort wiring, ZMQ
round-trip). Cross-process corruption-window closure to be verified by the
multi-P/D PD harness.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port upstream sgl-project/sglang #24416 (staging-free). Our transfer metrics
were computed from the full prompt length, so total_mb / speed_gb_s were
systematically over-reported on every prefix-cache hit and every CP shared-KV
per-rank page filter. Now each sender accumulates the actually-sent KV/state
indices and reports bytes = sent_pages * per-page item bytes via a new
KVTransferMetric returned by get_transfer_metric(); prefill consumes that
instead of estimating from len(origin_input_ids), and skips fake-bootstrap and
dummy-CP-rank senders (which transfer nothing).
Also route convert_to_duration() phase deltas through duration_between(), which
returns 0 when a phase timestamp is uninitialized, fixing nonsensical negative
durations in the time-stats log.
Unit-tested (test_kv_transfer_metrics.py, 10 cases) in the CUDA-13 container.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port the abort-correctness pair from upstream sgl-project/sglang:
- #24522: abort() now calls kv_mgr.update_status(room, Failed) so the
transfer worker and other pollers observe the abort, not only the local
conclude_state.
- #27372 (guard half): the transfer worker skips a chunk whose room has
already failed or been aborted, so it never transfers into KV pages that
may have been reclaimed. This matters more once per-layer async transfer
widens the abort-vs-in-flight window.
The decode->prefill ABORT/ABORT_ACK notification half of #27372 is deferred:
it interleaves with staging/tracing infra this branch lacks and needs the PD
test harness to verify, so porting it now would be unverified guesswork.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Port upstream sgl-project/sglang #23323. Our state_type in [swa, nsa]
transfer path only handled len(prefill_state_indices) < len(dst), and even
then clipped prefill_state_indices (a no-op when prefill is shorter), and
never handled the prefill > dst case. Now clip prefill when it is longer and
clip the dst state indices when it is shorter, so the extra-pool transfer
never reads past prefill_state_indices or writes past dst_state_indices.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Cache-hit GSM8K regressions only appeared after the second pass reused request-specific suffix pages, so this change adds fail-fast transfer validation, masks stale rectangular page-table tails, and extends CUDA/unit coverage across FP8 CP shared-KV write, load, top-k, and materialization paths. The temporary ledger records eliminated hypotheses to prevent re-debugging the same L2 and persistent-cache paths.\n\nConstraint: CP shared KV stores physical pages but scheduler-visible semantics must remain valid-token/page-bounded.\nConstraint: bs>1 FP8 prefill must preserve existing CP shared-KV fast paths without silent fallback.\nRejected: Blame raw HiCache L2 load without tests | L2 KV and index backup/load/materialize roundtrips pass on remote CUDA.\nRejected: Disable current/partial reuse broadly | hides the cache-hit contract regression and costs performance.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not weaken CP shared-KV fail-fast or rectangular-tail masking without rerunning second-pass cache-hit accuracy tests.\nTested: remote CUDA pytest for fused FP8 MLA store, fused persistent index store, L2-loaded FP8 KV materialize, L2-loaded index materialize, ragged top-k offset, TAI batched index MQA prepare.\nTested: local py_compile for touched test files and git diff --check.\nNot-tested: full second-pass GSM8K accuracy after these diagnostic tests; root cause remains under investigation.
Repeated GSM8K/cache-hit traffic can drain a CP owner lane while exact
allocation still succeeds. This changes L1/L2 free-room logic so the trigger
accounts for the pending allocation, then evicts a minimum target chunk or the
exact shortage, whichever is larger. CP load-back also performs best-effort
owner-lane free-room eviction before admitting tiny cache hits that would
otherwise skew one lane.
The commit also adds gated bs>1 prefill timing markers around
recv/bootstrap/batch/inflight polling so future hangs expose the scheduler
boundary instead of only surfacing as a watchdog shutdown. The post-fix repeat
GSM8K failure is recorded as an active regression to continue investigating,
not as old-process noise.
Constraint: Free-room policy must reduce repeated eviction and owner-lane starvation without reserving required+target pages after every allocation
Rejected: Evict to required+target availability | wastes L1/L2 residency under many short cache hits
Rejected: Treat free-room misses as fatal on load-back | exact capacity should remain the strict admission condition
Confidence: medium
Scope-risk: moderate
Directive: Do not remove the gated bs>1 timing markers until repeat-GSM8K hangs have a direct failing boundary
Tested: git diff --check
Tested: python -m py_compile for touched Python files
Tested: remote earlier test_cp_shared_kv_layout.py and test_cp_hicache_metadata.py passed 157 tests after this free-room formula
Not-tested: Two full repeat GSM8K runs after this commit; latest repeat still killed prefill and requires follow-up root-cause work
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.
CP shared-KV bs>1 exposed two separate overhead sources: HiCache load-back could thrash near capacity, and partial-current sync compose could all-reduce row-major page-table gaps between request prefixes. Keep the intended batch-plan path for bs=1, but make host/L1 free-room handling less reactive and teach the MLA/index sync compose path to use exact per-request prefix slot spans instead of one bounding span.\n\nThe exact-span path preserves the single-span IPC fast path for bs=1/single-span cases, while avoiding over-communication for heterogeneous cache-hit batches. The HiCache metadata tests cover host/L1 free-room propagation and load-back batching behavior.\n\nConstraint: bs=1 using CPSharedKVBatchPlan is the expected steady-state path and must not be treated as a regression.\nConstraint: Remote production-like validation runs inside g0034 container /sgl-workspace/sglang-tai.\nRejected: Disable batch-plan for bs=1 | user confirmed this is intended behavior and it would hide the actual bs>1 overhead.\nRejected: Keep one bounding prefix span for bs>1 | row-major page tables can include current/gap slots and inflate per-layer all-reduce work.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not replace exact prefix spans with a single row-major bounding span unless ETE data proves collective launch count dominates gap over-communication.\nTested: g0034 docker py_compile for changed Python/test 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: Full ETE bs>1 throughput after the exact-prefix-span change.\nNot-tested: CUDA kernel benchmark and live traffic replay.
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
Burst decode can sit in disaggregation prealloc/transfer queues while health probes still need an immediate scheduler-alive response, and EAGLE prebuilt metadata must outlive the synthetic prebuilt step until the first real decode consume. The transfer path also needs to send final aux metadata even when there are no new KV pages in the final chunk.
This keeps decode metadata slot ownership narrow instead of cloning EAGLE tensors, adds fail-safe cleanup for prebuilt exceptions/finished prebuilt requests, fixes final empty chunk metadata transfer, and records the investigation ledger for future debugging.
Constraint: bs=1 historically worked, so decode compute/sampling behavior must not be changed without direct evidence.
Rejected: Clone EAGLE metadata tensors on transfer commit | avoids lifetime issues but adds hot-path CPU/GPU memory traffic.
Rejected: Treat prefill AbortReq logs as root cause | current evidence shows they can be downstream of decode/router aborts.
Confidence: medium
Scope-risk: moderate
Directive: Do not move EAGLE metadata slot release earlier than first real decode result processing without proving the H2D/spec_info consume point changed.
Tested: g0034 docker py_compile for touched scheduler/disagg/test files; PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_req_to_token_pool.py test/registered/unit/managers/test_scheduler_health_check.py -> 24 passed
Not-tested: Fresh end-to-end burst traffic after restart; decode node_rank=1 persistent log capture.
Batch-planning metadata can now be attached even for single-request CP prefill, which routed FP8 flashmla_sparse through the batch cp_index path. That path used compact MQA-buffer row bases for score lookup but did not override the final ragged topk coordinate base consumed by attention, so topk indices could point at the wrong KV rows and produce low accept length or meaningless output.\n\nThis keeps ordinary long page tails out of compute padding, reserves compute padding for truly tiny suffixes, and makes cp_index RAGGED topk emit request-ragged offsets while preserving the compact buffer descriptors used for score materialization. The debug ledger records the rejected intermediate diagnoses and the confirmed coordinate-space failure.\n\nConstraint: CP shared-KV cache residency is page-granular, but attention/index compute must not consume synthetic long-tail rows.\nConstraint: FP8 CP prefill uses flashmla_sparse/RAGGED, where fused topk output is consumed directly as attention page_table_1.\nRejected: Disable current reuse or batch planning | would hide the regression and lose the intended bs>1 fast path.\nRejected: Treat all page tails as compute padding | regresses bs=1 semantics and can corrupt query/topk row contracts.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not change cp_index RAGGED topk offset handling without verifying score-buffer row bases and final attention KV coordinate space independently.\nTested: python -m py_compile on touched Python/test files; git diff --check; remote targeted ragged cp_index offset regression test; remote test_nsa_cp_utils.py; remote test_cp_shared_kv_runtime.py; user-reported ETE output recovered after restart.\nNot-tested: Agent-driven full ETE traffic run; broad multi-request bs>1 production soak.
CP shared KV current-only and partial-current paths were previously unsafe because owner-local current rows could be filled into dense page slots without synchronizing the current slots across CP ranks. This preserves the fast path and fixes the missing synchronization at the actual contract boundary: prefix slots remain materialized by the existing IPC/collective path, while freshly produced current slots are composed from current forward tensors and reduced only over the current slot ranges.\n\nThe change also keeps page-tail compute padding explicit, restores current-only index compose, fixes TAI index descriptor dtype consistency, and records the investigation ledger to avoid repeating discarded hypotheses.\n\nConstraint: CP shared KV must preserve current reuse and TAI materialize fast paths; blanket disable is not an acceptable fix.\nConstraint: Prefix slots must not be reduced twice after IPC/prefix materialization.\nRejected: Disable current-only MLA/index reuse | hides the owner-local composition bug and regresses the intended fast path.\nRejected: Disable TAI materialize globally | avoids symptoms without proving the byte/layout contract.\nConfidence: medium\nScope-risk: broad\nDirective: Do not remove current-slot range synchronization unless replacing it with an equivalent owner-aware P2P/IPC gather contract.\nTested: Local py_compile for touched files.\nTested: Remote g0034 py_compile for touched runtime/test files.\nTested: Remote test_cp_shared_kv_runtime.py: 111 passed, 5 warnings, 2 subtests passed.\nTested: Remote targeted current-slot compose tests and direct current-only index compose script.\nNot-tested: Full ETE output quality and decode accept len after restart.\nNot-tested: Full test_nsa_cp_utils.py collection on remote due incomplete installed sgl_kernel import/fake-op environment.
Mooncake previously truncated source prefill pages when the selected decode destination page list was shorter. That made an invalid prefill/decode page mapping continue as an incomplete KV transfer, which can surface later as decode garbage instead of the original mapping error.\n\nThis changes the transfer contract to require exact source/destination page-count equality and records compact page summaries in the fail-fast error. The helper is shared so the page-count contract can be unit-tested without constructing the transfer worker thread.\n\nConstraint: CP shared-KV page ownership requires a one-to-one prefill source page to decode destination page mapping.\nRejected: Keep warning-and-truncate | masks page-map corruption and can silently drop KV.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not restore source-page truncation; fix the mapping producer if this fail-fast triggers.\nTested: git diff --check; python -m py_compile python/sglang/srt/disaggregation/utils.py python/sglang/srt/disaggregation/mooncake/conn.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py\nNot-tested: pytest blocked locally by missing orjson; remote g0034 was unavailable during this pass.
CP shared-KV now uses CP-local current rows consistently across MLA/index current reuse, passes fp8 current-index K through the tai-kernel uint8 ABI, and clears the transient EAGLE CP-local hidden marker after draft capture. The disaggregation bootstrap also fingerprints the runtime source contract so prefill/decode mismatches fail fast instead of silently exchanging incompatible KV metadata.
Constraint: CP shared-KV batch paths flatten current K/V rows in CP-rank-local valid order, not global request order.
Constraint: tai-kernel current-index prepare validates current_index_k as uint8 bytes for fp8 payloads.
Rejected: Keep using global extend offsets for bs>1 current-index reuse | corrupts request-local bases once current_index_kv is CP-local.
Rejected: Infer CP-local EAGLE hidden semantics from tensor shape | static padding and bs>1 can make shape-based inference unsafe.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce forward_batch.out_cache_loc slicing in CP shared-KV current reuse without verifying CP-local owner-lane layout.
Tested: Remote container py_compile for touched runtime/test files.
Tested: Remote PYTHONPATH=python 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/disaggregation/test_common_conn_runtime_fingerprint.py (198 passed, 2 subtests passed).
Not-tested: Full remote ETE traffic after this commit; accept length and garbage-output recovery still require a fresh prefill/decode run.
Co-authored-by: OmX <omx@oh-my-codex.dev>
CP shared-KV bs>1 exposed three distinct padding domains: valid cache rows, CP page-tail compute rows, and MLP-sync flattened static padding. The previous implementation mixed these domains in direct-write and index top-k paths, so real requests failed when q/out_cache_loc lengths matched valid rows while metadata aliases described compute rows.\n\nThis change makes compute split strip only proven flattened static padding, keeps valid cache writes strict except for extend_num_tokens-proven static tails, marks CP-local EAGLE draft hidden state explicitly, and selects NSA index top-k query metadata by the actual q/weight row count.\n\nConstraint: CP shared-KV cache writes must never persist dummy page-tail or MLP static padding rows.\nConstraint: EAGLE draft hidden state can be CP-local before full CP metadata is visible in prepare_mlp_sync_batch.\nRejected: Use compute_padding_enabled as direct-write truncation proof | it silently accepts unknown out_cache_loc tails.\nRejected: Always consume compute q metadata in index top-k | actual q/weights can be valid-only after CP split.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not collapse valid rows, CP compute padding, and MLP static padding into one length condition; use explicit provenance.\nTested: remote py_compile for touched NSA files\nTested: remote targeted CP shared-KV padding/top-k regressions\nTested: remote pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -k 'not test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel' => 228 passed, 1 deselected, 5 warnings, 2 subtests passed\nNot-tested: full ETE replay after the final index top-k fix\nNot-tested: TAI current-index fast path dtype fallback
CP shared-KV prefill batching can present multiple requests to HiCache write-through in one forward. Keeping host free-room admission per request causes repeated small L2 evictions and defeats the free-room policy, while request metadata and radix attachment still need per-request ownership.\n\nThis change adds a batch prepare path that aggregates required host tokens by CP owner lane, performs one host admission/eviction step for the batch, and then reserves/submits each request independently. SessionAwareCache and the scheduler prefer the batch API when available.\n\nConstraint: Radix nodes, host slots, draft slots, rollback, and ack semantics remain per request.\nRejected: Merge reservations or radix nodes across requests | would complicate rollback and writing_check ack ownership.\nRejected: Add collective synchronization for host admission | local owner-lane logic already provides deterministic capacity accounting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce per-request host free-room eviction in the scheduler path without profiling host-full workloads.\nTested: local py_compile for touched Python files\nTested: remote g0034 docker PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 117 passed, 5 warnings\nNot-tested: full ETE bs>1 HiCache run under production traffic
The scheduler can now admit multi-request NSA in-seq CP shared-KV prefill batches only when the shared-KV bs>1 flag is explicitly enabled. The gate is still disabled by default and is scoped to CP shared-KV so ordinary CP is not widened accidentally.
Batch admission is bounded by optional request-count and page-aligned extend-token limits while real memory capacity remains allocator-owned. This keeps bf16 and fp8 on the same scheduler path because dtype differences are already reflected in KV pool token/page capacity.
Constraint: bs>1 runtime paths remain guarded by existing CP shared-KV fail-fast checks.
Constraint: Scheduler must not duplicate bf16/fp8 byte-level capacity estimation.
Rejected: Open the old CP gate unconditionally | ordinary CP would inherit an unverified shared-KV-specific batching path.
Rejected: Treat the extend-token cap as a hard per-request limit | a single large request could deadlock the scheduler.
Confidence: medium
Scope-risk: moderate
Directive: Keep CP shared-KV batching gated until ETE validates EAGLE accept length, output length, and HiCache load/backup behavior under real traffic.
Tested: local py_compile for server_args, schedule_policy, scheduler, prefill_adder tests, and server_args tests.
Tested: remote g0034 py_compile for the same files.
Tested: remote g0034 pytest target set: 5 passed for parser, parameter validation, default single-request CP gate, enabled bs>1 gate, and page-aligned extend cap.
Tested: remote g0034 pytest test_prefill_adder.py => 13 passed.
Not-tested: full server_args test file has an unrelated HuggingFace DNS/config-download failure in TestPrepareServerArgs.test_prepare_server_args.
Not-tested: ETE production traffic with --enable-cp-shared-kv-prefill-bs-gt1.
EAGLE draft shared-KV is supposed to mirror the target CP layout, so bs>1 must not fall back to legacy full-input or padded-hidden behavior when required batch metadata is missing or inconsistent. This change keeps the existing bs=1 compatibility path but makes batched CP draft fail fast on missing/mismatched spec hidden states, embedding pad metadata, or input embed shapes. The docs record the current W7 boundary: draft prefill follows target metadata, while scheduler admission and ETE remain gated.
Constraint: CP draft KV must mirror target layout and must not silently diverge under bs>1 shared-KV.
Rejected: Allow bs>1 to use the old full-input fallback | it can hide wrong owner/page metadata and corrupt accept length.
Confidence: medium
Scope-risk: moderate
Directive: Do not open the scheduler bs>1 CP gate until EAGLE accept length/output length are verified with this fail-fast path enabled.
Tested: Remote g0034 targeted EAGLE fail-fast unit test passed; remote full test/registered/unit/layers/test_nsa_cp_utils.py passed 70 tests.
Not-tested: EAGLE bs>1 ETE, because scheduler CP bs>1 admission gate remains closed.
Tiny extend requests can leave most CP lanes without query work, which has been tied to hangs and accept-length regressions. This change introduces a dual valid/compute metadata contract: forward paths may materialize compute-padded rows, while cache, current reuse, direct write, HiCache backup, and load remain valid/page based.
The implementation keeps radix/HiCache/device allocation on real page extents, filters dummy compute rows before MLA/index cache writes and current reuse, makes top-k/index consume compute rows while compacting valid rows, and opens tiny CP shared-KV in-seq split through compute padding. The accompanying plan document records the contract and P1-P7 evidence.
Constraint: CP shared KV and HiCache must stay page-granular; dummy compute rows must not allocate, write, backup, or load KV cache.
Constraint: Avoid silent fallback and avoid adding collectives on hot paths.
Rejected: Pad cache allocations to cp_size pages | would waste KV capacity and pollute radix/HiCache state.
Rejected: Keep tiny suffixes out of CP split | preserves the zero-lane behavior that compute padding is meant to remove.
Confidence: medium
Scope-risk: broad
Directive: Do not route compute-padded dummy rows into out_cache_loc, current reuse, HiCache reservation, or backup descriptors; keep valid/cache metadata explicit.
Tested: Remote g0034 container targeted P7 tests: 3 passed, 3 warnings.
Tested: Remote g0034 container full unit slice: 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_shared_kv_layout.py => 214 passed, 5 warnings, 2 subtests passed.
Tested: Local py_compile for touched P7 test file.
Not-tested: Latest CUDA/ETE traffic validation for dummy top-k rows, accept len, output len, and detokenizer hang behavior.
CP HiCache write reservations must stay per radix node, but the transfer descriptor does not need to be per request. This changes the layer-end hook to group pending write states for the same source and layer, so bs>1 prefill emits one target D2H descriptor and one draft D2H descriptor per layer while preserving per-node metadata, rollback, and ack semantics.\n\nConstraint: CP shared-KV HiCache metadata, host slots, and radix acknowledgements remain per request/node.\nConstraint: TAI direct transfer kernels already accept flattened page descriptors, so no tai-kernel change is required.\nRejected: Merge HiCache reservations or radix nodes | would complicate rollback and split handling.\nRejected: Add collective synchronization for grouped backup | grouping is local descriptor construction and must not add rank-level sync.\nConfidence: high\nScope-risk: moderate\nDirective: Keep target and draft source notifications separate; final ack must wait for both when draft HiCache is attached.\nTested: local py_compile for cache_controller.py and test_hicache_controller_cp.py\nTested: local git diff --check\nTested: remote pytest test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 85 passed, 3 warnings\nNot-tested: full ETE bs>1 CP HiCache replay with admission gate removed\nNot-tested: Nsight/throughput validation of reduced D2H submit count
The bs>1 path needs index top-k, shared-index prepare, current-index compact, and current-slot compose to consume flattened batch descriptors instead of falling back to per-request or per-segment Python/Torch work. This change wires SGLang to the new TAI batch prepare kernels, keeps fallback explicit, and records the remaining HiCache/load-backup gaps in the bs>1 workstream docs.
Constraint: CP shared-KV bs>1 must reuse fast paths rather than adding slow batch-only fallbacks
Constraint: No new collective operations were introduced
Rejected: Leave current-only cp_index as Python slice/cat | it keeps per-segment overhead in the short-extend bs>1 case
Rejected: Infer max segment lengths from CUDA descriptor tensors | .item() would add CPU synchronization on the hot path
Confidence: medium
Scope-risk: moderate
Directive: Do not remove the explicit fallback warnings without verifying the corresponding TAI symbols are present in production
Tested: local py_compile for touched SGLang files
Tested: remote g0034 test_nsa_cp_utils.py passed, 53 tests
Tested: remote g0034 test_fill_current_index_page_slots_uses_tai_kernel_when_available passed
Not-tested: full ETE bs>1 traffic with HiCache load/backup and draft/EAGLE enabled
PR #10 provides the W2 owner-lane allocation side needed by the current bs>1 CP shared-KV work: supported multi-request extends now derive page owners per request, preserve flattened request order, and use owner-lane allocation instead of legacy page allocation.\n\nThe merge applies cleanly on top of the W3/W4 target current-reuse work because the touched runtime surface is allocator/mem_cache only. Focused remote tests cover the new owner-lane planner/allocation behavior and the existing W4 target reuse regressions.\n\nConstraint: CP shared-KV page ownership is request-relative and page-granular.\nConstraint: Supported bs>1 CP shared-KV allocation must not silently fall back to legacy allocation.\nRejected: Treat multiple requests as one concatenated extend | would assign owners relative to the wrong request boundary.\nRejected: Keep multi_batch fallback | breaks downstream direct-write owner assumptions.\nConfidence: medium\nScope-risk: moderate\nDirective: Preserve request-order flattening when changing owner allocation; do not reorder by lane for convenience.\nTested: Local py_compile for allocator.py, common.py, cp_shared_kv_compute_owner.py, and test_cp_shared_kv_layout.py.\nTested: Local git diff --check --cached.\nTested: Remote g0034 py_compile for touched mem_cache files and test_cp_shared_kv_layout.py.\nTested: Remote g0034 pytest test_cp_shared_kv_layout.py => 34 passed, 3 warnings.\nTested: Remote g0034 pytest test_nsa_cp_utils.py test_cp_shared_kv_runtime.py => 157 passed, 5 warnings, 2 subtests passed.\nNot-tested: Full ETE/perf run; production HiCache load/back up with bs>1 allocation under sustained traffic.
Batch-size > 1 cache-hit traffic must not lose the current/partial-current reuse fast path. This change extends the target MLA/index sync-correct path to validate batched current suffix rows, compose page-aligned prefix spans, and route batched index top-k through current-only or partial-current reuse without falling back to scalar guards.\n\nThe implementation keeps page as the minimum cache unit: prefix cache coverage is page-aligned, while current suffix rows are sliced by valid extend lengths so padded tail rows are not exposed to attention. The index top-k batch path now mirrors the single-request contract: current-only reuses current index KV directly, partial-current materializes once and shares the dense buffer across request segments.\n\nConstraint: CP shared-KV supports in-seq zigzag only; round-robin remains fail-fast for shared-KV.\nConstraint: No new collectives are introduced; this is a sync correctness path, not a new communication scheme.\nRejected: Keep bs>1 current_index_kv fail-fast | disables the cache-hit path that W4 is meant to restore.\nRejected: Pad requests to a common batch length | violates the page-granular contract and exposes padded tails to consumers.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce batch_size_not_one or batch_gt1 current-reuse guards without proving an equivalent batched fast path exists.\nTested: Local py_compile for touched runtime/test files.\nTested: Remote g0034 pytest test_nsa_cp_utils.py test_cp_shared_kv_runtime.py => 157 passed, 5 warnings, 2 subtests passed.\nTested: Remote g0034 pytest test_cp_shared_kv_layout.py => 27 passed, 3 warnings.\nNot-tested: Full ETE with live traffic; CUDA kernel-backed batched descriptors; L1 prefetch, HiCache load/backup, and draft/EAGLE bs>1 current reuse.
