@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 01413e0cf5
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Fp8 paged attention update (#22222)

Signed-off-by: Xiao Yu <xiao.yu@amd.com> Signed-off-by: xiao-llm <xiao.yu.dc@outlook.com> Co-authored-by: Xiao Yu <xiao.yu@metamaterial.com> Co-authored-by: Xiao Yu <xiao.yu@amd.com> Co-authored-by: Bowen Bao <bowenbao@amd.com>

xiao-llm committed Sep 15, 2025 at 10:43 UTC 01413e0cf5a04da4049ffa38b6ff3df27ccabd06
5 files changed +259 -81
csrc/rocm/attention.cu
+247 -78
@@ -30,6 +30,10 @@
30 #define __HIP__GFX9__
31 #endif
32
33 +#if defined(__HIPCC__) && (defined(__gfx942__) || defined(__gfx950__))
34 + #define __HIP__FP8MFMA__
35 +#endif
36 +
37 #if defined(__HIPCC__) && (defined(__gfx1100__) || defined(__gfx1101__))
38 #define __HIP__GFX11__
39 #endif
@@ -51,6 +55,12 @@
55 #define MIN(a, b) ((a) < (b) ? (a) : (b))
56 #define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
57
58 +enum class MFMAType {
59 + F16 = 0,
60 + Fp8 = 1,
61 + Fp4 = 2,
62 +};
63 +
64 #if defined(__HIP__GFX9__)
65
66 #define GCN_MFMA_INSTR1 __builtin_amdgcn_mfma_f32_16x16x4f32
@@ -112,6 +122,21 @@ __device__ __forceinline__ floatx4 gcn_mfma16x16x16_instr(const _B16x4& inpA,
122 }
123 }
124
125 +template <typename T, int absz, int cbid, int blgp>
126 +__device__ __forceinline__ floatx4 gcn_mfma16x16x32_instr(const long& inpA,
127 + const long& inpB,
128 + const floatx4& inpC) {
129 + if constexpr (std::is_same<T, __hip_fp8_e4m3>::value) {
130 + return __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(inpA, inpB, inpC, absz,
131 + cbid, blgp);
132 + } else if constexpr (std::is_same<T, __hip_fp8_e5m2>::value) {
133 + return __builtin_amdgcn_mfma_f32_16x16x32_bf8_bf8(inpA, inpB, inpC, absz,
134 + cbid, blgp);
135 + } else {
136 + static_assert(false, "unsupported 8b dtype");
137 + }
138 +}
139 +
140 template <typename T>
141 __device__ __forceinline__ float to_float(const T& inp) {
142 if constexpr (std::is_same<T, _Float16>::value) {
@@ -256,12 +281,44 @@ __device__ __forceinline__ _B16x8 convert_b8x8_custom(const _B8x8 input) {
281 return ret;
282 }
283
284 +typedef union u64_cvt {
285 + half f16x4[4];
286 + int16_t b16x4[4];
287 + _B8x8 b8x8;
288 + _B16x4 b64;
289 + int64_t i64;
290 +} _T8x8;
291 +
292 +__device__ __forceinline__ _B8x8 convert_b16x8(const _B16x8& input,
293 + _T8x8& Mtemp) {
294 + _T8x8 Qtmp8x8;
295 +
296 + for (int i = 0; i < 2; i++) {
297 + floatx4 q_out = {0, 0, 0, 0};
298 + q_out = gcn_mfma16x16x16_instr<_Float16, 0, 0, 0>(Mtemp.b64, input.xy[i],
299 + q_out);
300 + Qtmp8x8.b16x4[i * 2] =
301 + __builtin_amdgcn_cvt_pk_fp8_f32(q_out[0], q_out[1], 0, false);
302 + Qtmp8x8.b16x4[i * 2 + 1] =
303 + __builtin_amdgcn_cvt_pk_fp8_f32(q_out[2], q_out[3], 0, false);
304 + }
305 + return Qtmp8x8.b8x8;
306 +}
307 +
308 +__device__ float warpReduceMax(float val) {
309 + for (int offset = warpSize / 2; offset > 0; offset /= 2) {
310 + val = max(
311 + val, __shfl_down(val, offset, WARP_SIZE)); // Using max() for reduction
312 + }
313 + return val;
314 +}
315 +
316 // grid (num_seqs, num_partitions,num_kv_heads)
317 // block (256)
318 // clang-format off
319 template <typename scalar_t, typename cache_t,
320 vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
264 - int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
321 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO, MFMAType MFMA_TYPE>
322 __global__
323 __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
324 const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -367,6 +424,10 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
424 const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
425
426 int kphysical_block_number[TLOOP];
427 + #if defined(__HIP__FP8MFMA__)
428 + float q_max = 0;
429 + float q_scale = 1.0;
430 + #endif
431
432 // fetch k physical block numbers
433 for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
@@ -416,6 +477,15 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
477 Qlocal[qkhe_depth][qkratio].xy[i] =
478 shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO]
479 [2 * qkratio + i];
480 + #if defined(__HIP__FP8MFMA__)
481 + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto &&
482 + MFMA_TYPE == MFMAType::Fp8) {
483 + scalar_t* qptr =
484 + reinterpret_cast<scalar_t*>(&Qlocal[qkhe_depth][qkratio].xy[i]);
485 + for (int k = 0; k < 4; k++)
486 + q_max = fmax(fabs(to_float<scalar_t>(qptr[k])), q_max);
487 + }
488 + #endif
489 }
490 }
491 }
@@ -515,6 +585,14 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
585 if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
586 // multiply by k_scale if fp8 kv cache
587 scale2 *= *k_scale;
588 + #if defined(__HIP__FP8MFMA__)
589 + q_max = warpReduceMax(q_max);
590 + constexpr float FP8_E4M3_SCALE_TARGET = 224.0f;
591 + if constexpr (MFMA_TYPE == MFMAType::Fp8) {
592 + q_scale = q_max > 0 ? FP8_E4M3_SCALE_TARGET / q_max : 1.0f;
593 + scale2 /= q_scale;
594 + }
595 + #endif
596 }
597
598 floatx4 d_out[TLOOP];
@@ -534,12 +612,41 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
612 auto Ktmp = Klocal[token_depth][qkhe_depth];
613 _B8x16 Ktmp8x16 = *reinterpret_cast<_B8x16*>(&Ktmp);
614 for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) {
537 - _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio];
538 - _B16x8 Klocaltmp = convert_b8x8_custom<scalar_t>(Ktmp8x8);
539 - for (int i = 0; i < 2; i++) {
540 - d_out[token_depth] = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
541 - Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i],
542 - d_out[token_depth]);
615 + if constexpr (MFMA_TYPE == MFMAType::F16) {
616 + _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio];
617 + _B16x8 Klocaltmp = convert_b8x8_custom<scalar_t>(Ktmp8x8);
618 + for (int i = 0; i < 2; i++) {
619 + d_out[token_depth] = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
620 + Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i],
621 + d_out[token_depth]);
622 + }
623 + } else {
624 + #if defined(__HIP__FP8MFMA__)
625 + _T8x8 Ktmp8x8, Qtmp8x8;
626 + Ktmp8x8.b8x8 = Ktmp8x16.xy[qkratio];
627 +
628 + for (int n = 0; n < 2; n++) {
629 + scalar_t* qptr = reinterpret_cast<scalar_t*>(
630 + &Qlocal[qkhe_depth][qkratio].xy[n]);
631 +
632 + Qtmp8x8.b16x4[n * 2] =
633 + vllm::fp8::scaled_vec_conversion<uint16_t, float2>(
634 + make_float2(to_float<scalar_t>(qptr[0]),
635 + to_float<scalar_t>(qptr[1])),
636 + q_scale);
637 + Qtmp8x8.b16x4[n * 2 + 1] =
638 + vllm::fp8::scaled_vec_conversion<uint16_t, float2>(
639 + make_float2(to_float<scalar_t>(qptr[2]),
640 + to_float<scalar_t>(qptr[3])),
641 + q_scale);
642 + }
643 +
644 + d_out[token_depth] =
645 + gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>(
646 + Ktmp8x8.i64, Qtmp8x8.i64, d_out[token_depth]);
647 + #else
648 + UNREACHABLE_CODE
649 + #endif
650 }
651 }
652 }
@@ -629,17 +736,36 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
736 // disable rtz conversion due to its impact on accuracy.
737 constexpr bool LOGITS_RTZ_CONVERSION = false;
738
739 + #if defined(__HIP__FP8MFMA__)
740 + int rowid_8x8 = rowid / 2;
741 + int offset = rowid % 2;
742 + #endif
743 +
744 // write logits to shared mem
745 for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
746 d_out[token_depth] *= inv_sum_scale;
635 - if constexpr (LOGITS_RTZ_CONVERSION) {
636 - // use rtz conversion for better performance, with negligible impact on
637 - // accuracy
638 - shared_logits[warpid][token_depth][lane16id][rowid] =
639 - from_floatx4_rtz<scalar_t>(d_out[token_depth]);
747 + if constexpr (MFMA_TYPE != MFMAType::Fp8) {
748 + if constexpr (LOGITS_RTZ_CONVERSION) {
749 + // use rtz conversion for better performance, with negligible impact on
750 + // accuracy
751 + shared_logits[warpid][token_depth][lane16id][rowid] =
752 + from_floatx4_rtz<scalar_t>(d_out[token_depth]);
753 + } else {
754 + shared_logits[warpid][token_depth][lane16id][rowid] =
755 + from_floatx4<scalar_t>(d_out[token_depth]);
756 + }
757 } else {
641 - shared_logits[warpid][token_depth][lane16id][rowid] =
642 - from_floatx4<scalar_t>(d_out[token_depth]);
758 + #if defined(__HIP__FP8MFMA__)
759 + // cast _B16x4* to _B8x8*
760 + _T8x8& logits_8x8 = *reinterpret_cast<_T8x8*>(
761 + &shared_logits[warpid][token_depth][lane16id][rowid_8x8]);
762 + logits_8x8.b16x4[offset * 2] = __builtin_amdgcn_cvt_pk_fp8_f32(
763 + d_out[token_depth][0], d_out[token_depth][1], 0, false);
764 + logits_8x8.b16x4[offset * 2 + 1] = __builtin_amdgcn_cvt_pk_fp8_f32(
765 + d_out[token_depth][2], d_out[token_depth][3], 0, false);
766 + #else
767 + UNREACHABLE_CODE
768 + #endif
769 }
770 }
771
@@ -692,19 +818,42 @@ __launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
818 _B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp);
819 for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) {
820 _B8x8 Vtmp8x8 = Vtmp8x16.xy[j];
695 - _B16x8 Vlocaltmp = convert_b8x8_custom<scalar_t>(Vtmp8x8);
696 - for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
697 - const int offset =
698 - rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
699 - j * ELEMS8_ELEMS4_RATIO + i;
700 - const int offset1 = offset % ROWS_PER_WARP;
701 - const int offset2 = offset / ROWS_PER_WARP;
702 - // output format is 16 qheads across 16 lanes, 16 head elems
703 - // spread across 4 rows
704 - tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
705 - Vlocaltmp.xy[i],
706 - shared_logits[vtoken_depth][offset2][lane16id][offset1],
707 - tmp_out);
821 + if constexpr (MFMA_TYPE == MFMAType::F16) {
822 + _B16x8 Vlocaltmp = convert_b8x8_custom<scalar_t>(Vtmp8x8);
823 + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
824 + const int offset =
825 + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
826 + j * ELEMS8_ELEMS4_RATIO + i;
827 + const int offset1 = offset % ROWS_PER_WARP;
828 + const int offset2 = offset / ROWS_PER_WARP;
829 + // output format is 16 qheads across 16 lanes, 16 head elems
830 + // spread across 4 rows
831 + tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
832 + Vlocaltmp.xy[i],
833 + shared_logits[vtoken_depth][offset2][lane16id][offset1],
834 + tmp_out);
835 + }
836 + } else {
837 + #if defined(__HIP__FP8MFMA__)
838 + for (int i = 0; i < ELEMS8_ELEMS4_RATIO / 2; i++) {
839 + const int offset =
840 + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
841 + j * ELEMS8_ELEMS4_RATIO + i;
842 + const int offset1 = (offset % ROWS_PER_WARP) / 2;
843 + const int offset2 = offset / ROWS_PER_WARP;
844 + // output format is 16 qheads across 16 lanes, 16 head elems
845 + // spread across 4 rows
846 + tmp_out = gcn_mfma16x16x32_instr<__hip_fp8_e4m3, 0, 0, 0>(
847 + reinterpret_cast<_T8x8*>(&Vtmp8x8)->i64,
848 + reinterpret_cast<_T8x8*>(
849 + &shared_logits[vtoken_depth][offset2][lane16id]
850 + [offset1])
851 + ->i64,
852 + tmp_out);
853 + }
854 + #else
855 + UNREACHABLE_CODE
856 + #endif
857 }
858 }
859 }
@@ -1570,7 +1719,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
1719 // clang-format off
1720 template <typename scalar_t, typename cache_t,
1721 vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
1573 - int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
1722 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
1723 + MFMAType MFMA_TYPE>
1724 __global__
1725 __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
1726 const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -2337,7 +2487,8 @@ __device__ __forceinline__ _B16x8 from_floatx8(const floatx8& inp) {
2487 // clang-format off
2488 template <typename scalar_t, typename cache_t,
2489 vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
2340 - int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
2490 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO,
2491 + MFMAType MFMA_TYPE>
2492 __global__
2493 __launch_bounds__(NUM_THREADS, 3) void paged_attention_ll4mi_QKV_mfma16_kernel(
2494 const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -2969,7 +3120,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3120 template <typename scalar_t, typename cache_t,
3121 vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
3122 int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
2972 - int GQA_RATIO>
3123 + int GQA_RATIO, MFMAType MFMA_TYPE>
3124 __global__
3125 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
3126 const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
@@ -3041,7 +3192,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3192 #define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \
3193 paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
3194 HEAD_SIZE, NTHR, ALIBI_ENABLED, \
3044 - GQA_RATIO> \
3195 + GQA_RATIO, MFMA_TYPE> \
3196 <<<grid, block, 0, stream>>>( \
3197 query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
3198 block_tables_ptr, seq_lens_ptr, query_start_loc_ptr, \
@@ -3069,7 +3220,7 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
3220
3221 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
3222 int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
3072 - bool ALIBI_ENABLED>
3223 + bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
3224 void paged_attention_custom_launcher(
3225 torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
3226 torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3225,7 +3376,7 @@ void paged_attention_custom_launcher(
3376
3377 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
3378 int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
3228 - bool ALIBI_ENABLED>
3379 + bool ALIBI_ENABLED, MFMAType MFMA_TYPE>
3380 void paged_attention_custom_launcher_navi(
3381 torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
3382 torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -3397,74 +3548,77 @@ void paged_attention_custom_launcher_navi(
3548 }
3549
3550 #define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, \
3400 - PSIZE, ALIBI_ENABLED) \
3551 + PSIZE, ALIBI_ENABLED, MFMA_TYPE) \
3552 if (!is_navi) { \
3553 paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3403 - OUTT, PSIZE, ALIBI_ENABLED>( \
3554 + OUTT, PSIZE, ALIBI_ENABLED, MFMA_TYPE>( \
3555 out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
3556 num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
3557 max_seq_len, alibi_slopes, k_scale, v_scale, fp8_out_scale); \
3558 } else { \
3408 - paged_attention_custom_launcher_navi< \
3409 - T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, ALIBI_ENABLED>( \
3559 + paged_attention_custom_launcher_navi<T, KVT, KV_DTYPE, BLK_SIZE, \
3560 + HEAD_SIZE, OUTT, PSIZE, \
3561 + ALIBI_ENABLED, MFMA_TYPE>( \
3562 out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
3563 num_kv_heads, scale, block_tables, seq_lens, query_start_loc, \
3564 max_seq_len, alibi_slopes, k_scale, v_scale); \
3565 }
3566
3567 #define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3416 - OUTT, PSIZE) \
3568 + OUTT, PSIZE, MFMA_TYPE) \
3569 if (alibi_slopes) { \
3570 CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \
3419 - true); \
3571 + true, MFMA_TYPE); \
3572 } else { \
3573 CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, OUTT, PSIZE, \
3422 - false); \
3574 + false, MFMA_TYPE); \
3575 }
3576
3577 #if defined(__HIPCC__) && defined(__gfx90a__)
3426 - #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE) \
3578 + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3579 + MFMA_TYPE) \
3580 if (fp8_out_scale) { \
3581 TORCH_CHECK(false, "fp8 out scale unsupported for gfx90a"); \
3582 } else { \
3583 CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \
3431 - 256); \
3584 + 256, MFMA_TYPE); \
3585 }
3586 #else
3434 - #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE) \
3587 + #define CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3588 + MFMA_TYPE) \
3589 if (fp8_out_scale) { \
3590 CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
3437 - uint8_t, 256); \
3591 + uint8_t, 256, MFMA_TYPE); \
3592 } else { \
3593 CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \
3440 - 256); \
3594 + 256, MFMA_TYPE); \
3595 }
3596 #endif
3597
3444 -#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE) \
3445 - switch (block_size) { \
3446 - case 16: \
3447 - CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE); \
3448 - break; \
3449 - case 32: \
3450 - CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE); \
3451 - break; \
3452 - default: \
3453 - TORCH_CHECK(false, "Unsupported block size: ", block_size); \
3454 - break; \
3455 - }
3456 -
3457 -#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE) \
3458 - switch (head_size) { \
3459 - case 64: \
3460 - CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64); \
3461 - break; \
3462 - case 128: \
3463 - CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128); \
3464 - break; \
3465 - default: \
3466 - TORCH_CHECK(false, "Unsupported head size: ", head_size); \
3467 - break; \
3598 +#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE, MFMA_TYPE) \
3599 + switch (block_size) { \
3600 + case 16: \
3601 + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 16, HEAD_SIZE, MFMA_TYPE); \
3602 + break; \
3603 + case 32: \
3604 + CALL_CUSTOM_LAUNCHER_OUT(T, KVT, KV_DTYPE, 32, HEAD_SIZE, MFMA_TYPE); \
3605 + break; \
3606 + default: \
3607 + TORCH_CHECK(false, "Unsupported block size: ", block_size); \
3608 + break; \
3609 + }
3610 +
3611 +#define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE, MFMA_TYPE) \
3612 + switch (head_size) { \
3613 + case 64: \
3614 + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 64, MFMA_TYPE); \
3615 + break; \
3616 + case 128: \
3617 + CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, 128, MFMA_TYPE); \
3618 + break; \
3619 + default: \
3620 + TORCH_CHECK(false, "Unsupported head size: ", head_size); \
3621 + break; \
3622 }
3623
3624 bool is_navi_gpu() {
@@ -3503,28 +3657,43 @@ void paged_attention(
3657 const std::optional<torch::Tensor>& alibi_slopes,
3658 const std::string& kv_cache_dtype, torch::Tensor& k_scale,
3659 torch::Tensor& v_scale,
3506 - const std::optional<torch::Tensor>& fp8_out_scale) {
3660 + const std::optional<torch::Tensor>& fp8_out_scale,
3661 + const std::string& mfma_type) {
3662 // clang-format on
3663 bool is_navi = is_navi_gpu();
3509 -
3664 const int head_size = query.size(2);
3665 if (kv_cache_dtype == "auto") {
3666 if (query.dtype() == at::ScalarType::Half) {
3513 - CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, _Float16,
3514 - vllm::Fp8KVCacheDataType::kAuto);
3667 + CALL_CUSTOM_LAUNCHER_BLK_HEAD(
3668 + _Float16, _Float16, vllm::Fp8KVCacheDataType::kAuto, MFMAType::F16);
3669 } else if (query.dtype() == at::ScalarType::BFloat16) {
3670 CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, __hip_bfloat16,
3517 - vllm::Fp8KVCacheDataType::kAuto);
3671 + vllm::Fp8KVCacheDataType::kAuto,
3672 + MFMAType::F16);
3673 } else {
3674 TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
3675 }
3676 } else if (kv_cache_dtype == "fp8" || kv_cache_dtype == "fp8_e4m3") {
3677 if (query.dtype() == at::ScalarType::Half) {
3523 - CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
3524 - vllm::Fp8KVCacheDataType::kFp8E4M3);
3678 + if (mfma_type == "fp8") {
3679 + CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
3680 + vllm::Fp8KVCacheDataType::kFp8E4M3,
3681 + MFMAType::Fp8);
3682 + } else {
3683 + CALL_CUSTOM_LAUNCHER_BLK_HEAD(_Float16, uint8_t,
3684 + vllm::Fp8KVCacheDataType::kFp8E4M3,
3685 + MFMAType::F16);
3686 + }
3687 } else if (query.dtype() == at::ScalarType::BFloat16) {
3526 - CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
3527 - vllm::Fp8KVCacheDataType::kFp8E4M3);
3688 + if (mfma_type == "fp8") {
3689 + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
3690 + vllm::Fp8KVCacheDataType::kFp8E4M3,
3691 + MFMAType::Fp8);
3692 + } else {
3693 + CALL_CUSTOM_LAUNCHER_BLK_HEAD(__hip_bfloat16, uint8_t,
3694 + vllm::Fp8KVCacheDataType::kFp8E4M3,
3695 + MFMAType::F16);
3696 + }
3697 } else {
3698 TORCH_CHECK(false, "Unsupported data type: ", query.dtype());
3699 }
csrc/rocm/ops.h
+2 -1
@@ -19,4 +19,5 @@ void paged_attention(
19 const std::optional<torch::Tensor>& query_start_loc, int64_t block_size,
20 int64_t max_seq_len, const std::optional<torch::Tensor>& alibi_slopes,
21 const std::string& kv_cache_dtype, torch::Tensor& k_scale,
22 - torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale);
22 + torch::Tensor& v_scale, const std::optional<torch::Tensor>& fp8_out_scale,
23 + const std::string& mfma_type);
csrc/rocm/torch_bindings.cpp
+2 -1
@@ -48,7 +48,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, rocm_ops) {
48 " Tensor? alibi_slopes,"
49 " str kv_cache_dtype,"
50 " Tensor k_scale, Tensor v_scale,"
51 - " Tensor? fp8_out_scale) -> ()");
51 + " Tensor? fp8_out_scale,"
52 + " str mfma_type) -> ()");
53 rocm_ops.impl("paged_attention", torch::kCUDA, &paged_attention);
54 }
55
vllm/_custom_ops.py
+2 -1
@@ -117,13 +117,14 @@ def paged_attention_rocm(
117 k_scale: torch.Tensor,
118 v_scale: torch.Tensor,
119 fp8_out_scale: Optional[torch.Tensor] = None,
120 + mfma_type: str = "fp8" if envs.VLLM_ROCM_FP8_MFMA_PAGE_ATTN else "f16",
121 ) -> None:
122 torch.ops._rocm_C.paged_attention(out, exp_sum, max_logits, tmp_out, query,
123 key_cache, value_cache, num_kv_heads,
124 scale, block_tables, seq_lens,
125 query_start_loc, block_size, max_seq_len,
126 alibi_slopes, kv_cache_dtype, k_scale,
126 - v_scale, fp8_out_scale)
127 + v_scale, fp8_out_scale, mfma_type)
128
129
130 def mla_decode_kvcache_cpu(
vllm/envs.py
+6
@@ -167,6 +167,7 @@ if TYPE_CHECKING:
167 VLLM_HAS_FLASHINFER_CUBIN: bool = False
168 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False
169 VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: bool = False
170 + VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
171 VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
172 VLLM_ALLREDUCE_USE_SYMM_MEM: bool = False
173 VLLM_TUNED_CONFIG_FOLDER: Optional[str] = None
@@ -1219,6 +1220,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1220 "VLLM_ENABLE_RESPONSES_API_STORE":
1221 lambda: bool(int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0"))),
1222
1223 + # If set, use the fp8 mfma in rocm paged attention.
1224 + "VLLM_ROCM_FP8_MFMA_PAGE_ATTN":
1225 + lambda: bool(int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0"))),
1226 +
1227 # Whether to use pytorch symmetric memory for allreduce
1228 "VLLM_ALLREDUCE_USE_SYMM_MEM":
1229 lambda: bool(int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "0"))),
@@ -1340,6 +1345,7 @@ def compute_hash() -> str:
1345 "VLLM_ROCM_QUICK_REDUCE_QUANTIZATION",
1346 "VLLM_ROCM_QUICK_REDUCE_CAST_BF16_TO_FP16",
1347 "VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB",
1348 + "VLLM_ROCM_FP8_MFMA_PAGE_ATTN",
1349 ]
1350 for key in environment_variables_to_hash:
1351 # if this goes out of sync with environment_variables,