92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
import torch
|
||
import torch_vacc
|
||
|
||
class BinCountTensorPooler:
|
||
"""
|
||
sampler类中（bin-count tensor）的缓冲器。
|
||
核心功能：为apply_penalties 提供创建好的bin-count, 为避免重复重头计算。
|
||
注意: 如果系统中的req_id未能确保在max_count的数量下的唯一性，要考虑是否适用该pooler
|
||
|
||
pooler = BinCountTensorPooler(15169, "cpu")
|
||
list_1 = pooler.request_tensors(['1','3','aad'])
|
||
bin_count_buffers = pooler.request_tensors(req_ids)
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
vocab_size: int,
|
||
device: torch.device,
|
||
max_cache_count: int = 20,
|
||
):
|
||
# 缓存容器：用列表维护请求ID与对应张量的顺序（FIFO淘汰）
|
||
self._cached_tensors: list[torch.Tensor] = []
|
||
self._cached_req_ids: list[str] = []
|
||
|
||
# 张量维度（+1 通常用于padding/特殊标记）
|
||
self.vocab_size: int = vocab_size + 1
|
||
self.device: torch.device = device
|
||
self.max_cache_count: int = max_cache_count
|
||
|
||
def request_tensors(self, req_ids: list[str]) -> list[torch.Tensor]:
|
||
"""
|
||
批量请求张量：已缓存的直接返回，未缓存的创建并加入缓存。
|
||
|
||
Args:
|
||
req_ids: 待请求的请求ID列表
|
||
|
||
Returns:
|
||
与req_ids顺序对应的bin-count tensor
|
||
"""
|
||
if not req_ids:
|
||
return [] # 空输入快速返回，避免无效循环
|
||
|
||
out_tensors = []
|
||
for req_id in req_ids:
|
||
if req_id not in self._cached_req_ids:
|
||
# 注册新的req, 增加cached-tensor
|
||
self._add_new_cache(req_id)
|
||
|
||
# 快速获取缓存索引（后续可优化为dict映射，提升性能）
|
||
cache_idx = self._cached_req_ids.index(req_id)
|
||
out_tensors.append(self._cached_tensors[cache_idx])
|
||
|
||
return out_tensors
|
||
|
||
def _add_new_cache(self, req_id: str) -> None:
|
||
"""
|
||
新增缓存项：创建张量并加入缓存，超出最大数量时按FIFO淘汰最早项。
|
||
私有方法：封装内部逻辑，避免外部直接调用
|
||
"""
|
||
# FIFO淘汰：缓存满时移除头部（最早加入的项）
|
||
if len(self._cached_req_ids) >= self.max_cache_count: # >= 更严谨（避免边界值问题）
|
||
self._cached_tensors.pop(0)
|
||
self._cached_req_ids.pop(0)
|
||
|
||
# 创建二进制计数张量（int32类型，初始全0）
|
||
bin_count_tensor = torch.zeros(
|
||
size=(1, self.vocab_size),
|
||
dtype=torch.int32,
|
||
device=self.device,
|
||
requires_grad=False, # 明确禁用梯度（计数张量无需反向传播）
|
||
)
|
||
|
||
# 同步添加ID和张量（保证两个列表索引一致）
|
||
self._cached_req_ids.append(req_id)
|
||
self._cached_tensors.append(bin_count_tensor)
|
||
|
||
def clear_cache(self) -> None:
|
||
"""清空所有缓存（可选扩展方法，方便外部手动清理）"""
|
||
self._cached_tensors.clear()
|
||
self._cached_req_ids.clear()
|
||
|
||
def get_cache_status(self) -> dict:
|
||
"""获取缓存状态（可选扩展方法，方便监控）"""
|
||
return {
|
||
"current_cache_count": len(self._cached_req_ids),
|
||
"max_cache_count": self.max_cache_count,
|
||
"cached_req_ids": self._cached_req_ids.copy(), # 返回副本避免外部修改
|
||
"tensor_shape": (1, self.vocab_size),
|
||
"tensor_dtype": torch.int32,
|
||
"device": str(self.device),
|
||
}
|