# Feature Request: Support Async-Stream Transfer for [AffineQuantizedTensor] (Fix diffusers group_offload device mismatch)

[lllllinux]

🐛 The Problem

When using diffusers Group Offloading feature (enable_group_offload) with use_stream=True alongside torchao quantization (via TorchAoConfig, which uses AffineQuantizedTensor), a RuntimeError: Expected all tensors to be on the same device error occurs. The quantized weight remains on the CPU while the input tensor is on CUDA.

Context:

• I raised the issue in diffusers: huggingface/diffusers#13281.
• diffusers maintainers merged a fix in huggingface/diffusers#13305.

However, the diffusers fix is a fallback. It detects tensor subclasses and forces them to use synchronous (blocking) transfer, skipping pin_memory() and record_stream. While this prevents the crash, it means torchao quantized models cannot benefit from the async-stream group offloading performance gains.

🔍 Root Cause Analysis

The issue stems from missing support for PyTorch's async transfer primitives in AffineQuantizedTensor (and its underlying implementation Float8AQTTensorImpl / PlainAQTTensorImpl):

1. pin_memory() is not implemented: AffineQuantizedTensor relies on the torch_dispatch fallback, which either raises NotImplementedError or silently returns a plain float CPU tensor (stripping the quantization subclass metadata). diffusers needs to pin the tensors in CPU memory for fast async DMA transfers.
2. non_blocking=True is stripped: In torchao/utils.py, the _get_to_kwargs method explicitly ignores the non_blocking argument. Consequently, when diffusers calls .to(device, non_blocking=True), torchao executes a synchronous .to(device).
3. is_pinned() is not implemented: Similar to pin_memory().

Note: the Float8Tensor implementation in workflows/float8/float8_tensor.py seems already implements aten.is_pinned and aten._pin_memory.

Suggested Direction

I’m not fully sure what the intended design direction for torchao is here, so why don't we just support pin_memory and non_blocking.

1. Supporting pin_memory / is_pinned

In torchao/dtypes/affine_quantized_tensor_ops.py:

@implements(aten.is_pinned.default)
def _(func, types, args, kwargs):
    return args[0].tensor_impl.is_pinned()

@implements(aten._pin_memory.default)
def _(func, types, args, kwargs):
    self = args[0]
    return self.__class__(
        self.tensor_impl.pin_memory(),
        self.block_size,
        self.shape,
        self.quant_min,
        self.quant_max,
        self.zero_point_domain,
        dtype=self.dtype,
    )

In layout implementations (e.g., torchao/dtypes/floatx/float8_layout.py):

@implements([aten.is_pinned.default])
def _(func, types, args, kwargs):
    return args[0].float8_data.is_pinned() and args[0].scale.is_pinned()

@implements([aten._pin_memory.default])
def _(func, types, args, kwargs):
    return Float8AQTTensorImpl(
        args[0].float8_data.pin_memory(),
        args[0].scale.pin_memory(),
        args[0].transposed,
        args[0]._layout,
    )

2. Propagate non_blocking in .to()

In torchao/utils.py (_get_to_kwargs):
Currently, _get_to_kwargs parses args but ignores non_blocking. It should return it so subclass .to() methods can use it:

def _get_to_kwargs(self, *args, **kwargs):
    # ... existing code ...
    device, dtype, non_blocking, _ = torch._C._nn._parse_to(*args, **kwargs)
    device = self.device if device is None else device
    dtype = self.dtype if dtype is None else dtype
    kwargs = {
        "device": device,
        "dtype": dtype,
        "non_blocking": non_blocking, # <-- Add this
    }
    return kwargs

In AffineQuantizedTensor.to and Float8AQTTensorImpl.to:
Extract non_blocking from the parsed kwargs and pass it down to inner tensor transfers.

def to(self, *args, **kwargs):
    parsed_kwargs = self._get_to_kwargs(*args, **kwargs)
    device = parsed_kwargs.pop("device")
    non_blocking = parsed_kwargs.pop("non_blocking", False) # <-- Default to False for BC
    return self.__class__(
        self.tensor_impl.to(device, non_blocking=non_blocking),
        # ... other args
    )

[namgyu-youn]

AffineQuantizedTensor (AQT) is triggered by Float8WeightOnlyConfig(version=1) which is an old workflow (will be deprecated).

Could you reproduce using Float8WeightOnlyConfig(version=2) instead? Float8Tensor already supports memory pin if I remember correctly.

[lllllinux]

i’ve tried Float8WeightOnlyConfig(version=2),the device mismatch error still exists

[jerryzh168]

would this fix your problem? huggingface/diffusers#13276