[PyTorch][torch.compile] Make quantizers opaque value objects by pggPL · Pull Request #7 · pggPL/TransformerEngine

pggPL · 2026-06-06T12:14:03Z

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 CustomRecipeinstance 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

kshitij12345

LGTM

…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>

…m-mem zero-copy (NVIDIA#3035) * Expert Parallelism: PyTorch wrapper + autograd ops with symm-mem zero-copy Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…NS_PER_RANK (NVIDIA#3150) * nccl with relax num_dispatch_tokens%64!=0 Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Skip EP tests/examples on nodes without NVLink Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

…VIDIA#3141) * Preserve fprop operands for dequantized backward override Signed-off-by: Evgeny <etsykunov@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add test_grouped_linear_backward_override_high_precision_forces_save_original_input test Signed-off-by: root <root@prenyx0017.a51.clusters.nvidia.com> --------- Signed-off-by: Evgeny <etsykunov@nvidia.com> Signed-off-by: root <root@prenyx0017.a51.clusters.nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root <root@prenyx0017.a51.clusters.nvidia.com>

* Make quantized-tensor __repr__ fake-safe under torch.compile Under torch.compile, TE quantized-tensor __repr__ methods are invoked on FakeTensors during AOT autograd's structured logging. The repr bodies call self._scale_inv.item() and/or self.dequantize() (which dispatches to the raw C++ op tex.dequantize), both of which access a FakeTensor's data pointer and raise: RuntimeError: Cannot access data pointer of Tensor (e.g. FakeTensor, FunctionalTensor) ... This was the sole cause of six fp8 failures in tests/pytorch/test_torch_compile.py. Fix: add one shared helper, safe_quantized_repr, in tensor/_quantization_helpers.py (a safe leaf module importing only torch) that builds a metadata-only repr string. Each data-touching __repr__ now wraps its existing body in a try/except and falls back to the helper when the data cannot be materialized. The eager (non-fake) repr output is unchanged; only a fallback path is added. Wrapped reprs: Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor and their *Storage counterparts. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make quantized __repr__ fallback universal, drop FakeTensor-specific logic Remove the FakeTensor-specific heuristic (_is_fake_data_access_error) and the warning path from safe_quantized_repr. The fallback is now a plain metadata-only repr triggered by any exception while materializing data, with each attribute access individually guarded so __repr__ never raises. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> --------- Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

…` with `total_recv_tokens_per_rank` placeholder (NVIDIA#3154) * versioning EP C configs Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Rename EP prepare token_counts to recv_tokens_per_expert Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Add total_recv_tokens_per_rank placeholder to nvte_ep_prepare Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Adapt PyTorch EP binding to versioned nvte_ep C config API Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> * Rename EP group config max_num_sms to num_comm_sms Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com> --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

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>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

_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>

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

# Conflicts: # transformer_engine/pytorch/tensor/float8_blockwise_tensor.py # transformer_engine/pytorch/tensor/float8_tensor.py # transformer_engine/pytorch/tensor/mxfp8_tensor.py # transformer_engine/pytorch/tensor/nvfp4_tensor.py

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>

… CUBLAS GGEMM heuristics (NVIDIA#3143) * support in grouped linear and relevant tests 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 * Unecessary details remove Removed details about FP8 current scaling methods. Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * fix grouped linear module's grouped tensor path Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * allow more current scaling use-cases.. block nvfp4+rht+single grouped weight being cuda graphable 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 * some minor comment fixing Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix heuristics Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * only dealyed scaling skip in failure comment Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> * address review comment 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 * fix for other 2 nvte APIs Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * fix m and n Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

@timmoon10

…n megatron. (NVIDIA#2898) * add support for THD CUDA graph Signed-off-by: HaochenYuan <haocheny@nvidia.com> * modify comment Signed-off-by: HaochenYuan <haocheny@nvidia.com> * address @timmoon10: drop FAv2-bwd alloc gate, rely on THD tail zero-fill Signed-off-by: HaochenYuan <haocheny@nvidia.com> * Support graph-safe MoE aux loss token count Signed-off-by: HaochenYuan <haocheny@nvidia.com> * add graph guard for one zero fill Signed-off-by: HaochenYuan <haocheny@nvidia.com> * remove redundant zero-fill Signed-off-by: HaochenYuan <haocheny@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename & remove prelude kernel Signed-off-by: HaochenYuan <haocheny@nvidia.com> * Update warp reduction function Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@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: HaochenYuan <haocheny@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

…ling (NVIDIA#3135) Implements grouped-tensor quantize for the FP8 1D (1x128) and 2D (128x128) block-scaling recipes in row-wise (RW), column-wise (CW) and BOTH quantization directions. A single CUDA kernel launch walks 128x128 tiles across every tensor in the group, with each CTA decoding its owning tensor from the device-side GroupedTensor metadata with (N, R, K) shapes. Supports SAME_BOTH_DIMS (all tensors identical) and VARYING_FIRST_DIM (constant K, varying R) shape representations. Three kernels share the dispatcher in group_quantize_blockwise_{1d,2d}: - group_block_scaled_1d_rw_kernel: RW-only dispatch; 8 threads/row, reads global memory directly into vec-16 registers; bypasses TMA since the shared-memory roundtrip and ptx::mbarrier do not buy anything without re-use in the CW path. - group_block_scaled_1d_tma_kernel: CW-only and BOTH dispatch. TMA bulk-load fills shared memory input cache. BOTH runs an RW pass (8 threads/row, vec-16 read from shared memory) then a CW pass; CW-only skips the RW pass. The CW pass uses 4 t/col with 32-row reg_data and two column passes in the BOTH instantiation (keeps the per-thread register footprint under the sm_90 3-CTAs/SM threshold) and 2 t/col with 64-row reg_data in the CW-only instantiation (avoids doubling the smem-load bank-conflict footprint that 4 t/col would introduce). - group_block_scaled_2d_tma_kernel: RW-only, CW-only and BOTH dispatch. TMA bulk-load fills shared memory input cache. Pass 1 stages 8 IVecs/thread in registers while computing the per-tile scalar amax. Pass 2 quantizes from registers, emits row-wise output, stages column-wise output to the shared memory transpose staging buffer, then drains smem_T to global memory. Per-expert scale offsets: - 1D RW: closed-form O(1) for both SAME_BOTH_DIMS and VARYING_FIRST_DIM (each M_i is a multiple of kTileDim=128, hence of kScaleColAlign=4, so DIVUP_TO_MULTIPLE collapses and the prefix sum reduces to a single tensor_offsets_ptr[tensor_id]/K load). - 2D CW: closed-form O(1) for SAME_BOTH_DIMS; CTA-cooperative warp-shuffle prefix sum for VARYING_FIRST_DIM (non-linear DIVUP_TO_MULTIPLE on blocks_y_t prevents a closed form). The cooperative reduction uses the existing warp_allreduce_sum helper from common/utils.cuh. Dequantize and bias-gradient (bgrad): - group_dequantize_fp8_blockwise.cuh: kernels for all four modes (1D/2D x rowwise/columnwise), inverting the per-expert layouts the quantize kernels write. - bgrad_group_quantize accepts Float8Block quantizers and computes dbias per-tile column-partial in-kernel (mirroring MXFP8); reduced per expert via the existing common::grouped_reduce_dbias. Scale constraints: the fused grouped FP8BS path supports only unconstrained FP32 scales (Float8BlockQuantizer::create_grouped_tensor rejects force_pow_2_scales=True). Power-of-2 scales remain available on the non-grouped/unfused split-quantize path used for Blackwell MXFP8 emulation. Tests: existing parametrized grouped quantize / dequantize / bgrad tests in test_grouped_tensor.py cover MXFP8, NVFP4, FP8 current scaling and the newly-added FP8 block scaling recipe. tests/cpp/operator/ test_cast_float8blockwise_grouped.cu adds 72 C++ unit-test cases over uniform/jagged shapes, all four (BD x direction) modes, K in {128, 256, 512}, and CUDA-graph capture coverage. Kernels are gated to Hopper (sm_90) at the host dispatcher (cuBlasLt grouped GEMM supports FP8 block-scaling only on Hopper). JAX integration is intentionally left out of scope and deferred to a follow-up PR. ResolvesNVIDIA#2525 Signed-off-by: Alp Dener <adener@nvidia.com>

…L2 Jax dist (NVIDIA#3159) * Keep the routing map format alive Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> * Fix incorrectly launched multi process EP tests in L2 Jax instead of L2 jax dist Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com> --------- Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com>

* skip tests on hopper Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> * Update qa/L1_pytorch_mcore_fsdp_integration/test.sh Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> --------- Signed-off-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: vthumbe1503 <vthumbe@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

…ests (NVIDIA#3174) fixing Blackwell skip condition for grouped FP8 block-scaling tests in C++ Signed-off-by: Alp Dener <adener@nvidia.com>

- 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>

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>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

- 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>

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>

_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>

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>

kshitij12345 approved these changes Jun 9, 2026
View reviewed changes

Comment threadtests/pytorch/test_torch_compile.py
Comment threadtransformer_engine/pytorch/dynamo/quantizer_opaque.py Outdated

pggPL force-pushed the remove_process_group_from_quantizers branch from e9097d6 to 948cd6dCompare June 16, 2026 12:23

pggPL requested a review from cyanguwa as a code owner June 16, 2026 12:23

pggPL force-pushed the remove_process_group_from_quantizers branch from b8c1bec to 6c9b986Compare June 16, 2026 14:56

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

pggPL force-pushed the make_qunatizers_opaque branch 2 times, most recently from adc65f6 to c7bbc83Compare June 29, 2026 07:33

pggPLand others added 7 commits June 29, 2026 11:25

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

pggPL changed the base branch from remove_process_group_from_quantizers to mainJune 29, 2026 09:26

pggPL closed this Jun 29, 2026

pggPL reopened this Jun 29, 2026

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

pggPLand others added 10 commits June 29, 2026 12:05

Reword opaque-flag comment: self-contained, no Linear reference
2c3c5df
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

phu0ngngand others added 26 commits June 30, 2026 19:45

Drop verbose comments around value-opaque flag stamping
9db604f
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Shorten amax_reduction_group check comment
fe5e5db
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

docs: document attention backend selection (NVIDIA#3142)
f721112
Signed-off-by: Santosh Bhavani <santosh.bhavani@live.com>

[Common] Blackwell skip condition for C++ grouped FP8 block-scaling t…
7cb8b31
…ests (NVIDIA#3174) fixing Blackwell skip condition for grouped FP8 block-scaling tests in C++ Signed-off-by: Alp Dener <adener@nvidia.com>

Drop redundant annotation-exclusion comments
dd96956
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Shorten _is_value_quantizer comment
447a4e1
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Fold annotation walk into _value_fields, trim comments
02a8fc9
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Simplify qualname-registry comment
5f55a51
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Reword qualname-registry comment for outside readers
02afbdc
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Merge remote-tracking branch 'upstream/main' into make_qunatizers_opaque
f4ddb51

Navigation Menu

[PyTorch][torch.compile] Make quantizers opaque value objects - #7

[PyTorch][torch.compile] Make quantizers opaque value objects#7
pggPL wants to merge 44 commits into
mainfrom
make_qunatizers_opaque

pggPL commented Jun 6, 2026•
edited
Loading

Uh oh!

kshitij12345 left a comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

Conversation

pggPL commented Jun 6, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

9 participants

pggPL commented Jun 6, 2026•
edited
Loading