refactor: register CUDA INT8 GEMM and epilogue extensions

src/core/cuda/convrot.cu

@@ -145,11 +145,6 @@ namespace sd_convrot {
         ggml_tensor* nodes[6];
         for (int j = 0; j < count; ++j)
             nodes[j] = ggml_graph_node(g, first + j);
-        // Native INT8 epilogues consume GEMM and GLU together; splitting them loses that fusion.
-        if (node->op == GGML_OP_GLU && nodes[count - 1]->op == GGML_OP_QUANTIZE_I8_CONVROT &&
-            node->src[0]->op == GGML_OP_MUL_MAT && node->src[0]->src[0]->type == GGML_TYPE_I8 &&
-            node->src[0]->src[2] && !node->src[0]->src[3])
-            return 0;
         Operation op;
         return describe(nodes, count, op) && supported(op, device_info(device)) ? count : 0;
     }

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

@@ -1,4 +1,7 @@
+#include <cstdlib>
+#include "convrot_kernel.h"
 #include "core/int8_ops.h"
+#include "i8_gemm.h"
 
 #include <cublas_v2.h>
 #include <cuda_runtime.h>
@@ -301,7 +304,7 @@ namespace sd_i8_cuda {
         int64_t n,
         int64_t rows) {
         const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
-        const int64_t row    = blockIdx.y;
+        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];
@@ -322,14 +325,16 @@ namespace sd_i8_cuda {
         const float4* weight_scales,
         const float4* bias,
         float4* dst,
-        int64_t n4) {
+        int64_t n4,
+        int64_t rows) {
         const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
-        if (output >= n4) {
+        const int64_t row    = blockIdx.y + int64_t(blockIdx.z) * 65535;
+        if (output >= n4 || row >= rows) {
             return;
         }
-        const int64_t index = (int64_t)blockIdx.y * n4 + output;
+        const int64_t index = row * n4 + output;
         const int4 v        = src[index];
-        const float scale   = scales[blockIdx.y];
+        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];
@@ -371,7 +376,7 @@ namespace sd_i8_cuda {
         cache.emplace(device, info);
         return info;
     }
-    static bool aligned(const ggml_tensor* tensor) {
+    static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) {
         if (!tensor)
             return true;
         size_t offset = 0;
@@ -379,14 +384,14 @@ namespace sd_i8_cuda {
             offset += tensor->view_offs;
             tensor = tensor->view_src;
         }
-        return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 4 == 0;
+        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 > 65535 || k % 4 || n % 4 || (op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256))))
+        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))
@@ -402,41 +407,105 @@ namespace sd_i8_cuda {
         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) && device_info(device).ordinal >= 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 <= device_info(device).shared)
+                return 3;
+        }
+        return 2;
+    }
     struct Plan {
         DeviceInfo device;
+        Kind kind;
         int64_t k, n, rows, padded;
-        size_t quant, scales, accum, partial, output;
+        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, void** state, size_t* workspace, void*) {
-        const auto* dst = nodes[0];
+    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(dst, &op);
+        const bool mat = sd_ggml_i8_describe(nodes[0], &op);
         auto* p        = new Plan{};
         *state         = p;
-        p->device      = device_info(ggml_backend_get_device(backend));
+        auto device    = ggml_backend_get_device(backend);
+        p->device      = 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 ? dst->ne[0] : 0;
-        p->rows        = mat ? ggml_nrows(dst) : ggml_nrows(dst->src[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);
-            p->accum  = reserve(total, size_t(p->n) * p->padded * sizeof(int32_t));
+#ifdef SD_CUDA_I8_CUTLASS
+            size_t gemm_bytes = 0;
+            // Custom GEMM previously used cuBLAS; retain that path when I32 accumulation can overflow.
+            if (op.weight_scale && !op.bias && (nodes[0]->op == GGML_OP_MUL_MAT || 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;
         p->partial                = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
-        p->output                 = reserve(total, ggml_nbytes(dst));
+        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*) {
-        delete static_cast<Plan*>(state);
+        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)
@@ -474,7 +543,54 @@ namespace sd_i8_cuda {
         }
         return cudaGetLastError();
     }
-    static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void* state, void*) {
+    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_t quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, cudaStream_t 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.ordinal]) {
+                auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, int(p.device.shared - 128));
+                if (error != cudaSuccess)
+                    return error;
+                initialized[p.device.ordinal] = true;
+            }
+        }
+        if (p.padded > p.rows) {
+            auto error = cudaMemsetAsync(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
+            if (error != cudaSuccess)
+                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 cudaGetLastError();
+    }
+    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));
         auto query    = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
@@ -483,12 +599,21 @@ namespace sd_i8_cuda {
             return GGML_STATUS_FAILED;
         auto stream   = static_cast<cudaStream_t>(context.stream);
         auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
-        auto* dst     = nodes[0];
+        auto* dst     = nodes[count - 1];
         SDI8Operation op;
-        const bool mat = sd_ggml_i8_describe(dst, &op);
+        const bool mat = sd_ggml_i8_describe(nodes[0], &op);
         bool alias     = false;
-        for (auto* input : dst->src)
-            alias |= overlaps(dst, input);
+        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_t error = cudaSuccess;
         if (!mat) {
@@ -496,10 +621,6 @@ namespace sd_i8_cuda {
             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 blas_query = reinterpret_cast<ggml_ext_cuda_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_blas_v1"));
-            auto handle     = blas_query ? static_cast<cublasHandle_t>(blas_query(execution->backend)) : nullptr;
-            if (!handle)
-                return GGML_STATUS_FAILED;
             auto* quant = reinterpret_cast<int8_t*>(scratch + p.quant);
             const float* scales;
             if (op.input->type == GGML_TYPE_F32) {
@@ -526,26 +647,64 @@ namespace sd_i8_cuda {
             }
             if (error != cudaSuccess || cudaGetLastError() != cudaSuccess)
                 return GGML_STATUS_FAILED;
-            auto* accum         = reinterpret_cast<int32_t*>(scratch + p.accum);
-            const int32_t alpha = 1, beta = 0;
-            if (cublasGemmEx(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) != CUBLAS_STATUS_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;
-            if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(output)) % 16 == 0) {
-                const dim3 grid(unsigned((p.n + 1023) / 1024), unsigned(p.rows));
-                dequantize_i32_rows_vec4_cuda<<<grid, 256, 0, stream>>>(reinterpret_cast<const int4*>(accum), scales, reinterpret_cast<const float4*>(ws), reinterpret_cast<const float4*>(bias), static_cast<float4*>(output), p.n / 4);
-            } else {
-                const dim3 grid(unsigned((p.n + 255) / 256), unsigned(p.rows));
-                dequantize_i32_rows_cuda<<<grid, 256, 0, stream>>>(accum, scales, ws, bias, static_cast<float*>(output), p.n, p.rows);
+            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_cuda_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_blas_v1"));
+                auto handle     = blas_query ? static_cast<cublasHandle_t>(blas_query(execution->backend)) : nullptr;
+                if (!handle)
+                    return GGML_STATUS_FAILED;
+                const int32_t alpha = 1, beta = 0;
+                if (cublasGemmEx(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) != CUBLAS_STATUS_SUCCESS)
+                    return GGML_STATUS_FAILED;
             }
-            error = cudaGetLastError();
+            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 = cudaGetLastError();
+            } 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 = cudaGetLastError();
+            }
+            if (error == cudaSuccess && prefix)
+                error = cudaMemcpyAsync(output, dst->src[0]->data, prefix, cudaMemcpyDeviceToDevice, stream);
         }
         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(ggml_backend_dev_t device) {
@@ -560,8 +719,27 @@ void sd_ggml_register_cuda_i8(ggml_backend_dev_t device) {
     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.cuda.override.v1";
+    desc.op        = GGML_OP_MUL_MAT;
+    desc.custom_op = nullptr;
+    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.cuda.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.name      = "sd.linear_i8.epilogue.cuda.custom.v1";
+    desc.op        = GGML_OP_CUSTOM;
+    desc.custom_op = sd_ggml_i8_op();
+    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.cuda.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));
+    desc.name      = "sd.quantize_i8_convrot.cuda.override.v1";
+    desc.op        = GGML_OP_QUANTIZE_I8_CONVROT;
+    desc.custom_op = nullptr;
+    GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
 }