diff --git a/README.md b/README.md index f600408..730d5bc 100644 --- a/README.md +++ b/README.md @@ -84,11 +84,11 @@ are co-located with each checkpoint and load automatically, so - **Backend isolation (Python framework)**: within the Python framework, all MLX code lives under `python/src/edge0/backends/mlx/`; the core logic (model specs, prerouter, streaming expert pool, server) depends - only on the backend facade (`backends/base.py`), so a new backend - implements the same facade (`backends/cuda/` is a reserved slot) with - zero changes to core code. The iOS / macOS / Android engines ship - platform-native stacks today — bringing every platform under one - access layer is exactly what the unified inference framework + only on the backend facade (`backends/base.py`), implemented by the + MLX backend and the optional Torch backend (`backends/cuda/`), so a new + backend needs zero changes to core code. The iOS / macOS / Android + engines ship platform-native stacks today — bringing every platform + under one access layer is exactly what the unified inference framework (see [Roadmap](#roadmap)) will deliver; - **Adapters as safetensors**: LoRA and prerouter weights are `.safetensors` files with provenance metadata (source, version, owner @@ -144,13 +144,17 @@ python examples/bench.py edge0-8b # via $EDGE0_8B_MODEL #### Requirements - **OS / hardware**: the MLX backend runs on macOS with Apple Silicon - (M1/M2/M3/M4). The CUDA backend is on the roadmap — no other - platforms are supported by the Python framework yet. + (M1/M2/M3/M4). The Torch backend supports CPU and Apple MPS. NVIDIA + CUDA results documented here were measured on a DGX Spark (GB10). + Choose a PyTorch wheel compatible with your OS, GPU, and CUDA + driver (see [PyTorch's install selector](https://pytorch.org/get-started/locally/)). - **Python**: 3.10+ (3.12 recommended). -- **MLX**: `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see +- **MLX backend**: `mlx==0.30.6` / `mlx-metal==0.30.6` with `mlx-lm==0.31.0` (see `python/pyproject.toml`). Garbled, mixed-language output on Apple A18 / A18 Pro means an older `mlx`: `pip install 'mlx==0.30.6' 'mlx-metal==0.30.6'` ([#8](https://github.com/Edge0-AI/Edge0/issues/8)). +- **CUDA backend**: install the PyTorch wheel for your system, then install + Edge0 with its optional `cuda` extra (which declares `torch>=2.14`). - **Memory**: ~2.9 GB peak active memory for `edge0-35b`, ~1.0 GB for `edge0-8b` (short contexts; see [Benchmark](#benchmark)). Add headroom for the OS, tokenizer, and long-context KV growth. @@ -166,6 +170,18 @@ cd python python3.12 -m venv .venv && .venv/bin/pip install -e '.[dev,fetch]' ``` +#### NVIDIA / CUDA backend (optional) + +Install the PyTorch wheel that matches your operating system, NVIDIA GPU, +and CUDA driver using [PyTorch's install selector](https://pytorch.org/get-started/locally/). +The CUDA backend was validated with PyTorch 2.14 and CUDA 13.0 on an NVIDIA +GB10. Then install Edge0 with the optional extra and select the backend: + +```bash +.venv/bin/pip install -e '.[cuda,dev,fetch]' +EDGE0_BACKEND=cuda .venv/bin/edge0 serve edge0-35b +``` + #### 2) Download a model The two tiers are published on Hugging Face and ModelScope — each repo @@ -279,6 +295,7 @@ runs this exact path). - [Architecture](docs/architecture.md) - [Attention](docs/attention.md) / [MoE](docs/moe.md) / [SSD streaming](docs/streaming.md) / [prerouter](docs/prerouter.md) - [Adding a model](docs/adding-a-model.md) +- [NVIDIA / CUDA support: investigation status](docs/nvidia.md) - [edge0-35b](docs/models/edge0-35b.md) / [edge0-8b](docs/models/edge0-8b.md) - Technical report: [The Other Half of the Memory Wall](https://arxiv.org/abs/2609.18063) ([PDF](paper/main.pdf)) diff --git a/docs/nvidia.md b/docs/nvidia.md new file mode 100644 index 0000000..59f7ad7 --- /dev/null +++ b/docs/nvidia.md @@ -0,0 +1,190 @@ +# NVIDIA / CUDA support — status + +`edge0` does **not** run on NVIDIA through MLX, at any version tested: +the first forward pass fails, differently at each version (the table +below). It does run on NVIDIA through the torch backend in +`backends/cuda/` (`EDGE0_BACKEND=cuda`): both edge0-8b and edge0-35b +generate text on a GB10, layer for layer within 5.3e-7 (8b) and 3.7e-6 +(35b, real weights) of MLX. The rest of this file is +that investigation, kept because it is what the next person attempting +this would otherwise repeat from scratch. + +## Environment + +- Hardware: DGX Spark, GB10 (`sm_121`, Grace-Blackwell, 128GB unified memory) +- Toolkit: CUDA 13.0 — use `mlx[cuda13]`, not `mlx[cuda12]`; the latter + ships NVRTC 12.9, which does not compile against the CUDA 13 headers + (`cuda_fp6.hpp` / `cuda_fp4.hpp`) on this platform. +- Tier tested: edge0-8b (Ling / Bailing hybrid) + +## The failure chain (three MLX versions, three distinct errors) + +Everything up to the model finishing construction works cleanly at +every version tested: install, checkpoint download, LoRA/prerouter +attach (`lora applied=153 not_found=0`, `prerouter installed: 16 heads`, +built in 0.7s). The failure is always at the first forward pass, and it +moves as the MLX version moves — meaning this is not one bug, it's the +edge of MLX's own CUDA-quantized-op support arriving in stages, with +`edge0`'s code (written and pinned against `0.30.4`) landing in a +different gap each time: + +| `mlx` version | Result | +|---|---| +| `0.30.4` (this repo's current pin) | `RuntimeError: QMM NYI` — quantized matmul has no CUDA implementation at all | +| `0.31.1` | `GatherQMM has no CUDA implementation` — the gather-variant specifically still missing | +| `0.32.0` / `0.32.2` | Both ops now present — clears `GatherQMM`, advances from `ling.py:181` to `ling.py:190`, then `IndexError: SmallVector out of range` inside `core.eval(logits)` | + +No CPU fallback was viable either: `mlx-cpu==0.30.4` fails to JIT on +g++13/aarch64, and `mlx-cpu==0.32.2` raises `There is no Stream(cpu, 3)`. + +The `0.32.x` failure is the most interesting one and the most worth a +second look by someone who wants to pursue this further: both required +CUDA kernels are confirmed present, the crash is downstream inside +`core.eval`, and it reproduced identically on two separate versions — +consistent with a real, narrow incompatibility between `edge0`'s +`ling.py` code path (written for `0.30.4`'s API) and something that +changed by `0.32.x`, not with a fundamentally unsupported operation. + +## A reliability caveat, independent of all of the above + +`mx.default_device()` reporting `Device(gpu, 0)` is **not** evidence the +GPU is usable — MLX is lazy and this call never touches the driver. It +reported `gpu` in every run above, including the ones where the GPU was +later confirmed dead. The real signal is whether `cuInit()` succeeds. +On our machine `cuInit()` started returning 999 a few minutes after +each boot while `nvidia-smi` still looked healthy. We first read that +as the driver dying after an aborted CUDA process. It was not: a +host-level cgroup device policy on that machine (unrelated to `edge0` +or MLX) denies `/dev/nvidia-uvm` and `/dev/nvidia-caps/*`, and the CUDA +runtime needs both. `nvidia-smi` goes through NVML on `/dev/nvidiactl`, +so it never notices. If `cuInit()` gives 999, first try opening +`/dev/nvidia-uvm` from the same shell. `EPERM` there means a device +policy, not a broken GPU. + +## Bottom line + +MLX's own CUDA backend does not close the gap at any version currently +available. The torch backend does: `EDGE0_BACKEND=cuda` runs both shipped +tiers on a GB10 and on CPUs, checked against MLX throughout (next +section). On the real weights, edge0-35b on the GB10 generates the same +32 greedy tokens as MLX run on its CPU device. + +## Torch backend: what exists and how it is checked + +MLX is the ground truth throughout: the checks run on Apple Silicon, where +both backends are available, and compare the torch side against MLX +(`tests/test_cuda_backend.py`, `tests/test_backend_parity.py`; the latter +runs each case under both backends in subprocesses). + +| Piece | Checked against | +|---|---| +| `quant.gather_qmm` (2/4/8-bit affine, all broadcast shapes the streaming layer uses) | `mx.gather_qmm` | +| `core` ops at their real call sites (routing, sampling), `nn.RMSNorm` / `gelu` | the same code on MLX: identical expert choices, identical sampler masks | +| `StreamingSwitchGLU`, every path (exact, whole-layer, hot, staged), on layer 1 of the real edge0-8b checkpoint (`EDGE0_8B_MODEL`) | MLX on the same inputs: 1.2-1.5% of output scale, about two bf16 ulps | +| `io.load_model` + `install_streaming_experts` on a small checkpoint in the exact published edge0-35b format | the source model on the same weights: 4e-7 relative, same argmax | +| `backends/cuda/_impl/bailing_hybrid.py` (torch port of the edge0-8b backbone) on the real checkpoint, every layer, chunked prefill + decode | the vendored MLX model on the MLX CPU device, float32: <= 1.5e-6 per layer, <= 1.7e-6 on the logits | +| the whole edge0-8b engine (`engine/ling.py` unchanged: LoRA, prerouter-staged decode, streaming, sampling) | the same engine on MLX: identical greedy tokens (`pytest -m slow`) | +| the same engine on Linux aarch64 (DGX Spark, torch on the CPU), 32 greedy tokens of a chat prompt | MLX on Apple Silicon, same checkpoint (same file hashes): identical 32 tokens; 1.2 GB peak anonymous memory | +| everything above with torch on an accelerator: `EDGE0_TORCH_DEVICE=mps` (Apple GPU), whole suite including the slow engine test, plus the 32-token run | the same references: all pass, identical 32 tokens. On a device a tensor left on the host fails loudly, as it would on CUDA; this is what found `load_model` leaving init-time buffers (the rotary `inv_freq`) on the host | +| **the edge0-8b backbone on a real GPU** — a GB10 (`sm_121`, DGX Spark), torch 2.14+cu130, float32, every layer fed MLX's input for that layer | MLX on the Apple CPU device: **5.3e-7 max per layer** (median 2.6e-7), 4.3e-7 on the logits, same argmax. TF32 off, `float32_matmul_precision=highest` | +| the whole edge0-8b engine on that GPU, 32 greedy tokens | MLX on Apple Silicon: 31 of 32 tokens identical, diverging at step 31. Not a GPU artifact: torch on the CPU differs from MLX by the same 4.5% median per-step logit distance (the staged prerouter path), and that step's top-2 margin is smaller than that noise | +| `backends/cuda/_impl/qwen3_5_moe.py` (torch port of the edge0-35b backbone), every layer, chunked prefill + decode, on a small model MLX wrote in the published format (bf16, 4-bit, 8-bit router and shared gate) | the vendored MLX model on the MLX CPU device, float32: <= 2.6e-7 per layer, <= 4.1e-7 on the logits | +| the whole edge0-35b engine (`engine/qwen.py` unchanged: streaming, staged decode, the class-level prerouter patch) on that small checkpoint | the same engine on MLX: identical greedy tokens, per-step logits within bf16 noise and tracking MLX *with* the prerouter (the prerouter moves them 5-11%) | +| **the edge0-35b backbone on the real 23 GB checkpoint**, every layer, chunked prefill + decode (`test_qwen35_port_matches_mlx_on_real_weights`, needs `EDGE0_35B_MODEL`) | the vendored MLX model on the MLX CPU device, float32: <= 2.4e-6 per GatedDeltaNet layer, <= 3.7e-6 per gated full-attention layer, <= 2.4e-6 on the logits, same argmax; 40 layers, no missing or unexpected tensor | +| the whole edge0-35b engine on the real weights (LoRA 310 targets, prerouter 33 heads, streamed experts), 32 greedy tokens | MLX on the MLX CPU device: **identical 32 tokens**. Against MLX on Metal, 29 of 32 -- and MLX-Metal disagrees with MLX-CPU at exactly those three positions, so the flip is its GPU precision | +| **the edge0-35b engine on a real GPU** — the same GB10, torch 2.14+cu130, as a Slurm job | **identical 32 tokens** to both MLX-CPU and torch on the Mac's CPU (18.4 s) | +| **the edge0-35b backbone on that GPU**, real weights, float32, every layer fed MLX's own input for that layer | MLX on the Apple CPU device: **1.6e-6 max per layer** (median 2.7e-8) — 5.4e-7 across the 30 GatedDeltaNet layers, 1.6e-6 across the 10 gated full-attention ones — 5.1e-7 on the logits, same argmax; 0 missing / 0 unexpected tensors, 40 streamed expert layers | + +Why the MLX *CPU* device: on some Apple GPUs MLX runs float32 matmul and +SDPA at reduced precision (an M5 Max measured 7.5e-4 from float64; MLX on +the CPU and torch both 2e-7). Against MLX on the GPU the port looks up to +1000x worse in the attention layers, all of it on the MLX side. + +`load_model` handles what the published checkpoints actually contain: MLX +quantization of nearly every linear and embedding (kept 4-bit resident via +`QuantizedLinear` / `QuantizedEmbedding`), 8-bit routers, and, for +edge0-35b, MLX's sanitize (conv1d stored `[C, k, 1]`, five RMSNorm kinds +stored as `w + 1`), which it undoes. + +## Speed, and where it goes (GB10, edge0-8b decode) + +The backend was written as a correctness reference, and it shows: the +default profile decodes at **316 ms a token (3.2 tok/s)** on a GB10. A +profile of three decode steps says the GPU is barely working -- 63 ms of +CUDA against 643 ms of CPU -- so this is launch and copy overhead, not +arithmetic: + +* **1681 host-to-device copies per token.** The streaming layer rebuilds + expert bundles from the mmap every step, which is the whole point on a + Mac with 36 GB and pointless on a 128 GB box. The shared LRU already + fixes it: `cache_slots=3072` (from the tier's 64) keeps every expert + resident after warm-up -> **206 ms a token (4.9 tok/s)**, 1.9 GB. A + config change, no code. +* **235 quantized-linear calls per token**, dequantizing the same + attention and lm_head weights every time (103 ms a token). + `EDGE0_TORCH_WEIGHT_CACHE=1` keeps the dequantized weight instead -> + **124 ms a token (8.1 tok/s)**, 6.0 GB resident. + +Together that is **2.5x** over the default profile, with the same greedy +tokens. The two paths are bit-for-bit identical on the CPU +(`test_quantized_linear_weight_cache_is_exact`); on CUDA, keeping the +whole weight changes cuBLAS's split against the chunked path, so they +differ within the engine's own noise -- teacher-forced they pick the same +token at all 32 steps and sit the same distance from MLX (4.8% vs 5.3% +median per-step), while free-running they can diverge a step earlier. + +### The syncs cost more than the arithmetic (Metal, edge0-8b decode) + +The same CPU-bound shape shows up on this Mac's GPU (`EDGE0_TORCH_DEVICE=mps`), +which makes the overhead measurable without a GB10. There the default +profile decoded at **688 ms a token**, and a profile blamed three torch +calls inside `gather_qmm`: `nonzero` (2.4 s), `tolist` (1.9 s) and +`unique` (1.4 s) of an 8.2 s run -- against **0.07 s** for the matmuls +they were arranging. None of them is arithmetic: each is a device->host +sync, and the loop paid one per expert to ask *which* experts a call +needs and *where* each one's rows are. + +Decode never needs to ask. One token routes to K experts, so the whole +call is a handful of rows: gather those rows' experts as one batch and +run a single `bmm`. No sync, and a repeated expert only costs a duplicate +dequantize. The per-expert loop still runs above that size (prefill), +where a dequantized copy per row is what blew memory up in the first +place -- now behind one host transfer of the index list instead of a +`nonzero` per expert. The cap is both a row count and a byte budget +(`EDGE0_TORCH_GATHER_BATCH`, `EDGE0_TORCH_GATHER_BYTES`), because one row +of a 35b expert is far bigger than an 8b one. + +| edge0-8b decode, Metal (M5 Max), greedy | ms/token | +| --- | --- | +| tier default (`cache_slots=64`) | 688 | +| + batched decode gather (this change) | 401 | +| + `cache_slots=3072` | 384 | +| + `EDGE0_TORCH_WEIGHT_CACHE=1` | 291 | + +**2.4x**, with the same 12 greedy tokens at every step. The GB10 numbers +above predate this and still need re-measuring: the Spark's GPU has been +fenced off since the 13-Sep reboot (`/dev/nvidia-uvm` EPERM, inside the +Slurm job as well as on the host). + +What is still on the table: a fused dequantize+matmul kernel (the +elementwise shift/mask/mul/add chain is most of the remaining CUDA time), +captured graphs or `torch.compile` for the per-step launch storm, and +bundles built directly on the device rather than copied per step. + +## What is left + +* **Performance.** Decode no longer pays a host sync per expert (see + above), but `gather_qmm` and the quantized linears still dequantize on + every call and `core.compile` is eager: this is a correctness + reference, not a fast path. + +## A stale assumption this also corrects + +`tests/test_repo_hygiene.py`'s module docstring says these hygiene tests +"can run on any platform — including the Linux CI job where MLX has no +wheels." That was true when written; it no longer is — MLX ships both a +CPU-only Linux wheel and CUDA-backed Linux wheels +(`manylinux_2_35_x86_64`/`aarch64`) as of at least `mlx==0.30.4`, the +exact version this repo already pins. Worth knowing regardless of +whether GPU runners are in reach for CI: the assumption that MLX-dependent +tests categorically cannot run in this repo's own CI is no longer correct. diff --git a/python/pyproject.toml b/python/pyproject.toml index 7d93e06..6850fe5 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -45,6 +45,9 @@ dev = [ fetch = [ "huggingface_hub>=0.23", ] +cuda = [ + "torch>=2.14", +] [project.scripts] edge0 = "edge0.cli:main" diff --git a/python/src/edge0/adapters/lora.py b/python/src/edge0/adapters/lora.py index 1103dd8..2866f70 100644 --- a/python/src/edge0/adapters/lora.py +++ b/python/src/edge0/adapters/lora.py @@ -23,11 +23,9 @@ import os -from edge0.backends import core +from edge0.backends import core, io from edge0.backends import nn -from edge0.backends.mlx.io import load_safetensors - class LoraLinear(nn.Module): """Wraps a Linear/QuantizedLinear and adds the low-rank delta.""" @@ -43,10 +41,10 @@ def __init__(self, base: nn.Module, lora_a: core.array, lora_b: core.array, def __call__(self, x: core.array) -> core.array: y = self.base(x) ad = self.lora_A.dtype - xd = x if x.dtype == ad else x.astype(ad) + xd = x if x.dtype == ad else core.astype(x, ad) # delta = (B @ A) applied on x: ((x @ A.T) @ B.T), shape [..., out] d = (xd @ self.lora_A.T) @ self.lora_B.T - return y + (self.lora_scale * d).astype(y.dtype) + return y + core.astype(self.lora_scale * d, y.dtype) def _resolve(model: nn.Module, key: str): @@ -83,7 +81,7 @@ def install_lora(model: nn.Module, adapters_path: str, "adapters; download them for this tier (README -> 'Getting " "the models & adapters') and place them in the model " "directory, or disable LoRA with lora="" / --no-lora.") - lora = load_safetensors(adapters_path) + lora = io.load_safetensors(adapters_path) if not lora: raise ValueError("adapters has no lora tensors") @@ -117,8 +115,8 @@ def install_lora(model: nn.Module, adapters_path: str, raise TypeError( f"lora target {key} is {type(mod).__name__}, not Linear") continue - a = d["A"].astype(dtype) - b = d["B"].astype(dtype) + a = core.astype(d["A"], dtype) + b = core.astype(d["B"], dtype) # re-bind the module attribute (e.g. attention.q_proj = LoraLinear) parts = key.split(".") owner = _resolve(model, ".".join(parts[:-1])) diff --git a/python/src/edge0/backends/__init__.py b/python/src/edge0/backends/__init__.py index 71e662c..4f404fc 100644 --- a/python/src/edge0/backends/__init__.py +++ b/python/src/edge0/backends/__init__.py @@ -26,7 +26,7 @@ only through the namespaces re-exported here. * ``core`` must provide (for the MLX backend these are mlx.core names): array, zeros, eye, arange, full, expand_dims, squeeze, reshape, - transpose, concatenate, stack, split, matmul, softmax, silu (via nn), + transpose, concatenate, stack, split, matmul, einsum, softmax, silu (via nn), sigmoid, erf, where, sum, cumsum, sort, topk, argpartition, take, take_along_axis, put_along_axis, astype, item, tolist, eval, compile, random.seed, random.categorical, and the dtypes float16 / float32 / @@ -49,12 +49,19 @@ quant, ) elif _BACKEND == "cuda": - raise ImportError( - "the CUDA backend is not implemented yet; set EDGE0_BACKEND=mlx " - "(or unset it)") + # Torch reference implementation. EDGE0_BACKEND=mlx on top of + # mlx[cuda*] is not a shortcut: no MLX release runs the shipped + # tiers on NVIDIA hardware (see docs/nvidia.md). + from edge0.backends.cuda.backend import ( # noqa: F401 + BackendImpl, + core, + io, + nn, + quant, + ) else: raise ImportError( - f"unknown backend {_BACKEND!r}; available: mlx") + f"unknown backend {_BACKEND!r}; available: mlx, cuda") backend = BackendImpl() diff --git a/python/src/edge0/backends/cuda/__init__.py b/python/src/edge0/backends/cuda/__init__.py new file mode 100644 index 0000000..215aaad --- /dev/null +++ b/python/src/edge0/backends/cuda/__init__.py @@ -0,0 +1,9 @@ +"""CUDA backend (torch reference implementation). + +Everything that touches ``torch`` for array/nn/quant ops lives under +this package, mirroring ``edge0/backends/mlx/``'s isolation rule. +""" + +from edge0.backends.cuda.backend import BackendImpl, core, io, nn, quant + +__all__ = ["BackendImpl", "core", "io", "nn", "quant"] diff --git a/python/src/edge0/backends/cuda/_impl/__init__.py b/python/src/edge0/backends/cuda/_impl/__init__.py new file mode 100644 index 0000000..6991ef5 --- /dev/null +++ b/python/src/edge0/backends/cuda/_impl/__init__.py @@ -0,0 +1 @@ +"""Torch ports of the vendored MLX model families (``backends/mlx/_impl``).""" diff --git a/python/src/edge0/backends/cuda/_impl/bailing_hybrid.py b/python/src/edge0/backends/cuda/_impl/bailing_hybrid.py new file mode 100644 index 0000000..c820cde --- /dev/null +++ b/python/src/edge0/backends/cuda/_impl/bailing_hybrid.py @@ -0,0 +1,653 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Torch port of the vendored Ling 3.0 backbone (``model_type: bailing_hybrid``). + +A line-by-line port of ``backends/mlx/_impl/bailing_hybrid.py`` -- see that +file for the architecture and its references -- with the same module tree, +parameter names and call signatures, so ``engine/ling.py``, the engine hooks +and the prerouter stager drive it unchanged. It replaces the checkpoint's +own ``modeling_bailing_moe_v3.py`` (remote code, ``fla``/Triton, per-expert +``ModuleList``) on the torch backend. + +Parameters are named as stored in the checkpoint, which differs from the +MLX module in one place: the MLX ``sanitize`` moves ``*_conv1d.weight`` +([C, 1, k], torch's depthwise layout) to ``*_conv1d.conv.weight`` in MLX's +[C, k, 1] layout; here the torch layout is used as stored. + +Checked against the MLX module on the real edge0-8b checkpoint, layer by +layer, in ``tests/test_backend_parity.py``. +""" + +from __future__ import annotations + +import inspect +import os +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + + +def _rope_interleave_torch(x: torch.Tensor, positions: torch.Tensor, + theta: float) -> torch.Tensor: + """``apply_rotary_pos_emb_interleave`` of ``modeling_bailing_moe_v3`` + (same formula as the MLX port): rotate consecutive pairs of the last + axis by ``positions * 1 / theta ** (arange(0, d, 2) / d)``.""" + B, H, L, D = x.shape + freqs = 1.0 / (theta ** (torch.arange(0, D, 2, dtype=torch.float32, + device=x.device) / D)) + angles = torch.outer(positions.to(torch.float32), freqs) + emb = torch.cat([angles, angles], dim=-1) + cos = torch.cos(emb)[None, None] + sin = torch.sin(emb)[None, None] + xi = x.reshape(B, H, L, D // 2, 2).transpose(-1, -2).reshape(B, H, L, D) + h = D // 2 + rh = torch.cat([-xi[..., h:], xi[..., :h]], dim=-1) + return xi * cos + rh * sin + + +def _rms_norm(x, weight, eps): + """``mx.fast.rms_norm``: normalize in float32, output in x's dtype.""" + xf = x.to(torch.float32) + y = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) + return (y * weight.to(torch.float32)).to(x.dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dims: int, eps: float = 1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dims)) + + def forward(self, x): + return _rms_norm(x, self.weight, self.eps) + + +def _linear(i, o, bias=False): + # the facade's Linear: accepts plain-tensor weight assignment, which + # prerouter/install.py relies on + from edge0.backends.cuda.nn import Linear + return Linear(i, o, bias=bias) + + +@dataclass +class ModelArgs: + model_type: str = "bailing_hybrid" + hidden_size: int = 1536 + num_hidden_layers: int = 24 + intermediate_size: int = 4608 + num_attention_heads: int = 16 + num_key_value_heads: int = 16 + head_dim: int = 128 + rms_norm_eps: float = 1e-6 + vocab_size: int = 157184 + max_position_embeddings: int = 131072 + tie_word_embeddings: bool = False + layer_group_size: int = 4 + first_k_dense_replace: int = 1 + short_conv_kernel_size: int = 4 + no_kda_lora: bool = True + kda_safe_gate: bool = True + kda_lower_bound: float = -5.0 + q_lora_rank: int | None = 256 + kv_lora_rank: int = 512 + qk_nope_head_dim: int = 128 + qk_rope_head_dim: int = 64 + v_head_dim: int = 128 + rope_theta: float = 6000000.0 + rope_interleave: bool = True + rope_scaling: dict | None = None + use_qkv_bias: bool = False + gated_attention_proj_granularity_type: str | None = "head_wise" + num_experts: int = 128 + num_experts_per_tok: int = 8 + num_shared_experts: int = 1 + moe_intermediate_size: int = 512 + moe_shared_expert_intermediate_size: int = 512 + n_group: int = 8 + topk_group: int = 4 + norm_topk_prob: bool = True + routed_scaling_factor: float = 2.5 + moe_router_enable_expert_bias: bool = True + prerouter_enabled: bool = False + prerouter_start_layer: int = 7 + prerouter_hidden: int = 512 + + @classmethod + def from_dict(cls, params: dict) -> "ModelArgs": + """Keep only this class's fields, like mlx-lm's BaseModelArgs.""" + names = inspect.signature(cls).parameters + return cls(**{k: v for k, v in params.items() if k in names}) + + @property + def qk_head_dim(self) -> int: + return self.qk_nope_head_dim + self.qk_rope_head_dim + + def is_mla_layer(self, idx: int) -> bool: + g = self.layer_group_size + full = self.num_hidden_layers // g * g + return (idx + 1) % g == 0 or idx >= full + + +# ---- caches (mlx-lm KVCache / ArraysCache, the parts the model uses) -------- + +class KVCache: + def __init__(self): + self.keys = None + self.values = None + self.offset = 0 + + def update_and_fetch(self, keys, values): + if self.keys is None: + self.keys, self.values = keys, values + else: + self.keys = torch.cat([self.keys, keys], dim=2) + self.values = torch.cat([self.values, values], dim=2) + self.offset += keys.shape[2] + return self.keys, self.values + + def make_mask(self, N: int): + return None if N == 1 else "causal" + + def size(self): + return self.offset + + +class ArraysCache: + def __init__(self, size: int): + self.cache = [None] * size + + def __setitem__(self, idx, value): + self.cache[idx] = value + + def __getitem__(self, idx): + return self.cache[idx] + + +def create_attention_mask(h, cache=None): + N = h.shape[1] + if cache is not None and hasattr(cache, "make_mask"): + return cache.make_mask(N) + return None if N == 1 else "causal" + + +def _sdpa(queries, keys, values, scale, mask): + """``mx.fast.scaled_dot_product_attention``. Its "causal" mask is + aligned to the END of the keys (query i of L sees keys up to + Lk - L + i); torch's ``is_causal`` aligns to the start, which is wrong + once a cache holds earlier tokens -- so build the mask explicitly.""" + if isinstance(mask, str) and mask == "causal": + L, Lk = queries.shape[-2], keys.shape[-2] + q_pos = torch.arange(Lk - L, Lk, device=queries.device)[:, None] + mask = torch.arange(Lk, device=queries.device)[None, :] <= q_pos + return F.scaled_dot_product_attention(queries, keys, values, + attn_mask=mask, scale=scale) + + +# ---- KDA --------------------------------------------------------------------- + +class ShortConv1d(nn.Module): + """Causal depthwise conv with silu and a rolling cache; weight in the + checkpoint's torch layout [channels, 1, ksize].""" + + def __init__(self, channels: int, kernel_size: int): + super().__init__() + self.kernel_size = kernel_size + self.weight = nn.Parameter(torch.zeros(channels, 1, kernel_size)) + + def forward(self, x, state=None): + B, T, C = x.shape + if state is None: + state = torch.zeros(B, self.kernel_size - 1, C, dtype=x.dtype, + device=x.device) + conv_input = torch.cat([state, x], dim=1) + out = F.conv1d(conv_input.transpose(1, 2), + self.weight.to(conv_input.dtype), groups=C) + out = F.silu(out.transpose(1, 2)) + if self.kernel_size == 1: + new_state = conv_input[:, :0, :] + else: + new_state = conv_input[:, -(self.kernel_size - 1):, :] + return out, new_state + + +def _kda_gate(f, A_log, dt_bias, *, safe_gate: bool, lower_bound: float): # noqa: N803 + f = f.to(torch.float32) + dt_bias.to(torch.float32) + a = torch.exp(A_log.to(torch.float32)) + if safe_gate: + return lower_bound * torch.sigmoid(a[..., None] * f) + return -a[..., None] * F.softplus(f) + + +def _kda_update(q, k, v, g_log, beta, state): + """mlx-lm ``gated_delta_ops`` with vectorized (per-channel) decay: + q, k [B, T, H, Dk], v [B, T, H, Dv], g_log [B, T, H, Dk], beta + [B, T, H], state [B, H, Dv, Dk] (float32). Sequential over T.""" + g = torch.exp(g_log) + if state is None: + B = q.shape[0] + state = torch.zeros(B, v.shape[-2], v.shape[-1], k.shape[-1], + dtype=torch.float32, device=q.device) + ys = [] + for t in range(q.shape[1]): + qt, kt, vt, gt, bt = q[:, t], k[:, t], v[:, t], g[:, t], beta[:, t] + state = state * gt[..., None, :] + kv_mem = (state * kt[..., None, :]).sum(dim=-1) # [B, H, Dv] + delta = (vt - kv_mem) * bt[..., None] + state = state + kt[..., None, :] * delta[..., None] + ys.append((state * qt[..., None, :]).sum(dim=-1)) + return torch.stack(ys, dim=1), state + + +class BailingKDA(nn.Module): + """Kimi Delta Attention with the Ling V3 safe gate.""" + + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.num_heads = args.num_attention_heads + self.head_dim = args.head_dim + self.proj_dim = self.num_heads * self.head_dim + self.conv_kernel = args.short_conv_kernel_size + self.safe_gate = args.kda_safe_gate + self.lower_bound = float(args.kda_lower_bound) + self.no_kda_lora = args.no_kda_lora + self.scale = float(self.head_dim) ** -0.5 + + hidden = args.hidden_size + self.q_proj = _linear(hidden, self.proj_dim) + self.k_proj = _linear(hidden, self.proj_dim) + self.v_proj = _linear(hidden, self.proj_dim) + self.q_conv1d = ShortConv1d(self.proj_dim, self.conv_kernel) + self.k_conv1d = ShortConv1d(self.proj_dim, self.conv_kernel) + self.v_conv1d = ShortConv1d(self.proj_dim, self.conv_kernel) + if self.no_kda_lora: + self.f_proj = _linear(hidden, self.proj_dim) + self.g_proj = _linear(hidden, self.proj_dim) + else: + self.f_a_proj = _linear(hidden, self.head_dim) + self.f_b_proj = _linear(self.head_dim, self.proj_dim) + self.g_a_proj = _linear(hidden, self.head_dim) + self.g_b_proj = _linear(self.head_dim, self.proj_dim) + self.b_proj = _linear(hidden, self.num_heads) + self.A_log = nn.Parameter(torch.zeros(self.num_heads)) + self.dt_bias = nn.Parameter(torch.zeros(self.proj_dim)) + self.o_norm = RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.o_proj = _linear(self.proj_dim, hidden) + + def forward(self, x, mask=None, cache: Any | None = None): + B, T, _ = x.shape + dtype = x.dtype + if cache is not None: + q_state, k_state, v_state, ssm_state = cache.cache + else: + q_state = k_state = v_state = ssm_state = None + + q_conv, q_state = self.q_conv1d(self.q_proj(x), q_state) + k_conv, k_state = self.k_conv1d(self.k_proj(x), k_state) + v_conv, v_state = self.v_conv1d(self.v_proj(x), v_state) + if cache is not None: + cache[0], cache[1], cache[2] = q_state, k_state, v_state + + q = q_conv.reshape(B, T, self.num_heads, self.head_dim) + k = k_conv.reshape(B, T, self.num_heads, self.head_dim) + v = v_conv.reshape(B, T, self.num_heads, self.head_dim) + qf, kf = q.to(torch.float32), k.to(torch.float32) + q = self.scale * qf / (torch.linalg.norm(qf, dim=-1, keepdim=True) + 1e-6) + k = kf / (torch.linalg.norm(kf, dim=-1, keepdim=True) + 1e-6) + + if self.no_kda_lora: + f = self.f_proj(x) + gate = self.g_proj(x) + else: + f = self.f_b_proj(self.f_a_proj(x)) + gate = self.g_b_proj(self.g_a_proj(x)) + f = f.reshape(B, T, self.num_heads, self.head_dim) + g = _kda_gate(f, self.A_log, + self.dt_bias.reshape(self.num_heads, self.head_dim), + safe_gate=self.safe_gate, lower_bound=self.lower_bound) + beta = torch.sigmoid(self.b_proj(x).to(torch.float32)) + + out, ssm_state = _kda_update(q, k, v.to(torch.float32), g, beta, + ssm_state) + if cache is not None: + cache[3] = ssm_state + + gate = gate.reshape(B, T, self.num_heads, self.head_dim) + out = self.o_norm(out.to(dtype)) * torch.sigmoid(gate) + return self.o_proj(out.reshape(B, T, -1)) + + +# ---- MLA --------------------------------------------------------------------- + +class BailingMLA(nn.Module): + """DeepSeek-style MLA plus the V3 head-wise output gate.""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.num_heads = args.num_attention_heads + self.qk_nope_head_dim = args.qk_nope_head_dim + self.qk_rope_head_dim = args.qk_rope_head_dim + self.qk_head_dim = args.qk_head_dim + self.v_head_dim = args.v_head_dim + self.kv_lora_rank = args.kv_lora_rank + self.q_lora_rank = args.q_lora_rank + self.scale = self.qk_head_dim ** -0.5 + self.gate_kind = args.gated_attention_proj_granularity_type + + hidden = args.hidden_size + bias = args.use_qkv_bias + if self.q_lora_rank is None: + self.q_proj = _linear(hidden, self.num_heads * self.qk_head_dim) + else: + self.q_a_proj = _linear(hidden, self.q_lora_rank, bias) + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=args.rms_norm_eps) + self.q_b_proj = _linear(self.q_lora_rank, + self.num_heads * self.qk_head_dim) + self.kv_a_proj_with_mqa = _linear( + hidden, self.kv_lora_rank + self.qk_rope_head_dim, bias) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=args.rms_norm_eps) + self.kv_b_proj = _linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)) + if self.gate_kind == "head_wise": + self.g_proj = _linear(hidden, self.num_heads) + elif self.gate_kind == "element_wise": + self.g_proj = _linear(hidden, self.num_heads * self.v_head_dim) + self.dense = _linear(self.num_heads * self.v_head_dim, hidden, bias) + if args.rope_scaling: + raise NotImplementedError( + "bailing_hybrid: rope_scaling is not supported " + f"(got {args.rope_scaling!r})") + self.rope_theta = args.rope_theta + + def forward(self, x, mask=None, cache: Any | None = None): + B, L, _ = x.shape + if self.q_lora_rank is None: + q = self.q_proj(x) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x))) + q = q.reshape(B, L, self.num_heads, self.qk_head_dim).transpose(1, 2) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + compressed = self.kv_a_proj_with_mqa(x) + kv_latent, k_pe = torch.split( + compressed, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + kv_latent = self.kv_a_layernorm(kv_latent) + k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(1, 2) + + kv = self.kv_b_proj(kv_latent) + kv = kv.reshape(B, L, self.num_heads, + self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2) + k_nope, values = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) + + offset = cache.offset if cache is not None else 0 + positions = torch.arange(L, device=x.device) + offset + q_pe = _rope_interleave_torch(q_pe, positions, self.rope_theta) + k_pe = _rope_interleave_torch(k_pe, positions, self.rope_theta) + k_pe = k_pe.expand(B, self.num_heads, L, self.qk_rope_head_dim) + + # RoPE computes in float32; MLX's concatenate promotes the same way + queries = torch.cat([q_nope.to(q_pe.dtype), q_pe], dim=-1) + keys = torch.cat([k_nope.to(k_pe.dtype), k_pe], dim=-1) + values = values.to(keys.dtype) + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + out = _sdpa(queries, keys, values, self.scale, mask) + out = out.transpose(1, 2).reshape(B, L, -1) + if self.gate_kind == "head_wise": + gate = torch.sigmoid(self.g_proj(x)) + out = out.reshape(B, L, self.num_heads, self.v_head_dim) + out = (out * gate[..., None]).reshape(B, L, -1) + elif self.gate_kind == "element_wise": + out = out * torch.sigmoid(self.g_proj(x)) + return self.dense(out) + + +# ---- MLP / MoE ----------------------------------------------------------------- + +class BailingMLP(nn.Module): + def __init__(self, args: ModelArgs, intermediate: int): + super().__init__() + self.gate_proj = _linear(args.hidden_size, intermediate) + self.up_proj = _linear(args.hidden_size, intermediate) + self.down_proj = _linear(intermediate, args.hidden_size) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +def _group_topk(scores, select, a: ModelArgs): + """Group-limited top-k shared by BailingGate and the prerouter path + (moe/routing.py's law): drop the n_group - topk_group groups with the + lowest top-2 sums, then take the top-k by selection score; the weights + are the raw scores of the chosen experts.""" + B = select.shape[:-1] + k_drop = a.n_group - a.topk_group + if k_drop > 0: + grouped = select.reshape(*B, a.n_group, a.num_experts // a.n_group) + group_scores = torch.topk(grouped, 2, dim=-1).values.sum(dim=-1) + drop = torch.argsort(group_scores, dim=-1, stable=True)[..., :k_drop] + masked = grouped.clone() + masked.scatter_(-2, drop[..., None].expand(*drop.shape, grouped.shape[-1]), + float("-inf")) + select = masked.reshape(*B, a.num_experts) + k = a.num_experts_per_tok + idx = torch.argsort(-select, dim=-1, stable=True)[..., :k] + w = torch.take_along_dim(scores, idx, dim=-1) + if a.norm_topk_prob: + w = w / (w.sum(dim=-1, keepdim=True) + 1e-20) + return idx, w * a.routed_scaling_factor + + +class BailingGate(nn.Module): + """Sigmoid-score router with noaux_tc expert-bias group selection.""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.weight = nn.Parameter(torch.zeros(args.num_experts, args.hidden_size)) + if args.moe_router_enable_expert_bias: + self.expert_bias = nn.Parameter(torch.zeros(args.num_experts)) + + def forward(self, x): + scores = torch.sigmoid((x @ self.weight.to(x.dtype).T).to(torch.float32)) + select = scores + if self.args.moe_router_enable_expert_bias: + select = scores + self.expert_bias.to(torch.float32) + return _group_topk(scores, select, self.args) + + +class BailingPrerouter(nn.Module): + """One-MoE-block-ahead expert selector (weights from the prerouter file, + installed by prerouter/install.py).""" + + def __init__(self, args: ModelArgs, hidden: int | None = None, + input_dim: int | None = None): + super().__init__() + self.args = args + self.hidden = hidden or args.prerouter_hidden + self.input_dim = input_dim or args.hidden_size + self.fc1 = _linear(self.input_dim, self.hidden) + self.fc2 = _linear(self.hidden, args.num_experts) + self.linear_init = _linear(self.input_dim, args.num_experts) + with torch.no_grad(): + self.linear_init.weight.zero_() + + def forward(self, x, topk_oh=None, prev_oh=None): + dtype = self.fc1.weight.dtype + feats = [x.to(dtype)] + if topk_oh is not None: + feats.append(topk_oh.to(dtype)) + if prev_oh is not None: + feats.append(prev_oh.to(dtype)) + xi = torch.cat(feats, dim=-1) + h = F.gelu(self.fc1(xi)) + # MLX promotes mixed dtypes (e.g. fp16 heads + a float32 zero + # linear_init) to the wider type; torch Linear needs one dtype. + a = self.fc2(h) + b = self.linear_init(xi.to(self.linear_init.weight.dtype)) + out = torch.promote_types(a.dtype, b.dtype) + return a.to(out) + b.to(out) + + +class BailingSparseMoE(nn.Module): + def __init__(self, args: ModelArgs, prerouter: BailingPrerouter | None = None, + use_prerouter: bool = False): + super().__init__() + self.args = args + self.gate = BailingGate(args) + self.prerouter = prerouter + self.use_prerouter = use_prerouter + self.last_topk = None + self.prev_topk_oh = None + self.feature_topk = os.environ.get("PREROUTER_FEATURE_TOPK", "teacher") + self.intra_mode = os.environ.get("PREROUTER_INTRA", "0") == "1" + # Routed experts come from the streaming installer + # (install_streaming_experts swaps a StreamingSwitchGLU in here); + # the checkpoint's stacked expert tensors are never made resident. + self.experts = None + self.shared_experts = BailingMLP( + args, args.moe_shared_expert_intermediate_size * args.num_shared_experts) + + def _select_from_logits(self, logits): + scores = torch.sigmoid(logits.to(torch.float32)) + return _group_topk(scores, scores, self.args) + + def forward(self, x, prev_prerouter_logits=None): + if self.use_prerouter and prev_prerouter_logits is not None: + idx, w = self._select_from_logits(prev_prerouter_logits) + if not self.intra_mode and self.feature_topk == "teacher": + feat_idx, _ = self.gate(x) + else: + feat_idx = idx + else: + idx, w = self.gate(x) + feat_idx = idx + self.last_topk = feat_idx if not self.intra_mode else None + if self.experts is None: + raise RuntimeError( + "routed experts not installed: run install_streaming_experts") + routed = self.experts(x, idx) + out = (routed * w[..., None].to(routed.dtype)).sum(dim=-2) + return out + self.shared_experts(x) + + +class BailingDecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.is_mla = args.is_mla_layer(layer_idx) + self.attention = BailingMLA(args) if self.is_mla else BailingKDA(args, layer_idx) + self.prerouter_enabled = bool(getattr(args, "prerouter_enabled", False)) + self.prerouter_start_layer = int(getattr(args, "prerouter_start_layer", 7)) + self.prerouter_logits = None + self.m_in_cache = None + self.after_attention_cb = None + self.has_prerouter = ( + self.prerouter_enabled + and layer_idx >= self.prerouter_start_layer - 1 + and layer_idx < args.num_hidden_layers - 1) + self.use_prerouter = ( + self.prerouter_enabled + and layer_idx >= self.prerouter_start_layer + and layer_idx < args.num_hidden_layers) + if layer_idx >= args.first_k_dense_replace: + self.mlp = BailingSparseMoE( + args, + prerouter=BailingPrerouter(args) if self.has_prerouter else None, + use_prerouter=self.use_prerouter) + else: + self.mlp = BailingMLP(args, args.intermediate_size) + self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def forward(self, x, mask=None, cache=None, prev_prerouter_logits=None): + h = x + self.attention(self.input_layernorm(x), mask, cache) + m_in = self.post_attention_layernorm(h) + self.m_in_cache = m_in if self.has_prerouter else None + self.prerouter_logits = None + if self.after_attention_cb is not None: + self.after_attention_cb(self.layer_idx, m_in) + if isinstance(self.mlp, BailingSparseMoE): + out = self.mlp(m_in, prev_prerouter_logits=prev_prerouter_logits) + else: + out = self.mlp(m_in) + return h + out + + +class BailingModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = nn.ModuleList( + [BailingDecoderLayer(args, i) for i in range(args.num_hidden_layers)]) + self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.first_mla_idx = next( + i for i in range(args.num_hidden_layers) if args.is_mla_layer(i)) + + def forward(self, inputs, cache=None, input_embeddings=None, + after_layer_cb=None, before_layer_cb=None, + async_eval_per_layer: bool = False, prerouter_cache=None): + h = input_embeddings if input_embeddings is not None \ + else self.word_embeddings(inputs) + if cache is None: + cache = [None] * len(self.layers) + mla_mask = create_attention_mask(h, cache[self.first_mla_idx]) + clip_v = float(os.environ.get("LING_HIDDEN_CLIP", "0")) + for li, (layer, c) in enumerate(zip(self.layers, cache)): + if before_layer_cb is not None: + before_layer_cb(li) + mask = mla_mask if layer.is_mla else None + prev = None + if (getattr(layer, "use_prerouter", False) + and prerouter_cache is not None and li in prerouter_cache): + prev = prerouter_cache[li] + h = layer(h, mask, c, prev) + if clip_v > 0: + h = torch.where(torch.isnan(h), torch.zeros_like(h), + h.clamp(-clip_v, clip_v)) + if after_layer_cb is not None: + after_layer_cb(li, h) + # async_eval_per_layer: an MLX lazy-graph hint; torch is eager. + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = BailingModel(args) + self.tie_word_embeddings = args.tie_word_embeddings + if not self.tie_word_embeddings: + self.lm_head = _linear(args.hidden_size, args.vocab_size) + + def forward(self, inputs, cache=None, input_embeddings=None, + after_layer_cb=None, before_layer_cb=None, + async_eval_per_layer: bool = False, prerouter_cache=None): + h = self.model(inputs, cache, input_embeddings, after_layer_cb, + before_layer_cb=before_layer_cb, + async_eval_per_layer=async_eval_per_layer, + prerouter_cache=prerouter_cache) + if self.tie_word_embeddings: + return F.linear(h, self.model.word_embeddings.weight) + return self.lm_head(h) + + def sanitize(self, weights: dict) -> dict: + """Drop MTP heads, like the MLX sanitize. Experts are stacked on + disk already, and conv weights are used in their stored layout.""" + return {k: v for k, v in weights.items() + if ".mtp_" not in k and "mtp." not in k} + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [KVCache() if layer.is_mla else ArraysCache(size=4) + for layer in self.model.layers] diff --git a/python/src/edge0/backends/cuda/_impl/qwen3_5_moe.py b/python/src/edge0/backends/cuda/_impl/qwen3_5_moe.py new file mode 100644 index 0000000..6c03b52 --- /dev/null +++ b/python/src/edge0/backends/cuda/_impl/qwen3_5_moe.py @@ -0,0 +1,507 @@ +"""Torch port of the vendored Qwen3.5-MoE backbone (edge0-35b). + +A line-by-line port of ``backends/mlx/_impl/qwen3_5_moe.py`` plus the +``qwen3_5.py`` / ``qwen3_next.py`` pieces it builds on (mlx-lm 0.31.0), +with the same module tree, parameter names and call signatures, so +``engine/qwen.py`` and the prerouter patch drive it unchanged. + +It computes the MLX way on the MLX checkpoint layout: RMSNorm weights are +the absolute multiplier (MLX's sanitize stores ``w + 1`` for the +zero-centred transformers norms) and ``conv1d.weight`` is [C, k, 1], so the +published edge0-35b tensors are used exactly as stored -- no conversion on +load. Checked against the MLX module in ``tests/test_backend_parity.py``. +""" + +from __future__ import annotations + +import inspect +import os +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from edge0.backends.cuda._impl.bailing_hybrid import ( + ArraysCache, + KVCache, + RMSNorm, + _linear, + _rms_norm, + _sdpa, + create_attention_mask, +) + + +def create_ssm_mask(h, cache=None): + """mlx-lm: only padded batches get an SSM mask; edge0 serves B=1.""" + return None + + +def _rope(x, dims: int, base: float, offset: int): + """``mx.fast.rope(traditional=False)``: rotate the first ``dims`` of the + last axis as two halves (i, i + dims/2) by ``pos * base**(-2i/dims)``; + the rest passes through.""" + L = x.shape[-2] + half = dims // 2 + inv_freq = base ** (-torch.arange(0, half, dtype=torch.float32, + device=x.device) * 2 / dims) + pos = torch.arange(offset, offset + L, dtype=torch.float32, device=x.device) + ang = pos[:, None] * inv_freq[None, :] + cos, sin = torch.cos(ang), torch.sin(ang) + xf = x.to(torch.float32) + x1, x2, rest = xf[..., :half], xf[..., half:dims], xf[..., dims:] + out = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos, rest], dim=-1) + return out.to(x.dtype) + + +@dataclass +class TextModelArgs: + model_type: str = "" + hidden_size: int = 4096 + intermediate_size: int = 14336 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + rms_norm_eps: float = 1e-6 + vocab_size: int = 151936 + num_key_value_heads: int = 8 + max_position_embeddings: int = 131072 + linear_num_value_heads: int = 64 + linear_num_key_heads: int = 16 + linear_key_head_dim: int = 192 + linear_value_head_dim: int = 128 + linear_conv_kernel_dim: int = 4 + tie_word_embeddings: bool = False + attention_bias: bool = False + head_dim: Optional[int] = None + full_attention_interval: int = 4 + num_experts: int = 0 + num_experts_per_tok: int = 0 + decoder_sparse_step: int = 1 + shared_expert_intermediate_size: int = 0 + moe_intermediate_size: int = 0 + norm_topk_prob: bool = True + rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = field( + default_factory=lambda: { + "type": "default", "mrope_section": [11, 11, 10], + "rope_theta": 100000, "partial_rotary_factor": 0.25}) + partial_rotary_factor: float = 0.25 + rope_theta: float = 100000.0 + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + + @classmethod + def from_dict(cls, params: dict): + names = inspect.signature(cls).parameters + return cls(**{k: v for k, v in params.items() if k in names}) + + def __post_init__(self): + if self.head_dim is None: + self.head_dim = self.hidden_size // self.num_attention_heads + if self.rope_parameters: + if ("type" not in self.rope_parameters + and "rope_type" in self.rope_parameters): + self.rope_parameters["type"] = self.rope_parameters.pop("rope_type") + self.partial_rotary_factor = self.rope_parameters.get( + "partial_rotary_factor", 0.25) + self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0) + self.rope_scaling = self.rope_parameters + + +class RMSNormGated(nn.Module): + """``Qwen3NextRMSNormGated``: rms_norm(x) * silu(gate), the gating in + float32 (``_precise_swiglu``), result in the input dtype.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(hidden_size)) + + def forward(self, hidden_states, gate=None): + x = _rms_norm(hidden_states, self.weight, self.eps) + if gate is None: + return x.to(hidden_states.dtype) + g = F.silu(gate.to(torch.float32)) + return (g * x.to(torch.float32)).to(hidden_states.dtype) + + +class RoPE(nn.Module): + """``nn.RoPE(dims, traditional=False, base)``: the only kind the + edge0-35b config asks for (rope_type "default").""" + + def __init__(self, dims: int, base: float, scale: float = 1.0): + super().__init__() + if scale != 1.0: + raise NotImplementedError("scaled RoPE is not ported") + self.dims, self.base = dims, base + + def forward(self, x, offset: int = 0): + return _rope(x, self.dims, self.base, offset) + + +def _initialize_rope(dims, base, scaling_config): + rope_type = "default" + if scaling_config is not None: + rope_type = scaling_config.get("type") or scaling_config.get( + "rope_type", "default") + if rope_type in ("default", "mrope"): + return RoPE(dims, base) + raise NotImplementedError(f"RoPE type {rope_type!r} is not ported") + + +class Attention(nn.Module): + """``Qwen3NextAttention``: gated full attention with GQA, q/k norms and + partial RoPE.""" + + def __init__(self, args: TextModelArgs): + super().__init__() + self.num_key_value_heads = args.num_key_value_heads + self.num_attention_heads = args.num_attention_heads + self.head_dim = args.head_dim + self.scale = self.head_dim ** -0.5 + bias = args.attention_bias + self.q_proj = _linear(args.hidden_size, + self.num_attention_heads * self.head_dim * 2, bias) + self.k_proj = _linear(args.hidden_size, + self.num_key_value_heads * self.head_dim, bias) + self.v_proj = _linear(args.hidden_size, + self.num_key_value_heads * self.head_dim, bias) + self.o_proj = _linear(self.num_attention_heads * self.head_dim, + args.hidden_size, bias) + self.q_norm = RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.rope = _initialize_rope( + int(self.head_dim * args.partial_rotary_factor), args.rope_theta, + args.rope_scaling) + + def forward(self, x, mask=None, cache: Optional[Any] = None): + B, L, _ = x.shape + q_out = self.q_proj(x).reshape(B, L, self.num_attention_heads, -1) + queries, gate = torch.split(q_out, q_out.shape[-1] // 2, dim=-1) + gate = gate.reshape(B, L, -1) + keys, values = self.k_proj(x), self.v_proj(x) + queries = self.q_norm(queries).transpose(1, 2) + keys = self.k_norm(keys.reshape(B, L, self.num_key_value_heads, -1)).transpose(1, 2) + values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(1, 2) + offset = cache.offset if cache is not None else 0 + queries = self.rope(queries, offset=offset) + keys = self.rope(keys, offset=offset) + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + # GQA as mx.fast.sdpa does it: query head h reads kv head h // rep + rep = self.num_attention_heads // self.num_key_value_heads + if rep > 1: + keys = keys.repeat_interleave(rep, dim=1) + values = values.repeat_interleave(rep, dim=1) + out = _sdpa(queries, keys, values, self.scale, mask) + out = out.transpose(1, 2).reshape(B, L, -1) + return self.o_proj(out * torch.sigmoid(gate)) + + +class MLP(nn.Module): + """``Qwen3NextMLP``.""" + + def __init__(self, dim, hidden_dim): + super().__init__() + self.gate_proj = _linear(dim, hidden_dim) + self.down_proj = _linear(hidden_dim, dim) + self.up_proj = _linear(dim, hidden_dim) + + def forward(self, x): + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +def _gated_delta_update(q, k, v, a, b, A_log, dt_bias, state=None): # noqa: N803 + """mlx-lm ``gated_delta_update`` on its ops path: beta = sigmoid(b), + g = exp(-exp(A_log) * softplus(a + dt_bias)) cast to a's dtype, one + decay scalar per value head, state in q's dtype, k heads repeated + (consecutively) to the value heads.""" + beta = torch.sigmoid(b) + g = torch.exp(-torch.exp(A_log.to(torch.float32)) + * F.softplus(a + dt_bias)).to(a.dtype) + B, T, Hk, Dk = q.shape + Hv, Dv = v.shape[-2:] + if state is None: + state = torch.zeros(B, Hv, Dv, Dk, dtype=q.dtype, device=q.device) + if Hv // Hk > 1: + q = q.repeat_interleave(Hv // Hk, dim=-2) + k = k.repeat_interleave(Hv // Hk, dim=-2) + ys = [] + for t in range(T): + qt, kt, vt, gt, bt = q[:, t], k[:, t], v[:, t], g[:, t], beta[:, t] + state = state * gt[..., None, None] + kv_mem = (state * kt[..., None, :]).sum(dim=-1) + delta = (vt - kv_mem) * bt[..., None] + state = state + kt[..., None, :] * delta[..., None] + ys.append((state * qt[..., None, :]).sum(dim=-1)) + return torch.stack(ys, dim=1), state + + +class GatedDeltaNet(nn.Module): + """``qwen3_5.GatedDeltaNet`` (split in_proj_qkv / z / b / a).""" + + def __init__(self, config: TextModelArgs): + super().__init__() + self.hidden_size = config.hidden_size + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + if self.num_v_heads % self.num_k_heads != 0: + raise ValueError("num_v_heads must be divisible by num_k_heads") + self.conv_kernel_size = config.linear_conv_kernel_dim + self.layer_norm_epsilon = config.rms_norm_eps + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = _DepthwiseConv(self.conv_dim, self.conv_kernel_size) + self.in_proj_qkv = _linear(self.hidden_size, self.key_dim * 2 + self.value_dim) + self.in_proj_z = _linear(self.hidden_size, self.value_dim) + self.in_proj_b = _linear(self.hidden_size, self.num_v_heads) + self.in_proj_a = _linear(self.hidden_size, self.num_v_heads) + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) + self.A_log = nn.Parameter(torch.zeros(self.num_v_heads)) + self.norm = RMSNormGated(self.head_v_dim, eps=self.layer_norm_epsilon) + self.out_proj = _linear(self.value_dim, self.hidden_size) + + def forward(self, inputs, mask=None, cache: Optional[Any] = None): + B, S, _ = inputs.shape + qkv = self.in_proj_qkv(inputs) + z = self.in_proj_z(inputs).reshape(B, S, self.num_v_heads, self.head_v_dim) + b = self.in_proj_b(inputs) + a = self.in_proj_a(inputs) + if cache is not None and cache[0] is not None: + conv_state = cache[0] + else: + conv_state = torch.zeros(B, self.conv_kernel_size - 1, self.conv_dim, + dtype=inputs.dtype, device=inputs.device) + if mask is not None: + qkv = torch.where(mask[..., None], qkv, 0) + conv_input = torch.cat([conv_state, qkv], dim=1) + if cache is not None: + cache[0] = conv_input[:, -(self.conv_kernel_size - 1):] + conv_out = F.silu(self.conv1d(conv_input)) + q, k, v = [ + t.reshape(B, S, h, d) for t, h, d in zip( + torch.split(conv_out, [self.key_dim, self.key_dim, self.value_dim], -1), + [self.num_k_heads, self.num_k_heads, self.num_v_heads], + [self.head_k_dim, self.head_k_dim, self.head_v_dim])] + state = cache[1] if cache is not None else None + inv_scale = k.shape[-1] ** -0.5 + ones = torch.ones(q.shape[-1], dtype=q.dtype, device=q.device) + q = (inv_scale ** 2) * _rms_norm(q, ones, 1e-6) + k = inv_scale * _rms_norm(k, ones, 1e-6) + out, state = _gated_delta_update(q, k, v, a, b, self.A_log, self.dt_bias, state) + if cache is not None: + cache[1] = state + out = self.norm(out, z) + return self.out_proj(out.reshape(B, S, -1)) + + +class _DepthwiseConv(nn.Module): + """MLX ``nn.Conv1d(groups=C, padding=0)`` on [B, T, C]; the weight in + MLX's layout [C, k, 1] as stored in the checkpoint.""" + + def __init__(self, channels: int, kernel_size: int): + super().__init__() + self.weight = nn.Parameter(torch.zeros(channels, kernel_size, 1)) + + def forward(self, x): + w = self.weight.transpose(1, 2).to(x.dtype) # [C, 1, k] + return F.conv1d(x.transpose(1, 2), w, groups=x.shape[-1]).transpose(1, 2) + + +class SparseMoeBlock(nn.Module): + """``Qwen3NextSparseMoeBlock``: softmax top-k router (precise softmax), + routed experts from the streaming installer, gated shared expert. The + 35b prerouter patches this class's ``__call__`` (prerouter/install.py), + as it does the MLX class.""" + + def __init__(self, args: TextModelArgs): + super().__init__() + dim = args.hidden_size + self.norm_topk_prob = args.norm_topk_prob + self.num_experts = args.num_experts + self.top_k = args.num_experts_per_tok + self.gate = _linear(dim, self.num_experts) + self.switch_mlp = None # install_streaming_experts swaps the twin in + self.shared_expert = MLP(dim, args.shared_expert_intermediate_size) + self.shared_expert_gate = _linear(dim, 1) + + def forward(self, x): + gates = self.gate(x) + gates = torch.softmax(gates.to(torch.float32), dim=-1).to(gates.dtype) + k = self.top_k + inds = torch.argsort(gates, dim=-1, stable=True)[..., -k:] + scores = torch.take_along_dim(gates, inds, dim=-1) + if self.norm_topk_prob: + scores = scores / scores.sum(dim=-1, keepdim=True) + if self.switch_mlp is None: + raise RuntimeError( + "routed experts not installed: run install_streaming_experts") + y = self.switch_mlp(x, inds) + y = (y * scores[..., None]).sum(dim=-2) + shared_y = self.shared_expert(x) + shared_y = torch.sigmoid(self.shared_expert_gate(x)) * shared_y + return y + shared_y + + +class DecoderLayer(nn.Module): + def __init__(self, args: TextModelArgs, layer_idx: int): + super().__init__() + self.is_linear = (layer_idx + 1) % args.full_attention_interval != 0 + if self.is_linear: + self.linear_attn = GatedDeltaNet(args) + else: + self.self_attn = Attention(args) + self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + if args.num_experts > 0: + self.mlp = SparseMoeBlock(args) + else: + self.mlp = MLP(args.hidden_size, args.intermediate_size) + + def forward(self, x, mask=None, cache=None): + if self.is_linear: + r = self.linear_attn(self.input_layernorm(x), mask, cache) + else: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + return h + self.mlp(self.post_attention_layernorm(h)) + + +class Qwen3_5TextModel(nn.Module): + def __init__(self, args: TextModelArgs): + super().__init__() + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = nn.ModuleList( + [DecoderLayer(args, i) for i in range(args.num_hidden_layers)]) + self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.ssm_idx = 0 + self.fa_idx = args.full_attention_interval - 1 + + def forward(self, inputs, cache=None, input_embeddings=None, + before_layer_cb=None, after_layer_cb=None, + async_eval_per_layer: bool = False): + hidden_states = input_embeddings if input_embeddings is not None \ + else self.embed_tokens(inputs) + if cache is None: + cache = [None] * len(self.layers) + fa_mask = create_attention_mask(hidden_states, cache[self.fa_idx]) + ssm_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx]) + # QWEN_HIDDEN_CLIP mirrors the vendored MLX qwen3_5: clamp each + # layer output to [-v, v] and zero NaN entries (0 = off). + clip_v = float(os.environ.get("QWEN_HIDDEN_CLIP", "0")) + for li, (layer, c) in enumerate(zip(self.layers, cache)): + if before_layer_cb is not None: + before_layer_cb(li) + mask = ssm_mask if layer.is_linear else fa_mask + hidden_states = layer(hidden_states, mask=mask, cache=c) + if clip_v > 0: + hidden_states = torch.where( + torch.isnan(hidden_states), torch.zeros_like(hidden_states), + hidden_states.clamp(-clip_v, clip_v)) + if after_layer_cb is not None: + after_layer_cb(li, hidden_states) + # async_eval_per_layer: an MLX lazy-graph hint; torch is eager. + return self.norm(hidden_states) + + +class TextModel(nn.Module): + def __init__(self, args: TextModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Qwen3_5TextModel(args) + if not args.tie_word_embeddings: + self.lm_head = _linear(args.hidden_size, args.vocab_size) + + def forward(self, inputs, cache=None, input_embeddings=None): + out = self.model(inputs, cache, input_embeddings=input_embeddings) + if self.args.tie_word_embeddings: + return F.linear(out, self.model.embed_tokens.weight) + return self.lm_head(out) + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [ArraysCache(size=2) if l.is_linear else KVCache() + for l in self.layers] + + def sanitize(self, weights: dict) -> dict: + """The MLX sanitize, on torch tensors: drop MTP; a raw transformers + checkpoint (conv1d [C, 1, k]) is moved to the MLX layout and its + zero-centred norms shifted by +1. The published edge0 checkpoints + are already in MLX form, so for them this changes nothing.""" + has_mtp = any("mtp." in k for k in weights) + unsanitized = any("conv1d.weight" in k and v.shape[-1] != 1 + for k, v in weights.items()) + shift = has_mtp or unsanitized + weights = {k: v for k, v in weights.items() if "mtp." not in k} + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + norm_keys = (".input_layernorm.weight", ".post_attention_layernorm.weight", + "model.norm.weight", ".q_norm.weight", ".k_norm.weight") + for k, v in list(weights.items()): + if "conv1d.weight" in k and v.shape[-1] != 1: + weights[k] = v.transpose(1, 2).contiguous() + if shift and k.endswith(norm_keys) and v.ndim == 1: + weights[k] = v + 1.0 + return weights + + +@dataclass +class ModelArgs: + model_type: str + text_config: dict + + @classmethod + def from_dict(cls, params: dict): + if "text_config" not in params: + return cls(model_type=params["model_type"], text_config=params) + return cls(model_type=params["model_type"], text_config=params["text_config"]) + + +class Model(nn.Module): + """``qwen3_5_moe.Model`` (via ``qwen3_5.Model``).""" + + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.language_model = TextModel(TextModelArgs.from_dict(args.text_config)) + + def forward(self, inputs, cache=None, input_embeddings=None): + return self.language_model(inputs, cache=cache, + input_embeddings=input_embeddings) + + def sanitize(self, weights: dict) -> dict: + new = {} + for key, value in weights.items(): + if key.startswith("vision_tower") or key.startswith("model.visual"): + continue + if key.startswith("model.language_model"): + key = key.replace("model.language_model", "language_model.model") + elif not key.startswith("language_model."): + key = "language_model." + key + new[key] = value + for l in range(self.language_model.args.num_hidden_layers): + prefix = f"language_model.model.layers.{l}.mlp" + gate_up_key = f"{prefix}.experts.gate_up_proj" + if gate_up_key in new: + gate_up = new.pop(gate_up_key) + mid = gate_up.shape[-2] // 2 + new[f"{prefix}.switch_mlp.gate_proj.weight"] = gate_up[..., :mid, :] + new[f"{prefix}.switch_mlp.up_proj.weight"] = gate_up[..., mid:, :] + new[f"{prefix}.switch_mlp.down_proj.weight"] = new.pop( + f"{prefix}.experts.down_proj") + return self.language_model.sanitize(new) + + @property + def layers(self): + return self.language_model.model.layers + + def make_cache(self): + return self.language_model.make_cache() diff --git a/python/src/edge0/backends/cuda/backend.py b/python/src/edge0/backends/cuda/backend.py new file mode 100644 index 0000000..54130bd --- /dev/null +++ b/python/src/edge0/backends/cuda/backend.py @@ -0,0 +1,17 @@ +"""CUDA backend: namespace assembly (torch-backed reference implementation).""" + +from __future__ import annotations + +from edge0.backends.cuda import core, io, nn, quant # noqa: F401 + + +class BackendImpl: + """Handle to the active backend implementation.""" + + name = "cuda" + description = "PyTorch / CUDA (reference implementation)" + + @property + def version(self) -> str: + import torch + return torch.__version__ diff --git a/python/src/edge0/backends/cuda/core.py b/python/src/edge0/backends/cuda/core.py new file mode 100644 index 0000000..4ce3da7 --- /dev/null +++ b/python/src/edge0/backends/cuda/core.py @@ -0,0 +1,332 @@ +"""CUDA backend: array ops namespace (mirrors ``mlx.core`` by name). + +The contract in ``edge0/backends/__init__.py`` documents this namespace +against MLX's naming, not PyTorch's idiomatic API -- so this module wraps +torch under MLX's names rather than exposing torch as-is. Framework code +(``streaming/layer.py``, ``moe/routing.py``, ``sampling.py``, ...) calls +``core.take_along_axis(...)``, not ``torch.gather(...)``. + +Known semantic gaps vs. MLX, flagged rather than silently papered over: + +* MLX arrays are copy-on-write / functionally immutable; every op here + that MLX documents as returning a new array (``put_along_axis``, ...) + is implemented with torch's out-of-place variant even where an + in-place ``_`` variant would be faster. ``streaming/layer.py``'s + incremental-stack path (sticky slots, see its module docstring) is + exactly the place that later wants an in-place + ``put_along_axis_``-style fast path once the CUDA staging design is + set -- deliberately not pre-empted here. +* ``eval`` / ``compile`` exist because MLX is lazy-by-default; torch is + eager, so ``eval`` is a no-op and ``compile`` wraps ``torch.compile`` + (itself opt-in-able, since it re-traces on shape changes -- exactly + the staged/exact path shape variability in ``streaming/layer.py``). +* ``take`` / ``take_along_axis`` follow numpy/MLX gather semantics + (index array shape composes with the source's remaining axes), not + ``torch.take``'s flat-index semantics -- implemented via + ``index_select`` + reshape. +""" + +from __future__ import annotations + +import os + +import torch + +# EDGE0_TORCH_DEVICE overrides the choice: "cpu" on a GPU machine, or "mps" +# on Apple Silicon, where a tensor left on the host fails loudly as it would +# on CUDA -- the only accelerator next to the MLX reference. +DEVICE = torch.device(os.environ.get("EDGE0_TORCH_DEVICE") + or ("cuda" if torch.cuda.is_available() else "cpu")) + +# ---- dtypes (contract: float16/float32/bfloat16/int32/uint32/int64) ------ + +float16 = torch.float16 +float32 = torch.float32 +bfloat16 = torch.bfloat16 +int32 = torch.int32 +uint32 = torch.uint32 +int64 = torch.int64 + + +def _as_tensor(x, dtype=None): + if isinstance(x, torch.Tensor): + t = x + else: + # numpy arrays (incl. non-writeable mmap views) and python lists + t = torch.as_tensor(x) + if dtype is not None and t.dtype != dtype: + t = t.to(dtype) + if t.device != DEVICE: + t = t.to(DEVICE, non_blocking=True) + return t + + +# ---- construction ---------------------------------------------------- + +def array(x, dtype=None): + """``mx.array`` equivalent: numpy/list/scalar -> device tensor. + + This is the single conversion point the streaming layer calls once + per expert bundle (see ``streaming/layer.py::_build``); the actual + host->device transfer for the SSD-offload path happens here. A + synchronous ``.to(DEVICE)`` is a correct placeholder -- the staging + strategy (pinned double-buffer vs. GPUDirect Storage) is the part + intentionally deferred, and replacing this function's body is the + entire diff either route needs; nothing in streaming/layer.py calls + torch directly. + """ + return _as_tensor(x, dtype) + + +def zeros(shape, dtype=float32): + return torch.zeros(shape, dtype=dtype, device=DEVICE) + + +def zeros_like(x): + return torch.zeros_like(x) + + +def eye(n, dtype=float32): + return torch.eye(n, dtype=dtype, device=DEVICE) + + +def arange(*args, dtype=None): + return torch.arange(*args, dtype=dtype, device=DEVICE) + + +def full(shape, value, dtype=float32): + return torch.full(shape, value, dtype=dtype, device=DEVICE) + + +# ---- shape ops ---------------------------------------------------------- + +def expand_dims(x, axes): + """numpy/MLX semantics: axes (negative ones too) index the OUTPUT, so + ``expand_dims(x[m, H], (-2, -3))`` is ``[m, 1, 1, H]``.""" + if isinstance(axes, int): + axes = (axes,) + out_ndim = x.ndim + len(axes) + for ax in sorted(a % out_ndim for a in axes): + x = torch.unsqueeze(x, ax) + return x + + +def squeeze(x, axis=None): + return torch.squeeze(x) if axis is None else torch.squeeze(x, axis) + + +def reshape(x, shape): + return torch.reshape(x, tuple(shape)) + + +def transpose(x, axes=None): + return x.permute(*axes) if axes is not None else x.t() + + +def concatenate(arrays, axis=0): + return torch.cat(list(arrays), dim=axis) + + +def stack(arrays, axis=0): + return torch.stack(list(arrays), dim=axis) + + +def split(x, indices_or_sections, axis=0): + # mx.split(x, [i, j], axis) takes SPLIT POINTS, like numpy.split -- + # torch.split takes SECTION SIZES. Convert when a list is given; an + # int section count (used e.g. in streaming/layer.py's + # split(x_gu, 2, axis=-1)) is already the torch convention. + if isinstance(indices_or_sections, int): + return list(torch.chunk(x, indices_or_sections, dim=axis)) + points = list(indices_or_sections) + sizes = [] + prev = 0 + n = x.shape[axis] + for p in points: + sizes.append(p - prev) + prev = p + sizes.append(n - prev) + return list(torch.split(x, sizes, dim=axis)) + + +# ---- math / reductions --------------------------------------------------- + +def matmul(a, b): + return torch.matmul(a, b) + + +def einsum(subscripts, *operands): + return torch.einsum(subscripts, *operands) + + +def softmax(x, axis=-1, precise=False): + """``precise=True`` accumulates in float32, as MLX does, and returns + the input dtype.""" + if precise: + return torch.softmax(x.float(), dim=axis).to(x.dtype) + return torch.softmax(x, dim=axis) + + +def sigmoid(x): + return torch.sigmoid(x) + + +def erf(x): + return torch.erf(x) + + +def where(cond, a, b): + return torch.where(cond, a, b) + + +def sum(x, axis=None, keepdims=False): + return torch.sum(x) if axis is None else torch.sum(x, dim=axis, keepdim=keepdims) + + +def cumsum(x, axis=None): + return torch.cumsum(x, dim=-1 if axis is None else axis) + + +def sort(x, axis=-1): + return torch.sort(x, dim=axis).values + + +def topk(x, k, axis=-1): + """The ``k`` largest VALUES in ascending order, like ``mx.topk`` + (``sampling.py`` reads the threshold from ``[..., :1]``). Unlike + ``torch.topk`` there are no indices.""" + vals = torch.topk(x, k, dim=axis).values + return torch.sort(vals, dim=axis).values + + +def argpartition(x, kth, axis=-1): + """Full index permutation with every element at or before ``kth`` + (negative counts from the end) no larger than the rest, like + ``mx.argpartition``. A stable sort satisfies that for any ``kth``; + call sites slice ``[..., :k]`` / ``[..., -k:]``. O(n log n) is fine + at expert counts (<= 256).""" + return torch.argsort(x, dim=axis, stable=True) + + +def take(a, indices, axis=None): + """numpy/MLX ``take`` semantics (index shape composes with the + source's remaining axes), NOT ``torch.take``'s flat-index semantics. + """ + if not isinstance(indices, torch.Tensor): + indices = torch.as_tensor(indices, device=DEVICE) + if axis is None: + flat = a.reshape(-1) + return flat[indices.reshape(-1)].reshape(indices.shape) + axis %= a.ndim + idx_flat = indices.reshape(-1).to(torch.long) + out = torch.index_select(a, axis, idx_flat) + return out.reshape( + tuple(a.shape[:axis]) + tuple(indices.shape) + tuple(a.shape[axis + 1:])) + + +def _along_axis_index(a, indices, axis): + """numpy/MLX ``*_along_axis`` broadcast ``indices`` against ``a`` on + every axis except ``axis``; torch.gather/scatter do not.""" + shape = list(a.shape) + shape[axis] = indices.shape[axis] + return torch.broadcast_to(indices.to(torch.long), shape) + + +def take_along_axis(a, indices, axis): + return torch.gather(a, axis, _along_axis_index(a, indices, axis)) + + +def put_along_axis(a, indices, values, axis): + """Functional (out-of-place) scatter, matching MLX's immutable-array + semantics -- see the module docstring re: an in-place fast path. + Indices broadcast against ``a`` and ``values`` against the indices + (``moe/routing.py`` masks whole groups with ``[..., k, 1]`` indices and + a scalar -inf). + """ + idx = _along_axis_index(a, indices, axis) + src = torch.as_tensor(values, dtype=a.dtype, device=a.device) + return torch.scatter(a, axis, idx, torch.broadcast_to(src, idx.shape)) + + +def index_add(a, indices, values): + """``a.at[indices].add(values)`` in MLX: out-of-place, rows given by + ``indices`` along axis 0, duplicates accumulate.""" + return a.index_add(0, indices.to(torch.long), values.to(a.dtype)) + + +def size(x) -> int: + """Number of elements (``x.size`` in MLX; a method in torch).""" + return x.numel() + + +def max(x, axis=None, keepdims=False): + if axis is None: + return torch.amax(x) + return torch.amax(x, dim=axis, keepdim=keepdims) + + +def maximum(a, b): + a = torch.as_tensor(a, device=DEVICE) + return torch.maximum(a, torch.as_tensor(b, dtype=a.dtype, device=a.device)) + + +def argmax(x, axis=None, keepdims=False): + return torch.argmax(x, dim=axis, keepdim=keepdims) + + +def set_cache_limit(limit): + """MLX buffer-cache cap; torch's caching allocator has no equivalent + knob here, so this is a no-op that returns the previous limit (0).""" + return 0 + + +def astype(x, dtype): + return x.to(dtype) + + +def item(x): + return x.item() + + +def tolist(x): + return x.tolist() + + +# ---- lazy-graph hooks (MLX-specific; no-ops / thin wraps under eager torch) + +def eval(*arrays): + """MLX forces lazy-graph materialization here; torch is eager, so + there is nothing to force. Kept as a no-op rather than removed so + call sites (``streaming/layer.py``, ``prerouter/*``) need no + backend-conditional code. + """ + return None + + +def compile(fn): + """Identity: runs eagerly. ``mx.compile`` is cheap to re-trace, but + ``torch.compile`` recompiles on every new shape, and the streaming + path changes shapes nearly every call (staged vs. exact, varying + expert counts). Revisit with profiling on real hardware.""" + return fn + + +class random: + """``mx.random`` submodule equivalent (contract: seed, categorical).""" + + @staticmethod + def seed(s): + torch.manual_seed(s) + + @staticmethod + def categorical(logits, axis=-1): + # mx.random.categorical draws ONE sample per row from unnormalized + # logits along `axis`; torch.multinomial wants probabilities on + # the LAST axis with shape [..., num_classes] -> 1 sample each. + if axis != -1 and axis != logits.ndim - 1: + logits = torch.movedim(logits, axis, -1) + probs = torch.softmax(logits.float(), dim=-1) + shape = probs.shape[:-1] + flat = probs.reshape(-1, probs.shape[-1]) + draw = torch.multinomial(flat, 1).reshape(shape) + return draw diff --git a/python/src/edge0/backends/cuda/io.py b/python/src/edge0/backends/cuda/io.py new file mode 100644 index 0000000..b466493 --- /dev/null +++ b/python/src/edge0/backends/cuda/io.py @@ -0,0 +1,464 @@ +"""CUDA backend: model / tokenizer / tensor-store loading. + +Contract (documented): ``load_model, load_tokenizer, open_tensor_store``. +Two more names are used in practice by call sites that import +``edge0.backends.mlx.io`` DIRECTLY instead of going through the generic +facade -- ``load_safetensors`` (``prerouter/install.py``, +``adapters/lora.py``) and ``open_shards`` (``engine/qwen.py``). Those +call sites need editing to import from ``edge0.backends.io`` once a +backend is selected; until then this module still implements both so +the CUDA-side equivalents exist and are testable on their own. +""" + +from __future__ import annotations + +import json +import struct + +import numpy as np +import torch + +from edge0.backends.base import TensorStore +from edge0.backends.cuda.core import DEVICE, _as_tensor + +_DTYPES = { + "F64": torch.float64, "F32": torch.float32, "F16": torch.float16, + "BF16": torch.bfloat16, "I64": torch.int64, "I32": torch.int32, + "I16": torch.int16, "I8": torch.int8, "U8": torch.uint8, + "BOOL": torch.bool, +} +_NP_DTYPES = { # for the raw numpy view before the device transfer + "F64": np.float64, "F32": np.float32, "F16": np.float16, + "I64": np.int64, "I32": np.int32, "I16": np.int16, "I8": np.int8, + "U8": np.uint8, "BOOL": np.bool_, + "U32": np.uint32, # MLX-quantized payloads (packed codes) +} + + +class SafeTensorsStore(TensorStore): + """mmap-backed safetensors store; ``get`` returns a device tensor. + + Structurally identical to ``backends/mlx/io.py::SafeTensorsStore`` + (same header parsing) -- the only change is the tail of ``get``: + numpy view -> device tensor instead of numpy view -> mx.array. BF16 + has no numpy dtype on most builds, so it is read as raw uint16 and + bit-cast via torch (``.view(torch.bfloat16)``), same trick the MLX + version uses via mlx's native bfloat16. + """ + + def __init__(self, path: str): + import mmap + self._path = path + with open(path, "rb") as f: + header_len = struct.unpack(" str: + return self._path + + def keys(self) -> list[str]: + return list(self._keys) + + def metadata(self) -> dict: + return dict(self._metadata) + + def get(self, name: str): + e = self._entries.get(name) + if e is None: + raise KeyError(f"{self.path}: no tensor {name!r}") + buf = np.frombuffer(self._mm, dtype=np.uint8, count=e["size"], + offset=e["offset"]) + if e["dtype"] == "BF16": + u16 = np.frombuffer(buf, dtype=np.uint16, + count=int(np.prod(e["shape"]))) + t = torch.from_numpy(u16.copy()).view(torch.bfloat16) + t = t.reshape(e["shape"]) + else: + arr = np.frombuffer(buf, dtype=_NP_DTYPES[e["dtype"]], + count=int(np.prod(e["shape"]))) + t = torch.from_numpy(arr.copy()).reshape(e["shape"]) + return t.to(DEVICE, non_blocking=True) + + def close(self): + self._mm.close() + self._file.close() + + def __del__(self): + try: + self.close() + except Exception: + pass + + +def open_shards(model_dir: str) -> list: + """Same contract as ``backends/mlx/io.py::open_shards`` -- reuses + ``streaming.mmap.SafetensorsMmap`` as-is, since that module is + already backend-agnostic (pure ``mmap`` + ``numpy``).""" + import glob + import os + from edge0.streaming.mmap import SafetensorsMmap + shards = [] + for path in sorted(glob.glob(os.path.join( + os.fspath(model_dir), "model*.safetensors"))): + shards.append(SafetensorsMmap(path)) + if not shards: + raise FileNotFoundError( + f"no model*.safetensors shards under {model_dir}") + return shards + + +def load_safetensors(path: str, dtype=None) -> dict: + """Load every tensor of a (small) safetensors file as device tensors + (adapter / prerouter weight files).""" + store = SafeTensorsStore(path) + try: + out = {} + for name in store.keys(): + t = store.get(name) + if dtype is not None and t.dtype != dtype: + t = t.to(dtype) + out[name] = t + return out + finally: + store.close() + + +def load_tokenizer(model_path, *, trust_remote_code=False): + """Identical to the MLX backend's implementation -- this already + goes through ``transformers.AutoTokenizer`` with no MLX dependency, + so it is genuinely backend-agnostic; duplicated here rather than + imported cross-backend so ``edge0.backends.cuda`` never imports + ``edge0.backends.mlx`` (keeps the two backends independently + installable).""" + from transformers import AutoTokenizer + return AutoTokenizer.from_pretrained( + model_path, local_files_only=True, + trust_remote_code=trust_remote_code) + + +_EXPERT_KEY_MARKERS = (".mlp.experts.", ".switch_mlp.") +"""Tensor-name substrings for the quantized MoE expert weights (the ones +streaming/layer.py streams from disk on demand — never meant to be +resident). Matches both tiers' REAL on-disk key naming, checked against +the actual published checkpoints (safetensors index/header), not +assumed: edge0-35b uses ``...mlp.switch_mlp.{gate,up,down}_proj``, +edge0-8b uses ``...mlp.experts.{gate,up,down}_proj`` — both already +pre-stacked ``[num_experts, ...]`` tensors (``WeightLayout.SEPARATE``), +regardless of what either tier's REFERENCE implementation does with +experts in memory (transformers' Bailing class uses a per-expert +``nn.ModuleList`` internally, but that in-memory shape is irrelevant +here since these keys are never loaded into it — see ``load_model``). +""" + +_KEY_PREFIX_STRIP = { + # architectures-string -> checkpoint key prefix to strip before + # matching transformers' own state_dict names. Confirmed against + # the REAL published checkpoints' safetensors index/header: + # edge0-35b's keys all carry "language_model." (MLX's own + # post-sanitize naming, from _impl/qwen3_5_moe.py wrapping + # TextModel under self.language_model) -- transformers' + # Qwen3_5MoeForCausalLM has no such wrapper, so every dense weight + # would silently fail to match without this. edge0-8b's keys + # already match transformers.BailingMoeV3's own naming with no + # prefix at all -- confirmed via the same check, not assumed + # identical by default (see _resolve_model_class's docstring). + "Qwen3_5MoeForConditionalGeneration": "language_model.", +} + + +def _resolve_model_class(model_path, raw_config: dict, *, + trust_remote_code=False): + """``config.json`` (+ directory, for trust_remote_code) -> a + ``(constructed transformers config, model_cls, needs_trust_remote_code, + key_prefix)`` tuple. + + Dispatches on ``architectures``/``auto_map``, NOT ``model_type`` -- + checked against the real ``Edge0/Edge0-8B-A1B-preview`` config.json + (downloaded directly, not assumed): it has no top-level + ``model_type`` field at all. A first version of this function keyed + on ``model_type`` would have silently mis-dispatched on it. + + * edge0-35b: the REAL published config.json's ``architectures`` is + ``["Qwen3_5MoeForConditionalGeneration"]`` -- the vision+text + wrapper, not plain ``Qwen3_5MoeForCausalLM`` (a first version of + this function assumed the latter). Checked directly: it has both + ``text_config`` AND ``vision_config`` keys. edge0 strips vision + entirely (confirmed in ``_impl/qwen3_5_moe.py::sanitize()``), so + this still resolves to ``Qwen3_5MoeForCausalLM`` (config class + ``Qwen3_5MoeTextConfig``) built from JUST ``config["text_config"]`` + -- the vision-wrapping class is never constructed. Every + checkpoint key carries a ``language_model.`` prefix (confirmed via + the real safetensors index) that the plain text-only class's own + attribute names don't have -- stripped via ``_KEY_PREFIX_STRIP``. + * edge0-8b (``BailingMoeV3ForCausalLM``): ships its OWN + ``modeling_bailing_moe_v3.py``/``configuration_bailing_moe_v3.py`` + co-located in the checkpoint repo (confirmed via the HF API file + listing, then downloaded and actually imported -- needs einops, + fla, triton as extra deps, not currently in edge0's own + dependency list). Experts are ``nn.ModuleList`` of per-expert MLP + modules in the reference implementation's OWN in-memory layout, + but the actual PUBLISHED CHECKPOINT stores them pre-stacked + (confirmed via the real safetensors header) -- the + ``nn.ModuleList`` is never populated from checkpoint data, it + gets swapped out by the streaming installer instead. No key + prefix needed -- confirmed via the same check, not assumed + identical to the 35b tier by default. + SECURITY NOTE, not a footnote: ``trust_remote_code=True`` runs + third-party Python shipped inside the checkpoint directory. That + is a real code-execution surface, not a formality -- worth a + deliberate decision (pin+review the exact modeling file once, + vendor it, or accept the risk per-checkpoint) before this path + is used unattended, e.g. in `edge0 serve`. + """ + archs = raw_config.get("architectures", []) + auto_map = raw_config.get("auto_map", {}) + + if any(a.startswith("Qwen3_5Moe") for a in archs): + from transformers import Qwen3_5MoeForCausalLM, Qwen3_5MoeTextConfig + text_cfg_dict = raw_config.get("text_config", raw_config) + config = Qwen3_5MoeTextConfig(**text_cfg_dict) + key_prefix = next( + (p for a, p in _KEY_PREFIX_STRIP.items() if a in archs), "") + return config, Qwen3_5MoeForCausalLM, False, key_prefix + + if any("Bailing" in a for a in archs) or "bailing" in str(auto_map).lower(): + if not trust_remote_code: + raise ValueError( + "this checkpoint requires remote model code; review its " + "Python files, then pass trust_remote_code=True explicitly") + from transformers import AutoConfig, AutoModelForCausalLM + config = AutoConfig.from_pretrained( + model_path, trust_remote_code=trust_remote_code) + return config, AutoModelForCausalLM, True, "" + + raise NotImplementedError( + f"cuda backend: unrecognized architectures={archs!r} -- no known " + f"transformers class for it (see this function's docstring for " + f"the two paths currently resolved)") + + +_QWEN35_SHIFTED_NORMS = ( + ".input_layernorm.weight", ".post_attention_layernorm.weight", + "model.norm.weight", ".q_norm.weight", ".k_norm.weight") +"""Norms that edge0's MLX ``qwen3_5.py::sanitize()`` stores as ``w + 1``: +transformers' Qwen3.5 RMSNorm computes ``x * (1 + w)``, the MLX one +``x * w``. Checked on the published edge0-35b bytes: input_layernorm mean +1.03, q_norm 1.33, model.norm 2.63 -- versus ~0 for a transformers +checkpoint. The gated ``linear_attn.norm`` is not shifted (mean 0.88).""" + + +def _undo_mlx_qwen35_sanitize(state: dict) -> None: + """Invert edge0's MLX Qwen3.5 sanitize on a loaded state dict, in + place -- only if the checkpoint went through it, detected the way the + sanitize itself detects the opposite: conv1d weights in MLX layout + ``[C, k, 1]`` instead of torch's ``[C, 1, k]``.""" + convs = [k for k in state if k.endswith("conv1d.weight")] + if not convs or not all(state[k].shape[-1] == 1 and state[k].shape[1] != 1 + for k in convs): + return + for k in convs: + state[k] = state[k].transpose(1, 2).contiguous() + for k in list(state): + if k.endswith(_QWEN35_SHIFTED_NORMS) and state[k].ndim == 1: + state[k] = state[k] - 1.0 + + +def _read_tensor(shard, meta, name): + import torch + raw = shard.raw(name) + count = int(np.prod(meta["shape"])) + if meta["dtype"] == "BF16": + u16 = np.frombuffer(raw, dtype=np.uint16, count=count) + t = torch.from_numpy(u16.copy()).view(torch.bfloat16) + else: + arr = np.frombuffer(raw, dtype=_NP_DTYPES[meta["dtype"]], count=count) + t = torch.from_numpy(arr.copy()) + return t.reshape(meta["shape"]) + + +def _params_on_meta(): + """Build a module with its PARAMETERS on the meta device (no memory) + but buffers real: non-persistent buffers such as rotary ``inv_freq`` + are computed at init and never come from the checkpoint, so building + everything under ``torch.device('meta')`` would leave them unusable.""" + import contextlib + + import torch + + @contextlib.contextmanager + def ctx(): + orig = torch.nn.Module.register_parameter + + def register_parameter(self, name, param): + if param is not None and param.device.type != "meta": + param = torch.nn.Parameter(param.to("meta"), + requires_grad=param.requires_grad) + orig(self, name, param) + + torch.nn.Module.register_parameter = register_parameter + try: + yield + finally: + torch.nn.Module.register_parameter = orig + return ctx() + + +def _install_quantized(model, state: dict, dtype) -> set: + """Replace every module whose weight is MLX-quantized in ``state`` + (``.scales`` present) with the torch equivalent, consuming the + three tensors. Linear / Embedding get quantized modules; anything else + (e.g. transformers' Qwen3.5 router, a custom module holding a plain + ``weight`` parameter) gets its weight dequantized in place. Returns the + module paths replaced (their tensors are already in place).""" + import torch + + from edge0.backends.cuda import nn as cnn + from edge0.backends.cuda.quant import _dequantize + + replaced = set() + for skey in [k for k in state if k.endswith(".scales")]: + path = skey[: -len(".scales")] + try: + mod = model.get_submodule(path) + except AttributeError: + continue # not part of this model: stays unexpected + w = state.pop(f"{path}.weight") + s = state.pop(skey) + b = state.pop(f"{path}.biases") + owner, _, attr = path.rpartition(".") + parent = model.get_submodule(owner) if owner else model + if isinstance(mod, torch.nn.Linear): + bias = state.pop(f"{path}.bias", None) + new = cnn.QuantizedLinear(w, s, b, mod.in_features, bias=bias) + elif isinstance(mod, torch.nn.Embedding): + new = cnn.QuantizedEmbedding(w, s, b, mod.embedding_dim, + dtype=dtype) + else: + in_features = mod.weight.shape[-1] + bits, group_size = cnn._quant_params(w, s, in_features) + dense = _dequantize(w, s, b, group_size, bits) + state[f"{path}.weight"] = dense.to(dtype or s.dtype) + continue + setattr(parent, attr, new.to(DEVICE)) + replaced.add(path) + return replaced + + +def load_model(model_path, lazy=True, strict=False, model_config=None, + get_model_classes=None, dtype=None, + trust_remote_code=False): + """Load an edge0 checkpoint into its transformers model class. + + The published checkpoints are MLX checkpoints: quantized throughout + (embeddings, lm_head, attention, shared experts, routers -- 392 + non-expert tensors for edge0-35b, 236 for edge0-8b) and, for edge0-35b, + run through MLX's sanitize. So: + + * modules whose weight has ``.scales`` become ``QuantizedLinear`` / + ``QuantizedEmbedding`` (4-bit payload stays resident); other + quantized params are dequantized in place; + * the MLX sanitize is undone where it was applied (Qwen3.5: conv1d + layout, ``w + 1`` norms -- see ``_undo_mlx_qwen35_sanitize``); + * routed-expert tensors (``_EXPERT_KEY_MARKERS``) are not loaded: the + streaming installer serves them from disk, and those parameters stay + on the meta device until it replaces them. + + ``dtype`` casts the dense floating-point tensors (default: as stored, + bf16). ``strict`` raises on non-expert parameters the checkpoint does + not provide and on checkpoint tensors the model has no place for. + ``trust_remote_code`` explicitly enables checkpoint-supplied Python on + the generic Transformers path. The engine path uses Edge0's own model + classes and does not load checkpoint-supplied Python. + """ + import json + import os + + import torch + + with open(os.path.join(os.fspath(model_path), "config.json")) as f: + raw_config = json.load(f) + + engine_path = get_model_classes is not None + if engine_path: + # The engines' path, same contract as mlx-lm's load_model: a + # vendored (Model, ModelArgs) pair built from config.json plus the + # model_config overrides; returns (model, config). + config = {**raw_config, **(model_config or {})} + model_cls, args_cls = get_model_classes(config=config) # as mlx-lm calls it + with _params_on_meta(): + model = model_cls(args_cls.from_dict(config)) + key_prefix = "" + else: + hf_config, model_cls, needs_trust_remote_code, key_prefix = \ + _resolve_model_class( + model_path, raw_config, + trust_remote_code=trust_remote_code) + with _params_on_meta(): + model = model_cls.from_config( + hf_config, trust_remote_code=trust_remote_code) \ + if needs_trust_remote_code else model_cls(hf_config) + + state = {} + skipped_expert_keys = [] + for shard in open_shards(model_path): + for name, meta in shard.entries.items(): + if any(m in name for m in _EXPERT_KEY_MARKERS): + skipped_expert_keys.append(name) + continue + mapped = name[len(key_prefix):] if key_prefix and \ + name.startswith(key_prefix) else name + state[mapped] = _read_tensor(shard, meta, name) + shard.close() + + if model_cls.__name__.startswith("Qwen3_5Moe"): + _undo_mlx_qwen35_sanitize(state) + if engine_path and hasattr(model, "sanitize"): + state = model.sanitize(state) + quantized = _install_quantized(model, state, dtype) + for k, t in state.items(): + if dtype is not None and t.is_floating_point(): + state[k] = t.to(dtype) + state[k] = state[k].to(DEVICE) + + missing, unexpected = model.load_state_dict(state, strict=False, + assign=True) + # Buffers computed at init and absent from the checkpoint (the rotary + # inv_freq) were built on the host; the checkpoint tensors went to + # DEVICE above. Not model.to(DEVICE): the never-loaded (streamed) + # expert parameters are still on meta and cannot be copied. + for mod in model.modules(): + for name, buf in mod._buffers.items(): + if buf is not None and not buf.is_meta and buf.device != DEVICE: + mod._buffers[name] = buf.to(DEVICE) + missing = [k for k in missing + if not any(m in k for m in _EXPERT_KEY_MARKERS) + and k.rpartition(".")[0] not in quantized] + if strict and (missing or unexpected): + raise RuntimeError( + f"load_model: {len(missing)} missing, {len(unexpected)} " + f"unexpected non-expert tensors: missing={missing[:5]} " + f"unexpected={unexpected[:5]}") + model.eval() + # MLX arrays carry no autograd state. Torch parameters default to + # requires_grad=True, and then every forward keeps its whole graph -- + # each dequantized expert included -- alive until the logits are + # dropped: on edge0-8b a 31-token prefill grew 1.2 GB/s past 120 GB. + model.requires_grad_(False) + model._edge0_skipped_expert_keys = skipped_expert_keys # for the streaming hook + model._edge0_load_report = {"missing": missing, "unexpected": unexpected} + return (model, config) if engine_path else model diff --git a/python/src/edge0/backends/cuda/model_specs.py b/python/src/edge0/backends/cuda/model_specs.py new file mode 100644 index 0000000..9f89124 --- /dev/null +++ b/python/src/edge0/backends/cuda/model_specs.py @@ -0,0 +1,47 @@ +"""MoESpec instances for the CUDA backend's two shipped tiers. + +Values are the SAME family facts already verified and shipped in +``models/edge0_35b/__init__.py`` / ``models/edge0_8b/__init__.py`` +(num_experts, top_k, router kind, quantization, ...) -- not +re-derived, since those are already checked against the real +checkpoints. What differs here is ``block_path``: the MLX specs' path +navigates MLX's OWN post-sanitize module tree (``language_model.model. +layers.{layer}.mlp``, from ``_impl/qwen3_5_moe.py`` wrapping +``TextModel`` under ``self.language_model``); ``transformers. +Qwen3_5MoeForCausalLM`` has no such wrapper, so its real attribute path +is one level shallower. ``key_template`` (which resolves SAFETENSORS +KEYS, not Python attributes) keeps the ``language_model.`` prefix, +because that prefix IS present in the real on-disk checkpoint -- +confirmed via the published ``model.safetensors.index.json``, not +assumed. edge0-8b needs no such split: its on-disk keys and +``transformers.BailingMoeV3ForCausalLM``'s attribute path already +agree (confirmed the same way), so its spec is byte-for-byte what +``models/edge0_8b/__init__.py`` already uses. +""" + +from __future__ import annotations + +from edge0.moe.spec import MoESpec, QuantSpec, RouterKind, WeightLayout + +QWEN35_MOE_SPEC = MoESpec( + num_experts=256, top_k=4, intermediate_size=512, + router=RouterKind.SOFTMAX_TOPK, norm_topk_prob=True, + shared_experts=1, + quant=QuantSpec(bits=4, group_size=64, mode="affine"), + layout=WeightLayout.SEPARATE, + key_template="language_model.model.layers.{layer}.mlp.switch_mlp", + block_path="model.layers.{layer}.mlp", + layer_path="model.layers.{layer}", +) + +BAILING_V3_MOE_SPEC = MoESpec( + num_experts=128, top_k=8, intermediate_size=512, + router=RouterKind.SIGMOID_GROUP, norm_topk_prob=True, + routed_scaling=2.5, n_group=8, topk_group=4, + shared_experts=1, + quant=QuantSpec(bits=4, group_size=64, mode="affine"), + layout=WeightLayout.SEPARATE, + key_template="model.layers.{layer}.mlp.experts", + block_path="model.layers.{layer}.mlp", + layer_path="model.layers.{layer}", +) diff --git a/python/src/edge0/backends/cuda/moe_blocks.py b/python/src/edge0/backends/cuda/moe_blocks.py new file mode 100644 index 0000000..b23f6a0 --- /dev/null +++ b/python/src/edge0/backends/cuda/moe_blocks.py @@ -0,0 +1,31 @@ +"""CUDA backend: adapters between StreamingSwitchGLU and transformers' MoE +blocks. + +The vendored MLX models call ``switch_mlp(x, inds)`` and weight the +per-expert outputs themselves -- the streaming twin's native interface. +transformers' blocks call their experts differently, so the twin is +wrapped before ``install_streaming_experts`` swaps it in (its ``wrap=``). +""" + +from __future__ import annotations + +import torch + + +class TransformersExpertsAdapter(torch.nn.Module): + """``experts(hidden [T, H], top_k_index [T, K], top_k_weights [T, K]) + -> [T, H]``, as ``Qwen3_5MoeSparseMoeBlock`` calls it, on top of a + twin that returns the raw per-expert outputs ``[T, K, H]``. + + Not usable for blocks that index or iterate ``self.experts`` as a + ``ModuleList`` (the edge0-8b checkpoint's own ``modeling_bailing_moe_v3.py`` + does): those need their forward replaced, not their experts. + """ + + def __init__(self, twin): + super().__init__() + object.__setattr__(self, "twin", twin) # not an nn.Module + + def forward(self, hidden_states, top_k_index, top_k_weights): + y = self.twin(hidden_states, top_k_index) # [T, K, H] + return (y * top_k_weights[..., None].to(y.dtype)).sum(dim=-2) diff --git a/python/src/edge0/backends/cuda/nn.py b/python/src/edge0/backends/cuda/nn.py new file mode 100644 index 0000000..1134ee9 --- /dev/null +++ b/python/src/edge0/backends/cuda/nn.py @@ -0,0 +1,148 @@ +"""CUDA backend: module-factory namespace (contract: Module, Linear, +RMSNorm, silu, gelu). + +``torch.nn`` already provides all four natively (``nn.Linear``, +``nn.functional.silu``/``gelu``); the only real gap is ``RMSNorm``, +which is standard ``torch.nn.RMSNorm`` since torch 2.4 -- re-exported +here under the contract's names rather than assuming call sites import +``torch.nn`` directly (see ``backends/__init__.py``'s enforcement that +framework code only ever imports ``edge0.backends.{core,nn,io,quant}``). +""" + +from __future__ import annotations + +import os + +import torch +import torch.nn as _tnn +import torch.nn.functional as F + +Module = _tnn.Module + + +class Linear(_tnn.Linear): + """``torch.nn.Linear`` that, like ``mlx.nn.Linear``, accepts a plain + tensor assigned to ``weight`` / ``bias`` (``prerouter/install.py`` does + ``head.fc1.weight = w``); torch itself insists on a Parameter. + + Parameters never require grad, as MLX arrays carry no autograd state: + modules built after ``load_model`` (the prerouter heads) would otherwise + make every forward that touches them record a graph.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.requires_grad_(False) + + def __setattr__(self, name, value): + if (name in ("weight", "bias") and isinstance(value, torch.Tensor) + and not isinstance(value, _tnn.Parameter)): + value = _tnn.Parameter(value, requires_grad=False) + super().__setattr__(name, value) + + +class RMSNorm(_tnn.RMSNorm): + """``mlx.nn.RMSNorm(dims, eps)`` signature on top of ``torch.nn.RMSNorm``. + + Subclassed rather than wrapped so the parameter keeps the name + ``weight``, as in MLX and the checkpoints. Numerics match + ``mx.fast.rms_norm`` (eps inside the sqrt); see + ``tests/test_cuda_backend.py``. + """ + + def __init__(self, dims: int, eps: float = 1e-5): + super().__init__(dims, eps=eps) + self.requires_grad_(False) + + +def _quant_params(weight, scales, in_features): + """(bits, group_size) of an MLX affine-quantized tensor, from shapes: + packed width = in * bits / 32, scales width = in / group_size. Works + for per-path overrides (the edge0-35b routers are 8-bit).""" + return (weight.shape[-1] * 32 // in_features, + in_features // scales.shape[-1]) + + +#: Dequantized weights kept per module instead of redone every call. +#: Off by default: the whole point of the quantized layers is that the +#: 4-bit payload is what stays resident. Worth turning on where memory is +#: plentiful: on a GB10 it takes the edge0-8b decode step from 206 ms to +#: 124 ms a token, for 4.1 GB more resident (EDGE0_TORCH_WEIGHT_CACHE=1). +CACHE_DEQUANTIZED = os.environ.get("EDGE0_TORCH_WEIGHT_CACHE", "") == "1" + + +class QuantizedLinear(_tnn.Module): + """``mlx.nn.QuantizedLinear`` layout (``weight`` packed uint32, + ``scales``, ``biases``) dequantized on the fly -- the 4-bit payload + stays resident, not a bf16 copy (unless ``CACHE_DEQUANTIZED``).""" + + def __init__(self, weight, scales, biases, in_features, bias=None): + super().__init__() + self.in_features = in_features + self.out_features = weight.shape[0] + self.bits, self.group_size = _quant_params(weight, scales, in_features) + self.register_buffer("weight", weight) + self.register_buffer("scales", scales) + self.register_buffer("biases", biases) + self.bias = None if bias is None else _tnn.Parameter(bias, False) + self._dequantized = None # (dtype, weight), see forward + + ROWS_PER_CHUNK = 4096 + + def forward(self, x): + """Dequantize ``ROWS_PER_CHUNK`` output rows at a time: the full + float weight of a large projection is never materialized (edge0-8b's + lm_head alone would be ~1 GB per call). Same arithmetic per + element as dequantizing everything first.""" + from edge0.backends.cuda.quant import _dequantize + if CACHE_DEQUANTIZED: + # Same weight every call: dequantize once. The chunking below + # exists to bound the transient, which a kept weight makes moot. + if self._dequantized is None or self._dequantized[0] != x.dtype: + w = _dequantize(self.weight, self.scales, self.biases, + self.group_size, self.bits).to(x.dtype) + self._dequantized = (x.dtype, w) + b = None if self.bias is None else self.bias.to(x.dtype) + return F.linear(x, self._dequantized[1], b) + outs = [] + for r in range(0, self.out_features, self.ROWS_PER_CHUNK): + sl = slice(r, r + self.ROWS_PER_CHUNK) + w = _dequantize(self.weight[sl], self.scales[sl], self.biases[sl], + self.group_size, self.bits).to(x.dtype) + b = None if self.bias is None else self.bias[sl].to(x.dtype) + outs.append(F.linear(x, w, b)) + return outs[0] if len(outs) == 1 else torch.cat(outs, dim=-1) + + +class QuantizedEmbedding(_tnn.Module): + """``mlx.nn.QuantizedEmbedding`` layout; only the looked-up rows are + dequantized. Output dtype: ``dtype``, else the scales' (the + checkpoint's).""" + + def __init__(self, weight, scales, biases, embedding_dim, dtype=None): + super().__init__() + self.num_embeddings = weight.shape[0] + self.embedding_dim = embedding_dim + self.bits, self.group_size = _quant_params(weight, scales, + embedding_dim) + self.out_dtype = dtype or scales.dtype + self.register_buffer("weight", weight) + self.register_buffer("scales", scales) + self.register_buffer("biases", biases) + + def forward(self, ids): + from edge0.backends.cuda.quant import _dequantize + # Gather on the int32 view of the packed rows: CUDA has no index + # kernel for uint32 ("index_cuda not implemented for UInt32"), and + # _dequantize reads the words as int32 anyway. + rows = _dequantize(self.weight.view(torch.int32)[ids], + self.scales[ids], + self.biases[ids], self.group_size, self.bits) + return rows.to(self.out_dtype) + + +def silu(x): + return F.silu(x) + + +def gelu(x): + return F.gelu(x) diff --git a/python/src/edge0/backends/cuda/quant.py b/python/src/edge0/backends/cuda/quant.py new file mode 100644 index 0000000..c7a8ff6 --- /dev/null +++ b/python/src/edge0/backends/cuda/quant.py @@ -0,0 +1,124 @@ +"""CUDA backend: quantized gather kernels (torch reference implementation). + +Semantics match ``mx.gather_qmm`` as checked against real MLX 0.30.4 +(Metal, the version this repo pins) in ``tests/test_cuda_backend.py``: + +* Packing: each uint32 word holds ``32 // bits`` codes, least significant + bits first, along the last (input) axis. Codes are unsigned; a group of + ``group_size`` codes dequantizes as ``w = code * scale + bias``. Scales + can be negative -- nothing here assumes otherwise. +* Broadcasting: the output batch shape is + ``broadcast(x.shape[:-2], rhs_indices.shape)`` and + ``out[b] = x[b] @ W[rhs_indices[b]].T``. Both call patterns in + ``streaming/layer.py`` rely on this: the unsorted path passes + ``x[..., 1, 1, D]`` against ``rhs_indices[..., K]``, the sorted path + passes ``x[T*K, 1, D]`` against ``rhs_indices[T*K]``. +* ``sorted_indices`` is a kernel hint in MLX; it never changes the values. + Ignored here. + +Every call still dequantizes the experts it touches; what this does avoid +is paying a device->host sync per expert to find out which those are (see +``gather_qmm``). +""" + +from __future__ import annotations + +import os + +import torch + +# Calls with at most this many (token, expert) rows dequantize the gathered +# experts as ONE batch, with no host sync at all -- decode (1 token, K +# experts per layer) lands here. Above it the per-expert loop is worth its +# syncs, because a dequantized copy per row is what blows up memory during +# prefill. Overridable for measurement. +BATCH_ROWS = int(os.environ.get("EDGE0_TORCH_GATHER_BATCH", "64")) +# ...and never more than this many bytes of dequantized float32 weight in +# flight, because one row of a 35b expert is far bigger than an 8b one. +BATCH_BYTES = int(os.environ.get("EDGE0_TORCH_GATHER_BYTES", str(256 << 20))) + + +def _dequantize(w: torch.Tensor, scales: torch.Tensor, biases: torch.Tensor, + group_size: int, bits: int) -> torch.Tensor: + """Packed ``[..., rows, in * bits / 32]`` uint32 -> float32 ``[..., rows, in]``.""" + per_word = 32 // bits + # int32 is enough: >> sign-extends, but the mask keeps only the low + # ``bits`` bits, which the sign bits never reach (int64 doubled the + # transient memory -- 1.9 GB of codes for edge0-8b's lm_head). + words = w.view(torch.int32) + shifts = torch.arange(per_word, device=w.device, dtype=torch.int32) * bits + codes = (words.unsqueeze(-1) >> shifts) & ((1 << bits) - 1) + codes = codes.reshape(*w.shape[:-1], w.shape[-1] * per_word) + grouped = codes.reshape(*codes.shape[:-1], -1, group_size).to(torch.float32) + deq = (grouped * scales.to(torch.float32).unsqueeze(-1) + + biases.to(torch.float32).unsqueeze(-1)) + return deq.reshape(codes.shape) + + +def gather_qmm(x, w, scales, biases, rhs_indices, transpose=True, + group_size=64, bits=4, mode="affine", + sorted_indices=False): + """Quantized matmul over a gathered subset of experts (``mx.gather_qmm``).""" + if mode != "affine" or bits not in (2, 4, 8): + raise NotImplementedError( + f"reference gather_qmm covers affine 2/4/8-bit only " + f"(got mode={mode!r}, bits={bits!r})") + idx = rhs_indices.to(torch.long) + bshape = torch.broadcast_shapes(x.shape[:-2], idx.shape) + M = x.shape[-2] + xf = x.to(torch.float32).expand(*bshape, *x.shape[-2:]).reshape(-1, M, x.shape[-1]) + flat = idx.expand(bshape).reshape(-1) + n_out = w.shape[-2] if transpose else w.shape[-1] * (32 // bits) + + rows_n = flat.numel() + deq_bytes = w[0].numel() * (32 // bits) * 4 + if rows_n <= BATCH_ROWS and rows_n * deq_bytes <= BATCH_BYTES: + # Decode-sized call. Asking WHICH experts are needed + # (unique/tolist) and WHERE each one's rows are (nonzero) costs a + # device->host sync each, and at this size those syncs dwarf the + # matmul -- measured 4.4 s of an 8.2 s edge0-8b decode on Metal, + # against 0.07 s for the matmuls themselves. Gather the rows' + # experts as one batch instead and run a single bmm: no sync, and + # a repeated expert only costs a duplicate dequantize. + deq = _dequantize(w[flat], scales[flat], biases[flat], group_size, bits) + out = torch.bmm(xf, deq.transpose(-1, -2) if transpose else deq) + return out.reshape(*bshape, M, n_out).to(x.dtype) + + out = torch.empty(flat.numel(), M, n_out, dtype=torch.float32, device=x.device) + # One distinct expert at a time: dequantizing a copy per (token, expert) + # pair peaked at several GB per MoE layer during prefill. An index past + # the last expert raises here; in MLX it silently reads out of bounds. + # ONE host transfer of the whole index list, then the row groups are + # built on the host -- a nonzero() per expert is a sync per expert. + groups: dict[int, list[int]] = {} + for row, e in enumerate(flat.tolist()): + groups.setdefault(e, []).append(row) + for e, rows_list in groups.items(): + rows = torch.as_tensor(rows_list, dtype=torch.long, device=x.device) + deq = _dequantize(w[e], scales[e], biases[e], group_size, bits) + out[rows] = torch.matmul(xf[rows], deq.T if transpose else deq) + return out.reshape(*bshape, M, n_out).to(x.dtype) + + +def gather_sort(x, indices): + """Sort token rows by expert id for ``gather_qmm(sorted_indices=True)``: + returns ``(x_sorted, indices_sorted, inv_order)``, as mlx-lm's + ``_gather_sort`` does.""" + m = indices.shape[-1] + flat = indices.reshape(-1).to(torch.long) + order = torch.argsort(flat, stable=True) + inv_order = torch.argsort(order) + return x.flatten(0, -3)[order // m], flat[order], inv_order + + +def scatter_unsort(x, inv_order, shape=None): + """Undo ``gather_sort``; ``shape`` re-splits the leading axis.""" + x = x[inv_order] + if shape is not None: + x = x.unflatten(0, tuple(shape)) + return x + + +def swiglu(up: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """SiLU gated activation: silu(gate) * up.""" + return torch.nn.functional.silu(gate) * up diff --git a/python/src/edge0/backends/mlx/backend.py b/python/src/edge0/backends/mlx/backend.py index d4f52b7..e84785f 100644 --- a/python/src/edge0/backends/mlx/backend.py +++ b/python/src/edge0/backends/mlx/backend.py @@ -1,16 +1,16 @@ """MLX backend: namespace assembly. -``core`` is mlx.core (arrays, ops, eval/compile, random), ``nn`` is -mlx.nn (module factory), ``io`` and ``quant`` are edge0 wrappers around -mlx-lm and mlx's quantized gather kernels. +``core`` is mlx.core (arrays, ops, eval/compile, random) plus the few +contract names MLX only has as array methods, ``nn`` is mlx.nn (module +factory), ``io`` and ``quant`` are edge0 wrappers around mlx-lm and mlx's +quantized gather kernels. """ from __future__ import annotations -import mlx.core as core # noqa: F401 import mlx.nn as nn # noqa: F401 -from edge0.backends.mlx import io, quant # noqa: F401 +from edge0.backends.mlx import core, io, quant # noqa: F401 class BackendImpl: diff --git a/python/src/edge0/backends/mlx/core.py b/python/src/edge0/backends/mlx/core.py new file mode 100644 index 0000000..659edbb --- /dev/null +++ b/python/src/edge0/backends/mlx/core.py @@ -0,0 +1,27 @@ +"""MLX backend: array ops namespace. + +Everything from ``mlx.core`` (same objects, so ``core.array``, +``core.compile`` ... are the real MLX ones), plus the few contract names +that MLX only offers as array methods or indexing syntax. Framework code +calls these instead of the method forms so it also runs on backends whose +arrays lack them (torch has no ``.astype`` or ``.at[]``, and its ``.size`` +is a method). +""" + +from __future__ import annotations + +from mlx.core import * # noqa: F401,F403 + + +def astype(x, dtype): + return x.astype(dtype) + + +def size(x) -> int: + """Number of elements (``x.size`` in MLX).""" + return x.size + + +def index_add(a, indices, values): + """``a.at[indices].add(values)``: out-of-place, duplicates accumulate.""" + return a.at[indices].add(values) diff --git a/python/src/edge0/backends/mlx/io.py b/python/src/edge0/backends/mlx/io.py index 242a4ea..363e636 100644 --- a/python/src/edge0/backends/mlx/io.py +++ b/python/src/edge0/backends/mlx/io.py @@ -33,7 +33,7 @@ def load_model(model_path, lazy=True, strict=False, model_config=None, get_model_classes=get_model_classes) -def load_tokenizer(model_path): +def load_tokenizer(model_path, *, trust_remote_code=False): """Load a HuggingFace tokenizer from a LOCAL directory only. mlx-lm 0.31's ``load_tokenizer`` funnels the path through its @@ -44,7 +44,8 @@ def load_tokenizer(model_path): """ from transformers import AutoTokenizer return AutoTokenizer.from_pretrained( - model_path, local_files_only=True, trust_remote_code=True) + model_path, local_files_only=True, + trust_remote_code=trust_remote_code) def open_shards(model_dir: str) -> list: diff --git a/python/src/edge0/backends/mlx/quant.py b/python/src/edge0/backends/mlx/quant.py index b1ca11d..b5a3b2c 100644 --- a/python/src/edge0/backends/mlx/quant.py +++ b/python/src/edge0/backends/mlx/quant.py @@ -3,6 +3,18 @@ from __future__ import annotations import mlx.core as mx +from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort + + +def gather_sort(x, indices): + """Sort token rows by expert id for ``gather_qmm(sorted_indices=True)``: + returns ``(x_sorted, indices_sorted, inv_order)`` (mlx-lm's helper).""" + return _gather_sort(x, indices) + + +def scatter_unsort(x, inv_order, shape=None): + """Undo ``gather_sort``; ``shape`` re-splits the leading axis.""" + return _scatter_unsort(x, inv_order, shape) def gather_qmm(x, w, scales, biases, rhs_indices, transpose=True, diff --git a/python/src/edge0/engine/base.py b/python/src/edge0/engine/base.py index c54ef16..6b9936f 100644 --- a/python/src/edge0/engine/base.py +++ b/python/src/edge0/engine/base.py @@ -27,6 +27,24 @@ from edge0.sampling import sample +def require_backend(tier: str, supported=("mlx",)) -> None: + """The engines drive a vendored model directly (per-layer callbacks, + caches, prerouter wiring), so a tier runs only on backends that ship + that model. Fail early and plainly elsewhere instead of pushing one + backend's arrays through another's ops.""" + from edge0.backends import backend + if backend.name not in supported: + raise NotImplementedError( + f"{tier}: the engine runs only on EDGE0_BACKEND=" + f"{' or '.join(supported)} today. The {backend.name} backend " + f"has the array/quant ops, the streaming layer and load_model, " + f"but no model port for this tier yet -- see docs/nvidia.md.") + + +def require_mlx_backend(tier: str) -> None: + require_backend(tier, ("mlx",)) + + class Edge0Engine: """Base streaming engine. Subclasses fill in the family hooks.""" diff --git a/python/src/edge0/engine/ling.py b/python/src/edge0/engine/ling.py index 936803c..4ab9a74 100644 --- a/python/src/edge0/engine/ling.py +++ b/python/src/edge0/engine/ling.py @@ -21,11 +21,8 @@ from edge0.backends import core -from edge0.backends.mlx._impl.bailing_hybrid import Model as BailingModel -from edge0.backends.mlx._impl.bailing_hybrid import ModelArgs as BailingArgs from edge0.backends import io -from edge0.backends.mlx.io import load_model, load_tokenizer -from edge0.engine.base import Edge0Engine +from edge0.engine.base import Edge0Engine, require_backend from edge0.engine.hooks import ( make_history_prefetch, make_prefill_before_layer, @@ -37,8 +34,14 @@ def _get_model_classes(config): - """mlx-lm class hook: serve the vendored bailing backbone.""" - return BailingModel, BailingArgs + """load_model class hook: serve the active backend's port of the bailing + backbone (imported here so the module itself imports without MLX).""" + from edge0.backends import backend + if backend.name == "cuda": + from edge0.backends.cuda._impl.bailing_hybrid import Model, ModelArgs + else: + from edge0.backends.mlx._impl.bailing_hybrid import Model, ModelArgs + return Model, ModelArgs def load_installed(model_dir: str, cfg): @@ -47,7 +50,8 @@ def load_installed(model_dir: str, cfg): Returns ``(model, model_config, shards, installs)``. """ - model, model_config = load_model( + require_backend("edge0-8b", ("mlx", "cuda")) + model, model_config = io.load_model( model_dir, lazy=True, strict=False, model_config={"model_type": "bailing_hybrid", "prerouter_enabled": cfg.prerouter is not None, @@ -145,7 +149,7 @@ def _build(self): opts = cfg.options if self._tok is None: try: - self._tok = load_tokenizer(cfg.model_dir) + self._tok = io.load_tokenizer(cfg.model_dir) except Exception: # noqa: BLE001 — tokenizer optional for CLI pass diff --git a/python/src/edge0/engine/qwen.py b/python/src/edge0/engine/qwen.py index b22a107..fc0618f 100644 --- a/python/src/edge0/engine/qwen.py +++ b/python/src/edge0/engine/qwen.py @@ -18,11 +18,8 @@ from edge0.backends import core -from edge0.backends.mlx._impl.qwen3_5_moe import Model as Qwen35Model -from edge0.backends.mlx._impl.qwen3_5_moe import ModelArgs as Qwen35Args from edge0.backends import io -from edge0.backends.mlx.io import load_model, load_tokenizer, open_shards -from edge0.engine.base import Edge0Engine +from edge0.engine.base import Edge0Engine, require_backend from edge0.engine.hooks import ( make_history_prefetch, make_intra_after_layer, @@ -34,8 +31,15 @@ def _get_model_classes(config): - """mlx-lm class hook: serve the vendored qwen3_5_moe backbone.""" - return Qwen35Model, Qwen35Args + """load_model class hook: serve the active backend's port of the + qwen3_5_moe backbone (imported here so the module itself imports + without MLX).""" + from edge0.backends import backend + if backend.name == "cuda": + from edge0.backends.cuda._impl.qwen3_5_moe import Model, ModelArgs + else: + from edge0.backends.mlx._impl.qwen3_5_moe import Model, ModelArgs + return Model, ModelArgs def load_installed(model_dir: str, cfg): @@ -46,11 +50,12 @@ def load_installed(model_dir: str, cfg): ``installs`` carries the layer maps and prerouter state the engine drives at the step boundary. """ - model, model_config = load_model( + require_backend("edge0-35b", ("mlx", "cuda")) + model, model_config = io.load_model( model_dir, lazy=True, strict=False, model_config={"model_type": "qwen3_5_moe"}, get_model_classes=_get_model_classes) - shards = open_shards(model_dir) + shards = io.open_shards(model_dir) spec = cfg.moe_spec opts = cfg.options # qwen config.json nests the text params under ``text_config``; mlx-lm @@ -135,7 +140,7 @@ def _build(self): opts = cfg.options if self._tok is None: try: - self._tok = load_tokenizer(cfg.model_dir) + self._tok = io.load_tokenizer(cfg.model_dir) except Exception: # noqa: BLE001 — tokenizer optional for CLI pass diff --git a/python/src/edge0/moe/routing.py b/python/src/edge0/moe/routing.py index 4c292bb..3cbddf1 100644 --- a/python/src/edge0/moe/routing.py +++ b/python/src/edge0/moe/routing.py @@ -38,7 +38,7 @@ def group_select_from_logits(logits, top_k: int, n_group: int, Returns ``(inds [..., k], scores [..., k])``. """ - scores = core.sigmoid(logits.astype(core.float32)) + scores = core.sigmoid(core.astype(logits, core.float32)) select = scores if expert_bias is not None: select = scores + expert_bias diff --git a/python/src/edge0/prerouter/heads.py b/python/src/edge0/prerouter/heads.py index 1368a98..9050f64 100644 --- a/python/src/edge0/prerouter/heads.py +++ b/python/src/edge0/prerouter/heads.py @@ -67,7 +67,7 @@ def __call__(self, h: core.array, executed_oh: core.array, prev_oh: core.array) -> core.array: feats = core.concatenate([h, executed_oh, prev_oh], axis=-1) if feats.dtype != self._dtype: - feats = feats.astype(self._dtype) + feats = core.astype(feats, self._dtype) # linear_init consumes the full concat features, same as training. return self.linear_init(feats) + self.fc2( gelu_erf(self.fc1(feats))) diff --git a/python/src/edge0/prerouter/install.py b/python/src/edge0/prerouter/install.py index 39976b6..553922f 100644 --- a/python/src/edge0/prerouter/install.py +++ b/python/src/edge0/prerouter/install.py @@ -17,11 +17,9 @@ from __future__ import annotations -from edge0.backends import core +from edge0.backends import core, io from edge0.backends import nn -from edge0.backends.mlx._impl.qwen3_next import Qwen3NextSparseMoeBlock -from edge0.backends.mlx.io import load_safetensors from edge0.moe.spec import MoESpec from edge0.prerouter.heads import PrerouterHead, topk_onehot from edge0.prerouter.spec import PrerouterSpec @@ -47,7 +45,7 @@ def _parse_weights(weights: dict[str, core.array], dtype): if parts[-2] not in ("fc1", "fc2", "linear_init"): continue if arr.dtype != dtype: - arr = arr.astype(dtype) + arr = core.astype(arr, dtype) heads.setdefault(owner, {})[f"{parts[-2]}.weight"] = arr return heads @@ -82,7 +80,7 @@ def install_prerouter( "the models & adapters') and place them in the model " "directory, or disable the prerouter with " "prerouter=None / --no-prerouter.") - weights = load_safetensors(pspec.weights_file) + weights = io.load_safetensors(pspec.weights_file) dtype = core.float16 if pspec.dtype == "fp16" else core.float32 head_weights = _parse_weights(weights, dtype) missing = [n for n in pspec.owner_layers(n_layers) if n not in head_weights] @@ -139,8 +137,7 @@ def install_prerouter( pg.linear_init.weight.shape, dtype=dtype) heads[owner] = pg - if pspec.patch_call and not getattr( - Qwen3NextSparseMoeBlock, _PATCHED_MARK, False): + if pspec.patch_call: _patch_qwen_consume() return state, heads @@ -152,7 +149,18 @@ def _patch_qwen_consume(): A class-level patch is required because implicit ``moe(x)`` calls look up the type, not the instance. Guarded per instance: only blocks with ``prerouter_enabled`` and single-token inputs (decode) route through - the prerouter; everything else takes the original router path.""" + the prerouter; everything else takes the original router path. + + The patched class is the active backend's port of the vendored block, + imported here so the module itself imports under any backend.""" + from edge0.backends import backend + if backend.name == "cuda": + from edge0.backends.cuda._impl.qwen3_5_moe import ( + SparseMoeBlock as Qwen3NextSparseMoeBlock) + else: + from edge0.backends.mlx._impl.qwen3_next import Qwen3NextSparseMoeBlock + if getattr(Qwen3NextSparseMoeBlock, _PATCHED_MARK, False): + return orig_call = Qwen3NextSparseMoeBlock.__call__ def prerouter_call(self, x): diff --git a/python/src/edge0/prerouter/stager.py b/python/src/edge0/prerouter/stager.py index c380ca5..22dfd9f 100644 --- a/python/src/edge0/prerouter/stager.py +++ b/python/src/edge0/prerouter/stager.py @@ -51,7 +51,7 @@ def _topk_onehot_pos(idx, num_experts: int, pos: int = -1): idx = idx[..., pos, :][..., None, :] # [B, 1, K] ar = core.arange(num_experts, dtype=idx.dtype) oh = (idx[..., None] == ar) # [B, 1, K, E] - return oh.sum(axis=-2).astype(core.float32) + return core.astype(oh.sum(axis=-2), core.float32) class PrerouterStager: @@ -158,7 +158,7 @@ def _logits(self, metas: list[int], inputs: list[core.array], feats = core.concatenate([core.stack(inputs, 0), core.stack(cur_ohs, 0), core.stack(prev_ohs, 0)], axis=-1) if feats.dtype != dtype: - feats = feats.astype(dtype) + feats = core.astype(feats, dtype) f2 = feats.reshape(len(metas), -1) h1 = core.einsum("ni,nij->nj", f2, w1) act = gelu_erf(h1) diff --git a/python/src/edge0/sampling.py b/python/src/edge0/sampling.py index 8d51aa7..ca0d410 100644 --- a/python/src/edge0/sampling.py +++ b/python/src/edge0/sampling.py @@ -10,7 +10,7 @@ def _mask_logits(logits, temperature, top_k, top_p): """Temperature / top-k / top-p truncation, shared by ``sample`` and any speculative verifier so both draw from the exact same distribution.""" - logits = logits.astype(core.float32) + logits = core.astype(logits, core.float32) if temperature > 0: logits = logits / temperature if top_k is not None and top_k > 0: @@ -21,10 +21,12 @@ def _mask_logits(logits, temperature, top_k, top_p): threshold = top[..., :1] logits = core.where(logits < threshold, float("-inf"), logits) if top_p is not None and top_p < 1.0: - sorted_vals = core.sort(logits, axis=-1)[..., ::-1] + # Descending sort as the negation of an ascending one (exact in + # IEEE arithmetic); torch cannot slice with a negative step. + sorted_vals = -core.sort(-logits, axis=-1) cum = core.cumsum(core.softmax(sorted_vals, axis=-1), axis=-1) cutoff = cum <= top_p - counts = core.sum(cutoff.astype(core.int32), axis=-1) + counts = core.sum(core.astype(cutoff, core.int32), axis=-1) k = core.maximum(counts, 1) threshold = core.take_along_axis( sorted_vals, (k - 1)[..., None], axis=-1) @@ -46,10 +48,10 @@ def sample(logits, temperature=0.7, top_k=None, top_p=None, penalty applies to them (vectorized, no per-token host syncs). """ if repetition_penalty != 1.0 and history: - logits = logits.astype(core.float32) + logits = core.astype(logits, core.float32) ids = core.array(sorted(set(int(x) for x in history)), dtype=core.int32) - if ids.size: + if core.size(ids): vals = core.take(logits, ids) logits[ids] = core.where( vals > 0, diff --git a/python/src/edge0/streaming/install.py b/python/src/edge0/streaming/install.py index 37b0e7e..c6115b6 100644 --- a/python/src/edge0/streaming/install.py +++ b/python/src/edge0/streaming/install.py @@ -18,6 +18,16 @@ from edge0.streaming.options import LayerOptions +def _swap_submodule(block, name: str, value) -> None: + """``setattr`` that also works on torch modules, which refuse to put a + non-Module where a child module is registered.""" + try: + setattr(block, name, value) + except TypeError: + delattr(block, name) + object.__setattr__(block, name, value) + + def install_streaming_experts( model, shards: list[SafetensorsMmap], @@ -26,6 +36,7 @@ def install_streaming_experts( num_layers: int | None = None, shared_cache: SharedExpertCache | None = None, prefetch_buffer: PrefetchBuffer | None = None, + wrap=None, ) -> list[StreamingSwitchGLU]: """Replace every routed-expert block in ``model`` with its streaming twin and return the twins in layer order. @@ -35,13 +46,19 @@ def install_streaming_experts( resolving. The original blocks are kept on the objects as ``_edge0_resident`` (used by the bit-parity tests); nothing else in the base model is touched. + + ``wrap(twin)`` adapts the twin to the block's own calling convention + before it is swapped in. The vendored MLX blocks call + ``switch_mlp(x, inds)`` / ``experts(x, idx)`` and weight the per-expert + outputs themselves, which is the twin's native interface, so the + default is no wrapping. """ if num_layers is None: n = 0 while True: try: spec.block_of(model, n) - except AttributeError: + except (AttributeError, IndexError): # past the last layer break n += 1 if n == 0: @@ -65,18 +82,19 @@ def install_streaming_experts( twin = StreamingSwitchGLU( shards, i, spec, options=opts, shared_cache=cache, prefetch_buffer=buf) - # keep the resident block reachable (bit-parity tests, debugging) - block._edge0_resident = block + # keep the resident block reachable (bit-parity tests, debugging); + # a plain attribute, not a child module -- as a child it would make + # the block its own descendant (torch's state_dict recurses forever) + object.__setattr__(block, "_edge0_resident", block) # swap the twin into the block, following the family convention: # qwen-style blocks call ``switch_mlp(x, inds)``; ling-style blocks # hold ``experts = SwitchGLU(...)`` and call ``experts(x, idx)``. # The original submodule is stashed for parity tests. - if hasattr(block, "switch_mlp"): - block._edge0_resident_switch = block.switch_mlp - block.switch_mlp = twin - else: - block._edge0_resident_switch = block.experts - block.experts = twin + name = "switch_mlp" if hasattr(block, "switch_mlp") else "experts" + # object.__setattr__: keep the stash out of torch's child registry + object.__setattr__(block, "_edge0_resident_switch", + getattr(block, name)) + _swap_submodule(block, name, wrap(twin) if wrap else twin) if opts.top_k is not None: res = getattr(block, "_edge0_resident", block) if hasattr(res, "top_k"): diff --git a/python/src/edge0/streaming/layer.py b/python/src/edge0/streaming/layer.py index ef1320b..b4fb5dd 100644 --- a/python/src/edge0/streaming/layer.py +++ b/python/src/edge0/streaming/layer.py @@ -46,8 +46,6 @@ from edge0.backends import nn import numpy as np -from mlx_lm.models.switch_layers import _gather_sort, _scatter_unsort - from edge0.moe.spec import MoESpec from edge0.streaming.cache import PrefetchBuffer, SharedExpertCache from edge0.streaming.mmap import SafetensorsMmap, u32_view @@ -190,7 +188,7 @@ def fn(x, w_u, s_u, b_u, w_g, s_g, b_g, w_d, s_d, b_d, _swiglu(x_up, x_gate), w_d, s_d, b_d, rhs_indices=local, transpose=True, group_size=gs, bits=bt, mode=md, sorted_indices=True) - return _scatter_unsort(z, inv_order, out_shape).squeeze(-2) + return quant.scatter_unsort(z, inv_order, out_shape).squeeze(-2) return fn def _make_moe_math_fused(): @@ -226,7 +224,7 @@ def fn(x, w_gu, s_gu, b_gu, w_d, s_d, b_d, _swiglu(x_up, x_gate), w_d, s_d, b_d, rhs_indices=local, transpose=True, group_size=gs, bits=bt, mode=md, sorted_indices=True) - return _scatter_unsort(z, inv_order, out_shape).squeeze(-2) + return quant.scatter_unsort(z, inv_order, out_shape).squeeze(-2) return fn self._moe_math = _make_moe_math() @@ -1118,7 +1116,7 @@ def __call__(self, x: core.array, indices: core.array | None = None) -> core.arr # to the resident top-N stack; misses (~3%) contribute zero from the # fast gather (overflow row) and are computed separately below via # the exact per-expert path, then scatter-added. Numerically exact. - if (self._hot_weights is not None and indices.size > 8 + if (self._hot_weights is not None and core.size(indices) > 8 and self._full_weights is None): with self._lock: flat = indices.reshape(-1).tolist() @@ -1151,7 +1149,7 @@ def __call__(self, x: core.array, indices: core.array | None = None) -> core.arr remapped = core.array(g_list, dtype=core.int32).reshape( ind_flat.shape) x3 = core.expand_dims(x_flat, (-2, -3)) - xs, local, inv_order = _gather_sort(x3, remapped) + xs, local, inv_order = quant.gather_sort(x3, remapped) x_up, x_gate = self._gather_gate_up(xs, w, local, True) z = quant.gather_qmm( _swiglu(x_up, x_gate), @@ -1159,7 +1157,7 @@ def __call__(self, x: core.array, indices: core.array | None = None) -> core.arr w[("down_proj", "biases")], rhs_indices=local, transpose=True, group_size=self.group_size, bits=self.bits, mode=self.mode, sorted_indices=True) - out = _scatter_unsort(z, inv_order) + out = quant.scatter_unsort(z, inv_order) # Miss contribution: switch_mlp's caller applies scores after # us, so only RAW per-expert outputs [T, k, H] at miss slots are # needed. Reuse _moe_math (same kernel as the staged path) with @@ -1216,7 +1214,7 @@ def _e1(t): ze = ze.squeeze(-2).squeeze(-2) # [m, H] flat_pos = core.array( [t * k_n + kk for t, kk in tk], dtype=core.int32) - out = out.at[flat_pos].add(ze[:, None, :]) + out = core.index_add(out, flat_pos, ze[:, None, :]) out = out.reshape(*ind_flat.shape, *out.shape[-2:]) out = out.squeeze(-2) # [T, k, H] if x.ndim == 3: # restore batch dim @@ -1228,7 +1226,7 @@ def _e1(t): # before_layer_cb load_full_layer / clear_full_layer), so the router # indices go straight into the sorted gather — same ops as the exact # path (bit-identical), no tolist / no bundles / no remap. - if self._full_weights is not None and indices.size > 8: + if self._full_weights is not None and core.size(indices) > 8: # Record usage so hot pins survive prefill (the whole-layer path # never touches _get_bundles, so without this _hot_counts stays # empty and PREFILL_PIN has nothing to select from). Also keep @@ -1244,7 +1242,7 @@ def _e1(t): self._last_prefill_topk = flat[-k:] w = self._full_weights x = core.expand_dims(x, (-2, -3)) - x, local, inv_order = _gather_sort(x, indices) + x, local, inv_order = quant.gather_sort(x, indices) x_up, x_gate = self._gather_gate_up(x, w, local, True) z = quant.gather_qmm( _swiglu(x_up, x_gate), @@ -1254,10 +1252,10 @@ def _e1(t): sorted_indices=True) # Matches the resident SwitchGLU contract: 4D [B, T, k, H] (the # caller's SparseMoeBlock sums over axis -2). NO batch restore - # here — _gather_sort/_scatter_unsort already keep [B, T, k, H] + # here — gather_sort/scatter_unsort already keep [B, T, k, H] # (verified against the deployment's whole-layer path, which # returns the same 4D shape). - return _scatter_unsort(z, inv_order, indices.shape).squeeze(-2) + return quant.scatter_unsort(z, inv_order, indices.shape).squeeze(-2) # Staged decode path (prerouter-style, single-token): consume the # fixed staged slots; router indices never leave the GPU (expert->slot @@ -1317,7 +1315,7 @@ def _build_asm(bundles, slot_of_list, exp_key): self._stats["staged_used"] += k return self._moe_math(x, *wargs, local2d) - if self._staged_mode and indices is not None and indices.size == self._staged_trigger: + if self._staged_mode and indices is not None and core.size(indices) == self._staged_trigger: self.wait_staged() st = self._staged_state if st is None: @@ -1337,7 +1335,7 @@ def _build_asm(bundles, slot_of_list, exp_key): for part in ("weight", "scales", "biases")) local2d = core.take(self._incr_slot_table, indices) with self._lock: - self._stats["staged_used"] += indices.size + self._stats["staged_used"] += core.size(indices) return self._moe_math(x, *wargs, local2d) else: bundles, slot_of_list, staged_exp = st @@ -1355,7 +1353,7 @@ def _build_asm(bundles, slot_of_list, exp_key): slot_of, wargs = asm local2d = core.take(slot_of, indices) with self._lock: - self._stats["staged_used"] += indices.size + self._stats["staged_used"] += core.size(indices) return self._moe_math(x, *wargs, local2d) # Exact decode path: resolve bundles for unique experts. @@ -1381,14 +1379,14 @@ def _build_asm(bundles, slot_of_list, exp_key): orig_x = x x = core.expand_dims(x, (-2, -3)) - do_sort = local2d.size >= 64 + do_sort = core.size(local2d) >= 64 inv_order = None if do_sort: if self._use_compile: - x_s, local, inv_order = _gather_sort(x, local2d) + x_s, local, inv_order = quant.gather_sort(x, local2d) return self._moe_math_sorted( x_s, *wargs, local, inv_order, indices.shape) - x, local, inv_order = _gather_sort(x, local2d) + x, local, inv_order = quant.gather_sort(x, local2d) elif self._use_compile: return self._moe_math(orig_x, *wargs, local2d) else: @@ -1421,7 +1419,7 @@ def _build_asm(bundles, slot_of_list, exp_key): transpose=True, group_size=self.group_size, bits=self.bits, mode=self.mode, sorted_indices=do_sort) if do_sort: - z = _scatter_unsort(z, inv_order, indices.shape) + z = quant.scatter_unsort(z, inv_order, indices.shape) return z.squeeze(-2) # ---- lifecycle -------------------------------------------------------- diff --git a/python/tests/backend_parity_worker.py b/python/tests/backend_parity_worker.py new file mode 100644 index 0000000..1f9e69e --- /dev/null +++ b/python/tests/backend_parity_worker.py @@ -0,0 +1,846 @@ +"""Run one shared-code case under whichever backend EDGE0_BACKEND selects. + +Invoked by tests/test_backend_parity.py as a subprocess, once per backend, +so the same framework code (routing, sampling, streaming) is exercised on +MLX and on the torch backend with identical inputs: + + EDGE0_BACKEND=cuda python tests/backend_parity_worker.py CASE IN.npz OUT.npz + +Inputs and outputs are numpy arrays in .npz files; bf16 never crosses the +boundary (cases upcast to float32 before returning). +""" + +from __future__ import annotations + +import sys + +import numpy as np + +from edge0.backends import core + + +# .tolist() is an array method on both backends (mlx.core has no tolist +# function, despite the contract listing one). +def _np(a): + return np.array(core.astype(a, core.float32).tolist(), dtype=np.float32) + + +def _np_int(a): + return np.array(a.tolist(), dtype=np.int64) + + +def _sorted_pair(inds, scores): + """Router outputs as (indices, scores) sorted by expert id per row -- + argpartition leaves the order within the top-k unspecified.""" + i, s = _np_int(inds), _np(scores) + order = np.argsort(i, axis=-1) + return np.take_along_axis(i, order, -1), np.take_along_axis(s, order, -1) + + +def case_select_from_logits(inp): + from edge0.moe.routing import select_from_logits + inds, scores = select_from_logits(core.array(inp["logits"]), top_k=4) + i, s = _sorted_pair(inds, scores) + return {"inds": i, "scores": s} + + +def case_group_select_from_logits(inp): + from edge0.moe.routing import group_select_from_logits + inds, scores = group_select_from_logits( + core.array(inp["logits"]), top_k=8, n_group=8, topk_group=4, + routed_scaling=2.5, expert_bias=core.array(inp["expert_bias"])) + i, s = _sorted_pair(inds, scores) + return {"inds": i, "scores": s} + + +def case_stacked_head_einsum(inp): + # The prerouter stager evaluates every head in ONE batch: feats from + # concatenate+stack, an astype to the head dtype, then einsum. + from edge0.prerouter.heads import gelu_erf + feats = core.concatenate([core.array(inp["a"]), core.array(inp["b"])], + axis=-1) + w1 = core.array(inp["w1"]) + feats = core.astype(feats, w1.dtype) + h1 = core.einsum("ni,nij->nj", feats, w1) + return {"h1": _np(h1), "act": _np(gelu_erf(h1))} + + +def case_mask_logits(inp): + from edge0.sampling import _mask_logits + out = {} + for name, kw in { + "topk": dict(temperature=0.7, top_k=20, top_p=None), + "topp": dict(temperature=0.9, top_k=None, top_p=0.8), + "both": dict(temperature=1.0, top_k=50, top_p=0.9), + "greedy": dict(temperature=0.0, top_k=None, top_p=None), + }.items(): + out[name] = _np(_mask_logits(core.array(inp["logits"]), **kw)) + return out + + +def _streaming_layer(model_dir, **opts): + """StreamingSwitchGLU over layer 1 of the real edge0-8b checkpoint.""" + import os + from edge0.moe.spec import MoESpec, QuantSpec, RouterKind, WeightLayout + from edge0.streaming.layer import StreamingSwitchGLU + from edge0.streaming.mmap import SafetensorsMmap + from edge0.streaming.options import LayerOptions + spec = MoESpec( + num_experts=128, top_k=8, intermediate_size=512, + router=RouterKind.SIGMOID_GROUP, norm_topk_prob=True, + routed_scaling=2.5, n_group=8, topk_group=4, shared_experts=1, + quant=QuantSpec(bits=4, group_size=64, mode="affine"), + layout=WeightLayout.SEPARATE, + key_template="model.layers.{layer}.mlp.experts", + block_path="model.layers.{layer}.mlp", + layer_path="model.layers.{layer}") + shards = [SafetensorsMmap(os.path.join(model_dir, "model.safetensors"))] + return StreamingSwitchGLU(shards, 1, spec, options=LayerOptions(**opts)) + + +def case_streaming(inp): + """Every StreamingSwitchGLU.__call__ path on real expert weights.""" + model_dir = str(inp["model_dir"]) + x1 = core.astype(core.array(inp["x1"]), core.bfloat16) # [1, H] + x16 = core.astype(core.array(inp["x16"]), core.bfloat16) # [16, H] + i1 = core.array(inp["i1"], dtype=core.int32) # [1, 8] + i16 = core.array(inp["i16"], dtype=core.int32) # [16, 8] + out = {} + for compiled in (True, False): + tag = "c" if compiled else "e" + layer = _streaming_layer(model_dir, use_compile=compiled) + out[f"exact_t1_{tag}"] = _np(layer(x1, i1)) # unsorted + out[f"exact_t16_{tag}"] = _np(layer(x16, i16)) # sorted + layer.load_full_layer() + out[f"full_t16_{tag}"] = _np(layer(x16, i16)) + layer.clear_full_layer() + # Hot stack holding every even expert: odd ones are misses and go + # through the exact scatter-add correction (core.index_add). + layer._hot_counts = {e: 1.0 for e in range(0, 128, 2)} + layer.load_hot_layer(n_hot=64) + layer.materialize_hot() + out[f"hot_t16_{tag}"] = _np(layer(x16, i16)) + layer.clear_hot_layer() + layer.close() + # Staged decode over a partial set: experts outside it are dropped. + staged = _streaming_layer(model_dir, use_compile=compiled, + staged=True, staged_n=8, staged_trigger=8) + staged.stage_experts([int(e) for e in inp["staged_set"]]) + staged.wait_staged() + out[f"staged_t1_{tag}"] = _np(staged(x1, i1)) + staged.close() + # Incremental stack (the staged_k4 preset since upstream #102): + # persistent [n+1, ...] tensors rewritten in place per changed slot. + # Two stagings with partial overlap exercise sticky slots, eviction + # back to the zero row and the write-back LRU. + incr = _streaming_layer(model_dir, use_compile=compiled, + staged=True, staged_n=8, staged_trigger=8, + staged_sync=True, incr_stack=True, + incr_writeback=True, warm_willneed=True) + assert incr._incr_mode + incr.stage_experts([int(e) for e in inp["staged_set"]]) + incr.wait_staged() + out[f"incr_t1_{tag}"] = _np(incr(x1, i1)) + incr.stage_experts([int(e) for e in inp["staged_set2"]]) + incr.wait_staged() + out[f"incr2_t1_{tag}"] = _np(incr(x1, i1)) + assert incr._stats["incr_writes"] > 8 + incr.close() + staged = _streaming_layer(model_dir, use_compile=compiled, + staged=True, staged_n=8, staged_trigger=8) + staged.stage_experts([int(e) for e in inp["staged_set2"]]) + staged.wait_staged() + out[f"staged2_t1_{tag}"] = _np(staged(x1, i1)) + staged.close() + return out + + +def case_bailing_port_parity(inp): + """torch backend only: the torch port of bailing_hybrid against the + vendored MLX model (run on the MLX CPU device -- full-precision float32; + some Apple GPUs run float32 matmul at ~7e-4), both loaded from the real + edge0-8b checkpoint, float32. Every torch layer is fed MLX's input to + that layer (per-layer error in isolation), and separately the whole + torch model runs free on the token ids with its own caches. A chunked + prefill (second chunk over a non-empty cache) and decode steps follow. + """ + import mlx.core as mx + import torch + + from edge0.backends.cuda._impl import bailing_hybrid as tb + from edge0.backends.cuda.core import DEVICE + from edge0.backends.cuda.io import load_model as t_load + from edge0.backends.cuda.model_specs import BAILING_V3_MOE_SPEC + from edge0.backends.mlx._impl import bailing_hybrid as mb + from edge0.backends.mlx.io import load_model as m_load + from edge0.streaming.install import install_streaming_experts + from edge0.streaming.mmap import SafetensorsMmap + + path = str(inp["model_dir"]) + over = {"model_type": "bailing_hybrid"} + mx.set_default_device(mx.cpu) + mm, _ = m_load(path, lazy=False, strict=False, model_config=over, + get_model_classes=lambda config: (mb.Model, mb.ModelArgs)) + mm.set_dtype(mx.float32) + tm, _ = t_load(path, strict=False, model_config=over, + get_model_classes=lambda config: (tb.Model, tb.ModelArgs), + dtype=torch.float32) + rep = tm._edge0_load_report + twins = install_streaming_experts( + tm, [SafetensorsMmap(f"{path}/model.safetensors")], + BAILING_V3_MOE_SPEC, num_layers=len(tm.model.layers)) + + def f32(a): + return np.array(a.astype(mx.float32)) if isinstance(a, mx.array) \ + else a.detach().float().cpu().numpy() + + def rel(a, b): + a, b = f32(a), f32(b) + return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-30)) + + ids = [int(i) for i in inp["ids"]] + steps = [ids[:7], ids[7:]] + [None] * int(inp["n_decode"]) + n = len(tm.model.layers) + m_cache, t_cache, t_free = mm.make_cache(), tm.make_cache(), tm.make_cache() + layer_err = np.zeros((len(steps), n)) + logit_err, m_arg, t_arg = [], [], [] + for s, chunk in enumerate(steps): + chunk = chunk if chunk is not None else [m_arg[-1]] + h = mm.model.word_embeddings(mx.array(chunk)[None]) + t_emb = tm.model.word_embeddings(torch.tensor(chunk, device=DEVICE)[None]) + m_mask = mb.create_attention_mask(h, m_cache[mm.model.first_mla_idx]) + t_mask = tb.create_attention_mask(t_emb, t_cache[tm.model.first_mla_idx]) + with torch.no_grad(): + for li in range(n): + ml, tl = mm.model.layers[li], tm.model.layers[li] + x_in = h + h = ml(x_in, m_mask if ml.is_mla else None, m_cache[li], None) + mx.eval(h) + t_out = tl(torch.from_numpy(f32(x_in)).to(DEVICE), + t_mask if tl.is_mla else None, t_cache[li], None) + layer_err[s, li] = rel(t_out, h) + mo = mm.lm_head(mm.model.norm(h))[0, -1] + to = tm.lm_head(tm.model(torch.tensor(chunk, device=DEVICE)[None], cache=t_free))[0, -1] + logit_err.append(rel(to, mo)) + m_arg.append(int(mx.argmax(mo).item())) + t_arg.append(int(torch.argmax(to).item())) + for t in twins: + if t is not None: + t.close() + return { + "layer_err": layer_err, "logit_err": np.array(logit_err), + "m_argmax": np.array(m_arg), "t_argmax": np.array(t_arg), + "is_mla": np.array([l.is_mla for l in tm.model.layers]), + "n_missing": np.array(len(rep["missing"])), + "unexpected": np.array(sorted({k.split(".")[3] + "." + k.split(".")[4] + for k in rep["unexpected"]}), dtype=str), + } + + +def _qwen35_tiny_ckpt(ckpt, seed): + """MLX builds a small edge0-35b-shaped model (3 GatedDeltaNet layers + 1 + gated full-attention layer, GQA, partial RoPE, 8 experts), casts it to + bf16, quantizes it with the model's own predicate (router and shared + gate at 8 bits) and saves it in its own format -- the published + checkpoint's format -- plus random prerouter heads in the prerouter + file's format. Returns the MLX model (on the MLX CPU device).""" + import json + import os + + import mlx.core as mx + import mlx.nn as mnn + from mlx.utils import tree_flatten, tree_map + + from edge0.backends.mlx._impl import qwen3_5_moe as mq + + mx.set_default_device(mx.cpu) + mx.random.seed(seed) + text = dict( + model_type="qwen3_5_moe_text", hidden_size=128, intermediate_size=256, + num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2, + head_dim=64, rms_norm_eps=1e-6, vocab_size=128, + linear_num_value_heads=4, linear_num_key_heads=2, + linear_key_head_dim=32, linear_value_head_dim=32, + linear_conv_kernel_dim=4, full_attention_interval=4, + num_experts=8, num_experts_per_tok=2, moe_intermediate_size=64, + shared_expert_intermediate_size=64, norm_topk_prob=True, + rope_parameters={"rope_type": "default", "rope_theta": 1e7, + "partial_rotary_factor": 0.25}) + config = {"model_type": "qwen3_5_moe", "text_config": text} + mm = mq.Model(mq.ModelArgs.from_dict(json.loads(json.dumps(config)))) + mm.update(tree_map(lambda p: p * 0.5, mm.parameters())) # keep O(1) + mm.set_dtype(mx.bfloat16) + pred = mm.quant_predicate + mnn.quantize(mm, group_size=64, bits=4, + class_predicate=lambda p, m: hasattr(m, "to_quantized") + and pred(p, m)) + os.makedirs(ckpt, exist_ok=True) + mx.save_safetensors(os.path.join(ckpt, "model.safetensors"), + dict(tree_flatten(mm.parameters()))) + # per-path overrides for the 8-bit layers, as in the published config + # (80 entries for edge0-35b: 40 layers x mlp.gate / shared_expert_gate); + # mlx-lm's loader reads the bits from here + quant = {"group_size": 64, "bits": 4} + for li in range(text["num_hidden_layers"]): + for leaf in ("mlp.gate", "mlp.shared_expert_gate"): + quant[f"language_model.model.layers.{li}.{leaf}"] = { + "group_size": 64, "bits": 8} + with open(os.path.join(ckpt, "config.json"), "w") as f: + json.dump(dict(config, quantization=quant), f) + f_in, hid = 128 + 2 * 8, 32 + heads = {} + for owner in range(0, 3): # start_layer=1 -> owners 0..n-2 + for name, shape in (("fc1", (hid, f_in)), ("fc2", (8, hid)), + ("linear_init", (8, f_in))): + heads[f"layers.{owner}.{name}.weight"] = ( + mx.random.normal(shape) * 0.3).astype(mx.float16) + mx.save_safetensors(os.path.join(ckpt, "prerouter.safetensors"), heads) + return mm + + +def case_qwen35_engine_generate(inp): + """Greedy-decode through the real Qwen35Engine (streaming, staged + decode, the class-level prerouter patch with trained-style heads) on + whichever backend EDGE0_BACKEND selects, on the small MLX-written + checkpoint. MLX runs on its CPU device.""" + import dataclasses + import os + + from edge0.config import GenerationConfig + from edge0.engine.qwen import Qwen35Engine + from edge0.models.edge0_35b import Qwen35Config + from edge0.streaming.options import LayerOptions + + ckpt = str(inp["ckpt_dir"]) + try: + import mlx.core as mx + mx.set_default_device(mx.cpu) + except ImportError: + pass + base = Qwen35Config._defaults(ckpt) + cfg = dataclasses.replace( + base, + moe_spec=dataclasses.replace(base.moe_spec, num_experts=8, top_k=2, + intermediate_size=64), + options=dataclasses.replace(LayerOptions.staged_k4(), staged_n=2, + staged_trigger=2, top_k=2, prefill_hot=4), + prerouter=dataclasses.replace( + base.prerouter, start_layer=1, hidden=32, + weights_file=(os.path.join(ckpt, "prerouter.safetensors") + if int(inp.get("prerouter", 1)) else "")), + prerouter_top_k=2, lora="") + from edge0.backends import core + engine = Qwen35Engine(ckpt, cfg) + # greedy by hand (what generate does with top_k=1) to keep every + # step's logits: token equality alone is too coarse on a small model + tokens, logits, graph = [], [], 0 + try: + engine.prefill([int(i) for i in inp["ids"]]) + lg = engine.next_logits() + for _ in range(int(inp["n"])): + # torch only: logits that carry an autograd graph keep every + # dequantized expert of the forward alive (MLX has no such state) + graph += int(getattr(lg, "grad_fn", None) is not None) + logits.append(_np(lg)) + tid = int(core.argmax(lg, axis=-1).item()) + tokens.append(tid) + lg = engine.step(tid) + finally: + engine.close() + return {"tokens": np.array(tokens, dtype=np.int64), + "logits": np.stack(logits), "graph": np.array(graph)} + + +def case_qwen35_make_ckpt(inp): + _qwen35_tiny_ckpt(str(inp["ckpt_dir"]), int(inp["seed"])) + return {"ok": np.array(1)} + + +def case_qwen35_port_parity(inp): + """torch backend only: the torch port of qwen3_5_moe against the vendored + MLX model. MLX builds a small edge0-35b-shaped model (3 GatedDeltaNet + layers + 1 gated full-attention layer, GQA, partial RoPE, 8 experts), + casts it to bf16 and quantizes it with the model's own predicate + (router and shared-expert gate at 8 bits), then saves it in its own + format -- the published checkpoint's format. The torch port loads that + file through the engine path, streams the experts from it, and is + compared layer by layer (teacher-forced) and free-running against the + MLX model on the MLX CPU device, both float32, over a chunked prefill + and decode steps.""" + import dataclasses + import os + + import mlx.core as mx + import torch + + from edge0.backends.cuda._impl import qwen3_5_moe as tq + from edge0.backends.cuda.core import DEVICE + from edge0.backends.cuda.io import load_model as t_load + from edge0.models.edge0_35b import Qwen35Config + from edge0.streaming.install import install_streaming_experts + from edge0.streaming.mmap import SafetensorsMmap + from mlx_lm.models.base import create_attention_mask as m_mask_fn + + ckpt = str(inp["ckpt_dir"]) + mm = _qwen35_tiny_ckpt(ckpt, int(inp["seed"])) + mm.set_dtype(mx.float32) + + tm, _ = t_load(ckpt, strict=False, model_config={"model_type": "qwen3_5_moe"}, + get_model_classes=lambda config: (tq.Model, tq.ModelArgs), + dtype=torch.float32) + rep = tm._edge0_load_report + spec = dataclasses.replace(Qwen35Config._defaults(ckpt).moe_spec, + num_experts=8, top_k=2, intermediate_size=64) + twins = install_streaming_experts( + tm, [SafetensorsMmap(os.path.join(ckpt, "model.safetensors"))], spec, + num_layers=4) + + def f32(a): + return np.array(a.astype(mx.float32)) if isinstance(a, mx.array) \ + else a.detach().float().cpu().numpy() + + def rel(a, b): + a, b = f32(a), f32(b) + return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-30)) + + ids = [int(i) for i in inp["ids"]] + steps = [ids[:6], ids[6:]] + [None] * int(inp["n_decode"]) + mlm, tlm = mm.language_model, tm.language_model + n = len(tlm.model.layers) + m_cache, t_cache, t_free = mlm.make_cache(), tlm.make_cache(), tlm.make_cache() + layer_err = np.zeros((len(steps), n)) + logit_err, m_arg, t_arg = [], [], [] + for s, chunk in enumerate(steps): + chunk = chunk if chunk is not None else [m_arg[-1]] + h = mlm.model.embed_tokens(mx.array(chunk)[None]) + t_emb = tlm.model.embed_tokens(torch.tensor(chunk, device=DEVICE)[None]) + m_fa = m_mask_fn(h, m_cache[mlm.model.fa_idx]) + t_fa = tq.create_attention_mask(t_emb, t_cache[tlm.model.fa_idx]) + with torch.no_grad(): + for li in range(n): + ml, tl = mlm.model.layers[li], tlm.model.layers[li] + x_in = h + h = ml(x_in, mask=None if ml.is_linear else m_fa, cache=m_cache[li]) + mx.eval(h) + t_out = tl(torch.from_numpy(f32(x_in)).to(DEVICE), + mask=None if tl.is_linear else t_fa, cache=t_cache[li]) + layer_err[s, li] = rel(t_out, h) + mo = mlm.lm_head(mlm.model.norm(h))[0, -1] + to = tlm.lm_head(tlm.model(torch.tensor(chunk, device=DEVICE)[None], cache=t_free))[0, -1] + logit_err.append(rel(to, mo)) + m_arg.append(int(mx.argmax(mo).item())) + t_arg.append(int(torch.argmax(to).item())) + for t in twins: + if t is not None: + t.close() + from edge0.backends.cuda import nn as cnn + return { + "layer_err": layer_err, "logit_err": np.array(logit_err), + "m_argmax": np.array(m_arg), "t_argmax": np.array(t_arg), + "is_linear": np.array([l.is_linear for l in tlm.model.layers]), + "n_missing": np.array(len(rep["missing"])), + "n_unexpected": np.array(len(rep["unexpected"])), + "gate_bits": np.array(tlm.model.layers[0].mlp.gate.bits), + "n_quantized": np.array(sum(isinstance(m, (cnn.QuantizedLinear, + cnn.QuantizedEmbedding)) + for m in tm.modules())), + } + + +def case_qwen35_hidden_clip(inp): + """torch backend only: QWEN_HIDDEN_CLIP (the per-layer clamp + NaN + scrub the engine turns on) in the torch port against the vendored MLX + model, through the full layer loop on the tiny MLX-written checkpoint. + Also reports how much the clip moved the output, so the test can + prove the clamp actually engaged.""" + import dataclasses + import os + + import mlx.core as mx + import torch + + from edge0.backends.cuda._impl import qwen3_5_moe as tq + from edge0.backends.cuda.core import DEVICE + from edge0.backends.cuda.io import load_model as t_load + from edge0.models.edge0_35b import Qwen35Config + from edge0.streaming.install import install_streaming_experts + from edge0.streaming.mmap import SafetensorsMmap + + ckpt = str(inp["ckpt_dir"]) + mm = _qwen35_tiny_ckpt(ckpt, int(inp["seed"])) + mm.set_dtype(mx.float32) + tm, _ = t_load(ckpt, strict=False, model_config={"model_type": "qwen3_5_moe"}, + get_model_classes=lambda config: (tq.Model, tq.ModelArgs), + dtype=torch.float32) + spec = dataclasses.replace(Qwen35Config._defaults(ckpt).moe_spec, + num_experts=8, top_k=2, intermediate_size=64) + twins = install_streaming_experts( + tm, [SafetensorsMmap(os.path.join(ckpt, "model.safetensors"))], spec, + num_layers=4) + ids = [int(i) for i in inp["ids"]] + out = {} + # The tiny model keeps its residual stream small: clip at a fraction of + # the largest per-layer output so the clamp is guaranteed to engage. + peak = [] + os.environ["QWEN_HIDDEN_CLIP"] = "0" + mx.eval(mm.language_model.model( + mx.array(ids)[None], + after_layer_cb=lambda li, h: peak.append(float(mx.abs(h).max())))) + clip_v = float(inp["clip_frac"]) * max(peak) + out["clip_v"] = np.array(clip_v) + for tag, v in (("off", "0"), ("on", repr(clip_v))): + os.environ["QWEN_HIDDEN_CLIP"] = v + m_h = mm.language_model.model(mx.array(ids)[None]) + mx.eval(m_h) + with torch.no_grad(): + t_h = tm.language_model.model(torch.tensor(ids, device=DEVICE)[None]) + out["m_" + tag] = np.array(m_h.astype(mx.float32)) + out["t_" + tag] = t_h.float().cpu().numpy() + os.environ.pop("QWEN_HIDDEN_CLIP") + for t in twins: + if t is not None: + t.close() + return out + + +def case_qwen35_real_parity(inp): + """torch backend only: the torch port of qwen3_5_moe against the vendored + MLX model, both loaded from the REAL edge0-35b checkpoint (four shards), + on the MLX CPU device in float32. Same shape as case_bailing_port_parity: + every torch layer is fed MLX's input to that layer, and separately the + whole torch model runs free with its own caches, over a chunked prefill + and decode steps.""" + import mlx.core as mx + import torch + + from edge0.backends.cuda._impl import qwen3_5_moe as tq + from edge0.backends.cuda.core import DEVICE + from edge0.backends.cuda.io import load_model as t_load + from edge0.backends.cuda.io import open_shards + from edge0.backends.mlx._impl import qwen3_5_moe as mq + from edge0.backends.mlx.io import load_model as m_load + from edge0.models.edge0_35b import Qwen35Config + from edge0.streaming.install import install_streaming_experts + from mlx_lm.models.base import create_attention_mask as m_mask_fn + + path = str(inp["model_dir"]) + over = {"model_type": "qwen3_5_moe"} + mx.set_default_device(mx.cpu) + mm, _ = m_load(path, lazy=False, strict=False, model_config=over, + get_model_classes=lambda config: (mq.Model, mq.ModelArgs)) + mm.set_dtype(mx.float32) + tm, _ = t_load(path, strict=False, model_config=over, + get_model_classes=lambda config: (tq.Model, tq.ModelArgs), + dtype=torch.float32) + rep = tm._edge0_load_report + mlm, tlm = mm.language_model, tm.language_model + n = len(tlm.model.layers) + twins = install_streaming_experts( + tm, open_shards(path), Qwen35Config._defaults(path).moe_spec, + num_layers=n) + + def f32(a): + return np.array(a.astype(mx.float32)) if isinstance(a, mx.array) \ + else a.detach().float().cpu().numpy() + + def rel(a, b): + a, b = f32(a), f32(b) + return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-30)) + + ids = [int(i) for i in inp["ids"]] + steps = [ids[:7], ids[7:]] + [None] * int(inp["n_decode"]) + m_cache, t_cache, t_free = mlm.make_cache(), tlm.make_cache(), tlm.make_cache() + layer_err = np.zeros((len(steps), n)) + logit_err, m_arg, t_arg = [], [], [] + for s, chunk in enumerate(steps): + chunk = chunk if chunk is not None else [m_arg[-1]] + h = mlm.model.embed_tokens(mx.array(chunk)[None]) + t_ids = torch.tensor(chunk, device=DEVICE)[None] + m_fa = m_mask_fn(h, m_cache[mlm.model.fa_idx]) + t_fa = tq.create_attention_mask(tlm.model.embed_tokens(t_ids), + t_cache[tlm.model.fa_idx]) + with torch.no_grad(): + for li in range(n): + ml, tl = mlm.model.layers[li], tlm.model.layers[li] + x_in = h + h = ml(x_in, mask=None if ml.is_linear else m_fa, + cache=m_cache[li]) + mx.eval(h) + t_out = tl(torch.from_numpy(f32(x_in)).to(DEVICE), + mask=None if tl.is_linear else t_fa, + cache=t_cache[li]) + layer_err[s, li] = rel(t_out, h) + mo = mlm.lm_head(mlm.model.norm(h))[0, -1] + to = tlm.lm_head(tlm.model(t_ids, cache=t_free))[0, -1] + logit_err.append(rel(to, mo)) + m_arg.append(int(mx.argmax(mo).item())) + t_arg.append(int(torch.argmax(to).item())) + for t in twins: + if t is not None: + t.close() + return { + "layer_err": layer_err, "logit_err": np.array(logit_err), + "m_argmax": np.array(m_arg), "t_argmax": np.array(t_arg), + "is_linear": np.array([bool(l.is_linear) for l in tlm.model.layers]), + "n_missing": np.array(len(rep["missing"])), + "n_unexpected": np.array(len(rep["unexpected"])), + "n_twins": np.array(sum(t is not None for t in twins)), + } + + +def case_engine_guard(inp): + """Every entry point imports without pulling MLX in, both engines + resolve to the torch ports, and a tier still refuses a backend it has + no port for.""" + import importlib + for mod in ("edge0", "edge0.engine", "edge0.engine.qwen", + "edge0.engine.ling", "edge0.cli", "edge0.prerouter.install", + "edge0.adapters.lora", "edge0.streaming.install"): + importlib.import_module(mod) + loaded_mlx = sorted(m for m in sys.modules + if m == "mlx" or m.startswith(("mlx.", "mlx_lm"))) + from edge0.engine.base import require_backend + try: + require_backend("some-tier", ("mlx",)) + error = "no error" + except NotImplementedError as e: + error = str(e) + from edge0.engine.ling import _get_model_classes as ling_classes + from edge0.engine.qwen import _get_model_classes as qwen_classes + return {"loaded_mlx": np.array(loaded_mlx, dtype=str), + "refusal": np.array(error), + "ling_model_module": np.array(ling_classes(config={})[0].__module__), + "qwen_model_module": np.array(qwen_classes(config={})[0].__module__)} + + +def case_engine_generate(inp): + """Greedy-decode through the real edge0-8b engine (LoRA, prerouter, + streaming, sampling) on whichever backend EDGE0_BACKEND selects.""" + from edge0 import AutoEngine + from edge0.config import GenerationConfig + n = int(inp["n"]) + engine = AutoEngine.from_pretrained(str(inp["model_dir"]), name="edge0-8b") + try: + ids = list(engine.encode_chat( + [{"role": "user", "content": str(inp["prompt"])}], think=False)) + out = engine.generate(ids, GenerationConfig( + temperature=0.0, top_k=1, top_p=1.0, max_new_tokens=n), + max_new_tokens=n) + finally: + engine.close() + return {"tokens": np.array(list(out), dtype=np.int64)} + + +def _tiny_qwen35(): + """A 4-layer transformers Qwen3.5-MoE (3 linear-attention layers, 1 full) + with every quantizable width a multiple of 64. Returns (config, model).""" + import torch + from transformers import Qwen3_5MoeForCausalLM, Qwen3_5MoeTextConfig + cfg = Qwen3_5MoeTextConfig( + vocab_size=128, hidden_size=128, num_hidden_layers=4, + num_attention_heads=4, num_key_value_heads=2, head_dim=32, + moe_intermediate_size=64, shared_expert_intermediate_size=64, + num_experts=8, num_experts_per_tok=2, + linear_num_key_heads=2, linear_num_value_heads=4, + linear_key_head_dim=16, linear_value_head_dim=16) + torch.manual_seed(0) + return cfg, Qwen3_5MoeForCausalLM(cfg).eval() + + +def _mlx_quantize(w, bits=4): + """MLX affine quantization (group 64), scales/biases rounded to bf16 as + on disk. Returns the on-disk tensors and the float32 weight they + dequantize to.""" + import mlx.core as mx + import torch + wq, s, b = mx.quantize(mx.array(w.detach().float().cpu().numpy()), + group_size=64, bits=bits) + s, b = s.astype(mx.bfloat16), b.astype(mx.bfloat16) + # bf16-valued scales, float32 arithmetic: with bf16 scales MLX would + # also round the dequantized weight to bf16 + deq = mx.dequantize(wq, s.astype(mx.float32), b.astype(mx.float32), + group_size=64, bits=bits) + + def bits16(a): + return torch.from_numpy(np.array(a.view(mx.uint16))).view(torch.bfloat16) + return ({"weight": torch.from_numpy(np.array(wq)), + "scales": bits16(s), "biases": bits16(b)}, + torch.from_numpy(np.array(deq))) + + +def case_load_model_qwen35_tiny(inp): + """torch backend only: write a small Qwen3.5-MoE checkpoint in the exact + on-disk format of the published edge0-35b -- ConditionalGeneration + config with text_config, language_model. prefix, every Linear and + Embedding 4-bit, the router 8-bit, experts as switch_mlp, norms stored + as w + 1, conv1d in MLX [C, k, 1] layout, the rest bf16 -- then + load_model + install_streaming_experts it and compare logits with the + source model running on the same effective weights.""" + import dataclasses + import json + import os + + import torch + from safetensors.torch import save_file + + from edge0.backends.cuda import nn as cnn + from edge0.backends.cuda.core import DEVICE + from edge0.backends.cuda.io import _QWEN35_SHIFTED_NORMS, load_model + from edge0.backends.cuda.model_specs import QWEN35_MOE_SPEC + from edge0.backends.cuda.moe_blocks import TransformersExpertsAdapter + from edge0.streaming.install import install_streaming_experts + from edge0.streaming.mmap import SafetensorsMmap + + cfg, ref = _tiny_qwen35() + I = cfg.moe_intermediate_size + ckpt = str(inp["ckpt_dir"]) + os.makedirs(ckpt, exist_ok=True) + P = "language_model." + disk, dense = {}, {} # on-disk tensors; reference weights + + def put(name, t): + disk[P + name] = t.contiguous() + + for name, t in ref.state_dict().items(): + t = t.detach().float() + mod_path, _, leaf = name.rpartition(".") + mod = ref.get_submodule(mod_path) + if name.endswith("mlp.experts.gate_up_proj"): + base = mod_path.replace(".experts", ".switch_mlp") + g, dg = _mlx_quantize(t[:, :I]) + u, du = _mlx_quantize(t[:, I:]) + for proj, q in (("gate_proj", g), ("up_proj", u)): + for part, v in q.items(): + put(f"{base}.{proj}.{part}", v) + dense[name] = torch.cat([dg, du], dim=1) + elif name.endswith("mlp.experts.down_proj"): + base = mod_path.replace(".experts", ".switch_mlp") + q, dense[name] = _mlx_quantize(t) + for part, v in q.items(): + put(f"{base}.down_proj.{part}", v) + elif leaf == "weight" and ( + isinstance(mod, (torch.nn.Linear, torch.nn.Embedding)) + or name.endswith("mlp.gate.weight")): + bits = 8 if name.endswith("mlp.gate.weight") else 4 + q, dense[name] = _mlx_quantize(t, bits=bits) + for part, v in q.items(): + put(f"{mod_path}.{part}", v) + else: + stored = t.to(torch.bfloat16) + dense[name] = stored.float() + if name.endswith(_QWEN35_SHIFTED_NORMS): + stored = (stored.float() + 1.0).to(torch.bfloat16) + if name.endswith("conv1d.weight"): + stored = stored.transpose(1, 2) # MLX [C, k, 1] + put(name, stored) + save_file(disk, os.path.join(ckpt, "model.safetensors")) + with open(os.path.join(ckpt, "config.json"), "w") as f: + json.dump({"architectures": ["Qwen3_5MoeForConditionalGeneration"], + "text_config": json.loads(cfg.to_json_string()), + "quantization": {"group_size": 64, "bits": 4, + "mode": "affine"}}, f) + ref.load_state_dict(dense) + + got = load_model(ckpt, strict=True, dtype=torch.float32) + spec = dataclasses.replace(QWEN35_MOE_SPEC, num_experts=cfg.num_experts, + top_k=cfg.num_experts_per_tok, + intermediate_size=I) + twins = install_streaming_experts( + got, [SafetensorsMmap(os.path.join(ckpt, "model.safetensors"))], + spec, wrap=TransformersExpertsAdapter) + out = { + "embed_type": np.array(type(got.model.embed_tokens).__name__), + "q_proj_type": np.array( + type(got.model.layers[3].self_attn.q_proj).__name__), + "lm_head_type": np.array(type(got.lm_head).__name__), + "router_dtype": np.array(str(got.model.layers[0].mlp.gate.weight.dtype)), + "router_bits": np.array(0), + "n_meta_params": np.array(sum( + p.is_meta for n, p in got.named_parameters())), + "n_quantized": np.array(sum( + isinstance(m, (cnn.QuantizedLinear, cnn.QuantizedEmbedding)) + for m in got.modules())), + } + with torch.no_grad(): + for key in ("ids6", "ids40"): + x = torch.from_numpy(inp[key].astype(np.int64)) + out[f"ref_{key}"] = ref(x).logits.float().numpy() + out[f"got_{key}"] = got(x.to(DEVICE)).logits.float().cpu().numpy() + for t in twins: + t.close() + return out + + +def case_install_qwen35_tiny(inp): + """torch backend only: a small transformers Qwen3_5MoeForCausalLM whose + experts are MLX-quantized into a real safetensors shard (named like the + edge0-35b checkpoint), streamed in with install_streaming_experts and + compared against the same model running its own dense experts.""" + import dataclasses + + import torch + from safetensors.torch import save_file + + from edge0.backends.cuda.core import DEVICE + from edge0.backends.cuda.model_specs import QWEN35_MOE_SPEC + from edge0.backends.cuda.moe_blocks import TransformersExpertsAdapter + from edge0.streaming.install import install_streaming_experts + from edge0.streaming.mmap import SafetensorsMmap + + cfg, model = _tiny_qwen35() + model.to(DEVICE) # where a user puts it; the streamed experts follow + E, K, I = cfg.num_experts, cfg.num_experts_per_tok, cfg.moe_intermediate_size + quantize = _mlx_quantize + + tensors = {} + for li in range(cfg.num_hidden_layers): + ex = model.model.layers[li].mlp.experts + prefix = f"language_model.model.layers.{li}.mlp.switch_mlp" + deq = {} + for proj, w in (("gate_proj", ex.gate_up_proj[:, :I]), + ("up_proj", ex.gate_up_proj[:, I:]), + ("down_proj", ex.down_proj)): + disk, deq[proj] = quantize(w) + for part, t in disk.items(): + tensors[f"{prefix}.{proj}.{part}"] = t.contiguous() + # the dense reference runs on exactly what the shard encodes + with torch.no_grad(): + ex.gate_up_proj.copy_(torch.cat([deq["gate_proj"], + deq["up_proj"]], dim=1)) + ex.down_proj.copy_(deq["down_proj"]) + shard_path = str(inp["shard_path"]) + save_file(tensors, shard_path) + + spec = dataclasses.replace(QWEN35_MOE_SPEC, num_experts=E, top_k=K, + intermediate_size=I) + out = {} + ids = {name: torch.from_numpy(inp[name].astype(np.int64)).to(DEVICE) + for name in ("ids6", "ids40")} # 12 pairs: unsorted; 80: sorted + with torch.no_grad(): + for name, x in ids.items(): + out[f"ref_{name}"] = model(x).logits.float().cpu().numpy() + twins = install_streaming_experts( + model, [SafetensorsMmap(shard_path)], spec, + wrap=TransformersExpertsAdapter) + out["n_twins"] = np.array(sum(t is not None for t in twins)) + out["experts_type"] = np.array( + type(model.model.layers[0].mlp.experts).__name__) + for name, x in ids.items(): + out[f"got_{name}"] = model(x).logits.float().cpu().numpy() + for t in twins: + t.close() + return out + + +CASES = {name[len("case_"):]: fn for name, fn in globals().items() + if name.startswith("case_")} + + +def main(argv): + case, inp_path, out_path = argv[1:4] + inp = dict(np.load(inp_path)) + np.savez(out_path, **CASES[case](inp)) + + +if __name__ == "__main__": + main(sys.argv) diff --git a/python/tests/test_backend_parity.py b/python/tests/test_backend_parity.py new file mode 100644 index 0000000..82b617b --- /dev/null +++ b/python/tests/test_backend_parity.py @@ -0,0 +1,338 @@ +"""Shared framework code must give the same answers on both backends. + +Each case in tests/backend_parity_worker.py runs twice in a subprocess -- +EDGE0_BACKEND=mlx and EDGE0_BACKEND=cuda -- on identical inputs, and the +outputs are compared here. MLX is the reference. Skipped unless both mlx +and torch are importable. +""" + +from __future__ import annotations + +import os +import pathlib +import subprocess +import sys + +import numpy as np +import pytest + +pytest.importorskip("mlx.core") +pytest.importorskip("torch") + +WORKER = pathlib.Path(__file__).with_name("backend_parity_worker.py") + + +def _run(case, inputs, tmp_path, backends=("mlx", "cuda")): + inp = tmp_path / f"{case}.in.npz" + np.savez(inp, **inputs) + results = {} + for backend in backends: + out = tmp_path / f"{case}.{backend}.npz" + env = dict(os.environ, EDGE0_BACKEND=backend) + proc = subprocess.run( + [sys.executable, str(WORKER), case, str(inp), str(out)], + env=env, capture_output=True, text=True) + assert proc.returncode == 0, ( + f"{case} failed on {backend}:\n{proc.stderr[-4000:]}") + results[backend] = dict(np.load(out)) + if len(backends) == 1: + return results[backends[0]] + return results["mlx"], results["cuda"] + + +def _logits(shape, seed=0, scale=3.0): + return (np.random.default_rng(seed).standard_normal(shape) * scale + ).astype(np.float32) + + +def test_select_from_logits(tmp_path): + ref, got = _run("select_from_logits", {"logits": _logits((6, 256))}, + tmp_path) + np.testing.assert_array_equal(got["inds"], ref["inds"]) + np.testing.assert_allclose(got["scores"], ref["scores"], rtol=1e-5, + atol=1e-6) + + +def test_group_select_from_logits(tmp_path): + rng = np.random.default_rng(1) + inputs = {"logits": _logits((6, 128), seed=1), + "expert_bias": (rng.standard_normal(128) * 0.1 + ).astype(np.float32)} + ref, got = _run("group_select_from_logits", inputs, tmp_path) + np.testing.assert_array_equal(got["inds"], ref["inds"]) + np.testing.assert_allclose(got["scores"], ref["scores"], rtol=1e-5, + atol=1e-6) + + +def _model_8b(): + path = os.environ.get("EDGE0_8B_MODEL") + if not path or not os.path.isfile(os.path.join(path, "model.safetensors")): + pytest.skip("set EDGE0_8B_MODEL to an edge0-8b checkpoint directory") + return path + + +def test_streaming_layer_real_experts(tmp_path): + rng = np.random.default_rng(3) + i16 = np.stack([rng.choice(128, 8, replace=False) for _ in range(16)]) + i1 = i16[:1] + inputs = { + "model_dir": np.array(_model_8b()), + "x1": rng.standard_normal((1, 1536)).astype(np.float32), + "x16": rng.standard_normal((16, 1536)).astype(np.float32), + "i1": i1.astype(np.int32), "i16": i16.astype(np.int32), + # half of token 0's experts plus unrelated ones: exercises drops + "staged_set": np.concatenate([i1[0, :4], [e for e in range(128) + if e not in i1[0]][:4]]).astype(np.int32), + # the other half of token 0's experts, two survivors of the first + # set and unrelated ones: incremental-stack churn + "staged_set2": np.concatenate([i1[0, 4:], i1[0, :2], + [e for e in range(128) + if e not in i1[0]][10:12]]).astype(np.int32), + } + ref, got = _run("streaming", inputs, tmp_path) + assert set(got) == set(ref) + for name in sorted(ref): + assert got[name].shape == ref[name].shape, name + # bf16 activations through three 4-bit matmuls: compare at bf16 + # resolution, relative to the output scale. + scale = np.abs(ref[name]).max() + np.testing.assert_allclose(got[name], ref[name], rtol=0, + atol=2e-2 * scale, err_msg=name) + # the staged run must actually drop the experts outside the staged set + assert not np.allclose(ref["staged_t1_e"], ref["exact_t1_e"]) + # The incremental stack is the same math as the rebuilt staged stack, + # before and after slot churn, on each backend. + for res in (ref, got): + for tag in ("c", "e"): + for a, b in (("incr", "staged"), ("incr2", "staged2")): + np.testing.assert_allclose( + res[f"{a}_t1_{tag}"], res[f"{b}_t1_{tag}"], rtol=0, + atol=1e-2 * np.abs(res[f"{b}_t1_{tag}"]).max(), + err_msg=f"{a} vs {b} ({tag})") + assert not np.allclose(ref["incr2_t1_e"], ref["incr_t1_e"]) + # Hot-stack prefill is documented as numerically exact, misses included + # (half the experts are misses here): same answer as the exact path, + # on each backend. Before the stack got its zero overflow row, misses + # gathered past the end of the stack. + for res in (ref, got): + for tag in ("c", "e"): + np.testing.assert_allclose(res[f"hot_t16_{tag}"], + res[f"exact_t16_{tag}"], rtol=0, + atol=1e-2 * np.abs(res[f"exact_t16_{tag}"]).max()) + + +def _model_35b(): + path = os.environ.get("EDGE0_35B_MODEL") + if not path or not os.path.isfile( + os.path.join(path, "model.safetensors.index.json")): + pytest.skip("set EDGE0_35B_MODEL to an edge0-35b checkpoint directory") + return path + + +def test_qwen35_port_matches_mlx_on_real_weights(tmp_path): + """The torch port of qwen3_5_moe on the real edge0-35b checkpoint: every + layer (GatedDeltaNet, gated full attention, routed + shared experts) and + the free-running logits within float32 noise of the MLX model, across a + chunked prefill and decode steps. Measured: <= 3.7e-6 per layer and + <= 2.4e-6 on the logits. + + Needs both models resident (about 22 GB); skipped without the checkpoint. + """ + from transformers import AutoTokenizer + path = _model_35b() + ids = AutoTokenizer.from_pretrained(path)( + "The capital of France is Paris. The capital of Italy is")["input_ids"] + res = _run("qwen35_real_parity", { + "model_dir": np.array(path), "ids": np.array(ids), + "n_decode": np.array(2)}, tmp_path, backends=("cuda",)) + assert int(res["n_missing"]) == 0 and int(res["n_unexpected"]) == 0 + assert int(res["n_twins"]) == len(res["is_linear"]) + assert res["is_linear"].any() and (~res["is_linear"]).any() + assert res["layer_err"].max() < 1e-5, res["layer_err"].max(axis=0) + assert res["logit_err"].max() < 1e-5, res["logit_err"] + np.testing.assert_array_equal(res["t_argmax"], res["m_argmax"]) + + +def test_bailing_port_matches_mlx_layer_by_layer(tmp_path): + """The torch port of bailing_hybrid on the real edge0-8b checkpoint: + every layer (KDA, MLA, dense MLP, routed + shared experts) and the + free-running logits within float32 noise of the MLX model, across a + chunked prefill and decode steps. Measured: <= 1.5e-6 per layer and + <= 1.7e-6 on the logits; the thresholds leave room for BLAS variation. + """ + from transformers import AutoTokenizer + path = _model_8b() + ids = AutoTokenizer.from_pretrained(path)( + "The capital of France is Paris. The capital of Italy is")["input_ids"] + res = _run("bailing_port_parity", { + "model_dir": np.array(path), "ids": np.array(ids), + "n_decode": np.array(2)}, tmp_path, backends=("cuda",)) + assert int(res["n_missing"]) == 0 + # the only checkpoint tensors without a home: the in-checkpoint copy + # of the prerouter, which edge0 loads from its own file + assert set(res["unexpected"]) <= {"mlp.pregate"}, res["unexpected"] + assert res["is_mla"].any() and (~res["is_mla"]).any() + assert res["layer_err"].max() < 1e-5, res["layer_err"].max(axis=0) + assert res["logit_err"].max() < 1e-5, res["logit_err"] + np.testing.assert_array_equal(res["t_argmax"], res["m_argmax"]) + + +def test_qwen35_port_matches_mlx_layer_by_layer(tmp_path): + """The torch port of qwen3_5_moe (edge0-35b) on a checkpoint MLX itself + wrote in the published format (bf16, 4-bit, 8-bit router and shared + gate, switch_mlp experts, MLX-layout conv1d and norms): every layer and + the free-running logits within float32 noise of the MLX model, across + a chunked prefill and decode steps. No 23 GB download needed.""" + rng = np.random.default_rng(7) + res = _run("qwen35_port_parity", { + "ckpt_dir": np.array(str(tmp_path / "ckpt")), "seed": np.array(7), + "ids": rng.integers(0, 128, 10), "n_decode": np.array(3)}, + tmp_path, backends=("cuda",)) + assert int(res["n_missing"]) == 0 and int(res["n_unexpected"]) == 0 + assert int(res["gate_bits"]) == 8 + assert res["is_linear"].any() and (~res["is_linear"]).any() + assert res["layer_err"].max() < 1e-5, res["layer_err"] + assert res["logit_err"].max() < 1e-5, res["logit_err"] + np.testing.assert_array_equal(res["t_argmax"], res["m_argmax"]) + + +def test_qwen35_port_hidden_clip_matches_mlx(tmp_path): + """QWEN_HIDDEN_CLIP, which the edge0-35b engine sets to 1000 by + default, clamps and NaN-scrubs every layer output the same way in the + torch port as in the vendored MLX model.""" + rng = np.random.default_rng(11) + res = _run("qwen35_hidden_clip", { + "ckpt_dir": np.array(str(tmp_path / "ckpt")), "seed": np.array(11), + "ids": rng.integers(0, 128, 9), "clip_frac": np.array(0.5)}, + tmp_path, backends=("cuda",)) + + def rel(a, b): + return float(np.abs(a - b).max() / (np.abs(b).max() + 1e-30)) + + assert rel(res["m_on"], res["m_off"]) > 1e-2 # the clip engaged + assert rel(res["t_off"], res["m_off"]) < 1e-5 + assert rel(res["t_on"], res["m_on"]) < 1e-5 + + +def test_qwen35_engine_matches_mlx_with_the_prerouter(tmp_path): + """The whole edge0-35b engine -- streaming experts, staged decode and the + class-level prerouter patch (patch_call=True) with heads in the real + file format -- on a small checkpoint MLX wrote in the published format + (LoRA off: no small adapter to match). Same greedy tokens as MLX, and + per-step logits that follow MLX *with* the prerouter: once predictions + exist (step 2 on) the prerouter moves the logits by 5-11% here, while + torch stays within ~1% of MLX -- so a torch path that silently fell back + to the plain router would sit much closer to MLX without it. Both + engines run in the checkpoint's bf16.""" + ckpt = tmp_path / "ckpt" + _run("qwen35_make_ckpt", {"ckpt_dir": np.array(str(ckpt)), + "seed": np.array(11)}, tmp_path, backends=("mlx",)) + inp = {"ckpt_dir": np.array(str(ckpt)), "n": np.array(10), + "ids": np.random.default_rng(12).integers(0, 128, 9)} + ref, got = _run("qwen35_engine_generate", dict(inp, prerouter=np.array(1)), + tmp_path) + (tmp_path / "no_prerouter").mkdir() + no_pr = _run("qwen35_engine_generate", dict(inp, prerouter=np.array(0)), + tmp_path / "no_prerouter", backends=("mlx",)) + np.testing.assert_array_equal(got["tokens"], ref["tokens"]) + # no autograd graph behind any step (it held every dequantized expert + # of the forward: the edge0-8b prefill grew past 120 GB on Linux) + assert got["graph"] == 0, got["graph"] + + def rel(a, b): + return np.abs(a - b).max(axis=-1) / np.abs(b).max(axis=-1) + to_ref = rel(got["logits"], ref["logits"]) + to_no_pr = rel(got["logits"], no_pr["logits"]) + assert to_ref.max() < 3e-2, to_ref # bf16 noise + acting = rel(ref["logits"], no_pr["logits"]) > 3e-2 # prerouter steps + assert acting.sum() >= 5, rel(ref["logits"], no_pr["logits"]) + assert (to_ref[acting] * 2 < to_no_pr[acting]).all(), (to_ref, to_no_pr) + + +def test_install_streaming_experts_into_transformers_qwen35(tmp_path): + """install_streaming_experts end to end on the torch backend: a real + transformers Qwen3.5-MoE model, the edge0-35b MoESpec paths, experts + streamed from a real safetensors shard. Same logits as the model's own + dense experts on the weights the shard encodes.""" + rng = np.random.default_rng(4) + res = _run("install_qwen35_tiny", { + "shard_path": np.array(str(tmp_path / "experts.safetensors")), + "ids6": rng.integers(0, 128, (1, 6)), + "ids40": rng.integers(0, 128, (1, 40)), + }, tmp_path, backends=("cuda",)) + assert int(res["n_twins"]) == 4 + assert str(res["experts_type"]) == "TransformersExpertsAdapter" + for name in ("ids6", "ids40"): + ref, got = res[f"ref_{name}"], res[f"got_{name}"] + np.testing.assert_allclose(got, ref, rtol=0, + atol=1e-5 * np.abs(ref).max(), err_msg=name) + + +def test_load_model_edge0_35b_format(tmp_path): + """cuda load_model on a checkpoint in the published edge0-35b on-disk + format (MLX-quantized throughout, MLX-sanitized norms and conv1d, + language_model. prefix), then streaming experts on top: same logits as + the source model on the weights the checkpoint encodes.""" + rng = np.random.default_rng(5) + res = _run("load_model_qwen35_tiny", { + "ckpt_dir": np.array(str(tmp_path / "ckpt")), + "ids6": rng.integers(0, 128, (1, 6)), + "ids40": rng.integers(0, 128, (1, 40)), + }, tmp_path, backends=("cuda",)) + assert str(res["embed_type"]) == "QuantizedEmbedding" + assert str(res["q_proj_type"]) == "QuantizedLinear" + assert str(res["lm_head_type"]) == "QuantizedLinear" + assert str(res["router_dtype"]) == "torch.float32" # 8-bit, dequantized + assert int(res["n_meta_params"]) == 0 # experts replaced by the twins + for key in ("ids6", "ids40"): + ref, got = res[f"ref_{key}"], res[f"got_{key}"] + np.testing.assert_allclose(got, ref, rtol=0, + atol=1e-5 * np.abs(ref).max(), err_msg=key) + assert (got.argmax(-1) == ref.argmax(-1)).all(), key + + +def test_engines_import_without_mlx_and_pick_the_backend_port(tmp_path): + res = _run("engine_guard", {"_": np.zeros(1)}, tmp_path, + backends=("cuda",)) + assert res["loaded_mlx"].size == 0, list(res["loaded_mlx"]) + assert str(res["refusal"]).startswith( + "some-tier: the engine runs only on EDGE0_BACKEND=mlx"), res["refusal"] + assert str(res["ling_model_module"]) == \ + "edge0.backends.cuda._impl.bailing_hybrid" + assert str(res["qwen_model_module"]) == \ + "edge0.backends.cuda._impl.qwen3_5_moe" + + +@pytest.mark.slow +def test_edge0_8b_engine_same_tokens_on_both_backends(tmp_path): + """The whole edge0-8b engine (LoRA, prerouter-staged decode, streaming + experts, sampling) greedy-decodes the same tokens on EDGE0_BACKEND=cuda + as on MLX. Slow: the torch path is a CPU reference here (~2 min).""" + ref, got = _run("engine_generate", { + "model_dir": np.array(_model_8b()), "n": np.array(8), + "prompt": np.array("Explain in two sentences why the sky is blue."), + }, tmp_path) + np.testing.assert_array_equal(got["tokens"], ref["tokens"]) + + +def test_stacked_head_einsum(tmp_path): + rng = np.random.default_rng(3) + ref, got = _run("stacked_head_einsum", { + "a": rng.standard_normal((33, 24)).astype(np.float16), + "b": rng.standard_normal((33, 8)).astype(np.float16), + "w1": rng.standard_normal((33, 32, 16)).astype(np.float32), + }, tmp_path) + for name in ref: + np.testing.assert_allclose(got[name], ref[name], rtol=1e-5, + atol=1e-5, err_msg=name) + + +def test_mask_logits(tmp_path): + ref, got = _run("mask_logits", {"logits": _logits((1000,), seed=2)}, + tmp_path) + for name in ref: + np.testing.assert_array_equal(np.isinf(got[name]), + np.isinf(ref[name]), err_msg=name) + keep = ~np.isinf(ref[name]) + np.testing.assert_allclose(got[name][keep], ref[name][keep], + rtol=1e-6, err_msg=name) diff --git a/python/tests/test_cuda_backend.py b/python/tests/test_cuda_backend.py new file mode 100644 index 0000000..4629a12 --- /dev/null +++ b/python/tests/test_cuda_backend.py @@ -0,0 +1,384 @@ +"""CUDA (torch reference) backend checked against real MLX. + +MLX is the ground truth here: every test builds inputs with ``mx.quantize`` +and compares the torch implementation against the MLX op it replaces. +Skipped unless both ``mlx`` and ``torch`` are importable (e.g. Apple +Silicon with torch installed); the torch side runs on CPU there. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core") +torch = pytest.importorskip("torch") + +from edge0.backends.cuda import quant as cq # noqa: E402 + +E, O, D = 8, 16, 128 + + +def _t(a): + if a.dtype == mx.bfloat16: # numpy has no bfloat16: move the raw bits + return torch.from_numpy(np.array(a.view(mx.uint16))).view(torch.bfloat16) + return torch.from_numpy(np.array(a)) + + +def _quantized(bits=4, group_size=64, dtype=mx.float32, shape=(E, O, D)): + mx.random.seed(0) + w = mx.random.normal(shape).astype(dtype) + wq, s, b = mx.quantize(w, group_size=group_size, bits=bits) + mx.eval(wq, s, b) + return wq, s, b + + +def _both(x, idx, wq, s, b, **kw): + """MLX reference against the torch port -- through BOTH of its paths. + + ``gather_qmm`` batches the gathered experts for decode-sized calls and + loops over the distinct ones above that, and every shape here is + decode-sized, so the loop would otherwise go untested. Forcing the cap + to 0 runs the same case through it; the two must agree.""" + ref = mx.gather_qmm(x, wq, s, b, rhs_indices=idx, **kw) + mx.eval(ref) + args = (_t(x), _t(wq), _t(s), _t(b), _t(idx)) + got = cq.gather_qmm(*args, **kw) + cap, cq.BATCH_ROWS = cq.BATCH_ROWS, 0 + try: + looped = cq.gather_qmm(*args, **kw) + finally: + cq.BATCH_ROWS = cap + # Same values, different summation order (bmm vs matmul): float32 + # tolerance, as in the MLX comparisons below. + np.testing.assert_allclose(looped.float().numpy(), got.float().numpy(), + rtol=1e-4, atol=1e-4) + return np.array(ref.astype(mx.float32)), got.float().numpy() + + +@pytest.mark.parametrize("x_shape,idx_shape", [ + ((3, D), (3,)), # default convention: broadcasts to (3, 3, O) + ((5, 1, 1, D), (5, 3)), # streaming/layer.py unsorted path + ((15, 1, D), (15,)), # streaming/layer.py sorted path (post _gather_sort) + ((2, 1, 4, D), (2, 3)), + ((1, 2, D), (4, 1)), +]) +def test_gather_qmm_matches_mlx(x_shape, idx_shape): + wq, s, b = _quantized() + mx.random.seed(1) + x = mx.random.normal(x_shape) + idx = mx.random.randint(0, E, idx_shape).astype(mx.uint32) + ref, got = _both(x, idx, wq, s, b, transpose=True, group_size=64, bits=4) + assert got.shape == ref.shape + np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4) + + +@pytest.mark.parametrize("bits", [2, 4, 8]) +def test_gather_qmm_bits(bits): + wq, s, b = _quantized(bits=bits) + x = mx.random.normal((4, 1, 1, D)) + idx = mx.random.randint(0, E, (4, 2)).astype(mx.uint32) + ref, got = _both(x, idx, wq, s, b, transpose=True, group_size=64, bits=bits) + np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4) + + +def test_gather_qmm_no_transpose(): + wq, s, b = _quantized(shape=(E, D, O * 4)) + x = mx.random.normal((3, 1, 1, D)) + idx = mx.random.randint(0, E, (3, 2)).astype(mx.uint32) + ref, got = _both(x, idx, wq, s, b, transpose=False, group_size=64, bits=4) + np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4) + + +def test_gather_qmm_sorted_flag_is_only_a_hint(): + wq, s, b = _quantized() + x = mx.random.normal((6, 1, D)) + idx = mx.array(sorted(np.random.default_rng(0).integers(0, E, 6).tolist()), + dtype=mx.uint32) + ref, got = _both(x, idx, wq, s, b, transpose=True, group_size=64, bits=4, + sorted_indices=True) + np.testing.assert_allclose(got, ref, rtol=1e-4, atol=1e-4) + + +def test_gather_qmm_bf16_checkpoint_dtypes(): + # Checkpoint scales/biases are bf16 and activations are bf16 in the + # real models; compare in float32 with a bf16-sized tolerance. + wq, s, b = _quantized(dtype=mx.bfloat16) + x = mx.random.normal((4, 1, 1, D)).astype(mx.bfloat16) + idx = mx.random.randint(0, E, (4, 2)).astype(mx.uint32) + ref, got = _both(x, idx, wq, s, b, transpose=True, group_size=64, bits=4) + np.testing.assert_allclose(got, ref, rtol=2e-2, atol=1e-1) + + +@pytest.mark.parametrize("shape,axes", [ + ((5, 7), (-2, -3)), ((5, 7), (0, 1)), ((5, 7), -1), ((2, 3, 4), (1, -1)), +]) +def test_expand_dims_matches_mlx(shape, axes): + from edge0.backends.cuda import core as cc + x = np.zeros(shape, dtype=np.float32) + assert tuple(cc.expand_dims(torch.from_numpy(x), axes).shape) == \ + tuple(mx.expand_dims(mx.array(x), axes).shape) + + +def test_rmsnorm_matches_mlx(): + from edge0.backends.cuda import nn as cnn + x = mx.random.normal((3, 64)) + w = mx.random.normal((64,)) + ref = np.array(mx.fast.rms_norm(x, w, 1e-6)) + norm = cnn.RMSNorm(64, eps=1e-6) + assert list(norm.state_dict()) == ["weight"] + norm.load_state_dict({"weight": _t(w)}) + got = norm(_t(x)).detach().numpy() + np.testing.assert_allclose(got, ref, rtol=1e-5, atol=1e-5) + + +def test_gelu_matches_mlx(): + import mlx.nn as mnn + from edge0.backends.cuda import nn as cnn + x = mx.random.normal((4, 32)) * 3 + np.testing.assert_allclose(cnn.gelu(_t(x)).numpy(), np.array(mnn.gelu(x)), + rtol=1e-5, atol=1e-5) + + +# ---- bailing_hybrid port, piece by piece (no checkpoint needed) ------------- + +def _bailing(): + from edge0.backends.cuda._impl import bailing_hybrid as tb + from edge0.backends.mlx._impl import bailing_hybrid as mb + return tb, mb + + +def _f32(shape, seed): + return np.random.default_rng(seed).standard_normal(shape).astype(np.float32) + + +def test_bailing_kda_recurrence_matches_mlx_gated_delta_ops(): + from mlx_lm.models.gated_delta import gated_delta_ops + tb, _ = _bailing() + B, T, H, Dk, Dv = 1, 6, 4, 16, 16 + q, k, v = _f32((B, T, H, Dk), 0), _f32((B, T, H, Dk), 1), _f32((B, T, H, Dv), 2) + g_log = -np.abs(_f32((B, T, H, Dk), 3)) # log-decay <= 0 + beta = 1 / (1 + np.exp(-_f32((B, T, H), 4))) + s0 = _f32((B, H, Dv, Dk), 5) + y_ref, s_ref = gated_delta_ops(mx.array(q), mx.array(k), mx.array(v), + mx.exp(mx.array(g_log)), mx.array(beta), + mx.array(s0)) + y, s = tb._kda_update(*(torch.from_numpy(a) for a in (q, k, v, g_log, beta, s0))) + np.testing.assert_allclose(y.numpy(), np.array(y_ref), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(s.numpy(), np.array(s_ref), rtol=1e-5, atol=1e-5) + + +def test_bailing_rope_interleave_matches_mlx(): + tb, mb = _bailing() + x = _f32((1, 4, 5, 64), 6) + pos = np.arange(5) + 7 # a cache offset of 7 + ref = mb._rope_interleave_torch(mx.array(x), mx.array(pos), 6e6) + got = tb._rope_interleave_torch(torch.from_numpy(x), torch.from_numpy(pos), 6e6) + np.testing.assert_allclose(got.numpy(), np.array(ref), rtol=1e-5, atol=1e-5) + + +def test_bailing_short_conv_with_state_matches_mlx(): + tb, mb = _bailing() + C, K = 8, 4 + w = _f32((C, 1, K), 7) # checkpoint layout + mc, tc = mb.ShortConv1d(C, K), tb.ShortConv1d(C, K) + mc.conv.weight = mx.array(np.swapaxes(w, 1, 2)) # what MLX's sanitize does + tc.weight.data = torch.from_numpy(w) + x1, x2 = _f32((1, 5, C), 8), _f32((1, 1, C), 9) # prefill, then one step + m1, ms = mc(mx.array(x1)) + m2, _ = mc(mx.array(x2), ms) + with torch.no_grad(): + t1, ts = tc(torch.from_numpy(x1)) + t2, _ = tc(torch.from_numpy(x2), ts) + for got, ref in ((t1, m1), (t2, m2)): + np.testing.assert_allclose(got.numpy(), np.array(ref), rtol=1e-5, atol=1e-5) + + +def test_bailing_gate_matches_mlx(): + tb, mb = _bailing() + ta = tb.ModelArgs(hidden_size=64, num_experts=128) + ma = mb.ModelArgs(hidden_size=64, num_experts=128) + w, bias, x = _f32((128, 64), 10), _f32((128,), 11) * 0.1, _f32((3, 5, 64), 12) + mg, tg = mb.BailingGate(ma), tb.BailingGate(ta) + mg.weight, mg.expert_bias = mx.array(w), mx.array(bias) + tg.weight.data, tg.expert_bias.data = torch.from_numpy(w), torch.from_numpy(bias) + # MLX on the CPU: float32 matmul on some Apple GPUs runs at reduced + # precision (~7e-4 from float64 on an M5 Max), the CPU path does not. + with mx.stream(mx.cpu): + mi, mw = mg(mx.array(x)) + mx.eval(mi, mw) + with torch.no_grad(): + ti, tw = tg(torch.from_numpy(x)) + mi, mw = np.array(mi), np.array(mw) + ti, tw = ti.numpy(), tw.numpy() + om, ot = np.argsort(mi, -1), np.argsort(ti, -1) # order within top-k is free + np.testing.assert_array_equal(np.take_along_axis(ti, ot, -1), + np.take_along_axis(mi, om, -1)) + np.testing.assert_allclose(np.take_along_axis(tw, ot, -1), + np.take_along_axis(mw, om, -1), rtol=1e-5) + + +def test_bailing_sdpa_causal_with_offset_matches_float64(): + """A second prefill chunk over a cache: query i of L must see keys up to + offset + i (end-aligned), which is what mx.fast's "causal" does and not + what torch's is_causal does.""" + tb, _ = _bailing() + L, off, D = 5, 7, 32 + q = np.random.default_rng(13).standard_normal((1, 2, L, D)) + k = np.random.default_rng(14).standard_normal((1, 2, off + L, D)) + v = np.random.default_rng(15).standard_normal((1, 2, off + L, D)) + s = np.einsum("bhqd,bhkd->bhqk", q, k) * D ** -0.5 + visible = np.arange(off + L)[None, :] <= (off + np.arange(L))[:, None] + s = np.where(visible, s, -np.inf) + p = np.exp(s - s.max(-1, keepdims=True)) + ref = np.einsum("bhqk,bhkd->bhqd", p / p.sum(-1, keepdims=True), v) + got = tb._sdpa(*(torch.from_numpy(a.astype(np.float32)) for a in (q, k, v)), + D ** -0.5, "causal") + np.testing.assert_allclose(got.numpy(), ref, rtol=1e-5, atol=1e-5) + + +# ---- qwen3_5_moe port, piece by piece (no checkpoint needed) ---------------- + +def _qwen(): + from edge0.backends.cuda._impl import qwen3_5_moe as tq + from edge0.backends.mlx._impl import qwen3_5 as mq5 + from edge0.backends.mlx._impl import qwen3_next as mqn + return tq, mq5, mqn + + +def test_qwen_gated_delta_update_matches_mlx(): + """Scalar per-head decay from (a, A_log, dt_bias), beta = sigmoid(b), + k heads repeated to the value heads -- vs mlx-lm's ops path.""" + from mlx_lm.models.gated_delta import gated_delta_update + tq, _, _ = _qwen() + B, T, Hk, Hv, Dk, Dv = 1, 6, 2, 4, 16, 16 + q, k, v = _f32((B, T, Hk, Dk), 20), _f32((B, T, Hk, Dk), 21), _f32((B, T, Hv, Dv), 22) + a, b = _f32((B, T, Hv), 23), _f32((B, T, Hv), 24) + A_log, dt_bias = _f32((Hv,), 25), _f32((Hv,), 26) + s0 = _f32((B, Hv, Dv, Dk), 27) * 0.1 + with mx.stream(mx.cpu): + y_ref, s_ref = gated_delta_update( + *(mx.array(t) for t in (q, k, v, a, b, A_log, dt_bias, s0)), + use_kernel=False) + mx.eval(y_ref, s_ref) + y, s = tq._gated_delta_update(*(torch.from_numpy(t) for t in + (q, k, v, a, b, A_log, dt_bias, s0))) + np.testing.assert_allclose(y.numpy(), np.array(y_ref), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(s.numpy(), np.array(s_ref), rtol=1e-5, atol=1e-5) + + +def test_qwen_partial_rope_matches_mlx_fast_rope(): + tq, _, _ = _qwen() + x = _f32((1, 4, 5, 256), 28) + for offset in (0, 7): + with mx.stream(mx.cpu): + ref = mx.fast.rope(mx.array(x), 64, traditional=False, base=1e7, + scale=1.0, offset=offset) + mx.eval(ref) + got = tq._rope(torch.from_numpy(x), 64, 1e7, offset) + np.testing.assert_allclose(got.numpy(), np.array(ref), rtol=1e-5, + atol=1e-5, err_msg=f"offset={offset}") + + +def test_qwen_gated_norm_and_conv_match_mlx(): + tq, _, mqn = _qwen() + w, x, z = _f32((16,), 29), _f32((1, 5, 4, 16), 30), _f32((1, 5, 4, 16), 31) + mn = mqn.Qwen3NextRMSNormGated(16) + mn.weight = mx.array(w) + tn = tq.RMSNormGated(16) + tn.weight.data = torch.from_numpy(w) + with mx.stream(mx.cpu): + ref = mn(mx.array(x), mx.array(z)) + mx.eval(ref) + with torch.no_grad(): + np.testing.assert_allclose(tn(torch.from_numpy(x), torch.from_numpy(z)).numpy(), + np.array(ref), rtol=1e-5, atol=1e-5) + import mlx.nn as mnn + C, K = 12, 4 + cw, cx = _f32((C, K, 1), 32), _f32((1, 9, C), 33) # MLX layout + mc = mnn.Conv1d(C, C, K, groups=C, bias=False, padding=0) + mc.weight = mx.array(cw) + tc = tq._DepthwiseConv(C, K) + tc.weight.data = torch.from_numpy(cw) + with mx.stream(mx.cpu): + cref = mc(mx.array(cx)) + mx.eval(cref) + with torch.no_grad(): + np.testing.assert_allclose(tc(torch.from_numpy(cx)).numpy(), np.array(cref), + rtol=1e-5, atol=1e-5) + + +def test_qwen_attention_with_cache_matches_mlx(): + """Gated GQA attention with partial RoPE, prefill then decode steps + through the KV cache (so RoPE offsets and the end-aligned causal mask + both matter).""" + from mlx.utils import tree_flatten + from mlx_lm.models.cache import KVCache as MKV + tq, mq5, _ = _qwen() + targs = dict(model_type="qwen3_5_moe_text", hidden_size=64, num_attention_heads=4, + num_key_value_heads=2, head_dim=32, rms_norm_eps=1e-6, + rope_parameters={"rope_type": "default", "rope_theta": 1e7, + "partial_rotary_factor": 0.25}) + ma = mq5.Attention(mq5.TextModelArgs.from_dict(dict(targs, + rope_parameters=dict(targs["rope_parameters"])))) + ta = tq.Attention(tq.TextModelArgs.from_dict(dict(targs, + rope_parameters=dict(targs["rope_parameters"])))) + rng = np.random.default_rng(34) + params = {k: (rng.standard_normal(v.shape) * 0.2).astype(np.float32) + for k, v in tree_flatten(ma.parameters())} + ma.load_weights([(k, mx.array(v)) for k, v in params.items()]) + ta.load_state_dict({k: torch.from_numpy(v) for k, v in params.items()}) + mc, tc = MKV(), tq.KVCache() + for step, L in enumerate((5, 3, 1, 1)): + x = _f32((1, L, 64), 40 + step) + with mx.stream(mx.cpu): + mask = "causal" if L > 1 else None + ref = ma(mx.array(x), mask, mc) + mx.eval(ref) + with torch.no_grad(): + got = ta(torch.from_numpy(x), "causal" if L > 1 else None, tc) + np.testing.assert_allclose(got.numpy(), np.array(ref), rtol=1e-4, + atol=1e-5, err_msg=f"step {step} (L={L})") + + +def test_swiglu_matches_mlx(): + import mlx.nn as mnn + up, gate = mx.random.normal((4, 32)), mx.random.normal((4, 32)) + ref = np.array(mnn.silu(gate) * up) + got = cq.swiglu(_t(up), _t(gate)).numpy() + np.testing.assert_allclose(got, ref, rtol=1e-5, atol=1e-5) + + +def test_quantized_linear_weight_cache_is_exact(): + """EDGE0_TORCH_WEIGHT_CACHE keeps the dequantized weight instead of + redoing it per call (2.5x on a GB10 decode step, with the expert cache). + It must be bit-for-bit the on-the-fly path, chunking included, so force + a projection wider than one chunk.""" + import mlx.nn as mnn + + from edge0.backends.cuda import nn as cnn + rows, cols = 5000, 256 # > ROWS_PER_CHUNK (4096) + w = mx.random.normal((rows, cols)) + ql = mnn.QuantizedLinear(cols, rows, bias=False, group_size=64, bits=4) + ql.weight, ql.scales, ql.biases = mx.quantize(w, group_size=64, bits=4) + x = mx.random.normal((3, cols)) + with mx.stream(mx.cpu): # MLX on this GPU runs f32 matmul low + ref = np.array(ql(x).astype(mx.float32)) + mx.eval(ref) + + def torch_out(cached): + layer = cnn.QuantizedLinear(_t(ql.weight), _t(ql.scales), + _t(ql.biases), cols) + old, cnn.CACHE_DEQUANTIZED = cnn.CACHE_DEQUANTIZED, cached + try: + with torch.no_grad(): + first = layer(_t(x)).float().numpy() + second = layer(_t(x)).float().numpy() # the cached call + finally: + cnn.CACHE_DEQUANTIZED = old + np.testing.assert_array_equal(first, second) + return first + + plain, cached = torch_out(False), torch_out(True) + np.testing.assert_array_equal(cached, plain) # exactness, not noise + np.testing.assert_allclose(plain, ref, rtol=1e-5, atol=1e-5) diff --git a/python/tests/test_cuda_core_take.py b/python/tests/test_cuda_core_take.py new file mode 100644 index 0000000..22df356 --- /dev/null +++ b/python/tests/test_cuda_core_take.py @@ -0,0 +1,30 @@ +"""``take`` on the torch backend follows numpy semantics, including negative axes.""" + +from __future__ import annotations + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from edge0.backends.cuda import core as cc # noqa: E402 + + +@pytest.mark.parametrize("axis", [0, 1, 2, -1, -2, -3]) +def test_take_matches_numpy_for_every_axis(axis): + rng = np.random.default_rng(0) + a = rng.standard_normal((4, 5, 6)).astype(np.float32) + indices = np.array([[3, 0], [2, 1]]) + out = cc.take(torch.as_tensor(a, device=cc.DEVICE), + torch.as_tensor(indices, device=cc.DEVICE), axis=axis) + expected = np.take(a, indices, axis=axis) + assert tuple(out.shape) == expected.shape + np.testing.assert_array_equal(out.cpu().numpy(), expected) + + +def test_take_without_axis_flattens_like_numpy(): + a = np.arange(12, dtype=np.float32).reshape(3, 4) + indices = np.array([[11, 0], [5, 6]]) + out = cc.take(torch.as_tensor(a, device=cc.DEVICE), + torch.as_tensor(indices, device=cc.DEVICE)) + np.testing.assert_array_equal(out.cpu().numpy(), np.take(a, indices)) diff --git a/python/tests/test_nvidia_backend_smoke.py b/python/tests/test_nvidia_backend_smoke.py new file mode 100644 index 0000000..7c5d6bd --- /dev/null +++ b/python/tests/test_nvidia_backend_smoke.py @@ -0,0 +1,99 @@ +"""Smoke test for MLX's GPU quantized-matmul support in isolation. + +Runs on any MLX GPU backend (Metal, or CUDA via ``mlx[cuda13]``). It does +NOT test edge0 end-to-end -- see docs/nvidia.md: edge0 does not run on +NVIDIA hardware with any MLX release tried (0.30.4 has no CUDA quantized +matmul; 0.31.1 lacks GatherQMM; 0.32.x has both but ``ling.py`` then hits +``IndexError`` in ``core.eval``). As of this writing the op below has +been checked on Metal only; on CUDA it has not run to completion yet. +""" + +from __future__ import annotations + +import pytest + +mx = pytest.importorskip("mlx.core", reason="mlx not installed") + + +def _has_gpu() -> bool: + # default_device() alone proves nothing: MLX is lazy and never touches + # the driver there. Evaluate something on the GPU for real. + try: + if mx.default_device().type != mx.DeviceType.gpu: + return False + mx.eval(mx.ones((8, 8)) @ mx.ones((8, 8))) + return True + except Exception: + return False + + +pytestmark = pytest.mark.skipif( + not _has_gpu(), + reason="requires an MLX GPU backend (Metal or CUDA) that can evaluate", +) + + +def test_default_device_is_gpu(): + assert mx.default_device().type == mx.DeviceType.gpu + + +@pytest.mark.skipif( + not mx.metal.is_available(), + reason="gather_qmm is verified on Metal only; the pinned MLX versions " + "do not complete it on CUDA (see docs/nvidia.md)", +) +def test_gather_qmm_matches_dense_reference(): + """The exact op edge0's streaming path depends on + (``backends/mlx/quant.py::gather_qmm``), checked against MLX's own + dense quantize/dequantize round-trip as ground truth -- no + cross-framework comparison needed since both sides are MLX here. + """ + import numpy as np + + rng = np.random.default_rng(0) + n_experts, out_features, in_features = 4, 32, 128 + group_size, bits = 64, 4 + + w = mx.array( + rng.standard_normal((n_experts, out_features, in_features)).astype(np.float32) + ) + wq, scales, biases = mx.quantize(w, group_size=group_size, bits=bits) + mx.eval(wq, scales, biases) + + # One row per selected expert: x is [2, 1, in] so the batch axis lines + # up with rhs_indices. A plain [2, in] x would broadcast against both + # indices and return [2, 2, out] (every row through every expert). + x = mx.array(rng.standard_normal((2, 1, in_features)).astype(np.float32)) + rhs_indices = mx.array([1, 3]) + + gathered = mx.gather_qmm( + x, wq, scales, biases, rhs_indices=rhs_indices, + transpose=True, group_size=group_size, bits=bits, + ).squeeze(-2) + mx.eval(gathered) + + w_deq = mx.dequantize(wq, scales, biases, group_size=group_size, bits=bits) + dense = mx.stack([ + x[i, 0] @ w_deq[int(rhs_indices[i])].T for i in range(x.shape[0]) + ]) + mx.eval(dense) + + rel_err = float(mx.linalg.norm(gathered - dense) / mx.linalg.norm(dense)) + assert rel_err < 1e-2, f"gather_qmm vs dense mismatch: rel_err={rel_err}" + + +@pytest.mark.skipif( + "EDGE0_NVIDIA_SMOKE_MODEL" not in __import__("os").environ, + reason="set EDGE0_NVIDIA_SMOKE_MODEL to a local checkpoint dir to run this", +) +def test_checkpoint_tokenizer_loads(): + """Tokenizer-only check on a real checkpoint. It does not build an + engine or generate: edge0 does not run end-to-end on MLX CUDA (see + docs/nvidia.md); the torch backend has its own checks.""" + import os + from edge0.backends import core, io + + model_path = os.environ["EDGE0_NVIDIA_SMOKE_MODEL"] + tokenizer = io.load_tokenizer(model_path) + ids = tokenizer.encode("The capital of France is") + assert len(ids) > 0 diff --git a/python/tests/test_remote_code_policy.py b/python/tests/test_remote_code_policy.py new file mode 100644 index 0000000..747995b --- /dev/null +++ b/python/tests/test_remote_code_policy.py @@ -0,0 +1,85 @@ +"""Keep checkpoint-supplied Python opt-in across backend loaders.""" + +from __future__ import annotations + +import ast +import pathlib + +ROOT = pathlib.Path(__file__).resolve().parents[1] + + +def _function(path: str, name: str) -> ast.FunctionDef: + tree = ast.parse((ROOT / path).read_text(encoding="utf-8")) + return next(node for node in tree.body + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + and node.name == name) + + +def _default_false(function: ast.FunctionDef, name: str) -> bool: + positional = function.args.args + positional_defaults = [None] * ( + len(positional) - len(function.args.defaults)) + positional_defaults.extend(function.args.defaults) + keyword_only = function.args.kwonlyargs + defaults = list(zip(positional, positional_defaults)) + defaults.extend(zip(keyword_only, function.args.kw_defaults)) + return any(arg.arg == name and isinstance(default, ast.Constant) + and default.value is False for arg, default in defaults) + + +def _keyword_value(function: ast.FunctionDef, method: str, + keyword: str) -> ast.expr | None: + for node in ast.walk(function): + if not isinstance(node, ast.Call): + continue + called_name = (node.func.attr if isinstance(node.func, ast.Attribute) + else node.func.id if isinstance(node.func, ast.Name) + else None) + if called_name == method: + for item in node.keywords: + if item.arg == keyword: + return item.value + return None + + +def test_tokenizer_remote_code_is_opt_in_for_both_backends(): + for path in ("src/edge0/backends/cuda/io.py", + "src/edge0/backends/mlx/io.py"): + loader = _function(path, "load_tokenizer") + assert _default_false(loader, "trust_remote_code") + value = _keyword_value(loader, "from_pretrained", "trust_remote_code") + assert isinstance(value, ast.Name) + assert value.id == "trust_remote_code" + + +def test_cuda_model_remote_code_is_explicitly_opt_in(): + loader = _function("src/edge0/backends/cuda/io.py", "load_model") + resolver = _function("src/edge0/backends/cuda/io.py", + "_resolve_model_class") + assert _default_false(loader, "trust_remote_code") + assert _default_false(resolver, "trust_remote_code") + + resolver_call = _keyword_value( + loader, "_resolve_model_class", "trust_remote_code") + assert isinstance(resolver_call, ast.Name) + assert resolver_call.id == "trust_remote_code" + + refusal = any( + isinstance(node, ast.If) + and isinstance(node.test, ast.UnaryOp) + and isinstance(node.test.op, ast.Not) + and isinstance(node.test.operand, ast.Name) + and node.test.operand.id == "trust_remote_code" + and any(isinstance(child, ast.Raise) + and isinstance(child.exc, ast.Call) + and isinstance(child.exc.func, ast.Name) + and child.exc.func.id == "ValueError" + for child in ast.walk(node)) + for node in ast.walk(resolver)) + assert refusal, "CUDA loader must refuse custom code unless opted in" + + for method in ("from_pretrained", "from_config"): + value = _keyword_value(resolver if method == "from_pretrained" + else loader, method, "trust_remote_code") + assert isinstance(value, ast.Name) + assert value.id == "trust_remote_code" diff --git a/python/tests/test_repo_hygiene.py b/python/tests/test_repo_hygiene.py index 9ea9b7c..7eeddb5 100644 --- a/python/tests/test_repo_hygiene.py +++ b/python/tests/test_repo_hygiene.py @@ -1,8 +1,9 @@ """Repository-hygiene guards for a clean open-source release. These tests import nothing from ``edge0`` (and nothing from ``mlx``), so -they can run on any platform — including the Linux CI job where MLX has no -wheels. They enforce the "no dependency on the author's machine" rule: +they can run on any platform, independent of whatever MLX wheels happen +to be available on the CI runner's platform. They enforce the "no +dependency on the author's machine" rule: * no personal / hardcoded absolute paths in shipped code, docs or scripts; * the MLX backend boundary is respected (``import mlx`` only inside