make scale shape 2d and match qdata shape in NVFP4Tensor

[vkuzo]

Summary:

Makes NVFP4Tensor._scale_e4m3 shape be consistent with NVFP4Tensor.shape. Specifically:

1. if the scales are not swizzled, contiguous data shape (M, K) corresponds to scale shape (M, K // 16)
2. if the scales are swizzled, contiguous data shape (M, K) corresponds to scale shape (ceil_div(M, 128) * 32, ceil_div(K, 64) * 16)

If we transpose axes 0 and 1, both qdata and scale now get transposed.

I want this because we need to reason about scales when combining 2d MoE weights into 3d weights, and having the scale shape match the data will make that reasoning easier (in a future PR).

Test Plan:

tests pass on a B200:

pytest test/prototype/mx_formats/ -s

also, a dense model quantized with torchao nvfp4 runs in vLLM before and after this PR

Reviewers:

Subscribers:

Tasks:

Tags:

[vkuzo]

Stack from ghstack (oldest at bottom):

• enable 3d weights for NVFP4Tensor #3109
• -> make scale shape 2d and match qdata shape in NVFP4Tensor #3108

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 08e9d13 with merge base 9368b28 ():
💚 Looks good so far! There are no failures yet. 💚

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

[vkuzo]

need to fix swizzle slicing tests