feat: support packed INT4 convrot models (#2276)

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

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,
-                                                x,
-                                                w,
-                                                weight_scale,
-                                                b,
-                                                int8_convrot ? int8_convrot_group_size : 0,
-                                                scale);
+#ifndef SD_USE_UPSTREAM_GGML
+            if (w4_convrot_kind) {
+                if (scale != 1.f)
+                    x = ggml_ext_scale(ctx->ggml_ctx, x, scale);
+                out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x,
+                                              w4_convrot_kind == 2 ? params["weight_codebook"] : weight_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;