| 448 | ) |
| 449 | |
| 450 | def insert_tasks_v1( |
| 451 | self, req_dicts: List[Request], num_running_requests: int, target_model_index_to_batch_id: dict = {} |
| 452 | ): |
| 453 | |
| 454 | if "caches" not in self.model_inputs: |
| 455 | self.initialize_kv_cache() |
| 456 | req_len = len(req_dicts) |
| 457 | self.model_inputs["num_running_requests"] = num_running_requests |
| 458 | self.model_inputs["running_requests_ids"] = range(num_running_requests) |
| 459 | if target_model_index_to_batch_id: |
| 460 | self.model_inputs.index_to_batch_id = dict(target_model_index_to_batch_id) |
| 461 | for i in range(req_len): |
| 462 | request = req_dicts[i] |
| 463 | logger.debug(f"{i}th request-{request.request_id}: {request}") |
| 464 | idx = self.model_inputs.get_index_by_batch_id(request.idx) |
| 465 | if request.task_type.value == RequestType.PREFILL.value: # prefill task |
| 466 | prefill_start_index = request.prefill_start_index |
| 467 | prefill_end_index = request.prefill_end_index |
| 468 | length = prefill_end_index - prefill_start_index |
| 469 | |
| 470 | input_ids = request.prompt_token_ids + request.output_token_ids |
| 471 | |
| 472 | self.model_inputs["input_ids_len"][idx] = length - 1 |
| 473 | self.model_inputs["pre_ids"][idx : idx + 1] = -1 |
| 474 | self.model_inputs["input_ids"][idx : idx + 1, : length - 1] = self.target_model_inputs["input_ids"][ |
| 475 | idx : idx + 1, 1:length |
| 476 | ] |
| 477 | self.model_inputs["input_ids_cpu"][idx : idx + 1, : length - 1] = self.target_model_inputs[ |
| 478 | "input_ids" |
| 479 | ][idx : idx + 1, 1:length].cpu() |
| 480 | encoder_block_num = len(request.block_tables) |
| 481 | self.model_inputs["encoder_block_lens"][idx : idx + 1] = encoder_block_num |
| 482 | self.model_inputs["block_tables"][idx : idx + 1, :] = -1 |
| 483 | self.model_inputs["block_tables"][idx : idx + 1, :encoder_block_num] = np.array( |
| 484 | request.block_tables, dtype="int32" |
| 485 | ) |
| 486 | self.model_inputs["stop_flags"][idx : idx + 1] = False |
| 487 | self.model_inputs["batch_drop"][idx : idx + 1] = False |
| 488 | |
| 489 | self.model_inputs["seq_lens_encoder"][idx : idx + 1] = length |
| 490 | self.model_inputs["seq_lens_decoder"][idx : idx + 1] = prefill_start_index |
| 491 | self.model_inputs["seq_lens_this_time_buffer"][idx : idx + 1] = length |
| 492 | self.model_inputs["step_idx"][idx : idx + 1] = ( |
| 493 | len(request.output_token_ids) if prefill_end_index >= len(input_ids) else 0 |
| 494 | ) |
| 495 | if self.enable_mm: |
| 496 | inputs = request.multimodal_inputs |
| 497 | self.model_inputs["attn_mask_offsets_full"][idx][0 : prefill_end_index - prefill_start_index] = ( |
| 498 | paddle.to_tensor( |
| 499 | inputs["attention_mask_offset"][prefill_start_index:prefill_end_index], dtype="int32" |
| 500 | ) |
| 501 | ) |
| 502 | self.model_inputs["attn_mask_offsets_decoder"][idx : idx + 1] = ( |
| 503 | inputs["attention_mask_offset"][prefill_end_index - 1] + 1 |
| 504 | ) |
| 505 | if ( |
| 506 | self.fd_config.scheduler_config.splitwise_role == "decode" |
| 507 | ): # In PD, we continue to decode after P generates first token |