Commit Graph

7374 Commits

Author SHA1 Message Date
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
4d2fdd14ac Keep kernel guard aligned with deployment env
The SYH DeepGEMM pick raised the global sglang-kernel runtime guard to 0.4.3, but this deployment and pyproject intentionally use sglang-kernel 0.4.0 while DeepGEMM compatibility is handled by the deep_gemm_wrapper import/call path and tai-provided runtime pieces. A global 0.4.3 guard blocks startup before those targeted checks can run.\n\nConstraint: Do not update pyproject.toml because local and remote environments intentionally differ.\nRejected: pip install sglang-kernel --force-reinstall | mutates the remote environment and conflicts with the pinned project environment.\nRejected: keep a global 0.4.3 guard | blocks the current CUDA deployment despite pyproject requiring 0.4.0.\nConfidence: medium\nScope-risk: narrow\nDirective: Do not raise this global guard without also updating the environment contract and pyproject together; feature-specific wheel/API checks belong near their import/call sites.\nTested: local py_compile for python/sglang/srt/entrypoints/engine.py; remote py_compile in g0034 container; remote installed sglang-kernel version confirmed as 0.4.0.\nNot-tested: full launch_server restart after the guard change.
2026-06-09 00:45:54 +08:00
laoyao0822
eabe0bbfad Preserve CP slot-remap reference paths after SYH pick
The SYH per-request page_inverse migration fixes bs>1 cache contamination, but the conflict resolution made several helper/reference paths require explicit request ids and broke existing unit coverage. This keeps production bs>1 callers on explicit req-id routing while allowing bs=1/reference helpers to infer or default request ids without reintroducing the hot-path batch-global inverse.

Constraint: Runtime bs>1 materialize paths must route through per-request page_inverse rows to avoid cross-request KV aliasing.

Constraint: The remote test container may still have single-row tai-kernel materialize helpers while production bs>1 requires the new req-id ABI.

Rejected: Revert to batch-global page_inverse | it is the GSM8K cache-hit corruption root cause.

Rejected: Update tests only | the helper API is still useful for bs=1/reference callers and py-level regression coverage.

Confidence: medium

Scope-risk: moderate

Directive: Do not remove explicit loc_req_id/current_req_id from production bs>1 call sites; default inference is for legacy/reference use only.

Tested: Local py_compile for cp_shared_kv_runtime.py and nsa_indexer.py.

Tested: Remote container py_compile for cp_shared_kv_runtime.py.

Tested: Remote container pytest: test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py => 263 passed, 5 warnings, 2 subtests passed.

Not-tested: Full ETE GSM8K/cache-hit run after SYH pick.
2026-06-09 00:19:15 +08:00
165b6a01f1 fix(nsa): correct bs>1 CP shared-KV KV per-request page_inverse + CP/EAGLE compute-padding
Two correctness fixes in the bs>1 CP shared-KV prefill path, validated end-to-end on B300 (GLM-5.1-FP8, EAGLE, fp8 KV, CP8):

1) Cross-request KV contamination: the value-keyed 1-D page_inverse was last-writer-wins, so two requests sharing a physical page (radix shared-prefix / current-reuse) aliased each other's KV. Replace with a per-request flat page_inverse[batch_rows*capacity] indexed req_id*capacity+lp, threaded through build/remap/fill_current/prefetch and the TAI/Triton kernels. req_id sourced via slot//pages_per_request (build), repeat_interleave (logical_locs), and a companion tensor through the same CP valid-split (current_locs).

2) CP+EAGLE compute-padding cache-write: forward_absorb_prepare's rebuild_cp_kv_cache all-gathers current k back to global rows under current-reuse, but the compute-padding branch fed that full k straight to select_cp_local_valid_rows_for_cache_write (which requires per-rank compute rows) -> fail-fast. Add cp_localize_current_kv_to_compute_rows to re-localize via the batch-plan compute split first (mirrors the non-padding branch).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 28efef91b05ede1a7072ee451e2aea39ecc3a5bd)
2026-06-09 00:01:04 +08:00
81a0191331 perf(deepgemm): CP-aware dense warmup M-grid for NSA prefill CP
Shrink the non-grouped (dense/attention) DeepGEMM warmup M grid by attn_cp_size under NSA prefill in-seq CP (per-rank M = tokens/cp), while keeping the grouped MoE GEMM grid full (deepep all-to-all re-gathers all tokens; topk==ep_size keeps MoE M ~= chunked). Gated by _cp_dense_warmup_divisor.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 672ef3a6609abd100e5e95b42016d17d0a2966e5)
2026-06-08 23:46:59 +08:00
134b88b9b4 fix(server_args): keep NSA prefill/decode backends in the same family for fp8
calculate_mla_kv_cache_dim only packs fp8 (override dim -> nsa_kv_cache_store_fp8=True, required by the flashmla_sparse dequant path) when NEITHER prefill nor decode backend is trtllm. The Blackwell-no-DP guard defaulted the unset side to trtllm, producing a mixed flashmla/trtllm pair that left fp8 KV unpacked and crashed flashmla_sparse ('kv must have dtype bf16'). Now an explicit flashmla choice on either side pulls the unset side to flashmla too.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit ed30312c3d599595e40ae99bc2832159b634cca6)
2026-06-08 23:46:59 +08:00
98fdae3375 fix(config): restore GlmMoeDsa index_topk_freq dropped by HF config load
GlmMoeDsaConfig.__init__ has no **kwargs, so transformers drops the raw
config.json fields the IndexCache/DSA path reads via getattr (index_topk_freq,
index_topk_pattern, index_skip_topk_offset) and may clobber qk_rope_head_dim —
only named fields like index_topk survive. Re-read them from config.json and
restore in get_config for GlmMoeDsaForCausalLM, matching upstream #27114.

Without this, setting index_topk_freq in the model config has NO effect
(getattr falls back to 1 -> IndexCache stays off). Verified on the B300 image:
get_config("/ssd/models/GLM-5.1-FP8") now returns index_topk_freq=4.

Refs: WI-2026-06-07-001

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 92ff8b47470c61462e252bc7ba428203ff3c324c)
2026-06-08 23:46:59 +08:00
ce3b5bc874 feat(nsa): port IndexCache (indexer topk sharing across layers)
Port upstream IndexCache for DeepSeek-V3.2 / GLM-5 (#21405 + #27114 gate fix) to
our diverged tree, matching upstream HEAD's merged form. The NSA indexer topk is
computed every `index_topk_freq` layers; `skip_topk` layers reuse the previous
layer's topk via prev_topk_indices threaded through the model forward. Default
index_topk_freq=1 => no sharing => zero behavior change until the model config opts in.

- forward_mla.py: gate the two indexer call sites
  (`if not skip_topk or (is_nextn and prev_topk_indices is None)`), else reuse
  prev_topk_indices; forward_absorb_core returns (output, topk_indices) when
  next_skip_topk is set (topk_indices already threaded prepare->core via inner_state).
- deepseek_v2.py: AttentionMLA.__init__ skip_topk/next_skip_topk setup
  (freq/pattern/offset; is_nextn=True/True); thread prev_topk_indices through
  forward/forward_prepare/op_core; DecoderLayer returns a 3-tuple (tuple-unpack
  placed AFTER the CP shared-KV finally); Model.forward loop threads topk_indices.
- deepseek_nextn.py: 3-tuple decoder unpack.
- server_args.py: port the #27114 guard - raise on --enable-two-batch-overlap with
  index-topk sharing (the TBO op path does not propagate topk across layers, so
  shared layers would run sparse attention with no indices).

All DecoderLayer.forward callers covered (Model loop, nextn, glm4_moe_lite /
mistral_large_3_eagle inherit, TBO op-path discards via op_core). Did NOT port
#24392 (orthogonal indexer-topk capture/output infra). Import-validated on the
B300 / torch-2.11 image. Enabling it requires setting `index_topk_freq` in the
model config (Zhipu-confirmed for GLM-5.1).

Refs: WI-2026-06-07-001

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit af2a41b8d396033e378b6bdd5f6baebf2a98c301)
2026-06-08 23:46:59 +08:00
4cf8885a74 perf(nsa): cache MQA-logits budget + drop indexer rope clones
Two upstream NSA-indexer perf ports (verified against our diverged tree), cutting
CPU launches / host syncs on the prefill critical path for GLM-5.1-FP8 on B300:

- #25299 (beaff00331): cache the MQA-logits chunking budget per device so prefill
  stops calling torch.cuda.mem_get_info (a host sync) on every indexer pass. The
  budget is computed once and capped by the mem_fraction_static serving headroom;
  a static guard is used (uncached) during cuda-graph capture, and the first real
  prefill caches the free-memory budget.
- #22232 (671fe73961): replace `dst = src.clone()` slice write-backs of the RoPE
  output with a data_ptr-guarded `dst.copy_(src)`. q_rope/k_rope are torch.split
  views of query/key, so when RoPE runs in place src/dst alias and the write-back
  is a redundant no-op (guard skips it); otherwise one copy instead of clone+copy.
  Saves an alloc+copy per q/k per indexer call. (Skipped the PR's AMD-only
  @torch.compile cleanup.)

Both verified to import on the B300 / torch-2.11 image. Deliberately NOT taken:
#21332 (un-force MHA one-shot on Blackwell — risky, per user), #23856 (torch.mm
indexer GEMM — per user).

Refs: WI-2026-06-07-001

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 6bc23945e4ff5c68b0618173856c7daf4653fa10)
2026-06-08 23:46:59 +08:00
db5ab48b19 feat(deepgemm): migrate to sgl-deep-gemm 0.1.2 wheel API
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)
2026-06-08 23:46:59 +08:00
laoyao0822
a32c639147 Remove temporary CP cache-hit trace instrumentation
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
2026-06-08 20:57:08 +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
fb5ccaff26 Preserve CP shared-KV cache-hit diagnostics
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.
2026-06-08 03:39:50 +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
b324407def Trace CP HiCache cache-hit state for GSM8K regressions
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)
2026-06-08 00:17:28 +08:00
laoyao0822
ff4717c2fb Keep L1 free-room eviction off the CP load-back hot path
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)
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
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
5791684864 feat(disagg): broaden lever A to full sequence incl cached prefix (high cache-hit)
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>
2026-06-07 09:51:05 +00:00
76dc60a1d0 test(disagg): finish() latency breakdown instrumentation (lever A A5 root-cause)
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>
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
246dbddac0 fix(disagg): per-layer completion can't hang + worker CUDA device + diagnostics
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>
2026-06-07 09:51:05 +00:00
501faa8f5b fix(disagg): tensor-safe prefix check + graceful per-layer register (lever A)
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>
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
74e22a0721 feat(disagg): register per-layer transfer manager notifier (lever A, A3-step1)
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>
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
2524fe4c9d fix(disagg/mooncake): mark abort Failed and skip aborted chunks in worker
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>
2026-06-07 09:51:04 +00:00
0cd5c38a9e fix(disagg/mooncake): clip CP/NSA state indices to the shorter side
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>
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
68defa3923 Expose CP shared-KV batch timing evidence
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.
2026-06-06 14:28:26 +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