Implementing H100 TMA multicast with cuda::ptx:: functions but its slower than 8 independent TMA operations fetching same tile in cluster

Accelerated Computing CUDA CUDA Programming and Performance
1

First, I initialized a mapped distributed shared memory:

// initialization of distributed shared memory barrier (mapped from one CTA)
if (threadIdx.x == 0 && blockRankInCluster == 0) {
    cuda::ptx::mbarrier_init(bar, 2); // tried 1 instead of 2 too
    cuda::ptx::fence_mbarrier_init(cuda::ptx::sem_release, cuda::ptx::scope_cluster);
    cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_release, cuda::ptx::scope_cluster, cuda::ptx::space_cluster, reinterpret_cast<uint64_t*>(bar), 1); // tried without this too
}
cluster.sync();

Then initiated TMA using tensor map:

if (threadIdx.x == 0 && blockRankInCluster == 0) {
    const uint16_t ctaMask = 0b1; // tried 0b11 
    cuda::ptx::cp_async_bulk_tensor(cuda::ptx::space_cluster, cuda::ptx::space_global, reinterpret_cast<void*>(s_mem), tensorMap, coords, reinterpret_cast<uint64_t*>(bar), ctaMask);                
}

After other asynchronous work, wait like this:

// cuda::ptx:: functions don't return any state. How can I know which state to start? Is state a token like non-cluster TMA operations use?
if (threadIdx.x == 0 && blockRankInCluster == 0) {
    while (!cuda::ptx::mbarrier_try_wait(cuda::ptx::sem_acquire, cuda::ptx::scope_cluster, reinterpret_cast<_CUDA_VSTD::uint64_t*>(bar), state))     {
    __nanosleep(10);
    };
}

but it doesn’t get out of this while-loop.

I used ctaMask = 0b1 in a two-block cluster launch (static define with cluster_dims(2, 1, 1)).

State is incremented after wait but its never reached due to infinite loop.

What can cause a try-wait to never return true value?

If this works, I will convert it to a multiple-CTA wait version to make synchronization more efficient. Currently only trying to run it correctly and have no idea how TMA knows the destinations for other CTAs smem. Can it infer their destinations from just the calling CTA’s destination parameter?

Tensor map encoding is correct, tile data on gmem is aligned, and works for non-cluster based TMA operations, distributed shared memory barrier and all the data pointers are aligned to 32 and 128 respectively. No error is returned. Its just running inside the while loop.

Also if someone can discuss the details about how mbarrier works, how a cluster initiates TMA from single and multiple CTAs for all CTAs, I appreciate. For example, more CTAs in cluster would require at least 1 arrival per CTA right? Then maybe try-wait per CTA instead of cluster.sync(). I guess arrive-count expected is (1 from TMA) +(1 for each thread calling try-wait).

I’m trying to do this because generating a single TMA request from single CTA and then distributing from that CTA to other CTAs in cluster (through distributed-smem copy manually) has a very limited bandwidth like 2-3 TB/s only (H100). Compared to the TMA bandwidth, its too low. So a multicast would make it potentially 5-10 times faster (or up to the peak shared-memory write bandwidth which is used as destination).

2

I completed the multi-cast like the following part but now it is slower than a normal TMA operation (non-multicast):

init:

if (threadIdx.x == 0 && blockRankInCluster == 0) {
    cuda::ptx::mbarrier_init(bar, 1);
    cuda::ptx::fence_mbarrier_init(cuda::ptx::sem_release, cuda::ptx::scope_cluster);
}
cluster.sync();

copy:

if (threadIdx.x == 0 && blockRankInCluster == 0) {
    const int32_t coords[2] = { x, y };
    const uint16_t ctaMask = 0b11111111;
    cuda::ptx::mbarrier_arrive_expect_tx(cuda::ptx::sem_release, cuda::ptx::scope_cluster, cuda::ptx::space_cluster, reinterpret_cast<uint64_t*>(bar), size);
    cuda::ptx::cp_async_bulk_tensor(cuda::ptx::space_cluster, cuda::ptx::space_global, reinterpret_cast<void*>(s_mem), tensorMap, coords, reinterpret_cast<uint64_t*>(bar), ctaMask);
}

wait:

if (threadIdx.x == 0 && blockRankInCluster == 0) {
    while (!cuda::ptx::mbarrier_try_wait(cuda::ptx::sem_acquire, cuda::ptx::scope_cluster, reinterpret_cast<_CUDA_VSTD::uint64_t*>(bar), state)) {
        __nanosleep(10);
    };
}

For example, in a non-multicast TMA operation where each CTA loads exact same global tile in same kernel, it has 8TB/s bandwidth while this multicast version has only about 2TB/s bandwidth which is even lower than a normal TMA (single CTA load) + manual broadcast to other CTAs using distributed smem.

Even without cluster.sync(), it is slower than non-multicast version.

Why is loading same tile 8 times not slower than loading only once(multicast)?

I checked incoming data, its correct for all CTAs of cluster. Just slower than independent TMA operations per CTA.

3

On H100, TMA multicast isn’t automatically a performance win over independent TMA.
Multicast requires cluster-level coordination and mbarrier costs, which can outweigh the benefit if the transfer size or reuse pattern doesn’t amortize that overhead.

4

Thank you Linzheng,

I was using single-source TMA (1 CTA to other CTAs in cluster for whole tile of data). I guess just 1 CTA can’t initiate other blocks TMA hardware to do this collectively. And probably loses the benefits of their interconnect bandwidth within cluster.

I mean, even if single TMA hardware usage isn’t bottleneck, maybe (not using all of ) cluster-interconnect is. I had assumed that nvidia unified that part too, similar to how warp-shuffle and smem throughput are unified.

5

That’s a reasonable intuition. In practice, a single “producer CTA” can still become the bottleneck — multicast doesn’t increase the issue rate, and you still pay a single TMA issue plus the cluster distribution/sync cost.
When the per-tile transfer is small (or coordination dominates), it’s not unusual to see higher aggregate throughput from multiple independent TMA calls.

1 Like
6

I just want to point out what seems to me to be a bug/misunderstanding in the code. You wrap the mbarrier operations in if (threadIdx.x == 0 && blockRankInCluster == 0), but when using multicast, each CTA in the cluster should use its own local mbarrier (see reference to PTX documentation below), since both data and mbarrier signal are multicast. So the mbarrier_init and mbarrier_arrive_expect_tx calls should happen in all CTAs, only cp_async_bulk_tensor with multicast should be limited to blockRankInCluster == 0.

The optional modifier .multicast::cluster allows copying of data from global memory to shared memory of multiple CTAs in the cluster. Operand ctaMask specifies the destination CTAs in the cluster such that each bit position in the 16-bit ctaMask operand corresponds to the %ctaid of the destination CTA. The source data is multicast to the same CTA-relative offset as dstMem in the shared memory of each destination CTA. The mbarrier signal is also multicast to the same CTA-relative offset as mbar in the shared memory of the destination CTA.

1 Like
7

Yes. Later I converted cuda::device::barrier_arrive_tx (and barrier.wait with token) to be called from each cta. Using single CTA to wait and then cluster-wide sync to get other CTAs in sync was a bad idea. But it didn’t change performance much. Still requires interleaved multicast to use full bandwidth of interconnect or to reduce overall latency of cluster multicast.

8

It can be tricky to benchmark this correctly as well. Cluster multicast should help you get higher overall bandwidth if you are getting close to the L2/HBM bandwidth limit in non-multicast situations. I can confirm that, for large transfers, pipelined and measured over a significant amount of time, non-multicast cp_async_bulk tops out at ~38B/cycle/SM on 80GB/132SM variants of H100, i.e. just under 5120B per cycle GPU-wide from L2. (Note that the 94GB cards have more L2 cache and bandwidth). Whereas with 4x cluster multicast you can get ~70B/cycle/SM. So it’s no longer maxing out L2 bandwidth and is limited by intra-GPC bandwidth instead. I’ve not seen an improvement going from 4x to 8x. (Note that cluster size 4 or 8 can no longer use all SMs)

However, this is throughput/bandwidth, not latency. One thing that’s less intuitive is that multicast will have higher latency, especially average latency.
If you’re not mulitcasting, the different CTAs don’t necessarily have to be in sync with each other. So if they’re individually loading the same block of data at slightly different times, what will happen is that one of them will arrive at the load first, and incur the latency penalty of any L2 cache miss. But when the other CTAs come around to loading the same data it might already be in L2, so their latency, as observed locally, can be much lower. Whereas a single multicast load will get all the same cache misses as the first non-multicast load. It’s thus important to pipeline your loads sufficiently to cover the worst case latency.

1 Like
9

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.