Commit Graph

2250 Commits

Author SHA1 Message Date
laoyao0822
cc908dd556 Reuse DeepGEMM JIT cache across restarts
New sgl-deep-gemm wheels persist cubins, but SGLang was still replaying the full warmup envelope after every restart. The wrapper also translated SGLANG_DG_* env after importing deep_gemm, so the wheel could observe stale cache settings.\n\nMove DeepGEMM JIT env preparation before the first import, bridge the current SGLANG_DG_USE_NVRTC spelling, replace dense M walks with category-specific sparse M grids, and add a SGLang-side warmup manifest under the DeepGEMM cache dir. A manifest hit skips only the SGLang replay loop; missing or stale cache entries still force warmup and refresh the manifest.\n\nConstraint: sgl-deep-gemm 0.1.2 removed the compile-mode API, so dense 1..m_max loops launch real kernels.\nConstraint: DeepGEMM cache keys include compiler flags such as the deep_gemm include path; path-stable environments are still required for cross-container cache hits.\nRejected: SGLANG_JIT_DEEPGEMM_PRECOMPILE=0 | bypasses warmup instead of making cache reuse correct.\nRejected: Enable per-process cache dirs by default | isolates corruption but defeats persistent cross-process cache reuse.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce dense DeepGEMM warmup for new wheels unless compile-mode support is verified on the target wheel.\nTested: Local py_compile for jit_cache.py configurer.py compile_utils.py model_runner.py\nTested: Local pytest test/registered/unit/layers/test_deep_gemm_jit_cache.py (6 passed)\nTested: Remote g0034 cjy-glm5-new py_compile for same files\nTested: Remote g0034 cjy-glm5-new pytest test/registered/unit/layers/test_deep_gemm_jit_cache.py (6 passed)\nTested: Remote env probe verified SGLANG_DG_CACHE_DIR and SGLANG_DG_USE_NVRTC bridge to DG_JIT_* before configurer completes\nNot-tested: Full decode cold-cache/hot-cache restart timing and manifest-hit ETE log validation
2026-06-11 02:55:04 +08:00
laoyao0822
3a43727216 Bound CP prefill batching by estimated temp memory
CP shared-KV bs>1 batching was only bounded by request count, extend tokens, and cached tokens. That left temporary GPU buffers such as MLA/index materialization, remap metadata, logits windows, and transfer descriptors implicit, and raw extend-token limits could exceed the active chunked-prefill budget.\n\nThis adds an explicit max-buffer-size admission gate with a CPU-only stream-aware estimator, wires it through PrefillAdder/Scheduler, performs a startup CUDA smoke allocation when configured, and reports the estimate in the scheduler admission benchmark. When chunked prefill is active, the effective CP extend-token limit is capped by the current chunk budget so the CP path does not advertise unreachable batch capacity or lift max-prefill-tokens too far.\n\nConstraint: Admission estimation must stay CPU-only on the scheduler hot path; CUDA allocation is limited to startup smoke checking.\nConstraint: Single oversized requests must still be allowed to run alone to avoid scheduler deadlock.\nRejected: Rely only on --max-prefill-tokens | it does not reliably bound the first oversized request and does not model cache-hit/load-back pressure.\nRejected: Let CP extend limit exceed chunked-prefill size | it creates an unreachable effective capacity and misleading budget lift.\nConfidence: medium\nScope-risk: moderate\nDirective: If bs>1 L1 prefetch is enabled later, update CPSharedKVPrefillBufferEstimatorContext.bs_gt1_l1_prefetch_enabled and include the live prefetch dense buffers in overlap windows.\nTested: local py_compile for touched files\nTested: local PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py (4 passed)\nTested: remote cjy-glm5-new targeted pytest for new server_args, PrefillAdder, estimator, and benchmark cases (10 passed)\nTested: remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_prefill_scheduler_admission_bench.py (29 passed before chunk cap, then test_prefill_adder.py 21 passed after chunk cap)\nNot-tested: full server_args suite because existing TestPrepareServerArgs tries to reach HuggingFace and fails under container DNS/network\nNot-tested: GLM5 ETE smoke with --cp-shared-kv-prefill-max-buffer-size
2026-06-11 01:33:28 +08:00
laoyao0822
9a9893e571 Model CP scheduler admission with cache-hit pressure
Add an offline benchmark that reuses PrefillAdder to model how L1 cached tokens, L2 HiCache hits, and extend tokens shape CP shared-KV prefill batch admission. The tool makes scheduler stop reasons and fake L2 load-back capacity pressure observable without starting a model.

Constraint: The benchmark must stay CPU/offline and avoid depending on CUDA execution or live services.

Constraint: L2 cached tokens are modeled as host_hit_length, so successful load-back both increases prefix_len and consumes fake L1 capacity.

Rejected: Build an ETE benchmark first | too slow for isolating scheduler admission behavior.

Rejected: Reimplement scheduler logic from scratch | would drift from PrefillAdder semantics.

Confidence: high

Scope-risk: narrow

Directive: Treat duration_us as Python admission overhead only; it is not an ETE latency metric.

Tested: Remote pytest test/registered/unit/managers/test_prefill_scheduler_admission_bench.py: 4 passed as part of 6 targeted tests.

Tested: Remote synthetic benchmark run with --cp-max-total-cached-tokens showed second 4096-token cached request stopped with OTHER.

Not-tested: Real traffic trace import from production logs.
2026-06-10 22:21:59 +08:00
laoyao0822
4f65d7a176 Bound CP prefill batches by cached-token pressure
High cache-hit CP shared-KV batches can have small extend tokens while still carrying substantial prefix/load-back work. Add a CP-specific cached-token admission limit so operators can bound that pressure independently from extend-token batching.

Constraint: The limit must not deadlock a single high-cache-hit request; it only stops adding additional requests to a non-empty batch.

Constraint: Cached tokens are counted after L2 load-back planning via prefix_len, so L1 hits and successful L2 hits share one scheduler budget.

Rejected: Reuse max_prefill_tokens | it limits generic input budget and does not represent cached-token work.

Rejected: Count only L1 prefix before load-back | would miss L2 hit pressure, which is one of the target cases.

Confidence: high

Scope-risk: moderate

Directive: Keep cached-token and extend-token limits separate; they bound different scheduler costs.

Tested: Remote pytest targeted cached-token PrefillAdder cases: 2 passed.

Tested: Remote pytest test/registered/unit/managers/test_prefill_adder.py: 18 passed.

Tested: Remote ServerArgs CP validation smoke: SERVER_ARGS_CACHED_LIMIT_OK.

Not-tested: Full ETE replay with a production cached-token limit value.
2026-06-10 22:21:46 +08:00
laoyao0822
342c552ab3 Decouple CP prefill batching from generic token cap
CP shared-KV bs>1 uses cp_shared_kv_prefill_max_total_extend_tokens as its grouping admission limit, but the generic max_prefill_tokens budget could still stop batching earlier. Raise only the legacy input-token admission budget for this CP path while keeping allocator-owned capacity checks unchanged.

Constraint: CP shared-KV bs>1 needs large cache-hit batches without relying on the generic max_prefill_tokens default.

Constraint: Allocator capacity must remain enforced by rem_total_tokens, cur_rem_tokens, and prepare_for_extend().

Rejected: Increase server-wide max_prefill_tokens | would change generic scheduler behavior and non-CP paths.

Confidence: high

Scope-risk: narrow

Directive: Do not use max_prefill_tokens as the CP shared-KV bs>1 grouping limit; use the CP-specific total extend token knob.

Tested: Local py_compile for schedule_policy.py and test_prefill_adder.py.

Tested: Remote pytest targeted CP PrefillAdder cases: 5 passed.

Tested: Remote pytest test/registered/unit/managers/test_prefill_adder.py: 16 passed.

Not-tested: Full ETE replay after this scheduler-only change.
2026-06-10 21:36:22 +08:00
3eef989739 Harden IndexCache for first e2e: validation, fail-fast, prefetch gating
Pre-e2e audit fixes for NSA index-topk sharing (index_topk_pattern):

- Validate the pattern at init: F/S charset only, must start with F
  (layer 0 has nothing to share from), and length must equal
  num_hidden_layers — a short pattern previously crashed deep in pool
  init with an obscure per-layer ValueError, a long one silently
  ignored tail characters. Warn when a pattern shadows a configured
  index_topk_freq (pattern takes precedence silently otherwise).
- Fail fast when a shared (S) layer receives prev_topk_indices=None:
  both gated indexer call sites previously fell through to omitting
  topk_indices entirely, running sparse attention without indices and
  silently corrupting output if threading were ever dropped.
- Reject pipeline parallelism with index-topk sharing (same rationale
  as the existing TBO guard): topk indices are not threaded across
  PP-stage boundaries, and each stage's loop resets them to None.
- Gate the CP shared-KV index prefetcher on the next layer having an
  index-cache slot: prefetching an S layer's nonexistent buffer
  tripped the inactive-layer fail-fast and disabled the prefetcher
  for the entire run on the first F->S transition.
- Log the resolved active index layer plan at pool init so an e2e run
  can confirm IndexCache is actually on.

Tests: pattern validation cases incl. the GLM-5 reference 78-char
pattern, next_skip == skip-of-next invariant; existing 'C' pattern
test updated to upstream 'F' charset.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 06:17:52 +00:00
5bec69f40a Drop redundant np.isin from CP-rank KV page filter
filter_kv_indices_for_cp_rank built rank_page_indices =
kv_indices[range_mask] and then ran np.isin(kv_indices,
rank_page_indices) — provably identical to range_mask itself (a value
is in the filtered subset iff it passes the same range test), at the
cost of an extra sort per chunk send under
SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER=1. Apply the range mask
directly: 30.5 -> 8.1 us per 1024-page chunk. Differential test pins
equivalence against a verbatim copy of the old logic.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 05:15:16 +00:00
e6864fbea7 Demote HiCache/CacheCtrl per-op logs from info to debug
The [HiCache-write]/[HiCache-load]/[HiCache-collective] and
[CacheCtrl-write] per-operation logs fired at INFO on the scheduler
thread during serving (every write-through, load-back, ack release,
and rate-limited collective summary), eating into stream gaps.
Demote all 29 to debug; one-time [HiCache-draft] init logs stay INFO
and genuine failures stay WARNING (write_backup CP FAILED after
deterministic retry).

Also repair TestHiCacheEvictLoggingLevels, which asserted markers
that no longer exist in this tree (stale list, failing on the
pristine branch): pin the current hot-path markers as debug-only and
extend the check to cache_controller.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 05:10:45 +00:00
87a22b17ce Fuse SamplingBatchInfo tensor construction into one pass
from_schedule_batch built temperature/top_p/top_k/min_p (+seed) with
4-5 separate list comprehensions and one synchronous H2D copy each,
plus 4 more passes for the is_all_greedy/need_* flags. Collect
everything in a single pass over reqs and upload the float params as
one pinned non-blocking H2D copy (disjoint device views of one
buffer; filter/merge only index and cat, producing fresh tensors, so
the shared buffer is safe), int32 top_k and optional int64 seeds as
their own pinned copies.

B300 (torch 2.11 cu130), scheduler-thread blocking time per call:
  bs=8: 38.8 -> 18.2 us (2.1x); bs=32: 1.9x; bs=200: 1.2x
CPU-only construction at bs=200: 2890 -> 1387 us (2.1x).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 05:04:39 +00:00
0d065a8ab0 Speed up Python bigram key conversion with C-level zip
The EAGLE bigram fallback built 64K-token keys with an index-based
Python comprehension (6.3 ms per call at 64K tokens). zip + islice
runs the pairing loop in C and avoids copying the shifted operand:
3.6-4.2 ms per call (~1.6x). Output is byte-identical (same tuples);
the tai-kernel fast path is unaffected.

Packing bigrams into int64s via numpy was measured and rejected: the
list<->array boxing makes it slower (4.1 ms) than zip until token ids
are numpy end-to-end, and it would change HiCache storage hash inputs.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 04:56:24 +00:00
8979c81a22 Make RadixKey slicing zero-copy to fix quadratic match_prefix
RadixKey.__getitem__ copied the token list on every slice, and the
match/insert tree walks re-slice the remaining key at every node hop,
making a single match O(len * hops) — quadratic for long cached
prefixes. Slices now return O(1) offset-based views over a shared
backing list; the key-match and child-key functions index the backing
list directly so views are never materialized on the hot path.

Keys stored in tree nodes are compacted at every store site (same cost
as the old copying slices), so lock-ref walks, eviction, splits, and
controller-thread reads never observe a key pinning a transient
backing list. Node-key compactness is enforced by a tree-walk test.

Microbenchmark (64K-token full hit, page_size=64):
  16-node path: 1.86 -> 0.87 ms/match (2.1x)
  128-node path: 9.80 -> 1.19 ms/match (8.2x), insert re-walk 6.9x

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-10 04:28:03 +00:00
laoyao0822
f78c414b64 Skip unused FP8 NSA K-cache dequant work
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
2026-06-10 07:00:43 +08:00
laoyao0822
af20334f4c Size CP HiCache host budget from compact index layers
The CP shared target+draft HiCache budget was still estimating NSA target
bytes per token as if every target layer had an index cache. After index
skip compacts L1 and L2 index buffers to active layers, that stale estimate
kept host/L2 token capacity artificially low even though the physical host
index allocation was already smaller.

Use the pool's index_active_layer_ids when estimating NSA HiCache bytes per
token, with the existing full-layer fallback for pools that do not expose
compact metadata. This keeps the shared target+draft capacity computation
consistent with the actual L2 host layout.

Constraint: Index skip compacts NSA index cache by active logical layers, but MLA KV remains full-layer.
Rejected: Keep full-layer HiCache estimate | preserves correctness but wastes the L2 capacity recovered by compact index allocation.
Confidence: high
Scope-risk: narrow
Directive: Keep HiCache byte estimates aligned with NSATokenToKVPool allocation metadata; otherwise target+draft token budgets will drift from real host memory use.
Tested: Remote container targeted pytest for active-index HiCache estimate and shared-budget tests -> 3 passed.
Tested: Remote container P3-P6 regression suite plus new capacity tests -> 203 passed, 2 subtests passed.
Not-tested: Full ETE replay/GSM8K after this budget correction.
Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-06-10 05:53:05 +08:00
laoyao0822
1ebde44e59 Compact skipped NSA index-cache state safely
Index skip reduces the number of target layers that own NSA index state,
but PD transfer and HiCache still assumed dense full-layer state buffers.
This change carries explicit state layer IDs through prefill/decode
registration, compacts device and host index buffers to active layers,
and maps logical layer IDs to compact slots on transfer paths.

The PD side fails fast when prefill/decode disagree on NSA state layer
identity instead of silently truncating or copying mismatched buffers.
Host direct tests now use the same CPU-index descriptor contract required
by the TAI cudaMemcpyBatchAsync path, and host registered memory is
unregistered on tensor finalization to avoid stale cudaHostRegister state
across CUDA tests.

Constraint: CP shared-KV with index_topk skip must keep target/draft state identity explicit before compacting buffers
Constraint: Direct HiCache TAI transfer rejects CUDA indices to avoid hidden D2H copies on the control path
Rejected: Keep full-layer L1/L2 index buffers | wastes the memory/bandwidth that index skip is meant to save
Rejected: Infer state buffer order by count only | can silently corrupt cache when active layer sets differ
Confidence: high
Scope-risk: moderate
Directive: Do not compact or reorder NSA state buffers without carrying logical layer IDs through PD registration and validating both sides
Tested: Remote container py_compile for touched runtime files
Tested: Remote container pytest: test_nsa_pool_host_unit.py, test_model_runner_kv_cache_mixin.py, test_cp_shared_kv_transfer_mapping.py, test_pd_state_layer_ids.py, test_cp_per_layer_transfer.py, test_cp_shared_kv_runtime.py -> 200 passed, 2 subtests passed
Not-tested: Full ETE GSM8K/replay after compacted P3-P6 changes
Co-authored-by: OmX <omx@oh-my-codex.dev>
2026-06-10 05:37:06 +08:00
laoyao0822
d21952b903 Reduce inactive NSA index-cache transfer safely
Centralize the IndexCache skip formula and thread the resulting active logical index layers into NSA KV pools. HiCache now skips only the indexer H2D/D2H payload for inactive target layers while preserving per-layer MLA KV transfer, keeping allocation shape unchanged for this phase.

Constraint: P0-P2 must not compact device or host allocation yet; prefill/decode state transfer still has no logical layer-id metadata.

Rejected: Recompute the skip formula separately in mem_cache | formula drift would corrupt cache or waste transfers when offset/pattern settings change.

Rejected: Skip whole-layer HiCache load/backup | MLA KV remains required for every attention layer.

Confidence: medium

Scope-risk: moderate

Directive: Before enabling compact state buffers or compact allocation, add layer-id metadata validation to PD transfer.

Tested: Local py_compile for touched files; remote pytest in g0034 container: test_nsa_index_layers.py and TestNSAIndexerPageIndices, 20 passed.

Not-tested: ETE replay/GSM8K with --nsa-index-topk-freq 4; PD state-transfer compaction remains unimplemented.
2026-06-10 04:28:26 +08:00
laoyao0822
a32050d1d8 Expose NSA index sharing as launch-time overrides
Index-topk sharing previously required editing model config or passing raw JSON overrides. Add first-class server args that merge into json_model_override_args before ModelConfig is read, preserving existing JSON overrides while letting launch scripts toggle index_topk_freq directly. Treat offset 0 as unset so wrappers can use 0 for default behavior without injecting a model-config override.

Constraint: Prefill and decode launch commands need the same effective model config without mutating /ssd model files.

Rejected: Require editing config.json | operationally fragile across g0034/g0035/g0036 model copies.

Rejected: Only use --json-model-override-args | works but is too error-prone for frequent launch-command tuning.

Confidence: high

Scope-risk: narrow

Directive: Keep these flags as model-config override shortcuts; apply them before any get_model_config() call.

Tested: Remote pytest in g0034 container for NSA index override parser tests: 2 passed.

Tested: py_compile for server_args.py and test_server_args.py.

Not-tested: Full prefill/decode ETE launch with --nsa-index-topk-freq enabled.
2026-06-10 03:24:33 +08:00
laoyao0822
24da983ff5 Enable CP HiCache direct transfers to use layer-page host layout
CP shared-KV HiCache transfers are per-layer, so the host layout should match the access pattern instead of forcing page-major strides through every layer. This adds a direct-only layer_page_first layout, routes per-layer KV and NSA index backup/load through the TAI LF<->LPF direct kernels, and keeps storage/page-buffer metadata paths fail-fast until their page-level contract is redesigned.\n\nThe direct controller keeps host indices in caller order for both page_first_direct and layer_page_first because the TAI direct path requires CPU index descriptors and owns descriptor coalescing. All-layer backup intentionally loops over per-layer direct kernels rather than using the sgl-kernel all-layer direct ABI.\n\nConstraint: layer_page_first is currently host-only CP HiCache; storage backends assume page-major contiguous page metadata.\nConstraint: TAI LPF direct kernels require CPU int64 page indices and complete page spans.\nRejected: silently fallback to SM copy when TAI LPF kernels are missing | that hides production performance regressions.\nRejected: support storage page metadata in this commit | LPF requires a layer-page-level storage contract, not a one-pointer-per-page contract.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not enable storage or kernel backend for layer_page_first without redesigning page-buffer metadata and adding remote ETE coverage.\nTested: local py_compile for touched runtime files.\nTested: remote py_compile in g0034 container for touched runtime files.\nTested: remote targeted pytest: 5 passed for parser/storage/layout/move_indices smoke coverage.\nNot-tested: full CP HiCache ETE with --hicache-mem-layout layer_page_first after this commit step.\nNot-tested: combined CUDA roundtrip tests in one pytest process; previous independent runs passed but combined run exposed a host-memory registration lifecycle issue.
2026-06-10 02:03:07 +08:00
laoyao0822
5e22279670 Carry CP-local hidden contract into EAGLE v2 prefill draft extend
SPEC_V2 builds an EAGLE draft input during target prefill. Under CP shared-KV, the target model may expose draft_hidden_states as a CP-local side channel before CP output collection. The v2 path was still passing global hidden_states and never marked the draft input as CP-local, unlike the non-v2 path.

Thread cp_local_hidden_states through the v2 _draft_extend_for_prefill helper and prefer draft_hidden_states when present. This preserves the semantic marker consumed by the static MLP-sync padding guard without changing the non-CP path.

Constraint: Absorb syh 5562937cf only; SPEC_V2 remains opt-in and broader SPEC_V2-on-CP validation is still separate.
Rejected: Infer CP-local hidden from tensor length | tensor length is ambiguous under static padding and bs>1 compute padding.
Confidence: high
Scope-risk: narrow
Directive: Keep EagleDraftInput.cp_local_hidden_states as an explicit semantic contract; do not replace it with shape-based inference.
Tested: Remote g0034 container red test failed before implementation with unexpected cp_local_hidden_states kwarg.
Tested: Remote g0034 container py_compile for eagle_worker_v2.py and test_nsa_cp_utils.py.
Tested: Remote g0034 container pytest target EAGLE marker tests: 2 passed.
Tested: Remote g0034 container pytest test_nsa_cp_utils.py -k eagle: 2 passed, 99 deselected.
Not-tested: Full ETE with SGLANG_ENABLE_SPEC_V2=1 on CP prefill.
2026-06-09 22:42:54 +08:00
laoyao0822
df2a3696cd Gate CP split validation to eligible extend forwards
CP shared-KV marks all forwards with uses_cp_shared_kv, but TARGET_VERIFY/decode style forwards may legitimately carry no extend prefix page-plan. The CP split helper previously ran the hard page-plan validator before checking context-parallel extend eligibility, so non-extend spec paths could fail before the later no-split decision.

Hoist the eligibility check and only validate page-plan metadata for real context-parallel EXTEND or shared-KV draft extend forwards. Add a regression that TARGET_VERIFY with no prefix metadata returns no-split instead of failing.

Constraint: Absorb syh 7fea88278 only; MQA logits chunk and per-forward budget changes are intentionally excluded.
Rejected: Relax the validator globally | real prefill page-plan violations must remain fail-fast.
Confidence: high
Scope-risk: narrow
Directive: Do not run CP shared-KV page-plan validation for non-context-parallel forward modes without proving those modes own extend_prefix_lens_cpu.
Tested: Remote g0034 container py_compile for utils/test file.
Tested: Remote g0034 container pytest test_nsa_cp_utils.py -k can_cp_split: 8 passed, 92 deselected.
Not-tested: Full ETE speculative non-deepep MoE path.
2026-06-09 19:55:22 +08:00
laoyao0822
81eb138a26 Avoid CP HiCache owner-lane planning stalls during load-back
CP HiCache load-back eviction planning previously recomputed per-owner page counts from node token tensors while scanning evictable leaves. Under shared-KV pressure this can put scheduler-side planning onto an expensive tensor-padding path and stall before load_back can complete.

This stores per-CP-size owner page counts on CP HiCache metadata and uses that CPU metadata for backed/resident CP nodes. Streaming abort handling also accepts int-like status codes so abort responses do not crash on .name/.value access. Temporary debug runbooks remain ignored.

Unnecessary prefill hot-path timing logs were removed before commit; owner-lane eviction now keeps warning-level output for slow planning, insufficient eviction, or remaining deficits only.

Constraint: CP shared-KV cache residency is page-owner based and already records page owners in CpHiCacheNodeMetadata.
Rejected: Keep verbose prefill/owner-lane timing logs | they proved the issue but add hot-path noise after validation.
Confidence: medium
Scope-risk: moderate
Directive: Do not reintroduce tensor-derived owner counting on CP HiCache backed nodes without measuring scheduler CPU/GPU sync cost.
Tested: python -m py_compile python/sglang/srt/mem_cache/hiradix_cache.py python/sglang/srt/entrypoints/openai/serving_base.py python/sglang/srt/entrypoints/openai/serving_chat.py python/sglang/srt/entrypoints/openai/serving_completions.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
Tested: git diff --check
Not-tested: Local pytest collection is blocked by missing starlette dependency.
Not-tested: Full ETE after log cleanup; previous pre-cleanup ETE replay reached 136098.82 prompt tok/s without killing prefill.
2026-06-09 06:55:39 +08:00
laoyao0822
50fde834ae Keep CP compute padding out of sparse MoE
CP shared-KV compute padding creates per-request lane slots, so valid rows are not a simple prefix/suffix mask. DeepEP MoE was still seeing dummy rows and using scalar non-padded semantics, which let padding participate in gate/topk and corrupted cache-hit tiny-extend inference.\n\nThe fix compacts CP-local valid rows before MoE dispatch and restores the compact output back to the compute-padded row layout before downstream layer communication. The local GSM8K investigation ledger is now removed from the tracked tree and ignored so future debug notes stay local.\n\nConstraint: CP shared-KV compute-padding layout must keep downstream communicator shapes stable.\nRejected: Disable bs>1/current reuse/cache-hit fast paths | hides the semantic bug and loses the intended performance path.\nRejected: Use num_token_non_padded for MoE under compute padding | valid rows are interleaved with dummy lane slots, not suffix-padded.\nConfidence: high\nScope-risk: moderate\nDirective: Do not feed compute-padded dummy rows into sparse MoE gate/topk; compact valid rows at the MoE boundary and restore shape afterward.\nTested: python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/models/deepseek_v2.py\nTested: remote focused CP utils tests passed, 4 tests.\nTested: remote GSM8K 50-question smoke accuracy 0.960; 200-question runs accuracy 0.955 and 0.965; full 1319-question run accuracy 0.952.\nNot-tested: Long-running production traffic beyond GSM8K after this commit.
2026-06-09 01:48:24 +08:00
laoyao0822
3698b0c22c Lock bs>1 CP compute-padding invariants
The GSM8K warm-cache investigation exposed several tiny-extend shapes where compute padding, valid-row selection, last-token collect, and all-gather rerange must stay request-slot aware. These tests pin those invariants without changing runtime behavior.

Constraint: bs>1 tiny extends use page-granular compute slots while only a subset of rows are semantically valid
Rejected: Rely on scalar non-padded token counts for these layouts | valid rows are not always a simple suffix mask
Confidence: medium
Scope-risk: narrow
Directive: Do not weaken these tests unless the replacement path proves equivalent request-slot and valid-row semantics
Tested: Remote pytest test_nsa_cp_utils.py as part of the 263-test CP regression suite
Not-tested: CUDA graph paths; prefill does not use cuda graph in the current CP setup
2026-06-08 20:56:48 +08:00
laoyao0822
b58513cba4 Preserve request-slot identity for CP cache-hit reuse
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
2026-06-08 20:56:31 +08:00
laoyao0822
2aa0b7313e Fix RAGGED CP cache-hit current KV reuse
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.
2026-06-08 03:39:34 +08:00
laoyao0822
8f23b06427 Lock fp8 CP cache-hit roundtrips across batch requests
The GSM8K cache-hit regression needs executable coverage for the fp8 paths that combine bs>1 page planning, valid-row writes, prefix materialization, and HiCache L2 reload.  These tests construct bs=5 page-aligned layouts and assert valid rows survive the persistent-page store/materialize and host roundtrip paths.

Constraint: Production runs use fp8_e4m3 CP shared KV with page_first_direct HiCache and bs>1 prefill.

Rejected: Validate with bf16-only tests | the observed regressions are fp8/cache-hit sensitive and bf16 coverage does not exercise scale/index byte layouts.

Confidence: medium

Scope-risk: narrow

Directive: Keep fp8 cache-hit tests aligned with the page-as-minimum-cache-unit contract.

Tested: python -m py_compile on changed runtime files.

Not-tested: CUDA/TAI tests not run locally; local pytest blocked before collection by missing orjson dependency.
(cherry picked from commit 6600f75fc8ee803610df0078bb880df808655635)
2026-06-08 00:17:28 +08:00
laoyao0822
0f9b445131 Keep compute-padding source offsets out of valid rerange output
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)
2026-06-08 00:17:28 +08:00
laoyao0822
c9a39ccdd2 Preserve decode suffix KV locations after cache hits
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)
2026-06-08 00:17:28 +08:00
b1cbacffae test(cp_shared_kv): load real sgl_kernel first to stop CPU-CI stub leaking
test_cp_shared_kv_runtime installs CPU-CI sgl_kernel stubs at import time via
sys.modules.setdefault + torch.library FRAGMENT defs. On a GPU box, when this
module is collected before the real sgl_kernel loads, the empty stub permanently
shadows it process-wide, so a later real-kernel test (e.g.
fast_topk_transform_ragged_fused in test_nsa_topk_transform) calls a None-returning
lambda and crashes. Importing the real sgl_kernel first makes setdefault keep the
real module (setattr only fills missing names) and the FRAGMENT defs hit the
already-registered path; on CPU-CI the import fails and stubbing proceeds as before.

Verified on GPU: test_cp_shared_kv_runtime alone 120 passed; combined with
test_nsa_topk_transform 125 passed (previously 1 failed / segfault on cleanup).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-07 10:15:35 +00:00
243f1b964c fix(cp_per_layer_transfer): gate per-layer path on a single non-dummy decode info
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>
2026-06-07 09:51:05 +00:00
e9354c41bc fix(cp_per_layer_transfer): use absolute chunk_page_start for chunked KV dst mapping
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>
2026-06-07 09:51:05 +00:00
beffd3a82f feat(disagg): unify per-layer transfer over chunked prefill (lever A)
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>
2026-06-07 09:51:05 +00:00
e12afe8ced perf(disagg): coarse-grained per-layer transfer + skip re-registration (lever A)
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>
2026-06-07 09:51:05 +00:00
5fdf439c7e fix(disagg): dedup per-layer enqueues so high-cache-hit can't hang (lever A)
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>
2026-06-07 09:51:05 +00:00
64f767bb42 fix(disagg): cover notifier-missed layers via post-forward fallback (lever A)
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>
2026-06-07 09:51:05 +00:00
ae18e3adc8 feat(disagg): wire per-layer overlapped transfer end-to-end (lever A, A3-step3)
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>
2026-06-07 09:51:05 +00:00
aa6acc9485 feat(disagg/mooncake): build_per_layer_context assembly (lever A, A3-step2)
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>
2026-06-07 09:51:05 +00:00
6bd5d5760c feat(disagg): per-layer block-address computation (lever A, A3-step2 core)
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>
2026-06-07 09:51:05 +00:00
7ac8067db1 feat(disagg): per-layer transfer manager (lever A, A2 orchestration)
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>
2026-06-07 09:51:04 +00:00
e0d47bfc41 feat(disagg): per-layer transfer context for overlap (lever A, A2 core)
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>
2026-06-07 09:51:04 +00:00
a09aec30a3 feat(disagg/mooncake): async per-layer KV transfer primitive (lever A foundation)
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>
2026-06-07 09:51:04 +00:00
9a91c73e09 feat(disagg/mooncake): add non-blocking async transfer wrapper bindings
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>
2026-06-07 09:51:04 +00:00
9a40e42384 fix(disagg/mooncake): notify prefill on decode-side abort
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>
2026-06-07 09:51:04 +00:00
c6c88c2617 fix(disagg): account KV transfer metrics by actually-sent bytes
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>
2026-06-07 09:51:04 +00:00
laoyao0822
f75ffff8d9 Protect CP shared-KV cache-hit correctness under batched FP8 reuse
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.
2026-06-07 13:26:49 +08:00
laoyao0822
b17976b60d Keep CP HiCache free room from collapsing under cache reuse
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
2026-06-07 01:23:34 +08:00
laoyao0822
81bfd4bfd3 Reduce CP shared-KV batch overhead without reverting bs1 planning
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.
2026-06-06 02:03:16 +08:00
laoyao0822
1b99de7459 Expose gated evidence for CP shared-KV bs>1 debugging
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
2026-06-05 23:48:08 +08:00
laoyao0822
6eea77e5e9 Stabilize disaggregated decode burst handoff
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.
2026-06-05 19:59:04 +08:00
laoyao0822
5eecc5f9c5 Preserve CP ragged topk coordinates under batch planning
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.
2026-06-05 07:45:29 +08:00
laoyao0822
c8510938b9 Repair CP current-slot composition instead of disabling reuse
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.
2026-06-05 06:25:06 +08:00