Skip to content

[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends - #22

Open
pggPL wants to merge 70 commits into
mainfrom
dpa_torch_compile
Open

[PyTorch][torch.compile] Support for DotProductAttention on flash and unfused backends#22
pggPL wants to merge 70 commits into
mainfrom
dpa_torch_compile

Conversation

@pggPL

@pggPL pggPL commented Jul 27, 2026

Copy link
Copy Markdown
Owner

Description

DotProductAttention.forward no longer carries @no_torch_dynamo, so torch.compile captures the whole module -- input unpacking, qkv layout, backend selection and the backend itself -- with fullgraph=True on the FlashAttention and UnfusedDotProductAttention backends.

FusedAttention stays an eager island (@no_torch_dynamo on its own forward): the graph break moves from the entry of DotProductAttention down 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

  • New feature (non-breaking change which adds functionality)
  • Bug fix (non-breaking change which fixes an issue)

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_metadata called get_fp8_recipe() unconditionally, which builds a default Recipe when 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_layout recognizes packed q/k/v from UntypedStorage.data_ptr and Tensor.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 through qkv_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 sets is_compiling() as well.

  • get_indices, which packs bshd/sbhd inputs 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_seqlens keyed its cache on torch.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_backend is 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 into FusedAttention, the resumed frame received the producing function instead of the enum member, and it reached fused_attn_fwd. get_attention_backend now returns the sub-backend as a plain int, which is baked into the graph as a literal; FusedAttnBackend.cast turns it back into a member, as fused_attn_fwd/bwd already did.

  • Backend selection is logged through the no-op logger while tracing, as get_attention_backend does: logging.Logger methods graph-break.

  • no_torch_dynamo takes a when predicate, 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:

why how to stay compiled
FP8 attention (fp8_dpa/fp8_mha) quantizers, fp8_meta mutation and Float8Tensor crossing the graph boundary, none of it covered by tests --
packed q/k/v passed as plain views recognizing them takes the data pointers dynamo cannot read declare the packing via qkv_layer/kv_layer
thd without max_seqlen_q/max_seqlen_kv deriving it reads the sequence lengths off cu_seqlens: a device synchronization, and a data-dependent value while tracing pass max_seqlen_q/max_seqlen_kv

The predicate is deliberately narrow for FP8: 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.

Without fullgraph, these calls work and warn. With fullgraph=True they 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):

  • the configuration matrix -- 14 configurations described with the same ModelConfig the eager attention tests use, against each backend that supports them, which get_available_attention_backends() decides. It covers bshd/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 shared cu_seqlens tensor for q and kv, and all four declared packed layouts. Compiled with fullgraph=True, compared against eager in forward and every input gradient;
  • CUDA graphs (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;
  • the compiled path around FusedAttention -- it is not compiled itself, but everything before and after the break is, and what crosses that break has to survive it;
  • the eager fallbacks -- that each fires exactly when it should, matches eager, and raises under 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.py and test_onnx_export.py show no new failures relative to main.

Not covered by this hardware and wanted from CI: FP8 and MXFP8 attention, the fused sub-backends, FlashAttention 3 and 4.

Known limitations

  • FusedAttention is not compiled; it is skipped from the matrix, and removing that skip is all that is needed once its forward traces. Registering tex.fused_attn_fwd/bwd as custom ops -- which is what lets flash-attn's kernels sit inside a graph -- would be the way there.
  • The unfused backend is compared to eager with a tolerance rather than exactly, for the reason above.

Checklist:

  • I have read and followed the contributing guidelines
  • The functionality is complete
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation
  • My changes generate no new warnings
  • I have added tests that prove my fix is effective or that my feature works
  • New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

@pggPL
pggPL requested a review from cyanguwa as a code owner July 27, 2026 09:13
@pggPL pggPL closed this Jul 27, 2026
@pggPL pggPL reopened this Jul 27, 2026
pggPL added 4 commits July 29, 2026 08:52
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>
@pggPL
pggPL force-pushed the dpa_torch_compile branch from 4a1c288 to 1cd6704 Compare July 29, 2026 06:56
@pggPL pggPL closed this Jul 29, 2026
@pggPL pggPL reopened this Jul 29, 2026
pggPL added 20 commits July 29, 2026 09:04
- 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>
xiuhu17 and others added 30 commits August 5, 2026 11:33
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.