[PyTorch][torch.compile] Make quantizers opaque value objects

[pggPL]

Description

Tensorless quantizers in TE (MXFP8, FP8 blockwise, FP8 current-scaling, NVFP4)
are fully described by a handful of plain, reproducible scalars — they hold no
live tensors and no process groups. This PR turns them into opaque value
objects so torch.compile can treat them as baked-in constants: two
quantizers with the same configuration become interchangeable, hashable, and
reconstructible inside an FX graph.

Quantizers that hold live state (delayed-scaling Float8Quantizer, which keeps
scale/amax tensors) and any user-defined quantizer keep the default
identity semantics, so the change is opt-in and backward compatible. On older
PyTorch builds without the opaque-object API the registration is a graceful
no-op.

Along the way this also un-breaks the existing test_torch_compile.py suite:
that file lived on main but was never wired into CI, and its
test_autocast_nested_custom case (nested te.autocast with multiple
CustomRecipe instances) was failing because of the CustomRecipe state-caching
bug fixed here. The file is now run in CI and passes.

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

• Add opt-in value-object identity to the base Quantizer
  (_value_fields / _value_key / __eq__ / __hash__). Returning None
  from _value_fields() (the default) keeps identity semantics.
• New module transformer_engine/pytorch/dynamo.py holding the
  torch.compile glue: __fx_repr__, value-key reconstruction and
  register_value_opaque_quantizer (gracefully no-op without PyTorch's
  opaque-object API).
• Register MXFP8Quantizer, Float8BlockQuantizer,
  Float8CurrentScalingQuantizer and NVFP4Quantizer as value opaque types
  (the deprecated amax_reduction_group is never part of the value).
• Fix CustomRecipe state caching in TransformerEngineBaseModule.set_meta_tensor:
  rebuild quantizers when the CustomRecipe instance changes (e.g. nested
  te.autocast regions) instead of reusing the first recipe's state, since
  every CustomRecipe shares the CustomRecipeState type but carries its own
  qfactory. This fixes the previously-failing test_autocast_nested_custom.
• Enable tests/pytorch/test_torch_compile.py in the L0_pytorch_unittest QA
  suite (it existed on main but was never run in CI), and add the quantizer
  value-object tests to it. Bringing it into CI required fixing the existing
  CustomRecipe torch.compile path: the qfactory now dispatches on
  QuantizerRole.tensor_type supplied by ToyLinear.get_quantizer_roles.
• Guard the value-object path against a stored amax reduction group: __fx_repr__
  already rejects any quantizer holding a process group, and __eq__ / __hash__
  now raise too. The group is excluded from the value key, so a stored group would
  otherwise compare/hash equal to a groupless quantizer and let torch.compile
  reuse a graph that skips the reduction. Pass the group per quantize call instead.

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

[pggPL]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR turns four tensorless quantizers (MXFP8Quantizer, Float8BlockQuantizer, Float8CurrentScalingQuantizer, NVFP4Quantizer) into hashable, equality-comparable value objects so torch.compile can bake them into compiled FX graphs as constants rather than breaking the graph. The change is opt-in: only classes explicitly registered via register_value_opaque_quantizer gain value semantics; all others keep identity-based __eq__/__hash__.

• New dynamo/quantizer_opaque.py: implements _rebuild_quantizer (bypasses __init__, then calls _rebuild_derived_state), _quantizer_fx_repr (generates eval-able FX code), and register_value_opaque_quantizer (validates annotated fields are plain value types, sets _value_field_names, attaches __fx_repr__, and calls PyTorch's opaque-object registration \u2014 gracefully no-op on older builds).
• NVFP4Quantizer restructuring: rht_matrix (a torch.Tensor) and amax_reduction_group removed from annotations; _rebuild_derived_state reconstructs rht_matrix from the cached get_rht_matrix; copy() now correctly propagates with_random_sign_mask via rht_matrix_random_sign_mask_t != 0.
• Tests: test_torch_compile.py is wired into CI and covers equality, hash, FX repr round-trip, and a fullgraph=True compile path; the kernel quantize round-trip verifies that derived state (including rht_matrix) is faithfully restored.

Confidence Score: 5/5

Safe to merge; the value-object path is strictly opt-in and the changes do not alter the quantization kernels or their inputs.

The core correctness paths — annotation validation at registration time, value-key construction, FX repr generation, and _rebuild_derived_state for NVFP4's derived tensor — are all well-guarded and covered by the new quantize kernel round-trip tests. The one identified gap (enum-to-int conversion keyed on the field name "dtype" rather than the field type) has no impact on any currently registered quantizer.

transformer_engine/pytorch/dynamo/quantizer_opaque.py and transformer_engine/pytorch/quantized_tensor.py — specifically the name-based enum conversion logic that would need updating before any new Enum-typed annotation is added to a registered quantizer.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/quantizer_opaque.py	New module implementing value-opaque quantizer registration for torch.compile; logic is correct for current quantizers but enum serialization is hardcoded to the "dtype" field name, creating a latent trap for future annotated enum fields.
transformer_engine/pytorch/quantized_tensor.py	Adds _value_fields, _value_key, eq, and hash to Quantizer base class; opt-in value semantics via _value_field_names class attribute; enum-to-int conversion in _value_key is name-specific ("dtype") rather than type-based.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Removes rht_matrix and amax_reduction_group from annotations; introduces _rebuild_derived_state hook that restores rht_matrix from the cached get_rht_matrix; copy() now passes with_random_sign_mask correctly.
transformer_engine/pytorch/tensor/float8_tensor.py	Removes amax_reduction_group annotation from Float8CurrentScalingQuantizer and adds register_value_opaque_quantizer call; copy() uses getattr with default None for the deprecated field.
transformer_engine/pytorch/tensor/mxfp8_tensor.py	Registers MXFP8Quantizer as a value-opaque type; minimal change, only dtype annotation is in the quantizer-specific fields.
transformer_engine/pytorch/tensor/float8_blockwise_tensor.py	Registers Float8BlockQuantizer as a value-opaque type; all annotated fields are plain value types.
tests/pytorch/test_torch_compile.py	Adds value-object tests (equality, hash, FX round-trip, fullgraph compile); tests are well-structured including quantize kernel round-trip to catch missing derived state.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A["register_value_opaque_quantizer(cls)"] --> B["_annotated_fields(): collect MRO annotations"]
    B --> C{"All fields are int/bool/float/str/Enum?"}
    C -- No --> D["Raise TypeError at import time"]
    C -- Yes --> E["cls._value_field_names = tuple(fields)"]
    E --> F["Attach __fx_repr__ = _quantizer_fx_repr"]
    F --> G{"torch opaque-object API available?"}
    G -- No --> H["Value semantics only; torch.compile graph-breaks"]
    G -- Yes --> I["register_opaque_type(cls, typ='value')"]
    I --> J["_VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__)"]
    K["quantizer.__fx_repr__()"] --> L["_value_key() raises if ProcessGroup stored"]
    L --> M["_rebuild_quantizer(cls, items)"]
    M --> N["cls.__new__(cls); set attrs via object.__setattr__"]
    N --> O{"_rebuild_derived_state exists?"}
    O -- Yes --> P["Restore rht_matrix (NVFP4 only)"]
    O -- No --> Q["Rebuilt quantizer ready"]
    P --> Q

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
flowchart TD
    A["register_value_opaque_quantizer(cls)"] --> B["_annotated_fields(): collect MRO annotations"]
    B --> C{"All fields are int/bool/float/str/Enum?"}
    C -- No --> D["Raise TypeError at import time"]
    C -- Yes --> E["cls._value_field_names = tuple(fields)"]
    E --> F["Attach __fx_repr__ = _quantizer_fx_repr"]
    F --> G{"torch opaque-object API available?"}
    G -- No --> H["Value semantics only; torch.compile graph-breaks"]
    G -- Yes --> I["register_opaque_type(cls, typ='value')"]
    I --> J["_VALUE_OPAQUE_QUALNAMES.add(cls.__qualname__)"]
    K["quantizer.__fx_repr__()"] --> L["_value_key() raises if ProcessGroup stored"]
    L --> M["_rebuild_quantizer(cls, items)"]
    M --> N["cls.__new__(cls); set attrs via object.__setattr__"]
    N --> O{"_rebuild_derived_state exists?"}
    O -- Yes --> P["Restore rht_matrix (NVFP4 only)"]
    O -- No --> Q["Rebuilt quantizer ready"]
    P --> Q

Loading

_{Reviews (10): Last reviewed commit: "Enable post-RHT amax in the NVFP4 value-..." | Re-trigger Greptile}

[kshitij12345]

Overall LGTM, would be good to resolve the inline comments before merging.

[pggPL]

/te-ci pytorch

[ptrendx]

There are some ways of getting that list without hardcoding it again, e.g. using inspect.get_annotations in the cases like we have where we provide all of the data members up front or via vars.

[pggPL]

I see It's quite messy, but it is result of mess in quantizers. There are some fields called in the constructor (and arguments names sometimes does not match the field name). Some fields are set like .internal or .block_len are not arguments to constructor. And some of them are excluded on purpose like amax_reduction_group or rht_matrix.

[ptrendx]

Shouldn't we then try to clean the quantizers code instead of adding to the messiness? I think the quantizers should advertise properly all of the fields that the value object would need.
I don't think internal is a problem necessarily (although we could just as well have it in the quantizer constructor, it seems that right now it is very easily discoverable), since you also do not put it in this function, and I assumed what you want here are just the constructor arguments from this subclass?
The ones that we want to "deprecate" in a sense like the amax_reduction_group should be removed from the annotated fields, right?

[pggPL]

/te-ci pytorch