Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[WIP][AMD][Kernel][Quantization] Add fp8 and int8 support for Triton FAv2 kernel #12534

Draft
wants to merge 21 commits into
base: main
Choose a base branch
from
+2,082 −649
Draft

[WIP][AMD][Kernel][Quantization] Add fp8 and int8 support for Triton FAv2 kernel #12534

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
21 commits
Select commit Hold shift + click to select a range
9e8bad6
TritonScaledMMLinearKernel implementation
rasmith Jan 21, 2025
daf9a71
Add regression test for rocm w8a8
rasmith Jan 22, 2025
9c11d5c
remote unused import
rasmith Jan 22, 2025
4e4d633
ruff
rasmith Jan 22, 2025
16cfc92
adding rocm fp8 support
rasmith Jan 28, 2025
5e0ada7
first successful execution
rasmith Jan 28, 2025
481ea60
remove prints
rasmith Jan 29, 2025
2de123f
Merge upstream
rasmith Jan 29, 2025
0369a6d
typo fix
rasmith Jan 29, 2025
4ff10e4
typo fix
rasmith Jan 29, 2025
b5f045d
typo fix
rasmith Jan 29, 2025
58ea057
typo fix
rasmith Jan 29, 2025
3ac3f08
typo fix
rasmith Jan 29, 2025
00cae8d
fix merge
rasmith Jan 29, 2025
a322dc6
fix compiler error
rasmith Jan 29, 2025
910d16d
Revert "Custom PA perf improvements (#222)"
rasmith Feb 1, 2025
29acd8e
still trying to fix
rasmith Feb 3, 2025
aac208d
compilation fixes:
rasmith Feb 4, 2025
0b8d7ae
seems to work now
rasmith Feb 4, 2025
69140e6
new rocm_flash_attn.py
rasmith Feb 7, 2025
07dabb5
Remove fp8_out_scale and apply FP8 FA fixes from Ilia
rasmith Feb 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions csrc/activation_kernels.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@
#include "cuda_compat.h"
#include "dispatch_utils.h"

#ifdef USE_ROCM
#include "quantization/fp8/amd/hip_float8.h"
#endif

namespace vllm {

template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
Expand Down
2 changes: 1 addition & 1 deletion csrc/attention/paged_attention_v1.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,4 +193,4 @@ void paged_attention_v1(
#undef WARP_SIZE
#undef MAX
#undef MIN
#undef DIVIDE_ROUND_UP
#undef DIVIDE_ROUND_UP
2 changes: 1 addition & 1 deletion csrc/attention/paged_attention_v2.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -203,4 +203,4 @@ void paged_attention_v2(
#undef WARP_SIZE
#undef MAX
#undef MIN
#undef DIVIDE_ROUND_UP
#undef DIVIDE_ROUND_UP
Loading