daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Validate aligned Vortex attention prototype

Browse source
This commit is contained in:
Daniel Maddern 2026-08-26 15:51:39 +07:00
parent c58e4c6d12
commit 67e05f517c
11 changed files with 1483 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
research/EXPERIMENT_REGISTRY.md
View file

@ -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

56
research/experiment_registry.json
View file

@ -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 <fixture> --output <report> --build-dir <build> --build-log <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",

11
research/vortex_exact_attention/DESIGN.md
View file

@ -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.

32
research/vortex_exact_attention/EXPERIMENT_LOG.md
View file

@ -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.

94
research/vortex_exact_attention/PHASE2B_ALIGNED_REPORT.md Normal file
View file

@ -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.

16
research/vortex_exact_attention/README.md
View file

@ -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.

10
research/vortex_exact_attention/kernels/README.md
View file

@ -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.

23
research/vortex_exact_attention/kernels/vea_b_numeric.cpp Normal file
View file

@ -0,0 +1,23 @@
#include <torch/extension.h>
#include <vector>
std::vector<torch::Tensor> 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<torch::Tensor> 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");
}

851
research/vortex_exact_attention/kernels/vea_b_numeric.cu Normal file
View file

@ -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 <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <cuda/barrier>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#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 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<uint32_t>(__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<uint32_t>(__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 <bool Capture>
__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 <bool Capture>
__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 <bool Capture>
__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<nv_bfloat162*>(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<nv_bfloat162*>(output + row * 128 + column) = packed;
}
}
}
}
template <bool Capture>
__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<uint4*>(shared)};
Smem128 k_slots[2] = {
{reinterpret_cast<uint4*>(shared + 16 * 1024)},
{reinterpret_cast<uint4*>(shared + 24 * 1024)}};
Smem64 v_slots[2] = {
{reinterpret_cast<uint4*>(shared + 32 * 1024)},
{reinterpret_cast<uint4*>(shared + 40 * 1024)}};
using BlockBarrier = cuda::barrier<cuda::thread_scope_block>;
BlockBarrier* barriers = reinterpret_cast<BlockBarrier*>(shared + 48 * 1024);
float* scale_slots = reinterpret_cast<float*>(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<const uint4*>(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<Capture>(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<uint4*>(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<uint4*>(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<Capture>(
v_slots[slot], reinterpret_cast<uint4*>(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<const uint4*>(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<const uint4*>(v + row * 192 + epoch * 64)[column];
}
(void)barriers[slot].arrive();
}
barriers[2].arrive_and_wait();
barriers[3].arrive_and_wait();
}
}
template <bool Capture>
__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<uint4*>(shared)};
Smem128 k_smem{reinterpret_cast<uint4*>(shared + 16 * 1024)};
Smem64 v_smem{reinterpret_cast<uint4*>(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<const uint4*>(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<const uint4*>(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<const uint4*>(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<Capture>(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<Capture>(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<torch::Tensor> 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<torch::Tensor>& outputs) {
return {outputs[1].data_ptr<float>(), outputs[2].data_ptr<float>(), outputs[3].data_ptr<float>(),
reinterpret_cast<uint32_t*>(outputs[4].data_ptr<int32_t>()), outputs[5].data_ptr<float>(),
reinterpret_cast<uint32_t*>(outputs[6].data_ptr<int32_t>()), outputs[7].data_ptr<float>(),
outputs[8].data_ptr<float>(), outputs[9].data_ptr<int64_t>()};
}
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 <typename Kernel>
std::array<int64_t, 8> 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<int64_t>(value.sharedSizeBytes), static_cast<int64_t>(value.localSizeBytes),
value.maxThreadsPerBlock, value.binaryVersion, value.ptxVersion,
static_cast<int64_t>(value.maxDynamicSharedSizeBytes), active};
}
} // namespace
std::vector<torch::Tensor> 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<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
const auto stream = at::cuda::getDefaultCUDAStream();
if (capture) {
vea_b_numeric_kernel<true><<<1, kThreads, kDynamicShared, stream>>>(
q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
reinterpret_cast<nv_bfloat16*>(outputs[0].data_ptr()), pointers(outputs));
} else {
vea_b_numeric_kernel<false><<<1, kThreads, kDynamicShared, stream>>>(
q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
reinterpret_cast<nv_bfloat16*>(outputs[0].data_ptr()), pointers(outputs));
}
CUDA_CHECK(cudaGetLastError());
return outputs;
}
std::vector<torch::Tensor> 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<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kReferenceShared));
sage_checkpoint_kernel<true><<<1, kReferenceThreads, kReferenceShared, at::cuda::getDefaultCUDAStream()>>>(
q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
reinterpret_cast<nv_bfloat16*>(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<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
Checkpoints empty{};
vea_b_numeric_kernel<false><<<1, kThreads, kDynamicShared, at::cuda::getDefaultCUDAStream()>>>(
q.data_ptr<int8_t>(), k.data_ptr<int8_t>(), reinterpret_cast<int8_t*>(v.data_ptr()),
q_scale.data_ptr<float>(), k_scale.data_ptr<float>(), v_scale.data_ptr<float>(),
reinterpret_cast<nv_bfloat16*>(output.data_ptr()), empty);
CUDA_CHECK(cudaGetLastError());
return output;
}
torch::Tensor numeric_kernel_attributes() {
CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<false>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
CUDA_CHECK(cudaFuncSetAttribute(vea_b_numeric_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kDynamicShared));
CUDA_CHECK(cudaFuncSetAttribute(sage_checkpoint_kernel<true>, cudaFuncAttributeMaxDynamicSharedMemorySize, kReferenceShared));
std::array<std::array<int64_t, 8>, 3> rows = {
attributes(vea_b_numeric_kernel<false>, kDynamicShared, kThreads),
attributes(vea_b_numeric_kernel<true>, kDynamicShared, kThreads),
attributes(sage_checkpoint_kernel<true>, kReferenceShared, kReferenceThreads)};
auto output = torch::empty({3, 8}, torch::TensorOptions().dtype(torch::kInt64));
auto access = output.accessor<int64_t, 2>();
for (int row = 0; row < 3; ++row)
for (int column = 0; column < 8; ++column) access[row][column] = rows[row][column];
return output;
}

133
research/vortex_exact_attention/tools/capture_phase2b_fixtures.py Normal file
View file

@ -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()

267
research/vortex_exact_attention/tools/run_phase2b.py Normal file
View file

@ -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()