[common/PyTorch] Fused attention: pass real tensor strides to cuDNN (opt-in, v2 C API)

[pggPL]

Problem

TE's cuDNN fused-attention integration loses tensor stride information at the pybind boundary: torch tensors are wrapped into NVTETensor (shape+ptr+dtype — no strides), and the C++ side reconstructs strides from the NVTE_QKV_Layout enum via hard-coded formulas (generateMatrixStrides). cudnn-frontend itself is stride-based (Tensor_attributes::set_stride; NVIDIA's own SDPA sample expresses bs3hd-interleaved QKV purely as strides).

Consequences of the enum channel:

• Python must detect the layout enum by inspecting data_ptr()/storage_offset() (get_qkv_layout) — untraceable under torch.compile (gb0156) and pure per-step CPU overhead;
• a mismatch between the declared enum and actual memory produces silently wrong results (no validation).

Change (2 commits)

1. Opt-in real-stride plumbing. The PyTorch extension reads the live at::Tensor strides of Q/K/V (and dO in backward), permutes them to cuDNN dim order [b,h,s,d], and the F16 arbitrary-seqlen backend uses them for the cuDNN graph set_stride instead of enum formulas — for dense (non-THD, non-paged) layouts. The graph/plan cache key (FADescriptor_v1) is extended with the stride values. Gated by NVTE_FUSED_ATTN_REAL_STRIDES=1 (default off → behavior unchanged).

2. Additive v2 C API. Strides travel as an explicit parameter, no hidden state:

typedef struct NVTEQKVStrides { const int64_t *q, *k, *v, *d_o; } NVTEQKVStrides;  // [b,h,s,d] each, NULL => enum-derived
void nvte_fused_attn_fwd_v2(..., NVTEQKVStrides qkv_strides, NVTETensor workspace, cudaStream_t stream);
void nvte_fused_attn_bwd_v2(..., NVTEQKVStrides qkv_strides, NVTETensor workspace, cudaStream_t stream);

nvte_fused_attn_fwd/bwd remain as thin wrappers calling _v2 with NULL strides — JAX and external callers are untouched byte-for-byte. (_v2 naming follows nvte_cublas_gemm_v2/nvte_quantize_v2.) Outputs (O, dQ/dK/dV) keep enum-derived strides. FP8/max-512 and THD/paged paths ignore the parameter (documented in the header).

Why this matters

With real strides, the enum no longer needs to encode memory geometry: q/k/v that are strided views into a packed buffer (bs3hd, sbh3d, kv-packed) work correctly even when declared as the plain separate layout. This removes the need for pointer-based layout detection in Python (torch.compile-hostile) and eliminates the silent-corruption failure mode.

Verification (RTX Ada sm89, bf16, F16 arbitrary-seqlen, fwd+bwd grads)

Case	Result
flag OFF, contiguous (determinism baseline)	bit-exact
flag ON, contiguous, declared separate	bit-exact vs baseline
flag ON, packed [b,s,3,h,d] views declared separate	bit-exact
flag OFF, packed views declared separate (status quo)	wrong (max diff 4.97) — documents the hazard
flag ON, sbh3d views declared sbhd_sbhd_sbhd	bit-exact
test_dot_product_attention (flag OFF and ON)	48 passed / 48 skipped
old symbols after wrapper conversion	exported; flag-OFF parity across full suite

New focused test: tests/pytorch/attention/test_fused_attn_real_strides.py (5 tests).

Notes / limitations

• THD/ragged and paged-KV intentionally keep the enum path (strides cannot express ragged memory; cu_seqlens/ragged offsets remain the channel there).
• dQ/dK/dV allocation still follows the enum layout (pre-existing behavior).
• Opt-in for now; flipping the default (and then removing Python-side layout detection) is follow-up work.

🤖 Generated with Claude Code

Flag lifecycle (commitment)

NVTE_FUSED_ATTN_REAL_STRIDES is a transitional rollout vehicle, scheduled for removal:

1. now: default OFF (byte-for-byte legacy behavior; opt-in for torch.compile users),
2. after bake time on CI/other archs: flip default to ON,
3. remove the flag entirely and make real strides unconditional for dense F16 layouts; delete the Python-side pointer-based layout detection (get_qkv_layout) at the same time.

The code is structured so step 3 is mechanical: the flag is read in exactly two places (the pybind glue in attention.cpp and one Python gate helper in dot_product_attention.py) — removal deletes those two checks and the detection path they guard. Nothing else depends on the flag.

Rationale for not shipping always-on immediately: today's detection also normalizes odd stride patterns via a .contiguous() retry; with detection gone those patterns reach cuDNN raw (loud plan-build failure, not corruption — but a new input class that deserves bake time), and the cuDNN support surface for general strides across archs/versions is broader than the enum set we could test on sm89.