diff --git a/python/fmha_sm100/api.py b/python/fmha_sm100/api.py index b02f747..8fd8d6c 100644 --- a/python/fmha_sm100/api.py +++ b/python/fmha_sm100/api.py @@ -765,13 +765,14 @@ def _fmha_sm100( output_maxscore: bool = True, output_o: bool = True, check_input_valid: bool = False, + page_table: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: if plan_info["MM-SA-Nv"]: return sparse_fmha(q=q, k=k, v=v, plan_info=plan_info, out=out, max_score=max_score, sm_scale=sm_scale, q_scale=q_scale, k_scale=k_scale, v_scale=v_scale, o_scale=o_scale, kv_indices=kv_indices, output_maxscore=output_maxscore, output_o=output_o, q_offset_override=q_offset_override, - kv_block_indexes=kv_block_indexes, check_input_valid=check_input_valid) + kv_block_indexes=kv_block_indexes, check_input_valid=check_input_valid, page_table=page_table) nnz_qo, num_qo_heads, head_dim_qk = q.shape @@ -1035,6 +1036,7 @@ def fmha_sm100( q_offset_override: Optional[Union[int, torch.Tensor]] = None, out: Optional[torch.Tensor] = None, max_score: Optional[torch.Tensor] = None, + page_table: Optional[torch.Tensor] = None, **kwargs ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Run dense, paged, or sparse SM100 FMHA using a precomputed plan. @@ -1054,9 +1056,15 @@ def fmha_sm100( Return value from ``fmha_sm100_plan`` for the same lengths, head layout, page size, and sparse/output mode. kv_indices : torch.Tensor, optional - Paged-KV physical page table, flattened across the batch. Required when - ``k`` and ``v`` use paged layout. Shape is ``[sum_pages]`` and dtype is + Paged-KV physical page table, flattened across the batch. Required for + generic paged paths. Shape is ``[sum_pages]`` and dtype is int32. + page_table : torch.Tensor, optional + Prebuilt two-dimensional physical page table with shape + ``[batch_size, max_pages_per_sequence]`` and dtype ``torch.int32``. + It must be on the same device as paged K/V. Sparse prefill uses it + directly instead of rebuilding it from ``kv_indices``; generic + sub-plans still require ``kv_indices``. kv_block_indexes : torch.Tensor, optional Sparse KV block indices from ``sparse_topk_select``. Shape ``[total_qo_len, num_kv_heads or num_qo_heads, kv_block_num]``, dtype @@ -1085,28 +1093,38 @@ def fmha_sm100( outputs are concatenated back into the original batch order. """ has_mixed_prefill, split, batch_size, decode, prefill = plan_info + if page_table is not None: + subplans = (decode, prefill) if has_mixed_prefill else (decode,) + if not any(plan["MM-SA-Nv"] for plan in subplans): + raise ValueError("page_table is only supported by sparse prefill") + if not has_mixed_prefill: - return _fmha_sm100(q, k, v, decode, out=out, max_score=max_score, kv_indices=kv_indices,kv_block_indexes=kv_block_indexes, q_offset_override=q_offset_override, **kwargs) + return _fmha_sm100(q, k, v, decode, out=out, max_score=max_score, kv_indices=kv_indices, page_table=page_table, kv_block_indexes=kv_block_indexes, q_offset_override=q_offset_override, **kwargs) else: decode_pack = decode.get("pack_factor", 1) decode_nnz = decode["qo_segment_offsets"][-1].item() // decode_pack - is_paged = kv_indices is not None + is_paged = kv_indices is not None or page_table is not None nnz_qo = q.shape[0] num_qo_heads = q.shape[1] q_decode = q[:decode_nnz] q_prefill = q[decode_nnz:] + decode_page_table = page_table[:split] if page_table is not None and decode["MM-SA-Nv"] else None + prefill_page_table = page_table[split:] if page_table is not None and prefill["MM-SA-Nv"] else None if is_paged: k_decode, v_decode = k, v k_prefill, v_prefill = k, v - if "kv_page_indptr" in decode: - kv_page_split = decode["kv_page_indptr"][-1].item() - else: - kv_page_split = decode["total_rows"] - decode_kv_indices = kv_indices[:kv_page_split] - prefill_kv_indices = kv_indices[kv_page_split:] + decode_kv_indices = None + prefill_kv_indices = None + if kv_indices is not None: + if "kv_page_indptr" in decode: + kv_page_split = decode["kv_page_indptr"][-1].item() + else: + kv_page_split = decode["total_rows"] + decode_kv_indices = kv_indices[:kv_page_split] + prefill_kv_indices = kv_indices[kv_page_split:] else: if "kv_segment_offsets" in decode: decode_kv_nnz = decode["kv_segment_offsets"][-1].item() @@ -1133,13 +1151,13 @@ def fmha_sm100( decode_out, decode_ms = _fmha_sm100( q_decode, k_decode, v_decode, decode, out=None, max_score=None, - kv_indices=decode_kv_indices, kv_block_indexes=decode_block_idx, + kv_indices=decode_kv_indices, page_table=decode_page_table, kv_block_indexes=decode_block_idx, q_offset_override=decode_qo_offset, **kwargs) prefill_out, prefill_ms = _fmha_sm100( q_prefill, k_prefill, v_prefill, prefill, out=None, max_score=None, - kv_indices=prefill_kv_indices, kv_block_indexes=prefill_block_idx, + kv_indices=prefill_kv_indices, page_table=prefill_page_table, kv_block_indexes=prefill_block_idx, q_offset_override=prefill_qo_offset, **kwargs) diff --git a/python/fmha_sm100/cute/interface.py b/python/fmha_sm100/cute/interface.py index d72b17a..dc505bd 100644 --- a/python/fmha_sm100/cute/interface.py +++ b/python/fmha_sm100/cute/interface.py @@ -52,6 +52,8 @@ _SUPPORTED_FWD_DTYPES = (torch.bfloat16, torch.float8_e4m3fn) _SUPPORTED_FWD_MMA_DTYPES = (torch.bfloat16, torch.float8_e4m3fn) _SUPPORTED_DECODE_QHEAD_PER_KV = 16 +# Row slices may start at any int32 boundary. +_PAGE_TABLE_ASSUMED_ALIGN_BYTES = 4 def _normalize_partial_dtype(partial_dtype: torch.dtype) -> torch.dtype: @@ -154,6 +156,31 @@ def _validate_cu_seqlens( raise ValueError(f"{name} must be contiguous") +def _validate_page_table( + page_table: torch.Tensor, + *, + device: torch.device, + batch: int, + page_size: int, + max_seqlen_k: int, +) -> None: + if page_table.device != device: + raise ValueError("page_table must be on the same device as q") + if page_table.dtype != torch.int32: + raise TypeError("page_table must be torch.int32") + if page_table.ndim != 2 or page_table.shape[0] != batch: + raise ValueError("page_table must have shape [B, max_num_pages_per_seq]") + if page_table.stride(-1) != 1: + raise ValueError("page_table must be contiguous in the last dimension") + + required_pages = (int(max_seqlen_k) + page_size - 1) // page_size + if page_table.shape[1] < required_pages: + raise ValueError( + f"page_table has {page_table.shape[1]} columns, " + f"but max_seqlen_k={max_seqlen_k} requires {required_pages}" + ) + + def _csr_row_capacity(k2q_row_ptr: torch.Tensor) -> int: return int(k2q_row_ptr.shape[1] - 1) @@ -170,6 +197,7 @@ def _validate_csr_varlen_inputs( cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, seqused_k: Optional[torch.Tensor], + max_seqlen_k: int, ) -> tuple[int, int]: if q.ndim != 3: raise ValueError("CSR sparse forward requires q to have shape [total_q, Hq, D]") @@ -257,14 +285,6 @@ def _validate_csr_varlen_inputs( if k.shape != (total_k, head_kv, q.shape[-1]) or v.shape != (total_k, head_kv, q.shape[-1]): raise ValueError("Sparse Attention k and v must match [total_k, Hkv, D]") else: - if page_table.device != q.device: - raise ValueError("page_table must be on the same device as q") - if page_table.dtype != torch.int32: - raise TypeError("page_table must be torch.int32") - if page_table.ndim != 2 or page_table.shape[0] != batch: - raise ValueError("page_table must have shape [B, max_num_pages_per_seq]") - if page_table.stride(-1) != 1: - raise ValueError("page_table must be contiguous in the last dimension") if k.ndim != 4 or v.ndim != 4: raise ValueError( "Sparse Page Attention requires k and v to have shape " @@ -292,6 +312,13 @@ def _validate_csr_varlen_inputs( raise ValueError("seqused_k must have shape [B]") if not seqused_k.is_contiguous(): raise ValueError("seqused_k must be contiguous") + _validate_page_table( + page_table, + device=q.device, + batch=batch, + page_size=page_size, + max_seqlen_k=max_seqlen_k, + ) if topK not in _SUPPORTED_SPARSE_TOPK: raise ValueError( f"CSR sparse forward supports topK in {_SUPPORTED_SPARSE_TOPK}, got {topK}" @@ -315,6 +342,7 @@ def _validate_csr_varlen_nvfp4_kv_inputs( cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, seqused_k: Optional[torch.Tensor], + max_seqlen_k: int, ) -> tuple[int, int]: if q.ndim != 3: raise ValueError("KVFP4 CSR sparse forward requires q to have shape [total_q, Hq, D]") @@ -389,14 +417,6 @@ def _validate_csr_varlen_nvfp4_kv_inputs( ) head_kv = int(k.shape[1]) required_scale_rows = int(k.shape[0]) * head_kv * page_size - if page_table.device != q.device: - raise ValueError("page_table must be on the same device as q") - if page_table.dtype != torch.int32: - raise TypeError("page_table must be torch.int32") - if page_table.ndim != 2: - raise ValueError("page_table must have shape [B, max_num_pages_per_seq]") - if page_table.stride(-1) != 1: - raise ValueError("page_table must be contiguous in the last dimension") if seqused_k is not None: if seqused_k.device != q.device: raise ValueError("seqused_k must be on the same device as q") @@ -425,10 +445,16 @@ def _validate_csr_varlen_nvfp4_kv_inputs( if cu_seqlens_k.shape != cu_seqlens_q.shape: raise ValueError("cu_seqlens_k must have shape [B + 1] matching cu_seqlens_q") batch = int(cu_seqlens_q.shape[0] - 1) - if page_table is not None and page_table.shape[0] != batch: - raise ValueError("page_table must have shape [B, max_num_pages_per_seq]") if seqused_k is not None and seqused_k.shape != (batch,): raise ValueError("seqused_k must have shape [B]") + if page_table is not None: + _validate_page_table( + page_table, + device=q.device, + batch=batch, + page_size=page_size, + max_seqlen_k=max_seqlen_k, + ) head_q = int(q.shape[1]) if head_q % head_kv != 0: raise ValueError("q.shape[1] must be divisible by Hkv") @@ -732,6 +758,7 @@ def sparse_atten_func( cu_seqlens_q, cu_seqlens_k, seqused_k, + max_seqlen_k, ) max_seqlen_q = int(max_seqlen_q) max_seqlen_k = int(max_seqlen_k) @@ -890,6 +917,7 @@ def sparse_atten_nvfp4_kv_func( cu_seqlens_q, cu_seqlens_k, seqused_k, + max_seqlen_k, ) total_q, head_q, dim = q.shape max_num_kv_blocks = _csr_row_capacity(k2q_row_ptr) @@ -1724,6 +1752,7 @@ def _call_sparse_forward_sm100_csr_varlen( partial_dtype, bool(causal), bool(paged_kv), + _PAGE_TABLE_ASSUMED_ALIGN_BYTES if paged_kv else None, bool(use_prepare_scheduler), page_size, bool(seqused_k is not None), @@ -1764,7 +1793,7 @@ def _call_sparse_forward_sm100_csr_varlen( else to_cute_tensor_kvouter(LSE_temperature_partial), to_cute_tensor_kvouter(Q_flat), None if Q_gather4_desc is None else to_cute_tensor_kvouter(Q_gather4_desc), - None if page_table is None else to_cute_tensor_kvouter(page_table), + None if page_table is None else to_cute_tensor_kvouter(page_table, assumed_align=_PAGE_TABLE_ASSUMED_ALIGN_BYTES), None if seqused_k is None else to_cute_tensor_kvouter(seqused_k), to_cute_tensor_kvouter(cu_seqlens_q), to_cute_tensor_kvouter(cu_seqlens_k), @@ -1916,6 +1945,7 @@ def _call_sparse_forward_sm100_csr_varlen_nvfp4_kv( partial_dtype, bool(causal), bool(paged_kv), + _PAGE_TABLE_ASSUMED_ALIGN_BYTES if paged_kv else None, bool(use_prepare_scheduler), page_size, bool(seqused_k is not None), @@ -1964,7 +1994,7 @@ def _call_sparse_forward_sm100_csr_varlen_nvfp4_kv( else to_cute_tensor_kvouter(LSE_temperature_partial), to_cute_tensor_kvouter(Q_flat), None if Q_gather4_desc is None else to_cute_tensor_kvouter(Q_gather4_desc), - None if page_table is None else to_cute_tensor_kvouter(page_table), + None if page_table is None else to_cute_tensor_kvouter(page_table, assumed_align=_PAGE_TABLE_ASSUMED_ALIGN_BYTES), None if seqused_k is None else to_cute_tensor_kvouter(seqused_k), to_cute_tensor_kvouter(cu_seqlens_q), to_cute_tensor_kvouter(cu_seqlens_k), diff --git a/python/fmha_sm100/cute/test_sparse_atten.py b/python/fmha_sm100/cute/test_sparse_atten.py index 21c777e..719ca97 100644 --- a/python/fmha_sm100/cute/test_sparse_atten.py +++ b/python/fmha_sm100/cute/test_sparse_atten.py @@ -2900,6 +2900,7 @@ def _build_sparse_nvfp4_kv_benchmark_context( paged: bool = False, page_size: int = BLK_KV, seqused_trim: int = 0, + page_table_storage_offset: int = 0, ) -> dict[str, object]: torch.random.manual_seed(seed) if paged: @@ -2958,6 +2959,17 @@ def _build_sparse_nvfp4_kv_benchmark_context( cu_seqlens_q = inputs["cu_seqlens_q"] cu_seqlens_k = inputs["cu_seqlens_k"] page_table = inputs["page_table"] if paged else None + if page_table is not None and page_table_storage_offset: + page_table_storage = torch.empty( + page_table.numel() + page_table_storage_offset, + dtype=page_table.dtype, + device=page_table.device, + ) + page_table_view = page_table_storage[page_table_storage_offset:].view_as( + page_table + ) + page_table_view.copy_(page_table) + page_table = page_table_view seqused_k = inputs["seqused_k"] if paged else None max_seqlen_q = int(inputs["max_seqlen_q"]) max_seqlen_k = int(inputs["max_seqlen_k"]) @@ -3059,7 +3071,57 @@ def run_csr() -> None: "backend_fns": backend_fns, "run_csr": run_csr, "paged": paged, + "page_table": page_table, + } + + +def test_page_table_validation_rejects_insufficient_capacity(): + page_table = torch.zeros((2, 1), dtype=torch.int32) + with pytest.raises(ValueError): + sparse_interface._validate_page_table( + page_table, + device=page_table.device, + batch=2, + page_size=128, + max_seqlen_k=256, + ) + + +def test_sparse_atten_nvfp4_kv_unaligned_page_table() -> None: + context = _build_sparse_nvfp4_kv_benchmark_context( + case_name="nvfp4_unaligned_page_table", + q2k_pattern="sink", + batch=1, + seqlen_q=64, + seqlen_k=4096, + head_kv=1, + qhead_per_kv=16, + dim=128, + topk=32, + blk_kv=128, + causal=True, + seed=42, + paged=True, + page_size=128, + page_table_storage_offset=1, + ) + page_table = context["page_table"] + assert page_table.data_ptr() % 16 == page_table.element_size() + + outputs = { + backend: run_forward() + for backend, run_forward in context["backend_fns"] } + torch.cuda.synchronize() + reference = outputs["bf16_prefill"] + actual = outputs["nvfp4_kv_prefill"] + for actual_tensor, reference_tensor in zip(actual, reference): + torch.testing.assert_close( + actual_tensor.float(), + reference_tensor.float(), + atol=2e-2, + rtol=2e-2, + ) def _run_sparse_benchmark_warmup(ctx: dict[str, object], *, warmup: int, sync_nvtx: bool) -> None: diff --git a/python/fmha_sm100/sparse_fmha_adapter.py b/python/fmha_sm100/sparse_fmha_adapter.py index 306b416..c6b1f0c 100644 --- a/python/fmha_sm100/sparse_fmha_adapter.py +++ b/python/fmha_sm100/sparse_fmha_adapter.py @@ -232,13 +232,8 @@ def _build_page_table( kv_lens = kv_segment_lens.tolist() pages_per_batch = [(int(kl) + page_size - 1) // page_size for kl in kv_lens] max_pages = max(pages_per_batch) - total = batch * max_pages - buf = torch.zeros(total + 4, dtype=torch.int32, device=kv_indices.device) - shift = ((-buf.data_ptr()) % 16) // 4 - page_table = buf[shift : shift + total].view(batch, max_pages) - assert page_table.data_ptr() % 16 == 0, ( - f"_build_page_table failed to align: buf=0x{buf.data_ptr():x} " - f"shift={shift} page_table=0x{page_table.data_ptr():x}" + page_table = torch.zeros( + (batch, max_pages), dtype=torch.int32, device=kv_indices.device ) offset = 0 for b in range(batch): @@ -266,6 +261,7 @@ def sparse_fmha( kv_block_indexes: Optional[torch.Tensor] = None, q_offset_override = None, check_input_valid: bool = False, + page_table: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, None]: """Run sparse prefill through ``sparse_atten_func`` using an FMHA-style API. @@ -290,7 +286,11 @@ def sparse_fmha( Accepted for FMHA API compatibility; only ``sm_scale`` is used by this backend. kv_indices : torch.Tensor, optional - Flattened physical page table with dtype int32. Required for paged KV. + Flattened physical page table with dtype int32. Required for paged KV + unless ``page_table`` is supplied. + page_table : torch.Tensor, optional + Prebuilt physical page table with shape ``[batch, max_pages]`` and dtype + int32. When supplied, it is used directly and ``kv_indices`` is ignored. output_maxscore : bool, optional Accepted for compatibility; sparse prefill returns no max-score tensor. output_o : bool, optional @@ -347,13 +347,10 @@ def sparse_fmha( is_paged = page_size > 0 and k.ndim == 4 - page_table = None - if is_paged: - - if kv_indices is not None: - page_table = _build_page_table( - kv_indices, kv_segment_lens, page_size, batch, - ) + if is_paged and page_table is None and kv_indices is not None: + page_table = _build_page_table( + kv_indices, kv_segment_lens, page_size, batch, + ) # build_k2q_csr(return_schedule=True) builds schedule using hardware SM count internally # (build_k2q_csr_native.cu), which ignores usable_SM_count. When SM-limited, skip its diff --git a/tests/regression/test_sparse_attn.py b/tests/regression/test_sparse_attn.py index 6e616b8..bd55fe9 100644 --- a/tests/regression/test_sparse_attn.py +++ b/tests/regression/test_sparse_attn.py @@ -14,6 +14,10 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "python")) from fmha_sm100.sparse_fmha_adapter import sparse_fmha as fmha_sm100, sparse_fmha_plan as fmha_sm100_plan +from fmha_sm100.api import ( + fmha_sm100 as dispatch_fmha_sm100, + fmha_sm100_plan as dispatch_fmha_sm100_plan, +) failed_cases = [] @@ -78,26 +82,65 @@ def sparse_ref(q_flat, k_pages, v_pages, qo_lens, kv_page_indptr, def run_sparse_flashinfer(q, k_pages, v_pages, qo_lens_list, original_kv_lens_list, qo_offsets_list, kv_indices, pages_per_batch, kv_block_indexes, kv_block_num, num_qo_heads, page_size, head_dim, device, - num_kv_splits=-1, dtype=torch.bfloat16): - """Run FlashInfer sparse attention.""" + num_kv_splits=-1, dtype=torch.bfloat16, + public_page_table_mode=None): + """Run sparse attention through the adapter or public dispatcher.""" num_kv_heads = k_pages.shape[1] batch_size = len(qo_lens_list) qo_segment_lens = torch.tensor(qo_lens_list, dtype=torch.int32) kv_segment_lens = torch.tensor(original_kv_lens_list, dtype=torch.int32) qo_offset_tensor = torch.tensor(qo_offsets_list, dtype=torch.int32) - plan_info = fmha_sm100_plan(qo_segment_lens, kv_segment_lens, + plan_fn = ( + dispatch_fmha_sm100_plan + if public_page_table_mode is not None + else fmha_sm100_plan + ) + plan_info = plan_fn(qo_segment_lens, kv_segment_lens, num_qo_heads, qo_offset=qo_offset_tensor, page_size=page_size, num_kv_splits=num_kv_splits, kv_block_num=kv_block_num, num_kv_heads=num_kv_heads, ) torch.cuda.synchronize() - out, _ = fmha_sm100( + page_table = None + if public_page_table_mode is not None: + page_table = torch.zeros( + batch_size, max(pages_per_batch), dtype=torch.int32, device=device + ) + offset = 0 + for batch_idx, num_pages in enumerate(pages_per_batch): + page_table[batch_idx, :num_pages] = kv_indices[offset:offset + num_pages] + offset += num_pages + + has_mixed, split, _, _, _ = plan_info + if public_page_table_mode == "precedence": + assert has_mixed, "expected public mixed dispatch" + prefill_page_table = page_table[split:] + assert prefill_page_table.data_ptr() % 16 != 0 + + kv_page_split = sum(pages_per_batch[:split]) + prefill_kv_indices = kv_indices[kv_page_split:] + rolled_kv_indices = prefill_kv_indices.roll(1) + assert not torch.equal(rolled_kv_indices, prefill_kv_indices) + kv_indices = kv_indices.clone() + kv_indices[kv_page_split:] = rolled_kv_indices + else: + assert not has_mixed, "expected non-mixed sparse dispatch" + kv_indices = None + + run_fn = ( + dispatch_fmha_sm100 + if public_page_table_mode is not None + else fmha_sm100 + ) + + out, _ = run_fn( q, k_pages, v_pages, plan_info=plan_info, sm_scale=1.0 / math.sqrt(head_dim), kv_indices=kv_indices, kv_block_indexes=kv_block_indexes, check_input_valid=True, + page_table=page_table, ) torch.cuda.synchronize() return out @@ -107,7 +150,8 @@ def _run_sparse_varlen(name, seed, batch_size, num_kv_heads, num_qo_heads, page_size=128, head_dim=128, shuffle_pages=False, qo_lens=None, original_kv_lens=None, qo_offsets=None, sparse_block_counts=None, max_sparse_blocks=16, - num_kv_splits=-1, dtype=torch.bfloat16): + num_kv_splits=-1, dtype=torch.bfloat16, + public_page_table_mode=None): """Generic sparse attention test with full control over parameters.""" torch.manual_seed(seed) random.seed(seed) @@ -211,6 +255,7 @@ def _run_sparse_varlen(name, seed, batch_size, num_kv_heads, num_qo_heads, kv_indices, pages_per_batch, kv_block_indexes, kv_block_num, num_qo_heads, page_size, head_dim, dev, num_kv_splits=num_kv_splits, dtype=dtype, + public_page_table_mode=public_page_table_mode, ) threshold = 0.9999 if dtype == torch.bfloat16 else 0.999 @@ -377,6 +422,27 @@ def _run_sparse_varlen(name, seed, batch_size, num_kv_heads, num_qo_heads, for dt in dtypes: all_pass &= _run_sparse_varlen(f"large_batch {dt}", 42, 8, 4, 4, qo_lens=[1]*8, original_kv_lens=[512]*8, dtype=dt) + print("\n=== 17. Direct page table ===") + for dt in dtypes: + all_pass &= _run_sparse_varlen( + f"public_mixed_unaligned_page_table {dt}", 42, 3, 1, 16, + qo_lens=[8, 17, 33], + original_kv_lens=[1024, 2176, 4224], + shuffle_pages=True, + max_sparse_blocks=32, + dtype=dt, + public_page_table_mode="precedence", + ) + + all_pass &= _run_sparse_varlen( + "public_sparse_page_table_only", 43, 1, 1, 16, + qo_lens=[64], + original_kv_lens=[4096], + shuffle_pages=True, + max_sparse_blocks=32, + public_page_table_mode="only", + ) + print() total = len(failed_cases) if all_pass: