Context

• Tensor subclass in question
• TP test in question

Summary

I am working on TP support for MXFP8WeightWrapperTensor subclass used for MXFP8 training in torchao: https://github.com/pytorch/ao/pull/3985

The basic design of this tensor subclass is the following:

• In __torch_function__ intercept linear ops or mm.default to execute an autograd function in its place, which does differentiable mxfp8 quantization + scaled_mm.
• For all other ops, continue to __torch_dispatch__ and behave like a regular tensor.

For TP, we wrap weights in the tensor subclass first, then wrap in DTensor when TP is applied, so it looks like:

DTensor(MXFP8WeightWrapperTensor(...))

When DTensor is not applied, the op override for linear in __torch_function__ works as expected. Logs:

(Pdb) toy_model_fp8(x_bf16)
[TORCH_FUNCTION] linear
mx_mm forward
[TORCH_FUNCTION] linear
mx_mm forward
[TORCH_FUNCTION] linear
mx_mm forward

However, when DTensor wraps the subclass, logging shows that linear is decomposed into __get__, t.default, mm.default, and the tranpose + mm ops are skipping __torch_function__ and going directly to __torch_dispatch__, which prevents us from intercepting our call and calling an autograd function in its place. Logs:

[TORCH_FUNCTION] __get__
[TORCH_DISPATCH]:  mm.default
[TORCH_DISPATCH]:  t.default
[TORCH_FUNCTION] __get__
[TORCH_DISPATCH]:  mm.default
[TORCH_DISPATCH]:  t.default
[TORCH_FUNCTION] __get__
[TORCH_DISPATCH]:  mm.default
[TORCH_DISPATCH]:  t.default

...
tp out comparison not close, shapes: tp=torch.Size([2, 128, 64]), global=torch.Size([2, 128, 64])
tp_out stats: min=-0.134765625, max=0.150390625, mean=0.000507354736328125
global_out stats: min=-0.1318359375, max=0.150390625, mean=0.000659942626953125
diff stats: min=0.0, max=0.010009765625, mean=0.0016021728515625
tp_out SQNR: 23.625

I don't know what __get__ is doing, but the fact that some linear ops are skipping __torch_function__ prevents us from intercepting and thus causes regular bf16 mm to be executed, which is not the desired behavior.

We can't intercept in torch dispatch because autograd runs at the torch_function level, it will not capture the backward pass we are defining in the custom autograd func.

Is there a way we can have DTensor send these mm ops through torch_function so we can intercept?

Repro

• On b200 devgpu, check out this torchao PR: https://github.com/pytorch/ao/pull/3985
• Run torchrun --nproc_per_node 2 test/prototype/mx_formats/test_mx_dtensor.py
• torch version: 2.12.0.dev20260305+cu130