Add host-side validation for paged-KV cp.async load configs

The cp.async (non-TMA) paged-KV path derives page_entry_per_thread =
n_block_size // <KV-load threads> in PagedKVManager.create. When tile_n is
not an exact multiple of the loader thread count, the trailing KV rows of
each n-block are silently dropped (or page_entry_per_thread hits 0),
producing wrong output or an opaque failure deep inside JIT compilation.

Validate this in _flash_attn_fwd before compilation: when page_table is set
and page_size != tile_n, require tile_n to be divisible by the KV-load
thread count (128 on SM90; 128/64 on SM100/SM110 depending on q_stage), and
raise a ValueError naming the offending values otherwise. Scoped to the
standard SM90/SM100/SM110 forward; the dedicated hd256 kernel and the MLA
(qv) path use their own loaders and keep their existing checks.

Add tests/cute/test_paged_kv_validation.py: a GPU-free test that runs under
FakeTensorMode with a cute.compile sentinel, asserting invalid configs raise
ValueError before compilation, valid configs clear the guard, and the TMA
path (page_size == tile_n) is unaffected.

Author: fxdv

flash_attn/cute/interface.py

@@ -602,6 +602,42 @@ def _flash_attn_fwd(
     use_dedicated_hd256_kernel = arch // 10 in [10, 11] and head_dim == 256 and head_dim_v == 256
     use_2cta_instrs = use_2cta_instrs or use_dedicated_hd256_kernel
 
+    # Host-side validation for the cp.async (non-TMA) paged-KV path. When
+    # page_size != tile_n the kernel loads the page table through PagedKVManager,
+    # which derives page_entry_per_thread = tile_n // <KV-load threads> via integer
+    # division (see PagedKVManager.create in paged_kv.py). If tile_n is not an exact
+    # multiple of the KV-load thread count, the trailing KV rows of every n-block are
+    # never fetched (page_entry_per_thread rounds down, or hits 0), silently
+    # corrupting the output. Fail fast here instead of deep inside JIT compilation.
+    # The dedicated hd256 kernel and the MLA (qv) path use their own KV loaders, so
+    # this guard is scoped to the standard SM90/SM100/SM110 forward.
+    if (
+        page_table is not None
+        and qv is None
+        and not use_dedicated_hd256_kernel
+        and arch // 10 in [9, 10, 11]
+        and page_size not in (None, tile_n)
+    ):
+        if arch // 10 == 9:
+            # SM90 loads KV with a single warp group (num_threads_per_warp_group).
+            paged_kv_load_threads = 128
+        else:
+            # SM100/SM110 cp.async path: the load warp group is softmax1 (4 warps =
+            # 128 threads) when q_stage == 1, otherwise the two trailing warps
+            # (64 threads). See FlashAttentionForwardSm100 load_warp_ids selection.
+            paged_kv_load_threads = 128 if q_stage == 1 else 64
+        if tile_n % paged_kv_load_threads != 0:
+            raise ValueError(
+                f"Paged KV with page_size ({page_size}) != n_block_size ({tile_n}) "
+                f"selects the cp.async KV-load path, which requires n_block_size "
+                f"(tile_n={tile_n}) to be a multiple of the KV-load thread count "
+                f"({paged_kv_load_threads}) so each thread loads a whole number of "
+                f"page-table entries; got tile_n={tile_n} % {paged_kv_load_threads} = "
+                f"{tile_n % paged_kv_load_threads}. Use page_size == n_block_size "
+                f"(={tile_n}) to take the TMA path, or pick a configuration whose "
+                f"tile_n is divisible by {paged_kv_load_threads}."
+            )
+
     if softcap is not None:
         assert score_mod is None, "softcap and score_mod cannot be used together"
         score_mod = utils.create_softcap_scoremod(softcap)

tests/cute/test_paged_kv_validation.py

@@ -0,0 +1,131 @@
+"""Host-side validation tests for paged-KV configurations.
+
+These tests do NOT require a GPU. They exercise the early input validation in
+``_flash_attn_fwd`` that rejects paged-KV configs whose ``n_block_size`` (tile_n)
+is not divisible by the cp.async KV-load thread count. Such configs would
+otherwise fail with an opaque error deep inside JIT compilation (or, worse,
+silently drop KV rows).
+
+The pattern mirrors the fast two-pass test workflow: we run under
+``FakeTensorMode`` (no GPU memory) and replace ``cute.compile`` with a sentinel
+so that reaching compilation is observable without a CUDA device. The kernel
+arch is forced via the private ``_arch`` argument so the selected forward path
+is deterministic on any host.
+"""
+
+import pytest
+
+# CUDA / CuTeDSL are Linux+NVIDIA only; skip cleanly where they are unavailable
+# (e.g. macOS dev hosts) instead of erroring at collection time.
+torch = pytest.importorskip("torch")
+pytest.importorskip("cutlass")
+
+from torch._subclasses.fake_tensor import FakeTensorMode  # noqa: E402
+
+from flash_attn.cute import interface  # noqa: E402
+from flash_attn.cute.interface import _flash_attn_fwd  # noqa: E402
+
+
+class _CompileReached(Exception):
+    """Raised by the stubbed ``cute.compile`` to mark that host validation passed."""
+
+
+def _install_compile_sentinel(monkeypatch):
+    """Replace cute.compile with a sentinel and reset the compile cache.
+
+    If host validation passes, control reaches ``cute.compile`` and the sentinel
+    fires; we treat that as "the config was accepted".
+    """
+
+    def _sentinel(*args, **kwargs):
+        raise _CompileReached
+
+    monkeypatch.setattr(interface.cute, "compile", _sentinel)
+    # Ensure we never short-circuit on a previously cached kernel.
+    monkeypatch.setattr(_flash_attn_fwd, "compile_cache", {})
+
+
+def _make_paged_inputs(
+    *,
+    batch_size=1,
+    seqlen_q=128,
+    num_head=4,
+    num_head_kv=4,
+    head_dim=128,
+    head_dim_v=128,
+    num_pages=4,
+    page_size=64,
+    max_num_pages_per_seq=4,
+    dtype=torch.bfloat16,
+):
+    """Build fake (meta) tensors for a paged-KV forward call."""
+    q = torch.empty(batch_size, seqlen_q, num_head, head_dim, dtype=dtype)
+    k = torch.empty(num_pages, page_size, num_head_kv, head_dim, dtype=dtype)
+    v = torch.empty(num_pages, page_size, num_head_kv, head_dim_v, dtype=dtype)
+    page_table = torch.zeros(batch_size, max_num_pages_per_seq, dtype=torch.int32)
+    return q, k, v, page_table
+
+
+@pytest.mark.parametrize("arch", [90, 100])
+def test_paged_kv_indivisible_tile_n_raises(monkeypatch, arch):
+    """tile_n not divisible by the KV-load thread count must raise ValueError.
+
+    tile_n is forced to 96 via ``tile_mn``. On SM90 the cp.async loader uses 128
+    threads; on SM100 (q_stage == 1 here, since seqlen_q == tile_m) it also uses
+    128. 96 % 128 != 0, so the config is rejected before compilation.
+    """
+    _install_compile_sentinel(monkeypatch)
+    with FakeTensorMode():
+        q, k, v, page_table = _make_paged_inputs(page_size=64)
+        with pytest.raises(ValueError, match=r"cp\.async KV-load path"):
+            _flash_attn_fwd(
+                q,
+                k,
+                v,
+                page_table=page_table,
+                tile_mn=(128, 96),  # force n_block_size = 96
+                _arch=arch,
+            )
+
+
+@pytest.mark.parametrize("arch", [90, 100])
+def test_paged_kv_divisible_tile_n_passes_guard(monkeypatch, arch):
+    """A valid paged config (tile_n divisible by load threads) clears the guard.
+
+    tile_n is forced to 128 (divisible by both 128 and 64), with page_size=64 so
+    the cp.async path is selected. The new validation must NOT raise; execution
+    should reach the stubbed ``cute.compile`` (``_CompileReached``).
+    """
+    _install_compile_sentinel(monkeypatch)
+    with FakeTensorMode():
+        q, k, v, page_table = _make_paged_inputs(page_size=64)
+        with pytest.raises(_CompileReached):
+            _flash_attn_fwd(
+                q,
+                k,
+                v,
+                page_table=page_table,
+                tile_mn=(128, 128),  # divisible by the KV-load thread count
+                _arch=arch,
+            )
+
+
+def test_paged_kv_tma_path_skips_guard(monkeypatch):
+    """When page_size == tile_n the TMA path is used and the guard is bypassed.
+
+    Even with tile_n=96 (not divisible by 128), page_size == tile_n means
+    PagedKVManager's cp.async loader is never used, so the config is accepted and
+    reaches compilation.
+    """
+    _install_compile_sentinel(monkeypatch)
+    with FakeTensorMode():
+        q, k, v, page_table = _make_paged_inputs(page_size=96, num_pages=4)
+        with pytest.raises(_CompileReached):
+            _flash_attn_fwd(
+                q,
+                k,
+                v,
+                page_table=page_table,
+                tile_mn=(128, 96),  # page_size == tile_n -> TMA path
+                _arch=90,
+            )