Prefetch.tensormap: elected lane vs full warp using CUTLASS/CuTe on SM90 and SM100

Accelerated Computing CUDA CUDA Programming and Performance
1

I would like to clarify the correctness and performance semantics of prefetch.tensormap on Hopper (SM90) and Blackwell (SM100). This is about prefetching the tensor-map descriptor, not cp.async.bulk.prefetch.tensor for tensor data.

NVIDIA CUTLASS provides these useful reference points (pinned to commit f614dc40e17fb3ddb1b7b474318b48a8d5d21a2c):

  • CuTe wrapper: cute::prefetch_tma_descriptor directly emits prefetch.tensormap without selecting a lane inside the wrapper.
  • SM100 GEMM caller: computes lane_predicate = cute::elect_one_sync() and guards descriptor prefetch with this predicate.
  • Collective comment: recommends issuing descriptor prefetch from a single thread for best performance.

The following device-side sketch isolates the two calling patterns using the official CuTe API. Elect=false is my comparison variant, not a claim that CUTLASS uses this pattern at the cited call site.

#include <cuda.h>
#include <cuda_runtime.h>
#include <cute/arch/cluster_sm90.hpp>
#include <cute/arch/copy_sm90_desc.hpp>

template <bool Elect>
__global__ void prefetch_descriptor(
    const __grid_constant__ CUtensorMap tensor_map) {
  if (threadIdx.x < 32) {
    if constexpr (Elect) {
      if (cute::elect_one_sync()) {
        cute::prefetch_tma_descriptor(&tensor_map);
      }
    } else {
      cute::prefetch_tma_descriptor(&tensor_map);
    }
  }
}

Assume a 1D launch with 128 threads per block, a valid tensor map encoded on the host, no descriptor updates, and a fully active first warp. Every participating lane therefore passes the same descriptor address. Host setup and actual TMA copies are omitted: this sketch expresses the calling-pattern question, not a standalone benchmark.

Questions:

  1. Are both variants supported for correctness? Is selecting one elected lane only a performance/code-generation recommendation for this particular instruction?
  2. For the full-warp, uniform-address variant, what determines the number of effective descriptor-prefetch requests? Are requests combined, repeated per lane, or handled by compiler-generated serialization? Is this specified or architecture/compiler dependent?
  3. Does the CUDA guidance about using an elected lane to avoid a compiler peeling loop for TMA copies also apply to descriptor prefetch.tensormap? Which SASS instructions or Nsight Compute counters would help distinguish the cases on SM90 and SM100?
  4. What is the sharing scope of the descriptor cache warmed by this instruction? Can a prefetch by one lane benefit subsequent TMA copies from other warps in the CTA/SM, and is any synchronization needed solely for this prefetch, separate from descriptor-update visibility and actual-copy completion?

PTX prefetch documentation

I have not measured a speed difference or compiled this sketch for a controlled SASS comparison. I am asking about the guaranteed semantics and the reason behind the official single-thread recommendation.

2

As far as I can tell the single thread recommendation is actually wrong and it’s (very slightly) better to use a single warp, at least based on the generated instructions. As often, it’s instructive to drop the code (in this case using inline PTX instead of CuTe) into Compiler Explorer and look at the generated SASS:

There is no peeling loop The SASS is basically the same between sm_90, sm_100, and sm_120. The elect variant uses a couple more SASS instructions to direct control flow but it’s actually unnecessary. UTMACCTL.PF is a uniform instruction, it doesn’t care about the warp active mask. And it’s executed once in both cases.

A prefetch executed multiple times, whether from different lanes within a warp (which, as shown, doesn’t happen) or from different warps (which would happen if you dropped the if (threadIdx.x < 32) condition) shouldn’t affect correctness – if data is already in cache it won’t be fetched again.

Re the cache scope, I don’t believe it’s specified and may change in future architectures. I’m sure it’s CTA at minimum, and possibly at the GPC level (given the multicast ability of TMA).
As I understand it, synchronisation for the purpose of descriptor update visibility may involve a full descriptor cache invalidation. In which case any prefetches before it would be pointless.