[Doc]: SDPA execute-time shape override cannot cross s_q == 1; heuristic engine flips between declared max_len 128 and 256

[YangXu1990uiuc]

Summary

Two properties of execute-time shape override (pygraph(is_override_shape_enabled=True) + execute(..., override_uids, override_shapes, override_strides)) on the unified SDPA forward engine are not documented and cost users a plan build (or an error) when hit:

1. The s_q == 1 boundary cannot be crossed. A graph built with s_q == 1 (decode class) rejects an override to s_q > 1 and vice versa with CUDNN_STATUS_NOT_SUPPORTED_INVALID_DYNAMIC_SHAPE. Callers have to keep separate graphs for the two classes (FlashInfer does: perf(prefill): build the cuDNN ragged graph once per length class via execute-time shape override; keep single-token GQA rows off cuDNN (NVBug 6783545) flashinfer-ai/flashinfer#5245).
2. Heuristic engine flips with the declared max sequence length. Between a declared max_len of 128 and 256 the mode-A heuristic switches engine (e.g. eng8 → eng10) on SM100 and SM107, independent of batch, head dims or Stats. Overriding a graph declared long to a short real shape runs the long-class engine: 2.0-2.3x slower on 4-token rows than a graph declared short (measured on B200). So callers should also class shapes as "<= 128" vs "> 128" when they build once and override.

Also worth stating in the docs: override forces the split-K factor to 1; for the cu_seq_len form this is already the case, for the per-batch seq_len form it costs 1.2-1.6x on low-occupancy shapes.

Ask

Document these in the SDPA / dynamic shape pages (and ideally surface (1) at build time rather than at execute).

_{note to self: claude::61d24ed2-7c90-4a97-9cbb-b91ae18136fd — "flashinfer frost prefill/GEMM enablement" · cwd /home/scratch.yanxu_libs/flashinfer}

[0z5a]

Hi @YangXu1990uiuc, I'd like to take this.

I'll start with the documentation + regression-coverage slice: document the s_q == 1 decode/prefill boundary, the declared-max-length sensitivity of heuristic selection, and the split-K limitation under execute-time overrides. I'll also add a minimal executable example/test that demonstrates the valid and invalid override cases without hard-coding the SM100/SM107 heuristic threshold as a cross-architecture rule.

If useful, I can then look at surfacing the s_q == 1 incompatibility earlier at build time as a separate follow-up rather than mixing that behavioral change into the documentation PR.

[0z5a]

Documentation + regression slice is up: #1148.

What it adds: an SDPA page section on declared-vs-overridden shapes, the s_q == 1 decode/prefill boundary (measured here: the backend answers CUDNN_STATUS_NOT_SUPPORTED_INVALID_DYNAMIC_SHAPE), the build-time plan-selection property, the split-KV caveat, and the distinction from the dynamic-shape kernel cache and CUDA graph capture. The SM100/SM107 numbers in your report (eng8 -> eng10 between declared 128/256, 2.0-2.3x on 4-token rows) are attributed to your platform and not generalized.

Two findings worth your call, both measured on an L20 (cuDNN 9.26, f16, B=1 H=4 D=64, plan eng8_k24=1_k27=0_k38=0_k40=3_k41=1) and pinned as strict xfails rather than papered over:

1. non-causal overrides off the plan's 64-row tile grid return wrong output with no error (96 -> max|dO| 8.1e-01, 160 -> 4.9e-01), while the same extents declared outright are correct;
2. overrides beyond the declared extent are not rejected (320 finite-but-wrong, 512 NaN).

Also: the frontend's own override-length check does not run on this path, so a short triple surfaces at variant-pack finalize as CUDNN_STATUS_BAD_PARAM. Happy to turn any of these into a follow-up if you want the frontend to reject them. The execute path also needs a frontend compiled against cuDNN >= 9.21 (a 9.19 build accepts the overrides and silently ignores them — that cost me a debugging round and is now written down).