[BUG] Unexpected Slower Performance in C++ CuTe (77_blackwell_fmha) vs Python CuTe DSL (fmha.py) on Jetson Thor

We observed significantly worse performance in the CUTLASS CuTe 77_blackwell_fmha example compared to the CuTe DSL examples/python/CuTeDSL/blackwell/fmha.py implementation on Jetson Thor.

Test Commands:
77_blackwell_fmha:

./77_blackwell_fmha --b=1 --h=16 --q=2225 --k=2225 --d=128 --mask=no --verbose --warmup_iterations=2 --iterations=10

fmha.py:

python3 fmha.py --iterations 10 --warmup_iterations 2 --skip_ref_check --mma_tiler_mn 128,128 --q_shape 1,2225,16,128 --k_shape 1,2225,16,128

Performance Results:
77_blackwell_fmha: 1,110,707 cycles
fmha.py: 922,724 cycles

Key Observations:

1. Local Memory Spilling:
ncu profiling reveals substantially higher Local Memory Spilling Requests in 77_blackwell_fmha compared to fmha.py.

2. Implementation Differences in softmax_step:
The softmax computation differs between implementations:

• fmha.py uses 4 fragments:

frg_cnt = 4
frg_tile = cute.size(tTMEM_LOADrS) // frg_cnt
tTMEM_LOADrS_frg = cute.logical_divide(tTMEM_LOADrS, cute.make_layout(frg_tile))
tTMEM_STORErS_x4_e_frg = cute.logical_divide(
    tTMEM_STORErS_x4_e, cute.make_layout(frg_tile)
)
for j in range(frg_cnt):
    for k in cutlass.range(
        cute.size(tTMEM_LOADrS_frg, mode=[0]), vectorize=True
    ):
        tTMEM_LOADrS_frg[k, j] = (
            tTMEM_LOADrS_frg[k, j] * scale + minus_row_max_scale
        )
        tTMEM_LOADrS_frg[k, j] = cute.math.exp2(
            tTMEM_LOADrS_frg[k, j], fastmath=True
        )

    s_vec = tTMEM_LOADrS_frg[None, j].load()
    tTMEM_STORErS_x4_e_frg[None, j].store(s_vec.to(self.q_dtype))

• 77_blackwell_fmha processes data without fragment:

CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tTMEM_LOADrS); i += 2) {
  float2 in = make_float2(
    tTMEM_LOADrS(i + 0),
    tTMEM_LOADrS(i + 1)
  );
  float2 out;
  cute::fma(out, scale_fp32x2, in, minus_row_max_scale_fp32x2);
  tTMEM_LOADrS(i + 0) = out.x;
  tTMEM_LOADrS(i + 1) = out.y;

  tTMEM_LOADrS(i+0) = ::exp2f(tTMEM_LOADrS(i+0));
  tTMEM_LOADrS(i+1) = ::exp2f(tTMEM_LOADrS(i+1));

  Array<ElementQK, kConversionsPerStep> in_conv;
  CUTLASS_PRAGMA_UNROLL
  for (int j = 0; j < kConversionsPerStep; j++) {
    in_conv[j] = tTMEM_LOADrS(i + j);
  }
  tTMEM_STORErS_x4_e[i / kConversionsPerStep] = convert(in_conv);


  if (i == size(tTMEM_LOADrS) - kReleasePipeCount) {
    order_s.arrive();
  }

  // this prevents register spills in fp16
  if constexpr (size<2>(tTMEM_STORErS_x4) == _2{}) {
    if (i == size(tTMEM_LOADrS) - 6) {
      copy(tiled_tmem_store, tTMEM_STORErS_x4(_, _, 0), tTMEM_STOREtS_x4(_, _, 0));
    }
  }
}

3. Failed Optimization Attempts:

• Porting fmha.py's fragment strategy to 77_blackwell_fmha did not reduce spilling or improve performance.

• Removing tiling from fmha.py increased spilling but still outperformed 77_blackwell_fmha (spilling remained lower than C++ version).

Despite similar algorithmic structures, why does 77_blackwell_fmha exhibit worse performance and higher spilling?

Start by reproducing both commands on Jetson Thor and profiling them with ncu. Compare the 77_blackwell_fmha entry point with examples/python/CuTeDSL/blackwell/fmha.py, focusing on softmax_step, fragment handling, and local-memory spilling. Done means identifying the cause of the performance gap and documenting a validated remediation or focused next step.