2.9 KiB
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:
- A compute-sanitizer quick check.
- Byte parity over one-, two-, three-, and four-tile boundary shapes.
- The real block-24 SHA and rotating timing harness.
- The absolute
<220 mscandidate 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.