[Feature] Add MiniCPM-SALA support (#30360)
Co-authored-by: Alex Nails <alex.nails@radixark.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
co-authored by
Alex Nails
Claude Opus 5
parent
d251fa2453
commit
092d85eb87
@@ -281,6 +281,7 @@ class ReqToTokenPool:
|
||||
)
|
||||
self.free_slots = list(range(1, self._alloc_size))
|
||||
self.req_generation = torch.zeros(self._alloc_size, dtype=torch.int64)
|
||||
self._aux_cache: Any = None
|
||||
|
||||
def write(self, indices, values):
|
||||
self.req_to_token[indices] = values
|
||||
@@ -324,12 +325,41 @@ class ReqToTokenPool:
|
||||
|
||||
def free(self, req: Req):
|
||||
assert req.req_pool_idx is not None, "request must have req_pool_idx"
|
||||
if self._aux_cache is not None:
|
||||
self._aux_cache.free(req.req_pool_idx)
|
||||
self.free_slots.append(req.req_pool_idx)
|
||||
req.req_pool_idx = None
|
||||
|
||||
def clear(self):
|
||||
self.free_slots = list(range(1, self._alloc_size))
|
||||
self.req_generation.zero_()
|
||||
if self._aux_cache is not None:
|
||||
self._aux_cache.clear()
|
||||
|
||||
def attach_aux_cache(self, aux_cache: Any) -> None:
|
||||
assert self._aux_cache is None
|
||||
self._aux_cache = aux_cache
|
||||
|
||||
def reset_aux_cache_allocator(self) -> None:
|
||||
if self._aux_cache is not None:
|
||||
self._aux_cache.reset_allocator()
|
||||
|
||||
def schedulable_token_capacity(self, physical_capacity: int) -> int:
|
||||
if self._aux_cache is None:
|
||||
return physical_capacity
|
||||
return self._aux_cache.dense_capacity
|
||||
|
||||
def alloc_aux_to_lengths(
|
||||
self,
|
||||
*,
|
||||
req_pool_indices_cpu: torch.Tensor,
|
||||
target_seq_lens_cpu: torch.Tensor,
|
||||
) -> None:
|
||||
if self._aux_cache is not None:
|
||||
self._aux_cache.alloc_to_lengths(
|
||||
req_pool_indices_cpu=req_pool_indices_cpu,
|
||||
target_seq_lens_cpu=target_seq_lens_cpu,
|
||||
)
|
||||
|
||||
|
||||
class MambaPool:
|
||||
|
||||
Reference in New Issue
Block a user