Stabilize EAGLE draft cache hits under CP HiCache
The failing runs showed EAGLE accept length collapsing when draft cache-hit suffixes used the new partial-current splice path. This keeps target partial-current reuse enabled, but returns EAGLE/NextN draft cache-hit suffixes to the previous full-materialize path with an explicit fallback warning until the draft splice path has value-level ETE proof.\n\nThe same change set also tightens the page-granular CP HiCache contract for scheduler-visible hits and makes the prefill-to-decode EAGLE handoff observable without cloning hot-path metadata. Exact non-page CP hits are floored to a page boundary for new scheduling decisions, while internal unfinished-request refresh keeps its exact accounting.\n\nConstraint: CP shared KV and HiCache operate at page granularity; exposing token-precise CP tails to scheduler-visible cache hits can force non-page partial materialization.\nConstraint: EAGLE/NextN draft has only one executable layer, so draft prefetch and draft partial-current splice need a separate correctness contract from target layers.\nRejected: Keep draft partial-current splice enabled | remote logs correlate it with avg accept length around 0.068 and median 0.\nRejected: Clone decode metadata tensors on transfer | slot ownership until process_prebuilt consumes them avoids extra hot-path copies.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable draft partial-current reuse without metadata/draft-KV value checks and ETE accept-length evidence.\nTested: g0034 container py_compile for touched modules.\nTested: g0034 container PYTHONPATH=python python -m pytest -q test/registered/unit/disaggregation/test_decode_queue_compaction.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 183 passed, 5 warnings, 2 subtests passed.\nNot-tested: Fresh ETE accept-length run after this exact commit; requires user-driven traffic restart.
This commit is contained in:
laoyao0822
@@ -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.
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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())
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user