Carry symbolic shapes through compile custom ops - #26
Open
shino16 wants to merge 115 commits into
Open
Conversation
…ompile 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>
…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>
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>
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>
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>
…__/__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>
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>
…fp4 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>
…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>
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>
…assthrough 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>
…paque 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>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
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>
Squashed PR NVIDIA#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>
…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>
…stants; fix SP memory leak; test suite hook-up Wrap CommOverlapCore pybind11 methods that return compile-time constants so torch.compile(fullgraph=True) can trace through them without graph breaks: - `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py; `_ub_is_fp8()` in gemm.py - `with_cublasmp()` → `ub_is_cublasmp()` in base.py All callers in linear.py, layernorm_linear.py, layernorm_mlp.py, base.py, gemm.py, userbuffers_backward_linear.py and userbuffers_forward_linear.py updated. Fix quantized grad_output not being freed early for column-parallel SP backward. Row-parallel SP already called clear_tensor_data(grad_output) to release the gathered tensor; column-parallel SP quantizes grad_output to Float8TensorStorage but never freed it before returning. Under torch.compile reduce-overhead this leaves 3 live pool tensors at recording end and triggers "Detected 3 tensor(s) in the cudagraph pool not tracked as outputs". Extend the existing clear_tensor_data guard to cover both parallel modes. Fix custom-recipe quantizer state being re-initialised on every forward call even when the recipe object has not changed. The existing early-exit for CustomRecipeState was missing an identity check on the recipe object, so any repeated call with the same recipe would bypass the early-return and rebuild quantizers unnecessarily. Add `if recipe_state.recipe is recipe: return` to restore the intended caching behaviour. Add test_torch_compile.py to L0_pytorch_unittest so the autocast and existing compile tests run in CI. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit bfce3a7)
for more information, see https://pre-commit.ci (cherry picked from commit afe364b)
ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit the no-roles warning, which graph-breaks under fullgraph=True. qfactory dispatches on role.tensor_type instead of a pre-baked string key. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit 22f80e4)
…tom_op Replace the low-level torch.library.Library (_TE_LIB.define/.impl + functional register_fake/register_autograd/register_torch_dispatch with lib=) with the standard torch.library.custom_op API, passing the dynamically built schema explicitly via schema=. register_fake/register_autograd/register_torch_dispatch are now methods on the returned CustomOpDef. Drops the TOR901 Library usage and is robust to re-registration (get_library_allowing_overwrite). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> (cherry picked from commit d590560)
…e sentinel Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Squashed PR NVIDIA#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>
…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>
test_python_alloc_matches_cpp_make_empty compared buffers the quantize kernel never writes: the scale-inv padding is allocated uninitialized by both paths, so the bit-exact comparison saw random bytes and failed on H100/B200 for fp8_blockwise. Zero every buffer before quantizing, so the comparison covers kernel output only. Also drop the param-level skips on the nvfp4 entries of _PROTO_QUANTIZERS and _VALUE_QUANTIZERS. is_fp8_available() and friends run at import time and go through torch.cuda.current_device(), so this module cannot be collected without CUDA at all and skipif(not torch.cuda.is_available()) never fires; the same goes for the torch.cuda.is_available() halves of the _hw_available() guards. Gating nvfp4 on nvfp4_available was also inconsistent with MXFP8 and blockwise, which are gated at runtime and only in the tests that run a kernel -- the allocation primitives themselves are pure Python and describe the layout on any HW. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
_linear_forward_impl_fake diverged from quantize_weight on the weight workspace in three ways: - it produced a new workspace only when update_ws was true, but the real cache-miss path returns (out, out) whenever cache=True, regardless of update_workspace; a first call with is_first_microbatch=False therefore lost the workspace and the "new_workspace" saved-weight alias; - it treated any non-None cached workspace as a hit, while the real path runs _is_weight_workspace_valid() first and falls through to a miss when the cached buffer layout no longer matches the quantizer's usage; - it kept quantizer.internal, so the descriptor resolved to a bare storage class, while the real path quantizes persistent workspaces with internal=False and caches wrapper tensors. On a cache hit the weightmat is now the workspace descriptor itself, and on a miss with cache_weight it is the same proto object returned as the new workspace, matching quantize_weight's aliasing. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Eager forward forces save_original_input=False for backward_override="dequantized", but the fake only handled "high_precision". With save_original_input=True and that override, the fake aliased the original input into saved-tensor slot 0 while eager saved a quantized input with rowwise-only usage, so the saved payload layout and the compiled backward setup disagreed. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The output proto's requires_grad considered only the input and the weight, so a frozen input and weight with a trainable bias described the output as non-differentiable while eager _Linear.apply produces a differentiable one. bias_requires_grad is already False when there is no bias. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
_linear_forward_impl_fake / _linear_backward_impl_fake, and the eager-side changes that existed only to support them (reading the requires_grad flags off LinearFwdArgs, the new_workspace/weight_workspace alias dedup and the _linear_setup_ctx signature carrying (out, new_weight_workspace)), have no caller in this PR: nothing registers them as a custom op's fake, so nothing exercises them here. They belong with the custom-op registration that consumes them. This PR is left as the TensorProto mechanism proper -- the proto, the storage flatten protocol and the pure-Python quantizer allocation hooks -- which the new tests do cover. linear.py returns to its upstream state. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Carried over from the TensorProto PR (NVIDIA#3153), where these impls used to live without a caller; they belong here, with the custom-op registration that consumes them. - Weight workspace: quantize_weight returns a fresh workspace on every cache miss with cache=True, not only when update_workspace is set; it discards a cached workspace that fails _is_weight_workspace_valid; and it quantizes persistent workspaces with internal=False so the cache holds wrapper tensors. The fake did none of the three. - backward_override="dequantized" forces save_original_input=False in the eager forward; the fake only handled "high_precision", so it aliased the original input where eager saves a rowwise-only quantized one. - The output's requires_grad ignored the bias, describing the output of a bias-only-trainable Linear as non-differentiable. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Quantized tensors now implement the wrapper-subclass flatten protocol, so nn.Module._apply moves them with torch.utils.swap_tensors instead of the `param.data = ...` path. The swap exchanges the parameter's entire __dict__: that is how the inner buffers reach the surviving object, but it also carries off everything attached to the parameter from the outside. TE relies on several such attributes: _high_precision_init_val and its two accessors (quantized_model_init(preserve_high_precision_init_val=True)), plus main_grad, grad_added_to_main_grad and overwrite_main_grad, which Megatron-Core attaches. They survived before only because `param.data = ...` is a no-op for a wrapper subclass -- the outer tensor is a zero-storage shell and the assignment never touched __dict__, so device moves silently did nothing at all. Snapshot the parameters' __dict__ before delegating to nn.Module._apply and restore the entries the swap dropped, rebinding bound accessors to the surviving parameter. Entries still present afterwards are the tensor's own state, where the post-swap value is the correct one. Covers the two test_sanity grouped-linear high-precision-init tests that broke on B200, and adds a direct test over .cuda() / .cpu() / .half(). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…en context The __tensor_flatten__ context stored the storage / wrapper class as a qualname string, which __tensor_unflatten__ then resolved through a global registry populated by __init_subclass__. The indirection only existed because OpaqueValueBundle could not carry a class object across the custom-op boundary. Teach the bundle to handle classes exactly as it already handles Enum (render by name, add the class to the FX graph globals) and store type(self) in the context directly. Quantizer._storage_metadata already returned the class, so create_metadata no longer has to stringify it. _all_quantized_tensor_subclasses now walks QuantizedTensor.__subclasses__() instead of filtering the registry. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…ps its quantizer to_tensor_proto reads the quantizer from _quantizer, an attribute of the real TE tensors; a TensorProto exposes it as quantizer. Re-describing a proto therefore produced a proto with quantizer=None. _proto_view converts weight_workspace to a TensorProto before the fake impl runs, and the fake Linear forward re-describes it on the FP8 weight-cache hit, so update_usage(rowwise_usage=True) hit a proto it believed to be unquantized: ValueError: update_usage called on a non-quantized TensorProto which dynamo surfaced as an Unsupported observed exception. This only fired with is_first_microbatch, the sole path feeding a cached workspace into the op, and made test_te_linear_compile_is_first_microbatch fail for both compile modes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… fake forward 6343317 made to_tensor_proto idempotent to stop the FP8 weight-cache hit from losing its quantizer. That papered over the actual defect: _proto_view has already turned weight_workspace into a TensorProto by the time the fake forward runs, so calling to_tensor_proto on it was a no-op that only ever had the chance to be wrong. Copy the proto with dataclasses.replace instead -- keeping the fresh-quantizer-copy semantics update_usage relies on -- and drop the idempotency branch, which now has no caller (the other call site is guarded by an isinstance check). _is_weight_workspace_valid goes with it: it dispatches on storage types (Float8TensorStorage, ...), so on a proto every branch fell through and it returned True unconditionally. The fake path therefore never detected a stale workspace; that gap is now stated in a comment rather than hidden behind a call that looked like it worked. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
The restore loop keys off "present after the swap": what survived is the tensor's own state, what did not is an externally attached annotation. That holds only as long as every declared buffer really is present afterwards. If one were not, the loop would quietly put the pre-move value back and splice a buffer from the old device (or from before a dtype conversion) into the moved parameter -- silently wrong numerics rather than a crash. Raise instead when a name from _FLATTEN_TENSOR_BUFFERS is about to be restored. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
… eager Registration of the opaque quantizer types, of OpaqueValueBundle and of the custom ops is best-effort so that TE stays importable on PyTorch builds without the opaque-object APIs. Failing silently makes the fallback resurface much later as an obscure error inside the compiled region, so warn once per process instead and point at the PyTorch version. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Keeps the Linear fake impls and the value-quantizer test parametrization on this branch; picks up the alloc-parity test from tensor_proto_mechanism. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
nn.Module._apply only assigns to self._parameters, never removes entries, so a missing parameter after it returns means something unexpected happened. Skipping it silently dropped every attribute attached to that parameter -- the failure this override exists to prevent. Match the buffer check and fail loudly. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Reverts 6e61d36. The check guarded a case that cannot arise today: the storages always set every declared buffer attribute, to None when unused, so the key is present whatever the usage flags say. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Each storage class listed its tensor buffers twice: once as a field
annotation, once as an (attribute, constructor kwarg) pair in
_FLATTEN_TENSOR_BUFFERS, in a different order and further down the file.
Adding a buffer meant remembering both.
Mark the field instead -- _scale_inv: Annotated[torch.Tensor,
Buffer("fp8_scale_inv")] -- and collect the declarations in
__init_subclass__, which already runs there for the storage registry.
_FLATTEN_TENSOR_BUFFERS survives as the derived attribute, so every consumer
is untouched, and the collected values are identical to the hand-written
tuples for all nine registered classes.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
"Buffer" collides with nn.Module's buffers, which are a different thing, and _FLATTEN_TENSOR_BUFFERS named a consumer (__tensor_flatten__) rather than the thing itself -- the list has four of them. PyTorch calls exactly this concept "inner tensors", which TensorProto.inner_names() already follows. Also drop the underscore from the two hooks every quantizer has to implement. They were the only members of the extension contract marked private, which is why the tests needed seven protected-access waivers to call them; the members nobody overrides (alloc_tensors, create_metadata) were public already. Buffer -> InnerTensor _FLATTEN_TENSOR_BUFFERS -> _INNER_TENSORS _describe_buffers -> inner_tensor_specs _storage_metadata -> storage_metadata Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
__tensor_flatten__ put the class qualname in the context and __tensor_unflatten__ looked it up in a module-level registry, populated from __init_subclass__. The indirection bought nothing: dynamo bakes the class object into the graph as a constant just as happily, which is what the custom-op branch already relies on. Store type(self) directly and drop _STORAGE_REGISTRY. __init_subclass__ stays for collecting the InnerTensor field annotations. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
…compile_rebased # Conflicts: # tests/pytorch/test_torch_compile.py # transformer_engine/pytorch/quantized_tensor.py # transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
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. Lost in a 3-way merge that took the HEAD side. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-Authored-By: Claude Opus 5 (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. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
_linear_forward_impl decided columnwise usage from inp.requires_grad, while
its fake twin used args.input_requires_grad -- a value read outside the graph
and baked into it as a constant. Under mode="reduce-overhead" cudagraph_trees
swaps the inputs for static placeholders while recording, and those carry
requires_grad=False, so the real impl quantized the weight without columnwise
data and returned one buffer fewer than the fake had promised. Every later
slot shifted by one and inductor's assert_size_stride landed on the 1-D
scale_inv where it expected the 2-D transpose:
AssertionError: wrong number of dimensions1 for op:
torch.ops.transformer_engine_compile.linear.default
Read both flags from LinearFwdArgs so the two impls decide from one source.
backward_needs_input is the same trap and changes with it, though it happens
not to fire today: the weight survives the swap as a static parameter.
Fixes test_te_linear_compiles[*-reduce-overhead] for every recipe whose
inner-tensor count depends on columnwise usage -- Float8CurrentScaling on
TN-capable archs, MXFP8/NVFP4/Float8BlockScaling on Blackwell.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
pggPL
changed the base branch from
linear_compile
to
linear_torch_compile_final_attempt
August 5, 2026 16:08
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
2 participants
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
Please include a brief summary of the changes, relevant motivation and context.
Fixes # (issue)
Type of change
Changes
Please list the changes introduced in this PR:
Checklist: