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 d0b71cc 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(

[drisspg]

Processed prompts: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 15.06it/s, est. speed input: 108.45 toks/s, output: 195.81 toks/s]
Prompt: 'Why is Pytorch 2.0 the best machine learning compiler?', Generated text: " Ok,\nI am five months after changing? This one's"
Prompt: 'Hello, my name is', Generated text: ' Samuel Rodriguez.\r\n\r\nIf you want to know some regulars of'
Prompt: 'The president of the United States is', Generated text: ' revered has taken old wine beefs, morgs.\n'
Prompt: 'The capital of France is', Generated text: ' ruling Tauresearchednow.com offers a regularsolutions'
Prompt: 'The future of AI is', Generated text: ' moving... I thought about implementing logic, physics are one transform'
[rank0]:[W524 10:43:22.504024005 ProcessGroupNCCL.cpp:1538] Warning: WARNING: destroy_process_group() was not called before progr

Per Tensor is broken, it has to do with how we are slicing

[drisspg]

Okay I verified on base PR tht at I am seeing weird behavior in VLLM I am going to make a separate issue for this

[jerryzh168]

is this per tensor? should we standardize on this? some other place are probably using -1 to represent this

[drisspg]

I think -1 is fine as well,