From 015bf0600876fd7cf2849411a0c99dcb02bcace2 Mon Sep 17 00:00:00 2001
From: "assistant-librarian[bot]"
 <210906412+assistant-librarian[bot]@users.noreply.github.com>
Date: Thu, 12 Feb 2026 09:44:51 -0800
Subject: [PATCH] add memsetasync for ck moe splitk (#4282)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Proposed changes

add memsetasync for ck moe splitk to fix

## Checklist

Please put an `x` into the boxes that apply. You can also fill these out
after creating the PR. If you're not sure, please don't hesitate to ask.

- [ ] I have added tests relevant to the introduced functionality, and
the unit tests are passing locally
- [ ] I have added the test to REGRESSION_TESTS list defined at the top
of CMakeLists.txt in tests/CMakeLists.txt, **IF** the test takes more
than 30 seconds to run.
- [ ] I have added inline documentation which enables the maintainers
with understanding the motivation
- [ ] I have removed the stale documentation which is no longer relevant
after this pull request
- [ ] (If this change is user-facing) I have added release notes which
provide the end users with a brief summary of the improvement from this
pull request
- [ ] I have run `clang-format` on all changed files
- [ ] Any dependent changes have been merged

## Discussion

If this is a relatively large or complex change, feel free to start a
discussion by explaining why you chose the solution you did and what
alternatives you considered



---
🔁 Imported from
[ROCm/composable_kernel#3630](https://github.com/ROCm/composable_kernel/pull/3630)
🧑‍💻 Originally authored by @lalala-sh

---------

Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: systems-assistant[bot] <systems-assistant[bot]@users.noreply.github.com>
Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com>
---
 .../impl/device_moe_gemm_blockscale.hpp       | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp
index 56d3b48547..12d28f572c 100644
--- a/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp
+++ b/include/ck/tensor_operation/gpu/device/impl/device_moe_gemm_blockscale.hpp
@@ -253,12 +253,12 @@ struct DeviceMoeGemmBlockScale
                         // rotating mem
                         rotating_mem.Next();
                         // clear c mem
-                        // if(arg_.KBatch > 1)
-                        //     hipGetErrorString(hipMemsetAsync(arg_.p_c_grid,
-                        //                                      0,
-                        //                                      arg_.M * arg_.N * sizeof(CDataType)
-                        //                                      * (IsInputGemm && IsSplitK ? 2 : 1),
-                        //                                      stream_config.stream_id_));
+                        if(arg_.KBatch > 1)
+                            hipGetErrorString(hipMemsetAsync(arg_.p_c_grid,
+                                                             0,
+                                                             arg_.M * arg_.N * sizeof(CDataType) *
+                                                                 (IsInputGemm && IsSplitK ? 2 : 1),
+                                                             stream_config.stream_id_));
                     };
 
                     ave_time = ck::utility::launch_and_time_kernel_with_preprocess<false>(
@@ -272,12 +272,12 @@ struct DeviceMoeGemmBlockScale
                 }
                 else
                 {
-                    // if(arg.KBatch > 1)
-                    //     hipGetErrorString(hipMemsetAsync(arg.p_c_grid,
-                    //                                      0,
-                    //                                      arg.M * arg.N * sizeof(CDataType) *
-                    //                                      (IsInputGemm && IsSplitK ? 2 : 1),
-                    //                                      stream_config.stream_id_));
+                    if(arg.KBatch > 1)
+                        hipGetErrorString(hipMemsetAsync(arg.p_c_grid,
+                                                         0,
+                                                         arg.M * arg.N * sizeof(CDataType) *
+                                                             (IsInputGemm && IsSplitK ? 2 : 1),
+                                                         stream_config.stream_id_));
 
                     ave_time = launch_and_time_kernel(
                         stream_config, kernel, dim3(gdx, gdy, gdz), dim3(BlockSize), 0, arg);