[FP8] Extend per-token-group quantization support to QuantFP8 #24342
+444
−61
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
Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
tahsintunan
requested review from
mgoin,
robertgshaw2-redhat,
tlrmchlsmth and
yewentao256
as code owners
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request successfully refactors FP8 quantization by extending QuantFP8 to support per-token-group quantization, unifying the quantization paths. The changes are well-structured, improving code clarity and maintainability. I've identified one potential issue where a group size of 1 is not handled correctly, which could lead to a runtime error. My review includes a suggestion to fix this.
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
This was referenced Sep 5, 2025
|
Hey @ProExpertProg , it's not ready for review yet. Mainly opened this draft PR for early feedback. I'll address the comments and also add the torch implementation |
|
Okay please keep us posted, really looking forward to this work! |
Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
ProExpertProg
requested changes
Sep 8, 2025
There was a problem hiding this comment.
This is moving in the right direction, but we should still try to move towards using QuantFP8 instead of a bunch of free function calls around. If you want to address that in a separate PR that's fine but then we shouldn't touch the MoE layers in this one at all and just add support for group quant to QuantFP8 here.
ProExpertProg
requested changes
Sep 9, 2025
|
@ProExpertProg understood. Then I’ll just revert the latest changes and keep only the torch impl in this PR. The refactoring seems more involved than I had imagined. This will also keep the PR short and easy to review. |
Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
tahsintunan
force-pushed
the
quantfp8-group
branch
from
db3dc5e to
2662be1
Compare
|
@tahsintunan can you post performance numbers for this? We might need to enable the custom path by default if the torch path is not fast enough |
This was referenced Sep 11, 2025
B200 results |
|
|
||
| if padded_dim != hidden_dim: | ||
| padding = padded_dim - hidden_dim | ||
| x = F.pad(x, (0, padding), mode='constant', value=0.0) |
There was a problem hiding this comment.
I wonder if there's a way to do this without padding - I worry the generated Triton kernel won't be able to eliminate the copy
There was a problem hiding this comment.
I'm realizing the padding is unlikely anyway as group shape will likely divide the hidden_size
There was a problem hiding this comment.
Yes, padding will only be used for non-standard dimensions, which should be rare.
Closed
|
@tahsintunan do you have a timeline for implementing the feedback? Based on the performance numbers we want to merge this as soon as feasible. I'd also be happy to finish your work and merge (you'll still be a coauthor on the commit) - let me know! Thanks for working on this |
|
@ProExpertProg I can take care of the remaining items by Monday. If you need this merged sooner, feel free to wrap it up. If not, I'll get it done by Monday. |
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
|
@tahsintunan yeah today or early tomorrow is totally fine! Thanks for taking this on. |
Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
ProExpertProg
approved these changes
Sep 15, 2025
ProExpertProg
added
the
ready
|
I ran the benchmark on H100 and found it was slower for 1x128 (the case we care about) |
mgoin
reviewed
Sep 16, 2025
Comment on lines
+36
to
+37
| # recompile for different shapes | ||
| fwd = torch.compile(fn, fullgraph=True, dynamic=False) |
There was a problem hiding this comment.
Why aren't we compiling with dynamic=True? I don't think we should be targeting shape specialization since we won't use that in practice
There was a problem hiding this comment.
In practice we specialize on all shapes except the first dim (num_tokens). The with_dyn_arg marks that shape as dynamic to fully simulate vLLM usage 👍
Comment on lines
59
to
96
| batch_size_range = [1, 16, 32, 64, 128] | ||
| seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096] | ||
| hidden_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096] | ||
| batch_sizes = [1, 16, 32, 64, 128] | ||
| group_shapes = [ | ||
| GroupShape.PER_TENSOR, | ||
| GroupShape.PER_TOKEN, | ||
| GroupShape(1, 64), | ||
| GroupShape(1, 128), | ||
| ] | ||
| column_major_scales = [True, False] | ||
|
|
||
| config_gen = itertools.product( | ||
| group_shapes, | ||
| column_major_scales, | ||
| batch_sizes, | ||
| hidden_sizes, | ||
| ) |
There was a problem hiding this comment.
Could we make some of these args? It takes a really long time by default to run all of this
There was a problem hiding this comment.
Can that be a follow-up it requires reworking the structure a lot (because currently this is passed to the function annotation).
There was a problem hiding this comment.
Addressed. Now you should be able to use something like this:
python3 benchmarks/kernels/bench_per_token_quant_fp8.py --hidden-sizes 1024 2048 4096 --batch-sizes 32 --group-sizes 128 --no-column-major
| assert not static, "Group quantization only supports dynamic mode" | ||
| self.group_size = group_shape.col | ||
| else: | ||
| assert group_shape in {GroupShape.PER_TOKEN, GroupShape.PER_TENSOR} |
There was a problem hiding this comment.
Can we add an assert that column_major_scales is False if non group?
|
|
||
| x_quant = x_quant.view(-1, padded_dim) | ||
| if padded_dim != hidden_dim: | ||
| x_quant = x_quant[..., :hidden_dim] |
There was a problem hiding this comment.
Should we make sure this is contiguous after stripping padding?
Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com>
Ran this on 5090 |
|
Merging for now since it is valid and isn't on by default, thanks for the nice work! |
1 task
FeiDaLI
pushed a commit
to FeiDaLI/vllm
that referenced
this pull request
Sep 25, 2025
…roject#24342) Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com>
charlifu
pushed a commit
to ROCm/vllm
that referenced
this pull request
Sep 25, 2025
…roject#24342) Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: charlifu <charlifu@amd.com>
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.
Purpose
Extends
QuantFP8to support per-token-group quantization and adds a torch implementation of the QuantFP8 group quantization.Addresses #24185
Changes
is_per_tensor(),is_per_token(),is_per_group()helper methods toGroupShapeQuantFP8to support arbitrary group sizes likeGroupShape(1, 128)test_fp8_quant_group.pyandbenchmark_quantfp8_group.pyTest Plan
Tested with existing test suite:
tests/kernels/quantization/test_fp8_quant.pytests/kernels/quantization/test_fp8_quant_group.py