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

Keep CP shared-KV cache hits page-aligned without short fallback

Browse Source 
CP shared KV and HiCache need a stable contract where physical cache coverage is page-aligned, while scheduler/radix-visible hit length remains the valid token length. This records the contract, adds page-aligned extent metadata, keeps owner assignment on actual tail pages instead of short-prefix fallback, and updates partial current reuse tests around tail-page masking.

Constraint: CP owner lanes operate on page units while scheduler and radix hit accounting must remain token-valid.

Rejected: Pad short suffixes to cp_size or 2*cp_size pages | wastes KV capacity and can turn a small tail into a much larger physical span.

Rejected: Silent direct-write or prefetch fallback | production fallback must be warning-visible for diagnosis.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce replicated short-radix fallback without checking docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md.

Tested: local py_compile for touched runtime, utility, owner, and unit-test files.

Tested: remote g0034 container three-file suite: 122 passed, 5 warnings.

Not-tested: full local pytest, blocked by missing runtime dependencies such as orjson.

Not-tested: CUDA E2E runtime for this commit.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 05:04:50 +08:00
parent 1e1807f2bc
commit c551623ca8
11 changed files with 1369 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/advanced_features/nsa_prefill_cp_batch_size_gt1_exploration.md
View File

@@ -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;

847
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md Normal file
View File

@@ -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.

20
docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md
View File

@@ -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

2
docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md
View File

@@ -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。

25
python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
View File

@@ -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)

171
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
View File

@@ -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(

113
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -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

14
python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py
View File

@@ -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:

82
test/registered/unit/layers/test_nsa_cp_utils.py
View File

@@ -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):

25
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
View File

@@ -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 (

195
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
View File

@@ -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):