From 5a40117a3741015d8d6f5f80966e98c42745609a Mon Sep 17 00:00:00 2001 From: Xinyu Jiang Date: Mon, 2 Nov 2026 04:55:30 +0000 Subject: [PATCH] [Bugfix] Skip empty inputs in custom all-gather Return from dispatchAllGather before the buffer lookup when the input is empty, and skip the zero-byte hipMemcpyAsync in all_gather_unreg, whose source pointer is null for an empty tensor. Signed-off-by: Xinyu Jiang Co-authored-by: Zhiyao Jiang --- csrc/include/custom_all_reduce.cuh | 6 ++++++ csrc/kernels/custom_all_reduce.cu | 7 +++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/csrc/include/custom_all_reduce.cuh b/csrc/include/custom_all_reduce.cuh index 9e72f961d84..2a2fcf5fa40 100644 --- a/csrc/include/custom_all_reduce.cuh +++ b/csrc/include/custom_all_reduce.cuh @@ -5130,6 +5130,12 @@ template void dispatchAllGather( hipStream_t stream, T* input, T* output, int size, int last_dim_size, int gather_dim) { + // Idle DP-attention groups can have empty shards. There is no payload to exchange, + // and launching zero blocks would leave an error for the next operator to report. + if(size == 0) + { + return; + } RankData* ptrs = get_buffer_RD(stream, input); auto d = 16 / sizeof(T); dim3 block(512); diff --git a/csrc/kernels/custom_all_reduce.cu b/csrc/kernels/custom_all_reduce.cu index 44f8bf40eae..a804dd9cdc8 100644 --- a/csrc/kernels/custom_all_reduce.cu +++ b/csrc/kernels/custom_all_reduce.cu @@ -521,8 +521,11 @@ void all_gather_unreg(fptr_t _fa, if(data_bytes > reg_bytes) throw std::runtime_error("registered buffer is too small to contain the input"); - HIP_CALL(hipMemcpyAsync((void*)reg_buffer, inp.data_ptr(), data_bytes, - hipMemcpyDeviceToDevice, stream)); + if(data_bytes != 0) + { + HIP_CALL(hipMemcpyAsync( + (void*)reg_buffer, inp.data_ptr(), data_bytes, hipMemcpyDeviceToDevice, stream)); + } _all_gather(_fa, (void*)reg_buffer, out.data_ptr(), inp.numel(), inp.dtype(), last_dim_size, dim); }