@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 848a6438ae
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 15, 21:00 UTC

[ROCm] Faster Custom Paged Attention kernels (#12348)

TJian committed Mar 4, 2025 at 01:24 UTC 848a6438aed2ef8d0de3d51e8ed2f970abf57853
6 files changed +1145 -447
.buildkite/run-amd-test.sh
-1
@@ -77,7 +77,6 @@ echo "Commands:$commands"
77 #ignore certain kernels tests
78 if [[ $commands == *" kernels "* ]]; then
79 commands="${commands} \
80 - --ignore=kernels/test_attention.py \
80 --ignore=kernels/test_attention_selector.py \
81 --ignore=kernels/test_blocksparse_attention.py \
82 --ignore=kernels/test_causal_conv1d.py \
benchmarks/kernels/benchmark_paged_attention.py
+51 -20
@@ -11,8 +11,9 @@ from vllm.platforms import current_platform
11 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser,
12 create_kv_caches_with_random)
13
14 -NUM_BLOCKS = 1024
14 +NUM_BLOCKS = 128 * 1024
15 PARTITION_SIZE = 512
16 +PARTITION_SIZE_ROCM = 256
17
18
19 @torch.inference_mode()
@@ -80,6 +81,12 @@ def main(
81 # Prepare for the paged attention kernel.
82 output = torch.empty_like(query)
83 if version == "v2":
84 + if current_platform.is_rocm():
85 + global PARTITION_SIZE
86 + if not args.custom_paged_attn:
87 + PARTITION_SIZE = 1024
88 + else:
89 + PARTITION_SIZE = PARTITION_SIZE_ROCM
90 num_partitions = ((max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE)
91 tmp_output = torch.empty(
92 size=(num_seqs, num_query_heads, num_partitions, head_size),
@@ -123,25 +130,46 @@ def main(
130 v_scale,
131 )
132 elif version == "v2":
126 - ops.paged_attention_v2(
127 - output,
128 - exp_sums,
129 - max_logits,
130 - tmp_output,
131 - query,
132 - key_cache,
133 - value_cache,
134 - num_kv_heads,
135 - scale,
136 - block_tables,
137 - seq_lens,
138 - block_size,
139 - max_seq_len,
140 - alibi_slopes,
141 - kv_cache_dtype,
142 - k_scale,
143 - v_scale,
144 - )
133 + if not args.custom_paged_attn:
134 + ops.paged_attention_v2(
135 + output,
136 + exp_sums,
137 + max_logits,
138 + tmp_output,
139 + query,
140 + key_cache,
141 + value_cache,
142 + num_kv_heads,
143 + scale,
144 + block_tables,
145 + seq_lens,
146 + block_size,
147 + max_seq_len,
148 + alibi_slopes,
149 + kv_cache_dtype,
150 + k_scale,
151 + v_scale,
152 + )
153 + else:
154 + ops.paged_attention_rocm(
155 + output,
156 + exp_sums,
157 + max_logits,
158 + tmp_output,
159 + query,
160 + key_cache,
161 + value_cache,
162 + num_kv_heads,
163 + scale,
164 + block_tables,
165 + seq_lens,
166 + block_size,
167 + max_seq_len,
168 + alibi_slopes,
169 + kv_cache_dtype,
170 + k_scale,
171 + v_scale,
172 + )
173 else:
174 raise ValueError(f"Invalid version: {version}")
175 torch.cuda.synchronize()
@@ -195,6 +223,9 @@ if __name__ == '__main__':
223 help="Data type for kv cache storage. If 'auto', will use model "
224 "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. "
225 "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)")
226 + parser.add_argument("--custom-paged-attn",
227 + action="store_true",
228 + help="Use custom paged attention")
229 args = parser.parse_args()
230 print(args)
231
csrc/rocm/attention.cu
+1084 -422
@@ -17,6 +17,7 @@
17 #include <torch/all.h>
18 #include <ATen/cuda/CUDAContext.h>
19 #include <c10/cuda/CUDAGuard.h>
20 +#include <hip/hip_fp8.h>
21 #include <hip/hip_bf16.h>
22 #include "cuda_compat.h"
23
@@ -50,6 +51,9 @@ using floatx4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;
51 using float16x4 =
52 __attribute__((__vector_size__(4 * sizeof(_Float16)))) _Float16;
53 typedef float16x4 _Half4;
54 +using float16x2 =
55 + __attribute__((__vector_size__(2 * sizeof(_Float16)))) _Float16;
56 +typedef float16x2 _Half2;
57 typedef struct _Half8 {
58 _Half4 xy[2];
59 } _Half8;
@@ -62,23 +66,17 @@ typedef struct _B16x8 {
66 } _B16x8;
67
68 using _B8x8 = uint2;
69 +using _B8x4 = int32_t; // used in builtins
70 +using bit8_t = uint8_t;
71
66 -////// Non temporal load stores ///////
67 -
68 -template <typename T>
69 -__device__ __forceinline__ T load(T* addr) {
70 - return addr[0];
71 -}
72 -
73 -template <typename T>
74 -__device__ __forceinline__ void store(T value, T* addr) {
75 - addr[0] = value;
76 -}
72 +typedef struct _B8x16 {
73 + _B8x8 xy[2];
74 +} _B8x16;
75
76 template <typename T, int absz, int cbid, int blgp>
79 -__device__ __forceinline__ floatx4 gcn_mfma_instr(const _B16x4& inpA,
80 - const _B16x4& inpB,
81 - const floatx4& inpC) {
77 +__device__ __forceinline__ floatx4 gcn_mfma4x4x4_instr(const _B16x4& inpA,
78 + const _B16x4& inpB,
79 + const floatx4& inpC) {
80 if constexpr (std::is_same<T, _Float16>::value) {
81 return __builtin_amdgcn_mfma_f32_4x4x4f16(inpA, inpB, inpC, absz, cbid,
82 blgp);
@@ -90,6 +88,21 @@ __device__ __forceinline__ floatx4 gcn_mfma_instr(const _B16x4& inpA,
88 }
89 }
90
91 +template <typename T, int absz, int cbid, int blgp>
92 +__device__ __forceinline__ floatx4 gcn_mfma16x16x16_instr(const _B16x4& inpA,
93 + const _B16x4& inpB,
94 + const floatx4& inpC) {
95 + if constexpr (std::is_same<T, _Float16>::value) {
96 + return __builtin_amdgcn_mfma_f32_16x16x16f16(inpA, inpB, inpC, absz, cbid,
97 + blgp);
98 + } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
99 + return __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(inpA, inpB, inpC, absz,
100 + cbid, blgp);
101 + } else {
102 + static_assert(false, "unsupported 16b dtype");
103 + }
104 +}
105 +
106 template <typename T>
107 __device__ __forceinline__ float to_float(const T& inp) {
108 if constexpr (std::is_same<T, _Float16>::value) {
@@ -121,17 +134,22 @@ __device__ __forceinline__ _B16x4 from_floatx4(const floatx4& inp) {
134 } t16;
135 _B16x4 ret;
136 if constexpr (std::is_same<T, _Float16>::value) {
124 - #pragma unroll
125 - for (int i = 0; i < 4; i++) {
126 - t16.f = (_Float16)inp[i];
127 - ret[i] = t16.u;
128 - }
129 - return ret;
137 + union h2cvt {
138 + __half2 h2[2];
139 + _B16x4 b16x4;
140 + } u;
141 + u.h2[0] = __float22half2_rn(make_float2(inp[0], inp[1]));
142 + u.h2[1] = __float22half2_rn(make_float2(inp[2], inp[3]));
143 + return u.b16x4;
144 } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
131 - #pragma unroll
145 for (int i = 0; i < 4; i++) {
133 - t16.b = __float2bfloat16(inp[i]);
134 - ret[i] = t16.u;
146 + union fcvt {
147 + uint32_t u32;
148 + float f32;
149 + } u;
150 + u.f32 = inp[i];
151 + u.u32 += 0x7fff + ((u.u32 >> 16) & 1); // BF16 RNE with no nan/inf check
152 + ret[i] = uint16_t(u.u32 >> 16);
153 }
154 return ret;
155 } else {
@@ -149,21 +167,25 @@ __device__ __forceinline__ _B16x4 addx4(const _B16x4& inp1,
167 } t1, t2, res;
168 _B16x4 ret;
169 if constexpr (std::is_same<T, _Float16>::value) {
152 - #pragma unroll
153 - for (int i = 0; i < 4; i++) {
154 - t1.u = inp1[i];
155 - t2.u = inp2[i];
156 - res.f = t1.f + t2.f;
157 - ret[i] = res.u;
158 - }
159 - return ret;
170 + union h2cvt {
171 + _B16x4 b16x4;
172 + __half2 h2[2];
173 + } u1, u2, s;
174 + u1.b16x4 = inp1;
175 + u2.b16x4 = inp2;
176 + s.h2[0] = u1.h2[0] + u2.h2[0];
177 + s.h2[1] = u1.h2[1] + u2.h2[1];
178 + return s.b16x4;
179 } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
161 - #pragma unroll
180 for (int i = 0; i < 4; i++) {
163 - t1.u = inp1[i];
164 - t2.u = inp2[i];
165 - res.b = t1.b + t2.b;
166 - ret[i] = res.u;
181 + union fcvt {
182 + float f32;
183 + uint32_t i32;
184 + } u1, u2, s;
185 + u1.i32 = uint32_t(inp1[i]) << 16;
186 + u2.i32 = uint32_t(inp2[i]) << 16;
187 + s.f32 = u1.f32 + u2.f32;
188 + ret[i] = uint16_t(s.i32 >> 16);
189 }
190 return ret;
191 } else {
@@ -171,53 +193,600 @@ __device__ __forceinline__ _B16x4 addx4(const _B16x4& inp1,
193 }
194 }
195
174 -template <typename T, vllm::Fp8KVCacheDataType KV_DTYPE>
175 -__device__ __forceinline__ _B16x8 scaled_convert_b8x8(const _B8x8 input,
176 - const float scale) {
177 - union alignas(16) {
178 - uint4 u4;
179 - _B16x8 u16x8;
180 - vllm::bf16_8_t b16x8;
181 - } tmp;
196 +__device__ __forceinline__ floatx4 to_float_fp8x4(const _B8x4& inp) {
197 + // From MI300+ platforms, we have v_cvt_pk_f32_fp8 instruction
198 + // to convert 2 packed fp8 to 2 packed fp32 values.
199 + // However, in MI200 platforms, we only have v_cvt_f32_fp8
200 + // to convert fp8 values individually. So we added
201 + // #else case for fewer instructions (# inst=2) in MI300+,
202 + // and fallback to
203 + // #if case for other platforms (# inst=4).
204 + #if defined(__gfx90a__)
205 + float4 f32x4 = vllm::fp8::vec_conversion<float4, uint32_t>(
206 + *reinterpret_cast<const uint32_t*>(&inp));
207 + return *reinterpret_cast<floatx4*>(&f32x4);
208 + #else // MI3xx+ optimized builtins
209 + const auto f0 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, false);
210 + const auto f1 = __builtin_amdgcn_cvt_pk_f32_fp8(inp, true);
211 + floatx4 ret;
212 + ret[0] = f0[0];
213 + ret[1] = f0[1];
214 + ret[2] = f1[0];
215 + ret[3] = f1[1];
216 + return ret;
217 + #endif
218 +}
219 +
220 +template <typename T>
221 +__device__ __forceinline__ _B16x4 from_floatx4_rtz(const floatx4& inp) {
222 + _B16x4 ret;
223 if constexpr (std::is_same<T, _Float16>::value) {
183 - tmp.u4 = vllm::fp8::scaled_convert<uint4, _B8x8, KV_DTYPE>(input, scale);
184 - return tmp.u16x8;
224 + union h2cvt {
225 + _Half2 h2[2];
226 + _B16x4 b16x4;
227 + } u;
228 + u.h2[0] = __builtin_amdgcn_cvt_pkrtz(inp[0], inp[1]);
229 + u.h2[1] = __builtin_amdgcn_cvt_pkrtz(inp[2], inp[3]);
230 + return u.b16x4;
231 } else if constexpr (std::is_same<T, __hip_bfloat16>::value) {
186 - tmp.b16x8 = vllm::fp8::scaled_convert<vllm::bf16_8_t, _B8x8, KV_DTYPE>(
187 - input, scale);
188 - return tmp.u16x8;
232 + for (int i = 0; i < 4; i++) {
233 + union fcvt {
234 + uint32_t i32;
235 + float f32;
236 + } u;
237 + u.f32 = inp[i];
238 + ret[i] = uint16_t(u.i32 >> 16);
239 + }
240 + return ret;
241 } else {
242 static_assert(false, "unsupported 16b dtype");
243 }
244 }
245
194 -///////////////////////////////////////
246 +template <typename T>
247 +__device__ __forceinline__ _B16x8 convert_b8x8_custom(const _B8x8 input) {
248 + union {
249 + _B8x8 b8x8;
250 + _B8x4 b8x4[2];
251 + } tmp;
252 + tmp.b8x8 = input;
253 + _B16x8 ret;
254 + for (int i = 0; i < 2; i++) {
255 + ret.xy[i] = from_floatx4_rtz<T>(to_float_fp8x4(tmp.b8x4[i]));
256 + }
257 + return ret;
258 +}
259 +
260 +// grid (num_seqs, num_partitions,num_kv_heads)
261 +// block (256)
262 +// clang-format off
263 +template <typename scalar_t, typename cache_t,
264 + vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
265 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED, int GQA_RATIO>
266 +__global__
267 +__launch_bounds__(NUM_THREADS, 5) void paged_attention_ll4mi_QKV_mfma16_kernel(
268 + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
269 + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
270 + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
271 + const int num_kv_heads,
272 + const float scale,
273 + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
274 + const int* __restrict__ context_lens, // [num_seqs]
275 + const int max_num_blocks_per_seq,
276 + const float* __restrict__ alibi_slopes, // [num_heads]
277 + const int q_stride,
278 + const int kv_block_stride,
279 + const int kv_head_stride,
280 + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
281 + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
282 + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
283 + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
284 + int max_ctx_blocks, const float* k_scale, const float* v_scale) {
285 + // clang-format on
286 + constexpr int NWARPS = NUM_THREADS / WARP_SIZE;
287 + const int warpid = threadIdx.x / WARP_SIZE;
288 + const int laneid = threadIdx.x % WARP_SIZE;
289 + const int lane4id = laneid % 4;
290 + const int lane16id = laneid % 16;
291 + const int rowid = laneid / 16;
292 +
293 + const int seq_idx = blockIdx.x;
294 + const int partition_idx = blockIdx.y;
295 +
296 + constexpr int T_PAR_SIZE = 256; // token partition size set to 256
297 +
298 + const int max_num_partitions = gridDim.y;
299 +
300 + const int context_len = context_lens[seq_idx];
301 +
302 + const int partition_start_token_idx =
303 + partition_idx * T_PAR_SIZE; // partition_size;
304 + // exit if partition is out of context for seq
305 + if (partition_start_token_idx >= context_len) {
306 + return;
307 + }
308 +
309 + constexpr int GQA_RATIO4 = DIVIDE_ROUND_UP(GQA_RATIO, 4);
310 +
311 + __shared__ float shared_qk_max[NWARPS][16 + 1];
312 + __shared__ float shared_exp_sum[NWARPS][16 + 1];
313 + // shared_logits is used for multiple purposes
314 + __shared__ _B16x4 shared_logits[NWARPS][4][16][4];
315 +
316 + // for QK mfma16x16, layout is QHead/Tokenx16 across every 16 lanes, 16 Bytes
317 + // HeadElements in each lane, 4x16B HeadElements across 4 rows of warp
318 + constexpr int ROWS_PER_WARP =
319 + WARP_SIZE / 16; // rows refers to 16 lanes; refer DDP (Data Parallel
320 + // Processing) terminology
321 + constexpr int CONTIGUOUS_KV_ELEMS_16B_LOAD =
322 + 16 / sizeof(cache_t); // 8 for 16 bit cache type, 16 for 8 bit types
323 + constexpr int QKHE_PER_FETCH =
324 + CONTIGUOUS_KV_ELEMS_16B_LOAD *
325 + ROWS_PER_WARP; // each fetch across a warp fetches these many elements
326 + constexpr int QK_SIZE_RATIO =
327 + sizeof(scalar_t) /
328 + sizeof(cache_t); // 1 for 16bit types, 2 for 8bit types
329 + constexpr int QKHELOOP = HEAD_SIZE / QKHE_PER_FETCH; // 4xQKHE_16B across
330 + // warp
331 +
332 + _B16x8 Qlocal[QKHELOOP]
333 + [QK_SIZE_RATIO]; // note that 16 contiguous elements of Q should
334 + // be fetched per lane for 8 bit cache types :
335 + // QK_SIZE_RATIO changes for this
336 +
337 + constexpr int CONTIGUOUS_SCALAR_ELEMS_16B = 16 / sizeof(scalar_t);
338 +
339 + constexpr int TOKENS_PER_WARP =
340 + T_PAR_SIZE /
341 + NWARPS; // sub partition of tokens per warp for qk calculation
342 + constexpr int TLOOP =
343 + TOKENS_PER_WARP /
344 + 16; // each mfma16x16x16 instruction processes 16 tokens
345 +
346 + // can be interpreted as B8x16 for 8 bit types
347 + _B16x8 Klocal[TLOOP][QKHELOOP];
348 +
349 + const int wg_start_head_idx = blockIdx.z * GQA_RATIO;
350 + const int wg_start_kv_head_idx = blockIdx.z;
351 + const int total_num_heads = gridDim.z * GQA_RATIO;
352 +
353 + // for QK mfma, tokens in multiples of TOKENS_PER_WARP are spread across warps
354 + // each mfma takes QH16xT16x16HE across warp
355 + // repeat mfmas across QKHELOOP dimension
356 + // output layout from QKmfma : QH16xT4x4 16 qheads across 16 lanes, 16 tokens
357 + // across 4 rows x 4 tokens per lane
358 +
359 + const int num_context_blocks = DIVIDE_ROUND_UP(context_len, BLOCK_SIZE);
360 + const int last_ctx_block = num_context_blocks - 1;
361 +
362 + const int* block_table_seq = block_tables + seq_idx * max_num_blocks_per_seq;
363 +
364 + int kphysical_block_number[TLOOP];
365 +
366 + // fetch k physical block numbers
367 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
368 + const int klocal_token_idx =
369 + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
370 + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
371 + const int kblock_idx = (kglobal_token_idx < context_len)
372 + ? kglobal_token_idx / BLOCK_SIZE
373 + : last_ctx_block;
374 + kphysical_block_number[token_depth] = block_table_seq[kblock_idx];
375 + }
376 +
377 + // fetch Q in shared across warps and then write to registers
378 + const int local_qhead_idx = 4 * warpid + rowid;
379 + const int global_qhead_idx = wg_start_head_idx + local_qhead_idx;
380 + const int64_t seq_idx64 = static_cast<int64_t>(seq_idx);
381 + const scalar_t* q_ptr =
382 + q + seq_idx64 * q_stride + global_qhead_idx * HEAD_SIZE;
383 +
384 + const int qhead_element = lane16id * CONTIGUOUS_SCALAR_ELEMS_16B;
385 + if ((local_qhead_idx < GQA_RATIO) && (qhead_element < HEAD_SIZE)) {
386 + const scalar_t* q_fetch_ptr = q_ptr + qhead_element;
387 + const _B16x8* q_fetch_ptr_16B =
388 + reinterpret_cast<const _B16x8*>(q_fetch_ptr);
389 + _B16x8 tmp = *q_fetch_ptr_16B;
390 + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
391 + const int offset1 =
392 + lane16id /
393 + 4; // 16 contiguous chunks of head elems are spread across 4x4lanes
394 + shared_logits[offset1][lane4id][local_qhead_idx][0] = tmp.xy[0];
395 + shared_logits[offset1][lane4id][local_qhead_idx][1] = tmp.xy[1];
396 + } else {
397 + for (int i = 0; i < 2; i++) {
398 + const int head_elem = lane16id * 2 + i; // element id in _B16x4 terms
399 + const int offset3 = head_elem % 4;
400 + const int offset2 = (head_elem / 4) % 4;
401 + const int offset1 = head_elem / 4 / 4;
402 + shared_logits[offset1][offset2][local_qhead_idx][offset3] = tmp.xy[i];
403 + }
404 + }
405 + }
406 + __syncthreads();
407 + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
408 + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) {
409 + for (int i = 0; i < 2; i++) {
410 + Qlocal[qkhe_depth][qkratio].xy[i] =
411 + shared_logits[qkhe_depth][rowid][lane16id % GQA_RATIO]
412 + [2 * qkratio + i];
413 + }
414 + }
415 + }
416 +
417 + constexpr int KX =
418 + 16 / sizeof(cache_t); // vLLM defines x as 16 Bytes of kv cache elements
419 + const cache_t* k_ptr = k_cache + wg_start_kv_head_idx * kv_head_stride;
420 +
421 + const int row_head_elem = rowid * CONTIGUOUS_KV_ELEMS_16B_LOAD;
422 + // fetch K values
423 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
424 + const int64_t kblock_number =
425 + static_cast<int64_t>(kphysical_block_number[token_depth]);
426 + const cache_t* k_ptr2 = k_ptr + kblock_number * kv_block_stride;
427 + const int klocal_token_idx =
428 + TOKENS_PER_WARP * warpid + token_depth * 16 + lane16id;
429 + const int kglobal_token_idx = partition_start_token_idx + klocal_token_idx;
430 + const int kphysical_block_offset = klocal_token_idx % BLOCK_SIZE;
431 + const cache_t* k_ptr3 = k_ptr2 + kphysical_block_offset * KX;
432 +
433 + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
434 + const int head_elem = row_head_elem + qkhe_depth * QKHE_PER_FETCH;
435 + const int offset1 = head_elem / KX;
436 + const int offset2 = head_elem % KX;
437 + const cache_t* k_fetch_ptr = k_ptr3 + offset1 * BLOCK_SIZE * KX + offset2;
438 + const _B16x8* k_fetch_ptr_16B =
439 + reinterpret_cast<const _B16x8*>(k_fetch_ptr);
440 + Klocal[token_depth][qkhe_depth] = *k_fetch_ptr_16B;
441 + }
442 + }
443 +
444 + float alibi_slope;
445 + if constexpr (ALIBI_ENABLED) {
446 + const int alibi_head_idx = wg_start_head_idx + lane16id;
447 + alibi_slope = (lane16id < GQA_RATIO) ? alibi_slopes[alibi_head_idx] : 0.f;
448 + }
449 +
450 + constexpr int VTOKENS_PER_LANE =
451 + TOKENS_PER_WARP / ROWS_PER_WARP; // 64/4 = 16 contiguous vtokens per lane
452 + constexpr int VBLOCKS_PER_LANE =
453 + 1; // assumes block size >=16, each lane can correspond to 1 block only
454 + constexpr int VTLOOP = NWARPS; // corresponds to tokens across warps
455 + constexpr int VTLANELOOP = DIVIDE_ROUND_UP(
456 + VTOKENS_PER_LANE,
457 + CONTIGUOUS_KV_ELEMS_16B_LOAD); // optimized for 16B fetches; assumes
458 + // minimum block size is 16
459 + constexpr int VHELOOP = HEAD_SIZE / 16 / NWARPS;
460 +
461 + int vphysical_block_number[VTLOOP][VBLOCKS_PER_LANE];
462 +
463 + // fetch v physical block numbers
464 + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
465 + for (int vblock_depth = 0; vblock_depth < VBLOCKS_PER_LANE;
466 + vblock_depth++) {
467 + const int vlocal_token_idx =
468 + vtoken_depth * VTOKENS_PER_LANE * ROWS_PER_WARP +
469 + rowid * VTOKENS_PER_LANE + vblock_depth * BLOCK_SIZE;
470 + // Safe to use an int32_t here assuming we are working with < 2 billion
471 + // tokens
472 + const int vglobal_token_idx =
473 + partition_start_token_idx + vlocal_token_idx;
474 + const int vblock_idx = (vglobal_token_idx < context_len)
475 + ? vglobal_token_idx / BLOCK_SIZE
476 + : last_ctx_block;
477 + vphysical_block_number[vtoken_depth][vblock_depth] =
478 + block_table_seq[vblock_idx];
479 + }
480 + }
481 +
482 + _B16x8 Vlocal[VTLOOP][VHELOOP][VTLANELOOP]; // this could be B8x16 too
483 +
484 + const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride +
485 + ((rowid * VTOKENS_PER_LANE) % BLOCK_SIZE);
486
196 -// grid (num_seqs, num_partitions,num_heads/gqa_ratio)
197 -// block (partition size)
487 + // v fetches are 16head elems across lanes x 16 tokens per lane
488 + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
489 + const int vhead_elem = vhe_depth * NWARPS * 16 + warpid * 16 + lane16id;
490 + const cache_t* v_ptr2 = v_ptr + vhead_elem * BLOCK_SIZE;
491 +
492 + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
493 + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
494 + const int vblock_depth = 0;
495 + const int64_t vblock_number = static_cast<int64_t>(
496 + vphysical_block_number[vtoken_depth][vblock_depth]);
497 + const cache_t* v_ptr3 = v_ptr2 + (vblock_number * kv_block_stride);
498 +
499 + const cache_t* v_fetch_ptr =
500 + v_ptr3 + vfetch_depth * CONTIGUOUS_KV_ELEMS_16B_LOAD;
501 + const _B16x8* v_fetch_ptr_16B =
502 + reinterpret_cast<const _B16x8*>(v_fetch_ptr);
503 + Vlocal[vtoken_depth][vhe_depth][vfetch_depth] = *v_fetch_ptr_16B;
504 + }
505 + }
506 + }
507 +
508 + // calculate post qk mfma scale
509 + float scale2 = scale;
510 + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
511 + // multiply by k_scale if fp8 kv cache
512 + scale2 *= *k_scale;
513 + }
514 +
515 + floatx4 d_out[TLOOP];
516 + // qk mfma
517 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
518 + d_out[token_depth] = {0};
519 + for (int qkhe_depth = 0; qkhe_depth < QKHELOOP; qkhe_depth++) {
520 + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
521 + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) {
522 + for (int i = 0; i < 2; i++) {
523 + d_out[token_depth] = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
524 + Klocal[token_depth][qkhe_depth].xy[i],
525 + Qlocal[qkhe_depth][qkratio].xy[i], d_out[token_depth]);
526 + }
527 + }
528 + } else { // kv cache dtype fp8
529 + auto Ktmp = Klocal[token_depth][qkhe_depth];
530 + _B8x16 Ktmp8x16 = *reinterpret_cast<_B8x16*>(&Ktmp);
531 + for (int qkratio = 0; qkratio < QK_SIZE_RATIO; qkratio++) {
532 + _B8x8 Ktmp8x8 = Ktmp8x16.xy[qkratio];
533 + _B16x8 Klocaltmp = convert_b8x8_custom<scalar_t>(Ktmp8x8);
534 + for (int i = 0; i < 2; i++) {
535 + d_out[token_depth] = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
536 + Klocaltmp.xy[i], Qlocal[qkhe_depth][qkratio].xy[i],
537 + d_out[token_depth]);
538 + }
539 + }
540 + }
541 + }
542 + d_out[token_depth] *= scale2;
543 + }
544 +
545 + const int qkout_token_idx =
546 + partition_start_token_idx + TOKENS_PER_WARP * warpid + rowid * 4;
547 +
548 + // apply alibi
549 + if constexpr (ALIBI_ENABLED) {
550 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
551 + const int local_token_idx = qkout_token_idx + token_depth * 16;
552 + const int alibi_offset = local_token_idx - context_len + 1;
553 + for (int i = 0; i < 4; i++) {
554 + d_out[token_depth][i] += alibi_slope * (alibi_offset + i);
555 + }
556 + }
557 + }
558 +
559 + // calculate qk_max and exp_sum per warp and write to shared memory
560 + float qk_max = -FLT_MAX;
561 + float exp_sum = 0.0f;
562 +
563 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
564 + const int local_token_idx = qkout_token_idx + token_depth * 16;
565 + for (int i = 0; i < 4; i++) {
566 + const float tmp = (local_token_idx + i < context_len)
567 + ? d_out[token_depth][i]
568 + : -FLT_MAX;
569 + qk_max = fmaxf(qk_max, tmp);
570 + }
571 + }
572 +
573 + for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) {
574 + qk_max = fmaxf(qk_max, __shfl_xor(qk_max, mask));
575 + }
576 +
577 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
578 + const int local_token_idx = qkout_token_idx + token_depth * 16;
579 + for (int i = 0; i < 4; i++) {
580 + const float tmp = (local_token_idx + i < context_len)
581 + ? __expf(d_out[token_depth][i] - qk_max)
582 + : 0.0f;
583 + d_out[token_depth][i] = tmp;
584 + exp_sum += tmp;
585 + }
586 + }
587 +
588 + for (int mask = WARP_SIZE / 2; mask >= 16; mask /= 2) {
589 + exp_sum += __shfl_xor(exp_sum, mask);
590 + }
591 +
592 + __syncthreads(); // sync before writing to shared mem
593 +
594 + float* shared_mem = reinterpret_cast<float*>(shared_logits);
595 + if (laneid < 16) {
596 + const int qk_max_offset = warpid * 16 + lane16id;
597 + shared_mem[qk_max_offset] = qk_max;
598 + const int exp_sum_offset = NWARPS * 16 + qk_max_offset;
599 + shared_mem[exp_sum_offset] = exp_sum;
600 + }
601 +
602 + __syncthreads();
603 +
604 + // calculate partition qk_max and exp_sum
605 + float partition_qk_max = -FLT_MAX;
606 + float warp_qk_max_exp[NWARPS];
607 + float partition_exp_sum = 0.0f;
608 +
609 + for (int w = 0; w < NWARPS; w++) {
610 + warp_qk_max_exp[w] = shared_mem[w * 16 + lane16id];
611 + partition_qk_max = fmaxf(partition_qk_max, warp_qk_max_exp[w]);
612 + }
613 +
614 + for (int w = 0; w < NWARPS; w++) {
615 + warp_qk_max_exp[w] = __expf(warp_qk_max_exp[w] - partition_qk_max);
616 + partition_exp_sum +=
617 + shared_mem[NWARPS * 16 + w * 16 + lane16id] * warp_qk_max_exp[w];
618 + }
619 +
620 + const float inv_sum_scale =
621 + __fdividef(1.f, partition_exp_sum + 1e-6f) * warp_qk_max_exp[warpid];
622 +
623 + __syncthreads();
624 +
625 + // disable rtz conversion due to its impact on accuracy.
626 + constexpr bool LOGITS_RTZ_CONVERSION = false;
627 +
628 + // write logits to shared mem
629 + for (int token_depth = 0; token_depth < TLOOP; token_depth++) {
630 + d_out[token_depth] *= inv_sum_scale;
631 + if constexpr (LOGITS_RTZ_CONVERSION) {
632 + // use rtz conversion for better performance, with negligible impact on
633 + // accuracy
634 + shared_logits[warpid][token_depth][lane16id][rowid] =
635 + from_floatx4_rtz<scalar_t>(d_out[token_depth]);
636 + } else {
637 + shared_logits[warpid][token_depth][lane16id][rowid] =
638 + from_floatx4<scalar_t>(d_out[token_depth]);
639 + }
640 + }
641 +
642 + // write out partition max_logits and exp_sum
643 + if (threadIdx.x < GQA_RATIO) {
644 + const int qhead_idx = lane16id;
645 + const int64_t offset = static_cast<int64_t>(seq_idx) *
646 + static_cast<int64_t>(total_num_heads) *
647 + static_cast<int64_t>(max_num_partitions) +
648 + (static_cast<int64_t>(wg_start_head_idx) +
649 + static_cast<int64_t>(qhead_idx)) *
650 + static_cast<int64_t>(max_num_partitions) +
651 + static_cast<int64_t>(partition_idx);
652 + max_logits[offset] = partition_qk_max;
653 + exp_sums[offset] = partition_exp_sum;
654 + }
655 +
656 + __syncthreads();
657 +
658 + constexpr int ELEMS8_ELEMS4_RATIO = 8 / 4;
659 + constexpr int ELEMS16_ELEMS8_RATIO = 16 / 8;
660 +
661 + _B16x4 outelems[VHELOOP];
662 + // Softmax V mfma
663 + // v layout: 16he across lanes x 16 tokens per lane
664 + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
665 + floatx4 tmp_out = {0};
666 +
667 + for (int vtoken_depth = 0; vtoken_depth < VTLOOP; vtoken_depth++) {
668 + if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
669 + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
670 + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
671 + const int offset = rowid * VTLANELOOP * ELEMS8_ELEMS4_RATIO +
672 + vfetch_depth * ELEMS8_ELEMS4_RATIO + i;
673 + const int offset1 = offset % ROWS_PER_WARP;
674 + const int offset2 = offset / ROWS_PER_WARP;
675 + // output format is 16 qheads across 16 lanes, 16 head elems spread
676 + // across 4 rows
677 + tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
678 + Vlocal[vtoken_depth][vhe_depth][vfetch_depth].xy[i],
679 + shared_logits[vtoken_depth][offset2][lane16id][offset1],
680 + tmp_out);
681 + }
682 + }
683 + // KV cache fp8
684 + } else {
685 + for (int vfetch_depth = 0; vfetch_depth < VTLANELOOP; vfetch_depth++) {
686 + _B16x8 Vtmp = Vlocal[vtoken_depth][vhe_depth][vfetch_depth];
687 + // reinterpret V format as 16 elements of 8bits
688 + _B8x16 Vtmp8x16 = *reinterpret_cast<_B8x16*>(&Vtmp);
689 + for (int j = 0; j < ELEMS16_ELEMS8_RATIO; j++) {
690 + _B8x8 Vtmp8x8 = Vtmp8x16.xy[j];
691 + _B16x8 Vlocaltmp = convert_b8x8_custom<scalar_t>(Vtmp8x8);
692 + for (int i = 0; i < ELEMS8_ELEMS4_RATIO; i++) {
693 + const int offset =
694 + rowid * ELEMS16_ELEMS8_RATIO * ELEMS8_ELEMS4_RATIO +
695 + j * ELEMS8_ELEMS4_RATIO + i;
696 + const int offset1 = offset % ROWS_PER_WARP;
697 + const int offset2 = offset / ROWS_PER_WARP;
698 + // output format is 16 qheads across 16 lanes, 16 head elems
699 + // spread across 4 rows
700 + tmp_out = gcn_mfma16x16x16_instr<scalar_t, 0, 0, 0>(
701 + Vlocaltmp.xy[i],
702 + shared_logits[vtoken_depth][offset2][lane16id][offset1],
703 + tmp_out);
704 + }
705 + }
706 + }
707 + }
708 + }
709 + // apply post Softmax V mfma v_scale
710 + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
711 + tmp_out *= *v_scale;
712 + }
713 + outelems[vhe_depth] = from_floatx4<scalar_t>(tmp_out);
714 + }
715 +
716 + __syncthreads();
717 +
718 + // store Softmax-V mfma output to shared mem
719 + for (int vhe_depth = 0; vhe_depth < VHELOOP; vhe_depth++) {
720 + // lane16 id head dimension; rowid head element dimension
721 + shared_logits[warpid][vhe_depth][lane16id][rowid] = outelems[vhe_depth];
722 + }
723 +
724 + __syncthreads();
725 +
726 + // write to tmp_out with coalesced writes after reading from shared mem
727 + if (warpid == 0) {
728 + _B16x8 vout[GQA_RATIO4];
729 + // each lane writes out 16Bytes of tmp_out along head elem dimension
730 + const int head_elem_idx = lane16id * 8;
731 + if (head_elem_idx < HEAD_SIZE) {
732 + for (int h = 0; h < GQA_RATIO4; h++) {
733 + const int local_head_idx = 4 * h + rowid;
734 + const int offset1 = (head_elem_idx / 16) % 4;
735 + const int offset2 = head_elem_idx / 16 / NWARPS;
736 + const int offset3 = (head_elem_idx / 4) % 4;
737 + for (int i = 0; i < 2; i++) {
738 + vout[h].xy[i] =
739 + shared_logits[offset1][offset2][local_head_idx][offset3 + i];
740 + }
741 + }
742 +
743 + const int64_t hsz_maxp_mult =
744 + static_cast<int64_t>(HEAD_SIZE * max_num_partitions);
745 + scalar_t* out_ptr = out + seq_idx * total_num_heads * hsz_maxp_mult +
746 + partition_idx * HEAD_SIZE;
747 + for (int h = 0; h < GQA_RATIO4; h++) {
748 + const int local_head_idx = 4 * h + rowid;
749 + if (local_head_idx < GQA_RATIO) {
750 + const int64_t out_head_idx =
751 + static_cast<int64_t>(wg_start_head_idx + local_head_idx);
752 + scalar_t* out_ptr2 = out_ptr + out_head_idx * hsz_maxp_mult;
753 + scalar_t* out_ptr3 = out_ptr2 + head_elem_idx;
754 + _B16x8* out_ptr_B16x8 = reinterpret_cast<_B16x8*>(out_ptr3);
755 + *out_ptr_B16x8 = vout[h];
756 + }
757 + }
758 + }
759 + }
760 +}
761 +
762 +// grid (num_seqs, num_partitions, num_kv_heads)
763 +// block (256 : partition size)
764 +// each WG handles 1 partition per sequence
765 +// clang-format off
766 template <typename scalar_t, typename cache_t,
199 - vllm::Fp8KVCacheDataType KV_DTYPE, int BLOCK_SIZE, int HEAD_SIZE,
200 - int NUM_THREADS,
767 + vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
768 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
769 int GQA_RATIO>
202 -__global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
203 - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
204 - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
205 - // head_size/x, block_size, x]
206 - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
207 - // head_size, block_size]
208 - const int num_kv_heads, const float scale,
209 - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
210 - const int* __restrict__ context_lens, // [num_seqs]
770 +__global__
771 +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
772 + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
773 + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
774 + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
775 + const int num_kv_heads,
776 + const float scale,
777 + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
778 + const int* __restrict__ context_lens, // [num_seqs]
779 const int max_num_blocks_per_seq,
212 - const float* __restrict__ alibi_slopes, // [num_heads]
213 - const int q_stride, const int kv_block_stride, const int kv_head_stride,
214 - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
215 - float* __restrict__ max_logits, // [num_seqs, num_heads,
216 - // max_num_partitions]
217 - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
218 - // head_size]
219 - scalar_t* __restrict__ final_out, // [num_seqs, num_heads, head_size]
220 - int max_ctx_blocks, const float* k_scale_ptr, const float* v_scale_ptr) {
780 + const float* __restrict__ alibi_slopes, // [num_heads]
781 + const int q_stride,
782 + const int kv_block_stride,
783 + const int kv_head_stride,
784 + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
785 + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
786 + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
787 + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
788 + int max_ctx_blocks, const float* k_scale, const float* v_scale) {
789 + // clang-format on
790 constexpr int NWARPS = NUM_THREADS / WARP_SIZE;
791 const int warpid = threadIdx.x / WARP_SIZE;
792 const int laneid = threadIdx.x % WARP_SIZE;
@@ -234,29 +803,37 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
803 if (partition_start_token_idx >= context_len) {
804 return;
805 }
237 - constexpr int QHLOOP =
238 - DIVIDE_ROUND_UP(GQA_RATIO, 4); // each 4 lanes fetch 4 different qheads,
239 - // total qheads =8, so qhloop is 2
806 + // every 4 lanes fetch 4 different qheads
807 + // qhloop = num loops over qhead dimension
808 + constexpr int QHLOOP = DIVIDE_ROUND_UP(GQA_RATIO, 4);
809 constexpr int GQA_RATIO4 = 4 * QHLOOP;
810 __shared__ float shared_qk_max[NWARPS][GQA_RATIO4 + 1];
811 __shared__ float shared_exp_sum[NWARPS][GQA_RATIO4 + 1];
812 _B16x8 Qlocal[QHLOOP];
813 constexpr int x = 16 / sizeof(scalar_t);
814 + // kheloop = num loops over head_size for 16Bytes of Q/dequantized K elements
815 constexpr int KHELOOP = HEAD_SIZE / x;
816 _B16x8 Klocal[KHELOOP];
817 _B8x8 Klocalb8[KHELOOP];
248 - constexpr int VHELOOP =
249 - HEAD_SIZE /
250 - WARP_SIZE; // v head_size dimension is distributed across lanes
251 - constexpr int VTLOOP = 8; // 16 separate 4xtokens across warp -> 16/2
252 - // 8xtokens
818 + // for SoftMax-V Gemm, V head_size dimension is distributed across warp
819 + // vheloop = num loops to cover v head size dimension
820 + constexpr int VHELOOP = HEAD_SIZE / WARP_SIZE;
821 + // softmax out has warp_size tokens across warp
822 + // vtloop = num loops to cover warp_size(64) tokens with 16Bytes of
823 + // dequantized V elements
824 + constexpr int VTLOOP = WARP_SIZE / 8;
825 + // num vblocks to cover warp_size(64) v elements
826 + constexpr int VBLOCKS = 8 * VTLOOP / BLOCK_SIZE;
827 + int vphysical_blocks[VBLOCKS];
828 _B16x8 Vlocal[VHELOOP][VTLOOP];
829 _B8x8 Vlocalb8[VHELOOP][VTLOOP];
255 - floatx4 dout[QHLOOP];
830 + floatx4 d_out[QHLOOP];
831 float qk_max[QHLOOP];
257 - #pragma unroll
832 +
833 + __shared__ _B16x4 vout_shared[QHLOOP][VHELOOP][WARP_SIZE][NWARPS + 1];
834 +
835 for (int h = 0; h < QHLOOP; h++) {
259 - dout[h] = {0};
836 + d_out[h] = {0};
837 qk_max[h] = -FLT_MAX;
838 }
839
@@ -278,25 +855,24 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
855 const int last_ctx_block = num_context_blocks - 1;
856
857 const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
281 -
858 + // token id within partition
859 const int local_token_idx = threadIdx.x;
860 + // token id within sequence
861 const int global_token_idx = partition_start_token_idx + local_token_idx;
862
863 + // fetch block number for k
864 const int block_idx = (global_token_idx < context_len)
865 ? global_token_idx / BLOCK_SIZE
866 : last_ctx_block;
288 - // fetch block number for q and k
289 - // int32 physical_block_number leads to overflow when multiplied with
290 - // kv_block_stride
867 +
868 + // fetch k physical block number
869 + // int32 physical_block_number leads to overflow when multiplied with
870 + // kv_block_stride
871 const int64_t physical_block_number =
872 static_cast<int64_t>(block_table[block_idx]);
873
874 // fetch vphysical block numbers up front
295 - constexpr int VBLOCKS = 8 * VTLOOP / BLOCK_SIZE;
296 - int vphysical_blocks[VBLOCKS];
297 -
875 const int warp_start_block_idx = warp_start_token_idx / BLOCK_SIZE;
299 - #pragma unroll
876 for (int b = 0; b < VBLOCKS; b++) {
877 const int vblock_idx = warp_start_block_idx + b;
878 const int vblock_idx_ctx =
@@ -304,12 +880,13 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
880 vphysical_blocks[b] = block_table[vblock_idx_ctx];
881 }
882
307 - // each 4 lanes fetch 8 helems, so warp fetches 8*16 = 128 helems
883 + // fetch q elements
884 + // every 4 lanes fetch 8 elems, so warp fetches 8*16 = 128 elems
885 const scalar_t* q_ptr =
886 q + seq_idx * q_stride + wg_start_head_idx * HEAD_SIZE;
887 const _B16x8* q_ptrh8 = reinterpret_cast<const _B16x8*>(q_ptr);
888 const int qhead_elemh8 = laneid / 4;
312 - #pragma unroll
889 +
890 for (int h = 0; h < QHLOOP - 1; h++) {
891 const int qhead_idx = h * 4 + lane4id;
892 Qlocal[h] = q_ptrh8[qhead_idx * HEAD_SIZE / 8 + qhead_elemh8];
@@ -323,22 +900,24 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
900 Qlocal[QHLOOP - 1].xy[1] = {0};
901 }
902
903 + // fetch k elements
904 const cache_t* k_ptr = k_cache + physical_block_number * kv_block_stride +
905 wg_start_kv_head_idx * kv_head_stride;
906
329 - const int physical_block_offset =
330 - local_token_idx % BLOCK_SIZE; // since x=half8, physical_block_offset
331 - // is already cast as _H8
907 + // physical_block_offset is already cast in terms of _B16x8
908 + const int physical_block_offset = local_token_idx % BLOCK_SIZE;
909 +
910 + // each K fetch is for 8 elements of cache_t which are later dequantized to
911 + // scalar_t for fp8
912 if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
913 const _B16x8* k_ptrh8 = reinterpret_cast<const _B16x8*>(k_ptr);
334 - #pragma unroll
914 for (int d = 0; d < KHELOOP; d++) {
915 Klocal[d] = k_ptrh8[d * BLOCK_SIZE + physical_block_offset];
916 }
917 } else {
918 + // vllm defines X as 16 Bytes of elements of cache_t
919 constexpr int X = 16 / sizeof(cache_t);
920 const cache_t* k_ptr2 = k_ptr + physical_block_offset * X;
341 - #pragma unroll
921 for (int d = 0; d < KHELOOP; d++) {
922 const int head_elem = d * 8;
923 const int offset1 = head_elem / X;
@@ -348,9 +927,9 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
927 }
928 }
929
930 + // optional alibi fetch
931 float alibi_slope[QHLOOP];
352 - if (alibi_slopes != nullptr) {
353 - #pragma unroll
932 + if constexpr (ALIBI_ENABLED) {
933 for (int h = 0; h < QHLOOP; h++) {
934 const int qhead_idx = h * 4 + lane4id;
935 alibi_slope[h] = (qhead_idx < GQA_RATIO)
@@ -360,10 +939,10 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
939 }
940
941 const cache_t* v_ptr = v_cache + wg_start_kv_head_idx * kv_head_stride;
942 + // fetch vcache in kv cache auto case
943 if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto) {
944 const _B16x8* v_ptrh8 = reinterpret_cast<const _B16x8*>(v_ptr);
945 // iterate over each v block
366 - #pragma unroll
946 for (int b = 0; b < VBLOCKS; b++) {
947 // int32 physical_block_number leads to overflow when multiplied with
948 // kv_block_stride
@@ -372,21 +951,20 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
951 const _B16x8* v_ptrh8b =
952 v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8;
953 // iterate over each head elem (within head_size)
375 - #pragma unroll
954 for (int h = 0; h < VHELOOP; h++) {
955 const int head_size_elem = h * WARP_SIZE + laneid;
956 const _B16x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
957 // iterate over all velems within block
380 - #pragma unroll
958 for (int d = 0; d < BLOCK_SIZE / 8; d++) {
959 Vlocal[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
960 }
961 }
962 }
386 - } else {
963 + } // if constexpr (KV_DTYPE == vllm::Fp8KVCacheDataType::kAuto)
964 + // fetch vcache in fp8 case
965 + else { // if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto)
966 const _B8x8* v_ptrh8 = reinterpret_cast<const _B8x8*>(v_ptr);
967 // iterate over each v block
389 - #pragma unroll
968 for (int b = 0; b < VBLOCKS; b++) {
969 // int32 physical_block_number leads to overflow when multiplied with
970 // kv_block_stride
@@ -395,164 +973,153 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
973 const _B8x8* v_ptrh8b =
974 v_ptrh8 + (vphysical_block_number * kv_block_stride) / 8;
975 // iterate over each head elem (within head_size)
398 - #pragma unroll
976 for (int h = 0; h < VHELOOP; h++) {
977 const int head_size_elem = h * WARP_SIZE + laneid;
978 const _B8x8* v_ptrh8be = v_ptrh8b + head_size_elem * BLOCK_SIZE / 8;
979 // iterate over all velems within block
403 - #pragma unroll
980 for (int d = 0; d < BLOCK_SIZE / 8; d++) {
405 - // Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
406 - const _B8x8 Vlocalb8 = v_ptrh8be[d];
407 - Vlocal[h][b * BLOCK_SIZE / 8 + d] =
408 - scaled_convert_b8x8<scalar_t, KV_DTYPE>(Vlocalb8, *v_scale_ptr);
981 + Vlocalb8[h][b * BLOCK_SIZE / 8 + d] = v_ptrh8be[d];
982 }
983 }
984 }
985 }
986
987 + #define QK_mfma(x) \
988 + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \
989 + Klocal[x] = convert_b8x8_custom<scalar_t>(Klocalb8[x]); \
990 + } \
991 + for (int h = 0; h < QHLOOP; h++) { \
992 + d_out[h] = gcn_mfma4x4x4_instr<scalar_t, 4, x, 0>( \
993 + Qlocal[h].xy[0], Klocal[x].xy[0], d_out[h]); \
994 + d_out[h] = gcn_mfma4x4x4_instr<scalar_t, 4, x, 0>( \
995 + Qlocal[h].xy[1], Klocal[x].xy[1], d_out[h]); \
996 + }
997 + // QK mfma with Q mfma block broadcast
998 + // Q values across head_size dimension stored across lanes
999 + // K values across head_size dimension are stored depthwise within lane
1000 + // Q broadcast with absz, cbid of mfma instruction
1001 + QK_mfma(0);
1002 + QK_mfma(1);
1003 + QK_mfma(2);
1004 + QK_mfma(3);
1005 + QK_mfma(4);
1006 + QK_mfma(5);
1007 + QK_mfma(6);
1008 + QK_mfma(7);
1009 + // below only needed for head size 128
1010 + if constexpr (KHELOOP > 8) {
1011 + QK_mfma(8);
1012 + QK_mfma(9);
1013 + QK_mfma(10);
1014 + QK_mfma(11);
1015 + QK_mfma(12);
1016 + QK_mfma(13);
1017 + QK_mfma(14);
1018 + QK_mfma(15);
1019 + }
1020 + #undef QK_mfma
1021 +
1022 + float scale2 = scale;
1023 if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
415 - #pragma unroll
416 - for (int d = 0; d < KHELOOP; d++) {
417 - Klocal[d] =
418 - scaled_convert_b8x8<scalar_t, KV_DTYPE>(Klocalb8[d], *k_scale_ptr);
419 - }
1024 + // post mfma scaling for fp8
1025 + scale2 *= *k_scale;
1026 }
1027
422 - #pragma unroll
1028 for (int h = 0; h < QHLOOP; h++) {
424 - dout[h] = gcn_mfma_instr<scalar_t, 4, 0, 0>(Qlocal[h].xy[0],
425 - Klocal[0].xy[0], dout[h]);
426 - dout[h] = gcn_mfma_instr<scalar_t, 4, 0, 0>(Qlocal[h].xy[1],
427 - Klocal[0].xy[1], dout[h]);
428 - dout[h] = gcn_mfma_instr<scalar_t, 4, 1, 0>(Qlocal[h].xy[0],
429 - Klocal[1].xy[0], dout[h]);
430 - dout[h] = gcn_mfma_instr<scalar_t, 4, 1, 0>(Qlocal[h].xy[1],
431 - Klocal[1].xy[1], dout[h]);
432 - dout[h] = gcn_mfma_instr<scalar_t, 4, 2, 0>(Qlocal[h].xy[0],
433 - Klocal[2].xy[0], dout[h]);
434 - dout[h] = gcn_mfma_instr<scalar_t, 4, 2, 0>(Qlocal[h].xy[1],
435 - Klocal[2].xy[1], dout[h]);
436 - dout[h] = gcn_mfma_instr<scalar_t, 4, 3, 0>(Qlocal[h].xy[0],
437 - Klocal[3].xy[0], dout[h]);
438 - dout[h] = gcn_mfma_instr<scalar_t, 4, 3, 0>(Qlocal[h].xy[1],
439 - Klocal[3].xy[1], dout[h]);
440 - dout[h] = gcn_mfma_instr<scalar_t, 4, 4, 0>(Qlocal[h].xy[0],
441 - Klocal[4].xy[0], dout[h]);
442 - dout[h] = gcn_mfma_instr<scalar_t, 4, 4, 0>(Qlocal[h].xy[1],
443 - Klocal[4].xy[1], dout[h]);
444 - dout[h] = gcn_mfma_instr<scalar_t, 4, 5, 0>(Qlocal[h].xy[0],
445 - Klocal[5].xy[0], dout[h]);
446 - dout[h] = gcn_mfma_instr<scalar_t, 4, 5, 0>(Qlocal[h].xy[1],
447 - Klocal[5].xy[1], dout[h]);
448 - dout[h] = gcn_mfma_instr<scalar_t, 4, 6, 0>(Qlocal[h].xy[0],
449 - Klocal[6].xy[0], dout[h]);
450 - dout[h] = gcn_mfma_instr<scalar_t, 4, 6, 0>(Qlocal[h].xy[1],
451 - Klocal[6].xy[1], dout[h]);
452 - dout[h] = gcn_mfma_instr<scalar_t, 4, 7, 0>(Qlocal[h].xy[0],
453 - Klocal[7].xy[0], dout[h]);
454 - dout[h] = gcn_mfma_instr<scalar_t, 4, 7, 0>(Qlocal[h].xy[1],
455 - Klocal[7].xy[1], dout[h]);
456 - if constexpr (KHELOOP > 8) {
457 - dout[h] = gcn_mfma_instr<scalar_t, 4, 8, 0>(Qlocal[h].xy[0],
458 - Klocal[8].xy[0], dout[h]);
459 - dout[h] = gcn_mfma_instr<scalar_t, 4, 8, 0>(Qlocal[h].xy[1],
460 - Klocal[8].xy[1], dout[h]);
461 - dout[h] = gcn_mfma_instr<scalar_t, 4, 9, 0>(Qlocal[h].xy[0],
462 - Klocal[9].xy[0], dout[h]);
463 - dout[h] = gcn_mfma_instr<scalar_t, 4, 9, 0>(Qlocal[h].xy[1],
464 - Klocal[9].xy[1], dout[h]);
465 - dout[h] = gcn_mfma_instr<scalar_t, 4, 10, 0>(Qlocal[h].xy[0],
466 - Klocal[10].xy[0], dout[h]);
467 - dout[h] = gcn_mfma_instr<scalar_t, 4, 10, 0>(Qlocal[h].xy[1],
468 - Klocal[10].xy[1], dout[h]);
469 - dout[h] = gcn_mfma_instr<scalar_t, 4, 11, 0>(Qlocal[h].xy[0],
470 - Klocal[11].xy[0], dout[h]);
471 - dout[h] = gcn_mfma_instr<scalar_t, 4, 11, 0>(Qlocal[h].xy[1],
472 - Klocal[11].xy[1], dout[h]);
473 - dout[h] = gcn_mfma_instr<scalar_t, 4, 12, 0>(Qlocal[h].xy[0],
474 - Klocal[12].xy[0], dout[h]);
475 - dout[h] = gcn_mfma_instr<scalar_t, 4, 12, 0>(Qlocal[h].xy[1],
476 - Klocal[12].xy[1], dout[h]);
477 - dout[h] = gcn_mfma_instr<scalar_t, 4, 13, 0>(Qlocal[h].xy[0],
478 - Klocal[13].xy[0], dout[h]);
479 - dout[h] = gcn_mfma_instr<scalar_t, 4, 13, 0>(Qlocal[h].xy[1],
480 - Klocal[13].xy[1], dout[h]);
481 - dout[h] = gcn_mfma_instr<scalar_t, 4, 14, 0>(Qlocal[h].xy[0],
482 - Klocal[14].xy[0], dout[h]);
483 - dout[h] = gcn_mfma_instr<scalar_t, 4, 14, 0>(Qlocal[h].xy[1],
484 - Klocal[14].xy[1], dout[h]);
485 - dout[h] = gcn_mfma_instr<scalar_t, 4, 15, 0>(Qlocal[h].xy[0],
486 - Klocal[15].xy[0], dout[h]);
487 - dout[h] = gcn_mfma_instr<scalar_t, 4, 15, 0>(Qlocal[h].xy[1],
488 - Klocal[15].xy[1], dout[h]);
489 - } // KHELOOP>8
490 - dout[h] *= scale;
1029 + d_out[h] *= scale2;
1030 }
492 - // transpose dout so that 4 token ids are in each lane, and 4 heads are across
493 - // 4 lanes
494 - #pragma unroll
1031 +
1032 + // transpose d_out so that 4 token ids are in each lane, and 4 heads are
1033 + // across 4 lanes
1034 for (int h = 0; h < QHLOOP; h++) {
1035 floatx4 tmp = {0};
497 - #pragma unroll
1036 for (int i = 0; i < 4; i++) {
1037 const float B = (lane4id == i) ? 1.0f : 0.0f;
500 - // const float A = (global_token_idx < context_len) ? dout[h][i] : 0.0f;
501 - tmp = __builtin_amdgcn_mfma_f32_4x4x1f32(dout[h][i], B, tmp, 0, 0, 0);
502 - // tmp = __builtin_amdgcn_mfma_f32_4x4x1f32(A, B, tmp, 0, 0, 0);
1038 + tmp = __builtin_amdgcn_mfma_f32_4x4x1f32(d_out[h][i], B, tmp, 0, 0, 0);
1039 }
504 - dout[h] = tmp;
1040 + d_out[h] = tmp;
1041 }
1042
1043 const int lane4_token_idx = 4 * (global_token_idx >> 2);
508 - const int alibi_offset = lane4_token_idx - context_len + 1;
509 - if (alibi_slopes != nullptr) {
510 - #pragma unroll
1044 +
1045 + if constexpr (ALIBI_ENABLED) {
1046 + const int alibi_offset = lane4_token_idx - context_len + 1;
1047 for (int h = 0; h < QHLOOP; h++) {
512 - #pragma unroll
1048 for (int i = 0; i < 4; i++) {
514 - dout[h][i] += alibi_slope[h] * (alibi_offset + i);
1049 + d_out[h][i] += alibi_slope[h] * (alibi_offset + i);
1050 }
1051 }
1052 }
1053
519 - #pragma unroll
1054 + const int bpermute_mask = 4 * (16 * ((laneid >> 2) % 4) + lane4id);
1055 +
1056 for (int h = 0; h < QHLOOP; h++) {
1057 qk_max[h] = -FLT_MAX;
522 - #pragma unroll
1058 for (int i = 0; i < 4; i++) {
1059 qk_max[h] = (lane4_token_idx + i < context_len)
525 - ? fmaxf(qk_max[h], dout[h][i])
1060 + ? fmaxf(qk_max[h], d_out[h][i])
1061 : qk_max[h];
1062 }
528 - #pragma unroll
529 - for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) {
530 - qk_max[h] = fmaxf(qk_max[h], __shfl_xor(qk_max[h], mask));
531 - }
1063 +
1064 + // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) {
1065 + // qk_max[h] = fmaxf(qk_max[h], __shfl_xor(qk_max[h], mask));
1066 + // }
1067 + // faster version of above code with dpp
1068 + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4"
1069 + : "=v"(qk_max[h])
1070 + : "v"(qk_max[h]), "v"(qk_max[h]));
1071 + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8"
1072 + : "=v"(qk_max[h])
1073 + : "v"(qk_max[h]), "v"(qk_max[h]));
1074 +
1075 + auto tmp = __builtin_amdgcn_ds_bpermute(
1076 + bpermute_mask, *reinterpret_cast<int*>(&qk_max[h]));
1077 + qk_max[h] = *reinterpret_cast<float*>(&tmp);
1078 + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:4"
1079 + : "=v"(qk_max[h])
1080 + : "v"(qk_max[h]), "v"(qk_max[h]));
1081 + asm("v_nop\n v_nop\n v_max_f32_dpp %0, %1, %2 row_ror:8"
1082 + : "=v"(qk_max[h])
1083 + : "v"(qk_max[h]), "v"(qk_max[h]));
1084 }
1085
1086 float exp_sum[QHLOOP];
535 - #pragma unroll
1087 for (int h = 0; h < QHLOOP; h++) {
1088 exp_sum[h] = 0.0f;
538 - #pragma unroll
1089 for (int i = 0; i < 4; i++) {
540 - dout[h][i] = (lane4_token_idx + i < context_len)
541 - ? __expf(dout[h][i] - qk_max[h])
542 - : 0.0f;
543 - exp_sum[h] += dout[h][i];
544 - }
545 - #pragma unroll
546 - for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) {
547 - exp_sum[h] += __shfl_xor(exp_sum[h], mask);
1090 + d_out[h][i] = (lane4_token_idx + i < context_len)
1091 + ? __expf(d_out[h][i] - qk_max[h])
1092 + : 0.0f;
1093 + exp_sum[h] += d_out[h][i];
1094 }
1095 + // for (int mask = WARP_SIZE / 2; mask >= 4; mask /= 2) {
1096 + // exp_sum[h] += __shfl_xor(exp_sum[h], mask);
1097 + // }
1098 + // faster version of above code with dpp
1099 + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4"
1100 + : "=v"(exp_sum[h])
1101 + : "v"(exp_sum[h]), "v"(exp_sum[h]));
1102 + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8"
1103 + : "=v"(exp_sum[h])
1104 + : "v"(exp_sum[h]), "v"(exp_sum[h]));
1105 +
1106 + auto tmp = __builtin_amdgcn_ds_bpermute(
1107 + bpermute_mask, *reinterpret_cast<int*>(&exp_sum[h]));
1108 + exp_sum[h] = *reinterpret_cast<float*>(&tmp);
1109 + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:4"
1110 + : "=v"(exp_sum[h])
1111 + : "v"(exp_sum[h]), "v"(exp_sum[h]));
1112 + asm("v_nop\n v_nop\n v_add_f32_dpp %0, %1, %2 row_ror:8"
1113 + : "=v"(exp_sum[h])
1114 + : "v"(exp_sum[h]), "v"(exp_sum[h]));
1115 }
1116
551 - #pragma unroll
552 - for (int h = 0; h < QHLOOP; h++) {
553 - const int head_idx = 4 * h + lane4id;
554 - shared_qk_max[warpid][head_idx] = qk_max[h];
555 - shared_exp_sum[warpid][head_idx] = exp_sum[h];
1117 + if (laneid < 4) {
1118 + for (int h = 0; h < QHLOOP; h++) {
1119 + const int head_idx = 4 * h + lane4id;
1120 + shared_qk_max[warpid][head_idx] = qk_max[h];
1121 + shared_exp_sum[warpid][head_idx] = exp_sum[h];
1122 + }
1123 }
1124 } // warp within context
1125
@@ -563,18 +1130,16 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
1130 max_logits + seq_idx * num_heads * max_num_partitions + partition_idx;
1131 float* exp_sums_ptr =
1132 exp_sums + seq_idx * num_heads * max_num_partitions + partition_idx;
566 - #pragma unroll
1133 + // calculate qk_max and exp_sums for partition
1134 for (int h = 0; h < QHLOOP; h++) {
1135 float global_qk_max = -FLT_MAX;
1136 float warp_qk_max[NWARPS];
1137 const int head_idx = 4 * h + lane4id;
571 - #pragma unroll
1138 for (int w = 0; w < NWARPS; w++) {
1139 warp_qk_max[w] = shared_qk_max[w][head_idx];
1140 global_qk_max = fmaxf(global_qk_max, warp_qk_max[w]);
1141 }
1142 float global_exp_sum = 0.0f;
577 - #pragma unroll
1143 for (int w = 0; w < NWARPS; w++) {
1144 global_exp_sum +=
1145 shared_exp_sum[w][head_idx] * __expf(warp_qk_max[w] - global_qk_max);
@@ -587,101 +1152,94 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
1152 }
1153 const float global_inv_sum_scale = __fdividef(1.f, global_exp_sum + 1e-6f) *
1154 __expf(qk_max[h] - global_qk_max);
590 - dout[h] *= global_inv_sum_scale;
1155 + d_out[h] *= global_inv_sum_scale;
1156 }
1157 + constexpr bool LOGITS_RTZ_CONVERSION = false;
1158 // logits[h] -> every 4 lanes hold 4 heads, each lane holds 4 tokens, there
1159 // are 4x16 tokens across warp
1160 _B16x4 logits[QHLOOP];
595 - #pragma unroll
1161 for (int h = 0; h < QHLOOP; h++) {
597 - logits[h] = from_floatx4<scalar_t>(dout[h]);
1162 + if constexpr (LOGITS_RTZ_CONVERSION) {
1163 + // use rtz for faster performance with no perceivable accuracy loss
1164 + logits[h] = from_floatx4_rtz<scalar_t>(d_out[h]);
1165 + } else {
1166 + logits[h] = from_floatx4<scalar_t>(d_out[h]);
1167 + }
1168 }
1169
600 - __shared__ _B16x4 vout_shared[QHLOOP][VHELOOP][WARP_SIZE][NWARPS + 1];
601 -
1170 if (warp_start_token_idx >= context_len) { // warp out of context
603 - #pragma unroll
1171 for (int qh = 0; qh < QHLOOP; qh++) {
605 - #pragma unroll
1172 for (int vh = 0; vh < VHELOOP; vh++) {
1173 vout_shared[qh][vh][laneid][warpid] = {0};
1174 }
1175 }
1176 } else { // warp in context
611 - // iterate across heads
612 - #pragma unroll
613 - for (int qh = 0; qh < QHLOOP; qh++) {
614 - // iterate over each v head elem (within head_size)
615 - #pragma unroll
616 - for (int vh = 0; vh < VHELOOP; vh++) {
617 - floatx4 acc = {0};
618 - // iterate over tokens
619 - acc = gcn_mfma_instr<scalar_t, 4, 0, 0>(logits[qh], Vlocal[vh][0].xy[0],
620 - acc);
621 - acc = gcn_mfma_instr<scalar_t, 4, 1, 0>(logits[qh], Vlocal[vh][0].xy[1],
622 - acc);
623 - acc = gcn_mfma_instr<scalar_t, 4, 2, 0>(logits[qh], Vlocal[vh][1].xy[0],
624 - acc);
625 - acc = gcn_mfma_instr<scalar_t, 4, 3, 0>(logits[qh], Vlocal[vh][1].xy[1],
626 - acc);
627 - acc = gcn_mfma_instr<scalar_t, 4, 4, 0>(logits[qh], Vlocal[vh][2].xy[0],
628 - acc);
629 - acc = gcn_mfma_instr<scalar_t, 4, 5, 0>(logits[qh], Vlocal[vh][2].xy[1],
630 - acc);
631 - acc = gcn_mfma_instr<scalar_t, 4, 6, 0>(logits[qh], Vlocal[vh][3].xy[0],
632 - acc);
633 - acc = gcn_mfma_instr<scalar_t, 4, 7, 0>(logits[qh], Vlocal[vh][3].xy[1],
634 - acc);
635 - acc = gcn_mfma_instr<scalar_t, 4, 8, 0>(logits[qh], Vlocal[vh][4].xy[0],
636 - acc);
637 - acc = gcn_mfma_instr<scalar_t, 4, 9, 0>(logits[qh], Vlocal[vh][4].xy[1],
638 - acc);
639 - acc = gcn_mfma_instr<scalar_t, 4, 10, 0>(logits[qh],
640 - Vlocal[vh][5].xy[0], acc);
641 - acc = gcn_mfma_instr<scalar_t, 4, 11, 0>(logits[qh],
642 - Vlocal[vh][5].xy[1], acc);
643 - acc = gcn_mfma_instr<scalar_t, 4, 12, 0>(logits[qh],
644 - Vlocal[vh][6].xy[0], acc);
645 - acc = gcn_mfma_instr<scalar_t, 4, 13, 0>(logits[qh],
646 - Vlocal[vh][6].xy[1], acc);
647 - acc = gcn_mfma_instr<scalar_t, 4, 14, 0>(logits[qh],
648 - Vlocal[vh][7].xy[0], acc);
649 - acc = gcn_mfma_instr<scalar_t, 4, 15, 0>(logits[qh],
650 - Vlocal[vh][7].xy[1], acc);
651 - vout_shared[qh][vh][laneid][warpid] = from_floatx4<scalar_t>(acc);
1177 + #define SV_mfma(x) \
1178 + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) { \
1179 + Vlocal[vh][x] = convert_b8x8_custom<scalar_t>(Vlocalb8[vh][x]); \
1180 + } \
1181 + for (int qh = 0; qh < QHLOOP; qh++) { \
1182 + acc[qh] = gcn_mfma4x4x4_instr<scalar_t, 4, 2 * x, 0>( \
1183 + logits[qh], Vlocal[vh][x].xy[0], acc[qh]); \
1184 + acc[qh] = gcn_mfma4x4x4_instr<scalar_t, 4, 2 * x + 1, 0>( \
1185 + logits[qh], Vlocal[vh][x].xy[1], acc[qh]); \
1186 + }
1187 +
1188 + for (int vh = 0; vh < VHELOOP; vh++) {
1189 + floatx4 acc[QHLOOP];
1190 + for (int qh = 0; qh < QHLOOP; qh++) {
1191 + acc[qh] = {0};
1192 + }
1193 + // SoftMax-V calculation
1194 + // logits -> token dimension is distributed across lanes
1195 + // Vlocal -> token dimension is depthwise within lane
1196 + // uses mfma instruction block broadcast for logits
1197 + SV_mfma(0);
1198 + SV_mfma(1);
1199 + SV_mfma(2);
1200 + SV_mfma(3);
1201 + SV_mfma(4);
1202 + SV_mfma(5);
1203 + SV_mfma(6);
1204 + SV_mfma(7);
1205 +
1206 + for (int qh = 0; qh < QHLOOP; qh++) {
1207 + if constexpr (KV_DTYPE != vllm::Fp8KVCacheDataType::kAuto) {
1208 + // post mfma v scale for fp8
1209 + acc[qh] *= *v_scale;
1210 + }
1211 + vout_shared[qh][vh][laneid][warpid] = from_floatx4<scalar_t>(acc[qh]);
1212 }
1213 }
1214 +
1215 + #undef SV_mfma
1216 } // warp in context
1217
1218 __syncthreads();
1219
1220 + // final write to tmp_out after vout accumulation
1221 if (warpid == 0) {
1222 _B16x4 vout[QHLOOP][VHELOOP];
1223 // iterate across heads
661 - scalar_t* out_ptr;
662 - int out_num_partitions;
663 - if (context_len > partition_size) {
664 - out_num_partitions = max_num_partitions;
665 - out_ptr = out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
666 - partition_idx * HEAD_SIZE;
667 - } else {
668 - out_num_partitions = 1;
669 - out_ptr = final_out + seq_idx * num_heads * HEAD_SIZE;
670 - }
671 - #pragma unroll
1224 for (int qh = 0; qh < QHLOOP; qh++) {
673 - // iterate over each v head elem (within head_size)
674 - #pragma unroll
1225 + // iterate over each v head elem (within head_size)
1226 for (int vh = 0; vh < VHELOOP; vh++) {
1227 vout[qh][vh] = {0};
677 - #pragma unroll
1228 for (int w = 0; w < NWARPS; w++) {
1229 vout[qh][vh] =
1230 addx4<scalar_t>(vout[qh][vh], vout_shared[qh][vh][laneid][w]);
1231 }
1232 + }
1233 + }
1234 +
1235 + scalar_t* out_ptr = out +
1236 + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
1237 + partition_idx * HEAD_SIZE;
1238 + const int out_num_partitions = max_num_partitions;
1239 + bit16_t* out_ptr_b16 = reinterpret_cast<bit16_t*>(out_ptr);
1240 + for (int qh = 0; qh < QHLOOP; qh++) {
1241 + for (int vh = 0; vh < VHELOOP; vh++) {
1242 const int head_size_elem = vh * WARP_SIZE + laneid;
683 - bit16_t* out_ptr_b16 = reinterpret_cast<bit16_t*>(out_ptr);
684 - #pragma unroll
1243 for (int i = 0; i < 4; i++) {
1244 const int head_idx = 4 * qh + i;
1245 if (head_idx < GQA_RATIO) {
@@ -692,15 +1250,15 @@ __global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
1250 }
1251 }
1252 }
695 - }
1253 + } // warpid == 0
1254 }
1255
1256 // Grid: (num_heads, num_seqs).
699 -template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
700 - int PARTITION_SIZE>
1257 +template <typename scalar_t, typename OUTT, int HEAD_SIZE, int NUM_THREADS,
1258 + int PARTITION_SIZE, int NPAR_LOOPS>
1259 __global__
1260 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
703 - scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size]
1261 + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size]
1262 const float* __restrict__ exp_sums, // [num_seqs, num_heads,
1263 // max_num_partitions]
1264 const float* __restrict__ max_logits, // [num_seqs, num_heads,
@@ -714,18 +1272,13 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1272 const int seq_idx = blockIdx.y;
1273 const int context_len = context_lens[seq_idx];
1274 const int num_partitions = DIVIDE_ROUND_UP(context_len, PARTITION_SIZE);
717 - if (num_partitions == 1) {
718 - // if num_partitions==1, main kernel will write to out directly, no work in
719 - // reduction kernel
720 - return;
721 - }
722 -
1275 constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
1276 const int warpid = threadIdx.x / WARP_SIZE;
1277 const int laneid = threadIdx.x % WARP_SIZE;
1278
1279 __shared__ float shared_global_exp_sum;
728 - __shared__ float shared_exp_sums[2 * WARP_SIZE];
1280 + // max num partitions supported is warp_size * NPAR_LOOPS
1281 + __shared__ float shared_exp_sums[NPAR_LOOPS * WARP_SIZE];
1282
1283 if (warpid == 0) {
1284 const float* max_logits_ptr = max_logits +
@@ -734,14 +1287,25 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1287
1288 // valid partition is the last valid partition in case threadid > num
1289 // partitions
737 - const int valid_partition =
738 - (threadIdx.x < num_partitions) ? threadIdx.x : num_partitions - 1;
739 - const int valid_partition2 = (WARP_SIZE + threadIdx.x < num_partitions)
740 - ? WARP_SIZE + threadIdx.x
741 - : num_partitions - 1;
742 - float reg_max_logit = max_logits_ptr[valid_partition];
743 - float reg_max_logit2 = max_logits_ptr[valid_partition2];
744 - float max_logit = fmaxf(reg_max_logit, reg_max_logit2);
1290 + int valid_partition[NPAR_LOOPS];
1291 + float reg_max_logit[NPAR_LOOPS];
1292 + const int last_valid_partition = num_partitions - 1;
1293 +
1294 + #pragma unroll
1295 + for (int i = 0; i < NPAR_LOOPS; i++) {
1296 + const int partition_no = i * WARP_SIZE + threadIdx.x;
1297 + valid_partition[i] =
1298 + (partition_no < num_partitions) ? partition_no : last_valid_partition;
1299 + }
1300 + #pragma unroll
1301 + for (int i = 0; i < NPAR_LOOPS; i++) {
1302 + reg_max_logit[i] = max_logits_ptr[valid_partition[i]];
1303 + }
1304 + float max_logit = reg_max_logit[0];
1305 + #pragma unroll
1306 + for (int i = 1; i < NPAR_LOOPS; i++) {
1307 + max_logit = fmaxf(max_logit, reg_max_logit[i]);
1308 + }
1309
1310 #pragma unroll
1311 for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
@@ -752,17 +1316,28 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1316 seq_idx * num_heads * max_num_partitions +
1317 head_idx * max_num_partitions;
1318
755 - float global_exp_sum = 0.0f;
756 - float rescaled_exp_sum = exp_sums_ptr[valid_partition];
757 - float rescaled_exp_sum2 = exp_sums_ptr[valid_partition2];
758 - rescaled_exp_sum *=
759 - (threadIdx.x < num_partitions) ? expf(reg_max_logit - max_logit) : 0.0f;
760 - rescaled_exp_sum2 *= (threadIdx.x + WARP_SIZE < num_partitions)
761 - ? expf(reg_max_logit2 - max_logit)
762 - : 0.0f;
763 - global_exp_sum += rescaled_exp_sum + rescaled_exp_sum2;
764 - shared_exp_sums[threadIdx.x] = rescaled_exp_sum;
765 - shared_exp_sums[threadIdx.x + WARP_SIZE] = rescaled_exp_sum2;
1319 + float rescaled_exp_sum[NPAR_LOOPS];
1320 + #pragma unroll
1321 + for (int i = 0; i < NPAR_LOOPS; i++) {
1322 + rescaled_exp_sum[i] = exp_sums_ptr[valid_partition[i]];
1323 + }
1324 + #pragma unroll
1325 + for (int i = 0; i < NPAR_LOOPS; i++) {
1326 + const int partition_no = i * WARP_SIZE + threadIdx.x;
1327 + rescaled_exp_sum[i] *= (partition_no < num_partitions)
1328 + ? expf(reg_max_logit[i] - max_logit)
1329 + : 0.0f;
1330 + }
1331 + float global_exp_sum = rescaled_exp_sum[0];
1332 + #pragma unroll
1333 + for (int i = 1; i < NPAR_LOOPS; i++) {
1334 + global_exp_sum += rescaled_exp_sum[i];
1335 + }
1336 + #pragma unroll
1337 + for (int i = 0; i < NPAR_LOOPS; i++) {
1338 + const int partition_no = i * WARP_SIZE + threadIdx.x;
1339 + shared_exp_sums[partition_no] = rescaled_exp_sum[i];
1340 + }
1341
1342 #pragma unroll
1343 for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
@@ -839,82 +1414,117 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1414 }
1415 }
1416
842 - if (num_partitions > MAX_NPAR) {
843 - idx = 0;
1417 + for (int p = 1; p < NPAR_LOOPS; p++) {
1418 + if (num_partitions > p * MAX_NPAR) {
1419 + idx = 0;
1420 #pragma unroll
845 - for (int j = MAX_NPAR * HEAD_SIZE; j < 2 * MAX_NPAR * HEAD_SIZE;
846 - j += HEAD_SIZE) {
847 - // lastj is last valid partition
848 - const int lastj_offset =
849 - (j < num_partition_offset) ? j : last_partition_offset;
850 - tmps[idx] = tmp_out_ptr[lastj_offset];
851 - idx++;
852 - }
1421 + for (int j = p * MAX_NPAR * HEAD_SIZE; j < (p + 1) * MAX_NPAR * HEAD_SIZE;
1422 + j += HEAD_SIZE) {
1423 + // lastj is last valid partition
1424 + const int lastj_offset =
1425 + (j < num_partition_offset) ? j : last_partition_offset;
1426 + tmps[idx] = tmp_out_ptr[lastj_offset];
1427 + idx++;
1428 + }
1429
1430 #pragma unroll
855 - for (int j = 0; j < MAX_NPAR; j++) {
856 - acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j + MAX_NPAR];
1431 + for (int j = 0; j < MAX_NPAR; j++) {
1432 + acc += to_float<scalar_t>(tmps[j]) * shared_exp_sums[j + p * MAX_NPAR];
1433 + }
1434 }
1435 }
1436
1437 const float inv_global_exp_sum =
1438 __fdividef(1.0f, shared_global_exp_sum + 1e-6f);
1439 acc *= inv_global_exp_sum;
863 - scalar_t* out_ptr =
864 - out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
865 - out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
1440 +
1441 + OUTT* out_ptr = out + static_cast<int64_t>(seq_idx) * num_heads * HEAD_SIZE +
1442 + static_cast<int64_t>(head_idx) * HEAD_SIZE;
1443 + if constexpr (std::is_same<OUTT, bit8_t>::value) {
1444 + out_ptr[threadIdx.x] =
1445 + __hip_cvt_float_to_fp8(acc, vllm::fp8::fp8_type::__default_saturation,
1446 + vllm::fp8::fp8_type::__default_interpret);
1447 + } else {
1448 + out_ptr[threadIdx.x] = from_float<scalar_t>(acc);
1449 + }
1450 }
1451
1452 #else // !defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1453
1454 +// clang-format off
1455 template <typename scalar_t, typename cache_t,
871 - vllm::Fp8KVCacheDataType KV_DTYPE, int BLOCK_SIZE, int HEAD_SIZE,
872 - int NUM_THREADS,
1456 + vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
1457 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
1458 int GQA_RATIO>
874 -__global__ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_kernel(
875 - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
876 - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads,
877 - // head_size/x, block_size, x]
878 - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads,
879 - // head_size, block_size]
880 - const int num_kv_heads, const float scale,
881 - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
882 - const int* __restrict__ context_lens, // [num_seqs]
1459 +__global__
1460 +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma16_kernel(
1461 + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
1462 + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
1463 + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
1464 + const int num_kv_heads,
1465 + const float scale,
1466 + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
1467 + const int* __restrict__ context_lens, // [num_seqs]
1468 const int max_num_blocks_per_seq,
1469 const float* __restrict__ alibi_slopes, // [num_heads]
885 - const int q_stride, const int kv_block_stride, const int kv_head_stride,
886 - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
887 - float* __restrict__ max_logits, // [num_seqs, num_heads,
888 - // max_num_partitions]
889 - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions,
890 - // head_size]
891 - scalar_t* __restrict__ final_out, // [num_seqs, num_heads, head_size]
1470 + const int q_stride,
1471 + const int kv_block_stride,
1472 + const int kv_head_stride,
1473 + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1474 + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
1475 + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
1476 + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
1477 + int max_ctx_blocks, const float* k_scale, const float* v_scale) {
1478 + UNREACHABLE_CODE
1479 +}
1480 +
1481 +template <typename scalar_t, typename cache_t,
1482 + vllm::Fp8KVCacheDataType KV_DTYPE, typename OUTT, int BLOCK_SIZE,
1483 + int HEAD_SIZE, int NUM_THREADS, bool ALIBI_ENABLED,
1484 + int GQA_RATIO>
1485 +__global__
1486 +__launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_QKV_mfma4_kernel(
1487 + const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size]
1488 + const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, head_size/x, block_size, x]
1489 + const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, head_size, block_size]
1490 + const int num_kv_heads,
1491 + const float scale,
1492 + const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
1493 + const int* __restrict__ context_lens, // [num_seqs]
1494 + const int max_num_blocks_per_seq,
1495 + const float* __restrict__ alibi_slopes, // [num_heads]
1496 + const int q_stride,
1497 + const int kv_block_stride,
1498 + const int kv_head_stride,
1499 + float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1500 + float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
1501 + scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, head_size]
1502 + OUTT* __restrict__ final_out, // [num_seqs, num_heads, head_size]
1503 int max_ctx_blocks, const float* k_scale, const float* v_scale) {
1504 UNREACHABLE_CODE
1505 }
1506
1507 // Grid: (num_heads, num_seqs).
897 -template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
898 - int PARTITION_SIZE>
1508 +template <typename scalar_t, typename OUTT, int HEAD_SIZE, int NUM_THREADS,
1509 + int PARTITION_SIZE, int NPAR_LOOPS>
1510 __global__
1511 __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
901 - scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size]
902 - const float* __restrict__ exp_sums, // [num_seqs, num_heads,
903 - // max_num_partitions]
904 - const float* __restrict__ max_logits, // [num_seqs, num_heads,
905 - // max_num_partitions]
906 - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads,
907 - // max_num_partitions, head_size]
1512 + OUTT* __restrict__ out, // [num_seqs, num_heads, head_size]
1513 + const float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions]
1514 + const float* __restrict__ max_logits, // [num_seqs, num_heads, max_num_partitions]
1515 + const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
1516 const int* __restrict__ context_lens, // [num_seqs]
1517 const int max_num_partitions) {
1518 UNREACHABLE_CODE
1519 }
1520 +// clang-format on
1521
1522 #endif // defined(__HIP__MI300_MI250__) TODO: Add NAVI support
1523
915 -#define LAUNCH_CUSTOM_ATTENTION(GQA_RATIO) \
916 - paged_attention_ll4mi_QKV_kernel<T, KVT, KV_DTYPE, BLOCK_SIZE, HEAD_SIZE, \
917 - NTHR, GQA_RATIO> \
1524 +#define LAUNCH_CUSTOM_ATTENTION_MFMA16(GQA_RATIO) \
1525 + paged_attention_ll4mi_QKV_mfma16_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
1526 + HEAD_SIZE, NTHR, ALIBI_ENABLED, \
1527 + GQA_RATIO> \
1528 <<<grid, block, 0, stream>>>( \
1529 query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
1530 block_tables_ptr, context_lens_ptr, max_num_blocks_per_seq, \
@@ -922,8 +1532,27 @@ __launch_bounds__(NUM_THREADS) void paged_attention_ll4mi_reduce_kernel(
1532 exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, max_ctx_blocks, \
1533 k_scale_ptr, v_scale_ptr);
1534
1535 +#define LAUNCH_CUSTOM_ATTENTION_MFMA4(GQA_RATIO) \
1536 + paged_attention_ll4mi_QKV_mfma4_kernel<T, KVT, KV_DTYPE, OUTT, BLOCK_SIZE, \
1537 + HEAD_SIZE, NTHR, ALIBI_ENABLED, \
1538 + GQA_RATIO> \
1539 + <<<grid, block, 0, stream>>>( \
1540 + query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, scale, \
1541 + block_tables_ptr, context_lens_ptr, max_num_blocks_per_seq, \
1542 + alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \
1543 + exp_sums_ptr, max_logits_ptr, tmp_out_ptr, out_ptr, max_ctx_blocks, \
1544 + k_scale_ptr, v_scale_ptr);
1545 +
1546 +#define LAUNCH_CUSTOM_REDUCTION(NPAR_LOOPS) \
1547 + paged_attention_ll4mi_reduce_kernel<T, OUTT, HEAD_SIZE, HEAD_SIZE, \
1548 + PARTITION_SIZE, NPAR_LOOPS> \
1549 + <<<reduce_grid, reduce_block, 0, stream>>>( \
1550 + out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, \
1551 + context_lens_ptr, max_num_partitions);
1552 +
1553 template <typename T, typename KVT, vllm::Fp8KVCacheDataType KV_DTYPE,
926 - int BLOCK_SIZE, int HEAD_SIZE, int PARTITION_SIZE = 512>
1554 + int BLOCK_SIZE, int HEAD_SIZE, typename OUTT, int PARTITION_SIZE_OLD,
1555 + bool ALIBI_ENABLED>
1556 void paged_attention_custom_launcher(
1557 torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
1558 torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
@@ -945,7 +1574,6 @@ void paged_attention_custom_launcher(
1574 ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
1575 : nullptr;
1576
948 - T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
1577 float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
1578 float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
1579 T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
@@ -956,109 +1584,143 @@ void paged_attention_custom_launcher(
1584 int* context_lens_ptr = context_lens.data_ptr<int>();
1585 const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
1586 const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
1587 + OUTT* out_ptr = reinterpret_cast<OUTT*>(out.data_ptr());
1588
1589 const int max_ctx_blocks = DIVIDE_ROUND_UP(max_context_len, BLOCK_SIZE);
1590 +
1591 + // partition size is fixed at 256 since both mfma4 and mfma16 kernels support
1592 + // it mfma4 kernel also supports partition size 512
1593 + constexpr int PARTITION_SIZE = 256;
1594 const int max_num_partitions =
1595 DIVIDE_ROUND_UP(max_context_len, PARTITION_SIZE);
1596 const int gqa_ratio = num_heads / num_kv_heads;
1597 assert(num_heads % num_kv_heads == 0);
1598 assert(head_size == HEAD_SIZE);
966 - assert(max_num_partitions <= 128);
1599
968 - constexpr int NTHR = PARTITION_SIZE;
1600 + constexpr int NTHR = 256;
1601 dim3 grid(num_seqs, max_num_partitions, num_kv_heads);
1602 dim3 block(NTHR);
1603 const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
1604 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
1605 +
1606 + // mfma4 kernel is faster than mfma16 for gqa_ratio <= 4
1607 switch (gqa_ratio) {
1608 case 1:
975 - LAUNCH_CUSTOM_ATTENTION(1);
1609 + LAUNCH_CUSTOM_ATTENTION_MFMA4(1);
1610 break;
1611 case 2:
978 - LAUNCH_CUSTOM_ATTENTION(2);
1612 + LAUNCH_CUSTOM_ATTENTION_MFMA4(2);
1613 break;
1614 case 3:
981 - LAUNCH_CUSTOM_ATTENTION(3);
1615 + LAUNCH_CUSTOM_ATTENTION_MFMA4(3);
1616 break;
1617 case 4:
984 - LAUNCH_CUSTOM_ATTENTION(4);
1618 + LAUNCH_CUSTOM_ATTENTION_MFMA4(4);
1619 break;
1620 case 5:
987 - LAUNCH_CUSTOM_ATTENTION(5);
1621 + LAUNCH_CUSTOM_ATTENTION_MFMA16(5);
1622 break;
1623 case 6:
990 - LAUNCH_CUSTOM_ATTENTION(6);
1624 + LAUNCH_CUSTOM_ATTENTION_MFMA16(6);
1625 break;
1626 case 7:
993 - LAUNCH_CUSTOM_ATTENTION(7);
1627 + LAUNCH_CUSTOM_ATTENTION_MFMA16(7);
1628 break;
1629 case 8:
996 - LAUNCH_CUSTOM_ATTENTION(8);
1630 + LAUNCH_CUSTOM_ATTENTION_MFMA16(8);
1631 break;
1632 case 9:
999 - LAUNCH_CUSTOM_ATTENTION(9);
1633 + LAUNCH_CUSTOM_ATTENTION_MFMA16(9);
1634 break;
1635 case 10:
1002 - LAUNCH_CUSTOM_ATTENTION(10);
1636 + LAUNCH_CUSTOM_ATTENTION_MFMA16(10);
1637 break;
1638 case 11:
1005 - LAUNCH_CUSTOM_ATTENTION(11);
1639 + LAUNCH_CUSTOM_ATTENTION_MFMA16(11);
1640 break;
1641 case 12:
1008 - LAUNCH_CUSTOM_ATTENTION(12);
1642 + LAUNCH_CUSTOM_ATTENTION_MFMA16(12);
1643 break;
1644 case 13:
1011 - LAUNCH_CUSTOM_ATTENTION(13);
1645 + LAUNCH_CUSTOM_ATTENTION_MFMA16(13);
1646 break;
1647 case 14:
1014 - LAUNCH_CUSTOM_ATTENTION(14);
1648 + LAUNCH_CUSTOM_ATTENTION_MFMA16(14);
1649 break;
1650 case 15:
1017 - LAUNCH_CUSTOM_ATTENTION(15);
1651 + LAUNCH_CUSTOM_ATTENTION_MFMA16(15);
1652 break;
1653 case 16:
1020 - LAUNCH_CUSTOM_ATTENTION(16);
1654 + LAUNCH_CUSTOM_ATTENTION_MFMA16(16);
1655 break;
1656 default:
1657 TORCH_CHECK(false, "Unsupported gqa ratio: ", gqa_ratio);
1658 break;
1659 }
1026 - // dim3 grid2(num_heads,num_seqs,head_size/HEAD_ELEMS_PER_WG);
1027 - // dim3 block2(1024);
1028 - // LAUNCH_CUSTOM_ATTENTION2;
1029 -
1030 - // reduction kernel is only required if max_context_len > partition size,
1031 - // otherwise main kernel writes directly to final output
1032 - // note there are cases with graphing where max_context_len is the max
1033 - // supported by graphing, not the actual max among all the sequences: in that
1034 - // case reduction kernel will still run but return immediately
1035 - if (max_context_len > PARTITION_SIZE) {
1036 - dim3 reduce_grid(num_heads, num_seqs);
1037 - dim3 reduce_block(head_size);
1038 - paged_attention_ll4mi_reduce_kernel<T, HEAD_SIZE, HEAD_SIZE, PARTITION_SIZE>
1039 - <<<reduce_grid, reduce_block, 0, stream>>>(
1040 - out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr,
1041 - context_lens_ptr, max_num_partitions);
1660 +
1661 + dim3 reduce_grid(num_heads, num_seqs);
1662 + dim3 reduce_block(head_size);
1663 + const int npar_loops = DIVIDE_ROUND_UP(max_num_partitions, WARP_SIZE);
1664 + // reduction kernel supports upto 8 NPAR_loops * 64 (warp_size) * 256
1665 + // (partition size) = 128K context length
1666 + switch (npar_loops) {
1667 + case 1:
1668 + LAUNCH_CUSTOM_REDUCTION(1);
1669 + break;
1670 + case 2:
1671 + LAUNCH_CUSTOM_REDUCTION(2);
1672 + break;
1673 + case 3:
1674 + LAUNCH_CUSTOM_REDUCTION(3);
1675 + break;
1676 + case 4:
1677 + LAUNCH_CUSTOM_REDUCTION(4);
1678 + break;
1679 + case 5:
1680 + LAUNCH_CUSTOM_REDUCTION(5);
1681 + break;
1682 + case 6:
1683 + LAUNCH_CUSTOM_REDUCTION(6);
1684 + break;
1685 + case 7:
1686 + LAUNCH_CUSTOM_REDUCTION(7);
1687 + break;
1688 + case 8:
1689 + LAUNCH_CUSTOM_REDUCTION(8);
1690 + break;
1691 + default:
1692 + TORCH_CHECK(false, "Unsupported npar_loops: ", npar_loops);
1693 + break;
1694 }
1695 }
1696
1045 -#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE) \
1046 - paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE>( \
1047 - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
1048 - num_kv_heads, scale, block_tables, context_lens, max_context_len, \
1697 +#define CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, PSIZE, \
1698 + ALIBI_ENABLED) \
1699 + paged_attention_custom_launcher<T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, T, \
1700 + PSIZE, ALIBI_ENABLED>( \
1701 + out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \
1702 + num_kv_heads, scale, block_tables, context_lens, max_context_len, \
1703 alibi_slopes, k_scale, v_scale);
1704
1051 -#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE) \
1052 - switch (block_size) { \
1053 - case 16: \
1054 - CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, 16, HEAD_SIZE); \
1055 - break; \
1056 - case 32: \
1057 - CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, 32, HEAD_SIZE); \
1058 - break; \
1059 - default: \
1060 - TORCH_CHECK(false, "Unsupported block size: ", block_size); \
1061 - break; \
1705 +#define CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, \
1706 + PSIZE) \
1707 + if (alibi_slopes) { \
1708 + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, PSIZE, true); \
1709 + } else { \
1710 + CALL_CUSTOM_LAUNCHER(T, KVT, KV_DTYPE, BLK_SIZE, HEAD_SIZE, PSIZE, false); \
1711 + }
1712 +
1713 +#define CALL_CUSTOM_LAUNCHER_BLK(T, KVT, KV_DTYPE, HEAD_SIZE) \
1714 + switch (block_size) { \
1715 + case 16: \
1716 + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, 16, HEAD_SIZE, 256); \
1717 + break; \
1718 + case 32: \
1719 + CALL_CUSTOM_LAUNCHER_ALIBI(T, KVT, KV_DTYPE, 32, HEAD_SIZE, 256); \
1720 + break; \
1721 + default: \
1722 + TORCH_CHECK(false, "Unsupported block size: ", block_size); \
1723 + break; \
1724 }
1725
1726 #define CALL_CUSTOM_LAUNCHER_BLK_HEAD(T, KVT, KV_DTYPE) \
@@ -1074,24 +1736,24 @@ void paged_attention_custom_launcher(
1736 break; \
1737 }
1738
1739 +// clang-format off
1740 void paged_attention(
1741 torch::Tensor& out, // [num_seqs, num_heads, head_size]
1742 torch::Tensor& exp_sums, // [num_seqs, num_heads, max_num_partitions]
1743 torch::Tensor& max_logits, // [num_seqs, num_heads, max_num_partitions]
1081 - torch::Tensor&
1082 - tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
1083 - torch::Tensor& query, // [num_seqs, num_heads, head_size]
1084 - torch::Tensor&
1085 - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
1086 - torch::Tensor&
1087 - value_cache, // [num_blocks, num_heads, head_size, block_size]
1088 - int64_t num_kv_heads, double scale,
1089 - torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
1090 - torch::Tensor& context_lens, // [num_seqs]
1744 + torch::Tensor& tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size]
1745 + torch::Tensor& query, // [num_seqs, num_heads, head_size]
1746 + torch::Tensor& key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
1747 + torch::Tensor& value_cache, // [num_blocks, num_heads, head_size, block_size]
1748 + int64_t num_kv_heads,
1749 + double scale,
1750 + torch::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq]
1751 + torch::Tensor& context_lens, // [num_seqs]
1752 int64_t block_size, int64_t max_context_len,
1753 const std::optional<torch::Tensor>& alibi_slopes,
1754 const std::string& kv_cache_dtype, torch::Tensor& k_scale,
1755 torch::Tensor& v_scale) {
1756 + // clang-format on
1757 const int head_size = query.size(2);
1758 if (kv_cache_dtype == "auto") {
1759 if (query.dtype() == at::ScalarType::Half) {
requirements-rocm.txt
+1 -1
@@ -11,4 +11,4 @@ peft
11 pytest-asyncio
12 tensorizer>=2.9.0
13 runai-model-streamer==0.11.0
14 -runai-model-streamer-s3==0.11.0
14 +runai-model-streamer-s3==0.11.0
\ No newline at end of file
tests/kernels/test_attention.py
+7 -1
@@ -25,6 +25,7 @@ MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512
25 # Reduce NUM_BLOCKS when it happens.
26 NUM_BLOCKS = 4321 # Arbitrary values for testing
27 PARTITION_SIZE = 512
28 +PARTITION_SIZE_ROCM = 256
29 # flshattF and tritonflashattF supported: {torch.float16, torch.bfloat16}
30 DTYPES = [
31 torch.half, torch.bfloat16, torch.float
@@ -146,6 +147,8 @@ def test_paged_attention(
147 or (version == "rocm" and head_size not in (64, 128))):
148 pytest.skip()
149
150 + global PARTITION_SIZE
151 +
152 current_platform.seed_everything(seed)
153 torch.set_default_device(device)
154 scale = float(1.0 / (head_size**0.5))
@@ -214,6 +217,9 @@ def test_paged_attention(
217 and block_size == BLOCK_SIZES[0]))
218
219 elif version in ("v2", "rocm"):
220 + if current_platform.is_rocm() and version == "rocm":
221 + PARTITION_SIZE = PARTITION_SIZE_ROCM
222 +
223 num_partitions = ((max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE)
224 assert PARTITION_SIZE % block_size == 0
225 num_seqs, num_heads, head_size = output.shape
@@ -432,4 +438,4 @@ def test_multi_query_kv_attention(
438 )
439 atol = get_default_atol(output) if current_platform.is_rocm() else 1e-3
440 rtol = get_default_rtol(output) if current_platform.is_rocm() else 1e-5
435 - torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
441 + torch.testing.assert_close(output, ref_output, atol=atol, rtol=rtol)
\ No newline at end of file
vllm/attention/backends/rocm_flash_attn.py
+2 -2
@@ -22,7 +22,7 @@ if TYPE_CHECKING:
22
23 logger = init_logger(__name__)
24
25 -_PARTITION_SIZE_ROCM = 512
25 +_PARTITION_SIZE_ROCM = 256
26 _GPU_ARCH = torch.cuda.get_device_properties("cuda").gcnArchName
27 _ON_NAVI = "gfx1" in _GPU_ARCH
28 _ON_MI250_MI300 = any(arch in _GPU_ARCH for arch in ["gfx90a", "gfx942"])
@@ -885,4 +885,4 @@ def _use_rocm_custom_paged_attention(qtype: torch.dtype, head_size: int,
885 and (qtype == torch.half or qtype == torch.bfloat16)
886 and (head_size == 64 or head_size == 128)
887 and (block_size == 16 or block_size == 32)
888 - and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768)
888 + and (gqa_ratio >= 1 and gqa_ratio <= 16) and max_seq_len <= 32768)
\ No newline at end of file