perf(cuda): fuse bias into INT8 CUTLASS epilogues (#2320)

docs/int8_convrot.md

@@ -94,7 +94,7 @@ The packed runtime activation tensor contains the I8 activation rows and their f
 
 - CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations.
 - Metal supports INT8 tensorwise and convrot on Apple silicon (Apple GPU family 7 or newer). It keeps weights in INT8 and rotates and quantizes activations on the GPU. On macOS/iOS 26 or newer, it probes Metal TensorOps `matmul2d` with INT8 inputs and INT32 cooperative results. It combines 1024-product partial sums using wrapping integer addition to preserve exact accumulation on older GPUs, then applies scales and bias in registers. Older systems or unavailable TensorOps pipelines use SIMD-group matrix multiplication: INT8 tiles are converted to FP16, and FP32 partial sums cover at most 256 products before being combined as integers, preserving the I32 accumulation contract. Convrot groups 1, 4, 16, 64, and 256 are supported. These kernels require the patched GGML Metal extension services and are enabled by `SD_METAL=ON`.
-- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
+- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled. For linear layers, it restores scales and adds any aligned bias before writing the result; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition.
 - AMD HIP devices in the CDNA, RDNA3 (including RDNA3.5), and RDNA4 families use the same INT8/H256 kernels with hipBLAS for I8 x I8 to I32 GEMM. Other AMD architectures fall back to CPU for INT8 matrix multiplication.
 - MUSA shares the CUDA-style H256 activation kernels. INT8 GEMM falls back to CPU because the MUSA backend does not advertise INT8 BLAS support.
 - Vulkan provides H256 activation quantization and INT8 matrix multiplication when the build and device support accelerated packed INT8 dot products. Unsupported configurations and GPU backends without these kernels use the backend scheduler to fall back to CPU.

src/core/cuda/i8_gemm.cu

@@ -176,9 +176,22 @@ namespace sd_i8_cutlass {
             dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)];
     }
 
+    template <class Fragment>
+    struct AddBias {
+        CUTLASS_DEVICE Fragment operator()(const Fragment& value, const Fragment& bias) const {
+            Fragment result;
+            // Prevent contraction with the preceding scale multiplication.
+            CUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < Fragment::kElements; ++i)
+                result[i] = __fadd_rn(value[i], bias[i]);
+            return result;
+        }
+    };
+
     template <class Config, bool StreamK, bool PackedW4 = false>
     struct DenseTypes {
         static constexpr bool kSplit = false;
+        static constexpr bool kBias  = !PackedW4;
         using Weight                 = std::conditional_t<PackedW4, cutlass::int4b_t, int8_t>;
         using Tile                   = typename Config::Tile;
         using Warp                   = typename Config::Warp;
@@ -190,7 +203,9 @@ namespace sd_i8_cutlass {
         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 Add                    = cutlass::epilogue::threadblock::VisitorCompute<AddBias, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+        using Biased                 = cutlass::epilogue::threadblock::Sm80EVT<Add, ScaleW, Row>;
+        using Epilogue               = cutlass::epilogue::threadblock::Sm80EVT<Store, std::conditional_t<kBias, Biased, ScaleW>>;
         using Kernel                 = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
             int8_t,
             cutlass::layout::RowMajor,
@@ -219,6 +234,7 @@ namespace sd_i8_cutlass {
     template <class Config, bool StreamK, bool Split = false>
     struct SwigluTypes {
         static constexpr bool kSplit = Split;
+        static constexpr bool kBias  = false;
         using Weight                 = int8_t;
         using Tile                   = typename Config::Tile;
         using Warp                   = typename Config::Warp;
@@ -260,23 +276,29 @@ namespace sd_i8_cutlass {
         bool split;
         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, const int8_t*, const float*) = nullptr;
+        ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*, const float*) = 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) {
+    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, const float* bias) {
         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}}};
+        typename Types::Epilogue::Arguments epilogue;
+        if constexpr (Types::kBias) {
+            // A negative zero default preserves signed-zero outputs.
+            epilogue = {{scales, {bias, -0.f, {}}, {}}, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+        } else {
+            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, reinterpret_cast<const typename Types::Weight*>(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, const int8_t* up, const float* up_scales) {
+    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, const int8_t* up, const float* up_scales, const float* bias) {
         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) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales)) %
+             reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out) | reinterpret_cast<uintptr_t>(up) | reinterpret_cast<uintptr_t>(up_scales) | reinterpret_cast<uintptr_t>(bias)) %
             16)
             return GGML_STATUS_FAILED;
         auto* paired = static_cast<float*>(scratch);
-        auto args    = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out);
+        auto args    = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out, bias);
         if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
             return GGML_STATUS_FAILED;
         if (!p.occupancy) {
@@ -312,7 +334,7 @@ namespace sd_i8_cutlass {
         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);
+        auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr);
         if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
             return false;
         size_t bytes = 0;
@@ -394,11 +416,13 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64
         return nullptr;
     return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true);
 }
-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) {
+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, const float* bias) {
-    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr);
+    if (bias && (plan->swiglu || plan->packed_w4))
+        return GGML_STATUS_FAILED;
+    return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr, bias);
 }
 ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) {
-    return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales);
+    return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales, nullptr);
 }
 void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) {
     delete plan;

src/core/cuda/i8_gemm.h

@@ -4,7 +4,7 @@
 #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, bool packed_w4 = false);
-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);
+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, const float* bias = nullptr);
 // n is the width of each branch; gate/up weights and scales remain separate.
 SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace);
 ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream);

src/core/cuda/int8.cu

@@ -447,7 +447,7 @@ namespace SD_CUDA_NAMESPACE {
         }
 #ifdef SD_CUDA_I8_CUTLASS
         static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) {
-            return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 &&
+            return op.weight_scale && op.weight->ne[0] <= INT_MAX / 16384 && aligned(op.bias, 16) &&
                    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)));
         }
@@ -531,7 +531,7 @@ namespace SD_CUDA_NAMESPACE {
 #ifdef SD_CUDA_I8_CUTLASS
                 size_t gemm_bytes = 0;
                 // Retain cuBLAS when I32 accumulation can overflow.
-                if (cutlass_compatible(op, dst) &&
+                if (cutlass_compatible(op, dst) && (!op.bias || p->kind == MAT) &&
                     (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);
@@ -819,7 +819,7 @@ namespace SD_CUDA_NAMESPACE {
 #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_at(p.gemm_offset), stream);
+                                                          intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch_at(p.gemm_offset), stream, bias);
                     if (status != GGML_STATUS_SUCCESS)
                         return status;
                 } else