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 90bac0e9a..f68088fb5 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 @@ -1983,3 +1983,208 @@ Correction: Guardrail: - Future diagnostics inside disaggregation helper queues should use scheduler-owned runtime state unless the queue constructor explicitly copies that state onto `self`. + +### C34. 2026-05-30 EAGLE accept debug run: partial accept collapse correlates with cache-hit decode, not missing draft-buffer registration + +Evidence from live logs under `/mnt/beegfs/cjy/log/` after enabling `SGLANG_EAGLE_ACCEPT_DEBUG=1` on prefill and both decode nodes: + +- Active files inspected: + - prefill: `sglang_cp_hicache_20260529_173736.log` + - decode node0: `decode0_20260529_174040.log` + - decode node1: `decode1_20260529_174053.log` +- Processes were still alive at 2026-05-30 01:53 CST. No new traceback/health-failure was found in the inspected tail. +- Startup audit now proves the simple “draft buffers are not registered” hypothesis is false for this run: target KV buffers, draft KV buffers, target state buffers, and draft state buffers are all visible in the disaggregation KV manager. +- Decode-side EAGLE debug parsing found `25009` logged samples and `876` decode batch stat lines. Logged samples are intentionally biased toward zero because every zero-accepted sample is logged, while non-zero samples are interval-sampled. +- Even accounting for that bias, the low-accept symptom is real in aggregate decode stats: decode batch `accept len` min `1.00`, average `1.412`, median `1.27`, and `168/876` windows at `<=1.05`. +- Low accept is not only a transfer-window effect: `#transfer-req: 0` has `141/765` windows at `accept len <= 1.05`; `#transfer-req: 1` has `22/93` such windows. +- Logged zero-accepted samples appear across cache-hit sizes, not only tiny prefix: top examples include cached `320`, `13312`, `14208`, `15680`, `62528`, `67968`, `74496`, `75520`, `88256`, and `98112`. +- Prefill in the same run heavily exercised CP current reuse: `forward_partial_current_reuse=2464`, with `used_prefetch=False=1432` and `used_prefetch=True=1032`. For small prefix `320` (`5` pages, below `min_prefix_pages=8`), async prefetcher creation was skipped but target partial-current sync compose still ran. +- Draft partial current reuse remained disabled and logged fallback for the small-prefix EAGLE/NextN path, e.g. `reason=draft_partial_current_reuse_disabled ... prefix_lens=[320] extend_lens=[10808]`. + +Implication: + +- The current root-cause target is narrower than “EAGLE disabled” or “draft KV buffers missing”. The suspicious path is cache-hit prefill with CP shared KV current reuse where target uses partial-current reuse but draft uses full materialize/fallback, followed by decode-side draft rejection. +- This still does not prove that partial-current reuse is the root cause: logged zero-accept samples also appear on large cached prefixes with async prefetch enabled. The next check must trace how prefill target hidden/topk and draft hidden/KV are captured and transferred to decode, especially when cached tokens are non-zero. +- Do not add more cache or prefetch behavior changes solely from accept-len logs. First inspect EAGLE prefill capture and disaggregation transfer metadata paths for target/draft position alignment. + +### C35. 2026-05-30 new EAGLE accept debug log: cache-hit runs still reject draft tokens after transfer drains + +Evidence from the current live run: + +- Active files: + - prefill: `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_173736.log` + - decode node0: `/mnt/beegfs/cjy/log/decode0_20260529_174040.log` + - decode node1: `/mnt/beegfs/cjy/log/decode1_20260529_174053.log` +- As of the first inspection, prefill and both decode launch processes were still alive and grep found no traceback, health failure, scheduler fatal error, or memory-leak exception in these three logs. +- Decode-side parsing across both decode logs found `3707` aggregate `Decode batch` windows with average accept length `1.400`, median `1.30`, min `1.00`, and `745` windows at `<=1.05`. This confirms the accept-length symptom remains real in the new run. +- The issue is not only an active transfer-window artifact: `#transfer-req: 0` windows still averaged `1.394` and had `609/3005` windows at `<=1.05`. +- Cache-hit remains the strongest visible correlate: windows classified near decode EAGLE debug samples had `cached0 n=26 avg=2.21 <=1.05=0`, but `cached_gt0 n=3681 avg=1.394 <=1.05=745`. Cached `320` windows averaged `1.484`, while larger cached prefixes still averaged around `1.38`. +- The collapse is not only tiny-extend: inferred `origin-cached` buckets all remain low (`1..63` avg `1.422`, `64..127` avg `1.314`, `2048..8191` avg `1.266`, `8192+` avg `1.430`). Therefore treating only sub-page or sub-CP-page suffixes is insufficient. +- Prefill CP shared-KV logs show current reuse is active in the same run: `forward_partial_current_reuse=8864`, `has_prefetcher=True=2608`, `has_prefetcher=False=6440`, `prefix_below_min=2864`. Small-prefix cases below the async prefetch threshold still use synchronous current compose. +- Explicit fallback counts are narrow: `CP_SHARED_KV_FALLBACK` is `432`, consisting of `184` MLA prefetch + `184` index prefetch fallbacks for `prefix_len=0`, and `64` `current_reuse` fallbacks for `draft_partial_current_reuse_disabled`. +- Recent tail also shows heavy owner-lane device eviction before large CP load-back when device lanes are nearly exhausted, e.g. load-back of `107520` tokens required evicting `135360` physical tokens to fix per-lane deficits. This is relevant to performance/device-cache churn, but it does not explain decode EAGLE rejection by itself. + +Implication: + +- The new logs strengthen, rather than weaken, the previous conclusion: accept length dropping toward `1` is a cache-hit-associated EAGLE draft rejection problem after transfer, not missing draft-buffer registration and not solely async-prefetch/tiny-suffix behavior. +- The next root-cause target is the prefill-to-decode EAGLE state contract on cache-hit paths: target/draft hidden-state capture, draft KV content for cached prefix/current suffix, and page-aligned CP local/global position trimming. The remaining suspicious asymmetry visible in logs is that target MLA uses partial-current reuse while draft cache-hit current reuse still explicitly falls back to full materialization. + +### C36. 2026-05-30 runtime-validation correction: growing log file is not proof of a new code version + +Correction: + +- A log file under `/mnt/beegfs/cjy/log/` continued to grow, but the active prefill PID was still `2586666` and the active decode PIDs were still `2446863` / `2208897`. That means the inspected evidence came from the already-running old process, not from a restarted process with any subsequent code changes. +- Therefore C35 must be treated as evidence for the old live runtime only. It is not validation of any newer local code or newer remote-synced code. + +Guardrail: + +- For future remote validation, do not infer a new runtime from log freshness or file name alone. Before using logs as validation evidence, record all three: + 1. process PID and start/elapsed time, + 2. runtime environment flags relevant to the change, + 3. code identity, preferably git commit or an explicit startup marker from the running checkout/container. +- If the process was not restarted after a code fix, report that explicitly and wait for/user-triggered traffic only after saying the new code is ready to restart. + +### C37. 2026-05-30 cache-hit EAGLE low-accept investigation narrowed to draft HiCache/materialize contract + +Evidence from the still-running process recorded in C36, not a restarted/new-code validation: + +- Decode samples with high cache-hit prefixes continue to show many `accepted_draft=0` entries. Examples include page-floored cache hits such as `origin=116609 cached=116544`, plus other concurrent high-hit requests (`origin=117775 cached=117760`, `81781/81728`, `79695/79680`, `93640/93632`). +- Aggregate parsing on the same logs still shows the symptom after transfer pressure drains: `#transfer-req:0` windows contain many `accept len <= 1.05`, while no-cache windows are materially healthier. +- Startup diagnostics already proved target/draft KV buffers and target/draft state buffers are registered in the disaggregation KV manager, so the simple “draft transfer surfaces are missing” hypothesis is rejected for this run. +- Prefill tail shows cache-hit target current reuse active with async prefetch (`used_prefetch=True`) and CP page-floored cached prefixes, followed by CP backup final ack and prepared backup logs. One inspected tail had `extend_lens=[8331]` but backup `logical_len=8330`; this is a candidate off-by-one/contract point to verify before changing behavior. + +Current narrowed hypotheses: + +1. Draft HiCache load/backup may not be fully symmetric with target under CP shared KV cache hit, even though draft buffers are registered. +2. Draft cache-hit materialization may diverge from target because target uses partial-current reuse/splice while draft currently disables partial-current reuse and falls back to full materialization. +3. Decode-side EAGLE prebuilt initialization may have an output-token off-by-one interaction with disaggregated cache-hit metadata, but this must be compared against the existing chunk-cache/prebuilt contract before modifying it. + +Guardrail: + +- Do not add another cache/prefetch behavior change from accept-length logs alone. First prove where draft prefix/current KV or transferred EAGLE state diverges from the target contract. + +### C38. 2026-05-30 anti-loop investigation ledger: checked points and do-not-repeat list + +Purpose: + +- This section is the read-first checkpoint for the current EAGLE/CP-HiCache low-accept investigation. Before re-scanning broad files, read C29-C38 and only inspect a point again if new runtime evidence contradicts it. + +Checked / rejected points: + +1. `omx explore` is currently unusable in this checkout/runtime: it fails before doing repository lookup with `Cargo.lock` version `4` unsupported by the installed Cargo. Use narrow shell searches until OMX is repaired; do not spend more investigation cycles on the same failure. +2. The active evidence in C35-C37 came from an old still-running process unless PID/start-time/code identity proves otherwise. Fresh log growth alone is not proof of a restarted runtime. +3. The simple hypothesis “draft KV/state transfer surfaces are not registered” is rejected for the EAGLE accept-debug run: startup diagnostics showed target KV, draft KV, target state, and draft state buffers visible in the disaggregation KV manager. +4. The observed `extend_lens=N+1` vs HiCache backup `logical_len=N` in EAGLE paths is currently explained by bigram radix keys: `convert_to_bigram_key([t0..tk])` has length `len(tokens)-1`. Do not treat this as an off-by-one root cause unless a decode-side contract requires token-count rather than bigram-count metadata. +5. CP HiCache target/draft host reservation is symmetric at the allocator level: `reserve_write_cp()` page-pads once, computes one owner pattern, allocates target host indices and draft host indices for the same physical owned positions, and stores both in `CpHiCacheNodeMetadata`. +6. CP HiCache target/draft backup is symmetric at the layer operation level: `submit_write_cp_layer()` submits target D2H and draft D2H on the same write stream and only queues final ack after all required target layers and all draft layers complete. +7. CP HiCache target/draft load-back is structurally symmetric for owned ranks: `load_cp()` replays the saved page-owner vector through `alloc_pages_with_owners()`, queues target `load_queue` and draft `draft_load_queue` against the same physical device indices, and `start_loading()` loads draft layer 0 before target layer 0 completes. This does not yet prove decode correctness, but it means “draft load queue is never populated” is not supported by code inspection. +8. Zero-owned CP ranks currently queue a zero-length target load op to keep load ACK rows aligned, but do not queue a draft zero-length op. Since they own no draft bytes this is likely OK, but if a future hang points at load ACK/node-id symmetry, this is the exact edge to revisit. +9. The remaining strongest code-level asymmetry is not host allocation/load/backup; it is MLA materialization/current-reuse: target cache-hit extend may use partial-current reuse, while EAGLE/NextN draft cache-hit extend intentionally returns false from `should_reuse_current_extend_kv()` with reason `draft_partial_current_reuse_disabled` and then falls through to full materialization. +10. Current non-prefetch partial-current compose is no longer a silent fallback for target: when no async prefetcher exists it does synchronous prefix materialize + current splice; if the prefix is not one positive page-aligned prefix, it fail-fast with `[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync]`. +11. Async MLA/index prefetch attention-window wait is intentionally deferred to consume-time in the current design; stale tests expecting attention-window wait/launch were updated earlier. Do not re-add attention-window waits without proving consume-time ordering is wrong. +12. The one-page allocator leak from repeated tiny EAGLE exact hits was already root-caused and fixed by freeing exact duplicate CP HiCache tail pages with `include_partial_tail=True` on duplicate insert. Do not revisit this unless the memory checker again reports an exact page-size accounting gap under the new code. + +Open checks still worth doing: + +1. Inspect the full decode traceback around the latest `Scheduler hit an exception` in decode logs before changing cache logic again. A concrete traceback can supersede the accept-length correlation hypotheses. +2. Trace the EAGLE cache-hit prefill-to-decode state contract: target hidden/topk capture, draft hidden capture, draft KV/state transfer, and decode prebuilt initialization. The key question is why cached-prefix decode rejects draft tokens even after transfer drains. +3. Decide whether draft cache-hit full materialization is semantically equivalent to target partial-current reuse. If not, draft needs an explicit same-layer current-reuse contract or a fail-fast guard, not another silent fallback. + +### C39. 2026-05-30 remote process stop invalidates connection-closed traceback as root-cause evidence + +Correction from operator: + +- The remote process observed around the latest decode exceptions was manually stopped by the operator. +- Therefore decode-side `RuntimeError: gloo ... Connection closed by peer` and paired `KVTransferError(...): Aborted by AbortReq` records in the old logs are downstream shutdown/abort symptoms, not a valid root-cause signal for the EAGLE accept-length collapse. + +Implication: + +- Do not use those shutdown tracebacks to drive cache or EAGLE code changes. +- Continue root-cause work from the stable evidence already recorded: cache-hit decode windows show low EAGLE accept length while draft KV/state transfer surfaces are registered, and the main unproven asymmetry remains cache-hit target partial-current reuse versus draft full materialization / draft current-reuse disablement. + +### C40. 2026-05-30 EAGLE/CP-shared-KV prefill contract checked points + +Purpose: + +- This section records the code facts checked after C39 so the next continuation does not repeat the same broad reads. + +Checked points: + +1. `_in_seq_collect_last_token()` is not currently an obvious padded-tail bug when NSA in-seq CP split is enabled. It uses `actual_token_count=sum(forward_batch.extend_seq_lens_cpu)` and `_get_in_seq_last_token_owner_and_offset(split_list, cp_size, actual_token_count)`, so the collected last hidden state is based on the valid extend length, not the padded physical tail. +2. `can_cp_split()` now skips NSA in-seq CP split for CP shared-KV cache-hit suffixes whose padded physical page count is below `cp_size`. When this happens, `forward_batch.nsa_cp_metadata` is never populated and `nsa_use_prefill_cp()` becomes false. The model then runs the unsplit/full-token path, while the NSA backend can still use CP shared-KV materialization because that gate is `forward_batch.uses_cp_shared_kv`, not `nsa_use_prefill_cp()`. +3. In the no-NSA-split path, MLA direct local write from `forward_mla._maybe_write_cp_shared_local_mla_kv()` is not used because it is called only under `nsa_use_prefill_cp()`. The later NSA backend save path still filters logical `out_cache_loc` to owned physical locs through `_maybe_filter_shared_mla_kv_write()` before writing. This is the expected fallback write route for short suffixes without in-seq CP metadata. +4. `ForwardMode.is_context_parallel_extend()` excludes `DRAFT_EXTEND` v1, but prefill-time EAGLE draft extension does not set the schedule batch to `DRAFT_EXTEND`; it keeps `ForwardMode.EXTEND` with `EagleDraftInput`. Therefore prefill draft can still create NSA CP metadata and hit the draft partial-current guard. Decode-time `forward_draft_extend_after_decode()` explicitly sets `ForwardMode.DRAFT_EXTEND`, so that later path is different and should not be conflated with prefill draft extend. +5. The logged `draft_partial_current_reuse_disabled` fallback is therefore specifically meaningful for prefill-time EAGLE draft cache-hit extends (`forward_mode=EXTEND`, `spec_info=EAGLE_DRAFT`), not for ordinary decode `DRAFT_EXTEND` after verification. +6. Target hidden capture for CP draft shared-KV is split-aware: target forward stores `forward_batch.draft_hidden_states` before the final CP collect. If CP split is active, this is CP-local hidden; if CP split is skipped, it is full-token hidden. Draft `_get_cp_local_spec_hidden_states()` accepts either shape: it returns local hidden unchanged or CP-splits full hidden when metadata exists. +7. `LogitsProcessor` compact-hidden behavior is compatible with CP narrow output: when the hidden tensor already has one row per request, it returns logits directly and returns hidden states only when `capture_hidden_mode=LAST`. The target CP shared-KV path bypasses the need for full hidden by using the separate `draft_hidden_states` side channel. + +Open after C40: + +- The remaining unproven point is still semantic, not a missing allocator/transfer surface: for prefill-time EAGLE cache-hit, target may splice current KV with partial-current reuse while draft currently full-materializes. Need prove whether those two produce equivalent draft logits/topk/hidden for the transferred prebuilt decode state. If not equivalent, implement draft same-layer current reuse or fail-fast that path; do not add silent fallback. +- Need a focused unit/instrumented check around prefill-time EAGLE draft cache-hit materialization, not another pass over HiCache target/draft host allocation. + +### C41. 2026-05-30 EAGLE disagg/cache contract checked points after manual stop + +Purpose: + +- Continue the C40 investigation without re-reading the same allocator and transfer paths. These points were checked after the operator noted the previous process had been manually stopped. + +Checked points: + +1. The newest prefill log at `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260529_184808.log` is not accept-length validation evidence. The active prefill process PID `2639412` started at `2026-05-30 02:48:08 CST`; the inspected tail only shows warmup and `The server is fired up and ready to roll!`, with no new decode traffic evidence in that file. +2. `send_kv_chunk()` transfers target/draft KV pages for `req_to_token[req_pool_idx, start_idx:end_idx]`, where the default final `end_idx` is `min(len(req.fill_ids), len(req.origin_input_ids))`. NSA state pages use the same `:seq_len` origin/fill span. This confirms prefill does not intentionally transfer generated output-token KV as part of the initial origin chunk; the first generated token is transferred through metadata (`output_id`, topk, hidden) and decode appends it on commit. +3. Decode pre-allocation mirrors that contract: `_pre_alloc()` allocates `origin_input_len` before metadata commit; `_commit_transfer_to_req()` appends the transferred `output_id` and stores transferred EAGLE `topk_p/topk_index/hidden_states_tensor`. This makes the prefill/decode origin-token transfer length consistent at the scheduler/metadata level. +4. EAGLE radix using one fewer key than token count is also the existing base radix contract, not a CP-only change. `RadixCache.cache_unfinished_req()` comments explicitly state that in EAGLE mode “bigram keys will only cache len - 1 kv indices”. Keep treating `logical_len = token_count - 1` as intended unless a draft-KV value mismatch is proven. +5. CP HiCache `load_cp()` internally allocates padded physical pages and returns only `visible_device_indices` sliced to each node's valid length. Therefore `HiRadixCache.load_back()` setting `loaded_node.value` using logical `host_len` offsets is correct for the returned tensor; the padded offset handling already happened inside `load_cp()`. +6. The remaining unproven correctness gap is therefore narrower: whether EAGLE/NextN draft prefill cache-hit attention should continue using full materialization from the draft pool while the target path uses partial-current compose, or whether the draft path must get the same current-suffix compose contract. This is not answered by allocator/load symmetry alone. + +Next check: + +- Inspect the draft full-materialize path against the partial-current compose path for prefill-time EAGLE cache hits. If full materialization can read stale/missing current suffix rows before draft attention, switch draft cache-hit prefill to the same explicit current-suffix compose contract or fail fast; do not add another silent fallback. + +### C42. 2026-05-30 draft cache-hit current-reuse contract change + +Root-cause update: + +- The previous draft guard in `should_reuse_current_extend_kv()` made prefill-time EAGLE/NextN draft cache-hit suffixes use full materialization from the persistent draft pool while the target path used prefix materialize + fresh current-suffix splice. +- That is the wrong contract boundary: disabling draft next-layer async prefetch is still correct because the draft model has no real next decoder layer, but disabling same-layer current suffix reuse changes the KV source visible to draft attention. +- The current fix removes the draft-specific current-reuse fallback. Target and prefill-time draft now share the same decision: current-only or cache-hit partial-current batches should materialize the page-aligned prefix and splice the freshly rebuilt current suffix. + +Guardrails: + +- This does not re-enable EAGLE/draft async prefetch. `cp_shared_kv_should_prefetch_next_layer()` still returns `False` for draft input. +- If a future hang reappears in draft MoE/DeepEP, investigate current-suffix rebuild/compose ordering directly; do not silently restore full-materialize fallback because that recreates the target/draft asymmetry. + +Verification status: + +- Local py_compile passed for the changed runtime/test files. +- Remote container verification on `g0034:/sgl-workspace/sglang-tai` passed: + - py_compile for `cp_shared_kv_runtime.py` and `test_cp_shared_kv_runtime.py` + - focused draft-current-reuse and draft-prefetch tests: `2 passed` + - full `test_cp_shared_kv_runtime.py`: `73 passed, 5 warnings, 2 subtests passed` + - related CP shared-KV/page-aligned suite: `test_nsa_cp_utils.py`, `test_cp_shared_kv_layout.py`, and `test_cp_shared_kv_runtime.py`: `128 passed, 5 warnings, 2 subtests passed` +- ETE/cache-hit accept-length validation is still not done for this change because the last inspected prefill log only showed warmup and the operator controls traffic generation. + +### C43. 2026-05-30 remote-runtime identity after draft current-reuse sync + +Checked point: + +- The changed runtime/test files were synced to `g0034:/mnt/beegfs/cjy/sglang-dev` and their SHA-256 hashes match the local checkout. +- The active prefill process on `g0034` was PID `2639412`, started at `2026-05-30 02:48:08 CST`, which predates this sync. Therefore that process cannot validate the draft current-reuse fix until it is restarted. +- The latest visible decode accept-debug log snippets are still evidence from the old runtime lineage. They remain useful for root-cause correlation, but not for validating C42. + +Guardrail: + +- Do not interpret any cache-hit accept-length numbers from the currently running old process as success or failure of C42. Require process start time after the sync, or an explicit code identity marker, before using ETE logs for this change. + +### C44. 2026-05-30 EAGLE prebuilt top-k metadata shape checked point + +Checked point: + +- The current launch uses `--speculative-algorithm EAGLE` without `--enable-multi-layer-eagle`, so the worker is `EAGLEWorkerV2`, not `MultiLayerEagleWorkerV2`. +- In this mode, prefill draft `_draft_extend_for_prefill()` only produces one draft top-k state for the next decode draft call. `decode_schedule_batch_mixin.process_prebuilt()` therefore intentionally slices `req.output_topk_p[:speculative_eagle_topk]` and does not multiply by `speculative_num_steps`. +- The decode debug log prints the full metadata buffer shape (for example `(16,)`) because the disaggregation metadata buffer is over-provisioned; process-prebuilt uses only the configured top-k prefix. This is not by itself evidence of a metadata shape bug. + +Rejected: + +- Do not change prebuilt `num_states` to `topk * speculative_num_steps` for the non-multi-layer EAGLE path unless the launch also enables `enable_multi_layer_eagle`. That would feed extra uninitialized/irrelevant metadata into `EAGLEWorkerV2.draft()`. 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 a3b972447..c0bd9569e 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,14 +978,13 @@ 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. - Partial current reuse appends the freshly computed suffix KV to the - materialized prefix buffer. That path is safe for the target model, but the - EAGLE/NextN draft layer has a different lifetime contract: it mirrors target - KV state and immediately enters DeepEP MoE after attention. A cache-hit - draft suffix observed in production can leave all ranks stuck in that MoE - collective after the partial-reuse attention path returns. Keep draft - cache-hit suffixes on the older full-materialize path until draft gets an - explicit same-layer reuse contract. + 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. """ if not cp_shared_kv_current_reuse_enabled(): @@ -993,21 +992,7 @@ def should_reuse_current_extend_kv(forward_batch) -> bool: current_only = is_current_only_extend_batch(forward_batch) partial_current = can_reuse_current_extend_kv(forward_batch) - if not (current_only or partial_current): - return False - - if cp_shared_kv_is_draft_input(forward_batch) and not current_only: - _log_current_reuse_fallback( - "draft_partial_current_reuse_disabled", - "CP shared KV current-reuse fallback (draft_partial_current_reuse): " - "cache-hit EAGLE/NextN draft uses full materialize instead of " - "partial current reuse. prefix_lens=%s extend_lens=%s", - getattr(forward_batch, "extend_prefix_lens_cpu", None), - getattr(forward_batch, "extend_seq_lens_cpu", None), - ) - return False - - return True + return current_only or partial_current def current_loc_remap_fast_path_args( 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 d1bd19e27..a31299ae4 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,7 @@ 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_disables_draft_cache_hit_suffix(self): + def test_should_reuse_current_extend_kv_allows_draft_cache_hit_suffix(self): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -593,15 +593,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True): - with self.assertLogs(runtime.logger.name, level="WARNING") as logs: - self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch)) - self.assertTrue( - any( - "draft_partial_current_reuse" in message - and "[CP_SHARED_KV_FALLBACK][current_reuse]" in message - for message in logs.output - ) - ) + self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) + self.assertEqual(runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS, {}) forward_batch.spec_info = TargetSpecInfo() forward_batch.cp_shared_kv_mla_prefetcher = object()