From 269f6eebba00ab9d774a7eb0b1f1ae465813a3e9 Mon Sep 17 00:00:00 2001
From: dylanyunlon <dylanyunlon@users.noreply.github.com>
Date: Wed, 5 Aug 2026 08:12:12 +0000
Subject: [PATCH] =?UTF-8?q?[CCCL-PORT]=20summary=5Fstatistics.cu=20transfo?=
 =?UTF-8?q?rm=5Freduce=20pattern=20=E2=86=92=20online=20softmax=20design?=
 =?UTF-8?q?=20doc?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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).
---
 paged_attn.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/paged_attn.py b/paged_attn.py
index ce83aa95..83ef4b0c 100644
--- a/paged_attn.py
+++ b/paged_attn.py
@@ -188,6 +188,28 @@ class PagedAttention:
 
         return output
 
+    # ================================================================
+    # CCCL Design Pattern: summary_statistics.cu transform_reduce
+    #
+    # CCCL packs {n, min, max, mean, M2, M3, M4} into one struct and
+    # computes ALL statistics in a single pass via transform_reduce.
+    # The binary_op merges two partial results (Welford parallel algo).
+    #
+    # Our online softmax is the same pattern:
+    #   accumulator = {m (running max), l (running sum_exp), o (running output)}
+    #   unary_op: score_tile → {max(tile), sum(exp(tile-max)), exp(tile-max) @ V}
+    #   binary_op: merge two accumulators with correction factor
+    #
+    # Key insight: kv_heads are INDEPENDENT — no cross-head dependency.
+    # Current code already batches via [kv_h, gqa, q_len, tile_sz] tensor ops.
+    # The CCCL pattern validates this is optimal: one matmul per tile across
+    # all heads simultaneously, not per-head iteration.
+    #
+    # Future optimization: if we ever get Triton/CUDA access, the binary_op
+    # merge step ({m,l,o} update) could be fused with the matmul via a
+    # custom epilogue — this is what FlashAttention-2/3 does at the CUDA level.
+    # ================================================================
+
     # paged_attention_v1 on BI-V100 fails for long contexts.
     # Route on actual sequence length (seq_lens.max()), not the max_seq_len
     # parameter which is inflated to max_model_len in CUDA graph mode.