Keep CP HiCache reuse page-safe without eviction log churn
Page-aligned CP shared KV can pad out_cache_loc beyond valid current rows, so current reuse now gates MLA composition on the valid extend rows and permits draft partial-current reuse once the TAI sparse-page capability check passes. The TAI current-slot path self-tests sparse pages before use and falls back to the torch reference when the installed kernel is stale. Eviction success and no-op diagnostics were also moved from INFO to DEBUG so owner-lane and host-admission churn does not flood production logs; true write failures remain WARNING. Constraint: CP shared KV uses page-aligned physical reservations where valid suffix rows can be shorter than padded out_cache_loc. Constraint: Production failure/fallback logs must remain visible, but hot successful eviction paths should not emit INFO per victim/rank. Rejected: Keep draft partial-current reuse disabled | would preserve avoidable full materialization on draft cache-hit suffixes. Rejected: Trust the TAI current-slot kernel unconditionally | stale kernels can corrupt sparse current-page composition. Confidence: medium Scope-risk: moderate Directive: Do not reintroduce INFO logging in eviction hot paths without rate limiting and runtime evidence. Tested: local py_compile for touched Python files Tested: local git diff --check Tested: remote container py_compile for touched Python files Tested: remote PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py::TestHiCacheEvictLoggingLevels::test_evict_hot_path_success_logs_are_debug_only test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 82 passed, 5 warnings, 2 subtests passed Not-tested: full ETE traffic after this commit; draft partial-current accept length still needs user-driven runtime validation
This commit is contained in:
laoyao0822
@@ -3365,3 +3365,251 @@ Verification on g0034 container:
|
||||
|
||||
Remaining risk:
|
||||
- This makes the all-layer direct/page-first backup API submit one TAI per-layer memcpy-batch op instead of one stale sgl-kernel all-layer op. It is correct and CUDA-13-safe, but it can increase CPU submit count on all-layer backup paths. Production forward-overlap paths already operate per-layer, so this is acceptable for the current branch goal.
|
||||
|
||||
### C81 — 2026-05-31 correction: do not reopen draft partial-current reuse before fixing root cause
|
||||
|
||||
Correction:
|
||||
|
||||
- Draft partial-current reuse must not be reopened merely by deleting the
|
||||
`draft_partial_current_reuse_disabled` guard.
|
||||
- The observed accept-length regression must be fixed first, then the default
|
||||
can be reopened without an env gate.
|
||||
- Until that root cause is fixed and verified, the safe contract remains:
|
||||
draft current-only reuse is allowed, draft cache-hit partial-current reuse is
|
||||
disabled with the explicit `draft_partial_current_reuse_disabled` warning.
|
||||
|
||||
Current known facts:
|
||||
|
||||
- Draft HiCache L2(host)→L1(device) load-back exists via `CacheController.load_cp()`
|
||||
and `CacheController.start_loading()` using `draft_load_queue` and
|
||||
`draft_mem_pool_host.load_to_device_per_layer(...)` when `has_draft_hicache` is
|
||||
true.
|
||||
- That L2→L1 load-back is separate from the CP shared-KV MLA/index next-layer
|
||||
prefetcher; the latter is still target-only because EAGLE/NextN has one
|
||||
executable draft layer.
|
||||
|
||||
Next target:
|
||||
|
||||
- Audit the draft partial-current compose path (`nsa_backend.py` plus
|
||||
`materialize_prefix_and_reuse_current_kv_page_slots`) for target/draft layout,
|
||||
layer-id, current-suffix placement, and page-slot mapping mismatches before
|
||||
reopening the guard.
|
||||
|
||||
### C82 — 2026-05-31 fixed-before-open path for draft partial-current reuse
|
||||
|
||||
Root-cause guard added before reopening draft partial-current reuse:
|
||||
|
||||
- The risky failure mode was not the draft predicate itself; it was accepting a
|
||||
stale TAI current-slot fill kernel whose masking logic assumed the current
|
||||
suffix covered a dense page range from first current page to last current page.
|
||||
- For sparse current pages, that stale kernel could mask unrelated prefix pages
|
||||
as suffix slack, corrupting partial-current compose while tensor shapes stayed
|
||||
valid.
|
||||
- `cp_shared_kv_runtime.py` now validates the installed TAI current-slot fill
|
||||
kernel once per CUDA device with a sparse-page self-test before using it. If
|
||||
the self-test fails, the TAI result is not used and the code falls back to the
|
||||
torch reference path, which preserves correctness.
|
||||
|
||||
After this capability fix:
|
||||
|
||||
- Draft cache-hit partial-current reuse is reopened directly, without an env
|
||||
gate.
|
||||
- Draft next-layer MLA/index prefetcher creation remains disabled; draft has one
|
||||
executable layer, so a future draft prefetch optimization should use an
|
||||
explicit same-layer contract.
|
||||
|
||||
Verification:
|
||||
|
||||
- Remote RED before runtime helper existed:
|
||||
`test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails`
|
||||
failed with missing `_tai_current_slot_fill_supports_sparse_pages`.
|
||||
- Remote GREEN after runtime helper:
|
||||
`test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails` → pass.
|
||||
- Remote installed-kernel evidence:
|
||||
`test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel`
|
||||
→ pass.
|
||||
- Remote tai-kernel sparse-page tests:
|
||||
`test_fill_current_token_kv_page_slots_handles_sparse_current_pages`,
|
||||
`test_fill_current_token_kv_page_slots_handles_nonzero_first_page_offset`, and
|
||||
`test_fill_current_token_kv_page_slots_masks_current_page_slack` → `3 passed`.
|
||||
|
||||
Remaining risk:
|
||||
|
||||
- Need ETE run with draft `branch=partial_current_sync` active to confirm decode
|
||||
aggregate accept remains healthy. The code-level stale-kernel corruption path
|
||||
is now blocked, but runtime traffic is still required before treating this as
|
||||
final performance/correctness validation.
|
||||
|
||||
### C83 — 2026-05-31 Mooncake P2P handshake receives HTTP traffic on port 17000
|
||||
|
||||
Observed remote log:
|
||||
|
||||
- `readString: too large length from socket: 3564087971010531152`
|
||||
- `SocketHandShakePlugin: failed to receive handshake message, malformed json`
|
||||
|
||||
Root-cause evidence:
|
||||
|
||||
- `3564087971010531152 == 0x31762f2054534f50`, little-endian bytes are
|
||||
`POST /v1`.
|
||||
- Tcpdump on g0034 captured the offending packet:
|
||||
`10.20.34.3:<ephemeral> > 10.20.32.34:17000` with HTTP payload
|
||||
`POST /v1/chat/completions HTTP/1.1` and `Host: 10.20.32.34:17000`.
|
||||
- In the same run, Mooncake P2P handshake selected port `17000` for one prefill
|
||||
scheduler process:
|
||||
`Transfer Engine RPC using P2P handshake, listening on 10.20.32.34:17000`.
|
||||
- The errors repeat every ~30 seconds and start before request traffic while the
|
||||
model is still loading, so this is not caused by HiCache per-layer backup
|
||||
payload corruption.
|
||||
|
||||
Current conclusion:
|
||||
|
||||
- This is a port/protocol collision: an external HTTP client/prober is sending
|
||||
OpenAI HTTP traffic to a Mooncake metadata/P2P handshake listener.
|
||||
- It is noisy and can mask real transfer errors, but it is not evidence that
|
||||
HiCache write/read payloads are corrupted.
|
||||
|
||||
Decision:
|
||||
|
||||
- Do not keep the temporary SGLang-side Mooncake port-range override.
|
||||
- The code change that set `MC_MAX_PRC_PORT=16999` by default was reverted on
|
||||
2026-05-31 at user request.
|
||||
- If this noise matters operationally, fix the external endpoint/prober or set
|
||||
Mooncake `MC_MIN_PRC_PORT`/`MC_MAX_PRC_PORT` explicitly in the launch
|
||||
environment instead of silently changing SGLang defaults.
|
||||
|
||||
### C84 — 2026-05-31 draft partial-current reuse also requires valid-row, not padded-loc, MLA gating
|
||||
|
||||
Finding:
|
||||
|
||||
- Reopening draft partial-current reuse still was not enough when `out_cache_loc`
|
||||
was page padded. The MLA gate in `nsa_backend.py` required
|
||||
`k.shape[0] == forward_batch.out_cache_loc.numel()` and the same for
|
||||
`k_rope`.
|
||||
- Under the page-aligned cache contract, `out_cache_loc` may cover the padded
|
||||
physical page extent while MLA projection tensors only contain the valid
|
||||
current suffix rows from `extend_seq_lens_cpu`.
|
||||
- That strict equality silently forces cache-hit draft suffixes (for example
|
||||
`extend_len=65`, padded locs `128`) away from partial-current reuse and back
|
||||
toward full materialization.
|
||||
|
||||
Fix direction:
|
||||
|
||||
- Validate current MLA tensors against valid current rows, not padded loc count.
|
||||
- Slice `k`, `k_rope`, and `out_cache_loc` to `extend_seq_lens_cpu[0]` before
|
||||
composing current rows. Padded tail locs are physical reservation/slack, not
|
||||
visible KV rows.
|
||||
- Keep draft next-layer MLA/index prefetcher disabled; this is only restoring the
|
||||
synchronous partial-current compose path for draft cache-hit suffixes.
|
||||
|
||||
Verification added:
|
||||
|
||||
- RED on remote before implementation:
|
||||
`test_current_extend_kv_rows_for_reuse_accepts_padded_out_cache_loc` failed
|
||||
because `current_extend_kv_rows_for_reuse` did not exist.
|
||||
- RED on remote before implementation:
|
||||
`test_mla_current_reuse_gate_accepts_padded_out_cache_loc` found the stale
|
||||
`k.shape[0] == forward_batch.out_cache_loc.numel()` gate.
|
||||
|
||||
Remaining runtime check:
|
||||
|
||||
- Remote GREEN after implementation:
|
||||
`test_should_reuse_current_extend_kv_enables_draft_partial_cache_hit_suffix`,
|
||||
`test_current_extend_kv_rows_for_reuse_accepts_padded_out_cache_loc`,
|
||||
`test_mla_current_reuse_gate_accepts_padded_out_cache_loc`,
|
||||
`test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails`, and
|
||||
`test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel`
|
||||
→ `5 passed`.
|
||||
- Remote broader runtime helper suite:
|
||||
`PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py`
|
||||
→ `81 passed, 5 warnings, 2 subtests passed`.
|
||||
- Still requires user-driven ETE traffic to verify that draft `partial_current_sync`
|
||||
appears without accept-length collapse.
|
||||
|
||||
### C85 — 2026-05-31 CP HiCache eviction path and overhead findings
|
||||
|
||||
Code path checked:
|
||||
|
||||
- Extend allocation under CP shared KV uses compute-owner page allocation in
|
||||
`alloc_paged_token_slots_extend()`.
|
||||
- If owner-lane allocation fails, `_evict_for_compute_owner_lanes()` computes
|
||||
per-owner deficits and calls `tree_cache.evict(EvictParams(...,
|
||||
owner_lane_deficits=...))`.
|
||||
- CP HiCache routes that to `_evict_cp_owner_lane_deficit_nodes()`, which plans
|
||||
victim nodes by owner-lane contribution and frees backed GPU nodes through
|
||||
`_evict_backuped()` -> `CacheController.evict_device()` ->
|
||||
`mem_pool_device_allocator.free()`.
|
||||
- Load-back can also evict device nodes through `_evict_cp_load_back_owner_lanes()`.
|
||||
- Host write admission evicts host nodes through
|
||||
`_reserve_write_cp_indices_no_collective()` ->
|
||||
`_evict_cp_host_for_write_admission()` -> `CacheController.evict_cp_host()`.
|
||||
|
||||
Remote log evidence from `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260530_184620.log`:
|
||||
|
||||
- `MemCache-evict _evict_for_compute_owner_lanes`: 3224 attempt lines across 8
|
||||
ranks, about 403 logical eviction cycles in ~18 minutes.
|
||||
- Average owner-lane request was ~18.5k tokens; average actual evicted was ~38k
|
||||
tokens, i.e. about ~19.5k token overshoot per rank-line.
|
||||
- Small requests were common: 1336 owner-lane evict lines requested less than
|
||||
4096 tokens, but eviction can still remove an entire victim node.
|
||||
- Host write-admission eviction appeared too: 1088 lines, about 136 logical
|
||||
host eviction cycles. Average local host freed was ~4.7k token slots.
|
||||
|
||||
Current conclusion:
|
||||
|
||||
- Eviction does not copy KV for already-backed nodes; it mainly frees allocator
|
||||
slots and updates radix/HiCache metadata. It should not consume PCIe bandwidth.
|
||||
- It can still be CPU/GPU-scheduler expensive because it is synchronous on the
|
||||
scheduler path, scans/heapifies evictable leaves, emits many INFO logs, and
|
||||
`PagedTokenToKVPoolAllocator.free()` uses `torch.unique(free_index // page_size)`
|
||||
per victim.
|
||||
- Per-owner capacity stats also do CUDA tensor reductions with `.item()` per
|
||||
owner lane, so failed owner-lane allocations can introduce host syncs.
|
||||
- The large overshoot is structural: eviction frees whole radix nodes, not the
|
||||
exact deficit. This can cause L1 churn even after removing the older `* cp_size`
|
||||
multiplier.
|
||||
|
||||
Candidate optimizations to consider:
|
||||
|
||||
1. Rate-limit or demote hot-path INFO eviction logs; current log volume alone is
|
||||
large enough to affect performance during heavy churn.
|
||||
2. Add a page-aligned fast path in `PagedTokenToKVPoolAllocator.free()` to use
|
||||
`free_index[::page_size] // page_size` when inputs are page-shaped, avoiding
|
||||
`torch.unique` for normal HiCache node eviction.
|
||||
3. Cache/maintain owner-lane free counts incrementally in the CP allocator to
|
||||
avoid repeated `.sum().item()` reductions during owner-lane admission.
|
||||
4. Improve victim selection to reduce overshoot for small deficits, or keep a
|
||||
small-node victim pool for small owner-lane shortages.
|
||||
|
||||
### C86 — 2026-05-31 Eviction hot-path logs demoted to DEBUG
|
||||
|
||||
User decision:
|
||||
|
||||
- Keep failure/fallback logs visible.
|
||||
- Demote successful/no-op eviction hot-path diagnostics from INFO to DEBUG.
|
||||
|
||||
Implemented scope:
|
||||
|
||||
- `mem_cache/common.py`
|
||||
- `evict_from_tree_cache()` capacity-trigger message is DEBUG.
|
||||
- `_evict_for_compute_owner_lanes()` no-evictable, attempt, and result messages are DEBUG.
|
||||
- `mem_cache/hiradix_cache.py`
|
||||
- CP load-back owner-lane eviction summary is DEBUG.
|
||||
- CP owner-lane no-victim / stale-victim / END summaries are DEBUG.
|
||||
- Deterministic CP host eviction success summary is DEBUG.
|
||||
- Generic `evict START/END`, `_evict_backuped`, `_evict_regular`, and host-slot eviction start are DEBUG.
|
||||
- `write_backup CP retry after deterministic host eviction` is DEBUG.
|
||||
- `write_backup CP FAILED after deterministic retry` is promoted to WARNING so true failures remain visible.
|
||||
|
||||
Not changed:
|
||||
|
||||
- Existing `warning`/`error` fallback, capacity failure, OOM, pin-release, and plan-divergence logs remain visible.
|
||||
- Eviction semantics are unchanged; this only changes log levels.
|
||||
|
||||
Verification:
|
||||
|
||||
- Added `TestHiCacheEvictLoggingLevels::test_evict_hot_path_success_logs_are_debug_only`.
|
||||
- RED on remote before implementation: failed because `MemCache-evict` success-path logs still used `logger.info`.
|
||||
- GREEN on remote after implementation:
|
||||
`PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py::TestHiCacheEvictLoggingLevels::test_evict_hot_path_success_logs_are_debug_only`
|
||||
→ `1 passed`.
|
||||
- Local and remote `py_compile` passed for touched Python files.
|
||||
|
||||
@@ -9,7 +9,9 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.nsa.utils import log_cp_draft_shared_kv_debug # noqa: F401
|
||||
from sglang.srt.layers.attention.nsa.utils import (
|
||||
log_cp_draft_shared_kv_debug,
|
||||
) # noqa: F401
|
||||
from sglang.srt.layers.dp_attention import get_attention_cp_group
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
@@ -85,11 +87,7 @@ def cp_shared_kv_mla_prefetch_min_prefix_pages(
|
||||
if page_size is not None and int(page_size) > 0:
|
||||
min_pages = max(
|
||||
min_pages,
|
||||
(
|
||||
_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS
|
||||
+ int(page_size)
|
||||
- 1
|
||||
)
|
||||
(_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS + int(page_size) - 1)
|
||||
// int(page_size),
|
||||
)
|
||||
return min_pages
|
||||
@@ -302,6 +300,110 @@ def _load_tai_materialize_kernels():
|
||||
return None
|
||||
|
||||
|
||||
@lru_cache(maxsize=16)
|
||||
def _tai_current_slot_fill_sparse_pages_self_test(
|
||||
device_type: str,
|
||||
device_index: int | None,
|
||||
) -> bool:
|
||||
"""Return whether TAI current-slot fill handles sparse current pages.
|
||||
|
||||
Older tai-kernel builds filled current rows correctly but masked every page
|
||||
between first and last current page as current-page slack. That silently
|
||||
hid unrelated prefix pages for cache-hit suffixes and was enough to corrupt
|
||||
CP shared-KV reuse. Validate the installed kernel once before allowing it
|
||||
on the hot path.
|
||||
"""
|
||||
|
||||
if device_type != "cuda":
|
||||
return True
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
|
||||
kernels = _load_tai_materialize_kernels()
|
||||
fill_kernel = (
|
||||
getattr(kernels, "fill_current_token_kv_page_slots_and_remap_locs", None)
|
||||
if kernels is not None
|
||||
else None
|
||||
)
|
||||
if fill_kernel is None:
|
||||
return False
|
||||
|
||||
device = torch.device(device_type, device_index)
|
||||
try:
|
||||
page_size = 4
|
||||
dense_kv = torch.zeros((16, 1), device=device, dtype=torch.float32)
|
||||
materialized_locs = torch.tensor(
|
||||
[[4, 5, 8, 9, 12, 13, 14, 15]],
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
current_kv = torch.tensor(
|
||||
[[10.0], [11.0], [12.0], [13.0]],
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
logical_locs = torch.tensor(
|
||||
[[20, 21, 40, 41, 100, 101, 102, 103]],
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
current_locs = torch.tensor(
|
||||
[20, 21, 100, 101],
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
page_inverse = torch.full((32,), -1, device=device, dtype=torch.long)
|
||||
page_inverse[0] = 0
|
||||
page_inverse[5] = 1
|
||||
page_inverse[10] = 2
|
||||
page_inverse[25] = 3
|
||||
|
||||
mixed_kv, mixed_locs, current_mask = fill_kernel(
|
||||
dense_kv,
|
||||
materialized_locs,
|
||||
current_kv,
|
||||
logical_locs,
|
||||
current_locs,
|
||||
page_inverse,
|
||||
page_size=page_size,
|
||||
mask_non_current_in_current_pages=True,
|
||||
)
|
||||
expected_locs = torch.tensor(
|
||||
[[4, 5, 8, 9, 12, 13, -1, -1]],
|
||||
device=device,
|
||||
dtype=torch.int64,
|
||||
)
|
||||
expected_mask = torch.tensor(
|
||||
[[True, True, False, False, True, True, False, False]],
|
||||
device=device,
|
||||
dtype=torch.bool,
|
||||
)
|
||||
expected_kv = torch.zeros_like(dense_kv)
|
||||
expected_kv[4:6] = current_kv[0:2]
|
||||
expected_kv[12:14] = current_kv[2:4]
|
||||
return bool(
|
||||
torch.equal(mixed_locs, expected_locs)
|
||||
and torch.equal(current_mask, expected_mask)
|
||||
and torch.equal(mixed_kv, expected_kv)
|
||||
)
|
||||
except Exception as exc:
|
||||
_log_tai_materialize_fallback(
|
||||
"fill_current_sparse_page_self_test_failed",
|
||||
"CP shared KV tai current-slot fill sparse-page self-test failed; "
|
||||
"falling back to torch reference. error=%s",
|
||||
exc,
|
||||
limit=1,
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _tai_current_slot_fill_supports_sparse_pages(device: torch.device) -> bool:
|
||||
return _tai_current_slot_fill_sparse_pages_self_test(
|
||||
device.type,
|
||||
device.index,
|
||||
)
|
||||
|
||||
|
||||
def _tai_materialize_runtime_enabled() -> bool:
|
||||
# Keep the debug path on the existing torch implementation. The debug path
|
||||
# intentionally preserves tensor summaries and value assertions used for
|
||||
@@ -719,6 +821,22 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs(
|
||||
)
|
||||
return None
|
||||
|
||||
if dense_kv_cache.is_cuda and not _tai_current_slot_fill_supports_sparse_pages(
|
||||
dense_kv_cache.device
|
||||
):
|
||||
_log_tai_materialize_fallback(
|
||||
"fill_current_sparse_page_unsupported",
|
||||
"CP shared KV tai current-slot fill kernel failed the sparse-page "
|
||||
"capability check; falling back to torch reference to avoid hiding "
|
||||
"prefix pages in partial-current reuse. page_size=%s current_rows=%s "
|
||||
"query_locs=%s",
|
||||
page_size,
|
||||
int(current_kv_cache.shape[0]),
|
||||
int(logical_locs.numel()),
|
||||
limit=1,
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
return fill_kernel(
|
||||
_contiguous_for_tai(dense_kv_cache),
|
||||
@@ -866,8 +984,10 @@ def fill_current_index_page_slots(
|
||||
)
|
||||
|
||||
current_pages = torch.div(current_locs, page_size, rounding_mode="floor")
|
||||
valid_pages = (current_locs >= 0) & (current_pages >= 0) & (
|
||||
current_pages < int(page_inverse.numel())
|
||||
valid_pages = (
|
||||
(current_locs >= 0)
|
||||
& (current_pages >= 0)
|
||||
& (current_pages < int(page_inverse.numel()))
|
||||
)
|
||||
safe_pages = torch.clamp(
|
||||
current_pages,
|
||||
@@ -1075,12 +1195,12 @@ def can_reuse_current_extend_kv(forward_batch) -> bool:
|
||||
def should_reuse_current_extend_kv(forward_batch) -> bool:
|
||||
"""Return whether MLA should splice current extend KV into materialized KV.
|
||||
|
||||
Current-only reuse is safe for both target and draft because there is no
|
||||
cached prefix to compose. Partial current reuse is currently a target-model
|
||||
contract only. EAGLE/NextN draft cache-hit suffixes keep using the older
|
||||
full-materialize path until the draft splice path has value-level ETE proof;
|
||||
the 2026-05-30 accept-length regression correlated with enabling that draft
|
||||
partial splice path.
|
||||
Current-only and partial-current reuse are both cache-layout operations: the
|
||||
cached page-aligned prefix is materialized from the shared KV cache and the
|
||||
current valid suffix is spliced from ``out_cache_loc``. Draft/NextN uses the
|
||||
same CP shared-KV layout contract as the target model. The stale TAI
|
||||
sparse-current-page corruption that made this unsafe is blocked at the
|
||||
current-slot fill capability check before any TAI result is used.
|
||||
"""
|
||||
|
||||
if not cp_shared_kv_current_reuse_enabled():
|
||||
@@ -1091,25 +1211,40 @@ def should_reuse_current_extend_kv(forward_batch) -> bool:
|
||||
return True
|
||||
|
||||
partial_current = can_reuse_current_extend_kv(forward_batch)
|
||||
if partial_current and cp_shared_kv_is_draft_input(forward_batch):
|
||||
prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
|
||||
extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
|
||||
_log_current_reuse_fallback(
|
||||
"draft_partial_current_reuse_disabled",
|
||||
"cache-hit EAGLE/NextN draft uses full materialize instead of "
|
||||
"partial current reuse. prefix_lens=%s extend_lens=%s",
|
||||
[int(x) for x in prefix_lens_cpu]
|
||||
if prefix_lens_cpu is not None
|
||||
else None,
|
||||
[int(x) for x in extend_lens_cpu]
|
||||
if extend_lens_cpu is not None
|
||||
else None,
|
||||
)
|
||||
return False
|
||||
|
||||
return current_only or partial_current
|
||||
|
||||
|
||||
def current_extend_kv_rows_for_reuse(
|
||||
forward_batch,
|
||||
*current_kv_tensors: torch.Tensor | None,
|
||||
) -> int | None:
|
||||
"""Return valid current rows if current KV tensors can be spliced.
|
||||
|
||||
``out_cache_loc`` may be page padded while MLA/index projections only produce
|
||||
valid-token rows. Partial-current reuse should therefore validate tensor
|
||||
rows against ``extend_seq_lens_cpu`` rather than the padded loc tensor length.
|
||||
"""
|
||||
|
||||
if not should_reuse_current_extend_kv(forward_batch):
|
||||
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:
|
||||
return None
|
||||
valid_current_rows = int(extend_seq_lens_cpu[0])
|
||||
if valid_current_rows <= 0:
|
||||
return None
|
||||
|
||||
out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
|
||||
if out_cache_loc is None or int(out_cache_loc.numel()) < valid_current_rows:
|
||||
return None
|
||||
|
||||
for tensor in current_kv_tensors:
|
||||
if tensor is None or int(tensor.shape[0]) < valid_current_rows:
|
||||
return None
|
||||
return valid_current_rows
|
||||
|
||||
|
||||
def current_loc_remap_fast_path_args(
|
||||
forward_batch,
|
||||
) -> tuple[int | None, int | None]:
|
||||
@@ -1366,7 +1501,9 @@ def _debug_assert_no_negative_tensor_values(
|
||||
)
|
||||
|
||||
|
||||
def build_dense_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
def build_dense_page_remap(
|
||||
logical_pages: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Build a dense per-call page remap for shared-KV runtime materialization."""
|
||||
dense_pages = logical_pages.clone()
|
||||
positive_mask = logical_pages > 0
|
||||
@@ -1389,17 +1526,25 @@ def remap_logical_pages_to_dense_pages(
|
||||
insert_positions = torch.searchsorted(unique_logical_pages, positive_pages)
|
||||
if cp_shared_kv_debug_enabled() and insert_positions.numel() > 0:
|
||||
if unique_logical_pages.numel() == 0:
|
||||
raise ValueError("unique_logical_pages is empty but logical_pages contains data")
|
||||
raise ValueError(
|
||||
"unique_logical_pages is empty but logical_pages contains data"
|
||||
)
|
||||
if torch.any(insert_positions >= unique_logical_pages.numel()):
|
||||
raise ValueError("logical_pages contains entries outside unique_logical_pages")
|
||||
raise ValueError(
|
||||
"logical_pages contains entries outside unique_logical_pages"
|
||||
)
|
||||
if not torch.equal(unique_logical_pages[insert_positions], positive_pages):
|
||||
raise ValueError("logical_pages contains entries outside unique_logical_pages")
|
||||
raise ValueError(
|
||||
"logical_pages contains entries outside unique_logical_pages"
|
||||
)
|
||||
|
||||
dense_pages[positive_mask] = insert_positions.to(dense_pages.dtype) + 1
|
||||
return dense_pages
|
||||
|
||||
|
||||
def build_slot_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
def build_slot_page_remap(
|
||||
logical_pages: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Build a fixed-shape page remap without dynamic CUDA output ops.
|
||||
|
||||
The compact remap path uses boolean compaction + unique/searchsorted. Those
|
||||
@@ -1618,9 +1763,7 @@ def build_shared_paged_buffer_slot_remap(
|
||||
physical_page_capacity=page_buffer.shape[0],
|
||||
)
|
||||
slot_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
|
||||
logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (
|
||||
layout.cp_size
|
||||
) + 1
|
||||
logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * (layout.cp_size) + 1
|
||||
page_inverse = build_slot_page_inverse_optimized(
|
||||
slot_logical_pages,
|
||||
logical_page_capacity=logical_page_capacity,
|
||||
@@ -1803,9 +1946,7 @@ def build_current_loc_remap(
|
||||
safe_query_locs = torch.where(
|
||||
valid_query, query_flat_long, torch.zeros_like(query_flat_long)
|
||||
)
|
||||
query_pages = torch.div(
|
||||
safe_query_locs, page_size, rounding_mode="floor"
|
||||
)
|
||||
query_pages = torch.div(safe_query_locs, page_size, rounding_mode="floor")
|
||||
query_offsets = torch.remainder(safe_query_locs, page_size)
|
||||
query_pages_in_range = query_pages < logical_page_capacity
|
||||
safe_query_pages = torch.clamp(query_pages, max=logical_page_capacity - 1)
|
||||
@@ -1822,7 +1963,9 @@ def build_current_loc_remap(
|
||||
row_values.to(compact_row_ids.dtype),
|
||||
torch.full_like(compact_row_ids.reshape(-1), -1),
|
||||
)
|
||||
return matched.reshape(query_locs.shape), compact_flat.reshape(query_locs.shape)
|
||||
return matched.reshape(query_locs.shape), compact_flat.reshape(
|
||||
query_locs.shape
|
||||
)
|
||||
finally:
|
||||
if cp_shared_kv_sort_nvtx_enabled():
|
||||
torch.cuda.nvtx.range_pop()
|
||||
@@ -1975,7 +2118,9 @@ def materialize_local_token_kv_pages(
|
||||
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
|
||||
torch.long
|
||||
)
|
||||
dense_page_ids = torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
|
||||
dense_page_ids = (
|
||||
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
|
||||
)
|
||||
page_offsets = torch.arange(page_size, device=kv_cache.device, dtype=torch.long)
|
||||
src_tokens = (owned_physical_pages[:, None] * page_size + page_offsets).reshape(-1)
|
||||
dst_tokens = (dense_page_ids[:, None] * page_size + page_offsets).reshape(-1)
|
||||
@@ -2236,9 +2381,9 @@ def token_page_copy_debug_checksum(
|
||||
if not torch.any(owned_mask):
|
||||
return "owned_pages=0"
|
||||
owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64)
|
||||
owned_physical_pages = layout.logical_pages_to_physical(
|
||||
owned_logical_pages
|
||||
).to(torch.long)
|
||||
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
|
||||
torch.long
|
||||
)
|
||||
dense_page_ids = (
|
||||
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
|
||||
)
|
||||
@@ -2270,7 +2415,9 @@ def materialize_local_paged_buffer(
|
||||
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
|
||||
torch.long
|
||||
)
|
||||
dense_page_ids = torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
|
||||
dense_page_ids = (
|
||||
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
|
||||
)
|
||||
dense_page_buffer[dense_page_ids] = page_buffer[owned_physical_pages]
|
||||
return dense_page_buffer
|
||||
|
||||
@@ -2369,9 +2516,9 @@ def paged_copy_debug_checksum(
|
||||
if not torch.any(owned_mask):
|
||||
return "owned_pages=0"
|
||||
owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64)
|
||||
owned_physical_pages = layout.logical_pages_to_physical(
|
||||
owned_logical_pages
|
||||
).to(torch.long)
|
||||
owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to(
|
||||
torch.long
|
||||
)
|
||||
dense_page_ids = (
|
||||
torch.nonzero(owned_mask, as_tuple=False).flatten().to(torch.long) + 1
|
||||
)
|
||||
|
||||
@@ -23,13 +23,13 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
cp_shared_kv_mla_prefetch_should_log_layer,
|
||||
cp_shared_kv_is_draft_input,
|
||||
cp_shared_kv_should_prefetch_next_layer,
|
||||
current_extend_kv_rows_for_reuse,
|
||||
current_loc_remap_fast_path_args,
|
||||
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,
|
||||
tensor_debug_summary,
|
||||
)
|
||||
@@ -1747,13 +1747,12 @@ class NativeSparseAttnBackend(
|
||||
mla_prefetcher = getattr(
|
||||
forward_batch, "cp_shared_kv_mla_prefetcher", None
|
||||
)
|
||||
can_reuse_current_kv = (
|
||||
should_reuse_current_extend_kv(forward_batch)
|
||||
and k is not None
|
||||
and k_rope is not None
|
||||
and k.shape[0] == forward_batch.out_cache_loc.numel()
|
||||
and k_rope.shape[0] == forward_batch.out_cache_loc.numel()
|
||||
current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse(
|
||||
forward_batch,
|
||||
k,
|
||||
k_rope,
|
||||
)
|
||||
can_reuse_current_kv = current_kv_rows_for_reuse is not None
|
||||
if cp_shared_kv_mla_prefetch_log_enabled():
|
||||
if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id):
|
||||
prefix_lens_cpu = getattr(
|
||||
@@ -1788,8 +1787,15 @@ class NativeSparseAttnBackend(
|
||||
else None,
|
||||
)
|
||||
if can_reuse_current_kv:
|
||||
current_kv_cache = _cat([k, k_rope], dim=-1)
|
||||
current_locs_for_reuse = forward_batch.out_cache_loc
|
||||
assert k is not None and k_rope is not None
|
||||
assert current_kv_rows_for_reuse is not None
|
||||
valid_current_rows = int(current_kv_rows_for_reuse)
|
||||
current_kv_cache = _cat(
|
||||
[k[:valid_current_rows], k_rope[:valid_current_rows]], dim=-1
|
||||
)
|
||||
current_locs_for_reuse = forward_batch.out_cache_loc[
|
||||
:valid_current_rows
|
||||
]
|
||||
logical_page_table_1 = page_table_1
|
||||
current_remap_page_size, current_remap_logical_page_capacity = (
|
||||
current_loc_remap_fast_path_args(forward_batch)
|
||||
|
||||
@@ -306,7 +306,7 @@ def evict_from_tree_cache(
|
||||
# Standard allocator
|
||||
available = allocator.available_size()
|
||||
if available < num_tokens:
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[MemCache-evict] evict_from_tree_cache: available=%d < num_tokens=%d deficit=%d, triggering eviction",
|
||||
available,
|
||||
num_tokens,
|
||||
@@ -343,7 +343,7 @@ def _evict_for_compute_owner_lanes(
|
||||
if isinstance(evictable_size, tuple):
|
||||
evictable_size = evictable_size[0]
|
||||
if evictable_size <= 0:
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size=%d <= 0, giving up",
|
||||
evictable_size,
|
||||
)
|
||||
@@ -355,7 +355,7 @@ def _evict_for_compute_owner_lanes(
|
||||
# load-back pressure.
|
||||
evict_tokens = max(allocator.page_size, deficit_pages * allocator.page_size)
|
||||
before_available = allocator.available_size()
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d: deficit_pages=%d evict_tokens=%d before_available=%d evictable_size=%d",
|
||||
attempt,
|
||||
deficit_pages,
|
||||
@@ -371,7 +371,7 @@ def _evict_for_compute_owner_lanes(
|
||||
)
|
||||
after_available = allocator.available_size()
|
||||
evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result: evicted=%d after_available=%d",
|
||||
attempt,
|
||||
evicted_tokens,
|
||||
|
||||
@@ -1220,7 +1220,7 @@ class HiRadixCache(RadixCache):
|
||||
num_evicted += self._evict_backuped(victim)
|
||||
|
||||
refreshed = self._refresh_cp_load_back_plan(plan)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-load] owner-lane device eviction before CP load-back: "
|
||||
"node_id=%d victims=%s num_evicted=%d required_by_owner=%s "
|
||||
"before_available_by_owner=%s before_deficit_by_owner=%s "
|
||||
@@ -1254,7 +1254,7 @@ class HiRadixCache(RadixCache):
|
||||
)
|
||||
eviction_plan = self._plan_cp_load_back_owner_lane_evictions(plan)
|
||||
if len(eviction_plan.victims) == 0:
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] owner-lane evict found no contributing victims: "
|
||||
"num_tokens=%d deficits=%s evictable_size=%d available_size=%d",
|
||||
params.num_tokens,
|
||||
@@ -1276,7 +1276,7 @@ class HiRadixCache(RadixCache):
|
||||
self._clear_pin(victim)
|
||||
|
||||
if not self._cp_device_leaf_is_load_back_victim(victim):
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] owner-lane evict victim no longer evictable: "
|
||||
"victim_id=%s deficits=%s",
|
||||
getattr(victim, "id", None),
|
||||
@@ -1319,7 +1319,7 @@ class HiRadixCache(RadixCache):
|
||||
if self._node_backuped(victim):
|
||||
num_evicted += self._evict_backuped(victim)
|
||||
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] owner-lane evict END: requested_tokens=%d "
|
||||
"deficits=%s victims=%s planned_freed_by_owner=%s "
|
||||
"remaining_deficit_by_owner=%s num_evicted=%d "
|
||||
@@ -1925,7 +1925,7 @@ class HiRadixCache(RadixCache):
|
||||
victim.host_value = None
|
||||
self._remove_host_leaf(victim)
|
||||
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] deterministic CP host eviction before write: "
|
||||
"node_id=%d phase=%s victims=%s local_freed=%d planned_freed=%s",
|
||||
node_id,
|
||||
@@ -1978,7 +1978,7 @@ class HiRadixCache(RadixCache):
|
||||
)
|
||||
)
|
||||
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-write] write_backup CP retry after deterministic host eviction: "
|
||||
"node_id=%d deficit_by_owner=%s",
|
||||
node_id,
|
||||
@@ -1991,7 +1991,7 @@ class HiRadixCache(RadixCache):
|
||||
if not isinstance(result, HiCacheWriteFailure):
|
||||
return result
|
||||
|
||||
logger.info(
|
||||
logger.warning(
|
||||
"[HiCache-write] write_backup CP FAILED after deterministic retry: "
|
||||
"node_id=%d len=%d needed_slots=%d",
|
||||
node_id,
|
||||
@@ -2800,7 +2800,7 @@ class HiRadixCache(RadixCache):
|
||||
]
|
||||
heapq.heapify(eviction_heap)
|
||||
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] evict START: num_tokens=%d heap_size=%d evictable_size=%d available_size=%d",
|
||||
num_tokens,
|
||||
len(eviction_heap),
|
||||
@@ -2864,7 +2864,7 @@ class HiRadixCache(RadixCache):
|
||||
self._evict_backuped(node)
|
||||
|
||||
self.update_eviction_metrics(num_evicted, start_time)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] evict END: num_tokens=%d num_evicted=%d num_locked_skipped=%d evictable_size_after=%d available_size_after=%d",
|
||||
num_tokens,
|
||||
num_evicted,
|
||||
@@ -2880,7 +2880,7 @@ class HiRadixCache(RadixCache):
|
||||
freed_len = self.cache_controller.evict_device(node.value)
|
||||
assert freed_len > 0
|
||||
self.evictable_size_ -= device_resident_len
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] _evict_backuped: node_id=%d num_evicted=%d physical_tokens=%d lock_ref=%d backed=%s",
|
||||
node.id,
|
||||
freed_len,
|
||||
@@ -2899,7 +2899,7 @@ class HiRadixCache(RadixCache):
|
||||
def _evict_regular(self, node: TreeNode):
|
||||
# evict a node not initiated write to host -- emit BlockRemoved
|
||||
num_evicted = self._node_device_resident_len(node)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] _evict_regular: node_id=%d num_evicted=%d",
|
||||
node.id,
|
||||
num_evicted,
|
||||
@@ -2943,7 +2943,7 @@ class HiRadixCache(RadixCache):
|
||||
return 0
|
||||
|
||||
leaves = list(self.evictable_host_leaves)
|
||||
logger.info(
|
||||
logger.debug(
|
||||
"[HiCache-evict] _evict_host_for_physical_slots: required_slots=%d sync=%s leaves=%d",
|
||||
required_host_slots,
|
||||
synchronize_across_ranks,
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
import inspect
|
||||
import re
|
||||
import sys
|
||||
import types
|
||||
import unittest
|
||||
@@ -97,6 +99,8 @@ from sglang.srt.managers.cache_controller import (
|
||||
HiCacheWriteFailure,
|
||||
HiCacheWriteReservation,
|
||||
)
|
||||
import sglang.srt.mem_cache.common as mem_cache_common
|
||||
import sglang.srt.mem_cache.hiradix_cache as hiradix_cache
|
||||
from sglang.srt.mem_cache.base_prefix_cache import (
|
||||
EvictParams,
|
||||
InsertParams,
|
||||
@@ -117,6 +121,57 @@ from sglang.test.test_utils import CustomTestCase
|
||||
register_cpu_ci(est_time=2, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class TestHiCacheEvictLoggingLevels(CustomTestCase):
|
||||
def assert_debug_marker(self, source: str, marker: str):
|
||||
self.assertIn(marker, source)
|
||||
self.assertRegex(
|
||||
source,
|
||||
r"logger\.debug\(\s*\n\s*\"" + re.escape(marker),
|
||||
msg=f"{marker} should be debug-only on the success/no-op hot path",
|
||||
)
|
||||
self.assertNotRegex(
|
||||
source,
|
||||
r"logger\.info\(\s*\n\s*\"" + re.escape(marker),
|
||||
msg=f"{marker} must not stay at INFO on the success/no-op hot path",
|
||||
)
|
||||
|
||||
def test_evict_hot_path_success_logs_are_debug_only(self):
|
||||
common_source = inspect.getsource(mem_cache_common)
|
||||
hiradix_source = inspect.getsource(hiradix_cache)
|
||||
|
||||
for marker in (
|
||||
"[MemCache-evict] evict_from_tree_cache:",
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes: evictable_size",
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d:",
|
||||
"[MemCache-evict] _evict_for_compute_owner_lanes attempt=%d result:",
|
||||
):
|
||||
self.assert_debug_marker(common_source, marker)
|
||||
|
||||
for marker in (
|
||||
"[HiCache-load] owner-lane device eviction before CP load-back: ",
|
||||
"[HiCache-evict] owner-lane evict found no contributing victims: ",
|
||||
"[HiCache-evict] owner-lane evict victim no longer evictable: ",
|
||||
"[HiCache-evict] owner-lane evict END: requested_tokens=%d ",
|
||||
"[HiCache-evict] deterministic CP host eviction before write: ",
|
||||
"[HiCache-write] write_backup CP retry after deterministic host eviction: ",
|
||||
"[HiCache-evict] evict START:",
|
||||
"[HiCache-evict] evict END:",
|
||||
"[HiCache-evict] _evict_backuped:",
|
||||
"[HiCache-evict] _evict_regular:",
|
||||
"[HiCache-evict] _evict_host_for_physical_slots:",
|
||||
):
|
||||
self.assert_debug_marker(hiradix_source, marker)
|
||||
|
||||
self.assertRegex(
|
||||
hiradix_source,
|
||||
r"logger\.warning\(\s*\n\s*\"\[HiCache-write\] write_backup CP FAILED after deterministic retry:",
|
||||
)
|
||||
self.assertNotRegex(
|
||||
hiradix_source,
|
||||
r"logger\.info\(\s*\n\s*\"\[HiCache-write\] write_backup CP FAILED after deterministic retry:",
|
||||
)
|
||||
|
||||
|
||||
class TestCpHiCacheImports(CustomTestCase):
|
||||
def test_cp_hicache_public_imports_without_sgl_kernel(self):
|
||||
import subprocess
|
||||
|
||||
@@ -54,9 +54,7 @@ for _name in ("flash_attn_varlen_func", "flash_attn_with_kvcache"):
|
||||
if not hasattr(flash_attn_stub, _name):
|
||||
setattr(flash_attn_stub, _name, lambda *args, **kwargs: None)
|
||||
|
||||
sgl_kernel_stub = sys.modules.setdefault(
|
||||
"sgl_kernel", types.ModuleType("sgl_kernel")
|
||||
)
|
||||
sgl_kernel_stub = sys.modules.setdefault("sgl_kernel", types.ModuleType("sgl_kernel"))
|
||||
if not hasattr(sgl_kernel_stub, "__path__"):
|
||||
sgl_kernel_stub.__path__ = []
|
||||
if not hasattr(sgl_kernel_stub, "flash_attn"):
|
||||
@@ -361,9 +359,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"sglang.srt.layers.attention.nsa.index_buf_accessor": index_accessor_stub
|
||||
},
|
||||
{"sglang.srt.layers.attention.nsa.index_buf_accessor": index_accessor_stub},
|
||||
):
|
||||
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
|
||||
|
||||
@@ -457,7 +453,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
self.assertIs(out, buffer)
|
||||
dist_all_reduce.assert_called_once()
|
||||
self.assertIs(dist_all_reduce.call_args.args[0], buffer)
|
||||
self.assertIs(dist_all_reduce.call_args.kwargs["group"], dummy_group.device_group)
|
||||
self.assertIs(
|
||||
dist_all_reduce.call_args.kwargs["group"], dummy_group.device_group
|
||||
)
|
||||
|
||||
def test_build_dense_page_remap_preserves_sentinels(self):
|
||||
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
|
||||
@@ -492,9 +490,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
)
|
||||
|
||||
current_locs = torch.tensor([100, 64, 256, 128], dtype=torch.int64)
|
||||
query_locs = torch.tensor(
|
||||
[[128, -1, 64], [512, 100, 256]], dtype=torch.int32
|
||||
)
|
||||
query_locs = torch.tensor([[128, -1, 64], [512, 100, 256]], dtype=torch.int32)
|
||||
|
||||
is_current, compact_rows = build_current_loc_remap(query_locs, current_locs)
|
||||
|
||||
@@ -569,7 +565,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
forward_batch.out_cache_loc = torch.arange(64, dtype=torch.int64)
|
||||
self.assertFalse(can_reuse_current_extend_kv(forward_batch))
|
||||
|
||||
def test_should_reuse_current_extend_kv_disables_draft_partial_cache_hit_suffix(
|
||||
def test_should_reuse_current_extend_kv_enables_draft_partial_cache_hit_suffix(
|
||||
self,
|
||||
):
|
||||
from sglang.srt.environ import envs
|
||||
@@ -595,13 +591,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
)
|
||||
|
||||
with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
|
||||
with self.assertLogs(runtime.logger.name, level="WARNING") as logs:
|
||||
self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch))
|
||||
self.assertIn(
|
||||
"[CP_SHARED_KV_FALLBACK][current_reuse]",
|
||||
logs.output[0],
|
||||
)
|
||||
self.assertIn("draft_partial_current_reuse_disabled", logs.output[0])
|
||||
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
|
||||
|
||||
forward_batch.spec_info = TargetSpecInfo()
|
||||
forward_batch.cp_shared_kv_mla_prefetcher = object()
|
||||
@@ -619,6 +609,73 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
forward_batch.cp_shared_kv_mla_prefetcher = None
|
||||
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
|
||||
|
||||
def test_current_extend_kv_rows_for_reuse_accepts_padded_out_cache_loc(self):
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
class DraftSpecInfo:
|
||||
def is_draft_input(self):
|
||||
return True
|
||||
|
||||
forward_batch = SimpleNamespace(
|
||||
forward_mode=_FakeExtendForwardMode(),
|
||||
batch_size=1,
|
||||
extend_prefix_lens_cpu=[40384],
|
||||
extend_seq_lens_cpu=[65],
|
||||
seq_lens_cpu=torch.tensor([40384 + 65], dtype=torch.int32),
|
||||
out_cache_loc=torch.arange(128, dtype=torch.int64),
|
||||
spec_info=DraftSpecInfo(),
|
||||
)
|
||||
k = torch.empty((65, 2, 4), dtype=torch.float32)
|
||||
k_rope = torch.empty((65, 2, 1), dtype=torch.float32)
|
||||
|
||||
with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True):
|
||||
self.assertEqual(
|
||||
runtime.current_extend_kv_rows_for_reuse(forward_batch, k, k_rope),
|
||||
65,
|
||||
)
|
||||
|
||||
self.assertIsNone(
|
||||
runtime.current_extend_kv_rows_for_reuse(
|
||||
forward_batch,
|
||||
k[:64],
|
||||
k_rope,
|
||||
)
|
||||
)
|
||||
|
||||
def test_mla_current_reuse_gate_accepts_padded_out_cache_loc(self):
|
||||
from pathlib import Path
|
||||
|
||||
source = (
|
||||
Path(__file__).resolve().parents[4]
|
||||
/ "python/sglang/srt/layers/attention/nsa_backend.py"
|
||||
).read_text()
|
||||
start = source.index(" current_kv_rows_for_reuse =")
|
||||
end = source.index(
|
||||
" if cp_shared_kv_mla_prefetch_log_enabled()", start
|
||||
)
|
||||
gate_source = source[start:end]
|
||||
|
||||
self.assertIn("current_extend_kv_rows_for_reuse", gate_source)
|
||||
self.assertNotIn(
|
||||
"k.shape[0] == forward_batch.out_cache_loc.numel()",
|
||||
gate_source,
|
||||
)
|
||||
self.assertNotIn(
|
||||
"k_rope.shape[0] == forward_batch.out_cache_loc.numel()",
|
||||
gate_source,
|
||||
)
|
||||
body_start = source.index(" if can_reuse_current_kv:", end)
|
||||
body_end = source.index(
|
||||
" logical_page_table_1 = page_table_1", body_start
|
||||
)
|
||||
body_source = "".join(source[body_start:body_end].split())
|
||||
|
||||
self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source)
|
||||
self.assertIn("k[:valid_current_rows]", body_source)
|
||||
self.assertIn("k_rope[:valid_current_rows]", body_source)
|
||||
self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
|
||||
|
||||
def test_runtime_fallback_helpers_use_standard_warning_marker(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
@@ -738,6 +795,53 @@ 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_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails(self):
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
class BadKernels:
|
||||
@staticmethod
|
||||
def fill_current_token_kv_page_slots_and_remap_locs(*args, **kwargs):
|
||||
raise AssertionError("stale TAI fill kernel should not be called")
|
||||
|
||||
with envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.override(True), patch.object(
|
||||
runtime,
|
||||
"_tai_current_slot_fill_supports_sparse_pages",
|
||||
return_value=False,
|
||||
), patch.object(
|
||||
runtime,
|
||||
"_load_tai_materialize_kernels",
|
||||
return_value=BadKernels,
|
||||
):
|
||||
result = runtime._try_tai_fill_current_kv_page_slots_and_remap_locs(
|
||||
dense_kv_cache=torch.zeros((16, 1), dtype=torch.float32),
|
||||
materialized_dense_locs=torch.tensor([[4, 5, 8, 9]], dtype=torch.int64),
|
||||
current_kv_cache=torch.ones((2, 1), dtype=torch.float32),
|
||||
logical_locs=torch.tensor([[20, 21, 40, 41]], dtype=torch.int64),
|
||||
current_locs=torch.tensor([20, 21], dtype=torch.int64),
|
||||
page_inverse=torch.tensor(
|
||||
[0, -1, -1, -1, -1, 1, -1, -1, -1, -1, 2],
|
||||
dtype=torch.long,
|
||||
),
|
||||
page_size=4,
|
||||
mask_non_current_in_current_pages=True,
|
||||
)
|
||||
|
||||
self.assertIsNone(result)
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required")
|
||||
def test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel(
|
||||
self,
|
||||
):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
|
||||
runtime._tai_current_slot_fill_sparse_pages_self_test.cache_clear()
|
||||
self.assertTrue(
|
||||
runtime._tai_current_slot_fill_supports_sparse_pages(
|
||||
torch.device("cuda", torch.cuda.current_device())
|
||||
)
|
||||
)
|
||||
|
||||
def test_materialize_prefix_and_reuse_current_kv_page_slots_without_prefetcher(
|
||||
self,
|
||||
):
|
||||
@@ -913,9 +1017,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
with patch.object(
|
||||
prefetch.torch.cuda, "current_stream", return_value=current_stream
|
||||
), patch.object(
|
||||
prefetcher, "launch_pending_reduce"
|
||||
) as launch_pending_reduce:
|
||||
), patch.object(prefetcher, "launch_pending_reduce") as launch_pending_reduce:
|
||||
prefetcher.wait_attention_window()
|
||||
|
||||
launch_pending_reduce.assert_not_called()
|
||||
@@ -1402,8 +1504,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
self.assertTrue(
|
||||
any(
|
||||
"token slot remap cache not reused (missing_cached_value)"
|
||||
in message
|
||||
"token slot remap cache not reused (missing_cached_value)" in message
|
||||
for message in logs.output
|
||||
)
|
||||
)
|
||||
@@ -1549,13 +1650,17 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
)
|
||||
self.assertIn("prefix_not_page_aligned", logger.call_args.args[1])
|
||||
|
||||
def test_mla_prefetch_min_prefix_pages_uses_cached_token_default_and_can_override(self):
|
||||
def test_mla_prefetch_min_prefix_pages_uses_cached_token_default_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_PREFIX_PAGES.clear()
|
||||
default_tokens = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_TOKENS.get()
|
||||
self.assertEqual(runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, default_tokens)
|
||||
self.assertEqual(
|
||||
runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, default_tokens
|
||||
)
|
||||
expected_pages = (default_tokens + 63) // 64
|
||||
self.assertEqual(
|
||||
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64),
|
||||
@@ -1844,9 +1949,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_tai_fused_mla_store_enabled", return_value=False
|
||||
), patch.object(
|
||||
runtime, "_load_tai_fused_mla_store_kernel"
|
||||
) as load_kernel:
|
||||
), patch.object(runtime, "_load_tai_fused_mla_store_kernel") as load_kernel:
|
||||
used = runtime.try_tai_fused_mla_store(
|
||||
token_to_kv_pool=FakePool(),
|
||||
layer=SimpleNamespace(layer_id=0),
|
||||
@@ -1962,7 +2065,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
self.inverse_calls = []
|
||||
self.remap_calls = []
|
||||
|
||||
def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
|
||||
def build_slot_page_inverse(
|
||||
self, slot_logical_pages, logical_page_capacity
|
||||
):
|
||||
self.inverse_calls.append((slot_logical_pages, logical_page_capacity))
|
||||
return torch.tensor([0, 1, 2, -1], dtype=torch.long)
|
||||
|
||||
@@ -2056,7 +2161,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=True
|
||||
), patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV materialize got logical token locs outside the physical pool",
|
||||
@@ -2068,7 +2175,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
def test_materialize_token_kv_skips_physical_pool_validation_when_debug_disabled(self):
|
||||
def test_materialize_token_kv_skips_physical_pool_validation_when_debug_disabled(
|
||||
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
|
||||
|
||||
@@ -2081,7 +2190,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=False
|
||||
), patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
):
|
||||
_, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
@@ -2100,7 +2211,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
logical_locs = torch.tensor([4, -1, 8], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "cp_shared_kv_debug_enabled", return_value=True):
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
|
||||
with patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
):
|
||||
_, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
@@ -2172,7 +2285,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
logical_locs = torch.tensor([8, 20, -1], dtype=torch.int64)
|
||||
remap_source_locs = torch.tensor([4, 8, 12, 16, 20], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
|
||||
with patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
):
|
||||
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
@@ -2223,7 +2338,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1)
|
||||
remap_source_locs = torch.tensor([4, 8, 12, 16, 20], dtype=torch.int64)
|
||||
|
||||
with patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
|
||||
with patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
):
|
||||
dense_kv_a, dense_locs_a = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=torch.tensor([4, 20], dtype=torch.int64),
|
||||
@@ -2254,7 +2371,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=True
|
||||
), patch.object(runtime, "_all_reduce_materialized_buffer", _identity_all_reduce):
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", _identity_all_reduce
|
||||
):
|
||||
with self.assertRaisesRegex(
|
||||
RuntimeError,
|
||||
"CP shared KV materialize got logical pages outside the physical page buffer",
|
||||
@@ -2401,7 +2520,9 @@ class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
|
||||
self.assertEqual(k_to_write.shape[0], 2)
|
||||
self.assertEqual(k_rope_to_write.shape[0], 2)
|
||||
|
||||
def test_index_write_filter_does_not_build_debug_summaries_when_debug_disabled(self):
|
||||
def test_index_write_filter_does_not_build_debug_summaries_when_debug_disabled(
|
||||
self,
|
||||
):
|
||||
from sglang.srt.layers.attention.nsa import nsa_indexer
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
@@ -2422,10 +2543,12 @@ class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
|
||||
"sglang.srt.environ.envs.SGLANG_DEBUG_CP_SHARED_KV.get",
|
||||
return_value=False,
|
||||
):
|
||||
physical_locs, key_to_write = nsa_indexer.Indexer._filter_shared_index_write(
|
||||
None,
|
||||
forward_batch,
|
||||
key,
|
||||
physical_locs, key_to_write = (
|
||||
nsa_indexer.Indexer._filter_shared_index_write(
|
||||
None,
|
||||
forward_batch,
|
||||
key,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(physical_locs.tolist(), [4, 8])
|
||||
@@ -2493,8 +2616,12 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
|
||||
self.remap_calls = []
|
||||
self.token_calls = []
|
||||
|
||||
def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
|
||||
self.page_inverse_calls.append((slot_logical_pages, logical_page_capacity))
|
||||
def build_slot_page_inverse(
|
||||
self, slot_logical_pages, logical_page_capacity
|
||||
):
|
||||
self.page_inverse_calls.append(
|
||||
(slot_logical_pages, logical_page_capacity)
|
||||
)
|
||||
return torch.tensor([0, 1, 2, -1, 3], dtype=torch.long)
|
||||
|
||||
def remap_logical_locs_to_slot_dense_locs(
|
||||
@@ -2567,7 +2694,9 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
class FakeTaiKernels:
|
||||
def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
|
||||
def build_slot_page_inverse(
|
||||
self, slot_logical_pages, logical_page_capacity
|
||||
):
|
||||
return torch.tensor([0, 1, 2, -1, 3], dtype=torch.long)
|
||||
|
||||
def remap_logical_locs_to_slot_dense_locs(
|
||||
|
||||
Reference in New Issue
Block a user