| 1 | #include "torch_utils.h" |
| 2 | #include "dispatch_utils.h" |
| 3 | |
| 4 | #include "../cuda_compat.h" |
| 5 | |
| 6 | #include "../quantization/w8a8/fp8/common.cuh" |
| 7 | #ifdef USE_ROCM |
| 8 | #include "../quantization/w8a8/fp8/amd/quant_utils.cuh" |
| 9 | #else |
| 10 | #include "../quantization/w8a8/fp8/nvidia/quant_utils.cuh" |
| 11 | #endif |
| 12 | |
| 13 | #ifdef USE_ROCM |
| 14 | #include <hip/hip_bf16.h> |
| 15 | typedef __hip_bfloat16 __nv_bfloat16; |
| 16 | #endif |
| 17 | |
| 18 | namespace vllm { |
| 19 | |
| 20 | // NOTE Be EXTRA careful with raw_kv_scalar_t, for __half and __nv_bfloat16 it's |
| 21 | // using u16 as the backing type. |
| 22 | template <typename qk_t, typename cos_sin_t, bool IS_NEOX, |
| 23 | typename raw_kv_scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt> |
| 24 | __global__ void concat_and_cache_mla_rope_fused_kernel( |
| 25 | const int64_t* __restrict__ positions, // [num_tokens] |
| 26 | qk_t* __restrict__ q_pe, // [num_tokens, num_q_heads, rot_dim] |
| 27 | qk_t* __restrict__ k_pe, // [num_tokens, rot_dim] |
| 28 | const qk_t* __restrict__ kv_c, // [num_tokens, kv_lora_rank] |
| 29 | const cos_sin_t* __restrict__ rope_cos_sin_cache, // [max_position, 2, |
| 30 | // rot_dim // 2] |
| 31 | const int rot_dim, const int64_t q_pe_stride_token, |
| 32 | const int64_t q_pe_stride_head, const int64_t k_pe_stride, |
| 33 | const int64_t kv_c_stride, const int num_q_heads, |
| 34 | cache_t* __restrict__ kv_cache, // [num_blocks, block_size, (kv_lora_rank + |
| 35 | // rot_dim)] |
| 36 | const int64_t* __restrict__ slot_mapping, // [num_tokens] |
| 37 | const int block_stride, const int entry_stride, const int kv_lora_rank, |
| 38 | const int block_size, const float* kv_cache_quant_scale) { |
| 39 | // Each thread block is responsible for one token. |
| 40 | const int64_t token_idx = blockIdx.x; |
| 41 | const int64_t slot_idx = slot_mapping[token_idx]; |
| 42 | // NOTE: slot_idx can be -1 if the token is padded |
| 43 | if (slot_idx < 0) { |
| 44 | return; |
| 45 | } |
| 46 | const int64_t pos = positions[token_idx]; |
| 47 | |
| 48 | const cos_sin_t* cos_sin_ptr = rope_cos_sin_cache + pos * rot_dim; |
| 49 | |
| 50 | const int embed_dim = rot_dim / 2; |
| 51 | |
| 52 | // Q ROPE |
| 53 | const int nq = num_q_heads * embed_dim; |
| 54 | for (int i = threadIdx.x; i < nq; i += blockDim.x) { |
| 55 | int head_idx = i / embed_dim; |
| 56 | int pair_idx = i % embed_dim; |
| 57 | |
| 58 | // NOTE: Would be nice to have interleaved sin/cos so we could just load |
| 59 | // both at the same time. |
| 60 | qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx)); |
| 61 | qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim)); |
| 62 | |
| 63 | qk_t* q_pe_head_ptr = |
| 64 | q_pe + token_idx * q_pe_stride_token + head_idx * q_pe_stride_head; |
| 65 | |
| 66 | int pair_idx_x, pair_idx_y; |
| 67 | if constexpr (IS_NEOX) { |
| 68 | // GPT-NeoX style rotary embedding. |
| 69 | pair_idx_x = pair_idx; |
| 70 | pair_idx_y = embed_dim + pair_idx; |
| 71 | } else { |
| 72 | // GPT-J style rotary embedding. |
| 73 | pair_idx_x = pair_idx * 2; |
| 74 | pair_idx_y = pair_idx * 2 + 1; |
| 75 | } |
| 76 | |
| 77 | qk_t x_src = q_pe_head_ptr[pair_idx_x]; |
| 78 | qk_t y_src = q_pe_head_ptr[pair_idx_y]; |
| 79 | |
| 80 | qk_t x_dst = x_src * cos - y_src * sin; |
| 81 | qk_t y_dst = y_src * cos + x_src * sin; |
| 82 | |
| 83 | q_pe_head_ptr[pair_idx_x] = x_dst; |
| 84 | q_pe_head_ptr[pair_idx_y] = y_dst; |
| 85 | } |
| 86 | |
| 87 | const int64_t block_idx = slot_idx / block_size; |
| 88 | const int64_t entry_idx = slot_idx % block_size; |
| 89 | |
| 90 | // K with 1 HEAD |
| 91 | for (int i = threadIdx.x; i < embed_dim; i += blockDim.x) { |
| 92 | int pair_idx = i; |
| 93 | |
| 94 | qk_t cos = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx)); |
| 95 | qk_t sin = static_cast<qk_t>(VLLM_LDG(cos_sin_ptr + pair_idx + embed_dim)); |
| 96 | |
| 97 | qk_t* k_pe_head_ptr = k_pe + token_idx * k_pe_stride; |
| 98 | |
| 99 | int pair_idx_x, pair_idx_y; |
| 100 | if constexpr (IS_NEOX) { |
| 101 | // GPT-NeoX style rotary embedding. |
| 102 | pair_idx_x = pair_idx; |
| 103 | pair_idx_y = embed_dim + pair_idx; |
| 104 | } else { |
| 105 | // GPT-J style rotary embedding. |
| 106 | pair_idx_x = pair_idx * 2; |
| 107 | pair_idx_y = pair_idx * 2 + 1; |
| 108 | } |
| 109 | |
| 110 | qk_t x_src = k_pe_head_ptr[pair_idx_x]; |
| 111 | qk_t y_src = k_pe_head_ptr[pair_idx_y]; |
| 112 | |
| 113 | qk_t x_dst = x_src * cos - y_src * sin; |
| 114 | qk_t y_dst = y_src * cos + x_src * sin; |
| 115 | |
| 116 | k_pe_head_ptr[pair_idx_x] = x_dst; |
| 117 | k_pe_head_ptr[pair_idx_y] = y_dst; |
| 118 | |
| 119 | // NOTE Why is this monster necessary? |
| 120 | // When K is of type float16, the actual template replacement for |
| 121 | // raw_kv_scalar_t with be u16. That's why it's used at the last moment |
| 122 | // otherwise CUDA ALU would break. |
| 123 | const raw_kv_scalar_t raw_x_value = |
| 124 | *reinterpret_cast<const raw_kv_scalar_t*>(&x_dst); |
| 125 | const raw_kv_scalar_t raw_y_value = |
| 126 | *reinterpret_cast<const raw_kv_scalar_t*>(&y_dst); |
| 127 | |
| 128 | cache_t* kv_cache_ptr = kv_cache + block_idx * block_stride + |
| 129 | entry_idx * entry_stride + kv_lora_rank; |
| 130 | |
| 131 | // MLA Cache Store |
| 132 | if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { |
| 133 | kv_cache_ptr[pair_idx_x] = raw_x_value; |
| 134 | kv_cache_ptr[pair_idx_y] = raw_y_value; |
| 135 | } else { |
| 136 | kv_cache_ptr[pair_idx_x] = |
| 137 | fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>( |
| 138 | raw_x_value, *kv_cache_quant_scale); |
| 139 | kv_cache_ptr[pair_idx_y] = |
| 140 | fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>( |
| 141 | raw_y_value, *kv_cache_quant_scale); |
| 142 | } |
| 143 | } |
| 144 | |
| 145 | // NOPE |
| 146 | for (int i = threadIdx.x; i < kv_lora_rank; i += blockDim.x) { |
| 147 | const qk_t* src_ptr = kv_c + token_idx * kv_c_stride + i; |
| 148 | const raw_kv_scalar_t src_value = |
| 149 | *reinterpret_cast<const raw_kv_scalar_t*>(src_ptr); |
| 150 | |
| 151 | cache_t* kv_cache_ptr = |
| 152 | kv_cache + block_idx * block_stride + entry_idx * entry_stride; |
| 153 | |
| 154 | if constexpr (kv_dt == Fp8KVCacheDataType::kAuto) { |
| 155 | kv_cache_ptr[i] = src_value; |
| 156 | } else { |
| 157 | kv_cache_ptr[i] = fp8::scaled_convert<cache_t, raw_kv_scalar_t, kv_dt>( |
| 158 | src_value, *kv_cache_quant_scale); |
| 159 | } |
| 160 | } |
| 161 | } |
| 162 | |
| 163 | } // namespace vllm |
| 164 | |
| 165 | #define CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED(RAW_KV_T, CACHE_T, KV_DTYPE) \ |
| 166 | do { \ |
| 167 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ |
| 168 | q_pe.scalar_type(), "qk_scalar_type", [&] { \ |
| 169 | using qk_t = scalar_t; \ |
| 170 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( \ |
| 171 | rope_cos_sin_cache.scalar_type(), \ |
| 172 | "rope_cos_sin_cache_scalar_type", [&] { \ |
| 173 | using cos_sin_t = scalar_t; \ |
| 174 | if (rope_is_neox) { \ |
| 175 | vllm::concat_and_cache_mla_rope_fused_kernel< \ |
| 176 | qk_t, cos_sin_t, true, RAW_KV_T, CACHE_T, KV_DTYPE> \ |
| 177 | <<<grid, block, 0, stream>>>( \ |
| 178 | positions.const_data_ptr<int64_t>(), \ |
| 179 | q_pe.mutable_data_ptr<qk_t>(), \ |
| 180 | k_pe.mutable_data_ptr<qk_t>(), \ |
| 181 | kv_c.const_data_ptr<qk_t>(), \ |
| 182 | rope_cos_sin_cache.const_data_ptr<cos_sin_t>(), \ |
| 183 | rot_dim, q_pe_stride_token, q_pe_stride_head, \ |
| 184 | k_pe_stride, kv_c_stride, num_q_heads, \ |
| 185 | reinterpret_cast<CACHE_T*>( \ |
| 186 | kv_cache.mutable_data_ptr()), \ |
| 187 | slot_mapping.const_data_ptr<int64_t>(), \ |
| 188 | block_stride, entry_stride, kv_lora_rank, \ |
| 189 | block_size, \ |
| 190 | kv_cache_quant_scale.const_data_ptr<float>()); \ |
| 191 | } else { \ |
| 192 | vllm::concat_and_cache_mla_rope_fused_kernel< \ |
| 193 | qk_t, cos_sin_t, false, RAW_KV_T, CACHE_T, KV_DTYPE> \ |
| 194 | <<<grid, block, 0, stream>>>( \ |
| 195 | positions.const_data_ptr<int64_t>(), \ |
| 196 | q_pe.mutable_data_ptr<qk_t>(), \ |
| 197 | k_pe.mutable_data_ptr<qk_t>(), \ |
| 198 | kv_c.const_data_ptr<qk_t>(), \ |
| 199 | rope_cos_sin_cache.const_data_ptr<cos_sin_t>(), \ |
| 200 | rot_dim, q_pe_stride_token, q_pe_stride_head, \ |
| 201 | k_pe_stride, kv_c_stride, num_q_heads, \ |
| 202 | reinterpret_cast<CACHE_T*>( \ |
| 203 | kv_cache.mutable_data_ptr()), \ |
| 204 | slot_mapping.const_data_ptr<int64_t>(), \ |
| 205 | block_stride, entry_stride, kv_lora_rank, \ |
| 206 | block_size, \ |
| 207 | kv_cache_quant_scale.const_data_ptr<float>()); \ |
| 208 | } \ |
| 209 | }); \ |
| 210 | }); \ |
| 211 | } while (false) |
| 212 | |
| 213 | // Executes RoPE on q_pe and k_pe, then writes k_pe and kv_c in the kv cache. |
| 214 | // q_pe and k_pe are modified in place. |
| 215 | // Replaces DeepseekScalingRotaryEmbedding.self.rotary_emb and |
| 216 | // concat_and_cache_mla. |
| 217 | void concat_and_cache_mla_rope_fused( |
| 218 | torch::stable::Tensor& positions, // [num_tokens] |
| 219 | torch::stable::Tensor& q_pe, // [num_tokens, num_q_heads, rot_dim] |
| 220 | torch::stable::Tensor& k_pe, // [num_tokens, rot_dim] |
| 221 | torch::stable::Tensor& kv_c, // [num_tokens, kv_lora_rank] |
| 222 | torch::stable::Tensor& rope_cos_sin_cache, // [max_position, rot_dim] |
| 223 | bool rope_is_neox, |
| 224 | torch::stable::Tensor& slot_mapping, // [num_tokens] or [num_actual_tokens] |
| 225 | torch::stable::Tensor& |
| 226 | kv_cache, // [num_blocks, block_size, (kv_lora_rank + rot_dim)] |
| 227 | const std::string& kv_cache_dtype, |
| 228 | torch::stable::Tensor& kv_cache_quant_scale) { |
| 229 | // NOTE(woosuk): In vLLM V1, query/key/position.size(0) can be different from |
| 230 | // slot_mapping.size(0) because of padding for CUDA graphs. |
| 231 | // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) |
| 232 | // because both include padding. |
| 233 | // In vLLM V1, however, key.size(0) can be larger than |
| 234 | // slot_mapping.size(0) since key includes padding for CUDA graphs, |
| 235 | // while slot_mapping does not. In this case, |
| 236 | // slot_mapping.size(0) represents the actual number of tokens |
| 237 | // before padding. |
| 238 | // For compatibility with both cases, we use slot_mapping.size(0) as |
| 239 | // the number of tokens. |
| 240 | const int64_t num_tokens = slot_mapping.size(0); |
| 241 | const int64_t num_padded_tokens = q_pe.size(0); |
| 242 | STD_TORCH_CHECK(num_padded_tokens >= num_tokens); |
| 243 | |
| 244 | const int num_q_heads = q_pe.size(1); |
| 245 | const int rot_dim = q_pe.size(2); |
| 246 | const int kv_lora_rank = kv_c.size(1); |
| 247 | |
| 248 | STD_TORCH_CHECK(positions.size(0) == num_padded_tokens); |
| 249 | STD_TORCH_CHECK(positions.dim() == 1); |
| 250 | STD_TORCH_CHECK(positions.scalar_type() == |
| 251 | torch::headeronly::ScalarType::Long); |
| 252 | |
| 253 | STD_TORCH_CHECK(q_pe.dim() == 3); |
| 254 | STD_TORCH_CHECK(q_pe.size(0) == num_padded_tokens); |
| 255 | STD_TORCH_CHECK(q_pe.size(1) == num_q_heads); |
| 256 | STD_TORCH_CHECK(q_pe.size(2) == rot_dim); |
| 257 | |
| 258 | STD_TORCH_CHECK(k_pe.dim() == 2); |
| 259 | STD_TORCH_CHECK(k_pe.size(0) == num_padded_tokens); |
| 260 | STD_TORCH_CHECK(k_pe.size(1) == rot_dim); |
| 261 | STD_TORCH_CHECK(k_pe.scalar_type() == q_pe.scalar_type()); |
| 262 | |
| 263 | STD_TORCH_CHECK(kv_c.dim() == 2); |
| 264 | STD_TORCH_CHECK(kv_c.size(0) == num_padded_tokens); |
| 265 | STD_TORCH_CHECK(kv_c.size(1) == kv_lora_rank); |
| 266 | STD_TORCH_CHECK(kv_c.scalar_type() == q_pe.scalar_type()); |
| 267 | |
| 268 | STD_TORCH_CHECK(rope_cos_sin_cache.size(1) == rot_dim); |
| 269 | STD_TORCH_CHECK(rope_cos_sin_cache.scalar_type() == q_pe.scalar_type()); |
| 270 | |
| 271 | STD_TORCH_CHECK(slot_mapping.size(0) == num_tokens); |
| 272 | STD_TORCH_CHECK(slot_mapping.scalar_type() == |
| 273 | torch::headeronly::ScalarType::Long); |
| 274 | |
| 275 | STD_TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + rot_dim); |
| 276 | STD_TORCH_CHECK(kv_cache.dim() == 3); |
| 277 | |
| 278 | STD_TORCH_CHECK(kv_cache_quant_scale.numel() == 1); |
| 279 | STD_TORCH_CHECK(kv_cache_quant_scale.scalar_type() == |
| 280 | torch::headeronly::ScalarType::Float); |
| 281 | |
| 282 | int64_t q_pe_stride_token = q_pe.stride(0); |
| 283 | int64_t q_pe_stride_head = q_pe.stride(1); |
| 284 | |
| 285 | int64_t k_pe_stride = k_pe.stride(0); |
| 286 | int64_t kv_c_stride = kv_c.stride(0); |
| 287 | |
| 288 | int block_size = kv_cache.size(1); |
| 289 | |
| 290 | int block_stride = kv_cache.stride(0); |
| 291 | int entry_stride = kv_cache.stride(1); |
| 292 | |
| 293 | int rope_block_size = std::min(num_q_heads * rot_dim / 2, 512); |
| 294 | int mla_block_size = kv_lora_rank; |
| 295 | int thread_block_size = |
| 296 | std::min(std::max(rope_block_size, mla_block_size), 512); |
| 297 | |
| 298 | dim3 grid(num_tokens, 1, 1); |
| 299 | dim3 block(thread_block_size, 1, 1); |
| 300 | |
| 301 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 302 | positions.get_device_index()); |
| 303 | const cudaStream_t stream = get_current_cuda_stream(); |
| 304 | |
| 305 | DISPATCH_BY_KV_CACHE_DTYPE(kv_c.scalar_type(), kv_cache_dtype, |
| 306 | CALL_CONCAT_AND_CACHE_MLA_ROPE_FUSED); |
| 307 | } |