jinchen62 · Oct 7, 2026 · Sep 30, 2026 · Oct 5, 2026 · Oct 5, 2026 · Oct 7, 2026
@@ -3777,6 +3777,7 @@ rocblaslt_status runContractionProblem(rocblaslt_handle                   handle
     }
     catch(const std::exception& e)
     {
+        log_error(__func__, e.what());
     }
     catch(...)
     {
@@ -4169,6 +4170,7 @@ rocblaslt_status makeArgument(rocblaslt_handle             handle,
     }
     catch(const std::exception& e)
     {
+        log_error(__func__, e.what());
     }
     catch(...)
     {

@@ -410,8 +410,9 @@ def _buildCustomKernelFromMetadata(kernelName, fullYaml, kernelConfig):
         grid = ["TilesXYBatchGSU", "One", "One"]
     else:
         # Version 0 kernels rely on hardware gridDim for tile decomposition,
-        # so the grid must be multi-dimensional.
-        grid = ["TilesX", "TilesY", "Batch"]
+        # so the grid must be multi-dimensional. Like generated kernels, they
+        # take their GSU slices along gridDim.y.
+        grid = ["TilesX", "TilesYGSU", "Batch"]
 
     # Workspace is left unset here and derived in
     # Solution._assignCustomKernelParameters, which sees the logic file's

@@ -1882,7 +1882,9 @@ def _assignCustomKernelParameters(state: dict) -> None:
 
     Raises:
       RuntimeError: If a macrotile component is supplied by neither the
-        CustomKernel block nor the consuming logic file.
+        CustomKernel block nor the consuming logic file, or if the solution
+        can split K (GlobalSplitU other than 0 or 1) but its grid launches one
+        GSU slice per tile.
     """
     ck = state["CustomKernel"]
 
@@ -1910,6 +1912,24 @@ def _assignCustomKernelParameters(state: dict) -> None:
     # tile and workspace sizing.
     ck["macrotile"] = macrotile
 
+    # A split-K kernel reduces into D only once every GSU slice of a tile has
+    # arrived, so a grid sized from tile counts alone would return success with D
+    # unwritten. GSU and persistent grids account for the split themselves.
+    # GlobalSplitU -1 lets the runtime pick a split above 1.
+    gsu = state.get("GlobalSplitU", 1)
+    grid = ck.get("grid", [])
+    oneSlicePerTile = all(
+      g in ("One", "TilesX", "TilesY", "Batch", "TilesXY", "TilesXYBatch") for g in grid)
+    if gsu not in (0, 1) and oneSlicePerTile:
+      raise RuntimeError(
+        f"Custom kernel '{ck.get('name', '?')}' runs with GlobalSplitU {gsu}, but its "
+        f"CustomKernel grid {grid} launches one GSU slice per tile; use TilesYGSU "
+        f"or TilesXYBatchGSU.")
+    # The same grid cannot honor a runtime GSU override either; clearing the flag
+    # rejects one during solution selection instead of failing at launch.
+    if oneSlicePerTile:
+      state.setdefault("InternalSupportParams", {})["SupportUserGSU"] = False
+
     # Derive _GlobalAccumulation from GlobalSplitUAlgorithm so the C++
     # runtime sees a non-zero sizeMapping.globalAccumulation for GSU>1
     # solutions.  Without this the legacy beta-only kernel

@@ -815,10 +815,11 @@ def test_build_from_metadata_numworkgroups_grid():
 
 
 def test_build_from_metadata_default_multidim_grid():
+    # Version 0 split-K kernels need one work group per GSU slice along Y.
     ck = _buildCustomKernelFromMetadata(
         "k", _kernel_yaml([_D_ARG]), {"MatrixInstruction": [16, 16, 16, 1]}
     )
-    assert ck["grid"] == ["TilesX", "TilesY", "Batch"]
+    assert ck["grid"] == ["TilesX", "TilesYGSU", "Batch"]
 
 
 def test_build_from_metadata_macrotile_name_fallback():

@@ -115,6 +115,66 @@ def test_assign_custom_kernel_params_default_depthu_is_not_a_tile():
         Solution._assignCustomKernelParameters(state)
 
 
+@pytest.mark.parametrize("grid", [
+    ["TilesX", "TilesYGSU", "Batch"],
+    ["TilesXYBatchGSU", "One", "One"],
+])
+def test_assign_custom_kernel_params_split_k_grid_is_accepted(grid):
+    state = _ck_state(
+        GlobalSplitU=16,
+        GlobalSplitUAlgorithm="MultipleBufferSingleKernel",
+        InternalSupportParams={"SupportUserGSU": True},
+    )
+    state["CustomKernel"]["grid"] = grid
+    Solution._assignCustomKernelParameters(state)
+    assert state["_GlobalAccumulation"] == "MultipleBufferSingleKernel"
+    assert state["InternalSupportParams"]["SupportUserGSU"] is True
+
+
+@pytest.mark.parametrize("over", [
+    {"GlobalSplitU": 16},
+    {"GlobalSplitU": -1},  # lets the runtime pick a split above 1
+    # generateCustomCall judges the grid alone, so a persistent kernel with a
+    # tile-count grid has to be rejected here too rather than at launch.
+    {"GlobalSplitU": 16, "TileProcessingStrategy": "StreamK"},
+])
+def test_assign_custom_kernel_params_split_k_without_gsu_grid_raises(over):
+    # A split-K kernel reduces into D only once every GSU slice has arrived, so a
+    # grid without a GSU term would launch one slice and leave D unwritten.
+    state = _ck_state(GlobalSplitUAlgorithm="MultipleBufferSingleKernel", **over)
+    state["CustomKernel"]["grid"] = ["TilesX", "TilesY", "Batch"]
+    with pytest.raises(RuntimeError, match="launches one GSU slice per tile"):
+        Solution._assignCustomKernelParameters(state)
+
+
+@pytest.mark.parametrize("gsu", [1, 0])  # 0: GSU disabled
+def test_assign_custom_kernel_params_grid_without_gsu_term_rejects_user_gsu(gsu):
+    # Such a grid launches one GSU slice per tile, so a runtime GSU override has to
+    # be turned away during solution selection rather than fail at launch.
+    state = _ck_state(GlobalSplitU=gsu, InternalSupportParams={"SupportUserGSU": True})
+    state["CustomKernel"]["grid"] = ["TilesX", "TilesY", "Batch"]
+    Solution._assignCustomKernelParameters(state)
+    assert state["InternalSupportParams"]["SupportUserGSU"] is False
+
+
+@pytest.mark.parametrize("strategy,grid", [
+    ("StreamK", ["StreamKWithBatch", "One", "One"]),
+    ("DataParallel", ["PersistentGrid", "One", "One"]),
+    ("None", ["PersistentNoBatch", "One", "One"]),
+])
+def test_assign_custom_kernel_params_persistent_keeps_user_gsu(strategy, grid):
+    # Persistent kernels distribute work through their own grid, so neither the
+    # GSU check nor the override flag applies to them.
+    state = _ck_state(
+        GlobalSplitU=16,
+        TileProcessingStrategy=strategy,
+        InternalSupportParams={"SupportUserGSU": True},
+    )
+    state["CustomKernel"]["grid"] = grid
+    Solution._assignCustomKernelParameters(state)
+    assert state["InternalSupportParams"]["SupportUserGSU"] is True
+
+
 def test_assign_custom_kernel_params_enable_mi_sets_wave_params():
     state = _ck_state()
     Solution._assignCustomKernelParameters(state)

@@ -205,6 +205,7 @@ namespace TensileLite
         PersistentGrid,
         PersistentWithBatch,
         PersistentNoBatch,
+        TilesYGSU,
         CustomGridSize_Count,
     };
 

@@ -159,6 +159,7 @@ namespace TensileLite
         case CustomGridSize::PersistentWithBatch: return "PersistentWithBatch";
         case CustomGridSize::PersistentNoBatch: return "PersistentNoBatch";
         case CustomGridSize::TilesXYBatchGSU:  return "TilesXYBatchGSU";
+        case CustomGridSize::TilesYGSU:        return "TilesYGSU";
         case CustomGridSize::CustomGridSize_Count:
             break;
         }
@@ -181,6 +182,7 @@ namespace TensileLite
             {"PersistentWithBatch", CustomGridSize::PersistentWithBatch},
             {"PersistentNoBatch", CustomGridSize::PersistentNoBatch},
             {"TilesXYBatchGSU",  CustomGridSize::TilesXYBatchGSU},
+            {"TilesYGSU",        CustomGridSize::TilesYGSU},
         };
 
         auto it = lookup.find(str);
@@ -2916,6 +2918,9 @@ namespace TensileLite
                 case CustomGridSize::TilesXYBatchGSU:
                     dim = tiles.x * tiles.y * tiles.z * (gsu > 0 ? gsu : 1);
                     break;
+                case CustomGridSize::TilesYGSU:
+                    dim = tiles.y * (gsu > 0 ? gsu : 1);
+                    break;
                 case CustomGridSize::PersistentWithBatch:
                 case CustomGridSize::StreamKWithBatch:
                     // generateCustomCall is only used for handwritten/external
@@ -2940,6 +2945,37 @@ namespace TensileLite
         assignGridSize(rv.numWorkGroups.y, customKernel.grid.y);
         assignGridSize(rv.numWorkGroups.z, customKernel.grid.z);
 
+        // A split-K kernel reduces into D only once every GSU slice of a tile has
+        // arrived, so a grid sized from tile counts alone would return success with
+        // D unwritten. GSU and persistent grids account for the split themselves.
+        auto tileCountOnly = [](CustomGridSize size) {
+            switch(size)
+            {
+            case CustomGridSize::One:
+            case CustomGridSize::TilesX:
+            case CustomGridSize::TilesY:
+            case CustomGridSize::Batch:
+            case CustomGridSize::TilesXY:
+            case CustomGridSize::TilesXYBatch:
+                return true;
+            default:
+                return false;
+            }
+        };
+        if(gsu > 1 && tileCountOnly(customKernel.grid.x) && tileCountOnly(customKernel.grid.y)
+           && tileCountOnly(customKernel.grid.z))
+            throw std::runtime_error(concatenate("Solution ",
+                                                 kernelName,
+                                                 " runs with GSU ",
+                                                 gsu,
+                                                 " but its custom-kernel grid [",
+                                                 customKernel.grid.x,
+                                                 ", ",
+                                                 customKernel.grid.y,
+                                                 ", ",
+                                                 customKernel.grid.z,
+                                                 "] launches one GSU slice per tile"));
+
         bool enableCluster = (sizeMapping.clusterDim.x > 1 || sizeMapping.clusterDim.y > 1);
         if(internalArgsSupport.persistentLoopArgsVersion == 1 && enableCluster)
         {

@@ -77,10 +77,13 @@ namespace
                           "mi350spx");
     }
 
-    uint32_t emittedSplitK(int16_t compiledGsu, int16_t runtimeGsu)
+    KernelInvocation splitKProbeCall(int16_t        compiledGsu,
+                                     int16_t        runtimeGsu,
+                                     CustomGridSize gridY = CustomGridSize::TilesYGSU)
     {
         ContractionSolution solution;
         configureProbeKernel(solution, {CustomArgType::uint32, CustomArgSemantic::SplitK});
+        solution.customKernel.grid.y      = gridY;
         solution.sizeMapping.globalSplitU = compiledGsu;
 
         auto problem = dummyProblem();
@@ -90,7 +93,12 @@ namespace
         ContractionInputs inputs;
         StreamKSettings   sk;
 
-        auto invocation = solution.generateCustomCall<false>(problem, inputs, device, sk);
+        return solution.generateCustomCall<false>(problem, inputs, device, sk);
+    }
+
+    uint32_t emittedSplitK(int16_t compiledGsu, int16_t runtimeGsu)
+    {
+        auto invocation = splitKProbeCall(compiledGsu, runtimeGsu);
 
         EXPECT_EQ(invocation.args.size(), sizeof(uint32_t));
         uint32_t splitK = 0;
@@ -141,6 +149,36 @@ TEST(CustomKernelTest, SplitKFallsBackToTheCompiledGsu)
     EXPECT_EQ(emittedSplitK(/*compiled*/ 16, /*runtime*/ 0), 4u);
 }
 
+// A split-K kernel reduces into D only once every GSU slice of a tile has arrived,
+// so TilesYGSU has to launch one row of work groups per slice, following the
+// runtime-effective GSU like the SplitK argument does.
+TEST(CustomKernelTest, TilesYGsuLaunchesEveryGsuSlice)
+{
+    auto const unsplit = splitKProbeCall(/*compiled*/ 1, /*runtime*/ 0).numWorkGroups;
+
+    auto const compiled = splitKProbeCall(/*compiled*/ 16, /*runtime*/ 0).numWorkGroups;
+    EXPECT_EQ(compiled.x, unsplit.x);
+    EXPECT_EQ(compiled.y, 16 * unsplit.y);
+    EXPECT_EQ(compiled.z, unsplit.z);
+
+    EXPECT_EQ(splitKProbeCall(/*compiled*/ 1, /*runtime*/ 4).numWorkGroups.y, 4 * unsplit.y);
+    EXPECT_EQ(splitKProbeCall(/*compiled*/ 0, /*runtime*/ 0).numWorkGroups.y, unsplit.y);
+    EXPECT_EQ(
+        splitKProbeCall(/*compiled*/ 1, /*runtime*/ 0, CustomGridSize::TilesY).numWorkGroups.y,
+        unsplit.y);
+}
+
+// A grid with no GSU term would launch one slice per tile and report success with
+// D unwritten, so a split-K launch on one has to fail instead.
+TEST(CustomKernelTest, GridWithoutGsuTermRejectsSplitK)
+{
+    EXPECT_THROW(splitKProbeCall(/*compiled*/ 16, /*runtime*/ 0, CustomGridSize::TilesY),
+                 std::runtime_error);
+    EXPECT_THROW(splitKProbeCall(/*compiled*/ 1, /*runtime*/ 4, CustomGridSize::TilesY),
+                 std::runtime_error);
+    EXPECT_NO_THROW(splitKProbeCall(/*compiled*/ 1, /*runtime*/ 0, CustomGridSize::TilesY));
+}
+
 // A custom kernel declares alpha and beta as 32-bit slots, so a narrower compute
 // type has to be widened before it is written or every argument after it shifts.
 // KernelArguments only widens an argument spelled exactly "alpha" or "beta".