[PG] Support full reduction ops (Product/Min/Max) and fix reduce kernel indexing bug by hhr2449 · Pull Request #1440 · kvcache-ai/Mooncake

hhr2449 · 2026-01-26T15:21:09Z

Description

Part of #1225

This PR implements support for reduction operations in the CUDA backend and fixes a indexing bug in the reduceKernel.

Key Changes:

Support for New Reduce Ops:
- Implemented PRODUCT, MIN, and MAX logic in reduceKernel
- Updated launchReduceKernel assertions to allow these operations.
Bugfix in reduceKernel:
- Fixed an indexing issue when initializing the accumulator (acc) in the CUDA kernel.
- Before: acc = src[elem_idx]
- After: acc = src[rank * numElements + elem_idx]
Unit Tests:
- Updated mooncake-wheel/tests/test_mooncake_backend.py to include test cases for all_reduce_product, all_reduce_min, and all_reduce_max.
- Ensured tests are running on CUDA device to trigger the kernel execution.

Type of Change

Types
- Bug fix
- New feature
  - Transfer Engine
  - Mooncake Store
  - Mooncake EP
  - Integration
  - P2P Store
  - Python Wheel
- Breaking change
- CI/CD
- Documentation update
- Other

How Has This Been Tested?

I ran the updated unit tests in a GPU environment to verify the correctness of the new reduction operations and the kernel logic.

Test Command: python3 -m unittest mooncake-wheel.tests.test_mooncake_backend
Results: Passed 5 tests (all_reduce_sum, all_reduce_product, all_reduce_min, all_reduce_max, all_gather).

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
I have updated the documentation.
I have added tests to prove my changes are effective.

…tests in test_mooncake_backend.py to cover new ops

gemini-code-assist · 2026-01-26T15:21:22Z

Summary of Changes

Hello @hhr2449, 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 significantly enhances the CUDA backend's distributed communication capabilities by introducing support for product, minimum, and maximum reduction operations. Concurrently, it addresses a crucial indexing error in the reduceKernel that previously led to incorrect accumulator initialization. The changes are thoroughly validated with new unit tests, ensuring robust and accurate collective operations.

Highlights

New Reduction Operations: Implemented support for PRODUCT, MIN, and MAX reduction operations within the CUDA backend's reduceKernel and updated the launchReduceKernel assertions to recognize these new operations.
Critical Bug Fix: Resolved an indexing bug in the reduceKernel where the accumulator (acc) was incorrectly initialized. The fix changes acc = src[elem_idx] to acc = src[rank * numElements + elem_idx] to correctly access elements across different ranks.
Expanded Unit Tests: Added new unit tests in test_mooncake_backend.py for all_reduce_product, all_reduce_min, and all_reduce_max to ensure the correctness of the newly implemented reduction operations and the bug fix, specifically targeting CUDA device execution.

🧠 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 successfully adds support for PRODUCT, MIN, and MAX reduction operations to the CUDA backend and fixes a critical indexing bug in the reduceKernel. The changes in the CUDA kernel are correct and the TORCH_CHECK has been updated accordingly. The accompanying Python tests are well-written and cover the new functionality. I've included a few minor suggestions to improve code style and adhere to PEP 8 guidelines in the test file.

gemini-code-assist · 2026-01-26T15:23:17Z

+        tensor = torch.tensor([2], dtype = torch.int32, device = "cuda")
+        dist.all_reduce(tensor, op = dist.ReduceOp.PRODUCT)


There are extra spaces around the = for keyword arguments. According to the PEP 8 style guide, there should be no spaces in this context. This issue is also present on lines 30-31 and 35-36.

Suggested change

tensor = torch.tensor([2], dtype = torch.int32, device = "cuda")

dist.all_reduce(tensor, op = dist.ReduceOp.PRODUCT)

tensor = torch.tensor([2], dtype=torch.int32, device="cuda")

dist.all_reduce(tensor, op=dist.ReduceOp.PRODUCT)

References

PEP 8 recommends not using spaces around the = sign for keyword arguments. ^(link)

codecov-commenter · 2026-01-26T15:43:25Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

UNIDY2002

Nice job!

…el indexing bug (kvcache-ai#1440)

hhr2449 added 3 commits January 23, 2026 22:22

test

1cf7ca5

Impelemented PRODUCT, MIN, MAX logic in reduce kerlel and Added unit …

e0fff75

…tests in test_mooncake_backend.py to cover new ops

Impelemented PRODUCT, MIN, MAX logic in reduce kerlel and Added unit …

5ae2e0a

…tests in test_mooncake_backend.py to cover new ops

hhr2449 requested review from ShangmingCai, UNIDY2002, stmatengss and ympcMark as code owners January 26, 2026 15:21

github-actions Bot added the run-ci label Jan 26, 2026

gemini-code-assist Bot reviewed Jan 26, 2026

View reviewed changes

style: fix clang-format errors

df32f57

stmatengss assigned UNIDY2002 and ympcMark Jan 26, 2026

UNIDY2002 approved these changes Jan 27, 2026

View reviewed changes

UNIDY2002 merged commit 0b8e1d6 into kvcache-ai:main Jan 27, 2026
16 checks passed

XucSh pushed a commit to XucSh/Mooncake that referenced this pull request Jan 27, 2026

[PG] Support full reduction ops (Product/Min/Max) and fix reduce kern…

9ab529c

…el indexing bug (kvcache-ai#1440)

JasonZhang517 pushed a commit to JasonZhang517/Mooncake that referenced this pull request Feb 9, 2026

[PG] Support full reduction ops (Product/Min/Max) and fix reduce kern…

cabaaa2

…el indexing bug (kvcache-ai#1440)

UNIDY2002 mentioned this pull request Mar 9, 2026

[RFC]: Future plans for improving Mooncake EP #1225

Closed

9 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[PG] Support full reduction ops (Product/Min/Max) and fix reduce kernel indexing bug#1440

[PG] Support full reduction ops (Product/Min/Max) and fix reduce kernel indexing bug#1440
UNIDY2002 merged 4 commits intokvcache-ai:mainfrom
hhr2449:main

hhr2449 commented Jan 26, 2026 •

edited by UNIDY2002

Loading

Uh oh!

gemini-code-assist Bot commented Jan 26, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jan 26, 2026

Uh oh!

codecov-commenter commented Jan 26, 2026

Uh oh!

UNIDY2002 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		tensor = torch.tensor([2], dtype = torch.int32, device = "cuda")
		dist.all_reduce(tensor, op = dist.ReduceOp.PRODUCT)

Conversation

hhr2449 commented Jan 26, 2026 • edited by UNIDY2002 Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist Bot commented Jan 26, 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 26, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Jan 26, 2026

Codecov Report

Uh oh!

UNIDY2002 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

hhr2449 commented Jan 26, 2026 •

edited by UNIDY2002

Loading