feat: add I2_S GGUF conversion for bitnet-b1.58-2B-4T and refactor T-MAC LUT path

- Add quantize_to_i2_s() for direct ternary-to-I2_S packing in conversion script
- Support offline-quantized models (uint8 packed weights + weight_scale)
- Fix weight_quant double-quantization bug for offline-quantized models
- Fix I2_S scale computation to use first nonzero absolute value
- Add I2_S ftype mapping and BitNetForCausalLM registration
- Refactor ggml-bitnet-lut T-MAC wrapper with proper mul_mat implementation
- Update llama.cpp submodule with I2_S ftype and 2B model type support

include/ggml-bitnet.h

@@ -14,6 +14,8 @@ typedef float bitnet_float_type;
 extern "C" {
 #endif
 
+struct ggml_compute_params;
+
 struct bitnet_tensor_extra {
     int lut_scales_size;
     int BK;
@@ -33,6 +35,7 @@ GGML_API size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, c
 GGML_API void ggml_bitnet_mul_mat_task_init(void * src1, void * qlut, void * lut_scales, void * lut_biases, int n, int k, int m, int bits);
 GGML_API void ggml_bitnet_mul_mat_task_compute(void * src0, void * scales, void * qlut, void * lut_scales, void * lut_biases, void * dst, int n, int k, int m, int bits);
 GGML_API void ggml_bitnet_transform_tensor(struct ggml_tensor * tensor);
+GGML_API void ggml_bitnet_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 GGML_API int ggml_bitnet_get_type_bits(enum ggml_type type);
 GGML_API void ggml_bitnet_set_n_threads(int n_threads);
 #if defined(GGML_BITNET_ARM_TL1)

include/kernel_config.ini

@@ -0,0 +1,21 @@
+[Kernels_0]
+m = 3200
+k = 8640
+bm = 160
+bk = 96
+bmm = 32
+
+[Kernels_1]
+m = 3200
+k = 3200
+bm = 320
+bk = 96
+bmm = 32
+
+[Kernels_2]
+m = 8640
+k = 3200
+bm = 320
+bk = 96
+bmm = 32
+

src/ggml-bitnet-lut.cpp

@@ -11,6 +11,7 @@
 
 #include "ggml-bitnet.h"
 #include "ggml-quants.h"
+#include "ggml-cpu-impl.h"
 
 #if defined(GGML_BITNET_ARM_TL1) || defined(GGML_BITNET_X86_TL2)
 #include "bitnet-lut-kernels.h"
@@ -19,16 +20,11 @@
 #if defined(GGML_BITNET_ARM_TL1)
 
 void ggml_bitnet_init(void) {
-    // LOG(INFO) << "ggml_bitnet_init";
-
     if (initialized) {
         return;
     }
     initialized = true;
 
-    // if (wrapper == nullptr) {
-    //     wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
-    // }
     if (bitnet_tensor_extras == nullptr) {
         bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
     }
@@ -36,26 +32,17 @@ void ggml_bitnet_init(void) {
 }
 
 void ggml_bitnet_free(void) {
-    // LOG(INFO) << "ggml_bitnet_free";
-
     if (!initialized) {
         return;
     }
     initialized = false;
 
-    // delete wrapper;
-    // wrapper = nullptr;
-    for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
-        // aligned_free(bitnet_tensor_extras[i].qweights);
-        // aligned_free(bitnet_tensor_extras[i].scales);
-    }
     delete[] bitnet_tensor_extras;
     bitnet_tensor_extras = nullptr;
 }
 
 static bool do_permutate(enum ggml_type type) {
     if (type == GGML_TYPE_TL1) {
-        // Add additional args to decide if permuted I2 or naive I2
         return false;
     } else {
         return true;
@@ -65,8 +52,7 @@ static bool do_permutate(enum ggml_type type) {
 bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
     if ((is_type_supported(src0->type)) &&
         src1->type == GGML_TYPE_F32 &&
-        dst->type == GGML_TYPE_F32 &&
-        src0->backend == GGML_BACKEND_TYPE_CPU) {
+        dst->type == GGML_TYPE_F32) {
         if (src1->ne[1] <= 1) {
             return true;
         }
@@ -79,10 +65,9 @@ size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const stru
     const size_t ne10 = src1->ne[0];
     const size_t ne11 = src1->ne[1];
     const int bits = ggml_bitnet_get_type_bits(src0->type);
-    
+
     size_t wsize = ne10 * ne11 * 15 * sizeof(int8_t) + 1 * ne11 * 2 * sizeof(bitnet_float_type);
     if (sizeof(bitnet_float_type) == 2) {
-        // Need fp32 to fp16 conversion
         wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
     }
     wsize = ((wsize - 1) / 64 + 1) * 64;
@@ -100,19 +85,61 @@ int ggml_bitnet_get_type_bits(enum ggml_type type) {
     }
 }
 
+void ggml_bitnet_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst) {
+    const struct ggml_tensor * src0 = dst->src[0];
+    const struct ggml_tensor * src1 = dst->src[1];
+
+    const size_t ne00 = src0->ne[0];
+    const size_t ne01 = src0->ne[1];
+    const size_t ne10 = src1->ne[0];
+    const size_t ne11 = src1->ne[1];
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+    const int bits = ggml_bitnet_get_type_bits(src0->type);
+
+    struct bitnet_tensor_extra * extra = (struct bitnet_tensor_extra *)src0->extra;
+    GGML_ASSERT(extra != nullptr);
+
+    char * wdata = (char *)params->wdata;
+    const size_t wsize_per_thread = ggml_bitnet_mul_mat_get_wsize(src0, src1, dst);
+
+    int8_t  * qlut  = (int8_t  *)(wdata);
+    bitnet_float_type * lut_scales = (bitnet_float_type *)(qlut + ne10 * ne11 * 15);
+    bitnet_float_type * lut_biases = (bitnet_float_type *)(lut_scales + ne11);
+
+    if (ith == 0) {
+        ggml_bitnet_mul_mat_task_init(
+            (void *)((char *)src1->data),
+            (void *)qlut,
+            (void *)lut_scales,
+            (void *)lut_biases,
+            ne10, ne00, ne11, bits);
+    }
+
+    // barrier
+    if (nth > 1) {
+        ggml_barrier(params->threadpool);
+    }
+
+    ggml_bitnet_mul_mat_task_compute(
+        (void *)extra->qweights,
+        (void *)extra->scales,
+        (void *)qlut,
+        (void *)lut_scales,
+        (void *)lut_biases,
+        (void *)((char *)dst->data),
+        ne10, ne00, ne11, bits);
+}
+
 #endif
 #if defined(GGML_BITNET_X86_TL2)
 void ggml_bitnet_init(void) {
-    // LOG(INFO) << "ggml_bitnet_init";
-
     if (initialized) {
         return;
     }
     initialized = true;
 
-    // if (wrapper == nullptr) {
-    //     wrapper = new BITNET::BITNETGeMMWrapper<bitnet_bitnet_float_type>();
-    // }
     if (bitnet_tensor_extras == nullptr) {
         bitnet_tensor_extras = new bitnet_tensor_extra[GGML_BITNET_MAX_NODES];
     }
@@ -120,19 +147,11 @@ void ggml_bitnet_init(void) {
 }
 
 void ggml_bitnet_free(void) {
-    // LOG(INFO) << "ggml_bitnet_free";
-
     if (!initialized) {
         return;
     }
     initialized = false;
 
-    // delete wrapper;
-    // wrapper = nullptr;
-    for (size_t i = 0; i < bitnet_tensor_extras_index; i++) {
-        // aligned_free(bitnet_tensor_extras[i].qweights);
-        // aligned_free(bitnet_tensor_extras[i].scales);
-    }
     delete[] bitnet_tensor_extras;
     bitnet_tensor_extras = nullptr;
 }
@@ -140,8 +159,7 @@ void ggml_bitnet_free(void) {
 bool ggml_bitnet_can_mul_mat(const struct ggml_tensor * src0, const struct ggml_tensor * src1, const struct ggml_tensor * dst) {
     if ((is_type_supported(src0->type)) &&
         src1->type == GGML_TYPE_F32 &&
-        dst->type == GGML_TYPE_F32 &&
-        src0->backend == GGML_BACKEND_TYPE_CPU) {
+        dst->type == GGML_TYPE_F32) {
         return true;
     }
     return false;
@@ -151,10 +169,9 @@ size_t ggml_bitnet_mul_mat_get_wsize(const struct ggml_tensor * src0, const stru
     const size_t ne01 = src0->ne[1];
     const size_t ne10 = src1->ne[0];
     const size_t ne11 = src1->ne[1];
-    
+
     size_t wsize = ne10 * ne11 * 11 * sizeof(int8_t) + 2 * ne11 * 2 * sizeof(bitnet_float_type);
     if (sizeof(bitnet_float_type) == 2) {
-        // Need fp32 to fp16 conversion
         wsize += std::max(ne10, ne01) * ne11 * sizeof(bitnet_float_type);
     }
     wsize = ((wsize - 1) / 64 + 1) * 64;
@@ -171,4 +188,5 @@ int ggml_bitnet_get_type_bits(enum ggml_type type) {
             return 0;
     }
 }
-#endif
+
+#endif

utils/convert-hf-to-gguf-bitnet.py

@@ -153,8 +153,12 @@ class Model(ABC):
             self.gguf_writer.add_expert_used_count(n_experts_used)
             logger.info(f"gguf: experts used count = {n_experts_used}")
 
-        self.gguf_writer.add_file_type(self.ftype)
-        logger.info(f"gguf: file type = {self.ftype}")
+        # Map ggml tensor type to llama ftype for general.file_type metadata
+        ftype_val = self.ftype
+        if self.ftype == gguf.GGMLQuantizationType.I2_S:
+            ftype_val = 40  # LLAMA_FTYPE_MOSTLY_I2_S (matches official model)
+        self.gguf_writer.add_file_type(ftype_val)
+        logger.info(f"gguf: file type = {ftype_val}")
 
     def write_tensors(self):
         block_count = self.hparams.get("n_layers", self.hparams.get("num_hidden_layers", self.hparams.get("n_layer")))
@@ -659,6 +663,70 @@ def preprocess_weights_tl2(
                              weight.shape[0]), mode='constant', constant_values=0)
     return weight
 
+def quantize_to_i2_s(w: np.ndarray, override_scale: float = None) -> np.ndarray:
+    """Quantize a float weight matrix to I2_S ternary format.
+
+    I2_S format: packed ternary bytes (4 values per byte) + 32-byte tail with f32 scale.
+    Dequantization: y = scale * ternary, where ternary in {-1, 0, +1}.
+
+    Args:
+        w: float weight tensor of shape (M, K)
+        override_scale: if provided, use this as the I2_S scale instead of computing from data.
+                       For offline-quantized BitNet models, this should be the weight_scale value.
+    """
+    M, K = w.shape
+    n = M * K
+    w_flat = w.flatten().astype(np.float32)
+
+    # Compute scale for I2_S dequantization
+    if override_scale is not None:
+        # override_scale is weight_scale from offline-quantized models ≈ mean(|original_bf16_weights|)
+        # Use it directly as the I2_S scale
+        scale = np.float32(override_scale)
+        # Weights are already ternary {-1, 0, 1}, use directly
+        q_float = w_flat
+    else:
+        # w_flat comes from weight_quant: values are ±scale or 0
+        # Use the first nonzero absolute value as scale (matches C quantize_i2_s)
+        nonzero = np.abs(w_flat[np.abs(w_flat) > 1e-6])
+        if len(nonzero) > 0:
+            scale = np.float32(nonzero[0])
+        else:
+            scale = np.float32(1e-5)
+        # Quantize to ternary {-1, 0, 1}
+        inv_scale = 1.0 / scale
+        q_float = np.round(w_flat * inv_scale).clip(-1, 1)
+
+    # Map ternary {-1, 0, 1} -> I2_S encoding {0, 1, 2}
+    q = np.ones(n, dtype=np.uint8)  # default to 1 (zero)
+    q[q_float > 0.5] = 2    # +1 -> 2
+    q[q_float < -0.5] = 0   # -1 -> 0
+
+    # Pack into I2_S layout: 128-value blocks, interleaved into 32 bytes
+    pad_len = (128 - n % 128) % 128
+    if pad_len:
+        q = np.pad(q, (0, pad_len), constant_values=1)
+
+    n_padded = len(q)
+    n_blocks = n_padded // 128
+    q = q.reshape(n_blocks, 4, 32)
+
+    packed = (q[:, 0, :].astype(np.uint8) << 6) | \
+             (q[:, 1, :].astype(np.uint8) << 4) | \
+             (q[:, 2, :].astype(np.uint8) << 2) | \
+             (q[:, 3, :].astype(np.uint8))
+    packed = packed.reshape(-1).astype(np.uint8)
+
+    # I2_S format: packed_bytes + 32-byte aligned tail (scale in first 4 bytes)
+    packed_size = n // 4
+    total_size = packed_size + 32
+    result = np.zeros(total_size, dtype=np.uint8)
+    result[:len(packed)] = packed[:packed_size]
+    result[packed_size:packed_size+4] = np.frombuffer(scale.tobytes(), dtype=np.uint8)
+
+    return result
+
+
 def transform_to_tl1(x: np.ndarray):
     scale = np.max(np.abs(x))
     # res = np.round(x / scale + 2).astype(np.uint8)
@@ -813,10 +881,17 @@ class LlamaModel(Model):
                             data = data.astype(np.float16)
                         data_qtype = gguf.GGMLQuantizationType.F16
 
-                if data_qtype is None:  # by default, convert to float32
-                    if data_dtype != np.float32:
-                        data = data.astype(np.float32)
-                    data_qtype = gguf.GGMLQuantizationType.F32
+                if data_qtype is None:  # by default
+                    # For I2_S/TL models, keep non-quantized 2D weights (e.g. embed) as F16 instead of F32
+                    if self.ftype in (gguf.GGMLQuantizationType.I2_S, gguf.GGMLQuantizationType.TL1, gguf.GGMLQuantizationType.TL2) \
+                       and n_dims >= 2 and not new_name.endswith("_norm.weight"):
+                        if data_dtype != np.float16:
+                            data = data.astype(np.float16)
+                        data_qtype = gguf.GGMLQuantizationType.F16
+                    else:
+                        if data_dtype != np.float32:
+                            data = data.astype(np.float32)
+                        data_qtype = gguf.GGMLQuantizationType.F32
 
                 shape = data_shape
                 # shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape
@@ -952,20 +1027,30 @@ class LlamaModel(Model):
                 raise ValueError(f"Unprocessed experts: {experts}")
 
 
-@Model.register("BitnetForCausalLM")
+@Model.register("BitnetForCausalLM", "BitNetForCausalLM")
 class BitnetModel(Model):
-    model_arch = gguf.MODEL_ARCH.BITNET
+    model_arch = gguf.MODEL_ARCH.BITNET_B158
 
     def set_vocab(self):
-        self._set_vocab_sentencepiece()
+        try:
+            self._set_vocab_sentencepiece()
+        except FileNotFoundError:
+            try:
+                self._set_vocab_llama_hf()
+            except (FileNotFoundError, TypeError):
+                self._set_vocab_gpt2()
         
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
 
         self.gguf_writer.add_vocab_size(self.hparams["vocab_size"])
 
-        self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR)
-        self.gguf_writer.add_rope_scaling_factor(1.0)
+        # rope dimension count (required for correct positional encoding)
+        if "head_dim" in self.hparams:
+            rope_dim = self.hparams["head_dim"]
+        else:
+            rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"]
+        self.gguf_writer.add_rope_dimension_count(rope_dim)
 
     def weight_quant(self, weight):
         dtype = weight.dtype
@@ -976,7 +1061,10 @@ class BitnetModel(Model):
 
     def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
         # quant weight to i2 (in fp16)
-        if name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight", 
+        # Skip weight_quant for offline-quantized models — weights are already ternary {-1,0,1}
+        # weight_quant would scale them to small floats (e.g. ±0.39), breaking quantize_to_i2_s
+        if not getattr(self, '_has_offline_quant', False) and \
+           name.endswith(("q_proj.weight", "k_proj.weight", "v_proj.weight",
                           "down_proj.weight", "up_proj.weight", "gate_proj.weight",
                           "o_proj.weight")):
             data_torch = self.weight_quant(data_torch)
@@ -986,15 +1074,40 @@ class BitnetModel(Model):
     def write_tensors(self):
         max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,")
 
+        # First pass: collect weight_scale tensors for offline-quantized models
+        scale_map = dict()
         for name, data_torch in self.get_tensors():
+            if name.endswith("weight_scale"):
+                data_torch = data_torch.to(torch.float32)
+                name = name.replace(".weight_scale", "")
+                scale_map[name] = data_torch
+
+        self._has_offline_quant = len(scale_map) > 0
+
+        for name, data_torch in self.get_tensors():
+            # skip weight_scale tensors
+            if name.endswith("weight_scale"):
+                continue
             # we don't need these
             if name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")):
                 continue
 
             old_dtype = data_torch.dtype
 
+            # Handle offline-quantized weights (uint8 packed with weight_scale)
+            if name.replace(".weight", "") in scale_map:
+                data_torch = data_torch.to(torch.uint8)
+                origin_shape = data_torch.shape
+                shift = torch.tensor([0, 2, 4, 6], dtype=torch.uint8).reshape((4, *(1 for _ in range(len(origin_shape)))))
+                data_torch = data_torch.unsqueeze(0).expand((4, *origin_shape)) >> shift
+                data_torch = data_torch & 3
+                data_torch = (data_torch.float() - 1).reshape((origin_shape[0] * 4, *origin_shape[1:]))
+                # For F16/F32 output: divide by weight_scale to get full float values
+                # For I2_S output: keep as ternary {-1,0,1}, scale is passed separately to quantize_to_i2_s
+                if self.ftype not in (gguf.GGMLQuantizationType.I2_S, gguf.GGMLQuantizationType.TL1, gguf.GGMLQuantizationType.TL2):
+                    data_torch = data_torch / scale_map[name.replace(".weight", "")].float()
             # convert any unsupported data types to float32
-            if data_torch.dtype not in (torch.float16, torch.float32):
+            elif data_torch.dtype not in (torch.float16, torch.float32):
                 data_torch = data_torch.to(torch.float32)
 
             # use the first number-like part of the tensor name as the block id
@@ -1048,7 +1161,13 @@ class BitnetModel(Model):
 
                 i2_scale = None
                 if self.ftype != gguf.GGMLQuantizationType.F32 and extra_f16 and not extra_f32:
-                    if self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2:
+                    if self.ftype == gguf.GGMLQuantizationType.I2_S and suit_i2:
+                        data_qtype = gguf.GGMLQuantizationType.I2_S
+                        # Use original weight_scale if available (offline-quantized models)
+                        orig_scale = scale_map.get(name.replace(".weight", ""))
+                        override_scale = orig_scale.item() if orig_scale is not None else None
+                        data = quantize_to_i2_s(data, override_scale=override_scale)
+                    elif self.ftype == gguf.GGMLQuantizationType.TL1 and suit_i2:
                         data, i2_scale = transform_to_tl1(data)
                         assert data.dtype == np.uint8
                         assert i2_scale.dtype == np.float32
@@ -1063,10 +1182,17 @@ class BitnetModel(Model):
                             data = data.astype(np.float16)
                         data_qtype = gguf.GGMLQuantizationType.F16
 
-                if data_qtype is None:  # by default, convert to float32
-                    if data_dtype != np.float32:
-                        data = data.astype(np.float32)
-                    data_qtype = gguf.GGMLQuantizationType.F32
+                if data_qtype is None:  # by default
+                    # For I2_S/TL models, keep non-quantized 2D weights (e.g. embed) as F16 instead of F32
+                    if self.ftype in (gguf.GGMLQuantizationType.I2_S, gguf.GGMLQuantizationType.TL1, gguf.GGMLQuantizationType.TL2) \
+                       and n_dims >= 2 and not new_name.endswith("_norm.weight"):
+                        if data_dtype != np.float16:
+                            data = data.astype(np.float16)
+                        data_qtype = gguf.GGMLQuantizationType.F16
+                    else:
+                        if data_dtype != np.float32:
+                            data = data.astype(np.float32)
+                        data_qtype = gguf.GGMLQuantizationType.F32
 
                 shape = data_shape
                 # shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape
@@ -1087,6 +1213,7 @@ class BitnetModel(Model):
 ftype_map = {
     "f32": gguf.GGMLQuantizationType.F32,
     "f16": gguf.GGMLQuantizationType.F16,
+    "i2_s": gguf.GGMLQuantizationType.I2_S,
     "tl1" : gguf.GGMLQuantizationType.TL1,
     "tl2" : gguf.GGMLQuantizationType.TL2,
 }