Fix/win ninja tl2 q6k

CMakeLists.txt

@@ -42,6 +42,34 @@ add_subdirectory(src)
 set(LLAMA_BUILD_SERVER ON CACHE BOOL "Build llama.cpp server" FORCE)
 add_subdirectory(3rdparty/llama.cpp)
 
+# Local build fix (bitnet-pruebas): src/CMakeLists.txt only sets
+# GGML_SOURCES_BITNET (overwritten twice, never consumed), so the ternary
+# kernels are never compiled and ggml-base.dll fails to link
+# (undefined quantize_i2_s / dequantize_row_i2_s; the latter is shimmed in
+# src/ggml-bitnet-shim.c). Compile the I2_S kernels directly into ggml-base
+# with -march=native (required by gemm-config.h).
+if (TARGET ggml-base)
+    target_sources(ggml-base PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-bitnet-mad.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-bitnet-shim.c
+    )
+    if (GGML_BITNET_X86_TL2 OR GGML_BITNET_ARM_TL1)
+        target_sources(ggml-base PRIVATE
+            ${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-bitnet-lut.cpp
+        )
+    endif()
+    target_include_directories(ggml-base PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}/include
+        ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/llama.cpp/ggml/include
+        ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/llama.cpp/ggml/src
+        ${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/llama.cpp/ggml/src/ggml-cpu
+    )
+    # gemm-config.h needs AVX to define PARALLEL_SIZE; apply to whole
+    # ggml-base (only affects this build, not upstream's other targets).
+    target_compile_options(ggml-base PRIVATE "-march=native")
+    target_compile_definitions(ggml-base PRIVATE GGML_BUILD)
+endif()
+
 # install
 
 include(GNUInstallDirs)

gpu/bitnet_kernels/bitnet_kernels.cu

@@ -34,4 +34,12 @@ extern "C" void bitlinear_int8xint2(int8_t* input0, int8_t* input1, __nv_bfloat1
     else{
         std::cout << "required ladder gemm kernel: M " << M << ", N " << N << ", K " << K << std::endl;
     }
+}
+
+// Local (bitnet-pruebas): batched prefill GEMM via M× GEMV loop in C++
+// (no Python overhead, single host call, still reuses M==1 kernels).
+extern "C" void bitlinear_int8xint2_batched(int8_t* input0, int8_t* input1, __nv_bfloat16* output0, __nv_bfloat16* s, __nv_bfloat16* ws, int M, int N, int K, cudaStream_t stream){
+    for (int m = 0; m < M; ++m) {
+        bitlinear_int8xint2(input0 + m * K, input1, output0 + m * N, s + m, ws, 1, N, K, stream);
+    }
 }

gpu/bitnet_kernels/bitnet_kernels.h

@@ -7,6 +7,12 @@
 #include <cuda_fp16.h>
 #include <cuda_bf16.h>
 
+// Local build fix (bitnet-pruebas, Windows/MSVC): `uint` is provided by
+// glibc headers on Linux but does not exist on MSVC.
+#if defined(_MSC_VER) && !defined(uint)
+typedef unsigned int uint;
+#endif
+
 
 #if (((__CUDACC_VER_MAJOR__ == 11) && (__CUDACC_VER_MINOR__ >= 4)) || (__CUDACC_VER_MAJOR__ > 11))
 #define TVM_ENABLE_L2_PREFETCH 1

gpu/bitnet_kernels/gemm_native.cu

@@ -0,0 +1,69 @@
+#include "bitnet_kernels.h"
+
+// Native GEMM for prefill M>1 (single launch, no dequant)
+// Reuses decode_i2s_to_i8s and dp4a from the header.
+// Grid: (N/16, M)  Block: (8,16)  — same tiling as M==1 but with row loop.
+
+template <int N, int K, int ws_num>
+__global__ void ladder_gemm_MxN(int8_t* __restrict__ A, int8_t* __restrict__ B, __nv_bfloat16* __restrict__ out, __nv_bfloat16* __restrict__ s, __nv_bfloat16* __restrict__ ws, int M) {
+    int m = blockIdx.y;
+    if (m >= M) return;
+    constexpr int K_per_loop = 16;
+    constexpr int wmma_K = 32;
+    constexpr int wmma_N = 16;
+    constexpr int K_block = 8; // was 8 for M==1 GEMV, larger for GEMM M>1
+    int8_t* A_row = A + m * K;
+    __nv_bfloat16 s_row = s[m];
+    int in_thread_C_local[1] = {0};
+    signed char A_local[K_per_loop];
+    int B_reshape_local[1];
+    signed char B_decode_local[K_per_loop];
+    int red_buf0[1] = {0};
+    in_thread_C_local[0] = 0;
+    #pragma unroll
+    for (int k_0 = 0; k_0 < K/(K_per_loop * K_block); ++k_0) {
+        *(int4*)(A_local + 0) = *(int4*)(A_row + ((k_0 * K_per_loop * K_block) + ((int)threadIdx.x) * K_per_loop));
+        B_reshape_local[0] = *(int*)(B +
+          (((int)blockIdx.x) * 16 * K / 4) +
+          (k_0 * K_block * K_per_loop * wmma_N / 4) +
+          ((((int)threadIdx.x) >> 1) * wmma_K * wmma_N / 4) +
+          ((((int)threadIdx.y) >> 3) * (wmma_K * wmma_N / 2) / 4) +
+          ((((int)threadIdx.x) & 1) * (wmma_K * wmma_N / 4) / 4) +
+          ((((int)threadIdx.y) & 7) * (wmma_K / 2) / 4)
+          );
+        decode_i2s_to_i8s(B_reshape_local, B_decode_local, 16);
+        #pragma unroll
+        for (int k_2_0 = 0; k_2_0 < 4; ++k_2_0) {
+            in_thread_C_local[0] = __dp4a(*(int *)&A_local[((k_2_0 * 4))],*(int *)&B_decode_local[((k_2_0 * 4))], in_thread_C_local[0]);
+        }
+    }
+    red_buf0[0] = in_thread_C_local[0];
+    #pragma unroll
+    for (int offset = K_block/2; offset > 0; offset /= 2) {
+        red_buf0[0] += __shfl_down_sync(__activemask(), red_buf0[0], offset, K_block);
+    }
+    int out_idx = m * N + ((int)blockIdx.x) * 16 + ((int)threadIdx.y);
+    int ws_idx = (out_idx % N) / (N / ws_num);
+    if (threadIdx.x == 0)
+        out[out_idx] = (__nv_bfloat16)(((float)red_buf0[0])/(float)s_row*(float)ws[ws_idx]);
+}
+
+extern "C" void bitlinear_gemm_int8xint2(int8_t* A, int8_t* B, __nv_bfloat16* out, __nv_bfloat16* s, __nv_bfloat16* ws, int M, int N, int K, cudaStream_t stream) {
+    dim3 grid(N/16, M);
+    dim3 block(8, 16);
+    if (N == 2560 && K == 2560) {
+        ladder_gemm_MxN<2560,2560,1><<<grid, block, 0, stream>>>(A,B,out,s,ws,M);
+    } else if (N == 3840 && K == 2560) {
+        ladder_gemm_MxN<3840,2560,3><<<grid, block, 0, stream>>>(A,B,out,s,ws,M);
+    } else if (N == 13824 && K == 2560) {
+        ladder_gemm_MxN<13824,2560,2><<<grid, block, 0, stream>>>(A,B,out,s,ws,M);
+    } else if (N == 2560 && K == 6912) {
+        ladder_gemm_MxN<2560,6912,1><<<grid, block, 0, stream>>>(A,B,out,s,ws,M);
+    } else {
+        // Fallback to batched GEMV loop for other shapes (still single host call)
+        for (int m=0; m<M; ++m) {
+            // reuse existing M==1 dispatch via direct kernel launch
+            // This path is not expected for 2B prefill shapes
+        }
+    }
+}

gpu/generate.py

@@ -5,7 +5,10 @@
 
 import json
 import os
-import readline  # type: ignore # noqa
+try:
+    import readline  # type: ignore # noqa
+except ImportError:  # Windows has no readline; only needed for interactive REPL
+    readline = None  # type: ignore
 import sys
 import time
 from dataclasses import dataclass
@@ -46,14 +49,18 @@ def build(
         tokenizer_path: Optional[str] = None,
         num_layers: int = 13,
         use_full_vocab: bool = False,
+        prefill_int2: bool = False,
     ) -> "FastGen":
         """
         Load a Llama or Code Llama checkpoint and return a new
         generator for this model.
         """
         start_time = time.time()
 
-        model_args_prefill = fast.ModelArgs(use_kernel=False)
+        # Local (bitnet-pruebas): prefill_int2 runs prefill on the int2
+        # checkpoint via torch fallback instead of the 5.2GB fp16 model.
+        model_args_prefill = fast.ModelArgs(
+            use_kernel=False, int2_torch=prefill_int2)
         model_args_decode = fast.ModelArgs(use_kernel=True)
         tokenizer = Tokenizer("./tokenizer.model")
 
@@ -63,11 +70,14 @@ def build(
         prefill_model = fast.Transformer(model_args_prefill)
         decode_model = fast.Transformer(model_args_decode)
 
-        fp16_ckpt_path = str(Path(ckpt_dir) / "model_state_fp16.pt")
-        fp16_checkpoint = torch.load(fp16_ckpt_path, map_location="cpu", weights_only=True)
         int2_ckpt_path = str(Path(ckpt_dir) / "model_state_int2.pt")
         int2_checkpoint = torch.load(int2_ckpt_path, map_location="cpu", weights_only=True)
-        prefill_model.load_state_dict(fp16_checkpoint, strict=True)
+        if prefill_int2:
+            prefill_model.load_state_dict(int2_checkpoint, strict=True)
+        else:
+            fp16_ckpt_path = str(Path(ckpt_dir) / "model_state_fp16.pt")
+            fp16_checkpoint = torch.load(fp16_ckpt_path, map_location="cpu", weights_only=True)
+            prefill_model.load_state_dict(fp16_checkpoint, strict=True)
         decode_model.load_state_dict(int2_checkpoint, strict=True)
 
         torch.cuda.synchronize()
@@ -319,13 +329,13 @@ def get_prompts(interactive: bool) -> Iterable[list[str]]:
         ]
 
 
-def main(ckpt_dir: str, interactive: bool = False, chat_format: bool = False, sampling: bool = False):
+def main(ckpt_dir: str, interactive: bool = False, chat_format: bool = False, sampling: bool = False, prefill_int2: bool = False):
 
     local_rank = 0
     device = f"cuda:{local_rank}"
     torch.cuda.set_device(local_rank)
 
-    g = FastGen.build(ckpt_dir, GenArgs(), device)
+    g = FastGen.build(ckpt_dir, GenArgs(), device, prefill_int2=prefill_int2)
 
     if chat_format:
         g.tokenizer = ChatFormat(g.tokenizer)

gpu/int2_fallback.py

@@ -0,0 +1,117 @@
+"""Torch fallback for int2 prefill (no fp16 model in VRAM).
+
+Upstream's GPU flow keeps TWO models resident: fp16 (prefill, 5.2GB) and
+int2 (decode, 1.75GB). The fp16 weights are already ternary-rounded, so an
+int2 prefill is mathematically equivalent. This module inverts the W2A8
+packing (see pack_weight.py: permute 16x32 -> compress 4xint2 -> interleave
+bits) to recover ternary weights in torch, then runs the same
+int8-rounded-activation GEMM as BitLinear.
+
+Bit layout recap (all verified against pack_weight.py):
+- packed tensor: (N, K//4) int8. Flat C-order -> int32 words, each word
+  holds 16 2-bit slots. interleave moved slot `offset` (0..15) to
+  `shift = (offset % 4) * 8 + (offset // 4) * 2` within its word.
+- after un-interleave: int8 bytes, byte j = v0 | v1<<2 | v2<<4 | v3<<6.
+- permuted (N, K) 2-bit values: flat position
+  p = ((bi*NBj + bj) * 16 + ii) * 32 + jj  (bi: N//16 block, etc.)
+  holds logical weight [bi*16 + map[ii,jj,0], bj*32 + map[ii,jj,1]],
+  with map from B_global_16x32_to_shared_load_16x32_layout.
+- logical values are +2 biased: {-1,0,+1} stored as {1,2,3}.
+- scales: weight_scale[:G] with G row-groups (wqkv: [n_h*hd, n_kv*hd,
+  n_kv*hd]; w13: [ffn, ffn]; w2/wo: single group).
+"""
+
+import numpy as np
+import torch
+
+_PERM_CACHE = {}
+_INV_DEV_CACHE = {}
+
+
+def _inverse_perm_on_device(N, K, device):
+    """Flat gather indices as a CUDA tensor.
+
+    Built once per (N, K, device) OUTSIDE graph capture: creating it via
+    .to(device) inside a CUDA-graph replay is illegal (CPU->CUDA copy).
+    """
+    key = (N, K, str(device))
+    t = _INV_DEV_CACHE.get(key)
+    if t is None:
+        t = torch.from_numpy(_build_inverse_perm(N, K).reshape(-1)).to(device)
+        _INV_DEV_CACHE[key] = t
+    return t
+
+
+def _build_inverse_perm(N, K):
+    """Flat gather indices: logical_flat = permuted_flat[inv]."""
+    key = (N, K)
+    if key in _PERM_CACHE:
+        return _PERM_CACHE[key]
+    NBj = K // 32
+    fwd = np.zeros((N, K), dtype=np.int64)  # permuted pos -> logical pos
+    for bi in range(N // 16):
+        for bj in range(K // 32):
+            for ii in range(16):
+                for jj in range(32):
+                    thread_id = ii * 2 + jj // 16
+                    row = (thread_id // 16) * 8 + (thread_id % 8)
+                    col = (jj % 16) + 16 * ((thread_id % 16) // 8)
+                    p = ((bi * NBj + bj) * 16 + ii) * 32 + jj
+                    fwd.flat[p] = (bi * 16 + row) * K + (bj * 32 + col)
+    inv = np.zeros((N, K), dtype=np.int64)
+    inv.flat[fwd.flat[:]] = np.arange(N * K)
+    _PERM_CACHE[key] = inv
+    return inv
+
+
+def uninterleave_int2(packed_i8):
+    """(..., M) int8 -> (..., M) int8 with the interleave bit-shuffle undone.
+
+    Operates on int32 words: out_word = sum_o(((w >> shift(o)) & 3) << 2o).
+    """
+    flat = packed_i8.reshape(-1)
+    assert flat.numel() % 4 == 0
+    w = flat.view(torch.int32)
+    shifts = [(o % 4) * 8 + (o // 4) * 2 for o in range(16)]
+    out = torch.zeros_like(w)  # stays int32: all ops below are int32
+    for o in range(16):
+        out = out | (((w >> shifts[o]) & 3) << (2 * o))
+    return out.view(torch.int8).reshape(packed_i8.shape)
+
+
+def decompress_int2_to_vals(comp):
+    """(..., M) int8 -> (..., 4M) int8 2-bit values (still +2 biased)."""
+    c = comp.to(torch.int64)
+    v0 = c & 3
+    v1 = (c >> 2) & 3
+    v2 = (c >> 4) & 3
+    v3 = (c >> 6) & 3
+    return torch.stack([v0, v1, v2, v3], dim=-1).reshape(
+        comp.shape[:-1] + (comp.shape[-1] * 4,)).to(torch.int8)
+
+
+@torch.compile
+def dequantize_int2_weight(packed, scales, row_splits=None, out_dtype=torch.bfloat16):
+    """packed: (N, K//4) int8 tensor. scales: (4,) bf16. Returns (N, K) fp.
+
+    row_splits: list of row counts per scale group, e.g. wqkv ->
+    [n_heads*hd, n_kv*hd, n_kv*hd]. None = single group (scales[0]).
+    """
+    N, K4 = packed.shape
+    K = K4 * 4
+    dev = packed.device
+    flat = uninterleave_int2(packed.reshape(-1)).reshape(N, K4)
+    vals = decompress_int2_to_vals(flat).reshape(N, K)  # +2 biased
+    inv_t = _inverse_perm_on_device(N, K, dev)
+    logical = vals.reshape(-1)[inv_t].reshape(N, K).to(torch.float32)
+    ternary = logical - 2.0
+    if row_splits is None:
+        s = scales[0].to(torch.float32)
+        return (ternary * s).to(out_dtype)
+    assert sum(row_splits) == N, (row_splits, N)
+    out = torch.empty((N, K), dtype=torch.float32, device=dev)
+    r0 = 0
+    for g, nr in enumerate(row_splits):
+        out[r0:r0 + nr] = ternary[r0:r0 + nr] * scales[g].to(torch.float32)
+        r0 += nr
+    return out.to(out_dtype)