Current limitation
PICache releases Mamba state slots one at a time. Each slot constructs a separate device tensor and calls the allocator independently, causing repeated kernel launches and repeated torch.cat work.
slot 1 --> tensor([1]) --> free() --> torch.cat
slot 2 --> tensor([2]) --> free() --> torch.cat
... ...
slot N --> tensor([N]) --> free() --> torch.cat
The slot IDs are already available together during cache removal, so this per-slot device work is avoidable.
Desired outcome
Collect all released Mamba slot IDs, construct one device tensor, and return them through one allocator call while preserving cache-removal and ownership semantics.
[slot 1, slot 2, ..., slot N]
|
v
one device tensor
|
v
one free()/torch.cat
Performance evidence
A focused CPU allocator benchmark, repeated seven times per batch size, measured:
| Released slots |
Per-slot release |
Batched release |
Improvement |
| 128 |
3.223 ms |
0.599 ms |
5.4x |
| 512 |
8.065 ms |
1.186 ms |
6.8x |
| 1,024 |
17.973 ms |
2.406 ms |
7.5x |
These measurements isolate allocator bookkeeping rather than end-to-end inference.
Current limitation
PICache releases Mamba state slots one at a time. Each slot constructs a separate device tensor and calls the allocator independently, causing repeated kernel launches and repeated
torch.catwork.The slot IDs are already available together during cache removal, so this per-slot device work is avoidable.
Desired outcome
Collect all released Mamba slot IDs, construct one device tensor, and return them through one allocator call while preserving cache-removal and ownership semantics.
Performance evidence
A focused CPU allocator benchmark, repeated seven times per batch size, measured:
These measurements isolate allocator bookkeeping rather than end-to-end inference.