h3-blackwell-runtime/SAGE2_BLACKWELL_DESIGN.md

# 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:

1. Adversarial short-shape comparison against the Sage2 mainloop.
2. Exact real block-24 attention output at 37,810 tokens.
3. Mainloop median below 220 ms before block integration.
4. Blocks 0, 24, and 49 exact after integration.
5. 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/`.