[BUG] SM90 cp.async + GMMA collectives: no fence.proxy.async between cp.async operand writes and wgmma reads #3721

Which component has the problem?

CUTLASS C++

Bug Report

Describe the bug

Summary. The SM90 cp.async + GMMA warp-specialized collective writes its shared-memory operands through the generic proxy and has wgmma.mma_async read them through the async proxy, with no fence.proxy.async on the path. Without the fence, PTX allows the wgmma to read operand bytes not yet visible to the async proxy, giving a silently wrong GEMM result; we have not run this on hardware, and current hardware may well make completed cp.async writes visible to the async proxy, though PTX does not state that. The fence can go on the producer side or the consumer side; see Options.

Cause

The collective include/cutlass/gemm/collective/sm90_mma_multistage_gmma_ss_warpspecialized.hpp (main 0b55a2f; unchanged since v4.4.0):

• fills the shared-memory operands with non-bulk cp.async, and with clear() for residue K-tiles (lines 274 and 284);
• publishes them with producer_commit(..., cutlass::arch::cpasync_barrier_arrive) (291, 310);
• then the consumer does consumer_wait (411, 436, 460), warpgroup_arrive(), and wgmma.mma_async, which reads those operands.

PTX ISA 9.4, §9.7.17.7.1 (the wgmma.fence page): "An async proxy fence must be used to establish an ordering between prior writes to shared memory matrices and subsequent reads of the same matrices in a wgmma.mma_async instruction." Several facts mean nothing else on the path supplies that fence:

• Non-bulk cp.async writes are generic-proxy operations (§9.7.10.28.3.1).
• wgmma reads its shared operands through the async proxy (§9.7.17.4).
• warpgroup_arrive() (wgmma.fence) orders register accesses.
• The implicit completion fence of §9.7.10.28.2 covers cp{.reduce}.async.bulk only.
• The cp.async.mbarrier.arrive arrive-on is ordered only with the thread's prior cp.async (§8.9.1.1), so it does not cover the clear() stores either.

Options

• Producer side: each producer thread waits for its own cp.async to complete, executes cutlass::arch::fence_view_async_shared(), then arrives on the barrier. The fence is executed by the writing thread, as in CUTLASS's own sm90_wgmma_transpose.hpp (lines 233–243) and the TMA epilogue (sm90_epilogue_tma_warpspecialized.hpp:736–737). It gives up the cp.async.mbarrier.arrive overlap; we have not measured the cost.
• Consumer side: one fence_view_async_shared() after each consumer_wait, before warpgroup_arrive(). This is cheaper, but it relies on a fence by a thread that only observed the writes. PTX says bi-directional proxy fences "take effect within a single thread" and that their effect "composes with other forms of synchronization according to the rules of the Memory Consistency Model" (§9.7.15.4); it does not state that this covers other threads' writes.

Scope

• Not observed on hardware: we have not run this on hardware, and current hardware may well make completed cp.async writes visible to the async proxy. PTX does not state that, though.
• For comparison, the Triton 3.5.0 and 3.7.0 PTX we compiled for an ordinary tl.dot matmul also had no fence between cp.async and wgmma; Triton 3.8.0 adds one.

Why there is no PR

Which placement is right is your call, given the producer side's cost and the consumer side's reliance on unstated composition. Tell us which you want and we'll write the PR.

Steps/Code to reproduce bug

We found this by checking compiled PTX against the PTX memory model with our own PTX-level checker. This instance (fp16, 64×64×64 tile, K = 320, one 384-thread CTA) was compiled with nvcc 12.6 for sm_90a:

#include "cutlass/cutlass.h"
#include "cutlass/gemm/collective/collective_builder.hpp"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/gemm/kernel/gemm_universal.hpp"
using namespace cute;
using BH12Tile = Shape<_64,_64,_64>;
using Cluster = Shape<_1,_1,_1>;
using H = cutlass::half_t;
using Main = cutlass::gemm::collective::CollectiveBuilder<
  cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp,
  H,cutlass::layout::RowMajor,8,H,cutlass::layout::ColumnMajor,8,
  float,BH12Tile,Cluster,cutlass::gemm::collective::StageCount<2>,
  cutlass::gemm::KernelCpAsyncWarpSpecialized>::CollectiveOp;
using Epi = cutlass::epilogue::collective::CollectiveBuilder<
  cutlass::arch::Sm90,cutlass::arch::OpClassTensorOp,BH12Tile,Cluster,
  cutlass::epilogue::collective::EpilogueTileAuto,float,float,
  void,cutlass::layout::RowMajor,8,H,cutlass::layout::RowMajor,8,
  cutlass::epilogue::NoSmemWarpSpecialized>::CollectiveOp;
using Kernel = cutlass::gemm::kernel::GemmUniversal<Shape<int,int,int,int>,Main,Epi>;
extern "C" __global__ __launch_bounds__(Kernel::MaxThreadsPerBlock)
void bh12_cpasync_gmma(H const* A,H const* B,H* D) {
  __shared__ __align__(128) char smem[Kernel::SharedStorageSize];
  Kernel::Params p{};
  p.mode = cutlass::gemm::GemmUniversalMode::kGemm;
  p.problem_shape = make_shape(64,64,320,1);
  p.mainloop.ptr_A=A; p.mainloop.ptr_B=B;
  p.mainloop.dA = make_stride(int64_t(320),_1{},int64_t(0));
  p.mainloop.dB = make_stride(int64_t(320),_1{},int64_t(0));
  p.epilogue.thread = typename Epi::ThreadEpilogueOp::Params(1.0f);
  p.epilogue.ptr_D=D;
  p.epilogue.dD=make_stride(int64_t(64),_1{},int64_t(0));
  Kernel{}(p,smem);
}

Its PTX has 8 cp.async and 8 wgmma.mma_async instructions and no fence.proxy (grep -c fence.proxy gives 0). Our checker finds the 16-byte cp.async that writes shared bytes [0, 16) unordered with the warpgroup's wgmma read of those bytes (it attributes the read to thread 256). The clear() path is from source reading; this instance has no residue.

Expected behavior

A fence.proxy.async.shared::cta orders the operand writes before the wgmma reads.

Environment details

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

Additional context

The _rs_ variant (sm90_mma_multistage_gmma_rs_warpspecialized.hpp) has the same pattern for the B operand, by source reading. There, A goes through ldmatrix into registers, and its TransposeB path already fences.