[pyTorch] Replace the make_empty implementation to use C++ implementation

[ptrendx]

Description

This PR unifies the implementation of the QuantizedTensor creation by using the C++ implementation of the create_tensor.

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

Please list the changes introduced in this PR:

• Replaced the Python implementations of the make_empty with the calls to C++ create_tensor

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

[ptrendx]

/te-ci L1 pytorch

[ptrendx]

/te-ci L1 pytorch

[ptrendx]

/te-ci pytorch L1

[greptile-apps]

Greptile Summary

This PR removes per-class Python make_empty overrides from all built-in quantizer types (Float8Quantizer, Float8CurrentScalingQuantizer, Float8BlockQuantizer, MXFP8Quantizer, NVFP4Quantizer) and replaces them with a single C++-backed implementation in the Quantizer base class. A new create_empty_quantized_tensor C++ function is exposed via pybind11 and a new resolve_device helper centralises device inference across all create_tensor overloads.

• Adds device and pin_memory parameters (defaulting to nullopt/false) to every create_tensor virtual declaration and implementation, maintaining backward compatibility at existing call sites.
• Introduces resolve_device which correctly infers the target device from an explicit argument, a pre-provided data tensor, or the current CUDA device — in that priority order.
• Python make_empty now delegates to tex.create_empty_quantized_tensor; a custom attribute guard is added to preserve a friendly error path for custom quantizers, though the attribute is undocumented.

Confidence Score: 5/5

The refactor is mechanically sound — all built-in quantizer paths are correctly unified through C++ with no correctness regressions.

All allocation sites are consistently updated; the new resolve_device helper properly handles all four device-inference cases. The only rough edge is an undocumented custom attribute that affects error message quality for niche custom-quantizer authors, not correctness of the built-in paths.

transformer_engine/pytorch/quantized_tensor.py — the undocumented custom attribute guard for custom quantizers.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/quantized_tensor.py	Base-class make_empty now routes through C++; the custom attribute guard is undocumented and risks replacing a clear Python error with an opaque C++ assertion for custom quantizer authors.
transformer_engine/pytorch/csrc/quantizer.cpp	Adds resolve_device helper and threads device/pin_memory through all create_tensor overloads; device inference from a pre-provided data tensor is correctly handled via the new helper.
transformer_engine/pytorch/csrc/extensions/cast.cpp	New create_empty_quantized_tensor function added cleanly before the anonymous namespace; delegates directly to the C++ quantizer's create_tensor with device and pin_memory.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Registers create_empty_quantized_tensor with pybind11 with correct argument names and types; no issues.
transformer_engine/pytorch/csrc/common.h	All create_tensor virtual declarations updated with device and pin_memory defaulting to nullopt/false; backward-compatible at call sites that don't supply them.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Removes Python make_empty from NVFP4Quantizer; old impl used torch.zeros for amax_* while C++ uses at::empty — deliberate per inline C++ comments that kernels zero those buffers.
transformer_engine/pytorch/tensor/mxfp8_tensor.py	Removes Python make_empty from MXFP8Quantizer; shape validation is preserved in C++ via NVTE_CHECK.
transformer_engine/pytorch/tensor/float8_blockwise_tensor.py	Removes Python make_empty from Float8BlockQuantizer; C++ now passes device kwarg which correctly flows through **kwargs to QuantizedTensor.__new__.
transformer_engine/pytorch/tensor/float8_tensor.py	Removes Python make_empty from Float8Quantizer and Float8CurrentScalingQuantizer; both now covered by the C++ path.

Sequence Diagram

sequenceDiagram
    participant PY as Python (Quantizer.make_empty)
    participant TEX as tex.create_empty_quantized_tensor (C++)
    participant CONV as convert_quantizer
    participant QC as QuantizerCpp::create_tensor
    participant RD as resolve_device
    participant PT as PyTorch (at::empty)
    participant PYTENSOR as Python Tensor __new__

    PY->>TEX: (self, shape, dtype, device, pin_memory)
    TEX->>CONV: convert Python quantizer to C++ Quantizer
    TEX->>QC: create_tensor(shape, te_dtype, device, pin_memory)
    QC->>RD: resolve_device(device_opt, data_opt)
    RD-->>QC: concrete at::Device
    QC->>PT: at::empty(..., opts with device + pin_memory)
    PT-->>QC: allocated at::Tensor(s)
    QC->>PYTENSOR: PyObject_Call(TensorClass, kwargs incl. device)
    PYTENSOR-->>QC: Python QuantizedTensor
    QC-->>TEX: (TensorWrapper, py::object)
    TEX-->>PY: QuantizedTensor

Loading

_{Reviews (13): Last reviewed commit: "Merge branch 'main' into pr_unify_make_e..." | Re-trigger Greptile}

[greptile-apps]

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[vthumbe1503]

Doing this in C++ itslef might be faster, since we are anyway going to call the QuantizedTensor.new method with requires_grad argument. Calling this from python for custom quantized tensor has severe python overheads

[vthumbe1503]

But I see it can get quite complicated since we might have to change the create_tensor API to accept the requires_grad argument.

[ptrendx]

/te-ci pytorch

[ptrendx]

/te-ci pytorch

[greptile-apps]

Additional Comments (1)

transformer_engine/pytorch/csrc/quantizer.cpp, line 596
device parameter shadowed by pre-existing local variable

The new device function parameter (line 562) is shadowed by the local variable at::Device device declared here. In C++, this is a scoping issue: the local declaration shadows the parameter, so kwargs["device"] on line 632 uses the local variable instead of the caller's argument.

This creates two problems:

1. If the build uses -Wshadow -Werror, this will be a compile error.
2. More critically, in the edge case where both with_data == false and with_transpose == false, the device will resolve to c10::cuda::current_device() rather than the device parameter passed by the caller. This causes the Python Float8Tensor object to report the wrong device.

Fix: Remove the shadowing local variable and use the parameter directly:

  // Construct Python FP8 tensor
  py::object out_py;
  py::object scale_inv_py = py::cast(scale_inv_tensor);
  py::object data_py = with_data ? py::cast(data_tensor) : py::none();
  py::object transpose_py = with_transpose ? py::cast(transpose_tensor) : py::none();
  if (internal) {
    // ...
    kwargs["quantizer"] = this->quantizer;

    py::tuple args(0);
    // ...
  } else {
    // ...
    kwargs["quantizer"] = this->quantizer;
    kwargs["device"] = py::cast(device);

The local at::Device device declaration (lines 593-596) should be removed entirely, as the device parameter already holds the requested device.

[ptrendx]

/te-ci pytorch

[ptrendx]

/te-ci pytorch

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[ptrendx]

/te-ci pytorch

[ptrendx]

/te-ci pytorch

[vthumbe1503]

LGTM

[vthumbe1503]

It might make sense for the tex API to accept requires_grad as an argument as well considering CPU overheads. Given that we are passing things like pin_memory, device which are torch::TensorOptions, we might as well add requires_grad in the picture which is also an attribute of torch::TensorOptions.

[ptrendx]

/te-ci pytorch