feat: support packed INT4 convrot models

docs/ggml_extensions.md

@@ -104,6 +104,16 @@ and backend `graph_plan_*` interfaces do not dispatch these extensions.
 
 ## Backend services
 
+Packed ComfyUI INT4 linear layers use `sd.linear_w4_convrot.v1`, built by
+`ggml_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the
+explicit output of `sd.quantize_i8_convrot.v1`; F32 callers get a quantization
+node during graph construction. Logical output dimensions are stored in the
+operator attributes, so backend copies of the packed activation do not need to
+retain the producer's op marker. W4A4 uses row scales; W4A8 additionally exposes
+the codebook, channel scales, and FP8 relative scales as tensor dependencies.
+CPU, CUDA/HIP/MUSA, and Vulkan register kernels through the same extension API
+as INT8. No ggml submodule changes are required.
+
 Services are discovered through `ggml_backend_reg_get_proc_address`:
 
 | Service | Purpose |

docs/int8_convrot.md

@@ -100,7 +100,46 @@ The packed runtime activation tensor contains the I8 activation rows and their f
 
 LoRA adapters are applied at runtime without modifying the INT8 weights. The INT8 convrot path computes the base linear output, while LoRA, LoHa, LoKr, and raw weight-difference adapters compute their output corrections from the original, unrotated activation and add them to the base output. `--lora-apply-mode auto` selects this path for models containing INT8 tensorwise weights. If `immediately` is requested, sd.cpp falls back to runtime application because merging an adapter would require dequantizing and rotating its weight update, then recalculating the per-row scales and requantizing the result.
 
-The shared CUDA/HIP/MUSA convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution.
+The shared CUDA/HIP/MUSA convrot activation paths support H64 and H256. Vulkan also supports H64 activation quantization; its INT8 matrix multiplication path currently requires H256.
+
+## Packed INT4 convrot
+
+ComfyUI `convrot_w4a4` and `asym_w4a8_int8` safetensors load directly, including
+checkpoints that mix packed INT4, INT8 tensorwise, and floating-point layers.
+Packed weights retain their checkpoint byte layout; they are not expanded or
+converted to GGUF Q4 blocks at load time.
+
+- `convrot_w4a4`: signed 4-bit weights, low nibble first, with F32 per-output
+  `weight_scale`. Activations use the same INT8 convrot quantization as above;
+  the format name does not select a native INT4 activation kernel.
+- `asym_w4a8_int8`: 4-bit indices into `weight_codebook` (16 F32 values), F32
+  `weight_s_channel`, and F8_E4M3 `weight_s_rel` per 16 input features. The
+  metadata must specify `group_size: 16`. CPU evaluates the floating-point
+  codebook; GPU kernels requantize it to INT8 with `amax / 127`, introducing a
+  small additional rounding error.
+
+Both formats imply convrot rotation even without a `convrot` flag; missing
+`convrot_groupsize` defaults to 256.
+
+Both formats reuse the convrot activation cache when the input scale is `1` and
+produce F32 outputs with optional bias. CPU provides a portable implementation.
+CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS;
+CUDA and supported AMD targets 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.
+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
+retain the integer-dot fallback. W4A4 retains integer dot products on HIP.
+
+LoRA application remains at runtime, using the original unrotated activations
+and logical weight dimensions. Packed weights and their codebook/scale tensors
+are excluded from weight type conversion. `SD_USE_UPSTREAM_GGML` builds reject
+these formats because they require the application extension API.
 
 ## Example
 

docs/model_config.md

@@ -60,6 +60,8 @@ static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage
 ```
 
 Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes.
+Use `logical_ne0()` for Linear input widths; packed INT4 weights store half of
+the logical width in `ne[0]`. Keep raw `ne` values for storage layout calculations.
 Do not load or parse tensor data for config detection.
 
 Detection should respect `prefix`. For nested weights, construct full names from

src/core/cpu/registration.cpp

@@ -5,6 +5,7 @@ namespace sd_cpu {
 
     void register_extensions(ggml_backend_dev_t device) {
         register_i8(device);
+        register_w4(device);
         register_rope(device);
         register_modulate(device);
     }

src/core/cpu/registration.h

@@ -9,6 +9,7 @@ namespace sd_cpu {
     void register_rope(ggml_backend_dev_t device);
     void register_modulate(ggml_backend_dev_t device);
     void register_i8(ggml_backend_dev_t device);
+    void register_w4(ggml_backend_dev_t device);
 }
 #endif
 

src/core/cpu/w4.cpp

@@ -0,0 +1,95 @@
+#include "core/cpu/registration.h"
+#include "core/ggml_extensions.h"
+
+#ifndef SD_USE_UPSTREAM_GGML
+#include <cstring>
+
+namespace sd_cpu {
+    struct W4Execution {
+        SDW4Operation op;
+        float* output;
+        int64_t rows;
+    };
+
+    static void multiply_w4(ggml_tensor*, int ith, int nth, void* data) {
+        auto& e         = *static_cast<W4Execution*>(data);
+        const auto& op  = e.op;
+        const int64_t k = op.input->ne[0], n = op.weight->ne[1];
+        const auto* w       = static_cast<const uint8_t*>(op.weight->data);
+        const auto* x       = static_cast<const int8_t*>(op.input->data);
+        const auto* xs      = reinterpret_cast<const float*>(x + k * ((e.rows + 3) / 4 * 4));
+        const auto* ws      = static_cast<const float*>(op.weight_scale->data);
+        const auto* bias    = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+        const auto* sc      = op.s_channel ? static_cast<const float*>(op.s_channel->data) : nullptr;
+        const auto* sr      = op.s_rel ? static_cast<const uint8_t*>(op.s_rel->data) : nullptr;
+        const auto to_float = ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float;
+        for (int64_t index = ith; index < n * e.rows; index += nth) {
+            const int64_t row = index / n, col = index % n;
+            const auto* packed = w + col * (k / 2);
+            float value;
+            if (op.kind == SDW4Kind::W4A4) {
+                uint32_t sum = 0;
+                for (int64_t i = 0; i < k; ++i) {
+                    const int nibble = (packed[i / 2] >> (4 * (i % 2))) & 15;
+                    sum += uint32_t(((nibble ^ 8) - 8) * int32_t(x[row * k + i]));
+                }
+                int32_t signed_sum;
+                std::memcpy(&signed_sum, &sum, sizeof(sum));
+                value = float(signed_sum) * xs[row] * ws[col];
+            } else {
+                float sum = 0.f;
+                for (int64_t g = 0; g < k / 16; ++g) {
+                    float relative;
+                    to_float(sr + col * (k / 16) + g, &relative, 1);
+                    const float scale = sc[col] * relative;
+                    for (int j = 0; j < 16; ++j) {
+                        const int64_t i = g * 16 + j;
+                        const int code  = (packed[i / 2] >> (4 * (i % 2))) & 15;
+                        sum += (ws[code] * scale) * float(x[row * k + i]);
+                    }
+                }
+                value = sum * xs[row];
+            }
+            e.output[index] = value + (bias ? bias[col] : 0.f);
+        }
+    }
+
+    static bool w4_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+        SDW4Operation op;
+        return sd_ggml_w4_describe(node, &op);
+    }
+    static ggml_status w4_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+        *bytes = ggml_nbytes(nodes[0]);
+        return GGML_STATUS_SUCCESS;
+    }
+    static ggml_status w4_execute(const ggml_ext_execution_context* ctx, ggml_tensor* const* nodes, int, void*, void*) {
+        auto* dst     = nodes[0];
+        W4Execution e = {};
+        if (!sd_ggml_w4_describe(dst, &e.op))
+            return GGML_STATUS_FAILED;
+        auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(ctx->backend));
+        auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+        if (!run)
+            return GGML_STATUS_FAILED;
+        e.rows      = ggml_nrows(dst);
+        e.output    = static_cast<float*>(ggml_backend_buffer_get_base(ctx->workspace));
+        auto status = run(ctx->backend, dst, multiply_w4, &e);
+        if (status == GGML_STATUS_SUCCESS)
+            std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+        return status;
+    }
+    void register_w4(ggml_backend_dev_t device) {
+        ggml_ext_kernel_desc d = {};
+        d.abi_version          = GGML_EXT_ABI_VERSION;
+        d.name                 = "sd.linear_w4_convrot.cpu.v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_w4_op();
+        d.priority             = 100;
+        d.supports             = w4_supports;
+        d.prepare              = w4_prepare;
+        d.execute              = w4_execute;
+        d.destroy              = [](void*, void*) {};
+        GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    }
+}
+#endif

src/core/cuda/CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW)
 
 function(sd_cuda_sources VENDOR LANGUAGE)
     set(SOURCES)
-    foreach(OP rope int8 convrot affine modulate bf16 rms_affine registration)
+    foreach(OP rope int8 w4 convrot affine modulate bf16 rms_affine registration)
         set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu")
         string(TOUPPER "${VENDOR}" UPPER_VENDOR)
         file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n")
@@ -49,6 +49,8 @@ if(SD_HIPBLAS)
         set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON)
     endif()
     set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math")
+    # W4A8 WMMA uses the wave64 fragment layout on both GFX11 and GFX12.
+    set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_hip.cu" APPEND PROPERTY COMPILE_OPTIONS "-mwavefrontsize64")
     target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas)
     target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas)
 endif()

src/core/cuda/int8.cu

@@ -386,12 +386,55 @@ namespace SD_CUDA_NAMESPACE {
             return info.int8_blas &&
                    ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
         }
+        template <bool Quantize>
+        static __global__ void convrot_h64(const float* src, int8_t* dst, float* partial, const float* scales, int64_t k) {
+            const int64_t group = blockIdx.x, groups = k / 64, row = group / groups;
+            const int tid = threadIdx.x;
+            __shared__ float values[64];
+            __shared__ float maxima[16];
+#pragma unroll
+            for (int j = 0; j < 4; ++j)
+                values[tid + j * 16] = src[group * 64 + tid + j * 16] * (1.f / 8.f);
+            __syncthreads();
+#pragma unroll
+            for (int stride = 1; stride < 64; stride *= 4) {
+                const int i   = (tid / stride) * 4 * stride + tid % stride;
+                const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+                values[i]              = convrot_h4_cuda(a, b, c, d, 0);
+                values[i + stride]     = convrot_h4_cuda(a, b, c, d, 1);
+                values[i + 2 * stride] = convrot_h4_cuda(a, b, c, d, 2);
+                values[i + 3 * stride] = convrot_h4_cuda(a, b, c, d, 3);
+                __syncthreads();
+            }
+            if constexpr (Quantize) {
+                const float inverse = scales[row] == 0.f ? 0.f : 1.f / scales[row];
+#pragma unroll
+                for (int j = 0; j < 4; ++j) {
+                    const int value                = __float2int_rn(values[tid + j * 16] * inverse);
+                    dst[group * 64 + tid + j * 16] = int8_t(max(-127, min(127, value)));
+                }
+            } else {
+                float amax = 0.f;
+#pragma unroll
+                for (int j = 0; j < 4; ++j)
+                    amax = fmaxf(amax, fabsf(values[tid + j * 16]));
+                maxima[tid] = amax;
+                __syncthreads();
+                for (int stride = 8; stride; stride >>= 1) {
+                    if (tid < stride)
+                        maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+                    __syncthreads();
+                }
+                if (!tid)
+                    partial[group] = maxima[0];
+            }
+        }
         static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
             int group;
-            if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0]))
+            if (!sd_ggml_convrot_describe(node, &group) || (group != 1 && group != 64 && group != 256) || !aligned(node) || !aligned(node->src[0]))
                 return false;
             const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]);
-            return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0;
+            return k <= INT_MAX && rows <= INT_MAX && (group == 1 || rows <= INT_MAX / (k / group)) && cuda_device_info(device).device >= 0;
         }
         enum Kind { MAT,
                     QUANTIZE,
@@ -421,7 +464,8 @@ namespace SD_CUDA_NAMESPACE {
                 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) &&
+                int group;
+                if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) &&
                     size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory)
                     return 3;
             }
@@ -430,6 +474,7 @@ namespace SD_CUDA_NAMESPACE {
         struct Plan {
             ggml_ext_gpu_device_info device;
             Kind kind;
+            int group;
             int64_t k, n, rows, padded;
             size_t quant, scales, accum, partial, output, gemm_offset;
             SDI8CudaGemmPlan* gemm = nullptr;
@@ -456,7 +501,10 @@ namespace SD_CUDA_NAMESPACE {
             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;
+            p->group       = mat ? op.group_size : 0;
+            if (!mat)
+                sd_ggml_convrot_describe(dst, &p->group);
+            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);
@@ -480,7 +528,8 @@ namespace SD_CUDA_NAMESPACE {
                 p->accum                = reserve(total, intermediate ? size_t(p->fast_swiglu ? p->n / 2 : p->n) * p->padded * sizeof(float) : 0);
             }
             const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory;
-            p->partial                = reserve(total, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0);
+            p->partial                = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float)
+                                                                                                                                       : 0);
             p->output                 = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst));
             *workspace                = total;
             return GGML_STATUS_SUCCESS;
@@ -506,6 +555,17 @@ namespace SD_CUDA_NAMESPACE {
                     return error;
             }
             const size_t shared = size_t(p.k) * sizeof(float);
+            if (p.group == 1) {
+                quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(src, dst, scales, p.k, p.rows);
+                return cuda_last_error();
+            }
+            if (p.group == 64) {
+                const auto groups = p.k / 64;
+                convrot_h64<false><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
+                reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
+                convrot_h64<true><<<unsigned(p.rows * groups), 16, 0, stream>>>(src, dst, partial, scales, p.k);
+                return cuda_last_error();
+            }
             if (shared + 128 <= p.device.shared_memory) {
                 static std::mutex mutex;
                 static std::unordered_map<int, bool> configured;

src/core/cuda/registration.cu

@@ -9,6 +9,7 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device);
 namespace SD_CUDA_NAMESPACE {
     void register_rope(ggml_backend_dev_t device);
     void register_i8(ggml_backend_dev_t device);
+    void register_w4(ggml_backend_dev_t device);
     void register_convrot(ggml_backend_dev_t device);
     void register_affine(ggml_backend_dev_t device);
     void register_modulate(ggml_backend_dev_t device);
@@ -25,6 +26,7 @@ namespace SD_CUDA_NAMESPACE {
         register_affine(device);
         register_modulate(device);
         register_i8(device);
+        register_w4(device);
         register_convrot(device);
         register_rope(device);
         register_bf16(device);