Skip to content

Conversation

@masahi
Copy link
Member

@masahi masahi commented Feb 15, 2024

Currently we rely on an ugly hack of abusing UsedMemory to signal the end of memory profiling, after which PooledAllocator switches to a different pooling mode (another hack). https://github.com/octoml/tvm/blob/for-mlc-serve-jan9/src/runtime/memory/pooled_allocator.h#L114-L116

In particular, this hack assumes that UsedMemory is called only once, after memory profiling. This blocks #210, which wants to call UsedMemory before / after memory-profiling inference.

This PR builds on https://github.com/octoml/tvm/pull/45 and explicitly toggles memory profiling start / stop. Please update to the latest for-mlc-serve-jan12 after this PR is merged.

@Lunderberg This should unblock your work.

@vinx13 @sunggg

@masahi masahi force-pushed the memory-profiling-cleanup branch from 721b6f5 to 4154bfe Compare February 15, 2024 09:01
@sunggg sunggg merged commit 6302271 into octoml:batch-serving Feb 15, 2024
@Lunderberg
Copy link

Thank you! With the changes here, all is running smoothly.

Lunderberg pushed a commit to Lunderberg/mlc-llm that referenced this pull request Feb 27, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants