Test reduce by littleyellowbicycle · Pull Request #307 · sgl-project/sgl-kernel-npu

littleyellowbicycle · 2026-01-07T12:32:44Z

No description provided.

…-npu into shmem_ver

Merge reduce scatter op

fix reduce scatter return

ag & rs fix and zero buff version update

gemini-code-assist · 2026-01-07T12:33:22Z

Summary of Changes

Hello @littleyellowbicycle, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a new ZCCL (Zero-Copy Collective Communication Library) module, providing optimized collective communication primitives for Ascend NPUs. It includes the core implementations for AllReduce, AllGather, and ReduceScatter operations, meticulously crafted for both host and device execution. The module incorporates advanced tiling strategies to efficiently manage data distribution and processing across multiple cores. Accompanying these implementations is a robust testing framework, complete with build scripts, data generation utilities, and performance analysis tools, ensuring the correctness and efficiency of the new collective operations. It's noted that the AllReduce and AllGather kernels currently utilize the same processing path for both small and large data, whereas ReduceScatter correctly differentiates its handling based on data size.

Highlights

New ZCCL Module: Introduction of a new zccl module for collective communication operations, including AllReduce, AllGather, and ReduceScatter.
Ascend NPU Optimization: Operations are tailored for Ascend NPUs, utilizing acl for device management and shmem_api for inter-process communication.
Tiling for Efficiency: Implements data tiling strategies (e.g., AllReduceTilingData, AllGatherTilingData, ReduceScatterTilingData) to optimize parallel processing across cores.
Comprehensive Test Suite: A robust testing framework is included, featuring CMake build configurations, shell scripts for execution and profiling, Python scripts for data generation and performance analysis, and C++ test cases with golden data verification.
Initial Data Size Handling: The AllReduce and AllGather kernels currently use the 'small data' processing path for both small and large data sizes, while ReduceScatter correctly dispatches to separate small and big data handling logic.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces new collective communication operations (AllReduce, AllGather, ReduceScatter) under a zccl module, along with corresponding tests. While the additions are significant, the current implementation contains several critical issues that must be addressed. These include use-after-free bugs due to missing stream synchronization before freeing resources, memory management errors such as leaks and incorrect deallocation, and logical flaws in the kernel implementations. I've provided detailed comments and suggestions to fix these problems.

gemini-code-assist · 2026-01-07T12:36:16Z

csrc/zccl/allgather/op_host/allgather.cpp

+    aclrtFreeHost(tiling_host.get());
+    aclrtFree(tiling_device_ptr);
+    shmem_free(gva);


The kernel is launched asynchronously, but resources (tiling_device_ptr, gva) are freed immediately without waiting for the kernel to complete. This is a use-after-free bug. You must synchronize the stream before freeing these resources. Also, aclrtFreeHost(tiling_host.get()) is incorrect as tiling_host is managed by a shared_ptr and was not allocated with aclrtMallocHost.

aclrtSynchronizeStream(stream); aclrtFree(tiling_device_ptr); shmem_free(gva);

gemini-code-assist · 2026-01-07T12:36:16Z

csrc/zccl/reduce_scatter/op_kernel/reduce_scatter_kernel.cpp

+    uint32_t rankSize = shmem_team_n_pes(teamId);
+    ZCCLDataType zcclDataType = static_cast<ZCCLDataType>(dataType);
+    size_t typeSize = getSizeFromTypeEnum(zcclDataType);
+    bool smallFlag = (totalLength >= BIG_DATA_SIZE / typeSize) ? false : true;


The logic for smallFlag seems to be inverted. totalLength * typeSize < BIG_DATA_SIZE should result in smallFlag being true. Currently, it will be false for small data.

bool smallFlag = (totalLength * typeSize < BIG_DATA_SIZE);

gemini-code-assist · 2026-01-07T12:36:16Z

csrc/zccl/all_reduce/op_host/all_reduce.cpp

+    tilingconfig->outputLastNumCore_ =
+        elements - ((tilingconfig->outputCorePerRank_) - 1) * (tilingconfig->outputCorePerRank_);


There appears to be a copy-paste error in the calculation of outputLastNumCore_. It's using tilingconfig->outputCorePerRank_ instead of tilingconfig->outputNumPerCore_ in the multiplication. This will lead to an incorrect tiling configuration.

tilingconfig->outputLastNumCore_ = elements - (tilingconfig->outputCorePerRank_ - 1) * tilingconfig->outputNumPerCore_;

gemini-code-assist · 2026-01-07T12:36:16Z

csrc/zccl/all_reduce/op_host/all_reduce.cpp

+    shmem_free(gva);
+    return 0;


The kernel is launched asynchronously on the stream. However, shmem_free(gva) is called immediately after, without waiting for the kernel to finish. This is a race condition that can lead to a use-after-free error. The same issue applies to deviceTilingConfig, which is a shared_ptr whose memory will be freed when it goes out of scope at the end of the function. You should synchronize the stream before freeing these resources.

aclrtSynchronizeStream(stream); shmem_free(gva);

gemini-code-assist · 2026-01-07T12:36:16Z

csrc/zccl/reduce_scatter/op_host/reduce_scatter.cpp

+    /* launch the kernel function via ACLRT_LAUNCH_KERNEL */
+    ACLRT_LAUNCH_KERNEL(ShmemReduceScatter)(blockDim, stream, inp, out, (uint8_t *)ptr,
+                                            fftsAddr, dataTypeNum, inpNumel, teamId, reduceOp);
+    shmem_free(ptr);


The kernel is launched asynchronously, but shmem_free(ptr) is called immediately without waiting for the kernel to finish. This can lead to a use-after-free error. You should synchronize the stream before freeing the memory.

ACLRT_LAUNCH_KERNEL(ShmemReduceScatter)(blockDim, stream, inp, out, (uint8_t *)ptr, fftsAddr, dataTypeNum, inpNumel, teamId, reduceOp); aclrtSynchronizeStream(stream); shmem_free(ptr);

gemini-code-assist · 2026-01-07T12:36:17Z

csrc/zccl/allgather/op_host/allgather_tiling_data.h

+    uint32_t output_last_num_core;
+};
+
+}  // namespace npu_kernel


The namespace here is npu_kernel, but it should be zccl to match the opening namespace sglang::zccl. This seems to be a copy-paste error.

Suggested change

} // namespace npu_kernel

} // namespace zccl

gemini-code-assist · 2026-01-07T12:36:17Z

csrc/zccl/all_reduce/op_kernel/all_reduce_kernel.cpp

+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#ifndef SGL_KERNEL_NPU_ALL_GATHER_KERNEL_H


The include guard SGL_KERNEL_NPU_ALL_GATHER_KERNEL_H seems to be a copy-paste error from the all_gather kernel. It should be specific to all_reduce.

#ifndef SGL_KERNEL_NPU_ALL_REDUCE_KERNEL_H

gemini-code-assist · 2026-01-07T12:36:17Z

csrc/zccl/all_reduce/op_host/all_reduce_tiling.h

+    uint32_t outputLastNumCore_;
+};
+
+}  // namespace npu_kernel


The namespace here is npu_kernel, but it should be zccl to match the opening namespace sglang::zccl. This seems to be a copy-paste error.

Suggested change

} // namespace npu_kernel

} // namespace zccl

gemini-code-assist · 2026-01-07T12:36:17Z

csrc/zccl/all_reduce/op_host/all_reduce.cpp

+    void* rawPtr = nullptr;
+    auto res = aclrtMalloc(&rawPtr, size, policy);
+    if (res != 0) {
+        return nullptr;  // 分配失败


This comment is in Chinese, while the rest of the codebase is in English. For consistency, please write comments in English.

return nullptr; // Allocation failed

gemini-code-assist · 2026-01-07T12:36:17Z

tests/csrc/zccl/reduce_scatter/scripts/data_gen.py

+
+
+def gen_random_data(size, dtype):
+    return np.ones_like(size, dtype=dtype)


Using np.ones_like for generating test data is not ideal for testing reduction operations like SUM, MAX, MIN, or PROD, as the results will be trivial. Using random data (as in the commented-out line) would provide more robust testing.

Suggested change

return np.ones_like(size, dtype=dtype)

return np.random.uniform(low=0.0, high=10.0, size=size).astype(dtype)

shadowxz109 and others added 21 commits December 18, 2025 10:37

add allgather using shmem

0a5e921

add allgather using shmem

e60e278

add allgather using shmem

5d33da6

add allgather using shmem

24485eb

Merge branch 'shmem_ver' of https://github.com/shadowxz109/sgl-kernel…

2867ece

…-npu into shmem_ver

add allgather using shmem

028f460

Merge branch 'shmem_ver' of https://github.com/shadowxz109/sgl-kernel…

405c3c0

…-npu into shmem_ver

add allgather using shmem

6b7fee2

Merge branch 'shmem_ver' of https://github.com/shadowxz109/sgl-kernel…

65bcc2c

…-npu into shmem_ver

reduce scatter op

2b40b73

Merge pull request #1 from zhangshuren416/merge

733becd

Merge reduce scatter op

fix return

aa1c331

Merge pull request #2 from zhangshuren416/zsr_dev_new

1948809

fix reduce scatter return

rename allgather

738964f

zero buff rs op

6c2cd1c

allgather zero buff op

dfff0e7

debug

4abf07d

allgather fix

814e29d

rs debug

1934d9f

Merge pull request #3 from zhangshuren416/zsr_dev_new

b111349

ag & rs fix and zero buff version update

fix bug

bc938b3

gemini-code-assist bot reviewed Jan 7, 2026

View reviewed changes

littleYellowBicycle added 2 commits January 12, 2026 14:36

add all_reduce

5860a25

fix data stream

dc76ca0

littleyellowbicycle force-pushed the test_reduce branch from 39929cc to dc76ca0 Compare January 12, 2026 06:36

littleYellowBicycle added 2 commits January 12, 2026 15:58

tests

b83362f

fix

3774397

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Test reduce#307

Test reduce#307
littleyellowbicycle wants to merge 25 commits intosgl-project:shmem_verfrom
shadowxz109:test_reduce

littleyellowbicycle commented Jan 7, 2026

Uh oh!

gemini-code-assist bot commented Jan 7, 2026

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

gemini-code-assist bot Jan 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		tilingconfig->outputLastNumCore_ =
		elements - ((tilingconfig->outputCorePerRank_) - 1) * (tilingconfig->outputCorePerRank_);



		def gen_random_data(size, dtype):
		return np.ones_like(size, dtype=dtype)

	return np.ones_like(size, dtype=dtype)
	return np.random.uniform(low=0.0, high=10.0, size=size).astype(dtype)

Conversation

littleyellowbicycle commented Jan 7, 2026

Uh oh!

gemini-code-assist bot commented Jan 7, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Jan 7, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants