diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index cb7a4de5d..d15d287c6 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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: > 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. diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 8810a854a..7eb21b872 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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 ) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index e972ad75d..9fa434007 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index b5bb8d6ca..fa0cbe4a3 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index b1b855d08..56b85f5c4 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index f5b0dfe23..6dc364f24 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 7720389ed..3ecaa2ffb 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -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(