Skip to content

[PyTorch] Declare Bias and the activations as custom ops - #30

Open
pggPL wants to merge 1 commit into
ops_fuser_compilefrom
ops_declare_ops
Open

[PyTorch] Declare Bias and the activations as custom ops#30
pggPL wants to merge 1 commit into
ops_fuser_compilefrom
ops_declare_ops

Conversation

@pggPL

@pggPL pggPL commented Aug 5, 2026

Copy link
Copy Markdown
Owner

Description

Builds on #29, which made an OperationFuser group holding one operation compile but converted no real
operation. This one converts Bias and the activation family, so they run through their custom ops
under torch.compile(fullgraph=True).

Nothing in the eager path changes: op_forward / op_backward drive the same implementations.

Type of change

  • Documentation change (change only to the documentation, either a fix or a new content)
  • Bug fix (non-breaking change which fixes an issue)
  • New feature (non-breaking change which adds functionality)
  • Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • Infra/Build change
  • Code refactoring

Changes

  • Bias and the activation family each declare their two argument containers and four compute
    classmethods; BasicOperation does the registration and the plumbing.
  • 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 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
    (_data, _scale_inv), rebuilt on the far side from the fake's TensorSpec. The subclass itself
    never 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 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 resulting metadata mismatch surfaces as an inductor
assertion on the sentinel's rank. That rule has to be static, which is why the activations keep their
input only when cache_quantized_input is set.

Testing

test_ops_custom_ops.py is 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 own
helpers), numerics against eager, and fullgraph=True, each with and without an FP8 output.

test_ops_hop_poc.py runs a three-operation pipeline whose backward walks a coarser grouping than its
forward -- 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:

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

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant