refactor: register INT8 projection and attention fusion

CMakeLists.txt

@@ -348,12 +348,15 @@ if (SD_CUDA)
         target_sources(${SD_LIB} PRIVATE src/core/cuda/rope.cu)
         set_source_files_properties(src/core/cuda/rope.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math")
         if (GGML_CUDA_CUTLASS_PATH)
-            target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu)
+            target_sources(${SD_LIB} PRIVATE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu)
-            set_source_files_properties(src/core/cuda/cute_attention.cu PROPERTIES
+            set_source_files_properties(src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu PROPERTIES
                 INCLUDE_DIRECTORIES "${GGML_CUDA_CUTLASS_PATH}/include"
                 COMPILE_OPTIONS "--use_fast_math;--extended-lambda;--expt-relaxed-constexpr")
             if (MSVC)
-                set_property(SOURCE src/core/cuda/cute_attention.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+                set_property(SOURCE src/core/cuda/cute_attention.cu src/core/cuda/i8_projection.cu APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus)
+            endif()
+            if (GGML_CUDA_FORCE_CUBLAS)
+                set_property(SOURCE src/core/cuda/i8_projection.cu APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_FORCE_CUBLAS)
             endif()
             set_property(SOURCE src/core/ggml_extensions.cpp APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION)
         endif()

docs/ggml_extensions.md

@@ -67,16 +67,20 @@ compute flags, external uses, output flags, and views; the initial implementatio
 accepts a conservative set of pure operators with only the final node exposed
 as the fused output. The rule remains responsible for mathematical equivalence,
 dtype, shape, layout, and hardware constraints. Internal views are accepted when
-their storage belongs to the fused subgraph. The final output must own storage,
+their storage belongs to the fused subgraph. In-place intermediates are accepted
-except for a terminal `CPY`, which retains its original destination and write
+only when their storage owner is an earlier node in the same match. External
+consumers and output flags still prevent eliminating those results. The final
+output must own storage, except for a terminal `CPY`, which retains its original destination and write
 semantics. The extension must handle input/output storage overlap, including
 allocator reuse, with scratch space when necessary. Other in-place writes are
-not eligible for extension fusion.
+not eligible for extension fusion when they target storage outside the match.
 
 An optional `optimize` callback runs before matching and allocation. Use
 `ggml_ext_graph_delay` to move a producer range after independent nodes. The
 helper rejects dependencies, unknown side effects, and writes that alias moved
-inputs or outputs. It can cross independent copies, including dtype conversions.
+inputs or outputs. It can cross independent copies and known in-place operators
+when their storage does not alias the moved inputs or outputs. Owning dtype
+conversions can also be moved or crossed.
 Reordering must preserve edges and tensor metadata. Fusion legality is checked
 again after reordering.
 
@@ -140,12 +144,21 @@ CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or
 private conversion helpers. The existing CUDA fusion and CuTe disable environment
 variables are honored.
 
-The combined INT8 GEMM epilogue feeding FP16 Q directly into attention still needs
+`src/core/cuda/i8_projection.cu` supplies the INT8 GEMM projection epilogues:
-to be externalized. The current extension can fuse RMSNorm/RoPE with attention,
+`GEMM -> Reshape -> RMSNorm -> Mul -> RoPE`, optional scaling and FP16 copy,
-while other INT8 GEMM and standalone projection fusions retain their native
+and K/V prefix concatenation with FP16 output. CuTe can consume a directly
-implementations. Native INT8 projection matchers do not recognize custom RoPE
+produced FP16 Q without materializing the FP32 GEMM or RoPE intermediates.
-yet, so those graphs currently run GEMM followed by the external RoPE/attention
+Independent affine operands, including broadcasted constants, are computed before
-path. Restoring their combined epilogues is part of the remaining migration.
+their normalization/projection
-Vulkan execution services and kernels also remain pending.
+chain when graph dependencies allow it. The attention optimizer can then move
+the complete five-node projection when safe.
+These rules accept both builtin and custom RoPE nodes. They preserve the original
+large-matrix selection thresholds, full interleaved head dimension 128, packed
+activation layout, and separate weight scales. External consumers prevent fusion;
+aliasing terminal copies use scheduler workspace. The CUDA projection kernels
+and matchers now live in sd.cpp. `GGML_CUDA_FORCE_CUBLAS` disables these epilogues.
+
+Ordinary INT8 GEMM, ConvRot quantization, and their other fusion rules still need
+to be externalized. Vulkan execution services and kernels also remain pending.
 This implementation does not yet externalize all model-specific kernels or add
 dynamic dtypes.

src/core/cuda/cute_attention.cu

@@ -44,6 +44,7 @@
 #include <mutex>
 #include <type_traits>
 #include <unordered_map>
+#include "core/cuda/i8_projection.h"
 #include "core/cuda/rope.h"
 #include "ggml-ext.h"
 
@@ -578,7 +579,14 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
     if (!fusion_enabled())
         return 0;
     int offset = 0;
-    if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
+    if (graph_node(graph, first)->op == GGML_OP_MUL_MAT) {
+        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))
@@ -601,6 +609,9 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo
         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;
             }
@@ -616,7 +627,8 @@ struct AttentionPlan {
 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 = nodes[0]->op == GGML_OP_RMS_NORM ? 3 : 0;
+    plan->index = nodes[0]->op == GGML_OP_MUL_MAT ? 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;
@@ -671,7 +683,10 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
     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 == 3) {
+        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) {
@@ -766,4 +781,8 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
     desc.preferred = nullptr;
     desc.optimize  = nullptr;
     GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+    desc.name     = "sd.i8_rope_cute_attention.cuda.v1";
+    desc.op       = GGML_OP_MUL_MAT;
+    desc.priority = 400;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
 }

src/core/cuda/i8_projection.cu

@@ -0,0 +1,389 @@
+#include "core/cuda/i8_projection.h"
+#include "core/cuda/rope.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) {
+    const auto* w  = mat->src[0];
+    const auto* x  = mat->src[1];
+    const auto* ws = mat->src[2];
+    if (mat->op != GGML_OP_MUL_MAT || mat->type != GGML_TYPE_F32 || mat->op_params[4] <= 0 || !w || !x || !ws || mat->src[3] ||
+        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_get(rope) ? ggml_ext_op_attributes(rope) : rope->op_params, 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_MUL_MAT;
+    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(device, &desc));
+}

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/rope.cu

@@ -235,6 +235,40 @@ static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, v
     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 || !sd_cuda_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 (mat->op == GGML_OP_MUL_MAT && mat->src[0]->type == GGML_TYPE_I8 &&
+                        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;
@@ -298,8 +332,9 @@ void sd_ggml_register_cuda_rope(ggml_backend_dev_t device) {
     GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
     desc.name     = "sd.rms_norm_mul_rope.cuda.v1";
     desc.op       = GGML_OP_RMS_NORM;
-    desc.priority = 200;
+    desc.priority = 250;
     desc.supports = nullptr;
     desc.match    = rope_match;
+    desc.optimize = rope_optimize;
     GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
 }

src/core/ggml_extensions.cpp

@@ -12,6 +12,7 @@
 void sd_ggml_register_cuda_rope(ggml_backend_dev_t device);
 #endif
 #ifdef SD_CUDA_CUTE_ATTENTION
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
 void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
 #endif
 
@@ -180,6 +181,7 @@ void sd_ggml_register_extensions(ggml_backend_t backend) {
         sd_ggml_register_cuda_rope(device);
 #endif
 #ifdef SD_CUDA_CUTE_ATTENTION
+        sd_ggml_register_cuda_i8_projection(device);
         sd_ggml_register_cuda_attention(device);
 #endif
         return;