mlc-ai / mlc-llm

Universal LLM Deployment Engine with ML Compilation
https://llm.mlc.ai/
Apache License 2.0
19.23k stars 1.58k forks source link

[C++] Invoke storage allocation for CUDA Graph explicitly #3042

Closed MasterJH5574 closed 14 hours ago

MasterJH5574 commented 16 hours ago

This PR adds a function that invokes the storage allocation function generated by CUDA Graph rewrite. With this function, we now manually trigger the storage allocation at initialization time.

The reason we need this is because that the storage allocation may contain CUDA IPC memory alloc that has to run through a Disco session. So when a function that needs CUDA graph storage allocation runs first outside a Disco session, there might be error caused if we did not initialize the allocation in advance.