[PyTorch] Split FusedAttnFunc into single-argument forward/backward helpers - #3480
[PyTorch] Split FusedAttnFunc into single-argument forward/backward helpers#3480pggPL wants to merge 2 commits into
Conversation
…elpers Behavior-neutral restructuring of the cuDNN fused attention autograd function, mirroring what NVIDIA#2967 did for Linear: FusedAttnFwdArgs / FusedAttnBwdArgs dataclasses, module-level _fused_attn_forward_impl, _fused_attn_setup_ctx and _fused_attn_backward_impl, and a thin FusedAttnFunc wrapper taking the differentiable tensors plus one args object instead of 37 positional arguments. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The bias / softmax_offset aux entries exist only when the tensor was passed, not merely when the bias / softmax type asks for one. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
|
/te-ci pytorch L1 |
Greptile SummaryThe PR restructures the cuDNN fused-attention autograd implementation into typed forward/backward argument containers and module-level helpers without intending to change behavior.
Confidence Score: 5/5The PR appears safe to merge with no changed-code defects identified. The refactored paths preserve saved-tensor handling, native auxiliary-pack ordering, the supported shared page-table model, backward gradient arity, and module importability. Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[FusedAttention.forward] --> B[Build FusedAttnFwdArgs]
B --> C[FusedAttnFunc.apply]
C --> D[_fused_attn_forward_impl]
D --> E[_fused_attn_setup_ctx]
E --> F[Save backward tensors and metadata]
F --> G[FusedAttnFunc.backward]
G --> H[Restore FusedAttnBwdArgs]
H --> I[_fused_attn_backward_impl]
I --> J[dQ, dK, dV, bias and offset gradients]
Reviews (1): Last reviewed commit: "[PyTorch] Match the fused attention aux ..." | Re-trigger Greptile |
| if t is not None | ||
| } | ||
| saved_from = tuple(sources.get(id(t)) for t in tensors_to_save) | ||
| tensors_to_save = tuple(None if src else t for t, src in zip(tensors_to_save, saved_from)) |
There was a problem hiding this comment.
We return bias and softmax_offset in fwd: same tensor memory but wrapped in different Pybind objects. Would id(t) be able to detect them and if not, would it fail to dedup the tensors here? Would checking the storage help?
def _alias_key(t: Any) -> Any:
# torch.library rejects op outputs that alias inputs, and pybind returns a
# fresh Python object for the same at::Tensor, so identity is not enough.
if isinstance(t, torch.Tensor):
return ("storage", t.untyped_storage().data_ptr(), t.storage_offset(),
tuple(t.shape), tuple(t.stride()))
return ("obj", id(t))
There was a problem hiding this comment.
Do we want to pass strict=True to the zips at ~1818 and ~1949 so a length drift would fail loudly rather than truncating?
| # --- FP8 --- | ||
| fp8: bool | ||
| fp8_meta: Optional[Dict[str, Any]] | ||
| quantizers: Optional[Any] |
There was a problem hiding this comment.
Would field annotations like Optional[Any] break custom-op registration?
| attn_scale: float = 1.0 | ||
| dropout_p: float = 0.0 | ||
| fast_zero_fill: bool = True | ||
| qkv_layout: str = "sbh3d" |
There was a problem hiding this comment.
We probably should leave the defaults of arguments like qkv_layout to None. This way, it can fail loudly rather than passing but incorrectly?
| fwd_args.attn_bias = attn_bias | ||
| fwd_args.softmax_offset = softmax_offset | ||
| out, max_logit, tensors_to_save_from_forward, ctx_attrs = _fused_attn_forward_impl(fwd_args) | ||
| bwd_args = FusedAttnBwdArgs() |
There was a problem hiding this comment.
Should we add if ctx is not None: here like in Linear:
| cu_seqlens_kv: torch.Tensor | ||
| cu_seqlens_q_padded: Optional[torch.Tensor] | ||
| cu_seqlens_kv_padded: Optional[torch.Tensor] | ||
| page_table: Optional[torch.Tensor] |
There was a problem hiding this comment.
Can we keep page_table_k and page_table_v separate even though we are assuming they are the same in our KV cache implementation? Users who subclass our KV cache manager may have different implementations for K/V page table. Thanks.
| cu_seqlens_kv_padded: Optional[torch.Tensor] | ||
| page_table: Optional[torch.Tensor] | ||
| packed_qkv: Optional[torch.Tensor] | ||
| packed_kv: Optional[torch.Tensor] |
There was a problem hiding this comment.
How does a mutable dataclass (with these tensor members) behave under torch.compile? can it still keep dynamo happy?
2 participants
Description
Behavior-neutral restructuring of
FusedAttnFunc, the cuDNN fused attentionautograd.Function, mirroring what #2967 did forLinearahead of #3053. Groundwork for #3472 (torch.compile support for the fused attention backend), which registers these helpers as a torch custom op throughdynamo/custom_op.py.FusedAttnFwdArgs/FusedAttnBwdArgsdataclasses replace the 37 positional arguments ofFusedAttnFunc.applyand the loosectx.*attributes._fused_attn_forward_impl,_fused_attn_setup_ctxand_fused_attn_backward_implhold the actual logic;FusedAttnFuncis a thin wrapper taking the differentiable tensors (q,k,v,attn_bias,softmax_offset) plus one args object, so its backward returns 6 grads instead of 36Nones.softmax_stats,rng_state,aux_bias,aux_softmax_offset) instead of a variable-length list.Nonein that slot and names its source inctx_attrs["saved_from"];_fused_attn_setup_ctxre-attaches it. This is what a custom op needs (it may not return its own inputs) and keeps the FP8 paths -- where aNonef16 slot means "not needed" -- exact.ctxattributes the backward never read (fp8_recipe,fp8_meta,is_output_fp8).No functional change intended. Tested on an RTX Ada with the fused backend forced:
tests/pytorch/attention/test_attention.py(test_dpa*,test_dot_product_attention,test_transformer_layer),test_kv_cache.py,test_cpu_offloading_v1.pyandtest_torch_compile.pyall pass. FP8 attention and context parallelism could not be exercised on that GPU and rely on CI.Type of change
Changes
attention/dot_product_attention/backends.py:FusedAttnFwdArgs,FusedAttnBwdArgs,_fused_attn_forward_impl,_fused_attn_setup_ctx,_fused_attn_backward_impl,_reload_qkv_layout;FusedAttnFuncreduced to a wrapper;FusedAttention.forwardbuilds the args object.Checklist: