| 1 | #include "torch_utils.h" |
| 2 | |
| 3 | #include "../cuda_compat.h" |
| 4 | #include "dispatch_utils.h" |
| 5 | |
| 6 | namespace vllm { |
| 7 | |
| 8 | template <typename scalar_t, typename cache_t, bool IS_NEOX> |
| 9 | inline __device__ void apply_token_rotary_embedding( |
| 10 | scalar_t* __restrict__ arr, const cache_t* __restrict__ cos_ptr, |
| 11 | const cache_t* __restrict__ sin_ptr, int rot_offset, int embed_dim, |
| 12 | const bool inverse) { |
| 13 | int x_index, y_index; |
| 14 | float cos_f, sin_f; |
| 15 | if (IS_NEOX) { |
| 16 | x_index = rot_offset; |
| 17 | y_index = embed_dim + rot_offset; |
| 18 | cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index)); |
| 19 | sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index)); |
| 20 | } else { |
| 21 | x_index = 2 * rot_offset; |
| 22 | y_index = 2 * rot_offset + 1; |
| 23 | cos_f = static_cast<float>(VLLM_LDG(cos_ptr + x_index / 2)); |
| 24 | sin_f = static_cast<float>(VLLM_LDG(sin_ptr + x_index / 2)); |
| 25 | } |
| 26 | if (inverse) { |
| 27 | sin_f = -sin_f; |
| 28 | } |
| 29 | const float x_f = static_cast<float>(arr[x_index]); |
| 30 | const float y_f = static_cast<float>(arr[y_index]); |
| 31 | arr[x_index] = static_cast<scalar_t>(x_f * cos_f - y_f * sin_f); |
| 32 | arr[y_index] = static_cast<scalar_t>(y_f * cos_f + x_f * sin_f); |
| 33 | } |
| 34 | |
| 35 | template <typename scalar_t, typename cache_t, bool IS_NEOX> |
| 36 | inline __device__ void apply_rotary_embedding( |
| 37 | scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads, |
| 38 | // head_size] or [num_tokens, num_heads, |
| 39 | // head_size] |
| 40 | scalar_t* __restrict__ key, // nullptr or |
| 41 | // [batch_size, seq_len, num_kv_heads, |
| 42 | // head_size] or [num_tokens, num_kv_heads, |
| 43 | // head_size] |
| 44 | const cache_t* cache_ptr, const int head_size, const int num_heads, |
| 45 | const int num_kv_heads, const int rot_dim, const int token_idx, |
| 46 | const int64_t query_stride, const int64_t key_stride, |
| 47 | const int64_t head_stride, const int64_t rope_dim_offset, |
| 48 | const bool inverse) { |
| 49 | const int embed_dim = rot_dim / 2; |
| 50 | const cache_t* cos_ptr = cache_ptr; |
| 51 | const cache_t* sin_ptr = cache_ptr + embed_dim; |
| 52 | |
| 53 | const int nq = num_heads * embed_dim; |
| 54 | for (int i = threadIdx.x; i < nq; i += blockDim.x) { |
| 55 | const int head_idx = i / embed_dim; |
| 56 | const int64_t token_head = |
| 57 | token_idx * query_stride + head_idx * head_stride + rope_dim_offset; |
| 58 | const int rot_offset = i % embed_dim; |
| 59 | apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>( |
| 60 | query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse); |
| 61 | } |
| 62 | |
| 63 | if (key != nullptr) { |
| 64 | const int nk = num_kv_heads * embed_dim; |
| 65 | for (int i = threadIdx.x; i < nk; i += blockDim.x) { |
| 66 | const int head_idx = i / embed_dim; |
| 67 | const int64_t token_head = |
| 68 | token_idx * key_stride + head_idx * head_stride + rope_dim_offset; |
| 69 | const int rot_offset = i % embed_dim; |
| 70 | apply_token_rotary_embedding<scalar_t, cache_t, IS_NEOX>( |
| 71 | key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim, inverse); |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | template <typename scalar_t, typename cache_t, bool IS_NEOX> |
| 77 | __global__ void rotary_embedding_kernel( |
| 78 | const int64_t* __restrict__ positions, // [batch_size, seq_len] or |
| 79 | // [num_tokens] |
| 80 | scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads, |
| 81 | // head_size] or [num_tokens, num_heads, |
| 82 | // head_size] |
| 83 | scalar_t* __restrict__ key, // nullptr or |
| 84 | // [batch_size, seq_len, num_kv_heads, |
| 85 | // head_size] or [num_tokens, num_kv_heads, |
| 86 | // head_size] |
| 87 | const cache_t* __restrict__ cos_sin_cache, // [max_position, rot_dim] |
| 88 | const int rot_dim, const int64_t query_stride, const int64_t key_stride, |
| 89 | const int64_t head_stride, const int num_heads, const int num_kv_heads, |
| 90 | const int head_size, const int64_t rope_dim_offset, const bool inverse) { |
| 91 | const int token_idx = blockIdx.x; |
| 92 | int64_t pos = positions[token_idx]; |
| 93 | const cache_t* cache_ptr = cos_sin_cache + pos * rot_dim; |
| 94 | |
| 95 | apply_rotary_embedding<scalar_t, cache_t, IS_NEOX>( |
| 96 | query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim, |
| 97 | token_idx, query_stride, key_stride, head_stride, rope_dim_offset, |
| 98 | inverse); |
| 99 | } |
| 100 | |
| 101 | } // namespace vllm |
| 102 | |
| 103 | void rotary_embedding( |
| 104 | torch::stable::Tensor& positions, // [batch_size, seq_len] or [num_tokens] |
| 105 | torch::stable::Tensor& |
| 106 | query, // [batch_size, seq_len, num_heads * head_size] or |
| 107 | // [num_tokens, num_heads * head_size] or |
| 108 | // [batch_size, seq_len, num_heads, head_size] or |
| 109 | // [num_tokens, num_heads, head_size] |
| 110 | std::optional<torch::stable::Tensor> key, |
| 111 | // null or |
| 112 | // [batch_size, seq_len, num_kv_heads * head_size] or |
| 113 | // [num_tokens, num_kv_heads * head_size] or |
| 114 | // [batch_size, seq_len, num_heads, head_size] or |
| 115 | // [num_tokens, num_heads, head_size] |
| 116 | int64_t head_size, |
| 117 | torch::stable::Tensor& cos_sin_cache, // [max_position, rot_dim] |
| 118 | bool is_neox, int64_t rope_dim_offset, bool inverse) { |
| 119 | // num_tokens = batch_size * seq_len |
| 120 | int64_t num_tokens = positions.numel(); |
| 121 | int positions_ndim = positions.dim(); |
| 122 | |
| 123 | // Make sure num_tokens dim is consistent across positions, query, and key |
| 124 | STD_TORCH_CHECK( |
| 125 | positions_ndim == 1 || positions_ndim == 2, |
| 126 | "positions must have shape [num_tokens] or [batch_size, seq_len]"); |
| 127 | if (positions_ndim == 1) { |
| 128 | STD_TORCH_CHECK( |
| 129 | query.size(0) == positions.size(0) && |
| 130 | (!key.has_value() || key->size(0) == positions.size(0)), |
| 131 | "query, key and positions must have the same number of tokens"); |
| 132 | } |
| 133 | if (positions_ndim == 2) { |
| 134 | STD_TORCH_CHECK( |
| 135 | query.size(0) == positions.size(0) && |
| 136 | (!key.has_value() || key->size(0) == positions.size(0)) && |
| 137 | query.size(1) == positions.size(1) && |
| 138 | (!key.has_value() || key->size(1) == positions.size(1)), |
| 139 | "query, key and positions must have the same batch_size and seq_len"); |
| 140 | } |
| 141 | |
| 142 | // Make sure head_size is valid for query and key |
| 143 | // hidden_size = num_heads * head_size |
| 144 | int query_hidden_size = query.numel() / num_tokens; |
| 145 | int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0; |
| 146 | STD_TORCH_CHECK(query_hidden_size % head_size == 0); |
| 147 | STD_TORCH_CHECK(key_hidden_size % head_size == 0); |
| 148 | |
| 149 | // Make sure query and key have consistent number of heads |
| 150 | int num_heads = query_hidden_size / head_size; |
| 151 | int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads; |
| 152 | STD_TORCH_CHECK(num_heads % num_kv_heads == 0); |
| 153 | |
| 154 | int rot_dim = cos_sin_cache.size(1); |
| 155 | int seq_dim_idx = positions_ndim - 1; |
| 156 | int64_t query_stride = query.stride(seq_dim_idx); |
| 157 | int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0; |
| 158 | |
| 159 | STD_TORCH_CHECK((rot_dim + rope_dim_offset) <= head_size); |
| 160 | // Determine head stride: for [*, heads, head_size] use stride of last dim; |
| 161 | // for flat [*, heads*head_size], heads blocks are contiguous of size |
| 162 | // head_size |
| 163 | int query_ndim = query.dim(); |
| 164 | int64_t head_stride = |
| 165 | (query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size; |
| 166 | |
| 167 | dim3 grid(num_tokens); |
| 168 | dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512)); |
| 169 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 170 | query.get_device_index()); |
| 171 | const cudaStream_t stream = get_current_cuda_stream(); |
| 172 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 173 | query.scalar_type(), "rotary_embedding", [&] { |
| 174 | using query_t = scalar_t; |
| 175 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 176 | cos_sin_cache.scalar_type(), "rotary_embedding_cache", [&] { |
| 177 | using cache_t = scalar_t; |
| 178 | if (is_neox) { |
| 179 | vllm::rotary_embedding_kernel<query_t, cache_t, true> |
| 180 | <<<grid, block, 0, stream>>>( |
| 181 | positions.const_data_ptr<int64_t>(), |
| 182 | query.mutable_data_ptr<query_t>(), |
| 183 | key.has_value() ? key->mutable_data_ptr<query_t>() |
| 184 | : nullptr, |
| 185 | cos_sin_cache.const_data_ptr<cache_t>(), rot_dim, |
| 186 | query_stride, key_stride, head_stride, num_heads, |
| 187 | num_kv_heads, head_size, rope_dim_offset, inverse); |
| 188 | } else { |
| 189 | vllm::rotary_embedding_kernel<query_t, cache_t, false> |
| 190 | <<<grid, block, 0, stream>>>( |
| 191 | positions.const_data_ptr<int64_t>(), |
| 192 | query.mutable_data_ptr<query_t>(), |
| 193 | key.has_value() ? key->mutable_data_ptr<query_t>() |
| 194 | : nullptr, |
| 195 | cos_sin_cache.const_data_ptr<cache_t>(), rot_dim, |
| 196 | query_stride, key_stride, head_stride, num_heads, |
| 197 | num_kv_heads, head_size, rope_dim_offset, inverse); |
| 198 | } |
| 199 | }); |
| 200 | }); |
| 201 | } |