feat: support packed INT4 convrot models (#2276)

Co-authored-by: noctrex <noctrex@gmail.com>

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);

src/core/cuda/w4.cu

@@ -0,0 +1,873 @@
+#include <climits>
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+    namespace sd_w4_cuda {
+        static __device__ __forceinline__ int shuffle_sum(int value, int lane, int width) {
+#if defined(SD_CUDA_VENDOR_HIP)
+            // ROCm LLVM 24 can crash lowering __shfl's wave_shuffle on RDNA2.
+            // These 8/16-lane groups stay within one ds_bpermute half-wave.
+            const int index = (lane & (width - 1)) + (__lane_id() & ~(width - 1));
+            return __builtin_amdgcn_ds_bpermute(index * 4, value);
+#else
+            return cuda_shuffle(value, lane, width);
+#endif
+        }
+
+        // Signed byte dots preserve the INT4 grid after the XOR-8 bias correction.
+        static __device__ __forceinline__ int dot4(int a, int b, int c) {
+#if defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__))
+            return __builtin_amdgcn_sudot4(true, a, true, b, c, false);
+#elif defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__))
+            return __builtin_amdgcn_sdot4(a, b, c, false);
+#elif defined(SD_CUDA_VENDOR_MUSA) || (!defined(SD_CUDA_VENDOR_HIP) && __CUDA_ARCH__ >= 610)
+            return __dp4a(a, b, c);
+#else
+            const auto* av = reinterpret_cast<const signed char*>(&a);
+            const auto* bv = reinterpret_cast<const signed char*>(&b);
+            for (int i = 0; i < 4; ++i)
+                c += int(av[i]) * int(bv[i]);
+            return c;
+#endif
+        }
+        static __device__ __forceinline__ float fp8_e4m3(uint8_t v) {
+            const int exponent = (v >> 3) & 15, mantissa = v & 7;
+            float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9)
+                                                                                                        : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7);
+            return v & 128 ? -result : result;
+        }
+        // A CTA shares decoded weights across 32 activation rows. Each double-buffered
+        // stage covers 64 K values; W4A8 restores the codebook scale per 16-value group.
+
+        static __device__ __forceinline__ uint32_t perm_bytes(uint32_t lo, uint32_t hi, uint32_t sel) {
+#if defined(SD_CUDA_VENDOR_HIP)
+            return __builtin_amdgcn_perm(hi, lo, sel);
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+            return __byte_perm(lo, hi, sel);
+#else
+            uint32_t r = 0;
+#pragma unroll
+            for (int i = 0; i < 4; ++i) {
+                const uint32_t b    = (sel >> (4 * i)) & 0xF;
+                const uint32_t byte = (b < 4 ? lo : hi) >> (8 * (b & 3)) & 0xFF;
+                r |= byte << (8 * i);
+            }
+            return r;
+#endif
+        }
+
+#if defined(SD_CUDA_VENDOR_HIP)
+#define PERM_SEL_EVEN 0x06040200u
+#define PERM_SEL_ODD 0x07050301u
+#else
+#define PERM_SEL_EVEN 0x6420u
+#define PERM_SEL_ODD 0x7531u
+#endif
+
+        static __device__ __forceinline__ void interleave_act_pair(
+            const int lo,
+            const int hi,
+            int* __restrict__ dst) {
+            dst[0] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_EVEN);
+            dst[1] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_ODD);
+        }
+
+#define MUL_MAT_W4_BLOCK 16
+#define MUL_MAT_W4_ROWS (2 * MUL_MAT_W4_BLOCK)
+#define MUL_MAT_W4_GPS 4
+#define MUL_MAT_W4_SLOTS (4 * MUL_MAT_W4_GPS)
+#define MUL_MAT_W4_OUT (MUL_MAT_W4_BLOCK * 4)
+
+        static __global__ void __launch_bounds__(MUL_MAT_W4_BLOCK* MUL_MAT_W4_BLOCK) mul_mat_w4a4_convrot_dp4a(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            int* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t packed_k) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4_OUT;
+            const uint32_t r0     = blockIdx.y * MUL_MAT_W4_ROWS + threadIdx.y;
+            const uint32_t r1     = r0 + MUL_MAT_W4_BLOCK;
+            const uint32_t lx     = threadIdx.x;
+            const uint32_t ly     = threadIdx.y;
+            const uint32_t idx    = ly * MUL_MAT_W4_BLOCK + lx;
+            const uint32_t dc     = idx & (MUL_MAT_W4_OUT - 1);
+            const uint32_t dw     = idx >> 6;
+
+            __shared__ int w_dec[2][MUL_MAT_W4_SLOTS * MUL_MAT_W4_OUT];
+            __shared__ int x_tile[2][MUL_MAT_W4_ROWS * MUL_MAT_W4_SLOTS];
+
+            const uint32_t w_words = packed_k / 4;
+            const uint32_t x_words = (packed_k * 2) / 4;
+            const uint32_t groups  = packed_k / 8;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            const uint32_t fl = lx & 7;
+            const uint32_t fr = (lx >> 3) ? r1 : r0;
+            const uint32_t tr = (lx >> 3) * MUL_MAT_W4_BLOCK + ly;
+            int my_sum        = 0;
+
+            {
+                const uint32_t ocol = o_base + dc;
+                int2 wv             = make_int2(0, 0);
+                if (ocol < n && dw * 2 < w_words) {
+                    wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + dw];
+                }
+                w_dec[0][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
+                w_dec[0][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+
+                const uint32_t gi = fl >> 1;
+                const uint32_t xb = gi * 4 + (fl & 1) * 2;
+                const int a       = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
+                const int b       = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
+                interleave_act_pair(a, b, &x_tile[0][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
+                my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
+            }
+            __syncthreads();
+
+            int acc0a = 0, acc0b = 0;
+            int acc1a = 0, acc1b = 0;
+            int acc2a = 0, acc2b = 0;
+            int acc3a = 0, acc3b = 0;
+            int acc4a = 0, acc4b = 0;
+            int acc5a = 0, acc5b = 0;
+            int acc6a = 0, acc6b = 0;
+            int acc7a = 0, acc7b = 0;
+
+            for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4_GPS) {
+                const uint32_t buf = (gs >> 2) & 1;
+
+                if (gs + MUL_MAT_W4_GPS < groups) {
+                    const uint32_t nbuf = buf ^ 1;
+                    const uint32_t ngs  = gs + MUL_MAT_W4_GPS;
+                    {
+                        const uint32_t ocol = o_base + dc;
+                        const uint32_t widx = ngs * 2 + dw * 2;
+                        int2 wv             = make_int2(0, 0);
+                        if (ocol < n && widx < w_words) {
+                            wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + ngs + dw];
+                        }
+                        w_dec[nbuf][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808;
+                        w_dec[nbuf][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808;
+                        const uint32_t gi                                         = ngs + (fl >> 1);
+                        const uint32_t xb                                         = gi * 4 + (fl & 1) * 2;
+                        const int a                                               = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0;
+                        const int b                                               = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0;
+                        interleave_act_pair(a, b, &x_tile[nbuf][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]);
+                        my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0);
+                    }
+                }
+#pragma unroll
+                for (uint32_t gl = 0; gl < MUL_MAT_W4_GPS; ++gl) {
+                    const int u0[4] = {
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 0],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 1],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 2],
+                        x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 3],
+                    };
+                    const int u1[4] = {
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 0],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 1],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 2],
+                        x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 3],
+                    };
+#pragma unroll
+                    for (uint32_t jp = 0; jp < 2; ++jp) {
+                        const int ue0     = jp == 0 ? u0[0] : u0[2];
+                        const int uo0     = jp == 0 ? u0[1] : u0[3];
+                        const int ue1     = jp == 0 ? u1[0] : u1[2];
+                        const int uo1     = jp == 0 ? u1[1] : u1[3];
+                        const uint32_t we = (gl * 4 + jp * 2) * MUL_MAT_W4_OUT;
+                        const uint32_t wo = we + MUL_MAT_W4_OUT;
+
+                        acc0a = dot4(w_dec[buf][we + lx], ue0, acc0a);
+                        acc1a = dot4(w_dec[buf][we + 16 + lx], ue0, acc1a);
+                        acc2a = dot4(w_dec[buf][we + 32 + lx], ue0, acc2a);
+                        acc3a = dot4(w_dec[buf][we + 48 + lx], ue0, acc3a);
+                        acc0b = dot4(w_dec[buf][wo + lx], uo0, acc0b);
+                        acc1b = dot4(w_dec[buf][wo + 16 + lx], uo0, acc1b);
+                        acc2b = dot4(w_dec[buf][wo + 32 + lx], uo0, acc2b);
+                        acc3b = dot4(w_dec[buf][wo + 48 + lx], uo0, acc3b);
+                        acc4a = dot4(w_dec[buf][we + lx], ue1, acc4a);
+                        acc5a = dot4(w_dec[buf][we + 16 + lx], ue1, acc5a);
+                        acc6a = dot4(w_dec[buf][we + 32 + lx], ue1, acc6a);
+                        acc7a = dot4(w_dec[buf][we + 48 + lx], ue1, acc7a);
+                        acc4b = dot4(w_dec[buf][wo + lx], uo1, acc4b);
+                        acc5b = dot4(w_dec[buf][wo + 16 + lx], uo1, acc5b);
+                        acc6b = dot4(w_dec[buf][wo + 32 + lx], uo1, acc6b);
+                        acc7b = dot4(w_dec[buf][wo + 48 + lx], uo1, acc7b);
+                    }
+                }
+                __syncthreads();
+            }
+
+            int v = my_sum;
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 4) % 8, 8);
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 2) % 8, 8);
+            v += shuffle_sum(v, (int(threadIdx.x) % 8 + 1) % 8, 8);
+            const int sum0 = shuffle_sum(v, 0, 16);
+            const int sum1 = shuffle_sum(v, 8, 16);
+
+            const int accv[2][4] = {
+                {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
+                {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
+            };
+            const int sumv[2] = {sum0, sum1};
+#pragma unroll
+            for (uint32_t t = 0; t < 2; ++t) {
+                const uint32_t rt = t == 0 ? r0 : r1;
+                if (rt < rows) {
+#pragma unroll
+                    for (uint32_t j = 0; j < 4; ++j) {
+                        const uint32_t oj = o_base + j * 16 + lx;
+                        if (oj < n) {
+                            y[(size_t)rt * n + oj] = accv[t][j] - 8 * sumv[t];
+                        }
+                    }
+                }
+            }
+        }
+
+#define MUL_MAT_W4A8_BLOCK 16
+#define MUL_MAT_W4A8_ROWS (2 * MUL_MAT_W4A8_BLOCK)
+#define MUL_MAT_W4A8_GPS 4
+#define MUL_MAT_W4A8_SLOTS (4 * MUL_MAT_W4A8_GPS)
+#define MUL_MAT_W4A8_OUT (MUL_MAT_W4A8_BLOCK * 4)
+
+#if defined(SD_CUDA_VENDOR_HIP) && defined(__HIP_DEVICE_COMPILE__)
+// __has_builtin checks the active device target's features, not just whether
+// the compiler recognizes the name. Prefer the non-replicated fragment layout.
+#if __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12)
+#define SD_W4_WMMA_GFX12
+#elif __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64)
+#define SD_W4_WMMA_GFX11
+#endif
+#endif
+
+// Keep each intrinsic paired with its fragment layout, including on newer GPUs.
+#if defined(SD_W4_WMMA_GFX11)
+        typedef int v4si32_w4a8 __attribute__((ext_vector_type(4)));
+
+        static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k,
+            const uint32_t r_base) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t idx    = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
+            const uint32_t wv     = idx >> 6;
+            const uint32_t lane   = idx & 63;
+            const uint32_t j      = lane & 15;
+            const uint32_t q      = lane >> 4;
+            const uint32_t ocol   = o_base + wv * 16 + j;
+            const uint32_t ra     = r_base + j;
+
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+            __shared__ float s_lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (idx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    s_lut[b] = fp8_e4m3(uint8_t(b));
+                }
+            }
+            __syncthreads();
+
+            const bool o_valid = ocol < n;
+            const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
+            const float scb    = s_cb;
+
+            float acc[4]           = {0.0f, 0.0f, 0.0f, 0.0f};
+            const v4si32_w4a8 zero = {};
+
+            int w0c = 0, w1c = 0, xc0 = 0, xc1 = 0, xc2 = 0, xc3 = 0, sc = 0;
+            int w0n = 0, w1n = 0, xn0 = 0, xn1 = 0, xn2 = 0, xn3 = 0, sn = 0;
+            {
+                const int* xr = xi + (size_t)ra * x_words;
+                xc0           = ra < rows ? xr[0] : 0;
+                xc1           = ra < rows ? xr[1] : 0;
+                xc2           = ra < rows ? xr[2] : 0;
+                xc3           = ra < rows ? xr[3] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words;
+                    w0c           = wp[0];
+                    w1c           = wp[1];
+                    sc            = s_rel[(size_t)ocol * groups];
+                }
+            }
+            if (groups > 1) {
+                const int* xr = xi + (size_t)ra * x_words + 4;
+                xn0           = ra < rows ? xr[0] : 0;
+                xn1           = ra < rows ? xr[1] : 0;
+                xn2           = ra < rows ? xr[2] : 0;
+                xn3           = ra < rows ? xr[3] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words + 2;
+                    w0n           = wp[0];
+                    w1n           = wp[1];
+                    sn            = s_rel[(size_t)ocol * groups + 1];
+                }
+            }
+
+            for (uint32_t g = 0; g < groups; ++g) {
+                v4si32_w4a8 bfrag = {};
+                float s           = 0.0f;
+                if (o_valid) {
+                    bfrag = {
+                        (lut[w0c & 0xFF] & 0xFFFF) | ((lut[(w0c >> 8) & 0xFF] & 0xFFFF) << 16),
+                        (lut[(w0c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w0c >> 24) & 0xFF] & 0xFFFF) << 16),
+                        (lut[w1c & 0xFF] & 0xFFFF) | ((lut[(w1c >> 8) & 0xFF] & 0xFFFF) << 16),
+                        (lut[(w1c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w1c >> 24) & 0xFF] & 0xFFFF) << 16),
+                    };
+
+                    s = scb * (s_chan * s_lut[sc & 0xFF]);
+                }
+                const v4si32_w4a8 afrag = {xc0, xc1, xc2, xc3};
+                const v4si32_w4a8 dfrag =
+                    __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64(true, afrag, true, bfrag, zero, false);
+
+                w0c = w0n;
+                w1c = w1n;
+                sc  = sn;
+                xc0 = xn0;
+                xc1 = xn1;
+                xc2 = xn2;
+                xc3 = xn3;
+                if (g + 2 < groups) {
+                    const int* xr = xi + (size_t)ra * x_words + (g + 2) * 4;
+                    xn0           = ra < rows ? xr[0] : 0;
+                    xn1           = ra < rows ? xr[1] : 0;
+                    xn2           = ra < rows ? xr[2] : 0;
+                    xn3           = ra < rows ? xr[3] : 0;
+                    if (o_valid) {
+                        const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
+                        w0n           = wp[0];
+                        w1n           = wp[1];
+                        sn            = s_rel[(size_t)ocol * groups + g + 2];
+                    }
+                }
+
+                acc[0] = fmaf((float)dfrag.x, s, acc[0]);
+                acc[1] = fmaf((float)dfrag.y, s, acc[1]);
+                acc[2] = fmaf((float)dfrag.z, s, acc[2]);
+                acc[3] = fmaf((float)dfrag.w, s, acc[3]);
+            }
+
+            if (o_valid) {
+#pragma unroll
+                for (uint32_t d = 0; d < 4; ++d) {
+                    const uint32_t r = r_base + 4 * d + q;
+                    if (r < rows) {
+                        float value = acc[d] * act_scales[r];
+                        if (bias != nullptr) {
+                            value += bias[ocol];
+                        }
+                        y[(size_t)r * n + ocol] = value;
+                    }
+                }
+            }
+        }
+#endif
+
+#if defined(SD_W4_WMMA_GFX12)
+        typedef int v4si32_w4a8_g12 __attribute__((ext_vector_type(4)));
+
+        static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma_gfx12(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k,
+            const uint32_t r_base) {
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t idx    = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x;
+            const uint32_t wv     = idx >> 6;
+            const uint32_t lane   = idx & 63;
+            const uint32_t j      = lane & 15;
+            const uint32_t h      = lane >> 4;
+            const uint32_t kh     = ((h & 1) << 1) | (h >> 1);
+            const uint32_t ocol   = o_base + wv * 16 + j;
+            const uint32_t ra     = r_base + j;
+
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+            __shared__ float s_lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (idx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8);
+                }
+#pragma unroll
+                for (int b = 0; b < 256; ++b) {
+                    s_lut[b] = fp8_e4m3(uint8_t(b));
+                }
+            }
+            __syncthreads();
+
+            const bool o_valid = ocol < n;
+            const float s_chan = o_valid ? s_channel[ocol] : 0.0f;
+            const float scb    = s_cb;
+
+            float acc[4]               = {0.0f, 0.0f, 0.0f, 0.0f};
+            const v4si32_w4a8_g12 zero = {};
+
+            int w0c = 0, w1c = 0, xc = 0, sc = 0;
+            int w0n = 0, w1n = 0, xn = 0, sn = 0;
+            {
+                const int* xr = xi + (size_t)ra * x_words;
+                xc            = ra < rows ? xr[kh] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words;
+                    w0c           = wp[0];
+                    w1c           = wp[1];
+                    sc            = s_rel[(size_t)ocol * groups];
+                }
+            }
+            if (groups > 1) {
+                const int* xr = xi + (size_t)ra * x_words;
+                xn            = ra < rows ? xr[4 + kh] : 0;
+                if (o_valid) {
+                    const int* wp = wi + (size_t)ocol * w_words + 2;
+                    w0n           = wp[0];
+                    w1n           = wp[1];
+                    sn            = s_rel[(size_t)ocol * groups + 1];
+                }
+            }
+
+            for (uint32_t g = 0; g < groups; ++g) {
+                int bfrag = 0;
+                float s   = 0.0f;
+                if (o_valid) {
+                    const int bwd = (kh & 2) ? w1c : w0c;
+                    const int bsh = (kh & 1) * 16;
+                    bfrag         = (lut[(bwd >> bsh) & 0xFF] & 0xFFFF) |
+                                    ((lut[(bwd >> (bsh + 8)) & 0xFF] & 0xFFFF) << 16);
+
+                    s = scb * (s_chan * s_lut[sc & 0xFF]);
+                }
+                const v4si32_w4a8_g12 dfrag =
+                    __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12(true, xc, true, bfrag, zero, false);
+
+                w0c = w0n;
+                w1c = w1n;
+                sc  = sn;
+                xc  = xn;
+                if (g + 2 < groups) {
+                    const int* xr = xi + (size_t)ra * x_words;
+                    xn            = ra < rows ? xr[(g + 2) * 4 + kh] : 0;
+                    if (o_valid) {
+                        const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2;
+                        w0n           = wp[0];
+                        w1n           = wp[1];
+                        sn            = s_rel[(size_t)ocol * groups + g + 2];
+                    }
+                }
+
+                acc[0] = fmaf((float)dfrag.x, s, acc[0]);
+                acc[1] = fmaf((float)dfrag.y, s, acc[1]);
+                acc[2] = fmaf((float)dfrag.z, s, acc[2]);
+                acc[3] = fmaf((float)dfrag.w, s, acc[3]);
+            }
+
+            if (o_valid) {
+#pragma unroll
+                for (uint32_t d = 0; d < 4; ++d) {
+                    const uint32_t r = r_base + 4 * kh + d;
+                    if (r < rows) {
+                        float value = acc[d] * act_scales[r];
+                        if (bias != nullptr) {
+                            value += bias[ocol];
+                        }
+                        y[(size_t)r * n + ocol] = value;
+                    }
+                }
+            }
+        }
+#endif
+
+        static __global__ void __launch_bounds__(MUL_MAT_W4A8_BLOCK* MUL_MAT_W4A8_BLOCK) mul_mat_w4a8_convrot(
+            const char* __restrict__ w,
+            const char* __restrict__ x,
+            const float* __restrict__ codebook,
+            const float* __restrict__ s_channel,
+            const uint8_t* __restrict__ s_rel,
+            const float* __restrict__ act_scales,
+            const float* __restrict__ bias,
+            float* __restrict__ y,
+            const uint32_t n,
+            const uint32_t rows,
+            const uint32_t k) {
+#if defined(SD_W4_WMMA_GFX11) || defined(SD_W4_WMMA_GFX12)
+            // Retain the 32-row launch grid shared with DP4A. Each WMMA wave
+            // computes a 16-row tile; the barrier protects the reused LUTs.
+            for (uint32_t tile = 0; tile < 2; ++tile) {
+#if defined(SD_W4_WMMA_GFX12)
+                mul_mat_w4a8_convrot_wmma_gfx12(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
+                                                blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
+#else
+                mul_mat_w4a8_convrot_wmma(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k,
+                                          blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK);
+#endif
+                __syncthreads();
+            }
+#else
+            const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT;
+            const uint32_t r0     = blockIdx.y * MUL_MAT_W4A8_ROWS + threadIdx.y;
+            const uint32_t r1     = r0 + MUL_MAT_W4A8_BLOCK;
+            const uint32_t lx     = threadIdx.x;
+            const uint32_t ly     = threadIdx.y;
+            const uint32_t idx    = ly * MUL_MAT_W4A8_BLOCK + lx;
+            const uint32_t dc     = idx & (MUL_MAT_W4A8_OUT - 1);
+            const uint32_t dw     = idx >> 6;
+
+            __shared__ int w_dec[2][MUL_MAT_W4A8_SLOTS * MUL_MAT_W4A8_OUT];
+            __shared__ int x_tile[2][MUL_MAT_W4A8_ROWS * MUL_MAT_W4A8_SLOTS];
+            __shared__ float s_dec[2][MUL_MAT_W4A8_GPS * MUL_MAT_W4A8_OUT];
+            __shared__ __align__(4) int8_t cb8[16];
+            __shared__ float s_cb;
+            __shared__ int lut[256];
+
+            const uint32_t w_words = k / 8;
+            const uint32_t x_words = k / 4;
+            const uint32_t groups  = k / 16;
+            const int* wi          = (const int*)w;
+            const int* xi          = (const int*)x;
+
+            if (ly == 0 && lx == 0) {
+                float amax = 0.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    amax = fmaxf(amax, fabsf(codebook[i]));
+                }
+
+                s_cb = amax / 127.0f;
+#pragma unroll
+                for (int i = 0; i < 16; ++i) {
+                    const int v = s_cb == 0.f ? 0 : __float2int_rn(codebook[i] / s_cb);
+                    cb8[i]      = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v));
+                }
+            }
+            __syncthreads();
+
+            lut[idx] = (cb8[idx & 0xF] & 0xFF) | ((cb8[idx >> 4] & 0xFF) << 8);
+            __syncthreads();
+
+#pragma unroll
+            for (uint32_t j = 0; j < 2; ++j) {
+                const uint32_t pw                              = dw * 2 + j;
+                const uint32_t ocol                            = o_base + dc;
+                const int wv                                   = (ocol < n && pw < w_words) ? wi[(size_t)ocol * w_words + pw] : 0;
+                w_dec[0][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
+                w_dec[0][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
+            }
+            {
+                const uint32_t ocol = o_base + dc;
+                float s             = 0.0f;
+                if (ocol < n && dw < groups) {
+                    s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + dw]);
+                }
+                s_dec[0][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
+            }
+            if (lx < 16) {
+                const uint32_t aidx = lx;
+                x_tile[0][ly * MUL_MAT_W4A8_SLOTS + lx] =
+                    (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
+                x_tile[0][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
+                    (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
+            }
+            __syncthreads();
+
+            float acc0 = 0.0f;
+            float acc1 = 0.0f;
+            float acc2 = 0.0f;
+            float acc3 = 0.0f;
+            float acc4 = 0.0f;
+            float acc5 = 0.0f;
+            float acc6 = 0.0f;
+            float acc7 = 0.0f;
+
+            for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4A8_GPS) {
+                const uint32_t buf = (gs >> 2) & 1;
+
+                if (gs + MUL_MAT_W4A8_GPS < groups) {
+                    const uint32_t nbuf = buf ^ 1;
+                    const uint32_t ngs  = gs + MUL_MAT_W4A8_GPS;
+#pragma unroll
+                    for (uint32_t j = 0; j < 2; ++j) {
+                        const uint32_t pw                                 = dw * 2 + j;
+                        const uint32_t ocol                               = o_base + dc;
+                        const uint32_t widx                               = ngs * 2 + pw;
+                        const int wv                                      = (ocol < n && widx < w_words) ? wi[(size_t)ocol * w_words + widx] : 0;
+                        w_dec[nbuf][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16);
+                        w_dec[nbuf][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16);
+                    }
+                    {
+                        const uint32_t ocol = o_base + dc;
+                        float s             = 0.0f;
+                        if (ocol < n && ngs + dw < groups) {
+                            s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + ngs + dw]);
+                        }
+                        s_dec[nbuf][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s;
+                    }
+                    if (lx < 16) {
+                        const uint32_t aidx = ngs * 4 + lx;
+                        x_tile[nbuf][ly * MUL_MAT_W4A8_SLOTS + lx] =
+                            (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0;
+                        x_tile[nbuf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] =
+                            (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0;
+                    }
+                }
+
+#pragma unroll
+                for (uint32_t gl = 0; gl < MUL_MAT_W4A8_GPS; ++gl) {
+                    const int u0[4] = {
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
+                        x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
+                    };
+                    const int u1[4] = {
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 0],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 1],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 2],
+                        x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 3],
+                    };
+                    int accg0 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u0[0], 0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u0[1], accg0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u0[2], accg0);
+                    accg0     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u0[3], accg0);
+                    int accg1 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u0[0], 0);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u0[1], accg1);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u0[2], accg1);
+                    accg1     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u0[3], accg1);
+                    int accg2 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u0[0], 0);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u0[1], accg2);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u0[2], accg2);
+                    accg2     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u0[3], accg2);
+                    int accg3 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u0[0], 0);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u0[1], accg3);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u0[2], accg3);
+                    accg3     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u0[3], accg3);
+                    int accg4 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u1[0], 0);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u1[1], accg4);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u1[2], accg4);
+                    accg4     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u1[3], accg4);
+                    int accg5 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u1[0], 0);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u1[1], accg5);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u1[2], accg5);
+                    accg5     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u1[3], accg5);
+                    int accg6 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u1[0], 0);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u1[1], accg6);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u1[2], accg6);
+                    accg6     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u1[3], accg6);
+                    int accg7 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u1[0], 0);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u1[1], accg7);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u1[2], accg7);
+                    accg7     = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u1[3], accg7);
+                    acc0      = fmaf((float)accg0, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc0);
+                    acc1      = fmaf((float)accg1, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc1);
+                    acc2      = fmaf((float)accg2, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc2);
+                    acc3      = fmaf((float)accg3, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc3);
+                    acc4      = fmaf((float)accg4, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc4);
+                    acc5      = fmaf((float)accg5, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc5);
+                    acc6      = fmaf((float)accg6, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc6);
+                    acc7      = fmaf((float)accg7, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc7);
+                }
+                __syncthreads();
+            }
+
+            const float accv[8] = {acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7};
+#pragma unroll
+            for (uint32_t t = 0; t < 2; ++t) {
+                const uint32_t rt = t == 0 ? r0 : r1;
+                if (rt < rows) {
+#pragma unroll
+                    for (uint32_t j = 0; j < 4; ++j) {
+                        const uint32_t oj = o_base + j * 16 + lx;
+                        if (oj < n) {
+                            float value = accv[t * 4 + j] * act_scales[rt];
+                            if (bias != nullptr) {
+                                value += bias[oj];
+                            }
+                            y[(size_t)rt * n + oj] = value;
+                        }
+                    }
+                }
+            }
+#endif
+        }
+
+        static __global__ void dequantize(const int32_t* accum, const float* xs, const float* ws, const float* bias, float* dst, int64_t n, int64_t total) {
+            const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+            if (i < total) {
+                float value = float(accum[i]) * xs[i / n] * ws[i % n];
+                dst[i]      = value + (bias ? bias[i % n] : 0.f);
+            }
+        }
+        static bool aligned(const ggml_tensor* t, size_t alignment) {
+            if (!t)
+                return true;
+            size_t offset = 0;
+            while (!t->data && t->view_src) {
+                offset += t->view_offs;
+                t = t->view_src;
+            }
+            return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
+        }
+        static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+            SDW4Operation op;
+            if (!sd_ggml_w4_describe(node, &op) || cuda_device_info(device).device < 0)
+                return false;
+            const int64_t k = op.input->ne[0], n = node->ne[0], rows = ggml_nrows(node);
+            if (k % 64 || k > INT_MAX || n > INT_MAX - 63 || rows > 65535 * 32 ||
+                ggml_nelements(node) > int64_t(INT_MAX) * 256 || !aligned(op.weight, 8) || !aligned(node, 4))
+                return false;
+            for (auto* t : node->src)
+                if (!aligned(t, 4))
+                    return false;
+            return true;
+        }
+        static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+            SDW4Operation op;
+            if (!sd_ggml_w4_describe(nodes[0], &op))
+                return GGML_STATUS_FAILED;
+            const auto* dst = nodes[0];
+            bool may_alias  = dst->data || dst->view_src;
+            if (dst->data && dst->buffer && !dst->view_src) {
+                may_alias = false;
+                for (const auto* input : dst->src) {
+                    if (!input)
+                        continue;
+                    size_t offset     = 0;
+                    const size_t size = ggml_nbytes(input);
+                    while (!input->data && input->view_src) {
+                        offset += input->view_offs;
+                        input = input->view_src;
+                    }
+                    if (!input->data)
+                        continue;
+                    const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(input->data) + offset;
+                    may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst));
+                }
+            }
+            *bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 ? 1 : 0) + (may_alias ? 1 : 0));
+            return GGML_STATUS_SUCCESS;
+        }
+        static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+            if (!src)
+                return false;
+            const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+            return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+        }
+        static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) {
+            auto* dst = nodes[0];
+            SDW4Operation op;
+            ggml_ext_gpu_context context;
+            if (!sd_ggml_w4_describe(dst, &op) || !cuda_context(e->backend, &context))
+                return GGML_STATUS_FAILED;
+            const auto stream = static_cast<CUDAStream>(context.stream);
+            const int64_t n = dst->ne[0], rows = ggml_nrows(dst), k = op.input->ne[0];
+            auto* scratch = e->workspace ? static_cast<float*>(ggml_backend_buffer_get_base(e->workspace)) : nullptr;
+            bool alias    = false;
+            for (auto* input : dst->src)
+                alias |= overlaps(dst, input);
+            auto* output     = alias ? scratch : static_cast<float*>(dst->data);
+            const auto* w    = static_cast<const char*>(op.weight->data);
+            const auto* x    = static_cast<const char*>(op.input->data);
+            const auto* xs   = reinterpret_cast<const float*>(x + k * ((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 dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16);
+            if (op.kind == SDW4Kind::W4A4) {
+                auto* accum = reinterpret_cast<int32_t*>(scratch + (alias ? ggml_nelements(dst) : 0));
+                mul_mat_w4a4_convrot_dp4a<<<grid, block, 0, stream>>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2));
+                if (cuda_last_error() != cuda_success)
+                    return GGML_STATUS_FAILED;
+                dequantize<<<unsigned((n * rows + 255) / 256), 256, 0, stream>>>(accum, xs, ws, bias, output, n, n * rows);
+            } else {
+                mul_mat_w4a8_convrot<<<grid, block, 0, stream>>>(w, x, ws,
+                                                                 static_cast<const float*>(op.s_channel->data), static_cast<const uint8_t*>(op.s_rel->data),
+                                                                 xs, bias, output, uint32_t(n), uint32_t(rows), uint32_t(k));
+            }
+            auto error = cuda_last_error();
+            if (error == cuda_success && alias)
+                error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+            return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+        }
+    }
+    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." SD_CUDA_NAME ".v1";
+        d.op                   = GGML_OP_CUSTOM;
+        d.custom_op            = sd_ggml_w4_op();
+        d.priority             = 100;
+        d.supports             = sd_w4_cuda::supports;
+        d.prepare              = sd_w4_cuda::prepare;
+        d.execute              = sd_w4_cuda::execute;
+        d.destroy              = [](void*, void*) {};
+        GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+    }
+}

src/core/ggml_extensions.cpp

@@ -301,6 +301,101 @@ ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int gro
     return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
 }
 
+struct W4Attributes {
+    int64_t shape[3];
+    SDW4Kind kind;
+    int32_t has_bias;
+};
+
+static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
+    if (!attributes || count < 3)
+        return false;
+    W4Attributes attr;
+    std::memcpy(&attr, attributes, sizeof(attr));
+    const bool codebook = attr.kind == SDW4Kind::W4A8;
+    if ((!codebook && attr.kind != SDW4Kind::W4A4) || (attr.has_bias != 0 && attr.has_bias != 1) ||
+        count != (codebook ? 5 : 3) + attr.has_bias)
+        return false;
+    for (int i = 0; i < count; ++i)
+        if (!inputs[i] || !ggml_is_contiguous(inputs[i]))
+            return false;
+    const auto* w = inputs[0];
+    const auto* x = inputs[1];
+    if (w->type != GGML_TYPE_I8 || w->ne[0] <= 0 || w->ne[0] > i8_size_limit / 2 ||
+        w->ne[1] <= 0 || w->ne[2] != 1 || w->ne[3] != 1)
+        return false;
+    const int64_t k = 2 * w->ne[0], n = w->ne[1];
+    int64_t rows = 1;
+    for (auto dim : attr.shape) {
+        if (dim <= 0 || dim > i8_size_limit / rows)
+            return false;
+        rows *= dim;
+    }
+    int64_t height;
+    if (!packed_height(k, rows, &height) || n > i8_size_limit / 4 / rows ||
+        x->type != GGML_TYPE_I8 || x->ne[0] != k || x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1 ||
+        !scale_matches(inputs[2], codebook ? 16 : n) ||
+        (attr.has_bias && !scale_matches(inputs[count - 1], n)))
+        return false;
+    if (codebook && (k % 16 || !scale_matches(inputs[3], n) || inputs[4]->type != GGML_TYPE_F8_E4M3 ||
+                     inputs[4]->ne[0] != k / 16 || inputs[4]->ne[1] != n || inputs[4]->ne[2] != 1 || inputs[4]->ne[3] != 1))
+        return false;
+    out->type  = GGML_TYPE_F32;
+    out->ne[0] = n;
+    std::memcpy(out->ne + 1, attr.shape, sizeof(attr.shape));
+    return true;
+}
+
+ggml_ext_op_t sd_ggml_w4_op() {
+    static const auto op = [] {
+        const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(W4Attributes), w4_infer};
+        auto result                 = ggml_ext_register_op(&desc);
+        GGML_ASSERT(result);
+        return result;
+    }();
+    return op;
+}
+
+bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) {
+    if (!node || ggml_ext_op_get(node) != sd_ggml_w4_op() || !ggml_is_contiguous(node))
+        return false;
+    W4Attributes attr;
+    std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr));
+    const bool codebook = attr.kind == SDW4Kind::W4A8;
+    const int count     = (codebook ? 5 : 3) + (attr.has_bias != 0);
+    ggml_ext_tensor_desc out;
+    if (!w4_infer(node->src, count, &attr, &out) || node->type != out.type ||
+        std::memcmp(node->ne, out.ne, sizeof(out.ne)))
+        return false;
+    *op = {node->src[0], node->src[1], node->src[2], codebook ? node->src[3] : nullptr,
+           codebook ? node->src[4] : nullptr, attr.has_bias ? node->src[count - 1] : nullptr, attr.kind};
+    return true;
+}
+
+ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* sc, ggml_tensor* sr, ggml_tensor* bias, SDW4Kind kind, int group) {
+    if (!w || !x || group < 0 || (group && !valid_group(group)))
+        return nullptr;
+    const int quant_group = group ? group : 1;
+    if (x->type == GGML_TYPE_F32)
+        x = ggml_quantize_i8_convrot(ctx, x, quant_group);
+    int packed_group;
+    if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group)
+        return nullptr;
+    const auto* logical    = x->src[0];
+    W4Attributes attr      = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0};
+    ggml_tensor* inputs[6] = {w, x, ws};
+    int count              = 3;
+    if (kind == SDW4Kind::W4A8) {
+        inputs[count++] = sc;
+        inputs[count++] = sr;
+    } else if (sc || sr) {
+        return nullptr;
+    }
+    if (bias)
+        inputs[count++] = bias;
+    return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &attr);
+}
+
 static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) {
     if (count != 3 || !attributes)
         return false;

src/core/ggml_extensions.h

@@ -39,6 +39,25 @@ bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
 ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size, ggml_tensor* input_scale = nullptr);
 ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
 
+enum class SDW4Kind : int32_t { W4A4 = 1,
+                                W4A8 = 2 };
+
+struct SDW4Operation {
+    const ggml_tensor* weight;
+    const ggml_tensor* input;
+    const ggml_tensor* weight_scale;
+    const ggml_tensor* s_channel;
+    const ggml_tensor* s_rel;
+    const ggml_tensor* bias;
+    SDW4Kind kind;
+};
+
+ggml_ext_op_t sd_ggml_w4_op();
+bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation);
+// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook
+// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node.
+ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* s_channel, ggml_tensor* s_rel, ggml_tensor* bias, SDW4Kind kind, int group_size);
+
 enum ggml_sage_attn_mode {
     GGML_SAGE_ATTN_AUTO = 0,
     GGML_SAGE_ATTN_2,

src/core/vulkan/CMakeLists.txt

@@ -53,6 +53,9 @@ sd_vulkan_shader(rms_modulate rms_modulate ON -O)
 sd_vulkan_shader(rms_rope_apply rms_rope_apply ON)
 sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16)
 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(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

@@ -11,7 +11,10 @@
 #include "mul_mat_i8_tensorwise.h"
 #include "mul_mat_i8_tensorwise_cm1.h"
 #include "mul_mat_i8_tensorwise_cm2.h"
+#include "mul_mat_w4a4_convrot.h"
+#include "mul_mat_w4a8_convrot.h"
 #include "quantize_i8_convrot.h"
+#include "quantize_i8_convrot_h64.h"
 #include "quantize_i8_convrot_register.h"
 #include "quantize_i8_convrot_subgroup.h"
 #include "rms_modulate_convrot.h"
@@ -111,6 +114,23 @@ namespace sd_vulkan_i8 {
                (cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] &&
                d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
     }
+    static bool w4_supports(ggml_backend_dev_t dev, const ggml_tensor* dst, void*) {
+        Device d;
+        SDW4Operation op;
+        if (!device_info(dev, d) || !sd_ggml_w4_describe(dst, &op) || !bounded(d, dst) ||
+            sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot) <= 4 || sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot) <= 4 ||
+            d.compute.shared_memory < (op.kind == SDW4Kind::W4A8 ? 16452u : 13312u) || d.info.max_push_constants < 16)
+            return false;
+        for (auto* t : dst->src)
+            if (t && (!bounded(d, t) || ggml_nbytes(t) > UINT32_MAX))
+                return false;
+        const uint64_t k = op.input->ne[0], n = dst->ne[0], rows = ggml_nrows(dst);
+        return k % 64 == 0 && k <= UINT32_MAX && n <= UINT32_MAX && rows <= UINT32_MAX &&
+               ggml_nelements(dst) <= UINT32_MAX &&
+               (k * ((rows + 3) / 4 * 4)) % d.info.buffer_alignment == 0 &&
+               (n + 63) / 64 <= d.info.max_group_count[0] && (rows + 31) / 32 <= d.info.max_group_count[1] &&
+               d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
+    }
     static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) {
         if (i + 2 >= count_nodes(graph) || !device.subgroup) {
             return false;
@@ -534,9 +554,12 @@ namespace sd_vulkan_i8 {
             status                     = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
         } else {
             ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
-            auto pipeline                      = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
+            int group;
-            bool subgroup                      = d.subgroup && k >= 1024;
+            if (!sd_ggml_convrot_describe(output, &group))
-            uint32_t width                     = k;
+                return GGML_STATUS_FAILED;
+            auto pipeline  = group == 64 ? SD_VK_PIPELINE(quantize_i8_convrot_h64, 3, QuantParameters) : SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
+            bool subgroup  = group == 256 && d.subgroup && k >= 1024;
+            uint32_t width = k;
             if (subgroup) {
                 pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters);
                 if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) {
@@ -653,6 +676,42 @@ namespace sd_vulkan_i8 {
         }
         return finish(e, d, output, dst, alias, status);
     }
+    static ggml_status multiply_w4(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* output) {
+        SDW4Operation op;
+        if (!sd_ggml_w4_describe(output, &op))
+            return GGML_STATUS_FAILED;
+        const uint32_t k = uint32_t(op.input->ne[0]), n = uint32_t(output->ne[0]), rows = uint32_t(ggml_nrows(output));
+        const size_t scale_offset                     = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
+        const ggml_ext_vulkan_binding xs              = {op.input, nullptr, scale_offset, size_t(rows) * 4};
+        std::vector<ggml_ext_vulkan_binding> bindings = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}};
+        if (op.kind == SDW4Kind::W4A8) {
+            bindings.push_back({op.s_channel});
+            bindings.push_back({op.s_rel});
+        }
+        bindings.push_back(xs);
+        bindings.push_back({op.bias ? op.bias : op.weight_scale});
+        ggml_ext_vulkan_binding dst;
+        bool alias;
+        if (!output_binding(e, d, output, bindings, dst, alias))
+            return GGML_STATUS_FAILED;
+        bindings.push_back(dst);
+        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};
+        auto status             = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
+        return finish(e, d, output, dst, alias, status);
+    }
+    static bool quant64_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
+        Device d;
+        int group;
+        return device_info(dev, d) && sd_ggml_convrot_describe(n, &group) && group == 64 &&
+               bounded(d, n) && bounded(d, n->src[0]) && n->ne[0] <= UINT32_MAX &&
+               ggml_nrows(n->src[0]) <= UINT32_MAX && d.info.max_group_invocations >= 16 && d.info.max_group_size[0] >= 16 &&
+               (uint64_t(n->ne[0]) * ((ggml_nrows(n->src[0]) + 3) / 4 * 4)) % d.info.buffer_alignment == 0;
+    }
     static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
         Device d;
         return device_info(dev, d) && can_quantize(d, n);
@@ -747,6 +806,8 @@ namespace sd_vulkan_i8 {
         Device d;
         if (!device_info(ggml_backend_get_device(e->backend), d))
             return GGML_STATUS_FAILED;
+        if (ggml_ext_op_get(nodes[0]) == sd_ggml_w4_op())
+            return multiply_w4(e, d, nodes[0]);
         return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count);
     }
 #undef SD_VK_PIPELINE
@@ -768,6 +829,13 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
     d.custom_op            = sd_ggml_convrot_op();
     d.supports             = quant_supports;
     GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.name     = "sd.quantize_i8_convrot.h64.vulkan.v1";
+    d.supports = quant64_supports;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+    d.name      = "sd.linear_w4_convrot.vulkan.v1";
+    d.custom_op = sd_ggml_w4_op();
+    d.supports  = w4_supports;
+    GGML_ASSERT(ggml_ext_register_kernel(device, &d));
     d.name      = "sd.linear_i8.vulkan.v1";
     d.op        = GGML_OP_CUSTOM;
     d.custom_op = sd_ggml_i8_op();

src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp

@@ -0,0 +1,189 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_integer_dot_product : require
+
+layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int 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 {
+    float data_d[];
+};
+
+shared int w_dec[2][16 * 64];
+shared int x_tile[2][32 * 16];
+shared int sum_p[2][16 * 8];
+
+void interleave_pair(int a, int b, uint buf, uint li) {
+    const uint a0       = uint(a);
+    const uint b0       = uint(b);
+    x_tile[buf][li + 0] = int((a0 & 0x000000FFu) | ((a0 >> 8) & 0x0000FF00u) | ((b0 << 16) & 0x00FF0000u) | ((b0 << 8) & 0xFF000000u));
+    x_tile[buf][li + 1] = int(((a0 >> 8) & 0x000000FFu) | ((a0 >> 16) & 0x0000FF00u) | ((b0 << 8) & 0x00FF0000u) | (b0 & 0xFF000000u));
+}
+
+void main() {
+    const uint out_base = gl_WorkGroupID.x * 64u;
+    const uint lx       = gl_LocalInvocationID.x;
+    const uint ly       = gl_LocalInvocationID.y;
+    const uint idx      = gl_LocalInvocationIndex;
+    const uint w_words  = p.k / 8u;
+    const uint a_words  = p.k / 4u;
+    const uint groups   = p.k / 16u;
+
+    const uint dc   = idx & 63u;
+    const uint dw   = idx >> 6u;
+    const uint ocol = out_base + dc;
+
+    const uint fl = lx & 7u;
+    const uint ft = lx >> 3u;
+    const uint fr = ft == 0u ? gl_WorkGroupID.y * 32u + ly : gl_WorkGroupID.y * 32u + ly + 16u;
+    const uint tr = ft * 16u + ly;
+    int my_sum    = 0;
+    int acc0a     = 0;
+    int acc0b     = 0;
+    int acc1a     = 0;
+    int acc1b     = 0;
+    int acc2a     = 0;
+    int acc2b     = 0;
+    int acc3a     = 0;
+    int acc3b     = 0;
+    int acc4a     = 0;
+    int acc4b     = 0;
+    int acc5a     = 0;
+    int acc5b     = 0;
+    int acc6a     = 0;
+    int acc6b     = 0;
+    int acc7a     = 0;
+    int acc7b     = 0;
+
+    {
+        [[unroll]] for (uint j = 0; j < 2; ++j) {
+            const uint pw                    = dw * 2 + j;
+            const uint widx                  = pw;
+            const uint wv                    = ocol < p.n && widx < w_words
+                                                   ? uint(data_w[ocol * w_words + widx])
+                                                   : 0u;
+            w_dec[0][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u);
+            w_dec[0][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u);
+        }
+        const uint gi = fl >> 1u;
+        const uint xb = gi * 4u + (fl & 1u) * 2u;
+        const int a   = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0;
+        const int b   = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0;
+        interleave_pair(a, b, 0, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u);
+        my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101)));
+    }
+    barrier();
+
+    for (uint gs = 0; gs < groups; gs += 4) {
+        const uint buf = (gs >> 2) & 1u;
+
+        if (gs + 4 < groups) {
+            const uint nbuf = buf ^ 1u;
+            const uint ngs  = gs + 4;
+            [[unroll]] for (uint j = 0; j < 2; ++j) {
+                const uint pw                       = dw * 2 + j;
+                const uint widx                     = ngs * 2 + pw;
+                const uint wv                       = ocol < p.n && widx < w_words
+                                                          ? uint(data_w[ocol * w_words + widx])
+                                                          : 0u;
+                w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u);
+                w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u);
+            }
+            const uint gi = ngs + (fl >> 1u);
+            const uint xb = gi * 4u + (fl & 1u) * 2u;
+            const int a   = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0;
+            const int b   = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0;
+            interleave_pair(a, b, nbuf, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u);
+            my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101)));
+        }
+
+        [[unroll]] for (uint gl = 0; gl < 4; ++gl) {
+            const int u0e  = x_tile[buf][ly * 16u + gl * 4u + 0u];
+            const int u0o  = x_tile[buf][ly * 16u + gl * 4u + 1u];
+            const int u0e2 = x_tile[buf][ly * 16u + gl * 4u + 2u];
+            const int u0o2 = x_tile[buf][ly * 16u + gl * 4u + 3u];
+            const int u1e  = x_tile[buf][(16u + ly) * 16u + gl * 4u + 0u];
+            const int u1o  = x_tile[buf][(16u + ly) * 16u + gl * 4u + 1u];
+            const int u1e2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 2u];
+            const int u1o2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 3u];
+            [[unroll]] for (uint jp = 0; jp < 2; ++jp) {
+                const int ue0 = jp == 0u ? u0e : u0e2;
+                const int uo0 = jp == 0u ? u0o : u0o2;
+                const int ue1 = jp == 0u ? u1e : u1e2;
+                const int uo1 = jp == 0u ? u1o : u1o2;
+                const uint we = (gl * 4u + jp * 2u) * 64u;
+                const uint wo = we + 64u;
+                acc0a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue0));
+                acc1a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue0));
+                acc2a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue0));
+                acc3a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue0));
+                acc0b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo0));
+                acc1b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo0));
+                acc2b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo0));
+                acc3b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo0));
+                acc4a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue1));
+                acc5a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue1));
+                acc6a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue1));
+                acc7a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue1));
+                acc4b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo1));
+                acc5b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo1));
+                acc6b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo1));
+                acc7b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo1));
+            }
+        }
+        barrier();
+    }
+
+    sum_p[ft][ly * 8u + fl] = my_sum;
+    barrier();
+    int sum0 = 0;
+    int sum1 = 0;
+    [[unroll]] for (uint i = 0; i < 8; ++i) {
+        sum0 += sum_p[0][ly * 8u + i];
+        sum1 += sum_p[1][ly * 8u + i];
+    }
+
+    const int accv[2][4] = {
+        {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b},
+        {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b},
+    };
+    const int sumv[2] = {sum0, sum1};
+    [[unroll]] for (uint t = 0; t < 2u; ++t) {
+        const uint rt = gl_WorkGroupID.y * 32u + ly + t * 16u;
+        if (rt < p.rows) {
+            [[unroll]] for (uint j = 0; j < 4; ++j) {
+                const uint oj = out_base + j * 16u + lx;
+                if (oj < p.n) {
+                    float value = float(accv[t][j] - 8 * sumv[t]) * data_as[rt] * data_ws[oj];
+                    if (p.has_bias != 0) {
+                        value += data_b[oj];
+                    }
+                    data_d[rt * p.n + oj] = value;
+                }
+            }
+        }
+    }
+}

src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp

@@ -0,0 +1,244 @@
+#version 460
+
+#extension GL_EXT_control_flow_attributes : require
+#extension GL_EXT_integer_dot_product : require
+
+layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint n;
+    uint rows;
+    uint k;
+    uint has_bias;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer W {
+    int data_w[];
+};
+layout(binding = 1, std430) readonly buffer A {
+    int data_a[];
+};
+layout(binding = 2, std430) readonly buffer CB {
+    float data_cb[];
+};
+layout(binding = 3, std430) readonly buffer SC {
+    float data_sc[];
+};
+layout(binding = 4, std430) readonly buffer SR {
+    uint data_sr[];
+};
+layout(binding = 5, std430) readonly buffer AS {
+    float data_as[];
+};
+layout(binding = 6, std430) readonly buffer B {
+    float data_b[];
+};
+layout(binding = 7, std430) writeonly buffer D {
+    float data_d[];
+};
+
+shared int w_dec[2][16 * 64];
+shared int x_tile[2][32 * 16];
+shared float s_dec[2][4 * 64];
+shared int lut[256];
+shared float s_lut[256];
+shared int cb8s[16];
+shared float s_cb;
+
+float f8_e4m3_to_f32(uint v) {
+    const uint sign = (v >> 7) & 1u;
+    const uint exp  = (v >> 3) & 0xFu;
+    const uint man  = v & 0x7u;
+    float result;
+    if (exp == 0xFu && man == 0x7u) {
+        result = uintBitsToFloat(0x7FC00000u);
+    } else if (exp == 0u) {
+        result = ldexp(float(man), -9);
+    } else {
+        result = ldexp(1.0 + float(man) / 8.0, int(exp) - 7);
+    }
+    return sign != 0u ? -result : result;
+}
+
+void main() {
+    const uint out_base = gl_WorkGroupID.x * 64u;
+    const uint lx       = gl_LocalInvocationID.x;
+    const uint ly       = gl_LocalInvocationID.y;
+
+    const uint row0     = gl_WorkGroupID.y * 32u + ly;
+    const uint row1     = row0 + 16u;
+    const uint idx      = gl_LocalInvocationIndex;
+    const uint w_words  = p.k / 8u;
+    const uint a_words  = p.k / 4u;
+    const uint groups   = p.k / 16u;
+    const uint sr_words = groups / 4u;
+
+    const uint dc   = idx & 63u;
+    const uint dw   = idx >> 6u;
+    const uint ocol = out_base + dc;
+    float sum0      = 0.0;
+    float sum1      = 0.0;
+    float sum2      = 0.0;
+    float sum3      = 0.0;
+    float sum4      = 0.0;
+    float sum5      = 0.0;
+    float sum6      = 0.0;
+    float sum7      = 0.0;
+
+    if (idx == 0u) {
+        float amax = 0.0;
+        for (uint i = 0; i < 16; ++i) {
+            amax = max(amax, abs(data_cb[i]));
+        }
+        s_cb = amax / 127.0;
+        for (uint i = 0; i < 16; ++i) {
+            const int v = (s_cb == 0.0 ? 0 : int(roundEven(data_cb[i] / s_cb)));
+            cb8s[i]     = v < -127 ? -127 : (v > 127 ? 127 : v);
+        }
+    }
+    barrier();
+
+    lut[idx]   = (cb8s[idx & 0xFu] & 0xFF) | (cb8s[idx >> 4] << 8);
+    s_lut[idx] = f8_e4m3_to_f32(idx);
+    barrier();
+
+    {
+        [[unroll]] for (uint j = 0; j < 2; ++j) {
+            const uint pw   = dw * 2 + j;
+            const uint widx = pw;
+            const uint wv   = ocol < p.n && widx < w_words
+                                  ? uint(data_w[ocol * w_words + widx])
+                                  : 0u;
+
+            w_dec[0][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16));
+            w_dec[0][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16));
+        }
+        {
+            const uint g = dw;
+            float s      = 0.0;
+            if (ocol < p.n && g < groups) {
+                const uint sr_word = data_sr[ocol * sr_words + g / 4u];
+                s                  = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu];
+            }
+            s_dec[0][dw * 64 + dc] = s_cb * s;
+        }
+        if (lx < 16u) {
+            const uint aidx                 = lx;
+            x_tile[0][ly * 16 + lx]         = row0 < p.rows && aidx < a_words
+                                                  ? data_a[row0 * a_words + aidx]
+                                                  : 0;
+            x_tile[0][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words
+                                                  ? data_a[row1 * a_words + aidx]
+                                                  : 0;
+        }
+    }
+    barrier();
+
+    for (uint gs = 0; gs < groups; gs += 4) {
+        const uint buf = (gs >> 2) & 1u;
+
+        if (gs + 4 < groups) {
+            const uint nbuf = buf ^ 1u;
+            const uint ngs  = gs + 4;
+            [[unroll]] for (uint j = 0; j < 2; ++j) {
+                const uint pw                       = dw * 2 + j;
+                const uint widx                     = ngs * 2 + pw;
+                const uint wv                       = ocol < p.n && widx < w_words
+                                                          ? uint(data_w[ocol * w_words + widx])
+                                                          : 0u;
+                w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16));
+                w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16));
+            }
+            {
+                const uint g = ngs + dw;
+                float s      = 0.0;
+                if (ocol < p.n && g < groups) {
+                    const uint sr_word = data_sr[ocol * sr_words + g / 4u];
+                    s                  = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu];
+                }
+                s_dec[nbuf][dw * 64 + dc] = s_cb * s;
+            }
+            if (lx < 16u) {
+                const uint aidx                    = ngs * 4 + lx;
+                x_tile[nbuf][ly * 16 + lx]         = row0 < p.rows && aidx < a_words
+                                                         ? data_a[row0 * a_words + aidx]
+                                                         : 0;
+                x_tile[nbuf][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words
+                                                         ? data_a[row1 * a_words + aidx]
+                                                         : 0;
+            }
+        }
+
+        float partial0 = 0.0;
+        float partial1 = 0.0;
+        float partial2 = 0.0;
+        float partial3 = 0.0;
+        float partial4 = 0.0;
+        float partial5 = 0.0;
+        float partial6 = 0.0;
+        float partial7 = 0.0;
+        {
+            [[unroll]] for (uint gl = 0; gl < 4; ++gl) {
+                int accg0 = 0;
+                int accg1 = 0;
+                int accg2 = 0;
+                int accg3 = 0;
+                int accg4 = 0;
+                int accg5 = 0;
+                int accg6 = 0;
+                int accg7 = 0;
+                [[unroll]] for (uint q = 0; q < 4; ++q) {
+                    const int a0 = x_tile[buf][ly * 16 + gl * 4 + q];
+                    const int a1 = x_tile[buf][(16u + ly) * 16 + gl * 4 + q];
+                    accg0 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a0));
+                    accg1 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a0));
+                    accg2 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a0));
+                    accg3 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a0));
+                    accg4 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a1));
+                    accg5 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a1));
+                    accg6 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a1));
+                    accg7 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a1));
+                }
+                const float s0 = s_dec[buf][gl * 64 + lx];
+                const float s1 = s_dec[buf][gl * 64 + 16u + lx];
+                const float s2 = s_dec[buf][gl * 64 + 32u + lx];
+                const float s3 = s_dec[buf][gl * 64 + 48u + lx];
+                partial0 += float(accg0) * s0;
+                partial1 += float(accg1) * s1;
+                partial2 += float(accg2) * s2;
+                partial3 += float(accg3) * s3;
+                partial4 += float(accg4) * s0;
+                partial5 += float(accg5) * s1;
+                partial6 += float(accg6) * s2;
+                partial7 += float(accg7) * s3;
+            }
+        }
+        sum0 += partial0;
+        sum1 += partial1;
+        sum2 += partial2;
+        sum3 += partial3;
+        sum4 += partial4;
+        sum5 += partial5;
+        sum6 += partial6;
+        sum7 += partial7;
+        barrier();
+    }
+
+    const float sumv[8] = {sum0, sum1, sum2, sum3, sum4, sum5, sum6, sum7};
+    [[unroll]] for (uint t = 0; t < 2u; ++t) {
+        const uint rt = t == 0u ? row0 : row1;
+        if (rt < p.rows) {
+            [[unroll]] for (uint j = 0; j < 4; ++j) {
+                const uint oj = out_base + j * 16u + lx;
+                if (oj < p.n) {
+                    float value = sumv[t * 4u + j] * data_as[rt];
+                    if (p.has_bias != 0) {
+                        value += data_b[oj];
+                    }
+                    data_d[rt * p.n + oj] = value;
+                }
+            }
+        }
+    }
+}

src/core/vulkan/shaders/quantize_i8_convrot_h64.comp

@@ -0,0 +1,115 @@
+#version 450
+
+#extension GL_EXT_control_flow_attributes : require
+
+#ifndef GROUP_SIZE
+#define GROUP_SIZE 256
+#endif
+#if GROUP_SIZE != 64 && GROUP_SIZE != 256
+#error unsupported GROUP_SIZE
+#endif
+#define GROUP_THREADS (GROUP_SIZE / 4)
+
+#if GROUP_SIZE == 256
+#define HADAMARD_NORM (1.0 / 16.0)
+#else
+#define HADAMARD_NORM (1.0 / 8.0)
+#endif
+
+layout(local_size_x = GROUP_THREADS, local_size_y = 1, local_size_z = 1) in;
+
+layout(push_constant) uniform parameter {
+    uint k;
+    uint rows;
+}
+p;
+
+layout(binding = 0, std430) readonly buffer A {
+    float data_a[];
+};
+layout(binding = 1, std430) writeonly buffer D {
+    uint data_d[];
+};
+layout(binding = 2, std430) writeonly buffer S {
+    float data_s[];
+};
+
+shared float values[GROUP_SIZE];
+shared float maxima[GROUP_THREADS];
+
+void transform_group(const uint row, const uint group) {
+    const uint tid         = gl_LocalInvocationID.x;
+    const uint base_offset = row * p.k + group * GROUP_SIZE;
+
+    [[unroll]] for (uint i = 0; i < 4; ++i) {
+        const uint index = tid + i * GROUP_THREADS;
+        values[index]    = data_a[base_offset + index] * HADAMARD_NORM;
+    }
+    barrier();
+
+    [[unroll]] for (uint stride = 1; stride < GROUP_SIZE; stride *= 4) {
+        const uint base = (tid / stride) * 4 * stride + tid % stride;
+        const uint i0   = base;
+        const uint i1   = i0 + stride;
+        const uint i2   = i1 + stride;
+        const uint i3   = i2 + stride;
+        const float a   = values[i0];
+        const float b   = values[i1];
+        const float c   = values[i2];
+        const float d   = values[i3];
+        values[i0]      = a + b + c - d;
+        values[i1]      = a + b - c + d;
+        values[i2]      = a - b + c + d;
+        values[i3]      = -a + b + c + d;
+        barrier();
+    }
+}
+
+uint quantize_value(const float value, const float inv_scale) {
+    const int quantized = clamp(int(roundEven(value * inv_scale)), -127, 127);
+    return uint(quantized) & 0xffu;
+}
+
+void quantize_row(const uint row) {
+    const uint tid    = gl_LocalInvocationID.x;
+    const uint groups = p.k / GROUP_SIZE;
+    float local_max   = 0.0;
+
+    for (uint group = 0; group < groups; ++group) {
+        transform_group(row, group);
+        [[unroll]] for (uint i = 0; i < 4; ++i) {
+            local_max = max(local_max, abs(values[tid + i * GROUP_THREADS]));
+        }
+        barrier();
+    }
+
+    maxima[tid] = local_max;
+    barrier();
+    [[unroll]] for (uint stride = GROUP_THREADS / 2; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            maxima[tid] = max(maxima[tid], maxima[tid + stride]);
+        }
+        barrier();
+    }
+
+    if (tid == 0) {
+        data_s[row] = maxima[0] / 127.0;
+    }
+    barrier();
+
+    const float row_scale = maxima[0] / 127.0;
+    const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
+    for (uint group = 0; group < groups; ++group) {
+        transform_group(row, group);
+        const uint index                                   = tid * 4;
+        const uint packed                                  = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
+        data_d[(row * p.k + group * GROUP_SIZE) / 4 + tid] = packed;
+        barrier();
+    }
+}
+
+void main() {
+    for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) {
+        quantize_row(row);
+    }
+}

src/model/adapter/ip_adapter.hpp

@@ -141,7 +141,7 @@ namespace IPAdapter {
                 }
                 auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight");
                 if (proj_in_iter != tensor_storage_map.end()) {
-                    embed_dim = proj_in_iter->second.ne[0];
+                    embed_dim = proj_in_iter->second.logical_ne0();
                 }
                 auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight");
                 if (proj_out_iter != tensor_storage_map.end()) {
@@ -168,7 +168,7 @@ namespace IPAdapter {
                 }
                 auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight");
                 if (proj_iter != tensor_storage_map.end()) {
-                    clip_dim = proj_iter->second.ne[0];
+                    clip_dim = proj_iter->second.logical_ne0();
                     out_dim  = proj_iter->second.ne[1];
                 }
                 num_tokens = out_dim / ctx_dim;

src/model/adapter/lora.hpp

@@ -1127,6 +1127,16 @@ public:
                                     ggml_tensor* output,
                                     const std::string& prefix,
                                     WeightAdapter::ForwardParams forward_params) override {
+        // Packed INT4 stores half the logical input width. Adapter shape checks
+        // and reshapes need the logical weight dimensions, without decoding it.
+        ggml_tensor logical_weight = *w;
+        if (forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_LINEAR &&
+            w->type == GGML_TYPE_I8 && x->ne[0] == 2 * w->ne[0]) {
+            logical_weight.ne[0] = x->ne[0];
+            for (int i = 1; i < GGML_MAX_DIMS; ++i)
+                logical_weight.nb[i] = logical_weight.nb[i - 1] * logical_weight.ne[i - 1];
+            w = &logical_weight;
+        }
         for (auto& lora_model : lora_models) {
             ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false);
             if (weight_diff != nullptr) {

src/model/common/block.hpp

@@ -319,7 +319,7 @@ protected:
             tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) {
             has_ip            = true;
             int64_t inner_dim = d_head * n_head;
-            int64_t ip_dim    = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0];
+            int64_t ip_dim    = tensor_storage_map.at(prefix + "to_k_ip.weight").logical_ne0();
             blocks["to_k_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
             blocks["to_v_ip"] = std::shared_ptr<GGMLBlock>(new Linear(ip_dim, inner_dim, false));
         }

src/model/common/ggml_block.hpp

@@ -148,6 +148,7 @@ protected:
     bool force_f32;
     bool force_prec_f32;
     bool has_weight_scale       = false;
+    int w4_convrot_kind         = 0;
     bool int8_convrot           = false;
     int int8_convrot_group_size = 0;
     float scale;
@@ -158,16 +159,21 @@ protected:
         has_weight_scale        = false;
         int8_convrot            = false;
         int8_convrot_group_size = 0;
+        auto weight_storage     = tensor_storage_map.find(prefix + "weight");
+        w4_convrot_kind         = weight_storage != tensor_storage_map.end() ? weight_storage->second.w4_convrot_kind : 0;
         enum ggml_type wtype    = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32);
-        if (in_features % ggml_blck_size(wtype) != 0 || force_f32) {
+        if (!w4_convrot_kind && (in_features % ggml_blck_size(wtype) != 0 || force_f32)) {
             wtype = GGML_TYPE_F32;
         }
-        params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features);
+        if (w4_convrot_kind) {
+            GGML_ASSERT(wtype == GGML_TYPE_I8 && in_features % 2 == 0);
+            GGML_ASSERT(weight_storage->second.ne[0] == in_features / 2 && weight_storage->second.ne[1] == out_features);
+        }
+        params["weight"] = ggml_new_tensor_2d(ctx, wtype, w4_convrot_kind ? in_features / 2 : in_features, out_features);
         if (bias) {
             enum ggml_type wtype = GGML_TYPE_F32;
             params["bias"]       = ggml_new_tensor_1d(ctx, wtype, out_features);
         }
-        auto weight_storage           = tensor_storage_map.find(prefix + "weight");
         const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
         auto weight_scale_storage     = tensor_storage_map.find(prefix + "weight_scale");
         if (weight_scale_storage != tensor_storage_map.end()) {
@@ -176,9 +182,16 @@ protected:
             params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
             has_weight_scale       = true;
         }
-        if (is_int8_tensorwise) {
+        if (w4_convrot_kind == 2) {
+            params["weight_codebook"]  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16);
+            params["weight_s_channel"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
+#ifndef SD_USE_UPSTREAM_GGML
+            params["weight_s_rel"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F8_E4M3, in_features / 16, out_features);
+#endif
+        }
+        if (is_int8_tensorwise || w4_convrot_kind) {
             GGML_ASSERT(wtype == GGML_TYPE_I8);
-            GGML_ASSERT(has_weight_scale);
+            GGML_ASSERT(w4_convrot_kind == 2 || has_weight_scale);
             int8_convrot            = weight_storage->second.int8_convrot;
             int8_convrot_group_size = weight_storage->second.int8_convrot_group_size;
         }
@@ -259,13 +272,31 @@ public:
                 }
             }
 #endif
-            out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+#ifndef SD_USE_UPSTREAM_GGML
-                                                x,
+            if (w4_convrot_kind) {
-                                                w,
+                if (scale != 1.f)
-                                                weight_scale,
+                    x = ggml_ext_scale(ctx->ggml_ctx, x, scale);
-                                                b,
+                out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x,
-                                                int8_convrot ? int8_convrot_group_size : 0,
+                                              w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale,
-                                                scale);
+                                              w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr,
+                                              w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr,
+                                              scale == 1.f ? b : nullptr, static_cast<SDW4Kind>(w4_convrot_kind),
+                                              int8_convrot ? int8_convrot_group_size : 0);
+                GGML_ASSERT(out);
+                if (scale != 1.f) {
+                    out = ggml_ext_scale(ctx->ggml_ctx, out, 1.f / scale);
+                    if (b)
+                        out = ggml_add_inplace(ctx->ggml_ctx, out, b);
+                }
+            } else
+#endif
+                out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
+                                                    x,
+                                                    w,
+                                                    weight_scale,
+                                                    b,
+                                                    int8_convrot ? int8_convrot_group_size : 0,
+                                                    scale);
             if (ctx->weight_adapter) {
                 WeightAdapter::ForwardParams forward_params;
                 forward_params.op_type               = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;

src/model/diffusion/boogu.hpp

@@ -72,10 +72,10 @@ namespace Boogu {
                 }
                 if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
                     int64_t patch_area = config.patch_size * config.patch_size;
-                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.in_channels = tensor_storage.logical_ne0() / patch_area;
                     config.hidden_size = tensor_storage.ne[1];
                 } else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) {
-                    config.instruction_feat_dim = tensor_storage.ne[0];
+                    config.instruction_feat_dim = tensor_storage.logical_ne0();
                     config.hidden_size          = tensor_storage.ne[1];
                 } else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
                     detected_head_dim = tensor_storage.ne[0];

src/model/diffusion/ernie_image.hpp

@@ -42,7 +42,7 @@ namespace ErnieImage {
                     config.in_channels = tensor_storage.ne[2];
                     config.hidden_size = tensor_storage.ne[3];
                 } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) {
-                    config.text_in_dim = tensor_storage.ne[0];
+                    config.text_in_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) {
                     detected_head_dim = tensor_storage.ne[0];
                 } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) {

src/model/diffusion/fibo2.hpp

@@ -31,12 +31,12 @@ namespace Fibo2 {
                 }
                 auto key = name.substr(prefix.size() + 1);
                 if (key == "x_embedder.weight") {
-                    config.channels = tensor.ne[0];
+                    config.channels = tensor.logical_ne0();
                     config.dim      = tensor.ne[1];
                 } else if (key == "perceiver.latent_queries") {
                     config.max_gist = tensor.ne[1];
                 } else if (key == "perceiver.layers.0.cross_attention.to_k.weight" || key == "perceiver.layers.0.cross_attention.to_kv.weight") {
-                    config.text_dim = tensor.ne[0];
+                    config.text_dim = tensor.logical_ne0();
                     config.fused_kv = key == "perceiver.layers.0.cross_attention.to_kv.weight";
                 } else if (key == "layers.0.attention.norm_q.weight") {
                     config.head_dim = tensor.ne[0];

src/model/diffusion/flux.hpp

@@ -158,7 +158,7 @@ namespace Flux {
                     }
                 }
                 if (ends_with(name, "txt_in.weight")) {
-                    config.context_in_dim = tensor_storage.ne[0];
+                    config.context_in_dim = tensor_storage.logical_ne0();
                     config.hidden_size    = tensor_storage.ne[1];
                 }
                 if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) {

src/model/diffusion/flux3.hpp

@@ -27,12 +27,12 @@ namespace Flux3 {
                 std::string key = name.substr(prefix.size() + 1);
                 if (key == "txt_in.weight") {
                     detected           = true;
-                    config.context_dim = tensor.ne[0];
+                    config.context_dim = tensor.logical_ne0();
                     config.hidden_size = tensor.ne[1];
                 } else if (key == "emb_in.image.weight") {
-                    config.channels = tensor.ne[0];
+                    config.channels = tensor.logical_ne0();
                 } else if (key == "single_blocks.0.mlp_out.weight") {
-                    config.mlp_dim = tensor.ne[0];
+                    config.mlp_dim = tensor.logical_ne0();
                 } else if (key == "ref_time_in.in_layer.weight") {
                     config.flow_map = true;
                 }

src/model/diffusion/krea2.hpp

@@ -92,11 +92,11 @@ namespace Krea2 {
                     continue;
                 }
                 if (ends_with(name, "first.weight") && tensor_storage.n_dims == 2) {
-                    config.in_channels  = tensor_storage.ne[0] / (config.patch_size * config.patch_size);
+                    config.in_channels  = tensor_storage.logical_ne0() / (config.patch_size * config.patch_size);
                     config.out_channels = config.in_channels;
                     config.features     = tensor_storage.ne[1];
                 } else if (ends_with(name, "blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) {
-                    detected_head_dim = tensor_storage.ne[0];
+                    detected_head_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) {
                     if (detected_head_dim > 0) {
                         config.heads = tensor_storage.ne[1] / detected_head_dim;
@@ -106,11 +106,11 @@ namespace Krea2 {
                         config.kv_heads = tensor_storage.ne[1] / detected_head_dim;
                     }
                 } else if (ends_with(name, "txtfusion.projector.weight") && tensor_storage.n_dims == 2) {
-                    config.text_layers = tensor_storage.ne[0];
+                    config.text_layers = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "txtfusion.layerwise_blocks.0.prenorm.scale") && tensor_storage.n_dims == 1) {
-                    config.text_dim = tensor_storage.ne[0];
+                    config.text_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) {
-                    detected_text_head_dim = tensor_storage.ne[0];
+                    detected_text_head_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) {
                     if (detected_text_head_dim > 0) {
                         config.text_heads = tensor_storage.ne[1] / detected_text_head_dim;

src/model/diffusion/lens.hpp

@@ -36,13 +36,13 @@ namespace Lens {
                     continue;
                 }
                 if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) {
-                    config.in_channels = tensor_storage.ne[0];
+                    config.in_channels = tensor_storage.logical_ne0();
                     int64_t inner_dim  = tensor_storage.ne[1];
                     if (config.attention_head_dim > 0) {
                         config.num_attention_heads = inner_dim / config.attention_head_dim;
                     }
                 } else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) {
-                    config.selected_layer_count = static_cast<int>(tensor_storage.ne[0] / config.joint_attention_dim);
+                    config.selected_layer_count = static_cast<int>(tensor_storage.logical_ne0() / config.joint_attention_dim);
                 } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) {
                     int64_t patch_area  = config.patch_size * config.patch_size;
                     config.out_channels = tensor_storage.ne[1] / patch_area;

src/model/diffusion/lingbot_video.hpp

@@ -63,16 +63,16 @@ namespace LingBotVideo {
                 }
 
                 if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) {
-                    int64_t patch_dim    = tensor_storage.ne[0];
+                    int64_t patch_dim    = tensor_storage.logical_ne0();
                     config.hidden_size   = tensor_storage.ne[1];
                     int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w;
                     if (patch_dim % patch_volume == 0) {
                         config.in_channels = patch_dim / patch_volume;
                     }
                 } else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
-                    config.text_dim = tensor_storage.ne[0];
+                    config.text_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
-                    config.freq_dim = tensor_storage.ne[0];
+                    config.freq_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) {
                     int64_t head_dim = tensor_storage.ne[0];
                     if (head_dim > 0) {

src/model/diffusion/llada_image.hpp

@@ -74,13 +74,13 @@ namespace LLaDAImage {
                 }
                 if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
                     int64_t patch_area = config.patch_size * config.patch_size;
-                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.in_channels = tensor_storage.logical_ne0() / patch_area;
                     config.hidden_size = tensor_storage.ne[1];
                 } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
-                    config.cap_feat_dim = tensor_storage.ne[0];
+                    config.cap_feat_dim = tensor_storage.logical_ne0();
                     config.hidden_size  = tensor_storage.ne[1];
                 } else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) {
-                    config.semantic_feat_dim = tensor_storage.ne[0];
+                    config.semantic_feat_dim = tensor_storage.logical_ne0();
                 } else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) {
                     detected_q_dim = tensor_storage.ne[1];
                 } else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) {

src/model/diffusion/ltxv.hpp

@@ -168,7 +168,7 @@ namespace LTXV {
             LTXAVConfig config;
             auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight");
             if (patchify_proj_iter != tensor_storage_map.end()) {
-                config.in_channels         = patchify_proj_iter->second.ne[0];
+                config.in_channels         = patchify_proj_iter->second.logical_ne0();
                 config.hidden_size         = patchify_proj_iter->second.ne[1];
                 int64_t video_heads        = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32);
                 auto attn_layout           = infer_attention_layout(config.hidden_size, video_heads);
@@ -178,7 +178,7 @@ namespace LTXV {
 
             auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight");
             if (audio_patchify_proj_iter != tensor_storage_map.end()) {
-                config.audio_in_channels         = audio_patchify_proj_iter->second.ne[0];
+                config.audio_in_channels         = audio_patchify_proj_iter->second.logical_ne0();
                 config.audio_hidden_size         = audio_patchify_proj_iter->second.ne[1];
                 config.audio_out_channels        = config.audio_in_channels;
                 int64_t audio_heads              = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32);
@@ -198,11 +198,11 @@ namespace LTXV {
 
             auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight");
             if (attn2_iter != tensor_storage_map.end()) {
-                config.cross_attention_dim = attn2_iter->second.ne[0];
+                config.cross_attention_dim = attn2_iter->second.logical_ne0();
             }
             auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight");
             if (audio_attn2_iter != tensor_storage_map.end()) {
-                config.audio_cross_attention_dim = audio_attn2_iter->second.ne[0];
+                config.audio_cross_attention_dim = audio_attn2_iter->second.logical_ne0();
             }
             if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) {
                 config.cross_attention_adaln = true;

src/model/diffusion/minimax_h3.hpp

@@ -73,16 +73,16 @@ namespace MiniMaxH3 {
             };
 
             if (const auto* weight = find("video_patch_proj.weight")) {
-                config.video_latent_channels = weight->ne[0] / 4;
+                config.video_latent_channels = weight->logical_ne0() / 4;
                 config.hidden_size           = weight->ne[1];
             }
             if (const auto* weight = find("audio_patch_proj.weight")) {
-                config.audio_latent_channels = weight->ne[0];
+                config.audio_latent_channels = weight->logical_ne0();
             }
             config.num_layers               = count_blocks(tensors, prefix + ".blocks.");
             config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks.");
             if (const auto* weight = find("blocks.0.attn.q_norm.weight")) {
-                config.attention_head_dim = weight->ne[0];
+                config.attention_head_dim = weight->logical_ne0();
             }
             if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) {
                 config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim);
@@ -91,14 +91,14 @@ namespace MiniMaxH3 {
                 config.ffn_hidden_size = weight->ne[1] / 2;
             }
             if (const auto* weight = find("condition_proj.weight")) {
-                config.text_dim = weight->ne[0];
+                config.text_dim = weight->logical_ne0();
             }
             if (const auto* table = find("adaln_t_table")) {
                 config.time_embed_dim   = table->ne[0];
                 config.adaln_curve_grid = table->ne[1];
             } else {
                 if (const auto* weight = find("time_embedder.proj_in.weight")) {
-                    config.timestep_input_dim     = weight->ne[0];
+                    config.timestep_input_dim     = weight->logical_ne0();
                     config.time_embed_hidden_size = weight->ne[1];
                 }
                 if (const auto* weight = find("time_embedder.proj_out.weight")) {

src/model/diffusion/minit2i.hpp

@@ -59,14 +59,14 @@ namespace MiniT2I {
                     config.pca_channels = tensor_storage.ne[2];
                     config.hidden_size  = tensor_storage.ne[3];
                 } else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) {
-                    config.txt_input_size  = tensor_storage.ne[0];
+                    config.txt_input_size  = tensor_storage.logical_ne0();
                     config.txt_hidden_size = tensor_storage.ne[1];
                 } else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) {
                     config.cond_vec_size = tensor_storage.ne[1];
                 } else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) {
                     int64_t inner3     = tensor_storage.ne[1];
                     int64_t inner      = inner3 / 3;
-                    config.hidden_size = tensor_storage.ne[0];
+                    config.hidden_size = tensor_storage.logical_ne0();
                     if (config.hidden_size == 768) {
                         config.num_heads = 12;
                         config.head_dim  = 64;
@@ -79,7 +79,7 @@ namespace MiniT2I {
                     }
                 } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) {
                     int64_t patch_area = config.patch_size * config.patch_size;
-                    config.hidden_size = tensor_storage.ne[0];
+                    config.hidden_size = tensor_storage.logical_ne0();
                     config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels;
                 } else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) {
                     config.prompt_length = tensor_storage.ne[1];

src/model/diffusion/mmdit.hpp

@@ -61,17 +61,17 @@ struct MMDiTConfig {
             } else if (name.find("y_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) {
                 has_weight_config      = true;
                 has_hidden_size        = true;
-                config.adm_in_channels = tensor_storage.ne[0];
+                config.adm_in_channels = tensor_storage.logical_ne0();
                 config.hidden_size     = tensor_storage.ne[1];
             } else if (name.find("context_embedder.weight") != std::string::npos && tensor_storage.n_dims == 2) {
                 has_weight_config               = true;
                 has_context_embed               = true;
-                config.context_size             = tensor_storage.ne[0];
+                config.context_size             = tensor_storage.logical_ne0();
                 config.context_embedder_out_dim = tensor_storage.ne[1];
             } else if (name.find("final_layer.linear.weight") != std::string::npos && tensor_storage.n_dims == 2) {
                 has_weight_config  = true;
                 has_hidden_size    = true;
-                config.hidden_size = tensor_storage.ne[0];
+                config.hidden_size = tensor_storage.logical_ne0();
                 int64_t patch_area = static_cast<int64_t>(config.patch_size) * config.patch_size;
                 if (patch_area > 0) {
                     config.out_channels = tensor_storage.ne[1] / patch_area;

src/model/diffusion/moss_image.hpp

@@ -32,16 +32,16 @@ namespace MOSSImage {
                 }
                 auto key = name.substr(prefix.size() + 1);
                 if (key == "video_embedder.weight") {
-                    config.channels = tensor.ne[0];
+                    config.channels = tensor.logical_ne0();
                     config.dim      = tensor.ne[1];
                 } else if (key == "text_embedder.proj.weight") {
-                    config.text_dim = tensor.ne[0];
+                    config.text_dim = tensor.logical_ne0();
                 } else if (key == "layers.0.proj.ffn_down.weight") {
-                    config.ffn_dim = tensor.ne[0];
+                    config.ffn_dim = tensor.logical_ne0();
                 } else if (key == "layers.0.q_norm.weight") {
                     config.head_dim = tensor.ne[0];
                 } else if (key == "time_embedding.0.weight") {
-                    config.freq_dim = tensor.ne[0];
+                    config.freq_dim = tensor.logical_ne0();
                 }
                 if (starts_with(key, "layers.")) {
                     depth = std::max<int64_t>(depth, std::stoi(key.substr(7)) + 1);

src/model/diffusion/pixart.hpp

@@ -51,10 +51,10 @@ namespace PixArt {
                 config.out_channels = w->ne[1] / (config.patch_size * config.patch_size);
             }
             if (auto w = find("y_embedder.y_proj.fc1.weight")) {
-                config.caption_channels = w->ne[0];
+                config.caption_channels = w->logical_ne0();
             }
             if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) {
-                config.cross_attention_dim = w->ne[0];
+                config.cross_attention_dim = w->logical_ne0();
             }
             if (auto w = find("blocks.0.mlp.fc1.weight")) {
                 config.ffn_dim = w->ne[1];

src/model/diffusion/qwen_image_2_1.hpp

@@ -24,14 +24,14 @@ namespace Qwen {
                 return it == weights.end() ? nullptr : &it->second;
             };
             if (auto w = find("img_in.weight")) {
-                config.in_channels = w->ne[0];
+                config.in_channels = w->logical_ne0();
                 config.hidden_size = w->ne[1];
             }
             if (auto w = find("proj_out.weight")) {
                 config.out_channels = w->ne[1];
             }
             if (auto w = find("txt_in.in_layer.weight")) {
-                config.context_dim = w->ne[0];
+                config.context_dim = w->logical_ne0();
             }
             if (auto w = find("transformer_blocks.0.attn.norm_q.weight")) {
                 config.head_dim = w->ne[0];

src/model/diffusion/sefi_image.hpp

@@ -31,7 +31,7 @@ namespace SefiImage {
                     continue;
                 }
                 if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) {
-                    config.timestep_guidance_in_dim = tensor_storage.ne[0];
+                    config.timestep_guidance_in_dim = tensor_storage.logical_ne0();
                     config.hidden_size              = tensor_storage.ne[1] * 2;
                 }
             }

src/model/diffusion/unet.hpp

@@ -104,13 +104,13 @@ struct UNetConfig {
         }
         if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) {
             if (time_embed->n_dims == 2) {
-                config.model_channels = static_cast<int>(time_embed->ne[0]);
+                config.model_channels = static_cast<int>(time_embed->logical_ne0());
                 config.time_embed_dim = static_cast<int>(time_embed->ne[1]);
             }
         }
         if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) {
             if (label_emb->n_dims == 2) {
-                config.adm_in_channels = static_cast<int>(label_emb->ne[0]);
+                config.adm_in_channels = static_cast<int>(label_emb->logical_ne0());
                 config.time_embed_dim  = static_cast<int>(label_emb->ne[1]);
             }
         }
@@ -124,7 +124,7 @@ struct UNetConfig {
                 continue;
             }
             if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) {
-                config.context_dim = static_cast<int>(tensor_storage.ne[0]);
+                config.context_dim = static_cast<int>(tensor_storage.logical_ne0());
                 break;
             }
         }

src/model/diffusion/z_image.hpp

@@ -62,10 +62,10 @@ namespace ZImage {
                 }
                 if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) {
                     int64_t patch_area = config.patch_size * config.patch_size;
-                    config.in_channels = tensor_storage.ne[0] / patch_area;
+                    config.in_channels = tensor_storage.logical_ne0() / patch_area;
                     config.hidden_size = tensor_storage.ne[1];
                 } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) {
-                    config.cap_feat_dim = tensor_storage.ne[0];
+                    config.cap_feat_dim = tensor_storage.logical_ne0();
                     config.hidden_size  = tensor_storage.ne[1];
                 } else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) {
                     detected_head_dim = tensor_storage.ne[0];

src/model/diffusion/z_image_l2p.hpp

@@ -18,7 +18,7 @@ namespace ZImageL2P {
             config.patch_size                          = 16;
             auto x_embedder                            = tensors.find(prefix + ".x_embedder.weight");
             if (x_embedder != tensors.end()) {
-                config.patch_size = static_cast<int>(std::lround(std::sqrt(static_cast<double>(x_embedder->second.ne[0] / config.in_channels))));
+                config.patch_size = static_cast<int>(std::lround(std::sqrt(static_cast<double>(x_embedder->second.logical_ne0() / config.in_channels))));
             }
             // L2P ships split to_q/to_k/to_v, which name conversion maps to qkv.weight plus
             // .1/.2 parts; the base detection only sees the q part and undercounts kv heads.

src/model/te/t5.hpp

@@ -57,7 +57,7 @@ struct T5Config {
         }
         if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) {
             if (q->n_dims == 2) {
-                config.model_dim  = q->ne[0];
+                config.model_dim  = q->logical_ne0();
                 int64_t inner_dim = q->ne[1];
                 config.inner_dim  = inner_dim;
                 // Flan-T5/T5 uses d_kv=64 for common sizes.
@@ -68,7 +68,7 @@ struct T5Config {
         }
         if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) {
             if (wi->n_dims == 2) {
-                config.model_dim = wi->ne[0];
+                config.model_dim = wi->logical_ne0();
                 config.ff_dim    = wi->ne[1];
             }
         }

src/model/upscaler/h3_latent_upscaler.hpp

@@ -75,7 +75,7 @@ namespace H3VUpsampler {
 
             const TensorStorage* embed_mid = find_weight("embed.2.weight");
             if (embed_mid != nullptr && embed_mid->n_dims >= 1) {
-                config.embed_dim = static_cast<int>(embed_mid->ne[0]);
+                config.embed_dim = static_cast<int>(embed_mid->logical_ne0());
             }
 
             std::set<int> in_res_idx, in_temporal_idx, out_res_idx, out_temporal_idx;

src/model_io/safetensors_io.cpp

@@ -109,8 +109,9 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
 
 struct ComfyQuantConfig {
     std::string format;
-    bool convrot   = false;
+    bool convrot          = false;
-    int group_size = 0;
+    int group_size        = 0;
+    int weight_group_size = 0;
 };
 
 static bool read_comfy_quant_config(std::ifstream& file,
@@ -140,6 +141,12 @@ static bool read_comfy_quant_config(std::ifstream& file,
         config.format             = json.value("format", "");
         config.convrot            = json.value("convrot", false);
         config.group_size         = json.value("convrot_groupsize", 0);
+        config.weight_group_size  = json.value("group_size", 0);
+        if (config.format == "convrot_w4a4" || config.format == "asym_w4a8_int8") {
+            // Packed W4 formats imply rotation, including exports without the INT8 convrot flag.
+            config.convrot    = true;
+            config.group_size = json.value("convrot_groupsize", 256);
+        }
     } catch (const std::exception&) {
         set_error(error, "parsing ComfyUI quantization metadata tensor failed: '" + tensor_name + "'");
         return false;
@@ -238,8 +245,8 @@ bool read_safetensors_file(const std::string& file_path,
             return false;
         }
 #ifdef SD_USE_UPSTREAM_GGML
-        if (config.format == "int8_tensorwise") {
+        if (config.format == "int8_tensorwise" || config.format == "convrot_w4a4" || config.format == "asym_w4a8_int8") {
-            set_error(error, "INT8 tensorwise/convrot is not supported by this ggml build (tensor '" + name + "')");
+            set_error(error, "ComfyUI INT8/packed INT4 is not supported by this ggml build (tensor '" + name + "')");
             return false;
         }
 #endif
@@ -340,29 +347,58 @@ bool read_safetensors_file(const std::string& file_path,
         if (ends_with(name, ".weight")) {
             const std::string module_name = name.substr(0, name.size() - std::string(".weight").size());
             auto config                   = comfy_quant_configs.find(module_name);
-            if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise") {
+            if (config != comfy_quant_configs.end() &&
+                (config->second.format == "int8_tensorwise" || config->second.format == "convrot_w4a4" || config->second.format == "asym_w4a8_int8")) {
+                const int w4_kind = config->second.format == "convrot_w4a4" ? 1 : config->second.format == "asym_w4a8_int8" ? 2
+                                                                                                                            : 0;
                 if (type != GGML_TYPE_I8) {
-                    set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'");
+                    set_error(error, "ComfyUI quantized weight is not I8: '" + name + "'");
                     return false;
                 }
+                if (w4_kind && (tensor_storage.n_dims != 2 || tensor_storage.ne[0] <= 0 ||
+                                tensor_storage.ne[0] > INT64_MAX / 2 || tensor_storage.ne[1] <= 0)) {
+                    set_error(error, "invalid packed INT4 weight shape: '" + name + "'");
+                    return false;
+                }
+                const int64_t input_features = tensor_storage.ne[0] * (w4_kind ? 2 : 1);
                 if (config->second.convrot) {
                     int group_size_remainder = config->second.group_size;
                     while (group_size_remainder > 1 && group_size_remainder % 4 == 0) {
                         group_size_remainder /= 4;
                     }
-                    if (group_size_remainder != 1 || tensor_storage.ne[0] % config->second.group_size != 0) {
+                    if (group_size_remainder != 1 || config->second.group_size > 256 || input_features % config->second.group_size != 0) {
                         set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'");
                         return false;
                     }
                 }
-                tensor_storage.is_int8_tensorwise      = true;
+                if (w4_kind) {
+                    auto matches = [&](const std::string& suffix, const std::string& expected_dtype, const std::vector<int64_t>& expected_shape) {
+                        auto entry = header_.find(module_name + suffix);
+                        return entry != header_.end() && entry->value("dtype", "") == expected_dtype &&
+                               (*entry)["shape"] == expected_shape;
+                    };
+                    const int64_t n  = tensor_storage.ne[1];
+                    const bool valid = w4_kind == 1
+                                           ? (matches(".weight_scale", "F32", {n}) || matches(".weight_scale", "F32", {n, 1}))
+                                           : config->second.weight_group_size == 16 && input_features % 16 == 0 &&
+                                                 matches(".weight_codebook", "F32", {16}) &&
+                                                 (matches(".weight_s_channel", "F32", {n}) || matches(".weight_s_channel", "F32", {n, 1})) &&
+                                                 matches(".weight_s_rel", "F8_E4M3", {n, input_features / 16});
+                    if (!valid) {
+                        set_error(error, "invalid packed INT4 scales/codebook for tensor '" + name + "'");
+                        return false;
+                    }
+                }
+                tensor_storage.is_int8_tensorwise      = w4_kind == 0;
+                tensor_storage.w4_convrot_kind         = w4_kind;
                 tensor_storage.int8_convrot            = config->second.convrot;
                 tensor_storage.int8_convrot_group_size = config->second.group_size;
             }
-        } else if (ends_with(name, ".weight_scale")) {
+        } else if (ends_with(name, ".weight_scale") || ends_with(name, ".weight_s_channel")) {
-            const std::string module_name = name.substr(0, name.size() - std::string(".weight_scale").size());
+            const std::string module_name = name.substr(0, name.rfind('.'));
             auto config                   = comfy_quant_configs.find(module_name);
-            if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise" &&
+            if (config != comfy_quant_configs.end() &&
+                (config->second.format == "int8_tensorwise" || config->second.format == "convrot_w4a4" || config->second.format == "asym_w4a8_int8") &&
                 tensor_storage.n_dims == 2 && tensor_storage.ne[0] == 1) {
                 tensor_storage.ne[0]  = tensor_storage.ne[1];
                 tensor_storage.ne[1]  = 1;

src/model_io/tensor_storage.h

@@ -22,6 +22,7 @@ struct TensorStorage {
     bool is_f64                 = false;
     bool is_i64                 = false;
     bool is_int8_tensorwise     = false;
+    int w4_convrot_kind         = 0;
     bool int8_convrot           = false;
     int int8_convrot_group_size = 0;
     int64_t ne[SD_MAX_DIMS]     = {1, 1, 1, 1, 1};
@@ -51,6 +52,10 @@ struct TensorStorage {
         return n;
     }
 
+    int64_t logical_ne0() const {
+        return ne[0] * (w4_convrot_kind ? 2 : 1);
+    }
+
     int64_t nbytes() const {
         return nelements() * ggml_type_size(type) / ggml_blck_size(type);
     }

src/model_loader.cpp

@@ -518,7 +518,7 @@ SDVersion ModelLoader::get_sd_version() const {
         }
         if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) {
             auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight");
-            if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) {
+            if (img_in != tensor_storage_map.end() && img_in->second.logical_ne0() == 128) {
                 return VERSION_MAGE_FLOW;
             }
             if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) {
@@ -593,7 +593,7 @@ SDVersion ModelLoader::get_sd_version() const {
                 }
             }
             auto connector = tensor_storage_map.find("embeddings_connectors.proj.weight");
-            if (connector != tensor_storage_map.end() && connector->second.ne[0] == 12800 &&
+            if (connector != tensor_storage_map.end() && connector->second.logical_ne0() == 12800 &&
                 connector->second.ne[1] == 5120 && tensor_storage_map.count("embeddings_connectors.sink_embedding")) {
                 return VERSION_P_H3_IMAGE;
             }
@@ -660,7 +660,7 @@ SDVersion ModelLoader::get_sd_version() const {
             tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) {
             has_output_block_71 = true;
             if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight") != std::string::npos) {
-                if (tensor_storage.ne[0] == 1024)
+                if (tensor_storage.logical_ne0() == 1024)
                     has_attn_1024 = true;
             }
         }
@@ -721,16 +721,16 @@ SDVersion ModelLoader::get_sd_version() const {
         if (has_distilled_guidance) {
             return VERSION_CHROMA;
         }
-        if (context_ebedding_weight.ne[0] == 3584) {
+        if (context_ebedding_weight.logical_ne0() == 3584) {
             return VERSION_LONGCAT;
         } else {
-            if (input_block_weight.ne[0] == 384) {
+            if (input_block_weight.logical_ne0() == 384) {
                 return VERSION_FLUX_FILL;
             }
-            if (input_block_weight.ne[0] == 128) {
+            if (input_block_weight.logical_ne0() == 128) {
                 return VERSION_FLUX_CONTROLS;
             }
-            if (input_block_weight.ne[0] == 196) {
+            if (input_block_weight.logical_ne0() == 196) {
                 return VERSION_FLEX_2;
             }
             return VERSION_FLUX;
@@ -1651,7 +1651,8 @@ bool ModelLoader::load_tensors(std::map<std::string, ggml_tensor*>& tensors,
 
 bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const {
     const std::string& name = tensor_storage.name;
-    if (tensor_storage.is_int8_tensorwise) {
+    if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind != 0 ||
+        ends_with(name, ".weight_codebook") || ends_with(name, ".weight_s_channel") || ends_with(name, ".weight_s_rel")) {
         return false;
     }
     if (type != GGML_TYPE_COUNT) {

src/model_manager_files.cpp

@@ -11,7 +11,7 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) {
            a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) &&
            a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 &&
            a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 &&
-           a.is_int8_tensorwise == b.is_int8_tensorwise && a.int8_convrot == b.int8_convrot &&
+           a.is_int8_tensorwise == b.is_int8_tensorwise && a.w4_convrot_kind == b.w4_convrot_kind && a.int8_convrot == b.int8_convrot &&
            a.int8_convrot_group_size == b.int8_convrot_group_size;
 }
 

src/pipeline/diffusion_engine.cpp

@@ -1019,7 +1019,7 @@ void StableDiffusionGGML::configure_weight_loading() {
     bool have_int8_tensorwise = false;
     bool have_quant_type_rule = false;
     for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) {
-        if (tensor_storage.is_int8_tensorwise) {
+        if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind) {
             have_int8_tensorwise = true;
         }
         if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) {
@@ -1051,7 +1051,7 @@ void StableDiffusionGGML::configure_weight_loading() {
     } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) {
         if (have_int8_tensorwise) {
             LOG_WARN(
-                "INT8 tensorwise weights do not support the immediately LoRA apply mode; "
+                "ComfyUI INT8/packed INT4 weights do not support the immediately LoRA apply mode; "
                 "using at_runtime instead");
             apply_lora_immediately = false;
         } else if (row_split_active()) {