[BUG] SM100 mixed TMA/cp.async GEMM: scheduler (CLC) pipeline does not count the epilogue-load warp, hangs when C is loaded #3782

Which component has the problem?

CUTLASS C++

Bug Report

Describe the bug

Summary. In the SM100 mixed TMA + cp.async GEMM kernel (include/cutlass/gemm/kernel/sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp, used by example 92), the scheduler (CLC) pipeline's consumer_arv_count leaves out the epilogue-load warp, which also consumes every scheduler response when the epilogue loads C. Each scheduler stage then receives more arrivals per phase than it expects, so a stage can be released before every consumer has read its response. Example 92, changed to pass beta = 1, hung in 10 of 10 runs on a B300; with #3781, which adds the missing term, it passed 10 of 10.

Cause

• consumer_arv_count (lines 616–624) counts NumMainloopTMALoadThreads + NumMainloopCpAsyncLoadThreads + NumEpilogueThreads + NumMMAThreads. In the static branch, which every example 92 grouped runner takes (its PersistentTileSchedulerSm100Group is not dynamic-persistent), that is 320; the compiled PTX has a literal 320 on the scheduler's empty barriers. The dynamic-persistent branch also adds NumSchedThreads and scales by cluster_size. Neither branch adds NumEpilogueLoadThreads (one warp).
• When collective_epilogue.is_producer_load_needed() holds, the epilogue-load warp calls scheduler.fetch_next_work on every tile (lines 949–964), and each of its threads releases the stage: PipelineAsync::consumer_release (include/cutlass/pipeline/sm90_pipeline.hpp, lines 1253–1257), reached through sm100_tile_scheduler_group.hpp (lines 312–317) and sm90_tile_scheduler_group.hpp (lines 596–610), with no elect. So in the static branch each phase of a stage's empty barrier receives 352 arrivals against the 320 it expects; in the dynamic branch the surplus is cluster_size × 32.
• The extra arrivals can complete a phase early. The scheduler warp can then refill clc_response[stage] before some consumer warp has read it, and the surplus carries into later phases until parity waits alias. We observed hangs; a wrong tile from the same mechanism is possible, but we have not seen one.
• Most SM100 kernels with an epilogue-load warp add NumEpilogueLoadThreads when is_epi_load_needed: for example sm100_gemm_array_tma_warpspecialized.hpp (lines 692–711, in both branches), sm100_gemm_tma_warpspecialized.hpp (lines 511–515) and sm100_implicit_gemm_tma_warpspecialized.hpp (lines 492–496).

Evidence

• One B300 (compute capability 10.3), CUDA 13.0, CUTLASS 4.8.0 (commit 0b55a2f, whose C++ sources are identical to the v4.8.0 tag), example 92 grouped with our harness change (below), 16 groups of 7168×256×512, each run with a 60 s timeout. Times are wall-clock per run, including setup and verification.
  • Unmodified: beta = 0 passed 3 of 3 (about 1.3–1.6 s each); beta = 1 hung in 10 of 10.
  • With Count the epilogue-load warp in the SM100 mixed TMA/cp.async GEMM's scheduler pipeline #3781: beta = 0 passed 3 of 3; beta = 1 passed 10 of 10 (about 1.5–2.1 s each).

Suggested fix

See #3781: the same conditional addition, in both branches.

Scope

• Affected: kernels built from sm100_gemm_mixed_tma_cpasync_warpspecialized.hpp whenever the epilogue loads C or another input. Example 92 sets beta = 0, so as shipped it does not take this path.
• sm100_gemm_cpasync_warpspecialized.hpp has the same omission (lines 464–465). Its epilogue-load warp also calls fetch_next_work on every tile (lines 673–690), so the same over-arrival applies there when a TMA epilogue loads C. Its unit tests use the NoSmem epilogue, which loads nothing, so they do not exercise it. We have not run it, so Count the epilogue-load warp in the SM100 mixed TMA/cp.async GEMM's scheduler pipeline #3781 leaves it alone; we're happy to add it.
• The kernel file, the example and the epilogue files are byte-identical on main (91843b7), where Count the epilogue-load warp in the SM100 mixed TMA/cp.async GEMM's scheduler pipeline #3781 applies cleanly. We have not run main itself.

Steps/Code to reproduce bug

Example 92 has no beta option. Apply our harness change (in the details below) to example 92. It adds --alpha, --beta and --dbeta, accepts compute capability 10.3, runs only the first runner (mixed TMA + cp.async, 1SM) and compares with a relative tolerance, since with beta ≠ 0 the reference and the kernel may round differently. Then build with CUDA 13.0 and run:

nvcc -std=c++17 -O3 -DNDEBUG --expt-relaxed-constexpr -gencode arch=compute_100f,code=sm_100f \
  -Iinclude -Itools/util/include -Iexamples/common \
  examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu -o ex92
for g in $(seq 0 15); do echo "$g 7168x256x512"; done > bench16.txt
timeout 60 ./ex92 --m=7168 --n=256 --k=512 --groups=16 --benchmark=bench16.txt --iterations=0 --beta=1

Expected: the run finishes and prints Disposition: Passed. Actual: it hangs until the timeout (10 of 10 runs).

Expected behavior

The kernel completes and matches the reference for any beta.

Environment details (please complete the following information):

• Environment location: Bare-metal. One NVIDIA B300 (compute capability 10.3), CUDA 13.0, CUTLASS 4.8.0 (commit 0b55a2f); the relevant files are unchanged on main (91843b7).

--- a/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu	2026-10-24 14:43:49
+++ b/examples/92_blackwell_moe_gemm/92_blackwell_moe_gemm_grouped.cu	2026-10-24 14:44:08
@@ -82,6 +82,9 @@
   bool verification;
 
   int m, n, k, groups;
+  float alpha, beta;
+  int64_t dbeta;
+  bool all_configs;
 
   int iterations;
 
@@ -92,6 +95,7 @@
     error(false),
     verification(true),
     m(2048), n(2048), k(2048), groups(1),
+    alpha(1.0f), beta(0.0f), dbeta(0), all_configs(false),
     iterations(10)
   { }
 
@@ -110,6 +114,10 @@
     cmd.get_cmd_line_argument("groups", groups, 1);
     cmd.get_cmd_line_argument("iterations", iterations, 10);
     cmd.get_cmd_line_argument("benchmark", benchmark_path);
+    cmd.get_cmd_line_argument("alpha", alpha, 1.0f);
+    cmd.get_cmd_line_argument("beta", beta, 0.0f);
+    cmd.get_cmd_line_argument("dbeta", dbeta, int64_t(0));
+    all_configs = cmd.check_cmd_line_flag("all_configs");
 
 
     if (cmd.check_cmd_line_flag("no_verif")) {
@@ -307,7 +315,9 @@
 
       // Check if output from CUTLASS kernel and reference kernel are equal or not
       // assume all M == maxM
-      bool passed = cutlass::reference::device::BlockCompareEqual(block_ref_D.get() + size_t(1) * i * maxN * maxM, block_D.get() + size_t(1) * i * maxN * maxM, M * N);
+      // Relative compare: with beta != 0 the reference and the kernel may round differently.
+      bool passed = cutlass::reference::device::BlockCompareRelativelyEqual(block_ref_D.get() + size_t(1) * i * maxN * maxM, block_D.get() + size_t(1) * i * maxN * maxM, size_t(M) * N,
+                      ElementD(0.01f), ElementD(0.5f));
       if (!passed) {
         return false;
       }
@@ -431,8 +441,12 @@
 
     // arguments.scheduler.max_swizzle_size = options.swizzle;
     
-    arguments.epilogue.thread.alpha = 1.0f;
-    arguments.epilogue.thread.beta = 0.0f;
+    arguments.epilogue.thread.alpha = options.alpha;
+    arguments.epilogue.thread.beta = options.beta;
+    if (options.dbeta != 0) {
+      // Batch stride for beta with a null beta_ptr: the scalar beta is used for every group.
+      arguments.epilogue.thread.dBeta = {_0{}, _0{}, options.dbeta};
+    }
 
     Gemm gemm_op;
 
@@ -470,7 +484,7 @@
 
     if (options.verification) {
       // Verify that the result is correct
-      bool passed = verify(problem_size, 1.0f, 0.0f);
+      bool passed = verify(problem_size, options.alpha, options.beta);
 
       std::cout << "  Disposition: " << (passed ? "Passed" : "Failed") << std::endl;
 
@@ -523,8 +537,8 @@
     return 0;
   }
   
-  if (!(props.major == 10 && props.minor == 0)) {
-    std::cerr << "This example requires a GPU of NVIDIA's Blackwell architecture (compute capability 100)." << std::endl;
+  if (!(props.major == 10 && (props.minor == 0 || props.minor == 3))) {
+    std::cerr << "This example requires a GPU of compute capability 10.0 or 10.3." << std::endl;
     return 0;
   }
 
@@ -563,7 +577,10 @@
 
   std::cout << "Running kernel with mixed TMA+CPASYNC load, 1SM:" << std::endl;
   ExampleRunner runner_mixed_tma_cpasync;
-  runner_mixed_tma_cpasync.run(options, hw_info);
+  bool ok = runner_mixed_tma_cpasync.run(options, hw_info);
+  if (!options.all_configs) {
+    return ok ? 0 : 2;
+  }
 
   std::cout << "\n\n\nRunning kernel with mixed TMA+CPASYNC load and 1x1 cluster:" << std::endl;
   ExampleRunner<