diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index ecdd19550..dd25f511d 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -2510,3 +2510,411 @@ Interpretation: - Split-prune removes a real overlapping-radix invariant violation, but it did not recover EAGLE accept length in this run. - The current accept collapse is more likely in the semantic correctness/content of EAGLE prebuilt state or draft/target KV alignment than in missing metadata lifetime or gross transfer failure. Next debugging should compare the actual prebuilt top-k token/content and draft input state against a non-HiCache or non-CP baseline, not just tensor shape/presence. + +### C56. 2026-05-30 accept-length collapse investigation: current checked facts + +Checked facts: + +- Current decode launch uses non-v2 `EAGLEWorker`, not `EAGLEWorkerV2`: `SpeculativeAlgorithm.create_worker()` selects V2 only when overlap scheduling is enabled, while the inspected decode server args have overlap disabled. V2-only hypotheses around `prepare_for_v2_draft()` are therefore not applicable to the latest run. +- Latest aggregate `Decode batch` accept length remains low but is not literally always 1: in the latest decode logs the average accept length is about `1.33`, median about `1.23-1.25`, with many windows at `<=1.05`. The per-request `EAGLE_ACCEPT_DEBUG` rows are intentionally zero-biased because zero-accept cases are logged much more often than non-zero samples, so they are useful for anomaly examples but not for true distribution estimates. +- Decode processes are not running HiCache/NSA-prefill-CP-shared-KV. The low accept issue is therefore more likely caused by prefill-transferred target/draft KV or prebuilt EAGLE state content than by decode-side HiCache behavior. +- Metadata lifetime is not the obvious failure: latest samples show `has_pd_hidden=True`, `pd_hidden_shape=(6144,)`, and `pd_topk_shape=(16,)`; missing hidden/top-k metadata is not the dominant failure mode. +- `process_prebuilt()` builds EAGLE prebuilt state from the transferred first output token, top-k, and hidden state, then calls `cache_unfinished_req(req)`. `fill_ids` is still the prompt/input tokens prepared before transfer; for EAGLE this is plausibly intentional because unfinished bigram cache should not cache the last row that depends on the generated output. This is a checked contract point, not a fix target unless content evidence proves otherwise. + +Next focused check: + +- Verify whether partial cache-hit prefill has symmetric current-suffix handling for MLA KV and NSA index KV. MLA current-suffix reuse was intentionally enabled after partial cache hits; index KV current-suffix reuse may still be restricted to current-only batches. If index materialization/top-k reads a stale or incomplete persistent index buffer for a cache-hit suffix, EAGLE top-k/hidden can be semantically bad even though all transferred tensors have correct shapes. + +### C57. 2026-05-30 accept-length collapse: target logs do not prove draft KV correctness + +Checked facts: + +- The current `SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH` samples are layer-probe based + and the active probe evidence is from target layer 2. EAGLE/NextN draft has a + single draft layer (`layer_id=0`), so these logs do not prove that draft + partial-current compose is correct or even which draft path is active. +- Latest prefill log does show many target cache-hit partial-current cases. For + small cache-hit prefixes such as `prefix_lens=[320]` (5 pages), the async + prefetcher is absent because the threshold is 8 pages, and target falls back + to the synchronous partial-current compose path. This is expected under the + current threshold policy and should not be re-read as “prefetch is missing”. +- NSA index current reuse is still current-only: `_can_reuse_current_index_kv()` + requires `is_current_only_extend_batch()`, so partial cache-hit requests do not + take the index-current-reuse fast path. However the indexer stores the current + suffix local index KV into the persistent pool before materializing the shared + index buffer for top-k, so the absence of an index partial-current fast path is + currently a performance risk more than a proven accept-length correctness root + cause. +- The remaining high-probability correctness surface is draft KV semantic/content + alignment across prefill cache-hit paths and decode prebuilt consumption: + decode itself is not running CP HiCache, so wrong draft KV transferred from + prefill can suppress EAGLE accept while metadata tensor shapes still look + healthy. + +Do-not-repeat: + +- Do not use target layer-2 MLA prefetch logs as evidence for EAGLE draft layer + correctness. Draft needs its own low-frequency probe or a direct KV/content + comparison. +- Do not treat missing async prefetcher on a 5-page prefix as a bug; it follows + the configured 8-page threshold. The correctness requirement is that sync + partial-current/full-current still works when prefetch is disabled. + +### C58. 2026-05-30 accept-length collapse: prebuilt/draft evidence gaps to close next + +Checked facts: + +- Disaggregated decode constructs the initial EAGLE running batch through + `get_new_prebuilt_batch()`: it first calls `req.init_next_round_input(tree_cache)`, + then `prepare_for_prebuilt()`, then `process_prebuilt()`. +- Because decode has its own radix tree, repeated prompts can enter prebuilt with + a non-zero decode-side `prefix_indices` even though KV has already been + preallocated/transferred for the full prompt. Current logs do not show this + `pre_len`, so we cannot yet correlate low accept length with decode-side + prefix reuse. +- `prepare_for_prebuilt()` builds `input_ids` from `fill_ids[pre_len:]`, but the + transferred `req_to_token` rows are full-prompt rows. If any later prebuilt or + EAGLE path consumes `out_cache_loc`, the correct transferred suffix rows start + at `pre_len`, not at zero. This is a concrete invariant risk, but current + non-overlap EAGLE mostly consumes the transferred req-to-token mapping and + prebuilt EAGLE state after `prepare_for_decode()`, so it still needs runtime + evidence before being treated as the accept-collapse root. +- The latest target-side CP shared-KV logs still do not identify the draft layer + path. We need a draft-layer (`is_draft_input`, layer 0) path log showing + whether EAGLE draft is using current-only, partial-current sync/prefetch, or + full materialize on cache-hit requests. + +Next instrumentation: + +- Under `SGLANG_EAGLE_ACCEPT_DEBUG`, add low-frequency decode prebuilt logs with + `pre_len`, `extend_input_len`, `fill_len`, `origin_len`, `output_len`, and + `cached_tokens`. +- Under `SGLANG_EAGLE_ACCEPT_DEBUG`, add low-frequency draft MLA path logs at + draft layer 0 with branch name, prefix/extend lengths, prefetcher presence, + current rows, KV rows, and page-table shape. Do not use the target layer-2 + MLA prefetch probe as draft evidence. + +### C59. 2026-05-30 latest EAGLE accept debug run: prebuilt prefix reuse is not active, draft prefill uses partial-current sync + +Checked runtime: + +- Latest prefill log: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_223451.log`. +- Latest decode logs: `/mnt/beegfs/cjy/log/decode0_20260529_223458.log`, `/mnt/beegfs/cjy/log/decode1_20260529_223500.log`. +- The new `[EAGLE_ACCEPT_DEBUG][prebuilt_prepare]` and `[EAGLE_ACCEPT_DEBUG][draft_mla_path]` probes are present in the latest run. + +Findings: + +- Decode prebuilt prefix reuse is not active in the inspected window: decode0 had `128` prebuilt debug rows and decode1 had `128`; all had `pre_len=0` even when transferred/cached tokens were large (`cached_max` around `106k-107k`). Therefore the suspected decode-side suffix offset from non-zero `pre_len` is not the current primary explanation. +- EAGLE accept remains low: latest decode logs are around `avg=1.35-1.38`, `median=1.26-1.29`, `p90=1.81-1.90`, with many windows at `<=1.05`. +- No server-side zero-output marker was found in the latest prefill/decode logs. +- Prefill draft MLA path is now directly observed. In the latest prefill log there are `128` draft-layer path rows: `48 current_only` and `80 partial_current_sync`. No draft prefetch path is used, as intended for draft. +- Partial-current draft examples include `prefix_lens=[320]`, large `extend_lens`, `can_current_reuse=True`, `has_prefetcher=False`, `branch=partial_current_sync`. This proves that the re-enabled draft current-suffix compose path is on the hot path for cache-hit EAGLE prefill. + +Next focused question: + +- Determine whether EAGLE draft partial-current compose uses the correct valid row/loc contract after `EagleDraftInput.prepare_for_extend()` shifts input ids left and appends the verified output token. If draft current rows are mapped to target-style `out_cache_loc` rows with an off-by-one semantic shift, the transferred draft KV can have healthy shapes but poor accept length. + +### C60. 2026-05-30 likely accept-collapse root: CP-local EAGLE hidden was padded then re-split as if full hidden + +Checked code path: + +- `EAGLEWorker.forward_target_extend()` stores target hidden for draft in `logits_output.draft_hidden_states` when `SGLANG_CP_DRAFT_SHARED_KV=1`. +- `DeepseekV3ForCausalLM.forward()` captures this side-channel before CP output collection, so the tensor is already CP-local. +- `EAGLEWorker.forward_draft_extend()` passes that CP-local tensor into `EagleDraftInput.hidden_states` for the NextN draft model. +- Before `DeepseekModelNextN.forward()` consumes it, `ForwardBatch.prepare_mlp_sync_batch()` pads `spec_info.hidden_states` to the global padded token count and saves the original tensor in `forward_batch.hidden_states_backup`. +- `DeepseekModelNextN._get_cp_local_spec_hidden_states()` previously checked `spec_hidden_states.shape[0] == full_num_tokens` and then ran `cp_split_and_rebuild_data(...)`. After the padding above, a CP-local tensor could look like a full-token tensor and get CP-split a second time. + +Why this matches the symptom: + +- It affects cache-miss and cache-hit EAGLE prefill, not only HiCache partial-current paths, matching the low accept length across runs. +- It preserves tensor shapes and transfer metadata, matching the earlier observation that hidden/top-k shapes were healthy while accept length remained bad. +- It is introduced by the CP-local draft-hidden optimization, not by decode prebuilt prefix reuse; latest logs show `pre_len=0` for all sampled decode prebuilt rows. + +Current local fix: + +- `DeepseekModelNextN._get_cp_local_spec_hidden_states()` now prefers `forward_batch.hidden_states_backup` when CP draft shared KV is enabled and MLP sync padding changed the visible hidden shape. +- The helper returns CP-local hidden directly, pads only a small local-padding gap, and only splits when the tensor is genuinely full-token hidden. +- A low-frequency `[EAGLE_ACCEPT_DEBUG][cp_draft_hidden]` probe records which branch was taken. + +Validation status: + +- Local `py_compile` for `python/sglang/srt/models/deepseek_nextn.py` passed. +- Needs remote sync/restart and ETE accept-length verification; no commit yet. + +### C61. 2026-05-30 restart after CP-local draft-hidden fix: code path is active but accept length is still low + +Checked runtime: + +- Latest prefill log: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260530_122138.log`. +- Latest decode logs: `/mnt/beegfs/cjy/log/decode0_20260530_122148.log`, `/mnt/beegfs/cjy/log/decode1_20260530_122154.log`. +- The `[EAGLE_ACCEPT_DEBUG][cp_draft_hidden]` probe is present in prefill, proving the restarted process loaded the CP-local hidden backup fix. + +Findings: + +- The CP-local hidden backup path is active. Example rows show `padded_shape=(10120, 6144)` and `backup_shape=(1280, 6144)`, followed by `using CP-local hidden directly`. +- EAGLE accept did **not** recover in this run: + - `decode0_20260530_122148.log`: `n=776`, `avg=1.357`, `median=1.2`, `last100=1.323`, `min=1.0`, `max=3.35`. + - `decode1_20260530_122154.log`: `n=754`, `avg=1.380`, `median=1.2`, `last100=1.418`, `min=1.0`, `max=3.05`. +- Decode prebuilt prefix reuse remains inactive in the sampled window: `128` prebuilt rows per decode log and all sampled `pre_len=0`, even with non-zero `cached_tokens`. +- No server-side zero-output marker or health/scheduler error was found in the latest prefill/decode logs. +- Draft path remains hot and conservative: prefill has `40 current_only` and `96 partial_current_sync` draft-layer debug rows, with no draft prefetcher. +- Only startup/empty-prefix fallback logs were observed (`prefix_len=0` for MLA/index prefetch). These are not evidence of the current accept collapse. + +Conclusion: + +- C60 fixed a real CP-local hidden contract bug, but it is not sufficient to restore accept length. The next root-cause surface is value/content correctness across the EAGLE prefill-to-decode boundary, not tensor shape or decode `pre_len` offset. + +Next focused checks: + +- Compare the EAGLE draft artifacts produced by prefill and consumed by decode at the request boundary: transferred draft KV content/row ordering, `output_topk_p`, `output_topk_index`, and draft hidden alignment. +- Keep the current shape/branch probes, but add only low-frequency value checksums if needed; avoid broad trace spam. + +### C62. 2026-05-30 accept-length collapse: active worker and metadata registration facts + +Checked facts: + +- The latest server args show `enable_multi_layer_eagle=False` and `disable_overlap_schedule=True`, so the active speculative path is the regular non-v2 `EAGLEWorker`. Multi-layer EAGLE top-k packing hypotheses are not active for this run. +- Prefill and decode manager startup logs both register `target_kv_bufs=78`, `draft_kv_bufs=1`, `total_kv_bufs=79`, `target_state_type=nsa`, and one appended draft NSA state buffer. This proves the draft KV/state buffer registration surface is present, but it does not prove values or row ordering are correct. +- Latest decode `prebuilt_prepare` logs remain `pre_len=0`; decode-side radix prefix reuse is not currently shifting the prebuilt input/off-cache mapping. +- The latest accept distribution is still low (`decode0 avg≈1.38`, `decode1 avg≈1.39`) while there are no health/scheduler errors or zero-output markers in the latest prefill/decode logs. + +Conclusion: + +- The next evidence gap is not worker selection, metadata shape, or missing buffer registration. We need value-level evidence at the EAGLE prefill-to-decode handoff: top-k token/probability, hidden checksum, and later draft KV content if metadata values match. + +Next instrumentation: + +- Add low-frequency `[EAGLE_ACCEPT_DEBUG][metadata_set]`, `[metadata_get]`, and `[prebuilt_state]` summaries under `SGLANG_EAGLE_ACCEPT_DEBUG`. +- Keep summaries bounded: request id/room, metadata slot, output id, cached tokens, top-k head, and hidden checksum/head only. Do not add per-layer or per-token spam. + +### C63. 2026-05-30 EAGLE handoff checksum instrumentation added + +Local change: + +- Added a bounded tensor digest helper in `sglang.srt.disaggregation.utils`. +- `MetadataBuffers.set_buf()` now logs `[EAGLE_ACCEPT_DEBUG][metadata_set]` for the prefill-side EAGLE handoff slot. +- Decode transfer commit now logs `[EAGLE_ACCEPT_DEBUG][metadata_get]` for the decode-side view of the same slot. +- `process_prebuilt()` now logs `[EAGLE_ACCEPT_DEBUG][prebuilt_state]` after stacking top-k and hidden into the initial decode EAGLE state. + +Purpose: + +- If `metadata_set`, `metadata_get`, and `prebuilt_state` match for the same request/room, the low accept length is unlikely to be caused by metadata corruption and the next target should be draft KV content/row ordering. +- If they diverge, the root is in metadata transfer/lifetime, not draft KV. + +Verification: + +- Local `py_compile` passed for `disaggregation/utils.py`, `disaggregation/decode.py`, `disaggregation/decode_schedule_batch_mixin.py`, `layers/attention/nsa_backend.py`, and `models/deepseek_nextn.py`. + +Not yet verified: + +- Needs remote sync/restart and one traffic run with `SGLANG_EAGLE_ACCEPT_DEBUG=1` to compare the actual handoff values. + +### C64. 2026-05-30 decode metadata lifetime audit: production free exists, stale aliases remain + +Checked facts: + +- `SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch()` calls + `new_batch.process_prebuilt(...)` and then releases each request's held decode + metadata slot in a `finally` block via `_free_decode_metadata_index_if_held`. +- The earlier suspicion that the decode metadata slot was never freed in + production is false; the production release path exists and is covered by + `test_get_new_prebuilt_batch_slices_waiting_queue_prefix`. +- The release helper currently only frees the allocator index and sets + `metadata_buffer_index=-1`. It does not clear `Req.output_topk_p`, + `Req.output_topk_index`, or `Req.hidden_states_tensor`, so a running request + can still retain stale views into a reusable metadata row after the slot has + been returned to the allocator. + +Risk: + +- The known prebuilt consumer copies top-k and hidden state into GPU tensors + before the slot is released, so this alias is not yet proven to be the + accept-length root cause. +- Keeping stale reusable-slot views on `Req` is still an unsafe lifecycle + contract: later debug, retraction, or future prebuilt paths can observe a + different request's EAGLE state after slot reuse. + +Planned correction: + +- Keep the no-clone transfer hot path, but make metadata release also detach the + request from reusable-slot views by clearing the three EAGLE handoff fields + after `process_prebuilt` has consumed them. +- Preserve the existing release timing; do not move release back to + `pop_transferred` and do not reintroduce transfer-time clones unless the new + handoff checksums prove process-prebuilt copy is not sufficient. + +### C65. 2026-05-30 latest remote run still lacks new metadata checksums + +Checked facts: + +- Remote files under `/mnt/beegfs/cjy/sglang-dev` match the local hashes for the + new `metadata_set`, `metadata_get`, and `prebuilt_state` instrumentation. +- The currently running logs (`sglang_cp_hicache_20260530_122138.log`, + `decode0_20260530_122148.log`, `decode1_20260530_122154.log`) contain the + older `draft_mla_path`, `cp_draft_hidden`, and `prebuilt_prepare` markers but + no `metadata_set`, `metadata_get`, or `prebuilt_state` markers. Therefore the + live process predates the latest handoff-checksum instrumentation. +- Current accept remains low in that old process: decode0 average around 1.38 + and decode1 average around 1.39 accepted length. + +Next evidence gate after restart: + +- Compare `metadata_set -> metadata_get -> prebuilt_state` for the same request. + If these match while accept remains low, the next root surface is draft KV + content/row ordering across prefill-to-decode transfer, not metadata lifetime + or shape. + +C64 correction applied: + +- `_free_decode_metadata_index_if_held` now clears `Req.output_topk_p`, + `Req.output_topk_index`, and `Req.hidden_states_tensor` after returning the + metadata slot to the allocator. +- The transfer hot path still stores views and does not clone; the request owns + the slot until `process_prebuilt` has synchronously consumed the tensors. +- Unit expectations were extended so future changes cannot reintroduce stale + reusable-slot aliases on long-lived decode requests. + +C64 verification note: + +- Local `py_compile` for the touched Python modules passed. +- Local pytest collection for `test_decode_queue_compaction.py` is blocked in + this environment by missing `orjson`; use the g0034 container for this test. + +### C66. 2026-05-30 metadata checksum debug guard compatibility + +Checked fact: + +- The remote decode queue compaction test uses a minimal fake `spec_algorithm` + with `is_none()` but without `is_eagle()`. The production decode transfer + path already gates EAGLE metadata ownership on `not spec_algorithm.is_none()`; + the new checksum-only debug branch used the narrower `is_eagle()` method and + therefore broke this test double without changing production semantics. + +Correction: + +- The `[EAGLE_ACCEPT_DEBUG][metadata_get]` guard now follows the existing + transfer ownership contract: emit only when speculative decoding is active + (`not spec_algorithm.is_none()`) and the debug sampler allows it. +- This is debug-only; it does not change metadata slot lifetime or request + scheduling behavior. + +Verification: + +- Local `py_compile` passed for `decode.py`, the metadata release helper file, + and the decode queue compaction test. +- Remote container verification on g0034 passed: + `PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py` + -> `11 passed, 5 warnings`. + +### C67. 2026-05-30 exact non-page CP cache hits still expose sub-page tails + +Finding: + +- C16/C54 floor CP radix hits when a request extends beyond an existing + non-page-aligned valid-tail node, but an exact key hit still exposes the full + token-precise tail. +- In normal prefill scheduling this is still an extend: the radix key is capped at + `input_len - 1`, so an exact EAGLE/bigram key hit with length `N` means the + model will compute at least the next current token while the CP shared-KV + backend sees `extend_prefix_len=N`. +- If `N` is not page-aligned, CP shared-KV current reuse either cannot use the + page-aligned prefix contract or must enter a mostly-zero/tiny split shape. + This matches the recurring symptom pattern where repeated cache-hit requests + are page-tail sized and EAGLE accept length collapses without transfer errors. + +Decision: + +- Treat CP HiCache radix cache as page-granular for scheduler-visible hits: an + exact non-page valid-tail hit is floored to the previous page boundary, just + like the extension case. +- The sub-page tail is sacrificed and recomputed by the new request. This costs + at most one page of reusable prefix, but keeps `prefix_indices`, + `cache_protected_len`, prepared backup start, and CP shared-KV materialization + on the same page-aligned contract. + +Regression first: + +- Added failing tests on g0034 before production changes: + - `test_cp_match_prefix_floors_exact_valid_tail_for_exact_key` + - `test_cp_insert_replaces_exact_valid_tail_from_page_boundary` +- Both failed under the old helper: match exposed `[0..5]` and insert reported + `prefix_len=6` for a 6-token node on 4-token pages. + +C67 implementation update: + +- Added `MatchPrefixParams.cp_floor_exact` so scheduler-visible CP prefix matches can + floor exact non-page valid-tail hits, while the internal `cache_unfinished_req()` + post-insert refresh can keep its own just-inserted exact tail and avoid + invalidating current-request accounting. +- `HiRadixCache._cp_floor_exact_valid_tail_extension_len()` now floors exact + non-page hits when `cp_floor_exact=True` or when the request key extends beyond + the child. Prepared-backup insertion enables the same exact-key floor so the + recomputed page-tail backup attaches to the replacement suffix node. +- `RadixCache.cache_unfinished_req()` calls its internal refresh match with + `cp_floor_exact=False`; this preserves the previous exact-duplicate accounting + fix for repeated tiny EAGLE requests. +- Existing CP host-hit expectations were updated from token-granular `6` to + page-granular `4` for a 6-token backed node on 4-token pages. + +C67 verification: + +```text +remote g0034 container: + PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_hicache_metadata.py \ + test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py + -> 109 passed, 5 warnings +``` + +### C68. 2026-05-30 accept-length collapse correlates with draft partial-current reuse + +Finding: + +- Latest available remote process is no longer running; `/mnt/beegfs/cjy/log` still contains the previous run as the newest evidence: + - Prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260530_122138.log` + - Decode: `/mnt/beegfs/cjy/log/decode0_20260530_122148.log`, `/mnt/beegfs/cjy/log/decode1_20260530_122154.log` +- Compared with the earlier baseline-ish run (`20260529_1605xx`), the failing run differs on the EAGLE/NextN draft MLA path: + - Earlier prefill `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_160555.log`: `draft_partial_current_reuse_disabled=64`, `branch=partial_current_sync=0`. + - Failing prefill `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260530_122138.log`: `draft_partial_current_reuse_disabled=0`, `branch=partial_current_sync=112`, `branch=current_only=48`. +- Decode accept distribution regressed at the same time: + - `decode0_20260529_160608.log`: `n=502`, `avg=1.324`, `median=1.18`, `last100=1.424`, `zero=0`. + - `decode1_20260529_160611.log`: `n=468`, `avg=1.348`, `median=1.20`, `last100=1.351`, `zero=0`. + - `decode0_20260530_122148.log`: `n=167806`, `avg=0.068`, `median=0.0`, `last100=0.01`, `zero=158223`. + - `decode1_20260530_122154.log`: `n=160622`, `avg=0.068`, `median=0.0`, `last100=0.011`, `zero=151281`. +- The failing logs predate the newest `metadata_set/get/prebuilt_state` checksum probes, so metadata handoff values are still unverified in a fresh process. + +Working hypothesis: + +- Target-model partial current reuse can remain enabled, but EAGLE/NextN draft partial current reuse is not yet a validated correctness contract. +- For draft cache-hit suffixes, the safer contract is the previous full-materialize path: direct write the local draft suffix into the draft pool, then materialize the full prefix+suffix from the pool/all-rank compose. This is slower than partial splice but matched the earlier non-collapsed accept distribution. + +Planned correction: + +- Keep draft current-only reuse enabled for empty-prefix batches. +- Disable only draft partial-current splice (`prefix_len > 0`) and emit an explicit `[CP_SHARED_KV_FALLBACK][current_reuse] reason=draft_partial_current_reuse_disabled` warning when this gate fires. +- Do not add new collectives; this restores the existing full-materialize path rather than inventing a new synchronization path. + +C68 implementation update: + +- `should_reuse_current_extend_kv()` now returns `True` immediately for current-only batches, preserving draft and target empty-prefix reuse. +- For partial-current batches, EAGLE/NextN draft inputs (`spec_info.is_draft_input()`) now return `False` and emit the explicit fallback warning: + `[CP_SHARED_KV_FALLBACK][current_reuse] reason=draft_partial_current_reuse_disabled`. +- Target partial-current reuse remains enabled. + +C68 verification: + +```text +remote g0034 container: + PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_should_reuse_current_extend_kv_disables_draft_partial_cache_hit_suffix + -> 1 passed + + PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py + -> 73 passed, 5 warnings, 2 subtests passed +``` + +Runtime validation still needed: + +- Start a fresh prefill/decode process with the synced code and `SGLANG_EAGLE_ACCEPT_DEBUG=1`. +- Expected prefill evidence: draft cache-hit suffix rows should log `draft_partial_current_reuse_disabled`; draft `branch=partial_current_sync` should disappear except for old logs. +- Expected decode evidence: accept distribution should move away from the `avg≈0.068 / median=0` collapse if C68 is the active root cause. diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ee0e283d7..741a92bd7 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -42,6 +42,8 @@ from sglang.srt.disaggregation.utils import ( ReqToMetadataIdxAllocator, TransferBackend, append_cp_draft_state_buffers, + eagle_accept_debug_should_log, + eagle_accept_debug_tensor_digest, get_kv_class, is_mla_backend, poll_and_all_reduce, @@ -1172,6 +1174,25 @@ class DecodeTransferQueue: return True # Case 3: Success - commit the transfer + if ( + not self.spec_algorithm.is_none() + and eagle_accept_debug_should_log("metadata_get") + ): + logger.warning( + "[EAGLE_ACCEPT_DEBUG][metadata_get] rid=%s room=%s idx=%s " + "output_id=%s cached_tokens=%s actual_room=%s topk_p=%s " + "topk_index=%s hidden=%s", + str(getattr(decode_req.req, "rid", ""))[:8], + decode_req.req.bootstrap_room, + idx, + int(output_id[0].item()), + int(cached_tokens[0].item()), + int(actual_room), + eagle_accept_debug_tensor_digest(output_topk_p[:1]), + eagle_accept_debug_tensor_digest(output_topk_index[:1]), + eagle_accept_debug_tensor_digest(output_hidden_states), + ) + decode_req.req.output_ids.append(output_id[0].item()) decode_req.req.cached_tokens = cached_tokens[0].item() if not self.spec_algorithm.is_none(): diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 342296585..30e10abd4 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -6,10 +6,16 @@ from typing import TYPE_CHECKING import torch +from sglang.srt.disaggregation.utils import ( + eagle_accept_debug_should_log, + eagle_accept_debug_tensor_digest, +) +from sglang.srt.environ import envs from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo logger = logging.getLogger(__name__) +_EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER = 0 if TYPE_CHECKING: from sglang.srt.managers.overlap_utils import FutureMap @@ -42,6 +48,7 @@ class ScheduleBatchDisaggregationDecodeMixin: for i, req in enumerate(reqs): req_pool_indices.append(req.req_pool_idx) + pre_len = len(req.prefix_indices) chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][ : req.extend_input_len ] @@ -51,7 +58,6 @@ class ScheduleBatchDisaggregationDecodeMixin: out_cache_loc[offset : offset + req.extend_input_len] = chunk offset += req.extend_input_len - pre_len = len(req.prefix_indices) seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1) seq_lens.append(seq_len) if len(req.output_ids) == 0: @@ -66,6 +72,35 @@ class ScheduleBatchDisaggregationDecodeMixin: pre_lens.append(pre_len) req.extend_logprob_start_len = 0 + if envs.SGLANG_EAGLE_ACCEPT_DEBUG.get(): + global _EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER + _EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER += 1 + counter = _EAGLE_ACCEPT_PREBUILT_DEBUG_COUNTER + if ( + counter <= 16 + or counter % 256 == 0 + or (pre_len > 0 and (counter <= 128 or counter % 128 == 0)) + ): + logger.warning( + "[EAGLE_ACCEPT_DEBUG][prebuilt_prepare] rid=%s " + "pre_len=%s extend_input_len=%s fill_len=%s origin_len=%s " + "output_len=%s seq_len=%s cached_tokens=%s " + "req_pool_idx=%s out_chunk_start=%s expected_suffix_start=%s " + "out_chunk_len=%s", + str(getattr(req, "rid", ""))[:8], + pre_len, + req.extend_input_len, + len(req.fill_ids), + len(req.origin_input_ids), + len(req.output_ids), + seq_len, + int(getattr(req, "cached_tokens", 0) or 0), + req.req_pool_idx, + 0, + pre_len, + int(chunk.numel()), + ) + extend_input_logprob_token_ids = None # Set fields @@ -161,6 +196,29 @@ class ScheduleBatchDisaggregationDecodeMixin: hidden_states_list = [req.hidden_states_tensor for req in self.reqs] hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device) + if eagle_accept_debug_should_log("prebuilt_state"): + req0 = self.reqs[0] if self.reqs else None + logger.warning( + "[EAGLE_ACCEPT_DEBUG][prebuilt_state] rid=%s bs=%s " + "num_states=%s output_ids=%s seq_lens=%s topk_p=%s " + "topk_index=%s hidden=%s", + str(getattr(req0, "rid", ""))[:8] if req0 is not None else None, + len(self.reqs), + num_states, + self.output_ids[: min(4, self.output_ids.numel())].detach() + .cpu() + .tolist(), + self.seq_lens[: min(4, self.seq_lens.numel())].detach() + .cpu() + .tolist(), + eagle_accept_debug_tensor_digest(topk_p[: min(1, topk_p.shape[0])]), + eagle_accept_debug_tensor_digest( + topk_index[: min(1, topk_index.shape[0])] + ), + eagle_accept_debug_tensor_digest( + hidden_states[: min(1, hidden_states.shape[0])] + ), + ) # local import to avoid circular import from sglang.srt.speculative.eagle_info import EagleDraftInput diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 4552a0090..c15fff9a3 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -2,6 +2,7 @@ from __future__ import annotations import os import random +import logging from collections import deque from contextlib import nullcontext from enum import Enum @@ -14,6 +15,8 @@ import torch.distributed as dist from sglang.srt.environ import envs from sglang.srt.utils import is_npu +logger = logging.getLogger(__name__) + if TYPE_CHECKING: from sglang.srt.disaggregation.base.conn import KVArgs from sglang.srt.disaggregation.common.conn import ( @@ -84,6 +87,63 @@ def poll_and_all_reduce_attn_cp_tp_group( # Metadata Buffers ######################### +_EAGLE_ACCEPT_DEBUG_COUNTERS = {} + + +def eagle_accept_debug_should_log( + key: str, + *, + first: int = 16, + every: int = 256, +) -> bool: + if not envs.SGLANG_EAGLE_ACCEPT_DEBUG.get(): + return False + count = _EAGLE_ACCEPT_DEBUG_COUNTERS.get(key, 0) + 1 + _EAGLE_ACCEPT_DEBUG_COUNTERS[key] = count + return count <= first or (every > 0 and count % every == 0) + + +def eagle_accept_debug_tensor_digest(tensor: Any, *, sample: int = 64) -> str: + """Small deterministic tensor summary for EAGLE handoff debugging. + + This intentionally samples only a prefix. It is enabled only under + SGLANG_EAGLE_ACCEPT_DEBUG and is for handoff equality checks, not for full + numerical validation. + """ + + if tensor is None: + return "None" + try: + t = torch.as_tensor(tensor) + except Exception as exc: # pragma: no cover - defensive debug helper + return f"unavailable({type(exc).__name__})" + shape = tuple(t.shape) + dtype = str(t.dtype).replace("torch.", "") + numel = int(t.numel()) + if numel == 0: + return f"shape={shape} dtype={dtype} numel=0" + + flat = t.detach().reshape(-1) + sample_count = min(sample, numel) + try: + sample_cpu = flat[:sample_count].cpu() + head_cpu = sample_cpu[: min(8, sample_count)] + head = head_cpu.tolist() + if sample_cpu.is_floating_point(): + checksum = float(sample_cpu.float().sum().item()) + abs_checksum = float(sample_cpu.float().abs().sum().item()) + return ( + f"shape={shape} dtype={dtype} sample={sample_count} " + f"sum={checksum:.6g} abs={abs_checksum:.6g} head={head}" + ) + checksum = int(sample_cpu.to(torch.int64).sum().item()) + return ( + f"shape={shape} dtype={dtype} sample={sample_count} " + f"sum={checksum} head={head}" + ) + except Exception as exc: # pragma: no cover - defensive debug helper + return f"shape={shape} dtype={dtype} numel={numel} digest_error={type(exc).__name__}" + def append_cp_draft_state_buffers( kv_args: Any, @@ -307,6 +367,27 @@ class MetadataBuffers: self.output_hidden_states[req.metadata_buffer_index].copy_( req.hidden_states_tensor ) + if eagle_accept_debug_should_log("metadata_set"): + logger.warning( + "[EAGLE_ACCEPT_DEBUG][metadata_set] rid=%s room=%s idx=%s " + "output_id=%s cached_tokens=%s topk=%s topk_p=%s " + "topk_index=%s hidden=%s", + str(getattr(req, "rid", ""))[:8], + getattr(req, "bootstrap_room", None), + req.metadata_buffer_index, + req.output_ids[0] if req.output_ids else None, + req.cached_tokens, + topk, + eagle_accept_debug_tensor_digest( + self.output_topk_p[req.metadata_buffer_index, :topk] + ), + eagle_accept_debug_tensor_digest( + self.output_topk_index[req.metadata_buffer_index, :topk] + ), + eagle_accept_debug_tensor_digest( + self.output_hidden_states[req.metadata_buffer_index] + ), + ) # Store bootstrap_room for validation on decode side self.bootstrap_room[req.metadata_buffer_index, 0] = ( req.bootstrap_room if req.bootstrap_room is not None else 0 diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index c0bd9569e..8a350fabe 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -978,20 +978,38 @@ def can_reuse_current_extend_kv(forward_batch) -> bool: def should_reuse_current_extend_kv(forward_batch) -> bool: """Return whether MLA should splice current extend KV into materialized KV. - The contract is model-role agnostic: target and prefill-time EAGLE/NextN - draft both write only the CP-owned suffix rows into the persistent pool, so - cache-hit attention should materialize the page-aligned prefix and splice the - freshly computed current suffix explicitly. Draft async prefetch remains - disabled elsewhere because NextN has no next layer to prefetch, but disabling - that optimization must not force draft cache-hit suffixes back to full - materialization from the pool. + Current-only reuse is safe for both target and draft because there is no + cached prefix to compose. Partial current reuse is currently a target-model + contract only. EAGLE/NextN draft cache-hit suffixes keep using the older + full-materialize path until the draft splice path has value-level ETE proof; + the 2026-05-30 accept-length regression correlated with enabling that draft + partial splice path. """ if not cp_shared_kv_current_reuse_enabled(): return False current_only = is_current_only_extend_batch(forward_batch) + if current_only: + return True + partial_current = can_reuse_current_extend_kv(forward_batch) + if partial_current and cp_shared_kv_is_draft_input(forward_batch): + prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + _log_current_reuse_fallback( + "draft_partial_current_reuse_disabled", + "cache-hit EAGLE/NextN draft uses full materialize instead of " + "partial current reuse. prefix_lens=%s extend_lens=%s", + [int(x) for x in prefix_lens_cpu] + if prefix_lens_cpu is not None + else None, + [int(x) for x in extend_lens_cpu] + if extend_lens_cpu is not None + else None, + ) + return False + return current_only or partial_current diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 7963ed774..e972ad75d 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from dataclasses import dataclass from enum import IntEnum, auto from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, TypeAlias @@ -68,6 +69,8 @@ if TYPE_CHECKING: _is_hip = is_hip() +logger = logging.getLogger(__name__) +_EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS: Dict[Tuple[int, int], int] = {} if _is_hip: from sglang.srt.layers.attention.nsa.triton_kernel import get_valid_kv_indices @@ -1737,6 +1740,10 @@ class NativeSparseAttnBackend( and topk_transform_method == TopkTransformMethod.PAGED ): assert forward_batch.cp_shared_kv_layout is not None + is_draft_mla_input = cp_shared_kv_is_draft_input(forward_batch) + eagle_draft_mla_branch = "not_selected" + eagle_draft_mla_used_prefetch = False + eagle_draft_mla_current_rows = None mla_prefetcher = getattr( forward_batch, "cp_shared_kv_mla_prefetcher", None ) @@ -1789,6 +1796,7 @@ class NativeSparseAttnBackend( ) if is_current_only_extend_batch(forward_batch): + eagle_draft_mla_branch = "current_only" current_mask, page_table_1 = build_current_loc_remap( logical_page_table_1, forward_batch.out_cache_loc, @@ -1846,8 +1854,11 @@ class NativeSparseAttnBackend( current_remap_logical_page_capacity=current_remap_logical_page_capacity, ) if prefetched_kv is not None: + eagle_draft_mla_branch = "partial_current_prefetch" + eagle_draft_mla_used_prefetch = True kv_cache, page_table_1 = prefetched_kv else: + eagle_draft_mla_branch = "partial_current_sync" prefix_lens_cpu = getattr( forward_batch, "extend_prefix_lens_cpu", None ) @@ -1969,7 +1980,9 @@ class NativeSparseAttnBackend( tensor_debug_checksum(k), tensor_debug_checksum(k_rope), ) + eagle_draft_mla_current_rows = int(current_kv_cache.shape[0]) else: + eagle_draft_mla_branch = "full_materialize" prefetched_kv = None if mla_prefetcher is not None: prefetched_kv = mla_prefetcher.consume( @@ -1978,6 +1991,8 @@ class NativeSparseAttnBackend( logical_locs=page_table_1, ) if prefetched_kv is not None: + eagle_draft_mla_branch = "full_materialize_prefetch" + eagle_draft_mla_used_prefetch = True kv_cache, page_table_1 = prefetched_kv else: slot_remap = get_or_build_shared_token_kv_slot_remap( @@ -1998,6 +2013,57 @@ class NativeSparseAttnBackend( nvtx_source="mla.full_materialize", nvtx_layer_id=layer.layer_id, ) + if ( + envs.SGLANG_EAGLE_ACCEPT_DEBUG.get() + and is_draft_mla_input + and int(layer.layer_id) == 0 + ): + debug_key = ( + int(forward_batch.cp_shared_kv_layout.cp_rank), + int(layer.layer_id), + ) + debug_count = ( + _EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS.get(debug_key, 0) + 1 + ) + _EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS[debug_key] = debug_count + if ( + debug_count <= 16 + or eagle_draft_mla_branch in ("full_materialize", "not_selected") + or debug_count % 256 == 0 + ): + prefix_lens_cpu = getattr( + forward_batch, "extend_prefix_lens_cpu", None + ) + extend_lens_cpu = getattr( + forward_batch, "extend_seq_lens_cpu", None + ) + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + logger.warning( + "[EAGLE_ACCEPT_DEBUG][draft_mla_path] cp_rank=%s " + "layer=%s count=%s branch=%s used_prefetch=%s " + "has_prefetcher=%s can_current_reuse=%s prefix_lens=%s " + "extend_lens=%s current_rows=%s kv_rows=%s " + "page_table_shape=%s out_cache_loc_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + debug_count, + eagle_draft_mla_branch, + eagle_draft_mla_used_prefetch, + mla_prefetcher is not None, + can_reuse_current_kv, + [int(x) for x in prefix_lens_cpu] + if prefix_lens_cpu is not None + else None, + [int(x) for x in extend_lens_cpu] + if extend_lens_cpu is not None + else None, + eagle_draft_mla_current_rows, + int(kv_cache.shape[0]) if kv_cache is not None else None, + tuple(page_table_1.shape) + if page_table_1 is not None + else None, + tuple(out_cache_loc.shape) if out_cache_loc is not None else None, + ) if mla_prefetcher is not None and cp_shared_kv_should_prefetch_next_layer( forward_batch, layer.layer_id ): diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index ec021482b..16e745f22 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -63,6 +63,13 @@ class SchedulerOutputProcessorMixin: ) allocator.free(idx) req.metadata_buffer_index = -1 + # The EAGLE handoff tensors are views into the reusable metadata slot. + # process_prebuilt consumes them into the batch before this helper is + # called; keeping the views afterwards makes the request observe a + # future transfer that reuses the same slot. + req.output_topk_p = None + req.output_topk_index = None + req.hidden_states_tensor = None def _maybe_log_eagle_accept_debug( self: Scheduler, diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 4462e18c3..66e09128e 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -37,6 +37,7 @@ class MatchPrefixParams: """Unified parameters for match_prefix across different cache types""" key: RadixKey + cp_floor_exact: bool = True # Mamba specific cow_mamba: bool = False diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index d5f7898f4..b1b855d08 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2109,7 +2109,7 @@ class HiRadixCache(RadixCache): child = node.children[child_key] prefix_len = self.key_match_fn(child.key, key) prefix_len = self._cp_floor_exact_valid_tail_extension_len( - child, prefix_len, len(key) + child, prefix_len, len(key), floor_exact_key=True ) if prefix_len <= 0: break @@ -2144,14 +2144,35 @@ class HiRadixCache(RadixCache): return prefix_len // self.page_size * self.page_size def _cp_floor_exact_valid_tail_extension_len( - self, child: TreeNode, prefix_len: int, request_len: int + self, + child: TreeNode, + prefix_len: int, + request_len: int, + *, + floor_exact_key: bool = False, ) -> int: + """Floor exact CP valid-tail hits to the previous physical page. + + CP HiCache may keep radix keys at scheduler-visible valid lengths while + the underlying target/draft pools and host reservations are page-owned. + Even when the incoming radix key exactly equals a non-page-aligned child + key, the scheduler will still compute the current suffix token + (`max_prefix_len = input_len - 1`). Exposing the sub-page tail as a + protected prefix lets later backup/current-reuse paths start inside a + page. Sacrifice that tail and let the new request re-own it. + + `floor_exact_key` is enabled for scheduler-visible prefix matching and + prepared-backup probing. Internal cache insertion refreshes can keep + exact sub-page tails for the current request so they do not immediately + invalidate their own just-inserted prefix. + """ + if ( not self._uses_cp_hicache or self.page_size <= 1 or prefix_len <= 0 or prefix_len != len(child.key) - or prefix_len >= request_len + or (prefix_len >= request_len and not floor_exact_key) or prefix_len % self.page_size == 0 ): return prefix_len @@ -3505,7 +3526,11 @@ class HiRadixCache(RadixCache): deferred_node = None try: - value, last_node = self._match_prefix_helper(self.root_node, key) + value, last_node = self._match_prefix_helper( + self.root_node, + key, + floor_exact_key=getattr(params, "cp_floor_exact", True), + ) except HiCachePendingBackupSplit as exc: value = [] last_node = exc.node.parent if exc.node.parent is not None else self.root_node @@ -3633,7 +3658,9 @@ class HiRadixCache(RadixCache): return matched_length - def _match_prefix_helper(self, node: TreeNode, key: RadixKey): + def _match_prefix_helper( + self, node: TreeNode, key: RadixKey, *, floor_exact_key: bool = True + ): node.last_access_time = time.monotonic() child_key = self.get_child_key_fn(key) value = [] @@ -3646,7 +3673,10 @@ class HiRadixCache(RadixCache): child.pin_expiry = time.monotonic() + child.pin_ttl raw_prefix_len = self.key_match_fn(child.key, key) prefix_len = self._cp_floor_exact_valid_tail_extension_len( - child, raw_prefix_len, len(key) + child, + raw_prefix_len, + len(key), + floor_exact_key=floor_exact_key, ) stop_after_page_floor = prefix_len != raw_prefix_len prune_stale_tail_after_split = stop_after_page_floor @@ -3765,7 +3795,10 @@ class HiRadixCache(RadixCache): node.priority = max(node.priority, priority) raw_prefix_len = self.key_match_fn(node.key, key) prefix_len = self._cp_floor_exact_valid_tail_extension_len( - node, raw_prefix_len, len(key) + node, + raw_prefix_len, + len(key), + floor_exact_key=prepared_cp_backup is not None, ) stop_after_page_floor = prefix_len != raw_prefix_len if prefix_len <= 0: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index f5f79e364..6765bd548 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -653,7 +653,9 @@ class RadixCache(BasePrefixCache): ) # The prefix indices could be updated, reuse it - match_result = self.match_prefix(MatchPrefixParams(key=radix_key)) + match_result = self.match_prefix( + MatchPrefixParams(key=radix_key, cp_floor_exact=False) + ) new_indices, new_last_node = ( match_result.device_indices, match_result.last_device_node, diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index f349d9194..5a0927ff8 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -63,6 +63,19 @@ logger = logging.getLogger(__name__) _is_cuda = is_cuda() _is_npu = is_npu() +_EAGLE_ACCEPT_CP_DRAFT_HIDDEN_DEBUG_COUNTS = {} + + +def _log_eagle_accept_cp_draft_hidden_debug(key: str, message: str, *args): + if not envs.SGLANG_EAGLE_ACCEPT_DEBUG.get(): + return + count = _EAGLE_ACCEPT_CP_DRAFT_HIDDEN_DEBUG_COUNTS.get(key, 0) + 1 + _EAGLE_ACCEPT_CP_DRAFT_HIDDEN_DEBUG_COUNTS[key] = count + if count <= 16 or count % 256 == 0: + logger.warning( + "[EAGLE_ACCEPT_DEBUG][cp_draft_hidden] " + message, + *args, + ) class DeepseekModelNextN(nn.Module): @@ -150,10 +163,75 @@ class DeepseekModelNextN(nn.Module): self._debug_cp_draft_shared_kv("fallback reason=missing_spec_hidden") return None + padded_spec_hidden_shape = tuple(spec_hidden_states.shape) + hidden_states_backup = getattr(forward_batch, "hidden_states_backup", None) + if ( + envs.SGLANG_CP_DRAFT_SHARED_KV.get() + and hidden_states_backup is not None + and hidden_states_backup.shape[0] != spec_hidden_states.shape[0] + ): + # ForwardBatch.prepare_mlp_sync_batch pads EagleDraftInput.hidden_states + # to the global padded token count before model.forward() runs. For + # CP-local draft, the target side-channel is already CP-local; if we + # let the padded tensor look like a full-token tensor, the branch + # below will CP-split it a second time and corrupt draft features. + spec_hidden_states = hidden_states_backup + _log_eagle_accept_cp_draft_hidden_debug( + "use_backup", + "using pre-pad CP-local hidden backup. cp_rank=%s " + "padded_shape=%s backup_shape=%s full_tokens=%s local_tokens=%s", + get_attention_cp_rank(), + padded_spec_hidden_shape, + tuple(spec_hidden_states.shape), + full_num_tokens, + local_num_tokens, + ) + if spec_hidden_states.shape[0] == local_num_tokens: + _log_eagle_accept_cp_draft_hidden_debug( + "local_direct", + "using CP-local hidden directly. cp_rank=%s shape=%s " + "full_tokens=%s local_tokens=%s", + get_attention_cp_rank(), + tuple(spec_hidden_states.shape), + full_num_tokens, + local_num_tokens, + ) return spec_hidden_states + if spec_hidden_states.shape[0] < local_num_tokens: + pad_rows = local_num_tokens - spec_hidden_states.shape[0] + if pad_rows <= max(get_attention_cp_size(), 1): + _log_eagle_accept_cp_draft_hidden_debug( + "local_pad", + "padding CP-local hidden to local token count. cp_rank=%s " + "shape=%s pad_rows=%s full_tokens=%s local_tokens=%s", + get_attention_cp_rank(), + tuple(spec_hidden_states.shape), + pad_rows, + full_num_tokens, + local_num_tokens, + ) + return torch.cat( + ( + spec_hidden_states, + spec_hidden_states.new_zeros( + pad_rows, *spec_hidden_states.shape[1:] + ), + ), + dim=0, + ) + if spec_hidden_states.shape[0] == full_num_tokens: + _log_eagle_accept_cp_draft_hidden_debug( + "full_split", + "splitting full hidden for CP-local draft. cp_rank=%s shape=%s " + "full_tokens=%s local_tokens=%s", + get_attention_cp_rank(), + tuple(spec_hidden_states.shape), + full_num_tokens, + local_num_tokens, + ) return cp_split_and_rebuild_data(forward_batch, spec_hidden_states) self._debug_cp_draft_shared_kv( diff --git a/test/registered/unit/disaggregation/test_decode_queue_compaction.py b/test/registered/unit/disaggregation/test_decode_queue_compaction.py index fb507afb1..6e743918d 100644 --- a/test/registered/unit/disaggregation/test_decode_queue_compaction.py +++ b/test/registered/unit/disaggregation/test_decode_queue_compaction.py @@ -428,6 +428,9 @@ class TestDecodeQueueCompaction(CustomTestCase): allocator = FakeAllocator() req = FakeReq("eagle", 3) req.metadata_buffer_index = 9 + req.output_topk_p = torch.ones((1,), dtype=torch.float32) + req.output_topk_index = torch.ones((1,), dtype=torch.int64) + req.hidden_states_tensor = torch.ones((4,), dtype=torch.float32) scheduler = SchedulerOutputProcessorMixin.__new__(SchedulerOutputProcessorMixin) scheduler.req_to_metadata_buffer_idx_allocator = allocator @@ -437,6 +440,9 @@ class TestDecodeQueueCompaction(CustomTestCase): self.assertEqual(allocator.freed, [9]) self.assertEqual(req.metadata_buffer_index, -1) + self.assertIsNone(req.output_topk_p) + self.assertIsNone(req.output_topk_index) + self.assertIsNone(req.hidden_states_tensor) def test_resume_retracted_reqs_compacts_queue_in_one_pass(self): prealloc_queue = DecodePreallocQueue.__new__(DecodePreallocQueue) @@ -689,6 +695,9 @@ class TestDecodeQueueCompaction(CustomTestCase): self.assertEqual(allocator.freed, [20, 21, 22]) for req in captured["reqs"]: self.assertEqual(req.metadata_buffer_index, -1) + self.assertIsNone(getattr(req, "output_topk_p", None)) + self.assertIsNone(getattr(req, "output_topk_index", None)) + self.assertIsNone(getattr(req, "hidden_states_tensor", None)) def test_get_new_prebuilt_batch_keeps_waiting_queue_when_no_capacity(self): scheduler = cast(Any, SimpleNamespace()) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index e502f5f58..f5b0dfe23 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -2596,9 +2596,12 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6))))) self.assertEqual(result.device_indices.tolist(), []) - self.assertEqual(result.host_hit_length, 6) + # CP HiCache exposes page-granular host hits. A non-page exact tail is + # floored and the tail is recomputed by the incoming request. + self.assertEqual(result.host_hit_length, 4) self.assertIs(result.last_device_node, root) - self.assertIs(result.last_host_node, node) + self.assertEqual(result.last_host_node.key.token_ids, [0, 1, 2, 3]) + self.assertIsNot(result.last_host_node, node) def test_cp_backed_tail_split_floors_to_page_boundary(self): cache = HiRadixCache.__new__(HiRadixCache) @@ -2811,6 +2814,47 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3]) self.assertNotIn((4, 5), result.last_device_node.children) + def test_cp_match_prefix_floors_exact_valid_tail_for_exact_key(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.device = "cpu" + cache.disable = False + cache.page_size = 4 + cache.ongoing_write_through = {} + cache.pending_host_backups = {} + cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) + cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) + cache.key_match_fn = lambda child_key, key: _key_match_paged( + child_key, key, page_size=4 + ) + cache.maybe_bigram_convert = lambda key: (key, None) + cache._update_leaf_status = lambda node: None + cache._update_host_leaf_status = lambda node: None + cache.enable_storage = False + cache.enable_kv_cache_events = False + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.host_len = 0 + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 149 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = torch.arange(6, dtype=torch.int64) + node.host_value = None + node.host_len = 0 + node.cp_hicache = None + root.children[(0, 1, 2, 3)] = node + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(6))))) + + self.assertEqual(result.device_indices.tolist(), [0, 1, 2, 3]) + self.assertEqual(result.host_hit_length, 0) + self.assertEqual(result.last_device_node.key.token_ids, [0, 1, 2, 3]) + self.assertNotIn((4, 5), result.last_device_node.children) + def test_cp_insert_floors_backed_tail_split_to_page_boundary(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = True @@ -2938,6 +2982,72 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertTrue(prepared.attached) self.assertIs(cache.pending_host_backups[148].node, new_tail) + def test_cp_insert_replaces_exact_valid_tail_from_page_boundary(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.disable = False + cache.is_eagle = False + cache.page_size = 4 + cache.pending_host_backups = {} + cache.ongoing_write_through = {} + cache.cache_controller = types.SimpleNamespace( + has_draft_hicache=False, + write_policy="write_through", + ) + cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4]) + cache.key_match_fn = lambda child_key, key: _key_match_paged( + child_key, key, page_size=4 + ) + cache.maybe_bigram_convert = lambda key, value=None: (key, value) + cache._update_leaf_status = lambda node: None + cache._update_host_leaf_status = lambda node: None + cache._inc_hit_count = lambda *args, **kwargs: None + cache._record_store_event = lambda node: None + cache.evictable_size_ = 0 + cache.protected_size_ = 0 + cache.enable_storage = False + cache.enable_kv_cache_events = False + cache.inc_node_lock_ref = lambda node: None + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 150 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = torch.arange(6, dtype=torch.int64) + node.host_len = 0 + node.cp_hicache = None + root.children[(0, 1, 2, 3)] = node + reservation = make_write_reservation( + torch.arange(4, 6, dtype=torch.int64), node_id=151, host_start=210 + ) + prepared = PreparedCpHiCacheBackup( + node_id=151, + reservation=reservation, + metadata=reservation.metadata, + logical_len=2, + ) + + result = cache.insert( + InsertParams( + key=RadixKey(list(range(6))), + value=torch.arange(6, dtype=torch.int64), + cp_hicache_prepared_backup=prepared, + ) + ) + + self.assertEqual(result.prefix_len, 4) + parent = root.children[(0, 1, 2, 3)] + self.assertEqual(parent.key.token_ids, [0, 1, 2, 3]) + new_tail = parent.children[(4, 5)] + self.assertIsNot(new_tail, node) + self.assertEqual(new_tail.key.token_ids, [4, 5]) + self.assertTrue(prepared.attached) + self.assertIs(cache.pending_host_backups[151].node, new_tail) + def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = False diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index a31299ae4..dfea1448c 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -569,7 +569,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64) self.assertFalse(can_reuse_current_extend_kv(forward_batch)) - def test_should_reuse_current_extend_kv_allows_draft_cache_hit_suffix(self): + def test_should_reuse_current_extend_kv_disables_draft_partial_cache_hit_suffix( + self, + ): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -593,8 +595,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True): - self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) - self.assertEqual(runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS, {}) + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch)) + self.assertIn( + "[CP_SHARED_KV_FALLBACK][current_reuse]", + logs.output[0], + ) + self.assertIn("draft_partial_current_reuse_disabled", logs.output[0]) forward_batch.spec_info = TargetSpecInfo() forward_batch.cp_shared_kv_mla_prefetcher = object()