snowllm 0.1.0__py3-none-any.whl

snowllm/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the SnowLLM project
+
+from ._capi import ABI_VERSION, SnowLLMError, synchronize
+from .geometry import ModelGeometry
+from ._version import __version__
+
+__all__ = ["ABI_VERSION", "ModelGeometry", "SnowLLMError", "__version__", "synchronize"]

snowllm/_capi.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the SnowLLM project
+
+import ctypes
+import importlib.util
+import os
+from ctypes import POINTER, c_char_p, c_float, c_int, c_int64, c_void_p
+
+__all__ = ["SnowLLMError", "build_geometry", "lib", "synchronize",
+           "assert_single_hip_runtime"]
+
+
+class SnowLLMError(RuntimeError):
+    pass
+
+
+class BuildGeometry(ctypes.Structure):
+    _fields_ = [
+        (n, c_int64)
+        for n in (
+            "hidden", "num_heads", "num_kv_heads", "head_size", "block_size", "vocab_size",
+            "qkv_proj_n", "qkv_q_head_stride", "qkv_off_scale", "qkv_off_k", "qkv_off_v",
+            "lin_num_k_heads", "lin_num_v_heads", "lin_head_k", "lin_head_v", "lin_key_dim",
+            "lin_value_dim", "lin_conv_dim", "lin_conv_k", "lin_conv_state", "lin_in_proj_n",
+            "lin_in_proj_n_pad",
+            "lin_in_proj_ba_n_pad", "lin_off_qkv", "lin_off_z", "lin_off_b", "lin_off_a",
+            "moe_num_experts", "moe_shared_expert", "moe_num_slabs", "moe_topk", "moe_topk_all",
+            "moe_inter", "sampling_max_k", "lm_head_decode_max_m",
+        )
+    ]
+
+
+def build_geometry() -> "BuildGeometry":
+    g = BuildGeometry()
+    lib.snowllm_model_config(ctypes.byref(g))
+    return g
+
+
+def _find_lib() -> str:
+    if env := os.environ.get("SNOWLLM_LIB"):
+        return env
+    try:
+        import snowllm_kernels
+    except ModuleNotFoundError as e:
+        if e.name != "snowllm_kernels":
+            raise
+        raise SnowLLMError(
+            "the snowllm-kernels package is not installed, so there are no kernels to call. "
+            "`pip install snowllm-kernels` (it carries libsnowllm.so and needs no compiler), or "
+            "point SNOWLLM_LIB at a libsnowllm.so you built yourself."
+        ) from e
+    lib = str(snowllm_kernels.LIB)
+    if not os.path.exists(lib):
+        raise SnowLLMError(
+            f"snowllm-kernels resolved to {os.path.dirname(snowllm_kernels.__file__)} but there is "
+            f"no libsnowllm.so at {lib}. If that is a source checkout, build it "
+            f"(`cmake --build build-release`); if it is an install, reinstall the wheel."
+        )
+    return lib
+
+
+def _preload_hip_runtime() -> None:
+    spec = importlib.util.find_spec("_rocm_sdk_core")
+    if spec is None or not spec.submodule_search_locations:
+        return
+    sdk = os.path.join(spec.submodule_search_locations[0], "lib", "libamdhip64.so.7")
+    if os.path.exists(sdk):
+        try:
+            ctypes.CDLL(sdk, mode=ctypes.RTLD_GLOBAL)
+        except OSError:
+            pass
+
+
+_preload_hip_runtime()
+_lib_path = _find_lib()
+lib = ctypes.CDLL(_lib_path)
+
+ABI_VERSION = 7
+
+
+def _check_abi() -> None:
+    try:
+        lib.snowllm_abi_version.restype = c_int
+        got = lib.snowllm_abi_version()
+    except AttributeError:
+        got = 0
+    if got != ABI_VERSION:
+        raise SnowLLMError(
+            f"ABI mismatch: this snowllm expects {ABI_VERSION}, "
+            f"{_lib_path} provides {got or 'no version at all'}. Rebuild the one that is stale."
+        )
+
+
+_check_abi()
+
+P, S = c_void_p, c_void_p
+_SIGS = [
+    ("snowllm_model_config", None, [POINTER(BuildGeometry)]),
+    ("snowllm_shuffle_bytes", c_int64, [c_int64]),
+    ("snowllm_prefill_row_quantum", c_int64, []),
+    ("snowllm_sampling_max_k", c_int64, []),
+    ("snowllm_error_string", c_char_p, [c_int]),
+    ("snowllm_synchronize", c_int, [S]),
+
+    ("snowllm_qkv_proj_shuffle_w", c_int, [P, P, S]),
+    ("snowllm_qkv_proj_scratch_bytes", c_int64, [c_int64]),
+    ("snowllm_qkv_proj_shuffle_w_fp8", c_int, [P, P, S]),
+    ("snowllm_attn_out_scale_oproj_shuffle_w", c_int, [P, P, S]),
+    ("snowllm_attn_out_scale_oproj_scratch_bytes", c_int64, [c_int64]),
+    ("snowllm_attn_out_scale_oproj_shuffle_w_fp8", c_int, [P, P, S]),
+    ("snowllm_linear_in_proj_shuffle_w", c_int, [P, P, S]),
+    ("snowllm_linear_out_proj_shuffle_w", c_int, [P, P, S]),
+    ("snowllm_linear_out_proj_shuffle_w_fp8", c_int, [P, P, S]),
+    ("snowllm_linear_in_proj_qz_shuffle_w_fp8", c_int, [P, P, S]),
+    ("snowllm_linear_in_proj_ba_shuffle_w", c_int, [P, P, S]),
+    ("snowllm_fused_linear_attn_workspace_bytes", c_int64, [c_int64]),
+    ("snowllm_mtp_fc_shuffle_w", c_int, [P, P, S]),
+    ("snowllm_mtp_fc_scratch_bytes", c_int64, [c_int64]),
+    ("snowllm_mtp_pre_fc", c_int, [P, P, P, P, P, c_int64, c_float, S]),
+    ("snowllm_mtp_fc", c_int, [P, P, P, P, c_int64, c_int, S]),
+    ("snowllm_moe_shuffle_gate_up", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_moe_shuffle_down", c_int, [P, P, c_int64, S]),
+    ("snowllm_moe_shuffle_gate_up_fused", c_int, [P, P, c_int64, S]),
+    ("snowllm_proj_scale_shuffle_fp8", c_int, [P, P, c_int64, c_int64, S]),
+    ("snowllm_moe_scale_shuffle_gate_up_fp8", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_moe_scale_shuffle_down_fp8", c_int, [P, P, c_int64, S]),
+    ("snowllm_moe_shuffle_gate_up_fp8", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_moe_shuffle_down_fp8", c_int, [P, P, c_int64, S]),
+    ("snowllm_moe_shuffle_router_bytes", c_int64, []),
+    ("snowllm_moe_shuffle_router", c_int, [P, P, S]),
+    ("snowllm_moe_workspace_bytes", c_int64, [c_int64]),
+    ("snowllm_lm_head_shuffle_weight", c_int, [P, P, S]),
+    ("snowllm_lm_head_rows_for", c_int64, [c_int64]),
+    ("snowllm_lm_head_pad_bytes", c_int64, [c_int64]),
+    ("snowllm_lm_head_scratch_bytes", c_int64, [c_int64]),
+    ("snowllm_lm_head", c_int, [P, P, P, c_int64, P, P, S]),
+    ("snowllm_paged_prefill_max_tokens", c_int64, []),
+    ("snowllm_paged_decode_workspace_size", c_int64, [c_int64, c_int64]),
+    ("snowllm_kv_pool_bytes", None, [c_int64, c_int, P, P]),
+    ("snowllm_kv_scale_bytes", None, [c_int64, P, P]),
+    ("snowllm_resolve_slots", c_int, [P, c_int64, P, P, P, c_int64, S]),
+    ("snowllm_kv_blocks_for", c_int64, [c_int64]),
+    ("snowllm_linear_state_slots", None, [c_int64, c_int64, c_int64, c_int64, P]),
+    ("snowllm_prefill_q_blocks", c_int64, [P, c_int64]),
+    ("snowllm_prefill_q_block_map", None, [P, c_int64, P]),
+    ("snowllm_moe_variant_force", None, [c_int]),
+    ("snowllm_moe_variant_name", c_char_p, [c_int]),
+    ("snowllm_paged_window_blocks", c_int64, [c_int64, c_int64]),
+    ("snowllm_paged_window_view", c_int, [P, c_int64, P, c_int64, c_int64, c_int64, P,
+                                          c_int64, P, S]),
+    ("snowllm_paged_decode_max_q_tokens", c_int64, []),
+    ("snowllm_paged_decode_num_slots", c_int64, [c_int64]),
+    ("snowllm_paged_decode_plan_elems", c_int64, [c_int64, c_int64]),
+
+    ("snowllm_gather_embedding", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_rmsnorm", c_int, [P, P, P, c_int64, c_float, S]),
+    ("snowllm_rmsnorm_residual", c_int, [P, P, P, P, c_int64, c_float, S]),
+    ("snowllm_rmsnorm_shuffled", c_int, [P, P, P, c_int64, c_float, S]),
+    ("snowllm_rmsnorm_residual_shuffled", c_int, [P, P, P, P, c_int64, c_float, S]),
+    ("snowllm_qkv_proj", c_int, [P, P, P, P, c_int64, c_int, S]),
+    ("snowllm_qkv_proj_fp8", c_int, [P, P, P, P, P, c_int64, c_int, S]),
+    ("snowllm_qk_norm", c_int, [P, c_int64, P, P, P, P, c_int64, c_float, S]),
+    ("snowllm_qk_norm_rope", c_int, [P, c_int64, P, P, P, P, P, P, c_int64, c_float, S]),
+    ("snowllm_attn_out_scale_oproj", c_int, [P, P, c_int64, P, P, P, c_int64, c_int, S]),
+    ("snowllm_attn_out_scale_oproj_fp8", c_int, [P, P, c_int64, P, P, P, P, c_int64, c_int, S]),
+    ("snowllm_rope_cos_sin", c_int, [P, P, P, P, c_int64, S]),
+    ("snowllm_rope_cos_sin_mrope", c_int, [P, P, P, P, c_int64, S]),
+    ("snowllm_rope_apply_q", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_rope_apply_k", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_reshape_and_cache", c_int, [P, P, c_int64, c_int64, P, P, P, c_int64, S]),
+    ("snowllm_reshape_and_cache_int8", c_int,
+     [P, P, c_int64, c_int64, P, P, P, P, P, c_int64, S]),
+    ("snowllm_paged_attn_prefill_int8", c_int,
+     [P, P, P, P, P, P, P, P, P, c_int64, c_int64, c_int64, c_float, S, P]),
+    ("snowllm_paged_attn_decode_int8", c_int,
+     [P, P, P, P, P, P, P, P, P, c_int64, c_int64, c_int64, P, c_float, c_int64, S]),
+    ("snowllm_paged_attn_prefill", c_int,
+     [P, P, P, P, P, P, P, c_int64, c_int64, c_int64, c_float, S, P]),
+    ("snowllm_paged_attn_decode", c_int,
+     [P, P, P, P, P, P, P, c_int64, c_int64, c_int64, P, c_float, c_int64, S]),
+    ("snowllm_paged_attn_decode_plan", c_int, [P, P, c_int64, c_int64, c_int, S]),
+    ("snowllm_fused_linear_attn", c_int,
+     [P, P, P, P, P, P, P,
+      P, P, P, P,
+      P, P, P, P,
+      P, P, P, P,
+      c_int64, c_int64, c_int, S]),
+    ("snowllm_fused_moe", c_int, [P, P, P, P, P, P, c_int64, S]),
+    ("snowllm_fused_moe_fp8", c_int, [P, P, P, P, P, P, P, P, c_int64, S]),
+    ("snowllm_lm_head_gemm_decode", c_int, [P, P, P, c_int64, S]),
+    ("snowllm_lm_head_gemm_prefill", c_int, [P, P, P, P, c_int64, S]),
+    ("snowllm_sampling_softmax", c_int, [P, P, c_int64, c_int64, P, S]),
+    ("snowllm_sampling_topk_topp", c_int, [P, P, P, c_int64, c_int64, c_int64, P, P, S]),
+    ("snowllm_sampling_multinomial", c_int, [P, P, P, P, c_int64, c_int64, S]),
+    ("snowllm_sampling_argmax", c_int, [P, P, c_int64, c_int64, S]),
+
+    ("snowllm_gemm_bf16_shuffle_b", c_int, [P, P, c_int64, c_int64, S]),
+    ("snowllm_vision_layernorm", c_int, [P, P, P, P, c_int64, c_int64, c_float, S]),
+    ("snowllm_vision_gemm_scratch_bytes", c_int64, [c_int64, c_int64]),
+    ("snowllm_vision_gemm_bias_act", c_int,
+     [P, P, P, P, P, P, c_int64, c_int64, c_int64, c_int, S]),
+    ("snowllm_vision_rope_qk", c_int, [P, P, P, c_int64, c_int64, c_int64, c_int64, S]),
+    ("snowllm_vision_attn", c_int, [P, P, P, P, c_int64, c_int64, c_int64, c_int64, S]),
+    ("snowllm_vision_supported_head_dims", c_int64, [P, c_int64]),
+    ("snowllm_vision_supported_ln_hidden", c_int64, [P, c_int64]),
+
+    ("snowllm_file_reader_create", P, [c_int, c_int64, c_int]),
+    ("snowllm_file_reader_destroy", None, [P]),
+    ("snowllm_file_reader_read", c_int, [P, c_char_p, c_int64, c_int64, P, S]),
+]
+_missing = []
+for _name, _res, _args in _SIGS:
+    try:
+        _fn = getattr(lib, _name)
+    except AttributeError:
+        _missing.append(_name)
+        continue
+    _fn.restype = _res
+    _fn.argtypes = _args
+if _missing:
+    raise SnowLLMError(f"{_lib_path} is missing {len(_missing)} entry points this binding needs:\n  "
+                       + "\n  ".join(_missing))
+
+
+def check(status: int, what: str) -> None:
+    if status != 0:
+        msg = lib.snowllm_error_string(status).decode()
+        raise SnowLLMError(f"{what}: hip status {status}: {msg}")
+
+
+def synchronize(stream: int = 0) -> None:
+    check(lib.snowllm_synchronize(stream), "synchronize")
+
+
+def assert_single_hip_runtime() -> None:
+    paths = set()
+    with open("/proc/self/maps") as f:
+        for line in f:
+            i = line.find("libamdhip64.so")
+            if i != -1:
+                paths.add(os.path.realpath(line.rsplit(" ", 1)[-1].strip()))
+    if len(paths) > 1:
+        raise SnowLLMError(
+            "two HIP runtimes are loaded in this process -- torch's device pointers are not valid "
+            "in snowllm's kernels:\n  " + "\n  ".join(sorted(paths))
+        )

snowllm/_version.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the SnowLLM project
+
+__version__ = "0.1.0"

snowllm/async_engine.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the SnowLLM project
+
+import asyncio
+import pathlib
+import time
+import traceback
+from collections.abc import AsyncIterator
+from concurrent.futures import ThreadPoolExecutor
+
+from .engine import Engine, Request, SamplingParams
+
+
+class AsyncEngine:
+    def __init__(self, engine: Engine, profile_dir: str | None = None,
+                 stats_interval: float = 0.0):
+        self.engine = engine
+        self.stats_interval = stats_interval
+        self._last_stats = (0.0, None)
+        self.profile_dir = pathlib.Path(profile_dir).expanduser() if profile_dir else None
+        if self.profile_dir:
+            self.profile_dir.mkdir(parents=True, exist_ok=True)
+        self._profiler = None
+        self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="snowllm-engine")
+        self._live: list[Request] = []
+        self._queues: dict[int, asyncio.Queue] = {}
+        self._sent: dict[int, int] = {}
+        self._wake = asyncio.Event()
+        self._task: asyncio.Task | None = None
+
+    def start(self) -> None:
+        self._task = asyncio.create_task(self._run())
+
+    async def start_profile(self) -> None:
+        if self.profile_dir is None:
+            raise RuntimeError("profiling is off; start the server with --profile-dir")
+        if self._profiler is not None:
+            raise RuntimeError("a profile is already running")
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(self._worker, self._profiler_start)
+
+    async def stop_profile(self) -> str:
+        if self._profiler is None:
+            raise RuntimeError("no profile is running")
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(self._worker, self._profiler_stop)
+
+    def _profiler_start(self) -> None:
+        import torch
+
+        self._profiler = torch.profiler.profile(
+            activities=[torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.CUDA],
+            with_stack=True,
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(str(self.profile_dir)),
+        )
+        self._profiler.start()
+
+    def _profiler_stop(self) -> str:
+        p, self._profiler = self._profiler, None
+        p.stop()
+        return str(self.profile_dir)
+
+    async def close(self) -> None:
+        if self._task:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+        self._worker.shutdown(wait=True)
+
+    def stats(self):
+        return self.engine.stats()
+
+    def submit(self, prompt: list[int], params: SamplingParams, rope_factor: float = 1.0,
+               **mm) -> Request:
+        r = self.engine.add(prompt, params, rope_factor=rope_factor, **mm)
+        self._live.append(r)
+        self._queues[id(r)] = asyncio.Queue()
+        self._sent[id(r)] = 0
+        self._wake.set()
+        return r
+
+    async def stream(self, r: Request) -> AsyncIterator[int]:
+        q = self._queues[id(r)]
+        try:
+            while True:
+                tok = await q.get()
+                if tok is None:
+                    return
+                yield tok
+        finally:
+            if not r.done:
+                self.abort(r)
+            self._queues.pop(id(r), None)
+
+    def abort(self, r: Request, reason: str = "abort") -> None:
+        r.done = True
+        r.finish_reason = r.finish_reason or reason
+        if r in self.engine.waiting:
+            self.engine.waiting.remove(r)
+            self._queues[id(r)].put_nowait(None)
+            self._close(r)
+
+    async def _run(self) -> None:
+        loop = asyncio.get_running_loop()
+        while True:
+            if not self.engine.waiting and not self.engine.running:
+                await self._wake.wait()
+                self._wake.clear()
+                continue
+            try:
+                await loop.run_in_executor(self._worker, self.engine.step)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                traceback.print_exc()
+                self._fail_all()
+                continue
+            self._dispatch()
+            self._log_stats()
+
+    def _log_stats(self) -> None:
+        if not self.stats_interval or self.engine.acct is None:
+            return
+        now = time.monotonic()
+        was_at, was = self._last_stats
+        if was is not None and now - was_at < self.stats_interval:
+            return
+        cur = self.engine.acct.snapshot()
+        self._last_stats = (now, cur)
+        if was is None:
+            return
+        dp_t = cur["prefill_tokens"] - was["prefill_tokens"]
+        dp_s = cur["prefill_seconds"] - was["prefill_seconds"]
+        dd_t = cur["decode_tokens"] - was["decode_tokens"]
+        dd_s = cur["decode_seconds"] - was["decode_seconds"]
+        dd_r = cur["decode_rows"] - was["decode_rows"]
+        st = self.engine.stats()
+        kv = 1.0 - st.free_kv_blocks / max(self.engine.blocks.total, 1)
+        print(f"[snowllm {time.strftime('%H:%M:%S')}] "
+              f"prefill {dp_t / dp_s if dp_s else 0.0:7.1f} tok/s | "
+              f"decode {dd_t / dd_s if dd_s else 0.0:7.1f} tok/s | "
+              f"AL {dd_t / dd_r if dd_r else 0.0:.3f} | "
+              f"running {st.running:3d} waiting {st.waiting:4d} | KV {kv * 100:4.1f}%", flush=True)
+
+    def _fail_all(self) -> None:
+        for r in self._live:
+            r.done, r.finish_reason = True, "error"
+        self._dispatch()
+
+    def _dispatch(self) -> None:
+        for r in list(self._live):
+            q = self._queues.get(id(r))
+            if q is None:
+                self._close(r)
+                continue
+            n = self._sent[id(r)]
+            for tok in r.out[n:]:
+                q.put_nowait(tok)
+            self._sent[id(r)] = len(r.out)
+            if r.done:
+                q.put_nowait(None)
+                self._close(r)
+
+    def _close(self, r: Request) -> None:
+        if r in self._live:
+            self._live.remove(r)
+        self._sent.pop(id(r), None)

snowllm/block_manager.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the SnowLLM project
+
+import torch
+
+from . import ops
+from .request import Request
+
+
+def pad_table(rows: list[list[int]]) -> torch.Tensor:
+    w = max((len(x) for x in rows), default=1) or 1
+    t = torch.zeros(len(rows), w, dtype=torch.int32)
+    for i, x in enumerate(rows):
+        t[i, : len(x)] = torch.tensor(x, dtype=torch.int32)
+    return t.to("cuda")
+
+
+def slot_mapping(bt: torch.Tensor, rows: list[tuple[int, int]]) -> torch.Tensor:
+    seq = torch.tensor([b for b, _ in rows], dtype=torch.int32, device="cuda")
+    pos = torch.tensor([p for _, p in rows], dtype=torch.int32, device="cuda")
+    return ops.resolve_slots(bt, seq, pos)
+
+
+def block_tables(batch: list[Request]) -> torch.Tensor:
+    return pad_table([r.blocks for r in batch])
+
+
+class BlockAllocator:
+    def __init__(self, num_blocks: int):
+        self.free = list(range(num_blocks))
+        self.total = num_blocks
+
+    def alloc(self, n: int) -> list[int] | None:
+        if n > len(self.free):
+            return None
+        return [self.free.pop() for _ in range(n)]
+
+    def release(self, blocks: list[int]) -> None:
+        self.free.extend(blocks)
+
+    def grow(self, r: Request, n: int = 1) -> bool:
+        need = ops.kv_blocks_for(r.num_cached + n)
+        while need > len(r.blocks):
+            more = self.alloc(1)
+            if more is None:
+                return False
+            r.blocks += more
+        return True
+
+
+class StateSlots:
+    def __init__(self, max_num_seqs: int, T: int, dummy: int):
+        self.T = T
+        self.free = list(range(max_num_seqs))
+        self._dummy = dummy
+
+    def take(self, r: Request) -> None:
+        r.slot = self.free.pop()
+        r.state_head = ops.linear_state_slots(r.slot, -1, 1, self.T)[0]
+
+    def give_back(self, r: Request) -> None:
+        self.free.append(r.slot)
+
+    def name(self, r: Request, rows: int) -> list[int]:
+        return ops.linear_state_slots(r.slot, r.state_head, rows, self.T)
+
+    def dummy_indices(self, n: int) -> torch.Tensor:
+        return torch.full((n,), self._dummy, dtype=torch.int32, device="cuda")