From be3833cc989d452f67f091d0e283b7e7a10d708f Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 20:47:15 -0400 Subject: [PATCH 1/4] python: centralize ctypes binder for model dylibs Add inference/c_binder.py exposing bind(family, abi). Each family declares a verb -> (argtypes, restype) ABI dict and the binder opens libsk.dylib (SK_DYLIB env / build/libsk.dylib), resolves sk__, and applies signatures. Optional symbols use optional() and are silently skipped when absent. Collapses each wrapper's _load() from ~25-35 lines of repetitive .argtypes/.restype assignments to a single bind() call plus a small ABI dict. Net Python LOC: -41. Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/inference/c_binder.py | 61 ++++++++++++++++++++++ SuperKittens/models/deepseek/deepseek.py | 31 +++++------ SuperKittens/models/gemma/gemma4/gemma4.py | 60 ++++++++------------- SuperKittens/models/mamba2/mamba2.py | 40 ++++++-------- SuperKittens/models/qwen/qwen.py | 45 ++++++---------- 5 files changed, 129 insertions(+), 108 deletions(-) create mode 100644 SuperKittens/inference/c_binder.py 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/models/deepseek/deepseek.py b/SuperKittens/models/deepseek/deepseek.py index 34a3bf3..3fe94e7 100644 --- a/SuperKittens/models/deepseek/deepseek.py +++ b/SuperKittens/models/deepseek/deepseek.py @@ -28,6 +28,8 @@ from pathlib import Path from dataclasses import dataclass, asdict +from SuperKittens.inference.c_binder import bind + # ─── C ABI ────────────────────────────────────────────────────────────────── @@ -88,26 +90,21 @@ 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 diff --git a/SuperKittens/models/gemma/gemma4/gemma4.py b/SuperKittens/models/gemma/gemma4/gemma4.py index 6c88098..890ae1c 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4.py +++ b/SuperKittens/models/gemma/gemma4/gemma4.py @@ -5,6 +5,8 @@ from pathlib import Path from dataclasses import dataclass +from SuperKittens.inference.c_binder import bind, optional + _VARIANT_TO_DIR = { "e2b": "gemma-4-E2B-it", @@ -51,48 +53,28 @@ 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 diff --git a/SuperKittens/models/mamba2/mamba2.py b/SuperKittens/models/mamba2/mamba2.py index 99e1ec3..1567d53 100644 --- a/SuperKittens/models/mamba2/mamba2.py +++ b/SuperKittens/models/mamba2/mamba2.py @@ -14,6 +14,8 @@ from pathlib import Path from typing import Optional +from SuperKittens.inference.c_binder import bind + @dataclass class Mamba2Config: @@ -79,32 +81,22 @@ class _CConfig(ctypes.Structure): class Mamba2Model: """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 + if lib_path is not None: + os.environ.setdefault("SK_DYLIB", lib_path) + self._lib = bind("mamba2", self._ABI) c_cfg = _CConfig( cfg.batch, cfg.seq_max, cfg.n_layers, cfg.d_model, cfg.intermediate, diff --git a/SuperKittens/models/qwen/qwen.py b/SuperKittens/models/qwen/qwen.py index 55031c8..615f064 100644 --- a/SuperKittens/models/qwen/qwen.py +++ b/SuperKittens/models/qwen/qwen.py @@ -13,6 +13,7 @@ from dataclasses import dataclass, asdict from SuperKittens.inference.generation import Model +from SuperKittens.inference.c_binder import bind, optional class _Config(ctypes.Structure): @@ -59,37 +60,25 @@ 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 From 1351878584595b909720e16bb1adeb3790f595eb Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 20:48:48 -0400 Subject: [PATCH 2/4] python: lift lifecycle methods onto Model base Move close/__enter__/__exit__/__del__/__repr__ to inference.generation.Model. Subclasses set _destroy_fn (and optional _repr_fields tuple); the base handles handle-nulling, _w_keep clearing, and a uniform repr template. Have Gemma4, DeepSeek, and Mamba2Model inherit Model so all four families share the lifecycle path. Drop the per-family duplicate methods. Net Python LOC across wrappers: -52 (counting the +36 added to generation.py base). Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/inference/generation.py | 36 ++++++++++++++++++++++ SuperKittens/models/deepseek/deepseek.py | 24 ++++----------- SuperKittens/models/gemma/gemma4/gemma4.py | 16 ++++------ SuperKittens/models/mamba2/mamba2.py | 11 ++----- SuperKittens/models/qwen/qwen.py | 20 +++--------- 5 files changed, 55 insertions(+), 52 deletions(-) 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 3fe94e7..2583fdd 100644 --- a/SuperKittens/models/deepseek/deepseek.py +++ b/SuperKittens/models/deepseek/deepseek.py @@ -29,6 +29,7 @@ from dataclasses import dataclass, asdict from SuperKittens.inference.c_binder import bind +from SuperKittens.inference.generation import Model # ─── C ABI ────────────────────────────────────────────────────────────────── @@ -179,9 +180,12 @@ 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() @@ -190,6 +194,7 @@ def __init__(self, config: Config | str | None = None): else: self.cfg = config lib = _load() + self._destroy_fn = lib.sk_deepseek_destroy self._cstruct = self.cfg._to_c() self._h = lib.sk_deepseek_create(ctypes.byref(self._cstruct)) if not self._h: @@ -334,20 +339,3 @@ def chat(self, text: str | list, *, max_new_tokens: int = 64) -> str: # 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 890ae1c..0e09a04 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4.py +++ b/SuperKittens/models/gemma/gemma4/gemma4.py @@ -6,6 +6,7 @@ from dataclasses import dataclass from SuperKittens.inference.c_binder import bind, optional +from SuperKittens.inference.generation import Model _VARIANT_TO_DIR = { @@ -175,14 +176,18 @@ def _to_cstruct(c: Gemma4Config) -> _Config: return cs -class Gemma4: +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") @@ -457,12 +462,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 1567d53..91973a8 100644 --- a/SuperKittens/models/mamba2/mamba2.py +++ b/SuperKittens/models/mamba2/mamba2.py @@ -15,6 +15,7 @@ from typing import Optional from SuperKittens.inference.c_binder import bind +from SuperKittens.inference.generation import Model @dataclass @@ -78,7 +79,7 @@ class _CConfig(ctypes.Structure): ] -class Mamba2Model: +class Mamba2Model(Model): """ctypes binding for libSuperKittens Mamba 2 C ABI.""" _ABI = { @@ -97,6 +98,7 @@ def __init__(self, cfg: Mamba2Config, lib_path: str | None = None): 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 = _CConfig( cfg.batch, cfg.seq_max, cfg.n_layers, cfg.d_model, cfg.intermediate, @@ -166,13 +168,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 615f064..55a613e 100644 --- a/SuperKittens/models/qwen/qwen.py +++ b/SuperKittens/models/qwen/qwen.py @@ -129,11 +129,15 @@ def _to_c(self) -> _Config: 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._destroy_fn = lib.sk_qwen_destroy self._cstruct = self.cfg._to_c() self._h = lib.sk_qwen_create(ctypes.byref(self._cstruct)) if not self._h: @@ -285,19 +289,3 @@ def decode_step(self) -> int: 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})") From cb199acf2254ea678f3d78924ff14e18f46d7522 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 20:50:11 -0400 Subject: [PATCH 3/4] python: dedup Config -> ctypes Structure conversion Add CtypesConfig mixin in inference/c_binder.py providing to_c(struct_cls) that copies dataclass fields into a Structure, skipping fields the Structure doesn't declare. All four family Config dataclasses now inherit it; per- family _to_c / _to_cstruct bodies collapse to a single call. Drops the unused asdict imports too. Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/models/deepseek/deepseek.py | 13 ++++------ SuperKittens/models/gemma/gemma4/gemma4.py | 28 +++------------------- SuperKittens/models/mamba2/mamba2.py | 11 +++------ SuperKittens/models/qwen/qwen.py | 13 ++++------ 4 files changed, 14 insertions(+), 51 deletions(-) diff --git a/SuperKittens/models/deepseek/deepseek.py b/SuperKittens/models/deepseek/deepseek.py index 2583fdd..61c450c 100644 --- a/SuperKittens/models/deepseek/deepseek.py +++ b/SuperKittens/models/deepseek/deepseek.py @@ -26,9 +26,9 @@ 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 +from SuperKittens.inference.c_binder import bind, CtypesConfig from SuperKittens.inference.generation import Model @@ -112,7 +112,7 @@ def _load(): # ─── Config ───────────────────────────────────────────────────────────────── @dataclass -class Config: +class Config(CtypesConfig): """Model dimensions. Use `Config.preset(name)` for known variants.""" n_layers: int = 60 d_model: int = 7168 @@ -169,11 +169,6 @@ 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 @@ -195,7 +190,7 @@ def __init__(self, config: Config | str | None = None): self.cfg = config lib = _load() self._destroy_fn = lib.sk_deepseek_destroy - self._cstruct = self.cfg._to_c() + 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 " diff --git a/SuperKittens/models/gemma/gemma4/gemma4.py b/SuperKittens/models/gemma/gemma4/gemma4.py index 0e09a04..9c7f006 100644 --- a/SuperKittens/models/gemma/gemma4/gemma4.py +++ b/SuperKittens/models/gemma/gemma4/gemma4.py @@ -5,7 +5,7 @@ from pathlib import Path from dataclasses import dataclass -from SuperKittens.inference.c_binder import bind, optional +from SuperKittens.inference.c_binder import bind, optional, CtypesConfig from SuperKittens.inference.generation import Model @@ -80,7 +80,7 @@ def _load(): @dataclass -class Gemma4Config: +class Gemma4Config(CtypesConfig): n_layers: int local_period: int d_model: int @@ -151,29 +151,7 @@ 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 + return c.to_c(_Config) class Gemma4(Model): diff --git a/SuperKittens/models/mamba2/mamba2.py b/SuperKittens/models/mamba2/mamba2.py index 91973a8..ef15b1b 100644 --- a/SuperKittens/models/mamba2/mamba2.py +++ b/SuperKittens/models/mamba2/mamba2.py @@ -14,12 +14,12 @@ from pathlib import Path from typing import Optional -from SuperKittens.inference.c_binder import bind +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 @@ -100,12 +100,7 @@ def __init__(self, cfg: Mamba2Config, lib_path: str | None = None): self._lib = bind("mamba2", self._ABI) self._destroy_fn = self._lib.sk_mamba2_destroy - 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, - ) + 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") diff --git a/SuperKittens/models/qwen/qwen.py b/SuperKittens/models/qwen/qwen.py index 55a613e..bfbfc79 100644 --- a/SuperKittens/models/qwen/qwen.py +++ b/SuperKittens/models/qwen/qwen.py @@ -10,10 +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 +from SuperKittens.inference.c_binder import bind, optional, CtypesConfig class _Config(ctypes.Structure): @@ -83,7 +83,7 @@ def _load(): @dataclass -class Config: +class Config(CtypesConfig): # Qwen3-32B dense defaults (per HF config.json). n_layers: int = 64 d_model: int = 5120 @@ -120,11 +120,6 @@ 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.""" @@ -138,7 +133,7 @@ def __init__(self, config: Config | str | None = None): else: self.cfg = config lib = _load() self._destroy_fn = lib.sk_qwen_destroy - self._cstruct = self.cfg._to_c() + 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)") From 50a6b07ac6b313601747a48fd3985701a57d3fd5 Mon Sep 17 00:00:00 2001 From: Alazar Manakelew Date: Wed, 13 May 2026 20:51:46 -0400 Subject: [PATCH 4/4] python: drop duplicate generate/prefill/decode_step in qwen + deepseek Both wrappers inherit Model and now use base.generate. Qwen had no own generate; remove the unused prefill/decode_step shims (no external callers). DeepSeek's duplicate generate/prefill/decode_step are deleted; add a base-compatible _forward and keep forward() as a thin int wrapper. DeepSeek's chat() now passes max_new_tokens by keyword (base.generate's arg is keyword-only) and forwards the tokenizer's eos_id. Gemma4 keeps its bf16 last_logits path / sample loop (numerics); Mamba2 has no generate yet and is untouched here. Co-Authored-By: Claude Opus 4.7 (1M context) --- SuperKittens/models/deepseek/deepseek.py | 41 ++++++------------------ SuperKittens/models/qwen/qwen.py | 9 ------ 2 files changed, 9 insertions(+), 41 deletions(-) diff --git a/SuperKittens/models/deepseek/deepseek.py b/SuperKittens/models/deepseek/deepseek.py index 61c450c..181beb7 100644 --- a/SuperKittens/models/deepseek/deepseek.py +++ b/SuperKittens/models/deepseek/deepseek.py @@ -262,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( @@ -274,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 @@ -330,7 +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) diff --git a/SuperKittens/models/qwen/qwen.py b/SuperKittens/models/qwen/qwen.py index bfbfc79..8d4ea3a 100644 --- a/SuperKittens/models/qwen/qwen.py +++ b/SuperKittens/models/qwen/qwen.py @@ -275,12 +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]) -