Sage2 Blackwell Mainloop

Decision

The adjacent NVFP4 producer fusions are complete. The next exact kernel experiment is now a retune of SageAttention 2.2.0's existing classic SM89 mainloop. Do not replace the algorithm with UMMA, FP8 Q/K, or FP4 attention.

The canonical GB10 profile is dominated by the attention kernel:

Complete block 24: 468.22 ms median.
Sage2 NHD path: 260.02 ms median.
SM89 attention mainloop: 238.81 ms median.
Q/K quantization: 7.68 ms median.
V quantization: 10.65 ms median.
Remaining framework overhead: 2.87 ms estimated.

The fresh real block-24 decomposition supersedes those internal timings:

Fused mainloop: 237.09 ms uninstrumented median.
K mean/smoothing: 2.37 ms median.
Q quantization: 3.80 ms median.
K subtract-mean quantization: 3.83 ms median.
V transpose/pad/permute: 5.08 ms median.
V scale/FP8 quantization: 5.53 ms median.
Manual decomposition versus public Sage2: byte-exact.

Q/K RMSNorm plus RoPE is 12.20 ms and is not the first target. NVFP4 ring dispatch is also excluded because it failed the controlled block gate.

Existing Contract

For the 37,810-token H3 shape, preserve these Sage2 inputs and outputs:

Batch 1, 56 heads, head dimension 128, non-causal attention.
Projection-strided BF16 NHD Q/K/V input.
Per-warp INT8 Q with 128-row blocks and 32-row warp groups.
Per-block INT8 K with 64-row blocks.
FP8 V and the existing per-channel scale layout.
Sage2's fp32+fp16 instantaneous-buffer accumulation order.
Contiguous BF16 NHD output with exact equality to Sage2 2.2.0.

The current mainloop is sageattention_sm89::qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf. It receives Q/K INT8 tensors, FP8 V, Q/K/V scales, and emits the final BF16 attention output.

Rejected Shortcut

SageAttention's Hopper SM90 WGMMA source was compiled experimentally for sm_121a. CUDA 13 ptxas rejects wgmma.fence, wgmma.mma_async, wgmma.commit_group, and wgmma.wait_group for SM121, including both INT8 and FP8 forms. A Hopper binary or source retarget is therefore not viable.

CUTLASS 4.6 SM120/121 exposes tcgen05/UMMA for F8/F6/F4 but not INT8. It therefore cannot preserve Sage2's INT8 QK operation. Changing Q/K to FP8 or FP4 would define a different attention algorithm and require a quality gate rather than exact parity.

Scheduler Diagnosis

The full-counter capture reports 255 registers/thread and 32 KiB dynamic shared memory for each 128-thread CTA. Registers limit residency to two CTAs, producing 16.83% achieved occupancy and only 0.46 eligible warps per scheduler. No warp is eligible in 63.53% of scheduler cycles. Fixed-latency dependency stalls are 2.01 cycles and math-pipe throttle is 1.24 cycles out of the 5.48 cycles between issued instructions.

The fused phase attribution is counter-based rather than separate wall time:

INT8 QK uses 37.77% of elapsed tensor-pipe capacity.
FP8 PV uses 37.77% of elapsed tensor-pipe capacity.
Scalar FMA and ALU pipelines use 15.01% and 8.81%, covering scale, online-softmax, conversion, normalization, and output work.
Combined tensor activity is 75.54%; the INT and FP sub-pipelines alternate, so neither individually exceeds 37.77%.
Memory throughput is 31.61%, L2 hit rate is 98.84%, and the warm off-chip request proxy is only 1.85 GB. This is not a DRAM bottleneck.
Shared accesses create 626,970,624 excessive wavefronts, 7.68% of all shared wavefronts.
Real Q and KV tail CTAs each cost less than 1 ms.

Completed P0

The source-identical classic mma.sync retune was completed without an accepted variant. The exact three-CTA threshold is 168 registers/thread; 170 remains at two CTAs. The 168-register cap causes 4.95 billion local spill requests and worsens no-eligible cycles from 63.13% to 78.79%. Narrowed scopes reduce static spills from 44/44 to 12/12 bytes but improve real-input latency by only 0.06%. In-place score reuse, early K prefetch, and independent softmax-chain interleaving are also neutral or slower.

The shared excess is produced by repeated V-staging LDGSTS.128 instructions. A 128-byte padded V layout leaves all 626,970,624 excessive wavefronts unchanged and provides no latency gain. This counter is therefore not an actionable layout target for the exact kernel.

The P0 gate was:

Adversarial short-shape comparison against the Sage2 mainloop.
Exact real block-24 attention output at 37,810 tokens.
Mainloop median below 220 ms before block integration.
Blocks 0, 24, and 49 exact after integration.
Alternating block benchmark must improve all three blocks before trajectory validation.

Use the Blackwell SageAttention3/CUTLASS code only as a scheduler reference. Its FP4 algorithm and output are not a correctness replacement for Sage2.

No variant exceeded the 3% mainloop gate, so block integration and deployment were intentionally skipped. Detailed evidence is in benchmarks/gb10-sage2-p0-register-scheduler-analysis.json.

Completed P1 And P2

P1 fused strided-NHD Q/K RMSNorm, split-half RoPE, and Q INT8 quantization. It was byte-exact, but complete-block gains were only 0.53-0.86%, below the 1% gate, so its runtime branch was removed.

P2 replaced Sage2's separate BF16 V transpose and FP8 quantization with an owned three-stage direct-NHD producer. It emits byte-identical FP8 V and FP32 scales for 13 boundary lengths. Canonical V preparation improves from 10.56 ms to 6.39 ms (39.48%), but the 4.17 ms saving is only 0.91% of the complete block and misses the 6.0 ms isolated go threshold. No block integration or trajectory validation was run. The validator-only implementation remains useful for a future owned attention backend; production Sage2 remains unchanged.

These results close the remaining preparation-only Sage2 opportunities. The next Sage-class optimization with plausible multi-percent block impact is a new warp-specialized QK/PV mainloop or a quality-gated Blackwell-native attention algorithm, not another launch-boundary fusion.

Completed P3

P3 tested the smallest exact temporal QK/PV specialization. Warp pairs A and B alternated INT8 QK plus online softmax against the other pair's prior-tile FP8 PV. Every warp retained its original query rows and private RS, RS_f8, RO, m, and d; no numerical state crossed warp boundaries. The candidate is compute-sanitizer clean and byte-exact across 13 short boundary cases and the real block-24 output SHA.

The schedule does not produce a material latency gain. In a rotating 50-sample comparison, baseline median is 245.44 ms and the candidate is 245.20 ms, a 0.10% improvement. It remains far above the <220 ms gate. Ptxas reports 254 registers/thread and 32/24-byte static store/load spills for the candidate, versus 255 registers and 24/24-byte spills for the clean baseline build. NCU and block integration were skipped because the latency gate failed.

This closes temporal warp separation as an exact Sage2 optimization on the current SM89 instruction stream. A multi-percent attention gain now requires a substantially new mainloop or a quality-gated Blackwell-native attention algorithm rather than another exact schedule rearrangement.

Successor Research

Vortex Exact Attention supersedes further source-identical Sage2 retuning. Its Phase 0 contract imports this document and retained P0-P3 evidence without claiming unrecoverable P0 source. Phase 1 evaluates only architectures that move or partition RS/RS_f8, m/d, and RO ownership. The recommended VEA-B design uses separate QK/softmax and PV/output warps around the still-supported exact INT8 mma.sync arithmetic on SM121. No kernel or production dispatch exists; projected latency ranges are not achieved results. See research/vortex_exact_attention/.

7.7 KiB Raw Blame History