Fix Per Row scaling for inference

[drisspg]

Stacked PRs:

• integration-vllm-test #2258
• Add a way to do power of 2 scaling #2256
• ->Fix Per Row scaling for inference #2253

---

Summary

In somewhere in the myriad of refactors we broke per-row scaling. This had no effect becuase dequant just so happend to work and a global scale, so we test now

[pytorch-bot]

🔗 Helpful Links

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

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

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

[Copilot]

Pull Request Overview

This PR restores per-row scaling support for float8 quantization and updates related APIs and tests.

• Add block_size support to quant parameter selection and dequantization for block-wise scaling
• Normalize granularity in the quantization API post-init
• Fix row-wise scale handling in float8 layouts and update tests for verifying per-row scale shapes

Reviewed Changes

Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
quant_primitives.py	Introduce block_size argument for block-wise scale compute and update dequant logic
quant_api.py	Normalize granularity tuple in __post_init__
float8_layout.py	Adjust row-wise scale transpose logic
affine_quantized_tensor.py	Pass block_size into choose_qparams_affine_float8
test_affine_quantized_float.py	Parametrize and verify per-row scale shapes in tests

Comments suppressed due to low confidence (4)

torchao/quantization/quant_primitives.py:1985

• The comment about supporting only tensorwise scaling is outdated since block-wise scaling is now implemented; please update or remove this comment to avoid confusion.

# only tensorwise scaling is supported for now:

torchao/quantization/quant_api.py:1579

• [nitpick] This normalization block appears duplicated from the class definition in the context excerpt; consider consolidating to avoid redundant logic.

activation_granularity, weight_granularity = _normalize_granularity(

torchao/dtypes/floatx/float8_layout.py:341

• Removing the unsqueeze(-1) before transpose changes the tensor's rank and may break row-wise scale alignment; restore the unsqueeze or adjust downstream logic to match expected dimensions.

w_scale = w_scale.T

torchao/dtypes/affine_quantized_tensor.py:465

• The variable block_size is used here but not defined or passed into this scope; ensure block_size is available or passed from the calling context.

scale = choose_qparams_affine_float8(

[danielvegamyhre]

was support for blockwise quantization intentionally included in this PR? this looks like the quant primitive support for it but presumably other things are still needed like scaling granularity (which only is PerTensor or PerRow right now), etc.

[drisspg]

This primitive, can/should be generic to block-size - which is what I did arguablly I hsould update the doc string a little, but its up to the caller to pass in the right block_size which in this case is only 2 difffernt forms today per-tensor and per-row