refactor: implement SageAttention through backend extensions

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
+Remaining model-specific CUDA paths still need auditing or migration. This
-need auditing or migration.
+implementation does not add dynamic dtypes.
-This implementation does not yet externalize all model-specific kernels or 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.

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)

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);

src/core/cuda/sage_attention.cu

@@ -0,0 +1,387 @@
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+#if __CUDACC_VER_MAJOR__ >= 12
+
+#include <algorithm>
+#include <climits>
+#include <cstring>
+#include <map>
+#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 <int D>
+    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 <int D, int TOKENS, bool SMOOTH>
+    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 <int D, bool KEY>
+    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 <int D>
+    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 <int D>
+    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 <int D, bool PLUS_PLUS>
+    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<int8_t*>(data + work.q);
+        auto* ki  = reinterpret_cast<int8_t*>(data + work.k);
+        auto* vi  = reinterpret_cast<int8_t*>(data + work.v);
+        auto* qs  = reinterpret_cast<float*>(data + work.q_scale);
+        auto* ks  = reinterpret_cast<float*>(data + work.k_scale);
+        auto* vs  = reinterpret_cast<float*>(data + work.v_scale);
+        auto* km  = reinterpret_cast<float*>(data + work.k_mean);
+        auto* out = reinterpret_cast<half*>(data + work.output);
+        sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
+        sage_quantize_per_thread<D, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), nullptr, qi, qs, lq);
+        sage_quantize_per_thread<D, true><<<dim3(padded / 64, hk * batch), 128, 0, stream>>>(
+            static_cast<const float*>(k->data), km, ki, ks, lk);
+        // The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++.
+        sage_value_scale<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f);
+        sage_quantize_value<D><<<dim3(padded / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(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<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), 256 * D, stream>>>(
+            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<float*>(dst->data), ggml_nelements(dst));
+    }
+
+    template <int D>
+    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<D, true>(stream, dst, data, work, scale);
+            return;
+        }
+        if (mode == GGML_SAGE_ATTN_2) {
+            sage_attn_fp8<D, false>(stream, dst, data, work, scale);
+            return;
+        }
+        auto* qi  = reinterpret_cast<int8_t*>(data + work.q);
+        auto* ki  = reinterpret_cast<int8_t*>(data + work.k);
+        auto* qs  = reinterpret_cast<float*>(data + work.q_scale);
+        auto* ks  = reinterpret_cast<float*>(data + work.k_scale);
+        auto* km  = reinterpret_cast<float*>(data + work.k_mean);
+        auto* out = reinterpret_cast<half*>(data + work.output);
+
+        sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
+        sage_quantize<D, 32, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), nullptr, qi, qs, lq);
+        sage_quantize<D, 64, true><<<dim3(GGML_PAD(lk, 64) / 64, hk * batch), 256, 0, stream>>>(
+            static_cast<const float*>(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<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>(
+            qi, ki, static_cast<half*>(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<float*>(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<const SageDevice*>(user_data);
+        SDSageAttentionParams params;
+        if (!sd_ggml_sage_attn_describe(op, &params) || 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<uintptr_t>(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], &params))
+            return GGML_STATUS_FAILED;
+        const auto mode = sage_mode(*static_cast<const SageDevice*>(user_data), params.mode);
+        auto* data      = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+        auto stream     = static_cast<cudaStream_t>(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<int, SageDevice> 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

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.

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 `<cuda/pipeline>` 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.

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<uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN,
+        typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean=false>
+__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<DTypeSVAccum, float>::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float");
+  static_assert(std::is_same<DTypeSVAccum, float>::value || std::is_same<DTypeSVAccum, half>::value, "DTypeSVAccum must be float or half");
+  static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::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<DTypeSVAccum, half>::value, "fuse_v_mean only supports half");
+  static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation
+
+  using DTypeOut2 = typename std::conditional<std::is_same<DTypeOut, half>::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<num_warps_q, num_warps_k>();
+  }
+  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<num_warps_q, num_warps_k>() * 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<num_warps_q, num_warps_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<num_warps_q, num_warps_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<DTypeSVAccum, float>::value)
+      {
+#pragma unroll
+        for (uint32_t k = 0; k < 8; k++)
+        {
+          RO[fq][fv][k] = 0.0f;
+        }
+      }
+      else if constexpr (std::is_same<DTypeSVAccum, half>::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<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
+  smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> 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<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> 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<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> 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<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
+  uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_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<num_warps_q, num_warps_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<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+  uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_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<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
+    &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<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+    &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<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+    &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        smem_K, RS, RQ, K_smem_offset_mma);
+    }
+    else
+    {
+      compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        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<DTypeSVAccum, float>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, sm_scale);
+    }
+    else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, sm_scale);
+    }
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+    }
+
+    uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+    RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+    }
+
+    __syncthreads();
+
+    // load K
+    load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+        smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    }
+    else
+    {
+      compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+        smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    }
+
+    __syncthreads();
+    // load V
+    load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        smem_K, RS, RQ, K_smem_offset_mma);
+    }
+    else
+    {
+      compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        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<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+    }
+    // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+    K_idx_lane_base += CTA_K;
+
+    if constexpr (std::is_same<DTypeSVAccum, float>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+    else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+    }
+
+    uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+    RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+    }
+
+    __syncthreads();
+
+    // load K with predicate
+    load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+        smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    }
+    else
+    {
+      compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+        smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    }
+
+    __syncthreads();
+    // load V with predicate
+    load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        smem_K, RS, RQ, K_smem_offset_mma);
+    }
+    else
+    {
+      compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        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<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+    }
+    // check out of bound in the last iter
+    apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+    K_idx_lane_base += CTA_K;
+
+    if constexpr (std::is_same<DTypeSVAccum, float>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+    else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+    }
+
+    uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+    RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+    }
+
+    // ensure V is ready
+    cp_async::wait_group<0>();
+    __syncthreads();
+
+    if constexpr (!use_inst_buffer)
+    {
+      compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+        smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    }
+    else
+    {
+      compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+        smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    }
+
+    __syncthreads();
+
+  }
+
+  // TODO: thread block sync mdo state for num_warps_k > 0
+
+  normalize_d<num_tiles_q, num_tiles_v, DenominatorAccumUnit>(RO, m, d);
+
+  // save the result
+  // if (get_warp_idx_k<num_warps_q, num_warps_k>() == 0)
+  // {
+
+  // convert half to bfloat16
+  if constexpr (std::is_same<DTypeSVAccum, half>::value && std::is_same<DTypeOut, nv_bfloat16>::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<num_warps_q, num_warps_k>() * 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<DTypeSVAccum, float>::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<DTypeOut, half>::value)
+          {
+            ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
+          }
+          else if constexpr (std::is_same<DTypeOut, nv_bfloat16>::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<DTypeSVAccum, half>::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<num_warps_q, num_warps_k>() + 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<num_warps_q, num_warps_k>() * 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<global_to_shared_line_lanes_O>(offset_O);
+    }
+
+    offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(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<num_warps_q, num_warps_k>();
+    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

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<uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN,
+        typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale=false, bool fuse_v_mean=false, bool use_pv_fp16_accu=false>
+__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<DTypeSVAccum, float>::value, "DTypeSVAccum must be float, half is WIP");
+  static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::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<num_warps_q, num_warps_k>();
+  }
+  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<num_warps_q, num_warps_k>() * 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<num_warps_q, num_warps_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<num_warps_q, num_warps_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<DTypeSVAccum, float>::value)
+      {
+#pragma unroll
+        for (uint32_t k = 0; k < 8; k++)
+        {
+          RO[fq][fv][k] = 0.0f;
+        }
+      }
+      else if constexpr (std::is_same<DTypeSVAccum, half>::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<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
+  smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> 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<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> 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<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> 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<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
+  uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_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<num_warps_q, num_warps_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<num_warps_q, num_warps_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<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+  uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_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<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
+    &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<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+    &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<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+    &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        smem_K, RS, RQ, K_smem_offset_mma);
+    }
+    else
+    {
+      compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        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<DTypeSVAccum, float>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, sm_scale);
+    }
+    else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, sm_scale);
+    }
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+    }
+
+    uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+    RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+    {
+      accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+    }
+
+    __syncthreads();
+
+    // load K without predicate
+    load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+    //   smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    if constexpr (!use_inst_buffer)
+    {
+      compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+        smem_V, RS_f8, RO, d);
+    }
+    else
+    {
+      if constexpr (!use_pv_fp16_accu){
+        compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+          smem_V, RS_f8, RO, d);
+      }
+      else{
+        compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+          smem_V, RS_f8, RO, d);
+      }
+    }
+    __syncthreads();
+    // load V
+    // for fp16:
+    // load_global_to_share                stride_seq_v
+    load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        smem_K, RS, RQ, K_smem_offset_mma);
+    }
+    else
+    {
+      compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        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<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+    }
+    // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+    K_idx_lane_base += CTA_K;
+
+    if constexpr (std::is_same<DTypeSVAccum, float>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+    else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+    }
+
+    uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+    RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+    {
+      accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+    }
+
+    __syncthreads();
+
+    // load K with predicate
+    load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+    //   smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    if constexpr (!use_inst_buffer)
+    {
+      compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+        smem_V, RS_f8, RO, d);
+    }
+    else
+    {
+      if constexpr (!use_pv_fp16_accu){
+        compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+          smem_V, RS_f8, RO, d);
+      }
+      else{
+        compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+          smem_V, RS_f8, RO, d);
+      }
+    }
+
+    __syncthreads();
+    // load V
+    // for fp16:
+    // load_global_to_share                stride_seq_v
+    load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+      &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        smem_K, RS, RQ, K_smem_offset_mma);
+    }
+    else
+    {
+      compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+        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<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+    }
+    apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+    K_idx_lane_base += CTA_K;
+
+    if constexpr (std::is_same<DTypeSVAccum, float>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+    else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+    {
+      update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
+    }
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+    {
+      accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+    }
+
+    uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+    RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+    if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+    {
+      accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+    }
+
+    // ensure V is ready
+    cp_async::wait_group<0>();
+    __syncthreads();
+
+    // for fp16:
+    // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+    //   smem_V, RS_f16, RO, d, V_smem_offset_mma);
+    if constexpr (!use_inst_buffer)
+    {
+      compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+        smem_V, RS_f8, RO, d);
+    }
+    else
+    {
+      if constexpr (!use_pv_fp16_accu){
+        compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+          smem_V, RS_f8, RO, d);
+      }
+      else{
+        compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+          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<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(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<num_warps_q, num_warps_k>() * 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<DTypeSVAccum, float>::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<DTypeOut, half>::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<DTypeSVAccum, half>::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<num_warps_q, num_warps_k>() + 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<num_warps_q, num_warps_k>() * 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<global_to_shared_line_lanes_O>(offset_O);
+    }
+
+    offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(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<num_warps_q, num_warps_k>();
+    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

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 <cuda_fp16.h>
+#include <cuda_pipeline_primitives.h>
+
+#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 <uint32_t num_warps_q, uint32_t num_warps_k>
+__device__ __forceinline__ uint32_t get_warp_idx_q()
+{
+  return get_warp_id() / num_warps_k;
+}
+
+template <uint32_t num_warps_q, uint32_t num_warps_k>
+__device__ __forceinline__ uint32_t get_warp_idx_k()
+{
+  return get_warp_id() % num_warps_k;
+}
+
+template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter,
+          uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
+__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset,
+                                                    const uint32_t &gmem_stride,
+                                                    const smem_t<swizzle_mode, stride> &smem)
+{
+  static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+  static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
+
+  constexpr uint32_t pack_size = std::is_same<T, half>::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<global_to_shared_line_lanes>(smem_offset);
+    }
+
+    smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(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 <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter,
+          uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
+__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset,
+                                                    const uint32_t &gmem_stride,
+                                                    const smem_t<swizzle_mode, stride> &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<T, half>::value || std::is_same<T, int8_t>::value);
+
+  constexpr uint32_t pack_size = std::is_same<T, half>::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<std::is_same<T, half>::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<global_to_shared_line_lanes>(smem_offset);
+    }
+
+    smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(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 <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter,
+          uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA>
+__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<swizzle_mode, stride> &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<global_to_shared_line_lanes>(smem_offset);
+    }
+
+    smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(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 <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner,
+          SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
+__device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride> &smem_Q, const smem_t<swizzle_mode, stride> &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<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+        }
+        else if constexpr (DTypeQK == DataType::kInt4)
+        {
+          mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(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<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
+        }
+        else if constexpr (DTypeQK == DataType::kInt4)
+        {
+          mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInplaceUpdate>(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 <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner,
+          SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
+__device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride> &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<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+      }
+      else if constexpr (DTypeQK == DataType::kInt4)
+      {
+        mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+      }
+    }
+  }
+}
+
+template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
+__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<DTypeQKAccum, float>::value)
+        {
+          RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
+        }
+        else if constexpr (std::is_same<DTypeQKAccum, half>::value)
+        {
+          RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
+        }
+      }
+    }
+  }
+}
+
+template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
+__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<DTypeQKAccum, float>::value)
+        {
+          RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
+        }
+        else if constexpr (std::is_same<DTypeQKAccum, half>::value)
+        {
+          RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
+        }
+      }
+    }
+  }
+}
+
+// for DTypeQKAccum float
+template <uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, bool use_half_o_scale, bool exp_offset, bool fuse_scale=false, typename DTypeSVAccum>
+__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<DTypeSVAccum, half>::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<DTypeSVAccum, float>::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<DTypeSVAccum, half>::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 <uint32_t num_tiles_q, uint32_t num_tiles_k, typename T>
+__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 <uint32_t num_tiles_q, uint32_t num_tiles_k>
+__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 <uint32_t num_tiles_q, uint32_t num_tiles_k>
+__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 <uint32_t num_tiles_q, uint32_t num_tiles_k>
+__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 <uint32_t num_tiles_q, uint32_t num_tiles_k, ComputeUnit compute_unit = ComputeUnit::kTensorCore, typename T>
+__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<T, float>::value && compute_unit == ComputeUnit::kCudaCore) ||
+                (std::is_same<T, uint32_t>::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 <uint32_t num_tiles_q, uint32_t num_tiles_k>
+__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 <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+          SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+__device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, stride> &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_warps_q, num_warps_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<DTypeSVAccum, float>::value)
+        {
+          mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+        {
+          mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV);
+        }
+      }
+    }
+  }
+}
+
+template <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+          SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width=4, typename T, typename DTypeSVAccum>
+__device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mode, stride> &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<DTypeSVAccum, float>::value)
+        {
+          mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::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 <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+          SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width=4, typename T, typename DTypeSVAccum>
+__device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<swizzle_mode, stride> &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<DTypeSVAccum, float>::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<mma::MMAMode::kInit>((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<mma::MMAMode::kInplaceUpdate>((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<uint32_t num_tiles_q, uint32_t num_tiles_v,
+       ComputeUnit compute_unit = ComputeUnit::kTensorCore, // compute unit for accumulate_d
+       typename DTypeQKAccum, typename DTypeSVAccum>
+__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<DTypeSVAccum, float>::value)
+        {
+          RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2];
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+        {
+          RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]);
+        }
+      }
+    }
+  }
+}
+
+template <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+          SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+__device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride> &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<num_warps_q, num_warps_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<DTypeSVAccum, float>::value)
+        {
+          mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::value)
+        {
+          // ! Not Implemented
+        }
+      }
+      offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+    }
+  }
+}
+
+template <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+          SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+__device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mode, stride> &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<num_warps_q, num_warps_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<DTypeSVAccum, float>::value)
+        {
+          mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::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<DTypeSVAccum, float>::value)
+        {
+          mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::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 <uint32_t num_warps_q, uint32_t num_warps_k,
+          uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+          SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+__device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<swizzle_mode, stride> &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<num_warps_q, num_warps_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<DTypeSVAccum, float>::value)
+        {
+          //mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+          mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInit>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::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<DTypeSVAccum, float>::value)
+        {
+          //mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+          mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInplaceUpdate>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
+        }
+        else if constexpr (std::is_same<DTypeSVAccum, half>::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

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 <cuda_runtime.h>
+
+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 <size_t n>
+__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 <PrefetchMode prefetch_mode, typename T>
+__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
+#ifdef CP_ASYNC_ENABLED
+  uint32_t smem_int_ptr = static_cast<uint32_t>(__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 <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+__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<uint32_t>(__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

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 <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+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

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+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 <typename T>
+__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X2_ENABLED
+  uint32_t smem_int_ptr = static_cast<uint32_t>(__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 <typename T>
+__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X4_ENABLED
+  uint32_t smem_int_ptr = static_cast<uint32_t>(__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 <typename T>
+__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X4_ENABLED
+  uint32_t smem_int_ptr = static_cast<uint32_t>(__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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<uint32_t*>(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 <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+__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

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+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<const int8_t&>(dst);
+}
+
+} // namespace sageattention
+
+#endif

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+
+#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 <SwizzleMode swizzle_mode, uint32_t stride>
+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 <typename T>
+  __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 <typename T>
+  __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 <uint32_t step_size>
+  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 <uint32_t step_size>
+  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 <uint32_t step_size>
+  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 <cp_async::SharedMemFillMode fill_mode, typename T>
+  __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<cp_async::PrefetchMode::kPrefetch, fill_mode>(
+        smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
+  }
+
+  template <typename T>
+  __device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr) const {
+    b128_t* smem_ptr = base + offset;
+    cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
+  }
+
+  template <typename T>
+  __device__ __forceinline__ void store_128b(const uint32_t &offset, T* gptr) const {
+    *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
+  }
+};
+
+} // namespace sageattention
+
+#endif

src/core/ggml_extensions.cpp

@@ -5,6 +5,7 @@
 #include "ggml-ext.h"
 
 #include <algorithm>
+#include <cmath>
 #include <cstring>
 #include <mutex>
 #include <unordered_set>
@@ -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(&params, 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, &params);
+}
+
 void sd_ggml_register_extensions(ggml_backend_t backend) {
     if (!backend) {
         return;

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