69 lines
2.9 KiB
Markdown
69 lines
2.9 KiB
Markdown
# 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:
|
|
|
|
```text
|
|
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
|
|
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:
|
|
|
|
```bash
|
|
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.
|