| 1 | #include "cpu_types.hpp" |
| 2 | |
| 3 | #include <algorithm> |
| 4 | |
| 5 | namespace cpu_utils { |
| 6 | |
| 7 | void eagle_prepare_inputs_padded_kernel_impl( |
| 8 | const torch::Tensor& cu_num_draft_tokens, |
| 9 | const torch::Tensor& valid_sampled_tokens_count, |
| 10 | const torch::Tensor& query_start_loc_gpu, |
| 11 | torch::Tensor& token_indices_to_sample, |
| 12 | torch::Tensor& num_rejected_tokens_gpu, const int64_t num_reqs) { |
| 13 | const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>(); |
| 14 | const int64_t* valid_count_ptr = |
| 15 | valid_sampled_tokens_count.data_ptr<int64_t>(); |
| 16 | const int32_t* query_loc_ptr = query_start_loc_gpu.data_ptr<int32_t>(); |
| 17 | int32_t* indices_out_ptr = token_indices_to_sample.data_ptr<int32_t>(); |
| 18 | int64_t* rejected_out_ptr = num_rejected_tokens_gpu.data_ptr<int64_t>(); |
| 19 | |
| 20 | #pragma omp parallel for |
| 21 | for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { |
| 22 | int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; |
| 23 | int64_t num_draft_tokens = cu_draft_ptr[req_idx] - start_idx; |
| 24 | int64_t num_valid_tokens = valid_count_ptr[req_idx]; |
| 25 | |
| 26 | int64_t num_rejected = 0; |
| 27 | if (num_draft_tokens > 0) { |
| 28 | num_rejected = num_draft_tokens + 1 - num_valid_tokens; |
| 29 | } |
| 30 | |
| 31 | int32_t q_last_tok_idx = query_loc_ptr[req_idx + 1] - 1; |
| 32 | int32_t index_to_sample = q_last_tok_idx - num_rejected; |
| 33 | |
| 34 | indices_out_ptr[req_idx] = index_to_sample; |
| 35 | rejected_out_ptr[req_idx] = num_rejected; |
| 36 | } |
| 37 | } |
| 38 | |
| 39 | void eagle_prepare_next_token_padded_kernel_impl( |
| 40 | const torch::Tensor& sampled_token_ids, |
| 41 | const torch::Tensor& discard_request_mask, |
| 42 | const torch::Tensor& backup_next_token_ids, torch::Tensor& next_token_ids, |
| 43 | torch::Tensor& valid_sampled_tokens_count, const int64_t vocab_size, |
| 44 | const int64_t num_sampled_tokens_per_req, const int64_t num_reqs) { |
| 45 | const int64_t* sampled_ids_ptr = sampled_token_ids.data_ptr<int64_t>(); |
| 46 | const bool* discard_mask_ptr = discard_request_mask.data_ptr<bool>(); |
| 47 | const int64_t* backup_ids_ptr = backup_next_token_ids.data_ptr<int64_t>(); |
| 48 | int64_t* next_ids_out_ptr = next_token_ids.data_ptr<int64_t>(); |
| 49 | int64_t* valid_count_out_ptr = valid_sampled_tokens_count.data_ptr<int64_t>(); |
| 50 | |
| 51 | const int64_t stride = sampled_token_ids.stride(0); |
| 52 | |
| 53 | #pragma omp parallel for |
| 54 | for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { |
| 55 | const int64_t* row_ptr = sampled_ids_ptr + req_idx * stride; |
| 56 | int64_t valid_count = 0; |
| 57 | int64_t last_valid_token = -1; |
| 58 | |
| 59 | for (int64_t pos = 0; pos < num_sampled_tokens_per_req; ++pos) { |
| 60 | int64_t token = row_ptr[pos]; |
| 61 | if (token != -1 && token < vocab_size) { |
| 62 | valid_count++; |
| 63 | last_valid_token = token; |
| 64 | } |
| 65 | } |
| 66 | |
| 67 | bool discard = discard_mask_ptr[req_idx]; |
| 68 | if (discard) { |
| 69 | next_ids_out_ptr[req_idx] = backup_ids_ptr[req_idx]; |
| 70 | valid_count_out_ptr[req_idx] = 0; |
| 71 | } else { |
| 72 | next_ids_out_ptr[req_idx] = |
| 73 | (valid_count > 0) ? last_valid_token : backup_ids_ptr[req_idx]; |
| 74 | valid_count_out_ptr[req_idx] = valid_count; |
| 75 | } |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | void eagle_step_slot_mapping_metadata_kernel_impl( |
| 80 | const torch::Tensor& positions, const torch::Tensor& block_table, |
| 81 | torch::Tensor& seq_lens, torch::Tensor& out_clamped_positions, |
| 82 | torch::Tensor& out_slot_mapping, const int64_t block_size, |
| 83 | const int64_t max_model_len, const int64_t PAD_ID) { |
| 84 | const int64_t batch_size = positions.size(0); |
| 85 | const int64_t input_batch_size = out_slot_mapping.size(0); |
| 86 | |
| 87 | const int64_t* pos_ptr = positions.data_ptr<int64_t>(); |
| 88 | const int32_t* bt_ptr = block_table.data_ptr<int32_t>(); |
| 89 | int32_t* seq_lens_ptr = seq_lens.data_ptr<int32_t>(); |
| 90 | int64_t* out_clamped_ptr = out_clamped_positions.data_ptr<int64_t>(); |
| 91 | int64_t* out_slot_ptr = out_slot_mapping.data_ptr<int64_t>(); |
| 92 | |
| 93 | const int64_t bt_stride = block_table.stride(0); |
| 94 | const int64_t n_blocks_per_req = block_table.size(1); |
| 95 | |
| 96 | #pragma omp parallel for |
| 97 | for (int64_t req_idx = 0; req_idx < input_batch_size; ++req_idx) { |
| 98 | if (req_idx >= batch_size) { |
| 99 | out_slot_ptr[req_idx] = PAD_ID; |
| 100 | continue; |
| 101 | } |
| 102 | |
| 103 | int64_t position = pos_ptr[req_idx]; |
| 104 | int64_t new_position = position + 1; |
| 105 | bool exceeds_max = new_position >= max_model_len; |
| 106 | int64_t clamped_position = exceeds_max ? 0 : new_position; |
| 107 | |
| 108 | out_clamped_ptr[req_idx] = clamped_position; |
| 109 | |
| 110 | int64_t block_number = clamped_position / block_size; |
| 111 | block_number = std::min(block_number, n_blocks_per_req - 1); |
| 112 | int32_t block_id = bt_ptr[req_idx * bt_stride + block_number]; |
| 113 | int64_t slot_id = block_id * block_size + (clamped_position % block_size); |
| 114 | out_slot_ptr[req_idx] = exceeds_max ? PAD_ID : slot_id; |
| 115 | |
| 116 | int32_t seq_len = seq_lens_ptr[req_idx]; |
| 117 | int32_t new_seq_len = exceeds_max ? 1 : (seq_len + 1); |
| 118 | new_seq_len = std::min(new_seq_len, static_cast<int32_t>(max_model_len)); |
| 119 | seq_lens_ptr[req_idx] = new_seq_len; |
| 120 | } |
| 121 | } |
| 122 | |
| 123 | void copy_and_expand_eagle_inputs_kernel_impl( |
| 124 | const torch::Tensor& target_token_ids, |
| 125 | const torch::Tensor& target_positions, const torch::Tensor& next_token_ids, |
| 126 | torch::Tensor& out_input_ids, torch::Tensor& out_positions, |
| 127 | torch::Tensor& out_is_rejected_token_mask, |
| 128 | torch::Tensor& out_is_masked_token_mask, |
| 129 | torch::Tensor& out_new_token_indices, |
| 130 | torch::Tensor& out_hidden_state_mapping, |
| 131 | const torch::Tensor& query_start_loc, const torch::Tensor& query_end_loc, |
| 132 | const int64_t padding_token_id, const int64_t parallel_drafting_token_id, |
| 133 | const int64_t total_input_tokens, |
| 134 | const int64_t num_padding_slots_per_request, const bool shift_input_ids) { |
| 135 | const int64_t num_reqs = query_end_loc.size(0); |
| 136 | |
| 137 | const int64_t* target_ids_ptr = target_token_ids.data_ptr<int64_t>(); |
| 138 | const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>(); |
| 139 | const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>(); |
| 140 | const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>(); |
| 141 | const int32_t* query_end_ptr = query_end_loc.data_ptr<int32_t>(); |
| 142 | |
| 143 | int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>(); |
| 144 | int64_t* out_pos_ptr = out_positions.data_ptr<int64_t>(); |
| 145 | bool* out_rej_mask_ptr = out_is_rejected_token_mask.data_ptr<bool>(); |
| 146 | bool* out_mask_ptr = out_is_masked_token_mask.data_ptr<bool>(); |
| 147 | int32_t* out_new_idx_ptr = out_new_token_indices.data_ptr<int32_t>(); |
| 148 | int32_t* out_hidden_map_ptr = out_hidden_state_mapping.data_ptr<int32_t>(); |
| 149 | |
| 150 | #pragma omp parallel for |
| 151 | for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { |
| 152 | int32_t q_start = query_start_ptr[req_idx]; |
| 153 | int32_t next_q_start = query_start_ptr[req_idx + 1]; |
| 154 | int32_t q_end = query_end_ptr[req_idx]; |
| 155 | |
| 156 | int64_t num_valid_tokens = |
| 157 | shift_input_ids ? (q_end - q_start) : (q_end - q_start + 1); |
| 158 | int64_t input_offset = shift_input_ids ? 1 : 0; |
| 159 | |
| 160 | int64_t out_start = q_start + req_idx * (num_padding_slots_per_request - |
| 161 | (shift_input_ids ? 1 : 0)); |
| 162 | int64_t num_rejected = next_q_start - q_end - 1; |
| 163 | int64_t total_output_tokens = |
| 164 | num_valid_tokens + num_padding_slots_per_request + num_rejected; |
| 165 | |
| 166 | int64_t start_pos = target_pos_ptr[q_start]; |
| 167 | int64_t bonus_token = next_ids_ptr[req_idx]; |
| 168 | |
| 169 | for (int64_t j = 0; j < total_output_tokens; ++j) { |
| 170 | int64_t out_idx = out_start + j; |
| 171 | bool is_valid = j < num_valid_tokens; |
| 172 | bool is_bonus = j == num_valid_tokens; |
| 173 | bool is_parallel = (j > num_valid_tokens) && |
| 174 | (j < num_valid_tokens + num_padding_slots_per_request); |
| 175 | bool is_rejected = j >= num_valid_tokens + num_padding_slots_per_request; |
| 176 | |
| 177 | int64_t in_idx = |
| 178 | std::min(static_cast<int64_t>(q_start + input_offset + j), |
| 179 | total_input_tokens - 1); |
| 180 | |
| 181 | int64_t token_id = padding_token_id; |
| 182 | if (is_valid) |
| 183 | token_id = target_ids_ptr[in_idx]; |
| 184 | else if (is_bonus) |
| 185 | token_id = bonus_token; |
| 186 | else if (is_parallel) |
| 187 | token_id = parallel_drafting_token_id; |
| 188 | |
| 189 | out_ids_ptr[out_idx] = token_id; |
| 190 | out_pos_ptr[out_idx] = is_rejected ? 0 : (start_pos + j); |
| 191 | out_rej_mask_ptr[out_idx] = is_rejected; |
| 192 | out_mask_ptr[out_idx] = is_parallel; |
| 193 | |
| 194 | if (is_bonus || is_parallel) { |
| 195 | int64_t new_token_local_idx = j - num_valid_tokens; |
| 196 | int64_t new_token_out_idx = |
| 197 | req_idx * num_padding_slots_per_request + new_token_local_idx; |
| 198 | out_new_idx_ptr[new_token_out_idx] = out_idx; |
| 199 | } |
| 200 | } |
| 201 | |
| 202 | if (shift_input_ids) { |
| 203 | int64_t n_input = next_q_start - q_start; |
| 204 | for (int64_t j = 0; j < n_input; ++j) { |
| 205 | out_hidden_map_ptr[q_start + j] = out_start + j; |
| 206 | } |
| 207 | } |
| 208 | } |
| 209 | } |
| 210 | |
| 211 | void copy_and_expand_dflash_inputs_kernel_impl( |
| 212 | const torch::Tensor& next_token_ids, const torch::Tensor& target_positions, |
| 213 | torch::Tensor& out_input_ids, torch::Tensor& out_context_positions, |
| 214 | torch::Tensor& out_query_positions, torch::Tensor& out_context_slot_mapping, |
| 215 | torch::Tensor& out_query_slot_mapping, torch::Tensor& out_token_indices, |
| 216 | const torch::Tensor& block_table, const torch::Tensor& query_start_loc, |
| 217 | const std::optional<torch::Tensor>& num_rejected_tokens, |
| 218 | const int64_t parallel_drafting_token_id, const int64_t block_size, |
| 219 | const int64_t num_query_per_req, const int64_t num_speculative_tokens, |
| 220 | const int64_t total_input_tokens, const bool has_num_rejected) { |
| 221 | const int64_t num_reqs = query_start_loc.size(0) - 1; |
| 222 | |
| 223 | const int64_t* next_ids_ptr = next_token_ids.data_ptr<int64_t>(); |
| 224 | const int64_t* target_pos_ptr = target_positions.data_ptr<int64_t>(); |
| 225 | const int32_t* block_table_ptr = block_table.data_ptr<int32_t>(); |
| 226 | const int32_t* query_start_ptr = query_start_loc.data_ptr<int32_t>(); |
| 227 | const int64_t* rejected_ptr = |
| 228 | has_num_rejected && num_rejected_tokens.has_value() |
| 229 | ? num_rejected_tokens.value().data_ptr<int64_t>() |
| 230 | : nullptr; |
| 231 | |
| 232 | int64_t* out_ids_ptr = out_input_ids.data_ptr<int64_t>(); |
| 233 | int64_t* out_ctx_pos_ptr = out_context_positions.data_ptr<int64_t>(); |
| 234 | int64_t* out_query_pos_ptr = out_query_positions.data_ptr<int64_t>(); |
| 235 | int64_t* out_ctx_slot_ptr = out_context_slot_mapping.data_ptr<int64_t>(); |
| 236 | int64_t* out_query_slot_ptr = out_query_slot_mapping.data_ptr<int64_t>(); |
| 237 | int32_t* out_token_idx_ptr = out_token_indices.data_ptr<int32_t>(); |
| 238 | |
| 239 | const int64_t block_table_stride = block_table.stride(0); |
| 240 | |
| 241 | #pragma omp parallel for |
| 242 | for (int64_t req_idx = 0; req_idx < num_reqs; ++req_idx) { |
| 243 | int32_t ctx_start = query_start_ptr[req_idx]; |
| 244 | int32_t ctx_end = query_start_ptr[req_idx + 1]; |
| 245 | int64_t num_ctx = ctx_end - ctx_start; |
| 246 | int64_t valid_ctx_end = ctx_end; |
| 247 | if (rejected_ptr != nullptr) { |
| 248 | valid_ctx_end -= rejected_ptr[req_idx]; |
| 249 | } |
| 250 | // Guard against out-of-bounds: ensure valid_ctx_end > ctx_start so that |
| 251 | // valid_ctx_end - 1 never reads before the request's context range. |
| 252 | valid_ctx_end = |
| 253 | std::max(valid_ctx_end, static_cast<int64_t>(ctx_start + 1)); |
| 254 | |
| 255 | int64_t last_pos = target_pos_ptr[valid_ctx_end - 1]; |
| 256 | |
| 257 | for (int64_t j = 0; j < num_ctx; ++j) { |
| 258 | int64_t ctx_idx = ctx_start + j; |
| 259 | int64_t ctx_pos_idx = std::min(ctx_idx, total_input_tokens - 1); |
| 260 | int64_t position = target_pos_ptr[ctx_pos_idx]; |
| 261 | int64_t block_num = position / block_size; |
| 262 | block_num = std::min(block_num, block_table_stride - 1); |
| 263 | int32_t block_id = |
| 264 | block_table_ptr[req_idx * block_table_stride + block_num]; |
| 265 | int64_t slot = block_id * block_size + (position % block_size); |
| 266 | |
| 267 | out_ctx_pos_ptr[ctx_idx] = position; |
| 268 | out_ctx_slot_ptr[ctx_idx] = slot; |
| 269 | } |
| 270 | |
| 271 | for (int64_t query_off = 0; query_off < num_query_per_req; ++query_off) { |
| 272 | int64_t query_out = req_idx * num_query_per_req + query_off; |
| 273 | int64_t position = last_pos + 1 + query_off; |
| 274 | int64_t block_num = position / block_size; |
| 275 | block_num = std::min(block_num, block_table_stride - 1); |
| 276 | int32_t block_id = |
| 277 | block_table_ptr[req_idx * block_table_stride + block_num]; |
| 278 | int64_t slot = block_id * block_size + (position % block_size); |
| 279 | |
| 280 | out_query_pos_ptr[query_out] = position; |
| 281 | out_query_slot_ptr[query_out] = slot; |
| 282 | out_ids_ptr[query_out] = |
| 283 | query_off == 0 ? next_ids_ptr[req_idx] : parallel_drafting_token_id; |
| 284 | |
| 285 | if (query_off > 0) { |
| 286 | int64_t sample_out_idx = |
| 287 | req_idx * num_speculative_tokens + (query_off - 1); |
| 288 | out_token_idx_ptr[sample_out_idx] = query_out; |
| 289 | } |
| 290 | } |
| 291 | } |
| 292 | } |
| 293 | |
| 294 | void rejection_greedy_sample_kernel_impl( |
| 295 | torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, |
| 296 | const torch::Tensor& draft_token_ids, const torch::Tensor& target_argmax, |
| 297 | const torch::Tensor& bonus_token_ids, |
| 298 | const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len) { |
| 299 | const int64_t batch_size = cu_num_draft_tokens.size(0); |
| 300 | |
| 301 | int64_t* out_ptr = output_token_ids.data_ptr<int64_t>(); |
| 302 | const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>(); |
| 303 | const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>(); |
| 304 | const int64_t* target_argmax_ptr = target_argmax.data_ptr<int64_t>(); |
| 305 | const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>(); |
| 306 | const bool* greedy_ptr = |
| 307 | is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr; |
| 308 | |
| 309 | const int64_t out_stride = output_token_ids.stride(0); |
| 310 | const int64_t bonus_stride = bonus_token_ids.stride(0); |
| 311 | |
| 312 | #pragma omp parallel for |
| 313 | for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { |
| 314 | if (greedy_ptr && !greedy_ptr[req_idx]) continue; |
| 315 | |
| 316 | int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; |
| 317 | int64_t end_idx = cu_draft_ptr[req_idx]; |
| 318 | int64_t num_draft_tokens = end_idx - start_idx; |
| 319 | |
| 320 | bool rejected = false; |
| 321 | for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { |
| 322 | int64_t target_id = target_argmax_ptr[start_idx + pos]; |
| 323 | out_ptr[req_idx * out_stride + pos] = target_id; |
| 324 | |
| 325 | if (draft_ids_ptr[start_idx + pos] != target_id) { |
| 326 | rejected = true; |
| 327 | break; |
| 328 | } |
| 329 | } |
| 330 | |
| 331 | if (!rejected) { |
| 332 | out_ptr[req_idx * out_stride + num_draft_tokens] = |
| 333 | bonus_ids_ptr[req_idx * bonus_stride]; |
| 334 | } |
| 335 | } |
| 336 | } |
| 337 | |
| 338 | void rejection_random_sample_kernel_impl( |
| 339 | torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, |
| 340 | const torch::Tensor& draft_token_ids, |
| 341 | const std::optional<torch::Tensor>& draft_probs, |
| 342 | const torch::Tensor& target_probs, const torch::Tensor& bonus_token_ids, |
| 343 | const torch::Tensor& recovered_token_ids, |
| 344 | const torch::Tensor& uniform_probs, |
| 345 | const std::optional<torch::Tensor>& is_greedy, const int64_t max_spec_len, |
| 346 | const int64_t vocab_size, const bool no_draft_probs) { |
| 347 | const int64_t batch_size = cu_num_draft_tokens.size(0); |
| 348 | |
| 349 | int64_t* out_ptr = output_token_ids.data_ptr<int64_t>(); |
| 350 | const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>(); |
| 351 | const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>(); |
| 352 | const float* draft_probs_ptr = |
| 353 | no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>(); |
| 354 | const float* target_probs_ptr = target_probs.data_ptr<float>(); |
| 355 | const int64_t* bonus_ids_ptr = bonus_token_ids.data_ptr<int64_t>(); |
| 356 | const int64_t* recovered_ids_ptr = recovered_token_ids.data_ptr<int64_t>(); |
| 357 | const float* uniform_probs_ptr = uniform_probs.data_ptr<float>(); |
| 358 | const bool* greedy_ptr = |
| 359 | is_greedy.has_value() ? is_greedy.value().data_ptr<bool>() : nullptr; |
| 360 | |
| 361 | const int64_t out_stride = output_token_ids.stride(0); |
| 362 | const int64_t bonus_stride = bonus_token_ids.stride(0); |
| 363 | const int64_t target_stride = target_probs.stride(0); |
| 364 | const int64_t draft_probs_stride = |
| 365 | no_draft_probs ? 0 : draft_probs.value().stride(0); |
| 366 | |
| 367 | #pragma omp parallel for |
| 368 | for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { |
| 369 | if (greedy_ptr && greedy_ptr[req_idx]) continue; |
| 370 | |
| 371 | int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; |
| 372 | int64_t end_idx = cu_draft_ptr[req_idx]; |
| 373 | int64_t num_draft_tokens = end_idx - start_idx; |
| 374 | |
| 375 | bool rejected = false; |
| 376 | for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { |
| 377 | int64_t token_idx = start_idx + pos; |
| 378 | int64_t draft_id = draft_ids_ptr[token_idx]; |
| 379 | |
| 380 | float p = target_probs_ptr[token_idx * target_stride + draft_id]; |
| 381 | float q = |
| 382 | no_draft_probs |
| 383 | ? 1.0f |
| 384 | : draft_probs_ptr[token_idx * draft_probs_stride + draft_id]; |
| 385 | float uniform_p = uniform_probs_ptr[token_idx]; |
| 386 | |
| 387 | float ratio = (q > 0.0f) ? (p / q) : 0.0f; |
| 388 | |
| 389 | if (ratio >= uniform_p) { |
| 390 | out_ptr[req_idx * out_stride + pos] = draft_id; |
| 391 | } else { |
| 392 | out_ptr[req_idx * out_stride + pos] = recovered_ids_ptr[token_idx]; |
| 393 | rejected = true; |
| 394 | break; |
| 395 | } |
| 396 | } |
| 397 | |
| 398 | if (!rejected) { |
| 399 | out_ptr[req_idx * out_stride + num_draft_tokens] = |
| 400 | bonus_ids_ptr[req_idx * bonus_stride]; |
| 401 | } |
| 402 | } |
| 403 | } |
| 404 | |
| 405 | void expand_kernel_impl(torch::Tensor& output, const torch::Tensor& input, |
| 406 | const torch::Tensor& cu_num_tokens, |
| 407 | const int64_t replace_from, const int64_t replace_to) { |
| 408 | const int64_t batch_size = cu_num_tokens.size(0); |
| 409 | const int64_t* cu_tokens_ptr = cu_num_tokens.data_ptr<int64_t>(); |
| 410 | |
| 411 | int64_t* out_ptr = output.data_ptr<int64_t>(); |
| 412 | const int64_t* in_ptr = input.data_ptr<int64_t>(); |
| 413 | |
| 414 | #pragma omp parallel for |
| 415 | for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { |
| 416 | int64_t start_idx = req_idx == 0 ? 0 : cu_tokens_ptr[req_idx - 1]; |
| 417 | int64_t end_idx = cu_tokens_ptr[req_idx]; |
| 418 | int64_t val = in_ptr[req_idx]; |
| 419 | |
| 420 | if (val == replace_from) { |
| 421 | val = replace_to; |
| 422 | } |
| 423 | |
| 424 | for (int64_t i = start_idx; i < end_idx; ++i) { |
| 425 | out_ptr[i] = val; |
| 426 | } |
| 427 | } |
| 428 | } |
| 429 | |
| 430 | void sample_recovered_tokens_kernel_impl( |
| 431 | torch::Tensor& output_token_ids, const torch::Tensor& cu_num_draft_tokens, |
| 432 | const torch::Tensor& draft_token_ids, |
| 433 | const std::optional<torch::Tensor>& draft_probs, |
| 434 | const torch::Tensor& target_probs, const torch::Tensor& inv_q, |
| 435 | const int64_t vocab_size, const bool no_draft_probs) { |
| 436 | const int64_t batch_size = cu_num_draft_tokens.size(0); |
| 437 | |
| 438 | int64_t* out_ptr = output_token_ids.data_ptr<int64_t>(); |
| 439 | const int64_t* cu_draft_ptr = cu_num_draft_tokens.data_ptr<int64_t>(); |
| 440 | const int64_t* draft_ids_ptr = draft_token_ids.data_ptr<int64_t>(); |
| 441 | const float* draft_probs_ptr = |
| 442 | no_draft_probs ? nullptr : draft_probs.value().data_ptr<float>(); |
| 443 | const float* target_probs_ptr = target_probs.data_ptr<float>(); |
| 444 | const float* inv_q_ptr = inv_q.data_ptr<float>(); |
| 445 | |
| 446 | const int64_t target_stride = target_probs.stride(0); |
| 447 | const int64_t draft_probs_stride = |
| 448 | no_draft_probs ? 0 : draft_probs.value().stride(0); |
| 449 | const int64_t inv_q_stride = inv_q.stride(0); |
| 450 | |
| 451 | #pragma omp parallel for |
| 452 | for (int64_t req_idx = 0; req_idx < batch_size; ++req_idx) { |
| 453 | int64_t start_idx = req_idx == 0 ? 0 : cu_draft_ptr[req_idx - 1]; |
| 454 | int64_t end_idx = cu_draft_ptr[req_idx]; |
| 455 | int64_t num_draft_tokens = end_idx - start_idx; |
| 456 | |
| 457 | const float* req_inv_q = inv_q_ptr + req_idx * inv_q_stride; |
| 458 | |
| 459 | for (int64_t pos = 0; pos < num_draft_tokens; ++pos) { |
| 460 | int64_t token_idx = start_idx + pos; |
| 461 | int64_t draft_id = draft_ids_ptr[token_idx]; |
| 462 | |
| 463 | const float* token_target_probs = |
| 464 | target_probs_ptr + token_idx * target_stride; |
| 465 | const float* token_draft_probs = |
| 466 | no_draft_probs ? nullptr |
| 467 | : (draft_probs_ptr + token_idx * draft_probs_stride); |
| 468 | |
| 469 | int64_t best_id = 0; |
| 470 | float best_val = -1.0f; |
| 471 | |
| 472 | for (int64_t v = 0; v < vocab_size; ++v) { |
| 473 | float prob = token_target_probs[v]; |
| 474 | if (no_draft_probs) { |
| 475 | if (v == draft_id) prob = 0.0f; |
| 476 | } else { |
| 477 | float diff = prob - token_draft_probs[v]; |
| 478 | prob = diff > 0.0f ? diff : 0.0f; |
| 479 | } |
| 480 | |
| 481 | float val = prob * req_inv_q[v]; |
| 482 | if (val > best_val) { |
| 483 | best_val = val; |
| 484 | best_id = v; |
| 485 | } |
| 486 | } |
| 487 | out_ptr[token_idx] = best_id; |
| 488 | } |
| 489 | } |
| 490 | } |
| 491 | |
| 492 | } // namespace cpu_utils |