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

87
python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py
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:

11
python/sglang/srt/layers/attention/nsa/utils.py
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

108
python/sglang/srt/layers/attention/nsa_backend.py
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,