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

Preserve FP8 CP shared-KV page contracts

Browse Source 
NSA FP8 CP shared-KV reuse must operate on packed page-slot rows, not bf16 compact rows. The change keeps current-only and partial-current reuse inside the page-aligned materialization contract, fails fast for non-page-aligned CP split inputs, and prevents FP8 FlashMLA-KV prefill from reaching incompatible in-seq CP metadata.

Constraint: NSA FP8 persistent MLA KV rows are packed 656-byte records and CP shared KV cache management is page-granular.\nConstraint: FlashMLA-KV prefill metadata is not CP-local after NSA in-seq splitting.\nRejected: Silently splice bf16 current rows into FP8 materialized cache | corrupts the packed cache layout.\nRejected: Keep FP8 FlashMLA-KV auto prefill under NSA CP | reaches num_splits shape errors after q-row splitting.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable FP8 FlashMLA-KV prefill for NSA in-seq CP until metadata is rebuilt after CP splitting or made CP-local.\nTested: Local git diff --check and py_compile for touched SGLang files.\nTested: Remote g0034 related unit sweep recorded in docs: test_nsa_cp_utils.py, test_cp_shared_kv_layout.py, test_cp_shared_kv_runtime.py, test_cp_hicache_metadata.py passed.\nNot-tested: Full FP8 ETE startup and performance run after this commit.
This commit is contained in:
laoyao0822
2026-06-01 03:33:44 +08:00
parent 46be97adc0
commit 6ef4face89
7 changed files with 810 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md
+304
View File
@@ -4290,3 +4290,307 @@ Risk / not yet covered:
capability, this can be tightened to fail-fast for the production GLM5 config.
- Remote CUDA verification is still required after syncing both `tai-kernel` and
`sglang-dev`; local CPU import is blocked by missing optional dependencies.
### C103 — 2026-06-01 FP8 MLA current reuse must splice packed cache rows, not bf16 compact rows
Problem found:
- NSA FP8 KV cache stores MLA rows in the packed persistent layout:
`[512 FP8 nope bytes][4 FP32 scales][64 BF16 rope values]`, i.e. `656` bytes
per token row.
- The CP shared-KV partial-current MLA path previously built fresh current rows
with `cat([k_nope, k_rope])`, i.e. the direct bf16/fp16 `576`-wide layout.
That is only valid for non-FP8 cache. In FP8 mode it would splice rows whose
dtype/row-width do not match the materialized prefix buffer.
Implementation direction:
- `tai-kernel` now provides `pack_quant_mla_kv(k_nope, k_rope) -> uint8[N,1,656]`
by reusing the fused quantized MLA store kernel with identity logical locs.
This is correctness-first and keeps the quantization/store contract in TAI
rather than rebuilding it with torch ops in SGLang.
- SGLang detects packed FP8 MLA KV cache by row width `656` plus `uint8` or
`float8_e4m3*` dtype, then packs current rows before current-only or
partial-current reuse.
- `fp8_e5m2` remains fail-fast for NSA CP shared-KV FP8 reuse; current model
config only supports `fp8_e4m3` for NSA.
- No silent bf16 fallback is allowed for FP8 current reuse. Missing TAI pack
capability raises `[CP_SHARED_KV_FAIL_FAST][fp8_current_pack]`.
IPC coverage:
- The descriptor-driven IPC materialize kernel is byte-count based. It already
covers FP8 MLA pages when `page_nbytes = page_size * 656`, so no separate IPC
copy kernel is required for packed FP8 rows.
- A CUDA test was added to lock the FP8 packed MLA page size through the IPC
slot-dense materialize path.
Verification status:
- Local and remote `py_compile` passed for touched SGLang and TAI files.
- Local source-level backend hook test passed. Local dynamic helper tests remain
blocked by missing local `orjson`; the remote container is authoritative for
those tests.
- Remote `g0034` SGLang targeted tests passed for FP8 current-pack success,
fail-fast-on-missing-TAI, and backend current-reuse hook detection.
- Remote `g0034` TAI CUDA tests passed for `pack_quant_mla_kv` identity layout
and FP8 packed MLA page IPC slot-dense materialize.
### C104 — 2026-06-01 current-only MLA reuse must still return page-slot KV buffers
Problem found:
- FP8 current-only CP shared-KV reuse packed the freshly computed current rows
into a compact tensor shaped like `[valid_tokens, 1, 656]` and assigned it
directly to `kv_cache`.
- `flashmla_kv` consumes MLA KV cache as pages and reshapes it with
`view(-1, page_size, 1, kv_cache_dim)`. A short current-only request such as
4 tokens therefore produced `4 * 656 = 2624` elements, which cannot be viewed
as full 64-token pages.
- This violated the page-aligned cache contract. FP8 does not get an exception:
even packed FP8 MLA rows must be stored and exposed through page-sized slots.
Implementation direction:
- The current-only branch now uses the same page-slot compose helper as
partial-current reuse, with `prefix_pages=0`.
- `prefix_pages=0` means there are no historical prefix pages to materialize or
reduce. It does **not** mean the output buffer has zero pages. The current
suffix still occupies `ceil(valid_tokens / page_size)` physical page slots.
- The compact current tensor remains only an input to the slot-fill step. The
attention-facing `kv_cache` is the composed page buffer, and returned locs are
remapped into that page buffer with non-current tail slack masked out.
Verification status:
- Local source-level regression test passed for the backend current-only branch:
it now requires `materialize_prefix_and_reuse_current_kv_page_slots`,
`prefix_pages=0`, and forbids `kv_cache = current_kv_cache`.
- Local `py_compile` passed for `nsa_backend.py` and `cp_shared_kv_runtime.py`.
- Remote `g0034` container `py_compile` passed for `nsa_backend.py` and
`cp_shared_kv_runtime.py`.
- Remote `g0034` targeted runtime tests passed for current-only page-slot
layout, padded-current reuse gating, FP8 current-pack success, and
fail-fast-on-missing-TAI.
- Local dynamic runtime helper tests are still blocked by missing optional local
dependency `orjson`; the remote container remains authoritative for those
tests.
### C105 — 2026-06-01 current-only tiny suffixes must also skip NSA in-seq CP split
Problem found:
- The short-page CP-split guard only handled radix-hit suffixes with a positive
page-aligned prefix. Current-only FP8 warmup/first-token paths with
`prefix_len=0` and `extend_len=64` still entered NSA in-seq CP split.
- With `cp_size=8` and `page_size=64`, that is only one physical page. The
page-aligned in-seq split produces mostly zero-token CP lanes. The downstream
FlashMLA KV path then sees inconsistent q-batch / metadata cardinality and the
kernel rejects `num_splits` with `num_splits must have shape (b+1)`.
Implementation direction:
- The page-unit guard now treats `prefix_len=0` as a valid page-aligned prefix
for the purpose of deciding whether CP split is safe.
- If `ceil(extend_len / page_size) < cp_size`, CP split is skipped for CP shared
KV regardless of whether the request is current-only or a radix hit.
- This preserves the page cache contract: we avoid splitting a single page over
eight owner lanes instead of falling back to token-balanced sub-page splits.
Verification status:
- Added a unit regression that `extend_len=64`, `prefix_len=0`, `page_size=64`,
`cp_size=8` does not use NSA in-seq CP split.
- Remote `g0034` container `py_compile` passed for `nsa/utils.py`.
- Remote `g0034` targeted tests passed for the new current-only skip guard and
adjacent short-page / one-page-per-rank CP-split behavior.
- Remote `g0034` combined targeted tests passed for the current-only page-slot
FP8 reuse checks plus the new CP-split guard.
### C106 — 2026-06-01 effective NSA impl must drive metadata and top-k layout
Finding:
- `forward_extend()` selects the effective NSA implementation from forward mode:
target-verify and draft-extend use `nsa_decode_impl`, while normal extend uses
`nsa_prefill_impl`.
- Two metadata/layout decisions still used global fields instead of that effective
implementation:
- `flashmla_metadata` construction in normal `init_forward_metadata()` was gated
only by `nsa_decode_impl == "flashmla_kv"`.
- `get_topk_transform_method()` selected ragged-vs-paged from
`nsa_prefill_impl` only.
Impact:
- The default Hopper FP8 setup (`prefill=flashmla_auto -> flashmla_kv`,
`decode=flashmla_kv`) did not expose this, but explicit backend combinations
could mismatch: e.g. decode `flashmla_kv` with prefill `flashmla_sparse` could
make target/draft use FlashMLA-KV attention with ragged top-k metadata.
- This is the same class of inconsistency as the recent FP8 failures: attention
kernel selection, top-k/index layout, and FlashMLA scheduler metadata must be
derived from one effective implementation for the current forward mode.
Correction:
- Added `_effective_nsa_impl_for_forward_batch()` and routed normal metadata
creation, forward-layer top-k transform selection, and indexer metadata through
it.
- CUDA-graph decode metadata remains decode-backend gated because that path is
decode/verify/draft graph state, not normal prefill metadata.
Verification:
- Remote g0034 container `py_compile` passed for `nsa_backend.py` and the runtime
test file.
- Remote targeted tests passed for effective top-k transform, FlashMLA metadata
source gating, and current-only page-slot reuse.
### C107 — 2026-06-01 non-page-aligned CP shared-KV prefixes must fail before CP split
Finding:
- Production radix/HiCache prefix matching is supposed to floor CP shared-KV
cache hits to page boundaries before scheduling. Evidence:
- base radix match floors keys to page size;
- HiRadix CP paths floor exact valid-tail hits and backed partial-tail splits;
- scheduler runtime checks require `cache_protected_len % page_size == 0`.
- The CP split builder still has a legacy token-balanced fallback when called
with a non-page-aligned `extend_prefix_len`. If production ever reaches that
branch, CP split would reintroduce sub-page split semantics instead of exposing
the upstream page-flooring bug.
Correction:
- `can_cp_split()` now fails fast for single-request CP shared-KV in-seq split
when `extend_prefix_lens_cpu[0]` is not page-aligned.
- This does not change the normal current-only path (`prefix_len=0`) or valid
page-aligned radix-hit path. It only makes an invalid production state explicit.
Verification status:
- Added unit coverage that a non-page-aligned CP shared-KV prefix raises
`[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix]` before CP split.
- Remote g0034 targeted test passed for non-page-aligned prefix fail-fast plus adjacent current-only and page-aligned CP-split cases.
### C108 — 2026-06-01 current-only index reuse tests must model the real CPU-length contract
Finding:
- A wider remote unit run exposed a stale test stub for NSA index current-only
reuse. The production gate is `is_current_only_extend_batch()`, which requires
`forward_mode.is_extend_without_speculative()`, zero prefix lengths, and
`seq_lens_cpu == extend_seq_lens_cpu`.
- The test passed `current_index_kv` but omitted those CPU metadata fields, so
the production helper correctly treated it as not current-only and attempted
shared index materialization.
Correction:
- The test stub now carries the same CPU-length contract as production current-only
batches. This keeps the assertion meaningful: current-only index reuse skips
shared index materialization only when the batch metadata proves that all
visible rows are current rows.
Verification:
- Remote g0034 container two-file unit suite passed:
`test_nsa_cp_utils.py` + `test_cp_shared_kv_runtime.py`
(`124 passed, 5 warnings, 2 subtests passed`).
- Remote g0034 container page-aligned/cache runtime suite passed:
`test_nsa_cp_utils.py` + `test_cp_shared_kv_layout.py` +
`test_cp_shared_kv_runtime.py`
(`149 passed, 5 warnings, 2 subtests passed`).
### C109 — 2026-06-01 FP8 HiCache target/draft budget split uses packed row bytes
Finding:
- CP HiCache shares `--hicache-size` across target and EAGLE/NextN draft when a
draft KV pool is attached.
- The split is byte-based but chooses a target token capacity first:
`target_pages = budget / ((target_bytes_per_token + draft_bytes_per_token) * page_size)`.
The remaining bytes go to draft, and draft capacity must be at least target
capacity.
- For NSA FP8, `NSATokenToKVPool.kv_cache_dim` is the packed persistent row width
(`656` uint8 bytes/token/layer), not the BF16 compact MLA width. HiCache host
pool construction passes the device `kv_cache_dim` as `override_kv_cache_dim`,
so `NSATokenToKVPoolHost.get_size_per_token()` and
`_estimate_hicache_size_per_token()` agree.
- GLM-5 EAGLE draft has one executable draft layer. Therefore under FP8 the
draft host allocation can be small in bytes (for example around 1/78 of the
target KV bytes) while still having equal or slightly larger token capacity.
A small draft GB log is not by itself a capacity bug.
Correction:
- No production code change was needed for the current target/draft HiCache size
split. Added regression tests to pin the FP8 packed-row byte contract and the
one-layer draft capacity implication.
Verification status:
- Remote g0034 container targeted FP8 HiCache budget tests passed
(`4 passed, 3 warnings`).
- Remote g0034 container full HiCache metadata unit file passed
(`102 passed, 5 warnings`).
### C110 — 2026-06-01 FP8 FlashMLA-KV prefill is incompatible with NSA in-seq CP split metadata
Finding:
- Remote g0034 prefill log `sglang_cp_hicache_20260531_183930.log` died on all
CP ranks with `RuntimeError: num_splits must have shape (b+1)` inside
`flash_mla_with_kvcache()`.
- Startup was Hopper FP8 KV cache with `nsa_prefill_backend=flashmla_auto`,
`nsa_decode_backend=flashmla_kv`, `--enable-nsa-prefill-context-parallel`, and
`--nsa-prefill-cp-mode in-seq-split`.
- The failing request was a normal target prefill/extend with about 40k input
tokens and no cache prefix. Auto dispatch selected `flashmla_kv` for FP8
prefill on Hopper.
- In in-seq CP split, the model input hidden states are split by
`prepare_input_dp_with_cp_dsa()`/`cp_split_and_rebuild_data()`, but normal NSA
metadata still builds FlashMLA-KV scheduler metadata from the unsplit full
`seqlens_expanded`. `_forward_flashmla_kv()` then reshapes local CP q rows as
`q=[local_q, 1, heads, dim]` while passing `num_splits` sized for the full q
row count, violating the FlashMLA-KV `(b+1)` contract.
- BF16 did not hit this because the default prefill implementation is
`flashmla_sparse`, which does not use FlashMLA-KV `num_splits` metadata.
Correction direction:
- Treat FP8 FlashMLA-KV as decode/verify/draft-only for NSA prefill CP until the
in-seq CP metadata is made FlashMLA-KV-aware.
- Auto FP8 prefill under NSA CP should dispatch to `flashmla_sparse` instead of
`flashmla_kv`; explicit invalid `flashmla_kv` prefill under active in-seq CP
should fail loudly rather than reaching the kernel shape error.
Verification target:
- Add unit coverage for FP8 `flashmla_auto` + NSA prefill CP selecting
`flashmla_sparse`.
- Sync to g0034 and rerun the remote unit slice before ETE restart.
C110 implementation note:
- `flashmla_auto` now selects `flashmla_sparse` for FP8 normal prefill whenever
NSA prefill CP is enabled and the forward mode is context-parallel extend.
This avoids constructing full-q FlashMLA-KV scheduler metadata before the
model's in-seq CP split narrows q rows.
- If an explicit configuration still routes an actual CP-split FP8 prefill layer
into `flashmla_kv`, `forward_extend()` raises
`[NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED]` before calling the kernel.
This is intentional: the alternative is the opaque kernel error
`num_splits must have shape (b+1)`.
C110 verification update:
- Remote g0034 container `py_compile` passed for `nsa_backend.py` and
`test_cp_shared_kv_runtime.py`.
- Remote targeted tests passed for FP8 auto prefill CP dispatch, explicit
FlashMLA-KV CP prefill fail-fast, effective top-k transform selection, and
FlashMLA metadata gating (`4 passed, 5 warnings`).
- Remote related unit sweep passed:
`test_nsa_cp_utils.py`, `test_cp_shared_kv_layout.py`,
`test_cp_shared_kv_runtime.py`, and `test_cp_hicache_metadata.py`
(`253 passed, 5 warnings, 2 subtests passed`).
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
+87
View File
@@ -477,6 +477,93 @@ def _load_tai_fused_mla_store_kernel():
return None
@lru_cache(maxsize=1)
def _load_tai_pack_quant_mla_kv_kernel():
try:
from tai_kernel.nsa_prefill import pack_quant_mla_kv
return pack_quant_mla_kv
except Exception as exc:
logger.warning(
"[CP_SHARED_KV_FAIL_FAST][fp8_current_pack] "
"CP shared KV FP8 current reuse requires tai-kernel "
"pack_quant_mla_kv; import failed. error=%s",
exc,
)
return None
def _float8_e4m3_dtypes() -> tuple[torch.dtype, ...]:
dtypes = []
for name in ("float8_e4m3fn", "float8_e4m3fnuz"):
dtype = getattr(torch, name, None)
if dtype is not None:
dtypes.append(dtype)
return tuple(dtypes)
def is_packed_fp8_mla_kv_cache(kv_cache: torch.Tensor) -> bool:
"""Return whether ``kv_cache`` uses the packed NSA FP8 MLA row layout."""
return (
kv_cache.ndim >= 2
and int(kv_cache.shape[-1]) == 656
and kv_cache.dtype in (torch.uint8, *_float8_e4m3_dtypes())
)
def pack_current_mla_kv_for_reuse(
k_nope: torch.Tensor,
k_rope: torch.Tensor,
*,
kv_cache: torch.Tensor,
) -> torch.Tensor:
"""Pack fresh current MLA rows to match a packed FP8 persistent KV cache.
CP shared-KV partial-current reuse splices freshly computed current suffix
rows into a materialized prefix buffer. For NSA FP8 KV cache, persistent
rows are stored as packed 656-byte records rather than the direct bf16
``[k_nope, k_rope]`` 576-wide layout. The splice source must therefore be
packed before it is copied into slot-dense pages.
"""
fp8_e5m2 = getattr(torch, "float8_e5m2", None)
if fp8_e5m2 is not None and kv_cache.dtype == fp8_e5m2:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][fp8_current_pack] "
"NSA CP shared KV supports fp8_e4m3 packed MLA KV only; "
"fp8_e5m2 current reuse is unsupported."
)
if not is_packed_fp8_mla_kv_cache(kv_cache):
raise ValueError(
"pack_current_mla_kv_for_reuse requires packed NSA FP8 MLA kv_cache "
f"with row width 656, got shape={tuple(kv_cache.shape)} dtype={kv_cache.dtype}"
)
kernel = _load_tai_pack_quant_mla_kv_kernel()
if kernel is None:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][fp8_current_pack] "
"TAI pack_quant_mla_kv is required for NSA FP8 CP shared-KV "
"partial-current reuse. Install/sync tai-kernel with the FP8 "
"current-pack kernel instead of falling back to bf16 compact rows."
)
packed_u8 = kernel(k_nope, k_rope)
expected_shape = (int(k_nope.shape[0]), 1, 656)
if packed_u8.dtype != torch.uint8 or tuple(packed_u8.shape) != expected_shape:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][fp8_current_pack] "
"TAI pack_quant_mla_kv returned an invalid layout. "
f"expected_shape={expected_shape} expected_dtype=torch.uint8 "
f"actual_shape={tuple(packed_u8.shape)} actual_dtype={packed_u8.dtype}"
)
if kv_cache.dtype == torch.uint8:
return packed_u8
return packed_u8.view(kv_cache.dtype)
@lru_cache(maxsize=1)
def _load_tai_index_mqa_prepare_kernel():
try:
python/sglang/srt/layers/attention/nsa/utils.py
+10 -1
View File
@@ -441,8 +441,17 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent(
return False
prefix_len = int(extend_prefix_lens_cpu[0])
if prefix_len <= 0 or prefix_len % page_size != 0:
if prefix_len < 0:
return False
if prefix_len % page_size != 0:
raise RuntimeError(
"[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] "
"CP shared KV NSA in-seq split requires a page-aligned prefix. "
"The radix/HiCache match path should floor cache hits to the "
"previous page boundary before CP split planning. "
f"prefix_len={prefix_len} extend_len={extend_len} "
f"page_size={page_size} cp_size={cp_size}"
)
padded_pages = ceil_div(extend_len, page_size)
return padded_pages < cp_size
python/sglang/srt/layers/attention/nsa_backend.py
+91 -17
View File
@@ -28,8 +28,10 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
filter_owned_logical_locs,
get_or_build_shared_token_kv_slot_remap,
is_current_only_extend_batch,
is_packed_fp8_mla_kv_cache,
materialize_prefix_and_reuse_current_kv_page_slots,
materialize_shared_token_kv_buffer,
pack_current_mla_kv_for_reuse,
tensor_debug_checksum,
tensor_debug_summary,
)
@@ -51,6 +53,7 @@ from sglang.srt.layers.attention.nsa.utils import (
is_nsa_enable_prefill_cp,
nsa_cp_round_robin_split_data,
nsa_cp_round_robin_split_q_seqs,
nsa_use_prefill_cp,
pad_nsa_cache_seqlens,
)
from sglang.srt.layers.attention.utils import (
@@ -670,7 +673,7 @@ class NativeSparseAttnBackend(
# Centralized dispatch: decide all strategies for this batch
self.set_nsa_prefill_impl(forward_batch)
topk_transform_method = self.get_topk_transform_method()
topk_transform_method = self.get_topk_transform_method(forward_batch)
# Batch indices selected when cp enabled: After splitting multiple sequences,
# a certain cp rank may not have some of these sequences.
# We use bs_idx_cpu to mark which sequences are finally selected by the current cp rank,
@@ -904,7 +907,8 @@ class NativeSparseAttnBackend(
cache_seqlens=nsa_cache_seqlens_int32,
seq_len_q=1,
)
if self.nsa_decode_impl == "flashmla_kv"
if self._effective_nsa_impl_for_forward_batch(forward_batch)
== "flashmla_kv"
else None
),
paged_mqa_schedule_metadata=paged_mqa_schedule_metadata,
@@ -1700,7 +1704,7 @@ class NativeSparseAttnBackend(
topk_indices = self._pad_topk_indices(topk_indices, q_nope.shape[0])
# NOTE(dark): here, we use page size = 1
topk_transform_method = self.get_topk_transform_method()
topk_transform_method = self.get_topk_transform_method(forward_batch)
if envs.SGLANG_NSA_FUSE_TOPK.get():
page_table_1 = topk_indices
else:
@@ -1827,9 +1831,16 @@ class NativeSparseAttnBackend(
assert k is not None and k_rope is not None
assert current_kv_rows_for_reuse is not None
valid_current_rows = int(current_kv_rows_for_reuse)
current_kv_cache = _cat(
[k[:valid_current_rows], k_rope[:valid_current_rows]], dim=-1
)
current_k_nope = k[:valid_current_rows]
current_k_rope = k_rope[:valid_current_rows]
if is_packed_fp8_mla_kv_cache(kv_cache):
current_kv_cache = pack_current_mla_kv_for_reuse(
current_k_nope,
current_k_rope,
kv_cache=kv_cache,
)
else:
current_kv_cache = _cat([current_k_nope, current_k_rope], dim=-1)
current_locs_for_reuse = forward_batch.out_cache_loc[
:valid_current_rows
]
@@ -1840,13 +1851,13 @@ class NativeSparseAttnBackend(
if is_current_only_extend_batch(forward_batch):
eagle_draft_mla_branch = "current_only"
current_mask, page_table_1 = build_current_loc_remap(
logical_page_table_1,
forward_batch.out_cache_loc,
page_size=current_remap_page_size,
logical_page_capacity=current_remap_logical_page_capacity,
)
if cp_shared_kv_debug_enabled():
current_mask, compact_current_rows = build_current_loc_remap(
logical_page_table_1,
forward_batch.out_cache_loc,
page_size=current_remap_page_size,
logical_page_capacity=current_remap_logical_page_capacity,
)
missing_current = (logical_page_table_1 >= 0) & (~current_mask)
if torch.any(missing_current):
bad_locs = logical_page_table_1[missing_current]
@@ -1862,11 +1873,32 @@ class NativeSparseAttnBackend(
forward_batch.cp_shared_kv_layout.cp_rank,
layer.layer_id,
tensor_debug_summary(current_locs_for_reuse),
tensor_debug_summary(page_table_1),
tensor_debug_summary(compact_current_rows),
tensor_debug_checksum(k),
tensor_debug_checksum(k_rope),
)
kv_cache = current_kv_cache
page_size = int(forward_batch.token_to_kv_pool.page_size)
slot_remap = get_or_build_shared_token_kv_slot_remap(
forward_batch,
kv_cache=kv_cache,
remap_logical_pages=metadata.real_page_table,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
)
kv_cache, page_table_1 = (
materialize_prefix_and_reuse_current_kv_page_slots(
kv_cache=kv_cache,
logical_locs=logical_page_table_1,
current_kv_cache=current_kv_cache,
current_locs=current_locs_for_reuse,
slot_remap=slot_remap,
layout=forward_batch.cp_shared_kv_layout,
page_size=page_size,
prefix_pages=0,
layer_id=layer.layer_id,
nvtx_source="mla.current_only_page_slots",
)
)
else:
extend_lens_cpu_for_current = getattr(
forward_batch, "extend_seq_lens_cpu", None
@@ -2165,6 +2197,19 @@ class NativeSparseAttnBackend(
v_head_dim=layer.v_head_dim,
)
elif nsa_impl == "flashmla_kv":
if (
self.nsa_kv_cache_store_fp8
and forward_batch.forward_mode.is_context_parallel_extend()
and nsa_use_prefill_cp(forward_batch)
):
raise RuntimeError(
"[NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED] "
"FP8 FlashMLA-KV prefill cannot be used after NSA "
"in-seq CP has split q rows without CP-local "
"FlashMLA metadata. Use flashmla_auto or "
"flashmla_sparse for nsa_prefill_backend under "
"NSA prefill CP."
)
if q_rope is not None:
q_all = concat_mla_absorb_q_general(q_nope, q_rope)
attn_output = self._forward_flashmla_kv(
@@ -2861,6 +2906,20 @@ class NativeSparseAttnBackend(
# Set MLA implementation only if not using MHA
if not self.use_mha and self.enable_auto_select_prefill_impl:
if self.nsa_kv_cache_store_fp8:
if (
forward_batch is not None
and forward_batch.forward_mode.is_context_parallel_extend()
and is_nsa_enable_prefill_cp()
):
# FlashMLA-KV metadata is decode-shaped: `num_splits` must
# match the q batch dimension after CP splitting. In NSA
# in-seq prefill CP the model hidden states are split later
# by `cp_split_and_rebuild_data()`, while this metadata is
# built before that split. Keep FP8 prefill CP on the
# sparse prefill kernel until FlashMLA-KV gets an explicit
# CP-local metadata contract.
self.nsa_prefill_impl = "flashmla_sparse"
return
if (
is_blackwell()
and forward_batch is not None
@@ -2877,15 +2936,30 @@ class NativeSparseAttnBackend(
# bf16 kv cache
self.nsa_prefill_impl = "flashmla_sparse"
def get_topk_transform_method(self) -> TopkTransformMethod:
def _effective_nsa_impl_for_forward_batch(
self, forward_batch: Optional[ForwardBatch] = None
) -> _NSA_IMPL_T:
forward_mode = getattr(forward_batch, "forward_mode", None)
if forward_mode is not None and (
forward_mode.is_decode_or_idle()
or forward_mode.is_target_verify()
or forward_mode.is_draft_extend(include_v2=True)
):
return self.nsa_decode_impl
return self.nsa_prefill_impl
def get_topk_transform_method(
self, forward_batch: Optional[ForwardBatch] = None
) -> TopkTransformMethod:
"""
SGLANG_NSA_FUSE_TOPK controls whether to fuse the topk transform into the topk kernel.
This method is used to select the topk transform method which can be fused or unfused.
"""
nsa_impl = self._effective_nsa_impl_for_forward_batch(forward_batch)
if (
# disable for MTP
self.nsa_kv_cache_store_fp8
and self.nsa_prefill_impl == "flashmla_sparse"
and nsa_impl == "flashmla_sparse"
):
topk_transform_method = TopkTransformMethod.RAGGED
else:
@@ -2901,7 +2975,7 @@ class NativeSparseAttnBackend(
)
return NSAIndexerMetadata(
attn_metadata=self.forward_metadata,
topk_transform_method=self.get_topk_transform_method(),
topk_transform_method=self.get_topk_transform_method(forward_batch),
paged_mqa_schedule_metadata=self.forward_metadata.paged_mqa_schedule_metadata,
force_unfused_topk=force_unfused,
validate_paged_topk=forward_batch.uses_cp_shared_kv,
test/registered/unit/layers/test_nsa_cp_utils.py
+72
View File
@@ -283,6 +283,62 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
):
self.assertFalse(can_cp_split(256, 8, True, forward_batch))
def test_can_cp_split_skips_current_only_when_page_units_do_not_cover_all_lanes(
self,
):
class Mode:
def is_context_parallel_extend(self):
return True
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_seq_lens_cpu=[64],
extend_prefix_lens_cpu=[0],
token_to_kv_pool=SimpleNamespace(page_size=64),
forward_mode=Mode(),
)
with (
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
return_value=True,
),
):
self.assertFalse(can_cp_split(64, 8, True, forward_batch))
def test_can_cp_split_fails_on_non_page_aligned_cp_shared_prefix(self):
class Mode:
def is_context_parallel_extend(self):
return True
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_seq_lens_cpu=[1024],
extend_prefix_lens_cpu=[65],
token_to_kv_pool=SimpleNamespace(page_size=64),
forward_mode=Mode(),
)
with (
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
return_value=True,
),
self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[cp_split_non_page_aligned_prefix\]",
),
):
can_cp_split(1089, 8, True, forward_batch)
def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self):
class Mode:
def is_context_parallel_extend(self):
@@ -653,10 +709,18 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
indexer._maybe_materialize_shared_index_buffer = fake_materialize
indexer._get_topk_ragged_with_cp = fake_get_topk
class Mode:
def is_extend_without_speculative(self):
return True
forward_batch = type(
"ForwardBatchStub",
(),
{
"forward_mode": Mode(),
"extend_prefix_lens_cpu": [0],
"extend_seq_lens_cpu": [5],
"seq_lens_cpu": torch.tensor([5], dtype=torch.int64),
"nsa_cp_metadata": NSAContextParallelMetadata(
kv_len_prev=5,
kv_len_next=9,
@@ -741,10 +805,18 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
indexer._maybe_materialize_shared_index_buffer = fake_materialize
indexer._get_topk_ragged_with_cp = fake_get_topk
class Mode:
def is_extend_without_speculative(self):
return True
forward_batch = type(
"ForwardBatchStub",
(),
{
"forward_mode": Mode(),
"extend_prefix_lens_cpu": [0],
"extend_seq_lens_cpu": [5],
"seq_lens_cpu": torch.tensor([5], dtype=torch.int64),
"nsa_cp_metadata": NSAContextParallelMetadata(
kv_len_prev=5,
kv_len_next=9,
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
+40
View File
@@ -113,6 +113,7 @@ from sglang.srt.mem_cache.hiradix_cache import (
PreparedCpHiCacheBackup,
_compute_shared_hicache_token_capacities,
)
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.radix_cache import RadixKey, TreeNode, _key_match_paged
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
from sglang.test.ci.ci_register import register_cpu_ci
@@ -191,6 +192,45 @@ class TestCpHiCacheImports(CustomTestCase):
class TestHiRadixCacheCPDraftHostPool(CustomTestCase):
def test_fp8_nsa_hicache_size_estimate_uses_packed_row_width(self):
pool = object.__new__(NSATokenToKVPool)
pool.store_dtype = torch.uint8
pool.kv_cache_dim = 656
pool.layer_num = 78
pool.index_head_dim = 128
pool.quant_block_size = 128
size_per_token = hiradix_cache._estimate_hicache_size_per_token(pool)
# FP8 NSA stores packed MLA rows as 656 uint8 bytes/token/layer plus
# the paged index buffer: 128 uint8 K bytes + 4 uint8 scale bytes.
self.assertEqual(size_per_token, (656 + 128 + 4) * 78)
def test_fp8_shared_budget_matches_target_and_one_layer_draft_capacity(self):
bytes_per_layer = 656 + 128 + 4
target_size_per_token = bytes_per_layer * 78
draft_size_per_token = bytes_per_layer
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
total_host_bytes=int(150 * 1e9),
target_size_per_token=target_size_per_token,
draft_size_per_token=draft_size_per_token,
page_size=64,
)
self.assertGreaterEqual(draft_tokens, target_tokens)
self.assertLessEqual(
target_tokens * target_size_per_token
+ draft_tokens * draft_size_per_token,
int(150 * 1e9),
)
# The draft pool is much smaller in bytes because GLM-5 EAGLE draft has
# one executable layer, but it still has at least target token capacity.
self.assertLess(
draft_tokens * draft_size_per_token,
target_tokens * target_size_per_token // 32,
)
def test_shared_budget_keeps_draft_at_least_target_capacity(self):
target_tokens, draft_tokens = _compute_shared_hicache_token_capacities(
total_host_bytes=1000,
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+206 -2
View File
@@ -672,10 +672,214 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
body_source = "".join(source[body_start:body_end].split())
self.assertIn("valid_current_rows=int(current_kv_rows_for_reuse)", body_source)
self.assertIn("k[:valid_current_rows]", body_source)
self.assertIn("k_rope[:valid_current_rows]", body_source)
self.assertIn("pack_current_mla_kv_for_reuse", body_source)
self.assertIn("forward_batch.out_cache_loc[:valid_current_rows]", body_source)
def test_flashmla_kv_current_only_reuse_keeps_page_slot_layout(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa_backend.py"
).read_text()
branch_start = source.index(" if is_current_only_extend_batch")
branch_end = source.index(" else:", branch_start)
branch_source = source[branch_start:branch_end]
self.assertIn(
"materialize_prefix_and_reuse_current_kv_page_slots", branch_source
)
self.assertIn("prefix_pages=0", branch_source)
self.assertNotIn("kv_cache = current_kv_cache", branch_source)
def test_nsa_backend_topk_transform_uses_effective_forward_impl(self):
from sglang.srt.layers.attention.nsa_backend import (
NativeSparseAttnBackend,
TopkTransformMethod,
)
class Mode:
def __init__(self, *, target=False, draft=False, decode=False):
self.target = target
self.draft = draft
self.decode = decode
def is_decode_or_idle(self):
return self.decode
def is_target_verify(self):
return self.target
def is_draft_extend(self, include_v2=False):
return self.draft
backend = NativeSparseAttnBackend.__new__(NativeSparseAttnBackend)
backend.nsa_kv_cache_store_fp8 = True
backend.nsa_prefill_impl = "flashmla_sparse"
backend.nsa_decode_impl = "flashmla_kv"
# Normal extend follows the prefill implementation and therefore needs
# ragged top-k for FP8 sparse prefill.
self.assertEqual(
backend.get_topk_transform_method(
SimpleNamespace(forward_mode=Mode())
),
TopkTransformMethod.RAGGED,
)
# Target verify / draft / decode follow the decode implementation. When
# decode is flashmla_kv, those paths need paged top-k even if prefill is
# currently flashmla_sparse.
for mode in (
Mode(target=True),
Mode(draft=True),
Mode(decode=True),
):
self.assertEqual(
backend.get_topk_transform_method(
SimpleNamespace(forward_mode=mode)
),
TopkTransformMethod.PAGED,
)
def test_flashmla_metadata_creation_uses_effective_forward_impl(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa_backend.py"
).read_text()
metadata_start = source.index(" flashmla_metadata=(")
metadata_end = source.index(" paged_mqa_schedule_metadata=", metadata_start)
metadata_source = source[metadata_start:metadata_end]
self.assertIn(
"_effective_nsa_impl_for_forward_batch(forward_batch)",
metadata_source,
)
self.assertNotIn(
'if self.nsa_decode_impl == "flashmla_kv"',
metadata_source,
)
def test_fp8_auto_prefill_cp_uses_sparse_not_flashmla_kv(self):
from sglang.srt.layers.attention.nsa_backend import NativeSparseAttnBackend
from sglang.srt.model_executor.forward_batch_info import ForwardMode
def make_backend():
backend = NativeSparseAttnBackend.__new__(NativeSparseAttnBackend)
backend.nsa_kv_cache_store_fp8 = True
backend.enable_auto_select_prefill_impl = True
backend.nsa_prefill_impl = "flashmla_auto"
return backend
forward_batch = SimpleNamespace(
forward_mode=ForwardMode.EXTEND,
seq_lens_cpu=torch.tensor([40392], dtype=torch.int32),
seq_lens_sum=40392,
extend_num_tokens=40392,
token_to_kv_pool=SimpleNamespace(dtype=torch.float8_e4m3fn),
hisparse_coordinator=None,
get_max_chunk_capacity=lambda: 65536,
)
with patch(
"sglang.srt.utils.get_device_sm", return_value=90
), patch(
"sglang.srt.utils.is_blackwell", return_value=False
), patch(
"sglang.srt.layers.attention.nsa_backend.is_nsa_enable_prefill_cp",
return_value=True,
):
backend = make_backend()
backend.set_nsa_prefill_impl(forward_batch)
self.assertEqual(backend.nsa_prefill_impl, "flashmla_sparse")
with patch(
"sglang.srt.utils.get_device_sm", return_value=90
), patch(
"sglang.srt.utils.is_blackwell", return_value=False
), patch(
"sglang.srt.layers.attention.nsa_backend.is_nsa_enable_prefill_cp",
return_value=False,
):
backend = make_backend()
backend.set_nsa_prefill_impl(forward_batch)
self.assertEqual(backend.nsa_prefill_impl, "flashmla_kv")
def test_explicit_fp8_flashmla_kv_cp_prefill_fails_before_kernel(self):
from pathlib import Path
source = (
Path(__file__).resolve().parents[4]
/ "python/sglang/srt/layers/attention/nsa_backend.py"
).read_text()
branch_start = source.index(' elif nsa_impl == "flashmla_kv":')
branch_end = source.index(' elif nsa_impl == "fa3":', branch_start)
branch_source = source[branch_start:branch_end]
self.assertIn("nsa_use_prefill_cp(forward_batch)", branch_source)
self.assertIn("NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED", branch_source)
self.assertLess(
branch_source.index("NSA_FP8_FLASHMLA_KV_CP_PREFILL_UNSUPPORTED"),
branch_source.index("self._forward_flashmla_kv("),
)
def test_pack_current_mla_kv_for_reuse_uses_tai_pack_for_fp8_cache(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
class FakeTaiPack:
def __init__(self):
self.calls = []
def __call__(self, k_nope, k_rope):
self.calls.append((k_nope, k_rope))
return torch.ones((k_nope.shape[0], 1, 656), dtype=torch.uint8)
fake_kernel = FakeTaiPack()
k_nope = torch.zeros((3, 1, 512), dtype=torch.bfloat16)
k_rope = torch.zeros((3, 1, 64), dtype=torch.bfloat16)
fp8_kv_cache = torch.empty((128, 1, 656), dtype=torch.float8_e4m3fn)
with patch.object(
runtime, "_load_tai_pack_quant_mla_kv_kernel", return_value=fake_kernel
):
packed = runtime.pack_current_mla_kv_for_reuse(
k_nope,
k_rope,
kv_cache=fp8_kv_cache,
)
self.assertEqual(len(fake_kernel.calls), 1)
self.assertIs(fake_kernel.calls[0][0], k_nope)
self.assertIs(fake_kernel.calls[0][1], k_rope)
self.assertEqual(tuple(packed.shape), (3, 1, 656))
self.assertEqual(packed.dtype, torch.float8_e4m3fn)
self.assertTrue(
torch.equal(
packed.view(torch.uint8),
torch.ones_like(packed.view(torch.uint8)),
)
)
def test_pack_current_mla_kv_for_reuse_fails_fast_when_fp8_pack_kernel_missing(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
k_nope = torch.zeros((1, 1, 512), dtype=torch.bfloat16)
k_rope = torch.zeros((1, 1, 64), dtype=torch.bfloat16)
fp8_kv_cache = torch.empty((128, 1, 656), dtype=torch.float8_e4m3fn)
with patch.object(
runtime, "_load_tai_pack_quant_mla_kv_kernel", return_value=None
), self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[fp8_current_pack\]",
):
runtime.pack_current_mla_kv_for_reuse(
k_nope,
k_rope,
kv_cache=fp8_kv_cache,
)
def test_runtime_fallback_helpers_use_standard_warning_marker(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime