[PyTorch][torch.compile] Make get_attention_backend traceable without graph breaks - #3189
Merged
pggPL merged 12 commits intoJul 23, 2026
Merged
Conversation
…reaks - Read NVTE_* env vars via os.environ.get instead of os.getenv so dynamo installs guards on the values (os.getenv reads are not guarded and would bake stale backend selections into compiled graphs). - Wrap tex.get_fused_attn_backend in a torch.compiler.assume_constant_result helper so the pybind call does not graph-break. - Mark get_device_compute_capability/get_cudnn_version with assume_constant_result for the same reason. - Use a no-op logger when compiling (logging.Logger methods graph-break) and skip debug-log blocks that call int()/str() on pybind enums. - Add test in tests/pytorch/test_torch_compile.py checking fullgraph=True tracing and recompilation on NVTE_* env var changes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Contributor
Greptile SummaryThis PR makes
Confidence Score: 5/5Safe to merge. The backend-selection logic is correct in all paths; the FusedAttnBackend IntEnum mirrors the C++ enum exactly with an import-time assert as a safety net; os.environ.get guards and assume_constant_result are applied consistently; and backward compatibility is preserved through cast() normalisation in fused_attn_fwd/bwd. All changes are in selection/dispatch logic with no tensor computation affected. The only findings are quality observations: warning/error logs are silently dropped during compilation (users lose fallback notices under torch.compile), a minor type annotation inaccuracy in ne, and three per-version env-var keys left out of the test pre-set. None of these affect correctness or introduce data-path regressions. transformer_engine/pytorch/attention/dot_product_attention/utils.py deserves a second look regarding the _NoOpLogger swallowing warning/error log levels. tests/pytorch/test_torch_compile.py is missing NVTE_FLASH_ATTN_V2/V3/V4 pre-sets which limits dynamo guard coverage for those keys. Important Files Changed
Reviews (8): Last reviewed commit: "[PyTorch] Address review: add info()/err..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comparing the pybind enum returned through assume_constant_result against module-level enum values generates guards dynamo cannot evaluate (crash when the comparison is true, i.e. when cuDNN rejects the config). The wrapper now returns a plain int, comparisons use precomputed int values, and the enum for callers is reconstructed by a second assume_constant_result helper that is never compared during tracing. Also document that os.environ.get (vs os.getenv) is intentional, and drop the guard_scalar specialization of numeric args: symbolic scalars (automatic dynamic) now graph break at the probe instead of forcing a full recompile per seqlen value; the test covers that path without fullgraph and checks the selection stays correct. A second test monkeypatches tex.get_fused_attn_backend to verify the baked result is trace-time-only and actually drives selection. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The pybind NVTE_Fused_Attn_Backend enum is not traceable by torch.compile: its C-implemented __eq__ cannot be traced, and a pybind enum instance baked through assume_constant_result produces guards dynamo cannot evaluate when compared against module-level enum values. FusedAttnBackend is now a plain python IntEnum generated at import time from tex.NVTE_Fused_Attn_Backend.__members__ (values always in sync with the C enum), and all remaining direct uses of the pybind enum on the python side are replaced with it. Name lookup (FusedAttnBackend["FP8"]) behaves the same as with the previous dict, and the backend value never crosses into a pybind call, so no boundary conversion is needed. This removes the previous int-based workaround in get_attention_backend (_fused_attn_backend_from_int and the precomputed int table): the assume_constant_result wrapper now simply returns the IntEnum and comparisons are traceable directly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Mirror how constants.DType migrated off the pybind enum: explicit IntEnum members pinned to the C values with an import-time sync assert, an __eq__ override comparing by integer value against NVTE_Fused_Attn_Backend (with matching __ne__/__hash__) so mixed comparisons stay equivalent regardless of the pybind11 version, and a cast() classmethod. fused_attn_fwd/bwd normalize their fused_attention_backend argument through cast(), so external callers still passing the pybind enum keep working. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
test_get_attention_backend_traceable_fp8 compiles the selection with fullgraph=True for AttentionParams(fp8=True) with a DelayedScaling(fp8_dpa) recipe, covering the FP8-only branch (run_config env reads, recipe filters, get_fp8_te_dtype) and checks that flipping NVTE_UnfusedDPA_Emulate_FP8 recompiles and keeps matching eager. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The is_compiling() guards around the available/selected-backend debug logs protected int() on the pybind enum, which crashed dynamo during tracing. With FusedAttnBackend now a python IntEnum, int() on it and str() on the flash-attn PkgVersion both trace cleanly (verified under fullgraph=True), so the logging blocks return to their upstream shape. The probe wrapper also reuses FusedAttnBackend.cast() and a shorter docstring. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Merge the three get_attention_backend tests into one covering: fullgraph tracing with the probe consulted at trace time only, env var flips (F16 and FP8), attention-param changes, and a forced No_Backend result driving the selection. The bitmask output now also encodes the fused sub-backend, which previously went unchecked. The probe wrapper takes layout/bias/mask/softmax as string keys and resolves the pybind enums internally, so every argument is a literal or a python enum - required for assume_constant_result(specialize_args=True) (pytorch#189042) to derive value guards once available. Scalars must stay concrete until then: the test pins specialize_int/float=True, because a symbolic scalar currently graph breaks at the probe and dynamo's resume then corrupts the returned fused backend (binds the wrapper function object instead of its result; surfaced by the sub-backend bits, minimal repro exists). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The probe-argument and static-scalar comments referenced assume_constant_result(specialize_args=True), which is not part of any released PyTorch; describe the current behavior only. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Drop the call-counting monkeypatch and its assertions; compiled-vs-eager output equality is what matters and already fails on stale selections. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…tion_backend_traceable # Conflicts: # tests/pytorch/test_torch_compile.py
Collaborator
Author
|
/te-ci pytorch |
pggPL
changed the title
cyanguwa
previously approved these changes
Jul 22, 2026
…tion_backend_traceable Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> # Conflicts: # transformer_engine/pytorch/attention/dot_product_attention/utils.py
Collaborator
Author
|
/te-ci pytorch |
shino16
approved these changes
Jul 23, 2026
shino16: _NoOpLogger does not subclass logging.Logger, so any traced logger.info()/logger.error() call under torch.compile would raise AttributeError. Add the two no-op methods for completeness. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
cyanguwa
approved these changes
Jul 23, 2026
9 tasks
pggPL
added a commit
to pggPL/TransformerEngine
that referenced
this pull request
Jul 29, 2026
Log arguments are evaluated regardless of the logger, so the no-op logger from NVIDIA#3189 is not enough here -- skip the whole block while tracing. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
9 tasks
pggPL
added a commit
that referenced
this pull request
Aug 17, 2026
… unfused backends (#3286) * [PyTorch][torch.compile] Support for DotProductAttention Make DotProductAttention itself traceable, so that torch.compile captures the whole module -- input unpacking, qkv layout, backend selection and the backend -- on the FlashAttention and UnfusedDotProductAttention backends. FusedAttention stays an eager island, as before. DotProductAttention.forward no longer carries @no_torch_dynamo. What that uncovered, and the fixes: - init_fp8_metadata unconditionally called get_fp8_recipe(), which builds a default Recipe when no autocast is active -- not traceable. It now returns early when FP8 is off everywhere, which is also what the rest of the function did in that case. - get_qkv_layout detects packed q/k/v from data pointers and storage offsets, neither of which dynamo can trace (and neither is meaningful on fake tensors). While tracing, the layout is built from the format instead; packed inputs are declared explicitly via qkv_layer/kv_layer, and non-contiguous q/k/v are made contiguous so the layout stays truthful. - get_full_cu_seqlens keyed its cache on torch.is_inference_mode_enabled(), a fundamental graph break; while tracing it builds the tensor directly. - get_padding_mask read sequence lengths on the host (a device sync, and data-dependent under torch.compile). It is now vectorized. - The fused sub-backend enum member does not survive a graph break: dynamo reconstructs the assume_constant_result value by re-emitting the call that produced it, so the resumed frame received the function itself instead of the enum -- which then reached fused_attn_fwd. get_attention_backend now keeps a plain int while tracing and casts back to FusedAttnBackend in eager mode only. - Backend selection logging is skipped while tracing (logging.Logger methods graph-break, and int() of the sub-backend enum is untraceable too). Tests: test_dpa_torch_compile covers 7 configurations (bshd/sbhd/thd, causal/no_mask/sliding window, GQA, cross attention, packed qkv_layer) against both backends with fullgraph=True, comparing forward and backward against eager; plus a CUDA-graphs (reduce-overhead) test and a test that the FusedAttention path stays correct around its graph break. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Delegate the no-FP8 case of DPA's init_fp8_metadata to the base class The override hand-rolled the state that TransformerEngineBaseModule already sets when FP8 is off everywhere (the three flags, fp8_checkpoint and fp8_initialized). Call super() instead, and read the flags straight off FP8GlobalStateManager.quantization_state, as the other hot paths do. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Run DPA eagerly for undeclared packed q/k/v instead of guessing the layout get_qkv_layout recognizes packed q/k/v by inspecting data pointers and storage offsets, which dynamo cannot trace. Rather than assume such inputs are three separate tensors -- a layout that would not describe memory -- detect the case from what is traceable (a strided tensor that does not own its whole storage) and run the whole module as an eager island, with a warning. Declaring the packing via qkv_layer/kv_layer keeps the call on the compiled path. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Log backend selection in eager only; shorten comments Log arguments are evaluated regardless of the logger, so the no-op logger from #3189 is not enough here -- skip the whole block while tracing. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fix lint: unused import and keyword-arg-before-vararg - utils.py no longer needs no_torch_dynamo: the eager island moved to dot_product_attention.py. - _needs_eager_dpa reads q/k/v out of *args/**kwargs instead of naming them before *args, which pylint flags (W1113). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Compare against eager under CUDA graphs, and tighten the tolerances The CUDA-graphs test only checked that the output was finite and that gradients existed. A replay that reuses stale buffers -- the failure mode that test is there to catch -- satisfies both, so it now compares against eager on fresh inputs every iteration. The remaining comparisons used the dtype defaults (rtol=1.6e-2 for bfloat16), which is far looser than what these paths deliver: compiled and eager agree bit-for-bit on every config and backend. FlashAttention now has to match exactly, as it runs the same kernel either way; the unfused backend is allowed a single bfloat16 rounding, since inductor may reassociate the softmax sums. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Support one cu_seqlens tensor for both q and kv under torch.compile Self attention on thd naturally passes the same tensor as cu_seqlens_q and cu_seqlens_kv, and flash-attn's varlen entry point forwards both to the same autograd.Function -- which dynamo refuses to trace ("duplicate tensor input"), failing the compilation under fullgraph=True. Hand the varlen call a copy of one of them while tracing. The clone is b + 1 int32 elements, only on that path, and only under torch.compile. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Apply the shared-cu_seqlens workaround to the FlashAttention v4 path too Pull the deduplication into a helper and use it wherever both cumulative sequence lengths are handed to flash-attn: the varlen path shared by v2 and v3 training, and v4's varlen kwargs. The v3 KV-cache path derives the kv lengths by subtraction, so its two arguments are distinct by construction. v4 is not installed in the environment this was tested in, so that call site is covered by construction rather than by a test; the helper is a no-op unless the two arguments are the same object. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Run FP8 attention eagerly under torch.compile FP8 attention brings quantizers, fp8_meta mutation and Float8Tensor across the graph boundary, and none of it is exercised by this PR's tests. Rather than trace it untested, route it to the eager island that packed q/k/v already use. The predicate is deliberately narrow: it asks the recipe for fp8_dpa/fp8_mha rather than whether an autocast is active, so FP8 GEMMs with attention in high precision -- the common training setup -- stay on the compiled path. The decorator now takes a predicate that returns the reason, so the warning names which of the two paths was taken. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Drive the compile tests off ModelConfig instead of a hand-picked intersection The configurations were chosen so that both backends could run all of them, which meant nothing backend-specific was ever exercised: no biases, no arbitrary masks, no MLA head dims, no sink softmax. Describe them with the ModelConfig the eager attention tests already use, and let get_available_attention_backends() decide which backend runs which, so a configuration only one backend supports is covered rather than avoided. Five such configurations are added. One of them immediately showed that the tolerance for the unfused backend was too tight: an off-by-one softmax differs from eager by two bfloat16 roundings on a single element out of 65536, because inductor reassociates the softmax sums. Only FlashAttention, which runs the same kernel either way, keeps the exact comparison. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Scale the unfused comparison to the tensor, and take tolerances from utils An off-by-one softmax exposed that comparing the unfused backend elementwise does not hold up: inductor's reassociation of the softmax sums produces an error proportional to those sums, which for gradient elements near zero is far outside any relative tolerance -- while being one bfloat16 rounding at the tensor's own scale. Derive the absolute tolerance from the reference tensor's magnitude, and take the dtype tolerances from tests/pytorch/utils.py rather than restating them. FlashAttention and FusedAttention keep the exact comparison. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Merge the per-backend compile tests, and give padding masks real padding The configurations built cu_seqlens with every sequence at the maximum length, so a padding mask was all-False: the padding branches ran, but on data that could not tell a difference. They now use varying sequence lengths, which immediately showed that FlashAttention on bshd/sbhd with a padding mask does not trace: it packs the tensors first, and get_indices built the index list with a Python comprehension over sequence lengths held in a GPU tensor -- a host synchronization per sequence in eager, untraceable while compiling. It is now built on the device, verified exhaustively against the previous implementation (10304 cases, 0 mismatches) and free of synchronization. The separate tests for one shared cu_seqlens tensor and for the fused backend folded into the main one: the first is a configuration, and the second is the backend axis, which now runs flash, fused and unfused with the graph break that FusedAttention's eager island implies. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Run eagerly when max_seqlen has to be derived from cu_seqlens On thd inputs without max_seqlen, DotProductAttention reads the sequence lengths off cu_seqlens to derive it. Unlike the padding mask and the packing indices, this one cannot be moved to the device: max_seqlen sizes tensors downstream, so it has to be a Python int, and reading it is a device synchronization -- a data-dependent value dynamo cannot trace. Route those calls to the eager island, so passing max_seqlen is what keeps a call compiled. Arguments are now looked up by name or position, so the predicate sees them whichever way the caller passed them. The two eager-fallback tests merge into one: they set up different inputs but assert the same thing -- a warning, results matching eager, and an error under fullgraph. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Share the run-and-compare code between the compile tests Both tests ran a callable, went backward, copied the gradients out and cleared them. That is now _run_and_capture(), with _assert_run_matches() comparing two of its results, which halves the CUDA-graphs test. The two stay separate tests: what CUDA graphs risk -- a replay reading stale buffers -- does not depend on the attention configuration, so folding the mode into the configuration matrix would multiply the cases without covering anything new, and would report a replay bug as a failure of some unrelated configuration. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Use the shared run-and-compare code in the remaining compile tests The eager-fallback test hand-rolled the same run, backward and gradient copy. The backend-level unfused test only checked that the output was finite and that gradients existed, which a CUDA-graph replay reading stale buffers also satisfies -- the weakness the DotProductAttention CUDA-graphs test was given an eager comparison for. It now compares against eager as well. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Read the predicate's arguments by name, not by hardcoded position The positions of forward's parameters were written out as a table of indices, which anyone inserting a parameter into a thirty-argument signature would silently invalidate: the predicate would go on reading whatever now sits at index 5 or 10 and decide the fallback on it, with nothing to notice. The decorator already has the function it wraps, so it takes the parameter names from its signature and hands the predicate the call's arguments keyed by name. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Move the eager-fallback decorator to jit.py Nothing about it is specific to attention: it takes a predicate, and runs the wrapped method as an eager island whenever that predicate gives a reason. It belongs next to no_torch_dynamo, which it builds on, rather than buried in the attention module where the modules that will need it next -- MultiheadAttention and TransformerLayer, both of which still have blockers of their own -- would not find it. What stays in the attention module is the predicate, which is where the attention-specific knowledge lives. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Skip FusedAttention rather than special-casing it in the compile tests The test compiled every backend with fullgraph=True except fused, which was passed fullgraph=False because its forward is an eager island. That put an assumption about one backend into the comparison logic. FusedAttention is skipped instead, next to the configurations a backend cannot run, with the reason spelled out. Everything else is compiled the same way, and supporting fused later is deleting the skip. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Cover the compiled path around FusedAttention FusedAttention is skipped from the backend matrix because it does not compile, but it is the default on Hopper and Blackwell, and this change does affect it: what used to be one eager island is now a traced prologue, a graph break at the backend, and a traced remainder. Whatever crosses that break has to survive it, which the sub-backend enum did not. One test covers it, without a fullgraph and outside the matrix. Verified to fail with the enum restored -- TypeError: int() argument must be ... not 'function', from cuDNN's entry point. Also rename _distinct_cu_seqlens: it returns its arguments untouched in eager and whenever they already differ, so it promised more than it does. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Log backend selection through the no-op logger, as get_attention_backend does Skipping the block with is_compiling() introduced a second mechanism for a problem this module already solves one way: get_attention_backend swaps in a no-op logger while tracing. The skip was needed while the sub-backend argument was an enum that did not survive tracing, and is not any more. The logger, which was private to the utils module, is now named without the underscore, since it is used from outside it. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Return the fused sub-backend as an int, in eager as well get_attention_backend cast it back to a FusedAttnBackend member unless it was tracing, which made its return type depend on the execution mode -- something every reader of that function then has to keep in mind, for no benefit inside transformer_engine: fused_attn_fwd/bwd call cast() on whatever they are given, and every comparison against a member works with an int on either side. It now returns an int always, and the docstring says so. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Rename eager_under_compile_if to fallback_to_eager_when The old name read as a run-on, and its '_if' suffix suggested a boolean condition where the argument actually returns the reason. The new one reads as a sentence where it is used and matches the warning the decorator emits. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fold the eager fallback into no_torch_dynamo as a `when` predicate It was a decorator of its own, which meant a second name for what is the same thing this module already does: disable dynamo for a frame. It is now the conditional form of no_torch_dynamo -- @no_torch_dynamo(when=predicate) -- so there is one decorator to find and one docstring to read. The predicate receives the call's arguments keyed by parameter name and returns the reason this call cannot be traced, or None to have it traced. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Cover the declared packed layouts beyond bs3hd Packing was pinned to one layout by an assert in the input builder, so of the twelve declarations DotProductAttention accepts -- qkv_layer or kv_layer, interleaved at -3 or -2, in each of the three formats -- exactly one was tested. The builder now takes what is packed and where it is interleaved, and derives the shapes and the layout string the way DotProductAttention does. Three configurations are added, so both packed inputs and both interleave dimensions are covered. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Let ONNX export keep its own path in get_qkv_layout The ONNX exporter runs through dynamo, so torch.compiler.is_compiling() is true during an export as well -- measured: inside torch.onnx.export(dynamo=True) both it and is_in_onnx_export_mode() are set. The compile branch was written first and took the export over: the layout came from the format instead of the detection ONNX relies on, tensors were no longer made contiguous, and packed views hit an assert where they used to be copied. Export mode is the narrower context, so it is checked first. This restores the previous behaviour for export while leaving the compiled path as it was. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Guard the assumption the argument binding rests on Arguments are matched to parameter names by position, which is only right while the signature has no *args in between -- and would misalign silently if one were added. It is checked when the decorator is applied. The docstring now also says that an argument the call left out is simply absent, so a predicate reading it with .get() sees its default only as long as that default is None. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Check that CUDA graphs were actually captured, and generalize a docstring fullgraph=True says dynamo did not break the graph, but inductor can still decline to capture -- a mutated input, a CPU scalar -- and run the thing normally, which would leave the CUDA-graphs test passing while measuring nothing. Its skip counter is asserted to be zero. Measured on this branch: nothing is skipped, the tree has one root, and interleaving the eager reference call between replays does not disturb it. The tolerance docstring described which backends are compared exactly by naming them; it now states the rule -- a backend whose kernel is the same either way -- so it still holds once FusedAttention traces. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Skip lazy compilation while already tracing Filling lazy_compile's closure cell is a side effect Dynamo rejects inside a higher-order op, so compiling a module whose autograd.Function backward makes the first call to a jit_fuser'd function fails with Unsupported: HOP: Unsafe side effect Attempted to mutate CellVariable() An earlier eager run hides this by filling the cell beforehand. The nested torch.compile is inlined into the outer graph either way, so calling the function directly while tracing is equivalent. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Run context parallel attention eagerly The context parallel implementation uses P2P communication and its own CUDA stream, neither of which dynamo traces, and FlashAttention reaches it from a forward that is no longer excluded from the graph. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Keep DotProductAttention compilable under a CUDA RNG states tracker Forking the tracker swaps the global CUDA generator state, which dynamo refuses to trace. With attention_dropout == 0 the dropout draws nothing from the generator, so the fork is a no-op -- use nullcontext instead and keep the call in the graph. Calls that do use dropout under a tracker run as an eager island. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Run checkpointed attention eagerly Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Support KV caching under torch.compile Wrap the KV cache kernels in custom ops: dynamo cannot trace pybind calls into transformer_engine_torch, so a decoding step used to break the graph. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Compile the fused sbh3d QKV split Wrap fa_prepare_fwd/bwd like the KV cache kernels, and cover the layout that reaches them: no configuration exercised it, compiled or eager. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Run FlashAttention 4 eagerly Version 4 registers no custom ops, unlike 2 and 3: it builds its kernels through the CUTLASS DSL as it runs. Disable dynamo on its entry points, so only the calls that reach it become eager islands. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Pin the sequence lengths backend selection bakes in A second sequence length makes dynamo's automatic dynamic turn max_seqlen symbolic, and _get_fused_attn_backend, being assume_constant_result, cannot take a symbolic argument -- so the call died once the length changed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Capture CUDA graphs with a KV cache The cache is mutated in place, so inductor declined to capture until its buffers were marked static. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Test generation against a KV cache A prefill and three single-token steps, each compared to eager, with and without CUDA graphs: the cache carries state, so a step that updates it wrongly only shows in the step after. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop the imports the custom ops made unused Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Rename _maybe_unshare_cu_seqlens to _unalias_cu_seqlens Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Split the get_full_cu_seqlens cache bypass per reason Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Skip get_fp8_recipe only while tracing in init_fp8_metadata The unconditional early return skipped main's new CustomRecipe local-recipes inference when quantization is off. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Skip the max_seqlen bake-in under ONNX export guard_scalar specializes user-declared dynamic sequence dims in the exported model; the exporter skips backend selection, which the bake-in is for. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Make undeclared non-contiguous q/k/v contiguous under compile The declared layout is only truthful for contiguous tensors. Eager repairs such inputs via its detect-and-retry loop; mirror that in the compiled branch instead of declaring a layout the memory contradicts. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
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
3 participants
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
This PR is part of an effort to enable
torch.compilesupport for attention: it makesget_attention_backendtraceable.It works on current PyTorch as-is — there is no need to wait for the two related PRs I opened in PyTorch. Those only improve the remaining rough edge: pytorch/pytorch#189042 (bigger and much more important; among other things it handles symbolic scalar arguments, which with stock dynamo must stay non-symbolic in the backend-selection probe — with static scalars, the common case, everything in this PR already works) and pytorch/pytorch#189027 (small, and possibly not needed anymore since this PR reads env vars via
os.environ.get, which stock dynamo already guards).get_attention_backendcurrently graph-breaks undertorch.compile, forcing backend selection to run as an eager island. This PR makes it fully traceable withfullgraph=True, with dynamo guards on theNVTE_*environment variables so that changing an env var triggers recompilation instead of silently reusing a stale backend selection.Type of change
Changes
Read
NVTE_*env vars inget_attention_backendviaos.environ.getinstead ofos.getenv: dynamo installs value guards onos.environaccesses, whileos.getenvreads are unguarded (the stdlib-global source is skipped) and would bake a stale backend selection into the compiled graph.Wrap the
tex.get_fused_attn_backendpybind call in atorch.compiler.assume_constant_resulthelper — the result depends only on the attention configuration, not on tensor values.Mark
get_device_compute_capabilityandget_cudnn_versionwithassume_constant_result(pybind/CUDA property calls are not traceable, and their results are constant for a process).Use a no-op logger when compiling (
logging.Loggermethods graph-break in dynamo) and skip the debug-log blocks that evaluateint()/str()on pybind enum values during tracing. Backend-selection debug logs are unchanged in eager mode.Add
test_get_attention_backend_traceabletotests/pytorch/test_torch_compile.py: compiles a function callingget_attention_backendwithfullgraph=True(any graph break fails the test) and flipsNVTE_FUSED_ATTN/NVTE_UNFUSED_ATTN/NVTE_FLASH_ATTNto verify the guards trigger recompilation and the result keeps matching eager.FusedAttnBackend(inpytorch/cpp_extensions/fused_attn.py) is now a pythonIntEnumgenerated at import time fromtex.NVTE_Fused_Attn_Backend.__members__, replacing the previous str->pybind-enum dict, and all direct python-side uses of the pybind enum are replaced with it. Rationale: the pybind enum is not traceable by dynamo (C__eq__), and a pybind enum baked throughassume_constant_resultproduces guards dynamo cannot evaluate when compared against module-level enum values (hard crash when cuDNN rejects a config). The transition follows the pattern used forconstants.DType: explicit members pinned to the C values with an import-time sync assert, an__eq__override so mixed comparisons withtex.NVTE_Fused_Attn_Backendstay equivalent regardless of the pybind11 version, and acast()classmethod;fused_attn_fwd/bwdnormalize their backend argument throughcast(), so callers still passing the pybind enum keep working. Name lookupFusedAttnBackend["FP8"]andint(...)behave as before. The remaining (minor) breaking surface isisinstancechecks against the pybind enum and dict-API such as.items()/iteration; no such uses exist inside TE and none were found in Megatron-LM/NeMo.Note: dynamo only guards
os.environkeys that exist at trace time; reads of absent keys are not guarded yet (upstream PyTorch limitation).Checklist: