(self, share_inputs)
| 145 | self.suffix_cache.add_active_response(req_id, token_array) |
| 146 | |
| 147 | def _run_impl(self, share_inputs): |
| 148 | |
| 149 | stop_flags_cpu = share_inputs["stop_flags"].cpu().numpy().flatten() |
| 150 | is_block_step_cpu = share_inputs["is_block_step"].cpu().numpy().flatten() |
| 151 | accept_tokens_cpu = share_inputs["accept_tokens"].cpu() |
| 152 | accept_num_cpu = share_inputs["accept_num"].cpu().numpy().flatten() |
| 153 | seq_lens_encoder = share_inputs["seq_lens_encoder"].cpu().numpy().flatten() |
| 154 | seq_lens_decoder = share_inputs["seq_lens_decoder"].cpu().numpy().flatten() |
| 155 | |
| 156 | draft_tokens_cpu = share_inputs["draft_tokens"].cpu() |
| 157 | seq_lens_this_time_cpu = share_inputs["seq_lens_this_time"].cpu() |
| 158 | |
| 159 | total_lens = seq_lens_encoder + seq_lens_decoder |
| 160 | batch_size = seq_lens_this_time_cpu.shape[0] |
| 161 | |
| 162 | for bid in range(batch_size): |
| 163 | req_id = self.idx_to_req_id.get(bid) |
| 164 | # 1. Stop condition has the highest priority |
| 165 | if stop_flags_cpu[bid]: |
| 166 | seq_lens_this_time_cpu[bid] = 0 |
| 167 | draft_tokens_cpu[bid, :] = -1 |
| 168 | if not is_block_step_cpu[bid]: |
| 169 | if req_id is not None and req_id in self.suffix_cache.active_requests: |
| 170 | self.stop_request(req_id) |
| 171 | continue |
| 172 | else: |
| 173 | seq_lens_this_time_cpu[bid] = 1 |
| 174 | draft_tokens_cpu[bid, 1:] = -1 |
| 175 | # 2. Skip some cases |
| 176 | num_tokens = total_lens[bid] |
| 177 | max_spec_tokens = min( |
| 178 | self.max_draft_token_num, |
| 179 | self.max_model_len - num_tokens - 1, |
| 180 | ) |
| 181 | if max_spec_tokens <= 1: |
| 182 | continue |
| 183 | if req_id is None: |
| 184 | continue |
| 185 | |
| 186 | # 3. Add accept tokens to context |
| 187 | acc_num = int(accept_num_cpu[bid]) |
| 188 | assert ( |
| 189 | acc_num > 0 |
| 190 | ), f"Request {req_id} (bid {bid}) must have at least one accepted token, but got {acc_num}." |
| 191 | if acc_num > 0: |
| 192 | token_ids = accept_tokens_cpu[bid, :acc_num] |
| 193 | ctx_start = seq_lens_decoder[bid] - acc_num |
| 194 | self.context_tokens[bid, ctx_start : ctx_start + acc_num] = token_ids |
| 195 | self.add_active_response(req_id, token_ids) |
| 196 | |
| 197 | # 4. Get context |
| 198 | start = max(0, num_tokens - self.max_tree_depth) |
| 199 | ctx = self.context_tokens[bid, start:num_tokens] |
| 200 | ctx = ctx[ctx >= 0] |
| 201 | if ctx.size == 0: |
| 202 | continue |
| 203 | if not ctx.flags["CONTIGUOUS"]: |
| 204 | ctx = np.ascontiguousarray(ctx, dtype=np.int32) |
nothing calls this directly
no test coverage detected