309 lines
13 KiB
Text
309 lines
13 KiB
Text
/*
|
|
* 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;
|
|
}
|