Feat (quant_tensor)!: convert QuantTensor to Tensor subclass by Giuseppe5 · Pull Request #1579 · Xilinx/brevitas

Giuseppe5 · 2026-08-14T12:07:55Z

Migrate QuantTensor to Tensor Subclasses

Summary

Migrate Brevitas quant tensors from NamedTuple containers to native torch.Tensor subclasses.
This preserves tensor behavior and autograd while storing quantization metadata as private attributes. The change also updates groupwise storage, runtime proxies, caching, pooling, export paths, and tests.

Changes

Convert QuantTensor and concrete quant tensor types to torch.Tensor subclasses.
Store metadata through _quant_tensor_metadata and private attributes such as _scale, _zero_point, and _bit_width.
Add shared reconstruction and tensor-operation support for set, detach, to, cpu, cuda, device checks, and PyTorch fallback dispatch.
Centralize generic arithmetic in QuantTensor; retain metadata-aware arithmetic for IntQuantTensor pairs.
Consolidate integer and floating-point metadata properties in IntMixin and FloatMixin.
Refactor groupwise quant tensors around shared compressed-storage handling.
Make groupwise shape operations operate on expanded values and return plain tensors.
Update metadata-only caching to use empty tensor storage instead of None.
Add metadata-based truncation for average-pooling paths, avoiding temporary quant-tensor reconstruction during tracing/export.
Update proxy and export code for the new metadata representation.
Correct float runtime metadata ordering during passthrough reconstruction.
Remove legacy fields and APIs.

TruncAvgPool Fix

Avoid reconstructing a temporary IntQuantTensor from cached metadata during tracing and export.
Pass the pooled tensor value and quantization metadata directly to the truncation proxy.
Preserve truncation behavior while avoiding Tensor-subclass reconstruction issues with FakeTensors.
Apply the same fix to TruncAvgPool2d and TruncAdaptiveAvgPool2d.

Breaking Changes

Quant tensors are no longer tuples or NamedTuple instances.
Tuple unpacking, _fields, and tuple-based serialization are no longer supported.
The .tensor property has been removed; use .value.
Legacy metadata fields have been removed:
- value_
- scale_
- zero_point_
- signed_t
- training_t
- saturating_t
Generic arithmetic with non-matching operands may now return a plain Tensor.
Unsupported PyTorch operations return plain tensors and may discard quantization metadata.
Metadata-only cached values are now empty tensors instead of None.

Tests

Added coverage for Tensor-subclass behavior, metadata reconstruction, autograd, device/dtype movement, arithmetic, groupwise storage, PyTorch dispatch, caching, and export-related paths.

nickfraser

My feeling is that we want to change the design pattern here. Current QuantTensor's methods are relatively complicated so that they may be used unmodified in the all the derived classes. My personal preference is to keep QuantTensor simple / readable and to override these methods in the derived classes through a Mixin or similar - unless there is a good reason to avoid this that I'm not seeing.

nickfraser · 2026-09-02T09:55:33Z

        if metadata_only:
            self.value = None
-            self.quant_tensor = quant_tensor.set(value=None)
+            self.quant_tensor = quant_tensor.set(value=torch.empty(0))


Any side-effects of this? Worth auditing anywhere we check for qt.value is None in the code or similar?

nickfraser

Quite a few comments - we need to be a bit cautious here. This PR is growing legs from what originally was a sideways movement from NamedTuple -> QuantTensor, with some cleanup included. We seem to be also changing:

some behaviour
some public APIs

The more we change, the more it feels like we should "do everything" in this PR, which would make it unwieldy and difficult to review. We might want to talk about scope here.

Essentially, we've deduped some code and left others duped, which leaves this PR is some "halfway" point between a proper refactor or not.

nickfraser · 2026-09-24T10:10:13Z

+        except Exception:
+            if enabled:
+                model.apply(lambda module: _override_cache_class(module, enabled=False))
+            raise


Does a bare raise like this still produce a human-readable error?

No, this must be changed.

nickfraser · 2026-09-24T10:11:12Z

+            raise
+        finally:
+            if not enabled:
+                model.apply(lambda module: _override_cache_class(module, enabled=False))


It's not clear to be if there is a change in behaviour here or if we're just guarding against future failures. Can you elaborate?

This is a new issues introduced with Tensor subclass instead of NamedTuple; it is a bit lengthy to explain here but the main issue is what happens when:

caching is enabled

and we try to perform operations like quant_tensor.set(value=new_value)

With NamedTuple this was not an issue.

nickfraser · 2026-09-24T10:14:09Z

                x.exponent_bit_width,
+                x.mantissa_bit_width,


Was this a bug before?

nickfraser · 2026-09-24T10:55:21Z

@@ -279,7 +279,7 @@ def forward(
            self._cached_act = cached_inp

        if self.is_quant_enabled:
-            if quant_input is None or isinstance(quant_input, Tensor):
+            if quant_input is None or not isinstance(quant_input, IntQuantTensor):


This logic changed here - it looks intentional, but thought I'd flag it.

Plain Tensors / QuantTensors will resolve to True and will be replaced with the cached value... Do we want this to occur for a QuantTensor that is not an IntQuantTensor?

This is only compatible with IntQuantTensor inputs (and it only produces IntQuantTensor weights, since this inherits from DecoupledWeightQuantWithInputProxyFromInjector which inherits from WeightQuantProxyFromInjector, which is the Int-specific implementation of weight proxy).

I'll add a specific assert to reject other QT

nickfraser · 2026-09-24T11:08:22Z

-    signed_t: Tensor
-    training_t: Tensor
-    dequant_shape: Optional[Tuple] = None
+class GroupwiseQuantTensorMixin:


Why are there so many methods defined here that are unrelated to Groupwise-specific logic. My hope was to avoid lots of duplicated logic my introducing the Mixins. Is there a reason why these should live in the Groupwise Mixin?

Methods that don't seem to be related to groupsize-specific logic:

signed

training

__torch_function__

device

view

reshape

flatten

transpose

permute

__add__

__mul__

__truediv__

A follow-up question is: should view, reshape, flatten, transpose, permute have groupwise-specific implementations, but we can leave that for a follow-up PR.

I've cleaned up a lot of them, but all the views have to be Groupwise specific.

For non-groupwise quanttensor we made the decision of allowing unsafe operations, i.e. you can flatten a tensor but the quant metadata will not receive the same view, even though a new QuantTensor is generated.

# Regular IntQuantTensor: unsafe reshape is allowed. qt = IntQuantTensor( value=torch.ones(2, 4), scale=torch.ones(1, 4), # Per-channel metadata. zero_point=torch.zeros(1, 4), bit_width=torch.tensor(8.), signed=True, training=False, ) flat = qt.view(-1) type(flat) # IntQuantTensor flat.value.shape # torch.Size([8]) flat.scale.shape # torch.Size([1, 4]) <- unchanged, no longer aligned with value

For groupwise quant tensor, generating a new QT with an unsafe view will cause issues everytime .value is called (or .scale, .zero_point), which is why we override and explicitly demote to normal Tensor.
If we preserved QT structure:

# GroupwiseIntQuantTensor: raw storage is compressed. qt = GroupwiseIntQuantTensor( value=compressed_value, # e.g. shape [2, 4, 8] scale=compressed_scale, zero_point=compressed_zp, group_size=8, group_dim=1, bit_width=torch.tensor(8.), signed=True, training=False, dequant_shape=(2, 32), # Logical expanded shape. ) flat = qt.view(-1) bad_qt.value # Expansion uses invalid storage/group geometry. bad_qt.scale # Same issue. bad_qt.zero_point # Same issue.

nickfraser · 2026-09-24T11:12:25Z


    @property
    def signed(self):
-        return self.signed_t.item()
+        return self._signed.item()


Other renaming of attributes have just been related to private values. Does the old version no longer exist? Do we want a breaking change like this with this PR?

I think so; The point is that we have two distinct ways to indicate a "private" attribute currently in master: value_ (NamedTuple doesn't like leading underscore), and training_t (since we store a tensor but we always retrieve a bool).

The idea is to have leading underscore variables as storage variables in all cases; I can also revert training_t/signed_t and leave them as an exception, but since we're changing stuff here (and the fact that this will be a breaking change is almost unavoidable anyway), I'd rather do it here.

Giuseppe5 · 2026-09-29T14:53:17Z


    @property
-    def signed(self):
-        return self.signed_t.item()
+    def exponent_bit_width(self):


These should go into floatmixin

Giuseppe5 · 2026-09-29T14:53:48Z

-            args = _unpack_quant_tensor(args)
-            kwargs = _unpack_quant_tensor(kwargs)
-            return func(*args, **kwargs)
+    def bit_width(self):


Maybe move to IntMixin

Giuseppe5 force-pushed the giuseppe/quant_tensor_subclass branch 3 times, most recently from 76ad16c to 52e0e98 Compare August 20, 2026 09:30

Giuseppe5 force-pushed the giuseppe/quant_tensor_subclass branch 3 times, most recently from 0b10189 to 5448b5a Compare August 26, 2026 08:45

Giuseppe5 self-assigned this Aug 26, 2026

Giuseppe5 force-pushed the giuseppe/quant_tensor_subclass branch from 48eed14 to 7ee8613 Compare September 1, 2026 12:46

nickfraser requested changes Sep 2, 2026

View reviewed changes

Giuseppe5 force-pushed the giuseppe/quant_tensor_subclass branch from 7ee8613 to b929a77 Compare September 2, 2026 13:56

Giuseppe5 requested a review from nickfraser September 8, 2026 09:32

Giuseppe5 force-pushed the giuseppe/quant_tensor_subclass branch from 5ed2434 to 1eb67b3 Compare September 9, 2026 13:51

Giuseppe5 mentioned this pull request Sep 21, 2026

Add squeeze/unsqueeze support to IntQuantTensor #1622

Draft

5 of 14 tasks

Giuseppe5 commented Sep 23, 2026

View reviewed changes

Comment thread src/brevitas/quant_tensor/base_quant_tensor.py

Giuseppe5 added 16 commits September 23, 2026 15:32

Feat (quant_tensor): convert from NamedTuple to Tensor subclass

e010eb4

fix notebook

2291a59

more fixes

2f617a1

added tests

2dfce1c

remove .tensor alias

3a77f66

fix tests

7a1db76

Feat (quant_tensor): support for slicing of metadata

bd85d70

Fix (quant_tensor): normalize negative group dimensions when slicing

bd9efcc

pre-commit

f7e3378

Fix: disable cached QuantTensor reconstruction during Dynamo export

abca05b

fix tests

937ca30

fix test 2

b0bd37c

Simplify interface

cfa7afc

coherent-ish names

f310102

fix tests

24b9485

Fix notebooks

c12eb10

Giuseppe5 added 5 commits September 23, 2026 15:33

simplify interface

3743dad

Review

f5715b8

cleanup

1d38ca0

Rename

83db66d

Last touches

6e813ea

Giuseppe5 force-pushed the giuseppe/quant_tensor_subclass branch from 45dd1bb to 6e813ea Compare September 23, 2026 14:34

removed useless method

aacc269

nickfraser requested changes Sep 24, 2026

View reviewed changes

Giuseppe5 added 3 commits September 24, 2026 13:44

cleanup add/mul/truediv

2372820

Cleanup torchfunc/signed/training

7f0fae8

cleanup

28c4541

Giuseppe5 requested a review from nickfraser September 24, 2026 14:50

Giuseppe5 commented Sep 29, 2026

View reviewed changes

Cleanup

ede737a

Giuseppe5 changed the title ~~Giuseppe/quant tensor subclass~~ Feat (quant_tensor)!: convert QuantTensor to Tensor subclass Sep 30, 2026

Conversation

Giuseppe5 commented Aug 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Migrate QuantTensor to Tensor Subclasses

Summary

Changes

TruncAvgPool Fix

Breaking Changes

Tests

Uh oh!

nickfraser left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

nickfraser left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nickfraser Sep 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Giuseppe5 commented Aug 14, 2026 •

edited

Loading

nickfraser Sep 24, 2026 •

edited

Loading