[PyTorch][torch.compile] DotProductAttention: declarative packed QKV/KV inputs

[pggPL]

Description

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.

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).

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• DotProductAttention.forward: new optional qkv_layer / kv_layer / qkv_interleave_dim arguments; packed inputs are resolved by _unpack_packed_qkv into zero-copy q/k/v views plus a declared layout, skipping get_qkv_layout entirely.
• get_qkv_layout 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: the original packed buffer is threaded down as packed_qkv/packed_kv and combine_and_quantize(combined_qkv=..., combined_kv=...) quantizes it directly instead of rebuilding it from views via combine_tensors (raw set_ with a silent adjacency assumption). Backward combine paths untouched.
• MultiheadAttention passes its fused projection output straight to DPA (self-attention: qkv_layer; cross-attention: kv_layer). The legacy sliced-views path is kept for RoPE, QK norm, KV caching, CPU offloading, GQA and FP8 projection outputs.
• Tests: _run_dot_product_attention gains a declarative_packed mode (packed buffer passed via qkv_layer/kv_layer, input grads read off the packed buffer); test_dpa_qkv_layout_declarative covers all 8 packed dense layouts on self- and cross-attention configs, test_dpa_qkv_layout_thd_declarative the 4 packed thd layouts (Hopper+). Verified on sm89 (bf16, fused+flash+unfused) with the existing test_dpa_qkv_layout matrix unchanged and green; lint clean.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adds a declarative packed QKV/KV input API to DotProductAttention, allowing callers to pass a fused projection buffer directly via new qkv_layer/kv_layer/qkv_interleave_dim arguments instead of slicing Q/K/V views that require pointer-based layout detection on every forward pass.

• _unpack_packed_qkv extracts zero-copy select() views and constructs the layout string (e.g. \"bs3hd\", \"thd_t2hd\") from the declaration, covering dense, thd, and FP8 DPA.
• MultiheadAttention passes its fused projection output straight to DPA when no per-tensor operation (RoPE, QK norm, KV caching, CPU offload) requires individual slices; GQA and FP8 QKV outputs fall back to the legacy sliced-views path.
• get_qkv_layout emits a DeprecationWarning (with stacklevel=2) when packed layouts are detected via pointer inspection, suppressed during CPU offload where migration is not possible.

Confidence Score: 5/5

The change is safe to merge: all layout construction paths are verified correct across dense, thd, and cross-attention formats, backward gradient routing through select-view autograd is sound, and existing callers are unaffected by the purely additive API.

The core _unpack_packed_qkv helper correctly handles every format/interleave combination. MHA packed paths gate correctly on QuantizedTensorStorage and num_queries_per_key_value. Gradient flow through select() views is handled properly. The only findings are two defensive assert statements in combine_and_quantize that should be raise ValueError.

No files require special attention; the two assert-vs-raise sites in utils.py are the only items worth a follow-up cleanup.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds _unpack_packed_qkv helper and three new optional forward args; layout construction logic and rank/stride validations are correct across all format/interleave combinations.
transformer_engine/pytorch/attention/multi_head_attention.py	Adds packed_dpa_eligible gate and wires packed buffers to DPA for self/cross-attention; FP8 QKV output and GQA self-attention correctly fall back to sliced-views path via QuantizedTensorStorage check and num_queries_per_key_value guard.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds deprecation warning in get_qkv_layout and combined_qkv/combined_kv fast-path in combine_and_quantize; dim-index derivation via string find is consistent with the declared layout construction.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Threads packed_qkv/packed_kv through FusedAttnFunc and FusedAttention; backward correctly returns None for both since gradients flow via the q/k/v select-view autograd path.
tests/pytorch/attention/test_attention.py	Adds declarative_packed mode to _run_dot_product_attention and two new test functions covering 8 dense layouts and 4 thd layouts; gradient extraction from packed buffer and backend comparison are correctly implemented.

_{Reviews (2): Last reviewed commit: "[PyTorch] Address review: warning stackl..." | Re-trigger Greptile}