[BUG] SM90 TMA GEMM mainloops rely on elect_one_sync() electing lane 0 (TMA issuer vs. PipelineTmaAsync is_leader) #3722

Which component has the problem?

CUTLASS C++

Bug Report

Describe the bug

Summary. Most SM90 TMA GEMM mainloop collectives work only if elect.sync elects lane 0, because two threads are chosen by different rules: the TMA issuer is the elected lane, while the producer's arrival comes from warp-group thread 0. PTX does not say which lane elect.sync elects, and if another lane were elected the kernel would hang. We suggest deriving the leader from the election in the mainloop warp, as in-tree code already does elsewhere.

Paths and lines are at main 0b55a2f.

Evidence

TMA issuer. These collectives call producer_acquire and issue the TMA copies only inside if (lane_predicate), where lane_predicate = cute::elect_one_sync(), that is, elect.sync with a full mask:

The non-warp-specialized sm90_mma_tma_gmma_ss.hpp (lines 365 and 504) and the SM90 implicit-GEMM conv collective (conv/collective/sm90_implicit_gemm_gmma_ss_warpspecialized.hpp:558–559, 612) have the same pattern.

Producer arrival. PipelineTmaAsync::producer_acquire performs arrive_and_expect_tx only when params_.is_leader is set (include/cutlass/pipeline/sm90_pipeline.hpp:533–535, and 555–557 in the token overload). This arrival is the stage's only producer arrival. The kernels set is_leader = warp_group_thread_idx == 0:

• sm90_gemm_tma_warpspecialized.hpp:324;
• sm90_gemm_tma_warpspecialized_cooperative.hpp:471;
• sm90_gemm_tma_warpspecialized_pingpong.hpp:463;
• the array cooperative and pingpong kernels, lines 522 and 531;
• sm90_mma_tma_gmma_ss.hpp:334.

The mainloop warp is warp 0 of the producer warp group, so warp-group thread 0 is its lane 0.

Cause

PTX ISA 9.4 (§9.7.15.15, elect.sync) promises only that the election is deterministic, with the same leader for the same membermask every time. It does not say which lane. If another lane were elected:

• that lane would call producer_acquire without being the leader, so nothing arrives on the full barrier;
• lane 0 would not enter the if at all;
• the TMA's complete-tx would only drive the transaction count negative.

The phase would never complete. Consumers would block in consumer_wait, then the producer in producer_acquire once every stage is used.

Suggested fix

Precedents for deriving the leader from the election:

• on the same PipelineTmaAsync, examples/88_hopper_fmha/kernel/fmha_kernel_tma_warpspecialized.hpp:205, 210 and sm120_gemm_tma_warpspecialized_cooperative_asymmetric_dma.hpp:443, 448;
• in the SM100 kernels, e.g. sm100_gemm_tma_warpspecialized.hpp:465.

For the SM90 kernels, that would be:

mainloop_pipeline_params.is_leader =
    warp_group_role == WarpGroupRole::Producer &&
    producer_warp_role == ProducerWarpRole::Mainloop && lane_predicate;

Restrict it to the mainloop warp. The MainloopAux warp in the blockwise kernels also has the Producer role and calls producer_acquire from all lanes. Giving it a leader would add a second arrive_and_expect_tx per stage.

This works because the kernel's elect_one_sync() and the collective's later one use the same full mask, and PTX guarantees the same leader for the same mask. Two more changes go with it:

• sm90_mma_tma_gmma_ss.hpp would need lane_predicate computed before is_leader is set;
• the debug-build brkpt checks in producer_acquire (sm90_pipeline.hpp:542–544 and 564–566), which fire when the leader is not lane 0, would need to go. They already assume lane 0 for the FMHA and SM120 kernels above.

The alternative is to gate the TMA issue in each affected collective on lane 0 instead of elect_one_sync().

Scope

• Unaffected: some collectives call producer_acquire from all 32 lanes and issue TMA only from the elected lane, so lane 0 always arrives:

  • the mixed-input collectives (sm90_mma_tma_gmma_rs_warpspecialized_mixed_input.hpp:684/694, and the array variant at 774/784);
  • the non-array FP8 blockwise collective (577/587–588).

  The epilogue's TMA loads are also unaffected, because one elected lane both posts the expect-tx and issues the copy.

• Current hardware: the kernels run today, so current hardware evidently elects lane 0 here. The code also states that intent:

  • sm90_mma_tma_gmma_ss.hpp:356 says "Set predicate for the lowest lane_id in the warp";
  • sm90_pipeline.hpp:286 says "We have exactly one thread elected in the Producer as the 'leader'";
  • CuTe's fallback without elect.sync returns (threadIdx.x % 32) == 0 (include/cute/arch/cluster_sm90.hpp:197).

Why there is no PR

The fix can go in the kernels (is_leader from the election) or in each affected collective (TMA issue gated on lane 0), so we'd rather follow your preference than guess. Tell us which and we'll write the PR.

Steps/Code to reproduce bug

Found by reading the source against the PTX ISA. There is no runtime reproduction.

Expected behavior

The lane that issues the TMA loads is the lane that arrives with the expected bytes.

Environment details

Source analysis; no runtime environment involved. Hopper (sm_90a).