perf(vulkan): accelerate packed INT4 with cooperative matrices

docs/int8_convrot.md

@@ -133,9 +133,26 @@ Other CUDA shapes and builds retain the packed byte-dot path. Set
 use packed byte-dot instructions, with portable byte arithmetic for other targets.
 Vulkan requires integer dot product support and at least 13312 bytes of shared
 memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input
-width to be divisible by 64; H64 and H256 are supported. Other shapes use CPU
-fallback. HIP W4A8 selects wave64 WMMA by the active compilation target's
-intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout.
+width to be divisible by 64; H64 and H256 are supported. When supported by the
+compiler and device, Vulkan uses NV cooperative matrix integer GEMM. W4A4
+decodes packed weights directly into cooperative matrices for aligned input
+and output widths. Complete activation tiles and the final partial tile use
+separate shader specializations. Other W4A4 shapes use double-buffered weight
+and activation tiles (20480 bytes of shared memory, or 30720 for larger tiles).
+W4A8 stages both weights and activations in shared memory in K=64 tiles,
+reusing them for each 16-value group's integer
+MMA and FP32 scale restoration. FP8 scales use a shared lookup table, and
+expanding projections visit activation tiles in groups of weight panels for
+cache reuse. Adjacent groups share a K=32 activation fragment, with
+complementary zero padding in the weights to preserve each group's scale.
+This path uses 17476 bytes of shared memory, or 22596 for larger tiles.
+Both paths use 256-thread workgroups and keep the model's
+weights packed; neither expands a full layer in device memory. Other Vulkan
+devices retain the packed dot-product path; set `GGML_VK_W4_FORCE_DP4A=1` to compare
+against it. FP32 accumulation order can differ between these paths.
+Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA by the active
+compilation target's intrinsic support, preferring the GFX12 fragment layout
+over the GFX11 layout.
 There is no GPU architecture allowlist: newer targets supporting either
 intrinsic automatically use its matching layout. Selection happens per device
 image, including multi-architecture builds; targets supporting neither intrinsic

src/core/vulkan/CMakeLists.txt

@@ -56,6 +56,9 @@ sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O)
 sd_vulkan_shader(quantize_i8_convrot_h64 quantize_i8_convrot_h64 ${SD_VK_DOT} -O -DGROUP_SIZE=64)
 sd_vulkan_shader(mul_mat_w4a4_convrot mul_mat_w4a4_convrot ${SD_VK_DOT} -O)
 sd_vulkan_shader(mul_mat_w4a8_convrot mul_mat_w4a8_convrot ${SD_VK_DOT} -O)
+sd_vulkan_shader(mul_mat_w4a4_convrot_cm2 mul_mat_w4a4_convrot_cm2 ${SD_VK_CM2} -O)
+sd_vulkan_shader(mul_mat_w4a4_convrot_decode_cm2 mul_mat_w4a4_convrot_decode_cm2 ${SD_VK_CM2} -O)
+sd_vulkan_shader(mul_mat_w4a8_convrot_cm2 mul_mat_w4a8_convrot_cm2 ${SD_VK_CM2} -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)

src/core/vulkan/int8.cpp

@@ -1,6 +1,7 @@
 #include <algorithm>
 #include <array>
 #include <climits>
+#include <cstdlib>
 #include <cstring>
 #include <vector>
 #include "core/ggml_extensions.h"
@@ -12,7 +13,10 @@
 #include "mul_mat_i8_tensorwise_cm1.h"
 #include "mul_mat_i8_tensorwise_cm2.h"
 #include "mul_mat_w4a4_convrot.h"
+#include "mul_mat_w4a4_convrot_cm2.h"
+#include "mul_mat_w4a4_convrot_decode_cm2.h"
 #include "mul_mat_w4a8_convrot.h"
+#include "mul_mat_w4a8_convrot_cm2.h"
 #include "quantize_i8_convrot.h"
 #include "quantize_i8_convrot_h64.h"
 #include "quantize_i8_convrot_register.h"
@@ -23,11 +27,17 @@
 #include "swiglu_quantize_i8_convrot_16384.h"
 
 namespace sd_vulkan_i8 {
+    // Keep these budgets in sync with the cooperative matrix shader tile layouts.
+    static constexpr uint32_t w4a4_shared       = 2 * (64 + 64) * (64 / 4 + 4) * 4;
+    static constexpr uint32_t w4a4_large_shared = 2 * (64 + 128) * (64 / 4 + 4) * 4;
+    static constexpr uint32_t w4a8_shared       = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 64 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4;
+    static constexpr uint32_t w4a8_large_shared = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 128 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4;
     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;
+        bool w4a4_cm2, w4a8_cm2;
         uint32_t cm_m, cm_n, cm_k;
     };
     static bool device_info(ggml_backend_dev_t dev, Device& d) {
@@ -55,6 +65,10 @@ namespace sd_vulkan_i8 {
                 d.cm2 = true;
                 d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0;
             }
+            if (prop.flexible && prop.workgroup_scope && prop.invocations == 256 && d.info.max_group_size[0] >= 256 && 64 % prop.m == 0 && 64 % prop.n == 0) {
+                d.w4a4_cm2 |= d.compute.shared_memory >= w4a4_shared && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_cm2) > 4;
+                d.w4a8_cm2 |= d.compute.shared_memory >= w4a8_shared && 32 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot_cm2) > 4;
+            }
         }
         return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4;
     }
@@ -698,9 +712,45 @@ namespace sd_vulkan_i8 {
         struct Parameters {
             uint32_t n, rows, k, bias;
         };
-        const Parameters pc     = {n, rows, k, op.bias ? 1u : 0u};
-        auto pipeline           = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters);
-        const uint32_t groups[] = {(n + 63) / 64, (rows + 31) / 32, 1};
+        const Parameters pc = {n, rows, k, op.bias ? 1u : 0u};
+        auto pipeline       = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters);
+        uint32_t tile_cols = 64, tile_rows = 32;
+        const bool large_w4a4  = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared;
+        const bool large_w4a8  = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared;
+        const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 &&
+                                 aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4;
+        const uint32_t constants[] = {64, large_w4a4 || large_w4a8 ? 128u : 64u, 64};
+        const char* force_dp4a     = std::getenv("GGML_VK_W4_FORCE_DP4A");
+        if (!(force_dp4a && std::atoi(force_dp4a)) && aligned(op.input, 16) &&
+            (op.kind == SDW4Kind::W4A4 ? d.w4a4_cm2 : d.w4a8_cm2 && aligned(op.weight, 8))) {
+            pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot_cm2, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot_cm2, 8, Parameters);
+            if (decode_w4a4) {
+                struct DecodeParameters {
+                    uint32_t n, rows, k, bias, row_offset;
+                };
+                DecodeParameters decoded    = {n, rows, k, pc.bias, 0};
+                uint32_t decode_constants[] = {64, 128, 128, 1};
+                pipeline                    = SD_VK_PIPELINE(mul_mat_w4a4_convrot_decode_cm2, 6, DecodeParameters);
+                pipeline.constants          = decode_constants;
+                pipeline.constant_count     = 4;
+                uint32_t groups[]           = {n / 64, rows / 128, 1};
+                auto status                 = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups);
+                if (status == GGML_STATUS_SUCCESS && rows % 128) {
+                    // Specialize the tail separately so bounds checks do not inhibit full-tile loads.
+                    decoded.row_offset  = rows / 128 * 128;
+                    decode_constants[1] = rows % 128 <= 64 ? 64 : 128;
+                    decode_constants[3] = 0;
+                    groups[1]           = 1;
+                    status              = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups);
+                }
+                return finish(e, d, output, dst, alias, status);
+            }
+            pipeline.constants      = constants;
+            pipeline.constant_count = 3;
+            tile_cols               = constants[0];
+            tile_rows               = constants[1];
+        }
+        const uint32_t groups[] = {(n + tile_cols - 1) / tile_cols, (rows + tile_rows - 1) / tile_rows, 1};
         auto status             = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
         return finish(e, d, output, dst, alias, status);
     }

src/core/vulkan/shaders/mul_mat_w4a4_convrot_cm2.comp

@@ -0,0 +1,88 @@
+#version 460
+#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(local_size_x = 256) in;
+layout(constant_id = 0) const uint BM = 64;
+layout(constant_id = 1) const uint BN = 64;
+layout(constant_id = 2) const uint BK = 64;
+const uint STRIDE                     = BK / 4 + 4;
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+layout(binding = 0, std430) readonly buffer W {
+    uint data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    uvec4 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 {
+    int32_t data_d[];
+};
+shared uint weights[2 * BM * STRIDE];
+shared uint activations[2 * BN * STRIDE];
+void load_tile(uint base, uint buf, uint row_base, uint col_base) {
+    const uint tid = gl_LocalInvocationIndex;
+    for (uint i = tid; i < BN * BK / 16; i += 256) {
+        const uint r = i / (BK / 16), t = i % (BK / 16);
+        const uvec4 value       = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0);
+        const uint offset       = (buf * BN + r) * STRIDE + t * 4;
+        activations[offset]     = value.x;
+        activations[offset + 1] = value.y;
+        activations[offset + 2] = value.z;
+        activations[offset + 3] = value.w;
+    }
+    for (uint i = tid; i < BM * BK / 8; i += 256) {
+        const uint c = i / (BK / 8), t = i % (BK / 8);
+        const uint v                                 = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 8) + base / 8 + t] : 0u;
+        const uint lo                                = (v & 15u) | ((v << 4) & 0xf00u) | ((v << 8) & 0xf0000u) | ((v << 12) & 0xf000000u);
+        const uint hi                                = ((v >> 16) & 15u) | ((v >> 12) & 0xf00u) | ((v >> 8) & 0xf0000u) | ((v >> 4) & 0xf000000u);
+        weights[(buf * BM + c) * STRIDE + 2 * t]     = (lo & 0x07070707u) | ((lo & 0x08080808u) * 31u);
+        weights[(buf * BM + c) * STRIDE + 2 * t + 1] = (hi & 0x07070707u) | ((hi & 0x08080808u) * 31u);
+    }
+}
+int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) {
+    if (row + c >= p.rows || col + r >= p.n)
+        return 0;
+    precise float value = float(v) * data_as[row + c] * data_ws[col + r];
+    if (p.has_bias != 0)
+        value += data_b[col + r];
+    return floatBitsToInt(value);
+}
+void main() {
+    const uint row = gl_WorkGroupID.y * BN, col = gl_WorkGroupID.x * BM;
+    coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+    load_tile(0, 0, row, col);
+    barrier();
+    for (uint base = 0; base < p.k; base += BK) {
+        const uint buf = (base / BK) & 1u;
+        if (base + BK < p.k)
+            load_tile(base + BK, buf ^ 1u, row, col);
+        coopmat<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+        coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> a;
+        coopMatLoad(w, weights, buf * BM * STRIDE, STRIDE, gl_CooperativeMatrixLayoutRowMajor);
+        coopMatLoad(a, activations, buf * BN * STRIDE, STRIDE, gl_CooperativeMatrixLayoutColumnMajor);
+        sum = coopMatMulAdd(w, a, sum);
+        barrier();
+    }
+    coopMatPerElementNV(sum, sum, dequantize, row, col);
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    dl                                                            = setTensorLayoutDimensionNV(dl, p.rows, p.n);
+    tensorViewNV<2, false, 1, 0> transpose                        = createTensorViewNV(2, false, 1, 0);
+    coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose);
+}

src/core/vulkan/shaders/mul_mat_w4a4_convrot_decode_cm2.comp

@@ -0,0 +1,92 @@
+#version 460
+#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
+#extension GL_EXT_shader_8bit_storage : require
+#extension GL_EXT_buffer_reference : 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(local_size_x = 256) in;
+layout(constant_id = 0) const uint BM   = 64;
+layout(constant_id = 1) const uint BN   = 64;
+layout(constant_id = 2) const uint BK   = 64;
+layout(constant_id = 3) const bool FULL = true;
+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 {
+    uvec4 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 {
+    int32_t data_d[];
+};
+layout(buffer_reference, std430, buffer_reference_align = 16) buffer PackedBlock {
+    uvec4 value;
+};
+int8_t decode_weight(const in PackedBlock block, const in uint block_coords[2], const in uint coord[2]) {
+    return int8_t(bitfieldExtract(int(block.value[coord[1] / 8]), int((coord[1] % 8) * 4), 4));
+}
+int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) {
+    if (row + c >= p.rows || col + r >= p.n)
+        return 0;
+    precise float value = float(v) * data_as[row + c] * data_ws[col + r];
+    if (p.has_bias != 0)
+        value += data_b[col + r];
+    return floatBitsToInt(value);
+}
+// The host selects complete weight tiles and a 256-aligned input width.
+void main() {
+    const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y;
+    const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x;
+    const uint panel = min(32u, cols), first_col = (linear / (panel * rows)) * panel;
+    const uint group_cols = min(panel, cols - first_col), within = linear - first_col * rows;
+    const uint row = p.row_offset + (within / group_cols) * BN, 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);
+    tensorLayoutNV<2> wl                                                     = createTensorLayoutNV(2);
+    wl                                                                       = setTensorLayoutBlockSizeNV(wl, 1, 32);
+    wl                                                                       = setTensorLayoutDimensionNV(wl, p.n, p.k & ~255u);
+    tensorLayoutNV<2> al                                                     = createTensorLayoutNV(2);
+    al                                                                       = setTensorLayoutDimensionNV(al, p.rows, p.k & ~255u);
+    tensorViewNV<2, false, 1, 0> at                                          = createTensorViewNV(2, false, 1, 0);
+    if (FULL) {
+        for (uint base = 0; base < p.k; base += 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(wl, col, BM, base, BK), decode_weight);
+            coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(al, row, BN, base, BK), at);
+            sum = coopMatMulAdd(w, a, sum);
+        }
+    } else {
+        tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tail = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+        tail                                                            = setTensorLayoutDimensionNV(tail, p.rows, p.k & ~255u);
+        for (uint base = 0; base < p.k; base += 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(wl, col, BM, base, BK), decode_weight);
+            coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(tail, row, BN, base, BK), at);
+            sum = coopMatMulAdd(w, a, sum);
+        }
+    }
+    coopMatPerElementNV(sum, sum, dequantize, row, col);
+    tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV);
+    dl                                                            = setTensorLayoutDimensionNV(dl, p.rows, p.n);
+    tensorViewNV<2, false, 1, 0> transpose                        = createTensorViewNV(2, false, 1, 0);
+    coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose);
+}