Repository navigation
IFU 2.20 - #756
Merged
Merged
IFU 2.20#756
Conversation
… unfused backends (#3286)
* [PyTorch][torch.compile] Support for DotProductAttention
Make DotProductAttention itself traceable, so that torch.compile captures
the whole module -- input unpacking, qkv layout, backend selection and the
backend -- on the FlashAttention and UnfusedDotProductAttention backends.
FusedAttention stays an eager island, as before.
DotProductAttention.forward no longer carries @no_torch_dynamo. What that
uncovered, and the fixes:
- init_fp8_metadata unconditionally called get_fp8_recipe(), which builds a
default Recipe when no autocast is active -- not traceable. It now returns
early when FP8 is off everywhere, which is also what the rest of the
function did in that case.
- get_qkv_layout detects packed q/k/v from data pointers and storage
offsets, neither of which dynamo can trace (and neither is meaningful on
fake tensors). While tracing, the layout is built from the format instead;
packed inputs are declared explicitly via qkv_layer/kv_layer, and
non-contiguous q/k/v are made contiguous so the layout stays truthful.
- get_full_cu_seqlens keyed its cache on torch.is_inference_mode_enabled(),
a fundamental graph break; while tracing it builds the tensor directly.
- get_padding_mask read sequence lengths on the host (a device sync, and
data-dependent under torch.compile). It is now vectorized.
- The fused sub-backend enum member does not survive a graph break: dynamo
reconstructs the assume_constant_result value by re-emitting the call that
produced it, so the resumed frame received the function itself instead of
the enum -- which then reached fused_attn_fwd. get_attention_backend now
keeps a plain int while tracing and casts back to FusedAttnBackend in
eager mode only.
- Backend selection logging is skipped while tracing (logging.Logger methods
graph-break, and int() of the sub-backend enum is untraceable too).
Tests: test_dpa_torch_compile covers 7 configurations (bshd/sbhd/thd,
causal/no_mask/sliding window, GQA, cross attention, packed qkv_layer)
against both backends with fullgraph=True, comparing forward and backward
against eager; plus a CUDA-graphs (reduce-overhead) test and a test that the
FusedAttention path stays correct around its graph break.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Delegate the no-FP8 case of DPA's init_fp8_metadata to the base class
The override hand-rolled the state that TransformerEngineBaseModule already
sets when FP8 is off everywhere (the three flags, fp8_checkpoint and
fp8_initialized). Call super() instead, and read the flags straight off
FP8GlobalStateManager.quantization_state, as the other hot paths do.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Run DPA eagerly for undeclared packed q/k/v instead of guessing the layout
get_qkv_layout recognizes packed q/k/v by inspecting data pointers and storage
offsets, which dynamo cannot trace. Rather than assume such inputs are three
separate tensors -- a layout that would not describe memory -- detect the case
from what is traceable (a strided tensor that does not own its whole storage)
and run the whole module as an eager island, with a warning. Declaring the
packing via qkv_layer/kv_layer keeps the call on the compiled path.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Log backend selection in eager only; shorten comments
Log arguments are evaluated regardless of the logger, so the no-op logger
from #3189 is not enough here -- skip the whole block while tracing.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Fix lint: unused import and keyword-arg-before-vararg
- utils.py no longer needs no_torch_dynamo: the eager island moved to
dot_product_attention.py.
- _needs_eager_dpa reads q/k/v out of *args/**kwargs instead of naming them
before *args, which pylint flags (W1113).
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Compare against eager under CUDA graphs, and tighten the tolerances
The CUDA-graphs test only checked that the output was finite and that
gradients existed. A replay that reuses stale buffers -- the failure mode
that test is there to catch -- satisfies both, so it now compares against
eager on fresh inputs every iteration.
The remaining comparisons used the dtype defaults (rtol=1.6e-2 for bfloat16),
which is far looser than what these paths deliver: compiled and eager agree
bit-for-bit on every config and backend. FlashAttention now has to match
exactly, as it runs the same kernel either way; the unfused backend is
allowed a single bfloat16 rounding, since inductor may reassociate the
softmax sums.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Support one cu_seqlens tensor for both q and kv under torch.compile
Self attention on thd naturally passes the same tensor as cu_seqlens_q and
cu_seqlens_kv, and flash-attn's varlen entry point forwards both to the same
autograd.Function -- which dynamo refuses to trace ("duplicate tensor input"),
failing the compilation under fullgraph=True.
Hand the varlen call a copy of one of them while tracing. The clone is
b + 1 int32 elements, only on that path, and only under torch.compile.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Apply the shared-cu_seqlens workaround to the FlashAttention v4 path too
Pull the deduplication into a helper and use it wherever both cumulative
sequence lengths are handed to flash-attn: the varlen path shared by v2 and
v3 training, and v4's varlen kwargs. The v3 KV-cache path derives the kv
lengths by subtraction, so its two arguments are distinct by construction.
v4 is not installed in the environment this was tested in, so that call site
is covered by construction rather than by a test; the helper is a no-op
unless the two arguments are the same object.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Run FP8 attention eagerly under torch.compile
FP8 attention brings quantizers, fp8_meta mutation and Float8Tensor across
the graph boundary, and none of it is exercised by this PR's tests. Rather
than trace it untested, route it to the eager island that packed q/k/v
already use.
The predicate is deliberately narrow: it asks the recipe for fp8_dpa/fp8_mha
rather than whether an autocast is active, so FP8 GEMMs with attention in
high precision -- the common training setup -- stay on the compiled path.
The decorator now takes a predicate that returns the reason, so the warning
names which of the two paths was taken.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Drive the compile tests off ModelConfig instead of a hand-picked intersection
The configurations were chosen so that both backends could run all of them,
which meant nothing backend-specific was ever exercised: no biases, no
arbitrary masks, no MLA head dims, no sink softmax. Describe them with the
ModelConfig the eager attention tests already use, and let
get_available_attention_backends() decide which backend runs which, so a
configuration only one backend supports is covered rather than avoided.
Five such configurations are added. One of them immediately showed that the
tolerance for the unfused backend was too tight: an off-by-one softmax
differs from eager by two bfloat16 roundings on a single element out of
65536, because inductor reassociates the softmax sums. Only FlashAttention,
which runs the same kernel either way, keeps the exact comparison.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Scale the unfused comparison to the tensor, and take tolerances from utils
An off-by-one softmax exposed that comparing the unfused backend elementwise
does not hold up: inductor's reassociation of the softmax sums produces an
error proportional to those sums, which for gradient elements near zero is
far outside any relative tolerance -- while being one bfloat16 rounding at
the tensor's own scale.
Derive the absolute tolerance from the reference tensor's magnitude, and take
the dtype tolerances from tests/pytorch/utils.py rather than restating them.
FlashAttention and FusedAttention keep the exact comparison.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Merge the per-backend compile tests, and give padding masks real padding
The configurations built cu_seqlens with every sequence at the maximum length,
so a padding mask was all-False: the padding branches ran, but on data that
could not tell a difference. They now use varying sequence lengths, which
immediately showed that FlashAttention on bshd/sbhd with a padding mask does
not trace: it packs the tensors first, and get_indices built the index list
with a Python comprehension over sequence lengths held in a GPU tensor -- a
host synchronization per sequence in eager, untraceable while compiling.
It is now built on the device, verified exhaustively against the previous
implementation (10304 cases, 0 mismatches) and free of synchronization.
The separate tests for one shared cu_seqlens tensor and for the fused backend
folded into the main one: the first is a configuration, and the second is the
backend axis, which now runs flash, fused and unfused with the graph break
that FusedAttention's eager island implies.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Run eagerly when max_seqlen has to be derived from cu_seqlens
On thd inputs without max_seqlen, DotProductAttention reads the sequence
lengths off cu_seqlens to derive it. Unlike the padding mask and the packing
indices, this one cannot be moved to the device: max_seqlen sizes tensors
downstream, so it has to be a Python int, and reading it is a device
synchronization -- a data-dependent value dynamo cannot trace. Route those
calls to the eager island, so passing max_seqlen is what keeps a call
compiled.
Arguments are now looked up by name or position, so the predicate sees them
whichever way the caller passed them.
The two eager-fallback tests merge into one: they set up different inputs but
assert the same thing -- a warning, results matching eager, and an error under
fullgraph.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Share the run-and-compare code between the compile tests
Both tests ran a callable, went backward, copied the gradients out and cleared
them. That is now _run_and_capture(), with _assert_run_matches() comparing two
of its results, which halves the CUDA-graphs test.
The two stay separate tests: what CUDA graphs risk -- a replay reading stale
buffers -- does not depend on the attention configuration, so folding the mode
into the configuration matrix would multiply the cases without covering
anything new, and would report a replay bug as a failure of some unrelated
configuration.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Use the shared run-and-compare code in the remaining compile tests
The eager-fallback test hand-rolled the same run, backward and gradient copy.
The backend-level unfused test only checked that the output was finite and
that gradients existed, which a CUDA-graph replay reading stale buffers also
satisfies -- the weakness the DotProductAttention CUDA-graphs test was given
an eager comparison for. It now compares against eager as well.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Read the predicate's arguments by name, not by hardcoded position
The positions of forward's parameters were written out as a table of indices,
which anyone inserting a parameter into a thirty-argument signature would
silently invalidate: the predicate would go on reading whatever now sits at
index 5 or 10 and decide the fallback on it, with nothing to notice.
The decorator already has the function it wraps, so it takes the parameter
names from its signature and hands the predicate the call's arguments keyed by
name.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Move the eager-fallback decorator to jit.py
Nothing about it is specific to attention: it takes a predicate, and runs the
wrapped method as an eager island whenever that predicate gives a reason. It
belongs next to no_torch_dynamo, which it builds on, rather than buried in the
attention module where the modules that will need it next -- MultiheadAttention
and TransformerLayer, both of which still have blockers of their own -- would
not find it.
What stays in the attention module is the predicate, which is where the
attention-specific knowledge lives.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Skip FusedAttention rather than special-casing it in the compile tests
The test compiled every backend with fullgraph=True except fused, which was
passed fullgraph=False because its forward is an eager island. That put an
assumption about one backend into the comparison logic.
FusedAttention is skipped instead, next to the configurations a backend cannot
run, with the reason spelled out. Everything else is compiled the same way, and
supporting fused later is deleting the skip.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Cover the compiled path around FusedAttention
FusedAttention is skipped from the backend matrix because it does not compile,
but it is the default on Hopper and Blackwell, and this change does affect it:
what used to be one eager island is now a traced prologue, a graph break at the
backend, and a traced remainder. Whatever crosses that break has to survive it,
which the sub-backend enum did not.
One test covers it, without a fullgraph and outside the matrix. Verified to
fail with the enum restored -- TypeError: int() argument must be ... not
'function', from cuDNN's entry point.
Also rename _distinct_cu_seqlens: it returns its arguments untouched in eager
and whenever they already differ, so it promised more than it does.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Log backend selection through the no-op logger, as get_attention_backend does
Skipping the block with is_compiling() introduced a second mechanism for a
problem this module already solves one way: get_attention_backend swaps in a
no-op logger while tracing. The skip was needed while the sub-backend argument
was an enum that did not survive tracing, and is not any more.
The logger, which was private to the utils module, is now named without the
underscore, since it is used from outside it.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Return the fused sub-backend as an int, in eager as well
get_attention_backend cast it back to a FusedAttnBackend member unless it was
tracing, which made its return type depend on the execution mode -- something
every reader of that function then has to keep in mind, for no benefit inside
transformer_engine: fused_attn_fwd/bwd call cast() on whatever they are given,
and every comparison against a member works with an int on either side.
It now returns an int always, and the docstring says so.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Rename eager_under_compile_if to fallback_to_eager_when
The old name read as a run-on, and its '_if' suffix suggested a boolean
condition where the argument actually returns the reason. The new one reads as
a sentence where it is used and matches the warning the decorator emits.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Fold the eager fallback into no_torch_dynamo as a `when` predicate
It was a decorator of its own, which meant a second name for what is the same
thing this module already does: disable dynamo for a frame. It is now the
conditional form of no_torch_dynamo -- @no_torch_dynamo(when=predicate) -- so
there is one decorator to find and one docstring to read.
The predicate receives the call's arguments keyed by parameter name and returns
the reason this call cannot be traced, or None to have it traced.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Cover the declared packed layouts beyond bs3hd
Packing was pinned to one layout by an assert in the input builder, so of the
twelve declarations DotProductAttention accepts -- qkv_layer or kv_layer,
interleaved at -3 or -2, in each of the three formats -- exactly one was
tested. The builder now takes what is packed and where it is interleaved, and
derives the shapes and the layout string the way DotProductAttention does.
Three configurations are added, so both packed inputs and both interleave
dimensions are covered.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Let ONNX export keep its own path in get_qkv_layout
The ONNX exporter runs through dynamo, so torch.compiler.is_compiling() is
true during an export as well -- measured: inside torch.onnx.export(dynamo=True)
both it and is_in_onnx_export_mode() are set. The compile branch was written
first and took the export over: the layout came from the format instead of the
detection ONNX relies on, tensors were no longer made contiguous, and packed
views hit an assert where they used to be copied.
Export mode is the narrower context, so it is checked first. This restores the
previous behaviour for export while leaving the compiled path as it was.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Guard the assumption the argument binding rests on
Arguments are matched to parameter names by position, which is only right
while the signature has no *args in between -- and would misalign silently if
one were added. It is checked when the decorator is applied.
The docstring now also says that an argument the call left out is simply
absent, so a predicate reading it with .get() sees its default only as long as
that default is None.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Check that CUDA graphs were actually captured, and generalize a docstring
fullgraph=True says dynamo did not break the graph, but inductor can still
decline to capture -- a mutated input, a CPU scalar -- and run the thing
normally, which would leave the CUDA-graphs test passing while measuring
nothing. Its skip counter is asserted to be zero. Measured on this branch:
nothing is skipped, the tree has one root, and interleaving the eager
reference call between replays does not disturb it.
The tolerance docstring described which backends are compared exactly by
naming them; it now states the rule -- a backend whose kernel is the same
either way -- so it still holds once FusedAttention traces.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Skip lazy compilation while already tracing
Filling lazy_compile's closure cell is a side effect Dynamo rejects inside
a higher-order op, so compiling a module whose autograd.Function backward
makes the first call to a jit_fuser'd function fails with
Unsupported: HOP: Unsafe side effect
Attempted to mutate CellVariable()
An earlier eager run hides this by filling the cell beforehand. The nested
torch.compile is inlined into the outer graph either way, so calling the
function directly while tracing is equivalent.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Run context parallel attention eagerly
The context parallel implementation uses P2P communication and its own CUDA
stream, neither of which dynamo traces, and FlashAttention reaches it from a
forward that is no longer excluded from the graph.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Keep DotProductAttention compilable under a CUDA RNG states tracker
Forking the tracker swaps the global CUDA generator state, which dynamo
refuses to trace. With attention_dropout == 0 the dropout draws nothing
from the generator, so the fork is a no-op -- use nullcontext instead and
keep the call in the graph. Calls that do use dropout under a tracker run
as an eager island.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Run checkpointed attention eagerly
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Support KV caching under torch.compile
Wrap the KV cache kernels in custom ops: dynamo cannot trace pybind calls
into transformer_engine_torch, so a decoding step used to break the graph.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Compile the fused sbh3d QKV split
Wrap fa_prepare_fwd/bwd like the KV cache kernels, and cover the layout
that reaches them: no configuration exercised it, compiled or eager.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Run FlashAttention 4 eagerly
Version 4 registers no custom ops, unlike 2 and 3: it builds its kernels
through the CUTLASS DSL as it runs. Disable dynamo on its entry points, so
only the calls that reach it become eager islands.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Pin the sequence lengths backend selection bakes in
A second sequence length makes dynamo's automatic dynamic turn max_seqlen
symbolic, and _get_fused_attn_backend, being assume_constant_result, cannot
take a symbolic argument -- so the call died once the length changed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Capture CUDA graphs with a KV cache
The cache is mutated in place, so inductor declined to capture until its
buffers were marked static.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Test generation against a KV cache
A prefill and three single-token steps, each compared to eager, with and
without CUDA graphs: the cache carries state, so a step that updates it
wrongly only shows in the step after.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Drop the imports the custom ops made unused
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Rename _maybe_unshare_cu_seqlens to _unalias_cu_seqlens
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Split the get_full_cu_seqlens cache bypass per reason
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Skip get_fp8_recipe only while tracing in init_fp8_metadata
The unconditional early return skipped main's new CustomRecipe
local-recipes inference when quantization is off.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Skip the max_seqlen bake-in under ONNX export
guard_scalar specializes user-declared dynamic sequence dims in the
exported model; the exporter skips backend selection, which the bake-in
is for.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Make undeclared non-contiguous q/k/v contiguous under compile
The declared layout is only truthful for contiguous tensors. Eager
repairs such inputs via its detect-and-retry loop; mirror that in the
compiled branch instead of declaring a layout the memory contradicts.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---------
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
…tests (#3398) Support synchronous collectives in HLO collective bytes assert tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com>
… (#3373) fix: SBHD reorder skip uses original shape instead of swapped tensor Use tensor.shape[seq_dim] instead of shape[seq_dim] when deciding whether a Striped SBHD case is large enough. Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com>
* Refactor GroupedLinear quantization dispatch Signed-off-by: Evgeny <etsykunov@nvidia.com> * Minor fixes Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename test alias Signed-off-by: Evgeny <etsykunov@nvidia.com> * fix lint Signed-off-by: Evgeny <etsykunov@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…348) * [PyTorch] Fix MXFP8 master weight cast when a rank owns no shard. Signed-off-by: Ritesh Patel <ripatel@nvidia.com> * [PyTorch] Destroy the process group created by the MXFP8 empty-shard test. Signed-off-by: Ritesh Patel <ripatel@nvidia.com> * [PyTorch] Add a 2-rank MXFP8 empty-shard test where one rank owns the master shard. Signed-off-by: Ritesh Patel <ripatel@nvidia.com> --------- Signed-off-by: Ritesh Patel <ripatel@nvidia.com> Co-authored-by: Evgeny Tsykunov <etsykunov@nvidia.com>
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com> Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
* [Common][PyTorch] Add QB router histogram paths Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [Common][PyTorch] Harden QB router validation Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [Common][PyTorch] Make QB bounds validation recoverable Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [Common][PyTorch] Validate QB bounds in common APIs Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [Common][PyTorch] Write QB cutoff in both modes Signed-off-by: Harry Zhou <hhanyu@nvidia.com> --------- Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
* Bump min compute_on version Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Bump to 0.11.1 Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Fix build issues with 26.08 JAX container images Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Update test to better cover change Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> * Remove backwards compat fix for jax 0.10.1 and 0.11.0 (see NVIDIA/TransformerEngine#3406) Signed-off-by: Fred Heinecke <fheinecke@nvidia.com> --------- Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
* Preserve quantized tensor subclasses on detach Signed-off-by: Dingqing Yang <dingqingy@nvidia.com> * Fix IdentityTensor detach alias assertion Signed-off-by: Dingqing Yang <dingqingy@nvidia.com> * Support quantized tensor subclasses in C++ bindings Signed-off-by: Dingqing Yang <dingqingy@nvidia.com> * Centralize quantized tensor detach Signed-off-by: Dingqing Yang <dingqingy@nvidia.com> --------- Signed-off-by: Dingqing Yang <dingqingy@nvidia.com>
…cs (#3409) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
* Surface distributed comm-overlap rank errors Comm-overlap launchers captured child stdout and stderr but raised only stderr, which could leave the originating rank context out of pytest and JUnit failures. Preserve the existing result predicates while attaching bounded tails of both streams, and record layer-worker exceptions so torchrun can report the rank-local traceback. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Isolate comm-overlap launcher environment The launcher mutated os.environ and then called os.unsetenv, which leaves Python's environment mapping unchanged. Conditional backend flags could therefore leak into later parameterized children. Build a child-only environment instead so each launch gets its intended overrides while preserving the parent process environment. 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 * Keep FA4 out of L1 distributed tests The moving PyTorch image installs FA4 by default, implicitly expanding an L1 suite that historically covered earlier attention backends. Export the FA4 selector at the suite boundary so every pytest and torchrun child retains the intended backend scope while dedicated attention suites own FA4 coverage. 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>
* [Common][PyTorch] Add SiTU-GLU activation Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Remove unsupported SiTU-GLU recompute option Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Reject unsupported scaled GLU recompute Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Complete grouped SiTU-GLU fusion coverage Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Consolidate grouped SiTU-GLU detection Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [Common] Normalize gated activation derivatives Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * Avoid implicit cast from bool to float Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * [Common] Fix scaled activation scale-gradient tests Signed-off-by: Harry Zhou <hhanyu@nvidia.com> --------- Signed-off-by: Harry Zhou <hhanyu@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
…nd JAX (#3396) * Add CI workflow building TE against torch 2.1 Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Remove MAX_JOBS=1 from pytorch 2.1 build workflow Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Use MAX_JOBS=2 in pytorch 2.1 workflow, unbounded build OOMs the runner Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Point TE build at apt cudnn9, torch 2.1 pip deps shadow it with cudnn 8.9 Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Avoid passing std::optional to at::get_generator_or_default, torch 2.1 needs c10::optional Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Gate NCCL EP in torch extension on nccl_ep_enabled, matching common CMake Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Pass comm streams as raw cudaStream_t handles, torch 2.1 pybind lacks c10::Stream caster Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Revert "Pass comm streams as raw cudaStream_t handles, torch 2.1 pybind lacks c10::Stream caster" This reverts commit 04b4b0f52843ede580a911059875fd802cee7cf1. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Revert "Avoid passing std::optional to at::get_generator_or_default, torch 2.1 needs c10::optional" This reverts commit 061834fc3484755a92ca82a553be8bae86fdba44. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Retarget CI workflow to torch 2.8 Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Rename workflow to Minimum supported PyTorch, parametrize torch version Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Add Minimum supported JAX job, rename workflow to minimum_versions Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Add packaging to min-jax job deps Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Skip NCCL EP in torch extension when torch lacks symm-mem headers; build min-torch job on sm90 Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fix min-versions jobs: pydantic for jax import, libcuda stub for sm90 import Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Extract torch symm-mem EP gate into nccl_ep_supported_by_torch helper Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Tidy nccl_ep_supported_by_torch Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Compile out EP zero-copy path when torch lacks symm-mem headers Per review: only the zero-copy path needs torch's c10d symm-mem (2.11); guard the includes with __has_include so the rest of EP builds on older torch, expose ep_zero_copy_supported() and raise in ep_bootstrap when zero_copy=True without support. Drop the build-time torch gate. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…in HLO (#3412) * [JAX] Fix counting of synced and wrapped collectives in HLO Signed-off-by: Alex Y. Chan <alechan@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alex Y. Chan <alechan@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* [Common] Make ptx.cuh arch agnoistic by moving arch specific helpers to a new header 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 * whoops Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * Update transformer_engine/common/util/ptx.cuh Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Kaining Zhong <44538064+kainzhong@users.noreply.github.com> --------- Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Kaining Zhong <44538064+kainzhong@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
* [Common/PyTorch] Fused grouped MXFP8 requantization Replace the group_dequantize -> group_quantize(columnwise) -> grouped_swizzle(rowwise scales) chain in group_requantize_inplace with a single kernel (NVTE_FUSED_GROUP_REQUANTIZE=0 restores the unfused path). Co-authored-by: Oleg Goncharov <ogoncharov@nvidia.com> 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 * fix zero tensor dim Signed-off-by: YangFei1990 <feiw@nvidia.com> * rename func; read less attrs 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: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Signed-off-by: Michał Marcinkiewicz <43240942+mmarcinkiewicz@users.noreply.github.com>
…uncher pass count (#3415) [PyTorch] Scale test_multi_process_ep outer timeout with launcher pass count Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> Co-authored-by: fheinecke <23390735+fheinecke@users.noreply.github.com>
…Fused Adam. (#3414) * Add the support of variable dtypes of params and grads in capturable Fused Adam. Signed-off-by: Przemek Tredak <ptredak@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add support in lower level JAX API for returning max logit and softmax aux to the user from TE JAX fused attn output Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add support for returning reduced per head max logit. Plumb max logit and softmax through the JAX fused attn primitives Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add max logit to JAX fused attn FFI and set it in the workspace Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Add first pass tests for max logit and softmax aux tensor outputs in JAX fused attn tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Reject aux returns with score_mod Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Handle SM120 max-logit layout Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Drop softmax aux return Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Modify static args in fused attn tests for jax Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Inline the choice of what is to be returned and remove redundant function for it Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Expose JAX max logit Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Support CP max logit Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Expand JAX max logit tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Remove JAX max logit integration tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Broaden JAX CP max logit tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Reduce JAX max logit across DP Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Simplify JAX max logit return Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Document JAX max logit reductions Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Refine JAX max logit tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Rename JAX max logit buffer Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Check JAX attention tensor pack capacity Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add distributed Muon optimizer Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Muon closure and reference test Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix Muon optimizer distributed API handling Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix Muon optimizer docs and params typing Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tensor-parallel Newton-Schulz wrapper Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Move Newton-Schulz wrapper into optimizers Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Use tensor-parallel Newton-Schulz in Muon Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Limit Muon branch to Newton-Schulz TP Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Run Newton-Schulz distributed cases in one launch Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Add single-GPU Newton-Schulz coverage Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Newton-Schulz compatibility and replicated coverage Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix clamped swiglu test Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
* Unify eager and non-eager EP dispatch into one fused prepare+dispatch cpp op * Trim eager EP dispatch and combine autograd operands to cut per-call Python overhead * Call the backend directly for eager EP dispatch, combine to cut per-call CPU overhead --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* Optimize GroupedLinear validation Signed-off-by: Evgeny <etsykunov@nvidia.com> * CI fix Signed-off-by: Evgeny <etsykunov@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
… (#3425) [JAX] Fix sync-tagged collective start classification Signed-off-by: Kshitij Janardan Lakhani <klakhani@nvidia.com>
* update_filter_fp8_thd_attention Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * add for fp8+thd debug Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * uncomment the configs Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Re-plumb FP8+THD ragged-offset support on top of merged main Mirrors the F16 arbitrary_seqlen ragged-offset pattern in the FP8 path: - Backend selector: enable FP8+THD for cuDNN >= 9.23 on sm >= 100 - fwd/bwd _impl: ragged detection, batch/seqlen bucketing, set_ragged_offset() on Q/K/V/O/dO/dQ/dK/dV/Stats, workspace allocation for ragged offsets, cu_seqlens_padded_to_offsets kernel - fwd/bwd dispatchers: accept num_tokens_q/kv, cu_seqlens_padded, compute max_batch/max_tokens, THD Stats shape 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 * Address FP8 THD review feedback Use cuDNN 9.23 as the FP8 THD ragged-offset gate and prefer int64 offsets, matching cuDNN guidance for the new path. Restrict FP8 THD backend selection to padding masks, align ragged offset tuple order with the F16 convention, and enable zero-fill for FP8 THD comparison tests. Suppress the forward FP8 graph-builder fn_size lint using the same local pattern already used by the backward builder, because refactoring the full graph construction is outside this review cleanup. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fix FP8 THD ragged version scope Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fix FP8 ragged offset kernel calls after main merge The shared conversion kernel now accepts RaggedOffsetMultipliers, so construct it in FP8 forward and backward instead of passing the removed scalar argument list. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fix FP8 direct-seqlen merge integration Keep the actual batch when cuDNN consumes user cu-seqlens because a bucketed batch would read past the buffers. Reuse the aligned fallback workspace and keep SM120 stats dense so allocation matches the graph layout. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Disable fast zero fill in FP8 THD comparison The optimized mha_fill path reads CUDA cu_seqlens from host C++ and segfaults for THD. A controlled A/B passed with False while the enabled path exited 139. Keep the comparison test on the safe path until a graph-safe zero-fill implementation lands. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Align FP8 THD MHA test token counts FP8 Linear flattens THD inputs to [t, h*d], so adjust generated sequence lengths before building cu_seqlens to make the total token count divisible by eight for both forward and backward. Keep fast zero fill disabled in the MHA helper to avoid the known host dereference path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Enable FP8 THD context parallel attention Remove obsolete FP8+THD+CP gates now that the fused-attention path supports this combination. Replace the unsafe host-side suffix calculation with stream-ordered zeroing, and preserve FP8 metadata and token-major THD layout in the all-gather path. 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 * Fix FP8 THD CP gradient padding P2P FP8 backward densely combines rank-local partial gradients, which can repopulate inter-sequence dK/dV padding after native zero initialization. Reconstruct every local padding interval from actual and padded cu-seqlens and clear dQ/dK/dV after reduction. Initially limit cleanup to FP8 because high-precision control runs kept padding zero; later commits can broaden that policy independently. 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 * Enable FP8 THD inference on Hopper with cuDNN 9.25 cuDNN 9.25 provides a working Hopper FP8 THD forward kernel, while 9.23 selects a plan that traps with an illegal instruction. Keep Hopper backward gated and preserve the existing cuDNN 9.23 requirement on Blackwell. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Cover FP8 THD CP inference on Hopper Reuse four existing Hopper cases for forward-only validation instead of adding a new Cartesian test axis. Blackwell retains forward-and-backward coverage, while Hopper runs the cuDNN 9.25-supported inference path. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Support bottom-right causal FP8 THD attention Admit the bottom-right causal mask only on SM100+, where the cuDNN FP8 backend supports it. Treat the mask as both causal and padding when constructing the frontend graph. Place delayed-scaling P2P half-gradients into their per-sequence THD halves so causal CP backward does not copy half-sized tensors into full buffers. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Fuse delayed FP8 THD half-gradient placement Extend the native THD gradient correction kernel with copy-and-zero operations for raw byte gradients. This replaces per-step Python index construction with one vectorized CUDA launch while preserving the inactive-half zeroing required by causal P2P backward. 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 * Limit FP8 THD attention to Blackwell Remove the Hopper forward-only selector and test specialization because this PR targets complete FP8 THD support, including backward. Apply THD gradient-padding cleanup regardless of FP8 state so enabling the cleanup does not change non-FP8 padding semantics. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Exercise default zero fill in FP8 THD tests Reserve seven positions in the final random THD sequence so aligning total tokens for cuBLAS only increases the sequence length without exceeding the configured maximum. Remove the temporary fast_zero_fill overrides now that the default path passes MHA and DPA backward validation. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Vectorize THD CP gradient padding cleanup Build padding masks with an on-device search instead of launching an operation per sequence, and reuse the KV mask for dK and dV. This removes batch-size-dependent overhead while preserving valid gradients. Add reference, mutation, and CUDA graph coverage for the helper. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Clean THD gradients before reverse A2A Per-step THD sequence metadata describes gradients in sequence order. Zero inter-sequence padding before reverse A2A converts those gradients to CP-rank order, preventing valid rank-local gradient rows from being cleared. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Address FP8 THD review feedback Remove redundant zero-fill and backend conditions while preserving strided THD gradient handling. Reset the shared delayed-scaling amax explicitly because ATen view zeroing does not update quantizer metadata. Co-locate the flattened THD padding mask with the dense padding-mask utilities and clarify raw-byte gradient correction semantics. 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 * Retry transient FP8 THD pool NaNs Long-lived context-parallel workers can retain state across heterogeneous cases, while the same FP8 THD case succeeds in a fresh worker. Reuse the existing one-retry mechanism only for the specific NaN assertion; persistent NaNs still fail on the fresh retry. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Normalize THD masks in FP8 attention tests THD inputs carry variable sequence lengths, so backend selection requires the padding-aware form of each logical mask. Copy each shared model config before translating the mask to avoid leaking mutations across parametrized cases. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Enable fused FP8 THD selection on SM90 cuDNN 9.23 adds FP8 THD forward support on Hopper. Relax the architecture gate to exercise SM90 support; the following commit narrows this to forward-only after validation showed Hopper backward must remain unavailable. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> * Limit SM90 fused FP8 THD to forward cuDNN supports fused FP8 THD forward on Hopper but rejects the backward path. Keep SM100+ eligible for training while allowing SM90 only for inference so unsupported training configurations skip during backend selection. 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 --------- Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…n all SMs (#3422) [Common] Default EP comm kernels to 32 SMs and keep shuffle/preprocess on all SMs Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* TE/JAX MoEBlock optimizations Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Fix lint Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> * Keep arch guard consistent in EP tests Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> --------- Signed-off-by: Jeremy Berchtold <jberchtold@nvidia.com> Co-authored-by: Teddy Do <tdophung@nvidia.com>
mxfp8 quant: skip fetching next tile Signed-off-by: qiyuw <qiyuw@nvidia.com>
…s and opaque config/params handles (#2964) * refactor nvte_get_fused_attn_backend with FE calls Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * replace code+string with string only Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up logic/comments/structure Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix compilation errors Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove handle from API; add bottom_right_diagonal Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add batch_size to API Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix jax binding Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * specify o_dtype for FP8s Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix BRCM and custom_fp8 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add o_format/etc to API and other tweaks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor tweaks for docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * replace with nvte_get_fused_attn_backend_v2 and add NVTEFusedAttnConfig Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix FP8 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add do_dtype and dqkv_dtype to API Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * replace with opaque handle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * repeat with fwd/bwd params Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove bucketed b/t_q/t_kv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * thread _v2 through, fix default scaling mode, make config specific to fwd/bwd, thread bias shapes through in jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reorder struct fields, consolidate APIs to derive, make_config, make_cache_key, fix Jax bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * simplify fused attn config/params wrappers via set_attr helper, add graph_debug, fix bias shape handling Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up derived fields, debug probe/exec graph mismatches Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * match fused attn availability probe to runtime for FP8 specs and per-step CP configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix bias for jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fused-attn: share graph cache across threads via mutex, require cuDNN 9.11, guard zero-size batch/token quantization, add graph-cache debug tracing Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Cache attention-backend selection keyed on (NVTE_* env, attention_params), drop temporary graph-debug instrumentation, revert min cudnn version update Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * use macros for attr_sizes[], cache_key_tuple(), and fprintf in cache debug Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "use macros for attr_sizes[], cache_key_tuple(), and fprintf in cache debug" This reverts commit aa34ccb356772e4854e8fc930a8d104d39cbcd0a. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove graph cache debug code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * review and clean up Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * guard against pre-scale bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix score mod Jax tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Mirror PyTorch NVTE_DEBUG logging in JAX fused-attn backend selection, document diagnostic semantics, add reject-path test for pre-scale bias Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tidy up on jax side Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix nvte_get_fused_attn_backend shim, docstring, bias/softmax pointers Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add device_id as a key Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add docstring for FP8 recipes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix doc/ipynb Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * avoid duplicate checks for fused backend and force to 0 for bias shape defaults Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * reduce ipynb diffs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add fused attn graph cache debug code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix kv cache probes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * deduplicate L0 pytest tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge with torch.compile PRs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove redundant change Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * group newly enabled SWA tests to tiers L0/L1 in Jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix torch.compile for get_backend Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * a cleaner way to make torch.compile work Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * skip fused attn checks for flash tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add tq/tkv to per-step cp configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix jax CI Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * temporary changes: cache debug, timers, single flight, is_probe, dry-run, still build plans in probes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "temporary changes: cache debug, timers, single flight, is_probe, dry-run, still build plans in probes" This reverts commit 8fdd81d291a96ef9aebc0136894b90cb040eed0c. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix probe/exec bias drift in keys Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix probe/exec cp drift in keys Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix MHA init to avoid probe/exec drift Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: graph cache and restructuring of impl Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert consolidating of derive() and make nvte_get_backend_v2 only read and not modify cfg Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * incorporate PR5 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * address review comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix CI failures Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more fixes: cache only genuine cuDNN refusals, bound the graph cache with LRU, drop cudnn.h from the public header, guard move-assignments against NULL Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: restructuring and polishing Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: clean up graph cache/debug Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: tidy up structure Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix compile warnings and misc changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: trim code/comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: trim code/comments v2 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor tweaks of notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * reorder FusedAttnFwd/BwdParams scalars to match FusedAttnConfig Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improvements/fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix jax merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.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: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix linter Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert to some main condition Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix thd in jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix CP AG gradient count Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add test for fused error message; thread t_q/kv to probe in tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add jax error message test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * swap input_batch for num_segments in jax get_backend Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks based on AI review Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * swap input_batch for num_segments in jax get_backend; more files Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * restrict MLA to sm100+ Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * drop token counts for bshd/sbhd as well Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweak to last commit regarding dropping token counts Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * UnsupportedByCudnn struct; distinguish multi-dev on single thread; record cache_graph after insert on winner thread Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * allow mixed THD+2BSHD on sm80 for kv cache Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add back fp8 sink bwd condition Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * recover 3 CP tests by removing brcm + thd constriction for p2p Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve jupiter notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve jupiter notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve jupiter notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve jupiter notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve jupiter notebook Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak test files Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks to common Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak docstrings Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert merge regressions to main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert merge regressions to main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert merge regressions to main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * pass window(-1,-1) for cp p2p Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweak to BRCM location Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix tq/tkv for jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix mixed_thd test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * make dtype fields in config local vars to reduce bloat Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * structurally fix batch_size/num_segments in jax probe/exec Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * narrow dedup for some tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * jax: reject unsupported configs before RNG/aux setup The FFI fwd/bwd paths queried the backend and then populated the RNG state and the aux tensor pack unconditionally. With NVTE_No_Backend that walked the FP8 branch of PopulateRngStateAsync, advanced the global RNG offset and launched a kernel for an attention that never runs, and the rejection message was discarded, so the failure only surfaced later from nvte_fused_attn_{fwd,bwd}_v2 without the diagnostic. Check the backend right after the query and surface the message. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * jax: query the fused attention backend through named setters GetFusedAttnBackendImpl took 36 positional arguments, including runs of same-typed heads, head dims and sequence lengths that could be transposed at a call site and still compile. Have each of the three call sites build a FusedAttnConfigWrapper with its named setters instead, and reduce GetFusedAttnBackendImpl to the query itself. The wrapper already names every field and owns the defaults, so it supplies the anti-swap property without a second struct mirroring FusedAttnConfig. The pybind entry point now reads like its PyTorch counterpart, and the two FFI handlers build the config in FUSED_ATTN_IMPL_COMMON_BLOCK, which already decodes the rest of the arguments they share, so the config cannot drift from the values it is derived from. GetFusedAttnBackendImpl stays separate from the pybind GetFusedAttnBackend because the FFI handlers run without the GIL and have no params object to pass. Their query is not a repeat of the one the Python layer ran: batch size is the segment count, which for ragged input is only resolved inside the handler. No behavior change: the fields the FFI paths never set are pinned to the values they previously picked up as struct defaults. return_max_logit is pinned to false, which is what the backward pass always passed, and the forward pass overrides it before querying. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * jax: construct FusedAttnHelper with keyword arguments The three production call sites passed the first sixteen fields positionally, including the is_training/batch_size and q_num_heads/kv_num_heads/q_max_seqlen/ kv_max_seqlen runs, where a transposition type-checks and yields a plausible backend answer rather than an error. Name them, convert the one remaining positional construction under tests/jax, and mark the dataclass kw_only so the convention cannot drift back. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tests/jax: assert the probed layout by name The positional index into the recorded FusedAttnHelper arguments had to move from 3 to 4 when batch_size was added. Now that the call site passes keywords, read qkv_layout out of the recorded kwargs instead. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * jax: probe the window the kernel actually runs with The fwd/bwd lowerings substitute cp_striped_window_size for the global window (CP + Ring P2P + THD + SWA), but the backend probe and the workspace-size query in abstract kept using config.window_size. The probe therefore validated, and the graph cache retained, a descriptor that never executes, while the executing descriptor was never checked. Add _FusedAttnConfig.effective_window_size and use it in both places, which also collapses the branch that was duplicated in the two lowerings. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * common: remove redundant headers Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add THD sink attn bwd restriction Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
fix(jax): interpolate the tensor-sequence-parallelism warning
The warning raised when tp_resource appears in lhs_specs is built from three
implicitly concatenated fragments, none of which carries the f prefix, so users
see the literal text {gsr.tp_resource} and {lhs_specs} instead of the values that
triggered the warning. Both names are locals defined a few lines above.
Signed-off-by: Tai An <antai12232931@outlook.com>
Co-authored-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com>
…3391) * feat: thread tanh logit softcapping through FlashAttention (FA2, opt-in FA3) Add a user `softcap` value (tanh logit softcapping, `softcap*tanh(x/softcap)`) to DotProductAttention so models like Gemma2 can run on the fused flash path instead of an unfused/FlexAttention kernel. - Add `softcap` to DotProductAttention (init+forward) and AttentionParams; thread it into the FA2 non-CP kwargs and all three context-parallel autograd functions (forward + ctx-saved backward). softcap=0.0 reproduces prior behavior. - get_attention_backend: when softcap != 0, disable FusedAttention/unfused and steer to FA2 -- disable FA3/FA4, and disable FA2 < 2.6.0 -- so the cap is never silently dropped (FA2 < 2.6.0) or hit at runtime as NotImplementedError (FA3/FA4). Also disable FA3 under context parallelism (its CP path hard-rejects nonzero softcap) so CP+softcap steers to FA2, which supports it, instead of crashing. - FA3 softcap opt-in: NVTE_FA3_SOFTCAP=1, Hopper (sm90) hd<=256, non-CP only, gated on a fail-closed signature probe (fa3_supports_softcap). Forward threads softcap into fa_3_optional_forward_kwargs; the existing Hopper autograd function carries it into backward automatically. Default off; unchanged behavior steers to FA2. - ONNX export: fail loudly (assert) rather than silently drop softcap -- export unconditionally force-selects UnfusedDotProductAttention, which has no softcap support, so this previously exported models with softcapping silently omitted. - Tests: test_softcap.py (FA2 fwd/bwd parity vs pure-PyTorch reference), wired into qa/L0_pytorch_unittest/test.sh. FA4 softcap opt-in is deliberately NOT included here -- see follow-up PR. On Blackwell (SM100), FA4's dedicated head_dim=256 forward kernel has no score_mod support at all (kernel constructor asserts `score_mod is None`), so there is currently no FA4 kernel path this could opt into; adding the scaffolding now would just be inert code with nothing to exercise. Addresses review findings: CP+FA3 softcap selection crash, ONNX silent drop, and the missing CI wiring for test_softcap.py. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: use raise instead of assert for the ONNX+softcap guard python -O / PYTHONOPTIMIZE strips assert statements, which would silently reopen the ONNX export softcap-drop bug the previous commit fixed (ONNX mode would again force-select UnfusedDotProductAttention with softcap silently omitted, with no error). Switch to an explicit if/raise ValueError, which survives optimized execution. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci (reapplied after a force-push rebase clobbered pre-commit.ci's original 19a21eb7 commit; same content, restored by hand) Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * refactor: move softcap reference into UnfusedDotProductAttention and fold test into test_attention.py UnfusedDotProductAttention now applies softcap * tanh(scores / softcap) to the already-scaled logits, matching how FlashAttention folds softmax_scale into its tanh argument, so it can serve as the softcap reference backend. Backend selection therefore no longer disqualifies unfused attention for softcap, and the ONNX-export guard is dropped since the export path force-selects unfused and torch.tanh is exportable. test_softcap.py is replaced by a model_configs_softcap dict and test_dpa_softcap in test_attention.py, which reuses test_dot_product_attention for backend sweeping. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): gate FA3 softcap on existing NVTE_FLASH_ATTN_V3 Drop the redundant NVTE_FA3_SOFTCAP opt-in. `use_flash_attention_3` already derives from NVTE_FLASH_ATTN_V3, so the existing flag governs the FA3 softcap path and NVTE_FLASH_ATTN_V3=0 disables it. Correctness stays established by the build-capability probe, head_dim <= 256, and the non-CP requirement. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): disable FlashAttention 4 for softcap FA4 exposes no softcap kwarg and its head_dim=256 kernel asserts score_mod is None, so there is no kernel to route the cap through. The FA4 call path in backends.py passes no softcap, so an FA4 selection with a nonzero softcap silently dropped the cap instead of failing closed. NVTE_FLASH_ATTN_V4 defaults to enabled, so this was reachable on SM100+ with flash-attn v4 installed and no context parallelism. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): disable FlashAttention 2 for softcap with dropout flash-attn rejects a nonzero softcap combined with nonzero dropout at dispatch: "Softcapping does not support dropout for now" in csrc/flash_attn/flash_api.cpp, present in mha_fwd and mha_varlen_fwd from v2.6.0 (the earliest version TE allows softcap on) onwards. Backend selection did not model this, so a softcap + attention-dropout config passed selection, routed to FA2, and crashed inside flash-attn. Dropout only reaches the kernel while training, since backends.py passes `self.attention_dropout if self.training else 0.0`, so the gate is on `attention_dropout != 0.0 and is_training` to avoid blocking valid inference configs. UnfusedDotProductAttention supports both softcap and dropout and stays available, so this steers rather than hard-fails. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: restore softcap dQ/dK/dV parity in the shared DPA harness test_dot_product_attention forced is_training=False whenever FusedAttention could not train a config, so that backends only available for inference could still be compared. softcap always disables FusedAttention, so test_dpa_softcap silently degraded to a forward-only comparison and the PR's backward-parity claim -- the FA2 softcap backward kernel included -- went untested. Add fwd_only_without_fused_attn (default True, so every other caller is byte-for-byte unchanged) and opt test_dpa_softcap out, which pairs FlashAttention against UnfusedDotProductAttention with is_training=True and restores the dgrad comparison. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: add softcap no-op and closed-form reference coverage Two gaps remained after folding test_softcap.py into test_attention.py. softcap=0.0 no-op: every model_configs_softcap entry uses a nonzero cap, so nothing asserted the backward-compatibility claim. The half that the PR actually changed is backend selection, and a filter that fired at 0.0 would silently remove FusedAttention and FA4 from other tests rather than fail one. test_dpa_softcap_zero_backend_selection asserts FusedAttention survives softcap=0.0 and is disabled by a nonzero cap. Unfused coverage and tanh's nonlinear region: test_dpa_softcap needs two TE backends, so it skips entirely without flash-attn even though UnfusedDotProductAttention now implements softcap and is the reference for everything else. It also cannot detect a dropped cap at all: 0.1 * randn inputs put the logits at O(1e-2), where the reference output moves by 9e-9 at cap=50 and 2e-4 at cap=0.01. test_dpa_softcap_vs_reference compares forward and dQ/dK/dV against a pure-PyTorch oracle one backend at a time, so it runs with unfused alone, and uses randn inputs so the cap moves the output by O(1). An assertion on that displacement keeps the test from going vacuous if the config drifts. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): align CP autograd backward arity with the softcap forward input Threading `softcap` through the three context-parallel autograd functions added a forward input to each without adding the matching gradient slot to the corresponding backward return tuple, leaving every CP backward one gradient short of its forward inputs. Because `softcap` sits mid-signature, the omission also shifted every later slot in AttnFuncWithCPAndQKVOA2A: `d_softmax_offset` was being returned in `softmax_type`'s position. Insert the missing slot at the `softcap` position in all three tuples. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test(pytorch): cover context parallelism with softcap No test exercised CP together with softcap, which is why a backward that returned one gradient fewer than its forward inputs went unnoticed in all three CP autograd functions. Thread softcap through the CP runner so it reaches DotProductAttention, and add one case per CP autograd function -- p2p, all_gather and a2a -- checking the softcapped forward and dgrad against the non-CP reference. The cap sits in tanh's nonlinear region so a path that dropped it diverges rather than matching a numerically linear reference. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): apply softcap before the additive bias, matching FlashAttention UnfusedDotProductAttention applied the tanh cap after adding post_scale_bias / ALiBi, computing cap(scale*QK + bias). FlashAttention-2 computes cap(scale*QK) + bias: its kernel softcaps immediately after the QK^T gemm and only then adds ALiBi, and it pre-divides alibi_slope by scale_softmax -- which softcapping sets to `softcap` -- so the bias deliberately lands outside the tanh (csrc/flash_attn/src/flash_fwd_kernel.h, mask.h, flash_api.cpp). ALiBi is the one bias type flash supports, so with softcap + ALiBi the unfused and flash paths returned different numerics depending only on whether a suitable flash-attn was installed. Defer the additive bias until after the cap so the two agree. pre_scale_bias is folded in before the scaling by construction and stays inside the cap; flash does not support it. softcap = 0.0 remains a bit-exact no-op for every bias type. Add two tests, both forcing UnfusedDotProductAttention: - test_dpa_softcap_bias_ordering pins cap(scores) + bias against cap(scores + bias), using post_scale_bias to drive the same branch ALiBi uses without needing slope machinery in the reference. - test_dpa_softcap_qk_layer_scaling covers softcap under NVTE_APPLY_QK_LAYER_SCALING, where the cap must be divided by layer_number; omitting that leaves an effective cap of softcap * layer_number. Both carry anti-vacuity asserts, and both were verified by mutation on an H100: reintroducing either bug makes the corresponding test fail, and the existing softcap suite still passes. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(test): skip the softcap no-op check when FusedAttention is unavailable test_dpa_softcap_zero_backend_selection asserted that FusedAttention is selectable at softcap=0.0. Whether it is available at all depends on the arch and on NVTE_ALLOW_NONDETERMINISTIC_ALGO, not on softcap: the A100 and L40 CI jobs fail the assert in their deterministic pass, where cuDNN offers no fused backend for this config, while the same test passes on H100 and B200. The test is about the softcap filter, so make the availability half a skip. The meaningful direction -- a nonzero softcap must disable FusedAttention -- still runs wherever a fused backend exists, and skipping keeps the comparison from being silently vacuous where one does not. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix(lint): silence missing-function-docstring on _IdentityWithMaskedGradient The Lint job fails on this branch with dot_product_attention.py:243:4: C0116: Missing function or method docstring dot_product_attention.py:248:4: C0116: Missing function or method docstring _IdentityWithMaskedGradient came from #3274 and is byte-identical to main; the same two errors fail Lint on that PR's own branch (desh/mixed-thd-pr-minimal), which merged before the check went green, so this branch inherited a red Lint by merging main. Not introduced here, but it blocks this PR. Use the same `# pylint: disable=missing-function-docstring` the other autograd Functions in this package use (e.g. context_parallel.py) rather than inventing docstrings. Drop this commit if #3274 is fixed upstream first. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat: expose softcap on MultiheadAttention and TransformerLayer; trim comments Review feedback from #3391. Thread softcap through MultiheadAttention and TransformerLayer, following the window_size pattern: a constructor argument stored on the module and a forward argument that falls back to it. TransformerLayer passes softcap via common_attention_kwargs, so it reaches cross-attention as well as self-attention, matching how softmax_type (also not mask-specific) is handled; scoping it to self-attention would silently leave decoder cross-attention uncapped. Resolution sits outside the thd_attention_policies branch at both levels, since softcap is not mask-specific. Drop "only supported by the FlashAttention and UnfusedDotProductAttention backends" from the softcap docstrings. It is inaccurate, since support depends on the FlashAttention version, and get_attention_backend() is the single place that should describe backend eligibility. The same wording appeared in the DotProductAttention class and forward docstrings, so all three now describe only what softcap does. Shorten the softcap comment in UnfusedDotProductAttention, keeping why the cap precedes the additive bias and why it is divided by layer_number, as each of those was a real bug. Remove the fail-loud block comment in the FA3 branch, whose selection criteria get_attention_backend() already documents, and keep a single line noting that the FA3 entry points are autograd functions and so the forward kwarg drives the backward too. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * fix: append softcap to the MHA and TransformerLayer forward signatures softcap was inserted mid-signature in MultiheadAttention.forward and TransformerLayer.forward, next to the other attention arguments. That shifts every parameter after it, so an external caller passing encoder_output or is_first_microbatch positionally would silently bind it to softcap instead. Both accept arbitrary objects, so it would be a wrong answer rather than a TypeError. Nothing in this repo calls either forward with enough positional arguments to hit it, but both classes are public API. Move softcap to the end of both forward signatures. Diffing against main confirms no existing parameter changes index in either. The constructors are unaffected and keep softcap grouped with window_size and bottom_right_diagonal, since it sits past every positional tuple used to construct these modules. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test: consolidate the softcap tests Review feedback from #3391. Cut model_configs_softcap from six configs to two: softcap_1_0 for a high cap through flash_attn_func, and softcap_3_1 with padding_causal for a low cap through flash_attn_varlen_func. Fold model_configs_softcap_reference from three to two by merging GQA into the causal config, since GQA is orthogonal to softcapping. Merge test_dpa_softcap_vs_reference, test_dpa_softcap_bias_ordering and test_dpa_softcap_qk_layer_scaling into one test parametrized over {plain, bias_outside_cap, qk_layer_scaling}. A helper returns the DPA kwargs, the forward kwargs and a right and a wrong reference closure per variant, so the shared work of forcing a backend, building inputs and asserting the two references differ by more than the tolerance happens once. The bias variant reuses _softcap_reference_attention through new bias and cap_includes_bias arguments rather than duplicating its GQA and causal handling. Both mutation signals are preserved. The special variants run on softcap_ref_1_0, the same shape the original margins were measured on, so capping the bias still separates the two references by about 0.98 and dropping the cap / layer_number division by about 0.44, against a threshold of 0.20. layer_number is hoisted to _SOFTCAP_QK_LAYER_NUMBER with a note that 3 is too small, since that value is load-bearing and would otherwise look arbitrary. Backward stays on the plain variant only, matching the previous coverage. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test: assert TransformerLayer forwards softcap to both attention modules Adds test_transformer_layer_softcap_plumbing. It builds a decoder TransformerLayer with a softcap, hooks both self_attention.core_attention and inter_attention.core_attention with a forward pre-hook, and asserts each sees the value, once from the constructor and once from a forward override. Only arrival is checked; the numerics are covered by the DotProductAttention tests. Cross-attention is the half worth asserting. It is reached through a separate call site from self-attention, so a refactor can drop the cap there while self-attention keeps working, and nothing else in the suite would notice. Main restructured both forward methods during this branch's life, so that is a live risk rather than a hypothetical one. Writing the test showed that passing softcap through common_attention_kwargs was not merely redundant but harmful: it pre-set each MultiheadAttention's self.softcap, so deleting the cross-attention pass-through still produced the right answer on the constructor path and the assertion could not fail. Removing that entry makes both assertions catch the regression, and leaves softcap matching how window_size is threaded, stored on self and passed per call rather than at construction. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test: set explicit backend state in the softcap plumbing test The test inherited NVTE_FLASH_ATTN / NVTE_FUSED_ATTN / NVTE_UNFUSED_ATTN from whichever softcap test ran before it, since none of them restore those vars. A nonzero cap also drops FusedAttention and FA4 in the backend filter, so a leftover NVTE_UNFUSED_ATTN=0 could leave no eligible backend at all and the forward would raise "No dot product attention backend is available" rather than fail an assertion. The outcome also depended on which parametrization of the preceding test happened to run last. Set the three variables explicitly, matching the convention of the other tests in this file, leaving flash and unfused both enabled so at least one backend is available with or without flash-attn. The assertions stay backend agnostic because the hook fires on DotProductAttention itself. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * chore(pytorch): tighten softcap comments and correct the FA4 attribution Halve the comment volume added to backends.py and utils.py, keeping only the facts that are not inferable from the code: that FA2 adds ALiBi outside the cap, and why the cap is divided by layer_number. Which backend supports what is already expressed by get_attention_backend and its debug messages. Also correct two inaccuracies found while auditing the surviving text. FA4 does implement softcap on both entry points TE calls, forward and backward, so it is disqualified because TE does not plumb the cap to it, not because the library lacks it; the log message said otherwise. And the FA3 signature probe detects whether the installed release exposes the kwarg at all, not whether the build set FLASHATTENTION_DISABLE_SOFTCAP, which still exposes it and rejects a nonzero cap at dispatch. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): probe softcap on all three FA3 entry points softcap is injected into fa_3_optional_forward_kwargs without regard to which entry point was selected, so the probe should cover every function the cap can reach. flash_attn_with_kvcache is the one chosen when inference_params is set, and it was validated only by proxy through the other two. No shipped FA3 build changes behaviour: all three entry points gained softcap in 2.7.3/hopper, which is TE's documented FA3 floor. Before that none of them had it. The symbol is already imported and None-guarded, so the probe cannot fire on an uninstalled FA3. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: exercise softcap at head_dim 128, not only 64 Cutting model_configs_softcap to two configs dropped head_dim 128 as a side effect, leaving every softcap config at 64. FA2 and FA3 compile a separate softcap kernel per head_dim bucket, and above 128 FA2 additionally forces the non-even-MN predicated path, so a bug in those kernels is invisible at 64. Widen the existing padding config instead of adding a third, so the count stays at two. Neither justification for the pair changes: head_dim plays no part in the flash_attn_func vs flash_attn_varlen_func dispatch, and the 0.01 cap rationale is head_dim invariant because softmax_scale cancels the sqrt(d) growth of the dot product. 128 rather than 256 because get_attention_backend disables FA3 for deterministic backward at head_dim >= 256, which would leave only UnfusedDotProductAttention and skip the comparison in the NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 CI pass. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: restore backend env vars in the softcap plumbing test The test set NVTE_FLASH_ATTN, NVTE_FUSED_ATTN and NVTE_UNFUSED_ATTN and never restored them, leaking NVTE_FUSED_ATTN=0 into every test that ran after it. Its own comment names that hazard as the reason it sets them explicitly, then reproduces it. Wrap the body in try/finally, restore the previous values, and invalidate the backend selection cache on the way out. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): reject a negative softcap instead of diverging silently tanh is odd, so UnfusedDotProductAttention's `cap * tanh(x / cap)` gives a bit-identical result for -c and +c, i.e. it caps at the absolute value. FlashAttention only enables capping under `softcap > 0.0`, so it applies no cap at all. get_attention_backend gates on `softcap != 0.0`, so a negative value passed every filter and the two backends silently disagreed, which is the exact failure mode the rest of this PR exists to prevent. Raise on a negative cap where softcap is resolved, which is the choke point every caller reaches, including MultiheadAttention and TransformerLayer. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * fix(pytorch): reject a non-finite softcap as well as a negative one The existing guard exists to stop the backends disagreeing silently about a cap they cannot all honour, but NaN slipped through it and caused exactly that disagreement: UnfusedDotProductAttention computes cap * tanh(x / cap) and returns NaN for every score, while FA2 and FA3 gate their kernels on softcap > 0.0, which is false for NaN, so they drop the cap and return ordinary uncapped attention. Infinity produces NaN on every backend instead, which at least fails loudly. Neither is a legitimate way to spell a disabled cap. The limit of c * tanh(x / c) as c grows is the identity, but evaluated at infinity it is inf * 0, which is NaN, so the intuition does not survive floating point. Fold both into the existing check with math.isfinite, matching how ScaledTanhSReLU and ScaledSiTUGLU validate the same shape of parameter. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> * test: pin the softcap validity guard The guard had no coverage, for either the negative case or the non-finite one. Parametrised over -1.0, inf and nan, exercising both the constructor value and the forward override, since both resolve on the same path. Mirrors test_scaled_tanh_srelu_rejects_bad_clamp_scale, which pins the same shape of check on ScaledTanhSReLU. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> --------- Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com>
* feat: mark GroupedLinear activations for paged stashing Signed-off-by: hongbinl <hongbinl@nvidia.com> * Clarify paged stash grouped tensor markers Signed-off-by: hongbinl <hongbinl@nvidia.com> --------- Signed-off-by: hongbinl <hongbinl@nvidia.com>
Replace TorchScript fallbacks with eager execution TorchScript is deprecated, and selecting it when torch.compile is disabled keeps an unnecessary legacy path. Reuse the existing decorator policy so enabled configurations use lazy torch.compile while disabled or unsupported configurations execute eagerly. Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
* Insert assert header Signed-off-by: Alex Y. Chan <alechan@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alex Y. Chan <alechan@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* cutlass commit update Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * change the namespace to different one 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 --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
…elpers (#3480) * [PyTorch] Split FusedAttnFunc into single-argument forward/backward helpers Behavior-neutral restructuring of the cuDNN fused attention autograd function, mirroring what #2967 did for Linear: FusedAttnFwdArgs / FusedAttnBwdArgs dataclasses, module-level _fused_attn_forward_impl, _fused_attn_setup_ctx and _fused_attn_backward_impl, and a thin FusedAttnFunc wrapper taking the differentiable tensors plus one args object instead of 37 positional arguments. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Match the fused attention aux pack split to the C++ conditions The bias / softmax_offset aux entries exist only when the tensor was passed, not merely when the bias / softmax type asks for one. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Make fused attention saved tensor sources explicit Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Preserve separate attention page tables and unset backward defaults Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…452) * [JAX] Add EP borrowed-comm path via XLA FFI collectives extension Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* Add MXFP8 support to cuBLASMp GEMM Signed-off-by: Almog Segal <asegal@nvidia.com> * Enable cuBLASMp MXFP8 overlap tests Signed-off-by: Almog Segal <asegal@nvidia.com> * Fix MXFP8 weight-gradient preparation with cuBLASMp Gather columnwise MXFP8 gradients from the original input for weight gradients, and keep Userbuffers-only external overlap out of the cuBLASMp backward path. Signed-off-by: Almog Segal <asegal@nvidia.com> * Require cuBLASMp 0.8.1 for MXFP8 support Signed-off-by: Almog Segal <asegal@nvidia.com> --------- Signed-off-by: Almog Segal <asegal@nvidia.com>
* [PyTorch] Fix mutable QB bounds in CUDA graphs Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Document mutable QB graph contract Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Require explicit QB bounds validation Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [Common][PyTorch] Validate QB bounds in CUDA graphs Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Allow trusted QB updates during graph capture Signed-off-by: Harry Zhou <hhanyu@nvidia.com> * [PyTorch] Clarify mutable QB graph bounds Signed-off-by: HarryZhou <2373256746@qq.com> --------- Signed-off-by: Harry Zhou <hhanyu@nvidia.com> Signed-off-by: HarryZhou <2373256746@qq.com> Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
…/2.20 # Conflicts: # build_tools/build_ext.py # build_tools/jax.py # build_tools/pytorch.py # build_tools/utils.py # pyproject.toml # qa/L0_jax_unittest/test.sh # qa/L2_jax_unittest/test.sh # tests/cpp/operator/CMakeLists.txt # tests/jax/test_distributed_fused_attn.py # tests/jax/test_fused_attn.py # tests/pytorch/attention/test_attention.py # tests/pytorch/attention/test_attention_with_cp.py # tests/pytorch/distributed/run_layer_with_overlap.py # tests/pytorch/distributed/test_cast_master_weights_to_fp8.py # tests/pytorch/distributed/test_comm_gemm_overlap.py # tests/pytorch/test_grouped_tensor.py # tests/pytorch/test_numerics.py # tests/pytorch/test_torch_compile.py # transformer_engine/__init__.py # transformer_engine/common/CMakeLists.txt # transformer_engine/common/fused_attn/fused_attn.cpp # transformer_engine/common/fused_attn/fused_attn_fp8.cu # transformer_engine/common/fused_attn/fused_attn_fp8.h # transformer_engine/common/fused_attn/utils.cu # transformer_engine/common/util/ptx.cuh # transformer_engine/common/util/pybind_helper.h # transformer_engine/jax/attention.py # transformer_engine/jax/cpp_extensions/attention.py # transformer_engine/jax/csrc/extensions.h # transformer_engine/jax/csrc/extensions/attention.cpp # transformer_engine/jax/csrc/extensions/pybind.cpp # transformer_engine/jax/flax/transformer.py # transformer_engine/jax/moe.py # transformer_engine/pytorch/__init__.py # transformer_engine/pytorch/attention/dot_product_attention/backends.py # transformer_engine/pytorch/attention/dot_product_attention/utils.py # transformer_engine/pytorch/csrc/extensions/attention.cpp # transformer_engine/pytorch/module/grouped_linear.py # transformer_engine/pytorch/module/linear.py # transformer_engine/pytorch/triton/cross_entropy.py
Micky774
requested review from
ipanfilo,
matthiasdiener,
wangye805 and
wenchenvincent
as code owners
October 7, 2026 19:30
Contributor
Author
|
For reference, here is a streamlined view of IFU-related changes |
ipanfilo
reviewed
Oct 13, 2026
| ) | ||
| install(FILES "${CMAKE_CURRENT_BINARY_DIR}/build_info.txt" DESTINATION "transformer_engine/") | ||
| endif() | ||
| if(NVTE_WITH_NCCL_EP) |
Collaborator
There was a problem hiding this comment.
Add USE_CUDA check here or explicitly set NVTE_WITH_NCCL_EP variable to FALSE if USE_ROCM
Contributor
|
One thing that concerns me: the changes assume MXFP8 scale swizzling is never done on ROCm, but on gfx1250 it should be. Perhaps we can gate these on gfx1250 (like swizzle_mxfp8_scale already does)? |
Contributor
Author
These kernels don't fuse the swizzle on any ROCm arch, same as the existing quantize kernels. They now flag the scales compact, so on gfx1250 the GEMM path does the swizzle, and CK grouped GEMM reorders compact scales itself. I fixed the misleading comments, and made the grouped pre-swizzle error on gfx1250 instead of writing the CUDA layout. |
ipanfilo
approved these changes
Oct 15, 2026
matthiasdiener
approved these changes
Oct 15, 2026
- Resolve build-system, common-layer, framework-binding, Python-layer and test/QA conflicts - Fix ROCm build of new upstream device code, build_tools imports, and test wiring broken by upstream renames; run new upstream tests in CI - Port fused grouped MXFP8 requantize and grouped scaled SwiGLU; flag their scales compact - Enable te.Linear torch.compile on ROCm - Disable FA2 softcap on ROCm; skip cuDNN-only and other unsupported test variants - Keep the fork's softmax_type aux condition in JAX attention - Guard CUDA-only fused-attn config and NCCL EP declarations - Drop the upstream minimum-versions workflow - Copyright header fixes Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Micky774
force-pushed
the
zain/ifu/2.20
branch
from
October 15, 2026 19:27
af4c320 to
e942d0d
Compare
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
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: