[PyTorch] DotProductAttention: declarative packed qkv/kv inputs

[pggPL]

Motivation

A fused QKV projection produces one packed buffer, but callers must slice it into query_layer/key_layer/value_layer views, which DotProductAttention then reverse-engineers via pointer-based layout detection (get_qkv_layout inspecting data pointers, strides and offsets on every forward). That detection graph-breaks under torch.compile and re-derives information the caller already knows.

This PR lets callers declare the packing instead. The declared layout is truthful by construction, so no detection runs on this path — including for thd and FP8 DPA.

API

DotProductAttention.forward gains three optional arguments (existing call sites unchanged):

• qkv_layer — fully packed QKV, e.g. [b,s,3,h,d] or [s,b,h,3,d] (thd: [t,3,h,d]).
• kv_layer — packed KV (e.g. [b,s,2,hg,d]), used with query_layer; supports GQA.
• qkv_interleave_dim (default -3) — where the 3/2 interleave sits (-3 or -2); explicit since shapes can be ambiguous.

Handling lives in a single helper (_unpack_packed_qkv): q/k/v come out as zero-copy select() views, the exact layout string (bs3hd, sbhd_sb2hd, th3d, ...) is built from the declaration, and inputs are strictly validated (mutual exclusion, rank/size at the interleave dim, stride(-1) == 1 so the declared layout cannot lie about memory, no inference_params).

get_qkv_layout now emits a DeprecationWarning when it detects a packed layout purely from pointers, pointing callers at the new arguments. Separate q/k/v never warn.

FP8: combine_and_quantize

For packed layouts, FP8 attention used to rebuild the packed buffer from the q/k/v views via combine_tensors (a raw set_ with a silent adjacency assumption). The original packed buffer is now threaded down as packed_qkv/packed_kv and combine_and_quantize(combined_qkv=..., combined_kv=...) quantizes it directly. Legacy call sites are unchanged (None defaults); backward combine paths untouched.

MultiheadAttention adoption

MHA hands its fused projection output straight to DPA: self-attention passes packed QKV as qkv_layer, cross-attention passes packed KV as kv_layer. The legacy sliced-views path is kept for RoPE, QK norm, KV caching, CPU offloading, GQA and FP8 projection outputs.

Tests

No new test file. _run_dot_product_attention gains a declarative_packed mode (packed buffer passed via qkv_layer/kv_layer, input grads read off the packed buffer), exercised by two small additions to the existing suite:

• test_dpa_qkv_layout_declarative — all 8 packed dense layouts x {self causal+bias, cross padding} configs, compared cross-backend like test_dpa_qkv_layout.
• test_dpa_qkv_layout_thd_declarative — the 4 packed thd layouts (Hopper+).

Verified on sm89 (bf16, fused+flash+unfused): declarative tests green, test_dpa_qkv_layout matrix unchanged and green, lint 10/10.

🤖 Generated with Claude Code