| 1 | #ifndef CONCAT_MLA_Q_CUH_ |
| 2 | #define CONCAT_MLA_Q_CUH_ |
| 3 | |
| 4 | #include "cuda_vec_utils.cuh" |
| 5 | |
| 6 | namespace vllm { |
| 7 | |
| 8 | // Concatenates ql_nope [num_tokens, num_heads, NOPE_DIM] and |
| 9 | // q_pe [num_tokens, num_heads, 64] |
| 10 | // into q_out [num_tokens, num_heads, NOPE_DIM+64]. |
| 11 | // Currently instantiated only for NOPE_DIM=512. |
| 12 | // Rope dim is hardcoded to 64 (DeepSeek V3.2 MLA) |
| 13 | template <typename DType, int NOPE_DIM> |
| 14 | __global__ void ConcatMLAQKernel( |
| 15 | DType* __restrict__ q_out, const DType* __restrict__ ql_nope, |
| 16 | const DType* __restrict__ q_pe, const int num_tokens, const int num_heads, |
| 17 | const int64_t out_stride_0, const int64_t out_stride_1, |
| 18 | const int64_t nope_stride_0, const int64_t nope_stride_1, |
| 19 | const int64_t pe_stride_0, const int64_t pe_stride_1) { |
| 20 | const int flat_warp_id = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; |
| 21 | if (flat_warp_id >= num_tokens * num_heads) return; |
| 22 | |
| 23 | const int token_id = flat_warp_id / num_heads; |
| 24 | const int head_id = flat_warp_id % num_heads; |
| 25 | const int lane_id = threadIdx.x & 31; |
| 26 | |
| 27 | constexpr bool use_256b = VLLM_256B_PTX_ENABLED; |
| 28 | constexpr int nope_vec_loads = |
| 29 | NOPE_DIM * sizeof(DType) / (VecTraits<use_256b>::ARCH_MAX_VEC_SIZE * 32); |
| 30 | |
| 31 | const DType* nope_src = |
| 32 | ql_nope + token_id * nope_stride_0 + head_id * nope_stride_1; |
| 33 | DType* nope_dst = q_out + token_id * out_stride_0 + head_id * out_stride_1; |
| 34 | |
| 35 | #pragma unroll |
| 36 | for (int i = 0; i < nope_vec_loads; i++) { |
| 37 | const int offset = i * 32 + lane_id; |
| 38 | if constexpr (use_256b) { |
| 39 | st256_cs(reinterpret_cast<u32x8_t*>(nope_dst) + offset, |
| 40 | ld256_cs(reinterpret_cast<const u32x8_t*>(nope_src) + offset)); |
| 41 | } else { |
| 42 | st128_cs(reinterpret_cast<int4*>(nope_dst) + offset, |
| 43 | ld128_cs(reinterpret_cast<const int4*>(nope_src) + offset)); |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | const int* rope_src = reinterpret_cast<const int*>( |
| 48 | q_pe + token_id * pe_stride_0 + head_id * pe_stride_1); |
| 49 | int* rope_dst = reinterpret_cast<int*>(q_out + token_id * out_stride_0 + |
| 50 | head_id * out_stride_1 + NOPE_DIM); |
| 51 | |
| 52 | st32_cs(rope_dst + lane_id, ld32_cs(rope_src + lane_id)); |
| 53 | } |
| 54 | |
| 55 | } // namespace vllm |
| 56 | |
| 57 | #endif // CONCAT_MLA_Q_CUH_ |