[PyTorch] torch.compile for an OperationFuser group holding one operation - #45
Open
pggPL wants to merge 28 commits into
Open
[PyTorch] torch.compile for an OperationFuser group holding one operation#45pggPL wants to merge 28 commits into
pggPL wants to merge 28 commits into
Conversation
register_custom_op now defines an operation's forward and backward as two independent two-tier custom ops and hands back both, leaving autograd to the caller. That is what lets a pipeline-level autograd.Function decide how the two are wired, and so group the forward and backward passes differently -- which is what ops.OperationFuser does. The variant that wires autograd itself keeps the old behaviour under register_custom_op_with_autograd, and is now built on the same registration: the pair is the primitive, autograd is what the other one adds. About two thirds of the two bodies were the same code before. BasicOperation gains the plumbing an operation needs to opt in: declare two argument containers and implement four compute classmethods, and __init_subclass__ registers the custom ops while op_forward / op_backward are written once in the base. compile_unsupported_reason lets an operation say why it cannot be compiled -- it sits here rather than on the args, as Linear has it, because in ops/ the compile boundary is the fuser group, not the operation. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
A group whose operations declare their compute halves now runs through their custom ops under torch.compile(fullgraph=True). The pipeline-level autograd.Function is traced as a higher-order op, which is what will later let its forward and backward walk different op groupings. Four side effects reached outside the higher-order op's scope and had to go: - OperationContext objects are created in the forward, but the backward is a separate subgraph, so writing to them there mutates an enclosing scope; the backward copies them into its own scope instead; - requires_grad_ on an output, which AOTAutograd's functionalization drops anyway -- autograd marks the outputs of an apply() itself; - _do_not_clear on inputs and outputs. They are gated on being traced rather than on using the custom ops. Under fullgraph there is no leaving the graph, so an unsupported operation does not fall back: the pipeline is traced either way and only the choice of implementation changes. The gate reports why a group runs eagerly through warn_compile_eager_fallback, which is safe to call from the traced region. Sequential builds its module groups outside the forward pass, since that constructs nn.Modules. Tested with a test-only operation, so the fuser's path does not depend on which real operations happen to declare their halves. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
An operation lists the forward kwargs it takes in fwd_kwarg_names. They are resolved into its args container like any other config, in the traced Python where Dynamo guards them, so they reach the custom op through the existing schema -- a value is guarded, a tensor is lifted into the graph, and a quantized one crosses as its inner buffers. An undeclared kwarg still sends the whole group to eager. That is not a schema limitation, as the old message implied: the kwargs that remain are the grouped operations' preallocated buffers, which the op writes to, and a custom op may not mutate a tensor from an enclosing scope. A kwarg carries no gradient. This matches the eager path, where kwargs never entered the autograd graph either, and is why only read-only ones are accepted. The fuser test helper now builds a separate model for the eager and the compiled pass. Previously both shared one model and the eager pass ran first to produce the reference, so the compiled pass was always traced on a model whose module groups, fusions and pre_first_fuser_forward had already run. Those paths are now traced as well. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
A value kwarg does not survive a second call. The other fields of an args container are read off the module and are constant across calls, so they are baked into the graph; a kwarg changes per call, and on the second value Dynamo hands over a symbolic scalar, which OpaqueValueBundle cannot carry -- it fails with AsPythonConstantNotImplementedError, not with a graph break. Measured on int and float alike; specialize_float=True cures only the float, and is global. The gate now takes tensor kwargs only, so a value sends the group to eager deterministically instead of failing on its second call. A 0-d tensor is the way to pass a scalar: it is a graph input, so it recompiles for no value at all. The test carries a quantized offset that changes on every call and confirms no recompilation, then adds a value kwarg to cover the gated path. That last call keeps its offset unquantized on purpose: a gated group runs the eager implementation, which is traced directly rather than hidden behind a custom op, and dequantize() graph-breaks there. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…4M3 (NVIDIA#3352) * [PyTorch] Decline fused grouped MLP when the backward format is not E4M3 The fused grouped MLP packs the incoming activation gradient by reinterpreting its storage as E4M3, conditioned only on NVFP4 and never on the FP8 format. Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit E5M2, so those bytes are read as the wrong format rather than converted, and every gradient out of the fusion is wrong. The forward pass is unaffected, so this shows up as a model that trains too slowly instead of one that fails. Fall back to the unfused ops when the recipe's backward format is not E4M3, and raise instead of reinterpreting if such a gradient reaches the kernel path. Signed-off-by: William <wilyan090@gmail.com> * [PyTorch] Gate the grouped MLP backward format check on MXFP8 fp8_format describes the FP8 formats of an MXFP8 recipe. NVFP4BlockScaling carries one too, pinned to E4M3, but its gradients are quantized to FP4 and the value says nothing about them, so testing it for an NVFP4 recipe reached the right answer for the wrong reason. Restrict the check to recipes where it means something. The runtime check at the pack site is unchanged: it reads the grad output quantizer's own dtype on the non-NVFP4 branch, not the recipe. Signed-off-by: William <wilyan090@gmail.com> * Test grouped MLP format fusion with real ops Signed-off-by: Przemek Tredak <ptredak@nvidia.com> --------- Signed-off-by: William <wilyan090@gmail.com> Signed-off-by: Przemek Tredak <ptredak@nvidia.com> Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
… torch version (NVIDIA#3466) [PyTorch] Resolve EP symm-mem window offset against both torch symm-mem layouts Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
* world_group() used a single hardcoded init_method="file:///tmp/rdzv",
shared across every world_size this test parametrizes over
([device_count(), 1, 2]), and neither world_group() nor
test_distributed_fuser_ops ever calls destroy_process_group() or
removes the file afterwards. A world_size=N run's leftover FileStore
content can then corrupt a differently-shaped world_size=M run that
reuses the same path in the same pytest session, surfacing as a
confusing NCCL bootstrap failure:
torch.distributed.DistBackendError: NCCL error ...
ncclOsSocketPollConnect: connect to <self> ... Connection refused,
exceeded error retry count after 35 attempts
instead of a clear rendezvous error. Confirmed 100% deterministic:
running the full file fresh (no pre-existing /tmp/rdzv) still fails
test_distributed_fuser_ops[2] every time, because the [4] parametrization
(which runs first) leaves /tmp/rdzv behind for [2] to trip over.
This was previously masked by older bundled NCCL (2.30.7), which
apparently tolerated the stale/mismatched FileStore well enough to
still succeed; a newer NCCL (2.31.2) surfaces it as a hard failure.
See the investigation writeup for the full comparison:
Fix: key the rendezvous path by world_size
(file:///tmp/rdzv_test_fusible_ops_{world_size}), and defensively
remove any pre-existing file at that path in test_distributed_fuser_ops
before launching each subprocess job, to also cover a leftover file
from an earlier crashed run of the same world_size.
Verified on a GB200 node
NCCL 2.31.2, the container that previously failed 2 of 3 parametrizations):
all 3 world_size parametrizations now pass, including two runs of the
full file back to back with no manual cleanup in between.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: Wei Wang <weiwan@nvidia.com>
* Use unique rendezvous files in fusible ops tests
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
---------
Signed-off-by: Wei Wang <weiwan@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
* Reduce grouped MLP fuser CPU overhead Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups. Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com> Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> * fix cutedsl wgrad crash Signed-off-by: tingyangk <tingyangk@nvidia.com> * resolve comments Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> --------- Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com> Signed-off-by: tingyangk <tingyangk@nvidia.com> Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com>
* [PyTorch] Ask cuDNN for a deterministic dprob under NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 The cuDNN grouped-GEMM dactivation backward that the CuTe DSL fused grouped MLP calls accumulates the scale gradient (dprob) with cross-CTA atomic adds, so its floating-point summation order follows the tile scheduler and varies run to run. Until now there was no way to switch that off, and NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 did not reach it: the run trained fine and was silently not reproducible. cuDNN frontend 1.28.0 (NVIDIA/cudnn-frontend#521) added a `deterministic` argument to grouped_gemm_dsrelu_wrapper_sm100 that parks each N-subtile's partial result in its own slot and sums the slots in a canonical order, for dprob and for dbias. Pass it from the TE flag. Passed as True or not at all, never as False. The wrapper's own default is None, which follows torch.use_deterministic_algorithms; sending an explicit False would override that and take determinism away from a caller who asked torch for it without setting the TE variable. The capability is reported per subclass rather than per environment variable, because grouped_gemm_dglu_wrapper_sm100 has no equivalent argument -- a GLU activation stays non-deterministic however new the installed front-end is. That case, and an SReLU op on a front-end older than 1.28.0, warn instead, once per distinct reason since the remedies differ. The warning is raised from where dprob is actually produced: with a unit activation scale the epilogue never runs its atomic accumulation, so there is nothing to make deterministic and nothing to warn about. Tests: TestGroupedMLPDeterminism covers the env-var parse, that only the SReLU op reports the capability and that it tracks the front-end version (no GPU or cuDNN needed for either), that the warning fires once per reason, and an MXFP8 end-to-end run under determinism for both SwiGLU and SReLU that checks numerics and pins which of the two arms warns. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Honor torch.use_deterministic_algorithms too, not just the env variable _deterministic_algorithms_required() copied the narrow check from transformer_engine.pytorch.triton.grouped_dbias_dscales, which reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. DotProductAttention takes the union instead -- the variable OR torch.use_deterministic_algorithms -- and that is the right precedent here. The two knobs answer different questions. The variable is set once in a job launcher, applies uniformly across ranks, and is the only one TE's C++ layer can read. The torch flag is the framework standard, is togglable at runtime, and is what a user who wants reproducibility usually reaches for; most have never heard of the variable. Keying on the variable alone left the torch flag half-honored. The SReLU path happened to come out right, but by delegation rather than by decision: TE passed nothing and the wrapper's own default read torch.are_deterministic_algorithms_enabled(). The GLU path did not -- TE stayed silent about an atomic dprob it cannot fix, for a user who had asked torch for reproducibility. That silence is the exact failure mode the warning exists to prevent, so it was the one case that most needed to warn. Passing the argument only as True, never as False, now needs a different justification than the one the first commit gave: with the union in place the two are equivalent, since the wrapper's default reads the same torch flag TE just read. The reason that survives is narrower and firmer -- the argument does not exist on the dGLU wrapper or on a front-end older than 1.28.0, where passing it at all, even as False, is a TypeError. Tests: the env-var parametrization becomes the two-knob truth table, including the row that motivates the change (torch flag set, NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -- the variable's default is the absence of a request, not a request for non-determinism, so the torch flag still wins). A fixture restores the process-global torch flag. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Test that dprob is actually bit-exact, not just within tolerance Review caught that nothing in the suite tested the property this change exists for. The end-to-end test runs the op once and checks numerics against a reference with rtol=0.125 / atol=0.25; reordering the same atomic adds moves dprob by about an ulp, so a run that is silently not reproducible passes it comfortably. The tolerance check proves the deterministic path is correct, which is worth keeping, but it cannot prove the path is deterministic. Add a second run. Same module, same inputs, grads cleared between passes, probs.grad compared with torch.equal. Three things the test has to get right to be worth having: * hidden_size 1024, not the 128 used elsewhere. dprob's reduction is over that extent and the tile is 256 wide, so 128 gives a single N-tile, one writer per token, and nothing to reorder -- the assertion would hold by construction and test nothing. * No bias. With an FC2 scale_bias the scale gradient is finished by the Triton grouped dbias/dscales kernel, which refuses to run under determinism, and probs.grad would stop being the dprob under test. * An assertion that the fusion happened, since dprob only comes from the cuDNN epilogue on the fused path. Skipped rather than xfailed on a front-end older than 1.28.0: there the kernel has no deterministic mode and is expected to vary, which is not a failure of this change. Weight gradients are deliberately left out of the comparison -- the CuTe DSL wgrad kernel has its own K-split atomics that this PR does not address. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Probe the dsrelu wrapper's signature instead of the frontend version _cudnn_frontend_supports_deterministic_dprob() gated on _cudnn_frontend_version_at_least("1.28.0"). That check is too coarse to answer the question it is asked, and would have raised at runtime on a build TE is actually run against. NVIDIA#521 merged after v1.27.0 was tagged, so `deterministic` ships in 1.28.0. But cudnn-frontend's develop branch has called itself 1.28.0 since shortly after that tag -- eleven days before the merge. Any front-end built from develop in that window reports 1.28.0 and does not accept the argument, so the version check passes, TE adds `deterministic=True` to the call, and the backward dies with TypeError: grouped_gemm_dsrelu_wrapper_sm100() got an unexpected keyword argument 'deterministic' This is not hypothetical, and not new. The same coarseness already bit use_single_group_runtime_offsets: a cuDNN reporting 1.27.0 that did not implement 1.27.0's arguments failed the identical way, in fuser_forward, before any backward code ran. Version numbers describe a release; they do not describe whatever happens to be installed. Ask the function instead. `"deterministic" in inspect.signature(...).parameters` is exact, cannot drift, and needs no maintenance when the release lands. The import is wrapped the way _grouped_gemm_dsrelu_backward_supported() already wraps it, so a missing cuDNN answers False rather than raising. Cached, since the call site runs every backward. This also removes the version constant from the code path entirely -- 1.28.0 now appears only in user-facing text, where a release number is the useful thing to say. Tests: a smoke test that the probe returns a bool without raising, with or without cuDNN installed, since reading a signature has more ways to fail than comparing two version strings. It deliberately does not assert which answer -- that depends on the installed front-end, and pinning it would only restate the implementation. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Revert the unrelated nccl-extensions submodule bump `git add -u` in the previous commit swept in a local 3rdparty/nccl-extensions pointer change that has nothing to do with this PR. Restore it to main's commit so the branch touches only the three files it means to. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Raise instead of warning, and cut the change down to what it needs Review asked for two things on the unsupported path: make it an error rather than a warning, and stop branching on self._cudnn_dact_func to pick a message. Both are right, and taking them removes most of the machinery this PR had accumulated. Raising matches what TE already does elsewhere: the Triton grouped dbias/dscales kernel refuses to run under determinism rather than running non-deterministically. It also matches what the variable documents -- "only deterministic algorithms are allowed" is not "prefer deterministic algorithms". A silently non-reproducible run is the failure this PR exists to prevent, so continuing past a request TE cannot honor was the wrong default. Checked that no existing determinism test hits this path: test_hybrid_quantization sets the variable for an attention recipe, and test_fusible_ops_with_userbuffers for linear ops. One message, no branch. The two cases did have different remedies, which is why the branch was there, but a single sentence states both facts -- "needs the scaled-SReLU activation and nvidia-cudnn-frontend 1.28.0 or later" -- without telling a SwiGLU user to go upgrade. What that let me delete: * _warn_nondeterministic_cudnn_dprob and its per-reason lru_cache, the two reason strings and the branch selecting them: 30 lines at the call site and above it, down to a single raise. * _cudnn_frontend_supports_deterministic_dprob as a standalone function. The probe now lives in GroupedMLP_CuTeGEMMUnary.grouped_gemm_dactivation_is_deterministic(), which reaches the wrapper through grouped_gemm_dactivation_kernel() -- the import and its ImportError handling already existed there, so folding it in dropped a duplicate import and an indirection. * The warn-once cache-clearing fixture in the tests, and the two tests that existed only to cover the warning. Tests: test_deterministic_dactivation_is_numerically_correct becomes test_determinism_either_runs_or_refuses -- it expects RuntimeError where the request cannot be honored and runs the full numerical check where it can, so both arms assert something either way. The bit-exactness and two-knob tests are unchanged in substance. Net: transformer_engine/pytorch/ops/fused/grouped_mlp.py goes from +106 to +68, all of it addition, no line of pre-existing code touched. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Apply suggestion from @vthumbe1503 Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * Match the feature-detection idiom main just landed The SiTU-GLU merge (NVIDIA#3402) brought _cudnn_frontend_supports_grouped_gemm_situglu() into this file, which asks inspect.signature(wrapper).parameters for the arguments it needs rather than comparing frontend versions -- the same conclusion this branch reached independently, now the house style. Two things to match. Guard the signature call with `except (TypeError, ValueError)`: a callable that is not introspectable answers "no" instead of raising out of a backward pass. I had left this out on the grounds that the wrapper is a plain undecorated function, which is true today but is not a property this code controls. And say "feature-detect" in the docstring summary, as the neighbor does. Also dropped the sentence about use_single_group_runtime_offsets from the docstring. The neighbor now demonstrates the pattern in the same file, so the cautionary tale is no longer what makes the choice legible. `import inspect` came in with the merge, so this branch no longer adds it. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Cut the comments down to the file's own register The new code carried multi-paragraph docstrings into a file whose 44 functions have a median docstring of one line. Measured before and after: grouped_mlp.py _deterministic_algorithms_required 10 -> 3 lines grouped_gemm_dactivation_is_deterministic (base) 5 -> 1 grouped_gemm_dactivation_is_deterministic (unary) 7 -> 1 test_grouped_mlp.py four new tests 4-6 -> 1-4 four inline comment blocks 2-3 -> 1 each Before this, the three new functions were the 2nd, 3rd and 5th longest docstrings in grouped_mlp.py; only fuse_grouped_mlp_ops, which has a full Parameters block, was longer. In the test file, 63 pre-existing tests have a median docstring of zero lines. Most of what came out was rationale, not explanation: why the union matches DotProductAttention, why feature detection beats a version compare, which cuDNN release window motivated it. That belongs in the commits that made those choices, where it already is, and it reads as noise next to _cudnn_frontend_supports_grouped_gemm_situglu -- the neighbor doing the very same feature detection in a one-line docstring with no rationale at all. What stayed is what the code cannot say itself: that the check sits inside the non-unit-scale branch because a unit scale produces no dprob; that hidden_size must exceed one N-tile or the bit-exactness test is vacuous; that bias would reroute probs.grad through Triton; that weight grads are excluded because wgrad has its own atomics. Each is now one line. No behavior change -- comments, docstrings and one local variable's reading order only. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Flatten the determinism check and the runs-or-refuses test Structural cleanups from the review pass. grouped_mlp.py: the check was nested two deep inside `if not unit_activation_scale`, and assigned deterministic_dactivation only to immediately test its own assignment. Hoisted to two flat statements right after unit_activation_scale is computed. `not unit_activation_scale and _deterministic_algorithms_required()` now says in the expression what the comment had to say in prose, and the separate `= False` initializer is gone. The local itself stays -- the kwargs dict is built about sixty lines further down. Also shortened the error: the tile-scheduler detail was not actionable, and "this activation's cuDNN dactivation kernel" is more accurate than naming the grouped-GEMM backward, since which kernel it is depends on the activation. test_grouped_mlp.py: fused_cls was derived from `activation` by a five-line conditional inside the test; it is now the second half of the parametrize pair. That also fixes the skip guard, which asked GroupedMLP_CuTeGEMMGLU.is_supported() on both parametrizations including the SReLU one -- the sibling test three functions down already gets this right. The _run closure existed only so an if/else could call it twice; a contextlib.nullcontext / pytest.raises choice removes the closure and the branch. nullcontext is used in ten test files here, so it is the local idiom rather than a new one. Not taken: dropping the `isinstance(..., bool)` assertion. It looks vacuous but it is the only coverage of the ImportError branch in the capability probe, which is the branch that runs on every machine without cuDNN -- including CI. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Close the second dprob producer, and stop discarding fp64 test tensors Two findings from the review pass. dprob has two producers in this backward, and the check only covered one. The cuDNN epilogue produces grad_scales at fuser_backward, and when scale_bias is set compute_grouped_dbias_dscales accumulates into it further down -- the Triton kernel that grouped_dbias_dscales.py documents as nondeterministic atomic adds. That kernel's own guard reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. So the hole opened exactly where this branch widened the trigger. With torch.use_deterministic_algorithms(True) and the variable unset -- the case the union exists to start honoring -- SReLU on a 1.28.0 front-end with scale_bias passed the new check, set deterministic=True, raised nothing, and then routed dprob through the nondeterministic path anyway. Env-var users were never exposed: the Triton guard fires for them. It was reachable only via the torch flag, which is to say only through what this branch added. The test picked bias=False and so never crossed it. scale_bias is computed ~130 lines earlier in the same scope, so the fix is to require both producers rather than one. Still one condition and one message, per review -- the message now lists all three requirements instead of two. Separately, the bit-exactness test built its tensors with make_reference_and_test_tensors and discarded the reference every time. That helper allocates an fp64 CPU companion, quantizes and dequantizes for MXFP8 representability, then copies back D2H with an implicit sync -- about 16 MB of host allocation across the two (1024, 1024) calls, for a test that compares run 1 against run 2 and never against a reference. Twelve of the file's other fifteen uses keep the reference; this one had no use for it. Plain uniform_ tensors instead. Also dropped a .item() sync for a token count already known in Python. Not taken, with reasons: * Hoisting _deterministic_algorithms_required into pytorch/utils.py so the Triton guard reads the same union. That is the deeper fix and it is correct, but broadening that guard changes behavior for callers this PR does not touch (ops/basic/grouped_linear.py, module/grouped_linear.py) -- users who set only the torch flag would start seeing RuntimeError where they now get silent nondeterminism. Worth doing deliberately, not as a side effect of this branch. * Extracting the signature-probe shared with _cudnn_frontend_supports_grouped_gemm_situglu. The overlap is about four lines and the two are not interchangeable; refactoring working code outside the diff to save them is not this PR's job. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Add the regression test for the scale_bias hole The previous commit fixed a real bug and shipped it with no test. Every test in the class used bias=False and every end-to-end one set the env var, so neither half of the bug was reachable: not scale_bias, and not the torch-flag-only trigger. Both halves are load-bearing. With the env var the Triton kernel raises on its own, so an env-var test would have passed before the fix as well as after and pinned nothing. Only torch.use_deterministic_algorithms with the variable unset reaches the state where this op's check said yes and the Triton reduction then ran nondeterministically. warn_only=True so torch's own enforcement cannot raise first and be mistaken for TE's refusal. are_deterministic_algorithms_enabled() still reports True in that mode -- the separate is_deterministic_algorithms_warn_only_enabled() getter exists precisely because the two are independent -- so the predicate under test sees what it should. Not executed: no GPU or torch on the machine this was written on. Formatting and syntax only, like the rest of the branch. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Make the bit-exactness test capable of failing Followed cudnn-frontend#521's own test work and found this test had the flaw its commit 88c7fab was written to fix, at the same config. That commit measured 16 launches per shape and found that at l=4 / [256]*4 / n=512 the NONDETERMINISTIC dprob is already bit-stable: the assertion cannot fail there, so a pass certifies nothing. It varies 15/15 at l=8 / [1024]*8 / n=2048. This test used l=4 / [256]*4 / n=1024 -- the vacuous shape, one power of two along n. Moved to the shape that actually varies. n > 256 was necessary but not sufficient, which is what the old comment got wrong. Spanning several N-tiles exercises the within-CTA subtile ordering; making the cross-CTA reduction unstable needs the larger token count and expert count too. Also took the rest of NVIDIA#521's discipline for these comparisons: * Repeat rather than compare a pair. The order determinism removes is set by the tile scheduler, so two runs can match by luck. Four by default, NVTE_TEST_DETERMINISM_REPEATS to raise it, matching that file's DETERMINISM_REPEATS. * Compare bytes, not values. torch.equal treats +0.0 and -0.0 as equal, and a change in reduction order produces exactly that; upstream's bitwise_bits views as uint8 for the same reason. * Assert the output is finite first, so a NaN run cannot be read as a determinism result. Not copied: asserting that the nondeterministic path *does* vary. It is the thing that makes the config meaningful, but as an assertion it is timing-dependent and would flake. Upstream settled this by measuring once and pinning the config; the comment now cites that measurement so the next person does not shrink the shape back. Not run yet: job 535935 is building the previous revision of this test. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Pick the bit-exactness config by measuring it, not by borrowing one Measured on GB300 across five shapes, determinism off, 8 launches each (job 538058), counting how many runs differ from run 0: l=8 tok/grp=1024 n=2048 2/7 max|d| 5.96e-08 l=8 tok/grp=1024 n=4096 2/7 max|d| 9.54e-07 l=16 tok/grp=1024 n=2048 7/7 max|d| 1.19e-07 l=8 tok/grp=2048 n=2048 5/7 max|d| 7.63e-06 l=4 tok/grp=512 n=8192 6/7 max|d| 1.91e-06 Moved to l=16, the only shape where every run differs, so the assertion cannot pass by luck. The previous choice, l=8, varies 2/7 -- an eight-run sample calls it stable often enough to be a poor detector, and an earlier control run (537313) did exactly that and reported 0/7 at this shape. I took that single sample as proof the config was vacuous and said so; it was a sampling artifact, and the shape does vary, just weakly. That earlier shape came from cudnn-frontend#521's own measurement, which was taken on its direct wrapper test. It does not transfer to TE's path -- different scheduler settings, different quantization -- so borrowing the number was the mistake underneath both errors. This config is measured through the fused grouped MLP itself. Two things the same job settled that are worth recording: * Without NVIDIA#521 the values genuinely move: 6e-08 to 8e-06 absolute across these shapes. Small, but nonzero every time, and the reason the refusal exists rather than a warning. * The refusal cannot be exercised against the stock 1.27.0 frontend on this image at all. TE's forward passes prob_tensor=None because _cudnn_frontend_version_at_least("1.27.0") reports optional-prob support that a stock 1.27.0 does not implement, so the op dies in fuser_forward with "prob_tensor is required" before any determinism code runs. Same version-gate-too-coarse failure this PR avoids for its own argument, on a gate it does not own. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> * Make the scale_bias half of the dprob check readable The condition was written as `not (A and not B)` with `scale_bias` as B, which gives a reader no way to tell why an FC2 bias flag decides whether a scale gradient is reproducible -- and the call that makes it relevant is ~250 lines further down. Same logic, named and nested: `dprob_is_deterministic` says what the conjunction means, and the comment names the mechanism instead of gesturing at it. dprob is finished by two kernels, not one -- the cuDNN dactivation epilogue writes it, and then, when scale_bias is set, fuser_backward hands it to compute_grouped_dbias_dscales as the `dscales` accumulator, which atomically adds into it. Its docstring is explicit: "Both outputs use fp32 atomic adds, so pre-populated tensors are accumulated into." No logic change; the truth table is identical. Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> --------- Signed-off-by: Zhiyu Li <zhiyul@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
* [PyTorch] Release warmup outputs after their last use Signed-off-by: Robin Zhang <robinz@nvidia.com> * [PyTorch] Release buffer-reuse capture temporaries Signed-off-by: Robin Zhang <robinz@nvidia.com> * [PyTorch] Release per-callable state on reset Signed-off-by: Robin Zhang <robinz@nvidia.com> * [PyTorch] Bundle per-callable lifecycle helpers Signed-off-by: Robin Zhang <robinz@nvidia.com> --------- Signed-off-by: Robin Zhang <robinz@nvidia.com>
) * Remove redundant runtime checks for activation recompute in MLP from _ScaledUnary class in activation.py Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Add warning for activation recompute in MLP outside fused path in _ScaledUnary class Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Add test for Scaled SReLU activation recompute warning outside fused MLP path Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Enhance activation recompute warning in ScaledSReLU: Update test to verify multiple warnings during backward passes and refactor warning mechanism to ensure it survives Dynamo tracing. Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> * Remove redundant logic to only warn once Signed-off-by: Tim Moon <tmoon@nvidia.com> * Remove unhelpful test Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Ravi Ghadia <rghadia@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com>
* fix: remove unreachable fused-attn backend skip and add regression test
Root cause: FusedAttnRunner._check_configs skipped with 'Unsupported
inputs combination or device compute capability.' unless the backend was
NVTE_F16_arbitrary_seqlen. The later elif re-testing
self.backend != NVTE_F16_arbitrary_seqlen could therefore never be
reached, so its skip message ('B1SS, BHSS and 11SS bias shapes are only
supported for the F16_arbitrary_seqlen backend') was dead code.
Fix: drop the dead elif arm. The padding-mask skip in the sibling if arm
is retained. A regression test locks the remaining behavior: a
non-1HSS post-scale-bias config (BiasShape._B1SS) that passes
_check_configs selects NVTE_F16_arbitrary_seqlen, proving the removal is
behaviorally invisible and the earlier guard is the sole gate.
Testing: not run locally - the JAX test stack (jax,
transformer_engine_jax) is not installed on this machine. The suite runs
in NVIDIA TransformerEngine CI. Contribution: tests/jax tests target
real GPU/cuDNN fused-attention kernels and skip otherwise; the new test
will follow that path via CI.
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
* Remove unnecessary test
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
---------
Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
…IDIA#3482) * Use the shmem alignment operator consistently Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Warn Linear argument documentation Signed-off-by: Evgeny <etsykunov@nvidia.com> * Update docstring, remove warning Signed-off-by: Evgeny <etsykunov@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com>
* Fine-grained recipe docs Signed-off-by: Evgeny <etsykunov@nvidia.com> * Update docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com> * resolve comments Signed-off-by: Evgeny <etsykunov@nvidia.com> * Rework heterogeneous quantization docs into mixed-format quantization; add per-recipe Quantizer sections Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Align mixed-format quantization diagrams with shared diagram-colors.css and dark mode Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Rename mixed-format quantization docs to fine-grained quantization recipes Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Simplify quantizer factory paragraph Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Generalize the fallback-path description Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com> Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Evgeny <etsykunov@nvidia.com>
The forward/backward pair was the registration primitive, returned as an _OpPair carrying every object the function had created, because the autograd-wired variant finished the registration outside it. Forward and backward are registered almost identically, so the primitive is now one op: _register_op returns a _RegisteredOp with the plan and the base/wrapper handles, and the two entry points register the pair as two calls. This also makes a forward-only op possible later without a placeholder backward. ForwardResult is gone: the autograd-free forward returns a plain (output, aux) tuple, symmetric with the backward. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… when available (NVIDIA#3495) [PyTorch] Use torch's register_custom_class API for opaque quantizers PyTorch renamed the opaque-object entry points: register_opaque_type -> register_custom_class, is_opaque_value_type -> is_opaque_constant_type, typ="value" -> typ="constant". The old names remain as thin wrappers that log a deprecation warning on every call, so each TE import now emits 21 warning lines (5 quantizer classes registered, 11 type checks) -- once per process, including every torch.compile worker. A 16-rank training job logs them from ~320 processes. Switch quantizer_opaque.py, custom_op.py and the torch.compile test to the new names. Both import sites already sit behind try/except that records the torch.compile custom-op path as disabled, so a PyTorch that predates the rename (torch < 2.14) keeps importing and running eagerly with the compile path off -- the same behaviour as a PyTorch without the opaque-object API at all. Signed-off-by: Michal Marcinkiewicz <michalm@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
* Add NVTE_PLUGIN support for plugin system on main branch Port plugin system changes from backend_rc2.14: - common/__init__.py: load plugin after framework extension init - dot_product_attention.py: override FlashAttention and get_attention_backend when plugin is active Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn> * Fix lint Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> --------- Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn> Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com> Co-authored-by: Xianduo Li <lixianduo@mail.nankai.edu.cn> Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com> Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
…LU (NVIDIA#3485) [PyTorch] Take the num_groups == 1 dense shortcuts only when the kernels agree At num_groups == 1 the fused grouped MLP may quantize and scale-swizzle the token buffer as ONE dense tensor over `tensor.shape[0]` rows and read it back the same way. That is coherent only if the cuDNN kernels also write their intermediates densely over `tensor.shape[0]`, which is what `use_single_group_runtime_offsets` asks of them -- they then derive M from the runtime tensor shapes rather than from the offsets. The two halves of that decision are currently gated differently. The request to the kernels is already guarded, because it can be refused: if supports_single_group_runtime_offsets: fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1 but its siblings -- `use_single_group_weight_swizzle`, `use_single_discrete_weight`, `use_single_group_dense_fc2`, `use_single_group_dense_dgrad`, the dense quantize and the plain wgrad GEMM -- are gated on `num_groups == 1` alone. So when the specialization is unavailable, `ScaledSReLU` at any version or any activation on a cuDNN frontend older than 1.27.0, the kernels pack their output to `sum(split_sizes)` rows while TE still reads everything densely. Because the columnwise swizzled MXFP8 scale layout is `[k/128][m/128][32][4][4]`, with the m-tile as an INNER stride, an m-extent mismatch misindexes every k-tile past the first: consumer k-tile j reads producer k-tile j*T_c/T_p. The HEAD rows -- the ones that do belong to the group -- pick up the wrong scale factors, and high k-tiles read scale memory the producer never wrote. The data is read correctly; only the scales move, so the error is order 1 and often NaN. Crucially this is value-independent: zeroing the padded tail does not prevent it, because the defect is in the indexing rather than in the padding's contents, so a caller cannot work around it. Callers may legally pass a token buffer padded past `sum(split_sizes)` -- an MoE expert-parallel fixed-capacity receive buffer always does, and `test_grouped_linear_cuda_graph_safe` documents the padded tail as "intentionally outside every group" and "uninitialized". Measured on B300 with MXFP8 and a native ScaledSReLU at num_groups == 1, input padded 512 -> 768 rows, FC1 main_grad against the unpadded result: before zeroed tail 406528 / 1048576 non-finite before poisoned tail 406528 / 1048576 non-finite (identical -- not the values) after either tail 0.0000e+00, bit-identical control num_groups=2 clean before and after Every corrected arm also sits at 1.70e-02 against an unfused reference, the MXFP8 fused-vs-unfused noise floor. Gate the whole shortcut family on the same predicate the kernels use, applying the guard TE already writes at the sites that were missing it. No new helper: the binding calls `_cudnn_frontend_supports_single_group_runtime_offsets` directly, as this file already does at three other sites, and the flag is named `use_dense_single_group` to match its neighbours. It deliberately does not claim the group spans the buffer -- that is unknowable here -- only that producer and consumer will agree on the extent. It is a host-side check on an activation type and a package version, so it costs nothing and stays CUDA graph capturable, and it leaves the GLU activations -- which do get the specialization -- byte-for-byte on their existing fast path. That matters: removing the shortcuts outright costs a shared expert +6% to +30% of the fused MLP step (+5-10% at hidden 4096 / ffn 14336, rising to +28-30% at 1024/4096 with few tokens), measured by CUDA-graph replay against num_groups >= 2 null controls, and Megatron's FusedSharedExpertMLP builds exactly this path. In the two grouped-quantize helpers the condition is just `use_dense_single_group`: `num_groups == 1` is already folded into that flag at both binding sites, and the `isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer))` test was dead because `fuse_grouped_mlp_ops` only builds this fusion for an MXFP8 or NVFP4 recipe. All the shortcuts have to move together: quantization layout and kernel indexing must agree, so enabling a subset pairs a dense consumer with a packed producer and yields NaN rather than a slowdown. Hence one flag threaded explicitly through the grouped-quantize helpers and `_compute_grad_params` rather than a predicate per site. What this does NOT do: where the specialization is available, producer and consumer agree and the residual exposure is that the plain wgrad GEMM still contracts over `logical_shape[0]` rows, summing any padding past `sum(split_sizes)` into the weight gradient. That part is value-dependent -- a zeroed tail contributes an all-zero outer product and is bit-exact, verified over 10 launches -- so a caller that zeroes its padding, or passes none, is unaffected. Closing that gap needs either a device-side masked zeroing of TE's own intermediate buffers or an explicit per-op dense-single-group contract; neither is attempted here. The condition cannot be evaluated at run time: `split_sizes` is device-resident, so testing `sum(split_sizes) == tensor.shape[0]` forces a device-to-host sync, and that test exists to keep this flow sync-free and graph-capturable. tests/pytorch/test_grouped_mlp.py is at baseline: 6 failed, 48 passed, the 6 being pre-existing nvfp4_rht and cuBLAS-version failures also present on main. Signed-off-by: Wentao Guo <wg0420@princeton.edu> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
…IA#3483) * Added device check that first dim is a multiple of 128 Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Fixed tensormap release-acquire protocol Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed error description Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Fix grouped quantize TMA descriptor updater launch Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * Fix shared-memory race in Grouped MXFP8 direct mapper Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
num_grad_inputs leaves BasicOperation: on the autograd-free path it only validated the length of backward_compute's result, and a default of 1 was wrong for every op with a parameter. register_custom_op takes it as an optional check. fwd_kwarg_names is read off resolve_fwd_args's keyword-only parameters instead of being declared twice. The fuser gate now matches the title: a group of several operations runs the eager implementations. Multi-op groups without fusions were slipping through. The saved_tensors reset in backward no longer skips compile; the contexts are copies local to that subgraph, so the write is allowed, and gating it only suggested a constraint that is not there. Drop test_te_ops_setup_context_saves_parameter. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
op_forward / compiled_op_forward and op_backward / compiled_op_backward share one body each, parametrized by the compute callable. Registration moves into _register_compile_ops. The base resolve_fwd_args no longer takes **kwargs, which the registration forbids on subclasses. Redundant checks and long comments removed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
forward_compute / backward_compute become forward_impl / backward_impl, matching module/linear.py and register_custom_op's parameters. The fixed resolve_fwd_args parameters are read off the base signature instead of a separate constant. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Skip tensor storage cleanup while tracing eager operation bodies. Reject delayed-scaling state in the fuser, including CustomRecipe, before compiled execution can omit its backward scale update. Add regression coverage for input preservation and built-in/custom delayed-scaling rejection after eager warmup. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Remove the test additions from 2f8582f at the requested scope. Keep both production fixes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
14 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
First step towards
torch.compile(fullgraph=True)support fortransformer_engine.pytorch.ops.The approach: keep the pipeline-level
_OperationFuserAutogradFunctionand let Dynamo trace it as a higher-order op, with each fusible operation calling its own custom op inside. That keeps the forward and backward fusion layouts independent (so the backward-only fusions survive), keepsOperationContextinside the traced scope, and bounds op registration to one entry per op class.This PR makes that work for an
OperationFusergroup holding one operation. It converts no real operation: the fuser's path is exercised by a test-only operation, so it does not depend on which operations happen to have a custom op. Converting the operations themselves is a follow-up.Type of change
What an operation implements
A
BasicOperationopts in by declaring two argument containers and implementing the methods below.op_forward/op_backwardare then inherited from the base class; the operation does not write any custom-op plumbing itself.Declarations
fwd_args_type,bwd_args_type: dataclasses. Their fields define the custom-op schema (tensors, quantized tensors, quantizers, plain values), exactly asLinearFwdArgs/LinearBwdArgsdo forte.Linear.Traced part (runs in the Dynamo-traced region, may read
self, module state andctx)resolve_fwd_args(input_, *, requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer, <kwargs>) -> fwd_args_type: gathers everything the forward needs into the container. Forward kwargs the operation accepts are declared as additional keyword-only parameters with defaults;fwd_kwarg_namesis derived from this signature.resolve_bwd_args(ctx, grad_output) -> bwd_args_type: rebuilds the backward's inputs from the saved context.setup_context(ctx, args, aux)(optional): decides what to save for backward. The default savesaux. An operation whose backward needs an input or a parameter saves it here fromargs, because a custom op may not return one of its own inputs.Custom op implementation (classmethods; read nothing but
args, mutate nothing)forward_impl(args) -> (output, aux): the forward.auxis a tuple of fresh tensors produced inside the op that the backward needs (norm statistics, a quantized copy of the input). Inputs and parameters are notaux.forward_fake(args) -> (output, aux): the same overTensorSpec, allocation-free.backward_impl(args) -> (grad_input, *grad_params): the backward.grad_inputmay beNone, meaning "grad_output, unchanged".backward_fake(args): itsTensorSpectwin.The
*_implmethods are called directly on the eager path and are the bodies of the registered custom ops on the compiled path, so there is one implementation of the math. Because they are custom-op bodies, they may not mutate their arguments: noclear_tensor_dataon inputs or saved tensors, no in-place writes to tensors from an enclosing scope.Derived by the base class
compile_ops: the registered(forward_fn, backward_fn)pair, one registration per class, under the lower-cased class name.fwd_kwarg_names: the extra keyword-only parameters ofresolve_fwd_args.compile_unsupported_reason():None, or why the operation must run its eager implementation. The base class rejects operations without a custom op and quantizers that are not value-opaque (delayed scaling).Changes
register_custom_opnow defines an operation's forward and backward as two independent custom ops and returns both, leaving autograd to the caller. The variant that wires autograd itself keeps the old behaviour underregister_custom_op_with_autograd. Both are built on a single-op primitive (_register_opreturning a_RegisteredOp), so the pair is two calls and a forward-only op is possible later. The autograd-free forward returns a plain(output, aux)tuple.BasicOperationgains the contract above;__init_subclass__registers the custom op once per class.OperationFuserdecides once per group whether to run its operations' custom ops. The group runs its eager implementations when: it holds more than one operation, either pass has a fusion, an operation has extra tensor inputs/outputs, a kwarg is undeclared or not a tensor, or an operation reportscompile_unsupported_reason. Only tensor kwargs are accepted because a value kwarg becomes a symbolic scalar on its second value, which the opaque value bundle cannot carry; a 0-d tensor is the way to pass a scalar.Side effects that had to go
All of them reached outside the higher-order op's scope:
OperationContextobjects are created in the forward, but the backward is a separate subgraph, so writing to them there mutates an enclosing scope. The backward copies them into its own scope.requires_grad_on an output: AOTAutograd's functionalization drops it anyway, and autograd marks the outputs of anapply()itself._do_not_clearon inputs and outputs.They are gated on being traced, not on using the custom ops. Under
fullgraph=Truethere is no leaving the graph, so an unsupported operation does not "fall back": the pipeline is traced either way and only the choice of implementation changes. The gate reports its reason throughwarn_compile_eager_fallback, which is safe to call from the traced region.Testing
test_torch_compile.pygains two test-only operations and five tests:fullgraph=True;test_torch_compile.py+test_fusible_ops.py(without grouped / userbuffers cases): 1616 passed, 1039 skipped, 1 xpassed. RTX Ada.Gated out and untouched: multi-operation groups, fused operations, operations with extra tensor inputs/outputs, grouped operations, userbuffers, delayed scaling, FP8 block scaling.
Checklist: