Validate paged-KV host config (tile_n divisible by KV-load thread count)

[fxdv]

What

Adds a host-side ValueError guard in _flash_attn_fwd for the cp.async paged-KV load path. It requires tile_n % <kv_load_threads> == 0 so that page_entry_per_thread = n_block_size // num_threads in PagedKVManager.create is exact. Without this, a misconfigured tile_n makes the page-table load loop run too few (or zero) iterations, so the trailing KV rows of every n-block are silently never fetched — producing wrong output or an opaque failure deep inside JIT compilation.

Scope

The guard fires only when:

• page_table is not None, and
• the cp.async path is selected (page_size not in (None, tile_n)), and
• the standard forward is used: non-MLA (qv is None), non-hd256 dedicated kernel, arch SM90/SM100/SM110.

The KV-load thread count is the loader's thread count, not the interface num_threads:

• SM90: 128 (the KV-load warp group, num_threads_per_warp_group).
• SM100/SM110: 128 when q_stage == 1 (load warps = softmax1, 4 warps), else 64 (load warps 14–15). This mirrors FlashAttentionForwardSm100's load_warp_ids selection and the num_load_threads passed to PagedKVManager.create.

The MLA (qv) path and the dedicated hd256 kernel use their own KV loaders and keep their existing checks, so they are excluded.

Why

Fail fast with a clear, actionable message for vLLM-style paging misconfigurations instead of returning silently-wrong results or surfacing an opaque compilation error.

Test plan

New GPU-free test tests/cute/test_paged_kv_validation.py (runs under FakeTensorMode with a cute.compile sentinel):

• invalid config (tile_n=96, cp.async path) → raises ValueError before compilation,
• valid config (tile_n=128, cp.async path) → clears the guard and reaches compilation,
• TMA path (page_size == tile_n) → guard bypassed even with tile_n=96.

Validation caveat

The change was validated statically (ruff and AST parse are clean) but the tests were not executed locally: the contributor host is Apple Silicon macOS without torch/CUDA. Running the test suite requires a Linux + CUDA host or CI.

[fxdv]

Hi @Johnsonms — would you have a moment to review this when convenient?

Small host-side validation for the cp.async paged-KV path: fails fast with a clear ValueError when tile_n is not divisible by the kernel KV-load thread count (128 on SM90; 128/64 on SM100/110 depending on q_stage), instead of silent corruption from a zero-iteration load loop.

Includes GPU-free tests in tests/cute/test_paged_kv_validation.py (FakeTensorMode + compile sentinel). Static validation only on my side — no CUDA host locally.

Happy to rebase on latest main if helpful. Thanks!