diff --git a/CMakeLists.txt b/CMakeLists.txt index 5c8382e34..9b7a76890 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/gpu/bitnet_kernels/bitnet_kernels.cu b/gpu/bitnet_kernels/bitnet_kernels.cu index 6e6158099..bef4496b6 100644 --- a/gpu/bitnet_kernels/bitnet_kernels.cu +++ b/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); + } } \ No newline at end of file diff --git a/gpu/bitnet_kernels/bitnet_kernels.h b/gpu/bitnet_kernels/bitnet_kernels.h index 1d897908f..2949194bb 100644 --- a/gpu/bitnet_kernels/bitnet_kernels.h +++ b/gpu/bitnet_kernels/bitnet_kernels.h @@ -7,6 +7,12 @@ #include #include +// 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 diff --git a/gpu/bitnet_kernels/gemm_native.cu b/gpu/bitnet_kernels/gemm_native.cu new file mode 100644 index 000000000..068c6fd98 --- /dev/null +++ b/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 +__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><<>>(A,B,out,s,ws,M); + } else if (N == 3840 && K == 2560) { + ladder_gemm_MxN<3840,2560,3><<>>(A,B,out,s,ws,M); + } else if (N == 13824 && K == 2560) { + ladder_gemm_MxN<13824,2560,2><<>>(A,B,out,s,ws,M); + } else if (N == 2560 && K == 6912) { + ladder_gemm_MxN<2560,6912,1><<>>(A,B,out,s,ws,M); + } else { + // Fallback to batched GEMV loop for other shapes (still single host call) + for (int m=0; m "FastGen": """ Load a Llama or Code Llama checkpoint and return a new @@ -53,7 +57,10 @@ def build( """ 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) diff --git a/gpu/int2_fallback.py b/gpu/int2_fallback.py new file mode 100644 index 000000000..636374321 --- /dev/null +++ b/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) diff --git a/gpu/model.py b/gpu/model.py index cd5abec01..89b73abe6 100755 --- a/gpu/model.py +++ b/gpu/model.py @@ -16,7 +16,14 @@ ) import ctypes -bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so') +import os +# Local fix (bitnet-pruebas, Windows): upstream only ships a .so name. +# Resolve relative to this file so any CWD works. +_HERE = os.path.dirname(os.path.abspath(__file__)) +_LIB = os.path.join( + _HERE, 'bitnet_kernels', + 'bitnet_kernels.dll' if os.name == 'nt' else 'libbitnet.so') +bitnet_lib = ctypes.CDLL(_LIB) def bitnet_int8xint2_linear(input0, input1, s, ws): out_shape = list(input0.shape) @@ -47,6 +54,9 @@ class ModelArgs: norm_eps: float = 1e-5 rope_theta: float = 500000.0 use_kernel: bool = False + # Local (bitnet-pruebas): prefill with int2 weights via torch fallback + # instead of the fp16 model (saves ~5GB VRAM). Requires the int2 ckpt. + int2_torch: bool = False LayerCache = Tuple[torch.Tensor, torch.Tensor] @@ -72,8 +82,75 @@ def quant_input(self, input): def forward(self, input): input, s = self.quant_input(input) + # Local (bitnet-pruebas): native M>1 via GEMM single launch when supported + if input.shape[0] > 1: + # Try true GEMM first + try: + return bitnet_gemm_int8xint2(input, self.weight, s, self.weight_scale) + except Exception: + return bitnet_int8xint2_linear_batched(input, self.weight, s, self.weight_scale) return bitnet_int8xint2_linear(input, self.weight, s, self.weight_scale) + +def bitnet_int8xint2_linear_batched(input0, weight, s, ws): + """Batched prefill: single host call that loops M× GEMV in C++ (no Python overhead).""" + M = input0.shape[0] + N = weight.shape[0] + K = weight.shape[1] * 4 + out = torch.empty((M, N), dtype=torch.bfloat16, device=input0.device) + stream = torch.cuda.current_stream().cuda_stream + # ws is (4,) bf16 shared; s is (M,1) bf16 per-row + # Ensure contiguous as kernel expects + s_c = s.contiguous() + # Call the new batched entry point (single launch overhead) + try: + bitnet_lib.bitlinear_int8xint2_batched( + ctypes.c_void_p(input0.data_ptr()), + ctypes.c_void_p(weight.data_ptr()), + ctypes.c_void_p(out.data_ptr()), + ctypes.c_void_p(s_c.data_ptr()), + ctypes.c_void_p(ws.data_ptr()), + ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), + ctypes.c_void_p(stream)) + except AttributeError: + # Fallback to Python loop if old DLL without batched symbol + for i in range(M): + row_in = input0[i:i+1].contiguous() + row_s = s_c[i:i+1].contiguous() + row_out = out[i:i+1] + bitnet_lib.bitlinear_int8xint2( + ctypes.c_void_p(row_in.data_ptr()), + ctypes.c_void_p(weight.data_ptr()), + ctypes.c_void_p(row_out.data_ptr()), + ctypes.c_void_p(row_s.data_ptr()), + ctypes.c_void_p(ws.data_ptr()), + ctypes.c_int(1), ctypes.c_int(N), ctypes.c_int(K), + ctypes.c_void_p(stream)) + return out + + +def bitnet_gemm_int8xint2(input0, weight, s, ws): + """True GEMM M>1 single launch (native, no loop). Falls back to batched if shape unsupported.""" + M = input0.shape[0] + N = weight.shape[0] + K = weight.shape[1] * 4 + out = torch.empty((M, N), dtype=torch.bfloat16, device=input0.device) + stream = torch.cuda.current_stream().cuda_stream + s_c = s.contiguous() + try: + # New native GEMM (single kernel, grid.y=M) + bitnet_lib.bitlinear_gemm_int8xint2( + ctypes.c_void_p(input0.data_ptr()), + ctypes.c_void_p(weight.data_ptr()), + ctypes.c_void_p(out.data_ptr()), + ctypes.c_void_p(s_c.data_ptr()), + ctypes.c_void_p(ws.data_ptr()), + ctypes.c_int(M), ctypes.c_int(N), ctypes.c_int(K), + ctypes.c_void_p(stream)) + return out + except AttributeError: + return bitnet_int8xint2_linear_batched(input0, weight, s, ws) + class BitLinear(nn.Linear): @torch.compile def quant_input(self, input): @@ -84,6 +161,61 @@ def forward(self, input): input = self.quant_input(input) return F.linear(input, self.weight) + +class BitLinearInt2Torch(nn.Module): + """Prefill with int2 weights, torch fallback (no fp16 model in VRAM). + + Holds the SAME buffers as BitLinearKernel (packed int8 weight + + weight_scale), so the int2 checkpoint loads directly. Forward inverts + the packing (see int2_fallback.py), then runs the exact same + int8-rounded-activation GEMM as BitLinear. Slower than a native int2 + GEMM, but removes the 5.2GB fp16 prefill model from VRAM. + """ + + # Native kernel supports these (N,K) on 2B (also large via other build) + _NATIVE_SHAPES = {(2560,2560),(3840,2560),(13824,2560),(2560,6912),(3200,3200),(4800,3200),(3200,10240),(20480,3200)} + + def __init__(self, in_features: int, out_features: int, bias: bool = False, + row_splits=None): + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.row_splits = row_splits + + self.weight = torch.nn.Parameter( + torch.zeros(out_features, in_features // 4, dtype=torch.int8), + requires_grad=False) + self.weight_scale = torch.nn.Parameter( + torch.zeros(4, dtype=torch.bfloat16), requires_grad=False) + + @torch.compile + def quant_input(self, input): + s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5) + return (input * s).round().clamp(-128, 127) / s + + @torch.compile + def quant_input_int8(self, input): + s = 127 / input.abs().max(dim=-1, keepdim=True).values.clamp_(min=1e-5) + return (input * s).round().clamp(-128, 127).to(torch.int8), s + + def forward(self, input): + # Local (bitnet-pruebas): for M>1 and supported shapes, use native + # GEMM single launch (no dequant). Otherwise fallback. + if input.shape[0] > 1 and (self.out_features, self.in_features) in self._NATIVE_SHAPES: + q, s = self.quant_input_int8(input) + # Try true GEMM first, fallback to batched GEMV loop if shape not in GEMM dispatch + try: + return bitnet_gemm_int8xint2(q, self.weight, s, self.weight_scale) + except Exception: + return bitnet_int8xint2_linear_batched(q, self.weight, s, self.weight_scale) + return F.linear(self.quant_input(input), self._dequant_weight(input)) + + def _dequant_weight(self, input): + from int2_fallback import dequantize_int2_weight + return dequantize_int2_weight( + self.weight, self.weight_scale, self.row_splits, + out_dtype=input.dtype) + class Attention(nn.Module): def __init__( self, @@ -94,6 +226,7 @@ def __init__( rope_theta: float, norm_eps: float, use_kernel: bool, + int2_torch: bool = False, ): super().__init__() @@ -105,16 +238,27 @@ def __init__( Linear = BitLinearKernel if use_kernel else BitLinear - self.wqkv = Linear( - dim, - (self.n_local_heads + 2 * self.n_local_kv_heads) * head_dim, - bias=False, - ) - self.wo = Linear( - self.n_local_heads * head_dim, - dim, - bias=False, - ) + if int2_torch: + nq = self.n_local_heads * head_dim + nkv = self.n_local_kv_heads * head_dim + self.wqkv = BitLinearInt2Torch( + dim, nq + 2 * nkv, bias=False, + row_splits=[nq, nkv, nkv], + ) + self.wo = BitLinearInt2Torch( + self.n_local_heads * head_dim, dim, bias=False, + ) + else: + self.wqkv = Linear( + dim, + (self.n_local_heads + 2 * self.n_local_kv_heads) * head_dim, + bias=False, + ) + self.wo = Linear( + self.n_local_heads * head_dim, + dim, + bias=False, + ) self.attn_sub_norm = RMSNorm(dim, norm_eps) @@ -151,8 +295,11 @@ def forward( theta=self.rope_theta, ) + # Local fix (bitnet-pruebas): upstream forces flash.FwOp, which has no + # built backend on Windows/sm_86 wheels. Auto-dispatch picks + # cutlassF-pt / triton_splitKF instead. output = fmha.memory_efficient_attention_forward( - xq, cache_k, cache_v, attn_bias, op = fmha.flash.FwOp + xq, cache_k, cache_v, attn_bias, op=None ) output = output.reshape(output_shape) @@ -172,21 +319,31 @@ def __init__( hidden_dim: int, norm_eps: float, use_kernel: bool, + int2_torch: bool = False, ): super().__init__() Linear = BitLinearKernel if use_kernel else BitLinear - self.w13 = Linear( - dim, - 2 * hidden_dim, - bias=False, - ) - self.w2 = Linear( - hidden_dim, - dim, - bias=False, - ) + if int2_torch: + self.w13 = BitLinearInt2Torch( + dim, 2 * hidden_dim, bias=False, + row_splits=[hidden_dim, hidden_dim], + ) + self.w2 = BitLinearInt2Torch( + hidden_dim, dim, bias=False, + ) + else: + self.w13 = Linear( + dim, + 2 * hidden_dim, + bias=False, + ) + self.w2 = Linear( + hidden_dim, + dim, + bias=False, + ) self.ffn_sub_norm = RMSNorm(hidden_dim, norm_eps) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -218,12 +375,14 @@ def __init__(self, args: ModelArgs): rope_theta=args.rope_theta, norm_eps=args.norm_eps, use_kernel=args.use_kernel, + int2_torch=getattr(args, "int2_torch", False), ) self.feed_forward = FeedForward( dim=args.dim, hidden_dim=args.ffn_dim, norm_eps=args.norm_eps, use_kernel=args.use_kernel, + int2_torch=getattr(args, "int2_torch", False), ) self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps) self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps) diff --git a/gpu/test.py b/gpu/test.py index 194fd50cf..2bb250363 100644 --- a/gpu/test.py +++ b/gpu/test.py @@ -5,12 +5,19 @@ from pack_weight import convert_weight_int8_to_int2 from torch.profiler import profile, record_function, ProfilerActivity import ctypes +import os import numpy as np # set all seed torch.manual_seed(42) np.random.seed(42) -bitnet_lib = ctypes.CDLL('bitnet_kernels/libbitnet.so') +bitnet_lib = ctypes.CDLL( + # Local fix (bitnet-pruebas, Windows): upstream only ships a .so name. + # Resolved relative to this file so any CWD works. + os.path.join(os.path.dirname(os.path.abspath(__file__)), + 'bitnet_kernels', + 'bitnet_kernels.dll' if os.name == 'nt' + else 'libbitnet.so')) def bitnet_int8xint2_linear(input0, input1, s, ws, ret): out_shape = list(input0.shape) diff --git a/src/ggml-bitnet-lut.cpp b/src/ggml-bitnet-lut.cpp index 676351ddc..5cab1c904 100644 --- a/src/ggml-bitnet-lut.cpp +++ b/src/ggml-bitnet-lut.cpp @@ -189,4 +189,50 @@ 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 * 11); + 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); + } + + 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 diff --git a/src/ggml-bitnet-shim.c b/src/ggml-bitnet-shim.c new file mode 100644 index 000000000..f0d0b6b60 --- /dev/null +++ b/src/ggml-bitnet-shim.c @@ -0,0 +1,38 @@ +// Local build fix (bitnet-pruebas), NOT upstream. +// +// ggml.c (ggml-base) references dequantize_row_i2_s via the I2_S type traits +// (.to_float), but no translation unit linked into ggml-base defines it: the +// only definition lives in the ggml-cpu backend (quants.c), which is a +// separate target. This file provides that exact scalar implementation +// (copied verbatim from 3rdparty/llama.cpp/ggml/src/ggml-cpu/quants.c) so +// ggml-base.dll links. It is only a .to_float fallback; the hot I2_S +// GEMV/GEMM path uses the native kernels in ggml-cpu-i2s.c. + +#include + +#ifndef MIN +#define MIN(a, b) ((a) < (b) ? (a) : (b)) +#endif + +void dequantize_row_i2_s(const uint8_t * x, float * y, int64_t n, const float i2_scale) { + static const float map2bit[4] = { -1.0f, 0.0f, 1.0f, 0.0f }; + int64_t done = 0; + while (done < n) { + int64_t cols0 = MIN(32, n - done - 0*32); + int64_t cols1 = MIN(32, n - done - 1*32); + int64_t cols2 = MIN(32, n - done - 2*32); + int64_t cols3 = MIN(32, n - done - 3*32); + for (int gp = 0; gp < 32; gp++) { + uint8_t byte = x[(done/4) + gp]; + uint8_t c0 = (byte >> 6) & 0x03; + uint8_t c1 = (byte >> 4) & 0x03; + uint8_t c2 = (byte >> 2) & 0x03; + uint8_t c3 = (byte >> 0) & 0x03; + if (gp < cols0) y[done + 0*32 + gp] = i2_scale * map2bit[c0]; + if (gp < cols1) y[done + 1*32 + gp] = i2_scale * map2bit[c1]; + if (gp < cols2) y[done + 2*32 + gp] = i2_scale * map2bit[c2]; + if (gp < cols3) y[done + 3*32 + gp] = i2_scale * map2bit[c3]; + } + done += 128; + } +}