diff --git a/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md b/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md index f0fd9817d..d4786df47 100644 --- a/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md +++ b/docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md @@ -326,7 +326,7 @@ Minimum tests before enabling true batch > 1: - two cache-hit requests with different prefix lengths; - page-aligned and non-page-aligned suffixes; -- short radix-hit suffix with replicated compute path; +- short radix-hit suffix with zero-owner-page segments; - HiCache host hit + load-back; - draft/MTP enabled with target/draft KV mirrored; - compare batch > 1 output against sequential batch-size-1 execution; 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 new file mode 100644 index 000000000..3e36b4b0a --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -0,0 +1,847 @@ +# NSA Prefill CP Page-Aligned Cache Contract + +This note records the contract we need before continuing MLA/index partial +reuse under CP shared KV + HiCache + EAGLE/MTP. + +The immediate trigger was a production hang where the prefill process stopped +after target-layer MLA partial current reuse: + +```text +/mnt/beegfs/cjy/sglang_cp_hicache_20260528_161408.log + +layer=2 +prefix_lens=[13184] +extend_lens=[37024] +page_size=64 +prefix_pages=206 +suffix_slots=579 +current_rows=37024 +kv_rows=87328 +used_prefetch=True +``` + +`suffix_slots * page_size = 579 * 64 = 37056`, but only `37024` current +rows were appended. The fast path consumed prefetched prefix pages and appended +only real current rows, so the last suffix page had `32` rows of page slack with +no explicit valid/invalid contract. + +The target design is **not** to fall back for these short or unaligned suffixes. +Instead, all physical cache paths should share one page-aligned contract while +logical token semantics stay unpadded. + +## Contract + +Use two lengths everywhere the path crosses a logical/physical boundary: + +```text +valid_tokens = real request/cache tokens +padded_pages = ceil(valid_tokens / page_size) +padded_tokens = padded_pages * page_size +padding_tokens = padded_tokens - valid_tokens +``` + +### Logical layer: always valid-token based + +The following must continue to use `valid_tokens`: + +- request token length, +- `seq_lens_cpu`, +- `extend_seq_lens_cpu`, +- radix tree key/logical node length, +- cache-hit logical length exposed to scheduler, +- decode position and max-context checks, +- EAGLE/MTP verification token count, +- attention-visible sequence length. + +Padding tokens must not become request tokens. + +### Physical layer: always page-aligned + +The following should use `padded_pages` / `padded_tokens`: + +- GPU KV pool allocation and owner-lane capacity, +- draft KV mirror allocation, +- HiCache host reservation, +- backup/load/evict accounting, +- disaggregated target/draft transfer page lists, +- CP shared-KV MLA/index dense materialization buffers, +- CP owner-lane capacity planning and snapshots. + +Padding rows may exist in temporary dense buffers or physical pages, but they +must be invisible to attention and radix semantics. + +### Visibility rule + +Every remap path must preserve this invariant: + +```text +logical loc maps to a dense row only if it is a real valid token. +padding/slack locs map to -1. +``` + +Padding KV rows can be zero-filled. Correctness comes from masking/remapping +them out, not from their value. + +### Kernel ownership + +Production hot paths should use `tai-kernel` for page-padded compose/remap and +copy-style operations. Python/Torch implementations are acceptable as CPU unit +test references and explicit warning fallbacks only; they should not be the +normal CUDA runtime path. + +For MLA partial-current reuse, the intended runtime primitive is: + +```text +preallocated slot-dense KV pages ++ current suffix KV copied into its slot-dense rows ++ zero padding rows left invisible +-> same dense KV buffer + mixed loc remap + current mask +``` + +This should be one `tai_kernel.nsa_prefill.cp_shared_kv_materialize` primitive, +not a sequence of `torch.sort/searchsorted/cat/zeros` operations in the +attention hot path. If the TAI primitive is unavailable, SGLang must emit a +visible warning fallback rather than silently taking the slow path. + +The TAI primitive should extend the existing materialize/remap kernel family: + +- copy each valid `current_locs[i]` row into `page_inverse[page] * page_size + offset`; +- build a current-page inverse from page-grouped current locs; +- keep real current locs visible by exact loc check; +- map non-current rows inside current pages to `-1`. + +This avoids the previous compact-current append path where tail-page slack had +slot-dense locs but no valid current KV row. + +## Root-cause record: unsafe partial-current prefetch + +The current MLA partial-current prefetch path has a narrower contract than the +page-level metadata it consumes. + +Evidence: + +- `cp_shared_kv_prefetch.py::consume_prefix_with_current()` says the prefetched + prefix is already materialized while current/suffix pages are not copied from + the shared pool; it then appends `current_kv_cache` directly. +- `cp_shared_kv_runtime.py::merge_materialized_and_current_kv()` only remaps + entries present in `current_locs`; all non-current entries keep the existing + `materialized_dense_locs`. +- `nsa_backend.py` calls `consume_prefix_with_current()` whenever an MLA + prefetcher exists on a mixed prefix/current batch, before falling back to the + older partial materialize path. + +That is safe only when the page-level suffix represented by the page table is +exactly covered by `current_locs`. It is not safe when: + +```text +ceil(extend_len / page_size) * page_size > extend_len +``` + +The older non-prefetch partial path is safer because it materializes all +non-current logical locations and only substitutes rows that match +`current_locs`. + +The fix should not be a permanent fallback. The intended fix is to make the +prefetch compose path create a page-padded current suffix buffer and map suffix +slack to `-1`. + +## Existing evidence in the codebase + +### Page-aligned split already treats tail as a page unit + +`build_page_aligned_in_seq_split_list()` computes: + +```python +full_pages = extend_len // page_size +tail_tokens = extend_len % page_size +num_page_units = full_pages + (1 if tail_tokens > 0 else 0) +``` + +The Phase 4 doc explicitly states that the tail unit is not split across CP +segments. However, the current `split_list` still sums to `extend_len`, not +`padded_tokens`. That is correct for compute tokens, but it means later cache +paths must not assume `split_list` size equals physical page coverage. + +### HiCache metadata currently requires page-aligned logical_len + +`CpHiCacheNodeMetadata.__post_init__()` rejects `logical_len % page_size != 0`, +and `reserve_write_cp()` rejects non-page-aligned `device_indices`. + +This is already a physical-cache assumption. If logical radix length remains +valid-token based, the metadata needs a clearer name/split: + +```text +valid_len # radix-visible logical length +padded_len # physical cache/host metadata length +page_owners # length padded_len / page_size +owned_positions # local positions inside padded_len +``` + +Without this split, node split/load checks can accidentally force logical cache +semantics to be page-aligned. + +### Owner-lane accounting is already page-derived + +`HiRadixCache._cp_owner_token_counts()` counts `page_size` tokens per owner page. +The no-collective capacity plan already describes a per-owner page ledger. This +means the capacity side is structurally ready for page-based accounting, but all +inputs to the ledger must be padded-page based. A valid-token scalar is not +enough for CP owner-lane admission. + +`cp_shared_kv_compute_owner.build_in_seq_page_compute_owners()` should therefore +return owners for every physical page unit, including short chunks where some +zigzag segments own zero pages. Falling back only because `num_pages < 2 * +cp_size` breaks the page-owner lane contract and can force later materialize or +HiCache paths back to token-balanced behavior. + +This owner assignment remains page-count based; `out_cache_loc` remains +valid-token based. `alloc_extend_compute_owner()` uses the owner list to choose +whole physical pages, then `alloc_extend_naive()` returns exactly the valid +extend token locs. The padded tail rows exist because the selected logical page +exists, not because fake token locs are appended to `out_cache_loc`. + +### Load must replay page owners + +`HiCacheController.load_cp()` reconstructs device allocation via +`alloc_pages_with_owners(page_owners)`. This confirms that load-back is already +page-pattern based and should stay physical-page based. + +## Impact surface + +### 1. CP split and metadata + +Files: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + +Needed changes: + +- keep `split_list` and `actual_token_count` valid-token based; +- add explicit padded page metadata to `NSAContextParallelMetadata` or an + attached helper object; +- keep `segment_page_starts/ends` page-based, including tail page coverage; +- add a helper that returns `valid_tokens`, `padded_tokens`, `padded_pages`, and + `padding_tokens` for the current extend/cache-hit suffix. + +Risk: + +- changing `split_list` to padded tokens would leak fake tokens into compute and + logits paths. Do not do that. + +### 2. MLA partial current reuse + +Files: + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` +- `python/sglang/srt/layers/attention/nsa_backend.py` + +Needed changes: + +- use the existing page-slot dense buffer for mixed prefix/current reuse; +- copy valid current suffix rows into their slot-dense page rows with + `tai-kernel` in the CUDA path; +- keep padding rows in the tail suffix page physically present but zero/unused; +- remap valid current locs to their slot-dense rows; +- map suffix slack/future locs to `-1`; +- keep existing full/current-only path behavior unchanged. + +Short-term safe gate if this cannot be implemented atomically: + +- WARNING fallback to the older partial materialize path when suffix slack is + detected. This is a stopgap only; the target contract is page-padded compose, + not fallback. + +### 3. NSA index partial reuse + +Files: + +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` + +Needed changes: + +- apply the same valid/padded contract as MLA; +- index page metadata can be reused across layers, but returned dense page/table + remaps must not expose padding locs as valid; +- current index reuse should not assume that all locations in a suffix page are + current tokens. + +Risk: + +- index is page-first by nature, so it is easy to accidentally treat the whole + tail page as visible. Tests must include non-page-aligned suffixes. + +### 4. HiCache host metadata and radix split + +Files: + +- `python/sglang/srt/mem_cache/hiradix_cache.py` +- `python/sglang/srt/managers/cache_controller.py` + +Needed changes: + +- distinguish radix-visible `valid_len` from physical `padded_len`; +- store `page_owners` for `padded_len / page_size` pages; +- keep split rules page-aware: a radix split inside an in-flight or host-backed + padded page must either split at a page boundary or keep the child pending + until the physical operation is committed; +- make `_node_host_len()` / host-hit accounting return valid length to the + scheduler but use padded length for load/evict/host free. + +Risk: + +- if `host_hit_length` becomes padded, scheduler may believe fake tokens were + cached. If physical metadata uses valid length, host reservation/load can + under-allocate the tail page. + +### 5. Capacity and owner-lane admission + +Files: + +- `python/sglang/srt/mem_cache/hiradix_cache.py` +- `python/sglang/srt/mem_cache/allocator.py` +- `python/sglang/srt/managers/cache_controller.py` + +Needed changes: + +- all required capacity should be expressed as pages or per-owner padded tokens; +- owner-lane vectors should be derived from padded `page_owners`, not valid + token counts; +- target and draft reservation must use the same padded page pattern; +- no hot-path collective should be introduced to discover the mismatch. + +Risk: + +- if target uses padded capacity and draft uses valid capacity, draft can miss + the tail page while target succeeds, causing cache-hit accept-rate corruption. + +### 6. Disaggregated transfer + +Files: + +- `python/sglang/srt/disaggregation/prefill.py` +- `python/sglang/srt/disaggregation/decode.py` +- `python/sglang/srt/disaggregation/mooncake/conn.py` +- `python/sglang/srt/disaggregation/nixl/conn.py` + +Needed changes: + +- transfer page lists remain page-based and may include the padded tail page; +- request metadata sent across disaggregation must keep valid token length; +- decode should not infer real prompt length from transferred padded page count; +- draft payload follows target pages, not independent draft token length. + +Risk: + +- page transfer count and logical request length diverge by up to + `page_size - 1` tokens. Any path that uses one as the other is wrong. + +### 7. EAGLE/MTP draft KV + +Files: + +- `python/sglang/srt/speculative/eagle_worker.py` +- `python/sglang/srt/speculative/eagle_info.py` +- `python/sglang/srt/managers/cache_controller.py` + +Needed changes: + +- draft physical pages mirror target padded pages; +- draft valid length mirrors target valid length; +- draft cannot independently choose a shorter valid-only capacity path; +- draft cache-hit prefetch/current reuse must use the same visibility mask. + +Risk: + +- draft and target can both be physically page-aligned while disagreeing on + which rows are valid. That still breaks verification/accept rate. + +## Missing work / correction ledger + +This section is the working checklist for the current implementation. Each item +records what is still missing or what existing code needs to be corrected before +we can treat the page-aligned contract as complete. + +### C1. Add an explicit page extent object + +Current state: + +- `build_page_aligned_in_seq_split_list()` computes `num_page_units` locally. +- Callers still pass raw lengths such as `extend_len`, `logical_len`, and + `required_slots` without naming whether they are valid-token or padded-page + quantities. + +Correction: + +- Add a small helper/data object, for example: + + ```text + PageAlignedCacheExtent( + valid_tokens, + padded_pages, + padded_tokens, + padding_tokens, + ) + ``` + +- Use it where the code crosses logical/physical boundaries. +- `valid_tokens=100, page_size=64` must produce `padded_pages=2` and + `padded_tokens=128`; it must not depend on `cp_size`. + +Tests: + +- Unit-test exact examples such as `0`, `64`, `100`, and `128` valid tokens. +- Add an assertion that `100` tokens with `page_size=64, cp_size=8` does not + become `1024` physical tokens. + +### C2. Extend CP split metadata without padding `split_list` + +Current state: + +- `split_list` is correct because it sums to valid tokens. +- `PageAlignedInSeqSplitInfo` and `NSAContextParallelMetadata` expose only + `page_aligned`, `page_size`, `extend_prefix_len`, and segment page ranges. +- There is no explicit field carrying padded suffix coverage. + +Correction: + +- Keep `split_list`, `actual_token_count`, and query compute valid-token based. +- Add explicit metadata such as `extend_valid_tokens`, `extend_padded_pages`, + `extend_padded_tokens`, and `extend_padding_tokens`. +- Keep `segment_page_starts/ends` page-based and include the tail page coverage. +- Do not change `cp_split_and_rebuild_data()` to consume fake padding tokens. + +Tests: + +- `extend_len=100, page_size=64, cp_size=8` should have: + - `sum(split_list) == 100`; + - `extend_padded_pages == 2`; + - `extend_padded_tokens == 128`; + - zero-token segments for ranks/zigzag segments without pages. + +### C3. Remove short-suffix fallback completely, but only as metadata padding + +Current state: + +- The old `num_pages < cp_size` / `num_pages < 2 * cp_size` fallback was removed + from the split and compute-owner helpers. +- `should_use_replicated_compute_for_short_radix_hit()` now returns `False`. +- `can_cp_split()` still has generic CP minimum-token checks; those are compute + feasibility checks, not page-padding rules. + +Correction: + +- Keep the principle: no fallback solely because a suffix has fewer than + `cp_size` or `2 * cp_size` pages. +- Padding means "cover the partial tail page physically"; it does not mean + "round up to cp_size pages". +- If a truly tiny batch cannot run through CP compute, the reason must be a + clearly named compute limitation, not `too_short_for_page_aligned`. + +Tests: + +- `extend_len=100, page_size=64, cp_size=8` should use two physical page units. +- `extend_len=3 pages, cp_size=8` should produce three compute owners, not + eight or sixteen. + +### C4. MLA partial-current prefetch compose is partially corrected + +Current state: + +- `CpSharedKVMlaPrefetcher.consume_prefix_with_current()` now uses a page-slot + compose helper instead of compact-appending current KV rows. +- `fill_current_kv_page_slots_and_remap_locs()` copies valid current rows into + their dense page slots and masks non-current locs in current pages to `-1`. +- A TAI primitive exists for the CUDA hot path, with a visible warning fallback + to the Torch reference when unavailable. + +Correction still needed: + +- Verify that every MLA mixed prefix/current path uses the page-slot compose + helper when a dense slot buffer already exists. +- Audit the non-prefetch fallback in `nsa_backend.py`: it still calls + `merge_materialized_and_current_kv()`, which is the compact-current helper. + That path may be correct only when the materialized loc table has already + masked suffix slack. This must be proven by tests or changed to the same + page-slot helper. +- Keep `merge_materialized_and_current_kv()` documented as compact-materialize + only; do not silently use it for page-slot prefetch compose. + +Tests: + +- Existing MLA slack test covers a tail page where only part of the suffix page + is valid. +- Add a non-prefetch partial-current test with the same tail-page shape before + changing the fallback path. +- CUDA execution of the TAI primitive must be verified remotely; do not test + CUDA locally. + +### C5. NSA index partial reuse still needs an equivalent padding contract + +Current state: + +- `CpSharedKVIndexPrefetcher.consume()` materializes suffix page slots into a + dense page buffer and returns dense pages/block tables. +- Index prefetch is page-first, so it can easily treat a whole tail page as + visible unless query/cache sequence lengths keep the tail rows masked. +- There is no dedicated unit test proving tail-page slack is invisible for NSA + index partial reuse. + +Correction: + +- Apply the same valid/padded distinction as MLA: + - page tables may include the padded tail page; + - valid sequence length must remain unpadded; + - any loc/table remap that can expose token rows must map padding rows to + invalid/sentinel entries or rely on a proven sequence-length mask. +- Current code inspection shows NSA index uses page ids plus explicit valid + sequence lengths: `GetKAndS` iterates `token_ids < seq_len`, and paged MQA + logits receive `seqlens_32` / `indexer_seq_lens`. Therefore tail-page slack + is not visible if those lengths remain valid-token based. +- The required immediate fix is a regression test that proves index prefetch keeps + the padded tail page in the dense page table while relying on valid + `indexer_seq_lens` to exclude slack. If a future path consumes index page + tables without an accompanying valid length, it must add an explicit + sentinel/remap or warning fallback before using the table. +- If index prefetch cannot prove this contract for a shape, emit a warning + fallback with a narrow reason instead of silently using a risky path. + +Tests: + +- Add an index partial reuse case with a non-page-aligned suffix: + - tail page included in dense pages; + - dense page table has only `ceil(valid_tokens / page_size)` physical pages, + not cp-size padding; + - valid sequence length remains the original unpadded token count so the + index kernels exclude tail slack. + +### C6. HiCache metadata conflates logical and physical lengths + +Current state: + +- `CpHiCacheNodeMetadata.logical_len` is required to be page-aligned. +- `reserve_write_cp()` rejects non-page-aligned `device_indices`. +- `_build_cp_load_back_plan()` and `load_cp()` use `node.host_len` as both the + logical cache-hit length and physical load length. + +Correction: + +- Split the metadata into: + + ```text + valid_len # radix/scheduler-visible tokens + padded_len # physical host/device cache span + page_owners # length padded_len / page_size + owned_positions # positions inside padded_len + ``` + +- Scheduler-facing cache hit length must stay `valid_len`. +- Host reservation, backup, load, eviction, and owner-lane capacity must use + `padded_len`. +- `load_cp()` should validate allocation length against padded physical length, + while returning enough information for the radix/scheduler path to expose only + valid tokens. + +Tests: + +- Metadata with `valid_len=100, padded_len=128, page_size=64`. +- Host hit reports `100` to scheduler but reserves/loads `128` physical slots. +- `page_owners` length is `2`. + +### C7. Radix page alignment currently floors partial tails + +Current state: + +- `page_align_keys()` truncates keys to a page boundary. +- `prepare_write_backup_for_req()` calls `page_align_keys()` before probing and + reserving CP HiCache backup. +- This can preserve correctness, but a short extend can be omitted from host + cache until it completes a full page. Repeated short suffixes can therefore + keep future hits on a fallback/partial path. + +Correction: + +- Do not append fake tokens to radix keys. +- Keep radix keys valid-token based, but attach physical padded metadata to the + node/backup record. +- If a radix split would land inside an in-flight padded physical page, either: + - delay/pending the split; or + - force the split to a page boundary and keep the valid-tail metadata attached + to the child. + +Tests: + +- Write a 100-token suffix and verify the radix-visible key length remains 100 + while physical metadata covers 128 tokens. +- Split a node around the tail page and verify valid/padded lengths remain + consistent. + +### C8. Owner-lane capacity must be padded-page based end to end + +Current state: + +- Compute-owner allocation returns one owner per physical page unit, including a + tail page. +- HiCache capacity planning counts `page_size` tokens per owner page. +- Some entry points still pass scalar valid lengths or page-aligned-only + `device_indices`. + +Correction: + +- Derive required capacity from padded `page_owners`, not valid token counts. +- Admission should use a per-owner vector built from the padded page list. +- Do not introduce an all-reduce/collective to discover capacity agreement; the + owner pattern should make the required vector deterministic. + +Tests: + +- A 100-token suffix with owners `[0, 1]` should require one page on owner 0 and + one page on owner 1. +- Target and draft required vectors must match when draft HiCache is enabled. + +### C9. Draft/EAGLE currently has safe fallbacks, not the final contract + +Current state: + +- Draft HiCache piggybacks target reservations/positions when + `draft_host_indices` are present. +- Draft partial-current reuse/prefetch has been disabled or pushed to older + materialization paths in some cases to avoid EAGLE hangs. + +Correction: + +- Draft physical pages must mirror target padded pages. +- Draft valid length must mirror target valid length. +- Draft cannot independently choose a valid-only capacity or prefetch path. +- Re-enable draft/EAGLE partial reuse only after the same page-slot visibility + mask is proven for draft. + +Tests: + +- Target/draft reservation with a non-page-aligned suffix uses identical + `owned_positions` and padded physical length. +- EAGLE one-layer path does not prefetch or consume an unmasked padded tail. + +### C10. Disaggregated transfer has not been audited for valid/padded split + +Current state: + +- Transfer APIs are page-oriented in several paths, but request metadata remains + token-length oriented. +- There is no current checklist proving that decode never infers prompt length + from transferred padded page count. + +Correction: + +- Transfer page lists may include the padded tail page. +- Request/bootstrap metadata must carry valid token length separately. +- Decode-visible prompt length and max-context checks must use valid length. +- Draft transfer follows target page pattern. + +Tests: + +- Send a non-page-aligned prompt through disaggregation: + - transferred pages include the padded tail page; + - decode sees the original valid length; + - draft and target transfer page counts match when draft KV is enabled. + +### C11. Fallback logging must be audited + +Current state: + +- MLA prefix misalignment now logs a warning fallback. +- Some prefetch misses use debug logs or normal fallback paths. +- Historical silent fallbacks made it hard to identify that async/per-layer or + page-aligned paths were not actually active. + +Correction: + +- Production-visible fallback from an intended hot path must emit + `[CP_SHARED_KV_FALLBACK]` or `[CP_HICACHE_FALLBACK]` at warning level with a + narrow reason. +- Debug-only logs are acceptable only for expected non-fallback states. +- Avoid per-layer spam for expected steady-state events; warn on path changes, + not every successful layer. + +Tests: + +- Unit-test representative fallback logs: + - MLA prefix not page-aligned; + - index prefetch consume miss after start layer; + - CP HiCache reservation/load capacity failure. + +### C12. Verification gaps before claiming completion + +Current state: + +- Python syntax checks pass for the edited SGLang files. +- Local pytest collection can be blocked by missing optional dependencies such + as `pybase64`; this workspace also hit `orjson` during targeted + `test_nsa_cp_utils.py` collection, then `starlette` after temporary stubbing. + `test_cp_shared_kv_layout.py` collection also hit `transformers` via + `sglang.test.test_utils`. +- CUDA kernels must not be tested locally. + +Correction: + +- Use a temporary local test stub only for optional import blockers such as + `pybase64` / `orjson`; do not commit stubs. +- If dependency stubbing starts expanding into unrelated web/server packages, + stop the dependency-stub path and use isolated module loading or + `py_compile` for the targeted helper. Do not spend the implementation loop + repeatedly rediscovering local environment gaps. +- Run CPU/unit tests for extent/split/owner/runtime helpers. +- Run CUDA/TAI kernel verification only on the remote CUDA environment. +- Record remote log evidence for: + - TAI page-slot compose being used; + - no warning fallback on the expected hot path; + - cache-hit request completing without detokenizer hang. + +Tests: + +- Targeted CPU pytest for: + - `test_nsa_cp_utils.py`; + - `test_cp_shared_kv_layout.py`; + - `test_cp_shared_kv_runtime.py`. +- Remote CUDA benchmark/runtime for the TAI materialize primitive and full ETE + prefill run. + +### C13. Remote source sync can expose mixed-version symbols + +Current state: + +- Remote container targeted new tests passed after scp sync. +- A wider remote run of: + + ```text + PYTHONPATH=python python -m pytest -q \ + test/registered/unit/layers/test_nsa_cp_utils.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_layout.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py + ``` + + produced 102 passes and 19 failures. +- The dominant failure is an import mismatch: + + ```text + nsa_indexer.py / nsa_backend.py import + cp_shared_kv_mla_prefetch_should_trace_tiny_extend + from cp_shared_kv_runtime.py, but the synced runtime file does not export it. + ``` + +Correction: + +- Treat this as a mixed-version/sync integrity problem before making logic + changes. +- Compare local and remote `cp_shared_kv_runtime.py`, `nsa_indexer.py`, and + `nsa_backend.py` symbol expectations. +- Either restore/export the expected helper or sync the matching caller files. +- Re-run the same remote wider suite after the symbol contract is repaired. + +Tests: + +- Import-level check: + `python - <<'PY' ... import cp_shared_kv_mla_prefetch_should_trace_tiny_extend`. +- Re-run the three-file remote pytest suite. + +### C14. Wider remote unit suite still exposes non-page-padding regressions + +Current state: + +- After directly overwriting remote `nsa_indexer.py` and `nsa_backend.py` to + remove the mixed-version import mismatch, remote py_compile passed. +- The three-file remote unit suite improved to: + + ```text + 114 passed, 7 failed, 5 warnings + ``` + +- Remaining failures are: + + 1. `test_local_out_cache_loc_logs_every_fallback_event` expects the old + fallback text `"CP shared KV direct-write fallback"`, while current code + emits the newer warning prefix + `[CP_SHARED_KV_FALLBACK][direct_write]`. + 2. MLA and index prefetch stream tests expect prefix materialization to run + inside the prefetch stream context, but current code records materialize on + the current stream and only the all-reduce on the prefetch stream. + 3. MLA and index `wait_attention_window()` tests expect the pending event to + be waited/launched. Current code returns without waiting in those unit + setups. + 4. `test_mla_prefetch_min_prefix_pages_uses_cached_token_default_and_can_override` + expects `SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-2` to return + `16` for `cp_size=4,page_size=64`, but current code returns the dynamic + token default of `8`. + +Correction: + +- Treat these as separate from the page-boundary extent work. +- The fallback-log test should be updated to the new warning prefix because C11 + defines `[CP_SHARED_KV_FALLBACK]` warning-level logging as the intended + production contract. +- Code inspection plus the Phase 8 prefetch doc confirm the current stream/wait + behavior is intentional: prefix materialize runs on the current stream because + MLA/index materialize kernels consume SM and should remain ordered; only the + async all-reduce is enqueued on the prefetch stream. `wait_attention_window()` + is a historical hook and must not wait; the real wait remains at next-layer + `consume(...)` immediately before use. The stale unit tests should assert this + deferred-consume model instead of forcing attention-tail waits. +- `SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-2` has no separate runtime + contract in docs or code; all negative values use the same dynamic default: + `max(cp_size, ceil(_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS / page_size))`. The + expected value changed from 16 to 8 after the default token threshold was reset + from 1024 to 512; the stale test expectation should be updated, not the code. + +Tests: + +- Re-run the exact failed tests after each fix, then the three-file remote + suite. + +## Testing requirements + +Add targeted tests before implementation: + +1. `extend_len % page_size != 0` MLA partial current reuse: + - prefix page-aligned, + - suffix has page slack, + - dense buffer includes padded suffix rows, + - slack locs remap to `-1`. + +2. Equivalent NSA index partial reuse case: + - tail page included in page table, + - valid current tokens remap, + - slack positions invalid. + +3. HiCache metadata: + - valid length and padded length differ, + - `page_owners` length is padded-page count, + - host hit exposed to scheduler is valid length. + +4. Capacity: + - per-owner required capacity is derived from padded pages; + - target/draft required counts match. + +5. Transfer: + - page list includes padded tail page; + - decode-visible sequence length remains valid length. + +## Implementation guidance + +1. Introduce a small helper/data object for page-aligned cache extents. Do not + pass raw `len`, `num_tokens`, or `required_slots` through new code without + naming whether it is valid or padded. +2. Update MLA prefetch compose first, because it is the observed failing path. +3. Update index partial reuse only after MLA has tests for slack masking. +4. Update HiCache metadata/capacity before enabling this contract for host-hit + load/backup paths. +5. Any temporary fallback must be a `WARNING` with a narrow reason and should be + removed after page-padded compose is complete. + +## Non-goals for this pass + +- Do not change request/logical token length semantics. +- Do not make `split_list` include fake tokens. +- Do not add a hot-path CP collective to reconcile capacity. +- Do not make draft KV manage an independent padded/valid contract. +- Do not silently fall back from the page-aligned contract. diff --git a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md index 3b310109c..91abcbb80 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md @@ -234,21 +234,17 @@ Phase 4 MVP 建议： fallback 到旧 token-average split ``` -当前实现采用保守 gate，但对 radix-hit suffix 放宽： +当前 page-aligned contract 已去掉 short-suffix fallback： ```text -如果 extend_prefix_len == 0 且 num_units < 2 * cp_size: - fallback 到旧 token-average split - -如果 extend_prefix_len > 0 且 prefix page-aligned: - 允许 cp_size <= num_units < 2 * cp_size，未覆盖的 segment 长度为 0 - -如果 extend_prefix_len > 0 且 prefix page-aligned 且 num_units < cp_size: - 不启用 CP split；保留 replicated compute（所有 rank 计算 short suffix）， - 但 compute-owner allocator 仍按 page owner 分配 logical page。 +如果 prefix page-aligned 且 extend_len > 0: + 使用 page-aligned split + 每个物理 page unit 只归属一个 zigzag segment + 没有 page 的 segment 长度为 0 + out_cache_loc 仍只包含 valid tokens，不追加 fake padding tokens ``` -原因是 CP 本身不适合无 cache 的短序列；但长 agent 上下文里 radix cache hit 很频繁，命中后 current suffix 可能只有少量 page。如果继续因为 suffix 太短 fallback，会导致 compute-owner allocation/write 在高频 radix-hit 路径失效。放宽后的约束仍然是：prefix 必须 page-aligned，实际 suffix page 仍不被切开。`cp_size <= num_units < 2 * cp_size` 时仍走 page-aligned CP split，每个 CP rank 至少拿到一个 page unit；`num_units < cp_size` 时避免构造 zero-token CP rank，改走 replicated compute，由已有 shared-KV write filter 只写本 rank 拥有的 page。 +原因是 CP 本身不适合无 cache 的短序列；但长 agent 上下文里 radix cache hit 很频繁，命中后 current suffix 可能只有少量 page。如果继续因为 suffix 太短 fallback，会导致 compute-owner allocation/write 在高频 radix-hit 路径失效。当前目标约束是：prefix 必须 page-aligned，实际 suffix page 不被切开。即使 `num_units < cp_size`，也保持 page-aligned CP split，使用 zero-token segment 表示没有 page 的 rank，由 shared-KV write filter 只写本 rank 拥有的 page。 --- @@ -469,7 +465,7 @@ test/registered/unit/attention/test_nsa_cp_page_aligned_split.py - MVP fallback 到旧 split，或显式返回 `page_aligned=False`。 6. too-short case - cache-miss: `num_units < 2 * cp_size` fallback 或 `page_aligned=False`。 - - radix-hit 且 prefix page-aligned: `cp_size <= num_units < 2 * cp_size` 允许；`num_units < cp_size` 不构造 CP split，改走 replicated compute，但 compute-owner allocation 仍可用。 + - radix-hit 且 prefix page-aligned: `num_units < 2 * cp_size` 也允许；没有 page 的 segment 使用 zero-token segment，compute-owner allocation 仍可用。 ### 7.2 Invariant tests diff --git a/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md b/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md index 412e60412..f307fe810 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md @@ -493,7 +493,7 @@ radix prefix 命中会让 current extend 从已有 logical pages 之后开始。 Phase 4/5 对 radix-hit short suffix 放宽 too-short gate： - 当 `extend_prefix_len > 0` 且 prefix page-aligned，并且 current suffix page 数至少为 `cp_size` 时，即使 page 数小于 `2 * cp_size`，仍生成 page-aligned split / compute-owner page owner list；未覆盖的第二段 zigzag segment 长度为 0。 -- 当 suffix page 数小于 `cp_size` 时，仍允许 compute-owner page allocation，但不启用 CP split。此时沿用 SGLang 原有 replicated compute 行为：所有 rank 计算 short suffix，shared-KV 的 MLA/index write filter 只保留本 rank owner page 的写入。这样避免 zero-token CP rank 的通信/kernel 边界问题，同时避免 radix-hit 高频短 suffix 回退到 legacy allocation。 +- 当前 page-aligned contract 已废弃 short suffix 的 replicated compute 特例。suffix page 数小于 `cp_size` 时，compute-owner page allocation 仍返回实际 page 的 owner list，CP split 使用 zero-token segments 表示没有 page 的 rank。这样保持 owner-lane / HiCache / materialize 的物理 page pattern，不再因为短 suffix 回退到 legacy allocation。 如果命中到 partial page，MVP fallback。 diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 76567bfe7..d8e89a260 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -18,11 +18,11 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_mla_prefetch_should_log_layer, filter_locs_mappable_to_physical_pool, filter_pages_mappable_to_physical_pool, + fill_current_kv_page_slots_and_remap_locs, get_or_build_shared_paged_buffer_slot_remap, get_or_build_shared_token_kv_slot_remap, materialize_local_paged_buffer_page_slots_into, materialize_local_token_kv_page_slots_into, - merge_materialized_and_current_kv, remap_logical_pages_to_slot_dense_pages, remap_logical_locs_to_slot_dense_locs_optimized, slot_range_to_page_slice, @@ -39,6 +39,14 @@ def _prefetch_log(message: str, *args) -> None: cp_shared_kv_mla_prefetch_log(message, *args) +def _mla_prefetch_fallback_log(reason: str, message: str, *args) -> None: + logger.warning( + "[CP_SHARED_KV_FALLBACK][mla_prefetch] reason=%s " + message, + reason, + *args, + ) + + def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None: logger.warning( "[CP_SHARED_KV_FALLBACK][index_prefetch] reason=%s " + message, @@ -384,8 +392,10 @@ class CpSharedKVMlaPrefetcher: return None extend_prefix_len = int(extend_prefix_lens_cpu[0]) if extend_prefix_len <= 0 or extend_prefix_len % page_size != 0: - _prefetch_log( - "create_skip reason=prefix_not_page_aligned prefix_len=%s page_size=%s", + _mla_prefetch_fallback_log( + "prefix_not_page_aligned", + "prefix length is zero or not page-aligned. " + "prefix_len=%s page_size=%s", extend_prefix_len, page_size, ) @@ -726,14 +736,15 @@ class CpSharedKVMlaPrefetcher: page_inverse=self.page_inverse, page_size=self.page_size, ) - mixed_kv_cache, mixed_locs, _ = merge_materialized_and_current_kv( - materialized_kv_cache=dense_kv_cache, + mixed_kv_cache, mixed_locs, _ = fill_current_kv_page_slots_and_remap_locs( + dense_kv_cache=dense_kv_cache, materialized_dense_locs=dense_locs, current_kv_cache=current_kv_cache, logical_locs=logical_locs, current_locs=current_locs, - page_size=current_remap_page_size, - logical_page_capacity=current_remap_logical_page_capacity, + page_inverse=self.page_inverse, + page_size=self.page_size, + mask_non_current_in_current_pages=True, ) remap_ms = _cpu_timing_ms(remap_cpu) total_ms = _cpu_timing_ms(consume_cpu) 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 56f74ff77..38c5fd0ab 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 @@ -642,6 +642,135 @@ def remap_logical_locs_to_slot_dense_locs_optimized( ) +def _try_tai_fill_current_kv_page_slots_and_remap_locs( + *, + dense_kv_cache: torch.Tensor, + materialized_dense_locs: torch.Tensor, + current_kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + current_locs: torch.Tensor, + page_inverse: torch.Tensor, + page_size: int, + mask_non_current_in_current_pages: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + if not _tai_materialize_runtime_enabled(): + return None + + kernels = _load_tai_materialize_kernels() + if kernels is None: + return None + fill_kernel = getattr( + kernels, + "fill_current_token_kv_page_slots_and_remap_locs", + None, + ) + if fill_kernel is None: + _log_tai_materialize_fallback( + "fill_current_missing", + "CP shared KV tai current-slot fill kernel is unavailable; " + "falling back to torch reference. Upgrade tai-kernel to keep this " + "hot path off PyTorch. 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), + _contiguous_for_tai(materialized_dense_locs), + _contiguous_for_tai(current_kv_cache), + _contiguous_for_tai(logical_locs), + _contiguous_for_tai(current_locs.reshape(-1)), + _contiguous_for_tai(page_inverse), + page_size=int(page_size), + mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages), + ) + except Exception as exc: + _log_tai_materialize_fallback( + "fill_current_failed", + "CP shared KV tai current-slot fill failed; falling back to torch " + "reference. error=%s page_size=%s current_rows=%s query_locs=%s", + exc, + page_size, + int(current_kv_cache.shape[0]), + int(logical_locs.numel()), + ) + return None + + +def fill_current_kv_page_slots_and_remap_locs( + *, + dense_kv_cache: torch.Tensor, + materialized_dense_locs: torch.Tensor, + current_kv_cache: torch.Tensor, + logical_locs: torch.Tensor, + current_locs: torch.Tensor, + page_inverse: torch.Tensor, + page_size: int, + mask_non_current_in_current_pages: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fill current suffix KV into preallocated dense page slots and remap locs. + + This is the CP shared-KV prefetch compose path. The dense buffer already has + page-aligned slots for prefix and suffix pages; current rows should be copied + into those suffix slots instead of appended to a second compact suffix. + """ + + tai_result = _try_tai_fill_current_kv_page_slots_and_remap_locs( + dense_kv_cache=dense_kv_cache, + materialized_dense_locs=materialized_dense_locs, + current_kv_cache=current_kv_cache, + logical_locs=logical_locs, + current_locs=current_locs, + page_inverse=page_inverse, + page_size=page_size, + mask_non_current_in_current_pages=mask_non_current_in_current_pages, + ) + if tai_result is not None: + return tai_result + + if dense_kv_cache.is_cuda: + _log_tai_materialize_fallback( + "fill_current_torch_reference_cuda", + "CP shared KV current-slot fill is using the torch reference on CUDA; " + "this is a fallback and should not be the steady-state hot path. " + "page_size=%s current_rows=%s query_locs=%s", + page_size, + int(current_kv_cache.shape[0]), + int(logical_locs.numel()), + limit=1, + ) + + current_dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( + current_locs.reshape(-1), + page_inverse=page_inverse, + page_size=page_size, + ) + valid_current_rows = current_dense_locs >= 0 + if torch.any(valid_current_rows): + dense_kv_cache[current_dense_locs[valid_current_rows].to(torch.long)] = ( + current_kv_cache[valid_current_rows] + ) + + current_mask, _ = build_current_loc_remap(logical_locs, current_locs) + mixed_locs = materialized_dense_locs + if mask_non_current_in_current_pages: + current_page_mask = build_current_page_mask( + logical_locs, + current_locs, + page_size=page_size, + ) + mixed_locs = torch.where( + current_page_mask & (~current_mask), + torch.full_like(materialized_dense_locs, -1), + materialized_dense_locs, + ) + return dense_kv_cache, mixed_locs, current_mask + + def _copy_tai_dense_slot_range_body( *, tai_dense_kv_cache: torch.Tensor, @@ -866,6 +995,43 @@ def current_loc_remap_fast_path_args( return page_size, logical_page_capacity +def build_current_page_mask( + query_locs: torch.Tensor, + current_locs: torch.Tensor, + *, + page_size: int, +) -> torch.Tensor: + """Return true for query locs that fall inside a current suffix page. + + Partial-current reuse can have a tail page where only the first N rows are + real current tokens. The page itself is part of the physical suffix, but + rows after the valid current tokens are padding/slack and must be invisible + to attention. + """ + + mask = torch.zeros_like(query_locs, dtype=torch.bool) + if page_size <= 1 or query_locs.numel() == 0 or current_locs.numel() == 0: + return mask + + query_flat = query_locs.reshape(-1).to(torch.long) + current_flat = current_locs.reshape(-1).to(torch.long) + current_pages = torch.unique( + torch.div(current_flat, page_size, rounding_mode="floor") + ) + if current_pages.numel() == 0: + return mask + + sorted_pages, _ = torch.sort(current_pages) + valid_query = query_flat >= 0 + safe_query = torch.where(valid_query, query_flat, torch.zeros_like(query_flat)) + query_pages = torch.div(safe_query, page_size, rounding_mode="floor") + insert_positions = torch.searchsorted(sorted_pages, query_pages) + safe_positions = torch.clamp(insert_positions, max=sorted_pages.numel() - 1) + in_range = insert_positions < sorted_pages.numel() + matched = valid_query & in_range & (sorted_pages[safe_positions] == query_pages) + return matched.reshape(query_locs.shape) + + def merge_materialized_and_current_kv( *, materialized_kv_cache: torch.Tensor, @@ -882,6 +1048,11 @@ def merge_materialized_and_current_kv( materialization path. Entries corresponding to current extend tokens are replaced with offsets into the appended ``current_kv_cache``. Non-current entries remain untouched, including ``-1`` invalid sentinels. + + This helper is for compact prefix materialization that appends a compact + current suffix. Prefetch paths with existing page slots must use + :func:`fill_current_kv_page_slots_and_remap_locs` instead, so the physical + page padding stays in the original dense slot layout. """ current_mask, current_rows = build_current_loc_remap( diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 76e8c89e8..08bc31cb4 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -173,6 +173,32 @@ def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens): return nsa_cache_seqlens +@dataclass(frozen=True) +class PageAlignedCacheExtent: + valid_tokens: int + padded_pages: int + padded_tokens: int + padding_tokens: int + + +def build_page_aligned_cache_extent( + *, valid_tokens: int, page_size: int +) -> PageAlignedCacheExtent: + if valid_tokens < 0: + raise ValueError(f"valid_tokens must be non-negative, got {valid_tokens}") + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + + padded_pages = ceil_div(valid_tokens, page_size) if valid_tokens > 0 else 0 + padded_tokens = padded_pages * page_size + return PageAlignedCacheExtent( + valid_tokens=valid_tokens, + padded_pages=padded_pages, + padded_tokens=padded_tokens, + padding_tokens=padded_tokens - valid_tokens, + ) + + @dataclass class PageAlignedInSeqSplitInfo: page_aligned: bool = False @@ -180,6 +206,10 @@ class PageAlignedInSeqSplitInfo: extend_prefix_len: int = 0 segment_page_starts: List[int] = None segment_page_ends: List[int] = None + extend_valid_tokens: int = 0 + extend_padded_pages: int = 0 + extend_padded_tokens: int = 0 + extend_padding_tokens: int = 0 @dataclass @@ -208,6 +238,10 @@ class NSAContextParallelMetadata: extend_prefix_len: int = 0 segment_page_starts: List[int] = None segment_page_ends: List[int] = None + extend_valid_tokens: int = 0 + extend_padded_pages: int = 0 + extend_padded_tokens: int = 0 + extend_padding_tokens: int = 0 def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]: @@ -233,6 +267,10 @@ def _fallback_page_aligned_split_info( extend_prefix_len=extend_prefix_len, segment_page_starts=[], segment_page_ends=[], + extend_valid_tokens=0, + extend_padded_pages=0, + extend_padded_tokens=0, + extend_padding_tokens=0, ) @@ -246,13 +284,10 @@ def build_page_aligned_in_seq_split_list( ) -> Tuple[List[int], PageAlignedInSeqSplitInfo]: """Build an in-seq split list whose real-token boundaries do not cut pages. - Phase 4 deliberately uses a conservative gate for cache-miss chunks: at - least `2 * cp_size` page units are required so every zigzag segment has at - least one page unit. For radix-hit suffixes with a page-aligned prefix, the - gate is relaxed to `cp_size` page units so every CP rank still receives at - least one local page while second zigzag segments may be empty. When the - gate does not hold, this helper falls back to the existing token-balanced - split and marks the result as not page-aligned. + The split remains valid-token based, but page metadata covers the physical + page units. Short chunks keep the page-aligned contract by assigning zero + valid tokens to surplus zigzag segments instead of falling back to + token-balanced splits that would poison later shared-KV/HiCache reuse. """ if extend_len < 0: @@ -271,14 +306,14 @@ def build_page_aligned_in_seq_split_list( if page_size <= 1 or extend_len <= 0 or extend_prefix_len % page_size != 0: return fallback_split, fallback_info + extent = build_page_aligned_cache_extent( + valid_tokens=extend_len, + page_size=page_size, + ) full_pages = extend_len // page_size tail_tokens = extend_len % page_size - num_page_units = full_pages + (1 if tail_tokens > 0 else 0) + num_page_units = extent.padded_pages cp_segment_num = cp_size * 2 - if num_page_units < cp_size or ( - num_page_units < cp_segment_num and extend_prefix_len == 0 - ): - return fallback_split, fallback_info base_units = num_page_units // cp_segment_num remainder_units = num_page_units % cp_segment_num @@ -315,6 +350,10 @@ def build_page_aligned_in_seq_split_list( extend_prefix_len=extend_prefix_len, segment_page_starts=segment_page_starts, segment_page_ends=segment_page_ends, + extend_valid_tokens=extent.valid_tokens, + extend_padded_pages=extent.padded_pages, + extend_padded_tokens=extent.padded_tokens, + extend_padding_tokens=extent.padding_tokens, ) @@ -356,50 +395,12 @@ def should_use_replicated_compute_for_short_radix_hit( ) -> bool: """Return whether a short radix-hit suffix should avoid CP splitting. - With CP shared KV, radix-hit suffixes can be page-aligned but shorter than - one page per CP rank. A page-aligned CP split would give some ranks zero - local tokens, which is unsafe for parts of the current CP collective/kernel - path. Instead, keep the original non-CP behavior: every rank computes the - short suffix, while shared-KV write filters persist only pages owned by the - local rank. + Kept as a compatibility hook for older callers. The page-aligned cache + contract no longer uses replicated compute for short suffixes: zero-length + CP segments are preferred over breaking the physical page-owner pattern. """ - if ( - forward_batch is None - or cp_size <= 0 - or not getattr(forward_batch, "uses_cp_shared_kv", False) - ): - return False - - extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) - extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) - if ( - extend_seq_lens_cpu is None - or extend_prefix_lens_cpu is None - or len(extend_seq_lens_cpu) != 1 - or len(extend_prefix_lens_cpu) != 1 - ): - return False - - token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None) - page_size = getattr(token_to_kv_pool, "page_size", None) - if page_size is None: - return False - page_size = int(page_size) - if page_size <= 1: - return False - - extend_len = int(extend_seq_lens_cpu[0]) - extend_prefix_len = int(extend_prefix_lens_cpu[0]) - if ( - extend_len <= 0 - or extend_prefix_len <= 0 - or extend_prefix_len % page_size != 0 - ): - return False - - num_page_units = ceil_div(extend_len, page_size) - return 0 < num_page_units < cp_size + return False def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): @@ -1177,6 +1178,10 @@ def prepare_input_dp_with_cp_dsa( extend_prefix_len=page_split_info.extend_prefix_len, segment_page_starts=page_split_info.segment_page_starts, segment_page_ends=page_split_info.segment_page_ends, + extend_valid_tokens=page_split_info.extend_valid_tokens, + extend_padded_pages=page_split_info.extend_padded_pages, + extend_padded_tokens=page_split_info.extend_padded_tokens, + extend_padding_tokens=page_split_info.extend_padding_tokens, ) return nsa_cp_metadata diff --git a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py index 8df8e6b6c..dd0930f60 100644 --- a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py +++ b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py @@ -24,9 +24,6 @@ def get_in_seq_page_compute_owner_unavailable_reason( full_pages = extend_len // page_size tail_tokens = extend_len % page_size num_page_units = full_pages + (1 if tail_tokens > 0 else 0) - if num_page_units < cp_size * 2 and extend_prefix_len == 0: - return "too_short_for_page_aligned" - return None @@ -39,12 +36,11 @@ def build_in_seq_page_compute_owners( ) -> Optional[List[int]]: """Return compute-owner CP rank for each newly allocated current page. - This mirrors the Phase 4 page-aligned `in-seq-split` segmentation for - normal CP chunks, but it only returns page-unit owners for the real extend - chunk. Short radix-hit suffixes with fewer pages than CP ranks are also - allowed: runtime keeps replicated compute for those chunks and the shared - KV write filters persist only locally owned pages. `None` means the batch - must stay on the legacy allocation/write path. + This mirrors the page-aligned `in-seq-split` segmentation and returns one + owner per physical page unit, including a tail page. Short chunks keep the + page-aligned owner pattern by assigning zero pages to surplus zigzag + segments instead of falling back to legacy allocation. `None` means the + batch must stay on the legacy allocation/write path. """ if cp_size <= 0: diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 79756b170..8b303c000 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -5,6 +5,8 @@ from unittest.mock import patch from sglang.srt.layers.attention.nsa.utils import ( NSAContextParallelMetadata, + PageAlignedCacheExtent, + build_page_aligned_cache_extent, _get_in_seq_last_token_owner_and_offset, build_page_aligned_in_seq_split_list, build_token_balanced_in_seq_split_list, @@ -22,6 +24,36 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=1, suite="stage-a-test-cpu") +class TestPageAlignedCacheExtent(unittest.TestCase): + def test_extent_uses_page_boundary_not_cp_size(self): + extent = build_page_aligned_cache_extent(valid_tokens=100, page_size=64) + + self.assertEqual(extent.valid_tokens, 100) + self.assertEqual(extent.padded_pages, 2) + self.assertEqual(extent.padded_tokens, 128) + self.assertEqual(extent.padding_tokens, 28) + + def test_extent_handles_empty_and_aligned_lengths(self): + self.assertEqual( + build_page_aligned_cache_extent(valid_tokens=0, page_size=64), + PageAlignedCacheExtent( + valid_tokens=0, + padded_pages=0, + padded_tokens=0, + padding_tokens=0, + ), + ) + self.assertEqual( + build_page_aligned_cache_extent(valid_tokens=128, page_size=64), + PageAlignedCacheExtent( + valid_tokens=128, + padded_pages=2, + padded_tokens=128, + padding_tokens=0, + ), + ) + + class TestNSAInSeqCPUtils(unittest.TestCase): def test_contiguous_valid_cp_query_count(self): from sglang.srt.layers.attention.nsa.nsa_indexer import ( @@ -118,6 +150,24 @@ class TestNSAInSeqCPUtils(unittest.TestCase): split_list, extend_prefix_len=0, extend_len=1100, page_size=64 ) + def test_page_aligned_split_exposes_padded_extent_without_padding_split_list(self): + split_list, info = build_page_aligned_in_seq_split_list( + total_len=100, + extend_len=100, + extend_prefix_len=0, + page_size=64, + cp_size=8, + ) + + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 100) + self.assertEqual(split_list[:2], [64, 36]) + self.assertEqual(split_list[2:], [0] * 14) + self.assertEqual(info.extend_valid_tokens, 100) + self.assertEqual(info.extend_padded_pages, 2) + self.assertEqual(info.extend_padded_tokens, 128) + self.assertEqual(info.extend_padding_tokens, 28) + def test_page_aligned_split_falls_back_when_prefix_is_not_page_aligned(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=1024, @@ -130,7 +180,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertFalse(info.page_aligned) self.assertEqual(split_list, build_token_balanced_in_seq_split_list(1024, 8)) - def test_page_aligned_split_falls_back_when_page_units_are_too_short(self): + def test_page_aligned_split_pads_zero_segments_when_page_units_are_short(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=512, extend_len=512, @@ -139,8 +189,13 @@ class TestNSAInSeqCPUtils(unittest.TestCase): cp_size=8, ) - self.assertFalse(info.page_aligned) - self.assertEqual(split_list, build_token_balanced_in_seq_split_list(512, 8)) + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 512) + self.assertEqual(split_list[:8], [64] * 8) + self.assertEqual(split_list[8:], [0] * 8) + self.assert_page_aligned_boundaries( + split_list, extend_prefix_len=0, extend_len=512, page_size=64 + ) def test_page_aligned_split_allows_radix_hit_suffix_with_one_page_per_rank(self): split_list, info = build_page_aligned_in_seq_split_list( @@ -159,7 +214,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): split_list, extend_prefix_len=54464, extend_len=512, page_size=64 ) - def test_page_aligned_split_falls_back_when_radix_hit_suffix_has_zero_rank(self): + def test_page_aligned_split_keeps_short_radix_hit_suffix_page_aligned(self): split_list, info = build_page_aligned_in_seq_split_list( total_len=256, extend_len=256, @@ -168,10 +223,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase): cp_size=8, ) - self.assertFalse(info.page_aligned) - self.assertEqual(split_list, build_token_balanced_in_seq_split_list(256, 8)) + self.assertTrue(info.page_aligned) + self.assertEqual(sum(split_list), 256) + self.assertEqual(split_list[:4], [64] * 4) + self.assertEqual(split_list[4:], [0] * 12) + self.assert_page_aligned_boundaries( + split_list, extend_prefix_len=54464, extend_len=256, page_size=64 + ) - def test_can_cp_split_uses_replicated_compute_for_short_radix_hit_suffix(self): + def test_can_cp_split_keeps_cp_for_short_radix_hit_suffix(self): class Mode: def is_context_parallel_extend(self): return True @@ -194,7 +254,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): return_value=True, ), ): - self.assertFalse(can_cp_split(256, 8, True, forward_batch)) + self.assertTrue(can_cp_split(256, 8, True, forward_batch)) def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self): class Mode: @@ -478,15 +538,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase): ) with self.assertLogs( - "sglang.srt.layers.attention.nsa.utils", level="INFO" + "sglang.srt.layers.attention.nsa.utils", level="WARNING" ) as cm: self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch)) self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch)) self.assertEqual(len(cm.output), 2) - self.assertIn("CP shared KV direct-write fallback", cm.output[0]) + self.assertIn("[CP_SHARED_KV_FALLBACK][direct_write]", cm.output[0]) self.assertIn("metadata is not page-aligned", cm.output[0]) - self.assertIn("CP shared KV direct-write fallback", cm.output[1]) + self.assertIn("[CP_SHARED_KV_FALLBACK][direct_write]", cm.output[1]) self.assertIn("metadata is not page-aligned", cm.output[1]) def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self): diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index f084e8578..6a8a31b25 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py @@ -150,7 +150,7 @@ class TestCPSharedPagedAllocator(CustomTestCase): [0, 0, 1, 1, 2, 2, 3, 3, 3, 3, 2, 2, 1, 1, 0, 0], ) - def test_compute_owner_page_assignment_falls_back_for_short_extend(self): + def test_compute_owner_page_assignment_keeps_page_aligned_short_extend(self): from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( build_in_seq_page_compute_owners, get_in_seq_page_compute_owner_unavailable_reason, @@ -163,17 +163,30 @@ class TestCPSharedPagedAllocator(CustomTestCase): cp_size=4, ) - self.assertIsNone(owners) - self.assertEqual( + self.assertEqual(owners, [0, 1, 2, 3, 3, 2, 1]) + self.assertIsNone( get_in_seq_page_compute_owner_unavailable_reason( extend_len=64 * 7, extend_prefix_len=0, page_size=64, cp_size=4, - ), - "too_short_for_page_aligned", + ) ) + def test_compute_owner_page_assignment_uses_tail_page_not_cp_padding(self): + from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( + build_in_seq_page_compute_owners, + ) + + owners = build_in_seq_page_compute_owners( + extend_len=100, + extend_prefix_len=0, + page_size=64, + cp_size=8, + ) + + self.assertEqual(owners, [0, 1]) + def test_compute_owner_page_assignment_allows_radix_hit_suffix_with_one_page_per_rank( self, ): @@ -199,7 +212,7 @@ class TestCPSharedPagedAllocator(CustomTestCase): ) ) - def test_compute_owner_page_assignment_allows_short_radix_hit_suffix_with_replicated_compute( + def test_compute_owner_page_assignment_keeps_short_radix_hit_suffix_page_aligned( self, ): from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( 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 80b6e94ce..134efadad 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 @@ -48,7 +48,7 @@ class _FakeExtendForwardMode: class TestCpSharedKVRuntimeHelpers(unittest.TestCase): - def test_mla_prefetch_materializes_and_reduces_on_prefetch_stream( + def test_mla_prefetch_materializes_on_current_stream_and_reduces_on_prefetch_stream( self, ): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -137,7 +137,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) self.assertEqual( calls, - [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], + [("materialize", "current"), ("reduce", "prefetch", "prefetch")], ) self.assertEqual(prefetch_stream.waited, ["current"]) @@ -145,11 +145,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual( calls, - [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], + [("materialize", "current"), ("reduce", "prefetch", "prefetch")], ) self.assertEqual(prefetch_stream.waited, ["current"]) - def test_index_prefetch_materializes_and_reduces_on_prefetch_stream( + def test_index_prefetch_materializes_on_current_stream_and_reduces_on_prefetch_stream( self, ): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -239,7 +239,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) self.assertEqual( calls, - [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], + [("materialize", "current"), ("reduce", "prefetch", "prefetch")], ) self.assertEqual(prefetch_stream.waited, ["current"]) @@ -247,7 +247,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual( calls, - [("materialize", "prefetch"), ("reduce", "prefetch", "prefetch")], + [("materialize", "current"), ("reduce", "prefetch", "prefetch")], ) self.assertEqual(prefetch_stream.waited, ["current"]) @@ -539,6 +539,41 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) self.assertEqual(mixed_locs.tolist(), [[4, 8, -1], [9, 7, -1]]) + def test_fill_current_kv_page_slots_masks_suffix_page_slack(self): + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + fill_current_kv_page_slots_and_remap_locs, + ) + + page_size = 4 + dense_kv = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1) + original_dense_kv = dense_kv.clone() + current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1) + logical_locs = torch.tensor([[4, 20, 21, 22, 23]], dtype=torch.int32) + # The prefetched dense buffer has physical rows for the whole suffix page, + # but only locs 20 and 21 are real current tokens. Rows for 22 and 23 are + # page slack and must not stay visible. + materialized_locs = torch.tensor([[4, 12, 13, 14, 15]], dtype=torch.int32) + current_locs = torch.tensor([20, 21], dtype=torch.int64) + page_inverse = torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64) + + mixed_kv, mixed_locs, current_mask = fill_current_kv_page_slots_and_remap_locs( + dense_kv_cache=dense_kv, + materialized_dense_locs=materialized_locs, + current_kv_cache=current_kv, + logical_locs=logical_locs, + current_locs=current_locs, + page_inverse=page_inverse, + page_size=page_size, + mask_non_current_in_current_pages=True, + ) + + expected_kv = original_dense_kv.clone() + expected_kv[12:14] = current_kv + self.assertEqual(int(mixed_kv.shape[0]), 16) + self.assertTrue(torch.equal(mixed_kv, expected_kv)) + self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]]) + self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]]) + def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -553,6 +588,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_stream = FakeCurrentStream() fake_event = object() dense_kv = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1) + original_dense_kv = dense_kv.clone() current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1) page_inverse = torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64) prefetcher = prefetch.CpSharedKVMlaPrefetcher( @@ -583,7 +619,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current( layer_id=1, kv_cache=torch.zeros((64, 1, 1), dtype=torch.float32), - logical_locs=torch.tensor([[4, 20], [21, 7]], dtype=torch.int32), + logical_locs=torch.tensor( + [[4, 20], [21, 7], [22, 23]], dtype=torch.int32 + ), current_kv_cache=current_kv, current_locs=torch.tensor([20, 21], dtype=torch.int64), ) @@ -591,11 +629,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(current_stream.events, [fake_event]) self.assertEqual(prefetcher.handles, {}) self.assertIsNone(prefetcher.pending_attention_handle) - self.assertTrue(torch.equal(mixed_kv[:16], dense_kv)) - self.assertTrue(torch.equal(mixed_kv[16:], current_kv)) - self.assertEqual(mixed_locs.tolist(), [[4, 16], [17, 7]]) + expected_kv = original_dense_kv.clone() + expected_kv[12:14] = current_kv + self.assertTrue(torch.equal(mixed_kv, expected_kv)) + self.assertEqual(mixed_locs.tolist(), [[4, 12], [13, 7], [-1, -1]]) - def test_mla_prefetch_attention_window_waits_on_pending_event(self): + def test_mla_prefetch_attention_window_defers_pending_event_wait(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -631,11 +670,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ): prefetcher.wait_attention_window() - self.assertEqual(current_stream.events, [fake_event]) - self.assertIsNone(prefetcher.pending_attention_handle) + self.assertEqual(current_stream.events, []) + self.assertIs(prefetcher.pending_attention_handle, handle) self.assertIs(prefetcher.handles[1], handle) - def test_mla_prefetch_attention_window_launches_pending_reduce_before_wait(self): + def test_mla_prefetch_attention_window_does_not_launch_pending_reduce(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -647,7 +686,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.events.append(event) current_stream = FakeCurrentStream() - fake_event = object() prefetcher = prefetch.CpSharedKVMlaPrefetcher( layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), page_size=4, @@ -666,21 +704,19 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetcher.handles[1] = handle prefetcher.pending_attention_handle = handle - def finish_reduce(): - handle.event = fake_event - with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream ), patch.object( - prefetcher, "launch_pending_reduce", side_effect=finish_reduce + prefetcher, "launch_pending_reduce" ) as launch_pending_reduce: prefetcher.wait_attention_window() - launch_pending_reduce.assert_called_once_with() - self.assertEqual(current_stream.events, [fake_event]) - self.assertIsNone(prefetcher.pending_attention_handle) + launch_pending_reduce.assert_not_called() + self.assertEqual(current_stream.events, []) + self.assertIs(prefetcher.pending_attention_handle, handle) + self.assertIsNone(handle.event) - def test_index_prefetch_attention_window_waits_on_pending_event(self): + def test_index_prefetch_attention_window_defers_pending_event_wait(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -715,10 +751,69 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ): prefetcher.wait_attention_window() - self.assertEqual(current_stream.events, [fake_event]) - self.assertIsNone(prefetcher.pending_attention_handle) + self.assertEqual(current_stream.events, []) + self.assertIs(prefetcher.pending_attention_handle, handle) self.assertIs(prefetcher.handles[1], handle) + def test_index_prefetch_tail_page_keeps_valid_length_unpadded(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeCurrentStream: + def __init__(self): + self.events = [] + + def wait_event(self, event): + self.events.append(event) + + page_size = 4 + valid_tokens = 10 + padded_pages = (valid_tokens + page_size - 1) // page_size + cp_size = 8 + page_buffer = torch.arange(0, 8 * 3, dtype=torch.uint8).view(8, 3) + dense_page_buffer = torch.zeros((padded_pages + 1, 3), dtype=torch.uint8) + dense_page_buffer[1] = page_buffer[1] + fake_event = object() + prefetcher = prefetch.CpSharedKVIndexPrefetcher( + layout=CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=0), + prefix_pages=1, + slot_logical_pages=torch.tensor([1, 2, 3], dtype=torch.int64), + page_inverse=torch.tensor([-1, 1, 2, 3], dtype=torch.int64), + dense_num_pages=padded_pages, + stream=object(), + ) + handle = prefetch.CpSharedKVIndexPrefetchHandle( + layer_id=1, + dense_page_buffer=dense_page_buffer, + prefix_rows=slice(1, 2), + event=fake_event, + ) + prefetcher.handles[1] = handle + prefetcher.pending_attention_handle = handle + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=FakeCurrentStream() + ), patch.object( + prefetch, "_all_reduce_materialized_buffer_range", _identity_all_reduce + ): + dense_pages_buffer, dense_pages = prefetcher.consume( + layer_id=1, + page_buffer=page_buffer, + logical_pages=torch.tensor([[1, 2, 3]], dtype=torch.int32), + ) + + self.assertEqual(valid_tokens, 10) + self.assertEqual(padded_pages, 3) + self.assertEqual(dense_pages.tolist(), [[1, 2, 3]]) + self.assertEqual(dense_pages.shape[1], padded_pages) + self.assertLess(dense_pages.shape[1], cp_size) + self.assertEqual(list(dense_pages_buffer.shape), [padded_pages + 1, 3]) + self.assertTrue(torch.equal(dense_pages_buffer[1], page_buffer[1])) + # Index kernels receive valid sequence lengths separately; the dense + # table includes the tail page, but no fake 12-token length is exposed. + self.assertLess(valid_tokens, padded_pages * page_size) + self.assertEqual(valid_tokens, 10) + def test_materialize_local_token_kv_pages(self): from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_dense_page_remap, @@ -1023,6 +1118,54 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True): self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled()) + def test_mla_prefetch_prefix_misalignment_logs_warning_fallback(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + hisparse_coordinator=None, + forward_mode=Mode(), + batch_size=1, + token_to_kv_pool=SimpleNamespace(page_size=64), + cp_shared_kv_layout=SimpleNamespace(cp_size=1, cp_rank=0), + extend_prefix_lens_cpu=[65], + ) + metadata = SimpleNamespace( + real_page_table=torch.tensor([1, 2], dtype=torch.int64), + page_table_1=torch.tensor([[64, 65]], dtype=torch.int32), + ) + + with patch.object( + prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True + ), patch.object( + prefetch, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + prefetch.torch.cuda, "is_available", return_value=True + ), patch.object( + prefetch, "_is_cuda_stream_capturing", return_value=False + ), patch.object( + prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True + ), patch.object( + prefetch.logger, "warning" + ) as logger: + result = prefetch.CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=True, + ) + + self.assertIsNone(result) + logger.assert_called_once() + self.assertIn( + "[CP_SHARED_KV_FALLBACK][mla_prefetch]", + logger.call_args.args[0], + ) + 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): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -1053,7 +1196,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(-2): self.assertEqual( runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(4, page_size=64), - 16, + max(4, expected_pages), ) with patch.object(runtime, "_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS", 2048):