[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends - #22
Open
pggPL wants to merge 70 commits into
Open
[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends#22pggPL wants to merge 70 commits into
pggPL wants to merge 70 commits into
Conversation
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>
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>
…ayout 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 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>
- 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>
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>
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>
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>
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>
…rsection 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>
…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>
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>
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>
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>
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>
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>
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>
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>
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>
…end 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>
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>
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>
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>
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>
…x::exp2f (NVIDIA#3262) Fix UE8M0 code 0 and 255 expansion in ptx::exp2f UE8M0 code 0 is 2^-127 and code 255 is NaN, but the exponent-field shift produced +0.0 and +Inf, so MXFP8 software dequantize zeroed every 1x32 block whose scale byte was 0. Mirror the special cases already present in exp2f_rcp and add a dequantize test with planted extreme scale codes. Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>
expose _BOOTSTRAPPED Signed-off-by: YangFei1990 <feiw@nvidia.com>
…_on (NVIDIA#3304) * Migrate EP collective-stream annotation to region-based compute_on * Constrain EP MoE output sharding to match activation input Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…rheads for MOE (NVIDIA#3238) * initial implementation for act grouped linear fusion Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments + cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comment + cleanup Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix CI + handle activation requires_grad=false and fusion optimizations for precomputed tensor offsets in backward Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * precomputed tensor offsets in grouped linear as well Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix lint Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * unecsaary based on op infra Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * better name Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * activation+group_quantize fusion via ops infra Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci ugly solution for dbias fusion Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert to have grouped_linear + activation fusion for now Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * get rid of fused ops now Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * improve comments Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add support for fused Q Up-Proj GEMM/RoPE/Quant. This commit add support for fusing the GEMM in the Q Up Proj step of DeepseekV3 training with the following RoPE and MXFP8 quantization operations. This uses a custom kernel from cudnn_frontend, and supports both 16-bit projection and mxfp8 projection. Signed-off-by: Chase Block <cblock@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove unused function from fused mla q uproj, add error handling Signed-off-by: Chase Block <cblock@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Adjust comment SBHD/BSHD Signed-off-by: Chase Block <cblock@nvidia.com> * Add cudnn_frontend version check; eliminate code dup in mxfp8_quantize_fast_path. Signed-off-by: Chase Block <cblock@nvidia.com> * Add missing docstrings Signed-off-by: Chase Block <cblock@nvidia.com> * Add unit test for q up proj fusion Signed-off-by: Chase Block <cblock@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add test for bwd path of fused gemm+rope+quant. Signed-off-by: Chase Block <cblock@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Consolidate fused q proj tests and add to unit test list Signed-off-by: Chase Block <cblock@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove backward test from fused mla q uproj. These tests really belong in Megatron. Signed-off-by: Chase Block <cblock@nvidia.com> --------- Signed-off-by: Chase Block <cblock@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com> Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
…IDIA#3244) * add support for row-wise quanted input for grouped gemm Signed-off-by: YangFei1990 <feiw@nvidia.com> * doc change Signed-off-by: YangFei1990 <feiw@nvidia.com> * implement for backward Signed-off-by: YangFei1990 <feiw@nvidia.com> * allow scaled_bias + frozen weights in prequant path Signed-off-by: YangFei1990 <feiw@nvidia.com> * allow bias + frozen weight path in prequantized input Signed-off-by: YangFei1990 <feiw@nvidia.com> * use .copy to create group tensor Signed-off-by: YangFei1990 <feiw@nvidia.com> * move the implementation to c++ layer Signed-off-by: YangFei1990 <feiw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * bug fix for operation ordering Signed-off-by: YangFei1990 <feiw@nvidia.com> * add comprehensive tests Signed-off-by: YangFei1990 <feiw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename to group_requantize Signed-off-by: YangFei1990 <feiw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename func with inplace tag Signed-off-by: YangFei1990 <feiw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor the requantize function Signed-off-by: YangFei1990 <feiw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fix for test Signed-off-by: YangFei1990 <feiw@nvidia.com> --------- Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Use a temp directory for CMake build Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Add fallback to exisiting build dir behavior when using an editable install Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> --------- Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
# Conflicts: # tests/pytorch/test_torch_compile.py
* fix: preserve FP8 recompute state for inner autocast Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: fix activation recompute test license header Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: run FP8 recompute coverage in PyTorch QA Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: compare inner FP8 autocast recompute numerics Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * refactor: rename FP8 recompute flag to _IN_ACTIVATION_RECOMPUTE_REGION The global no longer encodes FP8 state now that the FP8 gate lives in is_fp8_activation_recompute_enabled(); it only marks the checkpoint region. Rename it to match. The public getter keeps its name since it now returns the conjunction of the region flag and the current FP8 state. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * refactor: rename recompute phase global for symmetry _FP8_ACTIVATION_RECOMPUTE_PHASE carries no FP8 state either - in_fp8_activation_recompute_phase() returns it ungated and callers apply their own FP8 gate. Rename to _ACTIVATION_RECOMPUTE_PHASE to match _IN_ACTIVATION_RECOMPUTE_REGION. Both public getters keep their names. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: restore dropped FP8 recompute coverage and assert stash bookkeeping Moving the regression into test_numerics.py silently dropped two of the three original tests. Restore both: the non-FP8 negative case (which needs no FP8 hardware and is the direct guard for making the region flag FP8-agnostic) and the mixed FP8/non-FP8 region case. The inner-autocast test could not observe a missing stash: the outer autocast is disabled, so autocast_depth never returns to 0 with FP8 enabled, reduce_and_update_fp8_tensors is never called, and the forward scale stays at 1.0 - making stashed and unstashed recompute identical. Record the stash/restore of the forward scale and assert every stash is restored exactly once, and use an inner-autocast reference so the two runs differ only by recompute. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: make recompute stash assertions meaningful, guard bf16 The scale equality assertion was vacuous: record_restore read the scale after the real restore had already copied the stashed value into it, so it compared a tensor with itself, and every scale is 1.0 here anyway. Count stash and restore per module instead, and check a module was stashed before the real restore runs so a regression reports that rather than a KeyError from the recompute buffer lookup. Also skip the non-FP8 checkpoint test when bf16 is unavailable (it has no FP8 skipif, so it would otherwise run on pre-Ampere), and correct the is_fp8_activation_recompute_enabled docstring, which still claimed to return a bare global. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: assert the recompute reinstalls the stashed FP8 state Counting stashes and restores catches a missing stash but not a wrong one. The outputs cannot catch either: the forward scale never moves here (autocast_depth stays >= 1 so reduce_and_update_fp8_tensors never fires), and a wrong scale would only perturb FP8 rounding regardless, since scale_inv is derived from the same scale at cast time. Compare the state each restore installs against a clone captured at stash time, for both the scale and the amax history, and assert at least one restore actually moved the live state so the comparison is not vacuous. This is exact rather than tolerance-based and works for both reentrant and non-reentrant checkpointing. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: drop tautological recompute state assertions Comparing the restored state against a clone taken in the stash wrapper could not fail: the real stash clones the same tensors at the same instant, and the restore copies that entry back, so both sides derive from one snapshot. It exercised the stash/restore plumbing rather than anything this change touches. Keep the checks that do have teeth - that some module stashed, and that stashes and restores balance per module - and skip the bookkeeping for modules whose recipe is not delayed scaling, since the restore site is not gated on it while the stash site is. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: pin the changed predicate directly Replace test_checkpoint_without_fp8_does_not_save_fp8_recompute_state, which could not fail: with no autocast self.fp8 is False, so the stash site is unreachable before the region flag is ever consulted, and the assertion holds no matter what the getter returns. The mixed-region test already makes the same negative claim in an FP8-enabled session. Assert instead on is_fp8_activation_recompute_enabled() itself from inside a checkpointed callable that opens its own autocast: True in both phases inside the autocast, False outside it. The first half fails before the fix; the second pins the FP8 term against a later over-correction that drops it. Also move the fp8_meta key constant up to the other module constants. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> --------- Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>
…lint (NVIDIA#3329) [JAX] Gate collective-stream compute_on on import and silence pylint not-callable Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…dSrelu (NVIDIA#3328) fix cudnn version gurad Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…ons (NVIDIA#3291) * [Common] Fix pointer arithmatic to generate correct LDS/STS instructions Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove redundant comments Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * move helper function to common.h to avoid NVRTC issue Signed-off-by: Kaining Zhong <kainingz@nvidia.com> --------- Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…on (NVIDIA#3338) * mxfp8: add swizzled-scale fast path for cast-only quantization Produce GEMM-ready scales directly in the specialized cast-only kernels to avoid separate scale-swizzle launches while preserving generic fallbacks for unsupported shapes. Signed-off-by: qiyuw <qiyuw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * initializes padded scale entries Signed-off-by: qiyuw <qiyuw@nvidia.com> --------- Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Enable runtime resolution of CUDA header path for NVRTC Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/common/util/cuda_runtime.cpp Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: fheinecke <23390735+fheinecke@users.noreply.github.com> * Fix resolution not working properly when package is installed outside of site-packages directory Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * fix cuda include dir resolution when installing as an editable package Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * fix linter failure Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> --------- Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> Signed-off-by: fheinecke <23390735+fheinecke@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
…#3270) * EP Dispatch with MXFP8 * Separate EP dispatch-forward and combine-backward quant recipes * Simplify EP quantize path and guard zero-copy pool alloc under CUDA graph capture * Add eager-mode coverage for MXFP8 EP combine backward * Supply persistent symm-mem recv buffers for zero-copy 1F1B CUDA-graph capture * Enable MXFP8 EP combine backward under zero-copy * Release symm-mem pool in ep_finalize * Skip MXFP8 EP tests on pre-Blackwell devices Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…IA#3318) FusedAdam.step() skipped a param group with no parameters before touching its step counter, so a group that is empty on one data-parallel rank and populated on another stopped counting on the empty ranks. Since step is stored in param_groups it is checkpointed, and a rank that loads its shard from a rank where the group was empty resumes with a stale step and a wrong bias correction. Move the counter update above the empty-group skip. Empty groups have no parameter to read a device from, so the capturable tensor now falls back to the device of the optimizer scratch buffer. Fixes NVIDIA#1986 Signed-off-by: Aditya Singh <adisin650@gmail.com> Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
NVIDIA#3281) * Fix NVFP4 stochastic rounding on architectures without cvt.rs The four e2m1 stochastic rounding helpers gate on ARCH_HAS_STOCHASTIC_ROUNDING, which is sm_100 and sm_103, where the cvt.rs instruction exists. The other branch is NVTE_DEVICE_ERROR, which is printf plus assert(0), and release builds define NDEBUG, so on sm_120 and sm_121 the quantize returns all-zero FP4 data while printing once per thread. Replace that branch with a software e2m1 stochastic rounder that takes 8 random bits per element, the same as the instruction's rbits operand, and follows cvt.satfinite at the edges. Architectures with cvt.rs keep the asm path. Signed-off-by: David Kogan <davidkny22@gmail.com> * Make the software e2m1 stochastic rounder branchless The rounding ladder becomes a floor in units of the grid step at the jittered value, and the isnan branch folds into the final sign select, since fminf carries a NaN through to max_norm and the select keeps it positive, matching cvt.satfinite. Output is bit-identical to the previous form across all 2^32 float bit patterns at 16 rbits values and across every float with magnitude in [1, 8) at all 256 rbits values. An isolated microbenchmark on sm_121 runs 1.34x faster. Signed-off-by: David Kogan <davidkny22@gmail.com> --------- Signed-off-by: David Kogan <davidkny22@gmail.com>
…IA#3344) * fix swizzle scale output shape in variable shape case Signed-off-by: YangFei1990 <feiw@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: YangFei1990 <feiw@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
bump nccl-extensions commit Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…support (NVIDIA#3224) * improve device grouped linear Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * proper support for grouped bias Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes and improvements Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> # Conflicts: # tests/pytorch/test_grouped_mlp.py * fix unit test failure Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix group linear UT Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix group mlp UT by disabling wrong fallback Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix CI errors Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * rename toggle resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * don't always set optimize_for_gemm=True Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * test other fp8 recipes on hopper Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix test failure Signed-off-by: zhongboz <zhongboz@nvidia.com> * relax atol for bf16 sum errors Signed-off-by: zhongboz <zhongboz@nvidia.com> * fix CI Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * continue to fix edge cases Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * unify grouped tensor check Signed-off-by: zhongboz <zhongboz@nvidia.com> * special case handling for all empty moe inputs for fp8 CS Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * add another specific case guard Signed-off-by: zhongboz <zhongboz@nvidia.com> * fix cublas setup meta data issue for empty inputs Signed-off-by: zhongboz <zhongboz@nvidia.com> * mcore integration fixes Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * chore fix after rebase Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * chore Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * chore fix after rebase Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * resolve comments, enable save original input Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix cudnn version gurad Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * update benchmark script Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: zhongboz <zhongboz@nvidia.com> Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com>
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Squashed to single commit for review. Original PR: andrewwhitecdw#14 Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
* Gate FA4 and stabilize attention test imports FA4 can be installed on SM8x even though its current implementation rejects those GPUs. Disable selection and skip dedicated FA4 tests there so A100 and L40 use supported attention backends. FA4 and CUTLASS can also expose a generic utils package on sys.path. Prepend the Transformer Engine test helper directory in the context-parallel test so collection resolves the intended utilities. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Isolate FlashAttention CI backends Moving images can install FA4 alongside older FlashAttention generations, which mixes a shared Python namespace and can make context-parallel reference runs compile an unsupported backend. Isolate the L3 version matrix, keep current CP comparisons on FA2/FA3, and temporarily reject symmetric D512 FA4 on Blackwell until upstream kernel support is complete. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Prepend PyTorch test utility imports FA4 and its CUTLASS dependency expose a top-level utils module after Transformer Engine imports. Appending the test root can therefore bind these late imports to the installed module and fail collection. Give the repository helper precedence in the four test files that exhibited this ordering. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Recognize FA3 sliding-window CP support The all-gather and a2a guards use the FA2 package version check to recognize FlashAttention support, so an isolated FA3 run is rejected even though FA3 implements sliding-window attention. Accept the explicit FA3 backend in both guards. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Keep CP backend selection at CI boundaries The CP runner must honor an explicit generation selected by its caller, particularly the existing B200 L3 FA4 lane. Remove its internal V4 override, restore the L3 SM100 selection changed in 0f6c71e, and disable V4 only for the L1 suite that still targets FA2/FA3. This keeps per-generation L3 isolation intact without making the shared runner silently override directed coverage. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Align FlashAttention CI coverage by architecture Keep L0 on the mature FA2 path while L3 owns newer-generation coverage. Restrict H100 L3 to FA3 and B200 L3 to non-CP FA4 so unsupported H100 FA4 kernels and mislabeled Blackwell CP results do not obscure the intended signal. Make FA4-specific tests honor backend enablement to prevent silent fallback under an FA4 label. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Guard FA4 against incompatible CUTLASS installs Package metadata can report FA4 present even when a later dependency install leaves its transitive CUTLASS stack unusable. Reject the known b24/CUTLASS combination below the stable 4.6.2 release and treat a nested interface ImportError as an unavailable optional backend so unrelated Transformer Engine imports can continue. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> --------- Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Resolved conflict in backends.py: keep upstream's cutlass-dsl version guard and ImportError fallback for FlashAttention 4, with our no_torch_dynamo wrapping applied in the success branch. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The unconditional early return skipped main's new CustomRecipe local-recipes inference when quantization is off. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
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>
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>
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
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
DotProductAttention.forwardno longer carries@no_torch_dynamo, sotorch.compilecaptures the whole module -- input unpacking, qkv layout, backend selection and the backend itself -- withfullgraph=Trueon the FlashAttention and UnfusedDotProductAttention backends.FusedAttention stays an eager island (
@no_torch_dynamoon its own forward): the graph break moves from the entry ofDotProductAttentiondown to the backend call, so everything around it is now compiled.Follow-up to NVIDIA#3189 (traceable
get_attention_backend), NVIDIA#3200 (declarative packed QKV), NVIDIA#3201 (unfused DPA) and NVIDIA#3268 (device-side padding mask).Type of change
Changes
Removing the decorator exposed several things that could not be traced. Each is fixed at the source rather than worked around:
DotProductAttention.init_fp8_metadatacalledget_fp8_recipe()unconditionally, which builds a defaultRecipewhen no autocast is active -- and constructing one is not traceable. With FP8 off it now returns through the base class, which is where the rest of that function ended up anyway.get_qkv_layoutrecognizes packed q/k/v fromUntypedStorage.data_ptrandTensor.storage_offset, neither of which dynamo can trace, and neither of which means anything on the fake tensors used while tracing. While tracing, the layout comes from the format instead; packed inputs are declared throughqkv_layer/kv_layer([PyTorch][torch.compile] DotProductAttention: declarative packed QKV/KV inputs NVIDIA/TransformerEngine#3200). ONNX export keeps its own path -- the exporter runs through dynamo, so it setsis_compiling()as well.get_indices, which packsbshd/sbhdinputs for FlashAttention's varlen entry point, built the index list with a Python comprehension over sequence lengths held in a GPU tensor: a device synchronization per sequence in eager, and data-dependent while tracing. It is now built on the device, verified exhaustively against the previous implementation (10304 cases, 0 mismatches).get_full_cu_seqlenskeyed its cache ontorch.is_inference_mode_enabled(), a fundamental graph break. The cache only saves an allocation, so while tracing the tensor is built directly.The fused sub-backend enum does not survive a graph break.
_get_fused_attn_backendis decorated with@torch.compiler.assume_constant_result, and dynamo reconstructs such a value by re-emitting the call that produced it -- which is only valid inside the frame that made it. Crossing the break intoFusedAttention, the resumed frame received the producing function instead of the enum member, and it reachedfused_attn_fwd.get_attention_backendnow returns the sub-backend as a plain int, which is baked into the graph as a literal;FusedAttnBackend.castturns it back into a member, asfused_attn_fwd/bwdalready did.Backend selection is logged through the no-op logger while tracing, as
get_attention_backenddoes:logging.Loggermethods graph-break.no_torch_dynamotakes awhenpredicate, so a method can be an eager island for some calls only. The predicate receives the call's arguments keyed by parameter name and returns the reason a particular call cannot be traced.Falling back to eager
Three kinds of call cannot be compiled and run as an eager island instead, each with a warning naming what to do differently:
fp8_dpa/fp8_mha)fp8_metamutation andFloat8Tensorcrossing the graph boundary, none of it covered by testsqkv_layer/kv_layerthdwithoutmax_seqlen_q/max_seqlen_kvcu_seqlens: a device synchronization, and a data-dependent value while tracingmax_seqlen_q/max_seqlen_kvThe predicate is deliberately narrow for FP8: it asks the recipe for
fp8_dpa/fp8_mharather than whether an autocast is active, so FP8 GEMMs with attention in high precision -- the common training setup -- stay on the compiled path.Without
fullgraph, these calls work and warn. Withfullgraph=Truethey raise, which is what asking for a full graph means.Testing
tests/pytorch/test_torch_compile.py, 56 passed / 23 skipped on RTX Ada (sm89, flash-attn 2.7.4):ModelConfigthe eager attention tests use, against each backend that supports them, whichget_available_attention_backends()decides. It coversbshd/sbhd/thd, causal / no_mask / padding / sliding window / arbitrary masks, GQA, cross attention, alibi, post-scale bias, MLA head dims, off-by-one softmax, one sharedcu_seqlenstensor for q and kv, and all four declared packed layouts. Compiled withfullgraph=True, compared against eager in forward and every input gradient;mode="reduce-overhead") -- fresh inputs each iteration and compared against eager, since a replay reading stale buffers produces finite values and non-None gradients too; inductor's skip counter is asserted to be zero, so a declined capture cannot pass silently;fullgraph=True.Backend selection is invalidated right before every compiled call, so it is traced rather than served from the cache the eager call populated.
Comparison is exact for backends whose kernel is the same either way. The unfused backend is built from PyTorch ops that inductor fuses and reassociates, so its absolute tolerance is taken from the tensor's own scale -- an off-by-one softmax differs by two bfloat16 roundings on individual elements, which no elementwise relative tolerance covers near zero.
attention/test_attention.py,attention/test_kv_cache.py,test_gqa.pyandtest_onnx_export.pyshow no new failures relative tomain.Not covered by this hardware and wanted from CI: FP8 and MXFP8 attention, the fused sub-backends, FlashAttention 3 and 4.
Known limitations
tex.fused_attn_fwd/bwdas custom ops -- which is what lets flash-attn's kernels sit inside a graph -- would be the way there.Checklist:
🤖 Generated with Claude Code