| 1 | // N-gram embedding index kernel for LongCat-Flash (n-gram embedding variant). |
| 2 | // |
| 3 | // Adapted from SGLang: |
| 4 | // https://github.com/sgl-project/sglang/blob/main/python/sglang/jit_kernel/csrc/ngram_embedding.cuh |
| 5 | // |
| 6 | // For each position, computes the hashed n-gram embedding ids that index the |
| 7 | // concatenated embedder table. Integer tensors are int32 except ``row_indices`` |
| 8 | // (int64); the token table is ``[max_running_reqs, max_context_len]`` int32, |
| 9 | // where a negative entry marks an ignored token (e.g. an EOS boundary). |
| 10 | |
| 11 | #include "torch_utils.h" |
| 12 | |
| 13 | #include "ops.h" |
| 14 | |
| 15 | #include <cstdint> |
| 16 | |
| 17 | namespace vllm::ngram_embedding { |
| 18 | |
| 19 | constexpr int kBlockThreads = 256; |
| 20 | |
| 21 | __global__ void ComputeNGramIdsKernel( |
| 22 | int batch_size, int ne_n, int ne_k, |
| 23 | int* ne_weights, // [ne_n-1, ne_k, ne_n] |
| 24 | int* ne_mods, // [ne_n-1, ne_k] |
| 25 | int* exclusive_ne_embedder_size_sums, // [(ne_n-1)*ne_k + 1] |
| 26 | int* exclusive_req_len_sums, // [batch_size + 1] |
| 27 | int* ne_token_table, // [max_running_reqs, max_context_len] |
| 28 | int max_context_len, |
| 29 | const int64_t* __restrict__ row_indices, // [batch_size] |
| 30 | int* column_starts, // [batch_size] |
| 31 | int* n_gram_ids // [token_num, (ne_n-1)*ne_k] |
| 32 | ) { |
| 33 | const int req_id = blockIdx.x % batch_size; |
| 34 | const int config_id = (blockIdx.x - req_id) / batch_size; |
| 35 | // n and k are offset from their physical meaning: n = real_n - 2, k = real_k |
| 36 | // - 1 (they index into ne_weights / ne_mods). |
| 37 | const int k = config_id % ne_k; |
| 38 | const int n = (config_id - config_id % ne_k) / ne_k; |
| 39 | const int ne_weight_base_idx = n * ne_k * ne_n + k * ne_n; |
| 40 | const int ne_mod = ne_mods[n * ne_k + k]; |
| 41 | for (int i = exclusive_req_len_sums[req_id] + threadIdx.x; |
| 42 | i < exclusive_req_len_sums[req_id + 1]; i += blockDim.x) { |
| 43 | uint64_t n_gram_id = 0; |
| 44 | const int64_t current_token_offset = i - exclusive_req_len_sums[req_id]; |
| 45 | const int64_t req_token_table_index = |
| 46 | row_indices[req_id] * static_cast<int64_t>(max_context_len); |
| 47 | const int64_t current_token_table_index = |
| 48 | req_token_table_index + column_starts[req_id] + current_token_offset; |
| 49 | for (int j = 0; j < n + 2; j++) { |
| 50 | if (current_token_table_index - j < req_token_table_index) { |
| 51 | break; // outside this request's range |
| 52 | } |
| 53 | if (ne_token_table[current_token_table_index - j] < 0) { |
| 54 | break; // ignored token |
| 55 | } |
| 56 | const uint64_t term = |
| 57 | (uint64_t)ne_token_table[current_token_table_index - j] * |
| 58 | (uint64_t)ne_weights[ne_weight_base_idx + j]; |
| 59 | n_gram_id += term % ne_mod; |
| 60 | } |
| 61 | n_gram_id %= ne_mod; |
| 62 | n_gram_id += exclusive_ne_embedder_size_sums[n * ne_k + k]; |
| 63 | n_gram_ids[i * (ne_n - 1) * ne_k + n * ne_k + k] = (int)(n_gram_id); |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | } // namespace vllm::ngram_embedding |
| 68 | |
| 69 | void ngram_compute_n_gram_ids( |
| 70 | int64_t ne_n, int64_t ne_k, torch::stable::Tensor& ne_weights, |
| 71 | torch::stable::Tensor& ne_mods, |
| 72 | torch::stable::Tensor& exclusive_ne_embedder_size_sums, |
| 73 | torch::stable::Tensor& exclusive_req_len_sums, |
| 74 | torch::stable::Tensor& ne_token_table, torch::stable::Tensor& row_indices, |
| 75 | torch::stable::Tensor& column_starts, torch::stable::Tensor& n_gram_ids) { |
| 76 | const int batch_size = static_cast<int>(exclusive_req_len_sums.size(0) - 1); |
| 77 | const int max_context_len = static_cast<int>(ne_token_table.size(1)); |
| 78 | const int num_configs = (static_cast<int>(ne_n) - 1) * static_cast<int>(ne_k); |
| 79 | const int grid_size = num_configs * batch_size; |
| 80 | if (grid_size <= 0) return; |
| 81 | |
| 82 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 83 | ne_weights.get_device_index()); |
| 84 | const cudaStream_t stream = get_current_cuda_stream(); |
| 85 | vllm::ngram_embedding::ComputeNGramIdsKernel<<< |
| 86 | grid_size, vllm::ngram_embedding::kBlockThreads, 0, stream>>>( |
| 87 | batch_size, static_cast<int>(ne_n), static_cast<int>(ne_k), |
| 88 | ne_weights.mutable_data_ptr<int32_t>(), |
| 89 | ne_mods.mutable_data_ptr<int32_t>(), |
| 90 | exclusive_ne_embedder_size_sums.mutable_data_ptr<int32_t>(), |
| 91 | exclusive_req_len_sums.mutable_data_ptr<int32_t>(), |
| 92 | ne_token_table.mutable_data_ptr<int32_t>(), max_context_len, |
| 93 | row_indices.const_data_ptr<int64_t>(), |
| 94 | column_starts.mutable_data_ptr<int32_t>(), |
| 95 | n_gram_ids.mutable_data_ptr<int32_t>()); |
| 96 | } |