[FEAT] [AITER] [ROCm] integrate aiter bpreshuffle and ck ops#28837

Draft

vllmellm wants to merge 43 commits intovllm-project:mainfrom

EmbeddedLLM:refactor/aiter_integration

Contributor

vllmellm commented Nov 17, 2025 •

edited by github-actions bot

Loading

Purpose

This PR integrates the Aiter bpreshuffle and ck FP8 operators for ROCm, targeting Per-Token/Per-Channel (PTPC) quantization.

Adds AiterBpreshufflePerTokenFp8ScaledMMLinearKernel and AiterCKPerTokenFp8ScaledMMLinearKernel.
Hooks these kernels into the FP8ScaledMMLinearLayerConfig selection logic.
Triggers the bpreshuffle weight shuffling from compressed_tensors_w8a8_fp8.py.

Test Plan

Testing is performed by launching the vLLM server with specific environment variables to select each kernel and running lm_eval. Including lm_eval and benchmark

Test Result

vllmellm added 30 commits

October 28, 2025 16:26


          first try

974e682

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


          fix int8 path

e54e572

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


          clean up; fix quark path

c05027f

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


          update quark fp8 path; format

c089ea5

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


          Merge branch 'main' into refactor-fp8-linear

38825fc


          reduce logging boilerplate; update fp8 path

423e2a6

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


          reduce kernel init boilerplate

dd00106

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


          update ptpc path; bug fixes

7d36148

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


          revert input scale upper bounds

1f65cd5

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


          format; update fbgemm path

5fbe76b

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


          bug fix

e845035

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


          fix types; reduce boilerplate for int8

d92c23b

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


          Merge remote-tracking branch 'origin/main' into refactor-fp8-linear

8e8218e


          format

4ce0ba2

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


          update unit tests to use ScaledMMLinearKernels

dd5a70e

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


          update modelopt path

52ff537

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


          remove FP8LinearOps

b13c4bb

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


          add missing arg

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


          flash_infer missing out dtype bug fix

a8010c7

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


          prefer QuantKey over ScaledMMLinearQuantStrategy

f5e6cd9

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


          rename flash_infer.py to flashinfer.py

a76f7bb

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


          correct minimum capability req for channelwise torch

f10171c

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


          add missing kernels for cuda dispatch

fb72ec8

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


          reduce test boilerplate

93fb707

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


          fix quant key selection for ct; remove register_paramter calls; format

abf597e

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


          format

aaa0d55

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


          implement apply func in base FP8ScaledMMLinearKernel class

7fb4657

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


          add minimal documentation for torch scaled mm base class

56a05cd

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


          use for loops for fp8 linear layers init in tests

9ff9b44

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


          minor fixes

cfb476f

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

vllmellm added 12 commits

November 7, 2025 12:13


          force kernels for tests

e47d55b

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


          ensure static scales for ChannelWiseTorchScaledMMLinearKernel; remove…

edb6d43

… comment

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


          feat: Integrate AITER bpreshuffle and ck operators on top of fp8 refa…

45a3008

…ctor

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


          fix output padding for torch _scaled_mm

858765f

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


          Merge remote-tracking branch 'origin/refactor-fp8-linear' into refact…

45803e1

…or/aiter_integration


          Merge remote-tracking branch 'origin/main' into refactor-fp8-linear

4c596a0

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


          optional input scales

65ecf48

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


          remove maybe_create_device_identity

405d280

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


          Merge remote-tracking branch 'origin/main' into refactor-fp8-linear

9784a0c

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


          add CPU kernels; fix fp8 quant type selection

10eebd4

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


          Merge remote-tracking branch 'origin/refactor-fp8-linear' into refact…

686b3ec

…or/aiter_integration

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


          WIP: Integrate Aiter bpreshuffle and ck kernels

679a7cf

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

mergify bot added the nvidia label

mergify bot commented Nov 17, 2025

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

mergify bot added rocm needs-rebase labels

github-project-automation bot added this to NVIDIA


          sync upstream

30624ea

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

mergify bot removed the needs-rebase label

kliuae mentioned this pull request

[WIP][ROCm] Add AITER hipblaslt preshuffled gemm kernel #29981

Draft

5 tasks

mergify bot added the cpu label

mergify bot commented Dec 17, 2025

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

mergify bot added the needs-rebase label

tjtanaa mentioned this pull request

[ROCm] Change default settings for ROCm #33271

Closed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

tjtanaa Awaiting requested review from tjtanaa tjtanaa will be requested when the pull request is marked ready for review tjtanaa is a code owner

mgoin Awaiting requested review from mgoin mgoin will be requested when the pull request is marked ready for review mgoin is a code owner

robertgshaw2-redhat Awaiting requested review from robertgshaw2-redhat robertgshaw2-redhat will be requested when the pull request is marked ready for review robertgshaw2-redhat is a code owner

tlrmchlsmth Awaiting requested review from tlrmchlsmth tlrmchlsmth will be requested when the pull request is marked ready for review tlrmchlsmth is a code owner

yewentao256 Awaiting requested review from yewentao256 yewentao256 will be requested when the pull request is marked ready for review yewentao256 is a code owner

pavanimajety Awaiting requested review from pavanimajety pavanimajety will be requested when the pull request is marked ready for review pavanimajety is a code owner

At least 1 approving review is required to merge this pull request.

Labels

cpu needs-rebase nvidia rocm