[Aladoro]

The old reduction logic seemed to have mistakenly called cutlass.max rather than cute.arch.fmax during one of the in-place reductions:

...
 # Linear reduction is faster here, as well
  for j in cutlass.range_constexpr(cute.size(acc_qk_mn, mode=[1])):
      s_max[i] = cute.arch.fmax(s_max[i], acc_qk_mn[i, j])
      # old code: s_max[i] = cutlass.max(s_max[i], acc_qk_mn[i, j])
...

This seems to have been accidental (the kernel previously already uses the correct cute.arch.fmax for the first iteration), and seems very costly, emitting several operations rather than just nvmm fmax:

setp.le.f32
setp.nan.f32
selp.f32
selp.f32

Validation

Environment: NVIDIA H100 SxM, CUDA 12.9

Before the fix (default shapes/dtypes): 3884.84 us, 566.05 TFLOP/s
After: 3440.17 us, 639.22 TFLOP/s

[github-actions]

This PR has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this PR if it is no longer required. Otherwise, please respond with a comment indicating any updates. This PR will be labeled inactive-90d if there is no activity in the next 60 days.

[0z5a]

Measured this change on GH200 (fp16 in/out, fp32 accumulators, cold L2, 10 warmup + 50 timed iterations per measurement, 5 alternating rounds per variant, median shown). Both variants pass the same reference check (85/85 cases of the new FMHA test file in #3770, covering d32–d256, causal/sliding-window masks, persistent and non-persistent, GQA, bf16 and fp8), so the change is performance-only, as described.

q_shape (B, S, H, D)	mma_tiler_mn	persistent	mask	before: µs / TFLOPS	after: µs / TFLOPS	speedup
(4, 4096, 16, 128)	(64, 128)	yes	none	1028.6 / 534	906.5 / 606	1.135x
(4, 4096, 16, 128)	(64, 128)	yes	causal	622.9 / 441	553.5 / 497	1.125x
(4, 4096, 16, 128)	(64, 128)	no	none	1021.6 / 538	894.8 / 614	1.142x
(4, 4096, 32, 64)	(64, 128)	yes	none	1439.4 / 382	1171.1 / 469	1.229x
(4, 4096, 32, 64)	(64, 128)	yes	causal	1014.5 / 271	807.6 / 340	1.256x
(2, 8192, 8, 256)	(64, 64)	no	none	1694.0 / 649	1577.0 / 697	1.074x

Run-to-run spread (max − min) within each variant was at most 0.6%.

GPU: GH200 (cc 9.0) | driver 595.58 | CUDA 13.2 | nvidia-cutlass-dsl 4.8.1 | torch 2.14.1+cu130 | CUTLASS 91843b71 (+ this PR's one-line diff)