Skip to content

[Feature]: (Performance) Implement async two-phase MoE execution (begin/end) for CUDA backend #252

Description

@RaphyTwin

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    cudaBackend CUDA/NVIDIA

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions