79621cf8af
Previous _run_sdpa_fallback used Q-tiling but computed full attention weights
over the entire KV sequence per Q chunk:
attn_w = torch.softmax(Q_chunk @ K_full^T) → O(q_chunk × seq_len) memory
For seq_len=100K, kv_h=4, gqa=6, q_chunk=256:
[4, 6, 256, 100000] × 4B = 2.4 GB — causes OOM on BI-V100 (50GB/card, 4-way TP)
New version tiles BOTH Q and KV dimensions with online softmax:
For each Q chunk, iterate over KV tiles:
score = Q_chunk @ K_tile^T → O(q_chunk × kv_chunk) memory
{m, l, o} accumulator updated per tile (Flash Attention Algorithm 1)
Peak memory: [4, 6, 256, kv_chunk] × 4B where kv_chunk ≈ 8000 → ~48 MB
Architecture ported from CCCL source code:
- summary_statistics.cu: transform_reduce compound accumulator pattern
{n, min, max, mean, M2} maps to {m, l, o} online softmax state
- grid_even_share.cuh: adaptive tile sizing via _SCORE_BUDGET_BYTES
- agent_reduce.cuh: ConsumeFullTile vectorized load → GQA broadcast
- dispatch_reduce.cuh: two-path (single-tile vs multi-tile) dispatch
This is the same online softmax already used in paged_attn.py's
_forward_prefix_pytorch and _forward_decode_pytorch. Now xformers
fallback matches, giving consistent behavior across all attention paths.
Functional correctness: online softmax is mathematically equivalent to
torch.softmax — same output, different memory/compute schedule.
The {m, l, o} merge is the binary_op from CCCL's summary_stats_binary_op.