[PyTorch] Compile-friendly attention API: pointer-free qkv layout detection

[pggPL]

What

get_qkv_layout detects the qkv memory layout by inspecting storage — untyped_storage().data_ptr(), storage_offset(), stride(). Under torch.compile this graph-breaks hard (gb0156: Dynamo cannot trace UntypedStorage.data_ptr), and it was one of the main remaining breaks in DotProductAttention.forward.

This introduces a JAX-style, pointer-free way to determine the layout, gated so default eager behavior is unchanged.

How

Commit 1 — pointer-free layout detection (JAX-style):

• New get_qkv_layout_pointer_free(...): builds the layout string from qkv_format + packing info only — no data_ptr/storage_offset/stride. Same return contract as get_qkv_layout.
• DotProductAttention.forward now accepts packing signalled by argument presence (like JAX TE's QKVLayout): key_layer=None, value_layer=None → query_layer is qkv-packed [..,3,h,d]; value_layer=None → key_layer is kv-packed [..,2,h,d]; all three → separate. Packed tensors are unbound into views (traceable). An explicit qkv_layout: Optional[str] kwarg skips detection entirely.
• Gate: pointer-free path is used when packing is signalled, an explicit layout is passed, or torch.compiler.is_compiling(). Otherwise (default eager, 3 separate tensors) the original pointer-inspection path runs unchanged.

Commit 2 — gate inside get_qkv_layout itself:
The forward-level gate does not fire when an earlier graph break makes Dynamo skip the forward frame (forward then runs eagerly, is_compiling() is False there, and Dynamo compiles get_qkv_layout as its own frame — breaking on data_ptr). Checking is_compiling() inside get_qkv_layout folds to a constant during tracing and delegates to the pointer-free path, so a traced call never reaches pointer inspection.

Notes / limitations

• Under compile, 3 separate tensors get the plain separate layout rather than a pointer-detected packed layout — functionally correct for flash/fused attention; only the packed-storage micro-optimization is forgone.
• h3d/h2d interleaving cannot be signalled by argument presence (packing assumes 3/2 at dim -3, the JAX convention); pass an explicit qkv_layout for those.

Verification (workstation, RTX Ada, FlashAttention forced, bf16, bshd)

• data_ptr graph breaks through DotProductAttention.forward under torch.compile: 3 → 0.
• Compiled output bit-exact vs eager; qkv-packed and kv-packed inputs bit-exact vs equivalent separate tensors.
• Default eager 3-tensor call still routes through the original path (unchanged).

🤖 Generated with Claude Code