chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 Packages Projects Releases Wiki Activity

Keep current reuse independent of async CP prefetch

Browse Source 
Async MLA/index prefetch is a scheduling optimization, not the correctness contract for target current reuse. Tiny cache-hit suffixes can skip async prefetcher creation while target partial-current reuse still composes page-slot prefix materialization with current KV rows synchronously. CP HiCache radix/device accounting now treats retained valid-tail pages as physical page spans so allocator state stays consistent when logical cache keys are shorter than the retained page.

Constraint: CP shared KV ownership and HiCache residency are page-granular while request-visible cache lengths remain valid-token lengths.
Constraint: Async prefetch can hang or regress on large-prefix tiny-extend traffic and must not be required for current reuse.
Rejected: Treat missing prefetcher as fail-fast for target partial-current reuse | disabled useful current reuse and broke tiny-prefix/tiny-suffix traffic.
Rejected: Keep async prefetcher object with synchronous consume mode | conflates prefetch object existence with current-layer correctness and hides fallback semantics.
Confidence: medium
Scope-risk: moderate
Directive: Do not make current-only or target partial-current reuse depend on MLA/index prefetcher creation; prefetcher objects mean async next-layer work exists.
Tested: Remote g0034 container py_compile for touched modules.
Tested: Remote g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 73 passed.
Tested: Remote g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py -> 90 passed.
Not-tested: Latest full ETE traffic run with GLM-5.1 CP HiCache after this commit.
Not-tested: CUDA kernel-level performance impact of synchronous no-prefetch partial-current compose.
This commit is contained in:
laoyao0822
2026-05-29 19:34:27 +08:00
parent 2b1524bd8c
commit 2a9dfcca6f
9 changed files with 1067 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
+359
View File
@@ -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`
python/sglang/srt/environ.py
+1
View File
@@ -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)
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+47 -1
View File
@@ -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(
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+98
View File
@@ -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,
python/sglang/srt/layers/attention/nsa_backend.py
+63 -16
View File
@@ -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)
python/sglang/srt/mem_cache/hiradix_cache.py
+82 -18
View File
@@ -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)
python/sglang/srt/mem_cache/radix_cache.py
+19 -2
View File
@@ -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)
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+126
View File
@@ -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 = []
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+272 -11
View File
@@ -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