[Perf] Use np.ndarray instead of list[list[int]] to reduce GC overhead (#28245)

Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
This commit is contained in:
Jialin Ouyang
2025-11-11 10:28:47 -08:00
committed by GitHub
parent 76e4dcf225
commit 4228be7959
3 changed files with 17 additions and 10 deletions

View File

@@ -5,6 +5,7 @@ import random
from dataclasses import dataclass
from typing import TypeAlias
import numpy as np
import torch
from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
@@ -369,9 +370,9 @@ class MockEngineCore:
self.generated_logprobs_raw[req_idx][token_idx]
)
logprobs = LogprobsLists(
[logprobs_token_ids_],
[logprobs_],
[sampled_token_ranks_],
np.array([logprobs_token_ids_]),
np.array([logprobs_]),
np.array([sampled_token_ranks_]),
)
else:
logprobs = None