diff --git a/SuperKittens/inference/c_binder.py b/SuperKittens/inference/c_binder.py new file mode 100644 index 0000000..fee0d14 --- /dev/null +++ b/SuperKittens/inference/c_binder.py @@ -0,0 +1,61 @@ +"""Centralized ctypes binder for SK model family dylibs. + +Each family declares a verb -> (argtypes, restype) dict and calls bind(family, +abi). The binder opens libsk.dylib (SK_DYLIB env or build/libsk.dylib relative +to repo root), resolves each verb as ``sk__``, applies the +signatures, and returns the CDLL handle. Optional symbols are declared by +wrapping the (argtypes, restype) tuple with ``optional(...)``; missing optional +symbols are silently skipped instead of raising. +""" +from __future__ import annotations + +import ctypes +import os +from dataclasses import asdict +from pathlib import Path + + +def optional(argtypes, restype): + """Mark an ABI entry as optional (skipped if the symbol is absent).""" + return (argtypes, restype, True) + + +def _default_dylib_path() -> str: + return str(Path(__file__).resolve().parents[2] / "build" / "libsk.dylib") + + +_cache: dict[str, ctypes.CDLL] = {} + + +def bind(family: str, abi: dict) -> ctypes.CDLL: + """Open libsk.dylib and register ``sk__`` for each abi entry. + + abi values are ``(argtypes, restype)`` tuples; use ``optional(...)`` for + symbols that may be absent in older builds. + """ + dylib = os.environ.get("SK_DYLIB") or os.environ.get("SK_LIB") or _default_dylib_path() + lib = _cache.get(dylib) + if lib is None: + lib = ctypes.CDLL(dylib) + _cache[dylib] = lib + for verb, sig in abi.items(): + sym = f"sk_{family}_{verb}" + is_optional = len(sig) == 3 and sig[2] is True + argtypes, restype = sig[0], sig[1] + if is_optional and not hasattr(lib, sym): + continue + fn = getattr(lib, sym) + fn.argtypes = list(argtypes) + fn.restype = restype + return lib + + +class CtypesConfig: + """Mixin: copy dataclass fields into a ctypes.Structure via to_c().""" + + def to_c(self, struct_cls): + cs = struct_cls() + for k, v in asdict(self).items(): + if hasattr(cs, k): + setattr(cs, k, v) + return cs diff --git a/SuperKittens/inference/generation.py b/SuperKittens/inference/generation.py index 9c50f3b..05b3f3a 100644 --- a/SuperKittens/inference/generation.py +++ b/SuperKittens/inference/generation.py @@ -13,10 +13,46 @@ class Model: _last_logits() -> np.ndarray # (vocab,) fp16 of last position reset() # reset per-sequence cursor tokenizer # optional, attached by from_pretrained + + Lifecycle (close/__enter__/__exit__/__del__/__repr__) is provided here. + Subclasses set: + _destroy_fn — callable(handle) destroying the native handle (required for close) + _handle_attr — name of the instance attribute holding the handle (default "_h") + _repr_fields — optional tuple of (label, attr) for __repr__ (else uses cfg fields) """ cfg: Any tokenizer: Optional[Any] = None + _destroy_fn = None + _handle_attr: str = "_h" + _repr_fields: tuple = () + + def close(self) -> None: + h = getattr(self, self._handle_attr, None) + if h and self._destroy_fn is not None: + self._destroy_fn(h) + setattr(self, self._handle_attr, None) + if hasattr(self, "_w_keep"): + self._w_keep = None + + def __enter__(self): + return self + + def __exit__(self, *_): + self.close() + + def __del__(self): + try: + self.close() + except Exception: + pass + + def __repr__(self) -> str: + cls = type(self).__name__ + if self._repr_fields: + parts = [f"{label}={getattr(self.cfg, attr, '?')}" for label, attr in self._repr_fields] + return f"{cls}({', '.join(parts)})" + return f"{cls}(cfg={type(self.cfg).__name__})" def _forward(self, input_ids: np.ndarray) -> np.ndarray: raise NotImplementedError diff --git a/SuperKittens/models/deepseek/deepseek.py b/SuperKittens/models/deepseek/deepseek.py index 34a3bf3..181beb7 100644 --- a/SuperKittens/models/deepseek/deepseek.py +++ b/SuperKittens/models/deepseek/deepseek.py @@ -26,7 +26,10 @@ import ctypes, os import numpy as np from pathlib import Path -from dataclasses import dataclass, asdict +from dataclasses import dataclass + +from SuperKittens.inference.c_binder import bind, CtypesConfig +from SuperKittens.inference.generation import Model # ─── C ABI ────────────────────────────────────────────────────────────────── @@ -88,33 +91,28 @@ class _Weights(ctypes.Structure): _fields_ = [(n, ctypes.c_void_p) for n in _WEIGHT_FIELDS] +DEEPSEEK_ABI = { + "create": ([ctypes.POINTER(_Config)], ctypes.c_void_p), + "load_weights": ([ctypes.c_void_p, ctypes.POINTER(_Weights)], ctypes.c_int), + "forward": ([ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32)], ctypes.c_int), + "reset": ([ctypes.c_void_p], None), + "destroy": ([ctypes.c_void_p], None), +} + + _lib = None def _load(): global _lib - if _lib is not None: return _lib - dylib = os.environ.get("SK_DYLIB", - str(Path(__file__).resolve().parents[3] / "build" / "libsk.dylib")) - _lib = ctypes.CDLL(dylib) - _lib.sk_deepseek_create.argtypes = [ctypes.POINTER(_Config)] - _lib.sk_deepseek_create.restype = ctypes.c_void_p - _lib.sk_deepseek_load_weights.argtypes = [ctypes.c_void_p, ctypes.POINTER(_Weights)] - _lib.sk_deepseek_load_weights.restype = ctypes.c_int - _lib.sk_deepseek_forward.argtypes = [ctypes.c_void_p, - ctypes.POINTER(ctypes.c_int32), - ctypes.c_uint32, - ctypes.POINTER(ctypes.c_int32)] - _lib.sk_deepseek_forward.restype = ctypes.c_int - _lib.sk_deepseek_reset.argtypes = [ctypes.c_void_p] - _lib.sk_deepseek_reset.restype = None - _lib.sk_deepseek_destroy.argtypes = [ctypes.c_void_p] - _lib.sk_deepseek_destroy.restype = None + if _lib is None: + _lib = bind("deepseek", DEEPSEEK_ABI) return _lib # ─── Config ───────────────────────────────────────────────────────────────── @dataclass -class Config: +class Config(CtypesConfig): """Model dimensions. Use `Config.preset(name)` for known variants.""" n_layers: int = 60 d_model: int = 7168 @@ -171,20 +169,18 @@ def preset(cls, name: str) -> "Config": vocab_size=1024, seq_max=16, cache_max=64) raise ValueError(f"unknown DeepSeek preset: {name!r}") - def _to_c(self) -> _Config: - cs = _Config() - for k, v in asdict(self).items(): setattr(cs, k, v) - return cs - @property def dk(self) -> int: return self.qk_nope_dim + self.qk_rope_dim # ─── Main handle ──────────────────────────────────────────────────────────── -class DeepSeek: +class DeepSeek(Model): """Stateful DeepSeek V4 Flash inference handle.""" + _repr_fields = (("L", "n_layers"), ("D", "d_model"), ("H", "n_heads"), + ("E", "n_expert"), ("top_k", "top_k")) + def __init__(self, config: Config | str | None = None): if config is None: self.cfg = Config() @@ -193,7 +189,8 @@ def __init__(self, config: Config | str | None = None): else: self.cfg = config lib = _load() - self._cstruct = self.cfg._to_c() + self._destroy_fn = lib.sk_deepseek_destroy + self._cstruct = self.cfg.to_c(_Config) self._h = lib.sk_deepseek_create(ctypes.byref(self._cstruct)) if not self._h: raise RuntimeError("sk_deepseek_create failed (likely missing PSO or " @@ -265,9 +262,9 @@ def required_weights(self) -> tuple[str, ...]: return _WEIGHT_FIELDS # ─── inference ─── - def forward(self, input_ids) -> int: - """Forward pass over `seq` tokens; returns argmax token id.""" - ids = np.asarray(input_ids, dtype=np.int32).reshape(-1) + def _forward(self, input_ids: np.ndarray) -> np.ndarray: + """Model-base contract: int32 ids in, (batch,) int32 argmax out.""" + ids = np.ascontiguousarray(np.asarray(input_ids, dtype=np.int32)).reshape(-1) seq = ids.size // self.cfg.batch out = np.empty((self.cfg.batch,), dtype=np.int32) rc = _load().sk_deepseek_forward( @@ -277,36 +274,12 @@ def forward(self, input_ids) -> int: out.ctypes.data_as(ctypes.POINTER(ctypes.c_int32))) if rc: raise RuntimeError(f"forward failed: {rc}") self._last_token = int(out[0]) - return self._last_token - - def prefill(self, input_ids) -> int: - """Eats the prompt and returns the argmax of the last position.""" - self.reset() - return self.forward(input_ids) - - def decode_step(self) -> int: - """One decode step continuing from the last forward's argmax.""" - if self._last_token is None: - raise RuntimeError("decode_step called before prefill/forward") - return self.forward([self._last_token]) - - def generate(self, input_ids, max_new_tokens: int = 64, - *, stop_on_eos: bool = True) -> list[int]: - """Greedy decode: prefill, then loop. Stops at max_new_tokens or EOS - (when a tokenizer with an EOS id is attached and stop_on_eos=True). - - Sampling is currently argmax only; temperature / top-p / top-k will - plug in here once the GPU sampler kernel is wired.""" - out: list[int] = [self.prefill(input_ids)] - if stop_on_eos and self._tok and self._tok.is_eos(out[-1]): - return out - for _ in range(max_new_tokens - 1): - tok = self.decode_step() - out.append(tok) - if stop_on_eos and self._tok and self._tok.is_eos(tok): - break return out + def forward(self, input_ids) -> int: + """Backwards-compat wrapper returning the argmax token id.""" + return int(self._forward(input_ids)[0]) + # ── tokenizer + chat API ───────────────────────────────────────── def attach_tokenizer(self, tokenizer) -> "DeepSeek": """Attach a tokenizer (DeepSeekTokenizer or any object with @@ -333,24 +306,8 @@ def chat(self, text: str | list, *, max_new_tokens: int = 64) -> str: raise RuntimeError("attach_tokenizer() first") msgs = [{"role": "user", "content": text}] if isinstance(text, str) else text ids = self._tok.encode_chat(msgs) - out_ids = self.generate(ids, max_new_tokens=max_new_tokens) + out_ids = self.generate(ids, max_new_tokens=max_new_tokens, + eos_id=getattr(self._tok, "eos_id", None)) # Strip the prompt prefix; only return the newly generated portion. return self._tok.decode(out_ids[len(ids):] if len(out_ids) > len(ids) else out_ids) - # ─── lifecycle ─── - def close(self) -> None: - if self._h: - _load().sk_deepseek_destroy(self._h) - self._h = None - self._w_keep = None - - def __enter__(self): return self - def __exit__(self, *_): self.close() - def __del__(self): - try: self.close() - except Exception: pass - - def __repr__(self) -> str: - c = self.cfg - return (f"DeepSeek(L={c.n_layers}, D={c.d_model}, H={c.n_heads}, " - f"dk={c.dk}, dv={c.v_head_dim}, E={c.n_expert}, top_k={c.top_k})") diff --git a/SuperKittens/models/gemma/gemma4/gemma4.py b/SuperKittens/models/gemma/gemma4/gemma4.py index 6c88098..9c7f006 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4.py +++ b/SuperKittens/models/gemma/gemma4/gemma4.py @@ -5,6 +5,9 @@ from pathlib import Path from dataclasses import dataclass +from SuperKittens.inference.c_binder import bind, optional, CtypesConfig +from SuperKittens.inference.generation import Model + _VARIANT_TO_DIR = { "e2b": "gemma-4-E2B-it", @@ -51,53 +54,33 @@ class _Weights(ctypes.Structure): )] +GEMMA4_ABI = { + "create": ([ctypes.POINTER(_Config)], ctypes.c_void_p), + "load_weights": ([ctypes.c_void_p, ctypes.POINTER(_Weights)], ctypes.c_int), + "forward": ([ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32)], ctypes.c_int), + "reset": ([ctypes.c_void_p], None), + "destroy": ([ctypes.c_void_p], None), + "load_safetensors": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "load_safetensors_index": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "set_rope_tables": ([ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), + "get_last_logits": ([ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), + "set_dump_enabled": ([ctypes.c_void_p, ctypes.c_int], None), + "dump_layer": ([ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p], ctypes.c_int), +} + + _lib = None def _load(): global _lib - if _lib is not None: - return _lib - dylib = os.environ.get( - "SK_DYLIB", - str(Path(__file__).resolve().parents[4] / "build" / "libsk.dylib")) - _lib = ctypes.CDLL(dylib) - - _lib.sk_gemma4_create.argtypes = [ctypes.POINTER(_Config)] - _lib.sk_gemma4_create.restype = ctypes.c_void_p - - _lib.sk_gemma4_load_weights.argtypes = [ctypes.c_void_p, ctypes.POINTER(_Weights)] - _lib.sk_gemma4_load_weights.restype = ctypes.c_int - - _lib.sk_gemma4_forward.argtypes = [ - ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), - ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32)] - _lib.sk_gemma4_forward.restype = ctypes.c_int - - _lib.sk_gemma4_reset.argtypes = [ctypes.c_void_p] - _lib.sk_gemma4_reset.restype = None - - _lib.sk_gemma4_destroy.argtypes = [ctypes.c_void_p] - _lib.sk_gemma4_destroy.restype = None - - _lib.sk_gemma4_load_safetensors.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - _lib.sk_gemma4_load_safetensors.restype = ctypes.c_int - _lib.sk_gemma4_load_safetensors_index.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - _lib.sk_gemma4_load_safetensors_index.restype = ctypes.c_int - _lib.sk_gemma4_set_rope_tables.argtypes = [ctypes.c_void_p, - ctypes.c_void_p, ctypes.c_void_p, - ctypes.c_void_p, ctypes.c_void_p] - _lib.sk_gemma4_set_rope_tables.restype = ctypes.c_int - _lib.sk_gemma4_get_last_logits.argtypes = [ctypes.c_void_p, ctypes.c_void_p] - _lib.sk_gemma4_get_last_logits.restype = ctypes.c_int - - _lib.sk_gemma4_set_dump_enabled.argtypes = [ctypes.c_void_p, ctypes.c_int] - _lib.sk_gemma4_set_dump_enabled.restype = None - _lib.sk_gemma4_dump_layer.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] - _lib.sk_gemma4_dump_layer.restype = ctypes.c_int + if _lib is None: + _lib = bind("gemma4", GEMMA4_ABI) return _lib @dataclass -class Gemma4Config: +class Gemma4Config(CtypesConfig): n_layers: int local_period: int d_model: int @@ -168,39 +151,21 @@ def _preset(name: str) -> Gemma4Config: def _to_cstruct(c: Gemma4Config) -> _Config: - cs = _Config() - cs.batch = c.batch - cs.seq_max = c.seq_max - cs.cache_max = c.cache_max - cs.n_layers = c.n_layers - cs.local_period = c.local_period - cs.d_model = c.d_model - cs.n_int = c.n_int - cs.n_heads = c.n_heads - cs.n_kv_heads_local = c.n_kv_heads_local - cs.n_kv_heads_global = c.n_kv_heads_global - cs.head_dim_local = c.head_dim_local - cs.head_dim_global = c.head_dim_global - cs.window = c.window - cs.prope_p_pairs = c.prope_p_pairs - cs.vocab_size = c.vocab_size - cs.ple_dim = c.ple_dim - cs.has_ple = 1 if c.has_ple else 0 - cs.eps = c.eps - cs.final_logit_softcap = float(c.final_logit_softcap) - cs.use_double_wide_mlp = 1 if c.use_double_wide_mlp else 0 - cs.num_kv_shared_layers = int(c.num_kv_shared_layers) - return cs - - -class Gemma4: + return c.to_c(_Config) + + +class Gemma4(Model): """Stateful Gemma 4 inference handle. Holds KV cache between forwards.""" + _repr_fields = (("L", "n_layers"), ("D", "d_model"), ("H", "n_heads"), + ("n_int", "n_int")) + def __init__(self, variant_or_config): cfg = _preset(variant_or_config) if isinstance(variant_or_config, str) else variant_or_config self.cfg = cfg self._cstruct = _to_cstruct(cfg) lib = _load() + self._destroy_fn = lib.sk_gemma4_destroy self._h = lib.sk_gemma4_create(ctypes.byref(self._cstruct)) if not self._h: raise RuntimeError("sk_gemma4_create failed") @@ -475,12 +440,3 @@ def _get(key, default): print(f"[gemma4] hf-json attach failed: {e}") return m - def close(self): - if self._h: - _load().sk_gemma4_destroy(self._h) - self._h = None - self._w_keep = None - - def __del__(self): - try: self.close() - except Exception: pass diff --git a/SuperKittens/models/mamba2/mamba2.py b/SuperKittens/models/mamba2/mamba2.py index 99e1ec3..ef15b1b 100644 --- a/SuperKittens/models/mamba2/mamba2.py +++ b/SuperKittens/models/mamba2/mamba2.py @@ -14,9 +14,12 @@ from pathlib import Path from typing import Optional +from SuperKittens.inference.c_binder import bind, CtypesConfig +from SuperKittens.inference.generation import Model + @dataclass -class Mamba2Config: +class Mamba2Config(CtypesConfig): batch: int = 1 seq_max: int = 2048 n_layers: int = 24 @@ -76,42 +79,28 @@ class _CConfig(ctypes.Structure): ] -class Mamba2Model: +class Mamba2Model(Model): """ctypes binding for libSuperKittens Mamba 2 C ABI.""" + _ABI = { + "create": ([ctypes.POINTER(_CConfig)], ctypes.c_void_p), + "load_safetensors": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "forward": ([ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), + ctypes.c_uint32, ctypes.POINTER(ctypes.c_int)], ctypes.c_int), + "reset": ([ctypes.c_void_p], None), + "destroy": ([ctypes.c_void_p], None), + "dump_layer": ([ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_size_t], ctypes.c_int), + "get_last_logits": ([ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), + } + def __init__(self, cfg: Mamba2Config, lib_path: str | None = None): self.cfg = cfg - if lib_path is None: - lib_path = os.environ.get("SK_DYLIB") or os.environ.get("SK_LIB") - if lib_path is None: - here = Path(__file__).resolve() - lib_path = str(here.parents[3] / "build" / "libsk.dylib") - self._lib = ctypes.CDLL(lib_path) - - self._lib.sk_mamba2_create.argtypes = [ctypes.POINTER(_CConfig)] - self._lib.sk_mamba2_create.restype = ctypes.c_void_p - self._lib.sk_mamba2_load_safetensors.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - self._lib.sk_mamba2_load_safetensors.restype = ctypes.c_int - self._lib.sk_mamba2_forward.argtypes = [ - ctypes.c_void_p, ctypes.POINTER(ctypes.c_int), ctypes.c_uint32, - ctypes.POINTER(ctypes.c_int), - ] - self._lib.sk_mamba2_forward.restype = ctypes.c_int - self._lib.sk_mamba2_reset.argtypes = [ctypes.c_void_p] - self._lib.sk_mamba2_destroy.argtypes = [ctypes.c_void_p] - self._lib.sk_mamba2_dump_layer.argtypes = [ - ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p, ctypes.c_size_t, - ] - self._lib.sk_mamba2_dump_layer.restype = ctypes.c_int - self._lib.sk_mamba2_get_last_logits.argtypes = [ctypes.c_void_p, ctypes.c_void_p] - self._lib.sk_mamba2_get_last_logits.restype = ctypes.c_int - - c_cfg = _CConfig( - cfg.batch, cfg.seq_max, cfg.n_layers, cfg.d_model, cfg.intermediate, - cfg.n_heads, cfg.head_dim, cfg.state_size, cfg.n_groups, cfg.conv_kernel, - cfg.chunk_size, cfg.vocab_size, cfg.rms_eps, - cfg.time_step_min, cfg.time_step_max, cfg.tie_word_embeddings, - ) + if lib_path is not None: + os.environ.setdefault("SK_DYLIB", lib_path) + self._lib = bind("mamba2", self._ABI) + self._destroy_fn = self._lib.sk_mamba2_destroy + + c_cfg = cfg.to_c(_CConfig) self._h = self._lib.sk_mamba2_create(ctypes.byref(c_cfg)) if not self._h: raise RuntimeError("sk_mamba2_create returned NULL") @@ -174,13 +163,6 @@ def get_last_logits(self): def reset(self) -> None: self._lib.sk_mamba2_reset(self._h) - def __del__(self): - try: - if getattr(self, "_h", None): - self._lib.sk_mamba2_destroy(self._h) - except Exception: - pass - # Registry entry — kept import-light. SPEC = { diff --git a/SuperKittens/models/qwen/qwen.py b/SuperKittens/models/qwen/qwen.py index 55031c8..8d4ea3a 100644 --- a/SuperKittens/models/qwen/qwen.py +++ b/SuperKittens/models/qwen/qwen.py @@ -10,9 +10,10 @@ import ctypes, os import numpy as np from pathlib import Path -from dataclasses import dataclass, asdict +from dataclasses import dataclass from SuperKittens.inference.generation import Model +from SuperKittens.inference.c_binder import bind, optional, CtypesConfig class _Config(ctypes.Structure): @@ -59,42 +60,30 @@ class _Weights(ctypes.Structure): _fields_ = [(n, ctypes.c_void_p) for n in _WEIGHT_FIELDS] +QWEN_ABI = { + "create": ([ctypes.POINTER(_Config)], ctypes.c_void_p), + "load_weights": ([ctypes.c_void_p, ctypes.POINTER(_Weights)], ctypes.c_int), + "forward": ([ctypes.c_void_p, ctypes.POINTER(ctypes.c_int32), + ctypes.c_uint32, ctypes.POINTER(ctypes.c_int32)], ctypes.c_int), + "reset": ([ctypes.c_void_p], None), + "destroy": ([ctypes.c_void_p], None), + "load_safetensors": ([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "load_gguf": optional([ctypes.c_void_p, ctypes.c_char_p], ctypes.c_int), + "set_rope_tables": optional([ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), + "get_last_logits": optional([ctypes.c_void_p, ctypes.c_void_p], ctypes.c_int), +} + + _lib = None def _load(): global _lib - if _lib is not None: return _lib - dylib = os.environ.get("SK_DYLIB", - str(Path(__file__).resolve().parents[3] / "build" / "libsk.dylib")) - _lib = ctypes.CDLL(dylib) - _lib.sk_qwen_create.argtypes = [ctypes.POINTER(_Config)] - _lib.sk_qwen_create.restype = ctypes.c_void_p - _lib.sk_qwen_load_weights.argtypes = [ctypes.c_void_p, ctypes.POINTER(_Weights)] - _lib.sk_qwen_load_weights.restype = ctypes.c_int - _lib.sk_qwen_forward.argtypes = [ctypes.c_void_p, - ctypes.POINTER(ctypes.c_int32), - ctypes.c_uint32, - ctypes.POINTER(ctypes.c_int32)] - _lib.sk_qwen_forward.restype = ctypes.c_int - _lib.sk_qwen_reset.argtypes = [ctypes.c_void_p] - _lib.sk_qwen_reset.restype = None - _lib.sk_qwen_destroy.argtypes = [ctypes.c_void_p] - _lib.sk_qwen_destroy.restype = None - _lib.sk_qwen_load_safetensors.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - _lib.sk_qwen_load_safetensors.restype = ctypes.c_int - if hasattr(_lib, "sk_qwen_load_gguf"): - _lib.sk_qwen_load_gguf.argtypes = [ctypes.c_void_p, ctypes.c_char_p] - _lib.sk_qwen_load_gguf.restype = ctypes.c_int - if hasattr(_lib, "sk_qwen_set_rope_tables"): - _lib.sk_qwen_set_rope_tables.argtypes = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p] - _lib.sk_qwen_set_rope_tables.restype = ctypes.c_int - if hasattr(_lib, "sk_qwen_get_last_logits"): - _lib.sk_qwen_get_last_logits.argtypes = [ctypes.c_void_p, ctypes.c_void_p] - _lib.sk_qwen_get_last_logits.restype = ctypes.c_int + if _lib is None: + _lib = bind("qwen", QWEN_ABI) return _lib @dataclass -class Config: +class Config(CtypesConfig): # Qwen3-32B dense defaults (per HF config.json). n_layers: int = 64 d_model: int = 5120 @@ -131,21 +120,20 @@ def preset(cls, name: str) -> "Config": rope_freq_base=1_000_000.0, rope_n_ctx_orig=64) raise ValueError(f"unknown Qwen preset: {name!r}") - def _to_c(self) -> _Config: - cs = _Config() - for k, v in asdict(self).items(): setattr(cs, k, v) - return cs - class Qwen(Model): """Stateful Qwen3 (dense) inference handle.""" + _repr_fields = (("L", "n_layers"), ("D", "d_model"), ("H", "n_heads"), + ("hd", "head_dim"), ("n_int", "n_int")) + def __init__(self, config: Config | str | None = None): if config is None: self.cfg = Config() elif isinstance(config, str): self.cfg = Config.preset(config) else: self.cfg = config lib = _load() - self._cstruct = self.cfg._to_c() + self._destroy_fn = lib.sk_qwen_destroy + self._cstruct = self.cfg.to_c(_Config) self._h = lib.sk_qwen_create(ctypes.byref(self._cstruct)) if not self._h: raise RuntimeError("sk_qwen_create failed (missing PSO or wrong dims)") @@ -287,28 +275,3 @@ def chat(self, prompt, *, use_chat_template: bool = True, **gen_kwargs) -> str: out_ids = self.generate(np.array(ids, dtype=np.int32), eos_id=eos, **gen_kwargs) return self.tokenizer.decode(out_ids) - def prefill(self, input_ids) -> int: - self.reset() - return self.forward(input_ids) - - def decode_step(self) -> int: - if self._last_token is None: - raise RuntimeError("decode_step called before prefill/forward") - return self.forward([self._last_token]) - - def close(self) -> None: - if self._h: - _load().sk_qwen_destroy(self._h) - self._h = None - self._w_keep = None - - def __enter__(self): return self - def __exit__(self, *_): self.close() - def __del__(self): - try: self.close() - except Exception: pass - - def __repr__(self) -> str: - c = self.cfg - return (f"Qwen(L={c.n_layers}, D={c.d_model}, H={c.n_heads}/" - f"{c.n_kv_heads}KV, hd={c.head_dim}, n_int={c.n_int})")