make NVFP4Tensor handle per-expert outer scale

[vkuzo]

Summary:

NVFP4 MoE kernels usually do per-expert outer scale.
This PR extends NVFP4Tensor to support this.

Test Plan:

pytest test/prototype/mx_formats/test_nvfp4_tensor.py -s

[vkuzo]

Stack from ghstack (oldest at bottom):

• hook up hf + moe + nvfp4 script to lm_eval #4321
• wire up nvfp4 bmm #4320
• e2e example with HF model + MoE + torchao nvfp4 #4319
• hook up real nvfp4 grouped_gemm #4316
• -> make NVFP4Tensor handle per-expert outer scale #4315

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 1 New Failure

As of commit be9dc1b with merge base 0c8f44b ():

NEW FAILURE - The following job has failed:

• Run 1xH100 Tests / test (H100, linux.aws.h100, --pre torch torchvision torchaudio mslk --index-url https://download.... / linux-job (gh)
  RuntimeError: Command docker exec -t 5dd0e52871e98926ec9131b5dde934af7872e675cac6bc31722239d831223573 /exec failed with exit code 1

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

[jerryzh168]

not for this PR: I just checked the usage of self._orig_dtype seems it's duplicated with self.dtype now

[jerryzh168]

would be good to add a comment for attribute per_tensor_scale for NVFP4Tensor I think, also probably try to restrict this to weight only since it doesn't apply to activations?

[vkuzo]

would be good to add a comment for attribute per_tensor_scale for NVFP4Tensor I think

sure, fixed!

also probably try to restrict this to weight only since it doesn't apply to activations?

it could apply to activations (per-token-group)