nvfp4 tensor: switch to TorchAOBaseTensor

[vkuzo]

Summary:

Switches NVFP4Tensor to inherit from TorchAOBaseTensor, and removes the
custom tensor flatten/unflatten to use the base class version.

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• nvfp4 tensor: refactor weight-only vs dynamic quant #2790
• -> nvfp4 tensor: switch to TorchAOBaseTensor #2788

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ 5 Pending, 6 Unrelated Failures

As of commit b7d6069 with merge base 4463b79 ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio fbgemm-gpu-genai --index-url https... / linux-job (gh) (detected as infra flaky with no log or failing log classifier)
• Run 1xL4 Tests / test (SM-89, linux.g6.4xlarge.experimental.nvidia.gpu, --pre torch --index-url https://download.p... / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Run Regression Tests / test (CPU 2.6, linux.4xlarge, torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Run Regression Tests / test (CPU 2.7, linux.4xlarge, torch==2.7.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Run Regression Tests / test (CPU 2.8, linux.4xlarge, torch==2.8.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128
• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/wh... / linux-job (gh) (matched linux rule in flaky-rules.json)
  The process '/usr/bin/git' failed with exit code 128

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

[jerryzh168]

nit: remove underscores?

[vkuzo]

not related to current PR, we can do this in a future PR

[drisspg]

from a typehint / ide experience this kinda sucks

is it not possible to have the actual attributes typehinted on the class w/ AOBaseTensor?

[jerryzh168]

we can try to do that as an improvement for future maybe, not sure how to do that yet, maybe we can do some inspection

[jerryzh168]

not sure if we could change it, that's what we've been using in the past

[vkuzo]

I think it's worth changing, right now it's a list of attribute values without the names, which is way harder to debug than a key:value dictionary.