Problem
In the metal backend we have an elegant async execution model (coli_metal_moe_block_begin / _end), allowing the CPU/SSD to fetch missed experts while the gpu computes the resident ones. The CUDA backend (backend_cuda.h) currently only exposes synchronous blocking execution calls like coli_cuda_expert_group, missing out on parallelizing disk-I/O with GPU execution.
Proposed solution
Introduce a two-phase async API for the CUDA backend similar to Metal:
Implement coli_cuda_moe_block_begin and coli_cuda_moe_block_end.
Use non-default CUDA streams (cudaStream_t) to execute math kernels on resident experts on the gpu asynchronously.
While the gpu compute stream is busy, the host thread can continuously upload newly fetched experts to the GPU using asynchronous memory copies (cudaMemcpyAsync) on a separate copy stream.
Alternatives considered
No response
Scope and compatibility
No response
Problem
In the metal backend we have an elegant async execution model (coli_metal_moe_block_begin / _end), allowing the CPU/SSD to fetch missed experts while the gpu computes the resident ones. The CUDA backend (backend_cuda.h) currently only exposes synchronous blocking execution calls like coli_cuda_expert_group, missing out on parallelizing disk-I/O with GPU execution.
Proposed solution
Introduce a two-phase async API for the CUDA backend similar to Metal:
Implement coli_cuda_moe_block_begin and coli_cuda_moe_block_end.
Use non-default CUDA streams (cudaStream_t) to execute math kernels on resident experts on the gpu asynchronously.
While the gpu compute stream is busy, the host thread can continuously upload newly fetched experts to the GPU using asynchronous memory copies (cudaMemcpyAsync) on a separate copy stream.
Alternatives considered
No response
Scope and compatibility
No response