Stack from ghstack (oldest at bottom):

• -> support pinning for mx and nvfp4 tensors #4192

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4192

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 6e387dc with merge base c49a60e ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@claude review this please

Claude finished @andrewor14's task in 2m 4s —— View job

Code Review: Support pinning for mx and nvfp4 tensors

• Read changed files and PR diff
• Review implementation in mx_tensor.py
• Review implementation in nvfp4_tensor.py
• Review tests
• Post review feedback

Overall this is a clean, straightforward PR. The pattern follows existing conventions in the codebase (e.g., Float8Tensor, Int8Tensor, NF4Tensor all implement pinning the same way). Constructor argument ordering is correct for both tensor types. A few observations:

Implementation (mx_tensor.py:640-657)

The is_pinned and _pin_memory implementations for MXTensor are correct and minimal. The constructor call in mx_pin_memory correctly passes all 8 positional arguments in the right order matching MXTensor.__new__ at line 459.

Implementation (nvfp4_tensor.py:329-357)

The NVFP4Tensor implementation correctly handles the optional per_tensor_scale and act_per_tensor_scale fields in both is_pinned (checking them if non-None) and _pin_memory (conditionally pinning them). Constructor argument order matches NVFP4Tensor.__new__ at line 79.

Minor observations (non-blocking):

1. aten._pin_memory.default has a device kwarg — The ATen schema for _pin_memory is _pin_memory(Tensor self, Device? device=None) -> Tensor. The device arg is used for pinning to a specific CUDA device. This PR ignores it (doesn't forward kwargs to the inner .pin_memory() calls), but this is consistent with how every other tensor subclass in the repo handles it (Float8Tensor, Int8Tensor, NF4Tensor all ignore it too). So this is fine as-is — just noting for awareness.

2. MX test only covers float8_e4m3fn and float8_e5m2 — MXTensor also supports torch.float4_e2m1fn_x2. If FP4 MX tensors are used in diffusers offloading, it might be worth adding a parametrized case for that dtype too. But this is a nice-to-have, not a blocker.

Tests

Both test functions follow a solid pattern:

• Create on CUDA → move to CPU → verify not pinned → pin → verify pinned → verify inner tensors pinned → verify dequantization equivalence

The NVFP4 test additionally parametrizes over use_per_tensor_scale, which covers the optional field path. The assert not x_cpu.is_pinned() after pinning confirms that pin_memory() returns a new tensor rather than mutating in-place — good.

Verdict: LGTM. Clean implementation, correct constructor calls, good test coverage, and consistent with existing patterns in the codebase.