[FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968

vllmellm · 2025-03-17T15:33:00Z

Description

This PR integrates the Block-Scaled GEMM functionality from AITER into vLLM, and will allow any up-coming optimizations in AITER kernel to be directly used and evaluated within the vLLM framework.

Implementation

The gemm_a8w8_blockscale kernel from AITER has been added to /vllm/model_executor/layers/quantization/utils/fp8_utils.py. This kernel is:

Enabled only if VLLM_ROCM_USE_AITER and VLLM_ROCM_USE_AITER_LINEAR are both set to 1.
Suitable for DeepSeek models.

Testing

The integration has been verified through:

High-level integration tests with various models
Kernel function dispatch testing to ensure correct operation selection

[Updated] Performance

V1 Engine
Summary of Improvements
When comparing the performance with and without AITER Blockscaled GEMM FP8, the following improvements were observed:

Metric	Without AITER	With AITER	Improvement
Benchmark duration	142.67s	119.71s	16.1% faster
Request throughput	3.50 req/s	4.18 req/s	19.4% higher
Output token throughput	889.13 tok/s	1054.40 tok/s	18.6% higher
Total token throughput	4346.57 tok/s	5175.02 tok/s	19.1% higher
Mean inter-token latency	87.52 ms	73.41 ms	16.1% lower
Median inter-token latency	55.92 ms	45.66 ms	18.3% lower

Key Observations

Overall Speed: The benchmark completed 16.1% faster with AITER Blockscaled GEMM FP8.
Token Generation: Output token throughput increased by 18.6%, showing significantly faster token generation.
Latency Improvements: Both mean and median inter-token latency decreased, resulting in smoother token generation.
P99 Latency: The P99 inter-token latency improved from 2634.65ms to 2445.57ms, showing better worst-case performance.

[Updated] LM Eval accuracy

V1 Engine

Without AITER Block Scaled GEMM FP8

vllm (pretrained=deepseek-ai/DeepSeek-V3,tensor_parallel_size=8,max_model_len=32768,block_size=1,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9477	±	0.0061
		strict-match	5	exact_match	↑	0.9484	±	0.0061

With AITER Block Scaled GEMM FP8

vllm (pretrained=deepseek-ai/DeepSeek-V3,tensor_parallel_size=8,max_model_len=32768,block_size=1,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9469	±	0.0062
		strict-match	5	exact_match	↑	0.9477	±	0.0061

Old info

Throughput

SharedGPT Dataset

With AITER: Throughput: 5.42 requests/s, 2274.81 total tokens/s, 1087.69 output tokens/s
Without AITER: Throughput: 5.23 requests/s, 2196.01 total tokens/s, 1050.01 output tokens/s

Gain 3.5% in SharedGPT Dataset.

Accuracy Test GSM8K

vllm (pretrained=deepseek-ai/DeepSeek-V3,tensor_parallel_size=8,max_model_len=30000,trust_remote_code=True), gen_kwargs: (None), limit: None, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9484	±	0.0061
		strict-match	5	exact_match	↑	0.9492	±	0.0060

Environment Settings

Updates in Dockerfile.rocm_base

Added AITER Package:

AITER_BRANCH: 5a772490486fc9f8c0bc00f292ab94591fd9d28d

Additional Notes

When setting up AITER, it is crucial to use the command git clone --recursive. This is because the package depends on a third-party package (Composable Kernel).
For building and installing the AITER Python package, you must use the PREBUILD_KERNELS=1 flag along with the command python3 setup.py develop. This ensures that all kernels in the AITER package are built successfully.

The following branches were used as references for this integration:
https://github.com/ROCm/vllm/tree/aiter_upstream
https://github.com/ROCm/vllm/tree/aiter_integration_final
https://github.com/ROCm/vllm/tree/deepseek_v3_dev

This PR is part of a larger effort to integrate AITER kernels into vLLM for improved performance on the ROCm platform.

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

github-actions · 2025-03-17T15:33:10Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

… add AITER package Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

…eddedLLM/vllm into aiter-block-gemm-integration

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify · 2025-03-24T15:58:13Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vllmellm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

…gration Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

SageMoore

This generally looks fine. What models are you all using this kernel with? If there are any models that we would like to claim that this kernel supports, please just include lm_eval results in a comment on this PR.

mergify · 2025-03-31T15:29:52Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vllmellm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

…gration

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

mergify · 2025-04-22T10:03:44Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vllmellm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

houseroad · 2025-04-23T00:55:51Z

vllm/model_executor/layers/quantization/utils/fp8_utils.py

-                                       weight_scale,
-                                       block_size,
-                                       output_dtype=input.dtype)
+    # TODO is_shape_supported_by_cutlass is never used,


can we add the github id or create an issue for this for tracking purpose.

@houseroad We have cross-checked with main. It seems they have implemented and removed the comment. Thus, we have removed the TODO comment.

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

SageMoore

In general this looks fine, but let's iron out this "is cutlass supported logic" ironed out before landing.

SageMoore · 2025-05-12T20:33:32Z

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+                                  weight_scale: torch.Tensor,
+                                  input_2d: torch.Tensor) -> bool:
+    if current_platform.is_rocm():
+        # TODO this is never used, as cutlass_block_fp8_supported is False


If that's the case can we just return False? Or move the current_platform.is_rocm() check to apply_w8a8_block_fp8_linear?

@SageMoore I think we could.

Just make sure you have a look at #14397 as it describes how this is currently a bug.

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify · 2025-05-13T11:19:29Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @vllmellm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

SageMoore

Looks reasonable. Thanks for the contribution!

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: Yuqi Zhang <yuqizhang@google.com>

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Signed-off-by: minpeter <kali2005611@gmail.com>

vllmellm added 3 commits March 17, 2025 13:03

add aiter block-scaled gemm

961de9e

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

include AITER enable for rocm platforms in model end to end tests

45efca9

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

modify rocm docker

68a0479

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify bot added the ci/build label Mar 17, 2025

tjtanaa mentioned this pull request Mar 17, 2025

[Feature] [ROCm]: AITER Kernel Integration #14964

Open

53 tasks

remove aiter blockscale flag check from current platform

952e55e

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

vllmellm marked this pull request as ready for review March 18, 2025 05:33

vllmellm requested review from mgoin, robertgshaw2-redhat, DarkLight1337, ywang96 and tlrmchlsmth as code owners March 18, 2025 05:33

vllmellm added 5 commits March 18, 2025 06:48

keep packages versions in rocm docker_base file same as main and only…

70ae131

… add AITER package Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

Merge branch 'aiter-block-gemm-integration' of https://github.com/Emb…

a7025e0

…eddedLLM/vllm into aiter-block-gemm-integration

fix get envs variables in unit tests

a5c2e89

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

remove cascading logic in envs and use CK gemm w8a8 block from AITER

9507aac

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

revert back the unittest to their original format

7f17da7

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify bot added the needs-rebase label Mar 24, 2025

vllmellm added 2 commits March 24, 2025 16:01

Merge remote-tracking branch 'origin/main' into aiter-block-gemm-inte…

5e3e366

…gration Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

revert back unneccessary changes

31fefa3

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify bot removed the needs-rebase label Mar 24, 2025

revert back unneccassary changes

6bed514

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

SageMoore reviewed Mar 31, 2025

View reviewed changes

mergify bot added the needs-rebase label Mar 31, 2025

sync with main

df1ffc3

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

mergify bot removed the needs-rebase label Apr 16, 2025

wrap aiter kernel into direct register custom op

d6cce8f

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

tjtanaa added 2 commits April 22, 2025 07:52

Merge remote-tracking branch 'origin/main' into aiter-block-gemm-inte…

440617a

…gration

refactor the conditional checking logic to reduce checking overhead

55260a2

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

mergify bot added the needs-rebase label Apr 22, 2025

tjtanaa added 2 commits April 22, 2025 13:13

restore dispatcher abstraction andupdate unittest

e06cc3b

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

merge with main, remove AITER BLOCK SCALED GEMM flag from envs.py

f1a2044

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

mergify bot removed the needs-rebase label Apr 22, 2025

houseroad reviewed Apr 23, 2025

View reviewed changes

tjtanaa and others added 5 commits April 23, 2025 01:56

clean up TODO; it is resolved in main

41bf618

Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>

merge with main

65eaed1

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

fix pre-commit

fc97e94

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

fix pre-commit

7c5491f

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

fix pre-commit

a81cd66

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

SageMoore reviewed May 12, 2025

View reviewed changes

clean up use_cutlass logic

5943729

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify bot added the needs-rebase label May 13, 2025

fix pre-commit as main has changed linting rules

d0241c7

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

mergify bot removed the needs-rebase label May 13, 2025

SageMoore approved these changes May 13, 2025

View reviewed changes

DarkLight1337 added the ready ONLY add when PR is ready to merge/full CI is needed label May 13, 2025

DarkLight1337 approved these changes May 13, 2025

View reviewed changes

DarkLight1337 enabled auto-merge (squash) May 13, 2025 14:32

vllm-bot merged commit 40de1ef into vllm-project:main May 14, 2025
84 of 93 checks passed

tjtanaa deleted the aiter-block-gemm-integration branch May 16, 2025 16:28

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

[FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968

[FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968

Uh oh!

vllmellm commented Mar 17, 2025 •

edited by github-actions bot

Loading

Uh oh!

github-actions bot commented Mar 17, 2025

Uh oh!

mergify bot commented Mar 24, 2025

Uh oh!

SageMoore left a comment

Uh oh!

mergify bot commented Mar 31, 2025

Uh oh!

mergify bot commented Apr 22, 2025

Uh oh!

houseroad Apr 23, 2025

Uh oh!

tjtanaa Apr 23, 2025

Uh oh!

SageMoore left a comment

Uh oh!

SageMoore May 12, 2025

Uh oh!

tjtanaa May 13, 2025 •

edited

Loading

Uh oh!

ProExpertProg May 13, 2025

Uh oh!

mergify bot commented May 13, 2025

Uh oh!

SageMoore left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

[FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968

[FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968

Uh oh!

Conversation

vllmellm commented Mar 17, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Implementation

Testing

[Updated] Performance

[Updated] LM Eval accuracy

Throughput

Accuracy Test GSM8K

Environment Settings

Updates in Dockerfile.rocm_base

Additional Notes

Uh oh!

github-actions bot commented Mar 17, 2025

Uh oh!

mergify bot commented Mar 24, 2025

Uh oh!

SageMoore left a comment

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Mar 31, 2025

Uh oh!

mergify bot commented Apr 22, 2025

Uh oh!

houseroad Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

tjtanaa Apr 23, 2025

Choose a reason for hiding this comment

Uh oh!

SageMoore left a comment

Choose a reason for hiding this comment

Uh oh!

SageMoore May 12, 2025

Choose a reason for hiding this comment

Uh oh!

tjtanaa May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ProExpertProg May 13, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented May 13, 2025

Uh oh!

SageMoore left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

vllmellm commented Mar 17, 2025 •

edited by github-actions bot

Loading

tjtanaa May 13, 2025 •

edited

Loading