[PyTorch] torch.compile support for UnfusedDotProductAttention#19
Open
pggPL wants to merge 12 commits into
Open
[PyTorch] torch.compile support for UnfusedDotProductAttention#19pggPL wants to merge 12 commits into
pggPL wants to merge 12 commits into
Conversation
…le + CUDA graphs
Refactor TE custom kernels used by the unfused attention path so that
`torch.compile(fullgraph=True, mode="reduce-overhead")` can trace the
forward and backward and capture them into CUDA graphs without graph
breaks.
- softmax.py / softmax.cpp: register all `scaled_*_softmax_{forward,backward}`
kernels as `torch.library.custom_op`s with fake impls and an autograd
binding that mirrors the previous `torch.autograd.Function`s. The C++
backward kernels now allocate a fresh output buffer instead of writing
in-place into `output_grad`, so the ops no longer alias their inputs
(required by `torch.library.custom_op` and inductor cudagraph trees).
- utils.py: convert `ConvertTHDtoBSHD` / `ConvertBSHDtoTHD` to
`torch.library.custom_op`s, with thin wrapper classes that keep the
existing `.apply(...)` callsite syntax. Drop the
`int(cu_seqlens[-1].item())` from the hot path of `ConvertBSHDtoTHD.apply`
-- under `torch.compile` it created an unbacked SymInt, which made the
Inductor partitioner emit `None` placeholders for output buffers and
caused `cudagraph_trees` to assert. `num_tokens` is now passed in by
the caller as a regular (Sym)Int.
- backends.py: in the THD branch of unfused DPA, capture
`total_tokens_q = query_layer.shape[0]` before overwriting
`query_layer` with the BSHD form, and thread it back into
`ConvertBSHDtoTHD.apply` at the end of the forward.
- test_torch_compile.py: add `test_unfused_dpa_torch_compile`,
parametrized over qkv layouts (`bshd_bshd_bshd`, `sbhd_sbhd_sbhd`,
`thd_thd_thd`, `bs3hd`, `sbh3d`), that compiles
`UnfusedDotProductAttention.forward` directly with `fullgraph=True,
mode="reduce-overhead"` and runs forward+backward several times so the
CUDA graphs are recorded and replayed.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Made-with: Cursor
…to unfused_attention_torch_compile Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> # Conflicts: # tests/pytorch/test_torch_compile.py
…DotProductAttention
Make the FP8-emulation path (NVTE_UnfusedDPA_Emulate_FP8=1) of
UnfusedDotProductAttention traceable by torch.compile(fullgraph=True).
- backends.py: register the quantize+dequantize roundtrips used by
FP8EmulationFunc as torch.library custom ops
(te_fp8_emu::roundtrip_<QuantizerClass> and
te_fp8_emu::roundtrip_qkv_<QuantizerClass>) taking the quantizer as a
value-opaque argument, with fake impls for tracing. Ops are registered
only for the value-opaque quantizer classes
(Float8CurrentScalingQuantizer, MXFP8Quantizer); Float8Quantizer
(delayed scaling) carries scale/amax tensor state, is not
value-opaque, and deliberately keeps the plain eager path -- FP8
emulation with delayed scaling is not supported under torch.compile.
- backends.py: dispatch helpers `_fp8_emu_roundtrip{,_qkv}` key on
`type(quantizer).__qualname__` so they stay traceable for opaque
quantizer arguments; FP8EmulationFunc forward/backward now call them
(onnx_forward unchanged).
- backends.py: the joint q/k/v roundtrip clones any output whose
storage is shared with an input or another output, checking storage
identity directly -- the dequantized q/k/v can be views into one
combined buffer, and view metadata (`_base`) is not populated under
the torch-dispatch mode AOTAutograd runs custom ops with, so a
`_base`-guarded clone triggered the custom-op aliasing deprecation
warning under torch.compile.
- UnfusedDotProductAttention.forward: only query
FP8GlobalStateManager.get_fp8_recipe() when
fp8_meta["local_recipes"] is absent.
- test_torch_compile.py: add test_unfused_dpa_fp8_emulation_torch_compile
(current scaling + mxfp8, sbhd/bshd layouts; compiled fullgraph
forward+backward must match eager) and
test_unfused_dpa_fp8_emulation_delayed_scaling_eager guarding the
eager delayed-scaling path after the refactor.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…output=True With fp8_output=True the backend returns a Float8Tensor -- a tensor subclass that cannot cross a torch.compile graph boundary -- so the forward dispatches to a torch._dynamo.disable'd wrapper, the same mechanism DotProductAttention and FusedAttention use module-wide. With fp8_output=False the dispatcher is resolved at trace time and adds no graph break. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…cudagraphs) Parametrize test_unfused_dpa_fp8_emulation_torch_compile over compile mode (default, reduce-overhead), run 3 iterations so the CUDA graphs are recorded and replayed. The te_fp8_emu roundtrip ops for current scaling are pure (no mutated args), so inductor cudagraphs capture them; verified no cudagraph skips with TORCH_LOGS=cudagraphs. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…n; run FP8 as an eager island FP8 in the unfused backend (emulation and Float8Tensor output) is not supported under torch.compile: the forward dispatcher routes fp8=True and/or fp8_output=True to a torch._dynamo.disable'd wrapper, same as DotProductAttention does module-wide. Remove the FP8-emulation compile tests. The te_fp8_emu::* custom ops taking value-opaque quantizers stay as the eager implementation of FP8EmulationFunc. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…Func The ops existed solely to make the FP8-emulation path traceable by torch.compile; since FP8 in the unfused backend now always runs as an eager island, they are dead machinery (plus import-time registration and output clones the plain eager path never needed). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The tex softmax kernels take 'float scale_factor' directly. The 0-D tensor wrapping was a leftover of the old autograd.Function idiom, where the float had to be a tensor only to fit save_for_backward; the custom ops keep the scale on ctx as a plain attribute. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…the callsite) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…cate dict, silence W0613 - run black over the four changed files (earlier commits skipped pre-commit) - drop unused 'import os' in test_torch_compile.py - drop duplicated module-level _default_causal_mask dict in softmax.py - del unused 'output' arg in the conversion setup_context helpers Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL
changed the title
for more information, see https://pre-commit.ci
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Make the
UnfusedDotProductAttentionbackend traceable bytorch.compile(fullgraph=True, mode="reduce-overhead"), so the forward and backward can be captured into CUDA graphs without graph breaks.Scope:
torch.library.custom_ops with fake impls and autograd bindings; remove an unbacked-SymInt.item()from the hot path ofConvertBSHDtoTHD.fp8=True(emulation) and/orfp8_output=True(Float8Tensor output, a tensor subclass that cannot cross a graph boundary) the backend runs as an eager island — the forward dispatches to atorch._dynamo.disable'd wrapper, the same mechanismDotProductAttentionandFusedAttentionuse module-wide. FP8 attention always involves delayed scaling regardless of the recipe: S and dP are produced inside the kernel, so their amax cannot be known before quantization and they use delayed-scaling quantizers even under Float8CurrentScaling (seeDPA.init_fp8_metadata) — and delayed scaling (Float8Quantizer, tensor scale/amax state) is not supported under torch.compile.Type of change
Changes
softmax.py/softmax.cpp:scaled_*_softmax_{forward,backward}as custom ops; C++ backward kernels allocate a fresh output buffer instead of writing in-place intooutput_grad(custom ops and cudagraph trees forbid input aliasing).utils.py:ConvertTHDtoBSHD/ConvertBSHDtoTHDas custom ops;num_tokenspassed by the caller instead ofcu_seqlens[-1].item().backends.py:UnfusedDotProductAttention.forwarddispatches FP8 calls to an eager (dynamo-disabled) wrapper; the non-FP8 path is traced with no graph breaks.tests/pytorch/test_torch_compile.py:test_unfused_dpa_torch_compile(5 qkv layouts, fullgraph + reduce-overhead, fwd+bwd captured into CUDA graphs and replayed).Checklist:
🤖 Generated with Claude Code