Add support for float8 activation for Int4GroupwisePreshuffleTensor #2437
+239
−77
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2437
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Unrelated FailureAs of commit c13489e with merge base 2d61be8 ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following job failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
jerryzh168
added a commit
that referenced
this pull request
Jun 24, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
d6d3477 to
26517e8
Compare
This was referenced Jun 24, 2025
facebook-github-bot
added
the
CLA Signed
jerryzh168
added a commit
that referenced
this pull request
Jun 24, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
26517e8 to
d187f78
Compare
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
d187f78 to
2fcff42
Compare
jerryzh168
added a commit
that referenced
this pull request
Jun 24, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
added a commit
that referenced
this pull request
Jun 26, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
2fcff42 to
95856ed
Compare
jerryzh168
added a commit
that referenced
this pull request
Jun 27, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
95856ed to
1dec2cb
Compare
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
1dec2cb to
1645c79
Compare
jerryzh168
added a commit
that referenced
this pull request
Jun 27, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
added a commit
that referenced
this pull request
Jun 27, 2025
Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
1645c79 to
5e9e869
Compare
vkuzo
reviewed
Jul 2, 2025
| # making sure param.data is updated | ||
| assert param.data.packed_weight[0][0] != orig_value | ||
|
|
||
| def test_bmm(self): |
There was a problem hiding this comment.
unit tests are nice, feels like it would be good to also have an integration test to cover all of these ops in one go
There was a problem hiding this comment.
e2e tests can be the ones we run in vllm/sglang integration on real models I think?
vkuzo
reviewed
Jul 2, 2025
| @@ -2040,6 +2040,8 @@ class FbgemmConfig(AOBaseConfig): | |||
| weight_dtype (torch.dtype): weight dtype of the kernel | |||
| output_dtype (torch.dtype): output dtype of the kernel | |||
| group_size (int): The group size for weight | |||
| preshuffle (bool): whether preshuffle the weights or not | |||
| activation_dtype_for_int4 (str): the dtype for activation for int4 weight, either bf16 or fp8 | |||
There was a problem hiding this comment.
it's confusing to have both input_dtype and activation_dtype_for_int4, what are your thoughts
There was a problem hiding this comment.
yeah it's a bit confusing, although this is temporary, I'm deprecating this later in the stack: #2474
I'll remove this one and just use input_dtype
vkuzo
reviewed
Jul 2, 2025
| ) | ||
|
|
||
|
|
||
| class Float8ActivationInt4GroupwisePreshuffleTensor(TorchAOBaseTensor): |
There was a problem hiding this comment.
why not extend the existing Int4GroupwisePreshuffleTensor
There was a problem hiding this comment.
sorry this should be removed
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
c13fa2b to
cc359e6
Compare
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
cc359e6 to
b86939c
Compare
jerryzh168
commented
Jul 2, 2025
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
b86939c to
9ff26dd
Compare
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
9ff26dd to
b8466ab
Compare
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
b8466ab to
e666ef4
Compare
Summary: Added basic op support like linear and bmm, we have both float8 and bf16 in the same Tensor because it's the same dtype, only difference is whether the activation is quantized or not. Although there is some differneces in implementation: bf16 activaton: * group_scale * group_zero fp8 activation * group_scale * row_scale Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4
jerryzh168
force-pushed
the
jerryzh168/stack/4
branch
from
e666ef4 to
c13489e
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Stacked PRs:
transpose_inputfrom fbgemm configs #2422Add support for float8 activation for Int4GroupwisePreshuffleTensor
Summary:
Added basic op support like linear and bmm, we have both float8 and bf16 in the same Tensor
because it's the same dtype, only difference is whether the activation is quantized or not. Although
there is some differneces in implementation:
bf16 activaton:
fp8 activation
Test Plan:
python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py
Reviewers:
Subscribers:
Tasks:
Tags: