[PyTorch][torch.compile] Add TensorProto mechanism by pggPL · Pull Request #8 · pggPL/TransformerEngine

pggPL · 2026-06-06T14:52:17Z

Description

This PR introduces TensorProto — a data-free prototype of a tensor (or quantized tensor) that captures everything needed to reason about and rebuild a tensor without holding any storage: its logical shape/dtype and, for quantized tensors, the value-opaque quantizer defining the layout.

The key property is that TensorProto.create_tensor() materializes a quantized tensor purely in Python (via Quantizer.alloc_tensors + the storage's __tensor_unflatten__), so it traces under torch.compile(fullgraph=True) with no graph break — unlike make_empty, which goes through the opaque C++ tex.create_empty_quantized_tensor. This is the foundation for writing torch.library custom-op fake implementations of quantized ops.

This builds on the value-opaque quantizer work (so a TensorProto is itself safe to treat as a compile-time constant).

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

dynamo.py: Add TensorProto dataclass (shape, dtype, quantizer, requires_grad, device) with is_quantized, inner_names(), create_metadata() and create_tensor(), plus a to_tensor_proto() helper that builds a proto from a plain torch.Tensor or a QuantizedTensorStorage/QuantizedTensor.
quantized_tensor.py:
- Add the PyTorch wrapper-subclass flatten protocol (__tensor_flatten__ / __tensor_unflatten__) to QuantizedTensorStorage, driven by a per-class _FLATTEN_TENSOR_BUFFERS declaration of (attribute_name, constructor_kwarg) pairs.
- Add a _STORAGE_REGISTRY (populated via __init_subclass__) so __tensor_unflatten__ can resolve a concrete storage/wrapper class from its qualname inside an FX graph.
- Add pure-Python, traceable allocation hooks to Quantizer: alloc_tensors, create_metadata, and the opt-in overrides _describe_buffers, _storage_scalars, _resolve_storage_cls.
Quantizers: Implement the allocation hooks for Float8CurrentScalingQuantizer, MXFP8Quantizer and Float8BlockQuantizer.
Storage classes: Declare _FLATTEN_TENSOR_BUFFERS for Float8TensorStorage, MXFP8TensorStorage and Float8BlockwiseQTensorStorage.
ops/basic/basic_linear.py: Add allocation-free _functional_forward_fake / _functional_backward_fake that operate on TensorProto and return output/gradient protos, as a basis for custom-op fake impls (single-device only; TP/SP shape effects not yet modeled).
Tests: Add tests/pytorch/test_tensor_proto.py (CPU smoke tests for _describe_buffers/alloc_tensors/create_metadata, flatten round-trip, and to_tensor_proto) and torch.compile fullgraph tests in test_torch_compile.py.

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

Squashed PR #8 (tensor_proto_mechanism) onto the rebased base. Adds TensorProto (pure-Python, torch.compile-traceable quantized-tensor allocation via Quantizer.alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__), Linear fake fwd/bwd impls for the custom-op path, and tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

kshitij12345

Overall looks good

Would it be possible to reduce duplication between _linear_forward_impl_fake and _linear_forward_impl.

Squashed PR #8 (tensor_proto_mechanism) onto the rebased base. Adds TensorProto (pure-Python, torch.compile-traceable quantized-tensor allocation via Quantizer.alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__), Linear fake fwd/bwd impls for the custom-op path, and tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

shino16 · 2026-07-06T08:11:04Z

+    # Rebuilding from the derived proto matches the original tensor's structure.
+    assert _signature(proto.create_tensor(), proto.inner_names()) == _signature(
+        tensor, proto.inner_names()
+    )


What about adding basic tests of _linear_forward_impl_fake and _linear_backward_impl_fake?

shino16 · 2026-07-06T08:42:58Z

    # --- Differentiable tensors (also passed positionally to autograd) ---
    weight: TensorOrQuantized
-    inp: torch.Tensor
+    inp: TensorOrQuantized


As someone new to the codebase, I'm confused with the type annotations here. Can .weight and .inp be QuantizedTensorStorage?

shino16 · 2026-07-06T09:22:04Z

-    backward_needs_input = is_grad_enabled and weight.requires_grad
+    # Use the requires-grad flags captured into ``args`` at op-call time rather
+    # than the live tensors': the fake impl (``_linear_forward_impl_fake``) keys
+    # the number of FP8 inner buffers it emits off ``args.*_requires_grad``, so
+    # the real impl must agree to keep the custom-op output arity stable. Under
+    # ``torch.compile`` with CUDA-graph trees (``mode="reduce-overhead"``) the
+    # static graph inputs are detached during capture, so live
+    # ``weight.requires_grad`` / ``inp.requires_grad`` flip to False mid-capture
+    # and would otherwise diverge from the fake (schema/arity mismatch).
+    backward_needs_input = is_grad_enabled and args.weight_requires_grad


We may want to erase comments about the mismatch with the fake impl. This change looks more like a correctness fix on its own, but the comments made me suspect it was a haphazard patch to avoid crash.

shino16 · 2026-07-06T09:23:11Z

    # --- Differentiable tensors (also passed positionally to autograd) ---
    weight: TensorOrQuantized
-    inp: torch.Tensor
+    inp: TensorOrQuantized


As someone new to the codebase, I'm confused with the type annotations here. Why are they not just torch.Tensor even though we rely on Tensor attributes like .shape?

shino16 · 2026-07-06T09:41:16Z

+        elif new_weight_workspace is not None and wt_save is new_weight_workspace:
+            wt_alias = "new_workspace"
+        elif args.weight_workspace is not None and wt_save is args.weight_workspace:
+            wt_alias = "weight_workspace"


wt_save is guaranteed to be non-None.

Suggested change

elif new_weight_workspace is not None and wt_save is new_weight_workspace:

wt_alias = "new_workspace"

elif args.weight_workspace is not None and wt_save is args.weight_workspace:

wt_alias = "weight_workspace"

elif wt_save is new_weight_workspace:

wt_alias = "new_workspace"

elif wt_save is args.weight_workspace:

wt_alias = "weight_workspace"

shino16 · 2026-07-06T10:20:12Z

+import torch
+
+
+def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]:


Maybe you can use torch._prims_common.make_contiguous_strides_for, although it's a private API.

shino16 · 2026-07-06T10:38:27Z

+        if self.quantizer is None:
+            return


We could be strict and raise an error when self.quantizer is None.

shino16 · 2026-07-06T11:04:21Z

+        if self.quantizer is None:
+            return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)]
+        inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device)
+        return [inner[name] for name in self.inner_names()]


Does this have to be a list? We could return self.quantizer.alloc_tensors(...) directly since we make a dict with the same keys later.

shino16 · 2026-07-06T11:06:42Z

+    )
+
+    requires_grad = bool(getattr(tensor, "requires_grad", False))
+    if isinstance(tensor, QuantizedTensorStorage):


Do we need this if? The else path looks the same.

shino16 · 2026-07-06T11:40:17Z

+            return cls(**kwargs)
+        # Wrapper subclass: it also needs outer shape / dtype / device / stride.
+        fake_dtype = kwargs.get("fake_dtype")
+        device = next((t.device for t in inner_tensors.values() if t is not None), None)


Checking if t is not None may be too careful.

* disable 9.23.0/.1 for mxfp8 attention Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

…for QuantizedTensor (NVIDIA#3172) * fix torch function Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * make fp bs dequantize autograd aware + plus test for autograd flow of QuantizedTensor Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

…#3152) * [PyTorch] Make tensorless quantizers opaque value objects for torch.compile Give tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling, NVFP4) value-object semantics so torch.compile can treat them as baked-in constants: - Add opt-in value identity to the base Quantizer (_value_fields / _value_key / __eq__ / __hash__). Quantizers holding live tensors (delayed-scaling Float8Quantizer) and custom quantizers keep identity semantics. - New transformer_engine/pytorch/dynamo.py houses the torch.compile glue: __fx_repr__, value-key reconstruction and register_value_opaque_quantizer (gracefully a no-op on PyTorch builds without the opaque-object API). - Register the four tensorless quantizers as value opaque types. Also fix CustomRecipe state caching in TransformerEngineBaseModule: set_meta_tensor now rebuilds 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. Move the quantizer value-object tests into tests/pytorch/test_torch_compile.py and add that file to the L0 pytorch unittest QA suite. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Drop quantizer value registry; reconstruct via __fx_repr__ globals Follow-up to the value-opaque quantizer support: - Remove the module-level _QUANTIZER_VALUE_REGISTRY (qualname -> class) and _quantizer_from_value_key. __fx_repr__ now captures the quantizer class directly in the FX globals and reconstructs via _rebuild_quantizer(cls, items), matching how PyTorch's own value opaque types (e.g. DTensor placements) reconstruct themselves. This removes global mutable state and the qualname collision risk. - Consolidate the quantizer value-object tests in test_torch_compile.py down to two functions and exercise reconstruction through the public __fx_repr__ path instead of internal helpers. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Split dynamo.py into a dynamo/ package Replace the single dynamo.py module with a dynamo/ package so the torch.compile glue can grow with a clear responsibility split across the stacked branches. This branch owns the value-opaque quantizer layer. * dynamo/quantizer_opaque.py -- register_value_opaque_quantizer and helpers * dynamo/__init__.py -- re-exports the public API so callers keep importing from transformer_engine.pytorch.dynamo unchanged Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Raise in quantizer __fx_repr__ when a process group is stored A value-opaque quantizer must not carry live distributed state. Scan the quantizer attributes in __fx_repr__ and raise TypeError if any holds a torch.distributed.ProcessGroup (e.g. a non-None deprecated amax_reduction_group), so it cannot be silently baked into a torch.compile FX graph. Clarify the related comments accordingly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [PyTorch] Cover NVFP4 in quantizer value-object test NVFP4Quantizer is registered as a value-opaque quantizer but was missing from the value-semantics / __fx_repr__ round-trip test. Add it to _VALUE_QUANTIZERS (skipped without CUDA, which it needs to construct). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Reject a value quantizer that carries an amax reduction group in __eq__/__hash__ The amax reduction group is excluded from the value key, so a value quantizer that stored one would compare/hash equal to a groupless one and let torch.compile reuse a graph that skips the reduction. __eq__/__hash__ now raise (mirroring __fx_repr__, which already rejects any process-group-bearing quantizer). The group should be passed per quantize call, not stored on the quantizer. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Recognize value-opaque quantizers via a class flag Add is_value_opaque_quantizer() + the _te_compile_value_opaque flag stamped at registration, so dynamo-traced code can detect registered quantizers (and fall back to eager for unregistered ones). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Address review: narrow opaque-type except, add fullgraph test, fix nvfp4 value key - Narrow register_opaque_type except to (RuntimeError, TypeError): the API is already imported above, so ImportError/AttributeError there only mask real errors. - Add test_quantizer_value_object_fullgraph exercising torch.compile(fullgraph=True) end-to-end to verify opaque-type registration took effect. - Restore missing NVFP4Quantizer._with_random_sign_mask assignment required by _value_fields()/_value_key(). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Restore NVFP4 rht_matrix on value-key rebuild; assert quantize round-trip _rebuild_quantizer only restores value-key fields, so a reconstructed NVFP4Quantizer was missing the derived rht_matrix tensor (not hashable, so not in the value key) and failed at copy()/quantize time. Add a _rebuild_derived_state hook (called by _rebuild_quantizer) that NVFP4Quantizer uses to rebuild rht_matrix from _with_random_sign_mask (lru_cache -> cheap). Extend test_quantizer_value_object to also quantize with the original and the rebuilt quantizer and require bit-exact results (gated on HW support), so a field the kernel needs but the value key omits can no longer slip through. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Enforce process-group rejection in _value_key, not __fx_repr__; add test Move the ProcessGroup guard out of the (overridable) __fx_repr__ into Quantizer._value_key -- the single point every value-materialization path (__eq__/__hash__/__fx_repr__) goes through -- so a custom __fx_repr__ can no longer bypass it. Generalizes the old amax-only check to any field holding a ProcessGroup. Add a test that a value quantizer carrying a live group raises. Addresses review on NVIDIA#3152. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Strengthen fullgraph test: quantize/dequantize via a custom op, not passthrough Replace the trivial pass-through fullgraph test with one that drives each production quantizer through a minimal custom op (quantize + dequantize) under torch.compile(fullgraph=True) and compares to eager -- so the opaque-type registration is actually exercised inside the graph (a graph break would make fullgraph=True raise). Op registration sits right before the test. Also drop stale comments referencing the old __fx_repr__-side process-group guard. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Clarify comments: rht_matrix_random_sign_mask_t derivation; why the opaque flag - rht_matrix_random_sign_mask_t is a device-independent int derived from _with_random_sign_mask (the device only places a throwaway tensor); fix the misleading comment. - Explain why registration uses a class attribute, not a registry set: is_value_opaque_quantizer is traced inside the compile graph and dynamo can bake a getattr constant but cannot do 'type(q) in set' on the opaque class. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Reword opaque-flag comment: self-contained, no Linear reference Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Cover is_opaque_value_type with the import-safety guard too is_opaque_value_type(cls) sat between the import guard and the register_opaque_type guard, so on a partial/experimental opaque-object build it could raise RuntimeError/TypeError and crash TE import. Move it inside the same except so the 'registration never crashes import' promise holds for both calls. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Stamp value-opaque flag only after successful registration Move the _VALUE_OPAQUE_FLAG setattr to the end of register_value_opaque_quantizer, after register_opaque_type succeeds (or the type is already opaque). Previously the flag was set up front, so is_value_opaque_quantizer reported True even when the opaque-object API was missing or registration raised, since both paths are swallowed. Eager value semantics (__eq__/__hash__/__fx_repr__) are independent of the flag, so this only tightens the predicate to mean torch actually knows the type as opaque. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Drop verbose comments around value-opaque flag stamping Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Narrow value process-group check to amax_reduction_group _check_value_has_no_process_group ran on every guard eval (via __eq__/__hash__) and scanned all of vars(self) recursively. The only attribute that can hold a ProcessGroup is the deprecated amax_reduction_group, so check it directly (O(1)) and drop the _contains_process_group helper. Same guarantee, off the hot path. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Shorten amax_reduction_group check comment Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Drop trivial value-equality boilerplate from quantizer test Remove the a==b / hash / dict-key block that just exercised Python's own dict semantics; equality and hashing are still covered by the __fx_repr__ round-trip (rebuilt == a, hash match) and the bit-exact kernel check. other_kwargs is now unused, so drop it from the parametrization and both test signatures. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Address review comments: qualname registry, import and comment cleanups - Replace the class-attribute value-opaque flag with a module-level set of class qualnames: a set of class objects is untraceable under fullgraph=True (opaque classes have no equality rule in Dynamo), but the qualname constant-folds to a plain string; also avoids falsely reporting unregistered subclasses. - Register MXFP8Quantizer right after the class like the other quantizers. - Clarify the amax_reduction_group exclusion comment in Float8CurrentScalingQuantizer._value_fields. - Restore import order in test_torch_compile.py, import NVFP4Quantizer from transformer_engine.pytorch, drop unused Float8Quantizer import. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Derive quantizer value fields from class annotations Make the class annotations the single source of truth for what defines a quantizer's value, instead of hand-written per-class _value_fields lists: - Quantizer._value_fields is now derived from the annotations across the MRO (subsuming _BASE_VALUE_FIELDS); register_value_opaque_quantizer is the explicit opt-in and validates at import time that no annotated field is a tensor or process group. - Drop the four per-class _value_fields overrides. - Remove the deprecated amax_reduction_group annotation from Float8CurrentScalingQuantizer and NVFP4Quantizer (the attribute is still set for backward compatibility). - NVFP4: rename _with_random_sign_mask to with_random_sign_mask (annotated, matching the constructor argument), stop storing the derived rht_matrix_random_sign_mask_t in the value key and rebuild it together with rht_matrix in _rebuild_derived_state (lru-cached getters), which __init__ now also uses. copy() now propagates with_random_sign_mask. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Drop redundant annotation-exclusion comments Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Shorten _is_value_quantizer comment Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fold annotation walk into _value_fields, trim comments Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Simplify qualname-registry comment Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Reword qualname-registry comment for outside readers Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fix review findings: pickle compat, subclass opt-in, cached value fields - NVFP4: stop storing with_random_sign_mask; the annotated value field is the derived (deterministic, device-independent) rht_matrix_random_sign_mask_t and rht_matrix is rebuilt from it in _rebuild_derived_state. Quantizers pickled before this PR already carry the mask in __dict__, so old checkpoints keep working (a boolean back-fill default could lie for quantizers created with with_random_sign_mask=False). - Value semantics no longer leak into unregistered subclasses: register_value_opaque_quantizer stores the field tuple on the class and _value_fields looks it up in the class's own __dict__, so a subclass must register explicitly (it previously inherited value eq/hash that ignored its unannotated fields and skipped the annotation check). - The value-field tuple is computed once at registration instead of an MRO walk per __eq__/__hash__ call (these run per compiled-function invocation via the EQUALS_MATCH guard); registration validation and field derivation now share one annotation walk (Quantizer._annotated_fields). Drop the unreachable other._value_fields() branch and do the cheap type check first in __eq__. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Validate value-field annotations by resolved type, not annotation text Replace the substring blocklist with get_type_hints + an allowlist of value types (int/bool/float/str/enum): aliased tensor types no longer slip through and benign types whose name merely contains "Tensor" are no longer rejected. Runs once per class at import time, not in any hot path. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Drop the amax_reduction_group fixup from the generic rebuilder _rebuild_quantizer no longer back-fills the deprecated amax_reduction_group (and loses the field_names set that existed only for that check): a rebuilt quantizer deliberately lacks the attribute, so anything that genuinely needs it fails loudly instead of silently getting None. The only unconditional readers were the two copy() methods, which now tolerate the absent field; _canonicalized_amax_reduction_group (used by the kernel only when with_amax_reduction is set) still raises AttributeError on a rebuilt quantizer, which is the intended behavior. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Enable post-RHT amax in the NVFP4 value-object test factory The quantize kernel rejects with_rht=True without with_post_rht_amax=True (pre-RHT amax unsupported); mirror the recipe, which always sets both together. Unnoticed locally because the NVFP4 round-trip is skipped on non-NVFP4 hardware. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

* Remove cuDNN frontend git submodule Use nvidia-cudnn-frontend for the C++ headers and Python bindings. Keep the cuDNN library discovery helper in tree and update common, PyTorch, JAX, packaging, and test build paths. Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> * Fix JAX isolated build requirements Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com> --------- Signed-off-by: Vladimir Cherepanov <vcherepanov@nvidia.com>

shino16 · 2026-07-07T05:43:41Z

+            "(fsdp_group is not None); use FSDP2 or MCore FSDP."
+        )
+
+    weight = args.saved_weight if args.saved_weight is not None else args.weight_fp8


Note: Realistically args.saved_weight won't be None.

…behind cuBLAS 13.5+ (NVIDIA#3181) * Guard FP8 per-tensor scaling grouped GEMM on Hopper behind cuBLAS 13.5+ cuBLAS 13.4 has no SM90 grouped GEMM algorithms for the per-tensor (PER_BATCH_SCALAR_32F) FP8 scale mode, so the heuristic query fails with CUBLAS_STATUS_NOT_SUPPORTED and users get a cryptic "Unable to find suitable cuBLAS grouped GEMM algorithm" error (see NVIDIA#3176). The support is available starting with cuBLAS 13.5. - Add CUBLAS_FP8_TENSOR_SCALING_GROUPED_GEMM_HOPPER_VERSION (130500) and a runtime check with a clear error message on the FP8 tensor-scaling grouped GEMM path. - Enable the FP8Current grouped GEMM C++ tests on Hopper with cuBLAS 13.5+ (previously skipped as Blackwell-only, so this path had no C++ coverage on SM90). - Skip FP8 current scaling GroupedLinear tests on Hopper when cuBLAS < 13.5 instead of failing in algorithm selection. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * Fall back to legacy GroupedLinear path when fused grouped GEMM is unavailable Route around the fused GroupedTensor path instead of erroring out when the cuBLAS version is too old: grouped GEMM needs cuBLAS 13.3+ (13.4+ on Hopper), and FP8 per-tensor current scaling on Hopper needs 13.5+. Both the module (_is_grouped_tensor_path_supported) and the ops (_is_graph_safe_path_supported) checks now fall back to the legacy multi-stream path in these cases. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Squashed PR #8 (tensor_proto_mechanism) onto the rebased base. Adds TensorProto (pure-Python, torch.compile-traceable quantized-tensor allocation via Quantizer.alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__), Linear fake fwd/bwd impls for the custom-op path, and tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The cached FP8 weight is the same tensor returned as new_weight_workspace (cache miss) or passed in as weight_workspace (cache hit). A custom op may not return a tensor that aliases an input or another return, so mark those slots and reconstruct wt_save in _linear_setup_ctx instead of saving it twice. Mirrored in the fake impl so the saved-slot layout matches. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

NVFP4Quantizer._describe_buffers grouped each amax right after its scale (per-usage), diverging from NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (amax buffers last). The order is functionally irrelevant (buffers are consumed by name in alloc_tensors and reordered in TensorProto.inner_names), but aligning it makes describe/flatten agree and fixes test_to_tensor_proto_quantized[nvfp4]. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…upport - TensorProto.inner_names now raises if the quantizer describes buffer(s) absent from the storage's _FLATTEN_TENSOR_BUFFERS, instead of silently appending them. - Gate the nvfp4 proto-quantizer param on nvfp4_available so it skips on hardware without NVFP4 support rather than failing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

@staticmethods

…escribe_buffers Access NVFP4Quantizer @staticmethods (convert_shape_for_fp4, get_columnwise_shape) via the class instead of the instance. Under torch.compile, instance access of a @staticmethod on a value-opaque object crashes Dynamo guard generation with "'function' object has no attribute '__func__'" (pytorch/pytorch#182741). Temporary workaround until the PyTorch-side fix lands. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

pggPL force-pushed the make_qunatizers_opaque branch from 33e9d73 to d341eeb Compare June 16, 2026 15:21

pggPL requested a review from cyanguwa as a code owner June 16, 2026 15:21

pggPL force-pushed the tensor_proto_mechanism branch from 2cccc30 to 2e252f9 Compare June 16, 2026 15:31

pggPL force-pushed the tensor_proto_mechanism branch from 2e252f9 to ba92f5b Compare June 16, 2026 16:05

pggPL force-pushed the tensor_proto_mechanism branch from ba92f5b to b1273ea Compare June 16, 2026 16:12

kshitij12345 reviewed Jun 22, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/dynamo/tensor_proto.py Outdated

Comment thread tests/pytorch/test_torch_compile.py Outdated

pggPL force-pushed the make_qunatizers_opaque branch from e4a879b to adc65f6 Compare June 29, 2026 07:14

pggPL force-pushed the tensor_proto_mechanism branch from 85355a6 to c1e40b2 Compare June 29, 2026 07:16

pggPL force-pushed the make_qunatizers_opaque branch from adc65f6 to c7bbc83 Compare June 29, 2026 07:33

pggPL force-pushed the tensor_proto_mechanism branch from c1e40b2 to e760487 Compare June 29, 2026 07:34

pggPL force-pushed the make_qunatizers_opaque branch from c7bbc83 to f592cbb Compare June 29, 2026 09:26

pggPL force-pushed the tensor_proto_mechanism branch from e760487 to 50d5c21 Compare June 29, 2026 09:26

pggPL force-pushed the make_qunatizers_opaque branch from f592cbb to 945f62d Compare June 29, 2026 09:34

pggPL force-pushed the tensor_proto_mechanism branch from 50d5c21 to da709e7 Compare June 29, 2026 09:35

pggPL force-pushed the tensor_proto_mechanism branch 3 times, most recently from 5131ebc to 77831be Compare June 29, 2026 10:24

pggPL force-pushed the tensor_proto_mechanism branch from 77831be to 29e5245 Compare June 29, 2026 12:47

pggPL force-pushed the tensor_proto_mechanism branch from 29e5245 to 99c1377 Compare June 29, 2026 13:10

pggPL force-pushed the tensor_proto_mechanism branch from 99c1377 to afa86ff Compare June 29, 2026 13:29

pggPL force-pushed the tensor_proto_mechanism branch from afa86ff to 9e78a6c Compare June 29, 2026 13:30

pggPL force-pushed the tensor_proto_mechanism branch from 9e78a6c to 50c11cd Compare June 29, 2026 13:46

shino16 reviewed Jul 6, 2026

View reviewed changes

cyanguwa and others added 4 commits July 6, 2026 10:48

shino16 reviewed Jul 7, 2026

View reviewed changes

pggPL and others added 6 commits July 7, 2026 00:23

pggPL force-pushed the tensor_proto_mechanism branch from ff48e52 to e36cf6d Compare July 7, 2026 10:04

		import torch


		def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]:

Conversation

pggPL commented Jun 6, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

kshitij12345 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants