Skip to content

Move bitpacking.py to prototype and add uintx_utils.py#4152

Merged
jerryzh168 merged 32 commits into
mainfrom
gh/jerryzh168/68/head
Mar 25, 2026
Merged

Move bitpacking.py to prototype and add uintx_utils.py#4152
jerryzh168 merged 32 commits into
mainfrom
gh/jerryzh168/68/head

Conversation

@jerryzh168

@jerryzh168 jerryzh168 commented Mar 23, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

  • Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
    (original was deleted in the previous commit)
  • Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
    autoround and codebook features until they migrate to the new quantization design

… paths

- Remove PlainLayout class from dtypes/utils.py
- Delete torchao/dtypes/uintx/plain_layout.py
- Remove int8 weight and int8 dynamic activation dispatch from AQT dispatch table
- Remove AQT embedding dispatch (used PlainAQTTensorImpl)
- Change Int8WeightOnlyConfig default to version=2 (removes v1 AQT path)
- Change Int8DynamicActivationInt8WeightConfig default to version=2 (removes v1 AQT path)
- Remove PlainLayout from public exports
- Update tests to use v2 tensor types

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Mar 23, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⏳ No Failures, 3 Pending

As of commit 9d8b941 with merge base 1f90b4d (image):
💚 Looks good so far! There are no failures yet. 💚

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

…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
@jerryzh168 jerryzh168 added the module: not user facing Use this tag if you don't want this PR to show up in release notes label Mar 23, 2026
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]

@andrewor14 andrewor14 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are there any users for UintxTensor currently?

…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
@jerryzh168

Copy link
Copy Markdown
Contributor Author

Are there any users for UintxTensor currently?

I don't think so, this one is not very performant I think

@jerryzh168 jerryzh168 changed the base branch from gh/jerryzh168/68/base to main March 24, 2026 22:00
@jerryzh168

Copy link
Copy Markdown
Contributor Author

@claude review

@claude

claude Bot commented Mar 24, 2026

Copy link
Copy Markdown

Claude finished @jerryzh168's task in 1m 24s —— View job


Review of PR #4152

  • Gather context and read changed files
  • Review the diff against origin/main
  • Post review feedback

Summary: This PR moves bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/ and updates the import in uintx_utils.py accordingly. The changes are minimal and straightforward.

Findings

Bug: Broken test import

test/dtypes/test_bitpacking.py:10 still imports from the old path:

from torchao.dtypes.uintx.bitpacking import pack, pack_cpu, unpack, unpack_cpu

Since bitpacking.py was deleted from torchao/dtypes/uintx/ (by commit c98f09f earlier in the stack) and moved to torchao/prototype/dtypes/uintx/, this import will fail at runtime. The test file needs to be updated to import from the new location:

from torchao.prototype.dtypes.uintx.bitpacking import pack, pack_cpu, unpack, unpack_cpu

It's possible this is intentionally handled in another PR in the ghstack — but as it stands on this branch, the test is broken. Worth confirming.

Fix this →

The actual diff in this PR looks correct:

  • The rename of bitpacking.py is a clean move with no content changes (similarity index 100%).
  • The import update in uintx_utils.py:22 correctly changes from torchao.dtypes.uintx.bitpacking to torchao.prototype.dtypes.uintx.bitpacking.

Minor observations (non-blocking):

  • The __init__.py at torchao/prototype/dtypes/uintx/__init__.py exports nothing (__all__: list[str] = []). This is fine for now since uintx_utils.py and bitpacking.py are imported directly by path, but if these utilities are meant to be discoverable, you may want to add exports later.

Overall the change is clean — the only concern is the broken test import noted above.


@jerryzh168 jerryzh168 changed the base branch from main to gh/jerryzh168/68/base March 24, 2026 22:14
…ntx_utils.py"

- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
- Copy bitpacking.py from torchao/dtypes/uintx/ to torchao/prototype/dtypes/uintx/
  (original was deleted in the previous commit)
- Add uintx_utils.py with UintxTensor and UintxLayout classes needed by
  autoround and codebook features until they migrate to the new quantization design

[ghstack-poisoned]
@jerryzh168 jerryzh168 changed the base branch from gh/jerryzh168/68/base to main March 24, 2026 22:59
@jerryzh168

Copy link
Copy Markdown
Contributor Author

@claude review again

@claude

claude Bot commented Mar 24, 2026

Copy link
Copy Markdown

Claude finished @jerryzh168's task in 1m 0s —— View job


Re-review of PR #4152

  • Gather context and read changed files
  • Review the diff against origin/main
  • Post review feedback

Summary: This PR makes three changes:

  1. Renames torchao/dtypes/uintx/bitpacking.pytorchao/prototype/dtypes/uintx/bitpacking.py (100% similarity, no content changes)
  2. Updates the import in uintx_utils.py:22 to use the new torchao.prototype.dtypes.uintx.bitpacking path
  3. Updates the test import in test/dtypes/test_bitpacking.py:10 to use the new path

Findings

Previous issue resolved: The broken test import flagged in the last review has been fixed — test/dtypes/test_bitpacking.py:10 now correctly imports from torchao.prototype.dtypes.uintx.bitpacking.

No stale references: Searched the entire codebase for any remaining references to torchao.dtypes.uintx.bitpacking — none found. All references have been updated.

Import dependencies verified: The other imports in uintx_utils.py (PlainAQTTensorImpl from torchao.dtypes.uintx.plain_layout, Layout from torchao.dtypes.utils) point to files that exist on this branch.

No issues found. The changes are clean and correct. LGTM.


View job

@jerryzh168 jerryzh168 merged commit a927712 into main Mar 25, 2026
33 of 37 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: not user facing Use this tag if you don't want this PR to show up in release notes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants