[pculliton]

Summary

This PR introduces generic support for Static Range Quantization (SRQ) — static int8 activation quantization, intx (int1..int8) weight quantization, and optional static output quantize-dequantize scaling — via a dedicated IntxStaticActivationTensor subclass and Int8StaticActivationIntxWeightConfig.

(Note: Per review feedback, the IntxUnpackedToInt8Tensor embedding and row-select improvements have been split out into a separate PR.)

---

Motivation

Our primary motivation is to support pre-calibrated low-bit (int2 / int4 / int8) static range quantization (SRQ) models in TorchAO:

1. Dedicated IntxStaticActivationTensor subclass for intx weights (int1..int8) paired with calibrated static int8 activation scales (act_quant_scale, act_quant_zero_point) and optional output scaling (output_quant_scale).
2. Checkpoint serialization support (state_dict with weights_only=True and safetensors) storing calibration tensors directly on IntxStaticActivationTensor.

---

Key Changes

1. Core Tensor Subclass (torchao/quantization/quantize_/workflows/intx/intx_static_activation_tensor.py)

• Added IntxStaticActivationTensor to represent unpacked intx weights (torch.int1..torch.int8) with optional pre-calibrated per-tensor or per-channel act_quant_scale, act_quant_zero_point, and output_quant_scale.
• Dispatches CUDA per-channel (PerAxis(0)) linear layers with per-tensor activation scales (M > 16 and K, N % 8 == 0) to safe_int_mm with exact int32 affine zero-point corrections scaled once in fp32, routing smaller/unaligned shapes, grouped (PerGroup) weights, per-channel activation scales, and CPU execution through static activation quantize-dequantize + F.linear.
• Supports .to() (preserving calibration scale dtypes on device-only moves), aten.embedding.default / F.embedding, aten.select.int, and aten.slice.Tensor (including per-channel calibration scale slicing).

2. Configuration & Serialization (torchao/quantization/quant_api.py, torchao/prototype/safetensors/safetensors_utils.py)

• quant_api.py:
  • Added and registered Int8StaticActivationIntxWeightConfig (matching Int8StaticActivationInt8WeightConfig for in-memory per-layer calibration tensors, plus custom_scale / custom_zero_point handler support for QAT convert).
  • Updated _adjust_scale_dtype_in_intx_unpacked_tensor to support IntxStaticActivationTensor when weight_scale_dtype != weight.dtype.
• prototype/safetensors/safetensors_utils.py:
  • Registered IntxStaticActivationTensor in ALLOWED_CLASSES and ALLOWED_TENSORS_SUBCLASSES.

3. Unit Tests (test/quantization/quantize_/workflows/intx/test_intx_static_activation_tensor.py)

• Added a dedicated test suite TestIntxStaticActivationTensor covering:
  • CPU and CUDA numerical reference equivalence across PerAxis(0) and PerGroup(32), symmetric and asymmetric weights, nonzero act_quant_zero_point, output_quant_scale, bfloat16 / float32 scales, and small/large batch sizes (M=2 and M=32).
  • CUDA safe_int_mm fast-path vs. dequantize fallback selection.
  • weight_scale_dtype != weight.dtype, from_hp custom weight scale / zero point, and QATConfig(..., step="convert") range-learning conversion.
  • .to() device and dtype moves, calibration shape validation, and successive forward passes with varying input ranks (2D and 3D).
  • torch.compile(fullgraph=True) on CPU and CUDA (PerAxis(0) and PerGroup(32), M=2 and M=32).
  • Config metadata round-trip (config_from_dict(config_to_dict(config))), state_dict serialization (weights_only=True), safetensors flatten/unflatten, nn.Embedding, select, and slice.

---

Verification

Verified locally on CPU and CUDA:

pytest test/quantization/quantize_/workflows/intx/test_intx_static_activation_tensor.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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:

• PRs based before #202327 fail all Linux builds: can't find build-external-packages action

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

[meta-cla]

Hi @pculliton!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[meta-cla]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[metascroy]

Why not make this cuda usecase its own tensor subclass? It seems different enough from the rest of the intx tensor?

[pculliton]

Hmm, I think that would make sense, thanks! I'll split it up.

[pculliton]

@metascroy pushed 3f70ff5 splitting IntxStaticActivationTensor into its own subclass as suggested - please take a look when you get a chance!

[digantdesai]

Thanks @pculliton.

@metascroy get we please take a look at this once more?

[metascroy]

1. Please split the IntxUnpackedToInt8Tensor embedding improvements into their own PR.
2. The new static tensor should have its own test file.

AI disclosure: Scott requested the two changes above and explicitly asked Devmate to post the actionable findings from its review as inline comments. The AI-generated technical analysis is labeled and quoted in each inline comment.

Validation scope: targeted CPU reproductions used PR modules loaded in memory against local TorchAO dependencies. No full-suite or CUDA execution was performed; the CUDA zero-point example evaluates the branch's arithmetic on CPU. The added static-activation test checks output shapes, but does not check numerical correctness or invoke torch.compile.

[metascroy]

AI-assisted finding (Devmate), posted at Scott's request:

[P1] Apply static activation quantization in the fallback. CPU execution and CUDA execution with grouped weights call floating-point linear on the original input, so act_quant_scale has no effect on these paths. For example, with input values of 0.25, activation scale 1.0, and no bias, activation quantize-dequantize produces zero, but a targeted CPU check returned approximately 7.47 from this fallback. Quantize-dequantize the activation using the calibrated parameters before the floating-point linear, or explicitly reject unsupported configurations. Please add numerical reference tests for CPU and CUDA grouped-weight execution, not just output-shape assertions.

[metascroy]

AI-assisted finding (Devmate), posted at Scott's request:

[P1] Honor asymmetric activation and weight zero points. The config accepts asymmetric quantization, but activation quantization above hardcodes a zero point of zero instead of using act_quant_zero_point, and this matmul/rescaling omits the weight-zero-point correction. For asymmetric int4 unit weights of shape (2, 32), weight zero points of -8, unit input, and activation scale 1.0, evaluating this branch's arithmetic on CPU gives approximately 14.93 instead of the dequantized reference result 32. Implement the affine corrections, or reject/fall back for nonzero zero points. Please cover nonzero activation and weight zero points with CUDA numerical reference tests.

[metascroy]

AI-assisted finding (Devmate), posted at Scott's request:

[P2] Use a scale-dtype adjustment compatible with the new subclass. new_weight is an IntxStaticActivationTensor, but _adjust_scale_dtype_in_intx_unpacked_tensor explicitly asserts that its input is an IntxUnpackedToInt8Tensor. These are separate subclasses, so requesting a different scale dtype, such as float32 scales for bfloat16 weights, raises an assertion. Add support for the new type and a test exercising weight_scale_dtype != weight.dtype.

[metascroy]

AI-assisted finding (Devmate), posted at Scott's request:

[P2] Restore scalar calibration tensors during JSON deserialization. The new encoder converts a zero-dimensional tensor such as torch.tensor(0.5) into the JSON number 0.5, not a list. These list-only conversions therefore leave the restored scale as a Python float, and the quantization transform fails with AttributeError: 'float' object has no attribute 'detach'. This was reproduced with config_from_dict(config_to_dict(config)); scalar scales are also what the added test uses. Handle scalar tensor round-trips as well as lists, including the corresponding conversions in Int8StaticActivationInt8WeightConfig, and test that a restored config can actually quantize a module.

[metascroy]

AI-assisted finding (Devmate), posted at Scott's request:

[P2] Preserve embedding forward semantics in the sliced path. padding_idx does not force the stored padding row to zero during ordinary embedding forward; a loaded or explicitly supplied nonzero row must still be returned. This mask changes that result. The fast path also drops max_norm and norm_type entirely instead of forwarding them as the previous implementation did. A targeted check of the equivalent new static-tensor implementation with max_norm=1.0 returned row norms of approximately 5.28, versus 1.0 from the ordinary embedding reference. Preserve these options in both copies of the sliced implementation, or fall back when necessary. Add tests for a nonzero padding row and max-norm renormalization in the separate embedding PR.

[metascroy]

AI-assisted finding (Devmate), posted at Scott's request:

[P2] Normalize the activation scale locally instead of mutating the weight. The first forward replaces a scalar act_quant_scale with a tensor whose rank matches that input. Subsequent forwards of another rank no longer enter this branch and reuse the wrong scale shape. With per-tensor static activation quantization, a targeted check succeeded for input shape (1, 32) and then raised an assertion for (1, 1, 32) on the same weight. Keep the stored scalar unchanged, reshape a local scale for each call, and add a successive-forward test covering different input ranks.

[pculliton]

Thanks for the review @metascroy! I've updated the PR to address all feedback:

1. Split the IntxUnpackedToInt8Tensor embedding improvements out of this PR into a separate PR.
2. Added a dedicated test file with the requested tests.
3. Addressed all 6 P1 / P2 comments (static activation quantize-dequantize in the fallback linear path, affine zero-point corrections in the CUDA path, etc.).

Let me know what you think!

[metascroy]

Thanks @pculliton! Overall it looks good to me, but why do you need to extend Int8Tensor in this PR (it seems like an independent feature that could go in its own PR). More generally, does your IntxStaticActivationTensor with x=8 cover the usecases you had for the extension to Int8Tensor?

[pculliton]

Thanks @pculliton! Overall it looks good to me, but why do you need to extend Int8Tensor in this PR (it seems like an independent feature that could go in its own PR). More generally, does your IntxStaticActivationTensor with x=8 cover the usecases you had for the extension to Int8Tensor?

Thanks @metascroy! Yep, IntxStaticActivationTensor with x=8 covers our 8-bit SRQ use cases.

Good point on the Int8Tensor / Int8StaticActivationInt8WeightConfig changes - they were left over from the initial version of the PR before we split out IntxStaticActivationTensor. Two of the Devmate comments asked for fixes on those lines and I kinda forgot to remove them after I made the changes. :) I've pulled them out into the other PR.

Let me know what you think!

[vkuzo]

why is there a tensor here? this isn't expected in a config

[pculliton]

Thanks! Reverted torchao/core/config.py back to main and removed custom_weight_scale / custom_weight_zero_point (and the JSON list/scalar conversion) from Int8StaticActivationIntxWeightConfig.

For the calibration scales themselves (act_quant_scale, act_quant_zero_point, output_quant_scale), we kept Optional[torch.Tensor] = None on Int8StaticActivationIntxWeightConfig to match Int8StaticActivationInt8WeightConfig (and Float8StaticActivationFloat8WeightConfig) for in-memory per-layer quantization via FqnToConfig, while serialized checkpoints store them directly on IntxStaticActivationTensor in the state_dict.

[vkuzo]

@claude review

[claude]

Claude finished @vkuzo's task in 1m 51s —— View job

---

Review: SRQ via IntxStaticActivationTensor

• Gather diff/context
• Review tensor subclass
• Review config/serialization
• Review tests
• Post review

Overall the split into a dedicated subclass reads well, and the earlier Devmate findings look fixed. I checked the affine zero-point algebra in the CUDA path by hand and it is correct: it expands to s_x·s_w·[Σq_xq_w − zp_x·Σq_w − zp_w·Σq_x + K·zp_x·zp_w]. I did not run the tests: I couldn't fetch origin/main (that needs extra permissions), so I reviewed against the PR's base commit a195aaf.

---

1. Tensors in the config (expanding on @vkuzo's comment on torchao/core/config.py:142)

Int8StaticActivationInt8WeightConfig already sets a precedent for act_quant_scale living on the config. This PR goes further in two ways.

• Per-layer weight tensors on a config. custom_weight_scale and custom_weight_zero_point have the weight's shape, e.g. (N, K/group_size). That is per-layer, per-shape data, not configuration. A single config applied through quantize_(model, config) would fail or silently misapply on any layer with a different shape, so this only works with one FqnToConfig entry per layer. In that case the "config" is really a checkpoint. The transform also already accepts custom_scale / custom_zero_point kwargs (quant_api.py:1136), so the config fields duplicate that path. Suggestion: drop custom_weight_* from the config and keep only the handler kwargs. Even better, load pre-calibrated weights as tensor subclass data in the state_dict.

• The global ConfigJSONEncoder change is lossy and one-directional.

  • o.detach().cpu().tolist() drops dtype and shape information: a 0-d tensor becomes a scalar, and a (1,1) tensor becomes [[x]].
  • It emits a bare JSON value with no _type tag, so the decoder can't tell it was a tensor. That's why every field needs the hand-written isinstance(..., (list, int, float)) code in __post_init__ (quant_api.py:1085-1106).
  • The encoder change affects every config, while the decode fixup only exists for this one config.
  • Dtypes are not preserved: a bf16 scale comes back as fp32, and an int8 zero point comes back as int64.

  If tensors stay in the config, encode them like torch.dtype is encoded, as {"_type": "torch.Tensor", "_data": ..., "dtype": ..., "shape": ...}, and decode them centrally in config_from_dict. My preference is still to keep calibration data out of the config.

2. Precision of the CUDA fast path (intx_static_activation_tensor.py:388-412)

intermediate_dtype = torch.float if act_scale.dtype == torch.half else act_scale.dtype
y_dot_scaled = _int_scaled_matmul(...).to(torch.float32)

• If act_scale is bf16, _int_scaled_matmul returns int32_acc * bf16_scale in bf16. That happens after model.to(torch.bfloat16), because to() casts act_quant_scale to the model dtype at line 167.
• The zero-point corrections are then subtracted in fp32 from that already-rounded value. With asymmetric weights, Σq_x·q_w and zp_w·Σq_x are large and close to each other, so most of the result cancels. That turns bf16 rounding (about 3 significant digits) into a large relative error.
• The tests don't cover this because they always pass fp32 act_quant_scale.

Suggestion: apply the corrections in the integer domain, where they are exact, and scale once in fp32:

acc = safe_int_mm(tmp, w_t)                        # int32
acc -= zp_x * w_row_sums_int32                     # if act zp
acc -= zp_w * (x_row_sums_int32 - K * zp_x)        # if weight zp
y = acc.float() * s_x.float() * s_w.float()

This also skips the w_row_sums reduction over qdata on every forward. It is constant per weight, so it could be precomputed. Also, weight_tensor.zero_point is not None at line 402 is always true. Consider skipping the correction when the weights are symmetric (all-zero zp).

3. act_mapping_type is validated but never used

Int8StaticActivationIntxWeightConfig.act_mapping_type is asserted in __post_init__ but never passed to from_hp or used in dispatch. Whether activation quantization is asymmetric depends only on whether act_quant_zero_point is set. Remove the field, or make it meaningful, e.g. assert that act_quant_zero_point is None when the mapping is SYMMETRIC.

4. Silent int8 wrap of activation zero points

from_hp (line 285) and linear (line 356) cast act_quant_zero_point to int8 with .to(torch.int8). A calibrated zero point outside [-128, 127] wraps silently instead of failing; for example, a uint8-style calibrator producing 128 becomes -128. Weight zero points get the same cast at line 239. Add a range assert before casting.

5. aten.slice.Tensor issues (lines 543-581)

• Unhandled end: end=None, or sys.maxsize (which w[:, :] produces), is not clamped. end % block_size then raises TypeError or fails the assert. Clamp it with end = min(end, self.shape[dim]), as other TorchAO tensors do.

• Calibration tensors are not sliced: act_quant_scale, act_quant_zero_point and output_quant_scale are passed through unchanged.

  • If output_quant_scale is per-output-channel, slicing dim 0 leaves it with the wrong length.
  • If the activation scale is per-input-channel, slicing dim 1 has the same problem.

  Either slice them along the matching dim, or assert that they are per-tensor (numel() == 1) when slicing.

6. Scale broadcasting is permissive (lines 343-352, 442-448)

input_block_size = input.shape[i] // act_scale.shape[i] silently produces wrong blocks when the shapes don't divide evenly. A non-scalar act_scale that matches no dimension is also accepted. A per-token static scale (M,) is tied to one batch shape, which is probably not intended for SRQ. Suggestion: explicitly support only per-tensor and per-input-channel (K,) activation scales, and assert on anything else.

7. Leftovers from the split

• IntxUnpackedToInt8TensorActivationQuantization is newly exported from quantize_/workflows/__init__.py and added to ALLOWED_CLASSES in safetensors_utils.py. Nothing in this PR uses it, so it looks like a leftover from the earlier IntxUnpackedToInt8Tensor version.
• The swap_conv2d_1x1_to_linear lambda reformat (quant_api.py:241) is unrelated formatter churn.
• intx_packing_format: the assert at quant_api.py:1199 runs after quantization has finished. Move it into __post_init__, or drop the field, since only one value is accepted.

8. Smaller items

• version: int = 1: new configs elsewhere use version: int = 2, and the transform never checks it. Drop it, or follow the v2 convention.
• Fallback for some shapes: safe_int_mm falls back to a CPU matmul (.cpu()) when K or N is not a multiple of 8. test_asymmetric_unit_weight_zero_point hits this path (N=2). That's fine for a test, but it means the CUDA path is only fast for aligned shapes. Worth a comment, or falling back to the dequant path instead.
• Embedding fast path: max_norm on the F.embedding path now correctly falls back. padding_idx is ignored in the fast path, which matches forward semantics.
• Weight zero-point dtype: the __init__ assert now accepts int32/int64 weight zero points, but from_hp always converts to int8. Consider tightening the assert to int8 or float so the stored dtype stays canonical.

9. Tests

The coverage is good. Gaps worth adding:

• bf16 act_quant_scale, or model.to(torch.bfloat16) after quantization, on CUDA with asymmetric weights (catches Headers and gitignore #2).
• Slicing with end=None or a per-channel output_quant_scale (Test installation CI #5).
• Config round-trip that checks act_quant_scale.dtype / .shape are preserved, not only that the forward runs (Adding gpu quantization workflows and apis #1).

---

Summary: the numerics of the zero-point path are correct in exact arithmetic. The main items before merge:

1. The design question about tensors and per-layer weight data in the config, plus how they are serialized (Adding gpu quantization workflows and apis #1).
2. Moving the CUDA-path corrections into int32 to avoid bf16 cancellation (Headers and gitignore #2).
3. Cleaning up the unused act_mapping_type and the leftovers from the split (Move test.py #3, Adding installation and usage to readme #7).

[pculliton]

Thanks @vkuzo! Pushed an update addressing your comment and the items from the @claude review:

• Config & serialization (Adding gpu quantization workflows and apis #1, Move test.py #3, Adding installation and usage to readme #7, Documentation Updates #8): Reverted torchao/core/config.py and unrelated formatting in quant_api.py, removed unused IntxUnpackedToInt8TensorActivationQuantization exports, dropped custom_weight_* from Int8StaticActivationIntxWeightConfig (keeping custom_scale / custom_zero_point on the transform handler and from_hp), enforced act_quant_zero_point is None when act_mapping_type == SYMMETRIC, moved the intx_packing_format check to __post_init__, and set version = 2.
• CUDA fast path & numerics (Headers and gitignore #2, Move ao to torchao #4, Documentation Updates #8): Switched the CUDA per-channel path to safe_int_mm with exact int32 zero-point corrections scaled once in fp32, gated safe_int_mm on M > 16 and K, N % 8 == 0 (routing smaller/unaligned shapes to the dequantize path so torch.compile(fullgraph=True) works for decode shapes and avoids CPU fallback), preserved calibration scale dtypes on device-only .to(), and added [-128, 127] range checks before casting zero points to int8.
• Calibration shapes & slicing (Test installation CI #5, fixing imports for moved test #6, Documentation Updates #9): Restricted static scales to per-tensor (numel() == 1) or per-channel, clamped end and sliced per-channel calibration scales in aten.slice.Tensor, and expanded CPU/CUDA + torch.compile tests to cover the M > 16 fast path, bfloat16 .to(), QATConfig convert, safetensors, and slicing.

Please take a look when you get a chance!