diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 422624c2a..a84380782 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -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`). diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 2c531ee42..0e3a9c4e0 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -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: diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 29e33d703..81f005db3 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 54f5c6476..33e691adb 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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, diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 1028f1e94..24bb40e70 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 6dc364f24..630f9c180 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -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, diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 8af754fdd..d6067f481 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -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