From 67e05f517c97308610fa09cbc951066505ff00c2 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 26 Aug 2026 15:51:39 +0700 Subject: [PATCH] Validate aligned Vortex attention prototype --- research/EXPERIMENT_REGISTRY.md | 2 + research/experiment_registry.json | 56 ++ research/vortex_exact_attention/DESIGN.md | 11 +- .../vortex_exact_attention/EXPERIMENT_LOG.md | 32 + .../PHASE2B_ALIGNED_REPORT.md | 94 ++ research/vortex_exact_attention/README.md | 16 +- .../vortex_exact_attention/kernels/README.md | 10 +- .../kernels/vea_b_numeric.cpp | 23 + .../kernels/vea_b_numeric.cu | 851 ++++++++++++++++++ .../tools/capture_phase2b_fixtures.py | 133 +++ .../tools/run_phase2b.py | 267 ++++++ 11 files changed, 1483 insertions(+), 12 deletions(-) create mode 100644 research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md create mode 100644 research/vortex_exact_attention/kernels/vea_b_numeric.cpp create mode 100644 research/vortex_exact_attention/kernels/vea_b_numeric.cu create mode 100644 research/vortex_exact_attention/tools/capture_phase2b_fixtures.py create mode 100644 research/vortex_exact_attention/tools/run_phase2b.py diff --git a/research/EXPERIMENT_REGISTRY.md b/research/EXPERIMENT_REGISTRY.md index 11455f0..7eb2242 100644 --- a/research/EXPERIMENT_REGISTRY.md +++ b/research/EXPERIMENT_REGISTRY.md @@ -97,6 +97,8 @@ missing rather than inferred. | `vortex-exact-phase2a-invalid-slot-probe` | VEA-B invalid initial slot-indexing probe | `incomplete` | Deterministic repetition exposed K/V overwrite from incorrect two-slot indexing; its payload and 16.631 ms conclusions are invalid. | Initial durable probe report; experiment log | None; superseded by corrected probe | | `vortex-exact-phase2a-inline-named-barrier` | VEA-B inline named-barrier primitive | `architecture_rejected` | Corrected timing clears budget, but racecheck reports five hazards; do not use this primitive in the prototype. | Phase 2A sanitizer log and report | None required unless sanitizer semantics change | | `vortex-exact-phase2a-capability-decision` | VEA-B Phase 2A capability decision | `research_retained` | Mbarrier handoff passes resource, residency, deterministic, sanitizer, overlap, fixture, and synchronization gates. No attention speedup is claimed. | `research/vortex_exact_attention/PHASE2A_CAPABILITY_REPORT.md`; architecture model; durable probe/NCU/sanitizer/fixture artifacts | Executable exact attention parity and latency | +| `vortex-exact-phase2b-spill-variants` | VEA-B Phase 2B rejected spill variants | `architecture_rejected` | Register-lifetime, shared-RO, and setmaxnreg variants retained spills, changed FP32 checkpoints, or failed launch resources. | Phase 2B smoke/build artifacts; aligned report | None; superseded by selected zero-local variant | +| `vortex-exact-phase2b-aligned-prototype` | VEA-B Phase 2B fixed aligned numerical prototype | `research_retained` | Fixed `Q128 x KV192 x D128` kernel is checkpoint exact, zero-local, sanitizer-clean, overlapping, repeatable, and stably timed. No complete-H3 speedup is claimed. | `research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md`; authoritative JSON; NCU and sanitizer artifacts | Ragged tails, canonical blocks/trajectories, resident-service and production gates | ## Interpretation Notes diff --git a/research/experiment_registry.json b/research/experiment_registry.json index 4b71d03..752be00 100644 --- a/research/experiment_registry.json +++ b/research/experiment_registry.json @@ -938,6 +938,62 @@ "production_behavior": "No attention kernel and no dispatch; production Sage2 remains unchanged.", "source_recovery": "Capability source, provenance, commands, corrected report, ptxas, sanitizer, NCU, and fixture artifacts are retained." }, + { + "id": "vortex-exact-phase2b-spill-variants", + "name": "VEA-B Phase 2B rejected spill variants", + "family": "vortex-exact-attention", + "status": "architecture_rejected", + "hypothesis": "Compiler lifetime reduction or dynamic register redistribution can fit the exact ten-warp aligned kernel under the GB10 168-register launch ceiling without explicit retirement scheduling.", + "implementation_strategy": "Evaluate sequential PV fragments, union role state, role-separated loops, shared persistent RO, inner-pair fences, and setmaxnreg redistribution while requiring every numerical checkpoint and resource gate.", + "source_locations": ["research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md", "research/vortex_exact_attention/EXPERIMENT_LOG.md", "research/vortex_exact_attention/kernels/vea_b_numeric.cu"], + "active_source_location": null, + "commit_hash": null, + "benchmark_artifacts": [ + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/smoke-union.json", "exists": true}, + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/smoke-role-loops.json", "exists": true}, + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/smoke-shared-ro0.json", "exists": true} + ], + "profiler_artifacts": [], + "environment": {"gpu": "NVIDIA GB10", "cuda": "13.0", "driver": "580.173.02", "pytorch": "2.9.1+cu130", "triton": null, "container_image": "sha256:6d880d628334c981c3d155bf5244e65e26e22cc9273c80145f646eee3c3698c2", "commit_context": "isolated Phase 2B fixed-shape spill screening"}, + "metrics": {"initial_local_bytes_per_thread": 352, "union_local_bytes_per_thread": 208, "role_loop_local_bytes_per_thread": 40, "setmaxnreg_launch": "too many resources requested", "shared_ro_first_divergent_epoch": 1}, + "correctness_evidence": ["Register-only lifetime variants remained checkpoint exact but retained local memory.", "The shared persistent-RO variant retained final BF16 bytes but diverged at the epoch-1 FP32 RO checkpoint.", "The 200-register setmaxnreg variant was rejected by the runtime launch resource check."], + "decision_rationale": "Reject these variants. The selected retirement-slot design is exact, sanitizer-clean, and has zero local memory under the actual 168-register ceiling.", + "reproducer_commands": [], + "timestamp": "2026-08-26", + "evidence_missing": [], + "production_behavior": "Never integrated; all variants are isolated research artifacts.", + "source_recovery": "Durable smoke JSON/build logs and the Phase 2B report retain the failed-variant evidence; only the selected source remains active." + }, + { + "id": "vortex-exact-phase2b-aligned-prototype", + "name": "VEA-B Phase 2B fixed aligned numerical prototype", + "family": "vortex-exact-attention", + "status": "research_retained", + "hypothesis": "A ten-warp VEA-B CTA can preserve every Sage2 arithmetic checkpoint across three aligned K64 epochs while separating QK/softmax ownership from PV/output ownership and eliminating material spills.", + "implementation_strategy": "Build a four-warp Sage checkpoint oracle and ten-warp paired-owner kernel for prequantized Q128 by KV192 by D128; require byte-exact intermediates/output, deterministic repetition, ptxas/runtime/NCU resource gates, sanitizer-clean mbarriers, measured overlap, and stable timing.", + "source_locations": ["research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md", "research/vortex_exact_attention/kernels/vea_b_numeric.cpp", "research/vortex_exact_attention/kernels/vea_b_numeric.cu", "research/vortex_exact_attention/tools/capture_phase2b_fixtures.py", "research/vortex_exact_attention/tools/run_phase2b.py"], + "active_source_location": "research/vortex_exact_attention/kernels/vea_b_numeric.cu", + "commit_hash": null, + "benchmark_artifacts": [ + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/authoritative-batched-final.json", "exists": true}, + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/authoritative-final.json", "exists": true}, + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/fixtures/aligned-q128-kv192-seed73021/manifest.json", "exists": true} + ], + "profiler_artifacts": [ + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/vea-b-final.ncu-rep", "exists": true}, + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/memcheck-final.log", "exists": true}, + {"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/racecheck-final.log", "exists": true} + ], + "environment": {"gpu": "NVIDIA GB10", "cuda": "13.0", "driver": "580.173.02", "pytorch": "2.9.1+cu130", "triton": null, "container_image": "sha256:6d880d628334c981c3d155bf5244e65e26e22cc9273c80145f646eee3c3698c2", "commit_context": "isolated fixed aligned Phase 2B prototype; production dispatch unchanged"}, + "metrics": {"shape": {"batch": 1, "heads": 1, "q": 128, "kv": 192, "d": 128}, "epochs": 3, "registers_per_thread": 168, "local_bytes_per_thread": 0, "dynamic_shared_bytes": 51200, "resident_ctas_per_sm": 1, "repeat_iterations": 1000, "qk_pv_overlap_clocks": [49317, 46408], "batched_vea_p50_ms": 0.014425599575042724, "batched_sage2_p50_ms": 0.014871200174093246, "batched_vea_cv": 0.00873681739214492, "batched_speedup": 1.0308895721618005, "single_launch_vea_p50_ms": 0.016287999227643013, "single_launch_sage2_p50_ms": 0.019840000197291374, "single_launch_speedup": 1.2180747260608975, "memcheck_errors": 0, "racecheck_hazards": 0}, + "correctness_evidence": ["The recovered checkpoint oracle matches installed public Sage2 output byte-for-byte.", "VEA-B matches QK, m, d, FP8 probability, rescale, FP16 PV, FP32 RO, reciprocal, and final BF16 output checkpoints exactly.", "The final BF16 output SHA-256 is 75d4169c9ae0692b24a463a64b538f852acba5a5567404728c8af60c16fb8e1b across 1,000 repeated launches.", "Memcheck reports zero errors and racecheck reports zero hazards."], + "decision_rationale": "Phase 2B passes for the single fixed aligned fixture. The architecture advances to ragged-tail and complete-shape research only; the microbenchmark does not establish canonical H3 or production speedup.", + "reproducer_commands": ["python research/vortex_exact_attention/tools/run_phase2b.py --fixture-dir --output --build-dir --build-log --image sha256:6d880d628334c981c3d155bf5244e65e26e22cc9273c80145f646eee3c3698c2 --warmup 200 --iterations 500 --repeat 1000 --timing-batch 20"], + "timestamp": "2026-08-26", + "evidence_missing": ["Ragged-tail exactness", "Complete canonical attention", "Blocks 0/24/49", "Two-step and twelve-step trajectories", "Resident-service and production integration gates"], + "production_behavior": "No dispatch route; production Sage2 remains unchanged.", + "source_recovery": "Selected source, fixture capture, validator, report, ptxas logs, sanitizer logs, NCU report, and authoritative JSON are retained." + }, { "id": "gb10-post-fc2-production-profile-20260826", "name": "Authoritative GB10 post-FC2 production baseline and profile", diff --git a/research/vortex_exact_attention/DESIGN.md b/research/vortex_exact_attention/DESIGN.md index 093ffa3..098145c 100644 --- a/research/vortex_exact_attention/DESIGN.md +++ b/research/vortex_exact_attention/DESIGN.md @@ -2,11 +2,12 @@ ## Decision -Prototype VEA-B, the one-CTA Q128 paired-owner pipeline. Phase 2A capability -probes pass with CUDA block-scope mbarriers as the selected two-slot handoff. -VEA-A remains the fallback if the complete 320-thread kernel fails exactness, -resources, or latency. VEA-C is the higher-risk second prototype only after DSM -placement and handoff latency are measured. No attention result is achieved. +Retain VEA-B, the one-CTA Q128 paired-owner pipeline. Phase 2A capability probes +and the Phase 2B fixed aligned numerical prototype pass with CUDA block-scope +mbarriers as the selected two-slot handoff. VEA-A remains the fallback if a +complete ragged-tail kernel fails exactness, resources, or latency. VEA-C is the +higher-risk second prototype only after DSM placement and handoff latency are +measured. No complete H3 attention result is achieved. Every design is side-by-side research under `kernels/`. The production Sage2 path remains the reference and fallback; no dispatch integration is authorized. diff --git a/research/vortex_exact_attention/EXPERIMENT_LOG.md b/research/vortex_exact_attention/EXPERIMENT_LOG.md index ad8501b..2ac6c38 100644 --- a/research/vortex_exact_attention/EXPERIMENT_LOG.md +++ b/research/vortex_exact_attention/EXPERIMENT_LOG.md @@ -93,3 +93,35 @@ gates; the invalid `16.631 ms` result is not hardware evidence. Decision: VEA-B advances only to one isolated aligned-shape exact-attention prototype using mbarriers. No attention latency or speedup has been measured, and production dispatch remains unchanged. + +## 2026-08-26: Phase 2B Fixed Aligned Numerical Prototype + +Status: `research_retained` + +The selected ten-warp `Q128 x KV192 x D128` VEA-B kernel matches the recovered +Sage2 checkpoint oracle for final BF16 output, QK, `m`, `d`, FP8 probability, +rescale, FP16 PV, FP32 `RO`, and reciprocal values. The output SHA-256 is +`75d4169c9ae0692b24a463a64b538f852acba5a5567404728c8af60c16fb8e1b`. +It repeats exactly for 1,000 launches. + +Role-separated loops, ordered 8-column PV fragments, warp scheduling boundaries, +and a per-lane shared retirement slot reduce the fast and capture kernels to 168 +registers/thread with zero stack, spill traffic, and local bytes. The selected +alias is protected by an explicit warp phase boundary; memcheck reports zero +errors and racecheck reports zero hazards. Runtime and NCU show one ten-warp +CTA/SM and zero local spilling requests. QK/PV clocks overlap across both epoch +transitions. + +The final batched stability run reports VEA-B p50 `0.014426 ms`, CV `0.87%`, +and public Sage2 p50 `0.014871 ms`, a fixed-shape prototype ratio of `1.031x`. +The 2,000-sample single-launch run reports `0.016288 ms` versus `0.019840 ms`, +or `1.218x`, with VEA-B CV `5.95%`. + +Rejected spill variants are retained in durable probe/build artifacts: combined +ownership, sequential PV fragments, union role state, shared persistent `RO`, +and `setmaxnreg` redistribution. The latter fails launch resource checks; shared +persistent `RO` changes epoch-1 FP32 checkpoints. + +Decision: Phase 2B passes only for the isolated aligned fixture. No ragged-tail, +canonical-shape, H3 block, trajectory, resident-service, or production dispatch +claim is made. diff --git a/research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md b/research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md new file mode 100644 index 0000000..c69909a --- /dev/null +++ b/research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md @@ -0,0 +1,94 @@ +# Phase 2B Aligned Numerical Prototype + +## Decision + +VEA-B passes the isolated fixed-shape Phase 2B gate on NVIDIA GB10. The selected +kernel computes `B=1, H=1, Q=128, KV=192, D=128` from prequantized Q/K/V with +three increasing K64 epochs. It is retained as research only. Ragged tails, +preparation, complete H3 shapes, block integration, trajectories, and production +dispatch remain out of scope. + +## Exactness + +The four-warp checkpoint oracle reproduces the installed public SageAttention +2.2.0 output. The ten-warp VEA-B kernel matches the oracle byte-for-byte for: + +- final BF16 output, SHA-256 + `75d4169c9ae0692b24a463a64b538f852acba5a5567404728c8af60c16fb8e1b`; +- INT32-to-FP32 QK scores; +- per-epoch online-softmax `m` and `d`; +- FP8 E4M3 probability words and FP32 rescale values; +- FP16 PV partial words and FP32 `RO` after every epoch; +- final reciprocal values. + +The selected source also completes 1,000 consecutive fast launches with the +same output hash and no deadlock. + +## Selected Kernel + +- 320 threads: two K/V producer warps, four QK/softmax owners, and four + PV/output owners. +- Two ready, two free, and two K-ready CUDA block-scope mbarriers. +- Two K/V and probability slots with strict epoch reuse. +- Final reciprocal publication is folded into the final score-ready handoff. +- Probability and V fragments are processed as ordered 8-column MMA pairs. +- A per-lane volatile shared retirement slot aliases scale metadata only after a + warp phase boundary. It constrains ptxas scheduling without changing + arithmetic or allocating additional shared memory. + +Ptxas and runtime attributes report 168 registers/thread, zero stack frame, +zero spill stores/loads, zero local bytes/thread, 51.2 KiB dynamic shared +memory, and one resident CTA/SM. NCU independently reports zero local-memory +spilling requests, 168 registers/thread, one register/shared-memory-limited +block per SM, and ten theoretical active warps (`20.83%`). + +## Synchronization And Safety + +- Compute Sanitizer memcheck: `0 errors`. +- Compute Sanitizer racecheck: `0 hazards`, `0 errors`, `0 warnings`. +- QK epoch 1 overlaps PV epoch 0 by 49,317 clocks. +- QK epoch 2 overlaps PV epoch 1 by 46,408 clocks. + +The retirement alias initially produced racecheck WAR warnings because lanes +could overwrite scale metadata before sibling lanes completed their reads. The +selected variant adds an explicit warp phase boundary before the first write. + +## Timing + +The final 2,000-sample single-launch run measures VEA-B p50 `0.016288 ms` and +public Sage2 p50 `0.019840 ms`, a prototype-only `1.218x` ratio. Single-launch +CV is `5.95%`, reflecting a bimodal microsecond-scale distribution. + +The authoritative stability run uses 500 event samples with 20 identical +launches per sample and reports per-launch values. VEA-B p50 is `0.014426 ms`, +CV `0.87%`; public Sage2 p50 is `0.014871 ms`, CV `1.23%`; the ratio is `1.031x`. +This is fixed-shape screening evidence, not a complete-attention speedup claim. + +## Rejected Variants + +| Variant | Result | +| --- | --- | +| Initial combined ownership | Exact, 352 local bytes/thread | +| Sequential PV fragments | Exact, material spills remained | +| Union role state | Exact, 208 local bytes/thread | +| Role-separated epoch loops | Exact, 40 local bytes/thread | +| `setmaxnreg` redistribution | Compiled at 200 registers, rejected at launch for excessive resources | +| Shared persistent `RO` fragment | Final BF16 exact, FP32 `RO` diverged at epoch 1 and stack usage increased | +| Inner-pair warp fence | Exact, but increased the fast stack frame | +| Selected shared retirement alias | Exact, zero local memory; race-clean after warp phase boundary | + +## Evidence + +- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/authoritative-batched-final.json` +- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/authoritative-final.json` +- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/memcheck-final.log` +- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/probes/racecheck-final.log` +- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/vea-b-final.ncu-rep` +- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2b/fixtures/aligned-q128-kv192-seed73021/manifest.json` + +## Advancement Limit + +Phase 2B authorizes no production integration. The next phase must add ragged +tails and complete-shape work distribution, then pass canonical attention, +blocks 0/24/49, trajectory, resident-service, sanitizer, and complete-block +performance gates before any dispatch proposal. diff --git a/research/vortex_exact_attention/README.md b/research/vortex_exact_attention/README.md index 89dc16c..5c8bf9c 100644 --- a/research/vortex_exact_attention/README.md +++ b/research/vortex_exact_attention/README.md @@ -28,7 +28,10 @@ Such work belongs in a future `vortex_fast` backend with a separate contract. - Phase 1: design decision complete; VEA-B is the recommended prototype. - Phase 2A: capability probes pass with CUDA block-scope mbarriers; one isolated aligned-shape prototype is authorized. -- Attention kernel implementation: not started; capability kernels only. +- Phase 2B: the fixed aligned `Q128 x KV192 x D128` numerical prototype passes + exactness, resource, sanitizer, overlap, repeatability, and timing gates. +- Attention kernel implementation: isolated aligned-shape research prototype; + no ragged-tail or complete H3 implementation. - Production dispatch: unchanged; no `vortex_exact` route exists. ## Authoritative Baseline @@ -69,11 +72,13 @@ p95, block-level, and resident-service gates. - `DESIGN.md`: state-ownership architectures and selected prototype. - `PHASE2A_CAPABILITY_REPORT.md`: measured resource, handoff, overlap, sanitizer, NCU, and fixture decision. +- `PHASE2B_ALIGNED_REPORT.md`: exact aligned prototype, rejected variants, + resources, sanitizer, NCU, overlap, and timing decision. - `PERFORMANCE_MODEL.md`: resource and latency model. - `HARDWARE_CAPABILITIES.md`: target-specific hardware facts and unknowns. - `EXPERIMENT_LOG.md`: append-only research decisions. - `benchmarks/`: machine-readable models and fixture manifests. -- `kernels/`: isolated capability probes; no attention kernel. +- `kernels/`: isolated capability probes and aligned numerical prototype. - `tests/`: future parity and sanitizer harnesses. - `tools/`: project-local verification/modeling utilities. @@ -88,6 +93,13 @@ result. Phase 2A measures 96 registers/thread for the selected mbarrier handoff, blocks. These measurements advance VEA-B to an isolated prototype, not to a performance or production claim. +Phase 2B implements the fixed three-epoch aligned shape with 168 +registers/thread, zero local bytes, one ten-warp CTA/SM, exact Sage2 checkpoints, +clean memcheck/racecheck, and positive QK/PV overlap. Batched timing is stable +with a prototype-only `1.031x` p50 ratio to the equivalent public Sage2 launch. +These results do not cover preparation, ragged tails, complete H3 shapes, or +production integration. + ## Provenance The reference is the public SageAttention repository at the pinned commit above. diff --git a/research/vortex_exact_attention/kernels/README.md b/research/vortex_exact_attention/kernels/README.md index eb15125..fb478a9 100644 --- a/research/vortex_exact_attention/kernels/README.md +++ b/research/vortex_exact_attention/kernels/README.md @@ -1,7 +1,7 @@ # Kernels -This directory contains isolated VEA-B capability kernels for resource -allocation, bounded handoff, and concurrent INT8/FP8 issue. No attention kernel -exists. Future candidates remain isolated here until every exactness and -performance gate passes; this directory must not be imported by production -dispatch. +This directory contains isolated VEA-B capability kernels and the fixed aligned +Phase 2B numerical prototype. `vea_b_numeric.cu` supports only prequantized +`Q128 x KV192 x D128`; it is not a ragged-tail or complete H3 attention +backend. Candidates remain isolated here until every integration gate passes; +this directory must not be imported by production dispatch. diff --git a/research/vortex_exact_attention/kernels/vea_b_numeric.cpp b/research/vortex_exact_attention/kernels/vea_b_numeric.cpp new file mode 100644 index 0000000..c1c564c --- /dev/null +++ b/research/vortex_exact_attention/kernels/vea_b_numeric.cpp @@ -0,0 +1,23 @@ +#include + +#include + +std::vector launch_vea_b_numeric( + torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale, + bool capture); +torch::Tensor launch_vea_b_numeric_into( + torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale, + torch::Tensor output); +std::vector launch_sage_checkpoint( + torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale); +torch::Tensor numeric_kernel_attributes(); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("vea_b", &launch_vea_b_numeric, "Aligned VEA-B numerical prototype"); + module.def("vea_b_into", &launch_vea_b_numeric_into, "Aligned VEA-B numerical prototype into output"); + module.def("sage_checkpoint", &launch_sage_checkpoint, "Pinned Sage-layout checkpoint oracle"); + module.def("attributes", &numeric_kernel_attributes, "Numerical kernel attributes"); +} diff --git a/research/vortex_exact_attention/kernels/vea_b_numeric.cu b/research/vortex_exact_attention/kernels/vea_b_numeric.cu new file mode 100644 index 0000000..d5c1338 --- /dev/null +++ b/research/vortex_exact_attention/kernels/vea_b_numeric.cu @@ -0,0 +1,851 @@ +/* + * Isolated aligned-shape VEA-B numerical prototype. + * + * The MMA, online-softmax, FP8 conversion, PV accumulation, and output fragment + * mappings are adapted from SageAttention 2.2.0 at commit + * d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 (Apache-2.0), whose helpers record + * prior adaptation from FlashInfer 0.1.5. This prototype supports only + * B=H=1, Q=128, KV=192, D=128 and is not production code. + */ + +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +#define CUDA_CHECK(call) do { \ + cudaError_t error_ = (call); \ + TORCH_CHECK(error_ == cudaSuccess, cudaGetErrorString(error_)); \ +} while (0) + +namespace { + +constexpr int kQ = 128; +constexpr int kKv = 192; +constexpr int kD = 128; +constexpr int kEpochs = 3; +constexpr int kOwners = 4; +constexpr int kThreads = 320; +constexpr int kReferenceThreads = 128; +constexpr int kDynamicShared = 50 * 1024; +constexpr int kReferenceShared = 32 * 1024; +constexpr float kSoftmaxScale = 0.08838834764831845f; +constexpr float kLog2e = 1.44269504088896340736f; +constexpr float kFp8Offset = 8.807f; + +struct Checkpoints { + float* qk; + float* m; + float* d; + uint32_t* p; + float* rescale; + uint32_t* pv; + float* ro; + float* reciprocal; + int64_t* clocks; +}; + +__device__ __forceinline__ float ptx_exp2(float value) { + float result; + asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value)); + return result; +} + +__device__ __forceinline__ float ptx_rcp(float value) { + float result; + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(value)); + return result; +} + +__device__ __forceinline__ void unpack_half2(float* output, uint32_t value) { + const uint16_t low = value & 0xffff; + const uint16_t high = value >> 16; + asm("cvt.f32.f16 %0, %1;" : "=f"(output[0]) : "h"(low)); + asm("cvt.f32.f16 %0, %1;" : "=f"(output[1]) : "h"(high)); +} + +__device__ __forceinline__ uint32_t fp8x4(float* first, float* second) { + uint32_t output; + asm volatile( + "{ .reg .b16 lo; .reg .b16 hi;\n" + "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" + "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" + "mov.b32 %0, {lo, hi}; }" + : "=r"(output) + : "f"(first[0]), "f"(first[1]), "f"(second[0]), "f"(second[1])); + return output; +} + +__device__ __forceinline__ void ldmatrix_x4(uint32_t* output, const void* pointer) { + const uint32_t address = static_cast(__cvta_generic_to_shared(pointer)); + asm volatile( + "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];" + : "=r"(output[0]), "=r"(output[1]), "=r"(output[2]), "=r"(output[3]) + : "r"(address)); +} + +__device__ __forceinline__ uint2 ldmatrix_x2(const void* pointer) { + const uint32_t address = static_cast(__cvta_generic_to_shared(pointer)); + uint2 output; + asm volatile( + "ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];" + : "=r"(output.x), "=r"(output.y) + : "r"(address)); + return output; +} + +__device__ __forceinline__ void int8_mma_init(int32_t* c, uint32_t* a, uint32_t* b) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13};" + : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "r"(0), "r"(0), "r"(0), "r"(0)); +} + +__device__ __forceinline__ void int8_mma(int32_t* c, uint32_t* a, uint32_t* b) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " + "{%0,%1,%2,%3},{%4,%5,%6,%7},{%8,%9},{%10,%11,%12,%13};" + : "=r"(c[0]), "=r"(c[1]), "=r"(c[2]), "=r"(c[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "r"(c[0]), "r"(c[1]), "r"(c[2]), "r"(c[3])); +} + +__device__ __forceinline__ void fp8_mma_init(uint32_t* c, uint32_t* a, uint32_t* b) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "r"(0), "r"(0)); +} + +__device__ __forceinline__ void fp8_mma(uint32_t* c, uint32_t* a, uint32_t* b) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};" + : "=r"(c[0]), "=r"(c[1]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "r"(c[0]), "r"(c[1])); +} + +__device__ __forceinline__ void fp8_mma_init_scalar( + uint32_t& c0, uint32_t& c1, uint4 a, uint2 b) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};" + : "=r"(c0), "=r"(c1) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b.x), "r"(b.y), + "r"(0), "r"(0)); +} + +__device__ __forceinline__ void fp8_mma_scalar( + uint32_t& c0, uint32_t& c1, uint4 a, uint2 b) { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0,%1},{%2,%3,%4,%5},{%6,%7},{%8,%9};" + : "=r"(c0), "=r"(c1) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(b.x), "r"(b.y), + "r"(c0), "r"(c1)); +} + +struct Smem128 { + uint4* base; + __device__ __forceinline__ uint32_t offset(int row, int column) const { + return row * 8 + (column ^ (row & 7)); + } +}; + +struct Smem64 { + uint4* base; + __device__ __forceinline__ uint32_t offset(int row, int column) const { + return row * 4 + (column ^ ((row / 2) & 3)); + } +}; + +__device__ __forceinline__ void load_q_fragments( + Smem128 q_smem, int owner, int lane, uint32_t rq[2][4][4]) { + #pragma unroll + for (int inner = 0; inner < 4; ++inner) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + const int row = owner * 32 + fq * 16 + lane % 16; + const int column = lane / 16 + inner * 2; + ldmatrix_x4(rq[fq][inner], q_smem.base + q_smem.offset(row, column)); + } + } +} + +__device__ __forceinline__ void compute_qk( + Smem128 k_smem, int lane, uint32_t rq[2][4][4], int32_t rs[2][4][8]) { + #pragma unroll + for (int inner = 0; inner < 4; ++inner) { + #pragma unroll + for (int fk = 0; fk < 4; ++fk) { + uint32_t rk[4]; + const int row = lane % 8 + (lane / 16) * 8 + fk * 16; + const int column = (lane / 8) % 2 + inner * 2; + ldmatrix_x4(rk, k_smem.base + k_smem.offset(row, column)); + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + if (inner == 0) { + int8_mma_init(rs[fq][fk], rq[fq][inner], rk); + int8_mma_init(rs[fq][fk] + 4, rq[fq][inner], rk + 2); + } else { + int8_mma(rs[fq][fk], rq[fq][inner], rk); + int8_mma(rs[fq][fk] + 4, rq[fq][inner], rk + 2); + } + } + } + } +} + +__device__ __forceinline__ int qk_checkpoint_index( + int epoch, int owner, int lane, int fq, int fk, int element) { + return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 4 + fk) * 8 + element); +} + +__device__ __forceinline__ int state_checkpoint_index( + int epoch, int owner, int lane, int fq, int state) { + return ((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 2 + state); +} + +__device__ __forceinline__ int p_checkpoint_index( + int epoch, int owner, int lane, int fq, int fk, int element) { + return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 2 + fk) * 4 + element); +} + +__device__ __forceinline__ int pv_checkpoint_index( + int epoch, int owner, int lane, int fq, int fv, int element) { + return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 8 + fv) * 4 + element); +} + +__device__ __forceinline__ int ro_checkpoint_index( + int epoch, int owner, int lane, int fq, int fv, int element) { + return (((((epoch * 4 + owner) * 32 + lane) * 2 + fq) * 8 + fv) * 8 + element); +} + +template +__device__ __forceinline__ void softmax_update( + float rs[2][4][8], float m[2][2], float d[2][2], float old_scale[2][2], + uint32_t rs_f8[2][2][4], float scale, int epoch, int owner, int lane, + Checkpoints checkpoints) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int state = 0; state < 2; ++state) { + const float previous = m[fq][state]; + float tile_max = -5000000.0f; + #pragma unroll + for (int fk = 0; fk < 4; ++fk) { + const float local = fmaxf( + fmaxf(rs[fq][fk][state * 2], rs[fq][fk][state * 2 + 1]), + fmaxf(rs[fq][fk][state * 2 + 4], rs[fq][fk][state * 2 + 5])); + tile_max = fmaxf(tile_max, local); + } + tile_max = fmaf(tile_max, scale, -kFp8Offset); + tile_max = fmaxf(tile_max, __shfl_xor_sync(0xffffffff, tile_max, 1)); + tile_max = fmaxf(tile_max, __shfl_xor_sync(0xffffffff, tile_max, 2)); + m[fq][state] = fmaxf(m[fq][state], tile_max); + old_scale[fq][state] = ptx_exp2(previous - m[fq][state]); + d[fq][state] *= old_scale[fq][state]; + const float negative_m = -m[fq][state]; + #pragma unroll + for (int fk = 0; fk < 4; ++fk) { + rs[fq][fk][state * 2] = ptx_exp2(fmaf(rs[fq][fk][state * 2], scale, negative_m)); + rs[fq][fk][state * 2 + 1] = ptx_exp2(fmaf(rs[fq][fk][state * 2 + 1], scale, negative_m)); + rs[fq][fk][state * 2 + 4] = ptx_exp2(fmaf(rs[fq][fk][state * 2 + 4], scale, negative_m)); + rs[fq][fk][state * 2 + 5] = ptx_exp2(fmaf(rs[fq][fk][state * 2 + 5], scale, negative_m)); + } + } + } + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int fk = 0; fk < 4; ++fk) { + d[fq][0] += rs[fq][fk][0] + rs[fq][fk][1] + rs[fq][fk][4] + rs[fq][fk][5]; + d[fq][1] += rs[fq][fk][2] + rs[fq][fk][3] + rs[fq][fk][6] + rs[fq][fk][7]; + } + #pragma unroll + for (int fk = 0; fk < 2; ++fk) { + rs_f8[fq][fk][0] = fp8x4(rs[fq][fk * 2], rs[fq][fk * 2] + 4); + rs_f8[fq][fk][1] = fp8x4(rs[fq][fk * 2] + 2, rs[fq][fk * 2] + 6); + rs_f8[fq][fk][2] = fp8x4(rs[fq][fk * 2 + 1], rs[fq][fk * 2 + 1] + 4); + rs_f8[fq][fk][3] = fp8x4(rs[fq][fk * 2 + 1] + 2, rs[fq][fk * 2 + 1] + 6); + } + } + if constexpr (Capture) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int state = 0; state < 2; ++state) { + const int index = state_checkpoint_index(epoch, owner, lane, fq, state); + checkpoints.m[index] = m[fq][state]; + checkpoints.d[index] = d[fq][state]; + checkpoints.rescale[index] = old_scale[fq][state]; + } + #pragma unroll + for (int fk = 0; fk < 2; ++fk) { + #pragma unroll + for (int element = 0; element < 4; ++element) { + checkpoints.p[p_checkpoint_index(epoch, owner, lane, fq, fk, element)] = rs_f8[fq][fk][element]; + } + } + } + } +} + +template +__device__ __forceinline__ void compute_pv( + Smem64 v_smem, uint32_t rs_f8[2][2][4], float ro[2][8][8], + int epoch, int owner, int lane, Checkpoints checkpoints) { + const int row_base = lane % 8 + (lane / 16) * 8; + const int column_base = (lane / 8) % 2; + #pragma unroll + for (int fv = 0; fv < 8; ++fv) { + uint32_t instant[2][4]; + #pragma unroll + for (int fk = 0; fk < 2; ++fk) { + uint32_t rv[4]; + const uint32_t offset = v_smem.offset(row_base + fv * 16, column_base + fk * 2); + ldmatrix_x4(rv, v_smem.base + offset); + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + if (fk == 0) { + fp8_mma_init(instant[fq], rs_f8[fq][fk], rv); + fp8_mma_init(instant[fq] + 2, rs_f8[fq][fk], rv + 2); + } else { + fp8_mma(instant[fq], rs_f8[fq][fk], rv); + fp8_mma(instant[fq] + 2, rs_f8[fq][fk], rv + 2); + } + } + } + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int element = 0; element < 4; ++element) { + float pair[2]; + unpack_half2(pair, instant[fq][element]); + ro[fq][fv][element * 2] += pair[0]; + ro[fq][fv][element * 2 + 1] += pair[1]; + if constexpr (Capture) { + checkpoints.pv[pv_checkpoint_index(epoch, owner, lane, fq, fv, element)] = instant[fq][element]; + } + } + if constexpr (Capture) { + #pragma unroll + for (int element = 0; element < 8; ++element) { + checkpoints.ro[ro_checkpoint_index(epoch, owner, lane, fq, fv, element)] = ro[fq][fv][element]; + } + } + } + } +} + +template +__device__ __forceinline__ void compute_pv_from_shared( + Smem64 v_smem, const uint4* score_slot, float ro[2][8][8], + volatile float* retire_slot, int epoch, int owner, int lane, + Checkpoints checkpoints) { + const int row_base = lane % 8 + (lane / 16) * 8; + const int column_base = (lane / 8) % 2; + #pragma unroll + for (int fv = 0; fv < 8; ++fv) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int output_pair = 0; output_pair < 2; ++output_pair) { + uint32_t instant[2]; + #pragma unroll + for (int fk = 0; fk < 2; ++fk) { + const uint4 packed = score_slot[((owner * 4 + fq * 2 + fk) * 32 + lane)]; + const uint32_t offset = v_smem.offset( + row_base + fv * 16 + output_pair * 8, column_base + fk * 2); + const uint2 rv = ldmatrix_x2(v_smem.base + offset); + if (fk == 0) + fp8_mma_init_scalar(instant[0], instant[1], packed, rv); + else + fp8_mma_scalar(instant[0], instant[1], packed, rv); + } + #pragma unroll + for (int element = 0; element < 2; ++element) { + float pair[2]; + unpack_half2(pair, instant[element]); + const int ro_element = output_pair * 4 + element * 2; + ro[fq][fv][ro_element] += pair[0]; + ro[fq][fv][ro_element + 1] += pair[1]; + if constexpr (Capture) + checkpoints.pv[pv_checkpoint_index( + epoch, owner, lane, fq, fv, output_pair * 2 + element)] = instant[element]; + } + } + if constexpr (!Capture) { + #pragma unroll + for (int element = 0; element < 8; ++element) + *retire_slot = ro[fq][fv][element]; + } + __syncwarp(); + if constexpr (Capture) { + #pragma unroll + for (int element = 0; element < 8; ++element) + checkpoints.ro[ro_checkpoint_index(epoch, owner, lane, fq, fv, element)] = ro[fq][fv][element]; + } + } + } +} + +__device__ __forceinline__ void save_output( + nv_bfloat16* output, float ro[2][8][8], float reciprocal[2][2], + const float* v_scale, int owner, int lane) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int fv = 0; fv < 8; ++fv) { + float scales[4]; + const int base = (lane % 4) * 2 + fv * 16; + scales[0] = v_scale[base]; + scales[1] = v_scale[base + 1]; + scales[2] = v_scale[base + 8]; + scales[3] = v_scale[base + 9]; + #pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const int element = pair * 2; + const int state = (element % 4) / 2; + ro[fq][fv][element] *= reciprocal[fq][state]; + ro[fq][fv][element + 1] *= reciprocal[fq][state]; + const int scale_pair = pair >= 2 ? 2 : 0; + ro[fq][fv][element] *= scales[scale_pair]; + ro[fq][fv][element + 1] *= scales[scale_pair + 1]; + const int row = owner * 32 + fq * 16 + lane / 4 + 8 * state; + const int column = fv * 16 + (lane % 4) * 2 + (pair >= 2 ? 8 : 0); + const nv_bfloat162 packed = __float22bfloat162_rn( + make_float2(ro[fq][fv][element], ro[fq][fv][element + 1])); + *reinterpret_cast(output + row * 128 + column) = packed; + } + } + } +} + +__device__ __forceinline__ void save_output_streamed( + nv_bfloat16* output, float ro[2][8][8], const float* reciprocal, + const float* v_scale, int owner, int lane) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int fv = 0; fv < 8; ++fv) { + #pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const int element = pair * 2; + const int component = (element % 4) / 2; + const float normalizer = reciprocal[owner * 32 + fq * 16 + component * 8 + lane / 4]; + ro[fq][fv][element] *= normalizer; + ro[fq][fv][element + 1] *= normalizer; + const int column = fv * 16 + (lane % 4) * 2 + (pair >= 2 ? 8 : 0); + ro[fq][fv][element] *= v_scale[column]; + ro[fq][fv][element + 1] *= v_scale[column + 1]; + const int row = owner * 32 + fq * 16 + lane / 4 + 8 * component; + const nv_bfloat162 packed = __float22bfloat162_rn( + make_float2(ro[fq][fv][element], ro[fq][fv][element + 1])); + *reinterpret_cast(output + row * 128 + column) = packed; + } + } + } +} + +template +__global__ __launch_bounds__(kThreads, 1) void vea_b_numeric_kernel( + const int8_t* q, const int8_t* k, const int8_t* v, + const float* q_scale, const float* k_scale, const float* v_scale, + nv_bfloat16* output, Checkpoints checkpoints) { + extern __shared__ uint8_t shared[]; + Smem128 q_smem{reinterpret_cast(shared)}; + Smem128 k_slots[2] = { + {reinterpret_cast(shared + 16 * 1024)}, + {reinterpret_cast(shared + 24 * 1024)}}; + Smem64 v_slots[2] = { + {reinterpret_cast(shared + 32 * 1024)}, + {reinterpret_cast(shared + 40 * 1024)}}; + using BlockBarrier = cuda::barrier; + BlockBarrier* barriers = reinterpret_cast(shared + 48 * 1024); + float* scale_slots = reinterpret_cast(shared + 48 * 1024 + 128); + float* final_reciprocal = scale_slots + 2 * 128; + const int tid = threadIdx.x; + const int warp = tid / 32; + const int lane = tid % 32; + const bool qk_owner = warp >= 2 && warp <= 5; + const bool pv_owner = warp >= 6; + const int owner = qk_owner ? warp - 2 : (pv_owner ? warp - 6 : -1); + + if (tid < 6) { + const int count = tid < 4 ? 320 : 160; + init(barriers + tid, count); + } + for (int chunk = tid; chunk < 1024; chunk += kThreads) { + const int row = chunk / 8; + const int column = chunk % 8; + q_smem.base[q_smem.offset(row, column)] = reinterpret_cast(q + row * 128)[column]; + } + __syncthreads(); + + uint32_t rq[2][4][4]; + if (qk_owner) load_q_fragments(q_smem, owner, lane, rq); + __syncthreads(); + + if (qk_owner) { + float m[2][2] = {{-5000000.0f, -5000000.0f}, {-5000000.0f, -5000000.0f}}; + float d[2][2] = {{1.0f, 1.0f}, {1.0f, 1.0f}}; + #pragma unroll + for (int epoch = 0; epoch < kEpochs; ++epoch) { + const int slot = epoch & 1; + if (epoch >= 2) barriers[2 + slot].arrive_and_wait(); + barriers[4 + slot].arrive_and_wait(); + const uint64_t started = clock64(); + union ScoreStorage { int32_t integer[2][4][8]; float fp32[2][4][8]; } scores; + compute_qk(k_slots[slot], lane, rq, scores.integer); + const float dequant = q_scale[owner] * k_scale[epoch]; + const bool final_pair = epoch + 2 >= kEpochs; + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int fk = 0; fk < 4; ++fk) { + #pragma unroll + for (int element = 0; element < 8; ++element) { + const float converted = __int2float_rz(scores.integer[fq][fk][element]); + scores.fp32[fq][fk][element] = final_pair ? converted * dequant : converted; + if constexpr (Capture) { + checkpoints.qk[qk_checkpoint_index(epoch, owner, lane, fq, fk, element)] = scores.fp32[fq][fk][element]; + } + } + } + } + float old_scale[2][2]; + uint32_t probability[2][2][4]; + softmax_update(scores.fp32, m, d, old_scale, probability, + final_pair ? kSoftmaxScale * kLog2e : kSoftmaxScale * kLog2e * dequant, + epoch, owner, lane, checkpoints); + if ((lane & 3) == 0) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int state = 0; state < 2; ++state) + scale_slots[slot * 128 + owner * 32 + fq * 16 + state * 8 + lane / 4] = old_scale[fq][state]; + } + uint4* score_slot = reinterpret_cast(shared + slot * 8 * 1024); + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int fk = 0; fk < 2; ++fk) + score_slot[((owner * 4 + fq * 2 + fk) * 32 + lane)] = *reinterpret_cast(probability[fq][fk]); + if constexpr (Capture) { + if (lane == 0) { + checkpoints.clocks[((epoch * 8 + owner) * 2)] = started; + checkpoints.clocks[((epoch * 8 + owner) * 2) + 1] = clock64(); + } + } + if (epoch == kEpochs - 1) { + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int component = 0; component < 2; ++component) { + d[fq][component] += __shfl_xor_sync(0xffffffff, d[fq][component], 1); + d[fq][component] += __shfl_xor_sync(0xffffffff, d[fq][component], 2); + const float reciprocal = ptx_rcp(d[fq][component]); + if constexpr (Capture) + checkpoints.reciprocal[state_checkpoint_index(0, owner, lane, fq, component)] = reciprocal; + if ((lane & 3) == 0) + final_reciprocal[owner * 32 + fq * 16 + component * 8 + lane / 4] = reciprocal; + } + } + } + (void)barriers[slot].arrive(); + } + barriers[2].arrive_and_wait(); + barriers[3].arrive_and_wait(); + } else if (pv_owner) { + float ro[2][8][8]; + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int fv = 0; fv < 8; ++fv) + #pragma unroll + for (int element = 0; element < 8; ++element) ro[fq][fv][element] = 0.0f; + #pragma unroll + for (int epoch = 0; epoch < kEpochs; ++epoch) { + const int slot = epoch & 1; + barriers[slot].arrive_and_wait(); + const uint64_t started = clock64(); + #pragma unroll + for (int fq = 0; fq < 2; ++fq) { + #pragma unroll + for (int component = 0; component < 2; ++component) { + const float old_scale = scale_slots[slot * 128 + owner * 32 + fq * 16 + component * 8 + lane / 4]; + #pragma unroll + for (int fv = 0; fv < 8; ++fv) { + ro[fq][fv][component * 2] *= old_scale; + ro[fq][fv][component * 2 + 1] *= old_scale; + ro[fq][fv][component * 2 + 4] *= old_scale; + ro[fq][fv][component * 2 + 5] *= old_scale; + } + } + } + __syncwarp(); + compute_pv_from_shared( + v_slots[slot], reinterpret_cast(shared + slot * 8 * 1024), + ro, scale_slots + slot * 128 + owner * 32 + lane, + epoch, owner, lane, checkpoints); + if constexpr (Capture) { + if (lane == 0) { + checkpoints.clocks[((epoch * 8 + 4 + owner) * 2)] = started; + checkpoints.clocks[((epoch * 8 + 4 + owner) * 2) + 1] = clock64(); + } + } + (void)barriers[2 + slot].arrive(); + } + save_output_streamed(output, ro, final_reciprocal, v_scale, owner, lane); + } else if (warp == 0) { + #pragma unroll + for (int epoch = 0; epoch < kEpochs; ++epoch) { + const int slot = epoch & 1; + if (epoch >= 2) barriers[2 + slot].arrive_and_wait(); + for (int chunk = lane; chunk < 512; chunk += 32) { + const int row = chunk / 8; + const int column = chunk % 8; + k_slots[slot].base[k_slots[slot].offset(row, column)] = + reinterpret_cast(k + (epoch * 64 + row) * 128)[column]; + } + (void)barriers[4 + slot].arrive(); + (void)barriers[slot].arrive(); + } + barriers[2].arrive_and_wait(); + barriers[3].arrive_and_wait(); + } else { + #pragma unroll + for (int epoch = 0; epoch < kEpochs; ++epoch) { + const int slot = epoch & 1; + if (epoch >= 2) barriers[2 + slot].arrive_and_wait(); + for (int chunk = lane; chunk < 512; chunk += 32) { + const int row = chunk / 4; + const int column = chunk % 4; + v_slots[slot].base[v_slots[slot].offset(row, column)] = + reinterpret_cast(v + row * 192 + epoch * 64)[column]; + } + (void)barriers[slot].arrive(); + } + barriers[2].arrive_and_wait(); + barriers[3].arrive_and_wait(); + } +} + +template +__global__ __launch_bounds__(kReferenceThreads, 1) void sage_checkpoint_kernel( + const int8_t* q, const int8_t* k, const int8_t* v, + const float* q_scale, const float* k_scale, const float* v_scale, + nv_bfloat16* output, Checkpoints checkpoints) { + extern __shared__ uint8_t shared[]; + Smem128 q_smem{reinterpret_cast(shared)}; + Smem128 k_smem{reinterpret_cast(shared + 16 * 1024)}; + Smem64 v_smem{reinterpret_cast(shared + 24 * 1024)}; + const int tid = threadIdx.x; + const int owner = tid / 32; + const int lane = tid % 32; + for (int chunk = tid; chunk < 1024; chunk += kReferenceThreads) { + const int row = chunk / 8; + const int column = chunk % 8; + q_smem.base[q_smem.offset(row, column)] = reinterpret_cast(q + row * 128)[column]; + } + __syncthreads(); + uint32_t rq[2][4][4]; + load_q_fragments(q_smem, owner, lane, rq); + float m[2][2] = {{-5000000.0f, -5000000.0f}, {-5000000.0f, -5000000.0f}}; + float d[2][2] = {{1.0f, 1.0f}, {1.0f, 1.0f}}; + float ro[2][8][8]; + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int fv = 0; fv < 8; ++fv) + #pragma unroll + for (int element = 0; element < 8; ++element) ro[fq][fv][element] = 0.0f; + + #pragma unroll + for (int epoch = 0; epoch < kEpochs; ++epoch) { + for (int chunk = tid; chunk < 512; chunk += kReferenceThreads) { + const int row = chunk / 8; + const int column = chunk % 8; + k_smem.base[k_smem.offset(row, column)] = reinterpret_cast(k + (epoch * 64 + row) * 128)[column]; + const int vrow = chunk / 4; + const int vcolumn = chunk % 4; + v_smem.base[v_smem.offset(vrow, vcolumn)] = reinterpret_cast(v + vrow * 192 + epoch * 64)[vcolumn]; + } + __syncthreads(); + union ScoreStorage { int32_t integer[2][4][8]; float fp32[2][4][8]; } scores; + compute_qk(k_smem, lane, rq, scores.integer); + const float dequant = q_scale[owner] * k_scale[epoch]; + const bool final_pair = epoch + 2 >= kEpochs; + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int fk = 0; fk < 4; ++fk) + #pragma unroll + for (int element = 0; element < 8; ++element) { + const float converted = __int2float_rz(scores.integer[fq][fk][element]); + scores.fp32[fq][fk][element] = final_pair ? converted * dequant : converted; + if constexpr (Capture) + checkpoints.qk[qk_checkpoint_index(epoch, owner, lane, fq, fk, element)] = scores.fp32[fq][fk][element]; + } + float old_scale[2][2]; + uint32_t probability[2][2][4]; + softmax_update(scores.fp32, m, d, old_scale, probability, + final_pair ? kSoftmaxScale * kLog2e : kSoftmaxScale * kLog2e * dequant, + epoch, owner, lane, checkpoints); + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int fv = 0; fv < 8; ++fv) + #pragma unroll + for (int element = 0; element < 8; ++element) + ro[fq][fv][element] *= old_scale[fq][(element % 4) / 2]; + compute_pv(v_smem, probability, ro, epoch, owner, lane, checkpoints); + __syncthreads(); + } + float reciprocal[2][2]; + #pragma unroll + for (int fq = 0; fq < 2; ++fq) + #pragma unroll + for (int state = 0; state < 2; ++state) { + d[fq][state] += __shfl_xor_sync(0xffffffff, d[fq][state], 1); + d[fq][state] += __shfl_xor_sync(0xffffffff, d[fq][state], 2); + reciprocal[fq][state] = ptx_rcp(d[fq][state]); + if constexpr (Capture) checkpoints.reciprocal[state_checkpoint_index(0, owner, lane, fq, state)] = reciprocal[fq][state]; + } + save_output(output, ro, reciprocal, v_scale, owner, lane); +} + +std::vector make_outputs(torch::Device device) { + auto byte = torch::TensorOptions().device(device).dtype(torch::kInt32); + auto fp32 = torch::TensorOptions().device(device).dtype(torch::kFloat32); + auto output = torch::empty({1, 128, 1, 128}, torch::TensorOptions().device(device).dtype(torch::kBFloat16)); + auto qk = torch::empty({3, 4, 32, 2, 4, 8}, fp32); + auto m = torch::empty({3, 4, 32, 2, 2}, fp32); + auto d = torch::empty({3, 4, 32, 2, 2}, fp32); + auto p = torch::empty({3, 4, 32, 2, 2, 4}, byte); + auto rescale = torch::empty({3, 4, 32, 2, 2}, fp32); + auto pv = torch::empty({3, 4, 32, 2, 8, 4}, byte); + auto ro = torch::empty({3, 4, 32, 2, 8, 8}, fp32); + auto reciprocal = torch::empty({4, 32, 2, 2}, fp32); + auto clocks = torch::zeros({3, 8, 2}, torch::TensorOptions().device(device).dtype(torch::kInt64)); + return {output, qk, m, d, p, rescale, pv, ro, reciprocal, clocks}; +} + +Checkpoints pointers(std::vector& outputs) { + return {outputs[1].data_ptr(), outputs[2].data_ptr(), outputs[3].data_ptr(), + reinterpret_cast(outputs[4].data_ptr()), outputs[5].data_ptr(), + reinterpret_cast(outputs[6].data_ptr()), outputs[7].data_ptr(), + outputs[8].data_ptr(), outputs[9].data_ptr()}; +} + +void validate_inputs(torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale) { + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "inputs must be CUDA tensors"); + TORCH_CHECK(q.scalar_type() == torch::kInt8 && k.scalar_type() == torch::kInt8, "Q/K must be INT8"); + TORCH_CHECK(q.sizes() == torch::IntArrayRef({1, 128, 1, 128}), "Q shape must be [1,128,1,128]"); + TORCH_CHECK(k.sizes() == torch::IntArrayRef({1, 192, 1, 128}), "K shape must be [1,192,1,128]"); + TORCH_CHECK(v.numel() == 128 * 192 && v.element_size() == 1, "V must contain [1,128,1,192] FP8 bytes"); + TORCH_CHECK(q_scale.numel() == 4 && k_scale.numel() == 3 && v_scale.numel() == 128, "invalid scale shapes"); + TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "inputs must be contiguous"); +} + +template +std::array attributes(Kernel kernel, int shared_bytes, int threads) { + cudaFuncAttributes value{}; + CUDA_CHECK(cudaFuncGetAttributes(&value, kernel)); + int active = 0; + CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active, kernel, threads, shared_bytes)); + return {value.numRegs, static_cast(value.sharedSizeBytes), static_cast(value.localSizeBytes), + value.maxThreadsPerBlock, value.binaryVersion, value.ptxVersion, + static_cast(value.maxDynamicSharedSizeBytes), active}; +} + +} // namespace + +std::vector launch_vea_b_numeric( + torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale, + bool capture) { + validate_inputs(q, k, v, q_scale, k_scale, v_scale); + c10::cuda::CUDAGuard guard(q.device()); + auto outputs = make_outputs(q.device()); + CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared)); + CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared)); + const auto stream = at::cuda::getDefaultCUDAStream(); + if (capture) { + vea_b_numeric_kernel<<<1, kThreads, kDynamicShared, stream>>>( + q.data_ptr(), k.data_ptr(), reinterpret_cast(v.data_ptr()), + q_scale.data_ptr(), k_scale.data_ptr(), v_scale.data_ptr(), + reinterpret_cast(outputs[0].data_ptr()), pointers(outputs)); + } else { + vea_b_numeric_kernel<<<1, kThreads, kDynamicShared, stream>>>( + q.data_ptr(), k.data_ptr(), reinterpret_cast(v.data_ptr()), + q_scale.data_ptr(), k_scale.data_ptr(), v_scale.data_ptr(), + reinterpret_cast(outputs[0].data_ptr()), pointers(outputs)); + } + CUDA_CHECK(cudaGetLastError()); + return outputs; +} + +std::vector launch_sage_checkpoint( + torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale) { + validate_inputs(q, k, v, q_scale, k_scale, v_scale); + c10::cuda::CUDAGuard guard(q.device()); + auto outputs = make_outputs(q.device()); + CUDA_CHECK(cudaFuncSetAttribute(sage_checkpoint_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kReferenceShared)); + sage_checkpoint_kernel<<<1, kReferenceThreads, kReferenceShared, at::cuda::getDefaultCUDAStream()>>>( + q.data_ptr(), k.data_ptr(), reinterpret_cast(v.data_ptr()), + q_scale.data_ptr(), k_scale.data_ptr(), v_scale.data_ptr(), + reinterpret_cast(outputs[0].data_ptr()), pointers(outputs)); + CUDA_CHECK(cudaGetLastError()); + return outputs; +} + +torch::Tensor launch_vea_b_numeric_into( + torch::Tensor q, torch::Tensor k, torch::Tensor v, + torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale, + torch::Tensor output) { + validate_inputs(q, k, v, q_scale, k_scale, v_scale); + TORCH_CHECK(output.is_cuda() && output.scalar_type() == torch::kBFloat16 && + output.sizes() == torch::IntArrayRef({1, 128, 1, 128}) && output.is_contiguous(), + "output must be contiguous CUDA BF16 [1,128,1,128]"); + c10::cuda::CUDAGuard guard(q.device()); + CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared)); + Checkpoints empty{}; + vea_b_numeric_kernel<<<1, kThreads, kDynamicShared, at::cuda::getDefaultCUDAStream()>>>( + q.data_ptr(), k.data_ptr(), reinterpret_cast(v.data_ptr()), + q_scale.data_ptr(), k_scale.data_ptr(), v_scale.data_ptr(), + reinterpret_cast(output.data_ptr()), empty); + CUDA_CHECK(cudaGetLastError()); + return output; +} + +torch::Tensor numeric_kernel_attributes() { + CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared)); + CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared)); + CUDA_CHECK(cudaFuncSetAttribute(sage_checkpoint_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kReferenceShared)); + std::array, 3> rows = { + attributes(vea_b_numeric_kernel, kDynamicShared, kThreads), + attributes(vea_b_numeric_kernel, kDynamicShared, kThreads), + attributes(sage_checkpoint_kernel, kReferenceShared, kReferenceThreads)}; + auto output = torch::empty({3, 8}, torch::TensorOptions().dtype(torch::kInt64)); + auto access = output.accessor(); + for (int row = 0; row < 3; ++row) + for (int column = 0; column < 8; ++column) access[row][column] = rows[row][column]; + return output; +} diff --git a/research/vortex_exact_attention/tools/capture_phase2b_fixtures.py b/research/vortex_exact_attention/tools/capture_phase2b_fixtures.py new file mode 100644 index 0000000..ed6a0c3 --- /dev/null +++ b/research/vortex_exact_attention/tools/capture_phase2b_fixtures.py @@ -0,0 +1,133 @@ +"""Capture fixed aligned prequantized Sage2 fixtures for VEA-B Phase 2B.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import platform +from pathlib import Path + +import torch + + +SHAPE = {"batch": 1, "heads": 1, "q": 128, "kv": 192, "d": 128} +SEED = 73021 + + +def tensor_sha256(value: torch.Tensor) -> str: + data = value.detach().contiguous().view(torch.uint8).cpu().numpy() + return hashlib.sha256(memoryview(data)).hexdigest() + + +def file_sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def save(output_dir: Path, name: str, value: torch.Tensor) -> dict: + host = value.detach().contiguous().cpu() + path = output_dir / f"{name}.pt" + torch.save({"tensor": host}, path) + loaded = torch.load(path, map_location="cpu", weights_only=True)["tensor"] + if not torch.equal(host, loaded): + raise RuntimeError(f"reload mismatch for {name}") + return { + "path": str(path), + "shape": list(host.shape), + "dtype": str(host.dtype), + "tensor_sha256": tensor_sha256(host), + "file_sha256": file_sha256(path), + "size_bytes": path.stat().st_size, + "reload_verified": True, + } + + +def public_mainloop(q, k, v, q_scale, k_scale, v_scale): + import sageattention.core as sage_core + + output = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device) + sage_core.sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( + q, k, v, output, q_scale, k_scale, v_scale, + 0, 0, 2, SHAPE["d"] ** -0.5, 0, + ) + return output + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--image", required=True) + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + + generator = torch.Generator(device="cuda").manual_seed(SEED) + q_bf16 = torch.randn((1, 128, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator) + k_bf16 = torch.randn((1, 192, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator) + v_bf16 = torch.randn((1, 192, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator) + + import sageattention.core as sage_core + + with torch.inference_mode(): + k_mean = k_bf16.mean(dim=1, keepdim=True) + q_int8, q_scale, k_int8, k_scale = sage_core.per_warp_int8_cuda( + q_bf16, k_bf16, k_mean, BLKQ=128, WARPQ=32, BLKK=64, + tensor_layout="NHD", + ) + v_fp8, v_scale, _ = sage_core.per_channel_fp8( + v_bf16, tensor_layout="NHD", scale_max=2.25, smooth_v=False, + ) + public_output = public_mainloop(q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale) + torch.cuda.synchronize() + + tensors = { + "q_bf16": q_bf16, + "k_bf16": k_bf16, + "v_bf16": v_bf16, + "k_mean_bf16": k_mean, + "q_int8": q_int8, + "q_scale_fp32": q_scale, + "k_int8": k_int8, + "k_scale_fp32": k_scale, + "v_fp8_e4m3": v_fp8, + "v_scale_fp32": v_scale, + "public_output_bf16": public_output, + } + records = {name: save(args.output_dir, name, value) for name, value in tensors.items()} + manifest = { + "schema": "vortex-exact-phase2b-aligned-fixtures", + "version": 1, + "status": "captured_and_reload_verified", + "shape": SHAPE, + "epochs": 3, + "seed": SEED, + "scope": "prepared BF16 retained only for provenance; prototype inputs are prequantized tensors", + "reference": { + "sageattention_version": "2.2.0", + "sageattention_commit": "d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5", + "mainloop": "qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf", + "tensor_layout": "NHD", + "causal": False, + "qk_quant_granularity": "per_warp", + "v_scale_max": 2.25, + }, + "environment": { + "image": args.image, + "gpu": torch.cuda.get_device_name(), + "compute_capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "python": platform.python_version(), + }, + "tensors": records, + } + manifest_path = args.output_dir / "manifest.json" + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8") + print(json.dumps(manifest, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/research/vortex_exact_attention/tools/run_phase2b.py b/research/vortex_exact_attention/tools/run_phase2b.py new file mode 100644 index 0000000..3859528 --- /dev/null +++ b/research/vortex_exact_attention/tools/run_phase2b.py @@ -0,0 +1,267 @@ +"""Build and validate the fixed aligned VEA-B Phase 2B numerical prototype.""" + +from __future__ import annotations + +import argparse +import contextlib +import hashlib +import io +import json +import math +import os +import platform +import statistics +import sys +import time +from pathlib import Path + +import torch +from torch.utils.cpp_extension import load + + +PROJECT = Path(__file__).resolve().parents[1] +KERNELS = PROJECT / "kernels" +NAMES = ("output", "qk_scores", "m", "d", "probability_fp8", "rescale", "pv_fp16", "ro_fp32", "reciprocal", "clocks") +ATTRIBUTE_NAMES = ("registers_per_thread", "static_shared_bytes", "local_bytes_per_thread", "max_threads_per_block", "binary_version", "ptx_version", "max_dynamic_shared_bytes", "resident_ctas_per_sm") +KERNEL_NAMES = ("vea_b_fast", "vea_b_capture", "sage_checkpoint") + + +def tensor_sha256(value: torch.Tensor) -> str: + data = value.detach().contiguous().view(torch.uint8).cpu().numpy() + return hashlib.sha256(memoryview(data)).hexdigest() + + +def percentile(values: list[float], fraction: float) -> float: + ordered = sorted(values) + rank = (len(ordered) - 1) * fraction + low = int(rank) + high = min(low + 1, len(ordered) - 1) + weight = rank - low + return ordered[low] * (1 - weight) + ordered[high] * weight + + +def summarize(values: list[float]) -> dict: + mean = statistics.fmean(values) + return { + "count": len(values), "p10_ms": percentile(values, 0.10), + "p50_ms": percentile(values, 0.50), "p90_ms": percentile(values, 0.90), + "p95_ms": percentile(values, 0.95), "mean_ms": mean, + "stdev_ms": statistics.pstdev(values), + "coefficient_of_variation": statistics.pstdev(values) / mean if mean else 0.0, + "min_ms": min(values), "max_ms": max(values), + } + + +def benchmark(fn, warmup: int, iterations: int, batch: int) -> dict: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + values = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(batch): + fn() + end.record() + end.synchronize() + values.append(start.elapsed_time(end) / batch) + result = summarize(values) + result["launches_per_sample"] = batch + return result + + +def build_extension(build_dir: Path, build_log: Path): + build_dir.mkdir(parents=True, exist_ok=True) + stream = io.StringIO() + started = time.perf_counter() + with contextlib.redirect_stdout(stream), contextlib.redirect_stderr(stream): + extension = load( + name="vortex_exact_phase2b_numeric", + sources=[str(KERNELS / "vea_b_numeric.cpp"), str(KERNELS / "vea_b_numeric.cu")], + build_directory=str(build_dir), + extra_cflags=["-O3"], + extra_cuda_cflags=["-O3", "-lineinfo", "--ptxas-options=-v", "--maxrregcount=200", "-std=c++17"], + verbose=True, + ) + text = stream.getvalue() + sys.stderr.write(text) + build_log.parent.mkdir(parents=True, exist_ok=True) + build_log.write_text(text, encoding="utf-8") + return extension, time.perf_counter() - started + + +def load_fixture(directory: Path) -> tuple[dict, dict[str, torch.Tensor]]: + manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8")) + tensors = {} + for name, record in manifest["tensors"].items(): + value = torch.load(directory / Path(record["path"]).name, map_location="cpu", weights_only=True)["tensor"] + if tensor_sha256(value) != record["tensor_sha256"]: + raise RuntimeError(f"fixture hash mismatch: {name}") + tensors[name] = value.cuda() + return manifest, tensors + + +def public_launch(tensors: dict[str, torch.Tensor], output: torch.Tensor) -> torch.Tensor: + import sageattention.core as sage_core + + sage_core.sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( + tensors["q_int8"], tensors["k_int8"], tensors["v_fp8_e4m3"], output, + tensors["q_scale_fp32"], tensors["k_scale_fp32"], tensors["v_scale_fp32"], + 0, 0, 2, 128 ** -0.5, 0, + ) + return output + + +def comparison(actual: torch.Tensor, expected: torch.Tensor, name: str) -> dict: + equal = torch.equal(actual, expected) + result = { + "name": name, "equal": equal, + "actual_sha256": tensor_sha256(actual), + "expected_sha256": tensor_sha256(expected), + "different_elements": int(torch.count_nonzero(actual != expected).item()), + } + if not equal: + mismatch = torch.nonzero(actual != expected, as_tuple=False)[0].cpu().tolist() + index = tuple(mismatch) + result["first_mismatch_index"] = mismatch + result["first_divergent_epoch"] = mismatch[0] if actual.ndim and actual.shape[0] == 3 else None + result["actual_first"] = actual[index].item() + result["expected_first"] = expected[index].item() + if actual.is_floating_point(): + delta = actual.float() - expected.float() + result["max_abs"] = float(delta.abs().max().item()) + return result + + +def analyze_overlap(clocks: torch.Tensor) -> dict: + host = clocks.cpu() + transitions = [] + for epoch in range(2): + pv_start = int(host[epoch, 4:8, 0].min().item()) + pv_end = int(host[epoch, 4:8, 1].max().item()) + qk_start = int(host[epoch + 1, 0:4, 0].min().item()) + qk_end = int(host[epoch + 1, 0:4, 1].max().item()) + overlap = max(0, min(pv_end, qk_end) - max(pv_start, qk_start)) + transitions.append({"pv_epoch": epoch, "qk_epoch": epoch + 1, "overlap_clocks": overlap, "positive": overlap > 0}) + return {"transitions": transitions, "all_positive": all(item["positive"] for item in transitions)} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--fixture-dir", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--build-dir", type=Path, required=True) + parser.add_argument("--build-log", type=Path, required=True) + parser.add_argument("--image", required=True) + parser.add_argument("--warmup", type=int, default=100) + parser.add_argument("--iterations", type=int, default=500) + parser.add_argument("--repeat", type=int, default=1000) + parser.add_argument("--timing-batch", type=int, default=1) + parser.add_argument("--capture", choices=("none", "vea", "sage_checkpoint"), default="none") + parser.add_argument("--cuda-profiler-capture", action="store_true") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + extension, build_seconds = build_extension(args.build_dir, args.build_log) + manifest, tensors = load_fixture(args.fixture_dir) + inputs = ( + tensors["q_int8"], tensors["k_int8"], tensors["v_fp8_e4m3"], + tensors["q_scale_fp32"], tensors["k_scale_fp32"], tensors["v_scale_fp32"], + ) + attributes_tensor = extension.attributes() + attributes = { + kernel: {field: int(value) for field, value in zip(ATTRIBUTE_NAMES, row.tolist(), strict=True)} + for kernel, row in zip(KERNEL_NAMES, attributes_tensor, strict=True) + } + + if args.capture != "none": + if args.cuda_profiler_capture: + torch.cuda.cudart().cudaProfilerStart() + if args.capture == "vea": + output = torch.empty_like(tensors["public_output_bf16"]) + captured = extension.vea_b_into(*inputs, output) + else: + captured = extension.sage_checkpoint(*inputs)[0] + torch.cuda.synchronize() + if args.cuda_profiler_capture: + torch.cuda.cudart().cudaProfilerStop() + report = { + "status": "capture_only", "capture": args.capture, + "output_sha256": tensor_sha256(captured), "attributes": attributes, + } + else: + reference = extension.sage_checkpoint(*inputs) + actual = extension.vea_b(*inputs, True) + torch.cuda.synchronize() + public_expected = tensors["public_output_bf16"] + parity = [comparison(reference[0], public_expected, "sage_checkpoint_vs_public_output")] + parity.extend(comparison(actual[index], reference[index], name) for index, name in enumerate(NAMES[:-1])) + first_failure = next((item for item in parity if not item["equal"]), None) + + repeat_output = torch.empty_like(public_expected) + for _ in range(args.repeat): + extension.vea_b_into(*inputs, repeat_output) + torch.cuda.synchronize() + repeat_parity = comparison(repeat_output, public_expected, "repeated_vea_output") + + vea_output = torch.empty_like(public_expected) + public_output = torch.empty_like(public_expected) + vea_timing = benchmark( + lambda: extension.vea_b_into(*inputs, vea_output), + args.warmup, args.iterations, args.timing_batch) + public_timing = benchmark( + lambda: public_launch(tensors, public_output), + args.warmup, args.iterations, args.timing_batch) + torch.cuda.synchronize() + report = { + "status": "passed" if first_failure is None and repeat_parity["equal"] else "failed", + "parity": parity, + "first_failure": first_failure, + "repeat": {"iterations": args.repeat, "parity": repeat_parity}, + "overlap": analyze_overlap(actual[9]), + "timing": { + "vea_b": vea_timing, + "public_sage2": public_timing, + "prototype_speedup": public_timing["p50_ms"] / vea_timing["p50_ms"], + }, + "attributes": attributes, + "advancement": { + "byte_exact_final": all(item["equal"] for item in parity if item["name"] in ("output", "sage_checkpoint_vs_public_output")), + "all_intermediates_exact": all(item["equal"] for item in parity), + "no_deadlock_repeat": repeat_parity["equal"], + "registers_at_most_200": attributes["vea_b_fast"]["registers_per_thread"] <= 200, + "local_bytes_zero": attributes["vea_b_fast"]["local_bytes_per_thread"] == 0, + "one_ten_warp_cta_resident": attributes["vea_b_fast"]["resident_ctas_per_sm"] >= 1, + "qk_pv_overlap": analyze_overlap(actual[9])["all_positive"], + "stable_timing": vea_timing["coefficient_of_variation"] <= 0.05, + }, + } + + report.update({ + "schema": "vortex-exact-phase2b-aligned-prototype", "version": 1, + "build_seconds": build_seconds, "fixture_manifest": str(args.fixture_dir / "manifest.json"), + "environment": { + "image": args.image, "gpu": torch.cuda.get_device_name(), + "compute_capability": list(torch.cuda.get_device_capability()), + "torch": torch.__version__, "cuda": torch.version.cuda, + "python": platform.python_version(), "argv": sys.argv, + "environment_switches": {key: value for key, value in os.environ.items() if key.startswith(("CUDA_", "TORCH_", "MAX_JOBS"))}, + }, + "sources": { + "cpp_sha256": hashlib.sha256((KERNELS / "vea_b_numeric.cpp").read_bytes()).hexdigest(), + "cuda_sha256": hashlib.sha256((KERNELS / "vea_b_numeric.cu").read_bytes()).hexdigest(), + }, + "claims": {"fixed_aligned_shape_only": True, "production_changed": False, "h3_integration": False}, + }) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2), flush=True) + if report.get("status") == "failed": + raise SystemExit(2) + + +if __name__ == "__main__": + main()