[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 declarative packed QKV/KV inputs to DotProductAttention via three new optional arguments (qkv_layer, kv_layer, qkv_interleave_dim). When a caller provides the packed buffer directly, the new _unpack_packed_qkv helper derives zero-copy select() views for q/k/v and constructs the exact layout string from the declaration — bypassing the pointer-based get_qkv_layout detection entirely and enabling torch.compile compatibility.

• DotProductAttention.forward gains qkv_layer/kv_layer/qkv_interleave_dim; _unpack_packed_qkv validates inputs (mutual exclusion, rank, interleave-dim size, stride(-1)==1) and returns the declared layout, skipping get_qkv_layout on this path.
• MultiheadAttention passes the fused projection output directly to DPA via the new arguments when no per-tensor operation (RoPE, QK-norm, KV cache, CPU offload) is needed, gated by not isinstance(mixed_x_layer, QuantizedTensorStorage) for FP8 exclusion.
• combine_and_quantize accepts combined_qkv/combined_kv to quantize the original packed buffer directly instead of rebuilding it via combine_tensors; get_qkv_layout now emits a DeprecationWarning (with correct stacklevel=2) for pointer-detected packed layouts outside CPU-offload mode.

Confidence Score: 5/5

Safe to merge. All new code paths are additive: existing callers are unchanged, the packed path is guarded by explicit non-None checks, and the FP8 optimization in combine_and_quantize falls back to combine_tensors when no packed buffer is provided.

The layout-string construction (_packed_layout) is mathematically correct for all six format×interleave-dim combinations; validation in _unpack_packed_qkv covers mutual exclusion, rank, interleave size, and last-dimension stride; the MHA packed-eligible gate correctly excludes RoPE, QK-norm, inference, CPU offload, and FP8 projection outputs. The two assert statements in combine_and_quantize are defensively redundant (upstream validation already catches mismatches) and would only fire in abnormal use.

The two new assert statements in transformer_engine/pytorch/attention/dot_product_attention/utils.py (combine_and_quantize, case 1 and case 2) should be converted to ValueError raises for consistency with the rest of the validation layer.

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 args to DotProductAttention.forward; layout string construction and validation logic is correct for all format/interleave-dim combinations.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds combined_qkv/combined_kv shortcuts in combine_and_quantize and deprecation warning in get_qkv_layout; shape checks use assert (disabled at -O) rather than ValueError.
transformer_engine/pytorch/attention/multi_head_attention.py	Adds packed_dpa_eligible gate and passes packed buffers to DPA for both self- and cross-attention; QuantizedTensorStorage check correctly excludes FP8 outputs from the packed path.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Threads packed_qkv/packed_kv through FusedAttention and FusedAttnFunc for the FP8 optimization; backward correctly returns None gradients for these non-differentiable arguments.
tests/pytorch/attention/test_attention.py	New declarative tests cover all 8 dense and 4 thd packed layouts, comparing FusedAttention vs FlashAttention output; does not cross-validate the declarative path against the legacy pointer-detection path.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A["DotProductAttention.forward\n(query_layer, key_layer, value_layer,\nqkv_layer?, kv_layer?, qkv_interleave_dim)"]
A --> B["_unpack_packed_qkv()"]
B --> C{"qkv_layer or kv_layer?"}
C -- "neither (legacy)" --> D["return q, k, v unchanged\ndeclared_layout = None"]
C -- "qkv_layer provided" --> E["Validate: rank, shape[interleave_dim]==3,\nstride(-1)==1, q/k/v must be None"]
C -- "kv_layer provided" --> F["Validate: rank, shape[interleave_dim]==2,\nquery_layer required, k/v must be None"]
E --> G["q=select(dim,0), k=select(dim,1), v=select(dim,2)\nlayout = e.g. 'bs3hd'"]
F --> H["k=select(dim,0), v=select(dim,1)\nlayout = e.g. 'bshd_bs2hd'"]
D --> I{"Layout branch"}
G --> I
H --> I
I -- "declared_layout not None" --> J["Skip get_qkv_layout\nqkv_layout = declared_layout"]
I -- "FP8 tensors" --> K["get_qkv_layout on _data"]
I -- "normal" --> L["get_qkv_layout on q/k/v\n(may emit DeprecationWarning)"]
J --> M["Select backend"]
K --> M
L --> M
M -- "FusedAttention" --> N["FusedAttnFunc.forward\n(packed_qkv=qkv_layer, packed_kv=kv_layer)"]
M -- "FlashAttention" --> O["flash_attention.forward"]
M -- "Unfused" --> P["unfused_attention.forward"]
N --> Q{"FP8 quantization?"}
Q -- "Yes + packed buffer" --> R["combine_and_quantize(combined_qkv/combined_kv)\n→ quantize buffer directly"]
Q -- "Yes, rebuild" --> S["combine_and_quantize\n→ combine_tensors then quantize"]
Q -- "No" --> T["standard forward"]

Loading

_{Reviews (3): Last reviewed commit: "[PyTorch] Address review nits: clarify D..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch