[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:
cauphe
2026-08-24 02:25:16 -07:00
committed by GitHub
co-authored by Alex Nails Claude Opus 5
parent d251fa2453
commit 092d85eb87
44 changed files with 7055 additions and 111 deletions
@@ -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: