Validate Vortex attention capabilities
This commit is contained in:
Daniel Maddern
parent
ff8f7672b3
commit
c58e4c6d12
21 changed files with 1136 additions and 60 deletions
|
|
@ -322,10 +322,15 @@ the retained SageAttention 2.2.0 exactness contract at commit
|
|||
dispatch. Ten retained oracle artifacts match their recorded SHA-256 values.
|
||||
Phase 1 selects the VEA-B Q128 paired-owner architecture: four QK/softmax warps
|
||||
own `RS/RS_f8/m/d`, four separate PV warps own `RO`, and two producer warps own
|
||||
K/V staging. Its `180-207 ms` mainloop range is a heuristic screen, not achieved. VEA-A
|
||||
is the fallback and two-CTA DSM VEA-C is conditional on capability probes. No
|
||||
kernel exists. Missing self-contained canonical Q/K/V fixtures, checkpoint
|
||||
hash, and deployed Sage2 binary hash remain blockers. Production remains Sage2.
|
||||
K/V staging. Its `180-207 ms` mainloop range is a heuristic screen, not achieved.
|
||||
Phase 2A capability probes now pass with CUDA block-scope mbarriers: 96
|
||||
registers/thread for handoff, 139 for synthetic combined ownership, zero spills,
|
||||
one ten-warp CTA/SM, `0.076800 ms` barrier-only p50, and concurrent INT8/FP8
|
||||
progress in 48/48 blocks. Canonical Q/K/V/output fixtures are captured and
|
||||
reload-verified against the locked Sage2 output hash. The inline named-barrier
|
||||
primitive is rejected on racecheck. One isolated aligned-shape prototype is
|
||||
authorized; no attention kernel or attention speedup exists. Production remains
|
||||
Sage2.
|
||||
|
||||
The four-GEMM roofline selected FC2, and the guarded public split-K-1 schedule
|
||||
closed that target. Fresh NCU values at block 24 are `26.38 ms` QKV, `8.48 ms`
|
||||
|
|
|
|||
|
|
@ -517,9 +517,13 @@ P0-P3. The exact contract, fixture hashes, proof limits, and prior-work exclusio
|
|||
matrix are under `research/vortex_exact_attention/`. Phase 1 recommends VEA-B,
|
||||
a Q128 ten-warp pipeline with permanent disjoint QK/softmax and PV/output owners.
|
||||
Its `180-207 ms` mainloop and `205-235 ms` complete-attention ranges are
|
||||
heuristic screening models, not measured results. Register allocation, shared capacity, named
|
||||
barriers, and INT8/FP8 overlap must pass capability probes before kernel work.
|
||||
Production dispatch remains unchanged.
|
||||
heuristic screening models, not measured results. Phase 2A capability probes
|
||||
pass: selected block-scope mbarriers use 96 registers/thread and 51.2 KiB shared
|
||||
memory, permit one ten-warp CTA/SM, report zero spills and sanitizer errors, and
|
||||
measure `0.076800 ms` barrier-only p50 across 591 epochs. Synthetic INT8/FP8 role
|
||||
clocks overlap in 48/48 blocks. Canonical fixtures are captured and byte-locked.
|
||||
This authorizes only an isolated aligned-shape prototype; production dispatch
|
||||
remains unchanged and no attention speedup is claimed.
|
||||
|
||||
Materializing the optional BF16 LoRA input while these fused producers pack
|
||||
NVFP4 was also tested against the official Turbo-4 adapter. The isolated
|
||||
|
|
|
|||
|
|
@ -94,6 +94,9 @@ missing rather than inferred.
|
|||
| `gb10-post-fc2-production-profile-20260826` | Authoritative GB10 post-FC2 production baseline and profile | `production_accepted` | Exact `255.447 s` resident median and fresh NSYS/NCU ranking establish Sage2 as the dominant next investigation. | Post-FC2 baseline/profile summaries and raw reports | Matched SM120/SM100 profiles |
|
||||
| `vortex-exact-phase0-contract-import` | Vortex Exact Attention Phase 0 contract import | `research_retained` | Imported and verified retained exactness evidence without repeating closed Sage work; no kernel or dispatch exists. | `research/vortex_exact_attention/NUMERICAL_CONTRACT.md`; reference inventory | Self-contained canonical inputs, checkpoint and deployed binary hashes |
|
||||
| `vortex-exact-phase1-architecture-decision` | Vortex Exact Attention Phase 1 architecture decision | `research_retained` | Three structural ownership designs modeled; VEA-B recommended subject to capability probes. No speedup is claimed. | `research/vortex_exact_attention/DESIGN.md`; architecture model | Compiled resources, barrier/overlap probes, executable parity and timing |
|
||||
| `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 |
|
||||
|
||||
## Interpretation Notes
|
||||
|
||||
|
|
|
|||
|
|
@ -861,6 +861,83 @@
|
|||
"production_behavior": "No kernel and no dispatch; production Sage2 remains unchanged.",
|
||||
"source_recovery": "Design documents, machine-readable model, validator, and test are retained under research/vortex_exact_attention/."
|
||||
},
|
||||
{
|
||||
"id": "vortex-exact-phase2a-invalid-slot-probe",
|
||||
"name": "VEA-B invalid initial slot-indexing probe",
|
||||
"family": "vortex-exact-attention",
|
||||
"status": "incomplete",
|
||||
"hypothesis": "A two-slot 50 KiB shared-memory ring can preserve deterministic 591-epoch VEA-B producer-consumer order.",
|
||||
"implementation_strategy": "Compile and execute the first isolated ten-warp handoff and role-state probe on SM121.",
|
||||
"source_locations": ["research/vortex_exact_attention/EXPERIMENT_LOG.md", "research/vortex_exact_attention/PHASE2A_CAPABILITY_REPORT.md"],
|
||||
"active_source_location": null,
|
||||
"commit_hash": null,
|
||||
"benchmark_artifacts": [{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/probes/vea-b-capability-20260826.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 2A probe before slot-index correction"},
|
||||
"metrics": {"deterministic_repeat": false, "reported_inline_barrier_p50_ms": 16.631, "measurement_valid": false},
|
||||
"correctness_evidence": ["Deterministic repetition failed and exposed that two reserved 8 KiB K/V slots were indexed as one 16 KiB slot."],
|
||||
"decision_rationale": "Reject all payload and synchronization conclusions from this run. The corrected source and authoritative rerun supersede it.",
|
||||
"reproducer_commands": [],
|
||||
"timestamp": "2026-08-26",
|
||||
"evidence_missing": ["Valid two-slot indexing"],
|
||||
"production_behavior": "No attention kernel or dispatch integration.",
|
||||
"source_recovery": "The invalid report is retained; the corrected source is in research/vortex_exact_attention/kernels/."
|
||||
},
|
||||
{
|
||||
"id": "vortex-exact-phase2a-inline-named-barrier",
|
||||
"name": "VEA-B inline named-barrier primitive",
|
||||
"family": "vortex-exact-attention",
|
||||
"status": "architecture_rejected",
|
||||
"hypothesis": "Split arrive/sync inline PTX named barriers provide a valid bounded VEA-B handoff with lower overhead than block-scope mbarriers.",
|
||||
"implementation_strategy": "Run corrected 591-epoch two-slot payload and barrier probes, then require clean memcheck and racecheck.",
|
||||
"source_locations": ["research/vortex_exact_attention/kernels/vea_b_probe.cu", "research/vortex_exact_attention/PHASE2A_CAPABILITY_REPORT.md"],
|
||||
"active_source_location": null,
|
||||
"commit_hash": null,
|
||||
"benchmark_artifacts": [{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/probes/vea-b-capability-20260826-authoritative.json", "exists": true}],
|
||||
"profiler_artifacts": [{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/sanitizer/handoff-racecheck-20260826.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 corrected Phase 2A probe"},
|
||||
"metrics": {"barrier_p50_ms": 0.05612799897789955, "barrier_p95_ms": 0.05641600117087364, "budget_ms": 11.85445, "memcheck_errors": 0, "racecheck_hazards": 5},
|
||||
"correctness_evidence": ["Payload checksum and deterministic repetition pass, but Compute Sanitizer reports five shared-memory race hazards."],
|
||||
"decision_rationale": "Reject the primitive on the explicit zero-race-report gate. Select sanitizer-clean CUDA block-scope mbarriers instead.",
|
||||
"reproducer_commands": [],
|
||||
"timestamp": "2026-08-26",
|
||||
"evidence_missing": [],
|
||||
"production_behavior": "Not integrated.",
|
||||
"source_recovery": "The comparison implementation remains in the isolated capability probe only."
|
||||
},
|
||||
{
|
||||
"id": "vortex-exact-phase2a-capability-decision",
|
||||
"name": "VEA-B Phase 2A capability decision",
|
||||
"family": "vortex-exact-attention",
|
||||
"status": "research_retained",
|
||||
"hypothesis": "VEA-B can allocate disjoint QK and PV role state, sustain one ten-warp CTA/SM, perform a bounded sanitizer-clean 591-epoch handoff, and issue INT8 and FP8 MMA work concurrently on SM121.",
|
||||
"implementation_strategy": "Compile isolated role, handoff, and tensor-issue probes; measure 48-block 100-sample runs; run ptxas, Compute Sanitizer, and NCU; capture self-contained canonical Sage2 fixtures.",
|
||||
"source_locations": ["research/vortex_exact_attention/PHASE2A_CAPABILITY_REPORT.md", "research/vortex_exact_attention/kernels/vea_b_probe.cpp", "research/vortex_exact_attention/kernels/vea_b_probe.cu", "research/vortex_exact_attention/tools/run_capability_probes.py", "research/vortex_exact_attention/tools/capture_canonical_fixtures.py", "research/vortex_exact_attention/benchmarks/architecture_model.json"],
|
||||
"active_source_location": null,
|
||||
"commit_hash": null,
|
||||
"benchmark_artifacts": [
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/probes/vea-b-capability-20260826-authoritative.json", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/fixtures/canonical-20260826/manifest.json", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/build/vea-b-build-20260826-authoritative.log", "exists": true}
|
||||
],
|
||||
"profiler_artifacts": [
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/ncu/mbarrier-payload-20260826.ncu-rep", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/ncu/tensor-overlap-20260826.ncu-rep", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/sanitizer/mbarrier-payload-memcheck-20260826.log", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/sanitizer/mbarrier-payload-racecheck-20260826.log", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/sanitizer/tensor-memcheck-20260826.log", "exists": true},
|
||||
{"path": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/sanitizer/tensor-racecheck-20260826.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 Phase 2A capability probes; production dispatch unchanged"},
|
||||
"metrics": {"selected_handoff": "cuda_block_scope_mbarrier", "handoff_registers_per_thread": 96, "qk_registers_per_thread": 54, "pv_registers_per_thread": 138, "combined_registers_per_thread": 139, "local_bytes_per_thread": 0, "resident_ctas_per_sm": 1, "mbarrier_p50_ms": 0.07680000364780426, "mbarrier_p95_ms": 0.07887999713420868, "payload_p50_ms": 1.8749439716339111, "publication_errors": 0, "positive_overlap_blocks": 48, "overlap_blocks": 48, "canonical_output_sha256": "4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574", "achieved_attention_speedup": null},
|
||||
"correctness_evidence": ["Selected mbarrier payload returns [591, 4671090, 0, 128] and repeats deterministically.", "Selected handoff and tensor probes each report zero memcheck errors and zero racecheck hazards.", "Captured canonical Q/K/V/output tensors reload exactly and the Sage2 output matches the locked SHA-256."],
|
||||
"decision_rationale": "VEA-B passes Phase 2A and advances only to one isolated aligned-shape exact-attention prototype. Capability evidence is not an attention latency or production claim.",
|
||||
"reproducer_commands": ["python research/vortex_exact_attention/tools/validate_models.py", "python -m unittest discover -s research/vortex_exact_attention/tests -p test_*.py"],
|
||||
"timestamp": "2026-08-26",
|
||||
"evidence_missing": ["Executable aligned-shape exact attention parity", "Complete-kernel resource allocation", "Complete-kernel latency below 220 ms", "Canonical and trajectory integration"],
|
||||
"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": "gb10-post-fc2-production-profile-20260826",
|
||||
"name": "Authoritative GB10 post-FC2 production baseline and profile",
|
||||
|
|
|
|||
|
|
@ -2,10 +2,11 @@
|
|||
|
||||
## Decision
|
||||
|
||||
Prototype VEA-B, the one-CTA Q128 paired-owner pipeline, after its capability
|
||||
probes pass. VEA-A is the fallback if 320-thread role allocation fails. VEA-C is
|
||||
the higher-risk second prototype only after DSM placement and handoff latency
|
||||
are measured. These are projected designs, not implemented or achieved results.
|
||||
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.
|
||||
|
||||
Every design is side-by-side research under `kernels/`. The production Sage2
|
||||
path remains the reference and fallback; no dispatch integration is authorized.
|
||||
|
|
@ -95,9 +96,14 @@ the largest modeled 176-register role is 56,320 registers/CTA; compilation must
|
|||
remain at or below 200 registers/thread with zero material spills. One CTA gives
|
||||
ten active warps (`20.83%` of the 48-warp ceiling).
|
||||
|
||||
Four independent producer/consumer named-barrier pairs avoid a CTA-wide barrier
|
||||
inside the loop. K/V TMA barriers are separate. Projected synchronization cost
|
||||
is 6-15 ms.
|
||||
Two ready and two free block-scope mbarriers provide role-asymmetric two-slot
|
||||
reuse: six producer warps arrive without waiting on ready, four consumer warps
|
||||
wait and consume, then the roles reverse on free. Completion is CTA-scoped per
|
||||
slot; it is not a four-pair barrier claim. K/V TMA barriers are separate. The
|
||||
corrected Phase 2A 48-block, 591-epoch probe measures `0.076800 ms` p50 and
|
||||
`0.078880 ms` p95 for barriers alone, far below the `11.85445 ms` budget. The
|
||||
inline-PTX named-barrier alternative measures `0.056128 ms` p50 but is rejected
|
||||
because racecheck reports five hazards; mbarrier reports none.
|
||||
|
||||
Projected mainloop screening range: `180-207 ms`; conservative complete
|
||||
attention range with `25-28 ms` non-overlapped exact preparation is
|
||||
|
|
@ -155,10 +161,10 @@ boundary rather than a temporal schedule rearrangement.
|
|||
|
||||
## Prototype Order And Stop Conditions
|
||||
|
||||
1. Run register-allocation, shared-memory, named-barrier, and dual-pipe overlap
|
||||
probes without model inference.
|
||||
2. Reject VEA-B if it spills, cannot allocate 50-52 KiB predictably, or cannot
|
||||
demonstrate disjoint-role overlap.
|
||||
1. Phase 2A register-allocation, shared-memory, handoff, sanitizer, and dual-pipe
|
||||
probes are complete; see `PHASE2A_CAPABILITY_REPORT.md`.
|
||||
2. Retain CUDA block-scope mbarriers; do not use the rejected inline named
|
||||
barrier primitive.
|
||||
3. Implement VEA-B only for one aligned short shape and exact D=128 arithmetic.
|
||||
4. Consider VEA-C only after DSM feasibility; use VEA-A if B fails role/resource
|
||||
allocation.
|
||||
|
|
|
|||
|
|
@ -68,3 +68,28 @@ Phase commits:
|
|||
|
||||
- Phase 0 contract import: `c502842`
|
||||
- Phase 1 architecture decision: `b842bad`
|
||||
|
||||
## 2026-08-26: Phase 2A VEA-B Capability Decision
|
||||
|
||||
Status: `research_retained`
|
||||
|
||||
The corrected 320-thread VEA-B probes compile for SM121 with 96 registers/thread
|
||||
for handoff, 54 for QK role state, 138 for PV role state, and 139 for synthetic
|
||||
combined ownership. Ptxas and NCU report no spills, and 51.2 KiB shared memory
|
||||
permits one ten-warp CTA/SM.
|
||||
|
||||
The selected CUDA block-scope mbarrier handoff completes all 591 epochs with
|
||||
checksum `4,671,090`, zero publication errors, deterministic repetition,
|
||||
`0.076800 ms` barrier-only p50, and clean memcheck/racecheck. The faster inline
|
||||
named-barrier path is rejected because racecheck reports five hazards. INT8 and
|
||||
FP8 MMA role clocks overlap in 48/48 blocks. Canonical Q/K/V/output fixtures are
|
||||
self-contained and reload-verified; the output matches the locked Sage2 SHA-256.
|
||||
|
||||
The first full probe is invalid for payload conclusions because K/V slot bytes
|
||||
were indexed as one 16 KiB region instead of two 8 KiB regions. Deterministic
|
||||
repetition detected the bug. Corrected indexing passes checksum and sanitizer
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -41,24 +41,23 @@ handoff Blackwell-specific. FP8/FP4 QK would define `vortex_fast`, not
|
|||
`wgmma.wait_group` for `sm_121a`.
|
||||
- CUTLASS 4.6 executes SM120-family block-scaled narrow MMA on GB10, but that is
|
||||
not an exact replacement for Sage2 INT8 QK.
|
||||
- TMA producer/consumer scheduling is available, but role-specific register
|
||||
redistribution, named-barrier cost, and concurrent INT8/FP8 issue must be
|
||||
measured rather than assumed.
|
||||
- Phase 2A measures VEA-B role resources, one-CTA residency, bounded mbarrier
|
||||
handoff, and concurrent INT8/FP8 issue. TMA integration remains a complete
|
||||
prototype concern rather than a measured capability claim.
|
||||
|
||||
## Required Capability Probes Before Prototype
|
||||
## Capability Probe Status
|
||||
|
||||
1. Compile empty 256- and 320-thread role-specialized CTAs; record per-thread
|
||||
registers, aggregate register allocation, occupancy, and spill traffic.
|
||||
2. Test whether any supported register-allocation control can give producer and
|
||||
consumer warps different practical budgets on SM121.
|
||||
3. Measure one-CTA residency at 48, 52, 56, 64, and 68 KiB shared memory.
|
||||
4. Measure TMA plus named-barrier two-slot handoff for 4 and 8 KiB payloads.
|
||||
5. Verify simultaneous INT8 QK and FP8 PV warps produce actual overlap in pipe
|
||||
counters, not merely alternating aggregate utilization.
|
||||
6. Establish thread-block cluster/DSM support, placement, simultaneous
|
||||
residency, and remote shared-memory latency on GB10 before accepting VEA-C.
|
||||
7. Record sustained clocks and power behavior for 128-, 256-, and 320-thread
|
||||
CTAs.
|
||||
1. Complete: 320-thread QK, PV, combined-ownership, handoff, and tensor-issue
|
||||
resource probes with zero local spills.
|
||||
2. Complete: one 51.2 KiB ten-warp CTA/SM, confirmed by occupancy API and NCU.
|
||||
3. Complete: corrected two-slot 591-epoch mbarrier handoff, deterministic and
|
||||
sanitizer-clean, with 48-block p50/p95 measurements.
|
||||
4. Complete: simultaneous INT8 and FP8 warp progress in 48/48 blocks, plus an
|
||||
NCU tensor-issue capture.
|
||||
5. Remaining for the aligned prototype: integrated TMA staging, bank-aware
|
||||
payload layout, complete-kernel resources, sustained clocks, and exact output.
|
||||
6. Remaining before VEA-C: cluster/DSM support, placement, simultaneous
|
||||
residency, and remote shared-memory latency on GB10.
|
||||
|
||||
## Portability Policy
|
||||
|
||||
|
|
|
|||
|
|
@ -141,7 +141,7 @@ invalid V rows are zero; invalid Q rows suppress stores. Tail cost is below
|
|||
| Random short and ragged shapes | P1/P2/P3 retained JSON | Self-contained tensors plus exceptional values |
|
||||
| Prepared Q/K bytes and scales | P1 parity artifacts | Stored SHA-256 per tensor |
|
||||
| FP8 V bytes and scales | P2 randomized/repeat artifacts | Stored SHA-256 per tensor |
|
||||
| Block 24 attention | SHA `4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574` | Captured input Q/K/V and metadata |
|
||||
| Block 24 attention | SHA `4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574` | Captured and reload-verified Q/K/V/output plus metadata |
|
||||
| Blocks 0/24/49 | Strided-NHD and P1 artifacts | Self-contained boundary tensors |
|
||||
| Two-step latents | Strided-NHD contract | Re-capture under frozen environment |
|
||||
| Canonical 12-step latents | video/audio SHA below | Re-capture inputs and full environment |
|
||||
|
|
@ -155,13 +155,18 @@ audio 852005383770480a6503504e1ffec86dd1fb63a69c6400f92da18e39e0986de2
|
|||
|
||||
## Proof Limits And Blockers
|
||||
|
||||
- No self-contained canonical Q/K/V tensor fixture is retained.
|
||||
- No checkpoint SHA-256 is retained.
|
||||
- The deployed Sage2 extension binary hash is not retained.
|
||||
- Canonical Q/K/V/output tensors are retained in durable Phase 2A storage; all
|
||||
reload exactly and the output matches the locked Sage2 hash.
|
||||
- The checkpoint SHA-256 is
|
||||
`72fa9269ce551fb63ff42a32d9b46d0c122e84b4b2c511e22fa698287b088f70`.
|
||||
- The deployed Sage2 mainloop and fused-preparation binary SHA-256 values are
|
||||
`ea084d6a252524fb40ab5ede15583a22e87ec783bb9ec25119b87283141c6397`
|
||||
and `8139fa8ed9900c76fde4a76bb2d9f5b487b7196129df152ea7d091138a96318a`.
|
||||
- The exact production K-mean reduction tree and Comfy RMSNorm/RoPE binary are
|
||||
not source-frozen in this repository.
|
||||
- NaN, infinity, denormal, signed-zero, and tie behavior is not exhaustively
|
||||
proven by existing fixtures.
|
||||
|
||||
These gaps do not invalidate retained decisions, but they block executable
|
||||
prototype acceptance. They must be captured before comparing a new kernel.
|
||||
The remaining exceptional-value and source-freezing gaps do not invalidate the
|
||||
Phase 2A capability decision. They remain gates for executable attention parity,
|
||||
not for the isolated resource and handoff probes.
|
||||
|
|
|
|||
|
|
@ -1,7 +1,8 @@
|
|||
# Performance Model
|
||||
|
||||
All candidate values are projections. No Vortex kernel has been implemented or
|
||||
timed.
|
||||
Candidate attention latency values remain projections. Phase 2A capability
|
||||
kernels have measured resources, handoff cost, and concurrent issue, but no
|
||||
Vortex attention kernel has been implemented or timed.
|
||||
|
||||
## Baseline
|
||||
|
||||
|
|
@ -44,11 +45,18 @@ material local spilling rejects it.
|
|||
| Non-overlapped exact preparation | 25-28 ms | 25-28 ms | 25-28 ms | P1 entry 12.33 + K mean 2.37 + K quant 3.83 + P2 V 6.39 = 24.92 ms, plus margin |
|
||||
| Projected complete attention | `217-250 ms` | `205-235 ms` | `199-238 ms` | not achieved |
|
||||
|
||||
Phase 2A replaces the VEA-B synchronization assumption with capability evidence:
|
||||
the selected 48-block mbarrier path measures `0.076800 ms` p50 and `0.078880 ms`
|
||||
p95 for 591 barrier epochs. Its synthetic payload path measures `1.874944 ms`
|
||||
p50. These values clear the `11.85445 ms` budget but must not be substituted into
|
||||
the attention projection as if they were complete QK/softmax/PV work.
|
||||
|
||||
The screening range starts from `237.089 ms`, assumes a candidate-specific
|
||||
fraction of the measured `63.53%` no-eligible region becomes useful overlap,
|
||||
then adds the handoff budget. Scheduler percentages do not translate linearly
|
||||
to time, so this is a rejection model rather than a performance prediction.
|
||||
Capability microbenchmarks must replace every assumption before implementation.
|
||||
The remaining assumptions concern the complete exact attention implementation
|
||||
and its latency, not VEA-B's basic SM121 capability.
|
||||
|
||||
## Rejection Gates
|
||||
|
||||
|
|
|
|||
100
research/vortex_exact_attention/PHASE2A_CAPABILITY_REPORT.md
Normal file
100
research/vortex_exact_attention/PHASE2A_CAPABILITY_REPORT.md
Normal file
|
|
@ -0,0 +1,100 @@
|
|||
# Phase 2A VEA-B Capability Report
|
||||
|
||||
Date: 2026-08-26
|
||||
|
||||
Status: `advance_to_isolated_aligned_prototype`
|
||||
|
||||
This phase measures only the VEA-B ownership, resource, handoff, and concurrent
|
||||
issue capabilities. It does not implement attention, measure attention latency,
|
||||
or authorize production dispatch.
|
||||
|
||||
## Decision
|
||||
|
||||
VEA-B passes the Phase 2A capability gates with CUDA block-scope mbarriers as
|
||||
the two-slot handoff primitive. Proceed to one isolated aligned-shape attention
|
||||
prototype. Do not add `H3_ATTENTION=vortex_exact` or alter the Sage2 fallback.
|
||||
|
||||
The inline-PTX named-barrier primitive is rejected for the prototype. It is
|
||||
faster in the synthetic barrier-only probe, but Compute Sanitizer racecheck
|
||||
reports five shared-memory hazards because that path does not provide a
|
||||
sanitizer-recognized happens-before edge. The mbarrier path reports zero hazards
|
||||
and remains far below the synchronization budget.
|
||||
|
||||
## Authoritative Probe
|
||||
|
||||
Environment: NVIDIA GB10, SM121, CUDA 13.0, PyTorch 2.9.1+cu130, production
|
||||
image `sha256:6d880d628334c981c3d155bf5244e65e26e22cc9273c80145f646eee3c3698c2`.
|
||||
The run uses 48 blocks, 320 threads, 591 epochs, 20 warmups, and 100 measured
|
||||
iterations.
|
||||
|
||||
| Gate | Measurement | Result |
|
||||
| --- | ---: | --- |
|
||||
| Selected handoff registers | 96/thread | pass, limit 200 |
|
||||
| QK role registers | 54/thread | pass |
|
||||
| PV role registers | 138/thread | pass |
|
||||
| Combined ownership registers | 139/thread | pass |
|
||||
| Local memory / ptxas spills | 0 / 0 | pass |
|
||||
| Selected dynamic shared memory | 51,200 bytes | pass |
|
||||
| Resident ten-warp CTAs/SM | 1 | pass |
|
||||
| mbarrier-only p50 / p95 | `0.076800 / 0.078880 ms` | pass, budget `11.85445 ms` |
|
||||
| Synthetic payload p50 / p95 | `1.874944 / 1.879474 ms` | pass |
|
||||
| 591-epoch payload checksum | `4,671,090` | pass, zero publication errors |
|
||||
| Deterministic repeat | identical | pass |
|
||||
| INT8/FP8 positive clock overlap | 48/48 blocks | pass |
|
||||
| INT8/FP8 overlap p50 ratio | `0.9999929` | capability pass |
|
||||
| Combined issue p50 | `0.127072 ms` | below `0.224848 ms` serial p50 sum |
|
||||
|
||||
The overlap probe establishes that distinct VEA-B warp roles can make progress
|
||||
through INT8 and FP8 `mma.sync` loops concurrently. It is not a throughput model
|
||||
for the complete attention mainloop.
|
||||
|
||||
## Sanitizer And NCU
|
||||
|
||||
- Selected mbarrier payload: memcheck 0 errors; racecheck 0 hazards.
|
||||
- Tensor issue probe: memcheck 0 errors; racecheck 0 hazards.
|
||||
- Rejected inline named barrier: memcheck 0 errors; racecheck 5 hazards.
|
||||
- NCU selected handoff: 96 registers/thread, 51.2 KiB dynamic shared memory,
|
||||
shared-memory block limit 1, `20.91%` achieved occupancy, and zero local spill
|
||||
requests.
|
||||
- NCU tensor issue: 26 registers/thread and zero local spill requests.
|
||||
|
||||
The synthetic handoff payload has severe shared-load bank conflicts in the
|
||||
single-block NCU capture. Its deliberately simple byte access pattern is not an
|
||||
attention layout result, but the aligned prototype must choose and profile a
|
||||
bank-aware score/K/V layout rather than copy this access pattern unchanged.
|
||||
|
||||
## Canonical Fixtures
|
||||
|
||||
Self-contained BF16 NHD Q, K, V, and Sage2 output tensors were captured at
|
||||
`[1, 37810, 56, 128]`, saved, reloaded, and compared byte-for-byte. The output
|
||||
tensor SHA-256 is
|
||||
`4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574`,
|
||||
matching the locked public Sage2 oracle. Checkpoint and deployed Sage binary
|
||||
hashes also match the numerical contract.
|
||||
|
||||
## Invalid And Rejected Runs
|
||||
|
||||
The first full probe is invalid for payload conclusions. K/V allocation reserved
|
||||
two 8 KiB slots but indexed one 16 KiB region, allowing the next epoch to
|
||||
overwrite consumer data. Deterministic repetition exposed the defect; the
|
||||
corrected probe uses slot-relative 8 KiB indexing and passes deterministic,
|
||||
checksum, memcheck, and racecheck validation.
|
||||
|
||||
The original `16.631 ms` inline-barrier number from that invalid run is not
|
||||
retained as a hardware measurement. The corrected authoritative inline path is
|
||||
`0.056128 ms` p50, but remains rejected on the sanitizer gate.
|
||||
|
||||
## Durable Evidence
|
||||
|
||||
- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/probes/vea-b-capability-20260826-authoritative.json`
|
||||
- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/build/vea-b-build-20260826-authoritative.log`
|
||||
- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/fixtures/canonical-20260826/manifest.json`
|
||||
- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/sanitizer/`
|
||||
- `/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/ncu/`
|
||||
|
||||
## Next Gate
|
||||
|
||||
Implement only one aligned short-shape VEA-B prototype with the selected
|
||||
mbarrier handoff and exact D=128 arithmetic. It must compare against the captured
|
||||
prepared-byte fixtures before any canonical or integrated timing. Phase 2A does
|
||||
not validate the heuristic `180-207 ms` mainloop screen.
|
||||
17
research/vortex_exact_attention/PROVENANCE.md
Normal file
17
research/vortex_exact_attention/PROVENANCE.md
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
# Probe Provenance
|
||||
|
||||
The Phase 2A probes are original isolated capability microkernels. They do not
|
||||
contain the SageAttention mainloop or implement attention.
|
||||
|
||||
The inline MMA operand signatures are derived from SageAttention 2.2.0
|
||||
`csrc/mma.cuh` at commit `d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5`.
|
||||
That file is Apache-2.0 and itself records adaptation from FlashInfer 0.1.5.
|
||||
The local probe retains the required attribution in its source header.
|
||||
|
||||
Reference:
|
||||
|
||||
- https://github.com/thu-ml/SageAttention/blob/d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5/csrc/mma.cuh
|
||||
- https://github.com/thu-ml/SageAttention/blob/d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5/LICENSE
|
||||
|
||||
No CUTLASS, SageAttention, FlashAttention, or Sol source is copied into the
|
||||
Vortex backend. Production remains unchanged.
|
||||
|
|
@ -26,7 +26,9 @@ Such work belongs in a future `vortex_fast` backend with a separate contract.
|
|||
|
||||
- Phase 0: imported and verified retained Sage2 exactness specification.
|
||||
- Phase 1: design decision complete; VEA-B is the recommended prototype.
|
||||
- Kernel implementation: not started.
|
||||
- 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.
|
||||
- Production dispatch: unchanged; no `vortex_exact` route exists.
|
||||
|
||||
## Authoritative Baseline
|
||||
|
|
@ -65,11 +67,13 @@ p95, block-level, and resident-service gates.
|
|||
|
||||
- `NUMERICAL_CONTRACT.md`: imported exact Sage2 semantics and proof limits.
|
||||
- `DESIGN.md`: state-ownership architectures and selected prototype.
|
||||
- `PHASE2A_CAPABILITY_REPORT.md`: measured resource, handoff, overlap, sanitizer,
|
||||
NCU, and fixture 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/`: future isolated implementation; currently contains no kernel.
|
||||
- `kernels/`: isolated capability probes; no attention kernel.
|
||||
- `tests/`: future parity and sanitizer harnesses.
|
||||
- `tools/`: project-local verification/modeling utilities.
|
||||
|
||||
|
|
@ -78,7 +82,11 @@ warps own `RS/RS_f8/m/d`, four distinct PV warps own `RO`, and two producer
|
|||
warps own K/V staging. This differs structurally from P3 because complete state
|
||||
is transferred or partitioned rather than retained in every compute warp. Its
|
||||
`180-207 ms` mainloop range is a heuristic screening model, not a measured
|
||||
result; physical register allocation and overlap remain unverified.
|
||||
result. Phase 2A measures 96 registers/thread for the selected mbarrier handoff,
|
||||
139 for the synthetic combined-ownership probe, zero spills, one ten-warp CTA/SM,
|
||||
`0.076800 ms` barrier-only p50, and positive INT8/FP8 clock overlap in 48/48
|
||||
blocks. These measurements advance VEA-B to an isolated prototype, not to a
|
||||
performance or production claim.
|
||||
|
||||
## Provenance
|
||||
|
||||
|
|
|
|||
|
|
@ -22,7 +22,33 @@
|
|||
"formula": "baseline minus assumed recovered fraction of no-eligible region plus handoff budget",
|
||||
"exact_preparation_components_ms": {"p1_entry": 12.33, "k_mean": 2.37, "k_quant": 3.83, "p2_v": 6.39},
|
||||
"exact_preparation_ms": [25, 28],
|
||||
"unresolved": ["compiled registers", "named barrier latency", "INT8_FP8 overlap", "DSM placement and latency"]
|
||||
"unresolved": ["full attention kernel resources and latency", "DSM placement and latency"]
|
||||
},
|
||||
"phase2a_capability_measurements": {
|
||||
"status": "passed_for_isolated_aligned_prototype",
|
||||
"date": "2026-08-26",
|
||||
"candidate": "VEA-B",
|
||||
"geometry": {"threads": 320, "warps": 10, "epochs": 591, "dynamic_shared_bytes": 51200},
|
||||
"selected_handoff": "cuda_block_scope_mbarrier",
|
||||
"registers_per_thread": {"handoff": 96, "qk_role": 54, "pv_role": 138, "combined_ownership": 139},
|
||||
"local_bytes_per_thread": 0,
|
||||
"ptxas_spill_loads": 0,
|
||||
"ptxas_spill_stores": 0,
|
||||
"resident_ctas_per_sm": 1,
|
||||
"timing_ms": {
|
||||
"mbarrier_only_p50": 0.07680000364780426,
|
||||
"mbarrier_only_p95": 0.07887999713420868,
|
||||
"payload_p50": 1.8749439716339111,
|
||||
"payload_p95": 1.8794736385345459,
|
||||
"budget": 11.85445
|
||||
},
|
||||
"handoff_validation": {"checksum": 4671090, "publication_errors": 0, "deterministic_repeat": true},
|
||||
"tensor_issue": {"positive_overlap_blocks": 48, "total_blocks": 48, "p50_overlap_ratio": 0.999992875970649, "combined_p50_ms": 0.12707199901342392},
|
||||
"sanitizer": {"mbarrier_memcheck_errors": 0, "mbarrier_racecheck_hazards": 0, "tensor_memcheck_errors": 0, "tensor_racecheck_hazards": 0},
|
||||
"rejected_primitive": {"name": "inline_ptx_named_barrier", "racecheck_hazards": 5},
|
||||
"canonical_fixture": {"shape": [1, 37810, 56, 128], "dtype": "bfloat16", "output_sha256": "4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574", "reload_verified": true},
|
||||
"attention_kernel_implemented": false,
|
||||
"attention_latency_measured": false
|
||||
},
|
||||
"candidates": [
|
||||
{
|
||||
|
|
@ -43,14 +69,14 @@
|
|||
"id": "VEA-B", "name": "Q128 paired QK/PV owner pipeline", "recommended_order": 1,
|
||||
"geometry": {"q": 128, "k": 64, "d": 128, "threads": 320, "warps": 10},
|
||||
"ownership": {"qk_fragments": "warps_2_5", "m_d": "warps_2_5", "output_accumulation": "warps_6_9", "load_control": "warps_0_1"},
|
||||
"handoff": {"mechanism": "four_paired_two_slot_shared_rings_named_barriers", "payload": "fp8_probability_bytes_plus_fp32_ro_rescale_and_final_reciprocal", "bounded_tile_local": true, "consumer_order": "strict_epoch_order", "qk_and_pv_fragments_simultaneously_live": true, "same_warp": false},
|
||||
"handoff": {"mechanism": "two_slot_role_asymmetric_cuda_block_scope_mbarriers", "payload": "fp8_probability_bytes_plus_fp32_ro_rescale_and_final_reciprocal", "bounded_tile_local": true, "consumer_order": "strict_epoch_order", "qk_and_pv_fragments_simultaneously_live": true, "same_warp": false},
|
||||
"registers_per_thread": {"k_producer": [32, 48], "v_producer": [40, 56], "qk_softmax": [96, 128], "pv_output": [144, 176]},
|
||||
"logical_role_weighted_registers": 39000, "registers_per_cta_uniform_projection": 56320, "compiled_register_limit_per_thread": 200,
|
||||
"shared_kib": {"unaliased": [66, 68], "alias_saving": 16, "total": [50, 52]},
|
||||
"resident_ctas_per_sm_projected": 1, "synchronization_ms_projected": [6, 15],
|
||||
"recovered_no_eligible_fraction": [0.30, 0.42],
|
||||
"mainloop_ms_projected": [180, 207], "complete_attention_ms_projected": [205, 235],
|
||||
"capability_dependencies": ["compiled_registers_at_or_below_200", "zero_material_spills", "named_barrier_cost", "INT8_FP8_overlap"],
|
||||
"capability_dependencies": ["full_kernel_exactness", "full_kernel_compiled_resources", "full_kernel_latency"],
|
||||
"structural_difference": "QK/softmax and RO/PV have permanent disjoint owners; P1/P2 may feed slots but are not standalone boundaries"
|
||||
},
|
||||
{
|
||||
|
|
@ -67,5 +93,5 @@
|
|||
"structural_difference": "QK/softmax state and RO/PV state reside in separate CTAs and cross an explicit DSM byte handoff"
|
||||
}
|
||||
],
|
||||
"decision": {"recommended_prototype": "VEA-B", "fallback": "VEA-A", "conditional_second": "VEA-C", "kernel_implementation_started": false, "production_dispatch_changed": false}
|
||||
"decision": {"recommended_prototype": "VEA-B", "selected_handoff": "cuda_block_scope_mbarrier", "phase2a_capability_passed": true, "isolated_aligned_prototype_authorized": true, "fallback": "VEA-A", "conditional_second": "VEA-C", "kernel_implementation_started": false, "production_dispatch_changed": false}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"schema": "vortex-exact-reference-inventory",
|
||||
"version": 1,
|
||||
"version": 2,
|
||||
"verified_at": "2026-08-26",
|
||||
"reference": {
|
||||
"name": "SageAttention 2.2.0",
|
||||
|
|
@ -28,10 +28,20 @@
|
|||
{"path": "benchmarks/gb10-sage2-temporal-pair-short-20260825-p0.json", "size_bytes": 4348, "sha256": "e3c7b0f56125a358025184ca5c2b188d61999f8c3fdbd82fea41ee987b931f00"},
|
||||
{"path": "benchmarks/gb10-sage2-temporal-pair-timing-20260825-p0.json", "size_bytes": 5487, "sha256": "74b69b08873f31e98a0d40c501186e1f5853176c653381a1d4e65ea5d9a28552"}
|
||||
],
|
||||
"phase2a_canonical_fixture": {
|
||||
"manifest": "/home/daniel/StoryStudioAssets/H3-output/h3-blackwell-runtime/research/vortex_exact_attention/phase2a/fixtures/canonical-20260826/manifest.json",
|
||||
"shape": [1, 37810, 56, 128],
|
||||
"dtype": "bfloat16",
|
||||
"q_tensor_sha256": "1545e139b05b690e8ab3487dcce869cae74357479e42f8062e4734ce70f9d88a",
|
||||
"k_tensor_sha256": "67d383968d1dd299dc64023ce3edf622b9a84ad2b2a38d6595052bdffbd2f0a9",
|
||||
"v_tensor_sha256": "9a875d7210f328a1bce12e92ade95fcc93cd8b7f6446804308f1921d4ccc4f5b",
|
||||
"output_tensor_sha256": "4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574",
|
||||
"checkpoint_sha256": "72fa9269ce551fb63ff42a32d9b46d0c122e84b4b2c511e22fa698287b088f70",
|
||||
"sage_mainloop_binary_sha256": "ea084d6a252524fb40ab5ede15583a22e87ec783bb9ec25119b87283141c6397",
|
||||
"sage_fused_binary_sha256": "8139fa8ed9900c76fde4a76bb2d9f5b487b7196129df152ea7d091138a96318a",
|
||||
"reload_verified": true
|
||||
},
|
||||
"known_missing": [
|
||||
"deployed SageAttention extension binary hash",
|
||||
"checkpoint SHA-256",
|
||||
"self-contained canonical Q/K/V input tensors",
|
||||
"exact P0 candidate source patches",
|
||||
"complete historical P0 commands and environments"
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
# Kernels
|
||||
|
||||
No kernel exists. Phase 0 and Phase 1 are specification and design work only.
|
||||
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 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.
|
||||
|
|
|
|||
17
research/vortex_exact_attention/kernels/vea_b_probe.cpp
Normal file
17
research/vortex_exact_attention/kernels/vea_b_probe.cpp
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
torch::Tensor launch_handoff_probe(int64_t blocks, int64_t epochs, bool payload);
|
||||
torch::Tensor launch_mbarrier_handoff_probe(int64_t blocks, int64_t epochs, bool payload);
|
||||
std::vector<torch::Tensor> launch_tensor_issue_probe(int64_t blocks, int64_t iterations, int64_t mode);
|
||||
torch::Tensor launch_register_probe(int64_t blocks, int64_t iterations, int64_t role);
|
||||
torch::Tensor probe_kernel_attributes(int64_t dynamic_shared_bytes);
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
|
||||
module.def("handoff", &launch_handoff_probe, "VEA-B bounded handoff probe");
|
||||
module.def("handoff_mbarrier", &launch_mbarrier_handoff_probe, "VEA-B mbarrier handoff probe");
|
||||
module.def("tensor_issue", &launch_tensor_issue_probe, "VEA-B INT8/FP8 issue probe");
|
||||
module.def("register_probe", &launch_register_probe, "VEA-B role register probe");
|
||||
module.def("attributes", &probe_kernel_attributes, "VEA-B kernel attributes");
|
||||
}
|
||||
309
research/vortex_exact_attention/kernels/vea_b_probe.cu
Normal file
309
research/vortex_exact_attention/kernels/vea_b_probe.cu
Normal file
|
|
@ -0,0 +1,309 @@
|
|||
/*
|
||||
* Vortex Exact Attention isolated capability probes.
|
||||
*
|
||||
* The MMA operand signatures are adapted from SageAttention 2.2.0 csrc/mma.cuh
|
||||
* (commit d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5), Apache-2.0, which records
|
||||
* prior adaptation from FlashInfer 0.1.5. This file does not implement attention.
|
||||
*/
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda/barrier>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <array>
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
#define CUDA_CHECK(call) do { \
|
||||
cudaError_t error_ = (call); \
|
||||
TORCH_CHECK(error_ == cudaSuccess, cudaGetErrorString(error_)); \
|
||||
} while (0)
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kThreads = 320;
|
||||
constexpr int kProducerThreads = 192;
|
||||
constexpr int kSlotBytes = 8192;
|
||||
constexpr int kKvSlotBytes = 8192;
|
||||
constexpr int kDynamicSharedBytes = 50 * 1024;
|
||||
|
||||
__device__ __forceinline__ void ready_arrive(int slot) {
|
||||
if (slot == 0) asm volatile("bar.arrive 0, 320;" ::: "memory");
|
||||
else asm volatile("bar.arrive 1, 320;" ::: "memory");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void ready_sync(int slot) {
|
||||
if (slot == 0) asm volatile("bar.sync 0, 320;" ::: "memory");
|
||||
else asm volatile("bar.sync 1, 320;" ::: "memory");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void free_arrive(int slot) {
|
||||
if (slot == 0) asm volatile("bar.arrive 2, 320;" ::: "memory");
|
||||
else asm volatile("bar.arrive 3, 320;" ::: "memory");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void free_sync(int slot) {
|
||||
if (slot == 0) asm volatile("bar.sync 2, 320;" ::: "memory");
|
||||
else asm volatile("bar.sync 3, 320;" ::: "memory");
|
||||
}
|
||||
|
||||
template <bool UseMbarrier>
|
||||
__global__ __launch_bounds__(kThreads, 1)
|
||||
void vea_b_handoff_probe(int64_t* result, int epochs, bool payload) {
|
||||
extern __shared__ uint8_t shared[];
|
||||
uint8_t* scores = shared;
|
||||
uint8_t* k_slot = shared + 2 * kSlotBytes;
|
||||
uint8_t* v_slot = k_slot + 2 * kKvSlotBytes;
|
||||
using BlockBarrier = cuda::barrier<cuda::thread_scope_block>;
|
||||
BlockBarrier* barriers = reinterpret_cast<BlockBarrier*>(shared + 48 * 1024);
|
||||
const int tid = threadIdx.x;
|
||||
const bool producer = tid < kProducerThreads;
|
||||
int64_t checksum = 0;
|
||||
int completed_epoch = -1;
|
||||
|
||||
if constexpr (UseMbarrier) {
|
||||
if (tid < 4) init(barriers + tid, kThreads);
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (producer) {
|
||||
for (int epoch = 0; epoch < epochs; ++epoch) {
|
||||
const int slot = epoch & 1;
|
||||
if (epoch >= 2) {
|
||||
if constexpr (UseMbarrier) barriers[2 + slot].arrive_and_wait();
|
||||
else free_sync(slot);
|
||||
}
|
||||
if (payload) {
|
||||
if (tid < 64) {
|
||||
for (int index = tid; index < kKvSlotBytes; index += 64) {
|
||||
k_slot[slot * kKvSlotBytes + index] = static_cast<uint8_t>((epoch + index * 3) & 255);
|
||||
v_slot[slot * kKvSlotBytes + index] = static_cast<uint8_t>((epoch + index * 5) & 255);
|
||||
}
|
||||
} else {
|
||||
const int row = tid - 64;
|
||||
uint8_t* row_slot = scores + slot * kSlotBytes + row * 64;
|
||||
#pragma unroll
|
||||
for (int column = 0; column < 64; ++column) {
|
||||
row_slot[column] = static_cast<uint8_t>((epoch + row * 17 + column) & 255);
|
||||
}
|
||||
}
|
||||
}
|
||||
__threadfence_block();
|
||||
if constexpr (UseMbarrier) (void)barriers[slot].arrive();
|
||||
else ready_arrive(slot);
|
||||
completed_epoch = epoch;
|
||||
}
|
||||
if constexpr (UseMbarrier) barriers[2 + ((epochs - 1) & 1)].arrive_and_wait();
|
||||
else free_sync((epochs - 1) & 1);
|
||||
if (epochs > 1) {
|
||||
if constexpr (UseMbarrier) barriers[2 + ((epochs - 2) & 1)].arrive_and_wait();
|
||||
else free_sync((epochs - 2) & 1);
|
||||
}
|
||||
} else {
|
||||
const int consumer = tid - kProducerThreads;
|
||||
for (int epoch = 0; epoch < epochs; ++epoch) {
|
||||
const int slot = epoch & 1;
|
||||
if constexpr (UseMbarrier) barriers[slot].arrive_and_wait();
|
||||
else ready_sync(slot);
|
||||
if (payload) {
|
||||
const uint8_t* row_slot = scores + slot * kSlotBytes + consumer * 64;
|
||||
#pragma unroll
|
||||
for (int column = 0; column < 64; ++column) {
|
||||
const uint8_t expected = static_cast<uint8_t>((epoch + consumer * 17 + column) & 255);
|
||||
const uint8_t actual = row_slot[column];
|
||||
checksum += actual;
|
||||
if (actual != expected) atomicAdd(reinterpret_cast<unsigned long long*>(result + blockIdx.x * 4 + 2), 1ULL);
|
||||
}
|
||||
const int kv_index = consumer;
|
||||
const uint8_t expected_k = static_cast<uint8_t>((epoch + kv_index * 3) & 255);
|
||||
const uint8_t expected_v = static_cast<uint8_t>((epoch + kv_index * 5) & 255);
|
||||
const uint8_t actual_k = k_slot[slot * kKvSlotBytes + kv_index];
|
||||
const uint8_t actual_v = v_slot[slot * kKvSlotBytes + kv_index];
|
||||
if (actual_k != expected_k || actual_v != expected_v) {
|
||||
atomicAdd(reinterpret_cast<unsigned long long*>(result + blockIdx.x * 4 + 2), 1ULL);
|
||||
}
|
||||
checksum += actual_k + actual_v;
|
||||
}
|
||||
completed_epoch = epoch;
|
||||
__threadfence_block();
|
||||
if constexpr (UseMbarrier) (void)barriers[2 + slot].arrive();
|
||||
else free_arrive(slot);
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (tid == 0) result[blockIdx.x * 4] = completed_epoch + 1;
|
||||
if (tid == kProducerThreads) result[blockIdx.x * 4 + 1] = checksum;
|
||||
if (!producer && completed_epoch == epochs - 1) {
|
||||
atomicAdd(reinterpret_cast<unsigned long long*>(result + blockIdx.x * 4 + 3), 1ULL);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void int8_mma(int32_t* c, const uint32_t* a, const 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(uint32_t* c, const uint32_t* a, const 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]));
|
||||
}
|
||||
|
||||
__global__ __launch_bounds__(kThreads, 1)
|
||||
void vea_b_tensor_issue_probe(int64_t* values, int64_t* clocks, int iterations, int mode) {
|
||||
const int warp = threadIdx.x >> 5;
|
||||
const int lane = threadIdx.x & 31;
|
||||
const bool run_int = warp >= 2 && warp <= 5 && mode != 2;
|
||||
const bool run_fp = warp >= 6 && warp <= 9 && mode != 1;
|
||||
uint32_t a[4] = {0x01010101u + static_cast<uint32_t>(lane), 0x01010101u, 0x01010101u, 0x01010101u};
|
||||
uint32_t b[2] = {0x01010101u, 0x01010101u + static_cast<uint32_t>(lane)};
|
||||
int32_t ci[4] = {0, 0, 0, 0};
|
||||
uint32_t cf[2] = {0, 0};
|
||||
const uint64_t started = clock64();
|
||||
if (run_int) {
|
||||
for (int index = 0; index < iterations; ++index) int8_mma(ci, a, b);
|
||||
} else if (run_fp) {
|
||||
a[0] = a[1] = a[2] = a[3] = 0x38383838u;
|
||||
b[0] = b[1] = 0x38383838u;
|
||||
for (int index = 0; index < iterations; ++index) fp8_mma(cf, a, b);
|
||||
}
|
||||
const uint64_t finished = clock64();
|
||||
const int offset = (blockIdx.x * 10 + warp);
|
||||
if (lane == 0) {
|
||||
clocks[offset * 2] = static_cast<int64_t>(started);
|
||||
clocks[offset * 2 + 1] = static_cast<int64_t>(finished);
|
||||
values[offset] = run_int ? static_cast<int64_t>(ci[0]) : (run_fp ? static_cast<int64_t>(cf[0]) : 0);
|
||||
}
|
||||
}
|
||||
|
||||
template <int Role>
|
||||
__global__ __launch_bounds__(kThreads, 1)
|
||||
void vea_b_register_probe(float* output, int iterations) {
|
||||
const int warp = threadIdx.x >> 5;
|
||||
const int lane = threadIdx.x & 31;
|
||||
float total = 0.0f;
|
||||
if constexpr (Role == 0 || Role == 2) {
|
||||
if (warp >= 2 && warp <= 5) {
|
||||
int32_t rs[32];
|
||||
float m[8], d[8];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) rs[i] = i + lane;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) { m[i] = i + lane * 0.01f; d[i] = 1.0f + i; }
|
||||
for (int step = 0; step < iterations; ++step) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) rs[i] += (step + i) & 7;
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) { m[i] = fmaxf(m[i], step * 0.001f); d[i] = fmaf(d[i], 0.9999f, m[i]); }
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 32; ++i) total += static_cast<float>(rs[i]);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 8; ++i) total += m[i] + d[i];
|
||||
}
|
||||
}
|
||||
if constexpr (Role == 1 || Role == 2) {
|
||||
if (warp >= 6 && warp <= 9) {
|
||||
float ro[128];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 128; ++i) ro[i] = i * 0.001f + lane;
|
||||
for (int step = 0; step < iterations; ++step) {
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 128; ++i) ro[i] = fmaf(ro[i], 0.99999f, (step + i) * 0.00001f);
|
||||
}
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 128; ++i) total += ro[i];
|
||||
}
|
||||
}
|
||||
output[blockIdx.x * kThreads + threadIdx.x] = total;
|
||||
}
|
||||
|
||||
template <typename Kernel>
|
||||
std::array<int64_t, 8> attributes(Kernel kernel, int dynamic_shared_bytes) {
|
||||
cudaFuncAttributes value{};
|
||||
CUDA_CHECK(cudaFuncGetAttributes(&value, kernel));
|
||||
int active_blocks = 0;
|
||||
CUDA_CHECK(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&active_blocks, kernel, kThreads, dynamic_shared_bytes));
|
||||
return {value.numRegs, static_cast<int64_t>(value.sharedSizeBytes), static_cast<int64_t>(value.localSizeBytes),
|
||||
value.maxThreadsPerBlock, value.binaryVersion, value.ptxVersion,
|
||||
static_cast<int64_t>(value.maxDynamicSharedSizeBytes), active_blocks};
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
torch::Tensor launch_handoff_probe(int64_t blocks, int64_t epochs, bool payload) {
|
||||
TORCH_CHECK(blocks > 0 && epochs > 1, "blocks and epochs must be positive");
|
||||
c10::cuda::CUDAGuard guard(0);
|
||||
auto result = torch::zeros({blocks, 4}, torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA));
|
||||
CUDA_CHECK(cudaFuncSetAttribute(vea_b_handoff_probe<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicSharedBytes));
|
||||
vea_b_handoff_probe<false><<<blocks, kThreads, kDynamicSharedBytes, at::cuda::getDefaultCUDAStream()>>>(
|
||||
result.data_ptr<int64_t>(), static_cast<int>(epochs), payload);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
return result;
|
||||
}
|
||||
|
||||
torch::Tensor launch_mbarrier_handoff_probe(int64_t blocks, int64_t epochs, bool payload) {
|
||||
TORCH_CHECK(blocks > 0 && epochs > 1, "blocks and epochs must be positive");
|
||||
c10::cuda::CUDAGuard guard(0);
|
||||
auto result = torch::zeros({blocks, 4}, torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA));
|
||||
CUDA_CHECK(cudaFuncSetAttribute(vea_b_handoff_probe<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicSharedBytes));
|
||||
vea_b_handoff_probe<true><<<blocks, kThreads, kDynamicSharedBytes, at::cuda::getDefaultCUDAStream()>>>(
|
||||
result.data_ptr<int64_t>(), static_cast<int>(epochs), payload);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
return result;
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> launch_tensor_issue_probe(int64_t blocks, int64_t iterations, int64_t mode) {
|
||||
TORCH_CHECK(blocks > 0 && iterations > 0 && mode >= 0 && mode <= 2, "invalid tensor issue arguments");
|
||||
c10::cuda::CUDAGuard guard(0);
|
||||
auto options = torch::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA);
|
||||
auto values = torch::zeros({blocks, 10}, options);
|
||||
auto clocks = torch::zeros({blocks, 10, 2}, options);
|
||||
vea_b_tensor_issue_probe<<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(
|
||||
values.data_ptr<int64_t>(), clocks.data_ptr<int64_t>(), static_cast<int>(iterations), static_cast<int>(mode));
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
return {values, clocks};
|
||||
}
|
||||
|
||||
torch::Tensor launch_register_probe(int64_t blocks, int64_t iterations, int64_t role) {
|
||||
TORCH_CHECK(blocks > 0 && iterations > 0 && role >= 0 && role <= 2, "invalid register probe arguments");
|
||||
c10::cuda::CUDAGuard guard(0);
|
||||
auto output = torch::zeros({blocks, kThreads}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA));
|
||||
if (role == 0) vea_b_register_probe<0><<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(output.data_ptr<float>(), iterations);
|
||||
else if (role == 1) vea_b_register_probe<1><<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(output.data_ptr<float>(), iterations);
|
||||
else vea_b_register_probe<2><<<blocks, kThreads, 0, at::cuda::getDefaultCUDAStream()>>>(output.data_ptr<float>(), iterations);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
return output;
|
||||
}
|
||||
|
||||
torch::Tensor probe_kernel_attributes(int64_t dynamic_shared_bytes) {
|
||||
CUDA_CHECK(cudaFuncSetAttribute(
|
||||
vea_b_handoff_probe<false>, cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
static_cast<int>(dynamic_shared_bytes)));
|
||||
CUDA_CHECK(cudaFuncSetAttribute(
|
||||
vea_b_handoff_probe<true>, cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
static_cast<int>(dynamic_shared_bytes)));
|
||||
std::array<std::array<int64_t, 8>, 6> rows = {
|
||||
attributes(vea_b_handoff_probe<false>, static_cast<int>(dynamic_shared_bytes)),
|
||||
attributes(vea_b_handoff_probe<true>, static_cast<int>(dynamic_shared_bytes)),
|
||||
attributes(vea_b_tensor_issue_probe, 0),
|
||||
attributes(vea_b_register_probe<0>, 0),
|
||||
attributes(vea_b_register_probe<1>, 0),
|
||||
attributes(vea_b_register_probe<2>, 0),
|
||||
};
|
||||
auto output = torch::empty({6, 8}, torch::TensorOptions().dtype(torch::kInt64));
|
||||
auto accessor = output.accessor<int64_t, 2>();
|
||||
for (int row = 0; row < 6; ++row) for (int column = 0; column < 8; ++column) accessor[row][column] = rows[row][column];
|
||||
return output;
|
||||
}
|
||||
|
|
@ -10,5 +10,6 @@ Future test layers, in order:
|
|||
6. two-step and canonical 12-step video/audio latent parity;
|
||||
7. concurrent resident-service safety.
|
||||
|
||||
Existing retained outputs are oracles, not a substitute for the missing
|
||||
self-contained canonical input fixtures recorded in `NUMERICAL_CONTRACT.md`.
|
||||
Canonical Q/K/V/output fixtures are now captured and reload-verified in durable
|
||||
Phase 2A storage. Existing retained outputs remain additional oracles; executable
|
||||
attention parity has not started.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,143 @@
|
|||
"""Capture self-contained canonical Vortex Exact Attention Q/K/V fixtures."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import platform
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
PROJECT = Path(__file__).resolve().parents[1]
|
||||
ROOT = PROJECT.parents[1]
|
||||
sys.path.insert(0, str(ROOT / "tools"))
|
||||
|
||||
from profile_attention_path import prepare_qkv, representative_attention_inputs # noqa: E402
|
||||
|
||||
|
||||
EXPECTED_OUTPUT_SHA256 = "4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574"
|
||||
EXPECTED_CHECKPOINT_SHA256 = "72fa9269ce551fb63ff42a32d9b46d0c122e84b4b2c511e22fa698287b088f70"
|
||||
|
||||
|
||||
def sha256_file(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 tensor_sha256(value: torch.Tensor) -> str:
|
||||
data = value.detach().contiguous().view(torch.uint8).cpu().numpy()
|
||||
return hashlib.sha256(memoryview(data)).hexdigest()
|
||||
|
||||
|
||||
def save_tensor(output_dir: Path, name: str, value: torch.Tensor) -> dict:
|
||||
host = value.detach().contiguous().cpu()
|
||||
tensor_hash = tensor_sha256(host)
|
||||
path = output_dir / f"{name}.pt"
|
||||
torch.save({"tensor": host}, path)
|
||||
loaded = torch.load(path, map_location="cpu", weights_only=True)["tensor"]
|
||||
loaded_hash = tensor_sha256(loaded)
|
||||
if loaded_hash != tensor_hash or not torch.equal(loaded, host):
|
||||
raise RuntimeError(f"fixture reload verification failed for {name}")
|
||||
return {
|
||||
"path": str(path),
|
||||
"size_bytes": path.stat().st_size,
|
||||
"file_sha256": sha256_file(path),
|
||||
"tensor_sha256": tensor_hash,
|
||||
"shape": list(host.shape),
|
||||
"stride": list(host.stride()),
|
||||
"dtype": str(host.dtype),
|
||||
"contiguous": host.is_contiguous(),
|
||||
"reload_verified": True,
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument("--model-path", type=Path, default=Path("/models/minimax_h3_fl2va_pruned_nvfp4.safetensors"))
|
||||
parser.add_argument("--sage-mainloop-binary", type=Path, required=True)
|
||||
parser.add_argument("--sage-fused-binary", type=Path, required=True)
|
||||
parser.add_argument("--image", required=True)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
checkpoint_hash = sha256_file(args.model_path)
|
||||
if checkpoint_hash != EXPECTED_CHECKPOINT_SHA256:
|
||||
raise RuntimeError(f"checkpoint SHA-256 mismatch: {checkpoint_hash}")
|
||||
|
||||
workload = SimpleNamespace(
|
||||
model_path=str(args.model_path), device=args.device, width=1344, height=768,
|
||||
frames=124, steps=12, sampler_step=1, seed=440420, text_tokens=100,
|
||||
block_index=24, attention="sage2",
|
||||
)
|
||||
block, hidden, rotation, segments, metadata = representative_attention_inputs(workload)
|
||||
with torch.inference_mode():
|
||||
q, k, v, _ = prepare_qkv(block, hidden, rotation, None)
|
||||
from sageattention import sageattn
|
||||
|
||||
output = sageattn(q, k, v, tensor_layout="NHD", is_causal=False, smooth_k=False)
|
||||
torch.cuda.synchronize()
|
||||
output_hash = tensor_sha256(output)
|
||||
if output_hash != EXPECTED_OUTPUT_SHA256:
|
||||
raise RuntimeError(
|
||||
f"public Sage2 output mismatch: expected {EXPECTED_OUTPUT_SHA256}, got {output_hash}"
|
||||
)
|
||||
fixtures = {
|
||||
"q": save_tensor(args.output_dir, "q_prepared_bf16_nhd", q),
|
||||
"k": save_tensor(args.output_dir, "k_prepared_bf16_nhd", k),
|
||||
"v": save_tensor(args.output_dir, "v_bf16_nhd", v),
|
||||
"output": save_tensor(args.output_dir, "sage2_output_bf16_nhd", output),
|
||||
}
|
||||
|
||||
manifest = {
|
||||
"schema": "vortex-exact-canonical-fixtures",
|
||||
"version": 1,
|
||||
"status": "captured_and_reload_verified",
|
||||
"reference": {
|
||||
"sageattention_version": "2.2.0",
|
||||
"sageattention_commit": "d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5",
|
||||
"expected_output_sha256": EXPECTED_OUTPUT_SHA256,
|
||||
"actual_output_sha256": output_hash,
|
||||
"byte_exact": True,
|
||||
},
|
||||
"workload": metadata,
|
||||
"segments": segments,
|
||||
"checkpoint": {
|
||||
"path": str(args.model_path), "size_bytes": args.model_path.stat().st_size,
|
||||
"sha256": checkpoint_hash,
|
||||
},
|
||||
"sage_binaries": {
|
||||
"mainloop": {"path": str(args.sage_mainloop_binary), "size_bytes": args.sage_mainloop_binary.stat().st_size, "sha256": sha256_file(args.sage_mainloop_binary)},
|
||||
"fused_preparation": {"path": str(args.sage_fused_binary), "size_bytes": args.sage_fused_binary.stat().st_size, "sha256": sha256_file(args.sage_fused_binary)},
|
||||
},
|
||||
"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,
|
||||
"driver": torch.cuda.driver_version() if hasattr(torch.cuda, "driver_version") else None,
|
||||
"python": platform.python_version(),
|
||||
"argv": sys.argv,
|
||||
},
|
||||
"fixtures": fixtures,
|
||||
}
|
||||
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()
|
||||
288
research/vortex_exact_attention/tools/run_capability_probes.py
Normal file
288
research/vortex_exact_attention/tools/run_capability_probes.py
Normal file
|
|
@ -0,0 +1,288 @@
|
|||
"""Build and run isolated VEA-B SM121 capability probes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import contextlib
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
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"
|
||||
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 = ("handoff_named", "handoff_mbarrier", "tensor_issue", "qk_role", "pv_role", "combined_ownership")
|
||||
EPOCHS = 591
|
||||
CURRENT_MAINLOOP_MS = 237.089
|
||||
SYNC_BUDGET_MS = CURRENT_MAINLOOP_MS * 0.05
|
||||
|
||||
|
||||
def sha256_file(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||||
digest.update(chunk)
|
||||
return digest.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.0 - weight) + ordered[high] * weight
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict:
|
||||
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": statistics.mean(values),
|
||||
"min_ms": min(values), "max_ms": max(values),
|
||||
}
|
||||
|
||||
|
||||
def benchmark(fn, warmup: int, iterations: int) -> tuple[dict, object]:
|
||||
result = None
|
||||
for _ in range(warmup):
|
||||
result = 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()
|
||||
result = fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
values.append(start.elapsed_time(end))
|
||||
return summarize(values), 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_phase2a_probe",
|
||||
sources=[str(KERNELS / "vea_b_probe.cpp"), str(KERNELS / "vea_b_probe.cu")],
|
||||
build_directory=str(build_dir),
|
||||
extra_cflags=["-O3"],
|
||||
extra_cuda_cflags=["-O3", "-lineinfo", "--ptxas-options=-v"],
|
||||
verbose=True,
|
||||
)
|
||||
text = stream.getvalue()
|
||||
build_log.parent.mkdir(parents=True, exist_ok=True)
|
||||
build_log.write_text(text, encoding="utf-8")
|
||||
return extension, time.perf_counter() - started
|
||||
|
||||
|
||||
def expected_payload_checksum(epochs: int) -> int:
|
||||
total = 0
|
||||
for epoch in range(epochs):
|
||||
total += sum((epoch + column) & 255 for column in range(64))
|
||||
total += 2 * (epoch & 255)
|
||||
return total
|
||||
|
||||
|
||||
def validate_handoff(result: torch.Tensor, epochs: int, payload: bool) -> dict:
|
||||
host = result.cpu()
|
||||
expected_checksum = expected_payload_checksum(epochs) if payload else 0
|
||||
checks = {
|
||||
"all_producers_completed": bool(torch.all(host[:, 0] == epochs)),
|
||||
"all_checksums_match": bool(torch.all(host[:, 1] == expected_checksum)),
|
||||
"publication_errors": int(host[:, 2].sum().item()),
|
||||
"all_consumers_completed": bool(torch.all(host[:, 3] == 128)),
|
||||
"expected_checksum": expected_checksum,
|
||||
"observed_checksum_first_block": int(host[0, 1].item()),
|
||||
}
|
||||
checks["passed"] = all((
|
||||
checks["all_producers_completed"], checks["all_checksums_match"],
|
||||
checks["publication_errors"] == 0, checks["all_consumers_completed"],
|
||||
))
|
||||
return checks
|
||||
|
||||
|
||||
def analyze_clock_overlap(clocks: torch.Tensor) -> dict:
|
||||
host = clocks.cpu()
|
||||
ratios = []
|
||||
for block in range(host.shape[0]):
|
||||
int_start = int(host[block, 2:6, 0].min().item())
|
||||
int_end = int(host[block, 2:6, 1].max().item())
|
||||
fp_start = int(host[block, 6:10, 0].min().item())
|
||||
fp_end = int(host[block, 6:10, 1].max().item())
|
||||
overlap = max(0, min(int_end, fp_end) - max(int_start, fp_start))
|
||||
denominator = max(1, min(int_end - int_start, fp_end - fp_start))
|
||||
ratios.append(overlap / denominator)
|
||||
return {
|
||||
"blocks": len(ratios),
|
||||
"blocks_with_positive_overlap": sum(value > 0 for value in ratios),
|
||||
"positive_overlap_fraction": sum(value > 0 for value in ratios) / len(ratios),
|
||||
"p50_overlap_ratio": percentile(ratios, 0.50),
|
||||
"p10_overlap_ratio": percentile(ratios, 0.10),
|
||||
}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
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("--handoff-blocks", type=int, default=16576)
|
||||
parser.add_argument("--probe-blocks", type=int, default=128)
|
||||
parser.add_argument("--epochs", type=int, default=EPOCHS)
|
||||
parser.add_argument("--mma-iterations", type=int, default=8192)
|
||||
parser.add_argument("--warmup", type=int, default=3)
|
||||
parser.add_argument("--iterations", type=int, default=20)
|
||||
parser.add_argument("--capture", choices=("none", "handoff_barrier", "handoff_mbarrier", "handoff_payload", "handoff_mbarrier_payload", "tensor_overlap", "qk_role", "pv_role", "combined_ownership"), default="none")
|
||||
parser.add_argument("--cuda-profiler-capture", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
torch.manual_seed(440420)
|
||||
extension, build_seconds = build_extension(args.build_dir, args.build_log)
|
||||
attributes_tensor = extension.attributes(50 * 1024)
|
||||
attributes = {
|
||||
kernel: {name: int(value) for name, 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()
|
||||
with torch.cuda.nvtx.range(f"vortex_exact_{args.capture}"):
|
||||
if args.capture == "handoff_barrier":
|
||||
captured = extension.handoff(args.handoff_blocks, args.epochs, False)
|
||||
elif args.capture == "handoff_mbarrier":
|
||||
captured = extension.handoff_mbarrier(args.handoff_blocks, args.epochs, False)
|
||||
elif args.capture == "handoff_payload":
|
||||
captured = extension.handoff(args.handoff_blocks, args.epochs, True)
|
||||
elif args.capture == "handoff_mbarrier_payload":
|
||||
captured = extension.handoff_mbarrier(args.handoff_blocks, args.epochs, True)
|
||||
elif args.capture == "tensor_overlap":
|
||||
captured = extension.tensor_issue(args.probe_blocks, args.mma_iterations, 0)[0]
|
||||
else:
|
||||
role = {"qk_role": 0, "pv_role": 1, "combined_ownership": 2}[args.capture]
|
||||
captured = extension.register_probe(args.probe_blocks, 4, role)
|
||||
torch.cuda.synchronize()
|
||||
if args.cuda_profiler_capture:
|
||||
torch.cuda.cudart().cudaProfilerStop()
|
||||
report = {
|
||||
"status": "capture_only", "capture": args.capture,
|
||||
"captured_checksum": float(captured.float().sum().item()),
|
||||
"captured_values": captured.cpu().tolist(),
|
||||
"attributes": attributes, "build_seconds": build_seconds,
|
||||
}
|
||||
else:
|
||||
barrier_timing, barrier_result = benchmark(
|
||||
lambda: extension.handoff(args.handoff_blocks, args.epochs, False),
|
||||
args.warmup, args.iterations,
|
||||
)
|
||||
mbarrier_timing, mbarrier_result = benchmark(
|
||||
lambda: extension.handoff_mbarrier(args.handoff_blocks, args.epochs, False),
|
||||
args.warmup, args.iterations,
|
||||
)
|
||||
payload_timing, payload_result = benchmark(
|
||||
lambda: extension.handoff_mbarrier(args.probe_blocks, args.epochs, True),
|
||||
args.warmup, args.iterations,
|
||||
)
|
||||
payload_repeat = extension.handoff_mbarrier(args.probe_blocks, args.epochs, True)
|
||||
torch.cuda.synchronize()
|
||||
tensor_results = {}
|
||||
overlap_clocks = None
|
||||
for name, mode in (("overlap", 0), ("int8_only", 1), ("fp8_only", 2)):
|
||||
timing, result = benchmark(
|
||||
lambda mode=mode: extension.tensor_issue(args.probe_blocks, args.mma_iterations, mode),
|
||||
args.warmup, args.iterations,
|
||||
)
|
||||
tensor_results[name] = {"timing": timing, "checksum": int(result[0].sum().item())}
|
||||
if name == "overlap":
|
||||
overlap_clocks = result[1]
|
||||
register_results = {}
|
||||
for name, role in (("qk_role", 0), ("pv_role", 1), ("combined_ownership", 2)):
|
||||
timing, value = benchmark(
|
||||
lambda role=role: extension.register_probe(args.probe_blocks, 4, role),
|
||||
args.warmup, args.iterations,
|
||||
)
|
||||
register_results[name] = {"timing": timing, "checksum": float(value.sum().item())}
|
||||
report = {
|
||||
"status": "measured_capability_probe",
|
||||
"attributes": attributes,
|
||||
"handoff": {
|
||||
"selected_primitive": "cuda_block_scope_mbarrier",
|
||||
"barrier_only_full_grid": barrier_timing,
|
||||
"mbarrier_only_full_grid": mbarrier_timing,
|
||||
"payload_representative_blocks": payload_timing,
|
||||
"barrier_budget_ms": SYNC_BUDGET_MS,
|
||||
"named_barrier_under_budget": barrier_timing["p50_ms"] < SYNC_BUDGET_MS,
|
||||
"mbarrier_under_budget": mbarrier_timing["p50_ms"] < SYNC_BUDGET_MS,
|
||||
"barrier_validation": validate_handoff(barrier_result, args.epochs, False),
|
||||
"mbarrier_validation": validate_handoff(mbarrier_result, args.epochs, False),
|
||||
"payload_validation": validate_handoff(payload_result, args.epochs, True),
|
||||
"deterministic_repeat": bool(torch.equal(payload_result, payload_repeat)),
|
||||
},
|
||||
"tensor_issue": {
|
||||
**tensor_results,
|
||||
"clock_overlap": analyze_clock_overlap(overlap_clocks),
|
||||
"overlap_faster_than_serial_sum": tensor_results["overlap"]["timing"]["p50_ms"] < (
|
||||
tensor_results["int8_only"]["timing"]["p50_ms"] + tensor_results["fp8_only"]["timing"]["p50_ms"]
|
||||
),
|
||||
},
|
||||
"register_roles": register_results,
|
||||
"advancement": {
|
||||
"registers_at_most_200": attributes["combined_ownership"]["registers_per_thread"] <= 200,
|
||||
"local_bytes_zero": attributes["combined_ownership"]["local_bytes_per_thread"] == 0,
|
||||
"one_resident_ten_warp_cta": attributes["handoff_mbarrier"]["resident_ctas_per_sm"] >= 1,
|
||||
},
|
||||
"build_seconds": build_seconds,
|
||||
}
|
||||
properties = torch.cuda.get_device_properties(0)
|
||||
report.update({
|
||||
"schema": "vortex-exact-phase2a-capability-probes", "version": 1,
|
||||
"environment": {
|
||||
"image": args.image, "gpu": properties.name,
|
||||
"compute_capability": [properties.major, properties.minor],
|
||||
"multiprocessor_count": properties.multi_processor_count,
|
||||
"torch": torch.__version__, "cuda": torch.version.cuda,
|
||||
"python": sys.version, "argv": sys.argv,
|
||||
"environment_switches": {name: value for name, value in sorted(os.environ.items()) if name.startswith(("CUDA_", "TORCH_", "MAX_JOBS"))},
|
||||
},
|
||||
"sources": {
|
||||
"cpp_sha256": sha256_file(KERNELS / "vea_b_probe.cpp"),
|
||||
"cuda_sha256": sha256_file(KERNELS / "vea_b_probe.cu"),
|
||||
"build_log": str(args.build_log),
|
||||
},
|
||||
"configuration": {
|
||||
"threads": 320, "warps": 10, "producer_warps": 6,
|
||||
"consumer_warps": 4, "epochs": args.epochs,
|
||||
"dynamic_shared_bytes": 50 * 1024,
|
||||
"handoff_blocks": args.handoff_blocks, "probe_blocks": args.probe_blocks,
|
||||
"mma_iterations": args.mma_iterations,
|
||||
},
|
||||
"claims": {"attention_kernel_implemented": False, "production_changed": 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 __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -84,9 +84,32 @@ def validate_documents(inventory: dict, model: dict, root: Path = ROOT) -> None:
|
|||
c_shared = by_id["VEA-C"]["shared_kib"]
|
||||
if c_shared["cluster_total"] != c_shared["producer"] + c_shared["consumer"]:
|
||||
raise RuntimeError("VEA-C shared-memory arithmetic is inconsistent")
|
||||
phase2a = model["phase2a_capability_measurements"]
|
||||
if phase2a["status"] != "passed_for_isolated_aligned_prototype":
|
||||
raise RuntimeError("Phase 2A status must remain scoped to an isolated prototype")
|
||||
if phase2a["selected_handoff"] != "cuda_block_scope_mbarrier":
|
||||
raise RuntimeError("Phase 2A must retain the sanitizer-clean handoff")
|
||||
if max(phase2a["registers_per_thread"].values()) > 200:
|
||||
raise RuntimeError("Phase 2A measured register limit is exceeded")
|
||||
if phase2a["local_bytes_per_thread"] or phase2a["ptxas_spill_loads"] or phase2a["ptxas_spill_stores"]:
|
||||
raise RuntimeError("Phase 2A must retain zero local memory and spills")
|
||||
if phase2a["resident_ctas_per_sm"] < 1:
|
||||
raise RuntimeError("Phase 2A must retain one resident ten-warp CTA")
|
||||
if phase2a["timing_ms"]["mbarrier_only_p95"] >= phase2a["timing_ms"]["budget"]:
|
||||
raise RuntimeError("Phase 2A handoff exceeds its synchronization budget")
|
||||
sanitizer = phase2a["sanitizer"]
|
||||
if any(sanitizer.values()):
|
||||
raise RuntimeError("Phase 2A selected probes must remain sanitizer-clean")
|
||||
fixture = phase2a["canonical_fixture"]
|
||||
if not fixture["reload_verified"] or fixture["output_sha256"] != "4c666c20f5f8f651158a2ced33ccff08f3bada07665c595b99008d171db30574":
|
||||
raise RuntimeError("Phase 2A canonical fixture is not locked to Sage2")
|
||||
if phase2a["attention_kernel_implemented"] or phase2a["attention_latency_measured"]:
|
||||
raise RuntimeError("Phase 2A capability evidence must not claim attention results")
|
||||
decision = model["decision"]
|
||||
if decision["kernel_implementation_started"] or decision["production_dispatch_changed"]:
|
||||
raise RuntimeError("Phase 1 must remain design-only and isolated")
|
||||
raise RuntimeError("Vortex attention must remain unimplemented and isolated")
|
||||
if not decision["phase2a_capability_passed"] or not decision["isolated_aligned_prototype_authorized"]:
|
||||
raise RuntimeError("Phase 2A decision is inconsistent with retained measurements")
|
||||
|
||||
|
||||
def validate() -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue