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 5c46a8955..0c0231807 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 @@ -1033,6 +1033,281 @@ Completed C14: 123 passed, 5 warnings ``` +## C15: CP valid-tail radix nodes still need page-physical device accounting + +Finding: + +- Startup can fail immediately after disaggregated prefill warmup with: + + ```text + token_to_kv_pool_allocator memory leak detected! + max_total_num_tokens=833024, available_size=833024, + evictable_size=3, protected_size=0, session_held=0 + ``` + +- The observed `evictable_size=3` is consistent with EAGLE/bigram warmup: + a 4-token request is cached as a 3-token bigram radix key. Under CP HiCache + the radix key is valid-tail based, while the CP device allocator is + page-granular. +- `RadixCache.cache_finished_req()` still freed `kv_indices[len(keys):]`. + For CP + page allocator, freeing the 4th token of a 64-token page releases the + whole physical page even though the radix node still owns the first 3 valid + tokens. The allocator then reports the full pool available while radix still + reports an evictable 3-token node. +- HiRadix device accounting also still used `len(node.key)` / `len(node.value)` + in several places. With valid-tail nodes this is the scheduler-visible valid + length, not the allocator-visible physical page span. + +Correction: + +- CP HiCache device residency must account physical page spans while radix + matching and scheduler-visible prefix lengths remain valid-token based. +- Tail free for CP valid-tail keys must start at the next page boundary, not at + `len(keys)`, so partial retained pages are never released out from under a + radix node. +- Splits inside a CP physical page should be floored to the previous page + boundary (or rejected/deferred at zero) even when the node is not yet backed + by host metadata. Page is the minimum ownership unit for both GPU and host + HiCache. + +Tests: + +- Add a regression for EAGLE/bigram CP finished-cache tail free: retained + partial pages must not call allocator free on the in-page tail. +- Add CP page-physical accounting regressions for insert/lock/evict paths. +- Add a split-floor regression for an unbacked CP valid-tail node. + +Implemented C15: + +- Added `ceil_to_page_len()` and made CP `cache_finished_req()` free suffix + pages only from the next page boundary after the valid radix key. This keeps + a retained valid-tail node from releasing the physical page it still owns. +- Added `HiRadixCache._node_device_resident_len()` and routed CP insert, + single-node write locks, path locks, load-back locks, regular eviction, and + backed demotion through allocator-visible physical page lengths. +- Changed the CP split-floor policy to apply to unbacked valid-tail device + nodes too. Splits inside a physical page now floor to the previous page + boundary independent of host-backup state. + +Completed C15 tests: + +```text +remote g0034 container: + test_cp_hicache_metadata.py + 90 passed, 5 warnings + test_cp_hicache_load_back_owner_lanes.py + test_cp_shared_kv_layout.py + test_cp_shared_kv_runtime.py + 103 passed, 3 warnings +``` + +Not-tested C15: + +- `test_radix_cache_unit.py` collection currently fails in the remote container + before running tests because the `sgl_kernel::sgl_per_token_group_quant_8bit` + operator is not registered in that test import path. This is an environment / + test bootstrap issue, not a C15 assertion failure. + +## C16: Exact valid-tail hits can make the next backup start mid-page + +Finding: + +- Current C7/C15 code keeps exact CP valid-tail cache hits token-precise. That + means a request can hit, for example, 100 valid tokens on 64-token pages and + then start its next suffix at token 100. +- `prepare_write_backup_for_req()` derives the backup start as + `max(req.cache_protected_len, existing_prefix_len)`. If the protected prefix + is a valid tail, `start` can be inside a physical page. +- `_cp_required_host_tokens_by_rank()` / `reserve_write_cp()` then call + `pad_token_locs_to_page_boundary()` on `device_indices[start:]`. That helper + is only correct for spans that start at a page boundary; CPU tensors validate + this, but CUDA tensors intentionally skip the sync-heavy offset check. +- Therefore the production CUDA path can silently build a padded physical span + from a mid-page start, producing incorrect page-owner accounting and host + reservation metadata. + +Risk: + +- This is a correctness risk for repeated prompts where the previous request + leaves a non-page-aligned valid-tail node and the next request extends beyond + that exact tail. +- It also conflicts with the current design principle that CP HiCache owns cache + at page granularity and may sacrifice sub-page cache to keep management simple. + +Correction options: + +- Preferred under the current page-as-minimum-unit contract: expose only the + previous page boundary as `cache_protected_len` / `device_indices` for any hit + that will be extended, so new backup starts at a page boundary. +- More complex alternative: implement explicit partial-page sharing/refcounting + between a cached prefix tail and the new current suffix. This is intentionally + not the current direction. + +Tests needed: + +- Cache a 100-token valid-tail node, then issue a request that hits those 100 + tokens and extends by another suffix. Assert write reservation starts at 64, + not 100, or otherwise fails fast before mid-page padding. +- Add a CUDA/remote regression or source-level guard proving CUDA hot path cannot + call `pad_token_locs_to_page_boundary()` with a non-page-start span. + +## C17: Duplicate/no-insert frees are still token-granular under CP pages + +Finding: + +- C15 fixed the final tail free in `RadixCache.cache_finished_req()` by starting + CP tail free at the next page boundary. +- The duplicate/no-insert frees remain token-sliced: + + ```text + cache_finished_req: + kv_indices[req.cache_protected_len : new_prefix_len] + kv_indices[req.cache_protected_len : len(keys)] + + cache_unfinished_req: + kv_indices[req.cache_protected_len : new_prefix_len] + ``` + +- With CP page allocators, freeing any loc inside a page can release the whole + page. If the free range begins or ends inside a page, it can release a page + that still contains retained current/request KV or radix-owned valid-tail KV. + +Risk: + +- Remaining memory-accounting or allocator/radix divergence after the C15 startup + fix, especially on chunked/unfinished requests or insert-prefix duplicate + paths. + +Correction: + +- Audit every CP call to `token_to_kv_pool_allocator.free(kv_indices[a:b])`. + Under CP page mode, free ranges should either be page-aligned or deliberately + converted to whole pages that are no longer referenced by radix/request state. +- If a range is not page-aligned and cannot be proven safe, fail fast in tests + first, then floor/ceil according to ownership semantics. + +Tests needed: + +- Finished insert where `req.cache_protected_len` and `new_prefix_len` differ + inside the same page. +- `is_insert=False` finished path with a valid-tail key. +- `cache_unfinished_req()` duplicate-free path with a non-page-aligned protected + prefix. + +## C18: `cache_protected_len` now mixes valid-token and physical-page semantics + +Finding: + +- CP valid-tail matching can expose `req.cache_protected_len` as a valid-token + length. +- Some scheduler/runtime checker paths still assert or compute as if + `cache_protected_len` is page-aligned physical protection: + + ```text + scheduler_runtime_checker_mixin.py::_get_batch_uncached_size + schedule_batch.py::_evict_swa + ``` + +Risk: + +- Exact valid-tail cache hits may fail assertions or under/over-count uncached + tokens in memory self-check paths. +- Even when asserts do not fire in the current CP/NSA configuration, this is an + API ambiguity that can reappear when batch/chunked/SWA paths are combined. + +Correction: + +- Split request-visible fields conceptually: + - valid protected length for scheduler/logical prompt progress; + - physical protected length for allocator/page accounting. +- Until that refactor is done, CP HiCache should floor exposed protected length + to the nearest page boundary wherever allocator/page accounting consumes it. + +Tests needed: + +- Valid-tail CP hit through `Req.init_next_round_input()` and scheduler runtime + memory checker. +- Guard that SWA/page-aligned paths either remain unreachable for CP NSA HiCache + or receive a page-aligned protected length. + +## C19: Prefetcher creation must not gate target partial-current reuse + +Finding: + +- Remote ETE run `/mnt/beegfs/cjy/sglang_cp_hicache_20260529_001129.log` + died at `2026-05-29 00:17:24` with: + + ```text + [CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch] + reason=missing_prefetcher prefix_lens=[320] extend_lens=[10808] + ``` + +- The same log shows the async MLA prefetcher was intentionally not created: + + ```text + create_skip reason=prefix_below_min prefix_pages=5 min_prefix_pages=8 + ``` + +- A later run showed the inverse shape: `prefix_lens=[16320]` and + `extend_lens=[16]` passed the prefix gate and created async prefetchers even + though the current suffix was sub-page. So "short extend" did not request + prefetch directly; a large cache-hit prefix made the async prefetch gate pass. + +Corrected contract: + +- `CpSharedKVMlaPrefetcher` / `CpSharedKVIndexPrefetcher` mean async + next-layer prefetch only. If async prefetch is disabled, below threshold, or + otherwise not worthwhile, do not create a prefetcher object. +- Target MLA partial-current reuse must not depend on prefetcher existence. + When there is no prefetched prefix handle, `nsa_backend.py` synchronously + materializes the page-aligned prefix into the slot-dense buffer, all-reduces + that prefix range, then splices the current KV rows into their suffix page + slots with tail slack mapped to `-1`. +- Draft/EAGLE cache-hit partial-current reuse remains disabled until the draft + KV lifetime has its own same-layer contract. This does not affect target + partial-current reuse or cache-miss current-only reuse. +- The disabled path is only the old compact materialize/current merge fallback; + it can expose padded tail slack. The synchronous page-slot compose path is + the fallback-free current-layer contract when no async prefetcher exists. + +Tests: + +- Target cache-hit partial-current with `cp_shared_kv_mla_prefetcher=None` must + still return true from `should_reuse_current_extend_kv()`. +- `materialize_prefix_and_reuse_current_kv_page_slots()` covers the no-prefetcher + synchronous compose path and verifies tail slack maps to `-1`. +- Tiny-extend batches with large prefixes should skip async MLA/index prefetcher + creation before touching pool getters/remap builders. + +## C20: Prefetch disabled must not disable full current reuse + +Finding: + +- Cache-miss/current-only requests (`extend_prefix_lens_cpu=[0]` and + `seq_lens_cpu == extend_seq_lens_cpu`) do not need prefix materialization or a + prefetcher: they can reuse the full current forward KV directly. +- Conflating "no prefetcher" with "no current reuse" regresses warmup, + cache-miss, and tiny-prefix traffic by forcing unnecessary full materialize or + rebuild work. + +Risk: + +- Performance regression on no-cache or tiny-prefix traffic. This is especially + visible before HiCache warms up or when a prompt falls below the MLA prefetch + threshold (`prefix_pages < min_prefix_pages`, currently 8 pages by default for + `cp_size=8,page_size=64`). + +Correction: + +- Keep `should_reuse_current_extend_kv()` structured as: + - `current_only` is sufficient for full current reuse; + - target `partial_current and not current_only` is also sufficient, regardless + of prefetcher existence, because the backend has synchronous page-slot + compose; + - draft/EAGLE `partial_current and not current_only` remains rejected. +- Tests include both target partial-current/no-prefetcher and target + current-only/no-prefetcher coverage. + ## Testing requirements / status Targeted coverage needed for this contract: @@ -1082,3 +1357,87 @@ Targeted coverage needed for this contract: - Do not add a hot-path CP collective to reconcile capacity. - Do not make draft KV manage an independent padded/valid contract. - Do not silently fall back from the page-aligned contract. + +## C21: Large prefix with tiny extend must not enter MLA/index prefetch + +Finding: + +- The current prefetch creation gate is prefix-sized: it skips MLA/index + prefetch only when `prefix_pages < min_prefix_pages`. +- The 2026-05-29 remote hang did not violate that gate. The last batch before + detokenizer health failure had: + - `prefix_lens=[16320]` (`255` pages at `page_size=64`), so the prefix gate + allowed prefetch; + - `extend_lens=[16]`, less than one page; + - `has_mla=True has_index=True`, followed by no further forward-layer progress + and detokenizer timeout. +- Earlier batches with `prefix_lens=[320]` (`5` pages) correctly logged + `create_skip reason=prefix_below_min ... min_prefix_pages=8` and produced + `has_mla=False has_index=False`; those are not the observed hang pattern. +- In other words, short extend did not directly request prefetch. Cache hit + converted most of the request into a large cached prefix, and the large prefix + satisfied the prefetch gate while leaving only a sub-page current suffix. + +Root-cause analysis: + +- MLA/index prefetch is currently a next-layer optimization. It starts prefix + materialization/reduce on a separate stream from layer prepare hooks, then the + consumer waits on the event later. +- Partial-current reuse uses the MLA prefetcher as a page-slot compose contract: + prefetched prefix rows are reused and the fresh current suffix KV is inserted + into the padded tail page. +- With `extend_len < page_size`, the suffix page has exactly one CP owner lane + and many rank-local zero-owned cases. The current layer is also short enough + that async prefetch collective ordering and event waits are no longer hidden by + compute. A mismatch in next-layer async collective order or event completion + can present as a CUDA/NCCL wait rather than a Python exception. +- The observed symptom matches this: after `has_mla=True has_index=True` for + `prefix_lens=[16320], extend_lens=[16]`, no later forward probe appears and + the server only reports detokenizer heartbeat timeout. + +Risk: + +- Tiny cache-hit suffixes can pay prefetch launch/materialize/collective overhead + that cannot be amortized by the short extend. +- More importantly, sub-page suffixes exercise the still-sensitive padded-tail + partial-current path. Until that path is proven by ETE, allowing MLA/index + prefetch on `extend_len < page_size` is a correctness/hang risk. + +Correction: + +- Do not disable partial-current reuse just because `extend_len` is short. + Short cache-hit suffixes are exactly where avoiding suffix materialization is + valuable. +- Add a separate minimum-extend gate only for the async next-layer prefetch + strategy. When `extend_len` is below the async threshold, do **not** create an + MLA/index prefetcher object; a prefetcher object now means async next-layer + prefetch machinery exists. +- Target partial-current reuse is independent of prefetcher existence. Without + a prefetched prefix handle, the backend takes the synchronous page-slot compose + path: materialize/reduce the page-aligned prefix in the current layer and then + splice current KV rows into their padded suffix page slots. +- Full current reuse is also independent of prefetcher existence. Cache-miss + current-only batches keep using the forward KV directly even when async + prefetch is disabled, skipped by threshold, or not created. +- If partial-current reuse is intentionally rejected, for example the current + conservative draft/EAGLE cache-hit guard, that rejection must not disable + current-only/full-current reuse for other request shapes. +- Default the async gate to one page when `page_size` is known; allow an env + override to lower or raise the threshold for experiments. +- This must not force tiny cache-hit suffixes back to compact full + materialization. + +Tests: + +- Add runtime helper coverage for the min-extend threshold default and override. +- Add MLA/index prefetch creation tests where `prefix_pages >= 8` but + `extend_len < page_size`; both prefetchers should return `None` before + touching pool getters/remap builders, while current reuse remains available. +- Add a synchronous no-prefetch compose test proving sub-page current suffix rows + are filled while tail slack in the current page is masked. + +Verified: + +- Remote container `/sgl-workspace/sglang-tai` on `g0034`: + - `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`: `73 passed` + - `test/registered/unit/mem_cache/test_cp_hicache_metadata.py`: `90 passed` diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cf517603d..84f787bca 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -214,6 +214,7 @@ class Envs: SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX = EnvBool(False) SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS = EnvInt(512) SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1) + SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS = EnvInt(-1) SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False) SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False) SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index d8e89a260..f07517db6 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -14,6 +14,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, + cp_shared_kv_mla_prefetch_min_async_extend_tokens, cp_shared_kv_mla_prefetch_min_prefix_pages, cp_shared_kv_mla_prefetch_should_log_layer, filter_locs_mappable_to_physical_pool, @@ -430,6 +431,27 @@ class CpSharedKVMlaPrefetcher: ) return None + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) != 1: + _prefetch_log("create_skip reason=bad_extend_lens_metadata") + return None + extend_len = int(extend_seq_lens_cpu[0]) + min_async_extend_tokens = cp_shared_kv_mla_prefetch_min_async_extend_tokens( + page_size=page_size + ) + if extend_len < min_async_extend_tokens: + _prefetch_log( + "create_skip reason=extend_below_min cp_rank=%s cp_size=%s " + "extend_len=%s min_async_extend_tokens=%s prefix_pages=%s page_size=%s", + layout.cp_rank, + layout.cp_size, + extend_len, + min_async_extend_tokens, + prefix_pages, + page_size, + ) + return None + cp_group = get_attention_cp_group() if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: _prefetch_log( @@ -471,7 +493,8 @@ class CpSharedKVMlaPrefetcher: _prefetch_log( "create cp_rank=%s cp_size=%s prefix_pages=%s total_slots=%s " - "owned_prefix_pages=%s owned_total_pages=%s dense_pages=%s page_size=%s", + "owned_prefix_pages=%s owned_total_pages=%s dense_pages=%s page_size=%s " + "min_async_extend_tokens=%s extend_len=%s", layout.cp_rank, layout.cp_size, prefix_pages, @@ -480,6 +503,8 @@ class CpSharedKVMlaPrefetcher: owned_total_pages, remap.dense_num_pages, page_size, + min_async_extend_tokens, + extend_len, ) create_total_ms = _cpu_timing_ms(create_cpu) _prefetch_log( @@ -1195,6 +1220,27 @@ class CpSharedKVIndexPrefetcher: ) return None + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + if extend_seq_lens_cpu is None or len(extend_seq_lens_cpu) != 1: + _prefetch_log("index_create_skip reason=bad_extend_lens_metadata") + return None + extend_len = int(extend_seq_lens_cpu[0]) + min_extend_tokens = cp_shared_kv_mla_prefetch_min_async_extend_tokens( + page_size=page_size + ) + if extend_len < min_extend_tokens: + _prefetch_log( + "index_create_skip reason=extend_below_min cp_rank=%s cp_size=%s " + "extend_len=%s min_extend_tokens=%s prefix_pages=%s page_size=%s", + layout.cp_rank, + layout.cp_size, + extend_len, + min_extend_tokens, + prefix_pages, + page_size, + ) + return None + cp_group = get_attention_cp_group() if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: _index_prefetch_fallback_log( 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 4474e232b..df680775c 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 @@ -96,6 +96,28 @@ def cp_shared_kv_mla_prefetch_min_prefix_pages( return max(int(configured), 0) +def cp_shared_kv_mla_prefetch_min_async_extend_tokens( + *, page_size: int | None = None +) -> int: + """Minimum current-extend tokens required for async next-layer prefetch. + + This threshold gates creation of the async next-layer MLA/index prefetcher + only. It must not gate current-layer reuse: when no prefetcher exists, the + backend can still synchronously materialize prefix pages and splice current + KV rows for target partial-current reuse. + + A negative env value uses the page size as the dynamic default. Set the env + to 0 to allow async prefetch even for sub-page extends during experiments. + """ + + configured = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.get() + if configured < 0: + if page_size is not None and int(page_size) > 0: + return int(page_size) + return 0 + return max(int(configured), 0) + + def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None: if cp_shared_kv_mla_prefetch_log_enabled(): logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args) @@ -1945,6 +1967,82 @@ def materialize_local_token_kv_page_slots_into( dense_range.copy_(torch.where(owned_view, gathered, zero)) +def materialize_prefix_and_reuse_current_kv_page_slots( + *, + kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + current_kv_cache: torch.Tensor, + current_locs: torch.Tensor, + slot_remap: SharedTokenKVSlotRemap, + layout: CpSharedKVLayout, + page_size: int, + prefix_pages: int, + layer_id: int | None = None, + nvtx_source: str = "mla.partial_current_sync", +) -> tuple[torch.Tensor, torch.Tensor]: + """Synchronously compose prefix materialization with current KV rows. + + This is the non-prefetch partial-current path. It preserves the same + page-slot layout used by async prefetch compose, but does not require a + prefetcher object or an extra CUDA stream. Prefix pages are materialized and + reduced immediately; current rows are then inserted into their padded suffix + page slots and non-current tail slack is masked from the returned locs. + """ + + total_slots = int(slot_remap.slot_logical_pages.numel()) + if prefix_pages < 0 or prefix_pages > total_slots: + raise ValueError( + "Invalid CP shared KV partial-current prefix range: " + f"prefix_pages={prefix_pages} total_slots={total_slots}" + ) + + dense_kv_cache = kv_cache.new_zeros( + (slot_remap.dense_num_pages * page_size, *kv_cache.shape[1:]) + ) + materialize_local_token_kv_page_slots_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=slot_remap.slot_logical_pages, + layout=layout, + page_size=page_size, + start_slot=0, + end_slot=prefix_pages, + ) + + prefix_rows = slot_range_to_token_slice(page_size, 0, prefix_pages) + _all_reduce_materialized_buffer_range( + dense_kv_cache, + layout.cp_size, + prefix_rows.start, + prefix_rows.stop, + nvtx_source=nvtx_source, + nvtx_layer_id=layer_id, + nvtx_cp_rank=layout.cp_rank, + ) + + logical_locs = filter_locs_mappable_to_physical_pool( + logical_locs=logical_locs, + layout=layout, + physical_token_capacity=kv_cache.shape[0], + ) + dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=slot_remap.page_inverse, + page_size=page_size, + ) + mixed_kv_cache, mixed_locs, _ = fill_current_kv_page_slots_and_remap_locs( + dense_kv_cache=dense_kv_cache, + materialized_dense_locs=dense_locs, + current_kv_cache=current_kv_cache, + logical_locs=logical_locs, + current_locs=current_locs, + page_inverse=slot_remap.page_inverse, + page_size=page_size, + mask_non_current_in_current_pages=True, + ) + return mixed_kv_cache, mixed_locs + + def slot_range_to_token_slice( page_size: int, start_slot: int, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 40bee5d4b..bbdf2a0b8 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -26,6 +26,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( filter_owned_logical_locs, get_or_build_shared_token_kv_slot_remap, is_current_only_extend_batch, + materialize_prefix_and_reuse_current_kv_page_slots, materialize_shared_token_kv_buffer, should_reuse_current_extend_kv, tensor_debug_checksum, @@ -1850,23 +1851,69 @@ class NativeSparseAttnBackend( if extend_lens_cpu is not None else None ) - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_prefetch] " - "CP shared KV MLA partial-current reuse requires " - "page-slot prefetch compose. Compact " - "materialize/current merge fallback is disabled " - "because it can expose padded tail slack. " - f"reason={reason} " - f"cp_rank={forward_batch.cp_shared_kv_layout.cp_rank} " - f"layer_id={layer.layer_id} " - f"prefix_lens={prefix_lens} " - f"extend_lens={extend_lens} " - f"current_rows={int(current_kv_cache.shape[0])} " - f"logical_page_table_shape={tuple(logical_page_table_1.shape)} " - f"current_locs_shape={tuple(forward_batch.out_cache_loc.shape)} " - f"page_size={current_remap_page_size} " - f"logical_page_capacity={current_remap_logical_page_capacity}" + page_size = int(forward_batch.token_to_kv_pool.page_size) + if ( + prefix_lens_cpu is None + or len(prefix_lens_cpu) != 1 + or int(prefix_lens_cpu[0]) <= 0 + or int(prefix_lens_cpu[0]) % page_size != 0 + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync] " + "CP shared KV MLA partial-current sync compose " + "requires one positive page-aligned prefix. " + f"reason={reason} " + f"cp_rank={forward_batch.cp_shared_kv_layout.cp_rank} " + f"layer_id={layer.layer_id} " + f"prefix_lens={prefix_lens} " + f"extend_lens={extend_lens} " + f"current_rows={int(current_kv_cache.shape[0])} " + f"logical_page_table_shape={tuple(logical_page_table_1.shape)} " + f"current_locs_shape={tuple(forward_batch.out_cache_loc.shape)} " + f"page_size={page_size}" + ) + prefix_pages = int(prefix_lens_cpu[0]) // page_size + slot_remap = get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=metadata.real_page_table, + layout=forward_batch.cp_shared_kv_layout, + page_size=page_size, ) + kv_cache, page_table_1 = ( + materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=logical_page_table_1, + current_kv_cache=current_kv_cache, + current_locs=forward_batch.out_cache_loc, + slot_remap=slot_remap, + layout=forward_batch.cp_shared_kv_layout, + page_size=page_size, + prefix_pages=prefix_pages, + layer_id=layer.layer_id, + ) + ) + if ( + cp_shared_kv_mla_prefetch_log_enabled() + and cp_shared_kv_mla_prefetch_should_log_layer( + layer.layer_id + ) + ): + cp_shared_kv_mla_prefetch_log( + "forward_partial_current_sync_compose cp_rank=%s " + "layer=%s reason=%s prefix_lens=%s extend_lens=%s " + "prefix_pages=%s current_rows=%s kv_rows=%s " + "page_table_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + reason, + prefix_lens, + extend_lens, + prefix_pages, + int(current_kv_cache.shape[0]), + int(kv_cache.shape[0]), + tuple(page_table_1.shape), + ) if ( cp_shared_kv_mla_prefetch_log_enabled() and cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 381f18c03..fdb4f4518 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -46,6 +46,7 @@ from sglang.srt.mem_cache.radix_cache import ( RadixCache, RadixKey, TreeNode, + ceil_to_page_len, compute_node_hash_values, split_node_hash_value, ) @@ -2035,7 +2036,6 @@ class HiRadixCache(RadixCache): or self.page_size <= 1 or prefix_len <= 0 or prefix_len % self.page_size == 0 - or not self._node_backuped(child) ): return prefix_len return prefix_len // self.page_size * self.page_size @@ -2380,6 +2380,27 @@ class HiRadixCache(RadixCache): node.pin_expiry = 0.0 node.pin_ttl = 0 + def _node_device_resident_len(self, node: TreeNode) -> int: + """Allocator-visible device residency for a radix node. + + CP HiCache radix keys are valid-token lengths, but CP device KV + ownership is page-granular. Residency accounting therefore uses the + physical padded page span whenever CP HiCache is active, while normal + HiCache keeps historical token-count accounting. + """ + + value = getattr(node, "value", None) + if value is None: + return 0 + if not getattr(self, "_uses_cp_hicache", False): + return len(value) + + metadata = getattr(node, "cp_hicache", None) + padded_len = getattr(metadata, "padded_len", None) + if padded_len is not None: + return int(padded_len) + return ceil_to_page_len(len(value), getattr(self, "page_size", 1)) + def pin_prefix( self, token_ids: List[int], ttl_seconds: int = 300 ) -> Tuple[int, Optional[str]]: @@ -2455,10 +2476,11 @@ class HiRadixCache(RadixCache): delta = 0 while node != self.root_node: + resident_len = self._node_device_resident_len(node) if node.lock_ref == 0: - self.evictable_size_ -= len(node.key) - self.protected_size_ += len(node.key) - delta -= len(node.key) + self.evictable_size_ -= resident_len + self.protected_size_ += resident_len + delta -= resident_len node.lock_ref += 1 self._update_leaf_status(node) self._update_host_leaf_status(node) @@ -2473,10 +2495,11 @@ class HiRadixCache(RadixCache): delta = 0 while node != self.root_node: + resident_len = self._node_device_resident_len(node) if node.lock_ref == 1: - self.evictable_size_ += len(node.key) - self.protected_size_ -= len(node.key) - delta += len(node.key) + self.evictable_size_ += resident_len + self.protected_size_ -= resident_len + delta += resident_len node.lock_ref -= 1 self._update_leaf_status(node) self._update_host_leaf_status(node) @@ -2487,6 +2510,34 @@ class HiRadixCache(RadixCache): node = node.parent return DecLockRefResult(delta=delta) + def inc_node_lock_ref(self, node: TreeNode): + if self.disable: + return + if node == self.root_node: + return + resident_len = self._node_device_resident_len(node) + if node.lock_ref == 0: + self.evictable_size_ -= resident_len + self.protected_size_ += resident_len + node.lock_ref += 1 + self._update_leaf_status(node) + if hasattr(self, "evictable_host_leaves"): + self._update_host_leaf_status(node) + + def dec_node_lock_ref(self, node: TreeNode): + if self.disable: + return + if node == self.root_node: + return + resident_len = self._node_device_resident_len(node) + if node.lock_ref == 1: + self.evictable_size_ += resident_len + self.protected_size_ -= resident_len + node.lock_ref -= 1 + self._update_leaf_status(node) + if hasattr(self, "evictable_host_leaves"): + self._update_host_leaf_status(node) + def _update_host_leaf_status(self, node: TreeNode): if not node.evicted or node.lock_ref > 0: if node in self.evictable_host_leaves: @@ -2587,13 +2638,15 @@ class HiRadixCache(RadixCache): def _evict_backuped(self, node: TreeNode): # GPU -> CPU demotion: no BlockRemoved since block is still reachable via load_back - num_evicted = self.cache_controller.evict_device(node.value) - assert num_evicted > 0 - self.evictable_size_ -= num_evicted + device_resident_len = self._node_device_resident_len(node) + freed_len = self.cache_controller.evict_device(node.value) + assert freed_len > 0 + self.evictable_size_ -= device_resident_len logger.info( - "[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d lock_ref=%d backed=%s", + "[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d physical_tokens=%d lock_ref=%d backed=%s", node.id, - num_evicted, + freed_len, + device_resident_len, node.lock_ref, self._node_backuped(node), ) @@ -2603,11 +2656,11 @@ class HiRadixCache(RadixCache): # update leaf status for the parent because the node is evicted self._update_leaf_status(node.parent) self._update_host_leaf_status(node.parent) - return num_evicted + return device_resident_len def _evict_regular(self, node: TreeNode): # evict a node not initiated write to host -- emit BlockRemoved - num_evicted = len(node.value) + num_evicted = self._node_device_resident_len(node) logger.info( "[HiCache-evict] _evict_regular: node_id=%d num_evicted=%d", node.id, @@ -2618,6 +2671,18 @@ class HiRadixCache(RadixCache): self._delete_leaf(node) return num_evicted + def _delete_leaf(self, node: TreeNode): + key = self.get_child_key_fn(node.key) + v = node.parent.children.pop(key, None) + assert v == node, f"parent does not have child key, {key}" + + self.evictable_size_ -= self._node_device_resident_len(node) + if node in self.evictable_leaves: + self.evictable_leaves.remove(node) + self._update_leaf_status(node.parent) + if hasattr(self, "evictable_host_leaves"): + self._update_host_leaf_status(node.parent) + def _remove_host_leaf(self, node: TreeNode) -> TreeNode: parent = node.parent key = self.get_child_key_fn(node.key) @@ -3372,7 +3437,6 @@ class HiRadixCache(RadixCache): if ( self._uses_cp_hicache and self.page_size > 1 - and self._node_backuped(child) and prefix_len % self.page_size != 0 ): prefix_len = self._cp_floor_backed_partial_split_len( @@ -3474,7 +3538,7 @@ class HiRadixCache(RadixCache): # change the reference if the node is evicted # this often happens in the case of KV cache recomputation node.value = value[:prefix_len].clone() - self.evictable_size_ += len(node.value) + self.evictable_size_ += self._node_device_resident_len(node) self._update_leaf_status(node) self._update_host_leaf_status(node) # update parent status as a new leaf is added into device @@ -3495,7 +3559,7 @@ class HiRadixCache(RadixCache): new_node.priority = max(new_node.priority, priority) if new_node.evicted: new_node.value = value[:prefix_len].clone() - self.evictable_size_ += len(new_node.value) + self.evictable_size_ += self._node_device_resident_len(new_node) self._update_leaf_status(new_node) self._update_host_leaf_status(new_node) # update parent status as a new leaf is added into device @@ -3529,7 +3593,7 @@ class HiRadixCache(RadixCache): new_node.key = key new_node.value = value.clone() node.children[child_key] = new_node - self.evictable_size_ += len(value) + self.evictable_size_ += self._node_device_resident_len(new_node) self._update_leaf_status(node) self._update_leaf_status(new_node) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6a55209ce..554a90a39 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -117,6 +117,12 @@ def page_align_keys(key: list, page_size) -> list: return key[:page_aligned_len] +def ceil_to_page_len(length: int, page_size: int) -> int: + if page_size <= 1: + return length + return ((length + page_size - 1) // page_size) * page_size + + class TreeNode: counter = 0 @@ -526,8 +532,19 @@ class RadixCache(BasePrefixCache): if prepared_cp_backup is not None: req.cp_hicache_prepared_backup = None - # free the unaligned tail - self.token_to_kv_pool_allocator.free(kv_indices[len(keys) :]) + # Free the unaligned tail. + # + # CP HiCache radix keys are scheduler-visible valid lengths, while the + # backing allocator is page-granular. Freeing a loc inside the retained + # tail page releases the whole page and leaves radix accounting pointing + # at freed memory (observed as available_size=max with evictable_size>0 + # after EAGLE warmup). Therefore CP starts tail free at the next page + # boundary; non-CP keeps the historical token boundary because keys were + # already floored by page_align_keys above. + tail_free_start = len(keys) + if getattr(self, "_uses_cp_hicache", False): + tail_free_start = ceil_to_page_len(tail_free_start, self.page_size) + self.token_to_kv_pool_allocator.free(kv_indices[tail_free_start:]) # Remove req slot release the cache lock self.dec_lock_ref(req.last_node) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 4a1b9c18a..0d6a027ac 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -666,7 +666,133 @@ class FakeTokenAllocator: return "FakeTokenAllocator" +class RecordingTokenAllocator: + def __init__(self): + self.freed = [] + + def free(self, free_index): + if free_index.numel() > 0: + self.freed.append(free_index.detach().cpu().tolist()) + + +class RecordingDeviceAllocator: + def __init__(self): + self.freed = [] + + def free(self, free_index): + self.freed.append(free_index.detach().cpu().tolist()) + + class TestHiRadixCacheCPBackup(CustomTestCase): + def _minimal_cp_hiradix_cache(self, *, page_size=64): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable = False + cache._uses_cp_hicache = True + cache.page_size = page_size + cache.is_eagle = False + cache.enable_storage = False + cache.enable_kv_cache_events = False + cache.evictable_size_ = 0 + cache.protected_size_ = 0 + cache.evictable_leaves = set() + cache.evictable_host_leaves = set() + cache.get_child_key_fn = lambda key: tuple(key.token_ids[:page_size]) + cache.key_match_fn = lambda key0, key1: _key_match_paged( + key0, key1, page_size + ) + cache.cache_controller = types.SimpleNamespace(write_policy="write_back") + cache._record_store_event = lambda node: None + cache._record_remove_event = lambda node: None + + root = TreeNode() + root.key = RadixKey(token_ids=[], extra_key=None) + root.value = torch.empty((0,), dtype=torch.int64) + root.children = {} + root.parent = None + cache.root_node = root + return cache + + def test_cp_eagle_finished_cache_preserves_retained_tail_page(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache.disable_finished_insert = False + cache.disable = False + cache.is_eagle = True + cache.page_size = 64 + cache._uses_cp_hicache = True + cache.req_to_token_pool = types.SimpleNamespace( + req_to_token=torch.arange(128, dtype=torch.int64).view(1, 128) + ) + allocator = RecordingTokenAllocator() + cache.token_to_kv_pool_allocator = allocator + cache.insert = lambda params: types.SimpleNamespace(prefix_len=0) + cache.dec_lock_ref = lambda node: None + + req = types.SimpleNamespace( + origin_input_ids=[10, 11, 12, 13], + output_ids=[], + req_pool_idx=0, + extra_key=None, + cache_protected_len=0, + last_node=object(), + cp_hicache_prepared_backup=None, + pop_committed_kv_cache=lambda: 4, + ) + + cache.cache_finished_req(req) + + self.assertEqual(allocator.freed, []) + + def test_cp_valid_tail_device_accounting_uses_physical_page_span(self): + cache = self._minimal_cp_hiradix_cache(page_size=64) + + result = cache.insert( + InsertParams( + key=RadixKey(token_ids=[1, 2, 3], extra_key=None), + value=torch.arange(3, dtype=torch.int64), + ) + ) + + self.assertEqual(result.prefix_len, 0) + self.assertEqual(cache.evictable_size_, 64) + node = next(iter(cache.root_node.children.values())) + + cache.inc_node_lock_ref(node) + self.assertEqual(cache.evictable_size_, 0) + self.assertEqual(cache.protected_size_, 64) + + cache.dec_node_lock_ref(node) + self.assertEqual(cache.evictable_size_, 64) + self.assertEqual(cache.protected_size_, 0) + + def test_cp_valid_tail_regular_evict_reports_and_subtracts_physical_page(self): + cache = self._minimal_cp_hiradix_cache(page_size=64) + device_allocator = RecordingDeviceAllocator() + cache.cache_controller.mem_pool_device_allocator = device_allocator + cache.insert( + InsertParams( + key=RadixKey(token_ids=[1, 2, 3], extra_key=None), + value=torch.arange(3, dtype=torch.int64), + ) + ) + node = next(iter(cache.root_node.children.values())) + + num_evicted = cache._evict_regular(node) + + self.assertEqual(num_evicted, 64) + self.assertEqual(cache.evictable_size_, 0) + self.assertEqual(device_allocator.freed, [[0, 1, 2]]) + self.assertEqual(cache.root_node.children, {}) + + def test_cp_partial_split_floors_unbacked_valid_tail_to_page_boundary(self): + cache = self._minimal_cp_hiradix_cache(page_size=64) + node = TreeNode() + node.host_len = 0 + node.cp_hicache = None + + self.assertEqual(cache._cp_floor_backed_partial_split_len(node, 3), 0) + self.assertEqual(cache._cp_floor_backed_partial_split_len(node, 70), 64) + self.assertEqual(cache._cp_floor_backed_partial_split_len(node, 128), 128) + def test_session_aware_cache_forwards_cp_hicache_prepare(self): calls = [] 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 2dab40eba..0e1d04a81 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 @@ -7,16 +7,45 @@ from unittest.mock import patch import torch _sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT") -try: - _sgl_kernel_lib.define( - "moe_fused_gate(Tensor input_tensor, Tensor? bias, int num_expert_group, " - "int topk_group, int topk, int num_fused_shared_experts, " - "float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) " - "-> (Tensor, Tensor)" - ) -except RuntimeError as exc: - if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): - raise + + +def _define_sgl_kernel_stub(schema: str) -> None: + try: + _sgl_kernel_lib.define(schema) + except RuntimeError as exc: + if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower(): + raise + + +_define_sgl_kernel_stub( + "moe_fused_gate(Tensor input_tensor, Tensor? bias, int num_expert_group, " + "int topk_group, int topk, int num_fused_shared_experts, " + "float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) " + "-> (Tensor, Tensor)" +) +_define_sgl_kernel_stub( + "sgl_per_token_quant_fp8(Tensor input, Tensor output_q, Tensor output_s) -> ()" +) +_define_sgl_kernel_stub( + "sgl_per_token_group_quant_fp8(Tensor input, Tensor output_q, Tensor output_s, " + "int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()" +) +_define_sgl_kernel_stub( + "sgl_per_token_group_quant_8bit(Tensor input, Tensor output_q, Tensor output_s, " + "int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()" +) +_define_sgl_kernel_stub( + "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, " + "ScalarType out_dtype, Tensor? bias) -> Tensor" +) +_define_sgl_kernel_stub( + "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, " + "Tensor scales_b, ScalarType out_dtype) -> Tensor" +) +_define_sgl_kernel_stub( + "int8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, " + "ScalarType out_dtype, Tensor? bias) -> Tensor" +) flash_attn_stub = sys.modules.setdefault( "sgl_kernel.flash_attn", types.ModuleType("sgl_kernel.flash_attn") @@ -32,6 +61,72 @@ if not hasattr(sgl_kernel_stub, "__path__"): sgl_kernel_stub.__path__ = [] if not hasattr(sgl_kernel_stub, "flash_attn"): sgl_kernel_stub.flash_attn = flash_attn_stub +quantization_stub = sys.modules.setdefault( + "sgl_kernel.quantization", types.ModuleType("sgl_kernel.quantization") +) +if not hasattr(sgl_kernel_stub, "quantization"): + sgl_kernel_stub.quantization = quantization_stub +for _name in ( + "ggml_dequantize", + "ggml_moe_a8", + "ggml_moe_a8_vec", + "ggml_moe_get_block_size", + "ggml_mul_mat_a8", + "ggml_mul_mat_vec_a8", +): + if not hasattr(quantization_stub, _name): + setattr(quantization_stub, _name, lambda *args, **kwargs: None) +if not hasattr(sgl_kernel_stub, "sgl_per_token_quant_fp8"): + sgl_kernel_stub.sgl_per_token_quant_fp8 = lambda *args, **kwargs: None +if not hasattr(sgl_kernel_stub, "sgl_per_token_group_quant_fp8"): + sgl_kernel_stub.sgl_per_token_group_quant_fp8 = lambda *args, **kwargs: None +if not hasattr(sgl_kernel_stub, "sgl_per_token_group_quant_int8"): + sgl_kernel_stub.sgl_per_token_group_quant_int8 = lambda *args, **kwargs: None +for _name in ( + "concat_mla_absorb_q", + "gelu_and_mul", + "silu_and_mul", + "moe_align_block_size", + "moe_sum", + "moe_sum_reduce", + "moe_fused_gate", + "kimi_k2_moe_fused_gate", + "topk_softmax", + "topk_sigmoid", + "fast_topk_transform_fused", + "fast_topk_transform_ragged_fused", + "fast_topk_v2", + "fused_add_rmsnorm", + "gemma_fused_add_rmsnorm", + "gemma_rmsnorm", + "sgl_per_token_group_quant_8bit", + "fp8_blockwise_scaled_mm", + "fp8_scaled_mm", + "int8_scaled_mm", + "gptq_gemm", + "gptq_shuffle", + "qserve_w4a8_per_chn_gemm", + "qserve_w4a8_per_group_gemm", + "awq_dequantize", + "fused_experts", + "apply_shuffle_mul_sum", + "es_fp8_blockwise_scaled_grouped_mm", + "es_sm100_mxfp8_blockscaled_grouped_mm", + "es_sm100_mxfp8_blockscaled_grouped_quant", + "fp8_blockwise_scaled_grouped_mm", + "prepare_moe_input", + "shuffle_rows", + "cutlass_w4a8_moe_mm", + "get_cutlass_w4a8_moe_mm_data", + "merge_state_v2", + "cutlass_mla_decode", + "cutlass_mla_get_workspace_size", + "causal_conv1d_fwd", + "causal_conv1d_update", + "rmsnorm", +): + if not hasattr(sgl_kernel_stub, _name): + setattr(sgl_kernel_stub, _name, lambda *args, **kwargs: None) from sglang.test.ci.ci_register import register_cpu_ci @@ -252,7 +347,25 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(prefetch_stream.waited, ["current"]) def test_mla_pool_prefetch_getter_orders_layer_transfer_on_prefetch_stream(self): - from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool + index_accessor_stub = types.ModuleType( + "sglang.srt.layers.attention.nsa.index_buf_accessor" + ) + + class _IndexAccessorOp: + @staticmethod + def execute(*args, **kwargs): + raise AssertionError("index accessor is not used by this test") + + for _name in ("GetK", "GetS", "GetKAndS", "SetKAndS"): + setattr(index_accessor_stub, _name, _IndexAccessorOp) + + with patch.dict( + sys.modules, + { + "sglang.srt.layers.attention.nsa.index_buf_accessor": index_accessor_stub + }, + ): + from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool class FakeCounter: def __init__(self): @@ -483,6 +596,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) forward_batch.spec_info = TargetSpecInfo() + forward_batch.cp_shared_kv_mla_prefetcher = object() + self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) + + forward_batch.cp_shared_kv_mla_prefetcher = None self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) forward_batch.spec_info = DraftSpecInfo() @@ -490,6 +607,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): forward_batch.seq_lens_cpu = torch.tensor([56], dtype=torch.int32) self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) + forward_batch.spec_info = TargetSpecInfo() + forward_batch.cp_shared_kv_mla_prefetcher = None + self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch)) + def test_runtime_fallback_helpers_use_standard_warning_marker(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -609,6 +730,49 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]]) self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]]) + def test_materialize_prefix_and_reuse_current_kv_page_slots_without_prefetcher( + self, + ): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + page_size = 4 + layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0) + kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1) + logical_locs = torch.tensor([[4, 8, 20, 21, 22, 23]], dtype=torch.int64) + current_locs = torch.tensor([20, 21], dtype=torch.int64) + current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1) + remap_logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64) + slot_remap = runtime.build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=logical_locs, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=page_size, + ) + + with patch.object( + runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce + ): + mixed_kv, mixed_locs = ( + runtime.materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=logical_locs, + current_kv_cache=current_kv, + current_locs=current_locs, + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + prefix_pages=2, + ) + ) + + self.assertEqual(list(mixed_kv.shape), [16, 1, 1]) + self.assertTrue(torch.equal(mixed_kv[4:8], kv_cache[4:8])) + self.assertTrue(torch.equal(mixed_kv[8:12], kv_cache[8:12])) + self.assertTrue(torch.equal(mixed_kv[12:14], current_kv)) + self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]]) + def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -1240,6 +1404,103 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): 32, ) + def test_mla_prefetch_min_async_extend_tokens_defaults_to_one_page_and_can_override( + self, + ): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.clear() + self.assertEqual( + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=64), + 64, + ) + self.assertEqual( + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=None), + 0, + ) + + with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.override(0): + self.assertEqual( + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=64), + 0, + ) + + with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_EXTEND_TOKENS.override(128): + self.assertEqual( + runtime.cp_shared_kv_mla_prefetch_min_async_extend_tokens(page_size=64), + 128, + ) + + def test_mla_and_index_prefetch_skip_tiny_extend_even_with_large_prefix(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + hisparse_coordinator=None, + forward_mode=Mode(), + batch_size=1, + token_to_kv_pool=SimpleNamespace(page_size=64, start_layer=0), + cp_shared_kv_layout=SimpleNamespace(cp_size=8, cp_rank=0), + extend_prefix_lens_cpu=[16320], + extend_seq_lens_cpu=[16], + ) + metadata = SimpleNamespace( + real_page_table=torch.arange(256, dtype=torch.int64), + page_table_1=torch.zeros((1, 16336), dtype=torch.int32), + ) + stream = object() + kv_cache = torch.zeros((4096, 2), dtype=torch.float32) + remap = SimpleNamespace( + slot_logical_pages=torch.arange(1, 257, dtype=torch.int64), + page_inverse=torch.arange(0, 257, dtype=torch.int64), + dense_num_pages=257, + ) + + with patch.object( + prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True + ), patch.object( + prefetch, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + prefetch.torch.cuda, "is_available", return_value=True + ), patch.object( + prefetch, "_is_cuda_stream_capturing", return_value=False + ), patch.object( + prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True + ), patch.object( + prefetch, "get_attention_cp_group", return_value=SimpleNamespace(pynccl_comm=object()) + ), patch.object( + prefetch.torch.cuda, "Stream", return_value=stream + ), patch.object( + prefetch, "_prefetch_pool_get_key_buffer", return_value=kv_cache + ) as mla_getter, patch.object( + prefetch, "get_or_build_shared_token_kv_slot_remap", return_value=remap + ) as token_remap, patch.object( + prefetch, + "_prefetch_pool_get_index_buffer", + side_effect=AssertionError("index getter should not be reached"), + ) as index_getter: + mla_result = prefetch.CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=True, + ) + index_result = prefetch.CpSharedKVIndexPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=True, + ) + + self.assertIsNone(mla_result) + self.assertIsNone(index_result) + mla_getter.assert_not_called() + token_remap.assert_not_called() + index_getter.assert_not_called() + def test_fused_mla_store_uses_tai_kernel_when_enabled(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout