biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Change the name to vLLM (#150)

Browse Source
This commit is contained in:
Woosuk Kwon
2023-06-17 03:07:40 -07:00
committed by GitHub
parent e5464ee484
commit 0b98ba15c7
90 changed files with 342 additions and 339 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
csrc/activation_kernels.cu
View File

@@ -1,7 +1,7 @@
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
namespace cacheflow {
namespace vllm {
template<typename T>
__device__ __forceinline__ T silu(const T& x) {
@@ -22,7 +22,7 @@ __global__ void silu_and_mul_kernel(
}
}
} // namespace cacheflow
} // namespace vllm
void silu_and_mul(
torch::Tensor& out, // [num_tokens, d]
@@ -40,7 +40,7 @@ void silu_and_mul(
input.scalar_type(),
"silu_and_mul_kernel",
[&] {
cacheflow::silu_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
vllm::silu_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
d);

6
csrc/attention/attention_generic.cuh
View File

@@ -1,6 +1,6 @@
/*
* Adapted from https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
* Copyright (c) 2023, The CacheFlow team.
* 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");
@@ -19,7 +19,7 @@
#include <stdint.h>
namespace cacheflow {
namespace vllm {
// A vector type to store Q, K, V elements.
template<typename T, int VEC_SIZE>
@@ -61,4 +61,4 @@ inline __device__ void zero(T& dst) {
dst = tmp.raw;
}
} // namespace cacheflow
} // namespace vllm

8
csrc/attention/attention_kernels.cu
View File

@@ -1,6 +1,6 @@
/*
* 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 CacheFlow team.
* 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");
@@ -27,7 +27,7 @@
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
namespace cacheflow {
namespace vllm {
// Utility function for attention softmax.
template<int NUM_WARPS>
@@ -315,10 +315,10 @@ __global__ void single_query_cached_kv_attention_kernel(
}
}
} // namespace cacheflow
} // namespace vllm
#define LAUNCH_ATTENTION_KERNEL(T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS) \
cacheflow::single_query_cached_kv_attention_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS> \
vllm::single_query_cached_kv_attention_kernel<T, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS> \
<<<grid, block, shared_mem_size, stream>>>( \
out_ptr, \
query_ptr, \

6
csrc/attention/attention_utils.cuh
View File

@@ -1,6 +1,6 @@
/*
* 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 CacheFlow team.
* 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");
@@ -22,7 +22,7 @@
#include <float.h>
#include <type_traits>
namespace cacheflow {
namespace vllm {
// Q*K^T operation.
template<int THREAD_GROUP_SIZE, typename Vec, int N>
@@ -52,4 +52,4 @@ struct Qk_dot {
}
};
} // namespace cacheflow
} // namespace vllm

6
csrc/attention/dtype_bfloat16.cuh
View File

@@ -1,7 +1,7 @@
/*
* 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
* and https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
* Copyright (c) 2023, The CacheFlow team.
* 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");
@@ -25,7 +25,7 @@
#include <cuda_fp16.h>
#include <stdint.h>
namespace cacheflow {
namespace vllm {
// Define custom BF16 vector data types.
struct bf16_4_t {
@@ -420,4 +420,4 @@ inline __device__ void from_float(bf16_8_t& dst, Float8_ src) {
#endif
}
} // namespace cacheflow
} // namespace vllm

6
csrc/attention/dtype_float16.cuh
View File

@@ -1,7 +1,7 @@
/*
* 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
* and https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
* Copyright (c) 2023, The CacheFlow team.
* 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");
@@ -23,7 +23,7 @@
#include <stdint.h>
namespace cacheflow {
namespace vllm {
// FP16 vector types for Q, K, V.
template<>
@@ -441,4 +441,4 @@ inline __device__ Float8_ to_float(uint4 u) {
return tmp;
}
} // namespace cacheflow
} // namespace vllm

6
csrc/attention/dtype_float32.cuh
View File

@@ -1,7 +1,7 @@
/*
* 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
* and https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
* Copyright (c) 2023, The CacheFlow team.
* 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");
@@ -22,7 +22,7 @@
#include <stdint.h>
namespace cacheflow {
namespace vllm {
// Define custom FP32 vector data types.
struct Float4_ {
@@ -265,4 +265,4 @@ inline __device__ Float8_ to_float(Float8_ u) {
return u;
}
} // namespace cacheflow
} // namespace vllm

18
csrc/cache_kernels.cu
View File

@@ -46,7 +46,7 @@ void swap_blocks(
}
}
namespace cacheflow {
namespace vllm {
// Grid: (num_layers, num_pairs)
template<typename scalar_t>
@@ -77,7 +77,7 @@ __global__ void copy_blocks_kernel(
}
}
} // namespace cacheflow
} // namespace vllm
void copy_blocks(
std::vector<torch::Tensor>& key_caches,
@@ -129,7 +129,7 @@ void copy_blocks(
at::ScalarType::Half,
at::ScalarType::BFloat16,
key_caches[0].scalar_type(), "copy_blocks_kernel", ([&] {
cacheflow::copy_blocks_kernel<scalar_t><<<grid, block, 0, stream>>>(
vllm::copy_blocks_kernel<scalar_t><<<grid, block, 0, stream>>>(
key_cache_ptrs_tensor.data_ptr<int64_t>(),
value_cache_ptrs_tensor.data_ptr<int64_t>(),
block_mapping_tensor.data_ptr<int>(),
@@ -137,7 +137,7 @@ void copy_blocks(
}));
}
namespace cacheflow {
namespace vllm {
template<typename scalar_t>
__global__ void reshape_and_cache_kernel(
@@ -181,7 +181,7 @@ __global__ void reshape_and_cache_kernel(
}
}
} // namespace cacheflow
} // namespace vllm
void reshape_and_cache(
torch::Tensor& key, // [num_tokens, num_heads, head_size]
@@ -208,7 +208,7 @@ void reshape_and_cache(
key.scalar_type(),
"reshape_and_cache_kernel",
[&] {
cacheflow::reshape_and_cache_kernel<scalar_t><<<grid, block, 0, stream>>>(
vllm::reshape_and_cache_kernel<scalar_t><<<grid, block, 0, stream>>>(
key.data_ptr<scalar_t>(),
value.data_ptr<scalar_t>(),
key_cache.data_ptr<scalar_t>(),
@@ -223,7 +223,7 @@ void reshape_and_cache(
});
}
namespace cacheflow {
namespace vllm {
// Grid: (num_blocks, block_size).
template<typename scalar_t>
@@ -343,7 +343,7 @@ __global__ void gather_cached_kv_kernel_optimized(
}
}
} // namespace cacheflow
} // namespace vllm
void gather_cached_kv(
torch::Tensor& key, // [out] [num_tokens, num_heads, head_size]
@@ -370,7 +370,7 @@ void gather_cached_kv(
key.scalar_type(),
"gather_cached_kv_kernel_optimized",
[&] {
cacheflow::gather_cached_kv_kernel_optimized<scalar_t><<<grid, block, 0, stream>>>(
vllm::gather_cached_kv_kernel_optimized<scalar_t><<<grid, block, 0, stream>>>(
key.data_ptr<scalar_t>(),
value.data_ptr<scalar_t>(),
key_cache.data_ptr<scalar_t>(),

6
csrc/layernorm_kernels.cu
View File

@@ -3,7 +3,7 @@
#include "reduction_utils.cuh"
namespace cacheflow {
namespace vllm {
// TODO(woosuk): Further optimize this kernel.
template<typename scalar_t>
@@ -33,7 +33,7 @@ __global__ void rms_norm_kernel(
}
}
} // namespace cacheflow
} // namespace vllm
void rms_norm(
torch::Tensor& out, // [num_tokens, hidden_size]
@@ -52,7 +52,7 @@ void rms_norm(
input.scalar_type(),
"rms_norm_kernel",
[&] {
cacheflow::rms_norm_kernel<scalar_t><<<grid, block, 0, stream>>>(
vllm::rms_norm_kernel<scalar_t><<<grid, block, 0, stream>>>(
out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),

6
csrc/pos_encoding_kernels.cu
View File

@@ -1,7 +1,7 @@
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
namespace cacheflow {
namespace vllm {
template<typename scalar_t>
__global__ void rotary_embedding_neox_kernel(
@@ -46,7 +46,7 @@ __global__ void rotary_embedding_neox_kernel(
}
}
} // namespace cacheflow
} // namespace vllm
void rotary_embedding_neox(
torch::Tensor& positions, // [num_tokens]
@@ -70,7 +70,7 @@ void rotary_embedding_neox(
query.scalar_type(),
"rotary_embedding_neox",
[&] {
cacheflow::rotary_embedding_neox_kernel<scalar_t><<<grid, block, 0, stream>>>(
vllm::rotary_embedding_neox_kernel<scalar_t><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.data_ptr<scalar_t>(),

6
csrc/reduction_utils.cuh
View File

@@ -1,6 +1,6 @@
/*
* Adapted from https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/reduce_kernel_utils.cuh
* Copyright (c) 2023, The CacheFlow team.
* 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");
@@ -17,7 +17,7 @@
*/
#pragma once
namespace cacheflow {
namespace vllm {
template<typename T>
__inline__ __device__ T warpReduceSum(T val) {
@@ -48,4 +48,4 @@ __inline__ __device__ T blockReduceSum(T val) {
return val;
}
} // namespace cacheflow
} // namespace vllm