[Feature Request] GEMM+RS example

[ming-wei]

Hi,

It seems to me that example_gemm_rs.py is incomplete. It calls into gemm_rs_op, however the actual implementation (gemm_rs_producer_persistent/gemm_rs_producer_non_persistent) are missing/commented out.

It'd be great to provide a full example of GEMM+RS, since it is a common pattern in distributed LLM training/inference. Do we have a roadmap for this?

Thanks.

[chengyupku]

Hi ming-wei, some of the kernels are not fully implemented yet. We initially used the nvshmem API, but later discovered that it introduced some performance issues. We're currently replacing those parts with our own implementation using hand-written UVA (Unified Virtual Addressing) copy logic to address the problem.

We also recommend checking out our roadmap for more context.