refactor: move model-specific operators and fusions into backend extensions (#2253)

src/core/compute_workspace.cpp

@@ -6,9 +6,9 @@
 #include <unordered_map>
 #include <unordered_set>
 
+#include "core/ggml_compat.h"
 #include "core/ggml_extend_backend.h"
 #include "core/util.h"
-#include "ggml-impl.h"
 
 namespace sd {
     ComputeWorkspace::~ComputeWorkspace() {
@@ -69,7 +69,7 @@ namespace sd {
 
     bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) {
         GGML_ASSERT(!active_);
-        const size_t capacity = static_cast<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+        const size_t capacity = static_cast<size_t>(ggml_graph_n_nodes(graph) + ggml_graph_n_leafs(graph)) + 8;
         if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
             release();
             scheduler_          = make_scheduler(capacity);
@@ -101,11 +101,11 @@ namespace sd {
                 tensors.push_back(tensor);
             }
         };
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            visit(graph->nodes[i]);
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            visit(ggml_graph_node(graph, i));
         }
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            visit(graph->leafs[i]);
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            visit(ggml_graph_leaf(graph, i));
         }
         for (size_t i = 0; i < tensors.size(); ++i) {
             visit(tensors[i]->view_src);
@@ -120,9 +120,8 @@ namespace sd {
                 const auto* bytes = reinterpret_cast<const uint8_t*>(&value);
                 key.insert(key.end(), bytes, bytes + sizeof(value));
             };
-            append(graph->n_nodes);
-            append(graph->n_leafs);
-            append(graph->order);
+            append(ggml_graph_n_nodes(graph));
+            append(ggml_graph_n_leafs(graph));
             for (auto tensor : tensors) {
                 append(tensor);
                 append(tensor->type);
@@ -147,11 +146,11 @@ namespace sd {
                         append(usage);
                         append(type);
                     }
+                } else {
+                    const auto alignment_offset = reinterpret_cast<uintptr_t>(tensor->data) % ggml_backend_get_alignment(external);
+                    append(alignment_offset);
                 }
-                const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
-                const int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
-                                         ? graph->use_counts[index]
-                                         : 0;
+                const int32_t uses = ggml_graph_use_count(graph, tensor);
                 append(uses);
             }
             // Match the full allocation metadata, including topology and external placement.
@@ -187,28 +186,18 @@ namespace sd {
                     ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS);
                 }
                 copy->buffer = buffer;
-                copy->data   = reinterpret_cast<void*>(static_cast<uintptr_t>(1));
-                copy->extra  = nullptr;
-            }
-        }
-        auto copy_graph     = ggml_new_graph_custom(context, graph_size, false);
-        copy_graph->n_nodes = graph->n_nodes;
-        copy_graph->n_leafs = graph->n_leafs;
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            copy_graph->nodes[i] = copies.at(graph->nodes[i]);
-        }
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            copy_graph->leafs[i] = copies.at(graph->leafs[i]);
-        }
-        copy_graph->order = graph->order;
-        // Preserve fusion metadata, including uses outside a graph view, without rebuilding its node order.
-        for (const auto& entry : copies) {
-            const size_t source_index = ggml_hash_find(&graph->visited_hash_set, entry.first);
-            if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, source_index)) {
-                const size_t copy_index            = ggml_hash_insert(&copy_graph->visited_hash_set, entry.second);
-                copy_graph->use_counts[copy_index] = graph->use_counts[source_index];
+                // Capability checks inspect alignment even before weights are resident.
+                copy->data  = source->data ? source->data : reinterpret_cast<void*>(ggml_backend_get_alignment(external));
+                copy->extra = nullptr;
             }
         }
+        auto copy_graph = ggml_new_graph_custom(context, graph_size, false);
+        auto map_tensor = [](const ggml_tensor* tensor, void* user_data) -> ggml_tensor* {
+            const auto& mapped = *static_cast<const std::unordered_map<const ggml_tensor*, ggml_tensor*>*>(user_data);
+            const auto found   = mapped.find(tensor);
+            return found == mapped.end() ? nullptr : found->second;
+        };
+        GGML_ASSERT(ggml_graph_cpy_mapped(graph, copy_graph, map_tensor, &copies));
         Measurement result;
         auto scheduler = make_scheduler(graph_size);
         if (scheduler != nullptr) {

src/core/cpu/int8.cpp

@@ -0,0 +1,181 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+    static void transform(float* values, int group) {
+        const float scale = 1.f / std::sqrt(float(group));
+        for (int i = 0; i < group; ++i)
+            values[i] *= scale;
+        for (int stride = 1; stride < group; stride *= 4)
+            for (int base = 0; base < group; base += 4 * stride)
+                for (int j = 0; j < stride; ++j) {
+                    const int i   = base + j;
+                    const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+                    values[i]              = a + b + c - d;
+                    values[i + stride]     = a + b - c + d;
+                    values[i + 2 * stride] = a - b + c + d;
+                    values[i + 3 * stride] = -a + b + c + d;
+                }
+    }
+    struct I8Execution {
+        SDI8Operation op;
+        int64_t rows;
+        int8_t* quant;
+        float* scales;
+        float* output;
+    };
+    static void quantize(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e           = *static_cast<I8Execution*>(data);
+        const int64_t k   = e.op.input->ne[0];
+        const int group   = e.op.group_size;
+        const auto* input = static_cast<const float*>(e.op.input->data);
+        for (int64_t row = ith; row < e.rows; row += nth) {
+            const float* src = input + row * k;
+            int8_t* dst      = e.quant + row * k;
+            float amax       = 0.f, values[256];
+            auto visit       = [&](auto fn) { if (!group) { for (int64_t i = 0; i < k; ++i) fn(i, src[i]); }
+            else for (int64_t base = 0; base < k; base += group) { std::memcpy(values, src + base, group * sizeof(float)); transform(values, group); for (int j = 0; j < group; ++j) fn(base + j, values[j]); } };
+            visit([&](int64_t, float value) { amax = std::max(amax, std::fabs(value)); });
+            const float scale = amax / 127.f, inverse = scale == 0.f ? 0.f : 1.f / scale;
+            e.scales[row] = scale;
+            if (scale == 0.f)
+                std::memset(dst, 0, k);
+            else
+                visit([&](int64_t i, float value) { dst[i] = static_cast<int8_t>(std::max(-127L, std::min(127L, std::lrintf(value * inverse)))); });
+        }
+    }
+    static void multiply(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e         = *static_cast<I8Execution*>(data);
+        const auto& op  = e.op;
+        const int64_t k = op.weight->ne[0], n = op.weight->ne[1];
+        const auto* w    = static_cast<const int8_t*>(op.weight->data);
+        const auto* ws   = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+        const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+        for (int64_t index = ith; index < e.rows * n; index += nth) {
+            const int64_t row = index / n, column = index % n;
+            uint32_t sum = 0;
+            for (int64_t i = 0; i < k; ++i)
+                sum += uint32_t(int32_t(w[column * k + i]) * int32_t(e.quant[row * k + i]));
+            int32_t signed_sum;
+            std::memcpy(&signed_sum, &sum, sizeof(sum));
+            float value = float(signed_sum) * e.scales[row];
+            if (ws)
+                value *= ws[column];
+            if (bias)
+                value += bias[column];
+            e.output[index] = value;
+        }
+    }
+    static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+        if (!src)
+            return false;
+        const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+        return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+    }
+    static size_t quant_workspace(const ggml_tensor* node, const SDI8Operation& op) {
+        return op.input->type == GGML_TYPE_F32 ? (size_t(ggml_nelements(op.input)) + 3) / 4 * 4 + size_t(ggml_nrows(node)) * 4 : 0;
+    }
+    static bool i8_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDI8Operation op;
+        return sd_ggml_i8_describe(node, &op);
+    }
+    static bool convrot_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        int group;
+        return sd_ggml_convrot_describe(node, &group);
+    }
+    static ggml_status i8_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        SDI8Operation op;
+        if (!sd_ggml_i8_describe(nodes[0], &op))
+            return GGML_STATUS_FAILED;
+        *bytes = quant_workspace(nodes[0], op) + ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static ggml_status convrot_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void*, void*) {}
+    static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
+        return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
+    }
+    static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        I8Execution e = {};
+        if (!sd_ggml_i8_describe(dst, &e.op))
+            return GGML_STATUS_FAILED;
+        e.rows   = ggml_nrows(dst);
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
+        char* scratch           = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
+        const size_t quant_size = quant_workspace(dst, e.op);
+        if (e.op.input->type == GGML_TYPE_F32) {
+            e.quant     = reinterpret_cast<int8_t*>(scratch);
+            e.scales    = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
+            auto status = run(context->backend, dst, quantize, &e);
+            if (status != GGML_STATUS_SUCCESS)
+                return status;
+        } else {
+            e.quant  = static_cast<int8_t*>(e.op.input->data);
+            e.scales = e.op.packed ? reinterpret_cast<float*>(e.quant + e.op.input->ne[0] * ((e.rows + 3) / 4 * 4)) : static_cast<float*>(e.op.input_scale->data);
+        }
+        bool alias = false;
+        for (auto* input : dst->src)
+            alias |= overlaps(dst, input);
+        e.output    = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
+        auto status = run(context->backend, dst, multiply, &e);
+        if (status == GGML_STATUS_SUCCESS && alias)
+            std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+        return status;
+    }
+    static ggml_status convrot_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        I8Execution e = {};
+        if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
+            return GGML_STATUS_FAILED;
+        auto run = parallel(context->backend);
+        if (!run)
+            return GGML_STATUS_FAILED;
+        e.op.input      = dst->src[0];
+        e.rows          = ggml_nrows(e.op.input);
+        bool alias      = overlaps(dst, e.op.input);
+        e.quant         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
+        const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
+        e.scales    = reinterpret_cast<float*>(e.quant + k * padded);
+        auto status = run(context->backend, dst, quantize, &e);
+        if (status == GGML_STATUS_SUCCESS) {
+            std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
+            if (alias)
+                std::memcpy(dst->data, e.quant, ggml_nbytes(dst));
+        }
+        return status;
+    }
+    void register_i8(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.priority             = 100;
+        d.destroy              = destroy;
+        d.name                 = "sd.linear_i8.cpu.v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_i8_op();
+        d.supports             = i8_supports;
+        d.prepare              = i8_prepare;
+        d.execute              = i8_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+        d.name      = "sd.quantize_i8_convrot.cpu.v1";
+        d.op        = GGML_OP_CUSTOM;
+        d.custom_op = sd_ggml_convrot_op();
+        d.supports  = convrot_supports;
+        d.prepare   = convrot_prepare;
+        d.execute   = convrot_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+
+}
+#endif

src/core/cpu/registration.cpp

@@ -0,0 +1,12 @@
+#include "core/cpu/registration.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+namespace sd_cpu {
+
+    void register_extensions(ggml_backend_dev_t device) {
+        register_i8(device);
+        register_rope(device);
+    }
+
+}
+#endif

src/core/cpu/registration.h

@@ -0,0 +1,14 @@
+#ifndef __SD_CORE_CPU_REGISTRATION_H__
+#define __SD_CORE_CPU_REGISTRATION_H__
+
+#include "ggml-backend.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+namespace sd_cpu {
+    void register_extensions(ggml_backend_dev_t device);
+    void register_rope(ggml_backend_dev_t device);
+    void register_i8(ggml_backend_dev_t device);
+}
+#endif
+
+#endif

src/core/cpu/rope.cpp

@@ -0,0 +1,139 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+    static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        return sd_ggml_rope_describe(node);
+    }
+
+    struct RopeExecution {
+        const ggml_tensor* rope;
+        const ggml_tensor* norm;
+        const ggml_tensor* weight;
+    };
+
+    static void rope_compute(ggml_tensor* dst, int ith, int nth, void* user_data) {
+        const auto& execution = *static_cast<const RopeExecution*>(user_data);
+        const auto* rope      = execution.rope;
+        const auto* x         = execution.norm ? execution.norm->src[0] : rope->src[0];
+        const auto* pe        = rope->src[1];
+        int32_t interleaved;
+        std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
+        float eps = 0.0f;
+        if (execution.norm) {
+            std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
+        }
+        const int64_t half_dim = pe->ne[2];
+        const int64_t rows     = x->ne[1] * x->ne[2] * x->ne[3];
+        for (int64_t row = ith; row < rows; row += nth) {
+            const int64_t token = row % x->ne[2];
+            const int64_t head  = (row / x->ne[2]) % x->ne[1];
+            const int64_t batch = row / (x->ne[2] * x->ne[1]);
+            const char* input   = static_cast<const char*>(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3];
+            const char* pos     = static_cast<const char*>(pe->data) + token * pe->nb[3];
+            float* output       = static_cast<float*>(dst->data) + row * x->ne[0];
+            float scale         = 1.0f;
+            if (execution.norm) {
+                double sum = 0.0;
+                for (int64_t col = 0; col < x->ne[0]; ++col) {
+                    const float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+                    sum += static_cast<double>(v * v);
+                }
+                const float mean = static_cast<float>(sum / x->ne[0]);
+                scale            = 1.0f / std::sqrt(mean + eps);
+            }
+            auto value = [&](int64_t col) {
+                float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+                if (execution.norm) {
+                    const auto* weight = static_cast<const char*>(execution.weight->data);
+                    v                  = (v * scale) * *reinterpret_cast<const float*>(weight + col * execution.weight->nb[0]);
+                }
+                return v;
+            };
+            for (int64_t j = 0; j < half_dim; ++j) {
+                const int64_t i0 = interleaved ? 2 * j : j;
+                const int64_t i1 = interleaved ? 2 * j + 1 : j + half_dim;
+                const float a    = value(i0);
+                const float b    = value(i1);
+                const char* m    = pos + j * pe->nb[2];
+                const float c    = *reinterpret_cast<const float*>(m);
+                const float s    = *reinterpret_cast<const float*>(m + pe->nb[1]);
+                if (interleaved) {
+                    const float neg_s = *reinterpret_cast<const float*>(m + pe->nb[0]);
+                    const float c1    = *reinterpret_cast<const float*>(m + pe->nb[0] + pe->nb[1]);
+                    output[i0]        = a * c + b * neg_s;
+                    output[i1]        = a * s + b * c1;
+                } else {
+                    output[i0] = a * c - b * s;
+                    output[i1] = a * s + b * c;
+                }
+            }
+            for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) {
+                output[col] = value(col);
+            }
+        }
+    }
+
+    static int rope_fusion_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+        auto* view = const_cast<ggml_cgraph*>(graph);
+        if (first + 2 >= ggml_graph_n_nodes(view)) {
+            return 0;
+        }
+        const auto* norm = ggml_graph_node(view, first);
+        const auto* mul  = ggml_graph_node(view, first + 1);
+        const auto* rope = ggml_graph_node(view, first + 2);
+        if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] ||
+            norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) ||
+            mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul ||
+            ggml_ext_op_get(rope) != sd_ggml_rope_op() ||
+            !sd_ggml_rope_describe(rope) || !ggml_are_same_shape(norm, mul)) {
+            return 0;
+        }
+        const auto* weight = mul->src[0] == norm ? mul->src[1] : (mul->src[1] == norm ? mul->src[0] : nullptr);
+        if (!weight || weight->type != GGML_TYPE_F32 || weight->ne[0] != norm->ne[0] ||
+            weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) {
+            return 0;
+        }
+        return 3;
+    }
+
+    static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
+        auto device = ggml_backend_get_device(context->backend);
+        auto reg    = ggml_backend_dev_backend_reg(device);
+        auto run    = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        GGML_ASSERT(run);
+        RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
+        if (count == 3) {
+            execution.norm   = nodes[0];
+            execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
+        }
+        return run(context->backend, nodes[count - 1], rope_compute, &execution);
+    }
+
+    void register_rope(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.rope_apply.cpu.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_rope_op();
+        desc.priority             = 100;
+        desc.supports             = rope_supports;
+        desc.execute              = rope_execute;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+
+        desc.name      = "sd.rms_norm_mul_rope.cpu.v1";
+        desc.op        = GGML_OP_RMS_NORM;
+        desc.custom_op = nullptr;
+        desc.priority  = 200;
+        desc.supports  = nullptr;
+        desc.match     = rope_fusion_match;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    }
+
+}
+#endif

src/core/cuda/CMakeLists.txt

@@ -0,0 +1,84 @@
+cmake_policy(SET CMP0079 NEW)
+
+function(sd_cuda_sources VENDOR LANGUAGE)
+    set(SOURCES)
+    foreach(OP rope int8 convrot affine registration)
+        set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
+        string(TOUPPER "${VENDOR}" UPPER_VENDOR)
+        file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
+        set_source_files_properties("${SOURCE}" PROPERTIES LANGUAGE ${LANGUAGE})
+        list(APPEND SOURCES "${SOURCE}")
+    endforeach()
+    set(TARGET "sd-${VENDOR}")
+    add_library(${TARGET} OBJECT ${SOURCES})
+    target_include_directories(${TARGET} PRIVATE "${PROJECT_SOURCE_DIR}/src")
+    target_link_libraries(${TARGET} PRIVATE ggml)
+    set_target_properties(${TARGET} PROPERTIES POSITION_INDEPENDENT_CODE ON)
+    target_compile_features(${TARGET} PRIVATE cxx_std_17)
+    target_link_libraries(${SD_LIB} PRIVATE ${TARGET})
+    set(SD_CUDA_TARGET ${TARGET} PARENT_SCOPE)
+    set(SD_CUDA_SOURCES ${SOURCES} PARENT_SCOPE)
+endfunction()
+
+if(SD_CUDA)
+    enable_language(CUDA)
+    sd_cuda_sources(cuda CUDA)
+    target_sources(${SD_CUDA_TARGET} PRIVATE sage_attention.cu sol_attention.cu)
+    set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math")
+    if(SD_CUDA_CUTLASS_PATH)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
+        set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${PROJECT_SOURCE_DIR}/src/core/cuda")
+    endif()
+    set_target_properties(${SD_CUDA_TARGET} PROPERTIES CUDA_STANDARD 17 CUDA_STANDARD_REQUIRED ON)
+    target_link_libraries(${SD_CUDA_TARGET} PRIVATE CUDA::cudart CUDA::cublas)
+    target_link_libraries(${SD_LIB} PRIVATE CUDA::cudart CUDA::cublas)
+endif()
+
+if(SD_HIPBLAS)
+    list(APPEND CMAKE_PREFIX_PATH "$ENV{ROCM_PATH}" /opt/rocm /usr)
+    find_package(hip REQUIRED)
+    find_package(hipblas REQUIRED)
+    if(WIN32 OR CMAKE_CXX_COMPILER MATCHES "hipcc")
+        sd_cuda_sources(hip CXX)
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::device)
+    else()
+        sd_cuda_sources(hip HIP)
+        set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
+    endif()
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
+    target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
+    target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
+endif()
+
+if(SD_MUSA)
+    if(NOT MUSA_PATH)
+        if(DEFINED ENV{MUSA_PATH})
+            set(MUSA_PATH "$ENV{MUSA_PATH}")
+        elseif(EXISTS /opt/musa)
+            set(MUSA_PATH /opt/musa)
+        else()
+            set(MUSA_PATH /usr/local/musa)
+        endif()
+    endif()
+    list(APPEND CMAKE_MODULE_PATH "${MUSA_PATH}/cmake")
+    find_package(MUSAToolkit REQUIRED)
+    set(CMAKE_CXX_COMPILER "${MUSA_PATH}/bin/clang++")
+    sd_cuda_sources(musa CXX)
+    set(FLAGS -Od3 -fno-strict-aliasing -ffast-math -fsigned-char -x musa -mtgpu -fmusa-flush-denormals-to-zero)
+    if(NOT DEFINED MUSA_ARCHITECTURES)
+        set(MUSA_ARCHITECTURES "22;31")
+    endif()
+    foreach(ARCH ${MUSA_ARCHITECTURES})
+        list(APPEND FLAGS "--cuda-gpu-arch=mp_${ARCH}")
+    endforeach()
+    set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "${FLAGS}")
+    if(GGML_STATIC)
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart_static MUSA::mublas_static)
+        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart_static MUSA::mublas_static)
+    else()
+        target_link_libraries(${SD_CUDA_TARGET} PRIVATE MUSA::musart MUSA::mublas)
+        target_link_libraries(${SD_LIB} PRIVATE MUSA::musart MUSA::mublas)
+    endif()
+endif()

src/core/cuda/affine.cu

@@ -0,0 +1,183 @@
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include "ggml-ext.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_affine_cuda {
+        static bool disabled() {
+            static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+            return value;
+        }
+        static bool aligned(const ggml_tensor* tensor, size_t alignment) {
+            size_t offset = 0;
+            while (!tensor->data && tensor->view_src) {
+                offset += tensor->view_offs;
+                tensor = tensor->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+        }
+        static float param(const ggml_tensor* tensor, int index) {
+            float value;
+            std::memcpy(&value, tensor->op_params + index, sizeof(value));
+            return value;
+        }
+        static bool contiguous_f32(const ggml_tensor* tensor) {
+            return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
+        }
+        static bool bounded(const ggml_tensor* tensor) {
+            return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
+        }
+        static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (disabled() || first + 1 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* node = ggml_graph_node(g, first);
+            const auto* next = ggml_graph_node(g, first + 1);
+            if (node->op == GGML_OP_MUL) {
+                if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node))
+                    return 0;
+                const auto* x        = node->src[0];
+                const auto* gate     = node->src[1];
+                const auto* residual = next->src[0] == node ? next->src[1] : next->src[0];
+                return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) &&
+                               node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) &&
+                               ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] &&
+                               x->ne[0] % 4 == 0 && bounded(next)
+                           ? 2
+                           : 0;
+            }
+            if (node->op == GGML_OP_CONCAT) {
+                const bool permuted = next->op == GGML_OP_PERMUTE;
+                const int count     = permuted ? 3 : 2;
+                if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 ||
+                    node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) ||
+                    !contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1]))
+                    return 0;
+                if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] ||
+                                 next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] ||
+                                 next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3]))
+                    return 0;
+                const auto* dst = ggml_graph_node(g, first + count - 1);
+                return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 &&
+                               ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst)
+                           ? count
+                           : 0;
+            }
+            if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) {
+                const auto* src     = node->src[0];
+                const auto* permute = ggml_graph_node(g, first + 2);
+                const auto* dst     = ggml_graph_node(g, first + 3);
+                return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 &&
+                               aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 &&
+                               next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f &&
+                               permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 &&
+                               permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 &&
+                               dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 &&
+                               ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst)
+                           ? 4
+                           : 0;
+            }
+            return 0;
+        }
+        static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i >= count4)
+                return;
+            const float4 a = x[i], g = gate[i % n4], r = residual[i];
+            // The graph performs two FP32 operations; contracting them into FMA changes rounding.
+            dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
+                                 __fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
+        }
+        template <bool Permuted>
+        static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i >= count4)
+                return;
+            const int64_t col = i % columns4, row = i / columns4;
+            const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
+            const bool from_prefix   = token < prefix_tokens;
+            const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
+            const int64_t src_token  = from_prefix ? token : token - prefix_tokens;
+            const int64_t src_row    = Permuted ? (batch * src_tokens + src_token) * heads + head
+                                                : (batch * heads + head) * src_tokens + src_token;
+            const float4 value       = (from_prefix ? prefix : current)[src_row * columns4 + col];
+            dst[2 * i]               = __floats2half2_rn(value.x, value.y);
+            dst[2 * i + 1]           = __floats2half2_rn(value.z, value.w);
+        }
+        static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
+            const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
+            if (index >= count)
+                return;
+            const int64_t col = index % dim, row = index / dim;
+            const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
+            const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
+            dst[index / 2]     = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
+            dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+            *workspace = ggml_nbytes(nodes[count - 1]);
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void*, void*) {}
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context))
+                return GGML_STATUS_FAILED;
+            auto stream           = static_cast<CUDAStream>(context.stream);
+            const auto* node      = nodes[0];
+            const auto* a         = node->src[0];
+            const auto* b         = node->op == GGML_OP_CONT ? nullptr : node->src[1];
+            auto* dst             = nodes[count - 1];
+            const auto* residual  = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
+            const bool gate       = residual != nullptr;
+            const bool alias      = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
+                                    (residual && overlaps(dst, residual) && dst->data != residual->data);
+            void* output          = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+            const int64_t count4  = ggml_nelements(dst) / 4;
+            const unsigned blocks = unsigned((count4 + 255) / 256);
+            if (gate) {
+                gated_residual<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
+                                                           static_cast<const float4*>(residual->data), static_cast<float4*>(output), a->ne[0] / 4, count4);
+            } else if (node->op == GGML_OP_CONCAT) {
+                const bool permuted = count == 3;
+                auto kernel         = permuted ? concat_cast<true> : concat_cast<false>;
+                kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
+                                                   static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
+                                                   node->ne[permuted ? 1 : 2], count4);
+            } else {
+                scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
+                                                          a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
+            }
+            auto error = cuda_last_error();
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+
+    void register_affine(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.priority             = 200;
+        desc.match                = sd_affine_cuda::match;
+        desc.prepare              = sd_affine_cuda::prepare;
+        desc.destroy              = sd_affine_cuda::destroy;
+        desc.execute              = sd_affine_cuda::execute;
+        desc.name                 = "sd.gated_residual." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_MUL;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
+        desc.op   = GGML_OP_CONCAT;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
+        desc.op   = GGML_OP_CONT;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}

src/core/cuda/convrot.cu

@@ -0,0 +1,296 @@
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include "convrot_kernel.h"
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_convrot {
+
+        static bool fusion_disabled() {
+            static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+            return disabled;
+        }
+        static float param_f32(const ggml_tensor* tensor, int index) {
+            float value;
+            std::memcpy(&value, tensor->op_params + index, sizeof(value));
+            return value;
+        }
+        static bool f32(const ggml_tensor* tensor) {
+            return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
+        }
+        static bool vector(const ggml_tensor* tensor, int64_t k) {
+            return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k;
+        }
+        enum Kind { SCALE,
+                    NORM,
+                    RMS_MOD,
+                    SWIGLU,
+                    SILU,
+                    SIGMOID };
+        struct Operation {
+            Kind kind;
+            const ggml_tensor* input      = nullptr;
+            const ggml_tensor* up         = nullptr;
+            const ggml_tensor* modulation = nullptr;
+            const ggml_tensor* weight     = nullptr;
+            const ggml_tensor* shift      = nullptr;
+            int64_t k, rows, stride, up_stride, gate_offset = 0, up_offset = 0;
+            float eps = 0.f, scale = 1.f, bias = 0.f;
+        };
+        static bool describe(ggml_tensor* const* nodes, int count, Operation& op) {
+            int group;
+            const auto* first = nodes[0];
+            const auto* quant = nodes[count - 1];
+            if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
+                return false;
+            const auto* value = quant->src[0];
+            op.k              = value->ne[0];
+            op.rows           = ggml_nrows(value);
+            op.input          = first->src[0];
+            if (!f32(op.input) || !f32(first) || op.k > INT_MAX || op.rows > INT_MAX)
+                return false;
+            op.stride    = op.input->ne[0];
+            op.up_stride = op.k;
+            if (count == 2 && first->op == GGML_OP_SCALE && value == first) {
+                op.kind  = SCALE;
+                op.scale = param_f32(first, 0);
+                op.bias  = param_f32(first, 1);
+            } else if (count == 3 && first->op == GGML_OP_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && value == nodes[1]) {
+                op.kind       = NORM;
+                op.modulation = nodes[1]->src[1];
+                op.eps        = param_f32(first, 0);
+                if (op.k < 1024 || !vector(op.modulation, op.k))
+                    return false;
+            } else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first &&
+                       nodes[2]->op == GGML_OP_MUL && nodes[2]->src[0] == nodes[1] &&
+                       nodes[3]->op == GGML_OP_ADD && nodes[3]->src[0] == nodes[1] && nodes[3]->src[1] == nodes[2] &&
+                       nodes[4]->op == GGML_OP_ADD && nodes[4]->src[0] == nodes[3] && value == nodes[4]) {
+                op.kind       = RMS_MOD;
+                op.weight     = nodes[1]->src[1];
+                op.modulation = nodes[2]->src[1];
+                op.shift      = nodes[4]->src[1];
+                op.eps        = param_f32(first, 0);
+                if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !vector(op.shift, op.k))
+                    return false;
+            } else if (count == 2 && first->op == GGML_OP_GLU && ggml_get_glu_op(first) == GGML_GLU_OP_SWIGLU && value == first) {
+                op.kind = SWIGLU;
+                op.up   = first->src[1] ? first->src[1] : op.input;
+                if (!f32(op.up))
+                    return false;
+                op.up_stride = op.up->ne[0];
+                if (!first->src[1]) {
+                    if (first->op_params[1])
+                        op.gate_offset = op.k;
+                    else
+                        op.up_offset = op.k;
+                }
+            } else if (count == 3 && first->op == GGML_OP_UNARY && nodes[1]->op == GGML_OP_MUL && value == nodes[1] &&
+                       (nodes[1]->src[0] == first || nodes[1]->src[1] == first)) {
+                auto unary = ggml_get_unary_op(first);
+                if (unary != GGML_UNARY_OP_SILU && unary != GGML_UNARY_OP_SIGMOID)
+                    return false;
+                op.kind = unary == GGML_UNARY_OP_SILU ? SILU : SIGMOID;
+                op.up   = nodes[1]->src[0] == first ? nodes[1]->src[1] : nodes[1]->src[0];
+                if (!f32(op.up) || !ggml_are_same_shape(op.input, op.up))
+                    return false;
+            } else
+                return false;
+            for (int i = 0; i < count - 1; ++i) {
+                if (!f32(nodes[i]) || !ggml_are_same_shape(nodes[i], value))
+                    return false;
+                for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
+                    if (input == nodes[i])
+                        return false;
+            }
+            return true;
+        }
+        static bool supported(const Operation& op, const ggml_ext_gpu_device_info& device) {
+            const size_t extra = op.kind == NORM ? 384 : op.kind == RMS_MOD ? 256
+                                                                            : 128;
+            return device.device >= 0 && size_t(op.k) * sizeof(float) + extra <= device.shared_memory;
+        }
+        static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+            if (fusion_disabled())
+                return 0;
+            auto* g             = const_cast<ggml_cgraph*>(graph);
+            const int available = ggml_graph_n_nodes(g) - first;
+            const auto* node    = ggml_graph_node(g, first);
+            const int count     = node->op == GGML_OP_RMS_NORM ? 6 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3
+                                                                                                                           : 2;
+            if (available < count)
+                return 0;
+            ggml_tensor* nodes[6];
+            for (int j = 0; j < count; ++j)
+                nodes[j] = ggml_graph_node(g, first + j);
+            Operation op;
+            return describe(nodes, count, op) && supported(op, cuda_device_info(device)) ? count : 0;
+        }
+        static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+            if (fusion_disabled())
+                return;
+            for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) {
+                auto* quant = ggml_graph_node(graph, i);
+                int group;
+                if (!sd_ggml_convrot_describe(quant, &group) || group != 256)
+                    continue;
+                ggml_tensor* chain[6] = {};
+                int count             = 0;
+                auto* value           = quant->src[0];
+                if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) {
+                    auto* add    = value->src[0];
+                    auto* weight = add->src[0];
+                    if (weight->op == GGML_OP_MUL) {
+                        chain[0] = weight->src[0];
+                        chain[1] = weight;
+                        chain[2] = add->src[1];
+                        chain[3] = add;
+                        chain[4] = value;
+                        chain[5] = quant;
+                        count    = 6;
+                    }
+                } else if (value->op == GGML_OP_MUL) {
+                    for (int src = 0; src < 2; ++src)
+                        if (value->src[src]->op == GGML_OP_UNARY) {
+                            chain[0] = value->src[src];
+                            chain[1] = value;
+                            chain[2] = quant;
+                            count    = 3;
+                            break;
+                        }
+                }
+                Operation op;
+                if (!count || !describe(chain, count, op) || !supported(op, cuda_device_info(device)))
+                    continue;
+                int begin = -1;
+                for (int j = std::max(0, i - 64); j < i; ++j)
+                    if (ggml_graph_node(graph, j) == chain[0]) {
+                        begin = j;
+                        break;
+                    }
+                if (begin < 0)
+                    continue;
+                // Move the owning producer with its in-place consumers so alias checks retain their storage relationship.
+                for (int part = 1; part < count; ++part) {
+                    int next = -1;
+                    for (int j = begin + part; j <= i; ++j)
+                        if (ggml_graph_node(graph, j) == chain[part]) {
+                            next = j;
+                            break;
+                        }
+                    if (next < 0 || (next > begin + part && !ggml_ext_graph_delay(graph, begin, part, next)))
+                        break;
+                    begin = next - part;
+                }
+            }
+        }
+        struct Plan {
+            ggml_ext_gpu_device_info device;
+            Operation op;
+            bool configured = false;
+        };
+        static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+            auto* plan   = new Plan{};
+            *state       = plan;
+            plan->device = cuda_device_info(ggml_backend_get_device(backend));
+            if (!describe(nodes, count, plan->op) || !supported(plan->op, plan->device))
+                return GGML_STATUS_FAILED;
+            *workspace = ggml_nbytes(nodes[count - 1]);
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void* state, void*) {
+            delete static_cast<Plan*>(state);
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            if (!b)
+                return false;
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+        template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
+        static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) {
+            constexpr bool warp = Scaled || RMS;
+            constexpr int extra = Normalize ? 384 : RMS ? 256
+                                                        : 128;
+            if (!p.configured) {
+                auto status = cuda_set_shared_memory(quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
+                                                     int(p.device.shared_memory - extra));
+                if (status != cuda_success)
+                    return status;
+                p.configured = true;
+            }
+            const int threads    = Normalize ? 1024 : RMS                               ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
+                                                                                                            : op.k <= 8192   ? 512
+                                                                                                                             : 1024)
+                                                  : Swiglu && !Sigmoid && op.k >= 16384 ? 512
+                                                                                        : 256;
+            const int64_t padded = (op.rows + 3) / 4 * 4;
+            auto data            = [](const ggml_tensor* t) { return t ? static_cast<const float*>(t->data) : nullptr; };
+            quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
+                data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
+                op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
+            return cuda_last_error();
+        }
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+            auto& plan     = *static_cast<Plan*>(state);
+            const auto& op = plan.op;
+            auto* dst      = nodes[count - 1];
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != plan.device.device)
+                return GGML_STATUS_FAILED;
+            auto stream = static_cast<CUDAStream>(context.stream);
+            bool alias  = false;
+            for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
+                alias |= overlaps(dst, input);
+            auto* output         = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
+            const int64_t padded = (op.rows + 3) / 4 * 4;
+            if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success)
+                return GGML_STATUS_FAILED;
+            CUDAError error;
+            switch (op.kind) {
+                case SCALE:
+                    error = launch<false, false, true>(plan, op, output, stream);
+                    break;
+                case NORM:
+                    error = launch<false, true>(plan, op, output, stream);
+                    break;
+                case RMS_MOD:
+                    error = launch<false, false, false, false, true>(plan, op, output, stream);
+                    break;
+                case SWIGLU:
+                case SILU:
+                    error = launch<true>(plan, op, output, stream);
+                    break;
+                case SIGMOID:
+                    error = launch<true, false, false, true>(plan, op, output, stream);
+                    break;
+                default:
+                    return GGML_STATUS_FAILED;
+            }
+            if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success))
+                return GGML_STATUS_FAILED;
+            return GGML_STATUS_SUCCESS;
+        }
+    }
+
+    void register_convrot(ggml_backend_dev_t device) {
+        for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
+            ggml_ext_kernel_desc desc = {};
+            desc.abi_version          = GGML_EXT_ABI_VERSION;
+            const std::string name    = std::string("sd.convrot_fusion.") + ggml_op_name(op) + "." SD_CUDA_NAME ".v1";
+            desc.name                 = name.c_str();
+            desc.op                   = op;
+            desc.priority             = 225;
+            desc.optimize             = op == GGML_OP_SCALE ? sd_convrot::optimize : nullptr;
+            desc.match                = sd_convrot::match;
+            desc.prepare              = sd_convrot::prepare;
+            desc.execute              = sd_convrot::execute;
+            desc.destroy              = sd_convrot::destroy;
+            GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        }
+    }
+
+}

src/core/cuda/convrot_kernel.h

@@ -0,0 +1,251 @@
+#ifndef __SD_CORE_CUDA_CONVROT_KERNEL_H__
+#define __SD_CORE_CUDA_CONVROT_KERNEL_H__
+#include <cstdint>
+#include "runtime.h"
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_convrot {
+        static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1)
+                x += cuda_shuffle_xor(x, offset, 32);
+            return x;
+        }
+        static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
+#pragma unroll
+            for (int offset = 16; offset > 0; offset >>= 1) {
+                x.x += cuda_shuffle_xor(x.x, offset, 32);
+                x.y += cuda_shuffle_xor(x.y, offset, 32);
+            }
+            return x;
+        }
+        template <int Width>
+        static __device__ __forceinline__ float warp_reduce_max(float x) {
+#pragma unroll
+            for (int offset = Width / 2; offset > 0; offset >>= 1)
+                x = fmaxf(x, cuda_shuffle_xor(x, offset, Width));
+            return x;
+        }
+        static __device__ float2 block_sum(float2 x, float2* shared) {
+            x = warp_reduce_sum(x);
+            if (threadIdx.x % 32 == 0)
+                shared[threadIdx.x / 32] = x;
+            __syncthreads();
+            return warp_reduce_sum(shared[threadIdx.x % 32]);
+        }
+        static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+            a = index == 3 ? -a : a;
+            b = index == 2 ? -b : b;
+            c = index == 1 ? -c : c;
+            d = index == 0 ? -d : d;
+            return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+        }
+
+        template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false>
+        __global__ void quantize(
+            const T* src,
+            int8_t* dst,
+            float* scales,
+            int64_t k,
+            int64_t rows,
+            int64_t dst_stride,
+            const float* up,
+            int64_t src_stride,
+            int64_t up_stride,
+            const float* activation_scales = nullptr,
+            const float* weight_scales     = nullptr,
+            const float* modulation        = nullptr,
+            float eps                      = 0.0f,
+            float input_scale              = 1.0f,
+            float input_bias               = 0.0f,
+            const float* rms_weight        = nullptr,
+            const float* rms_shift         = nullptr) {
+            constexpr int group_size = 256;
+
+            const int64_t row = blockIdx.x;
+            const int tid     = threadIdx.x;
+            if (row >= rows) {
+                return;
+            }
+
+            extern __shared__ float values[];
+            __shared__ float maxima[32];
+            const T* src_row = src + row * src_stride;
+            int8_t* dst_row  = dst + row * dst_stride;
+
+            [[maybe_unused]] float rms_scale = 1.0f;
+            if constexpr (rms_modulation) {
+                __shared__ float sums[32];
+                // Preserve the unfused RMSNorm's 1024-thread reduction order with fewer physical warps.
+                for (int virtual_tid = tid; virtual_tid < 1024; virtual_tid += blockDim.x) {
+                    float partial = 0.0f;
+                    for (int64_t i = virtual_tid; i < k; i += 1024) {
+                        const float value = src_row[i];
+                        partial += value * value;
+                    }
+                    partial = warp_reduce_sum(partial);
+                    if (tid % 32 == 0) {
+                        sums[virtual_tid / 32] = partial;
+                    }
+                }
+                __syncthreads();
+                const float sum = warp_reduce_sum(sums[tid % 32]);
+                rms_scale       = rsqrtf(sum / k + eps);
+            }
+
+            if constexpr (swiglu || warp_rotation) {
+                const int lane = tid % 32;
+                for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+                    float v[8];
+#pragma unroll
+                    for (int j = 0; j < 8; ++j) {
+                        const int64_t i = group * group_size + j * 32 + lane;
+                        if constexpr (!swiglu) {
+                            float value = src_row[i];
+                            if constexpr (rms_modulation) {
+                                value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]);
+                                value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[i]);
+                            }
+                            if constexpr (scaled) {
+                                value = __fmaf_rn(value, input_scale, input_bias);
+                            }
+                            v[j] = value * (1.0f / 16.0f);
+                        } else {
+                            float gate_value, up_value;
+                            if constexpr (dequantize) {
+                                gate_value = (float)((const int32_t*)src_row)[i] * activation_scales[row];
+                                up_value   = (float)((const int32_t*)up)[row * up_stride + i] * activation_scales[row];
+                                gate_value *= weight_scales[i];
+                                up_value *= weight_scales[k + i];
+                            } else {
+                                gate_value = src_row[i];
+                                up_value   = up[row * up_stride + i];
+                            }
+                            const float activated = sigmoid_gate ? 1.0f / (1.0f + expf(-gate_value)) : (gate_value / (1.0f + expf(-gate_value)));
+                            v[j]                  = (activated * up_value) * (1.0f / 16.0f);
+                        }
+                    }
+                    // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+                    for (int stride = 1; stride <= 4; stride *= 4) {
+                        const int base = lane & ~(3 * stride);
+#pragma unroll
+                        for (int j = 0; j < 8; ++j) {
+                            const float a = cuda_shuffle(v[j], base, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+                            const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
+                            v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+                        }
+                    }
+#pragma unroll
+                    for (int j = 0; j < 8; j += 2) {
+                        const float a = cuda_shuffle(v[j], lane % 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+                        const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
+                        v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
+                        v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+                    }
+#pragma unroll
+                    for (int j = 0; j < 2; ++j) {
+                        const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+                        for (int t = 0; t < 4; ++t) {
+                            values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+                        }
+                    }
+                }
+                __syncthreads();
+            } else {
+                constexpr int group_threads   = group_size / 4;
+                constexpr int groups_per_wave = 1024 / group_threads;
+                if constexpr (normalize) {
+                    __shared__ float2 sums[32];
+                    float2 mean_var = make_float2(0.0f, 0.0f);
+                    for (int64_t i = tid; i < k; i += 1024) {
+                        const float value = src_row[i];
+                        mean_var.x += value;
+                        mean_var.y += value * value;
+                    }
+                    mean_var             = block_sum(mean_var, sums);
+                    const float mean     = mean_var.x / k;
+                    const float variance = mean_var.y / k - mean * mean;
+                    const float inv_std  = rsqrtf(variance + eps);
+                    for (int64_t i = tid; i < k; i += 1024) {
+                        const float normalized = __fmul_rn(__fsub_rn(src_row[i], mean), inv_std);
+                        values[i]              = __fmul_rn(__fmul_rn(normalized, modulation[i]), 1.0f / 16.0f);
+                    }
+                } else {
+                    for (int64_t i = tid; i < k; i += blockDim.x) {
+                        const float value = scaled ? __fmaf_rn(src_row[i], input_scale, input_bias) : src_row[i];
+                        values[i]         = value * (1.0f / 16.0f);
+                    }
+                }
+                __syncthreads();
+
+                const int group_lane = tid % group_threads;
+                const int wave_group = tid / group_threads;
+                const int64_t groups = k / group_size;
+                for (int64_t group_base = 0; group_base < groups; group_base += groups_per_wave) {
+                    const int64_t group = group_base + wave_group;
+#pragma unroll
+                    for (int stride = 1; stride < group_size; stride *= 4) {
+                        if (group < groups) {
+                            float* group_values = values + group * group_size;
+                            const int base      = (group_lane / stride) * 4 * stride + group_lane % stride;
+                            const int i0        = base;
+                            const int i1        = i0 + stride;
+                            const int i2        = i1 + stride;
+                            const int i3        = i2 + stride;
+                            const float a       = group_values[i0];
+                            const float b       = group_values[i1];
+                            const float c       = group_values[i2];
+                            const float d       = group_values[i3];
+                            group_values[i0]    = a + b + c - d;
+                            group_values[i1]    = a + b - c + d;
+                            group_values[i2]    = a - b + c + d;
+                            group_values[i3]    = -a + b + c + d;
+                        }
+                        __syncthreads();
+                    }
+                }
+            }
+
+            float local_max = 0.0f;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                local_max = fmaxf(local_max, fabsf(values[i]));
+            }
+            local_max      = warp_reduce_max<32>(local_max);
+            const int lane = tid & 31;
+            const int warp = tid >> 5;
+            if (lane == 0) {
+                maxima[warp] = local_max;
+            }
+            __syncthreads();
+
+            if (warp == 0) {
+                local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+                local_max = warp_reduce_max<32>(local_max);
+                if (lane == 0) {
+                    maxima[0] = local_max / 127.0f;
+                    if (scales != nullptr) {
+                        scales[row] = maxima[0];
+                    } else {
+                        *((float*)(dst_row + k)) = maxima[0];
+                    }
+                }
+            }
+            __syncthreads();
+
+            const float row_scale = maxima[0];
+            const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                int value  = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+                value      = max(-127, min(127, value));
+                dst_row[i] = (int8_t)value;
+            }
+        }
+
+    }
+}
+#endif

src/core/cuda/cute_attention.cu

@@ -0,0 +1,792 @@
+/*
+ * Adapted from FlashAttention v2.8.3:
+ * https://github.com/Dao-AILab/flash-attention/tree/v2.8.3/csrc/flash_attn/src
+ * Sources: kernel_traits.h, utils.h, softmax.h, flash_fwd_kernel.h.
+ * Copyright (c) 2023-2024, Tri Dao.
+ *
+ * BSD 3-Clause License
+ *
+ * Copyright (c) 2022, the respective contributors, as shown by the AUTHORS file.
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * * Redistributions of source code must retain the above copyright notice, this
+ *   list of conditions and the following disclaimer.
+ *
+ * * Redistributions in binary form must reproduce the above copyright notice,
+ *   this list of conditions and the following disclaimer in the documentation
+ *   and/or other materials provided with the distribution.
+ *
+ * * Neither the name of the copyright holder nor the names of its
+ *   contributors may be used to endorse or promote products derived from
+ *   this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+ * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+ * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+ * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+ * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+ * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cutlass/numeric_conversion.h>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cute/tensor.hpp>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+#include "core/cuda/i8_projection.h"
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+#include "ggml-ext.h"
+
+static __device__ __forceinline__ float sd_warp_max(float v) {
+    for (int mask = 16; mask > 0; mask >>= 1)
+        v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
+    return v;
+}
+namespace sd_fattn_cute {
+    using namespace cute;
+    template <bool A_in_regs = false, bool B_in_regs = false, typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename Tensor4, typename TiledMma, typename TiledCopyA, typename TiledCopyB, typename ThrCopyA, typename ThrCopyB>
+    __forceinline__ __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA, Tensor4 const& tCsB, TiledMma tiled_mma, TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) {
+        CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+        CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+        CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+        Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+        CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+        Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+        CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+        if (!A_in_regs) {
+            cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
+        }
+        if (!B_in_regs) {
+            cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+        }
+#pragma unroll
+        for (int i = 0; i < size<2>(tCrA); ++i) {
+            if (i < size<2>(tCrA) - 1) {
+                if (!A_in_regs) {
+                    cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
+                }
+                if (!B_in_regs) {
+                    cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+                }
+            }
+            cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+        }
+    }
+
+    template <typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename TiledMma, typename TiledCopy, typename ThrCopy>
+    __forceinline__ __device__ void gemm_rs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsB, TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, ThrCopy smem_thr_copy_B) {
+        CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+        CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+        CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+        Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+        CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+        cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+#pragma unroll
+        for (int i = 0; i < size<2>(tCrA); ++i) {
+            if (i < size<2>(tCrA) - 1) {
+                cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+            }
+            cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+        }
+    }
+
+    template <typename Layout>
+    __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
+        static_assert(decltype(size<0>(acc_layout))::value == 4);
+        static_assert(decltype(rank(acc_layout))::value == 3);
+        auto l = logical_divide(acc_layout, Shape<_2>{});
+        return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
+    };
+
+    template <typename MMA_traits, typename Layout>
+    __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
+        using X = Underscore;
+        static_assert(decltype(size<0>(acc_layout))::value == 4);
+        static_assert(decltype(rank(acc_layout))::value == 3);
+        constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
+        static_assert(mma_shape_K == 8 || mma_shape_K == 16);
+        if constexpr (mma_shape_K == 8) {
+            return acc_layout;
+        } else {
+            auto l = logical_divide(acc_layout, Shape<X, X, _2>{});
+            return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
+        }
+    };
+
+    template <typename TensorT>
+    __device__ __forceinline__ auto to_half(const TensorT& src) {
+        constexpr int count = decltype(size(src))::value;
+        auto dst            = make_tensor<cutlass::half_t>(src.layout());
+        cutlass::NumericArrayConverter<cutlass::half_t, float, count> convert;
+        *reinterpret_cast<cutlass::Array<cutlass::half_t, count>*>(dst.data()) =
+            convert(*reinterpret_cast<const cutlass::Array<float, count>*>(src.data()));
+        return dst;
+    }
+
+    template <bool Full, bool Clear, bool FullCols = true, typename Copy, typename Src, typename Dst, typename Id>
+    __device__ __forceinline__ void copy_rows(Copy copier, const Src& src, Dst& dst, const Id& id, int rows, int cols = 0) {
+#pragma unroll
+        for (int m = 0; m < size<1>(src); ++m) {
+#pragma unroll
+            for (int k = 0; k < size<2>(src); ++k) {
+                if ((Full || get<0>(id(0, m, k)) < rows) && (FullCols || get<1>(id(0, m, k)) < cols)) {
+                    cute::copy(copier, src(_, m, k), dst(_, m, k));
+                } else if constexpr (Clear) {
+                    cute::clear(dst(_, m, k));
+                }
+            }
+        }
+    }
+
+    template <int Rows>
+    struct Softmax {
+        decltype(make_tensor<float>(Shape<Int<Rows>>{})) maximum, sum;
+        template <bool First, typename Scores, typename Output>
+        __device__ __forceinline__ void update(Scores& acc_s, Output& acc_o) {
+            auto scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+            auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+            for (int m = 0; m < Rows; ++m) {
+                const float old_max = First ? -INFINITY : maximum(m);
+                float next_max      = First ? scores(m, 0) : fmaxf(old_max, scores(m, 0));
+#pragma unroll
+                for (int n = 1; n < size<1>(scores); ++n)
+                    next_max = fmaxf(next_max, scores(m, n));
+                next_max      = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 2));
+                next_max      = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 1));
+                maximum(m)    = next_max;
+                float row_sum = 0.f;
+                if constexpr (!First) {
+                    row_sum = sum(m);
+                    if (old_max != next_max) {
+                        const float factor = exp2f((old_max - next_max) * 1.4426950408889634f);
+                        row_sum *= factor;
+#pragma unroll
+                        for (int n = 0; n < size<1>(output); ++n)
+                            output(m, n) *= factor;
+                    }
+                }
+#pragma unroll
+                for (int n = 0; n < size<1>(scores); ++n) {
+                    const float value = exp2f(scores(m, n) * 1.4426950408889634f - next_max * 1.4426950408889634f);
+                    scores(m, n)      = value;
+                    row_sum += value;
+                }
+                // Defer the four-lane sum reduction until the final output normalization.
+                sum(m) = row_sum;
+            }
+        }
+        template <typename Output>
+        __device__ __forceinline__ void normalize(Output& acc_o) {
+            auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+            for (int m = 0; m < Rows; ++m) {
+                float total = sum(m);
+                total += __shfl_xor_sync(0xffffffff, total, 2);
+                total += __shfl_xor_sync(0xffffffff, total, 1);
+                const float inverse = 1.f / total;
+#pragma unroll
+                for (int n = 0; n < size<1>(output); ++n)
+                    output(m, n) *= inverse;
+            }
+        }
+    };
+
+    template <int BM, int BN, int D = 128, int Warps = 4>
+    struct Traits {
+        using Element                    = cutlass::half_t;
+        static constexpr int Threads     = Warps * 32;
+        using Mma                        = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
+                                                    Layout<Shape<Int<Warps>, _1, _1>>,
+                                                    Tile<Int<16 * Warps>, _16, _16>>;
+        static constexpr int SmemWidth   = D % 64 == 0 ? 64 : D % 32 == 0 ? 32
+                                                                          : 16;
+        static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2
+                                                                                 : 1;
+        using SmemAtom                   = decltype(composition(Swizzle<SwizzleBits, 3, 3>{},
+                                                                Layout<Shape<_8, Int<SmemWidth>>, Stride<Int<SmemWidth>, _1>>{}));
+        using QLayout                    = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BM>, Int<D>>{}));
+        using KVLayout                   = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BN>, Int<D>>{}));
+        using VtLayout                   = decltype(composition(KVLayout{}, make_layout(Shape<Int<D>, Int<BN>>{}, GenRowMajor{})));
+        using CopyAtom                   = Copy_Atom<SM75_U32x4_LDSM_N, Element>;
+        using CopyAtomT                  = Copy_Atom<SM75_U16x8_LDSM_T, Element>;
+        using GlobalLayout               = Layout<Shape<Int<Threads / (SmemWidth / 8)>, Int<SmemWidth / 8>>, Stride<Int<SmemWidth / 8>, _1>>;
+        using GlobalCopy                 = decltype(make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, Element>{},
+                                                                    GlobalLayout{},
+                                                                    Layout<Shape<_1, _8>>{}));
+        using OutputCopy                 = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+                                                                    GlobalLayout{},
+                                                                    Layout<Shape<_1, _8>>{}));
+        static constexpr int SharedBytes = (BM + 2 * BN) * D * sizeof(Element);
+    };
+
+    struct Params {
+        const cutlass::half_t *q, *k, *v;
+        cutlass::half_t* o;
+        int nq, nk, heads, kv_heads, dim;
+    };
+
+    template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
+    __global__ __launch_bounds__(Warps * 32) void kernel(Params p) {
+        using T       = Traits<BM, BN, D, Warps>;
+        using Element = typename T::Element;
+        extern __shared__ char storage[];
+        const int tid = threadIdx.x, m_block = blockIdx.x, batch = blockIdx.y, head = blockIdx.z;
+        const int64_t q_offset  = (int64_t(batch) * p.heads + head) * p.nq * p.dim;
+        const int64_t kv_offset = (int64_t(batch) * p.kv_heads + head / (p.heads / p.kv_heads)) * p.nk * p.dim;
+        auto mQ                 = make_tensor(make_gmem_ptr(p.q + q_offset), make_shape(p.nq, p.dim), make_stride(p.dim, _1{}));
+        auto mK                 = make_tensor(make_gmem_ptr(p.k + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
+        auto mV                 = make_tensor(make_gmem_ptr(p.v + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
+        auto gQ                 = local_tile(mQ, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+        auto gK                 = local_tile(mK, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
+        auto gV                 = local_tile(mV, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
+        auto sQ                 = make_tensor(make_smem_ptr(reinterpret_cast<Element*>(storage)), typename T::QLayout{});
+        auto sK                 = make_tensor(sQ.data() + size(sQ), typename T::KVLayout{});
+        auto sV                 = make_tensor(sK.data() + size(sK), typename T::KVLayout{});
+        auto sVt                = make_tensor(sV.data(), typename T::VtLayout{});
+        auto sVplain            = make_tensor(sV.data().get(), get_nonswizzle_portion(typename T::VtLayout{}));
+        typename T::GlobalCopy global_copy;
+        auto thread_copy = global_copy.get_thread_slice(tid);
+        auto tQgQ        = thread_copy.partition_S(gQ);
+        auto tQsQ        = thread_copy.partition_D(sQ);
+        auto tKgK        = thread_copy.partition_S(gK);
+        auto tKsK        = thread_copy.partition_D(sK);
+        auto tVgV        = thread_copy.partition_S(gV);
+        auto tVsV        = thread_copy.partition_D(sV);
+        auto tQcQ        = thread_copy.partition_S(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+        auto tKVcKV      = thread_copy.partition_S(make_identity_tensor(Shape<Int<BN>, Int<D>>{}));
+        typename T::Mma mma;
+        auto thread_mma = mma.get_thread_slice(tid);
+        auto rQ         = thread_mma.partition_fragment_A(sQ);
+        auto rK         = thread_mma.partition_fragment_B(sK);
+        auto rV         = thread_mma.partition_fragment_B(sVplain);
+        auto acc_o      = partition_fragment_C(mma, Shape<Int<BM>, Int<D>>{});
+        auto copyQ      = make_tiled_copy_A(typename T::CopyAtom{}, mma);
+        auto copyK      = make_tiled_copy_B(typename T::CopyAtom{}, mma);
+        auto copyV      = make_tiled_copy_B(typename T::CopyAtomT{}, mma);
+        auto threadQ    = copyQ.get_thread_slice(tid);
+        auto threadK    = copyK.get_thread_slice(tid);
+        auto threadV    = copyV.get_thread_slice(tid);
+        auto tSsQ       = threadQ.partition_S(sQ);
+        auto tSsK       = threadK.partition_S(sK);
+        auto tOsV       = threadV.partition_S(sVt);
+        auto score_id   = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<BN>>{}));
+        // Process the partial K/V tile first; subsequent tiles need no score masking.
+        int block = (p.nk - 1) / BN;
+        copy_rows<false, true, FullDim>(global_copy, tQgQ, tQsQ, tQcQ, p.nq - m_block * BM, p.dim);
+        auto last_k = tKgK(_, _, _, block);
+        copy_rows<false, true, FullDim>(global_copy, last_k, tKsK, tKVcKV, p.nk - block * BN, p.dim);
+        cute::cp_async_fence();
+        clear(acc_o);
+        Softmax<2 * size<1>(acc_o)> softmax;
+        auto step = [&](auto first_flag, int n) {
+            constexpr bool First = decltype(first_flag)::value;
+            auto acc_s           = partition_fragment_C(mma, Shape<Int<BM>, Int<BN>>{});
+            clear(acc_s);
+            cute::cp_async_wait<0>();
+            __syncthreads();
+            auto v_tile = tVgV(_, _, _, n);
+            copy_rows<!First, First, FullDim>(global_copy, v_tile, tVsV, tKVcKV, p.nk - n * BN, p.dim);
+            cute::cp_async_fence();
+            gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
+            if constexpr (First) {
+#pragma unroll
+                for (int i = 0; i < size(acc_s); ++i) {
+                    if (get<1>(score_id(i)) >= p.nk - n * BN)
+                        acc_s(i) = -INFINITY;
+                }
+            }
+            cute::cp_async_wait<0>();
+            __syncthreads();
+            // Do not commit an empty async group: the next iteration waits for this K tile.
+            if (n > 0) {
+                auto next_k = tKgK(_, _, _, n - 1);
+                copy_rows<true, false, FullDim>(global_copy, next_k, tKsK, tKVcKV, 0, p.dim);
+                cute::cp_async_fence();
+            }
+            softmax.template update<First>(acc_s, acc_o);
+            auto rP   = to_half(acc_s);
+            auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs<typename T::Mma>(rP.layout()));
+            gemm_rs(acc_o, tOrP, rV, tOsV, mma, copyV, threadV);
+        };
+        step(std::true_type{}, block);
+        for (--block; block >= 0; --block)
+            step(std::false_type{}, block);
+        softmax.normalize(acc_o);
+        auto rO      = to_half(acc_o);
+        auto sO      = make_tensor(sQ.data(), typename T::QLayout{});
+        auto copyO   = make_tiled_copy_C(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{}, mma);
+        auto threadO = copyO.get_thread_slice(tid);
+        auto tOrO    = threadO.retile_S(rO);
+        auto tOsO    = threadO.partition_D(sO);
+        cute::copy(copyO, tOrO, tOsO);
+        auto mO = make_tensor(make_gmem_ptr(p.o + int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim),
+                              make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{}));
+        auto gO = local_tile(mO, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+        typename T::OutputCopy output_copy;
+        auto global_threadO = output_copy.get_thread_slice(tid);
+        auto from_sO        = global_threadO.partition_S(sO);
+        auto to_gO          = global_threadO.partition_D(gO);
+        auto output_id      = global_threadO.partition_D(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+        __syncthreads();
+        auto output_regs = make_tensor<Element>(shape(to_gO));
+        cute::copy(output_copy, from_sO, output_regs);
+        copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+    }
+}
+
+static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+    a = index == 3 ? -a : a;
+    b = index == 2 ? -b : b;
+    c = index == 1 ? -c : c;
+    d = index == 0 ? -d : d;
+    return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+}
+
+static __global__ void quantize_attention(const half* src, int8_t* dst, float* scales, int64_t k, int64_t rows, float input_scale, float input_bias) {
+    const int64_t row = blockIdx.x;
+    const int tid     = threadIdx.x;
+    if (row >= rows)
+        return;
+    extern __shared__ float values[];
+    __shared__ float maxima[32];
+    const half* src_row = src + row * k;
+    int8_t* dst_row     = dst + row * k;
+    const int lane_id   = tid % 32;
+    {
+        const int lane           = lane_id;
+        constexpr int group_size = 256;
+        for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+            float v[8];
+#pragma unroll
+            for (int j = 0; j < 8; ++j) {
+                const float value = __fmaf_rn(__half2float(src_row[group * group_size + j * 32 + lane]), input_scale, input_bias);
+                v[j]              = value * (1.0f / 16.0f);
+            }
+            // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+            for (int stride = 1; stride <= 4; stride *= 4) {
+                const int base = lane & ~(3 * stride);
+#pragma unroll
+                for (int j = 0; j < 8; ++j) {
+                    const float a = __shfl_sync(0xffffffff, v[j], base, 32);
+                    const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
+                    const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
+                    const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
+                    v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+                }
+            }
+#pragma unroll
+            for (int j = 0; j < 8; j += 2) {
+                const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
+                const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
+                const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
+                const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
+                v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
+                v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+            }
+#pragma unroll
+            for (int j = 0; j < 2; ++j) {
+                const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+                for (int t = 0; t < 4; ++t) {
+                    values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+                }
+            }
+        }
+        __syncthreads();
+    }
+    float local_max = 0.0f;
+    for (int64_t i = tid; i < k; i += blockDim.x) {
+        local_max = fmaxf(local_max, fabsf(values[i]));
+    }
+    local_max      = sd_warp_max(local_max);
+    const int lane = tid & 31;
+    const int warp = tid >> 5;
+    if (lane == 0) {
+        maxima[warp] = local_max;
+    }
+    __syncthreads();
+
+    if (warp == 0) {
+        local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+        local_max = sd_warp_max(local_max);
+        if (lane == 0) {
+            maxima[0] = local_max / 127.0f;
+            if (scales != nullptr) {
+                scales[row] = maxima[0];
+            } else {
+                *((float*)(dst_row + k)) = maxima[0];
+            }
+        }
+    }
+    __syncthreads();
+
+    const float row_scale = maxima[0];
+    const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+    for (int64_t i = tid; i < k; i += blockDim.x) {
+        int value  = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+        value      = max(-127, min(127, value));
+        dst_row[i] = (int8_t)value;
+    }
+}
+
+template <typename T>
+static __global__ void copy_to_half(const char* src, half* dst, int64_t count, int64_t dim, int64_t rows, int64_t heads, size_t nb1, size_t nb2, size_t nb3, float scale) {
+    const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i < count) {
+        const int64_t row = i / dim;
+        const auto* p     = reinterpret_cast<const T*>(src + (row % rows) * nb1 +
+                                                       ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
+        dst[i]            = __float2half_rn(float(p[i % dim]) * scale);
+    }
+}
+
+static __global__ void copy_from_half(const half* src, float* dst, int64_t count, float scale, float bias, bool scaled) {
+    const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i < count)
+        dst[i] = scaled ? __fmaf_rn(__half2float(src[i]), scale, bias) : __half2float(src[i]);
+}
+
+static float param_f32(const ggml_tensor* node, int index) {
+    float value;
+    std::memcpy(&value, reinterpret_cast<const char*>(node->op_params) + index * sizeof(value), sizeof(value));
+    return value;
+}
+
+static int padded_dim(int64_t dim) {
+    return dim == 80 || dim == 96 || dim == 112 ? int(dim) : int((dim + 63) / 64 * 64);
+}
+
+struct CUDADevice {
+    int ordinal;
+    cudaDeviceProp properties;
+};
+
+static const CUDADevice* cuda_device(ggml_backend_dev_t device) {
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, CUDADevice> devices;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto found = devices.find(device);
+    if (found != devices.end())
+        return &found->second;
+    auto reg   = ggml_backend_dev_backend_reg(device);
+    auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+    if (!query)
+        return nullptr;
+    CUDADevice info = {};
+    info.ordinal    = query(device);
+    if (info.ordinal < 0 || cudaGetDeviceProperties(&info.properties, info.ordinal) != cudaSuccess)
+        return nullptr;
+    return &devices.emplace(device, info).first->second;
+}
+
+static bool attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+    static const bool disabled = std::getenv("GGML_CUDA_DISABLE_CUTE_FATTN") || std::getenv("GGML_CUDA_DISABLE_CUTLASS_FATTN");
+    const auto* info           = cuda_device(device);
+    const auto* q              = node->src[0];
+    const auto* k              = node->src[1];
+    const auto* v              = node->src[2];
+    if (disabled || !info || info->properties.major < 8 || !q || !k || !v ||
+        q->type != GGML_TYPE_F32 || node->type != GGML_TYPE_F32 ||
+        q->ne[0] < 32 || q->ne[0] > 256 || q->ne[0] % 8 ||
+        size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info->properties.sharedMemPerBlockOptin ||
+        q->ne[1] < 1 || k->ne[1] < 1 || k->ne[2] < 1 ||
+        node->src[3] || node->src[4] || param_f32(node, 1) != 0.0f || param_f32(node, 2) != 0.0f ||
+        q->ne[2] % k->ne[2] || q->ne[3] != k->ne[3] || !ggml_are_same_shape(k, v) ||
+        q->nb[0] != sizeof(float) || !ggml_is_contiguous(node) ||
+        (k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_F32) ||
+        (v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_F32) ||
+        k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type))
+        return false;
+    return k->ne[0] == q->ne[0] && q->ne[1] <= INT_MAX / q->ne[0] && k->ne[1] <= INT_MAX / q->ne[0] &&
+           q->ne[2] <= 65535 && q->ne[3] <= 65535 && ggml_nelements(q) <= INT_MAX && ggml_nelements(k) <= INT_MAX;
+}
+
+static bool attention_preferred(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+    const auto* q = node->src[0];
+    const auto* k = node->src[1];
+    if (q->ne[0] == 40 || q->ne[0] == 72)
+        return q->ne[1] >= 512 && k->ne[1] >= 512;
+    if (q->ne[0] == 80 || q->ne[0] == 96 || q->ne[0] == 112)
+        return false;
+    const int minimum = q->ne[0] == 128 ? 2048 : 4096;
+    return q->ne[1] >= minimum && k->ne[1] >= minimum;
+}
+
+static bool attention_selected(ggml_backend_dev_t device, const ggml_tensor* node) {
+    return node->op == GGML_OP_FLASH_ATTN_EXT && attention_supports(device, node, nullptr) &&
+           (attention_preferred(device, node, nullptr) || !ggml_ext_dev_supports_native_op(device, node));
+}
+
+static bool fusion_enabled() {
+    static const bool disabled = [] {
+        const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
+        return value && std::atoi(value) != 0;
+    }();
+    return !disabled;
+}
+
+static const ggml_tensor* graph_node(const ggml_cgraph* graph, int index) {
+    return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
+}
+
+static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, int first) {
+    const auto* attn     = graph_node(graph, first);
+    const auto* previous = attn;
+    const int total      = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
+    bool scaled          = false;
+    for (int count = 1; count < 8 && first + count < total; ++count) {
+        const auto* node = graph_node(graph, first + count);
+        int convrot_group;
+        if (node->src[0] != previous)
+            break;
+        if (node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE) {
+            if (node->view_offs || !ggml_is_contiguous(node) || ggml_nelements(node) != ggml_nelements(attn))
+                break;
+        } else if (node->op == GGML_OP_SCALE && !scaled) {
+            if (!ggml_is_contiguous(node) || !ggml_are_same_shape(previous, node))
+                break;
+            scaled = true;
+        } else if (sd_ggml_convrot_describe(node, &convrot_group)) {
+            if (convrot_group == 256 && previous->ne[0] % 256 == 0 &&
+                size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
+                return count + 1;
+            break;
+        } else
+            break;
+        previous = node;
+    }
+    if (first + 1 < total) {
+        const auto* next = graph_node(graph, first + 1);
+        if (next->op == GGML_OP_SCALE && next->src[0] == attn && ggml_is_contiguous(next) && ggml_are_same_shape(attn, next))
+            return 2;
+    }
+    return 1;
+}
+
+static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+    if (!fusion_enabled())
+        return 0;
+    int offset = 0;
+    if (ggml_ext_op_get(graph_node(graph, first)) == sd_ggml_i8_op()) {
+        if (first + 5 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
+            !sd_cuda_i8_rope_supports(device, graph_node(graph, first), graph_node(graph, first + 1),
+                                      graph_node(graph, first + 2), graph_node(graph, first + 3), graph_node(graph, first + 4)) ||
+            graph_node(graph, first + 5)->src[0] != graph_node(graph, first + 4))
+            return 0;
+        offset = 5;
+    } else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
+        if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
+            !sd_cuda::rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
+            graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
+            return 0;
+        offset = 3;
+    }
+    if (!attention_selected(device, graph_node(graph, first + offset)))
+        return 0;
+    const int count = offset + output_fusion(device, graph, first + offset);
+    return count > 1 ? count : 0;
+}
+
+static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+    if (!fusion_enabled())
+        return;
+    for (int i = 3; i < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)); ++i) {
+        const auto* attn = graph_node(graph, i);
+        if (!attention_selected(device, attn))
+            continue;
+        for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
+            if (attn->src[0] == graph_node(graph, j + 2) &&
+                sd_cuda::rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
+                if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) &&
+                    ggml_ext_graph_delay(graph, j - 2, 5, i))
+                    break;
+                ggml_ext_graph_delay(graph, j, 3, i);
+                break;
+            }
+        }
+    }
+}
+
+struct AttentionPlan {
+    int index;
+    size_t offsets[4];
+};
+
+static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) {
+    auto* plan  = new AttentionPlan{};
+    *state      = plan;
+    plan->index = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
+                                                                                                      : 0;
+    GGML_ASSERT(plan->index < count);
+    const auto* attn = nodes[plan->index];
+    size_t size      = 0;
+    for (int i = 0; i < 4; ++i) {
+        plan->offsets[i]   = size;
+        const size_t bytes = size_t(ggml_nelements(i < 3 ? attn->src[i] : attn)) * sizeof(half);
+        size += (bytes + 255) / 256 * 256;
+    }
+    *workspace_size = size;
+    return GGML_STATUS_SUCCESS;
+}
+
+template <int D, bool FullDim>
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+    auto kernel                = sd_fattn_cute::kernel<64, 64, D, FullDim>;
+    constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
+    static std::mutex mutex;
+    static std::unordered_map<int, bool> initialized;
+    {
+        std::lock_guard<std::mutex> lock(mutex);
+        if (!initialized[context.device]) {
+            auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes);
+            if (error != cudaSuccess)
+                return error;
+            initialized[context.device] = true;
+        }
+    }
+    kernel<<<dim3((p.nq + 63) / 64, batches, p.heads), 128, shared_bytes, static_cast<cudaStream_t>(context.stream)>>>(p);
+    return cudaGetLastError();
+}
+
+template <int D>
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+    return p.dim == D ? attention_launch<D, true>(context, p, batches) : attention_launch<D, false>(context, p, batches);
+}
+
+static ggml_status attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+    const auto& plan = *static_cast<AttentionPlan*>(state);
+    const auto* attn = nodes[plan.index];
+    auto device      = ggml_backend_get_device(execution->backend);
+    auto reg         = ggml_backend_dev_backend_reg(device);
+    auto query       = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+    ggml_ext_cuda_context context;
+    if (!query || !query(execution->backend, &context))
+        return GGML_STATUS_FAILED;
+    auto stream = static_cast<cudaStream_t>(context.stream);
+    auto* base  = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+    half* work[4];
+    for (int i = 0; i < 4; ++i)
+        work[i] = reinterpret_cast<half*>(base + plan.offsets[i]);
+    const half* inputs[3] = {work[0], work[1], work[2]};
+    for (int i = 0; i < 3; ++i) {
+        const auto* src   = attn->src[i];
+        const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
+        if (i == 0 && plan.index == 5) {
+            if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+        } else if (i == 0 && plan.index == 3) {
+            if (sd_cuda::rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+        } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
+            inputs[i] = static_cast<const half*>(src->data);
+        } else {
+            auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half<half> : copy_to_half<float>;
+            kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), work[i],
+                                                                          ggml_nelements(src), src->ne[0], src->ne[1], src->ne[2], src->nb[1], src->nb[2], src->nb[3], scale);
+        }
+        if (cudaGetLastError() != cudaSuccess)
+            return GGML_STATUS_FAILED;
+    }
+    const auto* q           = attn->src[0];
+    const auto* k           = attn->src[1];
+    sd_fattn_cute::Params p = {reinterpret_cast<const cutlass::half_t*>(inputs[0]),
+                               reinterpret_cast<const cutlass::half_t*>(inputs[1]), reinterpret_cast<const cutlass::half_t*>(inputs[2]),
+                               reinterpret_cast<cutlass::half_t*>(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])};
+    cudaError_t error       = cudaErrorInvalidValue;
+    switch (padded_dim(p.dim)) {
+        case 64:
+            error = attention_launch<64>(context, p, q->ne[3]);
+            break;
+        case 80:
+            error = attention_launch<80, true>(context, p, q->ne[3]);
+            break;
+        case 96:
+            error = attention_launch<96, true>(context, p, q->ne[3]);
+            break;
+        case 112:
+            error = attention_launch<112, true>(context, p, q->ne[3]);
+            break;
+        case 128:
+            error = attention_launch<128>(context, p, q->ne[3]);
+            break;
+        case 192:
+            error = attention_launch<192>(context, p, q->ne[3]);
+            break;
+        case 256:
+            error = attention_launch<256>(context, p, q->ne[3]);
+            break;
+    }
+    if (error != cudaSuccess)
+        return GGML_STATUS_FAILED;
+    float scale = 1.0f, bias = 0.0f;
+    for (int i = plan.index + 1; i < count; ++i) {
+        if (nodes[i]->op == GGML_OP_SCALE) {
+            scale = param_f32(nodes[i], 0);
+            bias  = param_f32(nodes[i], 1);
+        }
+    }
+    auto* dst = nodes[count - 1];
+    int convrot_group;
+    if (sd_ggml_convrot_describe(dst, &convrot_group)) {
+        const int64_t width = dst->src[0]->ne[0], rows = ggml_nrows(dst->src[0]), padded_rows = (rows + 3) / 4 * 4;
+        auto* packed = static_cast<int8_t*>(dst->data);
+        if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        if (cudaFuncSetAttribute(quantize_attention, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                 int(cuda_device(device)->properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        quantize_attention<<<rows, 256, width * sizeof(float), stream>>>(work[3], packed,
+                                                                         reinterpret_cast<float*>(packed + width * padded_rows), width, rows, scale, bias);
+    } else {
+        copy_from_half<<<(ggml_nelements(attn) + 255) / 256, 256, 0, stream>>>(work[3], static_cast<float*>(dst->data), ggml_nelements(attn), scale, bias, dst->op == GGML_OP_SCALE);
+    }
+    return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+static void attention_destroy(void* state, void*) {
+    delete static_cast<AttentionPlan*>(state);
+}
+
+void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
+    if (!cuda_device(device))
+        return;
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.cute_attention.cuda.v1";
+    desc.op                   = GGML_OP_FLASH_ATTN_EXT;
+    desc.priority             = 200;
+    desc.supports             = attention_supports;
+    desc.preferred            = attention_preferred;
+    desc.match                = attention_match;
+    desc.optimize             = attention_optimize;
+    desc.prepare              = attention_prepare;
+    desc.execute              = attention_execute;
+    desc.destroy              = attention_destroy;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    desc.name      = "sd.rms_rope_cute_attention.cuda.v1";
+    desc.op        = GGML_OP_RMS_NORM;
+    desc.priority  = 300;
+    desc.supports  = nullptr;
+    desc.preferred = nullptr;
+    desc.optimize  = nullptr;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    desc.name      = "sd.i8_rope_cute_attention.cuda.v1";
+    desc.op        = GGML_OP_CUSTOM;
+    desc.custom_op = sd_ggml_i8_op();
+    desc.priority  = 400;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}

src/core/cuda/i8_gated_residual.cu

@@ -0,0 +1,461 @@
+#include <cuda_runtime.h>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <unordered_map>
+#include "core/ggml_extensions.h"
+
+namespace sd_i8_gated {
+    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+        using Base::Base;
+
+        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+            const auto shape = tiled_shape();
+            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+                constexpr int band_width = 16;
+                if (index >= shape.m() * shape.n()) {
+                    return {shape.m(), shape.n(), get_batch_idx()};
+                }
+                const int band   = index / (shape.m() * band_width);
+                const int width  = min(band_width, shape.n() - band * band_width);
+                const int within = index - band * shape.m() * band_width;
+                return {within / width, band * band_width + within % width, get_batch_idx()};
+            }
+#endif
+            return Base::get_tile_offset(index);
+        }
+    };
+
+    template <int M, int N, int WM>
+    struct TileConfig {
+        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+    };
+
+    enum Tile {
+        TILE_128_256,
+        TILE_128_128,
+        TILE_64_128,
+        TILE_16_128,
+    };
+
+    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+    }
+
+    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+        if (rows < 32) {
+            return TILE_16_128;
+        }
+        const int64_t columns = (n + 127) / 128;
+        if (((rows + 15) / 16) * columns <= nsm) {
+            return TILE_16_128;
+        }
+        if (((rows + 63) / 64) * columns <= nsm) {
+            return TILE_64_128;
+        }
+        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
+        const int64_t square_tiles = ((rows + 127) / 128) * columns;
+        // A narrower tile costs more per output element, but can waste less of the last wave.
+        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+    }
+
+    template <class T>
+    struct MulRound {
+        CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
+            T result;
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < T::kElements; ++i) {
+                result[i] = __fmul_rn(a[i], b[i]);
+            }
+            return result;
+        }
+    };
+
+    template <class T>
+    struct AddRound {
+        CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
+            T result;
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < T::kElements; ++i) {
+                result[i] = __fadd_rn(a[i], b[i]);
+            }
+            return result;
+        }
+    };
+
+    template <class Config, bool scaled, bool stream_k>
+    struct KernelTypes {
+        using Tile         = typename Config::Tile;
+        using Warp         = typename Config::Warp;
+        using Swizzle      = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map          = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col          = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row          = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul          = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX       = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW       = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Residual     = cutlass::epilogue::threadblock::VisitorAuxLoad<Map, float, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using GateMul      = cutlass::epilogue::threadblock::VisitorCompute<MulRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Add          = cutlass::epilogue::threadblock::VisitorCompute<AddRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Scalar       = cutlass::epilogue::threadblock::VisitorScalarBroadcast<float>;
+        using ScaleOut     = cutlass::epilogue::threadblock::Sm80EVT<GateMul, ScaleW, Scalar>;
+        using Scaled       = std::conditional_t<scaled, ScaleOut, ScaleW>;
+        using Gated        = cutlass::epilogue::threadblock::Sm80EVT<GateMul, Scaled, Row>;
+        using WithResidual = cutlass::epilogue::threadblock::Sm80EVT<Add, Gated, Residual>;
+        using Output       = WithResidual;
+        using Store        = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Epilogue     = cutlass::epilogue::threadblock::Sm80EVT<Store, Output>;
+        using Kernel       = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+            int8_t,
+            cutlass::layout::RowMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            int8_t,
+            cutlass::layout::ColumnMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            float,
+            cutlass::layout::RowMajor,
+            4,
+            int,
+            float,
+            cutlass::arch::OpClassTensorOp,
+            cutlass::arch::Sm80,
+            Tile,
+            Warp,
+            cutlass::gemm::GemmShape<16, 8, 32>,
+            Epilogue,
+            Swizzle,
+            4,
+            cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+        using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+    };
+
+    struct DeviceInfo {
+        int ordinal        = -1;
+        int major          = 0;
+        int sms            = 0;
+        int threads_per_sm = 0;
+        size_t shared      = 0;
+    };
+    static DeviceInfo device_info(ggml_backend_dev_t device) {
+        static std::mutex mutex;
+        static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
+        std::lock_guard<std::mutex> lock(mutex);
+        auto found = cache.find(device);
+        if (found != cache.end())
+            return found->second;
+        DeviceInfo info;
+        auto reg             = ggml_backend_dev_backend_reg(device);
+        auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+        int ordinal          = query ? query(device) : -1;
+        cudaDeviceProp props = {};
+        if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
+            info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+        cache.emplace(device, info);
+        return info;
+    }
+    static float param_f32(const ggml_tensor* tensor, int index) {
+        float value;
+        std::memcpy(&value, tensor->op_params + index, sizeof(value));
+        return value;
+    }
+    static bool aligned(const ggml_tensor* tensor) {
+        if (!tensor)
+            return false;
+        size_t offset = 0;
+        while (tensor->view_src) {
+            offset += tensor->view_offs;
+            tensor = tensor->view_src;
+        }
+        return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
+    }
+    struct Operation {
+        SDI8Operation mat;
+        const ggml_tensor* gate;
+        const ggml_tensor* residual;
+        const ggml_tensor* prefix;
+        float scale;
+        int64_t rows, n, k;
+    };
+    static Operation describe(ggml_tensor* const* nodes, int count) {
+        Operation op = {};
+        sd_ggml_i8_describe(nodes[0], &op.mat);
+        const bool concat = nodes[count - 1]->op == GGML_OP_CONCAT;
+        const auto* add   = nodes[count - (concat ? 2 : 1)];
+        const auto* mul   = nodes[count - (concat ? 3 : 2)];
+        op.gate           = mul->src[1];
+        op.residual       = add->src[0] == mul ? add->src[1] : add->src[0];
+        op.prefix         = concat ? nodes[count - 1]->src[0] : nullptr;
+        op.scale          = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f;
+        op.rows           = ggml_nrows(nodes[0]);
+        op.n              = nodes[0]->ne[0];
+        op.k              = op.mat.weight->ne[0];
+        return op;
+    }
+    template <class Types>
+    static typename Types::Gemm::Arguments arguments(const Operation& op, float* output, bool metadata) {
+        const auto* x                            = metadata ? nullptr : static_cast<const int8_t*>(op.mat.input->data);
+        const auto* w                            = metadata ? nullptr : static_cast<const int8_t*>(op.mat.weight->data);
+        const auto* xs                           = metadata ? nullptr : op.mat.packed ? reinterpret_cast<const float*>(x + op.k * ((op.rows + 3) / 4 * 4))
+                                                                                      : static_cast<const float*>(op.mat.input_scale->data);
+        const auto* ws                           = metadata ? nullptr : static_cast<const float*>(op.mat.weight_scale->data);
+        const auto* gate                         = metadata ? nullptr : static_cast<const float*>(op.gate->data);
+        auto* residual                           = metadata ? nullptr : static_cast<float*>(op.residual->data);
+        typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+        typename Types::Scaled::Arguments scaled;
+        if constexpr (std::is_same<typename Types::Scaled, typename Types::ScaleW>::value)
+            scaled = scales;
+        else
+            scaled = {scales, {{op.scale}}, {}};
+        typename Types::Epilogue::Arguments epilogue = {{{scaled, {gate, 0.f, {}}, {}}, {residual, 0.f, {op.n, cute::_1{}, op.rows * op.n}}, {}}, {output, {op.n, cute::_1{}, op.rows * op.n}}};
+        return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0};
+    }
+    struct Plan {
+        DeviceInfo device;
+        size_t kernel_workspace                                                     = 0;
+        size_t output_offset                                                        = 0;
+        int occupancy                                                               = 0;
+        ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr;
+    };
+    template <class Types>
+    static ggml_status launch(Plan& plan, const Operation& op, float* output, void* scratch, cudaStream_t stream) {
+        using Kernel = typename Types::Kernel;
+        auto args    = arguments<Types>(op, output, false);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        if (plan.occupancy == 0) {
+            if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
+                cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+            int occupancy = 0;
+            if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&occupancy, cutlass::Kernel2<Kernel>,
+                                                                       Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
+                occupancy <= 0)
+                return GGML_STATUS_FAILED;
+            plan.occupancy = occupancy;
+        }
+        typename Kernel::Params params(args, plan.device.sms, plan.occupancy);
+        if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        cutlass::arch::synclog_setup();
+        cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+    template <class Config, bool Scaled, bool StreamK>
+    static bool configure(const Operation& op, Plan& plan) {
+        using Types  = KernelTypes<Config, Scaled, StreamK>;
+        using Kernel = typename Types::Kernel;
+        if (sizeof(typename Kernel::SharedStorage) > plan.device.shared)
+            return false;
+        auto args = arguments<Types>(op, nullptr, true);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return false;
+        size_t workspace = 0;
+        // Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement.
+        for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+            typename Kernel::Params params(args, plan.device.sms, occupancy);
+            const auto grid = params.get_grid_dims();
+            if (grid.y > 65535 || grid.z > 65535)
+                return false;
+            workspace = std::max(workspace, params.get_workspace_size());
+        }
+        plan.kernel_workspace = workspace;
+        plan.output_offset    = (workspace + 255) / 256 * 256;
+        plan.launch           = launch<Types>;
+        return true;
+    }
+    template <class Config>
+    static bool configure(const Operation& op, Plan& plan) {
+        if (op.scale != 1.f)
+            return op.rows >= 8192 ? configure<Config, true, true>(op, plan) : configure<Config, true, false>(op, plan);
+        return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
+    }
+    static bool configure(const Operation& op, Plan& plan) {
+        auto try_tile = [&](Tile tile) {
+            switch (tile) {
+                case TILE_128_256:
+                    return configure<TileConfig<128, 256, 64>>(op, plan);
+                case TILE_128_128:
+                    return configure<TileConfig<128, 128, 64>>(op, plan);
+                case TILE_64_128:
+                    return configure<TileConfig<64, 128, 32>>(op, plan);
+                case TILE_16_128:
+                    return configure<TileConfig<16, 128, 16>>(op, plan);
+            }
+            return false;
+        };
+        const auto tile = select_tile(op.rows, op.n, plan.device.sms);
+        if (try_tile(tile))
+            return true;
+        for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
+            if (fallback != tile && try_tile(fallback))
+                return true;
+        return false;
+    }
+    static bool fusion_disabled() {
+        static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+        return disabled;
+    }
+    static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+#ifndef SD_CUDA_FORCE_CUBLAS
+        if (fusion_disabled())
+            return;
+        const int total = ggml_graph_n_nodes(graph);
+        for (int i = 2; i + 1 < total; ++i) {
+            const auto* mul = ggml_graph_node(graph, i);
+            const auto* add = ggml_graph_node(graph, i + 1);
+            if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
+                continue;
+            const auto* producer = mul->src[0];
+            const bool scaled    = producer->op == GGML_OP_SCALE;
+            const auto* mat      = scaled ? producer->src[0] : producer;
+            SDI8Operation op;
+            if (!sd_ggml_i8_describe(mat, &op) || op.input->type != GGML_TYPE_I8 || !op.weight_scale || op.bias)
+                continue;
+            int scale_index = -1;
+            for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
+                if (scaled && ggml_graph_node(graph, j) == producer)
+                    scale_index = j;
+                if (ggml_graph_node(graph, j) != mat)
+                    continue;
+                if (scaled) {
+                    if (scale_index < 0 || (j + 1 < scale_index && !ggml_ext_graph_delay(graph, j, 1, scale_index)))
+                        break;
+                    j = scale_index - 1;
+                }
+                ggml_ext_graph_delay(graph, j, scaled ? 2 : 1, i);
+                break;
+            }
+        }
+#endif
+    }
+    static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+        return 0;
+#else
+        const int available = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) - first;
+        if (fusion_disabled() || available < 3)
+            return 0;
+        ggml_tensor* nodes[5] = {};
+        for (int i = 0; i < std::min(available, 5); ++i)
+            nodes[i] = ggml_graph_node(const_cast<ggml_cgraph*>(graph), first + i);
+        SDI8Operation mat;
+        if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias)
+            return 0;
+        const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0];
+        if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
+            ((rows + 15) / 16) > INT_MAX / ((n + 127) / 128))
+            return 0;
+        for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale})
+            if (!aligned(t))
+                return 0;
+        const bool scaled = nodes[1]->op == GGML_OP_SCALE;
+        int count         = scaled ? 4 : 3;
+        if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f)))
+            return 0;
+        const auto* mul = nodes[count - 2];
+        const auto* add = nodes[count - 1];
+        if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || mul->src[0] != nodes[count - 3] ||
+            (add->src[0] != mul && add->src[1] != mul))
+            return 0;
+        const auto* gate     = mul->src[1];
+        const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0];
+        if (!gate || !residual || gate->ne[0] != n || ggml_nelements(gate) != n || !aligned(gate) || !aligned(residual))
+            return 0;
+        for (const auto* t : {gate, residual, mul, add})
+            if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || (t != gate && !ggml_are_same_shape(nodes[0], t)))
+                return 0;
+        if (available > count) {
+            const auto* concat = nodes[count];
+            if (concat->op == GGML_OP_CONCAT && concat->src[1] == add && concat->op_params[0] == 1 &&
+                concat->type == GGML_TYPE_F32 && concat->ne[2] == 1 && concat->ne[3] == 1 && ggml_is_contiguous(concat) &&
+                concat->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(concat->src[0]))
+                ++count;
+        }
+        const auto* prefix = nodes[count - 1]->op == GGML_OP_CONCAT ? nodes[count - 1]->src[0] : nullptr;
+        for (int i = 0; i < count; ++i)
+            if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix)
+                return 0;
+        Plan plan;
+        plan.device = device_info(device);
+        if (plan.device.major < 8 || plan.device.sms <= 0)
+            return 0;
+        return configure(describe(nodes, count), plan) ? count : 0;
+#endif
+    }
+    static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+        auto* plan   = new Plan;
+        *state       = plan;
+        plan->device = device_info(ggml_backend_get_device(backend));
+        if (!configure(describe(nodes, count), *plan))
+            return GGML_STATUS_FAILED;
+        *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void* state, void*) {
+        delete static_cast<Plan*>(state);
+    }
+    static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+        if (!b)
+            return false;
+        const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+        return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+    }
+    static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+        auto& plan    = *static_cast<Plan*>(state);
+        const auto op = describe(nodes, count);
+        auto* dst     = nodes[count - 1];
+        auto reg      = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+        auto query    = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+        ggml_ext_cuda_context context;
+        if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
+            return GGML_STATUS_FAILED;
+        auto stream   = static_cast<cudaStream_t>(context.stream);
+        auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+        bool alias    = reinterpret_cast<uintptr_t>(dst->data) % 16 != 0;
+        for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) {
+            if (input == op.residual && !op.prefix && input->data == dst->data)
+                continue;
+            alias |= overlaps(dst, input);
+        }
+        auto* output              = alias ? scratch + plan.output_offset : static_cast<char*>(dst->data);
+        const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0;
+        auto status               = plan.launch(plan, op, reinterpret_cast<float*>(output + prefix_bytes), scratch, stream);
+        if (status != GGML_STATUS_SUCCESS)
+            return status;
+        if (prefix_bytes && cudaMemcpyAsync(output, op.prefix->data, prefix_bytes, cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        if (alias && cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
+            return GGML_STATUS_FAILED;
+        return GGML_STATUS_SUCCESS;
+    }
+}
+
+void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.i8_gated_residual.cuda.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_i8_op();
+    desc.priority             = 325;
+    desc.optimize             = sd_i8_gated::optimize;
+    desc.match                = sd_i8_gated::match;
+    desc.prepare              = sd_i8_gated::prepare;
+    desc.execute              = sd_i8_gated::execute;
+    desc.destroy              = sd_i8_gated::destroy;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}

src/core/cuda/i8_gemm.cu

@@ -0,0 +1,330 @@
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <unordered_map>
+#include "i8_gemm.h"
+
+namespace sd_i8_cutlass {
+    struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+        using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+        using Base::Base;
+
+        CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+            const auto shape = tiled_shape();
+            if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+                // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+                constexpr int band_width = 16;
+                if (index >= shape.m() * shape.n()) {
+                    return {shape.m(), shape.n(), get_batch_idx()};
+                }
+                const int band   = index / (shape.m() * band_width);
+                const int width  = min(band_width, shape.n() - band * band_width);
+                const int within = index - band * shape.m() * band_width;
+                return {within / width, band * band_width + within % width, get_batch_idx()};
+            }
+#endif
+            return Base::get_tile_offset(index);
+        }
+    };
+
+    template <int M, int N, int WM>
+    struct TileConfig {
+        using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+        using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+    };
+
+    enum Tile {
+        TILE_128_256,
+        TILE_128_128,
+        TILE_64_128,
+        TILE_16_128,
+    };
+
+    static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+        static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+        return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+    }
+
+    static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+        if (rows < 32) {
+            return TILE_16_128;
+        }
+        const int64_t columns = (n + 127) / 128;
+        if (((rows + 15) / 16) * columns <= nsm) {
+            return TILE_16_128;
+        }
+        if (((rows + 63) / 64) * columns <= nsm) {
+            return TILE_64_128;
+        }
+        const int64_t wide_tiles   = ((rows + 127) / 128) * ((n + 255) / 256);
+        const int64_t square_tiles = ((rows + 127) / 128) * columns;
+        // A narrower tile costs more per output element, but can waste less of the last wave.
+        const double wide_cost   = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+        const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+        return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+    }
+
+    struct DeviceInfo {
+        int ordinal        = -1;
+        int major          = 0;
+        int sms            = 0;
+        int threads_per_sm = 0;
+        size_t shared      = 0;
+    };
+    static DeviceInfo device_info(ggml_backend_dev_t device) {
+        static std::mutex mutex;
+        static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
+        std::lock_guard<std::mutex> lock(mutex);
+        auto found = cache.find(device);
+        if (found != cache.end())
+            return found->second;
+        DeviceInfo info;
+        auto reg             = ggml_backend_dev_backend_reg(device);
+        auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+        int ordinal          = query ? query(device) : -1;
+        cudaDeviceProp props = {};
+        if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
+            info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+        cache.emplace(device, info);
+        return info;
+    }
+    // The iterator pairs gate/up columns while reading the original checkpoint layout.
+    struct WeightLayout : cutlass::layout::PermuteBase {
+        int stride, half_n;
+        CUTLASS_HOST_DEVICE WeightLayout(cutlass::PitchLinearCoord extent, int64_t ld)
+            : stride(static_cast<int>(ld)), half_n(extent.strided() / 2) {}
+        CUTLASS_HOST_DEVICE WeightLayout(cutlass::MatrixCoord extent, int64_t ld)
+            : stride(static_cast<int>(ld)), half_n(extent.column() / 2) {}
+        CUTLASS_HOST_DEVICE int64_t operator()(cutlass::PitchLinearCoord c) const { return c.contiguous() + int64_t(c.strided() / 2 + (c.strided() % 2) * half_n) * stride; }
+        CUTLASS_HOST_DEVICE int64_t operator()(cutlass::MatrixCoord c) const { return c.row() + int64_t(c.column() / 2 + (c.column() % 2) * half_n) * stride; }
+    };
+
+    template <class Map>
+    struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+        using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Base::Base;
+        template <class CB>
+        struct Callbacks : CB {
+            CUTLASS_DEVICE Callbacks(CB const& base)
+                : CB(base) {}
+            CUTLASS_DEVICE void end_step(int step) {
+                auto src    = cute::filter(this->tC_rAux);
+                auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
+                CUTLASS_PRAGMA_UNROLL
+                for (int i = 0; i < cute::size(src); ++i) {
+                    auto coord = coords(i);
+                    if (cute::elem_less(coord, this->problem_shape)) {
+                        // Each vector contains two gate/up pairs; only their SwiGLU results are stored.
+                        const float4 a = *reinterpret_cast<const float4*>(&src(i));
+                        float2 value   = make_float2(__fmul_rn(a.x / (1.f + expf(-a.x)), a.y), __fmul_rn(a.z / (1.f + expf(-a.z)), a.w));
+                        int64_t n = cute::get<1>(this->problem_shape), m = cute::get<0>(this->problem_shape);
+                        int64_t offset                                                 = cute::get<0>(coord) * (n / 2) + cute::get<1>(coord) / 2 + cute::get<2>(coord) * m * (n / 2);
+                        *reinterpret_cast<float2*>(this->params_ptr->ptr_aux + offset) = value;
+                    }
+                }
+            }
+        };
+        template <class Shape>
+        CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
+            auto cb = Base::get_callbacks(tile, tid, shape);
+            return Callbacks<decltype(cb)>(cb);
+        }
+    };
+
+    static __global__ void pair_scales(const float* src, float* dst, int n) {
+        const int i = blockIdx.x * blockDim.x + threadIdx.x;
+        if (i < n)
+            dst[i] = src[i / 2 + (i % 2) * (n / 2)];
+    }
+
+    template <class Config, bool StreamK>
+    struct DenseTypes {
+        using Tile     = typename Config::Tile;
+        using Warp     = typename Config::Warp;
+        using Swizzle  = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map      = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col      = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row      = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul      = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX   = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW   = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Store    = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+        using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using Kernel   = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+            int8_t,
+            cutlass::layout::RowMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            int8_t,
+            cutlass::layout::ColumnMajor,
+            cutlass::ComplexTransform::kNone,
+            16,
+            float,
+            cutlass::layout::RowMajor,
+            4,
+            int,
+            float,
+            cutlass::arch::OpClassTensorOp,
+            cutlass::arch::Sm80,
+            Tile,
+            Warp,
+            cutlass::gemm::GemmShape<16, 8, 32>,
+            Epilogue,
+            Swizzle,
+            4,
+            cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+        using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+    };
+    template <class Config, bool StreamK>
+    struct SwigluTypes {
+        using Tile            = typename Config::Tile;
+        using Warp            = typename Config::Warp;
+        using Swizzle         = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+        using Map             = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+        using Col             = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+        using Row             = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+        using Mul             = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using ScaleX          = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+        using ScaleW          = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+        using Store           = SwigluStore<Map>;
+        using Epilogue        = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+        using GemmBase        = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
+        using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+        using Kernel          = std::conditional_t<StreamK,
+                                                   cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
+                                                   cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
+        using Gemm            = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+    };
+    struct Plan {
+        DeviceInfo device;
+        int64_t rows, n, k;
+        bool swiglu;
+        size_t kernel_bytes = 0, scratch_offset = 0;
+        int occupancy                                                                                                       = 0;
+        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t) = nullptr;
+    };
+    template <class Types>
+    static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
+        typename Types::ScaleW::Arguments scales     = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+        typename Types::Epilogue::Arguments epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+        return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, w, nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
+    }
+    template <class Types>
+    static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
+        using Kernel = typename Types::Kernel;
+        if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
+             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out)) %
+            16)
+            return GGML_STATUS_FAILED;
+        auto* paired = static_cast<float*>(scratch);
+        auto args    = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        if (!p.occupancy) {
+            if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
+                cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
+                return GGML_STATUS_FAILED;
+            if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&p.occupancy, cutlass::Kernel2<Kernel>, Kernel::kThreadCount,
+                                                                       sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
+                p.occupancy <= 0)
+                return GGML_STATUS_FAILED;
+        }
+        typename Kernel::Params params(args, p.device.sms, p.occupancy);
+        if (params.get_workspace_size() > p.kernel_bytes)
+            return GGML_STATUS_FAILED;
+        if (p.swiglu) {
+            pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, paired, int(p.n));
+            if (cudaGetLastError() != cudaSuccess)
+                return GGML_STATUS_FAILED;
+        }
+        if (params.init_workspace(static_cast<char*>(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess)
+            return GGML_STATUS_FAILED;
+        cutlass::arch::synclog_setup();
+        cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+    template <class Types>
+    static bool configure_kernel(Plan& p) {
+        using Kernel = typename Types::Kernel;
+        if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
+            return false;
+        auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
+        if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+            return false;
+        size_t bytes = 0;
+        for (int occupancy = 1; occupancy <= p.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+            typename Kernel::Params params(args, p.device.sms, occupancy);
+            const auto grid = params.get_grid_dims();
+            if (grid.y > 65535 || grid.z > 65535)
+                return false;
+            bytes = std::max(bytes, params.get_workspace_size());
+        }
+        p.kernel_bytes   = bytes;
+        p.scratch_offset = p.swiglu ? (size_t(p.n) * sizeof(float) + 255) / 256 * 256 : 0;
+        p.launch         = launch<Types>;
+        return true;
+    }
+    template <class Config>
+    static bool configure_tile(Plan& p) {
+        if (p.swiglu)
+            return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
+        return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
+    }
+    static bool configure(Plan& p) {
+        auto try_tile = [&](Tile tile) {
+            switch (tile) {
+                case TILE_128_256:
+                    return configure_tile<TileConfig<128, 256, 64>>(p);
+                case TILE_128_128:
+                    return configure_tile<TileConfig<128, 128, 64>>(p);
+                case TILE_64_128:
+                    return configure_tile<TileConfig<64, 128, 32>>(p);
+                case TILE_16_128:
+                    return configure_tile<TileConfig<16, 128, 16>>(p);
+            }
+            return false;
+        };
+        const auto tile = select_tile(p.rows, p.n, p.device.sms);
+        if (try_tile(tile))
+            return true;
+        for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
+            if (fallback != tile && try_tile(fallback))
+                return true;
+        return false;
+    }
+}
+struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+    return nullptr;
+#else
+    const auto info = sd_i8_cutlass::device_info(device);
+    if (info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) ||
+        rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128))
+        return nullptr;
+    auto* p   = new SDI8CudaGemmPlan;
+    p->device = info;
+    p->rows   = rows;
+    p->n      = n;
+    p->k      = k;
+    p->swiglu = swiglu;
+    if (!sd_i8_cutlass::configure(*p)) {
+        delete p;
+        return nullptr;
+    }
+    *workspace = p->scratch_offset + p->kernel_bytes;
+    return p;
+#endif
+}
+ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
+    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream);
+}
+void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
+    delete plan;
+}

src/core/cuda/i8_gemm.h

@@ -0,0 +1,9 @@
+#ifndef __SD_CORE_CUDA_I8_GEMM_H__
+#define __SD_CORE_CUDA_I8_GEMM_H__
+#include <cuda_runtime.h>
+#include "ggml-ext.h"
+struct SDI8CudaGemmPlan;
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace);
+ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream);
+void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan);
+#endif

src/core/cuda/i8_projection.cu

@@ -0,0 +1,394 @@
+#include "core/cuda/i8_projection.h"
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+
+#include <cuda_fp16.h>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+
+static __device__ __forceinline__ float warp_sum8(float value) {
+#pragma unroll
+    for (int offset = 4; offset > 0; offset >>= 1)
+        value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset, 8));
+    return value;
+}
+
+template <class Map, class Out, bool rms = true>
+struct I8RopeStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+    using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+    struct Arguments : Base::Arguments {
+        const float* weight;
+        const float* pe;
+        float eps;
+        float scale;
+        int64_t tokens;
+        half* half_output;
+        int64_t prefix_tokens;
+    };
+    using Params        = Arguments;
+    using SharedStorage = typename Base::SharedStorage;
+    const Params* extra;
+    template <class Shape>
+    static constexpr Params to_underlying_arguments(Shape const&, Arguments const& args, void*) { return args; }
+    CUTLASS_HOST_DEVICE I8RopeStore(Params const& p, SharedStorage const& s)
+        : Base(p, s), extra(&p) {}
+    template <class CB>
+    struct Callbacks : CB {
+        const Params* extra;
+        CUTLASS_DEVICE Callbacks(CB const& base, const Params* p)
+            : CB(base), extra(p) {}
+        CUTLASS_DEVICE float sum32(float4 v) {
+            float a = warp_sum8(__fmul_rn(v.x, v.x));
+            float b = warp_sum8(__fmul_rn(v.y, v.y));
+            float c = warp_sum8(__fmul_rn(v.z, v.z));
+            float d = warp_sum8(__fmul_rn(v.w, v.w));
+            return __fadd_rn(__fadd_rn(a, c), __fadd_rn(b, d));
+        }
+        CUTLASS_DEVICE void end_step(int step) {
+            static_assert(Map::Detail::kAccessWidth == 16);
+            auto src    = cute::filter(this->tC_rAux);
+            auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
+            static_assert(cute::size(src) == 4);
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < cute::size(src); i += 2) {
+                float4 va = *reinterpret_cast<const float4*>(&src(i)), vb = *reinterpret_cast<const float4*>(&src(i + 1));
+                float norm = 1.0f;
+                if constexpr (rms) {
+                    // Preserve the standalone RMSNorm reduction tree across four 32-element groups.
+                    float sum = __fadd_rn(sum32(va), sum32(vb));
+                    sum       = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 8));
+                    norm      = rsqrtf(sum / 128.f + extra->eps);
+                }
+                CUTLASS_PRAGMA_UNROLL
+                for (int j = 0; j < 2; ++j) {
+                    auto coord = coords(i + j);
+                    int row = cute::get<0>(coord), col = cute::get<1>(coord);
+                    int heads = cute::get<1>(this->problem_shape) / 128;
+                    if (cute::elem_less(coord, this->problem_shape)) {
+                        float4 value  = j ? vb : va;
+                        int local     = col % 128;
+                        int64_t token = row % extra->tokens, batch = row / extra->tokens, head = col / 128;
+                        float4 r = value;
+                        if constexpr (rms) {
+                            float4 weight = ((const float4*)extra->weight)[local / 4];
+                            value.x       = __fmul_rn(__fmul_rn(norm, value.x), weight.x);
+                            value.y       = __fmul_rn(__fmul_rn(norm, value.y), weight.y);
+                            value.z       = __fmul_rn(__fmul_rn(norm, value.z), weight.z);
+                            value.w       = __fmul_rn(__fmul_rn(norm, value.w), weight.w);
+                            float4 p0 = ((const float4*)extra->pe)[token * 64 + local / 2], p1 = ((const float4*)extra->pe)[token * 64 + local / 2 + 1];
+                            r.x = __fadd_rn(__fmul_rn(value.x, p0.x), __fmul_rn(value.y, p0.y));
+                            r.y = __fadd_rn(__fmul_rn(value.x, p0.z), __fmul_rn(value.y, p0.w));
+                            r.z = __fadd_rn(__fmul_rn(value.z, p1.x), __fmul_rn(value.w, p1.y));
+                            r.w = __fadd_rn(__fmul_rn(value.z, p1.z), __fmul_rn(value.w, p1.w));
+                        }
+                        int64_t offset = ((batch * heads + head) * (extra->tokens + extra->prefix_tokens) + token + extra->prefix_tokens) * 128 + local;
+                        if constexpr (std::is_same<Out, half>::value) {
+                            half2 a    = __floats2half2_rn(__fmul_rn(r.x, extra->scale), __fmul_rn(r.y, extra->scale));
+                            half2 b    = __floats2half2_rn(__fmul_rn(r.z, extra->scale), __fmul_rn(r.w, extra->scale));
+                            half2* dst = (half2*)(extra->half_output + offset);
+                            dst[0]     = a;
+                            dst[1]     = b;
+                        } else {
+                            *((float4*)(this->params_ptr->ptr_aux + offset)) = r;
+                        }
+                    }
+                }
+            }
+        }
+    };
+    template <class Shape>
+    CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
+        auto cb = Base::get_callbacks(tile, tid, shape);
+        return Callbacks<decltype(cb)>(cb, extra);
+    }
+};
+
+template <class Out, bool RMS = true>
+struct ProjectionKernel {
+    using Tile    = cutlass::gemm::GemmShape<128, 256, 64>;
+    using Warp    = cutlass::gemm::GemmShape<64, 64, 64>;
+    using Swizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
+    using Map     = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+    using Col     = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+    using Row     = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+    using Mul     = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+    using ScaleX  = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+    using ScaleW  = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+    using Store   = I8RopeStore<Map, Out, RMS>;
+    using EVT     = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+    using Kernel  = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+        int8_t,
+        cutlass::layout::RowMajor,
+        cutlass::ComplexTransform::kNone,
+        16,
+        int8_t,
+        cutlass::layout::ColumnMajor,
+        cutlass::ComplexTransform::kNone,
+        16,
+        float,
+        cutlass::layout::RowMajor,
+        4,
+        int,
+        float,
+        cutlass::arch::OpClassTensorOp,
+        cutlass::arch::Sm80,
+        Tile,
+        Warp,
+        cutlass::gemm::GemmShape<16, 8, 32>,
+        EVT,
+        Swizzle,
+        4,
+        cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+    using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+};
+
+static float param_f32(const ggml_tensor* tensor, int index) {
+    float result;
+    std::memcpy(&result, tensor->op_params + index, sizeof(result));
+    return result;
+}
+
+static bool device_supports(ggml_backend_dev_t device) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+    (void)device;
+    return false;
+#else
+    static std::mutex mutex;
+    static std::unordered_map<ggml_backend_dev_t, bool> cache;
+    std::lock_guard<std::mutex> lock(mutex);
+    auto found = cache.find(device);
+    if (found != cache.end())
+        return found->second;
+    auto reg             = ggml_backend_dev_backend_reg(device);
+    auto query           = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+    int ordinal          = query ? query(device) : -1;
+    cudaDeviceProp props = {};
+    size_t shared        = std::max({sizeof(typename ProjectionKernel<float>::Kernel::SharedStorage),
+                                     sizeof(typename ProjectionKernel<half>::Kernel::SharedStorage), sizeof(typename ProjectionKernel<half, false>::Kernel::SharedStorage)});
+    bool supported       = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared;
+    cache.emplace(device, supported);
+    return supported;
+#endif
+}
+
+static bool aligned(const ggml_tensor* tensor) {
+    size_t offset = 0;
+    while (tensor->view_src) {
+        offset += tensor->view_offs;
+        tensor = tensor->view_src;
+    }
+    return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
+}
+
+static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) {
+    SDI8Operation op;
+    if (!sd_ggml_i8_describe(mat, &op) || !op.packed || !op.weight_scale || op.bias)
+        return false;
+    const auto* w  = op.weight;
+    const auto* x  = op.input;
+    const auto* ws = op.weight_scale;
+    if (mat->type != GGML_TYPE_F32 ||
+        w->type != GGML_TYPE_I8 || x->type != GGML_TYPE_I8 || ws->type != GGML_TYPE_F32 ||
+        w->ne[2] != 1 || w->ne[3] != 1 || mat->ne[0] != w->ne[1] || ggml_nelements(ws) != mat->ne[0] ||
+        !ggml_is_contiguous(mat) || !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || !ggml_is_contiguous(ws) ||
+        !aligned(w) || !aligned(x) || !aligned(ws) || !aligned(mat) ||
+        view->op != GGML_OP_RESHAPE || view->src[0] != mat || view->view_offs != 0 || !ggml_is_contiguous(view) ||
+        view->ne[0] != 128 || view->ne[0] * view->ne[1] != mat->ne[0] || ggml_nelements(view) != ggml_nelements(mat))
+        return false;
+    const int64_t rows = ggml_nrows(mat), n = mat->ne[0], k = w->ne[0];
+    if (rows < 4096 || rows > INT_MAX || n < 4096 || n > INT_MAX || k < 4096 || k > INT_MAX || n % 256 || k % 64 || n / 256 > 65535 ||
+        x->ne[0] != k || x->ne[1] != (rows + 3) / 4 * 4 + (rows * 4 + k - 1) / k || x->ne[2] != 1 || x->ne[3] != 1)
+        return false;
+    return device_supports(device);
+}
+
+bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+    if (!projection_supports(device, mat, view) || norm->src[0] != view ||
+        !sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
+        return false;
+    int32_t interleaved;
+    std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
+    return interleaved != 0;
+}
+
+template <class Out, bool RMS = true>
+static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope, void* output, float scale, int64_t prefix_tokens) {
+    using P        = ProjectionKernel<Out, RMS>;
+    using Gemm     = typename P::Gemm;
+    const int rows = int(ggml_nrows(mat)), n = int(mat->ne[0]), k = int(mat->src[0]->ne[0]);
+    const auto* x  = static_cast<const int8_t*>(mat->src[1]->data);
+    const auto* w  = static_cast<const int8_t*>(mat->src[0]->data);
+    const auto* xs = reinterpret_cast<const float*>(x + int64_t(k) * ((int64_t(rows) + 3) / 4 * 4));
+    const auto* ws = static_cast<const float*>(mat->src[2]->data);
+    const auto* nw = RMS ? static_cast<const float*>(mul->src[1]->data) : nullptr;
+    const auto* pe = RMS ? static_cast<const float*>(rope->src[1]->data) : nullptr;
+    // The visitor forms FP32 addresses even when its only writes target the FP16 output.
+    float* address_space = std::is_same<Out, half>::value ? static_cast<float*>(mat->data) : static_cast<float*>(output);
+    if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
+         reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(nw) | reinterpret_cast<uintptr_t>(pe) |
+         reinterpret_cast<uintptr_t>(address_space)) %
+            16 ||
+        reinterpret_cast<uintptr_t>(output) % 8)
+        return cudaErrorInvalidValue;
+    typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}},
+                                      {{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
+    typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr,
+                                  int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0);
+    Gemm gemm;
+    if (gemm.can_implement(args) != cutlass::Status::kSuccess)
+        return cudaErrorInvalidValue;
+    if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess || gemm(stream) != cutlass::Status::kSuccess)
+        return cudaErrorUnknown;
+    return cudaGetLastError();
+}
+
+cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale) {
+    return projection_launch<half>(stream, nodes[0], nodes[1], nodes[2], nodes[3], nodes[4], output, scale, 0);
+}
+
+static const ggml_tensor* node_at(const ggml_cgraph* graph, int index) {
+    return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
+}
+
+static int concat_match(const ggml_cgraph* graph, int first, bool key) {
+    const int offset = key ? 5 : 2, count = key ? 7 : 5;
+    if (first + count > ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
+        return 0;
+    const auto* concat = node_at(graph, first + offset);
+    const auto* copy   = node_at(graph, first + count - 1);
+    if (concat->op != GGML_OP_CONCAT || concat->op_params[0] != (key ? 1 : 2) || concat->src[1] != node_at(graph, first + offset - 1) ||
+        copy->op != GGML_OP_CPY || copy->src[0] != node_at(graph, first + count - 2) || copy->type != GGML_TYPE_F16 ||
+        !ggml_is_contiguous(copy) || !ggml_are_same_shape(copy, copy->src[0]))
+        return 0;
+    for (int i = 0; i < 2; ++i)
+        if (!concat->src[i] || concat->src[i]->type != GGML_TYPE_F32 || !ggml_is_contiguous(concat->src[i]))
+            return 0;
+    const auto* view   = node_at(graph, first + 1);
+    const auto* prefix = concat->src[0];
+    if (prefix->ne[key ? 1 : 2] > INT_MAX - view->ne[2] || (ggml_nelements(prefix) + 255) / 256 > INT_MAX)
+        return 0;
+    if (!key) {
+        const auto* permute = node_at(graph, first + 3);
+        if (permute->op != GGML_OP_PERMUTE || permute->src[0] != concat || permute->op_params[0] != 0 || permute->op_params[1] != 2 ||
+            permute->op_params[2] != 1 || permute->op_params[3] != 3)
+            return 0;
+    }
+    return count;
+}
+
+static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+    static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION");return value&&std::atoi(value)!=0; }();
+    if (disabled || first + 4 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
+        return 0;
+    const auto* mat  = node_at(graph, first);
+    const auto* view = node_at(graph, first + 1);
+    const auto* norm = node_at(graph, first + 2);
+    const bool key   = norm->op == GGML_OP_RMS_NORM;
+    if (key) {
+        if (!sd_cuda_i8_rope_supports(device, mat, view, norm, node_at(graph, first + 3), node_at(graph, first + 4)))
+            return 0;
+        if (int count = concat_match(graph, first, true))
+            return count;
+        if (first + 6 < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph))) {
+            const auto* rope  = node_at(graph, first + 4);
+            const auto* scale = node_at(graph, first + 5);
+            const auto* copy  = node_at(graph, first + 6);
+            if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
+                copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+                return 7;
+        }
+        return 5;
+    }
+    return projection_supports(device, mat, view) ? concat_match(graph, first, false) : 0;
+}
+
+template <bool Aligned>
+static __global__ void prefix_copy(const float* input, half2* output, int64_t count4, int64_t prefix, int64_t tokens, int64_t heads, bool permuted) {
+    const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+    if (i >= count4)
+        return;
+    const int64_t col = i % 32, row = i / 32;
+    const int64_t token = row % prefix, head = (row / prefix) % heads, batch = row / (prefix * heads);
+    const int64_t source_row = permuted ? (batch * prefix + token) * heads + head : (batch * heads + head) * prefix + token;
+    const float* source      = input + source_row * 128 + col * 4;
+    float4 value;
+    if constexpr (Aligned)
+        value = *reinterpret_cast<const float4*>(source);
+    else
+        value = make_float4(source[0], source[1], source[2], source[3]);
+    const int64_t target   = ((batch * heads + head) * (prefix + tokens) + token) * 32 + col;
+    output[2 * target]     = __floats2half2_rn(value.x, value.y);
+    output[2 * target + 1] = __floats2half2_rn(value.z, value.w);
+}
+
+static ggml_status projection_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+    *workspace = ggml_nbytes(nodes[count - 1]);
+    return GGML_STATUS_SUCCESS;
+}
+static void projection_destroy(void*, void*) {}
+
+static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+    const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+    return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+}
+
+static ggml_status projection_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+    const auto* mat    = nodes[0];
+    const auto* view   = nodes[1];
+    const bool key     = nodes[2]->op == GGML_OP_RMS_NORM;
+    const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2];
+    const bool joined  = concat->op == GGML_OP_CONCAT;
+    auto* dst          = nodes[count - 1];
+    auto reg           = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+    auto query         = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+    ggml_ext_cuda_context context;
+    if (!query || !query(execution->backend, &context))
+        return GGML_STATUS_FAILED;
+    auto stream = static_cast<cudaStream_t>(context.stream);
+    bool alias  = reinterpret_cast<uintptr_t>(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0;
+    for (const auto* input : {mat->src[0], mat->src[1], mat->src[2], key ? nodes[3]->src[1] : nullptr, key ? nodes[4]->src[1] : nullptr, joined ? concat->src[0] : nullptr})
+        if (input)
+            alias |= overlaps(dst, input);
+    void* output   = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+    int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0;
+    float scale    = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f;
+    cudaError_t error;
+    if (dst->type == GGML_TYPE_F32)
+        error = projection_launch<float>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0);
+    else if (key)
+        error = projection_launch<half>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, prefix);
+    else
+        error = projection_launch<half, false>(stream, mat, view, nullptr, nullptr, nullptr, output, scale, prefix);
+    if (error == cudaSuccess && joined && prefix > 0) {
+        const auto* input    = concat->src[0];
+        const int64_t count4 = ggml_nelements(input) / 4;
+        auto kernel          = reinterpret_cast<uintptr_t>(input->data) % 16 == 0 ? prefix_copy<true> : prefix_copy<false>;
+        kernel<<<static_cast<unsigned int>((count4 + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(input->data), static_cast<half2*>(output), count4, prefix, view->ne[2], view->ne[1], !key);
+        error = cudaGetLastError();
+    }
+    if (error == cudaSuccess && alias)
+        error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+    return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device) {
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.i8_projection.cuda.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_i8_op();
+    desc.priority             = 350;
+    desc.match                = projection_match;
+    desc.prepare              = projection_prepare;
+    desc.execute              = projection_execute;
+    desc.destroy              = projection_destroy;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}

src/core/cuda/i8_projection.h

@@ -0,0 +1,11 @@
+#ifndef __SD_CORE_CUDA_I8_PROJECTION_H__
+#define __SD_CORE_CUDA_I8_PROJECTION_H__
+
+#include <cuda_runtime.h>
+#include "ggml-ext.h"
+
+bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale);
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
+
+#endif

src/core/cuda/int8.cu

@@ -0,0 +1,739 @@
+#include <cstdlib>
+#include "convrot_kernel.h"
+#include "core/ggml_extensions.h"
+#ifdef SD_CUDA_I8_CUTLASS
+#include "core/cuda/i8_gemm.h"
+#else
+struct SDI8CudaGemmPlan;
+#endif
+
+#include "runtime.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hipblas/hipblas.h>
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+#include <cublas_v2.h>
+#endif
+#include <algorithm>
+#include <climits>
+#include <mutex>
+#include <unordered_map>
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_i8_cuda {
+        template <int Width = 32>
+        static __device__ __forceinline__ float warp_reduce_max(float value) {
+#pragma unroll
+            for (int offset = Width / 2; offset > 0; offset >>= 1)
+                value = fmaxf(value, cuda_shuffle_xor(value, offset, Width));
+            return value;
+        }
+        static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) {
+            const int64_t row = blockIdx.x;
+            const int tid     = threadIdx.x;
+            if (row >= rows) {
+                return;
+            }
+
+            __shared__ float maxima[256];
+            const float* src_row = src + row * k;
+            int8_t* dst_row      = dst + row * k;
+            float local_max      = 0.0f;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                local_max = fmaxf(local_max, fabsf(src_row[i]));
+            }
+            maxima[tid] = local_max;
+            __syncthreads();
+
+            for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+                if (tid < stride) {
+                    maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                }
+                __syncthreads();
+            }
+
+            const float scale = maxima[0] / 127.0f;
+            if (tid == 0) {
+                scales[row] = scale;
+            }
+            if (scale == 0.0f) {
+                for (int64_t i = tid; i < k; i += blockDim.x) {
+                    dst_row[i] = 0;
+                }
+                return;
+            }
+
+            const float inv_scale = 1.0f / scale;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                int value  = __float2int_rn(src_row[i] * inv_scale);
+                value      = max(-127, min(127, value));
+                dst_row[i] = (int8_t)value;
+            }
+        }
+        static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+            a = index == 3 ? -a : a;
+            b = index == 2 ? -b : b;
+            c = index == 1 ? -c : c;
+            d = index == 0 ? -d : d;
+            return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+        }
+        static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) {
+            constexpr int group_size = 256;
+            const int64_t row        = blockIdx.x;
+            const int tid            = threadIdx.x;
+            if (row >= rows)
+                return;
+            extern __shared__ float values[];
+            __shared__ float maxima[32];
+            const float* src_row = src + row * k;
+            int8_t* dst_row      = dst + row * dst_stride;
+            {
+                const int lane = tid % 32;
+                for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+                    float v[8];
+#pragma unroll
+                    for (int j = 0; j < 8; ++j) {
+                        const int64_t i   = group * group_size + j * 32 + lane;
+                        const float value = src_row[i];
+                        v[j]              = value * (1.0f / 16.0f);
+                    }
+                    // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+                    for (int stride = 1; stride <= 4; stride *= 4) {
+                        const int base = lane & ~(3 * stride);
+#pragma unroll
+                        for (int j = 0; j < 8; ++j) {
+                            const float a = cuda_shuffle(v[j], base, 32);
+                            const float b = cuda_shuffle(v[j], base + stride, 32);
+                            const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+                            const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
+                            v[j]          = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+                        }
+                    }
+#pragma unroll
+                    for (int j = 0; j < 8; j += 2) {
+                        const float a = cuda_shuffle(v[j], lane % 16, 32);
+                        const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+                        const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+                        const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
+                        v[j]          = convrot_h4_cuda(a, b, c, d, lane / 16);
+                        v[j + 1]      = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+                    }
+#pragma unroll
+                    for (int j = 0; j < 2; ++j) {
+                        const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+                        for (int t = 0; t < 4; ++t) {
+                            values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+                        }
+                    }
+                }
+                __syncthreads();
+            }
+            float local_max = 0.0f;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                local_max = fmaxf(local_max, fabsf(values[i]));
+            }
+            local_max      = warp_reduce_max<32>(local_max);
+            const int lane = tid & 31;
+            const int warp = tid >> 5;
+            if (lane == 0) {
+                maxima[warp] = local_max;
+            }
+            __syncthreads();
+
+            if (warp == 0) {
+                local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+                local_max = warp_reduce_max<32>(local_max);
+                if (lane == 0) {
+                    maxima[0] = local_max / 127.0f;
+                    if (scales != nullptr) {
+                        scales[row] = maxima[0];
+                    } else {
+                        *((float*)(dst_row + k)) = maxima[0];
+                    }
+                }
+            }
+            __syncthreads();
+
+            const float row_scale = maxima[0];
+            const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+            for (int64_t i = tid; i < k; i += blockDim.x) {
+                int value  = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+                value      = max(-127, min(127, value));
+                dst_row[i] = (int8_t)value;
+            }
+        }
+
+        static __global__ void convrot_group_amax_i8_cuda(
+            const float* src,
+            float* partial_max,
+            int64_t k,
+            int64_t groups,
+            int64_t total_groups) {
+            const int64_t index = blockIdx.x;
+            const int tid       = threadIdx.x;
+            if (index >= total_groups) {
+                return;
+            }
+
+            __shared__ float values[256];
+            __shared__ float maxima[2];
+            const int64_t row   = index / groups;
+            const int64_t group = index - row * groups;
+            const float* input  = src + row * k + group * 256;
+            values[tid]         = input[tid] * (1.0f / 16.0f);
+            values[tid + 64]    = input[tid + 64] * (1.0f / 16.0f);
+            values[tid + 128]   = input[tid + 128] * (1.0f / 16.0f);
+            values[tid + 192]   = input[tid + 192] * (1.0f / 16.0f);
+            __syncthreads();
+
+#pragma unroll
+            for (int stride = 1; stride < 256; stride *= 4) {
+                const int base = (tid / stride) * 4 * stride + tid % stride;
+                const int i0   = base;
+                const int i1   = i0 + stride;
+                const int i2   = i1 + stride;
+                const int i3   = i2 + stride;
+                const float a  = values[i0];
+                const float b  = values[i1];
+                const float c  = values[i2];
+                const float d  = values[i3];
+                values[i0]     = a + b + c - d;
+                values[i1]     = a + b - c + d;
+                values[i2]     = a - b + c + d;
+                values[i3]     = -a + b + c + d;
+                __syncthreads();
+            }
+
+            float value = fmaxf(fabsf(values[tid]), fabsf(values[tid + 64]));
+            value       = fmaxf(value, fabsf(values[tid + 128]));
+            value       = fmaxf(value, fabsf(values[tid + 192]));
+            value       = warp_reduce_max<32>(value);
+            if ((tid & 31) == 0) {
+                maxima[tid >> 5] = value;
+            }
+            __syncthreads();
+            if (tid == 0) {
+                partial_max[index] = fmaxf(maxima[0], maxima[1]);
+            }
+        }
+
+        static __global__ void reduce_convrot_row_amax_i8_cuda(
+            const float* partial_max,
+            float* scales,
+            int64_t groups,
+            int64_t rows) {
+            const int64_t row = blockIdx.x;
+            const int tid     = threadIdx.x;
+            if (row >= rows) {
+                return;
+            }
+
+            __shared__ float maxima[256];
+            float value = 0.0f;
+            for (int64_t group = tid; group < groups; group += blockDim.x) {
+                value = fmaxf(value, partial_max[row * groups + group]);
+            }
+            maxima[tid] = value;
+            __syncthreads();
+            for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+                if (tid < stride) {
+                    maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                }
+                __syncthreads();
+            }
+            if (tid == 0) {
+                scales[row] = maxima[0] / 127.0f;
+            }
+        }
+
+        static __global__ void convrot_group_quantize_i8_cuda(
+            const float* src,
+            int8_t* dst,
+            const float* scales,
+            int64_t k,
+            int64_t groups,
+            int64_t total_groups,
+            int64_t dst_stride) {
+            const int64_t index = blockIdx.x;
+            const int tid       = threadIdx.x;
+            if (index >= total_groups) {
+                return;
+            }
+
+            __shared__ float values[256];
+            const int64_t row   = index / groups;
+            const int64_t group = index - row * groups;
+            const float* input  = src + row * k + group * 256;
+            int8_t* output      = dst + row * dst_stride + group * 256;
+            values[tid]         = input[tid] * (1.0f / 16.0f);
+            values[tid + 64]    = input[tid + 64] * (1.0f / 16.0f);
+            values[tid + 128]   = input[tid + 128] * (1.0f / 16.0f);
+            values[tid + 192]   = input[tid + 192] * (1.0f / 16.0f);
+            __syncthreads();
+
+#pragma unroll
+            for (int stride = 1; stride < 256; stride *= 4) {
+                const int base = (tid / stride) * 4 * stride + tid % stride;
+                const int i0   = base;
+                const int i1   = i0 + stride;
+                const int i2   = i1 + stride;
+                const int i3   = i2 + stride;
+                const float a  = values[i0];
+                const float b  = values[i1];
+                const float c  = values[i2];
+                const float d  = values[i3];
+                values[i0]     = a + b + c - d;
+                values[i1]     = a + b - c + d;
+                values[i2]     = a - b + c + d;
+                values[i3]     = -a + b + c + d;
+                __syncthreads();
+            }
+
+            const float row_scale = scales[row];
+            const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+            const int indices[4]  = {tid, tid + 64, tid + 128, tid + 192};
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                int value          = row_scale == 0.0f ? 0 : __float2int_rn(values[indices[i]] * inv_scale);
+                value              = max(-127, min(127, value));
+                output[indices[i]] = (int8_t)value;
+            }
+            if (group == 0 && tid == 0 && dst_stride != k) {
+                *((float*)(dst + row * dst_stride + k)) = row_scale;
+            }
+        }
+
+        static __global__ void dequantize_i32_rows_cuda(
+            const int32_t* src,
+            const float* scales,
+            const float* weight_scales,
+            const float* bias,
+            float* dst,
+            int64_t n,
+            int64_t rows) {
+            const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+            const int64_t row    = blockIdx.y + int64_t(blockIdx.z) * 65535;
+            if (output < n && row < rows) {
+                const int64_t index          = row * n + output;
+                const float activation_scale = scales[row];
+                float value                  = (float)src[index] * activation_scale;
+                if (weight_scales != nullptr) {
+                    value *= weight_scales[output];
+                }
+                if (bias != nullptr) {
+                    value += bias[output];
+                }
+                dst[index] = value;
+            }
+        }
+
+        static __global__ void dequantize_i32_rows_vec4_cuda(
+            const int4* src,
+            const float* scales,
+            const float4* weight_scales,
+            const float4* bias,
+            float4* dst,
+            int64_t n4,
+            int64_t rows) {
+            const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+            const int64_t row    = blockIdx.y + int64_t(blockIdx.z) * 65535;
+            if (output >= n4 || row >= rows) {
+                return;
+            }
+            const int64_t index = row * n4 + output;
+            const int4 v        = src[index];
+            const float scale   = scales[row];
+            float4 value        = make_float4(v.x * scale, v.y * scale, v.z * scale, v.w * scale);
+            if (weight_scales != nullptr) {
+                const float4 w = weight_scales[output];
+                value.x *= w.x;
+                value.y *= w.y;
+                value.z *= w.z;
+                value.w *= w.w;
+            }
+            if (bias != nullptr) {
+                const float4 b = bias[output];
+                value.x += b.x;
+                value.y += b.y;
+                value.z += b.z;
+                value.w += b.w;
+            }
+            dst[index] = value;
+        }
+
+        static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) {
+            if (!tensor)
+                return true;
+            size_t offset = 0;
+            while (tensor->view_src) {
+                offset += tensor->view_offs;
+                tensor = tensor->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+        }
+        static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            SDI8Operation op;
+            if (!sd_ggml_i8_describe(node, &op) || !aligned(node))
+                return false;
+            const int64_t k = op.weight->ne[0], n = op.weight->ne[1], rows = ggml_nrows(node);
+            if (k > INT_MAX || n > INT_MAX || rows > INT_MAX - 3 || k % 4 || n % 4 || (op.input->type == GGML_TYPE_F32 && op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256))))
+                return false;
+            for (auto* input : node->src)
+                if (!aligned(input))
+                    return false;
+            const auto info = cuda_device_info(device);
+            return info.int8_blas &&
+                   ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
+        }
+        static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            int group;
+            if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0]))
+                return false;
+            const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]);
+            return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0;
+        }
+        enum Kind { MAT,
+                    QUANTIZE,
+                    SWIGLU,
+                    SWIGLU_QUANTIZE,
+                    CONCAT };
+        static bool fusion_disabled() {
+            static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+            return disabled;
+        }
+        static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+            auto* g = const_cast<ggml_cgraph*>(graph);
+            if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g))
+                return 0;
+            const auto* mat = ggml_graph_node(g, first);
+            SDI8Operation op;
+            if (!supports(device, mat, nullptr) || !sd_ggml_i8_describe(mat, &op) || !op.weight_scale || op.bias)
+                return 0;
+            const auto* next = ggml_graph_node(g, first + 1);
+            if (next->op == GGML_OP_CONCAT && next->src[1] == mat && next->src[0] != mat && next->op_params[0] == 1 &&
+                next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 &&
+                next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0]))
+                return 2;
+            if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 ||
+                next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) ||
+                ggml_nelements(next) > int64_t(INT_MAX) * 256)
+                return 0;
+            if (first + 2 < ggml_graph_n_nodes(g)) {
+                const auto* quant = ggml_graph_node(g, first + 2);
+                if (quant->src[0] == next && convrot_supports(device, quant, nullptr) &&
+                    size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
+                    return 3;
+            }
+            return 2;
+        }
+        struct Plan {
+            ggml_ext_gpu_device_info device;
+            Kind kind;
+            int64_t k, n, rows, padded;
+            size_t quant, scales, accum, partial, output, gemm_offset;
+            SDI8CudaGemmPlan* gemm = nullptr;
+            bool fast_swiglu       = false;
+        };
+        static size_t reserve(size_t& total, size_t bytes) {
+            const size_t offset = total;
+            total += (bytes + 255) / 256 * 256;
+            return offset;
+        }
+        static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) {
+            const auto* dst = nodes[count - 1];
+            SDI8Operation op;
+            const bool mat = sd_ggml_i8_describe(nodes[0], &op);
+            auto* p        = new Plan{};
+            *state         = p;
+            auto device    = ggml_backend_get_device(backend);
+            p->device      = cuda_device_info(device);
+            p->kind        = !mat ? QUANTIZE : count == 3              ? SWIGLU_QUANTIZE
+                                           : count == 1                ? MAT
+                                           : dst->op == GGML_OP_CONCAT ? CONCAT
+                                                                       : SWIGLU;
+            p->k           = mat ? op.weight->ne[0] : dst->ne[0];
+            p->n           = mat ? nodes[0]->ne[0] : 0;
+            p->rows        = mat ? ggml_nrows(nodes[0]) : ggml_nrows(dst->src[0]);
+            p->padded      = (p->rows + 3) / 4 * 4;
+            size_t total   = 0;
+            if (mat) {
+                p->quant  = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0);
+                p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0);
+#ifdef SD_CUDA_I8_CUTLASS
+                size_t gemm_bytes = 0;
+                // Retain cuBLAS when I32 accumulation can overflow.
+                if (op.weight_scale && !op.bias && p->k <= INT_MAX / 16384 &&
+                    aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
+                    (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) &&
+                    (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) {
+                    if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) {
+                        p->gemm        = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes);
+                        p->fast_swiglu = p->gemm != nullptr;
+                    }
+                    if (!p->gemm)
+                        p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, false, &gemm_bytes);
+                }
+                p->gemm_offset = reserve(total, gemm_bytes);
+#endif
+                const bool intermediate = !p->gemm || p->kind == SWIGLU_QUANTIZE || (p->kind == SWIGLU && !p->fast_swiglu);
+                p->accum                = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
+            }
+            const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
+            p->partial                = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
+            p->output                 = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
+            *workspace                = total;
+            return GGML_STATUS_SUCCESS;
+        }
+        static void destroy(void* state, void*) {
+            auto* p = static_cast<Plan*>(state);
+#ifdef SD_CUDA_I8_CUTLASS
+            if (p)
+                sd_cuda_i8_gemm_free(p->gemm);
+#endif
+            delete p;
+        }
+        static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+            if (!b)
+                return false;
+            const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+            return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+        }
+        static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) {
+            if (p.padded > p.rows) {
+                auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                if (error != cuda_success)
+                    return error;
+            }
+            const size_t shared = size_t(p.k) * sizeof(float);
+            if (shared + 128 <= p.device.shared_memory) {
+                static std::mutex mutex;
+                static std::unordered_map<int, bool> configured;
+                {
+                    std::lock_guard<std::mutex> lock(mutex);
+                    if (!configured[p.device.device]) {
+                        auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
+                        if (error != cuda_success)
+                            return error;
+                        configured[p.device.device] = true;
+                    }
+                }
+                int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256
+                                                              : 1024;
+                quantize_convrot<<<unsigned(p.rows), threads, shared, stream>>>(src, dst, scales, p.k, p.rows, p.k);
+            } else {
+                const int64_t groups = p.k / 256, total_groups = p.rows * groups;
+                convrot_group_amax_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, partial, p.k, groups, total_groups);
+                reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
+                convrot_group_quantize_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, dst, scales, p.k, groups, total_groups, p.k);
+            }
+            return cuda_last_error();
+        }
+        template <bool Dequantize>
+        static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
+            const int64_t i    = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            const int64_t half = n / 2;
+            if (i >= rows * half)
+                return;
+            const int64_t row = i / half, col = i % half;
+            float gate, up;
+            if constexpr (Dequantize) {
+                gate = float(reinterpret_cast<const int32_t*>(src)[row * n + col]) * xs[row];
+                up   = float(reinterpret_cast<const int32_t*>(src)[row * n + half + col]) * xs[row];
+                gate *= ws[col];
+                up *= ws[half + col];
+            } else {
+                gate = src[row * n + col];
+                up   = src[row * n + half + col];
+            }
+            dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
+        }
+        template <bool Swiglu, bool Dequantize, bool Rotated>
+        static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
+            const int64_t hidden = p.n / 2;
+            auto kernel          = sd_convrot::quantize<Swiglu, Dequantize, false, false, Rotated>;
+            static std::mutex mutex;
+            static std::unordered_map<int, bool> initialized;
+            {
+                std::lock_guard<std::mutex> lock(mutex);
+                if (!initialized[p.device.device]) {
+                    auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
+                    if (error != cuda_success)
+                        return error;
+                    initialized[p.device.device] = true;
+                }
+            }
+            if (p.padded > p.rows) {
+                auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
+                if (error != cuda_success)
+                    return error;
+            }
+            const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
+                                                                                                           : 1024)
+                                                                  : 256;
+            kernel<<<unsigned(p.rows), threads, size_t(hidden) * sizeof(float), stream>>>(src, output,
+                                                                                          reinterpret_cast<float*>(output + hidden * p.padded), hidden, p.rows, hidden,
+                                                                                          Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr);
+            return cuda_last_error();
+        }
+        static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+            const auto& p = *static_cast<const Plan*>(state);
+            auto reg      = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+            ggml_ext_gpu_context context;
+            if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
+                return GGML_STATUS_FAILED;
+            auto stream   = static_cast<CUDAStream>(context.stream);
+            auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+            auto* dst     = nodes[count - 1];
+            SDI8Operation op;
+            const bool mat = sd_ggml_i8_describe(nodes[0], &op);
+            bool alias     = false;
+            for (int i = 0; i < count; ++i)
+                for (auto* input : nodes[i]->src) {
+                    bool internal = false;
+                    for (int j = 0; j < count; ++j)
+                        internal |= input == nodes[j];
+                    if (!internal)
+                        alias |= overlaps(dst, input);
+                }
+            // CUTLASS quantization epilogues read only the workspace intermediate after GEMM completes.
+            if (p.gemm && p.kind == SWIGLU_QUANTIZE)
+                alias = false;
+            void* output    = alias ? scratch + p.output : dst->data;
+            CUDAError error = cuda_success;
+            if (!mat) {
+                auto* quant  = static_cast<int8_t*>(output);
+                auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
+                error        = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch + p.partial));
+            } else {
+                auto* quant = reinterpret_cast<int8_t*>(scratch + p.quant);
+                const float* scales;
+                if (op.input->type == GGML_TYPE_F32) {
+                    auto* temporary_scales = reinterpret_cast<float*>(scratch + p.scales);
+                    scales                 = temporary_scales;
+                    if (op.group_size)
+                        error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
+                    else {
+                        if (p.padded > p.rows)
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                        if (error == cuda_success)
+                            quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(static_cast<const float*>(op.input->data), quant, temporary_scales, p.k, p.rows);
+                    }
+                } else {
+                    auto* data = static_cast<int8_t*>(op.input->data);
+                    scales     = op.packed ? reinterpret_cast<const float*>(data + p.k * p.padded) : static_cast<const float*>(op.input_scale->data);
+                    if (op.packed || p.rows == p.padded)
+                        quant = data;
+                    else {
+                        error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
+                        if (error == cuda_success)
+                            error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+                    }
+                }
+                if (error != cuda_success || cuda_last_error() != cuda_success)
+                    return GGML_STATUS_FAILED;
+                const auto* ws          = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+                const auto* bias        = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+                const size_t prefix     = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
+                auto* gemm_output       = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
+                auto* accum             = reinterpret_cast<int32_t*>(scratch + p.accum);
+                const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
+#ifdef SD_CUDA_I8_CUTLASS
+                if (p.gemm) {
+                    auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
+                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch + p.gemm_offset, stream);
+                    if (status != GGML_STATUS_SUCCESS)
+                        return status;
+                } else
+#endif
+                {
+                    auto blas_query = reinterpret_cast<ggml_ext_gpu_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1"));
+                    auto handle     = blas_query ? blas_query(execution->backend) : nullptr;
+                    if (!handle)
+                        return GGML_STATUS_FAILED;
+#if defined(SD_CUDA_VENDOR_MUSA)
+                    return GGML_STATUS_FAILED;
+#else
+                    const int32_t alpha = 1, beta = 0;
+#if defined(SD_CUDA_VENDOR_HIP)
+#if HIP_VERSION >= 60500000
+                    constexpr auto compute_type = HIPBLAS_COMPUTE_32I;
+#else
+                    constexpr auto compute_type = HIPBLAS_R_32I;
+#endif
+                    const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
+                                                      int(p.n), int(p.padded), int(p.k), &alpha,
+                                                      op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
+                                                      &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
+                    if (status != HIPBLAS_STATUS_SUCCESS)
+#else
+                    const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
+                                                     int(p.n), int(p.padded), int(p.k), &alpha,
+                                                     op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k),
+                                                     &beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+                    if (status != CUBLAS_STATUS_SUCCESS)
+#endif
+                        return GGML_STATUS_FAILED;
+#endif
+                }
+                const auto* values = reinterpret_cast<const float*>(accum);
+                if (p.kind == SWIGLU_QUANTIZE) {
+                    if (p.fast_swiglu)
+                        error = quantize_epilogue<false, false, true>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
+                    else if (p.gemm)
+                        error = quantize_epilogue<true, false, false>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
+                    else
+                        error = quantize_epilogue<true, true, false>(p, values, scales, ws, static_cast<int8_t*>(output), stream);
+                } else if (p.kind == SWIGLU && !p.fast_swiglu) {
+                    const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256);
+                    if (p.gemm)
+                        swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+                    else
+                        swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+                    error = cuda_last_error();
+                } else if (!p.gemm) {
+                    const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
+                    if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(gemm_output)) % 16 == 0) {
+                        dequantize_i32_rows_vec4_cuda<<<dim3(unsigned((p.n + 1023) / 1024), gy, gz), 256, 0, stream>>>(
+                            reinterpret_cast<const int4*>(accum), scales, reinterpret_cast<const float4*>(ws), reinterpret_cast<const float4*>(bias), reinterpret_cast<float4*>(gemm_output), p.n / 4, p.rows);
+                    } else {
+                        dequantize_i32_rows_cuda<<<dim3(unsigned((p.n + 255) / 256), gy, gz), 256, 0, stream>>>(accum, scales, ws, bias, gemm_output, p.n, p.rows);
+                    }
+                    error = cuda_last_error();
+                }
+                if (error == cuda_success && prefix)
+                    error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
+            }
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+
+    }
+
+    void register_i8(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.linear_i8." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_i8_op();
+        desc.priority             = 100;
+        desc.supports             = sd_i8_cuda::supports;
+        desc.prepare              = sd_i8_cuda::prepare;
+        desc.execute              = sd_i8_cuda::execute;
+        desc.destroy              = sd_i8_cuda::destroy;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name     = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1";
+        desc.priority = 240;
+        desc.supports = nullptr;
+        desc.match    = sd_i8_cuda::match;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.priority  = 100;
+        desc.match     = nullptr;
+        desc.name      = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
+        desc.custom_op = sd_ggml_convrot_op();
+        desc.supports  = sd_i8_cuda::convrot_supports;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}

src/core/cuda/registration.cu

@@ -0,0 +1,36 @@
+#include "runtime.h"
+
+#ifdef SD_CUDA_CUTE_ATTENTION
+void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
+void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+    void register_rope(ggml_backend_dev_t device);
+    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);
+    void register_sol_attention(ggml_backend_dev_t device);
+#endif
+
+    void register_extensions(ggml_backend_dev_t device) {
+        if (cuda_device_info(device).device < 0)
+            return;
+        register_affine(device);
+        register_i8(device);
+        register_convrot(device);
+        register_rope(device);
+#if !defined(SD_CUDA_VENDOR_HIP) && !defined(SD_CUDA_VENDOR_MUSA)
+        register_sage_attention(device);
+        register_sol_attention(device);
+#endif
+#ifdef SD_CUDA_CUTE_ATTENTION
+        sd_ggml_register_cuda_i8_projection(device);
+        sd_ggml_register_cuda_i8_gated_residual(device);
+        sd_ggml_register_cuda_attention(device);
+#endif
+    }
+}

src/core/cuda/registration.h

@@ -0,0 +1,13 @@
+#ifndef __SD_CORE_CUDA_REGISTRATION_H__
+#define __SD_CORE_CUDA_REGISTRATION_H__
+#include "ggml-backend.h"
+namespace sd_cuda {
+    void register_extensions(ggml_backend_dev_t device);
+}
+namespace sd_hip {
+    void register_extensions(ggml_backend_dev_t device);
+}
+namespace sd_musa {
+    void register_extensions(ggml_backend_dev_t device);
+}
+#endif

src/core/cuda/rope.cu

@@ -0,0 +1,337 @@
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <type_traits>
+
+namespace SD_CUDA_NAMESPACE {
+    static __device__ __forceinline__ float sd_warp_sum(float v) {
+        for (int mask = 16; mask > 0; mask >>= 1)
+            v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
+        return v;
+    }
+    template <int D, bool interleaved, typename T>
+    static __global__ void rms_norm_mul_rope_apply_fused(
+        const float* x,
+        const float* weight,
+        const float* pe,
+        T* dst,
+        int64_t heads,
+        int64_t tokens,
+        int64_t rows,
+        int64_t nb1,
+        int64_t nb2,
+        int64_t nb3,
+        int rot_half,
+        float eps,
+        float output_scale) {
+        const int lane    = threadIdx.x;
+        const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+        if (row >= rows) {
+            return;
+        }
+        const int64_t head  = row % heads;
+        const int64_t token = (row / heads) % tokens;
+        const int64_t batch = row / (heads * tokens);
+        const float* x_row  = x + head * nb1 + token * nb2 + batch * nb3;
+        float v[D / 32], sums[8] = {};
+#pragma unroll
+        for (int j = 0; j < D / 32; ++j) {
+            v[j]    = x_row[j * 32 + lane];
+            sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
+        }
+        // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
+        const float sum = __fadd_rn(
+            __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
+            __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
+        const float scale = rsqrtf(sum / float(D) + eps);
+        T* out            = dst + ((batch * heads + head) * tokens + token) * D;
+#pragma unroll
+        for (int j = 0; j < D / 32; ++j) {
+            const int col     = j * 32 + lane;
+            const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
+            float rotated     = value;
+            if constexpr (interleaved) {
+                const float other = cuda_shuffle_xor(value, 1, 32);
+                if (col < 2 * rot_half) {
+                    const float* m      = pe + 4 * (token * rot_half + col / 2);
+                    const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+                    rotated             = lane % 2 == 0
+                                              ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
+                                              : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
+                }
+            } else {
+                if (col < 2 * rot_half) {
+                    const int partner   = col < rot_half ? col + rot_half : col - rot_half;
+                    const float other   = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
+                    const float* m      = pe + 4 * (token * rot_half + col % rot_half);
+                    const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+                    rotated             = col < rot_half
+                                              ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
+                                              : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
+                }
+            }
+            if constexpr (std::is_same<T, half>::value) {
+                out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
+            } else {
+                out[col] = rotated;
+            }
+        }
+    }
+
+    struct RopeApplyStrides {
+        int64_t x[4];
+        int64_t pe[4];
+    };
+
+    template <bool interleaved>
+    static __global__ void rope_apply_f32(
+        const float* x,
+        const float* pe,
+        float* dst,
+        int64_t half_dim,
+        int64_t rot_half,
+        int64_t heads,
+        int64_t tokens,
+        int64_t pairs,
+        RopeApplyStrides strides) {
+        const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+        if (index >= pairs) {
+            return;
+        }
+        const int64_t j     = index % half_dim;
+        const int64_t row   = index / half_dim;
+        const int64_t token = row % tokens;
+        const int64_t head  = (row / tokens) % heads;
+        const int64_t batch = row / (tokens * heads);
+        const int64_t i0    = interleaved || j >= rot_half ? 2 * j : j;
+        const int64_t i1    = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half;
+        const float* x_row  = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3];
+        const float a       = x_row[i0 * strides.x[0]];
+        const float b       = x_row[i1 * strides.x[0]];
+        float* out          = dst + row * half_dim * 2;
+        if (j >= rot_half) {
+            out[i0] = a;
+            out[i1] = b;
+            return;
+        }
+        const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3];
+        const float c       = matrix[0];
+        const float s       = matrix[strides.pe[1]];
+        // Match the separate mul/add rounding of the unfused graph.
+        if (interleaved) {
+            out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]]));
+            out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]]));
+        } else {
+            out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
+            out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
+        }
+    }
+
+    static bool rope_interleaved(const ggml_tensor* rope) {
+        int32_t value;
+        std::memcpy(&value, ggml_ext_op_attributes(rope), sizeof(value));
+        return value != 0;
+    }
+
+    static float param_f32(const ggml_tensor* node, int index) {
+        float value;
+        std::memcpy(&value, node->op_params + index, sizeof(value));
+        return value;
+    }
+
+    static bool rope_is_node(const ggml_tensor* node) {
+        return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
+    }
+
+    static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
+        const auto* x  = rope->src[0];
+        const auto* pe = rope->src[1];
+        if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
+            x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
+            pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
+            rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
+            (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
+            return false;
+        }
+        for (int i = 0; i < 4; ++i) {
+            if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
+                return false;
+        }
+        return true;
+    }
+
+    bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+        if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !rope_is_node(rope) ||
+            mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
+            return false;
+        const auto* x      = norm->src[0];
+        const auto* weight = mul->src[1];
+        return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
+               x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
+               ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
+               weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
+               x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
+               x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
+               (ggml_nrows(x) + 3) / 4 <= INT_MAX;
+    }
+
+    template <typename T, bool Interleaved>
+    static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
+        const auto* x      = nodes[0]->src[0];
+        const auto* pe     = nodes[2]->src[1];
+        const int64_t rows = ggml_nrows(x);
+#define SD_ROPE(D)                                                                                                               \
+    case D:                                                                                                                      \
+        rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
+            static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data),                               \
+            static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows,                                               \
+            x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale);                            \
+        break
+        switch (x->ne[0]) {
+            SD_ROPE(32);
+            SD_ROPE(64);
+            SD_ROPE(96);
+            SD_ROPE(128);
+            SD_ROPE(160);
+            SD_ROPE(192);
+            SD_ROPE(224);
+            SD_ROPE(256);
+            default:
+                return cuda_invalid_value;
+        }
+#undef SD_ROPE
+        return cuda_last_error();
+    }
+
+    CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
+        if (fp16) {
+            return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
+                                              : rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
+        }
+        return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
+                                          : rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
+    }
+
+    static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+        static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
+        if (disabled)
+            return 0;
+        auto* g = const_cast<ggml_cgraph*>(graph);
+        if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
+                                                      ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
+            return 0;
+        if (first + 4 < ggml_graph_n_nodes(g)) {
+            const auto* rope  = ggml_graph_node(g, first + 2);
+            const auto* scale = ggml_graph_node(g, first + 3);
+            const auto* copy  = ggml_graph_node(g, first + 4);
+            if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
+                scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
+                ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+                return 5;
+        }
+        return 3;
+    }
+
+    static void rope_optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+        static const bool disabled = [] {
+            const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
+            return value && std::atoi(value) != 0;
+        }();
+        if (disabled)
+            return;
+        const int total = ggml_graph_n_nodes(graph);
+        for (int i = 2; i + 1 < total; ++i) {
+            const auto* mul  = ggml_graph_node(graph, i);
+            const auto* rope = ggml_graph_node(graph, i + 1);
+            if (mul->op != GGML_OP_MUL || !rope_is_node(rope) || rope->src[0] != mul ||
+                !mul->src[0] || mul->src[0]->op != GGML_OP_RMS_NORM)
+                continue;
+            const auto* norm = mul->src[0];
+            for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
+                if (ggml_graph_node(graph, j) != norm)
+                    continue;
+                if (j + 1 < i) {
+                    if (j >= 2) {
+                        const auto* mat  = ggml_graph_node(graph, j - 2);
+                        const auto* view = ggml_graph_node(graph, j - 1);
+                        if (ggml_ext_op_get(mat) == sd_ggml_i8_op() &&
+                            view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view &&
+                            ggml_ext_graph_delay(graph, j - 2, 3, i))
+                            break;
+                    }
+                    ggml_ext_graph_delay(graph, j, 1, i);
+                }
+                break;
+            }
+        }
+    }
+
+    static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+        *workspace = ggml_nbytes(nodes[count - 1]);
+        return GGML_STATUS_SUCCESS;
+    }
+
+    static void rope_destroy(void*, void*) {}
+
+    static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+        const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+        return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+    }
+
+    static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
+            return GGML_STATUS_FAILED;
+        auto stream      = static_cast<CUDAStream>(context.stream);
+        const auto* rope = nodes[count == 1 ? 0 : 2];
+        const auto* x    = count == 1 ? rope->src[0] : nodes[0]->src[0];
+        auto* dst        = nodes[count - 1];
+        const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
+        void* output     = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+        CUDAError error;
+        if (count > 1) {
+            error = rms_rope_launch(stream, nodes, output, count == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f);
+        } else {
+            RopeApplyStrides strides;
+            const auto* pe = rope->src[1];
+            for (int i = 0; i < 4; ++i) {
+                strides.x[i]  = x->nb[i] / 4;
+                strides.pe[i] = pe->nb[i] / 4;
+            }
+            const int64_t pairs = ggml_nelements(x) / 2;
+            auto kernel         = rope_interleaved(rope) ? rope_apply_f32<true> : rope_apply_f32<false>;
+            kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
+                                                                                       static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
+            error = cuda_last_error();
+        }
+        if (error == cuda_success && alias)
+            error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+        return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_rope(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.rope_apply." SD_CUDA_NAME ".v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_rope_op();
+        desc.priority             = 100;
+        desc.supports             = rope_supports;
+        desc.prepare              = rope_prepare;
+        desc.destroy              = rope_destroy;
+        desc.execute              = rope_execute;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+        desc.name      = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
+        desc.op        = GGML_OP_RMS_NORM;
+        desc.custom_op = nullptr;
+        desc.priority  = 250;
+        desc.supports  = nullptr;
+        desc.match     = rope_match;
+        desc.optimize  = rope_optimize;
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+
+}

src/core/cuda/rope.h

@@ -0,0 +1,11 @@
+#ifndef __SD_CORE_CUDA_ROPE_H__
+#define __SD_CORE_CUDA_ROPE_H__
+
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+    CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale);
+
+}
+#endif

src/core/cuda/runtime.h

@@ -0,0 +1,111 @@
+#ifndef __SD_CORE_CUDA_RUNTIME_H__
+#define __SD_CORE_CUDA_RUNTIME_H__
+
+#include "ggml-ext.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#define SD_CUDA_NAMESPACE sd_hip
+#define SD_CUDA_NAME "hip"
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <musa_fp16.h>
+#include <musa_runtime.h>
+#define SD_CUDA_NAMESPACE sd_musa
+#define SD_CUDA_NAME "musa"
+#else
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#define SD_CUDA_NAMESPACE sd_cuda
+#define SD_CUDA_NAME "cuda"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+#if defined(SD_CUDA_VENDOR_HIP)
+    using CUDAError                          = hipError_t;
+    using CUDAStream                         = hipStream_t;
+    static constexpr auto cuda_success       = hipSuccess;
+    static constexpr auto cuda_invalid_value = hipErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_HIP;
+    static CUDAError cuda_last_error() {
+        return hipGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return hipMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel, int) {
+        // HIP uses the device's shared-memory limit without CUDA's opt-in attribute.
+        return cuda_success;
+    }
+#elif defined(SD_CUDA_VENDOR_MUSA)
+    using CUDAError                          = musaError_t;
+    using CUDAStream                         = musaStream_t;
+    static constexpr auto cuda_success       = musaSuccess;
+    static constexpr auto cuda_invalid_value = musaErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_MUSA;
+    static CUDAError cuda_last_error() {
+        return musaGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return musaMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+        return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+    }
+#else
+    using CUDAError                          = cudaError_t;
+    using CUDAStream                         = cudaStream_t;
+    static constexpr auto cuda_success       = cudaSuccess;
+    static constexpr auto cuda_invalid_value = cudaErrorInvalidValue;
+    static constexpr auto cuda_kind          = GGML_EXT_GPU_CUDA;
+    static CUDAError cuda_last_error() {
+        return cudaGetLastError();
+    }
+    static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+        return cudaMemsetAsync(dst, value, bytes, stream);
+    }
+    static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+        return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream);
+    }
+    template <typename Kernel>
+    static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+        return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+    }
+#endif
+    template <typename T>
+    static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+        return __shfl(value, lane, width);
+#else
+        return __shfl_sync(0xffffffff, value, lane, width);
+#endif
+    }
+    template <typename T>
+    static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+        return __shfl_xor(value, lane, width);
+#else
+        return __shfl_xor_sync(0xffffffff, value, lane, width);
+#endif
+    }
+    static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
+        auto reg   = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+        auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
+        return query && query(backend, context) && context->kind == cuda_kind;
+    }
+    static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
+        ggml_ext_gpu_device_info info = {};
+        auto query                    = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
+        if (!query || !query(device, &info) || info.kind != cuda_kind)
+            return {cuda_kind, -1, 0, 0, false};
+        return info;
+    }
+}
+#endif

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,592 @@
+/*
+ * 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 "attn_utils.cuh"
+#include "math.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,622 @@
+/*
+ * 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 "attn_utils.cuh"
+#include "math.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,814 @@
+/*
+ * 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 "math.cuh"
+#include "mma.cuh"
+#include "numeric_conversion.cuh"
+#include "permuted_smem.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,166 @@
+/*
+ * 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,683 @@
+/*
+ * 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,148 @@
+/*
+ * 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,201 @@
+/*
+ * 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/cuda/sol_attention.cu

@@ -0,0 +1,409 @@
+// Sol-Attn diagonal-threshold algorithm:
+// https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+#if __CUDACC_VER_MAJOR__ >= 12
+
+#include <cuda_bf16.h>
+#include <algorithm>
+#include <climits>
+
+namespace sd_cuda {
+
+    static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+        for (int offset = 16; offset > 0; offset >>= 1)
+            x += __shfl_xor_sync(0xffffffff, x, offset);
+        return x;
+    }
+
+    static constexpr int SOL_DIM   = 128;
+    static constexpr int SOL_BLOCK = 64;
+
+    struct SolWorkspace {
+        size_t q, k, v, kc, vc, qm, stats, threshold, size;
+        int tokens, blocks;
+
+        explicit SolWorkspace(const ggml_tensor* dst) {
+            const auto* input  = dst->src[0];
+            const size_t heads = input->ne[2] * input->ne[3];
+            tokens             = GGML_PAD(input->ne[1], SOL_BLOCK);
+            blocks             = GGML_PAD(tokens / SOL_BLOCK, SOL_BLOCK);
+            size               = 0;
+            q                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            k                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            v                  = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+            kc                 = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+            vc                 = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+            qm                 = reserve(heads * (tokens / SOL_BLOCK) * SOL_DIM * sizeof(float));
+            stats              = reserve(heads * 2 * SOL_DIM * sizeof(float));
+            threshold          = reserve(heads * (tokens / SOL_BLOCK) * sizeof(float));
+        }
+
+        size_t reserve(size_t bytes) {
+            const size_t offset = size;
+            size += GGML_PAD(bytes, 256);
+            return offset;
+        }
+    };
+
+    static __global__ void sol_pack(const float* q, const float* k, const float* v, nv_bfloat16* qb, nv_bfloat16* kb, nv_bfloat16* vb, int length, int padded) {
+        const int index = blockIdx.x * blockDim.x + threadIdx.x;
+        if (index >= padded * SOL_DIM) {
+            return;
+        }
+        const int token   = index / SOL_DIM;
+        const int channel = index % SOL_DIM;
+        const size_t src  = size_t(blockIdx.y) * length * SOL_DIM + index;
+        const size_t dst  = size_t(blockIdx.y) * padded * SOL_DIM + index;
+        qb[dst]           = __float2bfloat16(token < length ? q[src] : 0.f);
+        kb[dst]           = __float2bfloat16(token < length ? k[src] : 0.f);
+        vb[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token] =
+            __float2bfloat16(token < length ? v[src] : 0.f);
+    }
+
+    static __global__ void sol_summaries(const nv_bfloat16* q, const nv_bfloat16* k, const nv_bfloat16* v, nv_bfloat16* kc, nv_bfloat16* vc, float* qm, int length, int padded, int padded_blocks) {
+        const int channel = threadIdx.x;
+        const int block   = blockIdx.x;
+        const int count   = min(SOL_BLOCK, max(0, length - block * SOL_BLOCK));
+        float sum_q = 0.f, sum_k = 0.f, sum_v = 0.f;
+        for (int i = 0; i < count; ++i) {
+            const int token    = block * SOL_BLOCK + i;
+            const size_t index = (size_t(blockIdx.y) * padded + token) * SOL_DIM + channel;
+            sum_q += __bfloat162float(q[index]);
+            sum_k += __bfloat162float(k[index]);
+            sum_v += __bfloat162float(v[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token]);
+        }
+        kc[(size_t(blockIdx.y) * padded_blocks + block) * SOL_DIM + channel] =
+            __float2bfloat16(count ? sum_k / count : 0.f);
+        vc[(size_t(blockIdx.y) * SOL_DIM + channel) * padded_blocks + block] = __float2bfloat16(sum_v);
+        if (count) {
+            qm[(size_t(blockIdx.y) * (padded / SOL_BLOCK) + block) * SOL_DIM + channel] = sum_q / count;
+        }
+    }
+
+    static __global__ void sol_statistics(const nv_bfloat16* kc, float* stats, int blocks, int padded_blocks) {
+        const int channel = threadIdx.x;
+        float sum = 0.f, sum_sq = 0.f;
+        for (int i = 0; i < blocks; ++i) {
+            const float value = __bfloat162float(kc[(size_t(blockIdx.x) * padded_blocks + i) * SOL_DIM + channel]);
+            sum += value;
+            sum_sq += value * value;
+        }
+        const float mean                                            = sum / blocks;
+        stats[size_t(blockIdx.x) * 2 * SOL_DIM + channel]           = mean;
+        stats[size_t(blockIdx.x) * 2 * SOL_DIM + SOL_DIM + channel] = fmaxf(sum_sq / blocks - mean * mean, 0.f);
+    }
+
+    static __global__ void sol_thresholds(const float* qm, const float* stats, float* threshold, int blocks, float scale, float tau) {
+        const int channel = threadIdx.x;
+        const float q     = qm[(size_t(blockIdx.y) * blocks + blockIdx.x) * SOL_DIM + channel];
+        float mean        = q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + channel];
+        float var         = q * q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + SOL_DIM + channel];
+        mean              = warp_reduce_sum(mean);
+        var               = warp_reduce_sum(var);
+        __shared__ float means[4], vars[4];
+        if (channel % 32 == 0) {
+            means[channel / 32] = mean;
+            vars[channel / 32]  = var;
+        }
+        __syncthreads();
+        if (channel == 0) {
+            mean = means[0] + means[1] + means[2] + means[3];
+            var  = vars[0] + vars[1] + vars[2] + vars[3];
+            threshold[size_t(blockIdx.y) * blocks + blockIdx.x] =
+                mean * scale + tau * sqrtf(fmaxf(var * scale * scale, 0.f) + 1.e-6f);
+        }
+    }
+
+    struct SolQTile {
+        uint32_t x[4];
+    };
+    struct SolKTile {
+        uint32_t x[2];
+    };
+    struct SolTile {
+        float x[4] = {};
+        static __device__ __forceinline__ int get_i(int i) {
+            return (i / 2) * 8 + threadIdx.x / 4;
+        }
+        static __device__ __forceinline__ int get_j(int i) {
+            return (threadIdx.x % 4) * 2 + i % 2;
+        }
+    };
+
+    static __device__ __forceinline__ void sol_load(SolQTile& tile, const nv_bfloat162* input, int stride) {
+        const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+        for (int i = 0; i < 4; ++i)
+            tile.x[i] = data[((i % 2) * 8 + threadIdx.x / 4) * stride + (i / 2) * 4 + threadIdx.x % 4];
+    }
+    static __device__ __forceinline__ void sol_load(SolKTile& tile, const nv_bfloat162* input, int stride) {
+        const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+        for (int i = 0; i < 2; ++i)
+            tile.x[i] = data[(threadIdx.x / 4) * stride + i * 4 + threadIdx.x % 4];
+    }
+    static __device__ __forceinline__ void sol_mma(SolTile& output, const SolQTile& q, const SolKTile& k) {
+#if __CUDA_ARCH__ >= 800
+        auto* dst = reinterpret_cast<uint32_t*>(output.x);
+        asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+            : "+r"(dst[0]), "+r"(dst[1]), "+r"(dst[2]), "+r"(dst[3])
+            : "r"(q.x[0]), "r"(q.x[1]), "r"(q.x[2]), "r"(q.x[3]), "r"(k.x[0]), "r"(k.x[1]));
+#else
+        asm("trap;");
+#endif
+    }
+
+    static __device__ __forceinline__ void sol_scores(const SolQTile (&q)[8], const nv_bfloat16* k, SolTile (&scores)[8], float scale) {
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+            scores[n] = SolTile{};
+#pragma unroll
+            for (int d = 0; d < 8; ++d) {
+                SolKTile kt;
+                sol_load(kt, reinterpret_cast<const nv_bfloat162*>(k + n * 8 * SOL_DIM + d * 16), SOL_DIM / 2);
+                sol_mma(scores[n], q[d], kt);
+            }
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                scores[n].x[i] *= scale;
+            }
+        }
+    }
+
+    template <bool PROXY>
+    static __device__ __forceinline__ void sol_update(SolTile (&scores)[8], SolTile (&output)[16], float (&row_max)[2], float (&row_sum)[2], nv_bfloat16* probabilities, const nv_bfloat16* values, int stride, int start, int length, const int* exact) {
+        float maximum[2] = {-INFINITY, -INFINITY};
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int column = n * 8 + SolTile::get_j(i);
+                const bool valid = PROXY ? ((start + column) * SOL_BLOCK < length && !exact[column])
+                                         : (start + column < length);
+                scores[n].x[i]   = valid ? scores[n].x[i] : -INFINITY;
+                maximum[i / 2]   = fmaxf(maximum[i / 2], scores[n].x[i]);
+            }
+        }
+        float alpha[2], sums[2] = {0.f, 0.f};
+#pragma unroll
+        for (int row = 0; row < 2; ++row) {
+            maximum[row]     = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 1));
+            maximum[row]     = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 2));
+            const float next = fmaxf(row_max[row], maximum[row]);
+            alpha[row]       = row_max[row] == next ? 1.f : exp2f(row_max[row] - next);
+            row_max[row]     = next;
+        }
+#pragma unroll
+        for (int n = 0; n < 8; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int column = n * 8 + SolTile::get_j(i);
+                const float p    = scores[n].x[i] == -INFINITY ? 0.f : exp2f(scores[n].x[i] - row_max[i / 2]);
+                const int count  = PROXY ? min(SOL_BLOCK, max(0, length - (start + column) * SOL_BLOCK)) : 1;
+                sums[i / 2] += p * count;
+                probabilities[SolTile::get_i(i) * SOL_BLOCK + column] = __float2bfloat16(p);
+            }
+        }
+#pragma unroll
+        for (int row = 0; row < 2; ++row) {
+            sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 1);
+            sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 2);
+            row_sum[row] = row_sum[row] * alpha[row] + sums[row];
+        }
+        __syncwarp();
+#pragma unroll
+        for (int n = 0; n < 16; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                output[n].x[i] *= alpha[i / 2];
+            }
+#pragma unroll
+            for (int d = 0; d < 4; ++d) {
+                SolQTile pt;
+                SolKTile vt;
+                sol_load(pt, reinterpret_cast<const nv_bfloat162*>(probabilities + d * 16), SOL_BLOCK / 2);
+                sol_load(vt, reinterpret_cast<const nv_bfloat162*>(values + n * 8 * stride + d * 16), stride / 2);
+                sol_mma(output[n], pt, vt);
+            }
+        }
+        __syncwarp();
+    }
+
+    static __global__ __launch_bounds__(128) void sol_forward(
+        const nv_bfloat16* q,
+        const nv_bfloat16* k,
+        const nv_bfloat16* v,
+        const nv_bfloat16* kc,
+        const nv_bfloat16* vc,
+        const float* thresholds,
+        float* dst,
+        int length,
+        int padded,
+        int padded_blocks,
+        int heads,
+        float scale) {
+#if __CUDA_ARCH__ >= 800
+        const int warp    = threadIdx.y;
+        const int lane    = threadIdx.x;
+        const int tid     = warp * 32 + lane;
+        const int blocks  = padded / SOL_BLOCK;
+        const int q_start = blockIdx.x * SOL_BLOCK;
+        const int q_len   = min(SOL_BLOCK, length - q_start);
+        const size_t base = size_t(blockIdx.y) * padded * SOL_DIM;
+        SolQTile qt[8];
+#pragma unroll
+        for (int d = 0; d < 8; ++d) {
+            sol_load(qt[d], reinterpret_cast<const nv_bfloat162*>(q + base + (q_start + warp * 16) * SOL_DIM + d * 16), SOL_DIM / 2);
+        }
+        SolTile output[16];
+        float row_max[2] = {-INFINITY, -INFINITY};
+        float row_sum[2] = {0.f, 0.f};
+        __shared__ nv_bfloat16 probabilities[SOL_BLOCK * SOL_BLOCK];
+        __shared__ float routes[4][SOL_BLOCK];
+        __shared__ int exact[SOL_BLOCK];
+        const float threshold = thresholds[size_t(blockIdx.y) * blocks + blockIdx.x];
+        for (int group = 0; group < blocks; group += SOL_BLOCK) {
+            SolTile scores[8];
+            sol_scores(qt, kc + (size_t(blockIdx.y) * padded_blocks + group) * SOL_DIM, scores, scale);
+#pragma unroll
+            for (int n = 0; n < 8; ++n) {
+#pragma unroll
+                for (int column = 0; column < 2; ++column) {
+                    float sum = scores[n].x[column] + scores[n].x[column + 2];
+                    sum += __shfl_xor_sync(0xffffffff, sum, 4);
+                    sum += __shfl_xor_sync(0xffffffff, sum, 8);
+                    sum += __shfl_xor_sync(0xffffffff, sum, 16);
+                    if (lane < 4) {
+                        routes[warp][n * 8 + 2 * lane + column] = sum;
+                    }
+                }
+            }
+            __syncthreads();
+            if (tid < SOL_BLOCK) {
+                const int block   = group + tid;
+                const float score = (routes[0][tid] + routes[1][tid] + routes[2][tid] + routes[3][tid]) / q_len;
+                exact[tid]        = block < blocks && (score > threshold || abs(int(blockIdx.x) - block) <= 1);
+            }
+            __syncthreads();
+            sol_update<true>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+                             vc + size_t(blockIdx.y) * padded_blocks * SOL_DIM + group,
+                             padded_blocks, group, length, exact);
+            for (int block = 0; block < SOL_BLOCK && group + block < blocks; ++block) {
+                if (!exact[block]) {
+                    continue;
+                }
+                const int start = (group + block) * SOL_BLOCK;
+                sol_scores(qt, k + base + start * SOL_DIM, scores, scale);
+                sol_update<false>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+                                  v + base + start, padded, start, length, nullptr);
+            }
+            __syncthreads();
+        }
+        const int head  = blockIdx.y % heads;
+        const int batch = blockIdx.y / heads;
+#pragma unroll
+        for (int n = 0; n < 16; ++n) {
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const int token = q_start + warp * 16 + SolTile::get_i(i);
+                if (token < length) {
+                    const int channel = n * 8 + SolTile::get_j(i);
+                    dst[((size_t(batch) * length + token) * heads + head) * SOL_DIM + channel] =
+                        __bfloat162float(__float2bfloat16(output[n].x[i] / row_sum[i / 2]));
+                }
+            }
+        }
+#else
+        asm("trap;");
+#endif
+    }
+
+    static bool sol_supports(ggml_backend_dev_t, const ggml_tensor* op, void*) {
+        SDSolAttentionParams params;
+        if (!sd_ggml_sol_attn_describe(op, &params))
+            return false;
+        const auto* q = op->src[0];
+        if (q->ne[2] * q->ne[3] > 65535 || GGML_PAD(q->ne[1], SOL_BLOCK) > INT_MAX / SOL_DIM)
+            return false;
+        for (const auto* tensor : {op->src[0], op->src[1], op->src[2]})
+            if (reinterpret_cast<uintptr_t>(tensor->data) & 3)
+                return false;
+        return !(reinterpret_cast<uintptr_t>(op->data) & 3);
+    }
+
+    static ggml_status sol_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+        *workspace = SolWorkspace(nodes[0]).size;
+        return GGML_STATUS_SUCCESS;
+    }
+    static void sol_destroy(void*, void*) {}
+
+    static ggml_status sol_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+        ggml_ext_gpu_context context;
+        if (!cuda_context(execution->backend, &context))
+            return GGML_STATUS_FAILED;
+        auto* dst = nodes[0];
+        SDSolAttentionParams params;
+        if (!sd_ggml_sol_attn_describe(dst, &params))
+            return GGML_STATUS_FAILED;
+        const auto* q = dst->src[0];
+        const SolWorkspace work(dst);
+        char* data       = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+        auto* qb         = reinterpret_cast<nv_bfloat16*>(data + work.q);
+        auto* kb         = reinterpret_cast<nv_bfloat16*>(data + work.k);
+        auto* vb         = reinterpret_cast<nv_bfloat16*>(data + work.v);
+        auto* kc         = reinterpret_cast<nv_bfloat16*>(data + work.kc);
+        auto* vc         = reinterpret_cast<nv_bfloat16*>(data + work.vc);
+        auto* qm         = reinterpret_cast<float*>(data + work.qm);
+        auto* stats      = reinterpret_cast<float*>(data + work.stats);
+        auto* threshold  = reinterpret_cast<float*>(data + work.threshold);
+        const int length = q->ne[1], heads = q->ne[2], batch_heads = heads * q->ne[3];
+        const int blocks  = work.tokens / SOL_BLOCK;
+        const float scale = params.scale * 1.4426950408889634f;
+        const auto stream = static_cast<cudaStream_t>(context.stream);
+        sol_pack<<<dim3((work.tokens * SOL_DIM + 255) / 256, batch_heads), 256, 0, stream>>>(
+            static_cast<const float*>(q->data), static_cast<const float*>(dst->src[1]->data), static_cast<const float*>(dst->src[2]->data),
+            qb, kb, vb, length, work.tokens);
+        sol_summaries<<<dim3(work.blocks, batch_heads), SOL_DIM, 0, stream>>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks);
+        sol_statistics<<<batch_heads, SOL_DIM, 0, stream>>>(kc, stats, blocks, work.blocks);
+        sol_thresholds<<<dim3(blocks, batch_heads), SOL_DIM, 0, stream>>>(qm, stats, threshold, blocks, scale, params.tau);
+        sol_forward<<<dim3(blocks, batch_heads), dim3(32, 4), 0, stream>>>(qb, kb, vb, kc, vc, threshold,
+                                                                           static_cast<float*>(dst->data), length, work.tokens, work.blocks, heads, scale);
+        return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+    }
+
+    void register_sol_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 || properties.major < 8 ||
+            cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
+            return;
+        const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
+        if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess ||
+            std::min(attributes.binaryVersion, attributes.ptxVersion) < 80)
+            return;
+        ggml_ext_kernel_desc desc = {};
+        desc.abi_version          = GGML_EXT_ABI_VERSION;
+        desc.name                 = "sd.sol_attn.cuda.v1";
+        desc.op                   = GGML_OP_CUSTOM;
+        desc.custom_op            = sd_ggml_sol_attn_op();
+        desc.priority             = 100;
+        desc.supports             = sol_supports;
+        desc.prepare              = sol_prepare;
+        desc.execute              = sol_execute;
+        desc.destroy              = sol_destroy;
+        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_sol_attention(ggml_backend_dev_t) {}
+}
+#endif

src/core/ggml_compat.cpp

@@ -0,0 +1,185 @@
+#include "core/ggml_compat.h"
+#include "core/ggml_extend_backend.h"
+
+#ifdef SD_GGML_GRAPH_COMPAT
+// Upstream ggml still requires its private graph layout for these operations.
+#include "ggml-impl.h"
+
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+int ggml_graph_n_leafs(const ggml_cgraph* graph) {
+    return graph->n_leafs;
+}
+
+ggml_tensor* ggml_graph_leaf(const ggml_cgraph* graph, int index) {
+    if (index < 0) {
+        index += graph->n_leafs;
+    }
+    GGML_ASSERT(index >= 0 && index < graph->n_leafs);
+    return graph->leafs[index];
+}
+
+void ggml_graph_add_leaf(ggml_cgraph* graph, ggml_tensor* tensor) {
+    GGML_ASSERT(tensor && graph->n_leafs < graph->size);
+    graph->leafs[graph->n_leafs++] = tensor;
+}
+
+int32_t ggml_graph_use_count(const ggml_cgraph* graph, const ggml_tensor* tensor) {
+    if (!graph || !tensor || !graph->visited_hash_set.size) {
+        return 0;
+    }
+    const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor);
+    return index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index)
+               ? graph->use_counts[index]
+               : 0;
+}
+
+bool ggml_graph_cpy_mapped(const ggml_cgraph* src, ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data) {
+    if (!src || !dst || src == dst || !map || dst->grads || dst->grad_accs ||
+        src->n_nodes > dst->size || src->n_leafs > dst->size) {
+        return false;
+    }
+    std::unordered_map<const ggml_tensor*, ggml_tensor*> mapping;
+    std::unordered_set<ggml_tensor*> mapped;
+    std::vector<const ggml_tensor*> sources;
+    bool valid = true;
+    auto remap = [&](const ggml_tensor* source) {
+        const auto found = mapping.find(source);
+        if (found != mapping.end()) {
+            return found->second;
+        }
+        auto* tensor = map(source, user_data);
+        if (tensor && !mapped.insert(tensor).second) {
+            valid = false;
+        }
+        mapping.emplace(source, tensor);
+        sources.push_back(source);
+        return tensor;
+    };
+    std::vector<ggml_tensor*> nodes(src->n_nodes), leafs(src->n_leafs);
+    for (int i = 0; i < src->n_nodes; ++i) {
+        if (!(nodes[i] = remap(src->nodes[i])) || !valid) {
+            return false;
+        }
+    }
+    for (int i = 0; i < src->n_leafs; ++i) {
+        if (!(leafs[i] = remap(src->leafs[i])) || !valid) {
+            return false;
+        }
+    }
+    std::vector<std::pair<ggml_tensor*, int32_t>> counts;
+    for (size_t i = 0; i < sources.size(); ++i) {
+        const auto* source = sources[i];
+        if (source->view_src) {
+            remap(source->view_src);
+        }
+        for (const auto* input : source->src) {
+            if (input) {
+                remap(input);
+            }
+        }
+        if (!valid) {
+            return false;
+        }
+        auto* tensor = mapping.at(source);
+        if (!tensor || !src->visited_hash_set.size) {
+            continue;
+        }
+        const size_t index = ggml_hash_find(&src->visited_hash_set, source);
+        if (index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, index)) {
+            continue;
+        }
+        if (counts.size() >= dst->visited_hash_set.size) {
+            return false;
+        }
+        counts.emplace_back(tensor, src->use_counts[index]);
+    }
+    ggml_graph_clear(dst);
+    dst->order = src->order;
+    dst->uid   = 0;
+    for (auto* tensor : nodes) {
+        ggml_graph_add_node(dst, tensor);
+    }
+    for (auto* tensor : leafs) {
+        ggml_graph_add_leaf(dst, tensor);
+    }
+    for (const auto& entry : counts) {
+        const size_t index     = ggml_hash_insert(&dst->visited_hash_set, entry.first);
+        dst->use_counts[index] = entry.second;
+    }
+    return true;
+}
+
+bool ggml_graph_cpy_use_counts(const ggml_cgraph* src, ggml_cgraph* dst) {
+    if (!src || !dst || !src->visited_hash_set.size || !dst->visited_hash_set.size) {
+        return false;
+    }
+    std::unordered_set<ggml_tensor*> represented;
+    for (int i = 0; i < dst->n_nodes; ++i) {
+        represented.insert(dst->nodes[i]);
+    }
+    for (int i = 0; i < dst->n_leafs; ++i) {
+        represented.insert(dst->leafs[i]);
+    }
+    size_t occupied = 0;
+    for (size_t i = 0; i < dst->visited_hash_set.size; ++i) {
+        occupied += ggml_bitset_get(dst->visited_hash_set.used, i);
+    }
+    std::vector<std::pair<ggml_tensor*, int32_t>> counts;
+    for (auto* tensor : represented) {
+        if (!tensor) {
+            return false;
+        }
+        const size_t source_index = ggml_hash_find(&src->visited_hash_set, tensor);
+        if (source_index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, source_index)) {
+            continue;
+        }
+        const size_t index = ggml_hash_find(&dst->visited_hash_set, tensor);
+        if (index == GGML_HASHSET_FULL || !ggml_bitset_get(dst->visited_hash_set.used, index)) {
+            if (++occupied > dst->visited_hash_set.size) {
+                return false;
+            }
+        }
+        counts.emplace_back(tensor, src->use_counts[source_index]);
+    }
+    for (const auto& entry : counts) {
+        const size_t index     = ggml_hash_find_or_insert(&dst->visited_hash_set, entry.first);
+        dst->use_counts[index] = entry.second;
+    }
+    return true;
+}
+
+ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend, ggml_cgraph* graph, int first, int last) {
+    if (!backend || !graph || first < 0 || last < first || last > graph->n_nodes) {
+        return GGML_STATUS_FAILED;
+    }
+    if (first == last) {
+        return GGML_STATUS_SUCCESS;
+    }
+    auto view = ggml_graph_view(graph, first, last);
+    return ggml_backend_graph_compute_async(backend, &view);
+}
+#endif
+
+#ifdef SD_USE_UPSTREAM_GGML
+#include "ggml-backend-impl.h"
+#endif
+
+ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void* ptr, size_t size, size_t max_tensor_size) {
+    auto buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
+#ifdef SD_USE_UPSTREAM_GGML
+    if (buffer && !buffer->context) {
+        auto reg = ggml_backend_dev_backend_reg(device);
+        if (reg && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
+            // Older Metal backends wrap failed mappings and cannot free their null context.
+            buffer->iface.free_buffer = nullptr;
+            ggml_backend_buffer_free(buffer);
+            return nullptr;
+        }
+    }
+#endif
+    return buffer;
+}

src/core/ggml_compat.h

@@ -0,0 +1,43 @@
+#ifndef __SD_CORE_GGML_COMPAT_H__
+#define __SD_CORE_GGML_COMPAT_H__
+
+#include "ggml-backend.h"
+
+#if __has_include("ggml-graph.h")
+#include "ggml-graph.h"
+#elif defined(SD_USE_UPSTREAM_GGML)
+#define SD_GGML_GRAPH_COMPAT
+
+int ggml_graph_n_leafs(const struct ggml_cgraph* graph);
+struct ggml_tensor* ggml_graph_leaf(const struct ggml_cgraph* graph, int index);
+// Like ggml_graph_add_node, appends metadata without traversing tensor dependencies.
+void ggml_graph_add_leaf(struct ggml_cgraph* graph, struct ggml_tensor* tensor);
+// Returns zero when the tensor has no recorded use count. Includes uses outside graph views.
+int32_t ggml_graph_use_count(const struct ggml_cgraph* graph, const struct ggml_tensor* tensor);
+
+typedef struct ggml_tensor* (*ggml_graph_tensor_map_fn)(const struct ggml_tensor* tensor, void* user_data);
+
+// Copies forward node/leaf order and recorded use counts, including external consumers.
+// The callback must be injective and may return NULL for dependencies outside the
+// node/leaf lists. It is called once per tensor in their dependency closure.
+// Tensor storage and edges are mapped by the caller.
+// The destination must have no gradients and enough capacity. Gradients are not copied.
+// Invalid mappings leave the destination unchanged. src and dst must be distinct.
+bool ggml_graph_cpy_mapped(const struct ggml_cgraph* src, struct ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data);
+
+// Copies recorded counts for tensors in dst's node/leaf lists. Other metadata is unchanged.
+// This preserves external consumers when constructing an inference graph segment.
+bool ggml_graph_cpy_use_counts(const struct ggml_cgraph* src, struct ggml_cgraph* dst);
+
+// Executes nodes [first, last) with the original graph's use counts. No graph ownership
+// is transferred. Tensor storage must already be allocated; invalid bounds fail.
+enum ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend,
+                                                        struct ggml_cgraph* graph,
+                                                        int first,
+                                                        int last);
+
+#else
+#error "Patched ggml requires ggml-graph.h; select SD_USE_UPSTREAM_GGML for upstream compatibility"
+#endif
+
+#endif

src/core/ggml_extend.h

@@ -4,6 +4,7 @@
 #include <cstdint>
 #include <vector>
 
+#include "core/ggml_extensions.h"
 #include "ggml-backend.h"
 #include "ggml.h"
 

src/core/ggml_extend_backend.cpp

@@ -1,4 +1,5 @@
 #include "core/ggml_extend_backend.h"
+#include "core/ggml_extensions.h"
 
 #include <algorithm>
 #include <cctype>
@@ -17,9 +18,8 @@
 #include <hip/hip_runtime_api.h>
 #endif
 
+#include "core/ggml_compat.h"
 #include "core/util.h"
-#include "ggml-backend-impl.h"
-#include "ggml-impl.h"
 #include "stable-diffusion.h"
 
 static std::string trim_copy(const std::string& value) {
@@ -450,24 +450,6 @@ bool sd_backend_is_cpu(ggml_backend_t backend) {
     return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU;
 }
 
-ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device,
-                                                          void* ptr,
-                                                          size_t size,
-                                                          size_t max_tensor_size) {
-    ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size);
-    if (buffer != nullptr && buffer->context == nullptr) {
-        ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device);
-        if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) {
-            // Metal can wrap a failed mapping in a non-null buffer. Its free callback also
-            // dereferences the missing context, so only release the outer buffer.
-            buffer->iface.free_buffer = nullptr;
-            ggml_backend_buffer_free(buffer);
-            return nullptr;
-        }
-    }
-    return buffer;
-}
-
 static int cuda_compiled_arch(ggml_backend_t backend) {
 #ifdef SD_USE_CUDA
     if (!sd_backend_is(backend, "CUDA")) {
@@ -596,7 +578,9 @@ bool sd_backend_supports_conv_implicit_gemm(ggml_backend_t backend, bool conv3d)
 
 ggml_backend_t sd_backend_cpu_init() {
     ggml_backend_load_all_once();
-    return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+    auto backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+    sd_ggml_register_extensions(backend);
+    return backend;
 }
 
 bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) {
@@ -641,23 +625,6 @@ void sd_backend_clear_buffer_cache(ggml_backend_t backend) {
     }
 }
 
-static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) {
-    ggml_cgraph cgraph = {
-        /*.size             =*/0,
-        /*.n_nodes          =*/i1 - i0,
-        /*.n_leafs          =*/0,
-        /*.nodes            =*/cgraph0->nodes + i0,
-        /*.grads            =*/nullptr,
-        /*.grad_accs        =*/nullptr,
-        /*.leafs            =*/nullptr,
-        /*.use_counts       =*/cgraph0->use_counts,
-        /*.visited_hash_set =*/cgraph0->visited_hash_set,
-        /*.order            =*/cgraph0->order,
-        /*.uid              =*/0,
-    };
-    return cgraph;
-}
-
 ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
                                                         ggml_cgraph* gf,
                                                         sd_graph_eval_callback_t callback_eval,
@@ -680,8 +647,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend,
             need = callback_eval(t, true, callback_eval_user_data);
         }
 
-        ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1);
-        status         = ggml_backend_graph_compute_async(backend, &gv);
+        status = ggml_backend_graph_compute_async_range(backend, gf, j0, j1 + 1);
         if (status != GGML_STATUS_SUCCESS) {
             break;
         }
@@ -1207,6 +1173,7 @@ ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) {
     }
 
     SDBackendHandle handle(backend);
+    sd_ggml_register_extensions(backend);
     backends_.emplace(actual_key, std::move(handle));
     return backend;
 }

src/core/ggml_extensions.cpp

@@ -0,0 +1,364 @@
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include "core/cpu/registration.h"
+#include "ggml-ext.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstring>
+#include <mutex>
+#include <unordered_set>
+#include <utility>
+
+#ifdef SD_BACKEND_EXTENSIONS
+#include "sd_backend_extensions.h"
+#endif
+
+static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) {
+    if (count != 2) {
+        return false;
+    }
+    const auto* x  = inputs[0];
+    const auto* pe = inputs[1];
+    if (x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || x->ne[0] % 2 != 0 ||
+        pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
+        pe->ne[3] != x->ne[2]) {
+        return false;
+    }
+    out->type  = GGML_TYPE_F32;
+    out->ne[0] = x->ne[0];
+    out->ne[1] = x->ne[2];
+    out->ne[2] = x->ne[1] * x->ne[3];
+    out->ne[3] = 1;
+    return true;
+}
+
+ggml_ext_op_t sd_ggml_rope_op() {
+    static const ggml_ext_op_t op = [] {
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.rope_apply.v1", sizeof(int32_t), rope_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return op;
+}
+
+bool sd_ggml_rope_describe(const ggml_tensor* node) {
+    if (node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node) || !node->src[0] || !node->src[1]) {
+        return false;
+    }
+    ggml_tensor* inputs[] = {node->src[0], node->src[1]};
+    ggml_ext_tensor_desc output;
+    return rope_infer(inputs, 2, nullptr, &output) &&
+           std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0;
+}
+
+struct I8Attributes {
+    int64_t shape[3];
+    int32_t group_size;
+    int32_t flags;
+};
+enum I8Flags { I8_PACKED       = 1,
+               I8_INPUT_SCALE  = 2,
+               I8_WEIGHT_SCALE = 4,
+               I8_BIAS         = 8 };
+
+static bool valid_group(int group) {
+    return group == 1 || group == 4 || group == 16 || group == 64 || group == 256;
+}
+static constexpr int64_t i8_size_limit = int64_t(std::min(uint64_t(INT64_MAX), uint64_t(SIZE_MAX)));
+static bool packed_height(int64_t k, int64_t rows, int64_t* height) {
+    if (k <= 0 || rows <= 0 || rows > i8_size_limit / 4)
+        return false;
+    const int64_t padded = (rows + 3) / 4 * 4;
+    const int64_t scales = (rows * 4 - 1) / k + 1;
+    if (padded > i8_size_limit / k || scales > i8_size_limit / k - padded)
+        return false;
+    *height = padded + scales;
+    return true;
+}
+static bool scale_matches(const ggml_tensor* scale, int64_t count) {
+    return !scale || (scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_nelements(scale) == count);
+}
+static bool i8_shape(const SDI8Operation& op, const int64_t* shape) {
+    const auto* w = op.weight;
+    const auto* x = op.input;
+    if (!w || !x || w->type != GGML_TYPE_I8 || w->ne[2] != 1 || w->ne[3] != 1 ||
+        !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || x->ne[0] != w->ne[0] ||
+        (op.group_size != 0 && (!valid_group(op.group_size) || w->ne[0] % op.group_size)))
+        return false;
+    int64_t rows = 1;
+    for (int i = 0; i < 3; ++i) {
+        if (shape[i] <= 0 || shape[i] > (INT64_MAX - 3) / rows)
+            return false;
+        rows *= shape[i];
+    }
+    const int64_t k = w->ne[0], n = w->ne[1];
+    int64_t height;
+    if (n <= 0 || n > (i8_size_limit - 4) / 4 || k <= 0 || k > i8_size_limit - 4 - n * 4 ||
+        rows > (i8_size_limit - 3) / (k + 4 + n * 4) || !packed_height(k, rows, &height))
+        return false;
+    if (op.packed) {
+        if (x->type != GGML_TYPE_I8 || !valid_group(op.group_size) || op.input_scale ||
+            x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1)
+            return false;
+    } else {
+        if (std::memcmp(shape, x->ne + 1, 3 * sizeof(int64_t)))
+            return false;
+        if (x->type == GGML_TYPE_I8) {
+            if (!op.input_scale || op.group_size != 0)
+                return false;
+        } else if (x->type != GGML_TYPE_F32 || op.input_scale)
+            return false;
+    }
+    return scale_matches(op.input_scale, rows) && scale_matches(op.weight_scale, w->ne[1]) && scale_matches(op.bias, w->ne[1]);
+}
+static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
+    I8Attributes attr;
+    std::memcpy(&attr, attributes, sizeof(attr));
+    if (count < 2 || (attr.flags & ~15))
+        return false;
+    SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & I8_PACKED) != 0};
+    int index        = 2;
+    for (auto entry : {std::make_pair(I8_INPUT_SCALE, &op.input_scale), std::make_pair(I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(I8_BIAS, &op.bias)}) {
+        if (attr.flags & entry.first) {
+            if (index >= count)
+                return false;
+            *entry.second = inputs[index++];
+        }
+    }
+    if (index != count || !i8_shape(op, attr.shape))
+        return false;
+    output->type  = GGML_TYPE_F32;
+    output->ne[0] = op.weight->ne[1];
+    std::memcpy(output->ne + 1, attr.shape, sizeof(attr.shape));
+    return true;
+}
+static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) {
+    int32_t group;
+    std::memcpy(&group, attributes, sizeof(group));
+    if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !valid_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group)
+        return false;
+    const int64_t k = inputs[0]->ne[0], rows = ggml_nrows(inputs[0]);
+    int64_t height;
+    if (!packed_height(k, rows, &height))
+        return false;
+    output->type  = GGML_TYPE_I8;
+    output->ne[0] = k;
+    output->ne[1] = height;
+    output->ne[2] = output->ne[3] = 1;
+    return true;
+}
+ggml_ext_op_t sd_ggml_i8_op() {
+    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(I8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
+    return op;
+}
+ggml_ext_op_t sd_ggml_convrot_op() {
+    static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.quantize_i8_convrot.v1", sizeof(int32_t), convrot_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }();
+    return op;
+}
+bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
+    if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node))
+        return false;
+    *op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
+    if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
+        I8Attributes attr;
+        std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
+        if (attr.flags & ~15 || std::memcmp(attr.shape, node->ne + 1, sizeof(attr.shape)))
+            return false;
+        op->group_size = attr.group_size;
+        op->packed     = (attr.flags & I8_PACKED) != 0;
+        int index      = 2;
+        if (attr.flags & I8_INPUT_SCALE)
+            op->input_scale = node->src[index++];
+        if (attr.flags & I8_WEIGHT_SCALE)
+            op->weight_scale = node->src[index++];
+        if (attr.flags & I8_BIAS)
+            op->bias = node->src[index++];
+        if (((attr.flags & I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & I8_BIAS) && !op->bias))
+            return false;
+    } else
+        return false;
+    return i8_shape(*op, node->ne + 1) && node->ne[0] == op->weight->ne[1];
+}
+bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
+    if (!node || !node->src[0] || !ggml_is_contiguous(node))
+        return false;
+    if (ggml_ext_op_get(node) == sd_ggml_convrot_op())
+        std::memcpy(group, ggml_ext_op_attributes(node), sizeof(int32_t));
+    else
+        return false;
+    ggml_ext_tensor_desc output;
+    return convrot_infer(node->src, 1, group, &output) && output.type == node->type && !std::memcmp(output.ne, node->ne, sizeof(output.ne));
+}
+ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group, ggml_tensor* xs) {
+    if (!w || !x)
+        return nullptr;
+    const auto* logical = x;
+    if (x->type == GGML_TYPE_I8 && !xs) {
+        int packed_group;
+        if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group)
+            return nullptr;
+        logical = x->src[0];
+    }
+    I8Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0};
+    ggml_tensor* inputs[5] = {w, x};
+    int count              = 2;
+    if (xs) {
+        attr.flags |= I8_INPUT_SCALE;
+        inputs[count++] = xs;
+    }
+    if (ws) {
+        attr.flags |= I8_WEIGHT_SCALE;
+        inputs[count++] = ws;
+    }
+    if (bias) {
+        attr.flags |= I8_BIAS;
+        inputs[count++] = bias;
+    }
+    return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr);
+}
+ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
+    int32_t attr = group;
+    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);
+}
+
+static bool sol_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+    if (count != 3 || !attributes)
+        return false;
+    SDSolAttentionParams params;
+    std::memcpy(&params, attributes, sizeof(params));
+    if (!std::isfinite(params.scale) || params.scale <= 0.f || !std::isfinite(params.tau))
+        return false;
+    const auto* q = inputs[0];
+    if (!q || q->ne[0] != 128)
+        return false;
+    for (int i = 0; i < count; ++i) {
+        const auto* t = inputs[i];
+        if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || !ggml_are_same_shape(q, t))
+            return false;
+    }
+    int64_t elements = 1;
+    for (int d = 0; d < GGML_MAX_DIMS; ++d) {
+        if (q->ne[d] <= 0 || q->ne[d] > INT64_MAX / int64_t(sizeof(float)) / elements)
+            return false;
+        elements *= q->ne[d];
+    }
+    *out = {GGML_TYPE_F32, {q->ne[0], q->ne[2], q->ne[1], q->ne[3]}};
+    return true;
+}
+
+ggml_ext_op_t sd_ggml_sol_attn_op() {
+    static const auto op = [] {
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.sol_attn.v1", sizeof(SDSolAttentionParams), sol_attn_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return op;
+}
+
+bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params) {
+    if (!node || ggml_ext_op_get(node) != sd_ggml_sol_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 (!sol_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_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau) {
+    ggml_tensor* inputs[]             = {q, k, v};
+    const SDSolAttentionParams params = {scale, tau};
+    return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, &params);
+}
+
+void sd_ggml_register_extensions(ggml_backend_t backend) {
+    if (!backend) {
+        return;
+    }
+    auto device    = ggml_backend_get_device(backend);
+    auto reg       = ggml_backend_dev_backend_reg(device);
+    const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr;
+    static std::mutex mutex;
+    static std::unordered_set<ggml_backend_dev_t> registered;
+    std::lock_guard<std::mutex> lock(mutex);
+    if (!registered.insert(device).second) {
+        return;
+    }
+    if (!cpu) {
+#ifdef SD_BACKEND_EXTENSIONS
+        for (auto register_backend : sd_backend_extensions)
+            register_backend(device);
+#endif
+        return;
+    }
+    sd_cpu::register_extensions(device);
+}
+
+ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) {
+    sd_ggml_register_extensions(backend);
+    ggml_tensor* inputs[] = {x, pe};
+    int32_t attributes    = interleaved ? 1 : 0;
+    return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes);
+}
+
+#else
+void sd_ggml_register_extensions(ggml_backend_t) {}
+#endif

src/core/ggml_extensions.h

@@ -0,0 +1,66 @@
+#ifndef __SD_CORE_GGML_EXTENSIONS_H__
+#define __SD_CORE_GGML_EXTENSIONS_H__
+
+#include "ggml-backend.h"
+
+void sd_ggml_register_extensions(ggml_backend_t backend);
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include "ggml-ext.h"
+
+ggml_ext_op_t sd_ggml_rope_op();
+bool sd_ggml_rope_describe(const ggml_tensor* node);
+// Apply precomputed rotary matrices, with shapes in ne[] order:
+// a: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens].
+// pe stores [[cos, -sin], [sin, cos]] for each pair and token, shared across heads/batches.
+// Returns contiguous F32 [head_dim, tokens, heads*batch]. Input strides are preserved.
+// Interleaved pairs are (2*i, 2*i+1); otherwise pairs are (i, i+rot_dim/2).
+// Channels from rot_dim to head_dim are copied unchanged.
+ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved);
+
+struct SDI8Operation {
+    const ggml_tensor* weight;
+    const ggml_tensor* input;
+    const ggml_tensor* input_scale;
+    const ggml_tensor* weight_scale;
+    const ggml_tensor* bias;
+    int group_size;
+    bool packed;
+};
+
+ggml_ext_op_t sd_ggml_i8_op();
+ggml_ext_op_t sd_ggml_convrot_op();
+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);
+
+struct SDSolAttentionParams {
+    float scale;
+    float tau;
+};
+
+ggml_ext_op_t sd_ggml_sol_attn_op();
+bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params);
+// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch].
+ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau);
+#endif
+
+#endif

src/core/ggml_graph_cut.cpp

@@ -16,7 +16,7 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 
-#include "ggml-impl.h"
+#include "core/ggml_compat.h"
 
 namespace sd::ggml_graph_cut {
 
@@ -35,8 +35,8 @@ namespace sd::ggml_graph_cut {
     static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) {
         GGML_ASSERT(gf != nullptr);
         GGML_ASSERT(tensor != nullptr);
-        for (int i = 0; i < gf->n_leafs; ++i) {
-            if (gf->leafs[i] == tensor) {
+        for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) {
+            if (ggml_graph_leaf(gf, i) == tensor) {
                 return i;
             }
         }
@@ -389,7 +389,7 @@ namespace sd::ggml_graph_cut {
                         return fail(segment_label + " references an unavailable cut node " +
                                     std::to_string(input.node_index));
                     }
-                } else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) {
+                } else if (input.leaf_index < 0 || input.leaf_index >= ggml_graph_n_leafs(gf)) {
                     return fail(segment_label + " references an invalid leaf " +
                                 std::to_string(input.leaf_index));
                 }
@@ -444,15 +444,15 @@ namespace sd::ggml_graph_cut {
 
     int leaf_count(ggml_cgraph* gf) {
         GGML_ASSERT(gf != nullptr);
-        return gf->n_leafs;
+        return ggml_graph_n_leafs(gf);
     }
 
     ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) {
         GGML_ASSERT(gf != nullptr);
-        if (leaf_index < 0 || leaf_index >= gf->n_leafs) {
+        if (leaf_index < 0 || leaf_index >= ggml_graph_n_leafs(gf)) {
             return nullptr;
         }
-        return gf->leafs[leaf_index];
+        return ggml_graph_leaf(gf, leaf_index);
     }
 
     ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) {
@@ -548,14 +548,14 @@ namespace sd::ggml_graph_cut {
     };
 
     std::vector<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
-        const size_t graph_size = static_cast<size_t>(graph->n_leafs) + graph->n_nodes;
+        const size_t graph_size = static_cast<size_t>(ggml_graph_n_leafs(graph)) + ggml_graph_n_nodes(graph);
         GraphLayoutTensors layout_tensors(graph_size);
         const auto& tensors = layout_tensors.tensors;
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            layout_tensors.add(graph->leafs[i]);
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            layout_tensors.add(ggml_graph_leaf(graph, i));
         }
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            layout_tensors.add(graph->nodes[i]);
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            layout_tensors.add(ggml_graph_node(graph, i));
         }
         for (size_t i = 0; i < tensors.size(); ++i) {
             layout_tensors.add(tensors[i]->view_src);
@@ -567,13 +567,13 @@ namespace sd::ggml_graph_cut {
         const size_t tensor_fields = 5 + 2 * GGML_MAX_DIMS + GGML_MAX_SRC +
                                      GGML_MAX_OP_PARAMS / sizeof(int32_t) + (include_bindings ? 2 : 0);
         signature.reserve(2 + graph_size + tensors.size() * tensor_fields);
-        signature.push_back(graph->n_nodes);
-        signature.push_back(graph->n_leafs);
-        for (int i = 0; i < graph->n_leafs; ++i) {
-            signature.push_back(layout_tensors.index(graph->leafs[i]));
+        signature.push_back(ggml_graph_n_nodes(graph));
+        signature.push_back(ggml_graph_n_leafs(graph));
+        for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) {
+            signature.push_back(layout_tensors.index(ggml_graph_leaf(graph, i)));
         }
-        for (int i = 0; i < graph->n_nodes; ++i) {
-            signature.push_back(layout_tensors.index(graph->nodes[i]));
+        for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+            signature.push_back(layout_tensors.index(ggml_graph_node(graph, i)));
         }
         for (auto tensor : tensors) {
             signature.push_back(tensor->op);
@@ -606,18 +606,18 @@ namespace sd::ggml_graph_cut {
                                    const Plan& plan,
                                    const std::vector<uint64_t>& layout) {
         GGML_ASSERT(gf != nullptr);
-        if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
+        if (plan.leaf_names.size() != static_cast<size_t>(ggml_graph_n_leafs(gf)) ||
             plan.layout != layout) {
             return false;
         }
-        for (int i = 0; i < gf->n_leafs; ++i) {
-            if (plan.leaf_names[i] != gf->leafs[i]->name) {
+        for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) {
+            if (plan.leaf_names[i] != ggml_graph_leaf(gf, i)->name) {
                 return false;
             }
         }
         size_t cut_index = 0;
-        for (int i = 0; i < gf->n_nodes; ++i) {
-            auto node = gf->nodes[i];
+        for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
+            auto node = ggml_graph_node(gf, i);
             if (is_graph_cut_tensor(node)) {
                 if (cut_index >= plan.cut_markers.size() ||
                     plan.cut_markers[cut_index].first != i || plan.cut_markers[cut_index].second != node->name) {
@@ -654,7 +654,7 @@ namespace sd::ggml_graph_cut {
             }
             return ggml_graph_node(gf, input_ref.node_index);
         }
-        if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) {
+        if (input_ref.leaf_index < 0 || input_ref.leaf_index >= ggml_graph_n_leafs(gf)) {
             return nullptr;
         }
         return leaf_tensor(gf, input_ref.leaf_index);
@@ -735,8 +735,7 @@ namespace sd::ggml_graph_cut {
         GGML_ASSERT(segment_graph != nullptr);
 
         for (ggml_tensor* leaf : leaves) {
-            GGML_ASSERT(segment_graph->n_leafs < segment_graph->size);
-            segment_graph->leafs[segment_graph->n_leafs++] = leaf;
+            ggml_graph_add_leaf(segment_graph, leaf);
         }
 
         for (int output_node_index : segment.output_node_indices) {
@@ -753,14 +752,7 @@ namespace sd::ggml_graph_cut {
         for (int node_idx : segment.internal_node_indices) {
             ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx));
         }
-        // Fusion checks need the original use counts, including consumers outside this segment.
-        for (ggml_tensor* tensor : represented) {
-            const size_t source_index = ggml_hash_find(&gf->visited_hash_set, tensor);
-            if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(gf->visited_hash_set.used, source_index)) {
-                const size_t index               = ggml_hash_insert(&segment_graph->visited_hash_set, tensor);
-                segment_graph->use_counts[index] = gf->use_counts[source_index];
-            }
-        }
+        GGML_ASSERT(ggml_graph_cpy_use_counts(gf, segment_graph));
         *graph_ctx_out = graph_ctx;
         return segment_graph;
     }
@@ -896,8 +888,8 @@ namespace sd::ggml_graph_cut {
             return plan;
         }
         plan.layout = graph_layout(gf, false);
-        for (int i = 0; i < gf->n_leafs; ++i) {
-            plan.leaf_names.emplace_back(gf->leafs[i]->name);
+        for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) {
+            plan.leaf_names.emplace_back(ggml_graph_leaf(gf, i)->name);
         }
         plan.compute_buffer_size =
             measure_graph_compute_buffer(backend, gf, params_tensor_set);

src/core/vulkan/CMakeLists.txt

@@ -0,0 +1,69 @@
+find_package(Vulkan COMPONENTS glslc REQUIRED)
+set(SD_VULKAN_GENERATED "${CMAKE_CURRENT_BINARY_DIR}/shaders")
+file(MAKE_DIRECTORY "${SD_VULKAN_GENERATED}")
+
+function(sd_vulkan_feature NAME EXTENSION)
+    set(SOURCE "${SD_VULKAN_GENERATED}/feature_${NAME}.comp")
+    set(EXTENSIONS "#extension ${EXTENSION} : require")
+    if(NAME STREQUAL "SD_VK_CM2")
+        string(PREPEND EXTENSIONS "#extension GL_KHR_cooperative_matrix : require\n")
+    endif()
+    file(WRITE "${SOURCE}" "#version 460\n${EXTENSIONS}\nlayout(local_size_x = 1) in;\nvoid main() {}\n")
+    execute_process(COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=vulkan1.3 "${SOURCE}"
+                    -o "${SD_VULKAN_GENERATED}/feature_${NAME}.spv"
+                    RESULT_VARIABLE RESULT OUTPUT_QUIET ERROR_QUIET)
+    if(RESULT EQUAL 0)
+        set(${NAME} ON PARENT_SCOPE)
+    else()
+        set(${NAME} OFF PARENT_SCOPE)
+        message(STATUS "sd.cpp Vulkan: ${EXTENSION} unavailable in glslc")
+    endif()
+endfunction()
+
+sd_vulkan_feature(SD_VK_DOT GL_EXT_integer_dot_product)
+sd_vulkan_feature(SD_VK_CM1 GL_KHR_cooperative_matrix)
+sd_vulkan_feature(SD_VK_CM2 GL_NV_cooperative_matrix2)
+
+function(sd_vulkan_shader NAME SHADER ENABLED)
+    set(HEADER "${SD_VULKAN_GENERATED}/${NAME}.h")
+    if(ENABLED)
+        set(SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${SHADER}.comp")
+        set(SPV "${SD_VULKAN_GENERATED}/${NAME}.spv")
+        set(ENV vulkan1.2)
+        if(NAME MATCHES "_cm2$")
+            set(ENV vulkan1.3)
+        endif()
+        add_custom_command(OUTPUT "${HEADER}"
+            COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=${ENV} ${ARGN} "${SOURCE}" -o "${SPV}"
+            COMMAND "${CMAKE_COMMAND}" "-DINPUT=${SPV}" "-DOUTPUT=${HEADER}" "-DNAME=${NAME}"
+                    -P "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake"
+            DEPENDS "${SOURCE}" "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake"
+            BYPRODUCTS "${SPV}" VERBATIM)
+    else()
+        file(WRITE "${HEADER}" "#include <cstdint>\nnamespace sd_vulkan_shaders { static constexpr uint32_t ${NAME}[] = {0}; }\n")
+    endif()
+    set(SD_VULKAN_HEADERS ${SD_VULKAN_HEADERS} "${HEADER}" PARENT_SCOPE)
+endfunction()
+
+# Preserve the native compilation policy: optimization changes RoPE reductions
+# and is disabled for KHR cooperative matrix shaders.
+sd_vulkan_shader(rope_apply rope_apply ON)
+sd_vulkan_shader(rms_rope_apply rms_rope_apply ON)
+sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16)
+sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O)
+sd_vulkan_shader(quantize_i8_convrot_subgroup quantize_i8_convrot_subgroup ${SD_VK_DOT} -O)
+sd_vulkan_shader(quantize_i8_convrot_register quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSINGLE_PASS=1)
+sd_vulkan_shader(swiglu_quantize_i8_convrot quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1)
+sd_vulkan_shader(swiglu_quantize_i8_convrot_16384 quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1 -DSINGLE_PASS=1 -DSTATIC_K=16384)
+sd_vulkan_shader(sigmoid_quantize_i8_convrot_6144 quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSIGMOID_MUL=1 -DSINGLE_PASS=1 -DSTATIC_K=6144)
+sd_vulkan_shader(modulate_convrot modulate_convrot ${SD_VK_DOT} -O)
+sd_vulkan_shader(rms_modulate_convrot modulate_convrot ${SD_VK_DOT} -O -DRMS_AFFINE=1)
+sd_vulkan_shader(mul_mat_i8_tensorwise mul_mat_i8_tensorwise ${SD_VK_DOT} -O)
+sd_vulkan_shader(mul_mat_i8_tensorwise_cm1 mul_mat_i8_tensorwise_cm1 ${SD_VK_CM1})
+sd_vulkan_shader(mul_mat_i8_tensorwise_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O)
+sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DSPLIT_SWIGLU=1)
+sd_vulkan_shader(mul_mat_i8_residual_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DRESIDUAL=1)
+add_custom_target(sd-vulkan-shaders DEPENDS ${SD_VULKAN_HEADERS})
+add_dependencies(${SD_LIB} sd-vulkan-shaders)
+target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp")
+target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}")

src/core/vulkan/int8.cpp

@@ -0,0 +1,806 @@
+#include <algorithm>
+#include <array>
+#include <climits>
+#include <cstring>
+#include <vector>
+#include "core/ggml_extensions.h"
+#include "ggml-ext-vulkan.h"
+#include "modulate_convrot.h"
+#include "mul_mat_i8_residual_cm2.h"
+#include "mul_mat_i8_swiglu_split_cm2.h"
+#include "mul_mat_i8_tensorwise.h"
+#include "mul_mat_i8_tensorwise_cm1.h"
+#include "mul_mat_i8_tensorwise_cm2.h"
+#include "quantize_i8_convrot.h"
+#include "quantize_i8_convrot_register.h"
+#include "quantize_i8_convrot_subgroup.h"
+#include "rms_modulate_convrot.h"
+#include "sigmoid_quantize_i8_convrot_6144.h"
+#include "swiglu_quantize_i8_convrot.h"
+#include "swiglu_quantize_i8_convrot_16384.h"
+
+namespace sd_vulkan_i8 {
+    struct Device {
+        const ggml_ext_vulkan_api* api;
+        ggml_ext_vulkan_device_info info;
+        ggml_ext_vulkan_compute_info compute;
+        bool subgroup, cm2, cm2_small;
+        uint32_t cm_m, cm_n, cm_k;
+    };
+    static bool device_info(ggml_backend_dev_t dev, Device& d) {
+        auto reg     = ggml_backend_dev_backend_reg(dev);
+        auto query   = reinterpret_cast<ggml_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1"));
+        auto compute = reinterpret_cast<ggml_ext_vulkan_compute_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_compute_v1"));
+        d            = {};
+        if (!query || !compute)
+            return false;
+        d.api        = query();
+        auto service = compute();
+        if (!d.api->device_info(dev, &d.info) || !service->properties(dev, &d.compute))
+            return false;
+        d.subgroup = d.info.subgroup_size == 32 && d.info.full_subgroups && d.info.subgroup_shuffle && d.compute.subgroup_arithmetic && d.info.max_group_invocations >= 256;
+        ggml_ext_vulkan_matrix_properties prop;
+        for (uint32_t i = 0; service->matrix_properties(dev, i, &prop); ++i) {
+            if (prop.a != GGML_TYPE_I8 || prop.b != GGML_TYPE_I8 || prop.accumulator != GGML_TYPE_I32 || prop.result != GGML_TYPE_I32 || prop.saturating || !prop.m || !prop.n || !prop.k)
+                continue;
+            if (!prop.flexible && !prop.workgroup_scope && !d.cm_m && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm1) > 4) {
+                d.cm_m = prop.m;
+                d.cm_n = prop.n;
+                d.cm_k = prop.k;
+            }
+            if (prop.flexible && prop.workgroup_scope && prop.invocations == 128 && 128 % prop.m == 0 && 64 % prop.n == 0 && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm2) > 4) {
+                d.cm2 = true;
+                d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0;
+            }
+        }
+        return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4;
+    }
+    static const ggml_tensor* node(const ggml_cgraph* graph, int i) {
+        return ggml_graph_node(const_cast<ggml_cgraph*>(graph), i);
+    }
+    static int count_nodes(const ggml_cgraph* graph) {
+        return ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
+    }
+    static int param_i32(const ggml_tensor* node, int i) {
+        return node->op_params[i];
+    }
+    static float param_f32(const ggml_tensor* node, int i) {
+        float value;
+        std::memcpy(&value, node->op_params + i, 4);
+        return value;
+    }
+    static bool is_convrot(const ggml_tensor* node) {
+        int group;
+        return sd_ggml_convrot_describe(node, &group);
+    }
+    static bool is_mat(const ggml_tensor* node) {
+        SDI8Operation op;
+        return sd_ggml_i8_describe(node, &op);
+    }
+    static bool aligned(const ggml_tensor* t, size_t alignment) {
+        size_t offset = 0;
+        while (!t->data && t->view_src) {
+            offset += t->view_offs;
+            t = t->view_src;
+        }
+        return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
+    }
+    static bool bounded(const Device& d, const ggml_tensor* t, bool descriptor_aligned = true) {
+        return t && ggml_nbytes(t) <= d.info.max_buffer_range - (d.info.buffer_alignment - 1) &&
+               aligned(t, descriptor_aligned ? d.info.buffer_alignment : 4);
+    }
+    static bool can_modulate(const Device& d, int64_t k, bool affine) {
+        return (k == 4096 || k == 6144 || (affine && k == 5376)) && d.info.max_group_invocations >= 512 && d.info.max_group_size[0] >= 512 && d.compute.shared_memory >= (k + 1024) * 4;
+    }
+    static bool can_quantize(const Device& d, const ggml_tensor* dst) {
+        int group;
+        if (!sd_ggml_convrot_describe(dst, &group) || group != 256 || !bounded(d, dst) || !bounded(d, dst->src[0]))
+            return false;
+        const auto* x = dst->src[0];
+        return x->ne[0] <= UINT32_MAX && ggml_nrows(x) <= UINT32_MAX && d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 256;
+    }
+    static bool can_matmul(const Device& d, const ggml_tensor* dst) {
+        SDI8Operation op;
+        if (!sd_ggml_i8_describe(dst, &op) || !op.packed || op.group_size != 256 || !op.weight_scale || !bounded(d, dst))
+            return false;
+        for (auto* t : {op.weight, op.input, op.weight_scale, op.bias})
+            if (t && !bounded(d, t))
+                return false;
+        const uint64_t cols = op.weight->ne[1], rows = ggml_nrows(dst);
+        return op.weight->ne[0] % 256 == 0 && op.weight->ne[0] <= UINT32_MAX && cols <= UINT32_MAX && rows <= UINT32_MAX &&
+               (cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] &&
+               d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
+    }
+    static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) {
+        if (i + 2 >= count_nodes(graph) || !device.subgroup) {
+            return false;
+        }
+        const ggml_tensor* sigmoid = node(graph, i);
+        const ggml_tensor* mul     = node(graph, i + 1);
+        const ggml_tensor* quant   = node(graph, i + 2);
+        if (sigmoid->op != GGML_OP_UNARY || ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID ||
+            mul->op != GGML_OP_MUL || mul->src[1] != sigmoid ||
+            !is_convrot(quant) || quant->src[0] != mul ||
+            sigmoid->ne[0] != 6144 || !ggml_are_same_shape(sigmoid, mul->src[0]) ||
+            !can_quantize(device, quant)) {
+            return false;
+        }
+        for (const ggml_tensor* tensor : {sigmoid->src[0], mul->src[0]}) {
+            if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) ||
+                tensor->view_offs % device.info.buffer_alignment != 0) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    static bool can_fuse_swiglu_convrot(const Device& device, const ggml_cgraph* graph, int i) {
+        if (i + 1 >= count_nodes(graph) || !device.subgroup) {
+            return false;
+        }
+        const ggml_tensor* glu   = node(graph, i);
+        const ggml_tensor* quant = node(graph, i + 1);
+        if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
+            glu->src[1] == nullptr || param_i32(glu, 1) != 0 ||
+            !is_convrot(quant) || quant->src[0] != glu ||
+            !can_quantize(device, quant) || !ggml_is_contiguous(glu->src[0]) || !ggml_is_contiguous(glu->src[1])) {
+            return false;
+        }
+        for (const ggml_tensor* tensor : {glu->src[0], glu->src[1]}) {
+            if (tensor->type != GGML_TYPE_F32 || tensor->view_offs % device.info.buffer_alignment != 0) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    static int match_modulate_convrot(const Device& device, const ggml_cgraph* graph, int quant_idx, int* indices) {
+        if (quant_idx < 1 || quant_idx >= count_nodes(graph)) {
+            return 0;
+        }
+        const ggml_tensor* quant = node(graph, quant_idx);
+        if (!is_convrot(quant) || !can_quantize(device, quant)) {
+            return 0;
+        }
+        const ggml_tensor* nodes[4]{};
+
+        const ggml_tensor* mod = nullptr;
+        const bool affine      = quant->src[0]->op == GGML_OP_ADD;
+        const int count        = affine ? 4 : 3;
+        if (affine) {
+            const ggml_tensor* shift = quant->src[0];
+            const ggml_tensor* add   = shift->src[0];
+            if (add->op != GGML_OP_ADD || add->src[1]->op != GGML_OP_MUL || add->src[1]->src[0] != add->src[0]) {
+                return 0;
+            }
+            nodes[0]                  = add->src[1];
+            nodes[1]                  = add;
+            nodes[2]                  = shift;
+            nodes[3]                  = quant;
+            mod                       = nodes[0]->src[1];
+            const ggml_tensor* vector = shift->src[1];
+            if (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) ||
+                vector->ne[0] != add->ne[0] || ggml_nelements(vector) != add->ne[0]) {
+                return 0;
+            }
+        } else {
+            const ggml_tensor* mul = quant->src[0];
+            if (mul->op != GGML_OP_MUL || mul->src[0]->op != GGML_OP_NORM) {
+                return 0;
+            }
+            nodes[0] = mul->src[0];
+            nodes[1] = mul;
+            nodes[2] = quant;
+            mod      = mul->src[1];
+        }
+        const ggml_tensor* x = nodes[0]->src[0];
+        if ((x->ne[0] != 4096 && x->ne[0] != 6144 && x->ne[0] != 5376) ||
+            !can_modulate(device, x->ne[0], affine) ||
+            x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || mod->type != GGML_TYPE_F32 ||
+            !ggml_is_contiguous(mod) || mod->ne[0] != x->ne[0] || ggml_nelements(mod) != x->ne[0]) {
+            return 0;
+        }
+        for (int j = 0; j < count; ++j) {
+            if (j < count - 1 && (nodes[j]->type != GGML_TYPE_F32 || !ggml_is_contiguous(nodes[j]) || !ggml_are_same_shape(x, nodes[j]))) {
+                return 0;
+            }
+        }
+        int next = count - 1;
+        for (int j = quant_idx; j >= std::max(0, quant_idx - 64) && next >= 0; --j) {
+            if (node(graph, j) == nodes[next]) {
+                indices[next--] = j;
+            }
+        }
+        return next < 0 ? count : 0;
+    }
+
+    static int match_rms_modulate_convrot(const Device& device, const ggml_cgraph* graph, int q, int* indices) {
+        if (!can_modulate(device, 5376, true) || match_modulate_convrot(device, graph, q, indices + 2) != 4) {
+            return 0;
+        }
+        const auto weighted = node(graph, indices[2])->src[0];
+        if (weighted->op != GGML_OP_MUL || weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 5376) {
+            return 0;
+        }
+        const auto norm   = weighted->src[0];
+        const auto x      = norm->src[0];
+        const auto weight = weighted->src[1];
+        if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || weight->type != GGML_TYPE_F32 ||
+            !ggml_is_contiguous(weight) || ggml_nelements(weight) != 5376 || weight->ne[0] != 5376) {
+            return 0;
+        }
+        int next = 1;
+        for (int j = indices[2] - 1; j >= std::max(0, indices[2] - 64) && next >= 0; --j) {
+            if (node(graph, j) == (next == 1 ? weighted : norm)) {
+                indices[next--] = j;
+            }
+        }
+
+        return next < 0 ? 6 : 0;
+    }
+
+    static int can_fuse_modulate_convrot(const Device& device, const ggml_cgraph* graph, int i) {
+        const ggml_op op = node(graph, i)->op;
+        if (op == GGML_OP_RMS_NORM) {
+            int indices[6];
+            if (match_rms_modulate_convrot(device, graph, i + 5, indices) != 6) {
+                return 0;
+            }
+            for (int j = 0; j < 6; ++j) {
+                if (indices[j] != i + j) {
+                    return 0;
+                }
+            }
+            return 6;
+        }
+        if (op != GGML_OP_NORM && op != GGML_OP_MUL) {
+            return 0;
+        }
+        const int count = op == GGML_OP_NORM ? 3 : 4;
+        int indices[4];
+        if (match_modulate_convrot(device, graph, i + count - 1, indices) != count) {
+            return 0;
+        }
+        for (int j = 0; j < count; ++j) {
+            if (indices[j] != i + j) {
+                return 0;
+            }
+        }
+        return count;
+    }
+
+    static bool match_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i, int* indices) {
+        if (i < 2 || i >= count_nodes(graph) || !device.cm2_small) {
+            return false;
+        }
+        const ggml_tensor* glu = node(graph, i);
+        if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
+            param_i32(glu, 1) != 0 || glu->src[1] == nullptr || glu->ne[0] % 128 != 0 ||
+            glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) ||
+            glu->view_offs % device.info.buffer_alignment != 0) {
+            return false;
+        }
+        const ggml_tensor* gate = glu->src[0];
+        const ggml_tensor* up   = glu->src[1];
+        if (gate == up || gate->src[1] != up->src[1] || !ggml_are_same_shape(gate, up)) {
+            return false;
+        }
+        for (auto mat : {gate, up}) {
+            if (!is_mat(mat) || mat->view_src != nullptr || mat->src[3] != nullptr ||
+                !can_matmul(device, mat)) {
+                return false;
+            }
+        }
+        int found = 0;
+        for (int j = std::max(0, i - 32); j < i; ++j) {
+            if (node(graph, j) == gate || node(graph, j) == up) {
+                indices[found++] = j;
+            }
+        }
+        if (found != 2) {
+            return false;
+        }
+        indices[2] = i;
+
+        return true;
+    }
+
+    static bool can_fuse_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i) {
+        int indices[3];
+        return match_i8_swiglu_split(device, graph, i + 2, indices) && indices[0] == i && indices[1] == i + 1;
+    }
+
+    static bool can_fuse_i8_swiglu(const Device& device, const ggml_cgraph* graph, int i) {
+        if (i + 1 >= count_nodes(graph) || !device.cm2_small) {
+            return false;
+        }
+        const ggml_tensor* mat = node(graph, i);
+        const ggml_tensor* glu = node(graph, i + 1);
+        return is_mat(mat) && can_matmul(device, mat) &&
+               mat->view_src == nullptr && mat->src[3] == nullptr && mat->ne[0] % 256 == 0 &&
+               glu->op == GGML_OP_GLU && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU &&
+               param_i32(glu, 1) == 0 && glu->src[0] == mat && glu->src[1] == nullptr &&
+               glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) &&
+               glu->view_offs % device.info.buffer_alignment == 0;
+    }
+
+    static int match_i8_residual(const Device& device, const ggml_cgraph* graph, int add_idx, int* indices) {
+        if (add_idx < 2 || add_idx >= count_nodes(graph) || !device.cm2) {
+            return 0;
+        }
+        const ggml_tensor* dst = node(graph, add_idx);
+        if (dst->op != GGML_OP_ADD) {
+            return 0;
+        }
+        const ggml_tensor* mul = dst->src[0]->op == GGML_OP_MUL ? dst->src[0] : dst->src[1];
+        if (mul->op != GGML_OP_MUL) {
+            return 0;
+        }
+        const ggml_tensor* scaled_mat = mul->src[0];
+        const bool scaled             = scaled_mat->op == GGML_OP_SCALE;
+        const ggml_tensor* mat        = scaled ? scaled_mat->src[0] : scaled_mat;
+        if (!is_mat(mat) || !can_matmul(device, mat) || mat->src[3] != nullptr ||
+            (scaled && param_f32(scaled_mat, 1) != 0.0f)) {
+            return 0;
+        }
+        const int count             = scaled ? 4 : 3;
+        const ggml_tensor* nodes[4] = {mat, scaled ? scaled_mat : mul, scaled ? mul : dst, dst};
+
+        const ggml_tensor* gate     = mul->src[1];
+        const ggml_tensor* residual = dst->src[0] == mul ? dst->src[1] : dst->src[0];
+        if (!ggml_are_same_shape(mat, dst) || !ggml_are_same_shape(residual, dst) ||
+            gate->ne[0] != mat->ne[0] || ggml_nelements(gate) != mat->ne[0]) {
+            return 0;
+        }
+        for (const ggml_tensor* tensor : {mat, scaled_mat, mul, gate, residual, dst}) {
+            if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) ||
+                tensor->view_offs % device.info.buffer_alignment != 0) {
+                return 0;
+            }
+        }
+        int next = count - 1;
+        for (int j = add_idx; j >= std::max(0, add_idx - 32) && next >= 0; --j) {
+            if (node(graph, j) == nodes[next]) {
+                indices[next--] = j;
+            }
+        }
+        return next < 0 ? count : 0;
+    }
+
+    static int can_fuse_i8_residual(const Device& device, const ggml_cgraph* graph, int i) {
+        if (i + 2 >= count_nodes(graph) || !is_mat(node(graph, i))) {
+            return 0;
+        }
+        const int count = node(graph, i + 1)->op == GGML_OP_SCALE ? 4 : 3;
+        int indices[4];
+        if (match_i8_residual(device, graph, i + count - 1, indices) != count) {
+            return 0;
+        }
+        for (int j = 0; j < count; ++j) {
+            if (indices[j] != i + j) {
+                return 0;
+            }
+        }
+        return count;
+    }
+
+    static int match_i8_row_concat(const Device& device, const ggml_cgraph* graph, int end, int* indices) {
+        const ggml_tensor* concat = node(graph, end);
+        if (concat->op != GGML_OP_CONCAT || param_i32(concat, 0) != 1 ||
+            concat->type != GGML_TYPE_F32 || concat->ne[2] != 1 || concat->ne[3] != 1 ||
+            !ggml_is_contiguous(concat) || !ggml_is_contiguous(concat->src[0]) ||
+            concat->view_offs % device.info.buffer_alignment != 0 ||
+            ggml_nbytes(concat->src[0]) % device.info.buffer_alignment != 0) {
+            return 0;
+        }
+        const ggml_tensor* tail = concat->src[1];
+        int count               = 0;
+        for (int j = end - 1; j >= std::max(0, end - 32); --j) {
+            if (node(graph, j) != tail) {
+                continue;
+            }
+            if (is_mat(tail) && can_matmul(device, tail)) {
+                indices[0] = j;
+                count      = 1;
+            } else {
+                count = match_i8_residual(device, graph, j, indices);
+            }
+            break;
+        }
+        if (!count) {
+            return 0;
+        }
+        indices[count++] = end;
+
+        return count;
+    }
+
+    static int can_fuse_i8_row_concat(const Device& device, const ggml_cgraph* graph, int i) {
+        if (!is_mat(node(graph, i))) {
+            return 0;
+        }
+        for (int end = i + 1; end < std::min(i + 5, count_nodes(graph)); ++end) {
+            int indices[5];
+            const int count = match_i8_row_concat(device, graph, end, indices);
+            if (count != end - i + 1) {
+                continue;
+            }
+            bool consecutive = true;
+            for (int j = 0; j < count; ++j) {
+                consecutive &= indices[j] == i + j;
+            }
+            if (consecutive) {
+                return count;
+            }
+        }
+        return 0;
+    }
+
+    struct QuantParameters {
+        uint32_t k, rows;
+    };
+    struct ModulateParameters {
+        uint32_t rows, x_offset, m_offset, h_offset;
+        float eps;
+        uint32_t k, w_offset;
+    };
+    struct MatParameters {
+        uint32_t n, rows, k, has_bias, row_offset;
+        float output_scale;
+        uint32_t group_cols;
+    };
+    static_assert(sizeof(QuantParameters) == 8 && sizeof(ModulateParameters) == 28 && sizeof(MatParameters) == 28, "INT8 push constant layout mismatch");
+    template <size_t N>
+    static ggml_ext_vulkan_pipeline shader(const char* name, const uint32_t (&spirv)[N], uint32_t bindings, uint32_t constants) {
+        return {name, spirv, sizeof(spirv), bindings, constants, nullptr, 0, 0, false};
+    }
+#define SD_VK_PIPELINE(NAME, BINDINGS, PARAMS) shader("sd." #NAME ".vulkan", sd_vulkan_shaders::NAME, BINDINGS, sizeof(PARAMS))
+    static ggml_ext_vulkan_binding range(ggml_ext_vulkan_binding binding, size_t offset, size_t size) {
+        binding.offset += offset;
+        binding.size = size;
+        return binding;
+    }
+    static bool overlap(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) {
+        return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size;
+    }
+    static bool output_binding(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const std::vector<ggml_ext_vulkan_binding>& inputs, ggml_ext_vulkan_binding& binding, bool& alias, const ggml_tensor* residual = nullptr) {
+        binding = {output};
+        alias   = false;
+        ggml_ext_vulkan_buffer_info dst, src;
+        if (!d.api->buffer_info(e->backend, &binding, &dst))
+            return false;
+        for (const auto& input : inputs) {
+            if (!d.api->buffer_info(e->backend, &input, &src))
+                return false;
+            // The residual epilogue reads each element before writing that same element.
+            if (residual && input.tensor == residual && dst.storage == src.storage && dst.offset == src.offset && dst.size == src.size)
+                continue;
+            alias |= overlap(dst, src);
+        }
+        if (alias)
+            binding = {nullptr, e->workspace, 0, ggml_nbytes(output)};
+        return true;
+    }
+    static ggml_status finish(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const ggml_ext_vulkan_binding& binding, bool alias, ggml_status status) {
+        if (status == GGML_STATUS_SUCCESS && alias) {
+            ggml_ext_vulkan_binding dst = {output};
+            return d.api->copy(e->backend, &binding, &dst);
+        }
+        return status;
+    }
+    static ggml_status quantize(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
+        const auto* output                          = nodes[count - 1];
+        const auto* first                           = nodes[0];
+        const bool modulate                         = count > 1 && (first->op == GGML_OP_NORM || first->op == GGML_OP_MUL || first->op == GGML_OP_RMS_NORM);
+        const bool rms                              = modulate && first->op == GGML_OP_RMS_NORM;
+        const bool affine                           = modulate && (rms || first->op == GGML_OP_MUL);
+        const bool sigmoid                          = count > 1 && first->op == GGML_OP_UNARY;
+        const auto* x                               = count == 1 ? output->src[0] : first->src[0];
+        const auto* up                              = count == 1 || modulate ? nullptr : sigmoid ? nodes[1]->src[0]
+                                                                                                 : first->src[1];
+        const int extra                             = rms ? 2 : 0;
+        const auto* mod                             = modulate ? nodes[extra + (affine ? 0 : 1)]->src[1] : nullptr;
+        const auto* shift                           = modulate ? (affine ? nodes[extra + 2]->src[1] : mod) : nullptr;
+        const auto* weight                          = rms ? nodes[1]->src[1] : x;
+        std::vector<ggml_ext_vulkan_binding> inputs = {{x}};
+        if (up)
+            inputs.push_back({up});
+        if (modulate) {
+            inputs.push_back({mod});
+            inputs.push_back({shift});
+            if (rms)
+                inputs.push_back({weight});
+        }
+        ggml_ext_vulkan_binding dst;
+        bool alias;
+        if (!output_binding(e, d, output, inputs, dst, alias))
+            return GGML_STATUS_FAILED;
+        const uint32_t k = uint32_t(x->ne[0]), rows = uint32_t(ggml_nrows(x));
+        const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
+        auto packed = range(dst, 0, scale_offset), scales = range(dst, scale_offset, size_t(rows) * 4);
+        const uint32_t groups[] = {std::min(rows, d.info.max_group_count[0]), 1, 1};
+        ggml_status status;
+        if (modulate) {
+            std::vector<ggml_ext_vulkan_binding> bindings = {{x}, {mod}, {shift}, packed, scales};
+            if (rms)
+                bindings.push_back({weight});
+            ModulateParameters pc = {rows, 0, 0, 0, affine && !rms ? 0.f : param_f32(first, 0), k, 0};
+            uint32_t* offsets[]   = {&pc.x_offset, &pc.m_offset, &pc.h_offset, nullptr, nullptr, &pc.w_offset};
+            for (size_t i = 0; i < bindings.size(); ++i)
+                if (offsets[i]) {
+                    ggml_ext_vulkan_buffer_info info;
+                    if (!d.api->buffer_info(e->backend, &bindings[i], &info))
+                        return GGML_STATUS_FAILED;
+                    *offsets[i] = info.shader_offset / 4;
+                }
+            auto pipeline              = rms ? SD_VK_PIPELINE(rms_modulate_convrot, 6, ModulateParameters) : SD_VK_PIPELINE(modulate_convrot, 5, ModulateParameters);
+            const uint32_t constants[] = {k, uint32_t(affine)};
+            pipeline.constants         = constants;
+            pipeline.constant_count    = 2;
+            status                     = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
+        } else {
+            ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
+            auto pipeline                      = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
+            bool subgroup                      = d.subgroup && k >= 1024;
+            uint32_t width                     = k;
+            if (subgroup) {
+                pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters);
+                if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) {
+                    pipeline                = SD_VK_PIPELINE(quantize_i8_convrot_register, 3, QuantParameters);
+                    pipeline.constants      = &width;
+                    pipeline.constant_count = 1;
+                }
+            }
+            if (up) {
+                subgroup = true;
+                pipeline = k == 16384 ? SD_VK_PIPELINE(swiglu_quantize_i8_convrot_16384, 4, QuantParameters) : SD_VK_PIPELINE(swiglu_quantize_i8_convrot, 4, QuantParameters);
+                if (sigmoid)
+                    pipeline = SD_VK_PIPELINE(sigmoid_quantize_i8_convrot_6144, 4, QuantParameters);
+            }
+            if (subgroup) {
+                pipeline.full_subgroups         = d.info.full_subgroups;
+                pipeline.required_subgroup_size = d.info.subgroup_size_control ? 32 : 0;
+            }
+            QuantParameters pc = {k, rows};
+            status             = d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups);
+        }
+        return finish(e, d, output, dst, alias, status);
+    }
+    static uint32_t group_columns(uint32_t cols, uint32_t k, uint32_t tile_cols, bool swiglu) {
+        if (!cols || !k || !tile_cols)
+            return 1;
+        const uint64_t col_tiles         = (uint64_t(cols) + tile_cols - 1) / tile_cols;
+        const uint64_t weight_tile_bytes = uint64_t(tile_cols) * k * (swiglu ? 2 : 1);
+        const uint64_t growth            = std::clamp(uint64_t(k) / cols, uint64_t(1), uint64_t(2));
+        return uint32_t(std::clamp(uint64_t(8) * 1024 * 1024 * growth / weight_tile_bytes, uint64_t(1), std::min(col_tiles, uint64_t(32))));
+    }
+    static ggml_status multiply(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
+        const auto* output      = nodes[count - 1];
+        const auto* concat      = output->op == GGML_OP_CONCAT ? output : nullptr;
+        const auto* tail        = concat ? concat->src[1] : output;
+        const auto* glu         = tail->op == GGML_OP_GLU ? tail : nullptr;
+        const auto* mat         = glu ? glu->src[0] : tail;
+        const ggml_tensor *gate = nullptr, *residual = nullptr;
+        float output_scale = 1.f;
+        if (tail->op == GGML_OP_ADD) {
+            const auto* mul = tail->src[0]->op == GGML_OP_MUL ? tail->src[0] : tail->src[1];
+            gate            = mul->src[1];
+            residual        = tail->src[0] == mul ? tail->src[1] : tail->src[0];
+            mat             = mul->src[0];
+            if (mat->op == GGML_OP_SCALE) {
+                output_scale = param_f32(mat, 0);
+                mat          = mat->src[0];
+            }
+        }
+        const auto* up_mat = glu ? glu->src[1] : nullptr;
+        SDI8Operation op, up;
+        if (!sd_ggml_i8_describe(mat, &op) || (up_mat && !sd_ggml_i8_describe(up_mat, &up)))
+            return GGML_STATUS_FAILED;
+        const uint32_t k = uint32_t(op.weight->ne[0]), n = uint32_t(op.weight->ne[1]), rows = uint32_t(ggml_nrows(mat));
+        const size_t scale_offset                   = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
+        std::vector<ggml_ext_vulkan_binding> inputs = {{op.weight}, {op.input}, {op.weight_scale}, {op.bias ? op.bias : op.weight_scale}};
+        if (gate) {
+            inputs.push_back({gate});
+            inputs.push_back({residual});
+        }
+        if (up_mat) {
+            inputs.push_back({up.weight});
+            inputs.push_back({up.weight_scale});
+        }
+        if (concat && ggml_nbytes(concat->src[0]))
+            inputs.push_back({concat->src[0]});
+        ggml_ext_vulkan_binding dst;
+        bool alias;
+        if (!output_binding(e, d, output, inputs, dst, alias, concat ? nullptr : residual))
+            return GGML_STATUS_FAILED;
+        auto destination                   = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst;
+        ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}};
+        auto dispatch                      = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows,
+                                                 std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) {
+            if (!dispatch_rows)
+                return GGML_STATUS_SUCCESS;
+            const uint32_t cols     = glu ? uint32_t(glu->ne[0]) : n;
+            MatParameters pc        = {n * (up_mat ? 2u : 1u), rows, k, op.bias ? 1u : 0u, row_offset, output_scale, group_columns(cols, k, tile_cols, glu != nullptr)};
+            pipeline.constants      = constants.begin();
+            pipeline.constant_count = uint32_t(constants.size());
+            const uint32_t groups[] = {(cols + tile_cols - 1) / tile_cols, (dispatch_rows + tile_rows - 1) / tile_rows, 1};
+            return d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups);
+        };
+        ggml_status status = GGML_STATUS_SUCCESS;
+        if (d.cm2) {
+            const bool large_k       = k % 4096 == 0 || (rows >= n && k >= n);
+            const uint32_t full_rows = (glu || n % 128 == 0) ? rows / 128 * 128 : 0;
+            auto pipeline            = residual ? SD_VK_PIPELINE(mul_mat_i8_residual_cm2, 8, MatParameters) : up_mat ? SD_VK_PIPELINE(mul_mat_i8_swiglu_split_cm2, 8, MatParameters)
+                                                                                                                     : SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm2, 6, MatParameters);
+            const uint32_t tile_cols = glu ? 64 : 128;
+            status                   = dispatch(pipeline, tile_cols, 128, {128, 1, glu ? 1u : 0u, tile_cols, glu || large_k ? 128u : 64u}, 0, full_rows);
+            if (status == GGML_STATUS_SUCCESS) {
+                const bool small         = d.cm2_small && (glu || n % 128 == 0) && rows - full_rows <= 32;
+                const uint32_t tail_cols = small || glu ? 64 : 128, tail_rows = small ? 32 : glu ? 128
+                                                                                                 : 64;
+                status = dispatch(pipeline, tail_cols, tail_rows, {tail_rows, small ? 2u : 0u, glu ? 1u : 0u, tail_cols, small ? 256u : 64u}, full_rows, rows - full_rows);
+            }
+        } else {
+            uint32_t coop_rows = 0;
+            if (d.cm_m && k % d.cm_k == 0 && n % (4 * d.cm_n) == 0) {
+                const uint32_t tile_rows        = 4 * d.cm_m;
+                coop_rows                       = rows / tile_rows * tile_rows;
+                auto pipeline                   = SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm1, 6, MatParameters);
+                pipeline.full_subgroups         = d.info.full_subgroups;
+                pipeline.required_subgroup_size = d.info.subgroup_size_control ? d.info.subgroup_size : 0;
+                status                          = dispatch(pipeline, 4 * d.cm_n, tile_rows, {d.info.subgroup_size, d.cm_m, d.cm_n, d.cm_k}, 0, coop_rows);
+            }
+            if (status == GGML_STATUS_SUCCESS)
+                status = dispatch(SD_VK_PIPELINE(mul_mat_i8_tensorwise, 6, MatParameters), 16, 16, {}, coop_rows, rows - coop_rows);
+        }
+        if (status == GGML_STATUS_SUCCESS && concat && ggml_nbytes(concat->src[0])) {
+            ggml_ext_vulkan_binding src = {concat->src[0]}, prefix = range(dst, 0, ggml_nbytes(concat->src[0]));
+            status = d.api->copy(e->backend, &src, &prefix);
+        }
+        return finish(e, d, output, dst, alias, status);
+    }
+    static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
+        Device d;
+        return device_info(dev, d) && can_quantize(d, n);
+    }
+    static bool mat_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
+        Device d;
+        return device_info(dev, d) && can_matmul(d, n);
+    }
+    static int match(ggml_backend_dev_t dev, const ggml_cgraph* graph, int first, void*) {
+        Device d;
+        if (!device_info(dev, d) || d.info.fusion_disabled)
+            return 0;
+        if (int count = can_fuse_modulate_convrot(d, graph, first))
+            return count;
+        if (int count = can_fuse_i8_row_concat(d, graph, first))
+            return count;
+        if (int count = can_fuse_i8_residual(d, graph, first))
+            return count;
+        if (can_fuse_i8_swiglu_split(d, graph, first))
+            return 3;
+        if (can_fuse_i8_swiglu(d, graph, first))
+            return 2;
+        if (can_fuse_sigmoid_convrot(d, graph, first))
+            return 3;
+        if (can_fuse_swiglu_convrot(d, graph, first))
+            return 2;
+        return 0;
+    }
+    static void optimize(ggml_backend_dev_t dev, ggml_cgraph* graph, void*) {
+        Device d;
+        if (!device_info(dev, d) || d.info.fusion_disabled || d.compute.graph_optimization_disabled)
+            return;
+        for (int last = 1; last < count_nodes(graph); ++last) {
+            int indices[6];
+            int count = match_rms_modulate_convrot(d, graph, last, indices);
+            if (!count)
+                count = match_modulate_convrot(d, graph, last, indices);
+            if (!count)
+                count = match_i8_row_concat(d, graph, last, indices);
+            if (!count)
+                count = match_i8_residual(d, graph, last, indices);
+            if (!count && match_i8_swiglu_split(d, graph, last, indices))
+                count = 3;
+            if (!count)
+                continue;
+            int first = indices[0];
+            for (int j = 1; j < count; ++j) {
+                if (!ggml_ext_graph_delay(graph, first, j, indices[j]))
+                    break;
+                first = indices[j] - j;
+            }
+        }
+    }
+    static bool may_alias(const ggml_tensor* output, const ggml_tensor* input) {
+        size_t offset     = 0;
+        const size_t size = ggml_nbytes(input);
+        while (!input->data && input->view_src) {
+            offset += input->view_offs;
+            input = input->view_src;
+        }
+        // Unallocated inputs get allocator storage, separate from a preallocated output.
+        if (!input->data)
+            return false;
+        if (!input->buffer)
+            return true;
+        if (output->buffer != input->buffer)
+            return false;
+        const auto dst = reinterpret_cast<uintptr_t>(output->data);
+        const auto src = reinterpret_cast<uintptr_t>(input->data) + offset;
+        return dst < src + size && src < dst + ggml_nbytes(output);
+    }
+    static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) {
+        const auto* output = nodes[count - 1];
+        // Fresh outputs cannot alias external inputs except the residual's elementwise reuse.
+        const bool fresh = !output->data && !output->view_src &&
+                           (output->op == GGML_OP_CUSTOM || output->op == GGML_OP_GLU || output->op == GGML_OP_CONCAT ||
+                            (output->op == GGML_OP_ADD && count > 1));
+        *bytes           = fresh ? 0 : ggml_nbytes(output);
+        if (output->data && output->buffer && !output->view_src) {
+            bool alias = false;
+            for (int i = 0; i < count; ++i)
+                for (const auto* src : nodes[i]->src)
+                    if (src && std::find(nodes, nodes + count, src) == nodes + count)
+                        alias |= may_alias(output, src);
+            if (!alias)
+                *bytes = 0;
+        }
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void*, void*) {}
+    static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) {
+        Device d;
+        if (!device_info(ggml_backend_get_device(e->backend), d))
+            return GGML_STATUS_FAILED;
+        return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count);
+    }
+#undef SD_VK_PIPELINE
+}
+
+void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
+    using namespace sd_vulkan_i8;
+    Device info;
+    if (!device_info(device, info))
+        return;
+    ggml_ext_kernel_desc d = {};
+    d.abi_version          = GGML_EXT_ABI_VERSION;
+    d.priority             = 100;
+    d.prepare              = prepare;
+    d.destroy              = destroy;
+    d.execute              = execute;
+    d.name                 = "sd.quantize_i8_convrot.vulkan.v1";
+    d.op                   = GGML_OP_CUSTOM;
+    d.custom_op            = sd_ggml_convrot_op();
+    d.supports             = quant_supports;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.name      = "sd.linear_i8.vulkan.v1";
+    d.op        = GGML_OP_CUSTOM;
+    d.custom_op = sd_ggml_i8_op();
+    d.supports  = mat_supports;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.priority  = 200;
+    d.supports  = nullptr;
+    d.match     = match;
+    d.optimize  = optimize;
+    d.name      = "sd.linear_i8.fused.vulkan.v1";
+    d.op        = GGML_OP_CUSTOM;
+    d.custom_op = sd_ggml_i8_op();
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.optimize  = nullptr;
+    d.custom_op = nullptr;
+    const struct {
+        ggml_op op;
+        const char* name;
+    } fusions[] = {
+        {GGML_OP_NORM, "sd.norm_convrot.vulkan.v1"},
+        {GGML_OP_RMS_NORM, "sd.rms_modulate_convrot.vulkan.v1"},
+        {GGML_OP_MUL, "sd.modulate_convrot.vulkan.v1"},
+        {GGML_OP_GLU, "sd.swiglu_convrot.vulkan.v1"},
+        {GGML_OP_UNARY, "sd.sigmoid_convrot.vulkan.v1"}};
+    for (const auto& fusion : fusions) {
+        d.op   = fusion.op;
+        d.name = fusion.name;
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+}

src/core/vulkan/rope.cpp

@@ -0,0 +1,226 @@
+#include "core/ggml_extensions.h"
+#include "ggml-ext-vulkan.h"
+#include "rms_rope_apply.h"
+#include "rms_rope_cast.h"
+#include "rope_apply.h"
+
+#include <algorithm>
+#include <climits>
+#include <cstring>
+
+namespace sd_vulkan_rope {
+    struct Parameters {
+        uint32_t half_dim, rot_half, heads, tokens, pairs;
+        uint32_t x_stride[4], pe_stride[4];
+        uint32_t x_offset, pe_offset, dst_offset, interleaved;
+    };
+    struct FusedParameters {
+        Parameters rope;
+        float eps;
+        uint32_t w_offset;
+        float output_scale;
+    };
+    static_assert(sizeof(Parameters) == 68 && sizeof(FusedParameters) == 80, "RoPE push constant layout mismatch");
+    static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) {
+        auto query = reinterpret_cast<ggml_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(
+            ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1"));
+        return query ? query() : nullptr;
+    }
+    static float param(const ggml_tensor* node, int i) {
+        float value;
+        std::memcpy(&value, node->op_params + i, sizeof(value));
+        return value;
+    }
+    static bool is_rope(const ggml_tensor* node) {
+        return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
+    }
+    static bool aligned(const ggml_tensor* tensor, size_t alignment) {
+        size_t offset = 0;
+        while (!tensor->data && tensor->view_src) {
+            offset += tensor->view_offs;
+            tensor = tensor->view_src;
+        }
+        return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+    }
+    static bool bounded(const ggml_tensor* node, const ggml_ext_vulkan_device_info& info) {
+        if (!node || !aligned(node, 4) || ggml_nbytes(node) > info.max_buffer_range - (info.buffer_alignment - 1))
+            return false;
+        for (int i = 0; i < 4; ++i)
+            if (node->nb[i] % 4 || node->nb[i] / 4 > UINT32_MAX || node->ne[i] > UINT32_MAX)
+                return false;
+        return true;
+    }
+    static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+        const auto* service = api(device);
+        ggml_ext_vulkan_device_info info;
+        if (!service || !service->device_info(device, &info) || info.max_group_size[0] < 256 || info.max_group_invocations < 256)
+            return false;
+        const auto* x  = node->src[0];
+        const auto* pe = node->src[1];
+        return x && pe && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 && node->type == GGML_TYPE_F32 &&
+               x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] > 0 && pe->ne[2] <= x->ne[0] / 2 &&
+               pe->ne[3] == x->ne[2] && node->ne[0] == x->ne[0] && node->ne[1] == x->ne[2] &&
+               node->ne[2] == x->ne[1] * x->ne[3] && node->ne[3] == 1 && ggml_is_contiguous(node) &&
+               bounded(x, info) && bounded(pe, info) && bounded(node, info) && ggml_nelements(node) / 2 <= UINT32_MAX &&
+               (uint64_t(ggml_nelements(node) / 2) + 255) / 256 <= uint64_t(info.max_group_count[0]) * info.max_group_count[1];
+    }
+    static bool fused_supports(ggml_backend_dev_t device, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+        ggml_ext_vulkan_device_info info;
+        if (!api(device)->device_info(device, &info) || info.fusion_disabled || !info.subgroup_shuffle || !info.full_subgroups ||
+            info.subgroup_size != 32 || !is_rope(rope) || !supports(device, rope, nullptr) ||
+            norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || mul->src[0] != norm || rope->src[0] != mul)
+            return false;
+        const auto* x      = norm->src[0];
+        const auto* weight = mul->src[1];
+        return x && weight && x->type == GGML_TYPE_F32 && x->ne[0] == 128 && x->nb[0] == 4 &&
+               norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_are_same_shape(x, norm) && ggml_are_same_shape(norm, mul) &&
+               weight->type == GGML_TYPE_F32 && weight->ne[0] == 128 && ggml_nelements(weight) == 128 &&
+               ggml_is_contiguous(weight) && bounded(weight, info) && bounded(x, info) && (x->ne[1] + 3) / 4 <= info.max_group_count[0];
+    }
+    static bool cast_supports(ggml_backend_dev_t device, const ggml_tensor* rope, const ggml_tensor* scale, const ggml_tensor* copy) {
+        ggml_ext_vulkan_device_info info;
+        api(device)->device_info(device, &info);
+        // Workspace writeback uses Vulkan's four-byte copy granularity.
+        return info.float16 && scale->op == GGML_OP_SCALE && scale->src[0] == rope && scale->type == GGML_TYPE_F32 &&
+               param(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale && copy->type == GGML_TYPE_F16 &&
+               ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(scale, copy);
+    }
+    static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+        auto* g = const_cast<ggml_cgraph*>(graph);
+        if (first + 2 >= ggml_graph_n_nodes(g) || !fused_supports(device, ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
+            return 0;
+        return first + 4 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), ggml_graph_node(g, first + 3), ggml_graph_node(g, first + 4)) ? 5 : 3;
+    }
+    static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+        for (int last = 2; last < ggml_graph_n_nodes(graph); ++last) {
+            const auto* output = ggml_graph_node(graph, last);
+            const auto* scale  = output->op == GGML_OP_CPY ? output->src[0] : nullptr;
+            const auto* rope   = scale && scale->op == GGML_OP_SCALE ? scale->src[0] : output;
+            if (!is_rope(rope))
+                continue;
+            const auto* mul  = rope->src[0];
+            const auto* norm = mul->src[0];
+            if (!norm || !fused_supports(device, norm, mul, rope))
+                continue;
+            const int count = output == rope ? 3 : cast_supports(device, rope, scale, output) ? 5
+                                                                                              : 0;
+            if (!count)
+                continue;
+            const ggml_tensor* nodes[] = {norm, mul, rope, scale, output};
+            int indices[5], next = count - 1;
+            for (int i = last; i >= std::max(0, last - 64) && next >= 0; --i)
+                if (ggml_graph_node(graph, i) == nodes[next])
+                    indices[next--] = i;
+            if (next >= 0)
+                continue;
+            int first = indices[0];
+            for (int i = 1; i < count; ++i) {
+                if (!ggml_ext_graph_delay(graph, first, i, indices[i]))
+                    break;
+                first = indices[i] - i;
+            }
+        }
+    }
+    static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) {
+        const auto* output = nodes[count - 1];
+        // RoPE outputs and owning casts cannot reuse allocator input storage.
+        *bytes = output->data || output->view_src ? ggml_nbytes(output) : 0;
+        return GGML_STATUS_SUCCESS;
+    }
+    static void destroy(void*, void*) {}
+    static bool overlaps(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) {
+        return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size;
+    }
+    static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+        auto device         = ggml_backend_get_device(execution->backend);
+        const auto* service = api(device);
+        ggml_ext_vulkan_device_info info;
+        if (!service || !service->device_info(device, &info))
+            return GGML_STATUS_FAILED;
+        const auto* rope                   = nodes[count == 1 ? 0 : 2];
+        const auto* x                      = count == 1 ? rope->src[0] : nodes[0]->src[0];
+        const auto* pe                     = rope->src[1];
+        const auto* output                 = nodes[count - 1];
+        ggml_ext_vulkan_binding bindings[] = {{x}, {pe}, {output}, {count == 1 ? nullptr : nodes[1]->src[1]}};
+        ggml_ext_vulkan_buffer_info buffers[4];
+        const int n_buffers = count == 1 ? 3 : 4;
+        for (int i = 0; i < n_buffers; ++i)
+            if (!service->buffer_info(execution->backend, bindings + i, buffers + i))
+                return GGML_STATUS_FAILED;
+        bool alias             = overlaps(buffers[2], buffers[0]) || overlaps(buffers[2], buffers[1]) || (count != 1 && overlaps(buffers[2], buffers[3]));
+        const auto destination = bindings[2];
+        if (alias) {
+            bindings[2] = {nullptr, execution->workspace, 0, ggml_nbytes(output)};
+            if (!service->buffer_info(execution->backend, bindings + 2, buffers + 2))
+                return GGML_STATUS_FAILED;
+        }
+        FusedParameters params = {};
+        auto& pc               = params.rope;
+        pc.half_dim            = uint32_t(x->ne[0] / 2);
+        pc.rot_half            = uint32_t(pe->ne[2]);
+        pc.heads               = uint32_t(x->ne[1]);
+        pc.tokens              = uint32_t(x->ne[2]);
+        pc.pairs               = uint32_t(ggml_nelements(x) / 2);
+        for (int i = 0; i < 4; ++i) {
+            pc.x_stride[i]  = uint32_t(x->nb[i] / 4);
+            pc.pe_stride[i] = uint32_t(pe->nb[i] / 4);
+        }
+        pc.x_offset   = buffers[0].shader_offset / 4;
+        pc.pe_offset  = buffers[1].shader_offset / 4;
+        pc.dst_offset = buffers[2].shader_offset / uint32_t(ggml_type_size(output->type));
+        std::memcpy(&pc.interleaved, ggml_ext_op_attributes(rope), 4);
+        uint32_t groups[3]                = {1, 1, 1};
+        ggml_ext_vulkan_pipeline pipeline = {};
+        pipeline.name                     = "sd.rope_apply.vulkan";
+        pipeline.spirv                    = sd_vulkan_shaders::rope_apply;
+        pipeline.spirv_size               = sizeof(sd_vulkan_shaders::rope_apply);
+        pipeline.binding_count            = n_buffers;
+        pipeline.push_constant_size       = sizeof(Parameters);
+        if (count == 1) {
+            const uint32_t n_groups = uint32_t((uint64_t(pc.pairs) + 255) / 256);
+            groups[0]               = std::min(n_groups, info.max_group_count[0]);
+            groups[1]               = (n_groups + groups[0] - 1) / groups[0];
+        } else {
+            pipeline.name                   = count == 5 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan";
+            pipeline.spirv                  = count == 5 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply;
+            pipeline.spirv_size             = count == 5 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply);
+            pipeline.push_constant_size     = sizeof(FusedParameters);
+            pipeline.full_subgroups         = true;
+            pipeline.required_subgroup_size = info.subgroup_size_control ? 32 : 0;
+            params.eps                      = param(nodes[0], 0);
+            params.w_offset                 = buffers[3].shader_offset / 4;
+            params.output_scale             = count == 5 ? param(nodes[3], 0) : 1.f;
+            groups[0]                       = (pc.heads + 3) / 4;
+            groups[1]                       = std::min(pc.tokens, info.max_group_count[1]);
+            groups[2]                       = std::min(uint32_t(x->ne[3]), info.max_group_count[2]);
+        }
+        auto status = service->dispatch(execution->backend, &pipeline, bindings, &params, groups);
+        if (status == GGML_STATUS_SUCCESS && alias)
+            status = service->copy(execution->backend, bindings + 2, &destination);
+        return status;
+    }
+}
+
+void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) {
+    if (!sd_vulkan_rope::api(device))
+        return;
+    ggml_ext_kernel_desc desc = {};
+    desc.abi_version          = GGML_EXT_ABI_VERSION;
+    desc.name                 = "sd.rope_apply.vulkan.v1";
+    desc.op                   = GGML_OP_CUSTOM;
+    desc.custom_op            = sd_ggml_rope_op();
+    desc.priority             = 100;
+    desc.supports             = sd_vulkan_rope::supports;
+    desc.prepare              = sd_vulkan_rope::prepare;
+    desc.execute              = sd_vulkan_rope::execute;
+    desc.destroy              = sd_vulkan_rope::destroy;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    desc.name      = "sd.rms_rope_apply.vulkan.v1";
+    desc.op        = GGML_OP_RMS_NORM;
+    desc.custom_op = nullptr;
+    desc.priority  = 200;
+    desc.supports  = nullptr;
+    desc.match     = sd_vulkan_rope::match;
+    desc.optimize  = sd_vulkan_rope::optimize;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+}

src/core/vulkan/shaders/modulate_convrot.comp

@@ -0,0 +1,160 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+
+layout(constant_id = 0) const uint K      = 4096;
+layout(constant_id = 1) const bool AFFINE = false;
+layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
+layout(binding = 0) readonly buffer X {
+    float data_x[];
+};
+layout(binding = 1) readonly buffer M {
+    float data_m[];
+};
+layout(binding = 2) readonly buffer H {
+    float data_h[];
+};
+layout(binding = 3) writeonly buffer D {
+    uint data_d[];
+};
+layout(binding = 4) writeonly buffer S {
+    float data_s[];
+};
+#ifdef RMS_AFFINE
+layout(binding = 5) readonly buffer W {
+    float data_w[];
+};
+#endif
+layout(push_constant) uniform parameter {
+    uint rows;
+    uint x_offset;
+    uint m_offset;
+    uint h_offset;
+    float eps;
+    uint k;
+    uint w_offset;
+}
+p;
+shared float values[K];
+shared vec2 sums[512];
+
+uint quantize(float value, float inv_scale) {
+    return uint(clamp(int(roundEven(value * inv_scale)), -127, 127)) & 0xffu;
+}
+
+#ifdef RMS_AFFINE
+// Keep RMS arithmetic outside the scope of precise on the quantized output.
+float rms_row(uint row, uint tid) {
+    float sum = 0.0;
+    [[unroll]] for (uint col = tid, iter = 0; iter < 12; col += 512, ++iter) {
+        float x = 0.0;
+        if (col < p.k) {
+            x           = data_x[p.x_offset + row * K + col];
+            values[col] = x;
+        }
+        sum += x * x;
+    }
+    sums[tid].x = sum;
+    barrier();
+    [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            sum += sums[tid + stride].x;
+            sums[tid].x = sum;
+        }
+        barrier();
+    }
+    return inversesqrt(sums[0].x / p.k + p.eps);
+}
+
+float rms_weight(float scale, float value, float weight) {
+    return scale * value * weight;
+}
+#endif
+
+void main() {
+    const uint tid = gl_LocalInvocationID.x;
+    for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
+#ifdef RMS_AFFINE
+        float inv_std = rms_row(row, tid);
+        float mean    = 0.0;
+#else
+        vec2 sum = vec2(0.0);
+        [[unroll]] for (uint col = tid; col < K; col += 512) {
+            const float x = data_x[p.x_offset + row * K + col];
+            values[col]   = x;
+            if (!AFFINE) {
+                sum.x += x;
+                sum.y = fma(x, x, sum.y);
+            }
+        }
+        float mean    = 0.0;
+        float inv_std = 1.0;
+        if (!AFFINE) {
+            sums[tid] = sum;
+            barrier();
+            [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) {
+                if (tid < stride)
+                    sums[tid] += sums[tid + stride];
+                barrier();
+            }
+            mean = sums[0].x / p.k;
+            // Match NORM's contraction even when precise propagates from the quantized output.
+            const float variance = fma(-mean, mean, sums[0].y / p.k);
+            inv_std              = inversesqrt(variance + p.eps);
+        }
+#endif
+        [[unroll]] for (uint col = tid; col < K; col += 512) {
+            precise float value;
+#ifdef RMS_AFFINE
+            values[col] = rms_weight(inv_std, values[col], data_w[p.w_offset + col]);
+#endif
+            if (AFFINE) {
+                precise float modulated = values[col] * data_m[p.m_offset + col];
+                value                   = values[col] + modulated;
+                value += data_h[p.h_offset + col];
+            } else {
+                value = (values[col] - mean) * inv_std;
+                value *= data_m[p.m_offset + col];
+            }
+            values[col] = value * (1.0 / 16.0);
+        }
+        barrier();
+        [[unroll]] for (uint stride = 1; stride < 256; stride *= 4) {
+            for (uint i = tid; i < K / 4; i += 512) {
+                const uint group          = i / 64;
+                const uint lane           = i % 64;
+                const uint base           = group * 256 + (lane / stride) * 4 * stride + lane % stride;
+                const float a             = values[base];
+                const float b             = values[base + stride];
+                const float c             = values[base + 2 * stride];
+                const float d             = values[base + 3 * stride];
+                values[base]              = a + b + c - d;
+                values[base + stride]     = a + b - c + d;
+                values[base + 2 * stride] = a - b + c + d;
+                values[base + 3 * stride] = -a + b + c + d;
+            }
+            barrier();
+        }
+        float maximum = 0.0;
+        [[unroll]] for (uint col = tid; col < K; col += 512) {
+            maximum = max(maximum, abs(values[col]));
+        }
+        sums[tid].x = maximum;
+        barrier();
+        [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) {
+            if (tid < stride) {
+                sums[tid].x = max(sums[tid].x, sums[tid + stride].x);
+            }
+            barrier();
+        }
+        const float scale     = sums[0].x / 127.0;
+        const float inv_scale = scale == 0.0 ? 0.0 : 1.0 / scale;
+        if (tid == 0) {
+            data_s[row] = scale;
+        }
+        for (uint i = tid; i < K / 4; i += 512) {
+            data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24);
+        }
+        barrier();
+    }
+}

src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp

@@ -0,0 +1,74 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_integer_dot_product : require
+
+layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+    uint row_offset;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    float data_d[];
+};
+
+shared int weights[16 * 8];
+shared int activations[16 * 8];
+
+void main() {
+    const uint output_index = gl_WorkGroupID.x * 16 + gl_LocalInvocationID.x;
+    const uint row          = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y;
+    const uint lx           = gl_LocalInvocationID.x;
+    const uint ly           = gl_LocalInvocationID.y;
+    const uint k_packed     = p.k / 4;
+    int sum                 = 0;
+
+    for (uint k_base = 0; k_base < k_packed; k_base += 8) {
+        if (ly < 8) {
+            weights[lx * 8 + ly] = output_index < p.n
+                                       ? data_w[output_index * k_packed + k_base + ly]
+                                       : 0;
+        }
+        if (lx < 8) {
+            activations[ly * 8 + lx] = row < p.rows
+                                           ? data_a[row * k_packed + k_base + lx]
+                                           : 0;
+        }
+        barrier();
+
+        [[unroll]]
+        for (uint i = 0; i < 8; ++i) {
+            sum += dotPacked4x8EXT(weights[lx * 8 + i], activations[ly * 8 + i]);
+        }
+        barrier();
+    }
+
+    if (output_index < p.n && row < p.rows) {
+        float value = float(sum) * data_as[row] * data_ws[output_index];
+        if (p.has_bias != 0) {
+            value += data_b[output_index];
+        }
+        data_d[row * p.n + output_index] = value;
+    }
+}

src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp

@@ -0,0 +1,106 @@
+#version 460
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_shader_8bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_KHR_cooperative_matrix : require
+#extension GL_KHR_memory_scope_semantics : require
+#extension GL_KHR_shader_subgroup_basic : require
+
+layout(constant_id = 0) const uint BLOCK_SIZE = 32;
+layout(constant_id = 1) const uint CM_M       = 16;
+layout(constant_id = 2) const uint CM_N       = 16;
+layout(constant_id = 3) const uint CM_K       = 16;
+const uint TILE_M                             = 4;
+const uint TILE_N                             = 4;
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+    uint row_offset;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int8_t data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int8_t data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 5, std430) writeonly buffer D {
+    float data_d[];
+};
+
+shared int32_t accumulators[1024];
+
+void main() {
+    const uint output_base = gl_WorkGroupID.x * CM_N * TILE_N;
+    const uint row_base    = gl_WorkGroupID.y * CM_M * TILE_M;
+    coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator> sums[TILE_M * TILE_N];
+    [[unroll]]
+    for (uint i = 0; i < TILE_M * TILE_N; ++i) {
+        sums[i] = coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator>(0);
+    }
+
+    for (uint k_base = 0; k_base < p.k; k_base += CM_K) {
+        coopmat<int8_t, gl_ScopeSubgroup, CM_M, CM_K, gl_MatrixUseA> activations[TILE_M];
+        coopmat<int8_t, gl_ScopeSubgroup, CM_K, CM_N, gl_MatrixUseB> weights[TILE_N];
+        [[unroll]]
+        for (uint tile_m = 0; tile_m < TILE_M; ++tile_m) {
+            coopMatLoad(activations[tile_m], data_a,
+                        (row_base + tile_m * CM_M) * p.k + k_base,
+                        p.k, gl_CooperativeMatrixLayoutRowMajor);
+        }
+        [[unroll]]
+        for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) {
+            coopMatLoad(weights[tile_n], data_w,
+                        (output_base + tile_n * CM_N) * p.k + k_base,
+                        p.k, gl_CooperativeMatrixLayoutColumnMajor);
+        }
+        [[unroll]]
+        for (uint tile_m = 0; tile_m < TILE_M; ++tile_m) {
+            [[unroll]]
+            for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) {
+                const uint index = tile_m * TILE_N + tile_n;
+                sums[index]      = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]);
+            }
+        }
+    }
+
+    [[unroll]]
+    for (uint tile_m = 0; tile_m < TILE_M; ++tile_m) {
+        [[unroll]]
+        for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) {
+            coopMatStore(sums[tile_m * TILE_N + tile_n], accumulators,
+                         0, CM_N, gl_CooperativeMatrixLayoutRowMajor);
+            controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup,
+                           gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
+
+            for (uint index = gl_SubgroupInvocationID; index < CM_M * CM_N; index += gl_SubgroupSize) {
+                const uint row          = row_base + tile_m * CM_M + index / CM_N;
+                const uint output_index = output_base + tile_n * CM_N + index % CM_N;
+                float value             = float(accumulators[index]) * data_as[row] * data_ws[output_index];
+                if (p.has_bias != 0) {
+                    value += data_b[output_index];
+                }
+                data_d[row * p.n + output_index] = value;
+            }
+            controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup,
+                           gl_StorageSemanticsShared, gl_SemanticsAcquireRelease);
+        }
+    }
+}

src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp

@@ -0,0 +1,172 @@
+#version 460
+
+#extension GL_EXT_shader_8bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require
+#extension GL_KHR_cooperative_matrix : require
+#extension GL_NV_cooperative_matrix2 : require
+#extension GL_KHR_memory_scope_semantics : require
+
+layout(constant_id = 0) const uint BN        = 128;
+layout(constant_id = 1) const uint ALIGNMENT = 0;
+layout(constant_id = 2) const bool SWIGLU    = false;
+layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
+layout(constant_id = 3) const uint BM = 128;
+layout(constant_id = 4) const uint BK = 64;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+    uint row_offset;
+    float output_scale;
+    uint group_cols;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int8_t data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int8_t data_a[];
+};
+layout(binding = 2, std430) readonly buffer WS {
+    float data_ws[];
+};
+layout(binding = 3, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 4, std430) readonly buffer B {
+    float data_b[];
+};
+// Store FP32 bit patterns without changing the integer accumulator fragment layout.
+layout(binding = 5, std430) writeonly buffer D {
+    int32_t data_d[];
+};
+#ifdef SPLIT_SWIGLU
+layout(binding = 6, std430) readonly buffer UW {
+    int8_t data_uw[];
+};
+layout(binding = 7, std430) readonly buffer US {
+    float data_us[];
+};
+#endif
+#ifdef RESIDUAL
+layout(binding = 6, std430) readonly buffer G {
+    float data_g[];
+};
+layout(binding = 7, std430) readonly buffer R {
+    float data_r[];
+};
+#endif
+int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
+    const uint r = row_base + col;
+    const uint c = col_base + row;
+    if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) {
+        return 0;
+    }
+    precise float value = float(acc) * data_as[r] * data_ws[c];
+    if (p.has_bias != 0) {
+        value += data_b[c];
+    }
+#ifdef RESIDUAL
+    value *= p.output_scale;
+    value *= data_g[c];
+    value += data_r[r * p.n + c];
+#endif
+    return floatBitsToInt(value);
+}
+
+int32_t dequantize_swiglu(const in uint row, const in uint col, const in int32_t gate_acc, const in int32_t up_acc, const in uint row_base, const in uint col_base) {
+    const uint r = row_base + col;
+    const uint c = col_base + row;
+    if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n / 2)) {
+        return 0;
+    }
+    precise float gate = float(gate_acc) * data_as[r] * data_ws[c];
+#ifdef SPLIT_SWIGLU
+    precise float up = float(up_acc) * data_as[r] * data_us[c];
+#else
+    precise float up = float(up_acc) * data_as[r] * data_ws[c + p.n / 2];
+#endif
+    return floatBitsToInt(gate / (1.0 + exp(-gate)) * up);
+}
+
+void main() {
+    const uint k_size                                                   = p.k & ~255u;
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> a_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> w_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    a_layout                                                            = setTensorLayoutDimensionNV(a_layout, p.rows, k_size);
+    w_layout                                                            = setTensorLayoutDimensionNV(w_layout, p.n, k_size);
+    d_layout                                                            = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n);
+    tensorViewNV<2, false, 1, 0> transpose                              = createTensorViewNV(2, false, 1, 0);
+    // Keep a weight panel in L2 while visiting the activation tiles.
+    const uint col_tiles                                                        = gl_NumWorkGroups.x;
+    const uint row_tiles                                                        = gl_NumWorkGroups.y;
+    const uint group_cols                                                       = clamp(p.group_cols, 1u, col_tiles);
+    const uint linear                                                           = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x;
+    const uint group                                                            = linear / (group_cols * row_tiles);
+    const uint first_col                                                        = group * group_cols;
+    const uint cols                                                             = min(col_tiles - first_col, group_cols);
+    const uint within_group                                                     = linear - group * group_cols * row_tiles;
+    const uint row                                                              = p.row_offset + (within_group / cols) * BN;
+    const uint col                                                              = (first_col + within_group % cols) * BM;
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum    = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum_up = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    if (ALIGNMENT != 0) {
+        tensorLayoutNV<2> a_full = createTensorLayoutNV(2);
+        tensorLayoutNV<2> w_full = createTensorLayoutNV(2);
+        a_full                   = setTensorLayoutDimensionNV(a_full, p.rows, k_size);
+        w_full                   = setTensorLayoutDimensionNV(w_full, p.n, k_size);
+        for (uint k = 0; k < p.k; k += BK) {
+            coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+            coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
+            coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_full, col, BM, k, BK));
+            if (ALIGNMENT == 1) {
+                coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(a_full, row, BN, k, BK), transpose);
+            } else {
+                coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(a_layout, row, BN, k, BK), transpose);
+            }
+            sum = coopMatMulAdd(w, a, sum);
+            if (SWIGLU) {
+#ifdef SPLIT_SWIGLU
+                coopMatLoadTensorNV(w, data_uw, 0, sliceTensorLayoutNV(w_full, col, BM, k, BK));
+#else
+                coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_full, col + p.n / 2, BM, k, BK));
+#endif
+                sum_up = coopMatMulAdd(w, a, sum_up);
+            }
+        }
+    } else {
+        for (uint k = 0; k < p.k; k += BK) {
+            coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+            coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
+            coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_layout, col, BM, k, BK));
+            coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(a_layout, row, BN, k, BK), transpose);
+            sum = coopMatMulAdd(w, a, sum);
+            if (SWIGLU) {
+#ifdef SPLIT_SWIGLU
+                coopMatLoadTensorNV(w, data_uw, 0, sliceTensorLayoutNV(w_layout, col, BM, k, BK));
+#else
+                coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_layout, col + p.n / 2, BM, k, BK));
+#endif
+                sum_up = coopMatMulAdd(w, a, sum_up);
+            }
+        }
+    }
+    if (SWIGLU) {
+        coopMatPerElementNV(sum, sum, dequantize_swiglu, sum_up, row, col);
+    } else {
+        coopMatPerElementNV(sum, sum, dequantize, row, col);
+    }
+    // The aligned dispatch contains only complete tiles and a 128-aligned output width.
+    if (ALIGNMENT == 1) {
+        tensorLayoutNV<2> d_full = createTensorLayoutNV(2);
+        d_full                   = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u);
+        coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_full, row, BN, col, BM), transpose);
+    } else {
+        coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row, BN, col, BM), transpose);
+    }
+}

src/core/vulkan/shaders/quantize_i8_convrot.comp

@@ -0,0 +1,105 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+
+layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint k;
+    uint rows;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer A {
+    float data_a[];
+};
+layout(binding = 1, std430) writeonly buffer D {
+    uint data_d[];
+};
+layout(binding = 2, std430) writeonly buffer S {
+    float data_s[];
+};
+
+shared float values[256];
+shared float maxima[64];
+
+void transform_group(const uint row, const uint group) {
+    const uint tid         = gl_LocalInvocationID.x;
+    const uint base_offset = row * p.k + group * 256;
+
+    [[unroll]]
+    for (uint i = 0; i < 4; ++i) {
+        const uint index = tid + i * 64;
+        values[index]    = data_a[base_offset + index] * (1.0 / 16.0);
+    }
+    barrier();
+
+    [[unroll]]
+    for (uint stride = 1; stride < 256; stride *= 4) {
+        const uint base = (tid / stride) * 4 * stride + tid % stride;
+        const uint i0   = base;
+        const uint i1   = i0 + stride;
+        const uint i2   = i1 + stride;
+        const uint i3   = i2 + stride;
+        const float a   = values[i0];
+        const float b   = values[i1];
+        const float c   = values[i2];
+        const float d   = values[i3];
+        values[i0]      = a + b + c - d;
+        values[i1]      = a + b - c + d;
+        values[i2]      = a - b + c + d;
+        values[i3]      = -a + b + c + d;
+        barrier();
+    }
+}
+
+uint quantize_value(const float value, const float inv_scale) {
+    const int quantized = clamp(int(roundEven(value * inv_scale)), -127, 127);
+    return uint(quantized) & 0xffu;
+}
+
+void quantize_row(const uint row) {
+    const uint tid    = gl_LocalInvocationID.x;
+    const uint groups = p.k / 256;
+    float local_max   = 0.0;
+
+    for (uint group = 0; group < groups; ++group) {
+        transform_group(row, group);
+        [[unroll]]
+        for (uint i = 0; i < 4; ++i) {
+            local_max = max(local_max, abs(values[tid + i * 64]));
+        }
+        barrier();
+    }
+
+    maxima[tid] = local_max;
+    barrier();
+    [[unroll]]
+    for (uint stride = 32; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            maxima[tid] = max(maxima[tid], maxima[tid + stride]);
+        }
+        barrier();
+    }
+
+    if (tid == 0) {
+        data_s[row] = maxima[0] / 127.0;
+    }
+    barrier();
+
+    const float row_scale = maxima[0] / 127.0;
+    const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
+    for (uint group = 0; group < groups; ++group) {
+        transform_group(row, group);
+        const uint index                            = tid * 4;
+        const uint packed                           = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
+        data_d[(row * p.k + group * 256) / 4 + tid] = packed;
+        barrier();
+    }
+}
+
+void main() {
+    for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
+        quantize_row(row);
+    }
+}

src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp

@@ -0,0 +1,140 @@
+#version 450
+
+#extension GL_EXT_shader_8bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_shuffle : require
+#extension GL_KHR_shader_subgroup_arithmetic : require
+#extension GL_EXT_control_flow_attributes : require
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+layout(push_constant) uniform parameter {
+    uint k;
+    uint rows;
+}
+p;
+layout(binding = 0, std430) readonly buffer A {
+    float data_a[];
+};
+layout(binding = 1, std430) writeonly buffer D {
+    int8_t data_d[];
+};
+layout(binding = 2, std430) writeonly buffer S {
+    float data_s[];
+};
+#if defined(SWIGLU) || defined(SIGMOID_MUL)
+layout(binding = 3, std430) readonly buffer B {
+    float data_b[];
+};
+#endif
+shared float maxima[8];
+#ifdef SINGLE_PASS
+#ifdef STATIC_K
+const uint K = STATIC_K;
+#else
+layout(constant_id = 0) const uint K = 4096;
+#endif
+#endif
+
+vec4 rotate(float a, float b, float c, float d) {
+    precise vec4 r = vec4(a + b + c - d, a + b - c + d,
+                          a - b + c + d, -a + b + c + d);
+    return r;
+}
+
+void transform_group(uint row, uint group, out float v[8]) {
+    const uint lane = gl_SubgroupInvocationID;
+    const uint base = row * p.k + group * 256;
+    [[unroll]] for (uint j = 0; j < 8; ++j) {
+        float value = data_a[base + lane + j * 32];
+#ifdef SWIGLU
+        precise float activated = value / (1.0 + exp(-value)) * data_b[base + lane + j * 32];
+        value                   = activated;
+#endif
+#ifdef SIGMOID_MUL
+        precise float gate      = 1.0 / (1.0 + exp(-value));
+        precise float activated = data_b[base + lane + j * 32] * gate;
+        value                   = activated;
+#endif
+        v[j] = value * (1.0 / 16.0);
+    }
+    [[unroll]] for (uint stride = 1; stride <= 4; stride *= 4) {
+        const uint first = lane & ~(3u * stride);
+        [[unroll]] for (uint j = 0; j < 8; ++j) {
+            vec4 r = rotate(subgroupShuffle(v[j], first), subgroupShuffle(v[j], first + stride),
+                            subgroupShuffle(v[j], first + 2 * stride), subgroupShuffle(v[j], first + 3 * stride));
+            v[j]   = r[(lane / stride) & 3u];
+        }
+    }
+    [[unroll]] for (uint j = 0; j < 8; j += 2) {
+        vec4 r   = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16),
+                          subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16));
+        v[j]     = r[lane / 16];
+        v[j + 1] = r[2 + lane / 16];
+    }
+    [[unroll]] for (uint j = 0; j < 2; ++j) {
+        vec4 r   = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]);
+        v[j]     = r.x;
+        v[j + 2] = r.y;
+        v[j + 4] = r.z;
+        v[j + 6] = r.w;
+    }
+}
+
+void main() {
+    for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
+        float local_max = 0.0;
+#ifdef SINGLE_PASS
+        float values[((K + 2047) / 2048) * 8];
+        [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
+            if (gl_SubgroupID + g * 8 >= K / 256)
+                continue;
+            float v[8];
+            transform_group(row, gl_SubgroupID + g * 8, v);
+            [[unroll]] for (uint j = 0; j < 8; ++j) {
+                values[g * 8 + j] = v[j];
+                local_max         = max(local_max, abs(v[j]));
+            }
+        }
+#else
+        for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
+            float v[8];
+            transform_group(row, group, v);
+            [[unroll]] for (uint j = 0; j < 8; ++j)
+                local_max = max(local_max, abs(v[j]));
+        }
+#endif
+        local_max = subgroupMax(local_max);
+        if (gl_SubgroupInvocationID == 0)
+            maxima[gl_SubgroupID] = local_max;
+        barrier();
+        if (gl_LocalInvocationID.x == 0) {
+            float row_max = 0.0;
+            for (uint j = 0; j < 8; ++j)
+                row_max = max(row_max, maxima[j]);
+            maxima[0]   = row_max;
+            data_s[row] = row_max / 127.0;
+        }
+        barrier();
+        const float row_scale = maxima[0] / 127.0;
+        const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
+#ifdef SINGLE_PASS
+        [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
+            if (gl_SubgroupID + g * 8 >= K / 256)
+                continue;
+            [[unroll]] for (uint j = 0; j < 8; ++j) {
+                data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127));
+            }
+        }
+#else
+        for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
+            float v[8];
+            transform_group(row, group, v);
+            [[unroll]] for (uint j = 0; j < 8; ++j) {
+                data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127));
+            }
+        }
+#endif
+        barrier();
+    }
+}

src/core/vulkan/shaders/rms_rope_apply.comp

@@ -0,0 +1,108 @@
+#version 450
+
+#extension GL_KHR_shader_subgroup_basic : require
+#extension GL_KHR_shader_subgroup_shuffle : require
+#extension GL_EXT_control_flow_attributes : require
+
+layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
+layout(binding = 0) readonly buffer X {
+    float data_x[];
+};
+layout(binding = 1) readonly buffer PE {
+    float data_pe[];
+};
+#ifdef CAST_F16
+#extension GL_EXT_shader_16bit_storage : require
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+layout(binding = 2) writeonly buffer D {
+    float16_t data_d[];
+};
+#else
+layout(binding = 2) writeonly buffer D {
+    float data_d[];
+};
+#endif
+layout(binding = 3) readonly buffer W {
+    float data_w[];
+};
+
+layout(push_constant) uniform parameter {
+    uint half_dim;
+    uint rot_half;
+    uint heads;
+    uint tokens;
+    uint pairs;
+    uint x_stride[4];
+    uint pe_stride[4];
+    uint x_offset;
+    uint pe_offset;
+    uint dst_offset;
+    uint interleaved;
+    float eps;
+    uint w_offset;
+    float output_scale;
+}
+pc;
+
+float shuffle_channel(vec4 values, uint channel) {
+    const uint lane     = channel & 31u;
+    const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane),
+                               subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane));
+    return shuffled[channel / 32];
+}
+
+void main() {
+    const uint head = gl_WorkGroupID.x * 4 + gl_SubgroupID;
+    if (head >= pc.heads) {
+        return;
+    }
+    const uint lane    = gl_SubgroupInvocationID;
+    const uint batches = pc.pairs / (64 * pc.heads * pc.tokens);
+    for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) {
+        for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) {
+            const uint base      = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
+            vec4 x               = vec4(data_x[base + lane], data_x[base + lane + 32],
+                                        data_x[base + lane + 64], data_x[base + lane + 96]);
+            precise vec4 squares = x * x;
+            precise float sum    = (squares.x + squares.z) + (squares.y + squares.w);
+            [[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) {
+                sum += subgroupShuffle(sum, lane ^ stride);
+            }
+            sum                   = subgroupShuffle(sum, 0);
+            const float scale     = inversesqrt(sum / 128.0 + pc.eps);
+            precise vec4 weighted = scale * x;
+            [[unroll]] for (uint i = 0; i < 4; ++i) {
+                weighted[i] *= data_w[pc.w_offset + lane + i * 32];
+            }
+            [[unroll]] for (uint i = 0; i < 4; ++i) {
+                const uint col      = lane + i * 32;
+                precise float value = weighted[i];
+                if (pc.interleaved != 0) {
+                    const float a = subgroupShuffle(value, lane & ~1u);
+                    const float b = subgroupShuffle(value, lane | 1u);
+                    if (col / 2 < pc.rot_half) {
+                        const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1];
+                        value             = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]];
+                    }
+                } else {
+                    const uint pair = col % pc.rot_half;
+                    const float a   = shuffle_channel(weighted, pair);
+                    const float b   = shuffle_channel(weighted, pair + pc.rot_half);
+                    if (col < 2 * pc.rot_half) {
+                        const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3];
+                        const float c     = data_pe[matrix];
+                        const float s     = data_pe[matrix + pc.pe_stride[1]];
+                        value             = col < pc.rot_half ? a * c - b * s : a * s + b * c;
+                    }
+                }
+#ifdef CAST_F16
+                value *= pc.output_scale;
+                value += 0.0;
+                data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = float16_t(value);
+#else
+                data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = value;
+#endif
+            }
+        }
+    }
+}

src/core/vulkan/shaders/rope_apply.comp

@@ -0,0 +1,66 @@
+#version 450
+
+layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
+
+layout(binding = 0) readonly buffer X {
+    float data_x[];
+};
+layout(binding = 1) readonly buffer PE {
+    float data_pe[];
+};
+layout(binding = 2) writeonly buffer D {
+    float data_d[];
+};
+
+layout(push_constant) uniform parameter {
+    uint half_dim;
+    uint rot_half;
+    uint heads;
+    uint tokens;
+    uint pairs;
+    uint x_stride[4];
+    uint pe_stride[4];
+    uint x_offset;
+    uint pe_offset;
+    uint dst_offset;
+    uint interleaved;
+}
+pc;
+
+void main() {
+    const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y;
+    if (index >= pc.pairs) {
+        return;
+    }
+    const uint j       = index % pc.half_dim;
+    const uint row     = index / pc.half_dim;
+    const uint token   = row % pc.tokens;
+    const uint head    = (row / pc.tokens) % pc.heads;
+    const uint batch   = row / (pc.tokens * pc.heads);
+    const uint i0      = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j : j;
+    const uint i1      = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j + 1 : j + pc.rot_half;
+    const uint x_row   = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
+    const uint out_row = pc.dst_offset + row * pc.half_dim * 2;
+    const float a      = data_x[x_row + i0 * pc.x_stride[0]];
+    const float b      = data_x[x_row + i1 * pc.x_stride[0]];
+    if (j >= pc.rot_half) {
+        data_d[out_row + i0] = a;
+        data_d[out_row + i1] = b;
+        return;
+    }
+    const uint matrix = pc.pe_offset + j * pc.pe_stride[2] + token * pc.pe_stride[3];
+    const float c     = data_pe[matrix];
+    const float s     = data_pe[matrix + pc.pe_stride[1]];
+    // Preserve the separate multiply/add rounding of the unfused graph.
+    precise float y0;
+    precise float y1;
+    if (pc.interleaved != 0) {
+        y0 = a * c + b * data_pe[matrix + pc.pe_stride[0]];
+        y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
+    } else {
+        y0 = a * c - b * s;
+        y1 = a * s + b * c;
+    }
+    data_d[out_row + i0] = y0;
+    data_d[out_row + i1] = y1;
+}