diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index eb29ed34..c196f0ef 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -399,10 +399,8 @@ The corresponding Vulkan shader sources and native fusion rules have been removed from ggml. Unsupported custom nodes retain scheduler fallback to another supported backend. -Remaining model-specific CUDA paths and Vulkan INT8/ConvRot implementations still -need auditing or migration. -This implementation does not yet externalize all model-specific kernels or add -dynamic dtypes. +Remaining model-specific CUDA paths still need auditing or migration. This +implementation does not add dynamic dtypes. `src/core/vulkan/int8.cpp` owns Vulkan INT8 GEMM, ConvRot, and their fusion rules. The Vulkan compute metadata service exposes shared-memory limits, subgroup @@ -412,3 +410,22 @@ SwiGLU, residual, row-concatenation, and modulation paths. Kernels use public Vulkan dispatch, buffer, copy, and workspace services without private backend headers. Vulkan currently handles packed group-256 activations with separate weight scales; unsupported custom nodes fall back to CPU. + +## SageAttention + +The existing `ggml_sage_attn` entry point is declared and implemented in sd.cpp's +common extension files. It constructs `sd.sage_attn.v1` with explicit Q/K/V inputs +and scale/mode attributes. ggml has no Sage-specific opcode, allocation rule, or +backend dispatch branch. + +`src/core/cuda/sage_attention.cu` and its `sageattention/` device headers own the +CUDA implementation. Quantization and temporary attention output use scheduler +workspace, rather than an oversized output tensor. The implementation obtains +the backend stream through the public GPU service and converts the final FP16 +output with a local CUDA kernel. It retains the SM80 compatibility path and the +SageAttention2/2++ toolkit and architecture requirements. Capability checks use +the application's compiled CUDA target as well as the physical device. + +CPU, HIP, MUSA, and Vulkan do not implement this NVIDIA-specific kernel. Existing +model-level capability checks retain flash/default attention for unsupported +layers; a custom Sage node itself does not silently change algorithms. diff --git a/ggml b/ggml index 9201d070..62da2f1b 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 9201d070ff57d41efbbd3f245bd80869c73f460a +Subproject commit 62da2f1b7e2ff90241dc48fb0fa3e9fd003706f5 diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 3d4509d6..9a610a49 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -23,6 +23,8 @@ endfunction() if(SD_CUDA) enable_language(CUDA) sd_cuda_sources(cuda CUDA) + target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu) + set_source_files_properties(sage_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math") if(GGML_CUDA_CUTLASS_PATH) set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS) diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu index 601890c7..f13d94f2 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -11,6 +11,9 @@ namespace SD_CUDA_NAMESPACE { void register_i8(ggml_backend_dev_t device); void register_convrot(ggml_backend_dev_t device); void register_affine(ggml_backend_dev_t device); +#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA) + void register_sage_attention(ggml_backend_dev_t device); +#endif void register_extensions(ggml_backend_dev_t device) { if (cuda_device_info(device).device < 0) @@ -19,6 +22,9 @@ namespace SD_CUDA_NAMESPACE { register_i8(device); register_convrot(device); register_rope(device); +#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA) + register_sage_attention(device); +#endif #ifdef SD_CUDA_CUTE_ATTENTION sd_ggml_register_cuda_i8_projection(device); sd_ggml_register_cuda_i8_gated_residual(device); diff --git a/src/core/cuda/sage_attention.cu b/src/core/cuda/sage_attention.cu new file mode 100644 index 00000000..46b9a08a --- /dev/null +++ b/src/core/cuda/sage_attention.cu @@ -0,0 +1,387 @@ +#include "core/ggml_extensions.h" +#include "runtime.h" + +#if __CUDACC_VER_MAJOR__ >= 12 + +#include +#include +#include +#include +#include "sageattention/attn.cuh" +#include "sageattention/attn_fp8.cuh" + +static constexpr int sage_cuda_version = __CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10; + +namespace sd_cuda { + + struct SageDevice { + int cc; + int compiled_cc; + }; + + static ggml_sage_attn_mode sage_mode(const SageDevice& device, ggml_sage_attn_mode requested) { + if (requested != GGML_SAGE_ATTN_AUTO) + return requested; + if (sage_cuda_version >= 12040 && device.cc >= 89 && device.compiled_cc >= 89) { + if (sage_cuda_version >= 12080 && device.cc != 90) + return GGML_SAGE_ATTN_2_PLUS_PLUS; + return GGML_SAGE_ATTN_2; + } + return GGML_SAGE_ATTN_AUTO; + } + + static __device__ __forceinline__ float warp_reduce_max(float x) { +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset)); + return x; + } + + static __global__ void sage_output(const half* input, float* output, int count) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < count) + output[i] = __half2float(input[i]); + } + + struct SageWorkspace { + size_t q; + size_t k; + size_t q_scale; + size_t k_scale; + size_t k_mean; + size_t v; + size_t v_scale; + size_t output; + size_t size; + + explicit SageWorkspace(const ggml_tensor* dst) { + const ggml_tensor* qt = dst->src[0]; + const ggml_tensor* kt = dst->src[1]; + size = 0; + q = reserve(ggml_nelements(qt)); + k = reserve(ggml_nelements(kt)); + q_scale = reserve(GGML_PAD(qt->ne[1], 128) / 32 * 8 * qt->ne[2] * qt->ne[3] * sizeof(float)); + k_scale = reserve(GGML_PAD(kt->ne[1], 64) / 64 * 4 * kt->ne[2] * kt->ne[3] * sizeof(float)); + k_mean = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float)); + v = reserve(kt->ne[0] * GGML_PAD(kt->ne[1], 64) * kt->ne[2] * kt->ne[3]); + v_scale = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float)); + output = reserve(ggml_nelements(dst) * sizeof(half)); + } + + size_t reserve(size_t bytes) { + const size_t offset = size; + size += GGML_PAD(bytes, 256); + return offset; + } + }; + + template + static __global__ void sage_key_mean(const float* k, float* mean, int length) { + const int channel = blockIdx.x * 32 + threadIdx.x % 32; + const int row = threadIdx.x / 32; + const size_t base = size_t(blockIdx.y) * length * D; + float sum = 0.0f; + for (int token = row; token < length; token += 8) { + sum += k[base + size_t(token) * D + channel]; + } + __shared__ float sums[256]; + sums[threadIdx.x] = sum; + __syncthreads(); + if (row == 0) { + for (int i = 1; i < 8; ++i) { + sum += sums[threadIdx.x + i * 32]; + } + mean[size_t(blockIdx.y) * D + channel] = sum / length; + } + } + + template + static __global__ void sage_quantize(const float* input, const float* mean, int8_t* output, float* scales, int length) { + constexpr int VALUES = D * TOKENS / 256; + const int head = blockIdx.y; + const int first = blockIdx.x * TOKENS * D; + const size_t base = size_t(head) * length * D; + float values[VALUES]; + float amax = 1e-7f; +#pragma unroll + for (int i = 0; i < VALUES; ++i) { + const int index = first + threadIdx.x + i * 256; + float value = 0.0f; + if (index < length * D) { + value = input[base + index]; + if constexpr (SMOOTH) { + value -= mean[size_t(head) * D + index % D]; + } + } + values[i] = value; + amax = fmaxf(amax, fabsf(value)); + } + amax = warp_reduce_max(amax); + __shared__ float maxima[8]; + if (threadIdx.x % 32 == 0) { + maxima[threadIdx.x / 32] = amax; + } + __syncthreads(); + if (threadIdx.x < 32) { + amax = threadIdx.x < 8 ? maxima[threadIdx.x] : 1e-7f; + amax = warp_reduce_max(amax); + if (threadIdx.x == 0) { + maxima[0] = amax; + scales[size_t(head) * gridDim.x + blockIdx.x] = amax / 127.0f; + } + } + __syncthreads(); + const float inverse_scale = 127.0f / maxima[0]; +#pragma unroll + for (int i = 0; i < VALUES; ++i) { + const int index = first + threadIdx.x + i * 256; + if (index < length * D) { + output[base + index] = int8_t(__float2int_rn(values[i] * inverse_scale)); + } + } + } + + template + static __global__ void sage_quantize_per_thread(const float* input, const float* mean, int8_t* output, float* scales, int length) { + constexpr int GROUPS = KEY ? 4 : 8; + constexpr int TOKENS = KEY ? 64 : 32; + constexpr int VALUES = D * TOKENS / GROUPS / 32; + const int group = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + const size_t base = size_t(blockIdx.y) * length * D; + float values[VALUES]; + float amax = 0.0f; +#pragma unroll + for (int i = 0; i < VALUES; ++i) { + const int index = lane + i * 32; + const int row = index / D; + const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group); + float value = token < length ? input[base + size_t(token) * D + index % D] : 0.0f; + if constexpr (KEY) { + if (token < length) { + value -= mean[size_t(blockIdx.y) * D + index % D]; + } + } + values[i] = value; + amax = fmaxf(amax, fabsf(value)); + } + const float scale = warp_reduce_max(amax) / 127.0f + 1e-7f; + if (lane == 0) { + scales[(size_t(blockIdx.y) * gridDim.x + blockIdx.x) * GROUPS + group] = scale; + } +#pragma unroll + for (int i = 0; i < VALUES; ++i) { + const int index = lane + i * 32; + const int row = index / D; + const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group); + if (token < length) { + const float value = values[i] / scale; + output[base + size_t(token) * D + index % D] = int8_t(__float2int_rz(value + copysignf(0.5f, value))); + } + } + } + + template + static __global__ void sage_value_scale(const half* v, float* scales, int length, float scale_max) { + const int channel = blockIdx.x * 32 + threadIdx.x % 32; + const int row = threadIdx.x / 32; + const size_t base = size_t(blockIdx.y) * length * D; + float amax = 1e-7f; + for (int token = row; token < length; token += 8) { + amax = fmaxf(amax, fabsf(__half2float(v[base + size_t(token) * D + channel]))); + } + __shared__ float maxima[256]; + maxima[threadIdx.x] = amax; + __syncthreads(); + if (row == 0) { + for (int i = 1; i < 8; ++i) { + amax = fmaxf(amax, maxima[threadIdx.x + i * 32]); + } + scales[size_t(blockIdx.y) * D + channel] = amax / scale_max; + } + } + + template + static __global__ void sage_quantize_value(const half* v, const float* scales, int8_t* output, int length, int padded) { + __shared__ uint8_t tile[32][D + 4]; + const int head = blockIdx.y; + const int first = blockIdx.x * 32; + const size_t base = size_t(head) * length * D; + for (int i = threadIdx.x; i < 32 * D; i += 256) { + const int token = first + i / D; + const int channel = i % D; + const float value = token < length ? __half2float(v[base + size_t(token) * D + channel]) / scales[size_t(head) * D + channel] : 0.0f; + // FP8 MMA consumes pairs of tokens in a permuted, transposed V tile. + const int row = i / D; + const int permuted = row / 16 * 16 + row % 16 / 8 * 2 + row % 8 / 2 * 4 + row % 2; + tile[permuted][channel] = __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3); + } + __syncthreads(); + for (int i = threadIdx.x; i < 32 * D; i += 256) { + output[(size_t(head) * D + i / 32) * padded + first + i % 32] = tile[i % 32][i / 32]; + } + } + + template + static void sage_attn_fp8(cudaStream_t stream, ggml_tensor* dst, char* data, const SageWorkspace& work, float scale) { + const ggml_tensor* q = dst->src[0]; + const ggml_tensor* k = dst->src[1]; + const ggml_tensor* v = dst->src[2]; + const int lq = q->ne[1], lk = k->ne[1], padded = GGML_PAD(lk, 64); + const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3]; + auto* qi = reinterpret_cast(data + work.q); + auto* ki = reinterpret_cast(data + work.k); + auto* vi = reinterpret_cast(data + work.v); + auto* qs = reinterpret_cast(data + work.q_scale); + auto* ks = reinterpret_cast(data + work.k_scale); + auto* vs = reinterpret_cast(data + work.v_scale); + auto* km = reinterpret_cast(data + work.k_mean); + auto* out = reinterpret_cast(data + work.output); + sage_key_mean<<>>(static_cast(k->data), km, lk); + sage_quantize_per_thread<<>>( + static_cast(q->data), nullptr, qi, qs, lq); + sage_quantize_per_thread<<>>( + static_cast(k->data), km, ki, ks, lk); + // The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++. + sage_value_scale<<>>(static_cast(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f); + sage_quantize_value<<>>(static_cast(v->data), vs, vi, lk, padded); + auto kernel = sageattention::qk_int_sv_f8_attn_kernel<128, 64, 32, 64, D, + sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerThread, + sageattention::QuantGranularity::kPerThread, float, true, half, + sageattention::ComputeUnit::kCudaCore, sageattention::MaskMode::kNone, false, true, false, PLUS_PLUS>; + kernel<<>>( + qi, ki, vi, out, nullptr, qs, ks, vs, nullptr, lq, lk, hq / hk, + D * lq * hq, D, D * lq, D * lk * hk, D, D * lk, + D * padded * hk, D * padded, padded, D * lq * hq, D * hq, D, scale); + sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast(dst->data), ggml_nelements(dst)); + } + + template + static void sage_attn(cudaStream_t stream, ggml_tensor* dst, char* data, ggml_sage_attn_mode mode, float scale) { + const ggml_tensor* q = dst->src[0]; + const ggml_tensor* k = dst->src[1]; + const ggml_tensor* v = dst->src[2]; + const int lq = q->ne[1], lk = k->ne[1]; + const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3]; + const SageWorkspace work(dst); + if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS) { + sage_attn_fp8(stream, dst, data, work, scale); + return; + } + if (mode == GGML_SAGE_ATTN_2) { + sage_attn_fp8(stream, dst, data, work, scale); + return; + } + auto* qi = reinterpret_cast(data + work.q); + auto* ki = reinterpret_cast(data + work.k); + auto* qs = reinterpret_cast(data + work.q_scale); + auto* ks = reinterpret_cast(data + work.k_scale); + auto* km = reinterpret_cast(data + work.k_mean); + auto* out = reinterpret_cast(data + work.output); + + sage_key_mean<<>>(static_cast(k->data), km, lk); + sage_quantize<<>>( + static_cast(q->data), nullptr, qi, qs, lq); + sage_quantize<<>>( + static_cast(k->data), km, ki, ks, lk); + + constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D; + auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D, + sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp, + sageattention::QuantGranularity::kPerWarp, float, false, half, + sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>; + kernel<<>>( + qi, ki, static_cast(v->data), out, nullptr, qs, ks, nullptr, + lq, lk, hq / hk, + D * lq * hq, D, D * lq, + D * lk * hk, D, D * lk, + D * lk * hk, D, D * lk, + D * lq * hq, D * hq, D, + scale); + sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast(dst->data), ggml_nelements(dst)); + } + + static bool sage_supports(ggml_backend_dev_t, const ggml_tensor* op, void* user_data) { + const auto& device = *static_cast(user_data); + SDSageAttentionParams params; + if (!sd_ggml_sage_attn_describe(op, ¶ms) || device.cc < 80 || device.compiled_cc < 80) + return false; + const auto mode = sage_mode(device, params.mode); + if (mode != GGML_SAGE_ATTN_AUTO && (sage_cuda_version < 12040 || device.cc < 89 || device.compiled_cc < 89)) + return false; + if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS && sage_cuda_version < 12080) + return false; + const auto* q = op->src[0]; + const auto* k = op->src[1]; + const auto* v = op->src[2]; + if (q->ne[0] != 64 && q->ne[0] != 128) + return false; + for (const auto* t : {q, k, v, op}) { + if (ggml_nelements(t) > INT_MAX || t->ne[2] * t->ne[3] > 65535 || + (reinterpret_cast(t->data) & 15)) + return false; + } + return k->ne[0] * GGML_PAD(k->ne[1], 64) * k->ne[2] * k->ne[3] <= INT_MAX; + } + + static ggml_status sage_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) { + *workspace = SageWorkspace(nodes[0]).size; + return GGML_STATUS_SUCCESS; + } + + static void sage_destroy(void*, void*) {} + + static ggml_status sage_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void* user_data) { + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context)) + return GGML_STATUS_FAILED; + SDSageAttentionParams params; + if (!sd_ggml_sage_attn_describe(nodes[0], ¶ms)) + return GGML_STATUS_FAILED; + const auto mode = sage_mode(*static_cast(user_data), params.mode); + auto* data = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + auto stream = static_cast(context.stream); + if (nodes[0]->src[0]->ne[0] == 64) + sage_attn<64>(stream, nodes[0], data, mode, params.scale); + else + sage_attn<128>(stream, nodes[0], data, mode, params.scale); + return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + + void register_sage_attention(ggml_backend_dev_t device) { + const auto info = cuda_device_info(device); + if (info.device < 0) + return; + int previous; + cudaDeviceProp properties; + cudaFuncAttributes attributes; + if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess || + cudaSetDevice(info.device) != cudaSuccess) + return; + const auto error = cudaFuncGetAttributes(&attributes, sage_output); + if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess) + return; + // PTX compiled for an older target retains that target's instruction guards after JIT. + static std::map devices; + auto& data = devices[info.device]; + data = {properties.major * 10 + properties.minor, std::min(attributes.binaryVersion, attributes.ptxVersion)}; + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.sage_attn.cuda.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_sage_attn_op(); + desc.priority = 100; + desc.supports = sage_supports; + desc.prepare = sage_prepare; + desc.destroy = sage_destroy; + desc.execute = sage_execute; + desc.user_data = &data; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + } + +} +#else +namespace sd_cuda { + void register_sage_attention(ggml_backend_dev_t) {} +} +#endif diff --git a/src/core/cuda/sageattention/LICENSE b/src/core/cuda/sageattention/LICENSE new file mode 100644 index 00000000..927221d6 --- /dev/null +++ b/src/core/cuda/sageattention/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2024 Jintao Zhang, Haofeng Huang + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/src/core/cuda/sageattention/README.md b/src/core/cuda/sageattention/README.md new file mode 100644 index 00000000..ea43a312 --- /dev/null +++ b/src/core/cuda/sageattention/README.md @@ -0,0 +1,28 @@ +# SageAttention CUDA device code + +Adapted from [thu-ml/SageAttention](https://github.com/thu-ml/SageAttention), +revision `d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5`, under Apache-2.0 (see LICENSE). + +`attn.cuh` contains the device kernel from +`csrc/qattn/qk_int_sv_f16_cuda_sm80.cu`; `attn_fp8.cuh` adapts +`csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh`. The remaining headers come from `csrc/` +and `csrc/qattn/attn_utils.cuh`. PyTorch includes and host wrappers were removed; +include paths and guards were adjusted, device helpers were placed in the +`sageattention` namespace, and `div_ceil` was renamed to avoid macro collisions. +Unused `` includes were removed so multi-architecture builds can +also compile targets below SM70; the backend rejects SageAttention below SM80. +Predicated loads now zero unused FP16 V lanes: masked probabilities cannot +suppress NaNs left in shared memory by previous kernels. + +The sd.cpp adapter in `../sage_attention.cu` supplies FP32 K smoothing and native CUDA +per-thread INT8 quantization, matching the token grouping in upstream +`sageattention/triton/quant_per_thread.py`. V is quantized per channel and +transposed, padded with zeros, and permuted according to +`csrc/fused/fused.cu`. Its FP8 range is 448 for SageAttention2 and 2.25 for +SageAttention2++. Zero channels use a positive scale to avoid division by zero. +SageAttention2 uses an FP32 instruction buffer; 2++ uses an FP16 instruction +buffer added to FP32 accumulators. The SM80 compatibility path retains per-warp +INT8 Q/K and FP16 V with FP32 accumulation. + +The adapter requests workspace and the current stream through public ggml extension APIs; it has no Python, PyTorch, +or Triton dependency. The public output is converted back to FP32. diff --git a/src/core/cuda/sageattention/attn.cuh b/src/core/cuda/sageattention/attn.cuh new file mode 100644 index 00000000..61b57580 --- /dev/null +++ b/src/core/cuda/sageattention/attn.cuh @@ -0,0 +1,686 @@ +/* + * Copyright (c) 2024 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH +#define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH + +#include "math.cuh" +#include "attn_utils.cuh" + +namespace sageattention { + +#define PACK_SIZE_QK 16 // as if it is int8 +#define PACK_SIZE_V 8 // fp16 +#define PACK_SIZE_O 8 // fp16 + +// treat as if int8 tensor core +#define MMA_QK_M 16 +#define MMA_QK_N 16 +#define MMA_QK_K 32 + +// fp16 tensor core +#define MMA_SV_M 16 +#define MMA_SV_N 16 +#define MMA_SV_K 16 + +template +__global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, half *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse, + float *__restrict__ Q_scale, float *__restrict__ K_scale, DTypeOut *__restrict__ V_mean, + const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, + const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, + const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, + const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, + const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, + float sm_scale) +{ + // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes. +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 + // compile time check + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4"); + static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(std::is_same::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float"); + static_assert(std::is_same::value || std::is_same::value, "DTypeSVAccum must be float or half"); + static_assert(std::is_same::value || std::is_same::value, "DTypeOut must be half or nv_bfloat16"); + static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64"); + static_assert(!fuse_v_mean || std::is_same::value, "fuse_v_mean only supports half"); + static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation + + using DTypeOut2 = typename std::conditional::value, half2, nv_bfloat162>::type; + + constexpr uint32_t num_warps_q = CTA_Q / WARP_Q; + constexpr uint32_t num_warps_k = CTA_K / WARP_K; + constexpr uint32_t num_warps = num_warps_q * num_warps_k; + constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M; + constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N; + constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K); + constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N; + + constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2); + constexpr uint32_t O_SMEM_STRIDE = head_dim; + constexpr uint32_t V_SMEM_STRIDE = head_dim; + + extern __shared__ int8_t smem[]; + + const uint32_t lane_id = get_lane_id(); + const uint32_t warp_id = get_warp_id(); + + // maximize L2 hit rate + const uint32_t batch_id = blockIdx.z; + const uint32_t bx = blockIdx.x; + const uint32_t num_qo_heads = gridDim.y; + const uint32_t head_id = blockIdx.y; + + // transfer to base 2 instead of base e with better numerical efficiency + sm_scale *= math::log2e; + + // RS holds the fragment of S + int32_t RS[num_tiles_q][num_tiles_k][8]; + DTypeSVAccum RO[num_tiles_q][num_tiles_v][8]; + float m[num_tiles_q][2]; // max + float d[num_tiles_q][2]; // denominator + + uint32_t q_scale_idx, k_scale_idx; + + if constexpr (Q_GRAN == QuantGranularity::kPerBlock) + { + const uint32_t num_block_q = gridDim.x; + q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx; + } + else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) + { + const uint32_t num_warp_block_q = gridDim.x * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q(); + } + else if constexpr (Q_GRAN == QuantGranularity::kPerThread) + { + const uint32_t num_warp_block_q = gridDim.x * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q() * 8 + lane_id / 4; + } + + if constexpr (K_GRAN == QuantGranularity::kPerBlock) + { + const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k; + } + else if constexpr (K_GRAN == QuantGranularity::kPerWarp) + { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k(); + } + else if constexpr (K_GRAN == QuantGranularity::kPerThread) + { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k() * 4 + lane_id % 4; + } + + constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) : (CTA_K / WARP_K) * 4; + + // initialize o, m, d +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + if constexpr (std::is_same::value) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RO[fq][fv][k] = 0.0f; + } + } + else if constexpr (std::is_same::value) + { +#pragma unroll + for (uint32_t k = 0; k < 4; k++) + { + ((int32_t*)RO[fq][fv])[k] = 0; + } + } + } + } +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) + { + m[fq][k] = -5000000.0f; + d[fq][k] = 1.0f; + } + } + + constexpr uint32_t K_smem_idx_offset = CTA_Q; + constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; + + constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_Q(smem); + smem_t smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); + constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE); + constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_O(smem); + + constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4; + constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4; + constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4; + + constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK); + constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V); + constexpr uint32_t V_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_V); + constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O); + constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O); + + int8_t *Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + int8_t *K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + half *V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V; + uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V); + + uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id % 16, lane_id / 16); + uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2); + uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 16, lane_id / 16); + + // for causal masking + uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; + uint32_t K_idx_lane_base = get_warp_idx_k() * WARP_K + 2 * (lane_id % 4); + + // for loading + uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + uint32_t V_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V; + + const uint32_t num_iterations = SAGE_DIV_CEIL( + mask_mode == MaskMode::kCausal + ? min(kv_len, (bx + 1) * CTA_Q) + : kv_len, + CTA_K); + + // load Q with predicate + load_global_to_share( + &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len); + cp_async::commit_group(); + cp_async::wait_group<0>(); + __syncthreads(); + + // for num_tiles_qk_inner = 1, we load all Qs in register + uint32_t RQ[num_tiles_q][4]; + if constexpr (num_tiles_qk_inner == 1) + { +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); + Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); + } + } + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + float q_scale = Q_scale[q_scale_idx]; + + float original_sm_scale = sm_scale; + float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset]; + + sm_scale = original_sm_scale * dequant_scale; + + // load V with predicate + load_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + K_load_idx_lane_base += CTA_K; + V_load_idx_lane_base += CTA_K; + +#pragma unroll + for (uint32_t iter = 1; iter < num_iterations - 1; iter++) + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]); + } + } + } + + // do not apply causal mask and out of bound mask for these iterations + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, sm_scale); + } + else if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d(RS_f16, d); + } + + __syncthreads(); + + // load K + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) + { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + else + { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + // load V + load_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V); + cp_async::commit_group(); + K_load_idx_lane_base += CTA_K; + V_load_idx_lane_base += CTA_K; + } + + // second last iter, apply causal mask + if (num_iterations > 1) + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) + { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + else if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d(RS_f16, d); + } + + __syncthreads(); + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) + { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + else + { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + // load V with predicate + load_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len); + cp_async::commit_group(); + K_load_idx_lane_base += CTA_K; + V_load_idx_lane_base += CTA_K; + } + + // last iter, apply causal mask and out of bound mask + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) + { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // check out of bound in the last iter + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + else if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d(RS_f16, d); + } + + // ensure V is ready + cp_async::wait_group<0>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) + { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + else + { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + + } + + // TODO: thread block sync mdo state for num_warps_k > 0 + + normalize_d(RO, m, d); + + // save the result + // if (get_warp_idx_k() == 0) + // { + + // convert half to bfloat16 + if constexpr (std::is_same::value && std::is_same::value) + { +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + ((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0])); + ((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1])); + ((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2])); + ((nv_bfloat162*)RO[fq][fv])[3] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[3])); + } + } + } + + // add v_mean + if constexpr (fuse_v_mean) + { + DTypeOut2 v_mean[2]; + DTypeOut *V_mean_lane_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + lane_id % 4 * 2; +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16)); + v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16)); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + ((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]); + ((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]); + ((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]); + ((DTypeOut2*)RO[fq][fv])[3] = __hadd2(((DTypeOut2*)RO[fq][fv])[3], v_mean[1]); + } + } + } + + // save the result to shared memory + uint32_t smem_O_row_base = get_warp_idx_q() * WARP_Q + lane_id / 4; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); + + if constexpr (std::is_same::value) + { + // convert RO to half + uint32_t RO_f16[4]; +#pragma unroll + for (uint32_t k = 0; k < 4; k++) + { + if constexpr (std::is_same::value) + { + ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); + } + else if constexpr (std::is_same::value) + { + ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); + } + } + + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1]; + + // ! permuted, make sure you know what you are doing + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2]; + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; + } + else if constexpr (std::is_same::value) + { + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1]; + + // ! permuted, make sure you know what you are doing + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[2]; + ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[3]; + } + } + } + + // ! do we need to sync here? + __syncwarp(); + + // shared memory to global memory + DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O; + uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O); + uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; + +#pragma unroll + for (uint32_t i = 0; i < O_smem_iters_col; i++) + { +#pragma unroll + for (uint32_t j = 0; j < O_smem_iters_row; j++) + { + if (O_load_idx_lane_base < qo_len) + { + smem_O.store_128b(offset_O, O_lane_ptr); + } + O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); + offset_O = smem_O.advance_offset_by_column(offset_O); + } + + offset_O = smem_O.advance_offset_by_row(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O)); + O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O)); + O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; + } + + if constexpr (return_lse) + { + uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q(); + float *lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; + uint32_t fq = (lane_id % 4) / 2; + uint32_t k = (lane_id % 4) % 2; + + if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) + { + lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); + } + } + + // } +#else + RUNTIME_ASSERT("SageAttention requires SM80 or newer"); +#endif +} + + +#undef PACK_SIZE_QK +#undef PACK_SIZE_V +#undef PACK_SIZE_O +#undef MMA_QK_M +#undef MMA_QK_N +#undef MMA_QK_K +#undef MMA_SV_M +#undef MMA_SV_N +#undef MMA_SV_K + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/attn_fp8.cuh b/src/core/cuda/sageattention/attn_fp8.cuh new file mode 100644 index 00000000..507f9e7f --- /dev/null +++ b/src/core/cuda/sageattention/attn_fp8.cuh @@ -0,0 +1,718 @@ +/* + * Copyright (c) 2024 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH +#define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH + +#include "math.cuh" +#include "attn_utils.cuh" + +namespace sageattention { + +#define PACK_SIZE_QK 16 // as if it is int8 +#define PACK_SIZE_V 16 // fp8 +#define PACK_SIZE_O 8 // fp16 + +// treat as if int8 tensor core +#define MMA_QK_M 16 +#define MMA_QK_N 16 +#define MMA_QK_K 32 + +// fp8 tensor core +#define MMA_SV_M 16 +#define MMA_SV_N 16 +#define MMA_SV_K 32 + +template +__global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, int8_t *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse, + float *__restrict__ Q_scale, float *__restrict__ K_scale, float *__restrict__ V_scale, float *__restrict__ V_mean, + const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, + const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, + const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, + const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, + const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, + float sm_scale) +{ + // Keep unsupported device passes from parsing architecture-specific intrinsics. +#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890 + // compile time check + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4"); + static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread"); + static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64"); + static_assert(std::is_same::value, "DTypeSVAccum must be float, half is WIP"); + static_assert(std::is_same::value || std::is_same::value, "DTypeOut must be half or nv_bfloat16"); + static_assert(CTA_K % 64 == 0); + static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation + + constexpr uint32_t num_warps_q = CTA_Q / WARP_Q; + constexpr uint32_t num_warps_k = CTA_K / WARP_K; + constexpr uint32_t num_warps = num_warps_q * num_warps_k; + constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M; + constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N; + constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K); + constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N; + + constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2); + constexpr uint32_t O_SMEM_STRIDE = head_dim; + // for fp16: head_dim + constexpr uint32_t V_SMEM_STRIDE = CTA_K; + + extern __shared__ int8_t smem[]; + + const uint32_t lane_id = get_lane_id(); + const uint32_t warp_id = get_warp_id(); + + // maximize L2 hit rate + const uint32_t batch_id = blockIdx.z; + const uint32_t bx = blockIdx.x; + const uint32_t num_qo_heads = gridDim.y; + const uint32_t head_id = blockIdx.y; + + // transfer to base 2 instead of base e with better numerical efficiency + sm_scale *= math::log2e; + + // RS holds the fragment of S + int32_t RS[num_tiles_q][num_tiles_k][8]; + DTypeSVAccum RO[num_tiles_q][num_tiles_v][8]; + float m[num_tiles_q][2]; // max + float d[num_tiles_q][2]; // denominator + + uint32_t q_scale_idx, k_scale_idx; + + if constexpr (Q_GRAN == QuantGranularity::kPerBlock) + { + const uint32_t num_block_q = gridDim.x; + q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx; + } + else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) + { + const uint32_t num_warp_block_q = gridDim.x * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q(); + } + else if constexpr (Q_GRAN == QuantGranularity::kPerThread) + { + const uint32_t num_warp_block_q = gridDim.x * num_warps_q; + q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q() * 8 + lane_id / 4; + } + + if constexpr (K_GRAN == QuantGranularity::kPerBlock) + { + const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k; + } + else if constexpr (K_GRAN == QuantGranularity::kPerWarp) + { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k(); + } + else if constexpr (K_GRAN == QuantGranularity::kPerThread) + { + const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); + k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k() * 4 + lane_id % 4; + } + + constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) : (CTA_K / WARP_K) * 4; + + // initialize o, m, d +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + if constexpr (std::is_same::value) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RO[fq][fv][k] = 0.0f; + } + } + else if constexpr (std::is_same::value) + { +#pragma unroll + for (uint32_t k = 0; k < 4; k++) + { + ((int32_t*)RO[fq][fv])[k] = 0; + } + } + } + } +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) + { + m[fq][k] = -5000000.0f; + d[fq][k] = 1.0f; + } + } + + constexpr uint32_t K_smem_idx_offset = CTA_Q; + constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; + + constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_Q(smem); + smem_t smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); + // for fp16: 32 + constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE); + constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; + smem_t smem_O(smem); + + constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4; + // for fp16: 32 + constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8; + // for fp16: 32 + constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 64) ? 8 : 4; + constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; + constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4; + + constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK); + constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK); + constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V); + // for fp16: CTA_K + constexpr uint32_t V_smem_iters_col = head_dim / (num_warps * global_to_shared_copy_lines_per_warp_V); + constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O); + constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O); + + int8_t *Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + int8_t *K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK; + // for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v + int8_t *V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V; + uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK); + uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V); + + uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id % 16, lane_id / 16); + uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2); + // for fp 16: + // uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k() * WARP_K + lane_id % 16, lane_id / 16); + uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(lane_id % 8 + (lane_id / 16) * 8, get_warp_idx_k() * WARP_K / PACK_SIZE_V + (lane_id / 8) % 2); + + // for causal masking + uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; + uint32_t K_idx_lane_base = get_warp_idx_k() * WARP_K + 2 * (lane_id % 4); + + // for loading + uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK; + + const uint32_t num_iterations = SAGE_DIV_CEIL( + mask_mode == MaskMode::kCausal + ? min(kv_len, (bx + 1) * CTA_Q) + : kv_len, + CTA_K); + + // load Q with predicate + load_global_to_share( + &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len); + cp_async::commit_group(); + cp_async::wait_group<0>(); + __syncthreads(); + + // for num_tiles_qk_inner = 1, we load all Qs in register + uint32_t RQ[num_tiles_q][4]; + if constexpr (num_tiles_qk_inner == 1) + { +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); + Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); + } + } + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + float q_scale = Q_scale[q_scale_idx]; + + float original_sm_scale = sm_scale; + float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset]; + + sm_scale = original_sm_scale * dequant_scale; + + // load V + // ! we assume that V is padded. If not, there might be illegal memory access or nan issue. + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + + K_load_idx_lane_base += CTA_K; + +#pragma unroll + for (uint32_t iter = 1; iter < num_iterations - 1; iter++) + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + float RS_f32[num_tiles_q][num_tiles_k][8]; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]); + } + } + } + + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, sm_scale); + } + else if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d_f8(RS_f8, d); + } + + __syncthreads(); + + // load K without predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) + { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } + else + { + if constexpr (!use_pv_fp16_accu){ + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } + else{ + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + __syncthreads(); + // load V + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + + K_load_idx_lane_base += CTA_K; + } + + // second last iter, apply causal mask + if (num_iterations > 1) + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) + { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + else if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d_f8(RS_f8, d); + } + + __syncthreads(); + + // load K with predicate + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len); + cp_async::commit_group(); + + dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset]; + sm_scale = original_sm_scale * dequant_scale; + + // ensure V is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) + { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } + else + { + if constexpr (!use_pv_fp16_accu){ + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } + else{ + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + + __syncthreads(); + // load V + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + K_load_idx_lane_base += CTA_K; + } + + // last iter, apply causal mask and out of bound mask + { + // ensure K is ready + cp_async::wait_group<1>(); + __syncthreads(); + + // compute QK^T + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk( + smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + compute_int_qk( + smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; + } + } + } + + if constexpr (mask_mode == MaskMode::kCausal) + { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + else if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d_f8(RS_f8, d); + } + + // ensure V is ready + cp_async::wait_group<0>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) + { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } + else + { + if constexpr (!use_pv_fp16_accu){ + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } + else{ + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + + __syncthreads(); + + } + + // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving. + + normalize_d(RO, m, d); + + // ! here we just implement the case for fp32 acumulation + if constexpr (fuse_v_scale) + { + float v_scale[4]; + float *V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4 ) * 2; +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + ((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16)); + ((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8)); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + RO[fq][fv][0] *= v_scale[0]; + RO[fq][fv][1] *= v_scale[1]; + RO[fq][fv][2] *= v_scale[0]; + RO[fq][fv][3] *= v_scale[1]; + RO[fq][fv][4] *= v_scale[2]; + RO[fq][fv][5] *= v_scale[3]; + RO[fq][fv][6] *= v_scale[2]; + RO[fq][fv][7] *= v_scale[3]; + } + } + } + + if constexpr (fuse_v_mean) + { + float v_mean[4]; + float *V_mean_base_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4 ) * 2; +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + ((float2*)v_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16)); + ((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8)); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + RO[fq][fv][0] += v_mean[0]; + RO[fq][fv][1] += v_mean[1]; + RO[fq][fv][2] += v_mean[0]; + RO[fq][fv][3] += v_mean[1]; + RO[fq][fv][4] += v_mean[2]; + RO[fq][fv][5] += v_mean[3]; + RO[fq][fv][6] += v_mean[2]; + RO[fq][fv][7] += v_mean[3]; + } + } + } + + // save the result to shared memory + uint32_t smem_O_row_base = get_warp_idx_q() * WARP_Q + lane_id / 4; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); + + if constexpr (std::is_same::value) + { + // convert RO to half + uint32_t RO_f16[4]; +#pragma unroll + for (uint32_t k = 0; k < 4; k++) + { + if constexpr (std::is_same::value) + { + ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); + } + else + { + ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); + } + } + + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1]; + + offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1); + ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2]; + ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; + } + else if constexpr (std::is_same::value) + { + // TODO: not implement + } + } + } + + // ! do we need to sync here? + __syncwarp(); + + // shared memory to global memory + DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O; + uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O); + uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; + +#pragma unroll + for (uint32_t i = 0; i < O_smem_iters_col; i++) + { +#pragma unroll + for (uint32_t j = 0; j < O_smem_iters_row; j++) + { + if (O_load_idx_lane_base < qo_len) + { + smem_O.store_128b(offset_O, O_lane_ptr); + } + O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); + offset_O = smem_O.advance_offset_by_column(offset_O); + } + + offset_O = smem_O.advance_offset_by_row(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O)); + O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O)); + O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; + } + + if constexpr (return_lse) + { + // ! this only works for num_tiles_q = 2 + uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q(); + float *lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; + uint32_t fq = (lane_id % 4) / 2; + uint32_t k = (lane_id % 4) % 2; + + if (lse_idx < qo_len) + { + lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); + } + } +#else + RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer"); +#endif +} + +#undef PACK_SIZE_QK +#undef PACK_SIZE_V +#undef PACK_SIZE_O +#undef MMA_QK_M +#undef MMA_QK_N +#undef MMA_QK_K +#undef MMA_SV_M +#undef MMA_SV_N +#undef MMA_SV_K + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/attn_utils.cuh b/src/core/cuda/sageattention/attn_utils.cuh new file mode 100644 index 00000000..32f4099a --- /dev/null +++ b/src/core/cuda/sageattention/attn_utils.cuh @@ -0,0 +1,998 @@ +/* + * Copyright (c) 2024 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_UTILS_CUH +#define SD_CORE_CUDA_SAGEATTENTION_ATTN_UTILS_CUH +#include +#include + +#include "cp_async.cuh" +#include "mma.cuh" +#include "permuted_smem.cuh" +#include "numeric_conversion.cuh" + +namespace sageattention { + +#define WARP_SIZE 32 + +#define S_FP8_OFFSET 8.807f +#define S_FP8_OFFSET_EXP 6680.8477f +#define S_FP8_OFFSET_EXP_INV 0.0022326917f + +#define SAGE_DIV_CEIL(M, N) (((M) + (N)-1) / (N)) + +enum class MaskMode { + kNone = 0, + kCausal = 1, +}; + +enum class DataType { + kHalf, + kInt8, + kInt4, + kE4M3, + kE5M2, +}; + +enum class QuantGranularity { + kPerTensor = 0, + kPerBlock = 1, + kPerWarp = 2, + kPerThread = 3, +}; + +enum class ComputeUnit { + kTensorCore, + kCudaCore, +}; + +__device__ __forceinline__ uint32_t get_warp_id() +{ + return threadIdx.y; +} + +__device__ __forceinline__ uint32_t get_lane_id() +{ + return threadIdx.x; +} + +template +__device__ __forceinline__ uint32_t get_warp_idx_q() +{ + return get_warp_id() / num_warps_k; +} + +template +__device__ __forceinline__ uint32_t get_warp_idx_k() +{ + return get_warp_id() % num_warps_k; +} + +template +__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset, + const uint32_t &gmem_stride, + const smem_t &smem) +{ + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + static_assert(std::is_same::value || std::is_same::value); + + constexpr uint32_t pack_size = std::is_same::value ? 8 : 16; + +#pragma unroll + for (uint32_t i = 0; i < smem_iters_col; i++) + { +#pragma unroll + for (uint32_t j = 0; j < smem_iters_row; j++) + { + smem.load_128b_async(smem_offset, *lane_ptr); + *lane_ptr += (global_to_shared_line_lanes * pack_size); + smem_offset = smem.advance_offset_by_column(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); + *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size)); + } + smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); + *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; +} + +// with predicate +template +__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset, + const uint32_t &gmem_stride, + const smem_t &smem, uint32_t base_idx, uint32_t max_len) +{ + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + static_assert(std::is_same::value || std::is_same::value); + + constexpr uint32_t pack_size = std::is_same::value ? 8 : 16; + +#pragma unroll + for (uint32_t i = 0; i < smem_iters_col; i++) + { +#pragma unroll + for (uint32_t j = 0; j < smem_iters_row; j++) + { + // Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN. + smem.load_128b_async::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len); + *lane_ptr += (global_to_shared_line_lanes * pack_size); + smem_offset = smem.advance_offset_by_column(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); + *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size)); + base_idx += global_to_shared_copy_lines_per_warp_per_iter; + } + smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); + *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; +} + +template +__device__ __forceinline__ void load_fp8_V_global_to_share(int8_t **lane_ptr, uint32_t &smem_offset, + const uint32_t &gmem_stride, + const smem_t &smem) +{ + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + constexpr uint32_t pack_size_fp8 = 16; + +#pragma unroll + for (uint32_t i = 0; i < smem_iters_col; i++) + { +#pragma unroll + for (uint32_t j = 0; j < smem_iters_row; j++) + { + smem.load_128b_async(smem_offset, *lane_ptr); + *lane_ptr += (global_to_shared_line_lanes * pack_size_fp8); + smem_offset = smem.advance_offset_by_column(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(smem_offset - (smem_iters_row * global_to_shared_line_lanes)); + *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size_fp8)); + } + smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride); + // for QK: *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; + *lane_ptr += CTA; // ! prevent underflow + *lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; +} + +template +__device__ __forceinline__ void compute_int_qk(const smem_t &smem_Q, const smem_t &smem_K, int32_t RS[][num_tiles_k][8], uint32_t &offset_Q, uint32_t &offset_K) +{ + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); + + uint32_t RQ[num_tiles_q][4]; + uint32_t RK[4]; + + // the first iteration, mma mode is kInit +#pragma unroll + for (uint32_t iter = 0; iter < 1; iter++) + { + // load RQ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); + offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); + } + // ! using permutation invariance + offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); + +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + // load RK + smem_K.ldmatrix_m8n8x4(offset_K, RK); + offset_K = smem_K.advance_offset_by_row<16>(offset_K); + + // mma +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (DTypeQK == DataType::kInt8) + { + mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + } + else if constexpr (DTypeQK == DataType::kInt4) + { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } + } + offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter); + } + + // following iteration, mma mode is kInplace +#pragma unroll + for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) + { + // load RQ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); + offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); + } + offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); + +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + // load RK + smem_K.ldmatrix_m8n8x4(offset_K, RK); + offset_K = smem_K.advance_offset_by_row<16>(offset_K); + + // mma +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (DTypeQK == DataType::kInt8) + { + mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + } + else if constexpr (DTypeQK == DataType::kInt4) + { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } + } + offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter); + } + + offset_Q -= (2 * num_tiles_qk_inner); + offset_K -= (2 * num_tiles_qk_inner); +} + +// for case when num_tiles_qk_inner = 1 +template +__device__ __forceinline__ void compute_int_qk(const smem_t &smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) +{ + static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); + static_assert(num_tiles_qk_inner == 1); + + uint32_t RK[4]; + + // mma mode is kInit +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + // load RK + smem_K.ldmatrix_m8n8x4(offset_K, RK); + offset_K = smem_K.advance_offset_by_row<16>(offset_K); + + // mma +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (DTypeQK == DataType::kInt8) + { + mma::mma_sync_m16n16k32_row_col_s8s8s32(RS[fq][fk], RQ[fq], RK); + } + else if constexpr (DTypeQK == DataType::kInt4) + { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } + } +} + +template +__device__ __forceinline__ void apply_causal_mask(const uint32_t &Q_idx_lane_base, const uint32_t &K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) +{ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2); + const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; + const bool out_of_boundary = (kv_idx > q_idx); + + if constexpr (std::is_same::value) + { + RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + } + else if constexpr (std::is_same::value) + { + RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); + } + } + } + } +} + +template +__device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t &K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t &kv_len) +{ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; + const bool out_of_boundary = (kv_idx >= kv_len); + + if constexpr (std::is_same::value) + { + RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + } + else if constexpr (std::is_same::value) + { + RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); + } + } + } + } +} + +// for DTypeQKAccum float +template +__device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float &sm_scale) +{ + static_assert(std::is_same::value || (!use_half_o_scale)); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) + { + // assign the smallest value possible + float m_prev = m[fq][k]; + float m_temp = -5000000.0f; +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]), + max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5])); + m_temp = max(m_temp, m_local); + } + + if constexpr (!fuse_scale) + { + if constexpr (exp_offset) + { + m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET); + } + else + { + m_temp *= sm_scale; + } + } + else if constexpr (exp_offset) + { + m_temp += (-S_FP8_OFFSET); + } + + // exchange element with the 4 threads in the row + m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x1)); // 0 exchange with 1, 2 exchange with 3 + m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x2)); // 0 exchange with 2, 1 exchange with 3 + + m[fq][k] = max(m[fq][k], m_temp); + + float o_scale = math::ptx_exp2(m_prev - m[fq][k]); + + // update denominator + d[fq][k] *= o_scale; + + half2 o_scale2; + if constexpr (use_half_o_scale) + { + o_scale2 = __floats2half2_rn(o_scale, o_scale); + } + + // update RO +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + if constexpr (std::is_same::value) + { + RO[fq][fv][k * 2 + 0] *= o_scale; + RO[fq][fv][k * 2 + 1] *= o_scale; + RO[fq][fv][k * 2 + 4] *= o_scale; + RO[fq][fv][k * 2 + 5] *= o_scale; + } + else if constexpr (std::is_same::value) + { + if constexpr (use_half_o_scale) + { + ((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2); + ((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2); + } + else + { + RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale); + RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale); + RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale); + RO[fq][fv][k * 2 + 5] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 5]) * o_scale); + } + } + } + + // raise RS to exponent + float negative_m = -m[fq][k]; +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + if constexpr (fuse_scale) + { + RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m); + RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m); + RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m); + RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m); + } + else + { + RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m)); + RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m)); + RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m)); + RS[fq][fk][k * 2 + 5] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 5], sm_scale, negative_m)); + } + } + } + } +} + +template +__device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) +{ + static_assert(sizeof(T) == 4); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + ((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]); + ((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]); + ((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]); + ((half2*)RS_16[fq][fk])[3] = __float22half2_rn(((float2*)RS[fq][fk])[3]); + } + } +} + +template +__device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) +{ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) + { + floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4); + floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6); + floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4); + floatx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 2, RS[fq][fk * 2 + 1] + 6); + } + } +} + +template +__device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) +{ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) + { + halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2); + halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3); + halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2); + halfx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3); + } + } +} + +template +__device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) +{ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) + { + e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]); + e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1); + e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2); + e4m3x4_to_halfx4(RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3, RS_8[fq][fk] + 3); + } + } +} + +template +__device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore)? 4 : 8], float d[][2]) +{ + // for compute unit cuda core, RS is float + // for compute unit tensor core, RS is packed half + static_assert((std::is_same::value && compute_unit == ComputeUnit::kCudaCore) || + (std::is_same::value && compute_unit == ComputeUnit::kTensorCore)); + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { + if constexpr (compute_unit == ComputeUnit::kTensorCore) + { + // full accumulate with tensor core + mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk])); + } + else if constexpr (compute_unit == ComputeUnit::kCudaCore) + { + // partial accumulate with cuda core + 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]; + } + } + } +} + +template +__device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) +{ +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) + { + mma::rowsum_f8f8f32(d[fq], RS[fq][fk]); + } + } +} + +template +__device__ __forceinline__ void compute_fp16_sv(const smem_t &smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) +{ + uint32_t smem_V_row_base = get_warp_idx_k() * (num_tiles_k * 16) + get_lane_id() % 16; + uint32_t smem_V_col_base = get_lane_id() / 16; +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2); + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV); + } + else if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV); + } + } + } + } +} + +template +__device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t &smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t &offset_V) +{ + static_assert(sizeof(T) == 4); + + // ! be sure you know what you are doing +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + else if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + } + + offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); + } + + // make offset_V their original value + offset_V -= (16 * num_tiles_k * stride); +} + +template +__device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t &smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t &offset_V) +{ + static_assert(sizeof(T) == 4); + static_assert(std::is_same::value); + + uint32_t RO_inst_buf[num_tiles_q][num_tiles_v][4]; + + // ! be sure you know what you are doing +#pragma unroll + for (uint32_t fk = 0; fk < 1; fk++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + } + + offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); + } + +#pragma unroll + for (uint32_t fk = 1; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } + } + + offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv); + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v)); + } + + // accumulate into RO +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x); + RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y); + RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x); + RO[fq][fv][3] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].y); + RO[fq][fv][4] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].x); + RO[fq][fv][5] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].y); + RO[fq][fv][6] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].x); + RO[fq][fv][7] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].y); + } + } + + // make offset_V their original value + offset_V -= (16 * num_tiles_k * stride); +} + +template +__device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) +{ + if constexpr (compute_unit == ComputeUnit::kCudaCore) + { + // accumulate_d performs partial accumulation with cuda core + // aggregate d +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) + { + d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3 + d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3 + } + } + } + + // divide O by d + float d_rcp[num_tiles_q][2]; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t k = 0; k < 2; k++) + { + // TODO: check m to prevent nan + d_rcp[fq][k] = math::ptx_rcp(d[fq][k]); + } + } + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + if constexpr (std::is_same::value) + { + RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2]; + } + else if constexpr (std::is_same::value) + { + RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]); + } + } + } + } +} + +template +__device__ __forceinline__ void compute_fp8_sv(const smem_t &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) +{ + uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; + // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; + uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) + { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV); + } + else if constexpr (std::is_same::value) + { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } +} + +template +__device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) +{ + uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; + // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; + uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; + + float RO_inst_buf[num_tiles_q][num_tiles_v][8]; + +#pragma unroll + for (uint32_t fk = 0; fk < 1; fk++) + { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + } + else if constexpr (std::is_same::value) + { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + +#pragma unroll + for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) + { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + } + else if constexpr (std::is_same::value) + { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; + RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; + RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; + RO[fq][fv][3] += RO_inst_buf[fq][fv][3]; + RO[fq][fv][4] += RO_inst_buf[fq][fv][4]; + RO[fq][fv][5] += RO_inst_buf[fq][fv][5]; + RO[fq][fv][6] += RO_inst_buf[fq][fv][6]; + RO[fq][fv][7] += RO_inst_buf[fq][fv][7]; + } + } +} + +template +__device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) +{ + uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; + // uint32_t smem_V_col_base = get_warp_idx_k() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; + uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; + + uint32_t RO_int32[num_tiles_q][num_tiles_v][4]; + +#pragma unroll + for (uint32_t fk = 0; fk < 1; fk++) + { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + //mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + } + else if constexpr (std::is_same::value) + { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + +#pragma unroll + for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) + { + uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); +#pragma unroll + for (uint32_t fv = 0; fv < num_tiles_v; fv++) + { + // load RV + uint32_t RV[4]; + // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); + smem_V.ldmatrix_m8n8x4(offset_V, RV); +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { + if constexpr (std::is_same::value) + { + //mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + } + else if constexpr (std::is_same::value) + { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + float RO_tmp_float[2]; +#pragma unroll + for(int i = 0; i < num_tiles_q; i++){ +#pragma unroll + for(int j = 0; j < num_tiles_v; j++){ + #pragma unroll + for(int k = 0; k < 4; k++){ + unpack_half2_from_uint32_to_float(RO_tmp_float, RO_int32[i][j][k]); + RO[i][j][k * 2 + 0] += RO_tmp_float[0]; + RO[i][j][k * 2 + 1] += RO_tmp_float[1]; + } + } + } + +// #pragma unroll +// for (uint32_t fq = 0; fq < num_tiles_q; fq++) +// { +// #pragma unroll +// for (uint32_t fv = 0; fv < num_tiles_v; fv++) +// { +// RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; +// RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; +// RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; +// RO[fq][fv][3] += RO_inst_buf[fq][fv][3]; +// RO[fq][fv][4] += RO_inst_buf[fq][fv][4]; +// RO[fq][fv][5] += RO_inst_buf[fq][fv][5]; +// RO[fq][fv][6] += RO_inst_buf[fq][fv][6]; +// RO[fq][fv][7] += RO_inst_buf[fq][fv][7]; +// } +// } +} + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/cp_async.cuh b/src/core/cuda/sageattention/cp_async.cuh new file mode 100644 index 00000000..51cf59d5 --- /dev/null +++ b/src/core/cuda/sageattention/cp_async.cuh @@ -0,0 +1,148 @@ +/* + * Copyright (c) 2024 by SageAttention team. + * + * This file is based on code from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/cp_async.cuh + * Copyright (c) 2023 by FlashInfer team. + * Small modifications made by SageAttention team, 2024 (e.g., renamed namespace). + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_CP_ASYNC_CUH +#define SD_CORE_CUDA_SAGEATTENTION_CP_ASYNC_CUH +#include + +namespace sageattention { + +namespace cp_async { + +enum class SharedMemFillMode { + kFillZero, // Fill zero to shared memory when predicate is false + kNoFill // Do not fill zero to shared memory when predicate is false +}; + +enum class PrefetchMode { + kNoPrefetch, // Do not fetch additional data from global memory to L2 + kPrefetch // Fetch additional data from global memory to L2 +}; + +#if (__CUDACC_VER_MAJOR__ >= 11) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) +#define CP_ASYNC_ENABLED +#endif +#endif + +/*! + * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted + * cp.async instructions to a group + */ +__device__ __forceinline__ void commit_group() { +#ifdef CP_ASYNC_ENABLED + asm volatile("cp.async.commit_group;\n" ::); +#endif +} + +/*! + * \brief Wrapper of PTX cp.async.wait_group instruction + * \tparam n Wait till most recent n groups are committed + */ +template +__device__ __forceinline__ void wait_group() { +#ifdef CP_ASYNC_ENABLED + asm volatile("cp.async.wait_group %0;\n" ::"n"(n)); +#endif +} + +/*! + * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from + * global memory to shared memory + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + */ +template +__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { +#ifdef CP_ASYNC_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(16)); + } else { + asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(16)); + } +#else + *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); +#endif +} + +/*! + * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from + * global memory to shared memory with predicate. + * \tparam prefetch_mode Whether to fetch additional data from global memory to L2 + * \tparam fill_mode Whether to fill zero to shared memory when predicate is false + * \tparam T Data type + * \param smem_ptr Pointer to shared memory + * \param gmem_ptr Pointer to global memory + * \param predicate Predicate value + * \note fill zero is slower than not fill zero + */ +template +__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) { +#ifdef CP_ASYNC_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + int src_in_bytes = predicate ? 16 : 0; + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); + } else { + asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr), + "l"(gmem_ptr), "n"(16), "r"(src_in_bytes)); + } + } else { + if constexpr (prefetch_mode == PrefetchMode::kPrefetch) { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n" + "}\n" ::"r"((int)predicate), + "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); + } else { + asm volatile( + "{\n" + " .reg .pred p;\n" + " setp.ne.b32 p, %0, 0;\n" + " @p cp.async.cg.shared.global [%1], [%2], %3;\n" + "}\n" ::"r"((int)predicate), + "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16)); + } + } +#else + if (predicate) { + *((uint4*)smem_ptr) = *((uint4*)gmem_ptr); + } else { + if constexpr (fill_mode == SharedMemFillMode::kFillZero) { + *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0); + } + } +#endif +} + +} // namespace cp_async + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/math.cuh b/src/core/cuda/sageattention/math.cuh new file mode 100644 index 00000000..3233a03f --- /dev/null +++ b/src/core/cuda/sageattention/math.cuh @@ -0,0 +1,162 @@ +/* + * Copyright (c) 2024 by SageAttention team. + * + * This file is based on code from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/math.cuh + * Copyright (c) 2023 by FlashInfer team. + * Small modifications made by SageAttention team, 2024 (e.g., renamed namespace). + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_MATH_CUH +#define SD_CORE_CUDA_SAGEATTENTION_MATH_CUH +#include +#include + +namespace sageattention { +#ifndef USHORT_TYPE +#define USHORT_TYPE +typedef unsigned short ushort; +#endif + +namespace math { + +// log2(e) +constexpr float log2e = 1.44269504088896340736f; +constexpr float log2e_recp = 1.0f / log2e; + +__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; } + +__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; } + +/*! + * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x + * \param x input + */ +__forceinline__ __device__ float ptx_exp2(float x) { + float y; + asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x) + * \param x input + */ +__forceinline__ __device__ float ptx_log2(float x) { + float y; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x + * \param x input + */ +__forceinline__ __device__ half2 ptx_exp2(half2 x) { + uint32_t y_u32; + uint32_t x_u32 = half2_as_uint32(x); + asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); + return uint32_as_half2(y_u32); +} + +/*! + * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x + * \param x input + */ +__forceinline__ __device__ half ptx_exp2(half x) { + ushort y_u16; + asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); + return __ushort_as_half(y_u16); +} + +/*! + * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x + * \param x input + */ +__forceinline__ __device__ float ptx_rcp(float x) { + float y; + asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle + * between threads in a warp. + * \param x The value in the source lane + * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta] + */ +__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) { + float y; + asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;" + : "=f"(y) + : "f"(x), "r"(lane_mask)); + return y; +} + +/*! + * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly + * shuffle between threads in a warp. + * \param x The value in the source lane + * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask] + */ +__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) { + return __shfl_xor_sync(0xffffffff, x, lane_mask); +} + +/*! + * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x) + * \param x input + */ +__forceinline__ __device__ float rsqrt(float x) { + float y; + asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x) + * \param x input + */ +__forceinline__ __device__ float tanh(float x) { + float y; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + +/*! + * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x) + * \param x input + */ +__forceinline__ __device__ half2 tanh(half2 x) { + uint32_t y_u32; + uint32_t x_u32 = half2_as_uint32(x); + asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32)); + return uint32_as_half2(y_u32); +} + +/*! + * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x) + * \param x input + */ +__forceinline__ __device__ half tanh(half x) { + ushort y_u16; + asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x))); + return __ushort_as_half(y_u16); +} + +} // namespace math + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/mma.cuh b/src/core/cuda/sageattention/mma.cuh new file mode 100644 index 00000000..03c4e69e --- /dev/null +++ b/src/core/cuda/sageattention/mma.cuh @@ -0,0 +1,729 @@ +/* + * Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/mma.cuh + * Copyright (c) 2023 by FlashInfer team. + * + * Modifications copyright (c) 2024 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_MMA_CUH +#define SD_CORE_CUDA_SAGEATTENTION_MMA_CUH +#include +#include +#include +#include + +namespace sageattention { + +namespace mma{ + +#if (__CUDACC_VER_MAJOR__ >= 11) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800)) +#define MMA_F16F16F32_M16N8K16_ENABLED +#define MMA_F16F16F16_M16N8K16_ENABLED +#define MMA_S8S8S32_M16N8K32_ENABLED +#define MMA_S4S4S32_M16N8K64_ENABLED +#endif +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 750)) +#define MMA_F16F16F32_M16N8K8_ENABLED +#define MMA_F16F16F16_M16N8K8_ENABLED +#define LDMATRIX_M8N8X2_ENABLED +#define LDMATRIX_M8N8X4_ENABLED +#endif +#endif + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890)) +#define MMA_F8F8F32_M16N8K16_ENABLED +#endif +#endif + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890)) +#define MMA_F8F8F16_M16N8K16_ENABLED +#endif +#endif + +#if defined(__CUDA_ARCH__) +#define RUNTIME_ASSERT(x) __brkpt() +#else + +#define RUNTIME_ASSERT(x) assert(0 && x) +#endif + +enum class MMAMode { + kInit = 0U, + kInplaceUpdate = 1U, +}; + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) { +#ifdef LDMATRIX_M8N8X2_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n" + : "=r"(R[0]), "=r"(R[1]) + : "r"(smem_int_ptr)); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory + * to fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) { +#ifdef LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) + : "r"(smem_int_ptr)); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from + * shared memory to fragment and transposes the fragment + * \tparam T data type of the fragment + * \param R pointer to the fragment + * \param smem_ptr pointer to the shared memory + */ +template +__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) { +#ifdef LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3]) + : "r"(smem_int_ptr)); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_F16F16F32_M16N8K16_ENABLED + // ! only support half dtype now + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_F16F16F32_M16N8K16_ENABLED + // ! only support half dtype now + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f16. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_F16F16F16_M16N8K16_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=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])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=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)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix + * multiplication, accumulated in f16. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_F16F16F16_M16N8K16_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=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])); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=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)); + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C[2]), "=r"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_S8S8S32_M16N8K32_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + 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};\n" + : "=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])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + 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};\n" + : "=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)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_S8S8S32_M16N8K32_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + 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};\n" + : "=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])); + 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};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), + "r"(C[6]), "r"(C[7])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + 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};\n" + : "=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)); + 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};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0), + "r"(0), "r"(0)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_S4S4S32_M16N8K64_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=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])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=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)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix + * multiplication, accumulated in int32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_S4S4S32_M16N8K64_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=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])); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), + "r"(C[6]), "r"(C[7])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=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)); + asm volatile( + "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0), + "r"(0), "r"(0)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix + * multiplication, accumulated in fp32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_F8F8F32_M16N8K16_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix + * multiplication, accumulated in fp16. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, + uint32_t* B) { + //uint32_t* C_uint32 = reinterpret_cast(C); +#ifdef MMA_F8F8F16_M16N8K16_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[0]), "=r"(C_uint32[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1])); + + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[2]), "=r"(C_uint32[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[0]), "=r"(C_uint32[1]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0)); + + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " + "{%0, %1}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, %9};\n" + : "=r"(C_uint32[2]), "=r"(C_uint32[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + + + +/*! + * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix + * multiplication, accumulated in fp32. + * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate + * \param C pointer to the accumulator + * \param A pointer to the fragment of matrix A + * \param B pointer to the fragment of matrix B + */ +template +__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, + uint32_t* B) { +#ifdef MMA_F8F8F32_M16N8K16_ENABLED + if constexpr (mma_mode == MMAMode::kInplaceUpdate) + { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), + "f"(C[2]), "f"(C[3])); + + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), + "f"(C[6]), "f"(C[7])); + } + else if constexpr (mma_mode == MMAMode::kInit) + { + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, %1, %2, %3}," + "{%4, %5, %6, %7}," + "{%8, %9}," + "{%10, %11, %12, %13};\n" + : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f), + "f"(0.f), "f"(0.f)); + } +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Use mma instructions to compute rowsum. + */ +__device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) { +#ifdef MMA_F16F16F32_M16N8K16_ENABLED + asm volatile( + "{\n" + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + "}\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320), // 1006648320 packs two 1.0f in half precision + "r"(1006648320), "f"(d[0]), "f"(d[1])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +/*! + * \brief Use mma instructions to compute rowsum. + */ +__device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) { +#ifdef MMA_F8F8F32_M16N8K16_ENABLED + asm volatile( + "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " + "{%0, _, %1, _}," + "{%2, %3, %4, %5}," + "{%6, %7}," + "{%8, 0., %9, 0.};\n" + : "=f"(d[0]), "=f"(d[1]) + : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504), // 943208504 packs four 1.0f in e4m3 + "f"(d[0]), "f"(d[1])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction"); +#endif +} + +} // namespace mma + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/numeric_conversion.cuh b/src/core/cuda/sageattention/numeric_conversion.cuh new file mode 100644 index 00000000..7049bacb --- /dev/null +++ b/src/core/cuda/sageattention/numeric_conversion.cuh @@ -0,0 +1,155 @@ +/* + * Copyright (c) 2024 by SageAttention team. + * + * Inspired by CUTLASS, https://github.com/NVIDIA/cutlass/blob/main/include/cutlass/numeric_conversion.h + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_NUMERIC_CONVERSION_CUH +#define SD_CORE_CUDA_SAGEATTENTION_NUMERIC_CONVERSION_CUH +#include +#include +#include +#include + +namespace sageattention { + +#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400) +#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890)) +#define FP8_CAST_ENABLED +#endif +#endif + +#if defined(__CUDA_ARCH__) +#define RUNTIME_ASSERT(x) __brkpt() +#else + +#define RUNTIME_ASSERT(x) assert(0 && x) +#endif + +__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) { + uint16_t h0 = source & 0xFFFF; + uint16_t h1 = (source >> 16) & 0xFFFF; + asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0)); + asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1)); +} + +__device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t *dest, float *source0, float *source1) +{ +#ifdef FP8_CAST_ENABLED + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".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};\n" \ + "}" \ + : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); +#endif +} + +__device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t *dest, float *source0, float *source1) +{ +#ifdef FP8_CAST_ENABLED + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" \ + "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); +#endif +} + +__device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) +{ +#ifdef FP8_CAST_ENABLED + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" \ + "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); +#endif +} + +__device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) +{ +#ifdef FP8_CAST_ENABLED + asm volatile( \ + "{\n" \ + ".reg .b16 lo;\n" \ + ".reg .b16 hi;\n" \ + "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" \ + "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" \ + "mov.b32 %0, {lo, hi};\n" \ + "}" \ + : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); +#endif +} + +__device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) +{ +#ifdef FP8_CAST_ENABLED + asm volatile( \ + "{\n" \ + ".reg .b16 lo, hi;\n" \ + "mov.b32 {lo, hi}, %2;\n" \ + "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \ + "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \ + "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); +#endif +} + +__device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) +{ +#ifdef FP8_CAST_ENABLED + asm volatile( \ + "{\n" \ + ".reg .b16 lo, hi;\n" \ + "mov.b32 {lo, hi}, %2;\n" \ + "cvt.rn.f16x2.e5m2x2 %0, lo;\n" \ + "cvt.rn.f16x2.e5m2x2 %1, hi;\n" \ + "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); +#else + RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); +#endif +} + +__device__ __forceinline__ int8_t float_to_int8_rn(float x) +{ + uint32_t dst; + asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); + return reinterpret_cast(dst); +} + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/permuted_smem.cuh b/src/core/cuda/sageattention/permuted_smem.cuh new file mode 100644 index 00000000..f44be703 --- /dev/null +++ b/src/core/cuda/sageattention/permuted_smem.cuh @@ -0,0 +1,202 @@ +/* + * Adapted from Flashinfer, https://github.com/flashinfer-ai/flashinfer/blob/v0.1.5/include/flashinfer/permuted_smem.cuh + * Copyright (c) 2023 by FlashInfer team. + * + * Modifications copyright (c) 2024 by SageAttention team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef SD_CORE_CUDA_SAGEATTENTION_PERMUTED_SMEM_CUH +#define SD_CORE_CUDA_SAGEATTENTION_PERMUTED_SMEM_CUH +#include +#include +#include + + +#include "cp_async.cuh" +#include "mma.cuh" + +namespace sageattention { + +enum class SwizzleMode { + k32B, // for k32B mode, a line of shared memory must have 32B (16 half value) + k64B, // for k64B mode, a line of shared memory must have 64B (32 half value) + k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B. +}; + +// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth +using b128_t = uint4; + +/*! + * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure + * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through + * permutation in the granularity of 128bit. + * + * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless, + * meaning it does not maintain any information about the current pointer position. The offset returnd by + * the struct can be used to access the shared memory through the provided interface. + * + * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j). + * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp. + * Under this assumption, we do not permute for k32B mode. + */ +template +struct smem_t { + // The base pointer. + b128_t* base; + // How many b128_t value a row contains + // uint32_t stride; + + __device__ __forceinline__ smem_t() : base(nullptr) {} + template + __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode"); + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode"); + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode"); + } else { + static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode"); + } + } + + /*! + * \brief Set the base pointer. + */ + template + __device__ __forceinline__ void set_base(T* new_base) { + base = (b128_t*)new_base; + } + + /*! + * \brief Compute the element offset given coordinates in a permuted shared memory. + * \param i The row index. + * \param j The column index. + */ + static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t &i, const uint32_t &j) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + return i * stride + (j ^ (i % 8)); + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + return i * stride + (j ^ ((i / 2) % 4)); + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + return i * stride + j; + } + } + + /*! + * \tparam step_size The step size to advance the offset in the permuted shared memory. + * \param offset The current offset. + */ + template + static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size % 8 == 0, + "Unsupported step size"); + return offset + step_size; + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(step_size == 4, "Unsupported step size"); + return offset + step_size; + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + static_assert(step_size == 2, "Unsupported step size"); + return offset + step_size; + } + } + + // ! use with care + template + static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset, const uint32_t &step_idx) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0, + "Unsupported step size"); + if constexpr (step_size == 2) { + return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8; + } else if constexpr (step_size == 4) { + return (offset ^ 0x4) + (step_idx % 2 == 1) * 8; + } else { + // step_size % 8 == 0 + return offset + step_size; + } + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(step_size == 2 || step_size == 4, "Unsupported step size"); + if constexpr (step_size == 2) { + return (offset ^ 0x2) + (step_idx % 2 == 1) * 4; + } else { + return offset + step_size; + } + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + return offset + step_size; + } + } + + template + static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t &offset) { + if constexpr (swizzle_mode == SwizzleMode::k128B) { + static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); + if constexpr (step_size == 4) { + return (offset ^ 0x4) + step_size * stride; + } else { + // step_size % 8 == 0 + return offset + step_size * stride; + } + } else if constexpr (swizzle_mode == SwizzleMode::k64B) { + static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size"); + if constexpr (step_size == 4) { + return (offset ^ 0x2) + step_size * stride; + } else { + // step_size % 8 == 0 + return offset + step_size * stride; + } + } else if constexpr (swizzle_mode == SwizzleMode::k32B) { + return offset + step_size * stride; + } + } + + __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t &offset, uint32_t* R) const { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x2(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t &offset, uint32_t* R) const { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4(R, smem_ptr); + } + + __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t &offset, uint32_t* R) const { + b128_t* smem_ptr = base + offset; + mma::ldmatrix_m8n8x4_trans(R, smem_ptr); + } + + template + __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr, bool predicate) const { + b128_t* smem_ptr = base + offset; + cp_async::pred_load_128b( + smem_ptr, reinterpret_cast(gptr), predicate); + } + + template + __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr) const { + b128_t* smem_ptr = base + offset; + cp_async::load_128b(smem_ptr, reinterpret_cast(gptr)); + } + + template + __device__ __forceinline__ void store_128b(const uint32_t &offset, T* gptr) const { + *reinterpret_cast(gptr) = *(base + offset); + } +}; + +} // namespace sageattention + +#endif diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index b659bb0d..e0214fe4 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -5,6 +5,7 @@ #include "ggml-ext.h" #include +#include #include #include #include @@ -223,6 +224,59 @@ ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int gro return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr); } +static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { + if (count != 3 || !attributes) + return false; + SDSageAttentionParams params; + std::memcpy(¶ms, attributes, sizeof(params)); + if (!std::isfinite(params.scale) || params.scale <= 0.f || params.mode < GGML_SAGE_ATTN_AUTO || params.mode > GGML_SAGE_ATTN_2_PLUS_PLUS) + return false; + for (int i = 0; i < count; ++i) { + const auto* t = inputs[i]; + if (!t || !ggml_is_contiguous(t) || t->type != (i == 2 ? GGML_TYPE_F16 : GGML_TYPE_F32)) + return false; + int64_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (t->ne[d] <= 0 || t->ne[d] > INT64_MAX / int64_t(sizeof(float)) / elements) + return false; + elements *= t->ne[d]; + } + } + const auto* q = inputs[0]; + const auto* k = inputs[1]; + if (q->ne[0] != k->ne[0] || !ggml_are_same_shape(k, inputs[2]) || q->ne[2] % k->ne[2] != 0 || q->ne[3] != k->ne[3]) + return false; + *out = {GGML_TYPE_F32, {q->ne[0], q->ne[2], q->ne[1], q->ne[3]}}; + return true; +} + +ggml_ext_op_t sd_ggml_sage_attn_op() { + static const auto op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.sage_attn.v1", sizeof(SDSageAttentionParams), sage_attn_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* params) { + if (!node || ggml_ext_op_get(node) != sd_ggml_sage_attn_op() || node->src[3] || !ggml_is_contiguous(node)) + return false; + ggml_ext_tensor_desc output; + const auto* attributes = ggml_ext_op_attributes(node); + if (!sage_attn_infer(node->src, 3, attributes, &output) || node->type != output.type || std::memcmp(node->ne, output.ne, sizeof(output.ne))) + return false; + std::memcpy(params, attributes, sizeof(*params)); + return true; +} + +ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode) { + ggml_tensor* inputs[] = {q, k, v}; + const SDSageAttentionParams params = {scale, mode}; + return ggml_ext_op_create(ctx, sd_ggml_sage_attn_op(), inputs, 3, ¶ms); +} + void sd_ggml_register_extensions(ggml_backend_t backend) { if (!backend) { return; diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index 208c717e..d069f10e 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -34,6 +34,23 @@ bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation); bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size); ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size, ggml_tensor* input_scale = nullptr); ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size); + +enum ggml_sage_attn_mode { + GGML_SAGE_ATTN_AUTO = 0, + GGML_SAGE_ATTN_2, + GGML_SAGE_ATTN_2_PLUS_PLUS, +}; + +struct SDSageAttentionParams { + float scale; + ggml_sage_attn_mode mode; +}; + +ggml_ext_op_t sd_ggml_sage_attn_op(); +bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* params); +// Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout. +// Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference. +ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode); #endif #endif