[siddhant-rajhans]

Description

Ops on a CPU-resident QuantizedTensor that have no quantized implementation (for example x_cpu + 0, or torch.testing.assert_close(x_cpu, ref)) fail with RecursionError for MXFP8 and NVFP4 tensors.

QuantizedTensor.__torch_dispatch__ dequantizes the arguments of such ops. MXFP8 and NVFP4 dequantize with tex.dequantize, which needs GPU buffers, so for a CPU tensor they first call tensor.to(device=...). Inside __torch_dispatch__ the autograd keys are excluded, so that call reaches the dispatcher as aten.to.dtype_layout instead of being decomposed into aten._to_copy. Only aten._to_copy was handled, so the call fell back into the dequantizing path and recursed. FP8 and FP8 block scaling are not affected because they dequantize the CPU buffers directly.

The fix follows #3279 by @janekb04.

Fixes #3294

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

• QuantizedTensor.__torch_dispatch__ handles the aten.to.* overloads the same way as aten._to_copy: the subclass is preserved and the internal buffers move to the target device.
• test_cpu_torch_ops in tests/pytorch/test_quantized_tensor.py, parametrized over all quantization formats: arithmetic and assert_close on a CPU tensor, .float(), and a round trip back to CUDA.

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

I don't have a Blackwell GPU, so I could not run the MXFP8 and NVFP4 cases locally. I reproduced the recursion with a minimal wrapper subclass on CPU (handling only aten._to_copy recurses on x + 0; handling the aten.to.* overloads as well gives the expected results), and the changed files pass black and pylint with the repo's configuration.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new actionable issue was found.

Summary

Prevents recursive dequantization when CPU-resident quantized tensors reach plain PyTorch operations.

• CPU quantized tensors now handle plain PyTorch operations without recursive dispatch.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[PyTorch operation on CPU quantized tensor] --> B[Dequantize arguments]
  B --> C[MXFP8 or NVFP4 moves buffers to CUDA]
  C --> D[Handle aten.to overload directly]
  D --> E[Rebuild quantized tensor with CUDA buffers]
  E --> F[Dequantize on CUDA]
  F --> G[Return dense result to CPU]

Loading

_{Reviews (6) · Last reviewed commit: "Merge branch 'main' into fix/cpu-quantiz..." · Reviewed by Greptile}

[vthumbe1503]

Thanks for the PR. There is a previously existing PR #3279 from @janekb04 to fix this problem. I think that PR also handles copy implementation for generic quantized_tensor. We also want to keep the test in test_quantized_tensor.py rather than introduce a new test file.

Could you please use that as reference and modify this PR? Although, I dont see fix/implementation for Float8Tensor and Float8BlockWiseTensor there. My guess is that base implementation in QuantizedTensor already handles the copy method for them. And this makes me wonder whether we even need the standalone implementations for MXFP8/NVFP4. And base QuantizedTensor's implementation should be able to handle it by itself for all QuantizedTensor.

[siddhant-rajhans]

Thanks for pointing me to #3279. I reworked the PR along those lines in f9380ee:

• The fix is now in the base class: QuantizedTensor.__torch_dispatch__ handles the aten.to.* overloads the same way as aten._to_copy, following [Pytorch] Fix RecursionError when dequantizing cpu-resident QuantizedTensor #3279.
• The test moved into test_quantized_tensor.py as test_cpu_torch_ops, parametrized over all quantization formats, and the separate test file is gone.

On your question: as you suspected, the MXFP8/NVFP4-specific changes aren't needed with the base-class fix, so I dropped them. The recursion comes from the base class. When the dequantizing fallback in __torch_dispatch__ moves a CPU tensor to the GPU with .to(), that call arrives as aten.to.dtype_layout (the decomposition into aten._to_copy is skipped inside dispatch), misses the _to_copy branch, and goes back into the fallback. FP8 and FP8 block scaling never reach that path: _FromFloat8Func has a CPU branch and the blockwise dequantization is plain tensor ops, while MXFP8 and NVFP4 need tex.dequantize on the GPU.

I don't have a Blackwell GPU, so I couldn't run the MXFP8/NVFP4 cases here (details in the description). Could you trigger CI when you get a chance?

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

LGTM