fix: support external page tables with 4-byte alignment

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)
 

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

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: