h3-blackwell-runtime/research/sage2_temporal_pair

Sage2 SM89 Temporal Pair Experiment

This is an isolated experiment for SageAttention commit d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5. It does not alter H3 runtime dispatch or install an extension into the runtime image.

Candidate

The patch adds a default-disabled kernel template option and enables it only in sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu in the candidate checkout. The four existing warps retain their original query rows and private RS, RS_f8, RO, m, and d state.

For each KV tile, warp pair A and pair B alternate these phases:

A: QK + scale + online softmax + score conversion | B: previous-tile PV
B: QK + scale + online softmax + score conversion | A: current-tile PV

K and V remain single-buffered in their existing distinct shared-memory regions. All warps cooperatively stage each next tile. Full-CTA barriers occur before either buffer is overwritten and after each cp.async.wait_group<0>(). There is no score or accumulator transfer. Pair B carries one RS_f8 fragment between loop iterations; the prologue fills and the epilogue drains that one fragment.

This differs from early-K prefetch because the intended overlap is INT QK from one pair with FP8 PV from the other pair. It differs from independent softmax chain interleaving because each warp's exact per-tile numerical sequence remains QK -> update_mdo -> accumulate_d -> RS_32_to_8 -> PV.

Run On Spark

From the project checkout on Spark:

bash research/sage2_temporal_pair/run_spark_experiment.sh

The script creates detached baseline and candidate worktrees, checks and applies the patch, builds separate _qattn_sm89 extensions for sm_121a, and runs:

1. A compute-sanitizer quick check.
2. Byte parity over one-, two-, three-, and four-tile boundary shapes.
3. The real block-24 SHA and rotating timing harness.
4. The absolute <220 ms candidate gate.

Useful controls:

RUN_SANITIZER=0 bash research/sage2_temporal_pair/run_spark_experiment.sh
RUN_NCU=1 bash research/sage2_temporal_pair/run_spark_experiment.sh
SAGE2_TEMPORAL_RUN_ID=manual-01 bash research/sage2_temporal_pair/run_spark_experiment.sh
MAX_JOBS=4 bash research/sage2_temporal_pair/run_spark_experiment.sh

The default working root is /home/daniel/aeon-spark-test/h3/sage2-temporal-pair. JSON and optional NCU reports are written under /home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/benchmarks. Build logs, including ptxas -v register and spill reports, remain under the run-specific builds directory.

Stop if the build fails, sanitizer reports an error, any short shape differs, the real output SHA differs from 4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574, candidate p50 is at least 220 ms, shared memory exceeds the baseline 32768 bytes, or candidate register/spill growth erases scheduler eligibility. Do not integrate or replace the deployed extension from this experiment.