269f6eebba
Source: cccl_upstream/thrust/examples/summary_statistics.cu
summary_statistics.cu demonstrates CCCL's core pattern: pack multiple
accumulation values into a single struct {n,min,max,mean,M2,M3,M4},
compute everything in ONE pass via thrust::transform_reduce with a
Welford parallel binary_op that merges two partial results.
Our Flash Attention online softmax is structurally identical:
accumulator = {m (running_max), l (running_sum_exp), o (running_output)}
unary_op: score_tile → {max, sum_exp, weighted_V}
binary_op: merge with correction factor exp(old_max - new_max)
Key validation: kv_heads are independent (no cross-head dependency),
so batching all heads in [kv_h, gqa, q_len, tile_sz] tensor ops is
the correct PyTorch equivalent of CCCL's transform_reduce approach.
This matches how dispatch_reduce.cuh handles multi-block results:
StableReductionOrder=false → atomic merge (one kernel)
StableReductionOrder=true → write partials, reduce in 2nd kernel
Our Python accumulator is the 'true' path (sequential merge per tile).