From d715b3aa1ea6af3f663eb6d3cd8f5b6bb15770e9 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 2 Jul 2026 15:31:26 -0400 Subject: [PATCH] Delete PagedAttention (#47361) Signed-off-by: mgoin --- CMakeLists.txt | 2 - .../kernels/benchmark_paged_attention.py | 131 +--- .../attention/attention_kernels.cuh | 667 ------------------ .../attention/paged_attention_v1.cu | 190 ----- .../attention/paged_attention_v2.cu | 202 ------ csrc/libtorch_stable/ops.h | 26 - csrc/libtorch_stable/torch_bindings.cpp | 30 - tests/kernels/attention/test_attention.py | 218 ++---- vllm/_custom_ops.py | 94 --- 9 files changed, 89 insertions(+), 1471 deletions(-) delete mode 100644 csrc/libtorch_stable/attention/attention_kernels.cuh delete mode 100644 csrc/libtorch_stable/attention/paged_attention_v1.cu delete mode 100644 csrc/libtorch_stable/attention/paged_attention_v2.cu diff --git a/CMakeLists.txt b/CMakeLists.txt index cf215a6f249..901f2be6bbb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -399,8 +399,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "csrc/libtorch_stable/sampler.cu" "csrc/libtorch_stable/topk.cu" "csrc/libtorch_stable/mamba/selective_scan_fwd.cu" - "csrc/libtorch_stable/attention/paged_attention_v1.cu" - "csrc/libtorch_stable/attention/paged_attention_v2.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels.cu" "csrc/libtorch_stable/cache_kernels_fused.cu" diff --git a/benchmarks/kernels/benchmark_paged_attention.py b/benchmarks/kernels/benchmark_paged_attention.py index b6a0b7ad8ca..f4249d6270a 100644 --- a/benchmarks/kernels/benchmark_paged_attention.py +++ b/benchmarks/kernels/benchmark_paged_attention.py @@ -19,13 +19,11 @@ from vllm.utils.torch_utils import ( logger = init_logger(__name__) NUM_BLOCKS = 128 * 1024 -PARTITION_SIZE = 512 PARTITION_SIZE_ROCM = 256 @torch.inference_mode() def main( - version: str, num_seqs: int, seq_len: int, num_query_heads: int, @@ -82,27 +80,20 @@ def main( # Prepare for the paged attention kernel. output = torch.empty_like(query) - if version == "v2": - if current_platform.is_rocm(): - global PARTITION_SIZE - if not args.custom_paged_attn and not current_platform.is_navi(): - PARTITION_SIZE = 1024 - else: - PARTITION_SIZE = PARTITION_SIZE_ROCM - num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE - tmp_output = torch.empty( - size=(num_seqs, num_query_heads, num_partitions, head_size), - dtype=output.dtype, - device=output.device, - ) - exp_sums = torch.empty( - size=(num_seqs, num_query_heads, num_partitions), - dtype=torch.float32, - device=output.device, - ) - max_logits = torch.empty_like(exp_sums) + num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM + tmp_output = torch.empty( + size=(num_seqs, num_query_heads, num_partitions, head_size), + dtype=output.dtype, + device=output.device, + ) + exp_sums = torch.empty( + size=(num_seqs, num_query_heads, num_partitions), + dtype=torch.float32, + device=output.device, + ) + max_logits = torch.empty_like(exp_sums) - def run_cuda_benchmark(num_iters: int, profile: bool = False) -> float: + def run_benchmark(num_iters: int, profile: bool = False) -> float: torch.accelerator.synchronize() if profile: torch.cuda.cudart().cudaProfilerStart() @@ -112,67 +103,26 @@ def main( k_scale = v_scale = torch.tensor(1.0, dtype=torch.float32, device=device) for _ in range(num_iters): - if version == "v1": - ops.paged_attention_v1( - output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - elif version == "v2": - if not args.custom_paged_attn: - ops.paged_attention_v2( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - else: - ops.paged_attention_rocm( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - None, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - else: - raise ValueError(f"Invalid version: {version}") + ops.paged_attention_rocm( + output, + exp_sums, + max_logits, + tmp_output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + None, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + ) torch.accelerator.synchronize() end_time = time.perf_counter() @@ -182,7 +132,6 @@ def main( # Warmup. print("Warming up...") - run_benchmark = run_cuda_benchmark run_benchmark(num_iters=3, profile=False) # Benchmark. @@ -195,12 +144,13 @@ def main( if __name__ == "__main__": logger.warning( - "This script benchmarks the paged attention kernel. " + "This script benchmarks the ROCm paged attention kernel. " "By default this is no longer used in vLLM inference." ) + if not current_platform.is_rocm(): + raise RuntimeError("This benchmark requires the ROCm platform.") parser = FlexibleArgumentParser(description="Benchmark the paged attention kernel.") - parser.add_argument("--version", type=str, choices=["v1", "v2"], default="v2") parser.add_argument("--batch-size", type=int, default=8) parser.add_argument("--seq-len", type=int, default=4096) parser.add_argument("--num-query-heads", type=int, default=64) @@ -208,7 +158,7 @@ if __name__ == "__main__": parser.add_argument( "--head-size", type=int, - choices=[64, 80, 96, 112, 120, 128, 192, 256], + choices=[64, 128], default=128, ) parser.add_argument("--block-size", type=int, choices=[16, 32], default=16) @@ -224,11 +174,7 @@ if __name__ == "__main__": choices=["auto", "fp8", "fp8_e5m2", "fp8_e4m3"], default="auto", help="Data type for kv cache storage. If 'auto', will use model " - "data type. CUDA 11.8+ supports fp8 (=fp8_e4m3) and fp8_e5m2. " - "ROCm (AMD GPU) supports fp8 (=fp8_e4m3)", - ) - parser.add_argument( - "--custom-paged-attn", action="store_true", help="Use custom paged attention" + "data type. ROCm (AMD GPU) supports fp8 (=fp8_e4m3)", ) args = parser.parse_args() print(args) @@ -236,7 +182,6 @@ if __name__ == "__main__": if args.num_query_heads % args.num_kv_heads != 0: raise ValueError("num_query_heads must be divisible by num_kv_heads") main( - version=args.version, num_seqs=args.batch_size, seq_len=args.seq_len, num_query_heads=args.num_query_heads, diff --git a/csrc/libtorch_stable/attention/attention_kernels.cuh b/csrc/libtorch_stable/attention/attention_kernels.cuh deleted file mode 100644 index c5f9a9876c3..00000000000 --- a/csrc/libtorch_stable/attention/attention_kernels.cuh +++ /dev/null @@ -1,667 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include - -#include "../../attention/attention_dtypes.h" -#include "attention_utils.cuh" -#include "../../cuda_compat.h" - -#ifdef USE_ROCM - #include - #include "../../quantization/w8a8/fp8/amd/quant_utils.cuh" -typedef __hip_bfloat16 __nv_bfloat16; -#else - #include "../../quantization/w8a8/fp8/nvidia/quant_utils.cuh" -#endif - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -namespace vllm { - -// Utility function for attention softmax. -template -inline __device__ float block_sum(float* red_smem, float sum) { - // Decompose the thread index into warp / lane. - int warp = threadIdx.x / WARP_SIZE; - int lane = threadIdx.x % WARP_SIZE; - - // Compute the sum per warp. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - sum += VLLM_SHFL_XOR_SYNC(sum, mask); - } - - // Warp leaders store the data to shared memory. - if (lane == 0) { - red_smem[warp] = sum; - } - - // Make sure the data is in shared memory. - __syncthreads(); - - // The warps compute the final sums. - if (lane < NUM_WARPS) { - sum = red_smem[lane]; - } - - // Parallel reduction inside the warp. -#pragma unroll - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) { - sum += VLLM_SHFL_XOR_SYNC(sum, mask); - } - - // Broadcast to other threads. - return VLLM_SHFL_SYNC(sum, 0); -} - -// TODO(woosuk): Merge the last two dimensions of the grid. -// Grid: (num_heads, num_seqs, max_num_partitions). -template // Zero means no partitioning. -__device__ void paged_attention_kernel( - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ out, // [num_seqs, num_heads, max_num_partitions, - // head_size] - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, // [num_heads] - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - const float* k_scale, const float* v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { - const int seq_idx = blockIdx.y; - const int partition_idx = blockIdx.z; - const int max_num_partitions = gridDim.z; - constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0; - const int seq_len = seq_lens[seq_idx]; - if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) { - // No work to do. Terminate the thread block. - return; - } - - const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE); - const int num_blocks_per_partition = - USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks; - - // [start_block_idx, end_block_idx) is the range of blocks to process. - const int start_block_idx = - USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0; - const int end_block_idx = - MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks); - const int num_blocks = end_block_idx - start_block_idx; - - // [start_token_idx, end_token_idx) is the range of tokens to process. - const int start_token_idx = start_block_idx * BLOCK_SIZE; - const int end_token_idx = - MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len); - const int num_tokens = end_token_idx - start_token_idx; - - constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1); - constexpr int NUM_THREAD_GROUPS = - NUM_THREADS / THREAD_GROUP_SIZE; // Note: This assumes THREAD_GROUP_SIZE - // divides NUM_THREADS - assert(NUM_THREADS % THREAD_GROUP_SIZE == 0); - constexpr int NUM_TOKENS_PER_THREAD_GROUP = - DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE); - constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE; - const int thread_idx = threadIdx.x; - const int warp_idx = thread_idx / WARP_SIZE; - const int lane = thread_idx % WARP_SIZE; - - const int head_idx = blockIdx.x; - const int num_heads = gridDim.x; - const int num_queries_per_kv = num_heads / num_kv_heads; - const int kv_head_idx = head_idx / num_queries_per_kv; - const float alibi_slope = - alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx]; - - // A vector type to store a part of a key or a query. - // The vector size is configured in such a way that the threads in a thread - // group fetch or compute 16 bytes at a time. For example, if the size of a - // thread group is 4 and the data type is half, then the vector size is 16 / - // (4 * sizeof(half)) == 2. - constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1); - using K_vec = typename Vec::Type; - using Q_vec = typename Vec::Type; - using Quant_vec = typename Vec::Type; - - constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE; - constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE; - - const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE; - const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE; - - // Load the query to registers. - // Each thread in a thread group has a different part of the query. - // For example, if the thread group size is 4, then the first thread in - // the group has 0, 4, 8, ... th vectors of the query, and the second thread - // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because - // q is split from a qkv tensor, it may not be contiguous. - const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE; - __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD]; -#pragma unroll - for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD; - i += NUM_THREAD_GROUPS) { - const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE; - q_vecs[thread_group_offset][i] = - *reinterpret_cast(q_ptr + vec_idx * VEC_SIZE); - } - __syncthreads(); // TODO(naed90): possible speedup if this is replaced with a - // memory wall right before we use q_vecs - - // Memory planning. - extern __shared__ char shared_mem[]; - // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy. - float* logits = reinterpret_cast(shared_mem); - // Workspace for reduction. - __shared__ float red_smem[2 * NUM_WARPS]; - - // x == THREAD_GROUP_SIZE * VEC_SIZE - // Each thread group fetches x elements from the key at a time. - constexpr int x = 16 / sizeof(cache_t); - float qk_max = -FLT_MAX; - - // Iterate over the key blocks. - // Each warp fetches a block of keys for each iteration. - // Each thread group in a warp fetches a key from the block, and computes - // dot product with the query. - const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq; - - // blocksparse specific vars - int bs_block_offset; - int q_bs_block_id; - if constexpr (IS_BLOCK_SPARSE) { - // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len, - // blocksparse_block_size); - q_bs_block_id = (seq_len - 1) / blocksparse_block_size; - if (blocksparse_head_sliding_step >= 0) - // sliding on q heads - bs_block_offset = - (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1; - else - // sliding on kv heads - bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) * - (-blocksparse_head_sliding_step) + - 1; - } - - for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; - block_idx += NUM_WARPS) { - // NOTE(woosuk): The block number is stored in int32. However, we cast it to - // int64 because int32 can lead to overflow when this variable is multiplied - // by large numbers (e.g., kv_block_stride). - // For blocksparse attention: skip computation on blocks that are not - // attended - if constexpr (IS_BLOCK_SPARSE) { - const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size; - const bool is_remote = - ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0); - const bool is_local = - (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks); - if (!is_remote && !is_local) { - for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { - const int physical_block_offset = - (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; - const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; - - if (thread_group_offset == 0) { - // NOTE(linxihui): assign very large number to skipped tokens to - // avoid contribution to the sumexp softmax normalizer. This will - // not be used at computing sum(softmax*v) as the blocks will be - // skipped. - logits[token_idx - start_token_idx] = -FLT_MAX; - } - } - continue; - } - } - const int64_t physical_block_number = - static_cast(block_table[block_idx]); - - // Load a key to registers. - // Each thread in a thread group has a different part of the key. - // For example, if the thread group size is 4, then the first thread in - // the group has 0, 4, 8, ... th vectors of the key, and the second thread - // has 1, 5, 9, ... th vectors of the key, and so on. - for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) { - const int physical_block_offset = - (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE; - const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; - K_vec k_vecs[NUM_VECS_PER_THREAD]; - -#pragma unroll - for (int j = 0; j < NUM_VECS_PER_THREAD; j++) { - const cache_t* k_ptr = - k_cache + physical_block_number * kv_block_stride + - kv_head_idx * kv_head_stride + physical_block_offset * x; - const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE; - const int offset1 = (vec_idx * VEC_SIZE) / x; - const int offset2 = (vec_idx * VEC_SIZE) % x; - - if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) { - k_vecs[j] = *reinterpret_cast( - k_ptr + offset1 * BLOCK_SIZE * x + offset2); - } else { - // Vector conversion from Quant_vec to K_vec. - Quant_vec k_vec_quant = *reinterpret_cast( - k_ptr + offset1 * BLOCK_SIZE * x + offset2); - k_vecs[j] = fp8::scaled_convert( - k_vec_quant, *k_scale); - } - } - - // Compute dot product. - // This includes a reduction across the threads in the same thread group. - float qk = scale * Qk_dot::dot( - q_vecs[thread_group_offset], k_vecs); - // Add the ALiBi bias if slopes are given. - qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0; - - if (thread_group_offset == 0) { - // Store the partial reductions to shared memory. - // NOTE(woosuk): It is required to zero out the masked logits. - const bool mask = token_idx >= seq_len; - logits[token_idx - start_token_idx] = mask ? 0.f : qk; - // Update the max value. - qk_max = mask ? qk_max : fmaxf(qk_max, qk); - } - } - } - - // Perform reduction across the threads in the same warp to get the - // max qk value for each "warp" (not across the thread block yet). - // The 0-th thread of each thread group already has its max qk value. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) { - qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); - } - if (lane == 0) { - red_smem[warp_idx] = qk_max; - } - __syncthreads(); - - // TODO(woosuk): Refactor this part. - // Get the max qk value for the sequence. - qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX; -#pragma unroll - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) { - qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask)); - } - // Broadcast the max qk value to all threads. - qk_max = VLLM_SHFL_SYNC(qk_max, 0); - - // Get the sum of the exp values. - float exp_sum = 0.f; - for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) { - float val = __expf(logits[i] - qk_max); - logits[i] = val; - exp_sum += val; - } - exp_sum = block_sum(&red_smem[NUM_WARPS], exp_sum); - - // Compute softmax. - const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f); - for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) { - logits[i] *= inv_sum; - } - __syncthreads(); - - // If partitioning is enabled, store the max logit and exp_sum. - if (USE_PARTITIONING && thread_idx == 0) { - float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions + partition_idx; - *max_logits_ptr = qk_max; - float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions + partition_idx; - *exp_sums_ptr = exp_sum; - } - - // Each thread will fetch 16 bytes from the value cache at a time. - constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE); - using V_vec = typename Vec::Type; - using L_vec = typename Vec::Type; - using V_quant_vec = typename Vec::Type; - using Float_L_vec = typename FloatVec::Type; - - constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE; - constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW; - constexpr int NUM_ROWS_PER_THREAD = - DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER); - - // NOTE(woosuk): We use FP32 for the accumulator for better accuracy. - float accs[NUM_ROWS_PER_THREAD]; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - accs[i] = 0.f; - } - - scalar_t zero_value; - zero(zero_value); - for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; - block_idx += NUM_WARPS) { - // NOTE(woosuk): The block number is stored in int32. However, we cast it to - // int64 because int32 can lead to overflow when this variable is multiplied - // by large numbers (e.g., kv_block_stride). - // For blocksparse attention: skip computation on blocks that are not - // attended - if constexpr (IS_BLOCK_SPARSE) { - int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size; - if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) && - !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) { - continue; - } - } - const int64_t physical_block_number = - static_cast(block_table[block_idx]); - const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE; - const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset; - L_vec logits_vec; - from_float(logits_vec, *reinterpret_cast(logits + token_idx - - start_token_idx)); - - const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride + - kv_head_idx * kv_head_stride; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE) { - const int offset = row_idx * BLOCK_SIZE + physical_block_offset; - V_vec v_vec; - - if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) { - v_vec = *reinterpret_cast(v_ptr + offset); - } else { - V_quant_vec v_quant_vec = - *reinterpret_cast(v_ptr + offset); - // Vector conversion from V_quant_vec to V_vec. - v_vec = fp8::scaled_convert(v_quant_vec, - *v_scale); - } - if (block_idx == num_seq_blocks - 1) { - // NOTE(woosuk): When v_vec contains the tokens that are out of the - // context, we should explicitly zero out the values since they may - // contain NaNs. See - // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472 - scalar_t* v_vec_ptr = reinterpret_cast(&v_vec); -#pragma unroll - for (int j = 0; j < V_VEC_SIZE; j++) { - v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value; - } - } - accs[i] += dot(logits_vec, v_vec); - } - } - } - - // Perform reduction within each warp. -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - float acc = accs[i]; -#pragma unroll - for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) { - acc += VLLM_SHFL_XOR_SYNC(acc, mask); - } - accs[i] = acc; - } - - // NOTE(woosuk): A barrier is required because the shared memory space for - // logits is reused for the output. - __syncthreads(); - - // Perform reduction across warps. - float* out_smem = reinterpret_cast(shared_mem); -#pragma unroll - for (int i = NUM_WARPS; i > 1; i /= 2) { - int mid = i / 2; - // Upper warps write to shared memory. - if (warp_idx >= mid && warp_idx < i) { - float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE]; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) { - dst[row_idx] = accs[i]; - } - } - } - __syncthreads(); - - // Lower warps update the output. - if (warp_idx < mid) { - const float* src = &out_smem[warp_idx * HEAD_SIZE]; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) { - accs[i] += src[row_idx]; - } - } - } - __syncthreads(); - } - - // Write the final output. - if (warp_idx == 0) { - scalar_t* out_ptr = - out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE; -#pragma unroll - for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) { - const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER; - if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) { - from_float(*(out_ptr + row_idx), accs[i]); - } - } - } -} - -// Grid: (num_heads, num_seqs, 1). -template -__global__ void paged_attention_v1_kernel( - scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size] - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, // [num_heads] - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - const float* k_scale, const float* v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { - paged_attention_kernel( - /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache, - v_cache, num_kv_heads, scale, block_tables, seq_lens, - max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride, - kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks, - blocksparse_vert_stride, blocksparse_block_size, - blocksparse_head_sliding_step); -} - -// Grid: (num_heads, num_seqs, max_num_partitions). -template -__global__ void paged_attention_v2_kernel( - float* __restrict__ exp_sums, // [num_seqs, num_heads, max_num_partitions] - float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const scalar_t* __restrict__ q, // [num_seqs, num_heads, head_size] - const cache_t* __restrict__ k_cache, // [num_blocks, num_kv_heads, - // head_size/x, block_size, x] - const cache_t* __restrict__ v_cache, // [num_blocks, num_kv_heads, - // head_size, block_size] - const int num_kv_heads, // [num_heads] - const float scale, - const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_blocks_per_seq, - const float* __restrict__ alibi_slopes, // [num_heads] - const int q_stride, const int kv_block_stride, const int kv_head_stride, - const float* k_scale, const float* v_scale, const int tp_rank, - const int blocksparse_local_blocks, const int blocksparse_vert_stride, - const int blocksparse_block_size, const int blocksparse_head_sliding_step) { - paged_attention_kernel( - exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale, - block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride, - kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank, - blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size, - blocksparse_head_sliding_step); -} - -// Grid: (num_heads, num_seqs). -template -__global__ void paged_attention_v2_reduce_kernel( - scalar_t* __restrict__ out, // [num_seqs, num_heads, head_size] - const float* __restrict__ exp_sums, // [num_seqs, num_heads, - // max_num_partitions] - const float* __restrict__ max_logits, // [num_seqs, num_heads, - // max_num_partitions] - const scalar_t* __restrict__ tmp_out, // [num_seqs, num_heads, - // max_num_partitions, head_size] - const int* __restrict__ seq_lens, // [num_seqs] - const int max_num_partitions) { - const int num_heads = gridDim.x; - const int head_idx = blockIdx.x; - const int seq_idx = blockIdx.y; - const int seq_len = seq_lens[seq_idx]; - const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE); - if (num_partitions == 1) { - // No need to reduce. Only copy tmp_out to out. - scalar_t* out_ptr = - out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE; - for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) { - out_ptr[i] = tmp_out_ptr[i]; - } - // Terminate the thread block. - return; - } - - constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE; - const int warp_idx = threadIdx.x / WARP_SIZE; - const int lane = threadIdx.x % WARP_SIZE; - - // Size: 2 * num_partitions. - extern __shared__ char shared_mem[]; - // Workspace for reduction. - __shared__ float red_smem[2 * NUM_WARPS]; - - // Load max logits to shared memory. - float* shared_max_logits = reinterpret_cast(shared_mem); - const float* max_logits_ptr = max_logits + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - float max_logit = -FLT_MAX; - for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) { - const float l = max_logits_ptr[i]; - shared_max_logits[i] = l; - max_logit = fmaxf(max_logit, l); - } - __syncthreads(); - - // Get the global max logit. - // Reduce within the warp. -#pragma unroll - for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask)); - } - if (lane == 0) { - red_smem[warp_idx] = max_logit; - } - __syncthreads(); - // Reduce across warps. - max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX; -#pragma unroll - for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) { - max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask)); - } - // Broadcast the max value to all threads. - max_logit = VLLM_SHFL_SYNC(max_logit, 0); - - // Load rescaled exp sums to shared memory. - float* shared_exp_sums = - reinterpret_cast(shared_mem + sizeof(float) * num_partitions); - const float* exp_sums_ptr = exp_sums + - seq_idx * num_heads * max_num_partitions + - head_idx * max_num_partitions; - float global_exp_sum = 0.0f; - for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) { - float l = shared_max_logits[i]; - float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit); - global_exp_sum += rescaled_exp_sum; - shared_exp_sums[i] = rescaled_exp_sum; - } - __syncthreads(); - global_exp_sum = block_sum(&red_smem[NUM_WARPS], global_exp_sum); - const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f); - - // Aggregate tmp_out to out. - const scalar_t* tmp_out_ptr = - tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE + - head_idx * max_num_partitions * HEAD_SIZE; - scalar_t* out_ptr = - out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; -#pragma unroll - for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) { - float acc = 0.0f; - for (int j = 0; j < num_partitions; ++j) { - acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] * - inv_global_exp_sum; - } - from_float(out_ptr[i], acc); - } -} - -} // namespace vllm - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/attention/paged_attention_v1.cu b/csrc/libtorch_stable/attention/paged_attention_v1.cu deleted file mode 100644 index 8fa41791593..00000000000 --- a/csrc/libtorch_stable/attention/paged_attention_v1.cu +++ /dev/null @@ -1,190 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "../torch_utils.h" -#include "attention_kernels.cuh" -#include "../../cuda_compat.h" - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE) \ - VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize( \ - ((void*)vllm::paged_attention_v1_kernel), \ - shared_mem_size); \ - vllm::paged_attention_v1_kernel \ - <<>>( \ - out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \ - scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq, \ - alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride, \ - k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks, \ - blocksparse_vert_stride, blocksparse_block_size, \ - blocksparse_head_sliding_step); - -// TODO(woosuk): Tune NUM_THREADS. -template -void paged_attention_v1_launcher( - torch::stable::Tensor& out, torch::stable::Tensor& query, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - int num_kv_heads, float scale, torch::stable::Tensor& block_tables, - torch::stable::Tensor& seq_lens, int max_seq_len, - const std::optional& alibi_slopes, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, - const int tp_rank, const int blocksparse_local_blocks, - const int blocksparse_vert_stride, const int blocksparse_block_size, - const int blocksparse_head_sliding_step) { - int num_seqs = query.size(0); - int num_heads = query.size(1); - int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: alibi_slopes is optional. - const float* alibi_slopes_ptr = - alibi_slopes - ? reinterpret_cast(alibi_slopes.value().data_ptr()) - : nullptr; - - T* out_ptr = reinterpret_cast(out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.mutable_data_ptr(); - int* seq_lens_ptr = seq_lens.mutable_data_ptr(); - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - - const int NUM_WARPS = NUM_THREADS / WARP_SIZE; - int padded_max_seq_len = - DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE; - int logits_size = padded_max_seq_len * sizeof(float); - int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float); - // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len - // Keep that in sync with the logic here! - int shared_mem_size = std::max(logits_size, outputs_size); - - dim3 grid(num_heads, num_seqs, 1); - dim3 block(NUM_THREADS); - const torch::stable::accelerator::DeviceGuard device_guard( - query.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - switch (head_size) { - // NOTE(woosuk): To reduce the compilation time, we only compile for the - // head sizes that we use in the model. However, we can easily extend this - // to support any head size which is a multiple of 16. - case 32: - LAUNCH_PAGED_ATTENTION_V1(32); - break; - case 64: - LAUNCH_PAGED_ATTENTION_V1(64); - break; - case 80: - LAUNCH_PAGED_ATTENTION_V1(80); - break; - case 96: - LAUNCH_PAGED_ATTENTION_V1(96); - break; - case 112: - LAUNCH_PAGED_ATTENTION_V1(112); - break; - case 120: - LAUNCH_PAGED_ATTENTION_V1(120); - break; - case 128: - LAUNCH_PAGED_ATTENTION_V1(128); - break; - case 192: - LAUNCH_PAGED_ATTENTION_V1(192); - break; - case 256: - LAUNCH_PAGED_ATTENTION_V1(256); - break; - default: - STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); - break; - } -} - -#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \ - paged_attention_v1_launcher( \ - out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \ - seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank, \ - blocksparse_local_blocks, blocksparse_vert_stride, \ - blocksparse_block_size, blocksparse_head_sliding_step); - -#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \ - if (is_block_sparse) { \ - CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \ - } else { \ - CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \ - } - -// NOTE(woosuk): To reduce the compilation time, we omitted block sizes -// 1, 2, 4, 64, 128, 256. -#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ - } - -void paged_attention_v1( - torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] - int64_t num_kv_heads, // [num_heads] - double scale, - torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::stable::Tensor& seq_lens, // [num_seqs] - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step) { - const bool is_block_sparse = (blocksparse_vert_stride > 1); - - DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, - CALL_V1_LAUNCHER_BLOCK_SIZE) -} - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/attention/paged_attention_v2.cu b/csrc/libtorch_stable/attention/paged_attention_v2.cu deleted file mode 100644 index 4e8e56ae05c..00000000000 --- a/csrc/libtorch_stable/attention/paged_attention_v2.cu +++ /dev/null @@ -1,202 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp - * Copyright (c) 2023, The vLLM team. - * Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "../torch_utils.h" -#include "attention_kernels.cuh" -#include "../../cuda_compat.h" - -#define MAX(a, b) ((a) > (b) ? (a) : (b)) -#define MIN(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b)) - -#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE) \ - vllm::paged_attention_v2_kernel \ - <<>>( \ - exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \ - value_cache_ptr, num_kv_heads, scale, block_tables_ptr, \ - seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride, \ - kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank, \ - blocksparse_local_blocks, blocksparse_vert_stride, \ - blocksparse_block_size, blocksparse_head_sliding_step); \ - vllm::paged_attention_v2_reduce_kernel \ - <<>>( \ - out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr, \ - max_num_partitions); - -template -void paged_attention_v2_launcher( - torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, - torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, - torch::stable::Tensor& query, torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, int num_kv_heads, float scale, - torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, - int max_seq_len, const std::optional& alibi_slopes, - torch::stable::Tensor& k_scale, torch::stable::Tensor& v_scale, - const int tp_rank, const int blocksparse_local_blocks, - const int blocksparse_vert_stride, const int blocksparse_block_size, - const int blocksparse_head_sliding_step) { - int num_seqs = query.size(0); - int num_heads = query.size(1); - int head_size = query.size(2); - int max_num_blocks_per_seq = block_tables.size(1); - int q_stride = query.stride(0); - int kv_block_stride = key_cache.stride(0); - int kv_head_stride = key_cache.stride(1); - - // NOTE: alibi_slopes is optional. - const float* alibi_slopes_ptr = - alibi_slopes - ? reinterpret_cast(alibi_slopes.value().data_ptr()) - : nullptr; - - T* out_ptr = reinterpret_cast(out.data_ptr()); - float* exp_sums_ptr = reinterpret_cast(exp_sums.data_ptr()); - float* max_logits_ptr = reinterpret_cast(max_logits.data_ptr()); - T* tmp_out_ptr = reinterpret_cast(tmp_out.data_ptr()); - T* query_ptr = reinterpret_cast(query.data_ptr()); - CACHE_T* key_cache_ptr = reinterpret_cast(key_cache.data_ptr()); - CACHE_T* value_cache_ptr = reinterpret_cast(value_cache.data_ptr()); - int* block_tables_ptr = block_tables.mutable_data_ptr(); - int* seq_lens_ptr = seq_lens.mutable_data_ptr(); - const float* k_scale_ptr = reinterpret_cast(k_scale.data_ptr()); - const float* v_scale_ptr = reinterpret_cast(v_scale.data_ptr()); - - const int NUM_WARPS = NUM_THREADS / WARP_SIZE; - int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE); - int logits_size = PARTITION_SIZE * sizeof(float); - int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float); - - // For paged attention v2 kernel. - dim3 grid(num_heads, num_seqs, max_num_partitions); - int shared_mem_size = std::max(logits_size, outputs_size); - // For paged attention v2 reduce kernel. - dim3 reduce_grid(num_heads, num_seqs); - int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float); - - dim3 block(NUM_THREADS); - const torch::stable::accelerator::DeviceGuard device_guard( - query.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - switch (head_size) { - // NOTE(woosuk): To reduce the compilation time, we only compile for the - // head sizes that we use in the model. However, we can easily extend this - // to support any head size which is a multiple of 16. - case 32: - LAUNCH_PAGED_ATTENTION_V2(32); - break; - case 64: - LAUNCH_PAGED_ATTENTION_V2(64); - break; - case 80: - LAUNCH_PAGED_ATTENTION_V2(80); - break; - case 96: - LAUNCH_PAGED_ATTENTION_V2(96); - break; - case 112: - LAUNCH_PAGED_ATTENTION_V2(112); - break; - case 120: - LAUNCH_PAGED_ATTENTION_V2(120); - break; - case 128: - LAUNCH_PAGED_ATTENTION_V2(128); - break; - case 192: - LAUNCH_PAGED_ATTENTION_V2(192); - break; - case 256: - LAUNCH_PAGED_ATTENTION_V2(256); - break; - default: - STD_TORCH_CHECK(false, "Unsupported head size: ", head_size); - break; - } -} - -#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE) \ - paged_attention_v2_launcher( \ - out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache, \ - num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \ - k_scale, v_scale, tp_rank, blocksparse_local_blocks, \ - blocksparse_vert_stride, blocksparse_block_size, \ - blocksparse_head_sliding_step); - -#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \ - if (is_block_sparse) { \ - CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true); \ - } else { \ - CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false); \ - } - -// NOTE(woosuk): To reduce the compilation time, we omitted block sizes -// 1, 2, 4, 64, 128, 256. -#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE) \ - switch (block_size) { \ - case 8: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE); \ - break; \ - case 16: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE); \ - break; \ - case 32: \ - CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE); \ - break; \ - default: \ - STD_TORCH_CHECK(false, "Unsupported block size: ", block_size); \ - break; \ - } - -void paged_attention_v2( - torch::stable::Tensor& out, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - exp_sums, // [num_seqs, num_heads, max_num_partitions] - torch::stable::Tensor& - max_logits, // [num_seqs, num_heads, max_num_partitions] - torch::stable::Tensor& - tmp_out, // [num_seqs, num_heads, max_num_partitions, head_size] - torch::stable::Tensor& query, // [num_seqs, num_heads, head_size] - torch::stable::Tensor& - key_cache, // [num_blocks, num_heads, head_size/x, block_size, x] - torch::stable::Tensor& - value_cache, // [num_blocks, num_heads, head_size, block_size] - int64_t num_kv_heads, // [num_heads] - double scale, - torch::stable::Tensor& block_tables, // [num_seqs, max_num_blocks_per_seq] - torch::stable::Tensor& seq_lens, // [num_seqs] - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step) { - const bool is_block_sparse = (blocksparse_vert_stride > 1); - DISPATCH_BY_KV_CACHE_DTYPE(query.scalar_type(), kv_cache_dtype, - CALL_V2_LAUNCHER_BLOCK_SIZE) -} - -#undef MAX -#undef MIN -#undef DIVIDE_ROUND_UP diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 6438daefcab..0daa024a767 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -452,32 +452,6 @@ torch::stable::Tensor gptq_gemm(torch::stable::Tensor a, void gptq_shuffle(torch::stable::Tensor q_weight, torch::stable::Tensor q_perm, int64_t bit); -void paged_attention_v1( - torch::stable::Tensor& out, torch::stable::Tensor& query, - torch::stable::Tensor& key_cache, torch::stable::Tensor& value_cache, - int64_t num_kv_heads, double scale, torch::stable::Tensor& block_tables, - torch::stable::Tensor& seq_lens, int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - -void paged_attention_v2( - torch::stable::Tensor& out, torch::stable::Tensor& exp_sums, - torch::stable::Tensor& max_logits, torch::stable::Tensor& tmp_out, - torch::stable::Tensor& query, torch::stable::Tensor& key_cache, - torch::stable::Tensor& value_cache, int64_t num_kv_heads, double scale, - torch::stable::Tensor& block_tables, torch::stable::Tensor& seq_lens, - int64_t block_size, int64_t max_seq_len, - const std::optional& alibi_slopes, - const std::string& kv_cache_dtype, torch::stable::Tensor& k_scale, - torch::stable::Tensor& v_scale, const int64_t tp_rank, - const int64_t blocksparse_local_blocks, - const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size, - const int64_t blocksparse_head_sliding_step); - // Cache ops (shared CUDA/ROCm) void swap_blocks(torch::stable::Tensor& src, torch::stable::Tensor& dst, int64_t block_size_in_bytes, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index a760df260ac..66a93a5ba13 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -598,33 +598,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor? initial_state_idx," "Tensor? cu_chunk_seqlen," "Tensor? last_chunk_indices) -> ()"); - - // Attention ops - // Compute the attention between an input query and the cached - // keys/values using PagedAttention. - ops.def( - "paged_attention_v1(" - " Tensor! out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); - - // PagedAttention V2. - ops.def( - "paged_attention_v2(" - " Tensor! out, Tensor! exp_sums, Tensor! max_logits," - " Tensor! tmp_out, Tensor query, Tensor key_cache," - " Tensor value_cache, int num_kv_heads, float scale," - " Tensor block_tables, Tensor seq_lens, int block_size," - " int max_seq_len, Tensor? alibi_slopes," - " str kv_cache_dtype, Tensor k_scale, Tensor v_scale," - " int tp_rank, int blocksparse_local_blocks," - " int blocksparse_vert_stride, int blocksparse_block_size," - " int blocksparse_head_sliding_step) -> ()"); } STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { @@ -747,9 +720,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { // Mamba kernels ops.impl("selective_scan_fwd", TORCH_BOX(&selective_scan_fwd)); - - ops.impl("paged_attention_v1", TORCH_BOX(&paged_attention_v1)); - ops.impl("paged_attention_v2", TORCH_BOX(&paged_attention_v2)); } STABLE_TORCH_LIBRARY_IMPL(_C, CPU, ops) { diff --git a/tests/kernels/attention/test_attention.py b/tests/kernels/attention/test_attention.py index 0bc2461463c..662e14a9bb8 100644 --- a/tests/kernels/attention/test_attention.py +++ b/tests/kernels/attention/test_attention.py @@ -21,16 +21,14 @@ MAX_SEQ_LEN = get_max_shared_memory_bytes() // FLOAT32_BYTES - 512 # There may not be enough gpu memory due to large NUM_BLOCKS. # Reduce NUM_BLOCKS when it happens. NUM_BLOCKS = 4321 # Arbitrary values for testing -PARTITION_SIZE = 512 PARTITION_SIZE_ROCM = 256 DTYPES = [torch.bfloat16] NUM_GEN_SEQS = [7] # Arbitrary values for testing NUM_PREFILL_SEQS = [3] # Arbitrary values for testing NUM_HEADS = [(32, 8), (40, 40), (64, 8)] # Arbitrary values for testing -# This should be sync with get_supported_head_sizes() in -# vllm.v1.attention.ops.paged_attn.PagedAttention -HEAD_SIZES = [32, 64, 80, 128, 256] +# Head sizes supported by the ROCm paged attention kernel. +HEAD_SIZES = [64, 128] BLOCK_SIZES = [16, 32] USE_ALIBI = [False, True] @@ -111,8 +109,8 @@ def ref_single_query_cached_kv_attention( output[i].copy_(out, non_blocking=True) -@pytest.mark.parametrize( - "version", ["v1", "v2"] if not current_platform.is_rocm() else ["v1", "v2", "rocm"] +@pytest.mark.skipif( + not current_platform.is_rocm(), reason="ROCm-only paged attention kernel" ) @pytest.mark.parametrize("num_seqs", NUM_GEN_SEQS) @pytest.mark.parametrize("num_heads", NUM_HEADS) @@ -125,7 +123,6 @@ def ref_single_query_cached_kv_attention( @pytest.mark.parametrize("device", CUDA_DEVICES) def test_paged_attention( kv_cache_factory, - version: str, num_seqs: int, num_heads: tuple[int, int], head_size: int, @@ -136,22 +133,11 @@ def test_paged_attention( seed: int, device: str, ) -> None: - if (kv_cache_dtype == "fp8" and head_size % 16) or ( - version == "rocm" and head_size not in (64, 128) + if current_platform.is_navi() and ( + kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi ): pytest.skip() - if ( - version == "rocm" - and current_platform.is_navi() - and ( - kv_cache_dtype == "fp8" or head_size != 128 or block_size != 16 or use_alibi - ) - ): - pytest.skip() - - global PARTITION_SIZE - set_random_seed(seed) torch.set_default_device(device) scale = float(1.0 / (head_size**0.5)) @@ -200,9 +186,46 @@ def test_paged_attention( # Call the paged attention kernel. output = torch.empty_like(query) - if version == "v1": - ops.paged_attention_v1( + num_partitions = (max_seq_len + PARTITION_SIZE_ROCM - 1) // PARTITION_SIZE_ROCM + assert PARTITION_SIZE_ROCM % block_size == 0 + num_seqs, num_heads, head_size = output.shape + tmp_output = torch.empty( + size=(num_seqs, num_heads, num_partitions, head_size), + dtype=output.dtype, + ) + exp_sums = torch.empty( + size=(num_seqs, num_heads, num_partitions), + dtype=torch.float32, + ) + max_logits = torch.empty_like(exp_sums) + ops.paged_attention_rocm( + output, + exp_sums, + max_logits, + tmp_output, + query, + key_cache, + value_cache, + num_kv_heads, + scale, + block_tables, + seq_lens, + None, + block_size, + max_seq_len, + alibi_slopes, + kv_cache_dtype, + k_scale, + v_scale, + ) + + opcheck( + torch.ops._rocm_C.paged_attention, + ( output, + exp_sums, + max_logits, + tmp_output, query, key_cache, value_cache, @@ -210,157 +233,18 @@ def test_paged_attention( scale, block_tables, seq_lens, + None, block_size, max_seq_len, alibi_slopes, kv_cache_dtype, k_scale, v_scale, - ) - - opcheck( - torch.ops._C.paged_attention_v1, - ( - output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - 0, - 0, - 0, - 64, - 0, - ), - cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]), - ) - - elif version in ("v2", "rocm"): - if current_platform.is_rocm() and version == "rocm": - PARTITION_SIZE = PARTITION_SIZE_ROCM - - num_partitions = (max_seq_len + PARTITION_SIZE - 1) // PARTITION_SIZE - assert PARTITION_SIZE % block_size == 0 - num_seqs, num_heads, head_size = output.shape - tmp_output = torch.empty( - size=(num_seqs, num_heads, num_partitions, head_size), - dtype=output.dtype, - ) - exp_sums = torch.empty( - size=(num_seqs, num_heads, num_partitions), - dtype=torch.float32, - ) - max_logits = torch.empty_like(exp_sums) - if version == "v2": - ops.paged_attention_v2( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - - opcheck( - torch.ops._C.paged_attention_v2, - ( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - 0, - 0, - 0, - 64, - 0, - ), - cond=(head_size == HEAD_SIZES[0] and block_size == BLOCK_SIZES[0]), - ) - - else: - ops.paged_attention_rocm( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - None, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - ) - - opcheck( - torch.ops._rocm_C.paged_attention, - ( - output, - exp_sums, - max_logits, - tmp_output, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - None, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - None, - "f16", - ), - cond=(head_size == 64 and block_size == BLOCK_SIZES[0]), - ) - - else: - raise AssertionError(f"Unknown version: {version}") + None, + "f16", + ), + cond=(head_size == 64 and block_size == BLOCK_SIZES[0]), + ) # Run the reference implementation. if kv_cache_dtype == "fp8": diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index a609035560e..2a32a38a695 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -117,100 +117,6 @@ if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"): # page attention ops -def paged_attention_v1( - out: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - num_kv_heads: int, - scale: float, - block_tables: torch.Tensor, - seq_lens: torch.Tensor, - block_size: int, - max_seq_len: int, - alibi_slopes: torch.Tensor | None, - kv_cache_dtype: str, - k_scale: torch.Tensor, - v_scale: torch.Tensor, - tp_rank: int = 0, - blocksparse_local_blocks: int = 0, - blocksparse_vert_stride: int = 0, - blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - torch.ops._C.paged_attention_v1( - out, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - tp_rank, - blocksparse_local_blocks, - blocksparse_vert_stride, - blocksparse_block_size, - blocksparse_head_sliding_step, - ) - - -def paged_attention_v2( - out: torch.Tensor, - exp_sum: torch.Tensor, - max_logits: torch.Tensor, - tmp_out: torch.Tensor, - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - num_kv_heads: int, - scale: float, - block_tables: torch.Tensor, - seq_lens: torch.Tensor, - block_size: int, - max_seq_len: int, - alibi_slopes: torch.Tensor | None, - kv_cache_dtype: str, - k_scale: torch.Tensor, - v_scale: torch.Tensor, - tp_rank: int = 0, - blocksparse_local_blocks: int = 0, - blocksparse_vert_stride: int = 0, - blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - torch.ops._C.paged_attention_v2( - out, - exp_sum, - max_logits, - tmp_out, - query, - key_cache, - value_cache, - num_kv_heads, - scale, - block_tables, - seq_lens, - block_size, - max_seq_len, - alibi_slopes, - kv_cache_dtype, - k_scale, - v_scale, - tp_rank, - blocksparse_local_blocks, - blocksparse_vert_stride, - blocksparse_block_size, - blocksparse_head_sliding_step, - ) - - def paged_attention_rocm( out: torch.Tensor, exp_sum: torch.Tensor,