Skip to content

Commit 96d9c8b

Browse files
authored
[Serve] add allocator in Storage as the upstream change (#1997)
The changes in apache/tvm#16750 modified the signature of the Storage, this pull request updates the caller code in mlc-llm to accommodate the new Storage class signature. Ran into build error w/o the change.
1 parent cc36324 commit 96d9c8b

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

cpp/serve/model.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,8 +63,8 @@ class ModelImpl : public ModelObj {
6363
memory::Allocator* allocator =
6464
memory::MemoryManager::GetOrCreateAllocator(device_host, memory::AllocatorType::kNaive);
6565
ICHECK_NOTNULL(allocator);
66-
token_ids_storage_ =
67-
memory::Storage(allocator->Alloc(device_host, {prefill_chunk_size_}, DataType::Int(32)));
66+
token_ids_storage_ = memory::Storage(
67+
allocator->Alloc(device_host, {prefill_chunk_size_}, DataType::Int(32)), allocator);
6868
this->logit_pos_arr_ = NDArray::Empty({max_num_sequence}, DataType::Int(32), device_host);
6969
}
7070

0 commit comments

Comments
 (0)