run_qwen3_4b.py

#!/usr/bin/env python3
"""
Unary model loader for Qwen3-4B-Thinking.
Loads converted weights and runs inference via unary_engine_v2.so
(c) 2026 OpenTransformers Ltd / Scott Bisset
"""
import ctypes, numpy as np, os, sys, json, time

def load_and_run(model_dir, prompt, max_tokens=128, temperature=0.0, top_p=0.9):
    # Load config
    config = json.load(open(os.path.join(model_dir, "config.json")))
    manifest = json.load(open(os.path.join(model_dir, "manifest.json")))
    n_planes = manifest["n_planes"]
    n_layers = config["num_hidden_layers"]
    hidden = config["hidden_size"]
    inter = config["intermediate_size"]
    n_heads = config["num_attention_heads"]
    n_kv_heads = config["num_key_value_heads"]
    head_dim = config.get("head_dim", hidden // n_heads)
    vocab = config["vocab_size"]
    rope_theta = config.get("rope_theta", 10000.0)
    has_attn_bias = 1 if config.get("attention_bias", False) else 0
    tie_embeddings = 1 if config.get("tie_word_embeddings", False) else 0

    print(f"Config: {n_layers}L, hidden={hidden}, inter={inter}, heads={n_heads}/{n_kv_heads}, vocab={vocab}")
    print(f"QK-Norm: yes, Tied embeddings: {'yes' if tie_embeddings else 'no'}, n_planes={n_planes}")

    # Load C engine
    engine_path = os.path.join(os.path.dirname(os.path.abspath(model_dir)), "unary_engine_v2.so")
    lib = ctypes.CDLL(engine_path)

    # Configure function signatures
    lib.model_alloc.restype = ctypes.c_void_p
    lib.model_alloc.argtypes = [
        ctypes.c_int,  # n_planes
        ctypes.c_int,  # hidden
        ctypes.c_int,  # inter
        ctypes.c_int,  # n_heads
        ctypes.c_int,  # n_kv_heads
        ctypes.c_int,  # head_dim
        ctypes.c_int,  # n_layers
        ctypes.c_int,  # vocab
        ctypes.c_float,  # rope_theta
        ctypes.c_int,  # has_attn_bias
        ctypes.c_int,  # tie_embeddings
    ]

    lib.forward_token.restype = ctypes.POINTER(ctypes.c_float)
    lib.forward_token.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_int]

    lib.generate.restype = ctypes.c_int
    lib.generate.argtypes = [
        ctypes.c_void_p,
        ctypes.POINTER(ctypes.c_int), ctypes.c_int,
        ctypes.POINTER(ctypes.c_int), ctypes.c_int,
        ctypes.c_float, ctypes.c_float, ctypes.c_int
    ]

    u16p = ctypes.POINTER(ctypes.c_uint16)
    f32p = ctypes.POINTER(ctypes.c_float)
    u64p = ctypes.POINTER(ctypes.c_uint64)

    lib.model_set_embed.argtypes = [ctypes.c_void_p, u16p]
    lib.model_set_final_norm.argtypes = [ctypes.c_void_p, f32p]
    lib.model_set_lm_head.argtypes = [ctypes.c_void_p, u16p, ctypes.c_int, ctypes.c_int]
    lib.layer_set_norms.argtypes = [ctypes.c_void_p, ctypes.c_int, f32p, f32p]
    lib.layer_set_bias.argtypes = [ctypes.c_void_p, ctypes.c_int, f32p, f32p, f32p]
    lib.layer_set_qk_norm.argtypes = [ctypes.c_void_p, ctypes.c_int, f32p, f32p]
    lib.layer_set_linears.argtypes = [
        ctypes.c_void_p, ctypes.c_int,
        # q: sign, planes, scales, out, in
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        # k
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        # v
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        # o
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        # gate
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        # up
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        # down
        u64p, u64p, f32p, ctypes.c_int, ctypes.c_int,
        ctypes.c_int,  # n_planes
    ]
    lib.model_reset_cache.argtypes = [ctypes.c_void_p]

    # Allocate model
    print("Allocating model...")
    model = lib.model_alloc(
        n_planes, hidden, inter, n_heads, n_kv_heads,
        head_dim, n_layers, vocab, rope_theta,
        has_attn_bias, tie_embeddings
    )

    # Keep references to prevent GC
    _refs = []

    def load_fp16(name):
        fname = name.replace(".", "_") + ".fp16"
        path = os.path.join(model_dir, fname)
        data = np.fromfile(path, dtype=np.uint16)
        _refs.append(data)
        return data.ctypes.data_as(u16p)

    def load_f32_from_fp16(name):
        fname = name.replace(".", "_") + ".fp16"
        path = os.path.join(model_dir, fname)
        data = np.fromfile(path, dtype=np.uint16)
        # Convert FP16 -> FP32
        f32 = data.view(np.float16).astype(np.float32)
        _refs.append(f32)
        return f32.ctypes.data_as(f32p)

    def load_unary(name):
        fname = name.replace(".", "_")
        sign = np.fromfile(os.path.join(model_dir, f"{fname}.sign"), dtype=np.uint64)
        planes = np.fromfile(os.path.join(model_dir, f"{fname}.planes"), dtype=np.uint64)
        scales = np.fromfile(os.path.join(model_dir, f"{fname}.scales"), dtype=np.float32)
        _refs.extend([sign, planes, scales])
        return (sign.ctypes.data_as(u64p), planes.ctypes.data_as(u64p),
                scales.ctypes.data_as(f32p))

    # Load embeddings
    print("Loading embeddings...")
    embed_ptr = load_fp16("model.embed_tokens.weight")
    lib.model_set_embed(model, embed_ptr)

    # Load final norm
    print("Loading final norm...")
    fnorm_ptr = load_f32_from_fp16("model.norm.weight")
    lib.model_set_final_norm(model, fnorm_ptr)

    # Load layers
    print(f"Loading {n_layers} layers...")
    for l in range(n_layers):
        prefix = f"model.layers.{l}"

        # Norms
        in_norm = load_f32_from_fp16(f"{prefix}.input_layernorm.weight")
        post_norm = load_f32_from_fp16(f"{prefix}.post_attention_layernorm.weight")
        lib.layer_set_norms(model, l, in_norm, post_norm)

        # QK-Norm
        q_norm = load_f32_from_fp16(f"{prefix}.self_attn.q_norm.weight")
        k_norm = load_f32_from_fp16(f"{prefix}.self_attn.k_norm.weight")
        lib.layer_set_qk_norm(model, l, q_norm, k_norm)

        # Linear layers
        q_s, q_p, q_sc = load_unary(f"{prefix}.self_attn.q_proj.weight")
        k_s, k_p, k_sc = load_unary(f"{prefix}.self_attn.k_proj.weight")
        v_s, v_p, v_sc = load_unary(f"{prefix}.self_attn.v_proj.weight")
        o_s, o_p, o_sc = load_unary(f"{prefix}.self_attn.o_proj.weight")
        g_s, g_p, g_sc = load_unary(f"{prefix}.mlp.gate_proj.weight")
        u_s, u_p, u_sc = load_unary(f"{prefix}.mlp.up_proj.weight")
        d_s, d_p, d_sc = load_unary(f"{prefix}.mlp.down_proj.weight")

        # Dims from manifest
        q_shape = manifest["unary"][f"{prefix}.self_attn.q_proj.weight"]
        k_shape = manifest["unary"][f"{prefix}.self_attn.k_proj.weight"]
        v_shape = manifest["unary"][f"{prefix}.self_attn.v_proj.weight"]
        o_shape = manifest["unary"][f"{prefix}.self_attn.o_proj.weight"]
        g_shape = manifest["unary"][f"{prefix}.mlp.gate_proj.weight"]
        u_shape = manifest["unary"][f"{prefix}.mlp.up_proj.weight"]
        d_shape = manifest["unary"][f"{prefix}.mlp.down_proj.weight"]

        lib.layer_set_linears(
            model, l,
            q_s, q_p, q_sc, q_shape[0], q_shape[1],
            k_s, k_p, k_sc, k_shape[0], k_shape[1],
            v_s, v_p, v_sc, v_shape[0], v_shape[1],
            o_s, o_p, o_sc, o_shape[0], o_shape[1],
            g_s, g_p, g_sc, g_shape[0], g_shape[1],
            u_s, u_p, u_sc, u_shape[0], u_shape[1],
            d_s, d_p, d_sc, d_shape[0], d_shape[1],
            n_planes
        )

        if (l + 1) % 6 == 0 or l == n_layers - 1:
            print(f"  Loaded layer {l+1}/{n_layers}")

    # Tokenize
    print("Tokenizing prompt...")
    from transformers import AutoTokenizer
    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    input_ids = tokenizer.encode(prompt)
    print(f"Prompt: {len(input_ids)} tokens")

    eos_token = config.get("eos_token_id", 151645)

    # Generate
    prompt_arr = (ctypes.c_int * len(input_ids))(*input_ids)
    out_arr = (ctypes.c_int * max_tokens)()

    print(f"\nGenerating (temp={temperature}, top_p={top_p})...")
    t0 = time.time()
    n_generated = lib.generate(
        model, prompt_arr, len(input_ids),
        out_arr, max_tokens,
        ctypes.c_float(temperature), ctypes.c_float(top_p),
        eos_token
    )
    dt = time.time() - t0

    out_ids = [out_arr[i] for i in range(n_generated)]
    text = tokenizer.decode(out_ids, skip_special_tokens=True)

    total_tokens = len(input_ids) + n_generated
    print(f"\n=== Output ({n_generated} tokens in {dt:.1f}s = {n_generated/dt:.1f} tok/s) ===")
    print(text)
    print(f"\nPrefill: {len(input_ids)} tokens, Decode: {n_generated} tokens")
    print(f"Total time: {dt:.1f}s, Speed: {total_tokens/dt:.1f} tok/s total, {n_generated/dt:.1f} tok/s decode")

    return text

if __name__ == "__main__":
    model_dir = sys.argv[1] if len(sys.argv) > 1 else "qwen3-4b-thinking-unary"
    prompt = sys.argv[2] if len(sys.argv) > 2 else "What is 2+2? Think step by step."
    max_tokens = int(sys.argv[3]) if len(sys.argv) > 3 else 64
    load_and_run(model_dir, prompt, max_tokens=max_tokens)