[PyTorch] Declare Bias and the activations as custom ops - #30
Open
pggPL wants to merge 1 commit into
Open
Conversation
Each declares its two argument containers and four compute classmethods, so BasicOperation registers its custom ops and the fuser can run it under torch.compile. Nothing in the eager path changes: op_forward / op_backward drive the same implementations. The activations share one implementation and dispatch to their per-class kernel through cls, so all ten subclasses get their own registered op without a factory. Their dispatch methods become staticmethods, which is what 'the compute half must not depend on an instance' means in practice. They are also the first operations here that hand back a quantized tensor: with a next-operation input quantizer the kernel writes FP8 directly, and the tensor crosses the op boundary as its inner buffers, rebuilt on the far side from the fake's TensorSpec. Two contracts the implementations had to respect. A custom op may not return one of its own inputs, so a backward that passes its gradient through returns None for that slot and the caller substitutes; cloning would cost a full-size copy on the common path. The same trick does not work for saved tensors -- the fake cannot see strides, so it cannot predict whether contiguous() will be a no-op, and the mismatch surfaces as an inductor assertion on the sentinel's rank. Hence the activations keep their input only when cache_quantized_input is set. test_ops_custom_ops.py is driven by a single list of operations: adding one means adding one entry. test_ops_hop_poc.py runs a three-operation pipeline whose backward walks a coarser grouping than its forward, which an operation with per-op autograd could not express. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
This was referenced Aug 5, 2026
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
Builds on #29, which made an
OperationFusergroup holding one operation compile but converted no realoperation. This one converts
Biasand the activation family, so they run through their custom opsunder
torch.compile(fullgraph=True).Nothing in the eager path changes:
op_forward/op_backwarddrive the same implementations.Type of change
Changes
Biasand the activation family each declare their two argument containers and four computeclassmethods;
BasicOperationdoes the registration and the plumbing.cls, so allten subclasses get their own registered op without a factory. Their dispatch methods become
staticmethods, which is what "the compute half must not depend on an instance" means in practice.
quantizer the kernel writes FP8 directly, and the tensor crosses the op boundary as its inner buffers
(
_data,_scale_inv), rebuilt on the far side from the fake'sTensorSpec. The subclass itselfnever crosses the schema.
Two contracts the implementations had to respect
A custom op may not return one of its own inputs, so a backward that passes its gradient through
returns
Nonefor that slot and the caller substitutes; cloning would cost a full-size copy on thecommon path.
The same trick does not work for saved tensors. The fake cannot see strides, so it cannot predict
whether
contiguous()will be a no-op, and the resulting metadata mismatch surfaces as an inductorassertion on the sentinel's rank. That rule has to be static, which is why the activations keep their
input only when
cache_quantized_inputis set.Testing
test_ops_custom_ops.pyis driven by a single list of operations -- adding one means adding one entry-- and checks fake-vs-real conformance (the flat
Tensor[]slot layout, via the framework's ownhelpers), numerics against eager, and
fullgraph=True, each with and without an FP8 output.test_ops_hop_poc.pyruns a three-operation pipeline whose backward walks a coarser grouping than itsforward -- something an operation with per-op autograd could not express -- and carries an FP8 tensor
from one operation to the next inside the traced region.
Full run: 1633 passed, 1187 skipped. Lint 10.00/10. RTX Ada, so the FP8 coverage is current scaling only.
Still gated out: multi-operation groups, fused operations,
BasicLinear, grouped operations,userbuffers, delayed scaling, FP8 block scaling.
Checklist: