diff --git a/README.md b/README.md index 561d5bad66..47f2af89ed 100644 --- a/README.md +++ b/README.md @@ -122,6 +122,7 @@ between MLX and CUDA, different quantizations, or different tensor-parallel rank | `--vision-max-images N` | With `--vision`, images across the full request history (default 4); other image limits still apply | Both | | `--vision-image-tokens N` | With `--vision`, the visual tokens a request's images share (default 4,096, up to 65,536); each image keeps at most 4,096 | CUDA Qwen | | `--alias` | Additional model IDs | Both | +| `--adapter NAME=FOLDER`, `--adapter-strength NAME=X` | Serve a LoRA adapter as model ID NAME beside the base model, from the same weights ([docs/adapters.md](docs/adapters.md)) | DeepSeek-V4.1 CUDA with `--parallel >=2` | | `--context N` | Prompt plus reply capacity | Both | | `--max-tokens N` | Default reply limit, 4096 | Both | | `--temperature`, `--top-p`, `--top-k`, `--min-p` | Sampling defaults; temperature zero is greedy | Both | diff --git a/docs/adapters.md b/docs/adapters.md new file mode 100644 index 0000000000..b22903b49a --- /dev/null +++ b/docs/adapters.md @@ -0,0 +1,132 @@ +# LoRA adapters on DeepSeek-V4.1 + +`tensorfold serve --adapter NAME=FOLDER` serves one or more LoRA adapters beside the base model from one set of +weights. The EXL3 weights are not touched: one TP2 allocation, one shared KV pool, the same context, MTP, vision and +Engram. Each adapter is a model id of its own, and every request picks the base model or an adapter. Requests for +different adapters (and for the base model) decode in the same rounds, in the same CUDA graphs. + +```bash +tensorfold serve DeepSeek-V4.1-Flash-EXL3 --parallel 4 --tp 2 --master \ + --adapter edit-a=/models/adapters/edit-a --adapter edit-b=/models/adapters/edit-b --adapter-strength edit-b=0.5 +``` + +```bash +curl http://127.0.0.1:8080/v1/chat/completions -d '{"model": "edit-a", "messages": [{"role": "user", "content": "Hi"}]}' +curl http://127.0.0.1:8080/v1/chat/completions -d '{"model": "DeepSeek-V4.1-Flash-EXL3", "adapter": "edit-b", ...}' +``` + +## What an adapter may edit + +An adapter edits the **output projection of attention** (`layers.N.attn.wo_b`) and the **shared expert's down +projection** (`layers.N.ffn.shared_experts.w2`) of backbone layers, at any rank, each module with its own rank and +scale. `delta = scale * B @ (A @ x)` is added to the module's fp32 partial sum before the cross-rank reduction the layer +already makes; no collective is added. + +Everything else is refused when the server starts, by name: routed experts, the router, the shared expert's gate and up +projections, other attention projections, the compressor and indexer, embeddings, the head, norms, hyper-connection +parameters, Engram, the vision tower, and every MTP/DSpark block (`mtp.*`). A file that is not a plain, always-on LoRA +pair (DoRA magnitudes, biases, full weights, `modules_to_save`, `fan_in_fan_out`, layer replication, and PEFT's +activated LoRA, whose `alora_invocation_tokens` make it edit only after those tokens) is refused too. Nothing is +dropped quietly. The drafter always runs the base model, so an adapted reply is still the adapted target's own: drafts +are verified against it, and a weaker draft only costs acceptance. + +## The folder + +PEFT's two files, and optionally a manifest: + +- `adapter_config.json`: `peft_type: LORA`, `r`, `lora_alpha`, and optionally `rank_pattern`, `alpha_pattern` and + `use_rslora`. The scale of a module is `alpha / rank` (`alpha / sqrt(rank)` with rsLoRA). +- `adapter_model.safetensors`: `layers.N.attn.wo_b.lora_A.weight` `[r, 8192]` and `lora_B.weight` `[5120, r]`, and + the same for `layers.N.ffn.shared_experts.w2` (`[r, 2304]` and `[5120, r]`). F32, F16 and BF16 are read; the prefixes + `base_model.model.`, `base_model.model.model.` and `model.` are accepted. +- `tensorfold_adapter.json` (optional) is a contract: `{"format": "tensorfold-adapter-1", "files": {"adapter_model.safetensors": + {"sha256": ...}}, "base": {"model_type": ..., "hidden_size": ..., "num_hidden_layers": ..., "config_sha256": ...}, + "modules": {"layers.10.attn.wo_b": {"rank": 2}, ...}, "provenance": {...}}`. What it declares must match the files + (the weights' hash, the base model, and exactly this set of modules at exactly these ranks), or the server does not + start; a member of another type than the one shown (`files`, `base` or `provenance` not an object) is refused the same + way. Without it the same checks run on what the files say; either way the normalized manifest is what the server + holds. + +At load, a rank component whose A row or B column is all zero is dropped (its product is exactly zero), and a module +with nothing left is left out; a value that is not finite is refused. + +`python3 tools/dsv41/adapter_check.py FOLDER MODEL_DIR --tp 2 --slots 2` runs the same checks offline, with no weights +loaded: it prints the modules, the table memory a rank would hold, and exits 1 with the reason where the server would +refuse. `--write-manifest` writes `tensorfold_adapter.json` from the files (keeping an existing `provenance`). + +## Tensor parallelism + +Both modules are row-parallel: a rank holds its slice of the input width (`wo_b`: its attention groups; `w2`: its half +of the shared expert's intermediate width), so **A is cut by columns and B is whole on every rank**. The ranks' partial +deltas sum in the reduction the layer already makes. Each rank reads the files itself and keeps its share in fp32 tables. + +The delta is `scale * B @ (A @ x)` summed in fp32 from the serialized factors (F16 and BF16 widen exactly). PEFT's own +forward rounds the rank-sized intermediate to the activation dtype; this one does not. + +The shared expert is usually packed with the routed experts, and its down projection then reads the hidden state only +after the expert kernels have scaled it by the matrix's input vector and rotated it (a 128-wide Hadamard). At load, A is +rewritten once to read that form (`A' = H (A / suh)`, float64 arithmetic, stored float32), so no hidden state is +rebuilt at run time. The delta of a packed shared expert is therefore computed from the fp16 value the expert kernels +store (relative error at most 2^-11 per element); a shared expert outside the pack reads its bf16 hidden state as is. + +On start every rank gathers a digest of the base config and each adapter's slot, name and identity key, and refuses to +serve if they differ. Slots follow the order of the `--adapter` flags, so every rank needs the same flags in the same +order: two adapters swapped is refused too. Each admission carries `[slot, key]` to the followers, and every step's +agreement check includes the slots of the live streams, so a rank that disagrees fails the step on every rank before any +collective. + +## Graphs and rows + +Slot 0 is the base model. A round's inputs carry one adapter slot a row (a sixth row of the graph's input buffer, read at +replay), padded rows copy their original's slot, and prompt chunks run one slot for all their rows. Two Triton kernels +add `scale * B @ (A @ x)` for a row's slot: one program per row and tile, fp32 sums in a fixed order, one tile shape for +every table. By construction a row's delta depends on its own input and its adapter's factors only: not on the other +rows, their adapters, the padding, the chunk it was prefilled in, or the largest rank of any other adapter. A row of +slot 0 is neither read nor written. The CPU tests check a row alone, reordered and padded against the whole batch, and +stock rows left untouched, under Triton's interpreter. The same on the compiled kernels, in prompt chunks, in captured +graphs and beside a larger rank is what the GPU tests (`tests/cuda/test_deepseek_v41_adapter_kernels.py`) and, on a +running server, `tools/dsv41/adapter_mix.py` check; they need a GPU. A layer no adapter edits runs nothing. The kernels +are built at start, before the graph warm-up. + +## Caches + +The KV pool is one pool, but a row computed under an adapter is that adapter's. Kept prompts (`TF_DS_KEEP`) and +spilled prompts (`TF_DS_SPILL_GIB`) carry the identity key of the adapter that made them: a prompt continues only a +kept or spilled prompt of the same key, and the spill's prefix hashes start with the key (the base model's hashes are +unchanged). The key changes with the weights file's bytes, the module set and every module's scale (including +`--adapter-strength`), never with the name. Adding this feature changes the engine identity hash of the spill +directory, so an existing spill directory is rebuilt once. + +## Memory + +Each rank holds, per edited module and layer, (held adapters + 1) slots of the largest rank any adapter has there +(rounded up to a multiple of 4), each `(local input width + output width) * 4` bytes per rank component. The tables are +made before the warm-up and counted by the allocator ceiling. Start refuses tables that would leave the host under +`TF_DS_MEM_FLOOR_GIB` plus one gibibyte of its available memory, before allocating them. `TF_DS_ADAPTER_MAX_RANK` +(default 256) bounds a module's rank. + +## Limits + +- `--parallel 2` or more is required (the single-stream engine has no per-row input). `TF_DS_KERNELS=0` and + `TF_DS_EXACT_MM=1`, the test modes, are refused with `--adapter`. +- Adapters are fixed at start: no loading or unloading at run time. +- Logprobs, as under `--parallel` for the base model, are not served. +- Correctness here means the adapted model is the base weights plus the serialized factors, run as written. It does not + claim that a checkpoint quantized again from merged weights would give the same activations; compare those separately. + +## Checking a running server + +Every adapter is a model id, so `python3 tools/bench_concurrent.py URL ADAPTER_ID --alone --serial` checks that an +adapter's concurrent replies equal their solo runs and drafted equals `"draft": false`, as for any model. +`python3 tools/dsv41/adapter_mix.py URL --models BASE,adapter-a,adapter-b --streams 4` checks the mixture: the ids rotate +over the streams, and each stream's token hash must be its own solo run's wherever it sits and whatever runs beside it. +To check that the flag changes nothing for the base model, compare the base model's hashes from a server with and +without `--adapter`. + +## Tests + +CPU: `tests/test_deepseek_v41_adapter_*.py` (folder validation and refusals, registry and identities, tensor-parallel +placement and the shared-expert fold in numpy, tables and hooks, the round columns, kept and spilled prompt scoping, the +admission ops, server and CLI selection, and the kernels under Triton's CPU interpreter). GPU: `tests/cuda/test_deepseek_v41_adapter_*.py` +(the compiled kernels, graph replay with changing slots, the fold against the real EXL3 expert kernels, and two ranks +over gloo). diff --git a/docs/api.md b/docs/api.md index d2f46ce3c3..8cb0a7bd54 100644 --- a/docs/api.md +++ b/docs/api.md @@ -148,6 +148,18 @@ asks for it and sends no `priority` of its own is treated as `priority: "backgro field always wins over the default. This lets one CUDA server answer to several ids at their usual priority while a background-only client (a batch extractor, say) gets one id that always yields, without sending the field itself. +### Adapters + +A server started with `--adapter NAME=FOLDER` (DeepSeek-V4.1 on CUDA, [adapters.md](adapters.md)) lists each adapter in +`/v1/models` as a model id of its own, with `parent` and `root` naming the base model. A request picks an adapter by +that `model` id, or by the `adapter` field (the adapter's name) beside any model id; replies name the model id asked +for. A request naming no adapter runs the base model: the adapter kernels never read or write its rows, so by +construction no adapter changes its output. To confirm on your hardware that it matches a server started without the +flag, compare the two servers' token hashes ([adapters.md](adapters.md#checking-a-running-server)). An `adapter` that +is not held or that contradicts the `model` id is a 400 before any stream opens, never the base model's answer. A +server started without `--adapter` ignores the field like any other unknown one. The Responses route forwards +`adapter`; the Anthropic Messages route picks an adapter by its `model` id only. + ## Messages and tools Developer messages use system-message semantics. Only leading system and developer messages merge, in order, diff --git a/src/tensorfold/cli.py b/src/tensorfold/cli.py index 58121f1b07..93bc01c4a2 100644 --- a/src/tensorfold/cli.py +++ b/src/tensorfold/cli.py @@ -16,7 +16,8 @@ from tensorfold import cli_args from tensorfold.server import stacks, thinking_notes from tensorfold.server.memory_budget import MEMORY_FRACTION -from tensorfold.serve_options import check as _check_serve_options, name_priority as _name_priority +from tensorfold.serve_options import adapter_specs as _adapter_specs, check as _check_serve_options +from tensorfold.serve_options import name_priority as _name_priority from tensorfold.serve_options import vision_options as _vision_options COMMANDS = ("serve", "pull", "models", "info", "update", "service", "tui", "plan") @@ -270,6 +271,9 @@ def _serve_cuda(args: argparse.Namespace, family: Any, model_dir: Path, context: if getattr(args, "checkpoint_slots", None) is not None and getattr(family.package, "CUDA_CHECKPOINT_SLOTS", False): options["checkpoint_slots"] = int(args.checkpoint_slots) served = args.name or (args.model.rstrip("/").split("/")[-1] if hub.is_repo_id(args.model) else model_dir.name) + adapters = _adapter_specs(args) + if adapters: + options["adapters"], options["adapter_reserved"] = adapters, [served, *args.alias] where = f", rank {args.rank} of 2" if args.tp == 2 else "" print(f"[tensorfold] loading {served}: {family.title} ({family.model_type}) on CUDA{where}", flush=True) from tensorfold.cuda import precision, prompt_precision diff --git a/src/tensorfold/cli_args.py b/src/tensorfold/cli_args.py index cd4fa5bb0d..b0c07b4f2d 100644 --- a/src/tensorfold/cli_args.py +++ b/src/tensorfold/cli_args.py @@ -49,6 +49,14 @@ def build_parser(handlers: dict[str, Callable[[argparse.Namespace], int]]) -> ar help="with --vision on CUDA Qwen checkpoints, the visual tokens a request's images share " "(default: 4096, at most 65536); each image keeps at most 4096") + adapters = serve.add_argument_group("adapters (DeepSeek-V4.1 on CUDA; docs/adapters.md)") + adapters.add_argument("--adapter", action="append", default=[], metavar="NAME=FOLDER", + help="serve the LoRA adapter in FOLDER (PEFT files, wo_b and the shared expert's w2) as model " + "id NAME beside the base model, both from one set of weights; repeatable; needs --parallel 2 " + "or more") + adapters.add_argument("--adapter-strength", action="append", default=[], metavar="NAME=NUMBER", + help="multiply adapter NAME's scale (default 1; a negative number inverts its edit)") + generation = serve.add_argument_group("generation (requests can override each of these)") generation.add_argument("--context", type=int, default=None, help="prompt plus reply window (default: model config; CUDA default/0: " diff --git a/src/tensorfold/cuda/http.py b/src/tensorfold/cuda/http.py index 85ce1e70fa..09226138d2 100644 --- a/src/tensorfold/cuda/http.py +++ b/src/tensorfold/cuda/http.py @@ -154,8 +154,11 @@ def do_GET(self): # vLLM's max_model_len: the prompt-plus-reply tokens a request may use (clients size their context by it) limit = getattr(app, "effective_context_window", None) or getattr(app, "native_context_window", None) extra = {"max_model_len": int(limit)} if limit else {} + adapters = getattr(app, "adapter_models", None) or {} # an adapter's entry names its base model self._json(200, {"object": "list", "data": [{"id": model_id, "object": "model", "owned_by": "tensorfold", - **extra} for model_id in app.model_ids]}) + **extra, **({"parent": app.served, "root": app.served} + if model_id in adapters else {})} + for model_id in app.model_ids]}) elif self.path.rstrip("/") in ("/health", "/v1/health"): self._json(200, {"status": "ok"} if getattr(getattr(app, "auth", None), "enabled", False) else health.of(app).snapshot(app)) diff --git a/src/tensorfold/cuda/scheduler.py b/src/tensorfold/cuda/scheduler.py index 4570d405a0..d5b4b1153b 100644 --- a/src/tensorfold/cuda/scheduler.py +++ b/src/tensorfold/cuda/scheduler.py @@ -76,12 +76,14 @@ def close(self, timeout: float | None = None, *, drain: bool = True) -> bool: def submit(self, prompt: list[int], count: int, sampling: Any, draft: bool, emit: Callable[[list[int]], bool | None], stop_eos: bool = True, *, vision: Any = None, - constraint: Any = None, background: bool = False, probabilities: Any = None) -> dict: - """Decode one request; ``emit`` runs on the calling thread and returns True to stop. Returns its stats.""" + constraint: Any = None, background: bool = False, probabilities: Any = None, + adapter: Any = None) -> dict: + """Decode one request; ``emit`` runs on the calling thread and returns True to stop. Returns its stats. + ``adapter``: the engine's reference to the adapter it runs under (None: the stock model).""" box: queue.Queue = queue.Queue() stream = Stream(list(prompt), max(1, count), sampling, draft=draft, stop_eos=stop_eos, vision=vision, - constraint=constraint, background=background, probabilities=probabilities) + constraint=constraint, background=background, probabilities=probabilities, adapter=adapter) cancel = [False] stream.emit = lambda new: (box.put(("tokens", new)), cancel[0])[1] self.waiting.put((stream, box)) diff --git a/src/tensorfold/cuda/server.py b/src/tensorfold/cuda/server.py index 590f48ef36..1a88f6ba24 100644 --- a/src/tensorfold/cuda/server.py +++ b/src/tensorfold/cuda/server.py @@ -59,6 +59,7 @@ class PreparedRequest: vision: Any = None grammar: Any = None # (spec, compiled grammar) of the request's response_format, or None think_budget: int = 0 # reply tokens before the server closes a think block the reply leaves open (0: no limit) + adapter: str | None = None # the adapter the request selects (its name), or None for the stock model def _native_context(model_dir: Path) -> int: @@ -94,6 +95,7 @@ def __init__(self, engine, model_dir: Path, served: str, *, default_thinking: bo self.served = served self.aliases = tuple(str(alias).strip() for alias in aliases if str(alias).strip()) self.background_ids = frozenset(background_ids) # --name-priority ID=background: a default for this id + self.adapter_models: dict[str, str] = {} # model id -> adapter name, for engines that serve adapters self.model_dir = Path(model_dir) self.tok = Tokenizer.from_file(str(model_dir / "tokenizer.json")) self.template = ChatTemplate(model_dir) @@ -124,11 +126,24 @@ def model_ids(self) -> list[str]: """The ids this endpoint answers to, as the MLX server lists them: ``--name`` first, then each ``--alias``.""" ids: list[str] = [] - for model_id in (self.served, *getattr(self, "aliases", ())): + for model_id in (self.served, *getattr(self, "aliases", ()), *getattr(self, "adapter_models", {})): if model_id and model_id not in ids: ids.append(model_id) return ids + def adapter_for(self, body: Any) -> str | None: + """The adapter a request selects: its model id when that is an adapter's, else its "adapter" field; None is the + stock model. Where the engine serves adapters, an unknown or contradictory choice is refused, never answered + by the stock model.""" + + registry = getattr(self.engine, "adapters", None) + if registry is None or not isinstance(body, dict): + return None + try: + return registry.pick(body, getattr(self, "adapter_models", None) or {}) + except LookupError as exc: + raise RequestError(str(exc)) from exc + def reply_model(self, body: Any) -> str: """The id a reply names: the one the request asked for when this endpoint answers to it, else ``--name``.""" @@ -148,6 +163,10 @@ def _check_fields(self, body: dict[str, Any]) -> str | None: return "this model's CUDA engine has no serial switch (\"draft\": false)" if not isinstance(body.get("messages", []), list): return "messages must be a list" + try: + self.adapter_for(body) + except RequestError as exc: + return str(exc) problem = grammar.refusal(body) # a malformed grammar field, or one beside a required call if problem is None and grammar.request_spec(body) and "constraint" not in inspect.signature( self.engine.generate).parameters: @@ -423,6 +442,7 @@ def prepare(self, body: dict[str, Any], chat: bool) -> PreparedRequest: if problem: raise refusal(problem) prepared = self._prepare(body, chat) + prepared.adapter = self.adapter_for(body) problem = self.check(body, prepared=prepared) if problem: raise refusal(problem) @@ -548,6 +568,10 @@ def on_tokens(new: list[int]) -> bool: options["probabilities"] = probabilities if takes_stop_eos: options["stop_eos"] = not prepared.ignore_eos + if prepared.adapter is not None: + if "adapter" not in inspect.signature(self.engine.generate).parameters: + raise RequestError("this model's engine does not serve adapters") + options["adapter"] = prepared.adapter shaped = prepared.grammar is not None or prepared.think_budget > 0 think_end = self.tok.token_to_id("") if chat and thinking and shaped else None budget = self._think_budget(prepared, think_end) diff --git a/src/tensorfold/cuda/streams.py b/src/tensorfold/cuda/streams.py index 5214c69e4a..d97c8fb52c 100644 --- a/src/tensorfold/cuda/streams.py +++ b/src/tensorfold/cuda/streams.py @@ -31,6 +31,7 @@ class Stream: drafts: list[int] = field(default_factory=list) # an MTP family's drafts for the next round stops: list[int] = field(default_factory=list) # prompt positions whose states the prefill keeps constraint: Any = None # the reply's grammar (tensorfold.engine.grammar), or None + adapter: Any = None # the adapter the reply runs (an engine's reference), or None background: bool = False # priority "background": after, and yielding to, the rest probabilities: Any = None carry: dict | None = None # the stats of the stream this one continues @@ -84,7 +85,7 @@ def continued(self) -> "Stream": """This stream again from its prompt, for later (as the Mac replays): what it sent is owed, not sent again.""" return Stream(self.prompt, self.count, self.sampling, draft=self.draft, stop_eos=self.stop_eos, emit=self.emit, - background=self.background, probabilities=self.probabilities, + background=self.background, probabilities=self.probabilities, adapter=self.adapter, carry=self.stats(), owed=[*self.out, *self.owed]) diff --git a/src/tensorfold/families/deepseek_v41/__init__.py b/src/tensorfold/families/deepseek_v41/__init__.py index 5f2928d2ab..d69e7787fb 100644 --- a/src/tensorfold/families/deepseek_v41/__init__.py +++ b/src/tensorfold/families/deepseek_v41/__init__.py @@ -12,6 +12,7 @@ MODEL_TYPES = ("deepseek_v41",) TITLE = "DeepSeek-V4.1-Flash" CUDA_VISION = True # --vision: the tower and image spans are this family's own (cuda/vision.py) +CUDA_ADAPTERS = True # --adapter: LoRA adapters on attention wo_b and the shared expert's w2 (docs/adapters.md) LANES = False MODELS = ("Mia-AiLab/DeepSeek-V4.1-Flash-EXL3-2.9bpw",) QUANT_METHODS = {"cuda": ("exl3",)} @@ -60,7 +61,8 @@ def cuda_engine(model_dir: str | Path, *, drafter: str = "", tp: int = 2, rank: engine = DsEngine(Path(model_dir), rank=int(rank), world=int(tp), master=master, port=int(master_port), drafts=drafts, context=options.get("context"), engram_dir=os.environ.get("TF_DS_ENGRAM") or None, vision=bool(options.get("vision")), vision_urls=bool(options.get("vision_urls")), - parallel=int(options.get("parallel") or 1)) + parallel=int(options.get("parallel") or 1), adapters=options.get("adapters"), + adapter_reserved=tuple(options.get("adapter_reserved") or ())) fp32_gemms() return engine diff --git a/src/tensorfold/families/deepseek_v41/cuda/app.py b/src/tensorfold/families/deepseek_v41/cuda/app.py index 5b168fb961..397bb82f16 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/app.py +++ b/src/tensorfold/families/deepseek_v41/cuda/app.py @@ -30,6 +30,9 @@ class DsApp(App): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) + registry = getattr(self.engine, "adapters", None) + if registry is not None: # --adapter: each adapter is a model id of its own + self.adapter_models = registry.model_map() # --vision-max-images: images take DsVision.split's path, not the shared image limits, so its own count # check gets the server's value (DsVision's default is 8) vision, images = getattr(self.engine, "vision", None), kwargs.get("vision_max_images") diff --git a/src/tensorfold/families/deepseek_v41/cuda/engine.py b/src/tensorfold/families/deepseek_v41/cuda/engine.py index 84f465bc0b..735a808e25 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/engine.py +++ b/src/tensorfold/families/deepseek_v41/cuda/engine.py @@ -50,6 +50,15 @@ def _drop_pages(fds: list) -> None: _drop_busy.release() +def _same_on_every_rank(flags: list[list[int]]) -> None: + """Refuse a start whose ranks gathered different words (each rank's row of them).""" + + if any(f != flags[0] for f in flags): + raise ValueError("--vision, --parallel, TF_DS_MARKOV, TF_DS_MARKOV_SPLIT and the --adapter files, names, " + "strengths and order must be the same on every rank (rank 0 and the workers run the same " + f"steps): {flags}") + + # drafts verified a round chosen from the confidence head's prefix survival and measured window costs (0: fixed k) ADAPTIVE = os.environ.get("TF_DS_ADAPTIVE", "0") == "1" # before serving: caches at full size, the Triton kernels prompt chunks use, every decode graph (each window size at each @@ -67,7 +76,8 @@ def _drop_pages(fds: list) -> None: class DsEngine: def __init__(self, model_dir: Path, *, rank: int, world: int, master: str, port: int, drafts: int = 3, context: int | None = None, engram_dir: str | None = None, vision: bool = False, - vision_urls: bool = False, parallel: int = 1) -> None: + vision_urls: bool = False, parallel: int = 1, adapters: list | None = None, + adapter_reserved: tuple = ()) -> None: from tensorfold.cuda.comm import NCCL from ..ops import compressed_token_map @@ -106,6 +116,8 @@ def __init__(self, model_dir: Path, *, rank: int, world: int, master: str, port: elif rank == 0: print("[tensorfold] WARNING: no Engram tables (TF_DS_ENGRAM): output will be degraded", flush=True) self.model = Model(self.w, Comm(nccl, world, rdma_bytes=int(RDMA_MB * (1 << 20)) // 64 * 64), eng) + # --adapter: the registry and every rank's share of the factors, before the pools and the warm-up size anything + self.adapters = self._load_adapters(adapters, adapter_reserved, parallel) if adapters else None # images (--vision): the tower runs on rank 0, which shares each image span's rows; every rank routes the span # with the gates' VL bias and keeps Engram out of it self.vision = None @@ -160,14 +172,12 @@ def __init__(self, model_dir: Path, *, rank: int, world: int, master: str, port: from . import markov nccl.barrier() # the Markov switches set the drafter's gathers (count, size) - mine = torch.tensor([int(self.vcfg is not None), int(parallel), int(markov.ON), int(markov.SPLIT)], + mine = torch.tensor([int(self.vcfg is not None), int(parallel), int(markov.ON), int(markov.SPLIT), + *(self.adapters.words() if self.adapters else [0, 0, 0, 0])], dtype=torch.int64, device="cuda") - every = torch.empty((world * 4,), dtype=torch.int64, device="cuda") + every = torch.empty((world * 8,), dtype=torch.int64, device="cuda") nccl.all_gather(mine, every) - flags = every.view(world, 4).tolist() - if any(f != flags[0] for f in flags): - raise ValueError("--vision, --parallel, TF_DS_MARKOV and TF_DS_MARKOV_SPLIT must be the same on every " - f"rank (rank 0 and the workers run the same steps): {flags}") + _same_on_every_rank(every.view(world, 8).tolist()) if WARM: self.warm() if self.concurrent: @@ -188,8 +198,51 @@ def __init__(self, model_dir: Path, *, rank: int, world: int, master: str, port: self._memory_ceiling() memtrace.start(rank) if rank == 0: + held = f", adapters {', '.join(self.adapters.names())}" if self.adapters else "" print(f"[tensorfold] DeepSeek-V4.1 engine ready: {world} rank(s), context {self.limit}, " - f"{'DSpark ' + str(drafts) + ' drafts' if self.drafter else 'serial decode'}", flush=True) + f"{'DSpark ' + str(drafts) + ' drafts' if self.drafter else 'serial decode'}{held}", flush=True) + + def _load_adapters(self, specs: list, reserved: tuple, parallel: int): + """The adapter registry (validated against this checkpoint) and this rank's factors on the GPU; every kernel + they launch is built here. Refused where a row's slot cannot reach the forward: one stream at a time, the + torch path (TF_DS_KERNELS=0) and the exact-matmul test mode.""" + + from . import lora_runtime + from . import model as _model + from .hostenv import meminfo_gib + from .lora_registry import AdapterRegistry + + if int(parallel) < 2: + raise ValueError("--adapter needs --parallel 2 or more: adapters ride the concurrent rounds' per-row inputs") + if not _model.KERNELS or _model.EXACT_MM: + raise ValueError("--adapter runs on the kernel path: TF_DS_KERNELS=0 and TF_DS_EXACT_MM=1 are test modes") + registry = AdapterRegistry.load(specs, self.w.cfg, self.model_dir, reserved) + floor = float(os.environ.get("TF_DS_MEM_FLOOR_GIB") or 2.0) + + def room() -> float | None: # GiB above the floor, a gigabyte kept for the warm-up + free = meminfo_gib("MemAvailable") + return None if free is None else free - floor - 1.0 + + runtime = lora_runtime.LoraRuntime.load(registry, self.w, self.rank, self.world, room=room) + self.model.lora = runtime + runtime.warm() + return registry + + def adapter_ok(self, ref) -> bool: + """Whether this engine holds exactly the adapter ``ref`` names (the stock model always).""" + + return ref is None or (self.adapters is not None and self.adapters.check(ref)) + + def adapter_ref(self, name: str | None): + """The reference a request's adapter name resolves to (None: the stock model); UnknownAdapter if not held.""" + + from .lora_registry import UnknownAdapter + + if name is None: + return None + if self.adapters is None: + raise UnknownAdapter(f"no adapter named {name!r}: this server was started without --adapter") + return self.adapters.ref(name) def _nccl_host_env(self) -> None: """NCCL's cuMem allocations off and its GID index set (hostenv.py); a value that the user set stays.""" @@ -347,15 +400,18 @@ def _share(self, values: list[int] | None) -> list[int]: return [int(v) for v in allv[:count].tolist()] def generate(self, prompt: list[int], max_tokens: int, sampling, on_tokens, draft: bool = True, - constraint=None, vision=None, **_: Any) -> dict[str, Any]: + constraint=None, vision=None, adapter: str | None = None, **_: Any) -> dict[str, Any]: from tensorfold.engine.grammar import pack stop_eos = bool(getattr(self.request, "stop_eos", True)) + ref = self.adapter_ref(adapter) if self.scheduler is not None: if vision is not None and getattr(vision, "spans", None) and self.tower is None: raise ValueError("image input needs the server started with --vision") return self.scheduler.submit(list(prompt), max_tokens, sampling, bool(draft), on_tokens, - stop_eos=stop_eos, vision=vision, constraint=constraint) + stop_eos=stop_eos, vision=vision, constraint=constraint, adapter=ref) + if ref is not None: + raise ValueError("adapters are served by the concurrent engine (--parallel 2 or more)") positions, rows = [], None if vision is not None and getattr(vision, "spans", None): if self.tower is None: @@ -412,22 +468,24 @@ def _sample(self, logits: torch.Tensor, positions: list[int], sampling) -> list[ return sample_rows(logits, positions, sampling) - def prefill(self, sc, dc, prompt: list[int], image: tuple | None = None) -> torch.Tensor: + def prefill(self, sc, dc, prompt: list[int], image: tuple | None = None, adapter: int = 0) -> torch.Tensor: """The prompt into ``sc`` (and the drafter's cache ``dc``, when given) in PREFILL_CHUNK chunks: the last row's logits. ``image``: (ascending positions of the prompt's image-span tokens, their rows [k, dim] bf16).""" - steps = self.prefill_steps(sc, dc, prompt, image) + steps = self.prefill_steps(sc, dc, prompt, image, adapter=adapter) while True: try: next(steps) except StopIteration as done: return done.value - def prefill_steps(self, sc, dc, prompt: list[int], image: tuple | None = None, start: int = 0, snap=None): + def prefill_steps(self, sc, dc, prompt: list[int], image: tuple | None = None, start: int = 0, snap=None, + adapter: int = 0): """``prefill`` one chunk a step (a generator: each ``next`` runs one chunk; its return value is the last row's logits), so a concurrent lane can decode between a long prompt's chunks. ``start`` (a chunk boundary): the caches already hold the prompt's rows before it, a kept prompt's (``multi``), and its chunks are a fresh - prefill's own from there; ``snap(end)`` runs after each chunk that ends on a chunk boundary.""" + prefill's own from there; ``snap(end)`` runs after each chunk that ends on a chunk boundary. ``adapter``: the + registry slot the prompt's rows run (0: stock).""" m = self.model assert start % PREFILL_CHUNK == 0 and start < len(prompt), (start, len(prompt)) @@ -469,7 +527,7 @@ def ahead(a: int) -> None: # the chunk at ``a``: its Engram if i1 > i0: block = (pos_dev[i0:i1] - s, rows[i0:i1]) tc = time.perf_counter() - out = m.forward(sc, ids, s, taps=taps, host_ids=host[s:e], replay=replay, image=block) + out = m.forward(sc, ids, s, taps=taps, host_ids=host[s:e], replay=replay, image=block, adapter=adapter) if ENGRAM_DONTNEED and m.engram is not None and _drop_busy.acquire(blocking=False): fds = sorted({f[0] for f in m.engram.files.values()}) threading.Thread(target=_drop_pages, args=(fds,), daemon=True).start() diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_format.py b/src/tensorfold/families/deepseek_v41/cuda/lora_format.py new file mode 100644 index 0000000000..ead4e4bb16 --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_format.py @@ -0,0 +1,364 @@ +"""A LoRA adapter folder read and validated for DeepSeek-V4.1: PEFT files plus an optional manifest, no torch. + +An adapter edits the output projection of attention (``wo_b``) and the shared expert's down projection (``w2``) of any +backbone layer, at any rank. Every other module, and anything that is not a plain LoRA pair, is refused by name: an edit +is served as the file says or not at all. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import os +import re +from dataclasses import dataclass, field +from pathlib import Path +from types import MappingProxyType +from typing import Any, Mapping + +import numpy as np + +FORMAT = "tensorfold-adapter-1" +CONFIG, WEIGHTS, MANIFEST = "adapter_config.json", "adapter_model.safetensors", "tensorfold_adapter.json" +WO_B, SHARED_W2 = "wo_b", "shared_w2" +KINDS = (WO_B, SHARED_W2) +MAX_RANK = int(os.environ.get("TF_DS_ADAPTER_MAX_RANK") or 256) # a module's rank, before zero components go +MAX_HEADER = 64 << 20 +DTYPES = {"F32": 4, "F16": 2, "BF16": 2} +TILE_J = 16 # the kernels' rank tile, fixed: no other adapter's rank can change how a row's sums are grouped +PAD_J = 4 # tables hold ranks in multiples of this (a 16-byte row of the fp32 factors) + +_PREFIXES = ("base_model.model.model.", "base_model.model.", "model.") +_KIND_NAMES = ((WO_B, re.compile(r"layers\.(\d+)\.attn\.wo_b")), + (SHARED_W2, re.compile(r"layers\.(\d+)\.ffn\.shared_experts\.w2"))) +_LORA = re.compile(r"(?P.+)\.lora_(?P[AB])\.weight") +_REFUSED = ("PEFT writes this field when the adapter is DoRA, has a bias, trains other modules, replicates layers or " + "is an activated LoRA; none of those is served") + + +class AdapterError(ValueError): + """The adapter folder cannot be served as written.""" + + +@dataclass(frozen=True) +class Dims: + """The base model's sizes an adapter's tensors must match.""" + + n_layers: int + dim: int + inter: int + o_in: int # o_groups * o_rank: wo_b's input width + + @classmethod + def of(cls, cfg: Any) -> Dims: + return cls(int(cfg.n_layers), int(cfg.dim), int(cfg.inter), int(cfg.o_groups) * int(cfg.o_rank)) + + def shape(self, kind: str) -> tuple[int, int]: + """(input width, output width) of a module kind's full matrix.""" + + return (self.o_in if kind == WO_B else self.inter), self.dim + + +@dataclass(frozen=True) +class ModuleSpec: + """One edited module as the file stores it. ``scale`` is alpha / r (alpha / sqrt(r) with rsLoRA).""" + + layer: int + kind: str + rank: int + alpha: float + scale: float + a_name: str + b_name: str + + @property + def name(self) -> str: + return module_name(self.layer, self.kind) + + +@dataclass(frozen=True) +class Factors: + """A module's factors in float32 with every all-zero rank component removed: delta = scale * B @ (A @ x).""" + + a: np.ndarray # [rank, in] + b: np.ndarray # [out, rank] + scale: float + + @property + def rank(self) -> int: + return int(self.a.shape[0]) + + +@dataclass(frozen=True) +class AdapterSpec: + path: str + weights: str + sha256: str # of the weights file + modules: tuple[ModuleSpec, ...] + manifest: Mapping[str, Any] = field(default_factory=dict) # the folder's tensorfold_adapter.json (read-only) + + def describe(self) -> dict[str, Any]: + """The normalized manifest: what this adapter edits, derived from the files (a declared one must agree).""" + + return {"format": FORMAT, "weights_sha256": self.sha256, + "modules": {m.name: {"rank": m.rank, "alpha": m.alpha, "scale": m.scale} for m in self.modules}} + + +def padded_rank(rank: int) -> int: + """The rank a table holds for a largest rank of ``rank``.""" + + return -(-max(int(rank), 1) // PAD_J) * PAD_J + + +def module_name(layer: int, kind: str) -> str: + return f"layers.{layer}.attn.wo_b" if kind == WO_B else f"layers.{layer}.ffn.shared_experts.w2" + + +def sha256_file(path: str | Path) -> str: + h = hashlib.sha256() + with open(path, "rb") as f: + for block in iter(lambda: f.read(16 << 20), b""): + h.update(block) + return h.hexdigest() + + +def read_header(path: str | Path) -> tuple[dict[str, dict], int]: + """(tensor entries, data start) of a safetensors file; every offset checked against the file's size.""" + + size = os.path.getsize(path) + with open(path, "rb") as f: + raw = f.read(8) + n = int.from_bytes(raw, "little") if len(raw) == 8 else 0 + if not 2 <= n <= MAX_HEADER or 8 + n > size: + raise AdapterError(f"{path}: not a safetensors file (header length {n})") + try: + header = json.loads(f.read(n)) + except ValueError as exc: + raise AdapterError(f"{path}: unreadable safetensors header ({exc})") from exc + if not isinstance(header, dict): + raise AdapterError(f"{path}: the safetensors header is not an object") + header.pop("__metadata__", None) + end, spans = size - 8 - n, [] + for name, e in header.items(): + try: + dtype, shape, (lo, hi) = e["dtype"], [int(v) for v in e["shape"]], [int(v) for v in e["data_offsets"]] + except (KeyError, TypeError, ValueError, OverflowError) as exc: # OverflowError: Infinity or 1e999 + raise AdapterError(f"{path}: tensor {name!r} has a malformed entry") from exc + if not isinstance(dtype, str) or dtype not in DTYPES: + raise AdapterError(f"{path}: tensor {name!r} is {dtype}; adapters are served as F32, F16 or BF16") + count = math.prod(shape) + if min(shape, default=1) < 0 or not 0 <= lo <= hi <= end or hi - lo != count * DTYPES[dtype]: + raise AdapterError(f"{path}: tensor {name!r} does not fit the file (shape {shape}, bytes {lo}:{hi})") + spans.append((lo, hi, name)) + spans.sort() + for (_, hi, a), (lo, _, b) in zip(spans, spans[1:]): + if lo < hi: + raise AdapterError(f"{path}: tensors {a!r} and {b!r} overlap") + return header, 8 + n + + +def read_tensor(path: str | Path, entry: Mapping[str, Any], base: int) -> np.ndarray: + """One tensor as float32 (BF16 and F16 widen exactly), read without the rest of the file.""" + + lo, hi = entry["data_offsets"] + raw = np.fromfile(path, dtype=np.uint8, count=hi - lo, offset=base + lo) + if raw.size != hi - lo: + raise AdapterError(f"{path}: short read") + shape = [int(v) for v in entry["shape"]] + if entry["dtype"] == "BF16": + out = (raw.view(" str: + for prefix in _PREFIXES: + if key.startswith(prefix): + return key[len(prefix):] + return key + + +def _classify(module: str) -> tuple[int, str]: + for kind, rx in _KIND_NAMES: + m = rx.fullmatch(module) + if m: + return int(m.group(1)), kind + raise AdapterError(f"{module}: not a module this engine can edit (only attention wo_b and the shared expert's w2 " + "of a backbone layer; MTP, routed-expert, router, embedding, head and norm edits are refused)") + + +def _real(value: Any) -> float | None: + """``value`` as a finite float; None for a bool, a string or anything else that is not a finite number.""" + + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + try: + number = float(value) + except OverflowError: # an integer no float can hold + return None + return number if math.isfinite(number) else None + + +def _object(value: Any, where: Path | str, what: str) -> Mapping[str, Any]: + """A JSON member that must be an object when it is present (absent: empty).""" + + if value is None: + return {} + if not isinstance(value, dict): + raise AdapterError(f"{where}: {what} must be an object") + return value + + +def _pattern(table: Any, module: str, what: str) -> Any: + """PEFT's rank_pattern / alpha_pattern value for a module: the first key that matches it, else None.""" + + for key, value in _object(table, CONFIG, what).items(): + try: + if re.fullmatch(rf"(?:.*\.)?(?:{key})", module) or re.fullmatch(rf"(?:.*\.)?(?:{key})", + "base_model.model." + module): + return value + except re.error as exc: + raise AdapterError(f"{CONFIG}: {what} key {key!r} is not a pattern ({exc})") from exc + return None + + +def read_config(folder: Path) -> dict[str, Any]: + try: + cfg = json.loads((folder / CONFIG).read_text()) + except (OSError, ValueError) as exc: + raise AdapterError(f"{folder}: {CONFIG} is missing or unreadable ({exc})") from exc + if not isinstance(cfg, dict) or str(cfg.get("peft_type", "")).upper() != "LORA": + raise AdapterError(f"{folder / CONFIG}: peft_type must be LORA") + for key in ("use_dora", "fan_in_fan_out", "lora_bias", "layer_replication", "modules_to_save", + "target_parameters", "trainable_token_indices"): + if cfg.get(key): + raise AdapterError(f"{folder / CONFIG}: {key} is set. {_REFUSED}") + # null is a plain LoRA; any other value is PEFT's activated LoRA, which edits only after its invocation tokens + if cfg.get("alora_invocation_tokens") is not None: + raise AdapterError(f"{folder / CONFIG}: alora_invocation_tokens is set. {_REFUSED}") + if str(cfg.get("bias", "none")) != "none": + raise AdapterError(f"{folder / CONFIG}: bias is {cfg['bias']!r}. {_REFUSED}") + for key in ("r", "lora_alpha"): + number = _real(cfg.get(key)) + if number is None or number <= 0: + raise AdapterError(f"{folder / CONFIG}: {key} must be a finite positive number") + return cfg + + +def _declared(folder: Path) -> dict[str, Any]: + path = folder / MANIFEST + if not path.is_file(): + return {} + try: + manifest = json.loads(path.read_text()) + except (OSError, ValueError) as exc: + raise AdapterError(f"{path}: unreadable ({exc})") from exc + if not isinstance(manifest, dict) or manifest.get("format") != FORMAT: + raise AdapterError(f"{path}: format must be {FORMAT!r}") + return manifest + + +def _check_manifest(manifest: Mapping[str, Any], folder: Path, sha: str, modules: tuple[ModuleSpec, ...], + dims: Dims, config_sha: str | None) -> None: + """A declared manifest is a contract: the file hash, the base model and the exact module set must agree.""" + + if not manifest: + return + files = _object(manifest.get("files"), folder / MANIFEST, "files") + want = _object(files.get(WEIGHTS), folder / MANIFEST, f"files.{WEIGHTS}").get("sha256") + if want is not None and want != sha: + raise AdapterError(f"{folder / MANIFEST}: {WEIGHTS} hashes to {sha}, the manifest says {want}") + base = _object(manifest.get("base"), folder / MANIFEST, "base") + _object(manifest.get("provenance"), folder / MANIFEST, "provenance") # --write-manifest copies it + for key, have in (("model_type", "deepseek_v41"), ("hidden_size", dims.dim), ("num_hidden_layers", dims.n_layers)): + if key in base and base[key] != have: + raise AdapterError(f"{folder / MANIFEST}: base {key} is {base[key]!r}, this model's is {have!r}") + if base.get("config_sha256") not in (None, config_sha): + raise AdapterError(f"{folder / MANIFEST}: it was made for another base config (config.json sha256 differs)") + declared = manifest.get("modules") + if declared is None: + return + if not isinstance(declared, dict): + raise AdapterError(f"{folder / MANIFEST}: modules must be an object") + found = {m.name: m.rank for m in modules} + given = {k: (v.get("rank") if isinstance(v, dict) else v) for k, v in declared.items()} + if given != found: + extra, lost = sorted(set(found) - set(given)), sorted(set(given) - set(found)) + wrong = sorted(k for k in set(found) & set(given) if found[k] != given[k]) + raise AdapterError(f"{folder / MANIFEST}: the module list and the tensors differ (not declared: {extra[:4]}, " + f"declared but absent: {lost[:4]}, other rank: {wrong[:4]})") + + +def inspect_folder(folder: str | Path, dims: Dims, *, config_sha: str | None = None) -> AdapterSpec: + """Validate an adapter folder against the base model's sizes. Raises AdapterError naming what is wrong.""" + + folder = Path(folder) + if not folder.is_dir(): + raise AdapterError(f"{folder}: not a directory") + cfg = read_config(folder) + weights = folder / WEIGHTS + if not weights.is_file(): + raise AdapterError(f"{folder}: {WEIGHTS} is missing") + manifest = _declared(folder) + header, _ = read_header(weights) + pairs: dict[str, dict[str, str]] = {} + for key in header: + m = _LORA.fullmatch(key) + if m is None: + raise AdapterError(f"{key}: not a LoRA A or B weight ({WEIGHTS} may hold only lora_A.weight and " + "lora_B.weight pairs; DoRA magnitudes, biases, embeddings and full weights are refused)") + module = _canonical(m["module"]) + if m["part"] in pairs.setdefault(module, {}): + raise AdapterError(f"{module}: two {m['part']} tensors") + pairs[module][m["part"]] = key + if not pairs: + raise AdapterError(f"{weights}: no LoRA tensors") + modules = [] + for module, ab in sorted(pairs.items()): + layer, kind = _classify(module) + if layer >= dims.n_layers: + raise AdapterError(f"{module}: layer {layer} of {dims.n_layers}") + if set(ab) != {"A", "B"}: + raise AdapterError(f"{module}: lora_{'B' if 'A' in ab else 'A'} is missing") + (a_in, b_out), (a_shape, b_shape) = dims.shape(kind), (header[ab["A"]]["shape"], header[ab["B"]]["shape"]) + if len(a_shape) != 2 or len(b_shape) != 2 or a_shape[1] != a_in or b_shape[0] != b_out or \ + a_shape[0] != b_shape[1]: + raise AdapterError(f"{module}: lora_A {a_shape} and lora_B {b_shape} do not fit a [{b_out}, {a_in}] matrix") + rank = int(a_shape[0]) + if not 1 <= rank <= MAX_RANK: + raise AdapterError(f"{module}: rank {rank} outside 1..{MAX_RANK} (TF_DS_ADAPTER_MAX_RANK)") + want = _pattern(cfg.get("rank_pattern"), module, "rank_pattern") + want = cfg["r"] if want is None else want + if want != rank: + raise AdapterError(f"{module}: the tensors have rank {rank}, {CONFIG} says {want}") + given = _pattern(cfg.get("alpha_pattern"), module, "alpha_pattern") + alpha = _real(cfg["lora_alpha"] if given is None else given) + if alpha is None: + raise AdapterError(f"{module}: alpha_pattern gives {given!r}, not a finite number") + scale = alpha / math.sqrt(rank) if cfg.get("use_rslora") else alpha / rank + modules.append(ModuleSpec(layer, kind, rank, alpha, scale, ab["A"], ab["B"])) + modules.sort(key=lambda m: (m.layer, m.kind)) + sha = sha256_file(weights) + spec = tuple(modules) + _check_manifest(manifest, folder, sha, spec, dims, config_sha) + return AdapterSpec(str(folder), str(weights), sha, spec, MappingProxyType(manifest)) + + +def read_factors(spec: AdapterSpec, module: ModuleSpec, parsed: tuple[dict, int] | None = None) -> Factors | None: + """The module's factors in float32, all-zero components dropped (their product is exactly zero); None when + nothing is left. A value that is not finite is refused: it would reach every reply. ``parsed``: read_header's + result for the weights file, when the caller reads several modules.""" + + header, base = parsed or read_header(spec.weights) + a = read_tensor(spec.weights, header[module.a_name], base) + b = read_tensor(spec.weights, header[module.b_name], base) + if not (np.isfinite(a).all() and np.isfinite(b).all()): + raise AdapterError(f"{module.name}: factors hold NaN or infinity") + live = np.flatnonzero(a.any(axis=1) & b.any(axis=0)) + if live.size == 0: + return None + return Factors(np.ascontiguousarray(a[live]), np.ascontiguousarray(b[:, live]), module.scale) diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_kernels.py b/src/tensorfold/families/deepseek_v41/cuda/lora_kernels.py new file mode 100644 index 0000000000..0bd56d9f62 --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_kernels.py @@ -0,0 +1,84 @@ +"""Row-invariant Triton kernels for the adapter delta: a row's bits never depend on the other rows or on their adapters. + +``delta[r] = scale[i] * B[i] @ (A[i] @ x[r])`` for the adapter slot ``i`` of row ``r``, summed in fp32 in a fixed order +over the input width and then over the rank, one program per row and tile. A row of slot 0 is not read or written, so a +stock row keeps its bits exactly. Slots, ranks and scales are read from device tables, and the slot of every row from a +device vector: a captured graph runs any mixture of adapters. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from .lora_tables import TILE_J, Tables + +TILE_K, TILE_N = 128, 128 + + +@triton.jit(do_not_specialize=["xs", "IDS"]) +def _project(X, xs, IDS, A, RANKS, T, K: tl.constexpr, R: tl.constexpr, BJ: tl.constexpr, BK: tl.constexpr): + """T[r, j] = sum_k A[slot, j, k] * x[r, k] for the rank tile of this program: fixed K order, fp32.""" + + r = tl.program_id(0) + jb = tl.program_id(1) + ad = tl.load(IDS + r) + rank = tl.load(RANKS + ad) + j = jb * BJ + tl.arange(0, BJ) + kk = tl.arange(0, BK) + live = j < rank + kend = tl.where(jb * BJ < rank, K, 0) + acc = tl.zeros((BJ,), dtype=tl.float32) + for k0 in range(0, kend, BK): + ok = (k0 + kk) < K + x = tl.load(X + r * xs + k0 + kk, mask=ok, other=0.0).to(tl.float32) + a = tl.load(A + (ad * R + j[:, None]) * K + (k0 + kk)[None, :], mask=live[:, None] & ok[None, :], other=0.0) + acc += tl.sum(a * x[None, :], axis=1) + tl.store(T + r * R + j, acc, mask=j < R) + + +@triton.jit(do_not_specialize=["os_", "IDS"]) +def _apply(OUT, os_, IDS, B, T, RANKS, SCALES, D: tl.constexpr, R: tl.constexpr, BN: tl.constexpr, BJ: tl.constexpr): + """out[r, n] += scale * sum_j B[slot, n, j] * T[r, j]: fixed rank order; a stock row's program does nothing.""" + + r = tl.program_id(0) + nb = tl.program_id(1) + ad = tl.load(IDS + r) + rank = tl.load(RANKS + ad) + if rank > 0: + nn = nb * BN + tl.arange(0, BN) + jj = tl.arange(0, BJ) + ok = nn < D + acc = tl.zeros((BN,), dtype=tl.float32) + for j0 in range(0, rank, BJ): + live = (j0 + jj) < rank + t = tl.load(T + r * R + j0 + jj, mask=live, other=0.0) + b = tl.load(B + (ad * D + nn[:, None]) * R + (j0 + jj)[None, :], mask=ok[:, None] & live[None, :], + other=0.0) + acc += tl.sum(b * t[None, :], axis=1) + cur = tl.load(OUT + r * os_ + nn, mask=ok, other=0.0) + tl.store(OUT + r * os_ + nn, cur + acc * tl.load(SCALES + ad), mask=ok) + + +def lowrank_add(out: torch.Tensor, x: torch.Tensor, ids: torch.Tensor, tab: Tables) -> torch.Tensor: + """out [n, D] fp32 += each row's adapter delta of x [n, K] (any float dtype, unit column stride); ``ids`` [n] int64 + holds each row's slot. Two launches, no host sync: safe inside a CUDA graph capture.""" + + n = out.shape[0] + if out.dtype != torch.float32 or out.dim() != 2 or out.shape[1] != tab.D or out.stride(1) != 1: + raise ValueError(f"out must be a float32 [rows, {tab.D}] tensor with unit column stride") + if x.dim() != 2 or x.shape != (n, tab.K) or x.stride(1) != 1 or x.dtype != tab.x_dtype: + raise ValueError(f"x must be a {tab.x_dtype} [{n}, {tab.K}] tensor with unit column stride, got " + f"{x.dtype} {tuple(x.shape)}") + if ids.dtype != torch.int64 or ids.dim() != 1 or ids.shape[0] != n or ids.stride(0) != 1: + raise ValueError("ids must be a contiguous int64 vector with a slot a row") + if n == 0: + return out + t = torch.empty((n, tab.R), dtype=torch.float32, device=out.device) + # one tile shape for every table: how a row's sums are grouped never depends on the largest rank of any adapter + _project[(n, triton.cdiv(tab.R, TILE_J))](x, x.stride(0), ids, tab.a, tab.ranks, t, K=tab.K, R=tab.R, BJ=TILE_J, + BK=TILE_K, num_warps=4) + _apply[(n, triton.cdiv(tab.D, TILE_N))](out, out.stride(0), ids, tab.b, t, tab.ranks, tab.scales, D=tab.D, + R=tab.R, BN=TILE_N, BJ=TILE_J, num_warps=4) + return out diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_place.py b/src/tensorfold/families/deepseek_v41/cuda/lora_place.py new file mode 100644 index 0000000000..9d7644441a --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_place.py @@ -0,0 +1,82 @@ +"""Where an adapter's factors live on a tensor-parallel rank, in numpy (no torch). + +Both edited modules are row-parallel: a rank holds a slice of the input width and adds a partial sum to the layer's +existing cross-rank reduction, so the factor A is cut by columns and B is whole on every rank. The packed shared expert +needs one more step: its down projection reads the hidden state only after the expert kernels have scaled and rotated it +(``fold_rotated``), so A is rewritten once, at load, to read that form. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import numpy as np + +from .lora_format import SHARED_W2, WO_B, Dims, Factors + +HAD = 128 # the EXL3 rotation's block width + + +def hadamard() -> np.ndarray: + """The 128 x 128 Walsh-Hadamard matrix in natural order divided by sqrt(128): symmetric, and its own inverse.""" + + i = np.arange(HAD) + both, par = i[:, None] & i[None, :], np.zeros((HAD, HAD), dtype=np.int64) + for bit in range(7): + par ^= (both >> bit) & 1 + return (1 - 2 * par) / math.sqrt(HAD) + + +def span(total: int, rank: int, world: int) -> tuple[int, int]: + """The columns [lo, hi) of an input width that rank ``rank`` of ``world`` holds.""" + + if world < 1 or not 0 <= rank < world or total % world or (total // world) % HAD: + raise ValueError(f"input width {total} does not split into {world} ranks of whole {HAD}-wide blocks") + per = total // world + return rank * per, (rank + 1) * per + + +def fold_rotated(a: np.ndarray, suh: np.ndarray) -> np.ndarray: + """A' with A' @ xd == A @ h, where xd = blockwise_hadamard(h * suh) is the form the expert kernels hand on: A' is + the blockwise Hadamard of A / suh (computed in float64, stored float32).""" + + suh = np.asarray(suh, dtype=np.float64) + if a.ndim != 2 or suh.shape != (a.shape[1],) or a.shape[1] % HAD: + raise ValueError(f"cannot fold factors {a.shape} with a scale vector {suh.shape}") + if not np.isfinite(suh).all() or not suh.all(): + raise ValueError("the expert's input scale vector holds zero or a value that is not finite") + z = (a.astype(np.float64) / suh[None, :]).reshape(a.shape[0], -1, HAD) @ hadamard() + return z.reshape(a.shape).astype(np.float32) + + +@dataclass(frozen=True) +class Placed: + """One rank's share of a module: ``a`` [rank, local input], ``b`` [output, rank], delta = scale * b @ (a @ x).""" + + a: np.ndarray + b: np.ndarray + scale: float + folded: bool # a reads the packed shared expert's rotated hidden, not the plain one + + @property + def rank(self) -> int: + return int(self.a.shape[0]) + + +def place(kind: str, factors: Factors, dims: Dims, rank: int, world: int, strength: float = 1.0, + suh: np.ndarray | None = None) -> Placed: + """This rank's slice of ``factors``. ``suh``: the packed shared expert's input scale vector for these columns.""" + + if kind not in (WO_B, SHARED_W2): + raise ValueError(f"no placement for {kind!r}") + width, out = dims.shape(kind) + if factors.a.shape[1] != width or factors.b.shape[0] != out: + raise ValueError(f"factors {factors.a.shape} / {factors.b.shape} do not fit a [{out}, {width}] matrix") + lo, hi = span(width, rank, world) + a = np.ascontiguousarray(factors.a[:, lo:hi]) + if suh is not None: + if kind != SHARED_W2: + raise ValueError("only the packed shared expert reads a rotated input") + a = fold_rotated(a, suh) + return Placed(a, np.ascontiguousarray(factors.b), float(factors.scale) * float(strength), suh is not None) diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_reference.py b/src/tensorfold/families/deepseek_v41/cuda/lora_reference.py new file mode 100644 index 0000000000..04b5cc7fcb --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_reference.py @@ -0,0 +1,21 @@ +"""The adapter delta in plain torch, a row at a time: the definition the Triton kernels are tested against.""" + +from __future__ import annotations + +import torch + + +def lowrank_add(out: torch.Tensor, x: torch.Tensor, ids: torch.Tensor, a: torch.Tensor, b: torch.Tensor, + ranks: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype = torch.float32) -> torch.Tensor: + """out[r] += scales[i] * (b[i, :, :k] @ (a[i, :k] @ x[r])) for every row r whose slot i is not 0 (k = ranks[i]); + other rows are not touched. ``a`` [slots, R, K], ``b`` [slots, D, R]; sums in ``dtype``; out is updated in place.""" + + for r in range(out.shape[0]): + i, k = int(ids[r]), 0 + if i: + k = int(ranks[i]) + if k == 0: + continue + t = a[i, :k].to(dtype) @ x[r].to(dtype) + out[r] += (float(scales[i]) * (b[i, :, :k].to(dtype) @ t)).to(out.dtype) + return out diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_registry.py b/src/tensorfold/families/deepseek_v41/cuda/lora_registry.py new file mode 100644 index 0000000000..9807c45bab --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_registry.py @@ -0,0 +1,184 @@ +"""The adapters a server holds: fixed at start, each with an identity that names its effect on every cached row. + +Slot 0 is the stock model and has no entry. An identity key is the same on every rank and in every process that loads +the same files at the same strength; it is what keeps a kept prompt, a spilled prompt or a graph input from meeting +another adapter's rows. +""" + +from __future__ import annotations + +import hashlib +import json +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Iterable, Mapping, Sequence + +from .lora_format import FORMAT, AdapterError, AdapterSpec, Dims, inspect_folder + +NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:/-]{0,95}") + + +class UnknownAdapter(LookupError): + """A request named an adapter this server does not hold.""" + + +@dataclass(frozen=True) +class AdapterRef: + """What travels with a request: the slot to run and the identity the slot must have on every rank.""" + + slot: int + key: str + + +@dataclass(frozen=True) +class AdapterEntry: + slot: int + name: str + key: str + strength: float + spec: AdapterSpec + + @property + def ref(self) -> AdapterRef: + return AdapterRef(self.slot, self.key) + + @property + def model_id(self) -> str: + return self.name + + +def adapter_key(spec: AdapterSpec, strength: float) -> str: + """32 hex digits that change with the factors' bytes, the module set and every module's effective scale.""" + + h = hashlib.sha256(f"{FORMAT}\0{spec.sha256}\0".encode()) + for m in spec.modules: + h.update(f"{m.name} {m.rank} {float(m.scale * strength).hex()}\0".encode()) + return h.hexdigest()[:32] + + +def slot_of(obj: Any) -> int: + """The registry slot of a stream or reference (0: stock).""" + + ref = obj if isinstance(obj, AdapterRef) or obj is None else getattr(obj, "adapter", None) + return 0 if ref is None else int(ref.slot) + + +def key_of(obj: Any) -> str: + """The identity key of a stream or reference ('' for stock).""" + + ref = obj if isinstance(obj, AdapterRef) or obj is None else getattr(obj, "adapter", None) + return "" if ref is None else str(ref.key) + + +class AdapterRegistry: + """Immutable after construction: entries, lookups and a digest every rank must agree on.""" + + def __init__(self, entries: Sequence[AdapterEntry], fingerprint: str = "") -> None: + self.entries = tuple(entries) + self.fingerprint = str(fingerprint) + for want, e in enumerate(self.entries, 1): + if e.slot != want: + raise ValueError(f"adapter {e.name!r} holds slot {e.slot}, expected {want}") + self._by_name = {e.name: e for e in self.entries} + if len(self._by_name) != len(self.entries): + raise ValueError("two adapters share a name") + + @classmethod + def load(cls, specs: Iterable[tuple[str, str, float]], cfg: Any, model_dir: str | Path, + reserved: Iterable[str] = ()) -> AdapterRegistry: + """Read and validate ``(name, folder, strength)`` specs against the base model. ``reserved``: model ids a + name may not take (--name and the aliases).""" + + dims, taken, entries = Dims.of(cfg), {str(r) for r in reserved}, [] + config = Path(model_dir) / "config.json" + sha = hashlib.sha256(config.read_bytes()).hexdigest() if config.is_file() else "" + for slot, (name, folder, strength) in enumerate(specs, 1): + if not NAME.fullmatch(str(name)): + raise AdapterError(f"adapter name {name!r}: letters, digits and . _ : / - only, 1 to 96 characters") + if name in taken: + raise AdapterError(f"adapter name {name!r} is already a model id of this server") + strength = float(strength) + if not math.isfinite(strength) or strength == 0.0: + raise AdapterError(f"adapter {name!r}: strength must be a finite number other than 0") + taken.add(name) + spec = inspect_folder(folder, dims, config_sha=sha or None) + entries.append(AdapterEntry(slot, name, adapter_key(spec, strength), strength, spec)) + return cls(entries, sha) + + # -- lookups ------------------------------------------------------------------------------------------------ + def __len__(self) -> int: + return len(self.entries) + + @property + def slots(self) -> int: + """Slots a table needs: the stock slot and one a adapter.""" + + return len(self.entries) + 1 + + def names(self) -> list[str]: + return [e.name for e in self.entries] + + def model_map(self) -> dict[str, str]: + """Model id -> adapter name, for the ids this registry adds to /v1/models.""" + + return {e.model_id: e.name for e in self.entries} + + def ref(self, name: str | None) -> AdapterRef | None: + """The reference for an adapter name (None: stock); UnknownAdapter when no such adapter is held.""" + + if name is None: + return None + entry = self._by_name.get(name) + if entry is None: + raise UnknownAdapter(f"no adapter named {name!r} (this server holds: {', '.join(self.names()) or 'none'})") + return entry.ref + + def entry(self, slot: int) -> AdapterEntry: + return self.entries[int(slot) - 1] + + def check(self, ref: AdapterRef | None) -> bool: + """Whether this registry holds exactly the adapter ``ref`` names (stock always).""" + + if ref is None: + return True + return 1 <= ref.slot <= len(self.entries) and self.entries[ref.slot - 1].key == ref.key + + def pick(self, body: Mapping[str, Any], models: Mapping[str, str]) -> str | None: + """The adapter name a request selects: its model id when that is an adapter's, else its "adapter" field. + The two may not name different adapters; an unknown name is an error, never a quiet fall back to stock.""" + + asked = body.get("model") + by_model = models.get(asked) if isinstance(asked, str) else None + field = body.get("adapter") + if field is not None and (not isinstance(field, str) or field not in self._by_name): + raise UnknownAdapter(f"adapter {field!r} is not held by this server (it holds: " + f"{', '.join(self.names()) or 'none'})") + if field is not None and by_model is not None and field != by_model: + raise UnknownAdapter(f"model {asked!r} is adapter {by_model!r}, but the request's adapter is {field!r}") + return field if field is not None else by_model + + # -- agreement ---------------------------------------------------------------------------------------------- + def digest(self) -> str: + """One hash of the base config, every slot, name and key: equal on two ranks exactly when they hold the same.""" + + body = json.dumps([self.fingerprint, [[e.slot, e.name, e.key] for e in self.entries]], separators=(",", ":")) + return hashlib.sha256(body.encode()).hexdigest() + + def words(self) -> list[int]: + """The digest as four 31-bit integers (what the ranks' start-up check gathers).""" + + raw = int(self.digest()[:32], 16) + return [(raw >> (31 * i)) & 0x7FFFFFFF for i in range(4)] + + def describe(self) -> list[dict[str, Any]]: + out = [] + for e in self.entries: + kinds: dict[str, int] = {} + for m in e.spec.modules: + kinds[m.kind] = kinds.get(m.kind, 0) + 1 + out.append({"slot": e.slot, "name": e.name, "key": e.key, "strength": e.strength, + "modules": kinds, "max_rank": max(m.rank for m in e.spec.modules), + "weights_sha256": e.spec.sha256}) + return out diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_runtime.py b/src/tensorfold/families/deepseek_v41/cuda/lora_runtime.py new file mode 100644 index 0000000000..a36e4ac8b0 --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_runtime.py @@ -0,0 +1,170 @@ +"""Adapter factors on one rank's device, and the hooks the model calls. Loaded once before warm-up, never written again. + +Every edited (layer, module) has one set of tables with a slot per adapter; a row's slot comes from a device vector, +so one captured graph serves any mixture of stock and adapter rows. The hooks add to the module's fp32 partial sum +before the layer's existing cross-rank reduction: no collective is added, and a layer no adapter edits runs nothing. +""" + +from __future__ import annotations + +import time +from typing import Any, Callable + +import torch + +from .lora_format import SHARED_W2, WO_B, AdapterError, Dims, read_factors, read_header +from .lora_place import place +from .lora_registry import AdapterRegistry +from .lora_tables import Tables + + +class LayerTables: + """The tables of one layer's edited modules (None: no adapter edits it).""" + + __slots__ = ("wo_b", "shared") + + def __init__(self) -> None: + self.wo_b: Tables | None = None + self.shared: Tables | None = None + + +class RowAdapters: + """One forward's rows and the adapter slot of each: ``rows`` is an int64 device vector, one slot a row.""" + + __slots__ = ("rt", "rows") + + def __init__(self, rt: LoraRuntime, rows: torch.Tensor) -> None: + self.rt, self.rows = rt, rows + + def wo_b(self, lay: Any, u: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """``out`` [n, d] fp32, the layer's wo_b partial for input ``u`` [n, local groups * o_rank] bf16, plus the + rows' adapter deltas.""" + + t = self.rt.layers.get(lay.idx) + if t is not None and t.wo_b is not None: + self.rt.add(out, u, self.rows, t.wo_b) + return out + + def shared(self, lay: Any, x: torch.Tensor, out: torch.Tensor, packed: bool) -> torch.Tensor: + """``out`` plus the shared expert's w2 delta. ``x``: the expert kernels' rotated hidden (fp16, ``packed``) or + the plain hidden (bf16) of a shared expert outside the pack.""" + + t = self.rt.layers.get(lay.idx) + tab = None if t is None else t.shared + if tab is not None: + if tab.folded != packed: + raise RuntimeError(f"layer {lay.idx}: the shared expert's factors were placed for the " + f"{'plain' if packed else 'rotated'} hidden state") + self.rt.add(out, x, self.rows, tab) + return out + + +class LoraRuntime: + """The tables of every edited layer on this rank. ``kernel`` is the function adding a delta (Triton by default).""" + + def __init__(self, registry: AdapterRegistry, layers: dict[int, LayerTables], device: str = "cuda", + kernel: Callable | None = None) -> None: + self.registry, self.layers, self.device, self.kernel = registry, layers, device, kernel + + @property + def nbytes(self) -> int: + return sum(tab.nbytes for t in self.layers.values() for tab in (t.wo_b, t.shared) if tab is not None) + + def add(self, out: torch.Tensor, x: torch.Tensor, rows: torch.Tensor, tab: Tables) -> torch.Tensor: + fn = self.kernel + if fn is None: + from .lora_kernels import lowrank_add as fn + + self.kernel = fn + return fn(out, x, rows, tab) + + def uniform(self, slot: int, n: int) -> RowAdapters | None: + """Rows [n] that all run adapter ``slot`` (a prompt chunk's); None for the stock slot: it runs no hook.""" + + if not slot: + return None + return RowAdapters(self, torch.full((n,), int(slot), dtype=torch.int64, device=self.device)) + + def bind(self, rows: torch.Tensor) -> RowAdapters: + """Rows whose slots a graph's input vector holds (read at replay, not at capture).""" + + return RowAdapters(self, rows) + + def warm(self) -> int: + """Build every kernel variant the model will launch (one launch each on zero rows), so no request compiles one. + Returns the number of table shapes warmed.""" + + seen = set() + for t in self.layers.values(): + for tab in (t.wo_b, t.shared): + if tab is None: + continue + key = (tab.K, tab.D, tab.R, tab.x_dtype) + if key in seen: + continue + seen.add(key) + out = torch.zeros((1, tab.D), dtype=torch.float32, device=self.device) + x = torch.zeros((1, tab.K), dtype=tab.x_dtype, device=self.device) + self.add(out, x, torch.zeros((1,), dtype=torch.int64, device=self.device), tab) + if self.device != "cpu": + torch.cuda.synchronize() + return len(seen) + + def describe(self) -> dict[str, Any]: + return {"adapters": self.registry.describe(), "layers": len(self.layers), "bytes": self.nbytes} + + @classmethod + def load(cls, registry: AdapterRegistry, weights: Any, rank: int, world: int, *, device: str = "cuda", + log: Callable[[str], None] = print, kernel: Callable | None = None, + room: Callable[[], float | None] | None = None) -> LoraRuntime: + """Read each adapter's factors, place this rank's share and upload it, one (layer, module) at a time (no + second copy of an adapter in host memory). ``room()``: GiB the host can still give, None if unknown; tables + that would not leave the floor that remains are refused before they are allocated.""" + + t0 = time.perf_counter() + dims = Dims.of(weights.cfg) + wanted: dict[tuple[int, str], list] = {} + for e in registry.entries: + for m in e.spec.modules: + wanted.setdefault((m.layer, m.kind), []).append((e, m)) + headers = {e.slot: read_header(e.spec.weights) for e in registry.entries} + layers: dict[int, LayerTables] = {} + total = dropped = 0 + for (layer, kind), items in sorted(wanted.items()): + if layer >= len(weights.layers): + raise AdapterError(f"layer {layer}: this engine holds {len(weights.layers)} layers") + lay = weights.layers[layer] + loaded = [] + for entry, module in items: + f = read_factors(entry.spec, module, headers[entry.slot]) + if f is None: + dropped += 1 + else: + loaded.append((entry, f)) + if not loaded: + continue + packed = kind == SHARED_W2 and getattr(lay, "shared", None) is None + suh = lay.experts.suh_d[lay.experts.count - 1].float().cpu().numpy() if packed else None + width, out = dims.shape(kind) + dtype = torch.float16 if packed else torch.bfloat16 + max_rank = max(f.rank for _, f in loaded) + needed = Tables.required_bytes(registry.slots, max_rank, width // world, out) + free = None if room is None else room() + if free is not None and needed > free * 2**30: + raise MemoryError(f"adapter factors for layer {layer} {kind} need {needed / 2**20:.0f} MiB, the " + f"host has {free:.2f} GiB to give above its floor") + tab = Tables.empty(registry.slots, max_rank, width // world, out, packed, dtype, device) + for entry, f in loaded: + p = place(kind, f, dims, rank, world, entry.strength, suh) + tab.put(entry.slot, p.a, p.b, p.scale) + total += tab.nbytes + slot = layers.setdefault(layer, LayerTables()) + if kind == WO_B: + slot.wo_b = tab + else: + slot.shared = tab + rt = cls(registry, layers, device, kernel) + names = ", ".join(f"{e.name} ({len(e.spec.modules)} modules)" for e in registry.entries) + log(f"[tensorfold] rank {rank}: adapters {names}: {len(layers)} layers, {total / 2**20:.1f} MiB of factors" + f"{f', {dropped} all-zero modules left out' if dropped else ''} in {time.perf_counter() - t0:.1f} s") + return rt diff --git a/src/tensorfold/families/deepseek_v41/cuda/lora_tables.py b/src/tensorfold/families/deepseek_v41/cuda/lora_tables.py new file mode 100644 index 0000000000..7a5d1e3e75 --- /dev/null +++ b/src/tensorfold/families/deepseek_v41/cuda/lora_tables.py @@ -0,0 +1,78 @@ +"""One module's factors for every adapter slot of one layer, as the kernels read them (torch only, no Triton).""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from .lora_format import PAD_J, TILE_J, padded_rank + +__all__ = ["PAD_J", "TILE_J", "Tables", "padded_rank"] + + +@dataclass +class Tables: + """Zero beyond a slot's rank; slot 0 (the stock model) is all zero and has rank 0.""" + + a: torch.Tensor # [slots, R, K] float32 + b: torch.Tensor # [slots, D, R] float32 + ranks: torch.Tensor # [slots] int32, 0 for the stock slot and for an adapter without this module + scales: torch.Tensor # [slots] float32 + folded: bool # a reads the packed shared expert's rotated hidden (lora_place.fold_rotated) + x_dtype: torch.dtype # the dtype of the activations this module reads + + @property + def slots(self) -> int: + return int(self.a.shape[0]) + + @property + def R(self) -> int: + return int(self.a.shape[1]) + + @property + def K(self) -> int: + return int(self.a.shape[2]) + + @property + def D(self) -> int: + return int(self.b.shape[1]) + + @property + def tile_j(self) -> int: + return TILE_J + + @property + def nbytes(self) -> int: + return sum(int(t.numel() * t.element_size()) for t in (self.a, self.b, self.ranks, self.scales)) + + @staticmethod + def required_bytes(slots: int, rank: int, k: int, d: int) -> int: + """Bytes for the padded fp32 factors, int32 ranks and fp32 scales, without allocating a tensor.""" + + return int(slots) * (padded_rank(rank) * (int(k) + int(d)) + 2) * 4 + + @classmethod + def empty(cls, slots: int, rank: int, k: int, d: int, folded: bool, x_dtype: torch.dtype, + device: str | torch.device = "cuda") -> Tables: + """All-zero tables for ``slots`` slots of up to ``rank`` components (rounded up to PAD_J).""" + + r = padded_rank(rank) + return cls(torch.zeros((slots, r, k), dtype=torch.float32, device=device), + torch.zeros((slots, d, r), dtype=torch.float32, device=device), + torch.zeros((slots,), dtype=torch.int32, device=device), + torch.zeros((slots,), dtype=torch.float32, device=device), folded, x_dtype) + + def put(self, slot: int, a, b, scale: float) -> None: + """Slot ``slot``'s factors: a [rank, K], b [D, rank] (numpy or torch float32), delta = scale * b @ (a @ x).""" + + a, b = torch.as_tensor(a, dtype=torch.float32), torch.as_tensor(b, dtype=torch.float32) + rank = int(a.shape[0]) + if slot < 1 or slot >= self.slots or rank > self.R or tuple(a.shape) != (rank, self.K) \ + or tuple(b.shape) != (self.D, rank): + raise ValueError(f"slot {slot}: factors {tuple(a.shape)} / {tuple(b.shape)} do not fit tables " + f"[{self.slots}, {self.R}, {self.K}] / [{self.D}]") + self.a[slot, :rank].copy_(a) + self.b[slot, :, :rank].copy_(b) + self.ranks[slot] = rank + self.scales[slot] = float(scale) diff --git a/src/tensorfold/families/deepseek_v41/cuda/model.py b/src/tensorfold/families/deepseek_v41/cuda/model.py index c33ee49303..2baf91537a 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/model.py +++ b/src/tensorfold/families/deepseek_v41/cuda/model.py @@ -414,8 +414,10 @@ def wo_a_rot(lay, n: int, device, head_dim: int) -> tuple: return suh, ga.buffers(n, device), lay.wo_a[0].k // head_dim -def wo_ab(lay, o: torch.Tensor | None = None, xh: list | None = None, fold: bool = True) -> torch.Tensor: - """Apply wo_a then wo_b with FP32 output; a dense wo_b skips the EXL3 input rotation.""" +def wo_ab(lay, o: torch.Tensor | None = None, xh: list | None = None, fold: bool = True, + adapter=None) -> torch.Tensor: + """Apply wo_a then wo_b with FP32 output; a dense wo_b skips the EXL3 input rotation. ``adapter`` (lora_runtime. + RowAdapters): the rows' wo_b adapter deltas are added to the partial sum.""" ga = _group(lay, "_tf_wo_a", lay.wo_a) dense = hasattr(lay.wo_b, "weight") @@ -436,8 +438,10 @@ def wo_ab(lay, o: torch.Tensor | None = None, xh: list | None = None, fold: bool c += wo.n ga.rotated(xh, outs, rot=rot if fold else None) if not fold: # wo_b's input rotated by its own rot_many launch - return mm(lay.wo_b, u, F32) - return gb.rotated(xb, out_dtypes=[F32])[0] + out = mm(lay.wo_b, u, F32) + else: + out = gb.rotated(xb, out_dtypes=[F32])[0] + return out if adapter is None else adapter.wo_b(lay, u, out) @dataclass @@ -741,6 +745,7 @@ def __init__(self, w, comm: Comm, engram: Engram | None = None): c = self.cfg self.Hl = c.n_heads // w.world self.scratch: dict = {} + self.lora = None # lora_runtime.LoraRuntime, when the engine serves adapters self.rope = RopeTables(c) self._zero = torch.zeros((1,), dtype=torch.int64, device="cuda") if not KERNELS and KV_QUANT: @@ -1009,9 +1014,10 @@ def kv_source_update(self, lay, x: torch.Tensor, sc: SeqCache, start: int, share store_rows(sc.comp[lay.idx], groups, lat, 16, True) def attention_k(self, lay, x: torch.Tensor, sc: SeqCache, start: int, shared: dict, pos: torch.Tensor, - floor: int = 0, kv_done: bool = False): + floor: int = 0, kv_done: bool = False, adapter=None): """The same math as ``attention`` with the row-independent Triton kernels. ``floor``: no window key before - this position (bounded replay); ``kv_done``: the kv-source update already ran for these rows.""" + this position (bounded replay); ``kv_done``: the kv-source update already ran for these rows; ``adapter`` + (lora_runtime.RowAdapters): the rows' wo_b adapter deltas.""" c = self.cfg n = x.shape[0] @@ -1110,7 +1116,9 @@ def attention_k(self, lay, x: torch.Tensor, sc: SeqCache, start: int, shared: di if not ring_mode: keep = min(n, R) ring[pos[-keep:] % R] = kv[-keep:] - return mm(lay.wo_b, wo_a_out(lay, o), F32) + u = wo_a_out(lay, o) + out = mm(lay.wo_b, u, F32) + return out if adapter is None else adapter.wo_b(lay, u, out) def _neg(self, n: int) -> torch.Tensor: t = self.scratch.get(("neg", n)) @@ -1121,10 +1129,11 @@ def _neg(self, n: int) -> torch.Tensor: # -- MoE ------------------------------------------------------------------------------------------------------- def moe(self, lay, x: torch.Tensor, topk: int | None = None, img: torch.Tensor | None = None, - shared_side: bool = False) -> torch.Tensor: + shared_side: bool = False, adapter=None) -> torch.Tensor: """fp32 partial [n, d] of the MoE. ``shared_side`` (the decode bodies, which prefetch this layer's mHC mix and gate into L2 first): with SHARED_OVERLAP a decode window's shared expert runs on a side stream (_shared_side). - (Without the gate in L2 the side stream's DRAM traffic slows the gate's matmul by more than it overlaps.)""" + (Without the gate in L2 the side stream's DRAM traffic slows the gate's matmul by more than it overlaps.) + ``adapter`` (lora_runtime.RowAdapters): the rows' adapter deltas of the shared expert's w2.""" c = self.cfg n = x.shape[0] topk = topk or c.topk @@ -1167,10 +1176,14 @@ def moe(self, lay, x: torch.Tensor, topk: int | None = None, img: torch.Tensor | s = self._moe_scratch(lay, n, slots) if EXACT_MM: return self._moe_exact(lay, x, pick, wts) - shared = self._shared_apart(lay, x) if separate_shared else None + shared = self._shared_apart(lay, x, adapter) if separate_shared else None out = exl3_experts.routed(x.contiguous(), pick.contiguous(), wts, lay.experts, s, None, n, limit=c.swiglu_limit, act_mode=exl3_experts.ACT_F32, add=shared, before_down=side_done) + if adapter is not None and not separate_shared: + # the packed shared expert's rotated hidden is still in the scratch of the call that computed it + scratch = s if side_done is None else self.scratch[("moe_side", slots, lay.experts.count)] + adapter.shared(lay, scratch.xd[:n * slots].view(n, slots, -1)[:, slots - 1], out, packed=True) return out # fp32 partial [n, d] def _moe_scratch(self, lay, n: int, slots: int) -> "exl3_experts.Scratch": @@ -1187,7 +1200,7 @@ def _moe_scratch(self, lay, n: int, slots: int) -> "exl3_experts.Scratch": self.scratch[skey] = s return s - def _shared_apart(self, lay, x: torch.Tensor) -> torch.Tensor: + def _shared_apart(self, lay, x: torch.Tensor, adapter=None) -> torch.Tensor: """Run a shared expert that has a dense or other-format projection outside the routed EXL3 pack.""" w1, w3, w2 = lay.shared @@ -1203,7 +1216,8 @@ def _shared_apart(self, lay, x: torch.Tensor) -> torch.Tensor: h = K.clamped_swiglu(g.contiguous(), u.contiguous(), limit) else: h = (F.silu(g.float().clamp(max=limit)) * u.float().clamp(-limit, limit)).to(BF16) - return mm(w2, h, F32) + out = mm(w2, h, F32) + return out if adapter is None else adapter.shared(lay, h, out, packed=False) def _shared_side(self, lay, x: torch.Tensor, n: int, slots: int) -> torch.cuda.Event: """SHARED_OVERLAP: the shared expert of x's n rows on the side stream, from now on: exl3_experts.routed with @@ -1265,15 +1279,18 @@ def mat(ptrs, k2s, suh, svh, e, k, n): @torch.inference_mode() def forward(self, sc: SeqCache, ids: torch.Tensor, start: int, all_logits: bool = False, taps: list | None = None, host_ids: list[int] | None = None, replay: int | None = None, - image: tuple | None = None): + image: tuple | None = None, adapter: int = 0): """Rows ids [n] at positions start.. of one sequence -> fp32 logits [n or 1, V] (both ranks the same). ``image``: (rows [k] long, embeddings [k, dim] bf16), this block's image-span rows: their embeddings replace the token's, the MoE gate picks their experts with its VL bias and Engram leaves them untouched (DeepSeek's - image_mask); ``host_ids`` then holds a negative id at each of them (no n-gram spans one).""" + image_mask); ``host_ids`` then holds a negative id at each of them (no n-gram spans one). ``adapter``: the + registry slot every row of this block runs (0: the stock model).""" c, w = self.cfg, self.w n = ids.shape[0] assert start == sc.length, (start, sc.length) + if adapter and (self.lora is None or not KERNELS): + raise ValueError("adapters run on the kernel path of an engine started with --adapter") sc.tokens[start:start + n] = ids emb = w.embed[ids].to(BF16) img = None @@ -1291,7 +1308,7 @@ def forward(self, sc: SeqCache, ids: torch.Tensor, start: int, all_logits: bool hashes = self.engram.hashes(sc.host.view(), start, n) # [n, L, cols] host shared: dict = {} if KERNELS: - return self._forward_k(sc, ids, start, all_logits, taps, h, pre, hashes, shared, replay, img) + return self._forward_k(sc, ids, start, all_logits, taps, h, pre, hashes, shared, replay, img, adapter) for lay in w.layers: if hashes is not None and lay.engram_wkv is not None: h = self.engram_apply(lay, h, hashes[:, c.engram_layers.index(lay.idx)], img) @@ -1316,7 +1333,8 @@ def forward(self, sc: SeqCache, ids: torch.Tensor, start: int, all_logits: bool g = self.comm.gather(local) # [world, n, V / world] return g.permute(1, 0, 2).reshape(local.shape[0], -1) - def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, replay=None, img=None): + def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, replay=None, img=None, + adapter: int = 0): """``replay`` (decoder SWA bounded replay, CED's prefill): the decoder layers run only for rows at positions >= replay, their window truncated there; the first decoder layer's kv source still covers every row.""" @@ -1332,6 +1350,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r comb = torch.empty((n, c.hc, c.hc), dtype=F32, device=dev) floor, kv_done = 0, False h_alt = None # (switch "hc_pf") the streams' second buffer + arow = self.lora.uniform(adapter, n) if adapter else None # every row of the block runs one adapter self.taps_start = start for lay in w.layers: if replay is not None and lay.idx == c.n_layers // 2: @@ -1356,6 +1375,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r pre_f = torch.empty((n, c.hc), dtype=F32, device=dev) post = torch.empty((n, c.hc), dtype=F32, device=dev) comb = torch.empty((n, c.hc, c.hc), dtype=F32, device=dev) + arow = self.lora.uniform(adapter, n) if adapter else None self.taps_start = start floor, kv_done = replay, True if hashes is not None and lay.engram_wkv is not None: @@ -1375,7 +1395,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r comb, part) with _T("attn_r%d" % (2 if lay.comp_wkv is not None else 1 if lay.idx_wq_b is not None else 0)): pa = self.attention_k(lay, x, sc, start, shared, pos, floor=floor, - kv_done=kv_done and lay.idx == c.n_layers // 2) + kv_done=kv_done and lay.idx == c.n_layers // 2, adapter=arow) with _T("gather"): g = self.comm.gather(pa) with _T("hc"): @@ -1392,7 +1412,7 @@ def _forward_k(self, sc, ids, start, all_logits, taps, h, pre, hashes, shared, r K.hc_pre(h, fn, scale, base, pre_a, lay.ffn_norm, c.eps, c.hc_eps, c.hc_iters, x, pre_f, post, comb, part) with _T("moe"): - pm = self.moe(lay, x, img=img) + pm = self.moe(lay, x, img=img, adapter=arow) with _T("gather"): g = self.comm.gather(pm) with _T("hc"): diff --git a/src/tensorfold/families/deepseek_v41/cuda/multi.py b/src/tensorfold/families/deepseek_v41/cuda/multi.py index 38ebb42711..ae8db0b7c3 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/multi.py +++ b/src/tensorfold/families/deepseek_v41/cuda/multi.py @@ -54,6 +54,7 @@ from ..ops import HostIds from . import constrain, memtrace +from .lora_registry import AdapterRef, key_of, slot_of from .model import RAW # "conf", "fixed", "even": the most drafts a stream verifies a round by how many streams decode ("5,5,3,3": up to the @@ -444,12 +445,15 @@ class Kept: """A finished stream's prompt kept in the window: positions [base, base + size) hold the compressed and indexer rows (and token ids) of its first ``top`` prompt tokens; ``snaps`` its slot's rings and compressor inputs at each kept chunk boundary; ``host`` the prompt's host ids to ``top`` (Engram's n-grams); ``keys`` (rank 0 only) the - prompt's ids to ``top`` with each image span's positions keyed by its picture; ``tick`` when it was last used.""" + prompt's ids to ``top`` with each image span's positions keyed by its picture; ``tick`` when it was last used; + ``adapter`` the identity key of the adapter whose forward made the rows ('' for the stock model): only a prompt + of the same adapter continues it.""" def __init__(self, eid: int, base: int, size: int, top: int, host: list, snaps: dict, replay: bool, keys, - tick: int) -> None: + tick: int, adapter: str = "") -> None: self.eid, self.base, self.size, self.top, self.host = eid, base, size, top, host self.snaps, self.replay, self.keys, self.tick, self.hits = snaps, replay, keys, tick, 0 + self.adapter = adapter self.disk: dict | None = None # its spilled manifest (TF_DS_SPILL_GIB), when written @@ -695,9 +699,11 @@ def _all(self, ok: bool) -> bool: def _shape(self) -> list: return [self.next_id, list(self.free), list(self.extents.gaps), - [[s.sid, s.st.index, s.st.sc.length, len(s.out), bool(s.done)] for s in self.streams.values()], - [[s.sid, s.st.index, s.filled] for s in self.filling], - [[k.eid, k.base, k.size, k.top, k.tick, sorted(k.snaps)] for k in self.kept.values()]] + [[s.sid, s.st.index, s.st.sc.length, len(s.out), bool(s.done)] + ([slot_of(s)] if slot_of(s) else []) + for s in self.streams.values()], + [[s.sid, s.st.index, s.filled] + ([slot_of(s)] if slot_of(s) else []) for s in self.filling], + [[k.eid, k.base, k.size, k.top, k.tick, sorted(k.snaps)] + ([k.adapter] if k.adapter else []) + for k in self.kept.values()]] # -- kept prompts --------------------------------------------------------------------------------------------- def _replay(self) -> bool: @@ -726,7 +732,7 @@ def _match(self, s: Stream, keys) -> list | None: limit = len(s.prompt) - (self.m.cfg.window if replay else 1) best = None for k in self.kept.values(): - if k.replay != replay or k.keys is None: + if k.replay != replay or k.keys is None or k.adapter != key_of(s): # (only the same adapter's rows) continue n = min(len(k.keys), len(keys)) diff = np.flatnonzero(k.keys[:n] != keys[:n]) @@ -884,14 +890,14 @@ def _keep(self, s: Stream) -> None: self.extents.give(s.base + size, s.size - size) keys = getattr(s, "keys", None) k = Kept(self.next_kept, s.base, size, top, s.st.sc.host.copy(top), {b: s.snaps[b] for b in marks}, - self._replay(), None if keys is None else keys[:top], self.ticks) + self._replay(), None if keys is None else keys[:top], self.ticks, key_of(s)) self.kept[k.eid] = k self.next_kept += 1 if self.spill is not None: # write-through: its rows past what its lineage holds with memtrace.phase("spill-save"): k.disk = self.spill.persist(f"{self.spill_nonce or 'run'}-{k.eid:06d}", base=k.base, top=k.top, replay=k.replay, snaps=k.snaps, host=k.host, - lineage=getattr(s, "lineage", None), keys=k.keys) + lineage=getattr(s, "lineage", None), keys=k.keys, adapter=k.adapter) s.st.sc = None s.snaps = None @@ -910,7 +916,7 @@ def _covered(self, done: list[Stream]) -> list[int]: top = max(marks) for k in self.kept.values(): if (k.eid not in drops and k.keys is not None and k.top <= top and k.replay == self._replay() - and set(k.snaps) <= marks and np.array_equal(k.keys, keys[:k.top])): + and k.adapter == key_of(s) and set(k.snaps) <= marks and np.array_equal(k.keys, keys[:k.top])): drops.append(k.eid) return drops @@ -923,6 +929,8 @@ def admit(self, s: Stream) -> None: if s.probabilities is not None: raise ValueError("logprobs are not served with --parallel on DeepSeek-V4.1") + if s.adapter is not None and not self.e.adapter_ok(s.adapter): + raise ValueError("this request names an adapter the engine does not hold") room = self.extents.total - len(s.prompt) - self.max_rows - 8 - 4 if room < 1: raise ValueError(f"a prompt of {len(s.prompt)} tokens leaves no room in the {self.extents.total}-token " @@ -939,13 +947,23 @@ def admit(self, s: Stream) -> None: reuse = self._match(s, s.keys) if KEEP and self.kept else None if self.spill is not None and s.keys is not None: # a spilled boundary past every kept prompt's far = self.spill.best(s.keys, len(s.prompt), window=self.m.cfg.window, replay=self._replay(), - above=int(reuse[1]) if reuse else 0) + above=int(reuse[1]) if reuse else 0, adapter=key_of(s)) if far is not None: reuse = ["disk", far[0], far[1]] + # (a stock request's op is the one it always was: the adapter's slot and key ride after the grammar) self._send(["admit", list(s.prompt), s.count, _pack(s.sampling), bool(s.draft), bool(s.stop_eos), index, - positions, reuse, pack_grammar(s.constraint)]) + positions, reuse, pack_grammar(s.constraint)] + self._adapter_plan(s)) self._admit(s, index, positions, reuse) + def _adapter_plan(self, s: Stream, check: bool = False) -> list: + """[] for the stock model, else [[slot, key]] (with ``check``, [slot, key, whether this rank holds that adapter]: + a rank that does not makes the step's digests differ, and the step fails on every rank).""" + + ref = s.adapter + if ref is None: + return [] + return [[int(ref.slot), str(ref.key)] + ([bool(self.e.adapter_ok(ref))] if check else [])] + def _need(self, s: Stream) -> int: """A stream's positions: its prompt, its reply, a round's rows and its extent's scratch row (a group long).""" @@ -961,7 +979,8 @@ def _admit(self, s: Stream, index: int, positions: list[int], reuse: list | None self._step(True) try: self._agree("admission", [self._shape(), list(s.prompt), s.count, _pack(s.sampling), bool(s.draft), - bool(s.stop_eos), index, positions, reuse, pack_grammar(s.constraint)]) + bool(s.stop_eos), index, positions, reuse, pack_grammar(s.constraint)] + + self._adapter_plan(s, check=True)) self.ticks += 1 need = self._need(s) disk = bool(reuse) and reuse[0] == "disk" @@ -1038,8 +1057,9 @@ def snap(end: int) -> None: if end in marks: s.snaps[end] = self._snapshot(index) + held = {} if s.adapter is None else {"adapter": slot_of(s)} # (a stock prompt's call is as it was) s.steps = e.prefill_steps(slot.sc, slot.dc if s.draft else None, list(s.prompt), image, start=cut, - snap=snap) + snap=snap, **held) s.filled, s.prefill_s, s.cached = cut, 0.0, cut self.filling.append(s) except Exception as exc: @@ -1384,7 +1404,8 @@ def mark(): self._stream_failed(s, exc) window = [s.pending] + drafts sc.host.set(P, window) - windows.append((s.st.index, s.base, s.size, P, window, sc.host.view())) + windows.append((s.st.index, s.base, s.size, P, window, sc.host.view()) + + (() if s.adapter is None else (slot_of(s),))) kept.append((s, P, drafts)) _ht("drafts back + plan") logits, taps = self.runner.forward(windows) @@ -1392,7 +1413,7 @@ def mark(): absorb_done = None if EAGER_ABSORB and taps is not None: eager, r = [], 0 - for (s, P, _), (_, _, _, _, window, _) in zip(kept, windows): + for (s, P, _), (_, _, _, _, window, *_) in zip(kept, windows): if s.draft: eager.append((s.st.index, taps[r:r + len(window)], P)) r += len(window) @@ -1402,7 +1423,7 @@ def mark(): done, r0, tokens = [], 0, 0 t_absorb = 0.0 absorbs = [] - for (s, P, drafts), (_, _, _, _, window, _) in zip(kept, windows): + for (s, P, drafts), (_, _, _, _, window, *_) in zip(kept, windows): n = len(window) if not s.done: # (a stream that failed before the forward: its rows, no tokens) try: @@ -1598,9 +1619,10 @@ def follow(self, link: Link) -> None: kind = op[0] try: if kind == "admit": - _, prompt, count, sampling, draft, stop_eos, index, positions, reuse, packed = op + _, prompt, count, sampling, draft, stop_eos, index, positions, reuse, packed, *held = op s = Stream(list(prompt), int(count), _unpack(sampling), draft=bool(draft), stop_eos=bool(stop_eos), - constraint=constrain.follow(self.e, packed)) + constraint=constrain.follow(self.e, packed), + adapter=AdapterRef(int(held[0][0]), str(held[0][1])) if held else None) s.emit = lambda new: None s.keys = None self._admit(s, int(index), list(positions), reuse) diff --git a/src/tensorfold/families/deepseek_v41/cuda/rounds.py b/src/tensorfold/families/deepseek_v41/cuda/rounds.py index 6f32616fa4..0cb9ec0771 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/rounds.py +++ b/src/tensorfold/families/deepseek_v41/cuda/rounds.py @@ -80,9 +80,11 @@ def __init__(self, model: Model, pool: PoolCache, rows: int, bucket: int, taps: self.out = out # (logits, taps) buffers the runner's graphs share (SHARED_OUT) c = model.cfg dev = "cuda" - self.inputs = torch.zeros((5, rows), dtype=torch.long, device=dev) # one buffer, so one copy fills it - # ids, positions, slots, and the row's stream's extent (base, end) in positions - self.ids, self.pos, self.slot, self.base, self.end = self.inputs.unbind(0) + # one buffer, so one copy fills it: ids, positions, slots, the row's stream's extent (base, end) in positions + # and, in an engine that serves adapters, the row's adapter slot (read at replay by the adapter kernels) + self.inputs = torch.zeros((6 if model.lora is not None else 5, rows), dtype=torch.long, device=dev) + self.ids, self.pos, self.slot, self.base, self.end, *spare = self.inputs.unbind(0) + self.adapter = model.lora.bind(spare[0]) if spare else None self.end.fill_(1) self.table = SeqCache(cap=pool.cap) # RoPE tables for any pool position self.e_in = {} @@ -217,7 +219,10 @@ def _attention(self, lay, x, shared, cos, sin, xh=None): o = K.sparse_attn(q, lay.sink, ring, m._zero, True, comp, cidx, pos, hd ** -0.5, c.window, wbase=wbase, cbase=cbase, ring_rows=RS) K.rope_heads(o, cos, sin, pos, rd, inverse=True) - out = mm(lay.wo_b, wo_a_out(lay, o), F32) + u = wo_a_out(lay, o) + out = mm(lay.wo_b, u, F32) + if self.adapter is not None: + out = self.adapter.wo_b(lay, u, out) l2_fork(m, lay, "moe") # the MoE's mHC mix, gate, shared expert return out @@ -260,15 +265,18 @@ def _attention2(self, lay, x, shared, cos, sin, xh=None): suh, xh, gh = wo_a_rot(lay, n, x.device, hd) K.sparse_attn(q, lay.sink, ring, m._zero, True, comp, cidx, pos, hd ** -0.5, c.window, wbase=wbase, cbase=cbase, ring_rows=RS, rot=(cos, sin, rd, suh, xh[0], gh)) - out = wo_ab(lay, xh=xh, fold=K.on("rot_wob")) + out = wo_ab(lay, xh=xh, fold=K.on("rot_wob"), adapter=self.adapter) else: o = K.sparse_attn(q, lay.sink, ring, m._zero, True, comp, cidx, pos, hd ** -0.5, c.window, wbase=wbase, cbase=cbase, ring_rows=RS) K.rope_heads(o, cos, sin, pos, rd, inverse=True) if K.on("rot_wob"): - out = wo_ab(lay, o) # wo_a's epilogue rotates for wo_b + out = wo_ab(lay, o, adapter=self.adapter) # wo_a's epilogue rotates for wo_b else: - out = mm(lay.wo_b, wo_a_out(lay, o), F32) + u = wo_a_out(lay, o) + out = mm(lay.wo_b, u, F32) + if self.adapter is not None: + out = self.adapter.wo_b(lay, u, out) l2_fork(m, lay, "moe") # the MoE's mHC mix, gate, shared expert return out @@ -438,7 +446,7 @@ def _stretch(self, k: int): fn, scale, base = lay.hc_ffn K.hc_pre(h, fn, scale, base, pre_a, lay.ffn_norm, c.eps, c.hc_eps, c.hc_iters, x, pre_f, post, comb, part) - y = m.moe(lay, x, shared_side=moe_side()) # the gate is in L2 by then + y = m.moe(lay, x, shared_side=moe_side(), adapter=self.adapter) # the gate is in L2 by then self._l2_next(li, end) K.hc_post(m.comm.gather(y), h, post, comb, h) pre, pre_f = pre_f, pre @@ -541,7 +549,7 @@ def pre_mix(h, params, pre_in, norm, pre_out, rot=None): join() pending = m.comm.gather(y) h = pre_mix(h, lay.hc_ffn, pre_a, lay.ffn_norm, pre_f) - y = m.moe(lay, x, shared_side=moe_side()) # the gate is in L2 by then + y = m.moe(lay, x, shared_side=moe_side(), adapter=self.adapter) # the gate is in L2 by then self._l2_next(li, end) join() pending = m.comm.gather(y) @@ -567,11 +575,16 @@ def capture(self, pool=None) -> None: self.graphs = graphs self.graph = graphs[0] - def set(self, ids: list[int], pos: list[int], slots: list[int], base: list[int], end: list[int]) -> None: + def set(self, ids: list[int], pos: list[int], slots: list[int], base: list[int], end: list[int], + adapters: list[int] | None = None) -> None: + cols = [ids, list(pos), slots, base, end] + if self.adapter is not None: + cols.append(adapters if adapters is not None else [0] * len(ids)) if ONE_COPY: # pinned: the copy runs in stream order, nothing waits for it - self.inputs.copy_(torch.tensor([ids, list(pos), slots, base, end], dtype=torch.long).pin_memory(), - non_blocking=True) + self.inputs.copy_(torch.tensor(cols, dtype=torch.long).pin_memory(), non_blocking=True) return + if self.adapter is not None: + self.adapter.rows.copy_(torch.tensor(cols[5], dtype=torch.long), non_blocking=False) self.ids.copy_(torch.tensor(ids, dtype=torch.long), non_blocking=False) self.pos.copy_(torch.tensor(pos, dtype=torch.long), non_blocking=False) self.slot.copy_(torch.tensor(slots, dtype=torch.long), non_blocking=False) @@ -579,7 +592,7 @@ def set(self, ids: list[int], pos: list[int], slots: list[int], base: list[int], self.end.copy_(torch.tensor(end, dtype=torch.long), non_blocking=False) def run(self, ids: list[int], pos: list[int], slots: list[int], base: list[int], end: list[int], - e_rows, e_start: dict | None = None) -> torch.Tensor: + e_rows, e_start: dict | None = None, adapters: list[int] | None = None) -> torch.Tensor: """``e_rows``: layer -> its Engram rows on the GPU, or a callable giving them (called just before the stretch that starts at that layer, so the read overlaps the stretches before it). ``e_start``: layer -> a callable that starts that layer's read, called once the stretch before the one that needs the rows is enqueued (before the @@ -595,7 +608,7 @@ def start(k: int) -> None: # the reads the stretch after s from .multi import _ht _ht("fwd prep") - self.set(ids, pos, slots, base, end) + self.set(ids, pos, slots, base, end, adapters) _ht("fwd set") start(-1) for k, (first, _) in enumerate(self.stretches): @@ -664,29 +677,34 @@ def engram_touch(self, tails: list[list[int]]) -> None: _ht(f"touch {i} submitted") def forward(self, windows: list[tuple], replay: bool = True) -> tuple[torch.Tensor, torch.Tensor | None]: - """``windows``: each stream's (slot, extent base, extent size, first position, token ids, host ids so far): - rows in that order. Returns logits [R, V] and taps [R, 3 d] (the rows in window order; with SHARED_OUT views of - the runner's pair, valid until its next forward); ``replay=False`` + """``windows``: each stream's (slot, extent base, extent size, first position, token ids, host ids so far[, + adapter slot]): rows in that order. Returns logits [R, V] and taps [R, 3 d] (the rows in window order; with + SHARED_OUT views of the runner's pair, valid until its next forward); ``replay=False`` only captures a missing graph (warm-up) and returns (None, None).""" m = self.m - ids, pos, slots, base, end = [], [], [], [], [] + ids, pos, slots, base, end, adapters = [], [], [], [], [], [] hashes = [] - for slot, b0, size, p0, toks, host in windows: + for slot, b0, size, p0, toks, host, *more in windows: n = len(toks) ids += toks pos += range(p0, p0 + n) slots += [slot] * n base += [b0] * n end += [b0 + size] * n + adapters += [more[0] if more else 0] * n if m.engram is not None: hashes.append(m.engram.hashes(host, p0, n)) + if any(adapters) and m.lora is None: + raise ValueError("a stream asks for an adapter, and this engine was started without one") R = len(ids) b = self.bucket(max(p + 1 for p in pos)) Rg = R if self.graphs is not None: # also a capture-only call: it makes the key a round uses Rg = padded_rows(R, sizes_for(b, self.sizes, self.sizes_long)) - for col in (ids, pos, slots, base, end): # the extra rows copy the last row (graph_keys.ROUND_ROWS) + # the extra rows copy the last row (graph_keys.ROUND_ROWS): its adapter too, or the copy would write + # other values into the cache rows its original writes + for col in (ids, pos, slots, base, end, adapters): col += col[-1:] * (Rg - R) e_rows = e_start = None if m.engram is not None and replay: @@ -712,7 +730,7 @@ def forward(self, windows: list[tuple], replay: bool = True) -> tuple[torch.Tens g = self.graphs.get(key) if self.graphs is not None else None if g is None: g = RoundDecoder(m, self.pool, Rg, b, True, out=self.out) - g.set(ids, pos, slots, base, end) + g.set(ids, pos, slots, base, end, adapters) if self.graphs is not None: if self.graph_pool is None: self.graph_pool = torch.cuda.graph_pool_handle() @@ -723,5 +741,5 @@ def forward(self, windows: list[tuple], replay: bool = True) -> tuple[torch.Tens self.graphs[key] = g if not replay: return None, None - out = g.run(ids, pos, slots, base, end, e_rows, e_start) + out = g.run(ids, pos, slots, base, end, e_rows, e_start, adapters) return out[:R], (None if g.taps is None else g.taps[:R]) diff --git a/src/tensorfold/families/deepseek_v41/cuda/spill.py b/src/tensorfold/families/deepseek_v41/cuda/spill.py index d7cb899b24..e5a20368c1 100644 --- a/src/tensorfold/families/deepseek_v41/cuda/spill.py +++ b/src/tensorfold/families/deepseek_v41/cuda/spill.py @@ -149,6 +149,7 @@ def _write_all(fd: int, data: memoryview) -> None: _NAME = re.compile(r"[0-9A-Za-z_-]{1,64}") +_ADAPTER = re.compile(r"[0-9a-f]{0,64}") # an adapter's identity key; '' is the stock model _FILE = re.compile(r"(seg|snap)-[0-9A-Za-z_-]{1,64}-\d{1,9}\.st") @@ -160,7 +161,8 @@ def _valid(m: dict) -> bool: """A manifest names only files of this directory's own shapes (a tampered one could name a path elsewhere).""" files = [seg.get("file", "") for seg in m.get("segments", [])] + [r.get("file", "") for r in m["snaps"].values()] - return bool(_NAME.fullmatch(str(m.get("name", "")))) and all(_FILE.fullmatch(str(f)) for f in files) + return (bool(_NAME.fullmatch(str(m.get("name", "")))) and all(_FILE.fullmatch(str(f)) for f in files) + and _ADAPTER.fullmatch(str(m.get("adapter", ""))) is not None) def engine_id(model_dir, *, world: int, pool, dpool=None) -> str: @@ -308,11 +310,21 @@ def _write_snap(self, fname: str, snap) -> None: self._save(fname, tensors, {"n": len(snap), "engine": self.engine}) @staticmethod - def _hashes(keys, boundaries) -> dict: + def _seed(adapter: str = ""): + """A prefix hash's start: empty for the stock model, so its hashes are what they always were; an adapter's + identity key first, so no prefix of one adapter can equal a prefix of another.""" + + h = hashlib.sha256() + if adapter: + h.update(b"adapter\0" + adapter.encode() + b"\0") + return h + + @classmethod + def _hashes(cls, keys, boundaries, adapter: str = "") -> dict: """Each boundary's prefix hash: sha256 of the prompt's ids (int64, image spans keyed by picture) to it.""" k64 = np.ascontiguousarray(keys, dtype=np.int64) - h, prev, out = hashlib.sha256(), 0, {} + h, prev, out = cls._seed(adapter), 0, {} for b in sorted(int(x) for x in boundaries): if b > len(k64): break @@ -322,11 +334,12 @@ def _hashes(keys, boundaries) -> dict: return out def persist(self, name: str, *, base: int, top: int, replay: bool, snaps: dict, host, lineage: dict | None, - keys=None) -> dict | None: + keys=None, adapter: str = "") -> dict | None: """Write a kept prompt at pool positions [base, base + top): the rows past what ``lineage`` (the kept or spilled prompt it continued, to where it continued) already holds on this disk, the snapshots not already there, then its manifest. Returns the manifest, None when the prompt is too short or a write failed (the - prompt stays in the window either way).""" + prompt stays in the window either way). ``adapter``: the identity key of the adapter that made the rows + ('' for the stock model); the manifest names it, and only a prompt of the same adapter reads it back.""" top = int(top) if top < self.min_top or not snaps: @@ -352,8 +365,10 @@ def persist(self, name: str, *, base: int, top: int, replay: bool, snaps: dict, wrote += self._path(fname).stat().st_size m = {"format": FORMAT, "engine": self.engine, "name": name, "top": top, "replay": bool(replay), "segments": segs, "snaps": refs, "created": time.time()} + if adapter: + m["adapter"] = str(adapter) if keys is not None: - m["hashes"] = self._hashes(keys, [int(b) for b in refs]) + m["hashes"] = self._hashes(keys, [int(b) for b in refs], adapter) self._write_json(self._entry(name), m) self.manifests[name] = m self._index_add(m) @@ -399,10 +414,11 @@ def _index_add(self, m: dict) -> None: if ref is not None: self.index.setdefault(hx, []).append((m["name"], int(b), bool(ref["full"]))) - def best(self, keys, plen: int, *, window: int, replay: bool, above: int = 0) -> tuple[str, int] | None: + def best(self, keys, plen: int, *, window: int, replay: bool, above: int = 0, + adapter: str = "") -> tuple[str, int] | None: """Rank 0: (manifest, boundary) of the furthest spilled boundary this prompt shares, past ``above``: a boundary snapshot below the prompt's last row (with replay prefill, a window below), a full one anywhere below - its last row; the most recently used of equals.""" + its last row; the most recently used of equals. Only manifests of the same ``adapter`` key count.""" if not self.index: return None @@ -412,7 +428,7 @@ def best(self, keys, plen: int, *, window: int, replay: bool, above: int = 0) -> if not cands: return None k64 = np.ascontiguousarray(keys, dtype=np.int64) - h, prev, best = hashlib.sha256(), 0, None + h, prev, best = self._seed(adapter), 0, None for b in cands: if b > len(k64): break @@ -422,6 +438,8 @@ def best(self, keys, plen: int, *, window: int, replay: bool, above: int = 0) -> m = self.manifests.get(name) if bb != b or m is None or bool(m.get("replay")) != bool(replay) or b > (lim_f if full else lim_b): continue + if str(m.get("adapter", "")) != adapter: # (the hashes are seeded by it too: belt and braces) + continue t = self._mtime(name) if best is None or (b, t) > (best[1], best[2]): best = (name, b, t) diff --git a/src/tensorfold/serve_options.py b/src/tensorfold/serve_options.py index 3316d8e91a..a2ac795b39 100644 --- a/src/tensorfold/serve_options.py +++ b/src/tensorfold/serve_options.py @@ -4,9 +4,13 @@ import argparse import inspect +import math +import re from pathlib import Path from typing import Any +_ADAPTER_NAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._:/-]{0,95}") + def _served_name(args: argparse.Namespace) -> str: """``--name``, else a repo id's last segment or the model path's name; no I/O (matches ``cli._serve_cuda``).""" @@ -31,9 +35,51 @@ def name_priority(args: argparse.Namespace) -> dict[str, str]: return parsed +def adapter_specs(args: argparse.Namespace) -> list[tuple[str, str, float]]: + """``--adapter NAME=FOLDER`` and ``--adapter-strength NAME=NUMBER`` as (name, folder, strength) in the order + given (strength 1 unless named). No I/O: the folders are read when the engine starts.""" + + folders: dict[str, str] = {} + for entry in getattr(args, "adapter", None) or []: + name, sep, folder = str(entry).partition("=") + name, folder = name.strip(), folder.strip() + if not sep or not name or not folder: + raise ValueError(f"--adapter takes NAME=FOLDER, not {entry!r}") + if not _ADAPTER_NAME.fullmatch(name): + raise ValueError(f"--adapter name {name!r}: letters, digits and . _ : / - only, 1 to 96 characters") + if name in folders: + raise ValueError(f"--adapter {name} is given twice") + folders[name] = folder + strengths: dict[str, float] = {} + for entry in getattr(args, "adapter_strength", None) or []: + name, sep, value = str(entry).partition("=") + try: + number = float(value) + except ValueError: + number = math.nan + if not sep or not math.isfinite(number) or number == 0.0: + raise ValueError(f"--adapter-strength takes NAME=NUMBER (finite, not 0), not {entry!r}") + if name.strip() not in folders: + raise ValueError(f"--adapter-strength names {name.strip()!r}, which no --adapter defines") + strengths[name.strip()] = number + return [(name, folder, strengths.get(name, 1.0)) for name, folder in folders.items()] + + def check(args: argparse.Namespace, family: Any, backend: str, config_dir: Any = None) -> None: """Refuse KV cache, draft rule, image, share, slot and precision options the backend or family can't serve.""" + adapters = adapter_specs(args) + if adapters: + if backend != "cuda" or not getattr(family.package, "CUDA_ADAPTERS", False): + raise ValueError(f"--adapter serves LoRA adapters from the DeepSeek-V4.1 CUDA engine; {family.title} on " + f"{'CUDA' if backend == 'cuda' else 'MLX'} has none") + if _cuda_streams(getattr(args, "parallel", "auto")) < 2: + raise ValueError("--adapter needs --parallel 2 or more: each round's rows carry their own adapter") + taken = {_served_name(args), *(str(a).strip() for a in getattr(args, "alias", None) or () if str(a).strip())} + clash = sorted(name for name, _, _ in adapters if name in taken) + if clash: + raise ValueError(f"--adapter {', '.join(clash)} is already a model id of this server (--name or --alias)") + if getattr(args, "vision_urls", False) and not getattr(args, "vision", False): raise ValueError("--vision-urls needs --vision") images = getattr(args, "vision_max_images", None) diff --git a/src/tensorfold/server/responses_translate.py b/src/tensorfold/server/responses_translate.py index 023f860bf8..7ea0111e00 100644 --- a/src/tensorfold/server/responses_translate.py +++ b/src/tensorfold/server/responses_translate.py @@ -15,7 +15,7 @@ # fields a chat completion reads as they are: OpenAI's and this server's own (top_k, min_p, seed, draft, ...) PASSED = ("model", "temperature", "top_p", "top_k", "min_p", "seed", "stream", "parallel_tool_calls", "stop", "draft", - "thinking_budget", "ignore_eos", "priority", "return_token_ids", "chat_template_kwargs") + "thinking_budget", "ignore_eos", "priority", "return_token_ids", "chat_template_kwargs", "adapter") REFUSED = {"background": "background responses are not supported: send the request and wait for it", "conversation": "conversations are not supported: send previous_response_id or the items", "prompt": "prompt templates are not supported: send input and instructions", diff --git a/tests/adapter_fakes.py b/tests/adapter_fakes.py new file mode 100644 index 0000000000..c3a1cd7099 --- /dev/null +++ b/tests/adapter_fakes.py @@ -0,0 +1,103 @@ +"""Adapter folders for the DeepSeek-V4.1 adapter tests: PEFT files written without torch, small model sizes.""" + +from __future__ import annotations + +import json +import re +import struct +from pathlib import Path +from types import SimpleNamespace + +import numpy as np + +# a small model: 4 layers, hidden 256, shared expert width 512, wo_b input 1024 (two ranks hold half of each) +CFG = SimpleNamespace(n_layers=4, dim=256, inter=512, o_groups=4, o_rank=256) +WO_B = "layers.{}.attn.wo_b" +SHARED = "layers.{}.ffn.shared_experts.w2" +PREFIX = "base_model.model." + + +def bf16_bytes(x: np.ndarray) -> bytes: + """float32 -> bfloat16 bytes (round to nearest even).""" + + u = np.ascontiguousarray(x, dtype="> 16) & 1)) >> 16).astype(" bytes: + if dtype == "BF16": + return bf16_bytes(x) + return np.ascontiguousarray(x, dtype=" None: + """``tensors``: name -> (dtype word, shape, raw bytes).""" + + header, off = {"__metadata__": {"format": "pt"}}, 0 + for name, (dtype, shape, raw) in tensors.items(): + header[name] = {"dtype": dtype, "shape": list(shape), "data_offsets": [off, off + len(raw)]} + off += len(raw) + head = json.dumps(header).encode() + head += b" " * (-len(head) % 8) + with open(path, "wb") as f: + f.write(struct.pack(" Path: + """An adapter folder: ``modules`` maps a module name (``layers.1.attn.wo_b``) to (A [r, in], B [out, r]).""" + + folder.mkdir(parents=True, exist_ok=True) + tensors = {} + for name, (a, b) in modules.items(): + tensors[f"{prefix}{name}.lora_A.weight"] = (dtype, list(a.shape), encode(a, dtype)) + tensors[f"{prefix}{name}.lora_B.weight"] = (dtype, list(b.shape), encode(b, dtype)) + for name, spec in (extra or {}).items(): + tensors[name] = spec + write_safetensors(folder / "adapter_model.safetensors", tensors) + ranks = {name: int(a.shape[0]) for name, (a, _) in modules.items()} + rank = min(ranks.values(), default=1) # PEFT's default rank; the others come in a rank_pattern + pattern = {re.escape(name): k for name, k in ranks.items() if k != rank} + cfg = {"peft_type": "LORA", "r": rank, "lora_alpha": rank, "bias": "none", "target_modules": ["wo_b", "w2"], + "base_model_name_or_path": "base", **({"rank_pattern": pattern} if pattern else {}), **(config or {})} + (folder / "adapter_config.json").write_text(json.dumps(cfg)) + if manifest is not None: + (folder / "tensorfold_adapter.json").write_text(json.dumps(manifest)) + return folder + + +def standard(folder: Path, seed: int = 0, ranks: tuple[int, int] = (2, 3), dtype: str = "F32", **kw) -> Path: + """Layer 1's wo_b at ranks[0] and layer 2's shared w2 at ranks[1].""" + + rng = np.random.default_rng(seed) + wo_in = CFG.o_groups * CFG.o_rank + mods = {WO_B.format(1): factors(rng, CFG.dim, wo_in, ranks[0]), + SHARED.format(2): factors(rng, CFG.dim, CFG.inter, ranks[1])} + return write_adapter(folder, mods, dtype=dtype, **kw) + + +def model_dir(root: Path) -> Path: + root.mkdir(parents=True, exist_ok=True) + (root / "config.json").write_text(json.dumps({"model_type": "deepseek_v41", "hidden_size": CFG.dim})) + return root + + +def full_config() -> dict: + """A config.json the family's Cfg reads, at the small sizes of CFG.""" + + return {"model_type": "deepseek_v41", "vocab_size": 1024, "hidden_size": CFG.dim, "moe_intermediate_size": CFG.inter, + "num_hidden_layers": CFG.n_layers, "num_attention_heads": 4, "head_dim": 64, "qk_rope_head_dim": 16, + "q_lora_rank": 64, "o_lora_rank": CFG.o_rank, "o_groups": CFG.o_groups, "sliding_window": 128, + "rms_norm_eps": 1e-6, "n_routed_experts": 8, "num_experts_per_tok": 2, "routed_scaling_factor": 1.5, + "swiglu_limit": 10.0, "compress_ratios": [0, 4, 128, 0], "kv_source_layer_ids": [1, 2], + "index_source_layer_ids": [1], "index_n_heads": 4, "index_head_dim": 32, "index_topk": 16, "hc_mult": 4, + "hc_sinkhorn_iters": 20, "hc_eps": 1e-6, "rope_theta": 10000.0, "compress_rope_theta": 160000.0} diff --git a/tests/adapter_kernel_checks.py b/tests/adapter_kernel_checks.py new file mode 100644 index 0000000000..4305c5d622 --- /dev/null +++ b/tests/adapter_kernel_checks.py @@ -0,0 +1,134 @@ +"""Checks of the adapter delta kernels against the plain-torch definition, run on a GPU or, for the logic, by Triton's CPU +interpreter (tests/test_deepseek_v41_adapter_kernels_cpu.py). Each check raises AssertionError.""" + +from __future__ import annotations + +import torch + +from tensorfold.families.deepseek_v41.cuda import lora_kernels as lk +from tensorfold.families.deepseek_v41.cuda import lora_reference as ref +from tensorfold.families.deepseek_v41.cuda.lora_tables import Tables + + +def make_tables(device, ranks, k, d, x_dtype, seed=0, scales=None, junk=False): + """Tables of slots 1..len(ranks); a rank of 0 leaves the slot without this module. ``junk``: nonzero values past + each slot's rank, which the kernels must never read.""" + + g = torch.Generator().manual_seed(seed) + t = Tables.empty(len(ranks) + 1, max(max(ranks), 1), k, d, False, x_dtype, device) + for slot, r in enumerate(ranks, 1): + if r: + t.put(slot, torch.randn((r, k), generator=g) * 0.05, torch.randn((d, r), generator=g) * 0.05, + (scales or [1.0] * len(ranks))[slot - 1]) + if junk: + t.a[slot, r:] = 7.0 + t.b[slot, :, r:] = 7.0 + return t + + +def rand(device, shape, dtype, seed): + g = torch.Generator().manual_seed(seed) + return torch.randn(shape, generator=g).to(dtype).to(device) + + +def run(out, x, ids, tab): + return lk.lowrank_add(out.clone(), x, ids, tab) + + +def want(out, x, ids, tab): + return ref.lowrank_add(out.clone().double(), x, ids, tab.a, tab.b, tab.ranks, tab.scales, dtype=torch.float64) + + +def check_matches_the_definition(device, x_dtype, k=256, d=384, ranks=(1, 3, 5), n=6, seed=1): + tab = make_tables(device, ranks, k, d, x_dtype, seed, scales=[1.0, 0.25, -2.0][:len(ranks)]) + x, out = rand(device, (n, k), x_dtype, seed + 1), rand(device, (n, d), torch.float32, seed + 2) + ids = torch.tensor([(i % (len(ranks) + 1)) for i in range(n)], dtype=torch.int64, device=device) + got, exp = run(out, x, ids, tab), want(out, x, ids, tab) + err = (got.double() - exp).abs().max().item() + assert err <= 2e-5 * max(1.0, exp.abs().max().item()), f"max error {err}" + return err + + +def check_stock_rows_keep_their_bits(device, x_dtype, k=256, d=384): + tab = make_tables(device, (2, 4), k, d, x_dtype, 3) + x, out = rand(device, (5, k), x_dtype, 4), rand(device, (5, d), torch.float32, 5) + ids = torch.tensor([0, 2, 0, 1, 0], dtype=torch.int64, device=device) + got = run(out, x, ids, tab) + for r in (0, 2, 4): + assert torch.equal(got[r], out[r]), f"stock row {r} changed" + assert not torch.equal(got[1], out[1]) and not torch.equal(got[3], out[3]) + + +def check_a_row_does_not_depend_on_the_others(device, x_dtype, k=384, d=256): + """Row by row, in a different order, beside other adapters and padded by copies: the same delta bits.""" + + tab = make_tables(device, (2, 3, 1), k, d, x_dtype, 6) + n = 7 + x, out = rand(device, (n, k), x_dtype, 7), rand(device, (n, d), torch.float32, 8) + ids = torch.tensor([3, 1, 0, 2, 1, 3, 2], dtype=torch.int64, device=device) + whole = run(out, x, ids, tab) + for r in range(n): # a row alone + alone = run(out[r:r + 1], x[r:r + 1].contiguous(), ids[r:r + 1].contiguous(), tab) + assert torch.equal(alone[0], whole[r]), f"row {r} differs alone" + perm = torch.tensor([6, 2, 4, 0, 5, 1, 3], device=device) # another order + again = run(out[perm], x[perm].contiguous(), ids[perm].contiguous(), tab) + assert torch.equal(again, whole[perm]), "a reordered batch differs" + pad = torch.cat([torch.arange(n, device=device), torch.full((9,), n - 1, device=device)]) # padded with copies + padded = run(out[pad], x[pad].contiguous(), ids[pad].contiguous(), tab) + assert torch.equal(padded[:n], whole) and torch.equal(padded[n:], whole[n - 1:].expand(9, -1)), "padding changed bits" + + +def check_nothing_past_a_slots_rank_is_read(device, x_dtype, k=256, d=256): + clean = make_tables(device, (2, 5), k, d, x_dtype, 9) + dirty = make_tables(device, (2, 5), k, d, x_dtype, 9, junk=True) + x, out = rand(device, (4, k), x_dtype, 10), rand(device, (4, d), torch.float32, 11) + ids = torch.tensor([1, 2, 1, 2], dtype=torch.int64, device=device) + assert torch.equal(run(out, x, ids, clean), run(out, x, ids, dirty)) + + +def check_a_strided_view_of_the_hidden_state_is_read_in_place(device, k=256, d=128, slots=7): + """The packed shared expert's rotated hidden is one slot of [rows, slots, K] fp16.""" + + tab = make_tables(device, (3,), k, d, torch.float16, 12) + big = rand(device, (5, slots, k), torch.float16, 13) + out = rand(device, (5, d), torch.float32, 14) + ids = torch.tensor([1, 0, 1, 1, 0], dtype=torch.int64, device=device) + view = big[:, slots - 1] + assert view.stride() == (slots * k, 1) + assert torch.equal(run(out, view, ids, tab), run(out, view.contiguous(), ids, tab)) + + +def check_a_slot_without_the_module_adds_nothing(device, x_dtype, k=256, d=256): + tab = make_tables(device, (0, 3), k, d, x_dtype, 15) + x, out = rand(device, (3, k), x_dtype, 16), rand(device, (3, d), torch.float32, 17) + ids = torch.tensor([1, 2, 1], dtype=torch.int64, device=device) + got = run(out, x, ids, tab) + assert torch.equal(got[0], out[0]) and torch.equal(got[2], out[2]) and not torch.equal(got[1], out[1]) + + +def check_refusals(device): + tab = make_tables(device, (2,), 256, 128, torch.bfloat16, 18) + x, out = rand(device, (2, 256), torch.bfloat16, 19), rand(device, (2, 128), torch.float32, 20) + ids = torch.tensor([1, 0], dtype=torch.int64, device=device) + for args in ((out.to(torch.float16), x, ids), (out, x.to(torch.float16), ids), (out, x, ids.int()), + (out, x[:, :128], ids), (out, x, ids[:1]), (out[:, :64], x, ids)): + try: + lk.lowrank_add(args[0], args[1], args[2], tab) + except ValueError: + continue + raise AssertionError("a malformed call was accepted") + + +def run_all(device) -> dict: + results = {} + for dt in (torch.bfloat16, torch.float16, torch.float32): + name = str(dt).split(".")[1] + results[f"definition_{name}"] = check_matches_the_definition(device, dt) + check_stock_rows_keep_their_bits(device, dt) + check_a_row_does_not_depend_on_the_others(device, dt) + check_nothing_past_a_slots_rank_is_read(device, dt) + check_a_slot_without_the_module_adds_nothing(device, dt) + results["definition_wide_rank"] = check_matches_the_definition(device, torch.bfloat16, ranks=(17, 33), n=3, seed=21) + check_a_strided_view_of_the_hidden_state_is_read_in_place(device) + check_refusals(device) + return results diff --git a/tests/cuda/test_deepseek_v41_adapter_kernels.py b/tests/cuda/test_deepseek_v41_adapter_kernels.py new file mode 100644 index 0000000000..a0b8176e6e --- /dev/null +++ b/tests/cuda/test_deepseek_v41_adapter_kernels.py @@ -0,0 +1,90 @@ +"""The adapter delta kernels compiled for the GPU: the plain-torch definition, independent rows (alone, reordered, padded), +production shapes and ranks, and a captured CUDA graph that replays any mixture of adapters from its input vector.""" + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device", allow_module_level=True) + +import adapter_kernel_checks as checks + +from tensorfold.families.deepseek_v41.cuda import lora_kernels as lk + +DEV = "cuda" + + +def test_the_compiled_kernels_pass_every_check_of_the_cpu_interpreter_run(): + errors = checks.run_all(DEV) + assert max(errors.values()) < 1e-4, errors + + +# one rank's shapes: wo_b reads 4096 columns (4 groups of 1024) and writes 5120; the shared w2 reads 1152 +@pytest.mark.parametrize("k", [4096, 1152]) +@pytest.mark.parametrize("ranks", [(1,), (2, 3), (16, 7), (17, 64), (33, 1, 128)]) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16, torch.float16]) +def test_production_shapes_and_ranks_match_the_definition(k, ranks, x_dtype): + assert checks.check_matches_the_definition(DEV, x_dtype, k=k, d=5120, ranks=ranks, n=16, seed=len(ranks)) < 1e-4 + + +def test_a_prompt_chunk_of_many_rows_matches_the_definition_and_keeps_every_row_independent(): + tab = checks.make_tables(DEV, (4, 9), 4096, 5120, torch.bfloat16, 30) + n = 2048 + x, out = checks.rand(DEV, (n, 4096), torch.bfloat16, 31), checks.rand(DEV, (n, 5120), torch.float32, 32) + ids = (torch.arange(n, device=DEV) % 3).to(torch.int64) + whole = checks.run(out, x, ids, tab) + sample = torch.tensor([0, 1, 2, 3, 700, 1023, 2047], device=DEV) + exp = checks.want(out[sample], x[sample], ids[sample], tab) + assert (whole[sample].double() - exp).abs().max().item() < 1e-4 + for r in sample.tolist(): # a decode-sized call gives the chunk's bits + alone = checks.run(out[r:r + 1], x[r:r + 1].contiguous(), ids[r:r + 1].contiguous(), tab) + assert torch.equal(alone[0], whole[r]) + stock = ids == 0 + assert torch.equal(whole[stock], out[stock]) + + +@pytest.mark.parametrize("rows", [1, 3, 8, 16]) +def test_a_captured_graph_replays_any_mixture_of_adapters_from_its_input_vector(rows): + tab = checks.make_tables(DEV, (2, 5, 1), 1152, 5120, torch.float16, 40) + x = checks.rand(DEV, (rows, 1152), torch.float16, 41) + out0 = checks.rand(DEV, (rows, 5120), torch.float32, 42) + ids = torch.zeros((rows,), dtype=torch.int64, device=DEV) # the graph's input, like RoundDecoder's + out = out0.clone() + + def body(): + out.copy_(out0) + lk.lowrank_add(out, x, ids, tab) + + side = torch.cuda.Stream() + with torch.cuda.stream(side): + body() # the warm-up the captures follow + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + body() + g = torch.Generator().manual_seed(43) + for _ in range(6): + ids.copy_(torch.randint(0, 4, (rows,), generator=g).to(DEV)) + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, checks.run(out0, x, ids, tab)), ids.tolist() # the eager call's bits + ids.zero_() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(out, out0) # all stock: nothing added + + +def test_another_adapters_larger_rank_does_not_change_a_rows_bits(): + """The same adapter in tables made wider by a second adapter gives the same bits: the kernels' tile is fixed and + read only a slot's rank, so a server's other adapters never move an adapter's output.""" + + k, d = 1152, 512 + narrow = checks.make_tables(DEV, (3,), k, d, torch.bfloat16, 50) + wide = checks.make_tables(DEV, (3, 40), k, d, torch.bfloat16, 50) + wide.put(1, narrow.a[1, :3], narrow.b[1, :, :3], float(narrow.scales[1])) + x, out = checks.rand(DEV, (4, k), torch.bfloat16, 51), checks.rand(DEV, (4, d), torch.float32, 52) + ids = torch.ones((4,), dtype=torch.int64, device=DEV) + assert narrow.R < wide.R + assert torch.equal(checks.run(out, x, ids, narrow), checks.run(out, x, ids, wide)) diff --git a/tests/cuda/test_deepseek_v41_adapter_shared.py b/tests/cuda/test_deepseek_v41_adapter_shared.py new file mode 100644 index 0000000000..068a1990b6 --- /dev/null +++ b/tests/cuda/test_deepseek_v41_adapter_shared.py @@ -0,0 +1,112 @@ +"""The packed shared expert's w2 adapter reads what the EXL3 expert kernels leave for the down projection. The factors are +folded at load to read that rotated, scaled fp16 hidden state (lora_place.fold_rotated); these tests run the real expert +kernels on synthetic layers and check both the premise (the scratch holds Hadamard(hidden * suh_d) for the shared slot, +decode and prompt paths) and the delta computed from it.""" + +import math + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device", allow_module_level=True) + +from test_exl3_experts import _hadamard, _layer, _picks, _rot + +from tensorfold.cuda.exl3 import experts +from tensorfold.families.deepseek_v41.cuda import lora_format as lf +from tensorfold.families.deepseek_v41.cuda import lora_kernels as lk +from tensorfold.families.deepseek_v41.cuda import lora_place as lp +from tensorfold.families.deepseek_v41.cuda.lora_tables import Tables + +D, I, E, TOPK = 512, 256, 10, 3 # E routed experts, and the shared one last in the pack +SLOTS = TOPK + 1 +CB = 2 + + +def scratch_for(ex, rows): + prompt = rows >= experts.EXACT_ROWS # Model._moe_scratch: prompt chunks take a scratch of their own + return experts.Scratch(ex, rows=rows if prompt else max(rows, experts.EXACT_ROWS), slots=SLOTS, prompt=prompt) + + +def hidden_of_the_shared_expert(x, mats): + """The shared expert's SwiGLU hidden state [rows, I] in float64, from the decoded weights.""" + + gate, up, _ = mats + H = _hadamard(x.device) + + def lin(v, m): + t, suh, svh = m + return _rot(_rot(v * suh.double(), H) @ experts.dequant(t, CB).double(), H) * svh.double() + + out = [] + for r in range(x.shape[0]): + gg, uu = lin(x[r].double(), gate[-1]), lin(x[r].double(), up[-1]) + out.append(gg * torch.sigmoid(gg) * uu) + return torch.stack(out) + + +@pytest.mark.parametrize("rows", [1, 4, 16, 63, 64, 130]) +def test_the_scratch_holds_the_rotated_scaled_hidden_state_of_the_shared_slot(rows): + ex, mats = _layer(E + 1, D, I, [(4, 4, 4)] * (E + 1), CB, seed=3) + g = torch.Generator().manual_seed(rows) + x = torch.randn((rows, D), generator=g).to(torch.bfloat16).cuda() + sel, w = _picks(E, rows, TOPK, g, shared=True) + assert int(sel[0, -1]) == E == ex.count - 1 # the last pack entry is the shared slot's pick + s = scratch_for(ex, rows) + experts.routed(x, sel, w, ex, s, None, rows) + xd = s.xd[:rows * SLOTS].view(rows, SLOTS, I)[:, SLOTS - 1] + act = hidden_of_the_shared_expert(x, mats) + want = _rot(act * ex.suh_d[E].double(), _hadamard(x.device)) + assert xd.dtype == torch.float16 + assert (xd.double() - want).abs().max().item() <= 2e-3 * want.abs().max().item() + + +@pytest.mark.parametrize("rows", [1, 6, 64, 100]) +def test_the_folded_factors_add_the_delta_of_the_plain_hidden_state(rows): + ex, mats = _layer(E + 1, D, I, [(4, 4, 4)] * (E + 1), CB, seed=4) + g = torch.Generator().manual_seed(40 + rows) + x = torch.randn((rows, D), generator=g).to(torch.bfloat16).cuda() + sel, w = _picks(E, rows, TOPK, g, shared=True) + s = scratch_for(ex, rows) + out = experts.routed(x, sel, w, ex, s, None, rows) + xd = s.xd[:rows * SLOTS].view(rows, SLOTS, I)[:, SLOTS - 1] + act = hidden_of_the_shared_expert(x, mats) + + rng = np.random.default_rng(5) + dims = lf.Dims(n_layers=1, dim=D, inter=I, o_in=I) + f = lf.Factors((rng.standard_normal((5, I)) * 0.1).astype(np.float32), + (rng.standard_normal((D, 5)) * 0.1).astype(np.float32), 0.5) + placed = lp.place(lf.SHARED_W2, f, dims, 0, 1, strength=2.0, suh=ex.suh_d[E].float().cpu().numpy()) + assert placed.folded + tab = Tables.empty(3, 5, I, D, True, torch.float16, "cuda") + tab.put(2, placed.a, placed.b, placed.scale) + ids = (torch.arange(rows, device="cuda") % 3 == 2).to(torch.int64) * 2 # every third row runs slot 2 + got = lk.lowrank_add(out.clone(), xd, ids, tab).double() + + delta = placed.scale * (act @ torch.from_numpy(f.a).double().cuda().T) @ torch.from_numpy(f.b).double().cuda().T + want = out.double() + delta * (ids != 0).double()[:, None] + scale = delta.abs().max().item() + assert (got - want).abs().max().item() <= 4e-3 * scale # fp16 storage of the scratch + keep = ids == 0 + assert torch.equal(got[keep].float(), out[keep]) # the rows of the base model: its bits + + +def test_the_delta_is_the_same_for_a_row_in_a_small_window_and_in_a_prompt_chunk(): + """The scratch of a prompt chunk and of a decode window hold the same rotated hidden state for a row: fp16 apart.""" + + ex, _ = _layer(E + 1, D, I, [(4, 4, 4)] * (E + 1), CB, seed=6) + g = torch.Generator().manual_seed(7) + rows = 128 + x = torch.randn((rows, D), generator=g).to(torch.bfloat16).cuda() + sel, w = _picks(E, rows, TOPK, g, shared=True) + big = scratch_for(ex, rows) + experts.routed(x, sel, w, ex, big, None, rows) + small = scratch_for(ex, 8) + experts.routed(x[:8].contiguous(), sel[:8].contiguous(), w[:8].contiguous(), ex, small, None, 8) + a = big.xd[:rows * SLOTS].view(rows, SLOTS, I)[:8, SLOTS - 1].double() + b = small.xd[:8 * SLOTS].view(8, SLOTS, I)[:, SLOTS - 1].double() + assert (a - b).abs().max().item() <= 4e-3 * b.abs().max().item() + assert math.isfinite(float(a.sum())) diff --git a/tests/cuda/test_deepseek_v41_adapter_tp.py b/tests/cuda/test_deepseek_v41_adapter_tp.py new file mode 100644 index 0000000000..9f488ed241 --- /dev/null +++ b/tests/cuda/test_deepseek_v41_adapter_tp.py @@ -0,0 +1,107 @@ +"""Two ranks on one GPU over gloo, as the Nemotron TP test runs them: each rank places its share of an adapter, adds its +partial delta with the compiled kernels, and the ranks' sum (the reduction the layer already makes) equals the unsharded +delta of the whole matrix. For wo_b and for the packed shared expert (folded factors on each rank's own scale vector).""" + +import socket + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("needs a CUDA device", allow_module_level=True) + +ROWS = 5 +DIMS = (4, 256, 512, 1024) # layers, dim, shared width, wo_b width +SLOTS_OF_ROWS = [0, 1, 2, 1, 0] + + +def _free_port() -> int: + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def problem(kind): + """Both ranks and the test build the same numbers from the same seeds.""" + + from tensorfold.families.deepseek_v41.cuda import lora_format as lf + + dims = lf.Dims(*DIMS) + width, out = dims.shape(kind) + rng = np.random.default_rng(11) + mods = [lf.Factors((rng.standard_normal((r, width)) * 0.1).astype(np.float32), + (rng.standard_normal((out, r)) * 0.1).astype(np.float32), s) for r, s in ((2, 1.0), (5, 0.5))] + x = rng.standard_normal((ROWS, width)).astype(np.float32) + suh = (rng.uniform(0.5, 2.0, width) * np.where(rng.random(width) < 0.5, -1.0, 1.0)).astype(np.float16) + return dims, mods, x, suh.astype(np.float64) + + +def _worker(rank: int, port: int, kind: str, out) -> None: + import torch.distributed as dist + + from tensorfold.families.deepseek_v41.cuda import lora_kernels as lk + from tensorfold.families.deepseek_v41.cuda import lora_place as lp + from tensorfold.families.deepseek_v41.cuda.lora_tables import Tables + + torch.cuda.set_device(0) + dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=2) + dims, mods, x, suh = problem(kind) + width, dim = dims.shape(kind) + lo, hi = lp.span(width, rank, 2) + packed = kind == "shared_w2" + if packed: # what the expert kernels leave for the down projection + z = (x[:, lo:hi].astype(np.float64) * suh[lo:hi]).reshape(ROWS, -1, 128) @ lp.hadamard() + local = torch.from_numpy(z.reshape(ROWS, hi - lo)).to(torch.float16).cuda() + else: + local = torch.from_numpy(x[:, lo:hi]).to(torch.bfloat16).cuda() + tab = Tables.empty(3, 5, hi - lo, dim, packed, local.dtype, "cuda") + for slot, f in enumerate(mods, 1): + p = lp.place(kind, f, dims, rank, 2, suh=suh[lo:hi] if packed else None) + tab.put(slot, p.a, p.b, p.scale) + ids = torch.tensor(SLOTS_OF_ROWS, dtype=torch.int64, device="cuda") + partial = lk.lowrank_add(torch.zeros((ROWS, dim), dtype=torch.float32, device="cuda"), local, ids, tab) + total = partial.cpu() + dist.all_reduce(total) # the reduction the layer's gather already makes + out.put((rank, partial.cpu().numpy(), total.numpy())) + dist.barrier() + dist.destroy_process_group() + + +def run_pair(kind): + import torch.multiprocessing as mp + + ctx = mp.get_context("spawn") + out = ctx.Queue() + port = _free_port() + procs = [ctx.Process(target=_worker, args=(r, port, kind, out)) for r in (0, 1)] + for p in procs: + p.start() + got = {r: (part, total) for r, part, total in (out.get(timeout=600) for _ in procs)} + for p in procs: + p.join(timeout=120) + assert p.exitcode == 0 + return got + + +def reference(kind): + dims, mods, x, _ = problem(kind) + bf16 = torch.from_numpy(x).to(torch.bfloat16).float().numpy().astype(np.float64) + out = np.zeros((ROWS, dims.dim)) + for i, slot in enumerate(SLOTS_OF_ROWS): + if slot: + f = mods[slot - 1] + out[i] = f.scale * (f.b.astype(np.float64) @ (f.a.astype(np.float64) @ (x if kind == "shared_w2" else bf16)[i])) + return out + + +@pytest.mark.parametrize("kind", ["wo_b", "shared_w2"]) +def test_two_ranks_partial_deltas_sum_to_the_whole_matrices_delta(kind): + got = run_pair(kind) + assert np.array_equal(got[0][1], got[1][1]) # both ranks hold the same sum + want = reference(kind) + tol = 4e-3 if kind == "shared_w2" else 1e-4 # fp16 storage of the rotated hidden state + assert np.abs(got[0][1] - want).max() <= tol * np.abs(want).max() + stock = [i for i, s in enumerate(SLOTS_OF_ROWS) if s == 0] + assert all(not got[r][0][stock].any() for r in (0, 1)) # a stock row: no partial on either rank diff --git a/tests/test_deepseek_v41_adapter_engine.py b/tests/test_deepseek_v41_adapter_engine.py new file mode 100644 index 0000000000..b3e391995b --- /dev/null +++ b/tests/test_deepseek_v41_adapter_engine.py @@ -0,0 +1,108 @@ +"""The engine's adapter entry points without a GPU: what start refuses, what a name resolves to, and what generate does +with an adapter when no concurrent scheduler is there to carry it.""" + +from types import SimpleNamespace + +import pytest +from adapter_fakes import CFG, model_dir, standard + +pytest.importorskip("torch") +pytest.importorskip("triton") + +from tensorfold.families.deepseek_v41.cuda import engine as eng # noqa: E402 +from tensorfold.families.deepseek_v41.cuda import lora_registry as lr # noqa: E402 +from tensorfold.families.deepseek_v41.cuda import lora_runtime # noqa: E402 +from tensorfold.families.deepseek_v41.cuda import model as _model # noqa: E402 + + +def bare(tmp_path, adapters=None): + e = eng.DsEngine.__new__(eng.DsEngine) + e.w, e.rank, e.world, e.model_dir = SimpleNamespace(cfg=CFG, layers=[]), 0, 2, model_dir(tmp_path / "m") + e.model, e.adapters, e.scheduler = SimpleNamespace(lora=None), adapters, None + e.request = SimpleNamespace(stop_eos=True) + return e + + +def specs(tmp_path): + return [("alpha", str(standard(tmp_path / "a", 1)), 1.0)] + + +def test_start_refuses_adapters_where_a_rows_slot_cannot_reach_the_forward(tmp_path, monkeypatch): + e = bare(tmp_path) + with pytest.raises(ValueError, match="--parallel 2 or more"): + e._load_adapters(specs(tmp_path), (), 1) + monkeypatch.setattr(_model, "KERNELS", False) + with pytest.raises(ValueError, match="kernel path"): + e._load_adapters(specs(tmp_path), (), 4) + monkeypatch.setattr(_model, "KERNELS", True) + monkeypatch.setattr(_model, "EXACT_MM", True) + with pytest.raises(ValueError, match="kernel path"): + e._load_adapters(specs(tmp_path), (), 4) + assert e.model.lora is None + + +def test_start_validates_loads_and_warms_and_hands_back_the_registry(tmp_path, monkeypatch): + seen = {} + + class Runtime: + def warm(self): + seen["warmed"] = True + + def load(registry, weights, rank, world, room=None, **kw): + seen.update(registry=registry, rank=rank, world=world, room=room) + return Runtime() + + monkeypatch.setattr(_model, "KERNELS", True) + monkeypatch.setattr(_model, "EXACT_MM", False) + monkeypatch.setattr(lora_runtime.LoraRuntime, "load", staticmethod(load)) + e = bare(tmp_path) + registry = e._load_adapters(specs(tmp_path), ("base",), 2) + assert registry.names() == ["alpha"] and seen["registry"] is registry and seen["warmed"] + assert (seen["rank"], seen["world"]) == (0, 2) and isinstance(e.model.lora, Runtime) + assert seen["room"] is None or isinstance(seen["room"](), float) + with pytest.raises(Exception, match="already a model id"): + bare(tmp_path / "again")._load_adapters(specs(tmp_path / "again"), ("alpha",), 2) + + +def test_ranks_started_with_other_adapters_or_another_order_are_refused_and_the_message_says_so(tmp_path): + a, b, base = str(standard(tmp_path / "a", 1)), str(standard(tmp_path / "b", 2)), model_dir(tmp_path / "m") + + def flags(*held): # what a rank gathers: vision, parallel, Markov x2, the digest + return [0, 4, 0, 0, *lr.AdapterRegistry.load(list(held), CFG, base, ("base",)).words()] + + both = flags(("x", a, 1.0), ("y", b, 1.0)) + eng._same_on_every_rank([both, flags(("x", a, 1.0), ("y", b, 1.0))] * 2) # the same flags: nothing raised + for other in (flags(("y", b, 1.0), ("x", a, 1.0)), # the order alone + flags(("x", a, 1.0), ("y", b, 0.5)), flags(("x", a, 1.0))): # a strength; a missing adapter + with pytest.raises(ValueError, match="--adapter files, names, strengths and order") as seen: + eng._same_on_every_rank([both, other]) + assert str([both, other]) in str(seen.value) + + +def test_a_name_resolves_to_a_reference_and_a_missing_one_is_an_error(tmp_path): + registry = lr.AdapterRegistry.load(specs(tmp_path), CFG, model_dir(tmp_path / "m"), ("base",)) + e = bare(tmp_path, registry) + ref = e.adapter_ref("alpha") + assert ref == registry.ref("alpha") and e.adapter_ref(None) is None + assert e.adapter_ok(ref) and e.adapter_ok(None) and not e.adapter_ok(lr.AdapterRef(1, "0" * 32)) + with pytest.raises(lr.UnknownAdapter, match="no adapter named"): + e.adapter_ref("beta") + without = bare(tmp_path / "plain") + assert without.adapter_ok(None) and not without.adapter_ok(ref) + with pytest.raises(lr.UnknownAdapter, match="without --adapter"): + without.adapter_ref("alpha") + + +def test_generate_hands_the_reference_to_the_scheduler_and_refuses_it_without_one(tmp_path): + registry = lr.AdapterRegistry.load(specs(tmp_path), CFG, model_dir(tmp_path / "m"), ("base",)) + e = bare(tmp_path, registry) + got = {} + e.scheduler = SimpleNamespace(submit=lambda *a, **k: got.update(k) or {"ok": 1}) + e.tower = None + assert e.generate([1, 2], 4, None, lambda new: None, adapter="alpha") == {"ok": 1} + assert got["adapter"] == registry.ref("alpha") + e.generate([1, 2], 4, None, lambda new: None) + assert got["adapter"] is None # the base model: the reference is None + e.scheduler = None + with pytest.raises(ValueError, match="concurrent engine"): + e.generate([1, 2], 4, None, lambda new: None, adapter="alpha") diff --git a/tests/test_deepseek_v41_adapter_format.py b/tests/test_deepseek_v41_adapter_format.py new file mode 100644 index 0000000000..cb03993cd8 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_format.py @@ -0,0 +1,343 @@ +"""DeepSeek-V4.1 adapter folders: what is read, what is refused, and that nothing is dropped quietly (no torch, no GPU).""" + +import json +import struct + +import numpy as np +import pytest +from adapter_fakes import CFG, PREFIX, SHARED, WO_B, factors, standard, write_adapter, write_safetensors + +from tensorfold.families.deepseek_v41.cuda import lora_format as lf + +DIMS = lf.Dims.of(CFG) +WO_IN = CFG.o_groups * CFG.o_rank + + +def rng(seed=0): + return np.random.default_rng(seed) + + +def test_a_standard_adapter_is_read_with_each_modules_rank_and_scale(tmp_path): + spec = lf.inspect_folder(standard(tmp_path / "a", ranks=(2, 3), config={"lora_alpha": 8}), DIMS) + assert [(m.layer, m.kind, m.rank) for m in spec.modules] == [(1, "wo_b", 2), (2, "shared_w2", 3)] + assert [m.scale for m in spec.modules] == [8 / 2, 8 / 3] + assert spec.describe()["modules"]["layers.2.ffn.shared_experts.w2"] == {"rank": 3, "alpha": 8.0, "scale": 8 / 3} + + +def test_rank_and_alpha_patterns_give_each_module_its_own_scale(tmp_path): + r = rng() + mods = {WO_B.format(0): factors(r, CFG.dim, WO_IN, 2), WO_B.format(3): factors(r, CFG.dim, WO_IN, 4)} + cfg = {"r": 2, "lora_alpha": 2, "rank_pattern": {r"layers\.3\.attn\.wo_b": 4}, + "alpha_pattern": {r"layers\.3\.attn\.wo_b": 1}} + spec = lf.inspect_folder(write_adapter(tmp_path / "a", mods, config=cfg), DIMS) + assert [(m.rank, m.alpha, m.scale) for m in spec.modules] == [(2, 2.0, 1.0), (4, 1.0, 0.25)] + + +def test_rslora_scales_by_the_square_root_of_the_rank(tmp_path): + spec = lf.inspect_folder(standard(tmp_path / "a", ranks=(4, 4), config={"use_rslora": True, "lora_alpha": 4, + "r": 4}), DIMS) + assert [m.scale for m in spec.modules] == [2.0, 2.0] + + +@pytest.mark.parametrize("prefix", [PREFIX, "base_model.model.model.", "model.", ""]) +def test_the_common_prefixes_name_the_same_module(tmp_path, prefix): + mods = {WO_B.format(2): factors(rng(), CFG.dim, WO_IN, 2)} + spec = lf.inspect_folder(write_adapter(tmp_path / "a", mods, prefix=prefix), DIMS) + assert [(m.layer, m.kind) for m in spec.modules] == [(2, "wo_b")] + + +@pytest.mark.parametrize("dtype", ["F32", "F16", "BF16"]) +def test_every_served_dtype_widens_exactly(tmp_path, dtype): + folder = standard(tmp_path / "a", dtype=dtype) + spec = lf.inspect_folder(folder, DIMS) + header, base = lf.read_header(spec.weights) + for m in spec.modules: + a = lf.read_tensor(spec.weights, header[m.a_name], base) + assert a.dtype == np.float32 and a.shape == (m.rank, DIMS.shape(m.kind)[0]) + if dtype == "BF16": # bf16 values are the high half of the float32 + raw = np.fromfile(spec.weights, dtype=np.uint8) + lo, hi = header[spec.modules[0].a_name]["data_offsets"] + want = (raw[base + lo:base + hi].view("= {"definition_bfloat16", "definition_float16", "definition_float32", "definition_wide_rank"} + assert max(errors.values()) < 1e-3 diff --git a/tests/test_deepseek_v41_adapter_mix_tool.py b/tests/test_deepseek_v41_adapter_mix_tool.py new file mode 100644 index 0000000000..fa0eb2c168 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_mix_tool.py @@ -0,0 +1,99 @@ +"""tools/dsv41/adapter_mix.py against a stand-in server: equal hashes pass, a reply that depends on its neighbours (or on +drafting) is reported and fails the run, and the model ids rotate over the streams.""" + +import hashlib +import importlib.util +import json +import pathlib +import threading +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +TOOL = pathlib.Path(__file__).resolve().parents[1] / "tools" / "dsv41" / "adapter_mix.py" + + +def tool(): + spec = importlib.util.spec_from_file_location("adapter_mix_under_test", TOOL) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def serve(leaky: bool = False, no_hash: bool = False): + """A server whose reply hash is a function of (model, prompt, seed): and, when ``leaky``, of how many requests + were in flight.""" + + live, lock, seen = [0], threading.Lock(), [] + + class Handler(BaseHTTPRequestHandler): + def log_message(self, *args): + pass + + def do_POST(self): + body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + seen.append((body["model"], body.get("draft", True))) + with lock: + live[0] += 1 + crowd = live[0] + try: + import time + + time.sleep(0.05) + with lock: + crowd = max(crowd, live[0]) + key = f"{body['model']}|{body['prompt']}|{body['seed']}" + (f"|{crowd}" if leaky else "") + sha = hashlib.sha256(key.encode()).hexdigest()[:16] + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + self.wfile.write(b'data: {"choices": [{"text": "x"}]}\n\n') + last = {"choices": [], "usage": {"completion_tokens": 3}} + if not no_hash: + last["tensorfold"] = {"token_sha": sha} + self.wfile.write(b"data: " + json.dumps(last).encode() + b"\n\ndata: [DONE]\n\n") + finally: + with lock: + live[0] -= 1 + + httpd = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + return httpd, seen + + +def run(module, port, monkeypatch, capsys, *argv): + monkeypatch.setattr("sys.argv", ["adapter_mix.py", f"http://127.0.0.1:{port}", *argv]) + code = module.main() + return code, [json.loads(line) for line in capsys.readouterr().out.splitlines()] + + +def test_equal_hashes_pass_and_every_id_takes_every_place_in_turn(monkeypatch, capsys): + httpd, seen = serve() + try: + code, lines = run(tool(), httpd.server_port, monkeypatch, capsys, "--models", "base,a,b", "--streams", "4", + "--tokens", "8") + finally: + httpd.shutdown() + assert code == 0 and len(lines) == 3 + assert [line["models"] for line in lines] == [["base", "a", "b", "base"], ["a", "b", "base", "a"], + ["b", "base", "a", "b"]] + assert all(all(line["concurrent_equals_solo"]) and all(line["solo_equals_serial"]) for line in lines) + assert {m for m, _ in seen} == {"base", "a", "b"} and (("a", False) in seen) # serial runs ask draft: false + + +def test_a_reply_that_depends_on_the_streams_beside_it_fails_the_run(monkeypatch, capsys): + httpd, _ = serve(leaky=True) + try: + code, lines = run(tool(), httpd.server_port, monkeypatch, capsys, "--models", "base,a", "--streams", "3", + "--tokens", "8") + finally: + httpd.shutdown() + assert code == 1 and any(not all(line["concurrent_equals_solo"]) for line in lines) + + +def test_a_reply_without_a_hash_is_an_error_not_a_pass(monkeypatch, capsys): + httpd, _ = serve(no_hash=True) + try: + code, lines = run(tool(), httpd.server_port, monkeypatch, capsys, "--models", "base", "--streams", "2", + "--tokens", "8") + finally: + httpd.shutdown() + assert code == 1 and "no token_sha" in lines[0]["errors"][0] and not any(lines[0]["concurrent_equals_solo"]) diff --git a/tests/test_deepseek_v41_adapter_model.py b/tests/test_deepseek_v41_adapter_model.py new file mode 100644 index 0000000000..dab07ebd37 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_model.py @@ -0,0 +1,169 @@ +"""The hooks in the DeepSeek-V4.1 model code: wo_b's delta is computed from the very input wo_b consumed, and the shared +expert's from the rotated hidden state still in the expert scratch (packed) or the plain hidden state (apart). The model +functions run on CPU tensors with the GPU kernels and expert extension replaced by stand-ins that write known values.""" + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") + +import torch.nn.functional as F # noqa: E402 + +from tensorfold.families.deepseek_v41.cuda import lora_reference as ref # noqa: E402 +from tensorfold.families.deepseek_v41.cuda import lora_runtime as rt # noqa: E402 +from tensorfold.families.deepseek_v41.cuda import model as M # noqa: E402 +from tensorfold.families.deepseek_v41.cuda.lora_tables import Tables # noqa: E402 + +D, K_IN, I, E, TOPK = 8, 6, 8, 5, 2 +SLOTS = TOPK + 1 +G = torch.Generator().manual_seed(0) + + +def kernel(out, x, ids, tab): + return ref.lowrank_add(out, x, ids, tab.a, tab.b, tab.ranks, tab.scales, dtype=torch.float64) + + +def tables(k, x_dtype, folded, seed): + g = torch.Generator().manual_seed(seed) + t = Tables.empty(3, 3, k, D, folded, x_dtype, "cpu") + t.put(1, torch.randn((2, k), generator=g), torch.randn((D, 2), generator=g), 0.5) + t.put(2, torch.randn((3, k), generator=g), torch.randn((D, 3), generator=g), -1.0) + return t + + +def runtime(idx, wo_b=None, shared=None): + layer = rt.LayerTables() + layer.wo_b, layer.shared = wo_b, shared + return rt.LoraRuntime(None, {idx: layer}, "cpu", kernel) + + +def expected(tab, x, slots, out): + want = out.clone().double() + for i, s in enumerate(slots): + if s: + k = int(tab.ranks[s]) + want[i] += float(tab.scales[s]) * (tab.b[s, :, :k].double() @ (tab.a[s, :k].double() @ x[i].double())) + return want + + +# -- wo_b --------------------------------------------------------------------------------------------------------- +class FakeGroup: + """Exl3Group stand-in: writes known values into wo_a's output slices.""" + + def __init__(self, values): + self.values = values + + def rotated(self, xh, outs, rot=None): + c = 0 + for o in outs: + o.copy_(self.values[:, c:c + o.shape[1]]) + c += o.shape[1] + return outs + + +def test_wo_ab_adds_the_delta_of_the_input_wo_b_consumed(monkeypatch): + n = 4 + u_vals = torch.randn((n, K_IN), generator=G).to(torch.bfloat16) + monkeypatch.setattr(M, "_group", lambda lay, key, layers: FakeGroup(u_vals)) + w = torch.randn((D, K_IN), generator=G).to(torch.bfloat16) + lay = SimpleNamespace(idx=3, wo_a=[SimpleNamespace(n=3), SimpleNamespace(n=3)], wo_b=SimpleNamespace(weight=w)) + xh = [torch.zeros((n, 5), dtype=torch.float16)] + base = M.wo_ab(lay, xh=xh, fold=False) + assert torch.equal(base, u_vals.float() @ w.float().t()) # without an adapter: as it was + tab = tables(K_IN, torch.bfloat16, False, 1) + slots = [0, 1, 2, 1] + got = M.wo_ab(lay, xh=xh, fold=False, adapter=runtime(3, wo_b=tab).bind(torch.tensor(slots))) + torch.testing.assert_close(got.double(), expected(tab, u_vals, slots, base), rtol=1e-5, atol=1e-5) + assert torch.equal(got[0], base[0]) # the stock row: the same bits + assert M.wo_ab(lay, xh=xh, fold=False, adapter=runtime(9, wo_b=tab).bind(torch.tensor(slots))).equal(base) # layer 9: none + + +# -- the MoE ------------------------------------------------------------------------------------------------------ +class FakeScratch: + def __init__(self, ex, rows, slots, prompt=False): + self.rows, self.slots, self.prompt = rows, slots, prompt + self.xd = torch.zeros((rows * slots, I), dtype=torch.float16) + + +def fake_experts(written): + """routed() leaves each (row, slot)'s down-projection input in the scratch, as the real one does.""" + + def routed(x, pick, wts, ex, s, out, n, limit=None, act_mode=None, add=None, before_down=None): + s.xd[:n * SLOTS] = torch.arange(n * SLOTS * I, dtype=torch.float16).view(n * SLOTS, I) + written.append(s.xd[:n * SLOTS].clone()) + base = torch.zeros((n, D)) if add is None else add.clone() + return base + + return SimpleNamespace(routed=routed, Scratch=FakeScratch, ACT_F32=1, EXACT_ROWS=64) + + +def fake_kernels(): + def route(logits, bias, topk, scale, shared_id, pick, wts): + pick[:, :topk] = torch.arange(topk, dtype=torch.int32) + pick[:, topk] = shared_id + wts.fill_(1.0) + + return SimpleNamespace(on=lambda name: False, DECODE_ROWS=16, rowmm=lambda x, w: x.float() @ w.float().t(), + route=route) + + +def moe_model(monkeypatch, written): + monkeypatch.setattr(M, "K", fake_kernels()) + monkeypatch.setattr(M, "exl3_experts", fake_experts(written)) + monkeypatch.setattr(M, "KERNELS", True) + monkeypatch.setattr(M, "EXACT_MM", False) + monkeypatch.setattr(M, "SHARED_OVERLAP", False) + m = M.Model.__new__(M.Model) + m.cfg, m.scratch = SimpleNamespace(topk=TOPK, route_scale=1.0, swiglu_limit=10.0), {} + return m + + +def packed_layer(idx): + return SimpleNamespace(idx=idx, shared=None, experts=SimpleNamespace(count=E), gate_w=torch.zeros((E, D)), + gate_b=torch.zeros((E,)), gate_b_vl=None) + + +def test_the_packed_shared_experts_delta_reads_its_slots_rotated_hidden_state_in_the_scratch(monkeypatch): + written = [] + m = moe_model(monkeypatch, written) + n = 5 + x = torch.randn((n, D), generator=G).to(torch.bfloat16) + lay = packed_layer(2) + base = m.moe(lay, x) + assert torch.equal(base, torch.zeros((n, D))) # no adapter: routed()'s result as is + tab = tables(I, torch.float16, True, 2) + slots = [2, 0, 1, 0, 2] + got = m.moe(lay, x, adapter=runtime(2, shared=tab).bind(torch.tensor(slots))) + shared_rows = written[-1].view(n, SLOTS, I)[:, SLOTS - 1] # the last slot of each row + torch.testing.assert_close(got.double(), expected(tab, shared_rows, slots, base), rtol=1e-5, atol=1e-3) + assert torch.equal(got[1], base[1]) and torch.equal(got[3], base[3]) + other = m.moe(packed_layer(7), x, adapter=runtime(2, shared=tab).bind(torch.tensor(slots))) + assert torch.equal(other, base) # a layer no adapter edits + + +def apart_layer(idx): + w = [torch.randn((I, D), generator=G).to(torch.bfloat16), torch.randn((I, D), generator=G).to(torch.bfloat16), + torch.randn((D, I), generator=G).to(torch.bfloat16)] + lay = packed_layer(idx) + lay.shared = tuple(SimpleNamespace(weight=t) for t in w) + lay.experts = SimpleNamespace(count=E - 1) + return lay, w + + +def test_a_shared_expert_apart_adds_the_delta_of_its_plain_hidden_state(monkeypatch): + written = [] + m = moe_model(monkeypatch, written) + n = 4 + x = torch.randn((n, D), generator=G).to(torch.bfloat16) + lay, (w1, w3, w2) = apart_layer(1) + base = m.moe(lay, x) + h = (F.silu((x.float() @ w1.float().t()).to(torch.bfloat16).float().clamp(max=10.0)) + * (x.float() @ w3.float().t()).to(torch.bfloat16).float().clamp(-10.0, 10.0)).to(torch.bfloat16) + torch.testing.assert_close(base, h.float() @ w2.float().t()) # (routed() here adds only ``add``) + tab = tables(I, torch.bfloat16, False, 3) + slots = [1, 2, 0, 2] + got = m.moe(lay, x, adapter=runtime(1, shared=tab).bind(torch.tensor(slots))) + torch.testing.assert_close(got.double(), expected(tab, h, slots, base), rtol=1e-4, atol=1e-3) + assert torch.equal(got[2], base[2]) diff --git a/tests/test_deepseek_v41_adapter_multi.py b/tests/test_deepseek_v41_adapter_multi.py new file mode 100644 index 0000000000..55b56fd529 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_multi.py @@ -0,0 +1,263 @@ +"""Kept and spilled prompts never meet another adapter's rows, and the lane's steps carry the adapter to every rank, +through the real MultiDecoder admission, finish and follow paths on CPU tensors (no GPU, kernels or checkpoint).""" + +import json +import sys +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("safetensors") + +from tests.test_deepseek_v41_spill import CHUNK, LAYERS, RAW, _fill, _pool, _prompt, _rows + +KEY_A, KEY_B = "a" * 32, "b" * 32 + + +@pytest.fixture +def mods(monkeypatch, tmp_path): + """multi.py and spill.py with the CUDA model and engine modules stubbed (PREFILL_CHUNK, REPLAY, RAW).""" + + monkeypatch.setenv("TF_DS_KEEP", "1") + model = ModuleType("tensorfold.families.deepseek_v41.cuda.model") + model.RAW = RAW + engine = ModuleType("tensorfold.families.deepseek_v41.cuda.engine") + engine.PREFILL_CHUNK, engine.REPLAY = CHUNK, False + monkeypatch.setitem(sys.modules, model.__name__, model) + monkeypatch.setitem(sys.modules, engine.__name__, engine) + for name in ("tensorfold.families.deepseek_v41.cuda.multi", "tensorfold.families.deepseek_v41.cuda.spill"): + monkeypatch.delitem(sys.modules, name, raising=False) + import tensorfold.families.deepseek_v41.cuda.multi as multi + import tensorfold.families.deepseek_v41.cuda.spill as spill + + return multi, spill, tmp_path / "spill" + + +def ref(multi, slot, key): + from tensorfold.families.deepseek_v41.cuda.lora_registry import AdapterRef + + return AdapterRef(slot, key) + + +def decoder(multi, spill, root, *, nonce="run1", held=("a", "b")): + """A MultiDecoder whose fake prefill makes rows that depend on the adapter slot, as a real adapter's do.""" + + pool, calls = _pool(), [] + + def pool_view(p, slot, base, size): + return SimpleNamespace(base=base, size=size, slot=slot, length=0, host=multi.HostIds()) + + def prefill_steps(sc, dc, prompt, image=None, start=0, snap=None, adapter=0): + sc.length = start + sc.host.truncate(start) + first = min(len(prompt), (start // CHUNK + 1) * CHUNK) + ends = [first] + list(range(first + CHUNK, len(prompt), CHUNK)) + if ends[-1] != len(prompt): + ends.append(len(prompt)) + calls.append((start, adapter)) + shifted = [t + 100000 * adapter for t in prompt] + for s, e in zip([start] + ends[:-1], ends): + _fill(pool, sc.base, sc.slot, shifted, s, e) + sc.host.set(s, prompt[s:e]) + if snap is not None and e % CHUNK == 0: + snap(e) + if e < len(prompt): + yield e + return "logits" + + holds = {(1, KEY_A) if "a" in held else None, (2, KEY_B) if "b" in held else None} + e = SimpleNamespace(rank=0, world=1, drafter=None, eos=(1,), replay_mode=False, prefill_steps=prefill_steps, + _sample=lambda last, pos, sampling: [7], + adapter_ok=lambda r: r is None or (r.slot, r.key) in holds) + md = object.__new__(multi.MultiDecoder) + md.e, md.pool, md.dpool, md.cap = e, pool, None, pool.cap + md.m = SimpleNamespace(cfg=SimpleNamespace(window=128, compress_ratios=[4, 128, 1, 1], n_layers=LAYERS), + pool_view=pool_view, comm=None) + md.slots = [multi.Slot(i, None) for i in range(pool.slots)] + md.extents, md.free, md.max_rows, md.eos = multi.Extents(pool.cap), list(range(pool.slots)), 6, (1,) + md.streams, md.filling, md.chunk_s, md.since_fill, md.next_id = {}, [], 0.0, 0.0, 0 + md.link = md.watch = md.round_end = None + md.kept, md.next_kept, md.ticks, md.keep_stats = {}, 0, 0, {} + md.spill = spill.Spill(root, engine="testengine", rank=0, pool=pool, ratios={0: 4, 1: 128}, budget=1 << 30, + max_age=3600.0, min_top=4096, device="cpu") + md.spill_nonce = nonce + return md, calls + + +def serve(multi, md, prompt, adapter=None): + s = multi.Stream(list(prompt), 1, None, draft=False, stop_eos=False, adapter=adapter) + s.emit = lambda new: None + md.admit(s) + while s in md.filling: + md._fill(s) + md.finish([s]) + return s + + +def test_a_prompt_continues_only_a_kept_prompt_of_the_same_adapter(mods): + multi, spill, root = mods + md, calls = decoder(multi, spill, root) + a, b = ref(multi, 1, KEY_A), ref(multi, 2, KEY_B) + p = _prompt(9000, 1) + serve(multi, md, p, a) + (k,) = md.kept.values() + assert k.adapter == KEY_A + longer = p + _prompt(2000, 2) + for other in (None, b): # the stock model and another adapter: a fresh fill + calls.clear() + s = serve(multi, md, longer, other) + assert s.cached == 0 and calls == [(0, 0 if other is None else 2)] + calls.clear() + s = serve(multi, md, longer, a) # its own adapter: continues + assert s.cached == 8192 and calls == [(8192, 1)] + + +def test_a_stock_prompt_is_never_continued_by_an_adapter_and_the_reverse(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root) + p = _prompt(9000, 3) + serve(multi, md, p, None) + assert [k.adapter for k in md.kept.values()] == [""] + assert serve(multi, md, p + [5, 6], ref(multi, 1, KEY_A)).cached == 0 + assert serve(multi, md, p + [5, 6, 7], None).cached == 8192 + + +def test_a_kept_prompt_is_not_covered_by_a_finishing_prompt_of_another_adapter(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root) + a, b = ref(multi, 1, KEY_A), ref(multi, 2, KEY_B) + p = _prompt(9000, 4) + serve(multi, md, p, a) + serve(multi, md, p, b) # the same ids, another adapter: both are kept + assert sorted(k.adapter for k in md.kept.values()) == [KEY_A, KEY_B] + serve(multi, md, p + [9] * 3000, a) # the same adapter's longer prompt covers its own + assert sorted(k.adapter for k in md.kept.values()) == [KEY_A, KEY_B] + assert len(md.kept) == 2 + + +def test_spilled_prompts_are_read_back_by_their_own_adapter_only(mods): + multi, spill, root = mods + a, b = ref(multi, 1, KEY_A), ref(multi, 2, KEY_B) + p = _prompt(20000, 6) + md1, _ = decoder(multi, spill, root) + serve(multi, md1, p, a) + (k,) = md1.kept.values() + assert k.disk is not None and k.disk["adapter"] == KEY_A + want = _rows(md1.pool, k.base, k.top) + for other in (None, b): # a new run: an empty window, the same disk + md, calls = decoder(multi, spill, root, nonce="run" + ("0" if other is None else "b")) + s = serve(multi, md, p + [1, 2, 3], other) + assert s.cached == 0 and md.keep_stats == {} and calls[0][0] == 0 + md, calls = decoder(multi, spill, root, nonce="runa") + s = multi.Stream(list(p) + [1, 2, 3], 1, None, draft=False, stop_eos=False, adapter=a) + s.emit = lambda new: None + md.admit(s) + assert s.cached == k.top and md.keep_stats == {"disk": 1} + assert all(torch.equal(x, y) for x, y in zip(want, _rows(md.pool, s.base, k.top))) + + +def test_a_stock_spilled_prompt_has_no_adapter_in_its_manifest_and_the_hashes_it_always_had(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root) + p = _prompt(20000, 7) + serve(multi, md, p, None) + (k,) = md.kept.values() + assert "adapter" not in k.disk + import hashlib + + keys = np.asarray(p, dtype=np.int64) + plain = hashlib.sha256(keys[:k.top].tobytes()).hexdigest() + assert k.disk["hashes"][str(k.top)] == plain + assert spill.Spill._hashes(keys, [k.top], KEY_A)[str(k.top)] != plain + + +def test_a_manifest_naming_an_adapter_in_a_form_that_is_no_key_is_dropped_at_start(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root) + serve(multi, md, _prompt(20000, 8), ref(multi, 1, KEY_A)) + (entry,) = md.spill.dir.glob("entry-*.json") + manifest = json.loads(entry.read_text()) + manifest["adapter"] = "../../etc" + entry.write_text(json.dumps(manifest)) + md2, _ = decoder(multi, spill, root, nonce="run2") + assert md2.spill.manifests == {} + + +def test_a_stock_admission_is_the_op_it_always_was_and_an_adapters_rides_after_the_grammar(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root) + sent = [] + md.link = SimpleNamespace(send=sent.append) + serve(multi, md, _prompt(100, 9), None) + serve(multi, md, _prompt(100, 9), ref(multi, 2, KEY_B)) + stock, held = [op for op in sent if op[0] == "admit"] + assert len(stock) == 10 and len(held) == 11 and held[10] == [2, KEY_B] and held[:10] == stock[:10] + + +def test_a_follower_builds_the_stream_with_the_adapter_the_op_names_and_fills_under_it(mods): + multi, spill, root = mods + md, calls = decoder(multi, spill, root) + p = _prompt(100, 10) + ops = [["admit", p, 1, None, False, False, 0, [], None, [], [2, KEY_B]], None] + link = SimpleNamespace(receive=lambda: ops.pop(0)) + with pytest.raises(RuntimeError, match="step link closed"): + md.follow(link) + (s,) = md.filling + assert (s.adapter.slot, s.adapter.key) == (2, KEY_B) + while s in md.filling: + md._fill(s) + assert calls == [(0, 2)] + + +def test_a_stock_follower_admission_still_parses_and_fills_stock(mods): + multi, spill, root = mods + md, calls = decoder(multi, spill, root) + ops = [["admit", _prompt(100, 11), 1, None, False, False, 0, [], None, []], None] + link = SimpleNamespace(receive=lambda: ops.pop(0)) + with pytest.raises(RuntimeError, match="step link closed"): + md.follow(link) + (s,) = md.filling + assert s.adapter is None and md.slots[0].sc is not None + while s in md.filling: + md._fill(s) + assert calls == [(0, 0)] + + +def test_the_plan_every_rank_agrees_on_names_the_adapter_and_whether_this_rank_holds_it(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root, held=("a",)) + stock = multi.Stream([1], 1, None) + held, lost = (multi.Stream([1], 1, None, adapter=ref(multi, s, k)) for s, k in ((1, KEY_A), (2, KEY_B))) + assert md._adapter_plan(stock) == [] and md._adapter_plan(stock, check=True) == [] + assert md._adapter_plan(held) == [[1, KEY_A]] + assert md._adapter_plan(held, check=True) == [[1, KEY_A, True]] + assert md._adapter_plan(lost, check=True) == [[2, KEY_B, False]] # a rank without it plans differently + + +def test_a_request_for_an_adapter_the_engine_does_not_hold_fails_before_anything_is_sent(mods): + multi, spill, root = mods + md, _ = decoder(multi, spill, root, held=("a",)) + sent = [] + md.link = SimpleNamespace(send=sent.append) + s = multi.Stream(_prompt(100, 12), 1, None, adapter=ref(multi, 2, KEY_B)) + s.emit = lambda new: None + with pytest.raises(ValueError, match="does not hold"): + md.admit(s) + assert sent == [] and md.filling == [] and md.free == list(range(len(md.slots))) + + +def test_the_shape_each_round_checks_names_adapter_streams_and_is_unchanged_for_stock(mods, monkeypatch): + multi, spill, root = mods + monkeypatch.setattr(multi, "ADMIT_GIB", 0.0) # (no memory gate: two streams are live) + md, _ = decoder(multi, spill, root) + stock = multi.Stream(_prompt(100, 13), 1, None) + stock.emit = lambda new: None + md.admit(stock) + entry = md._shape()[4][0] + assert entry == [stock.sid, 0, 0] # a stock stream's entry has no slot in it + held = multi.Stream(_prompt(100, 13), 1, None, adapter=ref(multi, 1, KEY_A)) + held.emit = lambda new: None + md.admit(held) + assert md._shape()[4] == [entry, [held.sid, 1, 0, 1]] diff --git a/tests/test_deepseek_v41_adapter_place.py b/tests/test_deepseek_v41_adapter_place.py new file mode 100644 index 0000000000..a1b96bbc75 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_place.py @@ -0,0 +1,114 @@ +"""Where an adapter's factors live on a rank: the cut, the sum over ranks, and the shared expert's rotated hidden state +(numpy only: the algebra the kernels are told to run).""" + +import numpy as np +import pytest +from adapter_fakes import CFG + +from tensorfold.families.deepseek_v41.cuda import lora_format as lf +from tensorfold.families.deepseek_v41.cuda import lora_place as lp + +DIMS = lf.Dims.of(CFG) + + +def make(kind, rank, seed=0): + rng = np.random.default_rng(seed) + width, out = DIMS.shape(kind) + return lf.Factors((rng.standard_normal((rank, width)) * 0.1).astype(np.float32), + (rng.standard_normal((out, rank)) * 0.1).astype(np.float32), 0.75) + + +def test_the_hadamard_matrix_is_symmetric_orthonormal_and_in_natural_order(): + h = lp.hadamard() + assert h.shape == (128, 128) + np.testing.assert_allclose(h, h.T, atol=0) + np.testing.assert_allclose(h @ h, np.eye(128), atol=1e-12) + assert h[0, 0] > 0 and np.all(h[0] == h[0, 0]) and np.all(h[:, 0] == h[0, 0]) # the all-ones row + np.testing.assert_allclose(h[1], h[0] * np.where(np.arange(128) & 1, -1, 1)) # Sylvester: bit 0 alternates + + +def test_a_rank_holds_its_own_whole_blocks_of_the_input(): + assert lp.span(512, 0, 2) == (0, 256) and lp.span(512, 1, 2) == (256, 512) and lp.span(512, 0, 1) == (0, 512) + for total, rank, world in ((512, 2, 2), (512, 0, 3), (192, 0, 2), (512, -1, 2), (512, 0, 0)): + with pytest.raises(ValueError): + lp.span(total, rank, world) + + +@pytest.mark.parametrize("kind", [lf.WO_B, lf.SHARED_W2]) +@pytest.mark.parametrize("world", [1, 2]) +def test_the_ranks_partial_deltas_add_up_to_the_whole_matrices_delta(kind, world): + f = make(kind, 3) + width, out = DIMS.shape(kind) + x = np.random.default_rng(1).standard_normal((5, width)).astype(np.float32) + whole = 0.75 * (x @ f.a.T) @ f.b.T + parts = 0 + for rank in range(world): + p = lp.place(kind, f, DIMS, rank, world) + lo, hi = lp.span(width, rank, world) + assert p.a.shape == (3, width // world) and p.b.shape == (out, 3) and not p.folded and p.scale == 0.75 + parts = parts + p.scale * (x[:, lo:hi] @ p.a.T) @ p.b.T + np.testing.assert_allclose(parts, whole, rtol=1e-5, atol=1e-5) + + +def test_strength_multiplies_the_scale_and_nothing_else(): + f = make(lf.WO_B, 2) + p, q = lp.place(lf.WO_B, f, DIMS, 0, 2), lp.place(lf.WO_B, f, DIMS, 0, 2, strength=-2.0) + assert q.scale == -2.0 * p.scale + np.testing.assert_array_equal(p.a, q.a) + np.testing.assert_array_equal(p.b, q.b) + + +def rotated_hidden(act, suh): + """What the expert kernels leave for the down projection: blockwise Hadamard of (act * suh), in fp16.""" + + h = lp.hadamard() + z = (act.astype(np.float64) * suh).reshape(act.shape[0], -1, 128) @ h + return z.reshape(act.shape).astype(np.float16) + + +@pytest.mark.parametrize("world", [1, 2]) +def test_the_folded_factors_read_the_rotated_hidden_state_as_the_plain_ones_read_the_hidden_state(world): + f = make(lf.SHARED_W2, 4, seed=5) + rng = np.random.default_rng(6) + act = rng.standard_normal((6, DIMS.inter)).astype(np.float32) + suh = np.where(rng.random(DIMS.inter) < 0.5, -1.0, 1.0) * rng.uniform(0.5, 2.0, DIMS.inter) + suh = suh.astype(np.float16).astype(np.float64) + want = 0.75 * (act @ f.a.T) @ f.b.T + got = 0 + for rank in range(world): + lo, hi = lp.span(DIMS.inter, rank, world) + xd = rotated_hidden(act[:, lo:hi], suh[lo:hi]) # fp16, as the kernels store it + p = lp.place(lf.SHARED_W2, f, DIMS, rank, world, suh=suh[lo:hi]) + assert p.folded and p.a.dtype == np.float32 + got = got + p.scale * (xd.astype(np.float32) @ p.a.T) @ p.b.T + scale = np.abs(want).max() + assert np.abs(got - want).max() <= 3e-3 * scale # fp16 storage of xd is the only difference + + +def test_folding_is_exact_without_the_fp16_rounding(): + rng = np.random.default_rng(7) + a = rng.standard_normal((3, 256)).astype(np.float32) + suh = rng.uniform(0.5, 2.0, 256) * np.where(rng.random(256) < 0.5, -1, 1) + act = rng.standard_normal((4, 256)) + xd = (act * suh).reshape(4, 2, 128) @ lp.hadamard() + np.testing.assert_allclose(xd.reshape(4, 256) @ lp.fold_rotated(a, suh).T.astype(np.float64), act @ a.T, + rtol=2e-5, atol=2e-5) + + +def test_folding_refuses_a_scale_vector_it_cannot_divide_by(): + a = np.ones((2, 128), dtype=np.float32) + for bad in (np.zeros(128), np.full(128, np.inf), np.full(128, np.nan), np.ones(64)): + with pytest.raises(ValueError): + lp.fold_rotated(a, bad) + with pytest.raises(ValueError): + lp.fold_rotated(np.ones((2, 100), dtype=np.float32), np.ones(100)) + + +def test_only_the_shared_expert_reads_a_rotated_input_and_the_factors_must_fit(): + f = make(lf.WO_B, 2) + with pytest.raises(ValueError, match="only the packed shared expert"): + lp.place(lf.WO_B, f, DIMS, 0, 2, suh=np.ones(256)) + with pytest.raises(ValueError, match="do not fit"): + lp.place(lf.SHARED_W2, f, DIMS, 0, 2) + with pytest.raises(ValueError, match="no placement"): + lp.place("wq_b", f, DIMS, 0, 2) diff --git a/tests/test_deepseek_v41_adapter_registry.py b/tests/test_deepseek_v41_adapter_registry.py new file mode 100644 index 0000000000..efe57aad28 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_registry.py @@ -0,0 +1,128 @@ +"""The adapter registry: identities, slots, what a request may select, and what two ranks must agree on (no torch).""" + +import shutil +from types import SimpleNamespace + +import pytest +from adapter_fakes import CFG, model_dir, standard + +from tensorfold.families.deepseek_v41.cuda import lora_format as lf +from tensorfold.families.deepseek_v41.cuda import lora_registry as lr + + +@pytest.fixture +def base(tmp_path): + return model_dir(tmp_path / "model") + + +def load(base, specs, reserved=("base",)): + return lr.AdapterRegistry.load(specs, CFG, base, reserved) + + +def test_slots_follow_the_order_given_and_zero_is_the_stock_model(tmp_path, base): + reg = load(base, [("one", str(standard(tmp_path / "a", 1)), 1.0), ("two", str(standard(tmp_path / "b", 2)), 1.0)]) + assert reg.names() == ["one", "two"] and reg.slots == 3 and len(reg) == 2 + assert [e.slot for e in reg.entries] == [1, 2] and reg.entry(2).name == "two" + assert reg.model_map() == {"one": "one", "two": "two"} + assert reg.ref(None) is None and reg.ref("two") == lr.AdapterRef(2, reg.entry(2).key) + + +def test_an_unknown_adapter_is_an_error_not_the_stock_model(tmp_path, base): + reg = load(base, [("one", str(standard(tmp_path / "a")), 1.0)]) + with pytest.raises(lr.UnknownAdapter, match="no adapter named 'nope'"): + reg.ref("nope") + + +def test_the_key_names_the_factors_and_their_scales_not_the_label(tmp_path, base): + a = str(standard(tmp_path / "a", 1)) + shutil.copytree(a, tmp_path / "same") + keys = [load(base, [(n, f, s)]).entry(1).key for n, f, s in + [("x", a, 1.0), ("y", str(tmp_path / "same"), 1.0), ("x", a, 2.0), ("x", str(standard(tmp_path / "c", 3)), 1.0)]] + assert keys[0] == keys[1] # the same files at the same strength: the same rows + assert len({keys[0], keys[2], keys[3]}) == 3 # another strength or other factors: other rows + assert all(len(k) == 32 and set(k) <= set("0123456789abcdef") for k in keys) + + +def test_one_changed_byte_of_the_weights_changes_the_key(tmp_path, base): + a = standard(tmp_path / "a", 1) + before = load(base, [("x", str(a), 1.0)]).entry(1).key + path = a / lf.WEIGHTS + raw = bytearray(path.read_bytes()) + raw[-1] ^= 1 + path.write_bytes(bytes(raw)) + assert load(base, [("x", str(a), 1.0)]).entry(1).key != before + + +def test_the_digest_two_ranks_compare_follows_every_slot_name_and_key(tmp_path, base): + a, b = str(standard(tmp_path / "a", 1)), str(standard(tmp_path / "b", 2)) + ref = load(base, [("one", a, 1.0), ("two", b, 1.0)]) + assert load(base, [("one", a, 1.0), ("two", b, 1.0)]).digest() == ref.digest() + assert load(base, [("two", b, 1.0), ("one", a, 1.0)]).digest() != ref.digest() # slots swapped + assert load(base, [("one", a, 1.0), ("two", b, 0.5)]).digest() != ref.digest() # a strength + assert load(base, [("one", a, 1.0)]).digest() != ref.digest() + other = model_dir(tmp_path / "other") + (other / "config.json").write_text('{"model_type": "deepseek_v41", "hidden_size": 256, "x": 1}') + assert load(other, [("one", a, 1.0), ("two", b, 1.0)]).digest() != ref.digest() # another base config + assert len(ref.words()) == 4 and all(0 <= w < 2**31 for w in ref.words()) + assert ref.words() == load(base, [("one", a, 1.0), ("two", b, 1.0)]).words() + + +def test_a_reference_is_good_only_where_the_same_adapter_holds_its_slot(tmp_path, base): + a, b = str(standard(tmp_path / "a", 1)), str(standard(tmp_path / "b", 2)) + reg = load(base, [("one", a, 1.0), ("two", b, 1.0)]) + swapped = load(base, [("two", b, 1.0), ("one", a, 1.0)]) + assert reg.check(None) and reg.check(reg.ref("one")) + assert not reg.check(swapped.ref("one")) # the same name, another slot and key + assert not reg.check(lr.AdapterRef(3, "0" * 32)) and not reg.check(lr.AdapterRef(0, "")) + + +def pick_registry(tmp_path, base): + reg = load(base, [("one", str(standard(tmp_path / "a", 1)), 1.0), ("two", str(standard(tmp_path / "b", 2)), 1.0)]) + return reg, reg.model_map() + + +def test_a_request_picks_its_adapter_by_model_id_or_field(tmp_path, base): + reg, models = pick_registry(tmp_path, base) + assert reg.pick({"model": "base"}, models) is None and reg.pick({}, models) is None + assert reg.pick({"model": "one"}, models) == "one" + assert reg.pick({"model": "base", "adapter": "two"}, models) == "two" + assert reg.pick({"model": "two", "adapter": "two"}, models) == "two" + assert reg.pick({"model": "no such model"}, models) is None # an unknown model id was always the default + + +@pytest.mark.parametrize("body", [{"adapter": "nope"}, {"adapter": 3}, {"adapter": ""}, {"model": "one", "adapter": "two"}]) +def test_an_unknown_or_contradictory_choice_is_refused(tmp_path, base, body): + reg, models = pick_registry(tmp_path, base) + with pytest.raises(lr.UnknownAdapter): + reg.pick(body, models) + + +def test_names_and_strengths_that_cannot_be_served_are_refused(tmp_path, base): + a = str(standard(tmp_path / "a")) + for bad in ("", "-x", "a b", "a=b", "x" * 97): + with pytest.raises(lf.AdapterError, match="letters, digits"): + load(base, [(bad, a, 1.0)]) + with pytest.raises(lf.AdapterError, match="already a model id"): + load(base, [("base", a, 1.0)]) + with pytest.raises(lf.AdapterError, match="already a model id"): + load(base, [("one", a, 1.0), ("one", a, 1.0)]) + for bad in (0.0, float("nan"), float("inf")): + with pytest.raises(lf.AdapterError, match="strength"): + load(base, [("one", a, bad)]) + assert load(base, [("neg", a, -1.0)]).entry(1).strength == -1.0 # an inverted edit is a valid strength + + +def test_a_stream_or_a_reference_gives_its_slot_and_key(): + ref = lr.AdapterRef(2, "ab" * 16) + assert (lr.slot_of(ref), lr.key_of(ref)) == (2, "ab" * 16) + assert (lr.slot_of(None), lr.key_of(None)) == (0, "") + stream = SimpleNamespace(adapter=ref) + assert (lr.slot_of(stream), lr.key_of(stream)) == (2, "ab" * 16) + assert (lr.slot_of(SimpleNamespace(adapter=None)), lr.key_of(SimpleNamespace())) == (0, "") + + +def test_the_description_names_what_each_adapter_edits(tmp_path, base): + reg = load(base, [("one", str(standard(tmp_path / "a", ranks=(2, 3))), 1.0)]) + d = reg.describe()[0] + assert d["slot"] == 1 and d["modules"] == {"wo_b": 1, "shared_w2": 1} and d["max_rank"] == 3 + assert d["weights_sha256"] == reg.entry(1).spec.sha256 diff --git a/tests/test_deepseek_v41_adapter_rounds.py b/tests/test_deepseek_v41_adapter_rounds.py new file mode 100644 index 0000000000..1dff0b33b8 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_rounds.py @@ -0,0 +1,100 @@ +"""A concurrent round's inputs carry each row's adapter slot: in window order, padded with copies of the last row, and +absent from a stock round (the round runner on CPU with fake graphs, no GPU).""" + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("triton") + +from tensorfold.families.deepseek_v41.cuda import graph_keys, rounds + +SIZES = graph_keys.row_sizes(16, graph_keys.ROUND_ROWS) +LONG = graph_keys.row_sizes(16, graph_keys.ROUND_ROWS_LONG) +CAP = 1 << 20 + + +class FakeGraph: + taps = None + + def __init__(self, rows): + self.rows, self.runs, self.sets, self.captured = rows, [], [], False + + def set(self, *cols): + self.sets.append(cols) + + def capture(self, pool): + self.captured = True + + def run(self, *args): + self.runs.append(args) + return torch.zeros((self.rows, 4)) + + +ADAPTERS = object() # (an engine that serves adapters: only whether lora is set matters here) + + +def runner(lora=ADAPTERS, graphs=None): + r = rounds.RoundRunner.__new__(rounds.RoundRunner) + r.m, r.pool = SimpleNamespace(engram=None, lora=lora), SimpleNamespace(cap=CAP) + r.graphs, r.graph_pool, r.captures, r.warmed, r.out = graphs if graphs is not None else {}, object(), 0, True, None + r.sizes, r.sizes_long = SIZES, LONG + return r + + +def window(slot, p0, n, *adapter): + return (slot, slot * 4096, 4096, p0, list(range(n)), None) + adapter + + +def test_the_columns_follow_the_windows_and_the_padding_copies_the_last_rows_adapter(): + g = FakeGraph(6) + r = runner(graphs={(6, 1024): g}) + logits, taps = r.forward([window(0, 10, 3, 0), window(1, 20, 2, 2)]) # 5 rows run in the 6-row graph + assert logits.shape == (5, 4) and taps is None + ids, pos, slots, base, end, _, _, adapters = g.runs[0] + assert adapters == [0, 0, 0, 2, 2, 2] and slots == [0, 0, 0, 1, 1, 1] + assert pos == [10, 11, 12, 20, 21, 21] and ids == [0, 1, 2, 0, 1, 1] # (the pad row is the last row again) + assert base == [0, 0, 0, 4096, 4096, 4096] and end == [4096, 4096, 4096, 8192, 8192, 8192] + + +def test_a_window_without_an_adapter_is_slot_zero_and_a_stock_round_has_only_zeros(): + g = FakeGraph(4) + r = runner(graphs={(4, 1024): g}) + r.forward([window(0, 5, 2), window(1, 9, 2)]) + assert g.runs[0][-1] == [0, 0, 0, 0] + g2 = FakeGraph(4) + runner(lora=None, graphs={(4, 1024): g2}).forward([window(0, 5, 2), window(1, 9, 2)]) + assert g2.runs[0][-1] == [0, 0, 0, 0] # an engine without adapters + + +def test_a_row_of_an_adapter_in_an_engine_without_adapters_is_refused_not_run_as_stock(): + r = runner(lora=None, graphs={(2, 1024): FakeGraph(2)}) + with pytest.raises(ValueError, match="started without one"): + r.forward([window(0, 5, 2, 1)]) + + +def test_a_round_that_captures_a_graph_sets_it_up_with_the_same_columns(monkeypatch): + made = [] + monkeypatch.setattr(rounds, "RoundDecoder", lambda m, pool, rows, bucket, taps, out=None: made.append( + FakeGraph(rows)) or made[-1]) + r = runner() # (warmed: a capture now is counted) + r.forward([window(0, 5, 1, 1), window(1, 9, 3, 2)]) # 4 rows: a 4-row graph + (g,) = made + assert g.captured and g.rows == 4 and g.sets[0][-1] == [1, 2, 2, 2] and g.runs[0][-1] == [1, 2, 2, 2] + assert r.captures == 1 and (4, 1024) in r.graphs + + +def test_without_graphs_the_round_runs_unpadded(): + g = FakeGraph(5) + r = runner() + r.graphs = None + r.sizes = SIZES + made = [] + r.forward = rounds.RoundRunner.forward.__get__(r) + rounds.RoundDecoder, saved = (lambda m, pool, rows, bucket, taps, out=None: made.append(g) or g), rounds.RoundDecoder + try: + r.forward([window(0, 5, 2, 1), window(1, 9, 3, 2)]) + finally: + rounds.RoundDecoder = saved + assert g.runs[0][-1] == [1, 1, 2, 2, 2] diff --git a/tests/test_deepseek_v41_adapter_runtime.py b/tests/test_deepseek_v41_adapter_runtime.py new file mode 100644 index 0000000000..aaf43aa7f8 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_runtime.py @@ -0,0 +1,239 @@ +"""Adapter factors as tables on a rank, and the hooks the model calls, on CPU tensors: each rank's share, the row-wise +sum over slots, the shared expert's two input forms, and that a stock row keeps its bits.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +from adapter_fakes import CFG, SHARED, WO_B, factors, model_dir, standard, write_adapter + +torch = pytest.importorskip("torch") + +from tensorfold.families.deepseek_v41.cuda import lora_format as lf +from tensorfold.families.deepseek_v41.cuda import lora_place as lp +from tensorfold.families.deepseek_v41.cuda import lora_reference as ref +from tensorfold.families.deepseek_v41.cuda import lora_registry as lr +from tensorfold.families.deepseek_v41.cuda import lora_runtime as rt +from tensorfold.families.deepseek_v41.cuda.lora_tables import TILE_J, Tables, padded_rank + +DIMS = lf.Dims.of(CFG) +WO_IN = DIMS.o_in + + +def kernel(out, x, ids, tab): + return ref.lowrank_add(out, x, ids, tab.a, tab.b, tab.ranks, tab.scales, dtype=torch.float64) + + +def weights(packed=(2,), world=2, seed=0): + """Layers whose shared expert is packed with the routed ones (has suh_d) or apart (``shared`` set).""" + + g = torch.Generator().manual_seed(seed) + layers = [] + for i in range(CFG.n_layers): + suh = (torch.rand((5, CFG.inter // world), generator=g) + 0.5) * torch.where( + torch.rand((5, CFG.inter // world), generator=g) < 0.5, -1.0, 1.0) + layers.append(SimpleNamespace(idx=i, shared=None if i in packed else ("w1", "w3", "w2"), + experts=SimpleNamespace(count=5, suh_d=suh.to(torch.float16)))) + return SimpleNamespace(cfg=CFG, layers=layers) + + +@pytest.fixture +def registry(tmp_path): + base = model_dir(tmp_path / "model") + a = standard(tmp_path / "a", seed=1, ranks=(2, 3)) + rng = np.random.default_rng(2) + b = write_adapter(tmp_path / "b", {WO_B.format(1): factors(rng, CFG.dim, WO_IN, 4), + WO_B.format(3): factors(rng, CFG.dim, WO_IN, 1), + SHARED.format(2): factors(rng, CFG.dim, CFG.inter, 2)}) + return lr.AdapterRegistry.load([("a", str(a), 1.0), ("b", str(b), 0.5)], CFG, base, ("model",)) + + +def load(registry, rank=0, world=2, packed=(2,), **kw): + return rt.LoraRuntime.load(registry, weights(packed, world), rank, world, device="cpu", log=lambda s: None, + kernel=kernel, **kw) + + +def file_factors(registry, slot, name): + e = registry.entry(slot) + m = next(m for m in e.spec.modules if m.name == name) + return lf.read_factors(e.spec, m) + + +def test_tables_hold_ranks_in_fours_and_the_kernels_tile_is_one_constant(): + assert [padded_rank(r) for r in (0, 1, 2, 3, 4, 5, 16, 17, 33, 256)] == [4, 4, 4, 4, 4, 8, 16, 20, 36, 256] + assert TILE_J == 16 and Tables.empty(2, 3, 128, 128, False, torch.float16, "cpu").tile_j == TILE_J + assert Tables.empty(2, 200, 128, 128, False, torch.float16, "cpu").tile_j == TILE_J # (not a function of the rank) + + +def test_tables_hold_exactly_what_was_put_and_zero_elsewhere(): + t = Tables.empty(3, 5, 256, 128, False, torch.bfloat16, "cpu") + assert Tables.required_bytes(3, 5, 256, 128) == t.nbytes + assert (t.slots, t.R, t.K, t.D, t.tile_j) == (3, 8, 256, 128, 16) and int(t.ranks.sum()) == 0 + a, b = np.ones((5, 256), np.float32), np.full((128, 5), 2.0, np.float32) + t.put(2, a, b, 0.5) + assert t.ranks.tolist() == [0, 0, 5] and t.scales.tolist() == [0.0, 0.0, 0.5] + assert float(t.a[2, :5].sum()) == 5 * 256 and float(t.a[2, 5:].abs().sum()) == 0 and float(t.a[0].abs().sum()) == 0 + for bad in (0, 3, -1): + with pytest.raises(ValueError): + t.put(bad, a, b, 1.0) + with pytest.raises(ValueError): + t.put(1, np.ones((9, 256), np.float32), np.ones((128, 9), np.float32), 1.0) # past the tables' rank + with pytest.raises(ValueError): + t.put(1, np.ones((5, 128), np.float32), b, 1.0) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_each_rank_holds_its_cut_of_every_slot_and_the_stock_slot_holds_nothing(registry, rank): + r = load(registry, rank) + assert sorted(r.layers) == [1, 2, 3] + wo1 = r.layers[1].wo_b + assert wo1.slots == 3 and wo1.R == 4 and wo1.K == WO_IN // 2 and wo1.D == CFG.dim and not wo1.folded + assert wo1.x_dtype == torch.bfloat16 and wo1.ranks.tolist() == [0, 2, 4] + lo, hi = lp.span(WO_IN, rank, 2) + fa, fb = file_factors(registry, 1, "layers.1.attn.wo_b"), file_factors(registry, 2, "layers.1.attn.wo_b") + np.testing.assert_array_equal(wo1.a[1, :2].numpy(), fa.a[:, lo:hi]) + np.testing.assert_array_equal(wo1.b[1, :, :2].numpy(), fa.b) + np.testing.assert_array_equal(wo1.a[2, :4].numpy(), fb.a[:, lo:hi]) + assert float(wo1.a[0].abs().sum()) == 0 and float(wo1.b[0].abs().sum()) == 0 + assert wo1.scales.tolist() == [0.0, fa.scale, fb.scale * 0.5] # strength 0.5 on the second + assert r.layers[3].wo_b.ranks.tolist() == [0, 0, 1] and r.layers[3].shared is None # only adapter b edits layer 3 + assert r.layers[1].shared is None and r.nbytes > 0 + + +def test_the_packed_shared_expert_gets_folded_factors_and_a_plain_one_does_not(registry): + packed, apart = load(registry, 1, packed=(2,)), load(registry, 1, packed=()) + w = weights((2,), 2) + tab = packed.layers[2].shared + assert tab.folded and tab.x_dtype == torch.float16 + f = file_factors(registry, 1, "layers.2.ffn.shared_experts.w2") + lo, hi = lp.span(CFG.inter, 1, 2) + want = lp.fold_rotated(f.a[:, lo:hi], w.layers[2].experts.suh_d[4].float().numpy()) # the last expert: shared + np.testing.assert_array_equal(tab.a[1, :3].numpy(), want) + plain = apart.layers[2].shared + assert not plain.folded and plain.x_dtype == torch.bfloat16 + np.testing.assert_array_equal(plain.a[1, :3].numpy(), f.a[:, lo:hi]) + + +def test_the_hooks_add_each_rows_own_delta_and_a_stock_row_keeps_its_bits(registry): + r = load(registry, 0) + g = torch.Generator().manual_seed(3) + lay = SimpleNamespace(idx=1) + u = torch.randn((5, WO_IN // 2), generator=g).to(torch.bfloat16) + out = torch.randn((5, CFG.dim), generator=g) + rows = torch.tensor([0, 1, 2, 0, 1]) + got = r.bind(rows).wo_b(lay, u, out.clone()) + tab = r.layers[1].wo_b + for i, slot in enumerate(rows.tolist()): + if slot == 0: + assert torch.equal(got[i], out[i]) # not touched: the same bits + continue + k = int(tab.ranks[slot]) + t = tab.a[slot, :k].double() @ u[i].double() + want = out[i].double() + float(tab.scales[slot]) * (tab.b[slot, :, :k].double() @ t) + torch.testing.assert_close(got[i].double(), want, rtol=1e-6, atol=1e-6) + assert not torch.equal(got[1], out[1]) and not torch.equal(got[2], out[2]) + + +def test_the_hooks_run_nothing_for_a_layer_no_adapter_edits(registry): + r = load(registry, 0) + seen = [] + r.kernel = lambda *a: seen.append(a) + out = torch.zeros((2, CFG.dim)) + rows = r.bind(torch.tensor([1, 2])) + assert rows.wo_b(SimpleNamespace(idx=0), torch.zeros((2, WO_IN // 2), dtype=torch.bfloat16), out) is out + assert rows.shared(SimpleNamespace(idx=0), torch.zeros((2, CFG.inter // 2)), out, packed=True) is out + assert rows.wo_b(SimpleNamespace(idx=2), torch.zeros((2, WO_IN // 2), dtype=torch.bfloat16), out) is out # no wo_b + assert seen == [] + + +def test_the_two_ranks_partials_add_up_to_the_merged_weight_result(registry): + ranks = [load(registry, k) for k in (0, 1)] + g = torch.Generator().manual_seed(4) + lay = SimpleNamespace(idx=1) + u = torch.randn((4, WO_IN), generator=g).to(torch.bfloat16) + slots = torch.tensor([1, 2, 0, 2]) + total = torch.zeros((4, CFG.dim), dtype=torch.float64) + for k, r in enumerate(ranks): + half = u[:, k * WO_IN // 2:(k + 1) * WO_IN // 2].contiguous() + total += r.bind(slots).wo_b(lay, half, torch.zeros((4, CFG.dim))).double() + for i, slot in enumerate(slots.tolist()): + if slot == 0: + assert float(total[i].abs().sum()) == 0 + continue + f = file_factors(registry, slot, "layers.1.attn.wo_b") + merged = (f.b.astype(np.float64) @ f.a.astype(np.float64)) * f.scale * registry.entry(slot).strength + np.testing.assert_allclose(total[i].numpy(), merged @ u[i].double().numpy(), rtol=1e-5, atol=1e-6) + + +def rotated(act, suh): + z = (act.double() * suh.double()).reshape(act.shape[0], -1, 128) @ torch.from_numpy(lp.hadamard()) + return z.reshape(act.shape).to(torch.float16) + + +def test_both_shared_expert_forms_give_the_same_delta(registry): + packed, apart = load(registry, 0, packed=(2,)), load(registry, 0, packed=()) + suh = weights((2,), 2).layers[2].experts.suh_d[4] + g = torch.Generator().manual_seed(5) + act = torch.randn((3, CFG.inter // 2), generator=g).to(torch.bfloat16) + slots = torch.tensor([2, 0, 1]) + lay = SimpleNamespace(idx=2) + plain = apart.bind(slots).shared(lay, act, torch.zeros((3, CFG.dim)), packed=False) + folded = packed.bind(slots).shared(lay, rotated(act.float(), suh.float()), torch.zeros((3, CFG.dim)), packed=True) + assert float(folded[1].abs().sum()) == 0 and float(plain[1].abs().sum()) == 0 + scale = plain.abs().max() + assert (folded - plain).abs().max() <= 4e-3 * scale # fp16 storage of the rotated hidden state + assert float(plain[0].abs().sum()) > 0 and float(plain[2].abs().sum()) > 0 + + +def test_a_shared_expert_read_in_the_wrong_form_is_an_error_not_a_wrong_delta(registry): + r = load(registry, 0, packed=(2,)) + with pytest.raises(RuntimeError, match="placed for the rotated"): + r.bind(torch.tensor([1])).shared(SimpleNamespace(idx=2), torch.zeros((1, CFG.inter // 2)), + torch.zeros((1, CFG.dim)), packed=False) + + +def test_modules_whose_factors_are_all_zero_are_left_out(tmp_path): + base = model_dir(tmp_path / "model") + rng = np.random.default_rng(6) + a, b = factors(rng, CFG.dim, WO_IN, 2) + folder = write_adapter(tmp_path / "z", {WO_B.format(0): (a, 0 * b), WO_B.format(2): (a, b)}) + reg = lr.AdapterRegistry.load([("z", str(folder), 1.0)], CFG, base) + logged = [] + r = rt.LoraRuntime.load(reg, weights(), 0, 2, device="cpu", log=logged.append, kernel=kernel) + assert sorted(r.layers) == [2] and "1 all-zero modules left out" in logged[0] + + +def test_tables_that_would_not_fit_the_host_are_refused_before_they_are_allocated(registry, monkeypatch): + with monkeypatch.context() as patch: + def unexpected_allocation(*args, **kwargs): + pytest.fail("factor tables allocated before the memory check") + + patch.setattr(Tables, "empty", unexpected_allocation) + with pytest.raises(MemoryError, match="the host has 0.00 GiB"): + load(registry, 0, room=lambda: 0.0) + assert load(registry, 0, room=lambda: 64.0).nbytes > 0 + assert load(registry, 0, room=lambda: None).nbytes > 0 # unknown: no check + + +def test_an_engine_with_fewer_layers_than_the_adapter_edits_is_refused(registry): + w = weights() + w.layers = w.layers[:2] + with pytest.raises(lf.AdapterError, match="holds 2 layers"): + rt.LoraRuntime.load(registry, w, 0, 2, device="cpu", log=lambda s: None, kernel=kernel) + + +def test_uniform_rows_are_one_slot_and_stock_prompts_run_no_hook(registry): + r = load(registry, 0) + assert r.uniform(0, 7) is None + rows = r.uniform(2, 5) + assert rows.rows.dtype == torch.int64 and rows.rows.tolist() == [2] * 5 and rows.rt is r + + +def test_warm_runs_each_table_shape_once_on_zero_rows(registry): + r = load(registry, 0) + seen = [] + r.kernel = lambda out, x, ids, tab: seen.append((tuple(x.shape), x.dtype, tab.R, int(ids.sum()), float(out.sum()))) + n = r.warm() + assert n == len(seen) == 2 # layers 1 and 3's wo_b share a shape (R=4); layer 2's shared w2 is packed (fp16) + assert {(s[0][1], s[1], s[2]) for s in seen} == {(WO_IN // 2, torch.bfloat16, 4), (CFG.inter // 2, torch.float16, 4)} + assert all(s[0][0] == 1 and s[3] == 0 and s[4] == 0 for s in seen) diff --git a/tests/test_deepseek_v41_adapter_server.py b/tests/test_deepseek_v41_adapter_server.py new file mode 100644 index 0000000000..e6c4d1f7e6 --- /dev/null +++ b/tests/test_deepseek_v41_adapter_server.py @@ -0,0 +1,211 @@ +"""Adapter selection on the CUDA server: model ids, the adapter field, refusals before a stream opens, and a server +without adapters answering exactly as it did (fake engine, no GPU).""" + +import json + +import pytest + +pytest.importorskip("jinja2") + +from adapter_fakes import CFG, model_dir, standard + +from tensorfold import cli, serve_options +from tensorfold.cuda import server +from tensorfold.families.deepseek_v41.cuda import lora_registry as lr +from tests.test_cuda_admission import http_server +from tests.test_cuda_server_errors import HI, Engine, app_for, request + +NOT_PASSED = "" + + +class AdapterEngine(Engine): + """The fake engine, taking an adapter name as the DeepSeek-V4.1 engine does.""" + + def __init__(self, registry=None): + super().__init__() + self.adapters = registry + self.seen = [] # the adapter argument of each call + + def generate(self, prompt, max_tokens, sampling, on_tokens, draft=True, adapter=NOT_PASSED): + self.seen.append(adapter) + return super().generate(prompt, max_tokens, sampling, on_tokens, draft) + + +def served(tmp_path): + base = model_dir(tmp_path / "model") + reg = lr.AdapterRegistry.load([("alpha", str(standard(tmp_path / "a", 1)), 1.0), + ("beta", str(standard(tmp_path / "b", 2)), 1.0)], CFG, base, ("fake-cuda",)) + (tmp_path / "app").mkdir() + app = app_for(tmp_path / "app") + app.engine = AdapterEngine(reg) + app.adapter_models = reg.model_map() + return app + + +def models(port): + import http.client + + connection = http.client.HTTPConnection("127.0.0.1", port, timeout=10) + try: + connection.request("GET", "/v1/models") + return json.loads(connection.getresponse().read())["data"] + finally: + connection.close() + + +def test_models_lists_each_adapter_beside_the_base_model_and_names_its_parent(tmp_path): + app = served(tmp_path) + with http_server(app) as port: + data = models(port) + assert [m["id"] for m in data] == ["fake-cuda", "alpha", "beta"] + assert "parent" not in data[0] and [(m["parent"], m["root"]) for m in data[1:]] == [("fake-cuda", "fake-cuda")] * 2 + + +def test_a_request_naming_an_adapters_model_id_runs_that_adapter_and_the_reply_names_it(tmp_path): + app = served(tmp_path) + with http_server(app) as port: + status, _, text = request(port, {"messages": HI, "max_tokens": 8, "model": "beta"}) + assert status == 200 and json.loads(text)["model"] == "beta" + assert app.engine.seen == ["beta"] + + +def test_the_adapter_field_selects_an_adapter_beside_the_base_model_id(tmp_path): + app = served(tmp_path) + with http_server(app) as port: + status, _, text = request(port, {"messages": HI, "max_tokens": 8, "model": "fake-cuda", "adapter": "alpha"}) + assert status == 200 and json.loads(text)["model"] == "fake-cuda" + status, _, _ = request(port, {"prompt": "Hi", "max_tokens": 8, "adapter": "beta"}, chat=False) + assert status == 200 + status, _, _ = request(port, {"messages": HI, "max_tokens": 8, "model": "alpha", "adapter": "alpha"}) + assert status == 200 + assert app.engine.seen == ["alpha", "beta", "alpha"] + + +def test_the_base_model_runs_without_an_adapter_argument_at_all(tmp_path): + app = served(tmp_path) + with http_server(app) as port: + for body in ({"messages": HI, "max_tokens": 8}, {"messages": HI, "max_tokens": 8, "model": "fake-cuda"}, + {"messages": HI, "max_tokens": 8, "model": "unlisted"}): + assert request(port, body)[0] == 200 + assert app.engine.seen == [NOT_PASSED] * 3 # the keyword is not sent for the base model + + +@pytest.mark.parametrize("extra", [{"adapter": "nope"}, {"adapter": ""}, {"adapter": 7}, {"adapter": ["alpha"]}, + {"model": "alpha", "adapter": "beta"}]) +def test_an_unknown_or_contradictory_adapter_is_refused_before_any_stream_opens(tmp_path, extra): + app = served(tmp_path) + with http_server(app) as port: + for stream in (False, True): + status, kind, text = request(port, {"messages": HI, "max_tokens": 8, "stream": stream, **extra}) + assert status == 400 and kind == "application/json", text + assert "adapter" in json.loads(text)["error"]["message"] + assert app.engine.calls == [] + + +def test_a_server_without_adapters_ignores_the_field_as_it_ignores_any_unknown_one(tmp_path): + app = app_for(tmp_path) # the plain fake engine: no adapters attribute + with http_server(app) as port: + status, _, text = request(port, {"messages": HI, "max_tokens": 8, "adapter": "anything"}) + assert status == 200 and json.loads(text)["model"] == "fake-cuda" + app = served(tmp_path / "again") + app.engine.adapters = None # an engine started without --adapter + with http_server(app) as port: + assert request(port, {"messages": HI, "max_tokens": 8, "adapter": "alpha"})[0] == 200 + assert app.engine.seen == [NOT_PASSED] + + +def test_an_engine_that_cannot_take_an_adapter_is_not_asked_to(tmp_path): + app = served(tmp_path) + app.engine = Engine() # no adapter parameter on generate + app.engine.adapters = served(tmp_path / "x").engine.adapters + with http_server(app) as port: + status, _, text = request(port, {"messages": HI, "max_tokens": 8, "adapter": "alpha"}) + assert status in (400, 500) and "does not serve adapters" in text and app.engine.calls == [] + + +def test_the_prepared_request_carries_the_choice(tmp_path): + app = served(tmp_path) + assert app.prepare({"messages": HI, "model": "alpha"}, True).adapter == "alpha" + assert app.prepare({"messages": HI}, True).adapter is None + assert app.model_ids == ["fake-cuda", "alpha", "beta"] + + +# -- the command line --------------------------------------------------------------------------------------------- +def args(*argv): + return cli.build_parser().parse_args(["serve", "/m", "--backend", "cuda", *argv]) + + +def test_the_flags_parse_to_names_folders_and_strengths_in_order(): + specs = serve_options.adapter_specs(args("--adapter", "b=/x/b", "--adapter", "a=/x/a", "--adapter-strength", "a=0.5")) + assert specs == [("b", "/x/b", 1.0), ("a", "/x/a", 0.5)] + assert serve_options.adapter_specs(args()) == [] + + +@pytest.mark.parametrize("argv,match", [ + (["--adapter", "nofolder"], "NAME=FOLDER"), (["--adapter", "=/x"], "NAME=FOLDER"), (["--adapter", "a="], "NAME=FOLDER"), + (["--adapter", "a b=/x"], "letters, digits"), (["--adapter", "a=/x", "--adapter", "a=/y"], "twice"), + (["--adapter", "a=/x", "--adapter-strength", "a"], "NAME=NUMBER"), + (["--adapter", "a=/x", "--adapter-strength", "a=zero"], "NAME=NUMBER"), + (["--adapter", "a=/x", "--adapter-strength", "a=0"], "NAME=NUMBER"), + (["--adapter", "a=/x", "--adapter-strength", "a=nan"], "NAME=NUMBER"), + (["--adapter", "a=/x", "--adapter-strength", "b=2"], "no --adapter defines"), + (["--adapter-strength", "a=2"], "no --adapter defines"), +]) +def test_malformed_adapter_flags_are_refused(argv, match): + with pytest.raises(ValueError, match=match): + serve_options.adapter_specs(args(*argv)) + + +def family(adapters=True): + from types import SimpleNamespace + + return SimpleNamespace(title="Fam", model_type="deepseek_v41", package=SimpleNamespace(CUDA_ADAPTERS=adapters)) + + +def test_adapters_need_a_family_that_serves_them_the_cuda_backend_and_parallel_streams(): + ok = args("--adapter", "a=/x", "--parallel", "4", "--name", "base") + serve_options.check(ok, family(), "cuda") + with pytest.raises(ValueError, match="has none"): + serve_options.check(ok, family(False), "cuda") + with pytest.raises(ValueError, match="MLX has none"): + serve_options.check(ok, family(), "mlx") + with pytest.raises(ValueError, match="--parallel 2 or more"): + serve_options.check(args("--adapter", "a=/x", "--name", "base"), family(), "cuda") + with pytest.raises(ValueError, match="--parallel 2 or more"): + serve_options.check(args("--adapter", "a=/x", "--parallel", "1", "--name", "base"), family(), "cuda") + + +def test_an_adapter_may_not_take_the_base_models_name_or_an_alias(): + for argv in (["--adapter", "base=/x", "--name", "base"], ["--adapter", "other=/x", "--alias", "other"]): + with pytest.raises(ValueError, match="already a model id"): + serve_options.check(args(*argv, "--parallel", "2", "--name", "base") if "--name" not in argv + else args(*argv, "--parallel", "2"), family(), "cuda") + + +def test_serve_options_without_adapters_check_exactly_as_before(): + serve_options.check(args("--name", "base"), family(False), "cuda") + + +def test_serve_hands_the_adapters_and_the_taken_model_ids_to_the_engine(tmp_path, monkeypatch): + from types import SimpleNamespace + + seen = {} + package = SimpleNamespace(CUDA_ADAPTERS=True, cuda_engine=lambda *a, **k: seen.update(k) or SimpleNamespace(max_len=8192)) + family = SimpleNamespace(title="Test family", model_type="test", package=package) + monkeypatch.setattr(server, "App", lambda *a, **k: SimpleNamespace(effective_context_window=8192)) + monkeypatch.setattr(server, "serve", lambda *a: None) + base = ["serve", str(tmp_path), "--backend", "cuda", "--no-drafts", "--name", "base", "--alias", "chat"] + parsed = cli.build_parser().parse_args(base + ["--parallel", "2", "--adapter", "x=/x", "--adapter-strength", "x=0.5"]) + assert cli._serve_cuda(parsed, family, tmp_path, 8192) == 0 + assert seen["adapters"] == [("x", "/x", 0.5)] and seen["adapter_reserved"] == ["base", "chat"] + seen.clear() + assert cli._serve_cuda(cli.build_parser().parse_args(base), family, tmp_path, 8192) == 0 + assert "adapters" not in seen and "adapter_reserved" not in seen # no flag: the engine is called as before + + +def test_the_responses_route_forwards_the_adapter_field_and_the_model_id_to_the_chat_it_runs(): + from tensorfold.server.responses_translate import translate + + chat = translate({"model": "alpha", "input": "hi", "adapter": "beta"}, None).chat + assert chat["model"] == "alpha" and chat["adapter"] == "beta" + assert "adapter" not in translate({"model": "alpha", "input": "hi"}, None).chat diff --git a/tests/test_deepseek_v41_adapter_tool.py b/tests/test_deepseek_v41_adapter_tool.py new file mode 100644 index 0000000000..3c57dac05f --- /dev/null +++ b/tests/test_deepseek_v41_adapter_tool.py @@ -0,0 +1,86 @@ +"""tools/dsv41/adapter_check.py: the operator's offline check of an adapter folder against a checkpoint (no torch).""" + +import importlib.util +import json +import pathlib +import struct + +import pytest +from adapter_fakes import CFG, SHARED, WO_B, factors, full_config, standard, write_adapter + +from tensorfold.families.deepseek_v41.cuda import lora_format as lf + +TOOL = pathlib.Path(__file__).resolve().parents[1] / "tools" / "dsv41" / "adapter_check.py" + + +def run(monkeypatch, capsys, *argv): + spec = importlib.util.spec_from_file_location("adapter_check_under_test", TOOL) + tool = importlib.util.module_from_spec(spec) + spec.loader.exec_module(tool) + monkeypatch.setattr("sys.argv", ["adapter_check.py", *map(str, argv)]) + code = tool.main() + out = capsys.readouterr() + return code, out.out, out.err + + +@pytest.fixture +def model(tmp_path): + (tmp_path / "model").mkdir() + (tmp_path / "model" / "config.json").write_text(json.dumps(full_config())) + return tmp_path / "model" + + +def test_it_prints_the_modules_and_the_memory_a_rank_would_hold(tmp_path, model, monkeypatch, capsys): + folder = standard(tmp_path / "a", ranks=(2, 3)) + code, out, _ = run(monkeypatch, capsys, folder, model, "--tp", "2", "--slots", "2") + assert code == 0 and "ok: 2 modules" in out and "layers.1.attn.wo_b: rank 2" in out + code, out, _ = run(monkeypatch, capsys, folder, model, "--json", "--tp", "2", "--slots", "1") + report = json.loads(out) + # slots+1 slots, rank rounded up to 4, float32: wo_b reads half of 1024 and writes 256, w2 half of 512 and writes 256 + assert report["bytes_per_rank"] == 2 * 4 * (1024 // 2 + 256) * 4 + 2 * 4 * (512 // 2 + 256) * 4 + assert set(report["modules"]) == {WO_B.format(1), SHARED.format(2)} + + +def test_a_written_manifest_is_accepted_by_the_server_and_its_provenance_is_kept(tmp_path, model, monkeypatch, capsys): + folder = standard(tmp_path / "a", manifest={"format": lf.FORMAT, "provenance": {"made_by": "a tool"}}) + code, _, _ = run(monkeypatch, capsys, folder, model, "--write-manifest") + assert code == 0 + written = json.loads((folder / lf.MANIFEST).read_text()) + assert written["provenance"] == {"made_by": "a tool"} and written["modules"][WO_B.format(1)] == {"rank": 2} + assert lf.inspect_folder(folder, lf.Dims.of(CFG), config_sha=written["base"]["config_sha256"]).manifest == written + assert run(monkeypatch, capsys, folder, model)[0] == 0 # and it checks out again + + +def test_what_the_server_would_refuse_is_refused_with_the_reason_and_exit_one(tmp_path, model, monkeypatch, capsys): + bad = write_adapter(tmp_path / "bad", {"mtp.0.attn.wo_b": factors(__import__("numpy").random.default_rng(0), + CFG.dim, 1024, 2)}) + code, out, err = run(monkeypatch, capsys, bad, model) + assert code == 1 and out == "" and "not a module this engine can edit" in err + nan = write_adapter(tmp_path / "nan", {WO_B.format(1): (factors(__import__("numpy").random.default_rng(1), + CFG.dim, 1024, 2)[0], + __import__("numpy").full((CFG.dim, 2), float("nan"), + dtype="float32"))}) + assert run(monkeypatch, capsys, nan, model)[0] == 1 + + +@pytest.mark.parametrize("name,manifest,config,match", [ + ("base", {"format": lf.FORMAT, "base": "deepseek_v41"}, None, "base must be an object"), + ("files", {"format": lf.FORMAT, "files": {lf.WEIGHTS: "abc"}}, None, + "files.adapter_model.safetensors must be an object"), + ("provenance", {"format": lf.FORMAT, "provenance": ["a"]}, None, "provenance must be an object"), + ("activated", None, {"alora_invocation_tokens": [5, 6]}, "alora_invocation_tokens is set"), +]) +def test_a_folder_of_the_wrong_shape_is_refused_with_the_reason_not_a_traceback(tmp_path, model, monkeypatch, capsys, + name, manifest, config, match): + folder = standard(tmp_path / name, manifest=manifest, config=config) + before = sorted((p.name, p.read_bytes()) for p in folder.iterdir()) + code, out, err = run(monkeypatch, capsys, folder, model, "--write-manifest") + assert code == 1 and out == "" and err.startswith("refused: ") and match in err + assert sorted((p.name, p.read_bytes()) for p in folder.iterdir()) == before # and nothing was written + + +def test_a_weights_header_that_is_not_an_object_is_refused_with_the_reason(tmp_path, model, monkeypatch, capsys): + folder = standard(tmp_path / "a") + (folder / lf.WEIGHTS).write_bytes(struct.pack(" int: + """Float32 table bytes one rank holds for this adapter alone: slots a module, the rank rounded up to fours.""" + + total = 0 + for m in spec.modules: + width, out = dims.shape(m.kind) + total += (slots + 1) * lf.padded_rank(m.rank) * (width // world + out) * 4 + return total + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("adapter") + ap.add_argument("model") + ap.add_argument("--tp", type=int, default=2) + ap.add_argument("--slots", type=int, default=1) + ap.add_argument("--write-manifest", action="store_true", help="write tensorfold_adapter.json from the files") + ap.add_argument("--json", action="store_true") + args = ap.parse_args() + model = pathlib.Path(args.model) + cfg = Cfg.read(model) + dims = lf.Dims.of(cfg) + config = model / "config.json" + sha = hashlib.sha256(config.read_bytes()).hexdigest() + try: + spec = lf.inspect_folder(args.adapter, dims, config_sha=sha) + for module in spec.modules: # a value that is not finite is refused at load + lf.read_factors(spec, module) + except lf.AdapterError as exc: + print(f"refused: {exc}", file=sys.stderr) + return 1 + report = spec.describe() + report["base"] = {"model_type": "deepseek_v41", "hidden_size": dims.dim, "num_hidden_layers": dims.n_layers, + "config_sha256": sha} + report["bytes_per_rank"] = rank_bytes(spec, dims, args.tp, args.slots) + if args.write_manifest: + manifest = {"format": lf.FORMAT, "base": report["base"], "files": {lf.WEIGHTS: {"sha256": spec.sha256}}, + "modules": {m.name: {"rank": m.rank} for m in spec.modules}, + "provenance": dict(spec.manifest.get("provenance") or {})} + (pathlib.Path(args.adapter) / lf.MANIFEST).write_text(json.dumps(manifest, indent=2) + "\n") + if args.json: + print(json.dumps(report, indent=2)) + return 0 + print(f"ok: {len(spec.modules)} modules, weights sha256 {spec.sha256[:16]}") + for m in spec.modules: + print(f" {m.name}: rank {m.rank}, alpha {m.alpha:g}, scale {m.scale:g}") + print(f"float32 tables on a rank of {args.tp}, {args.slots} adapter(s) held: " + f"{report['bytes_per_rank'] / 2**20:.1f} MiB") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/dsv41/adapter_mix.py b/tools/dsv41/adapter_mix.py new file mode 100644 index 0000000000..2247a3001c --- /dev/null +++ b/tools/dsv41/adapter_mix.py @@ -0,0 +1,90 @@ +"""Exactness of mixed adapter rounds on a running server: each stream of a concurrent burst, whatever model id it asks +for (the base model or an adapter) and wherever the ids sit among the streams, must give the token-id hash (token_sha) +of its own solo run, and each solo run that of its "draft": false run. The ids rotate over the streams, so every +stream takes every id in turn. One JSON line a rotation; exit 1 on a mismatch or a failed request. Standard library only. + + python3 adapter_mix.py http://127.0.0.1:8080 --models BASE,adapter-a,adapter-b --streams 4 --tokens 128 +""" + +import argparse +import json +import sys +import threading +import urllib.request + +PROMPTS = ["Write a short Python function that computes the Fibonacci sequence and explain it.", + "Explain how matrix multiplication uses a GPU in plain English, then give a small numerical example.", + "List the planets of the solar system in order from the sun, one line each, with one fact about each.", + "Summarize the plot of a heist story set on a space station in about five sentences."] + + +def ask(base: str, model: str, prompt: str, seed: int, tokens: int, temperature: float, draft: bool = True, + gate: threading.Barrier | None = None) -> dict: + body = {"model": model, "prompt": prompt, "max_tokens": tokens, "temperature": temperature, "stream": True, + "stream_options": {"include_usage": True}, "ignore_eos": True, "seed": seed} + if temperature > 0: + body.update(top_k=20, top_p=0.95) + if not draft: + body["draft"] = False + req = urllib.request.Request(base + "/v1/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + try: + if gate is not None: + gate.wait() + sha = None + with urllib.request.urlopen(req, timeout=1800) as resp: + for raw in resp: + line = raw.decode().strip() + if line.startswith("data:") and line != "data: [DONE]": + sha = (json.loads(line[5:]).get("tensorfold") or {}).get("token_sha") or sha + return {"sha": sha, "error": None if sha else "the reply carried no token_sha"} + except Exception as exc: # noqa: BLE001 - a failed request is reported + return {"sha": None, "error": f"{type(exc).__name__}: {exc}"} + + +def burst(base: str, models: list[str], seeds: list[int], tokens: int, temperature: float) -> list[dict]: + gate = threading.Barrier(len(models)) + out: list[dict] = [{} for _ in models] + + def run(i: int) -> None: + out[i] = ask(base, models[i], PROMPTS[i % len(PROMPTS)], seeds[i], tokens, temperature, gate=gate) + + threads = [threading.Thread(target=run, args=(i,)) for i in range(len(models))] + for t in threads: + t.start() + for t in threads: + t.join() + return out + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + p.add_argument("base") + p.add_argument("--models", required=True, help="model ids, a comma list (the base model's and the adapters')") + p.add_argument("--streams", type=int, default=4) + p.add_argument("--tokens", type=int, default=128) + p.add_argument("--seed", type=int, default=1234) + p.add_argument("--temperature", type=float, default=1.0) + args = p.parse_args() + ids = [m.strip() for m in args.models.split(",") if m.strip()] + seeds = [args.seed + i for i in range(args.streams)] + bad = 0 + for turn in range(len(ids)): + models = [ids[(i + turn) % len(ids)] for i in range(args.streams)] + solo = [ask(args.base, models[i], PROMPTS[i % len(PROMPTS)], seeds[i], args.tokens, args.temperature) + for i in range(args.streams)] + serial = [ask(args.base, models[i], PROMPTS[i % len(PROMPTS)], seeds[i], args.tokens, args.temperature, + draft=False) for i in range(args.streams)] + together = burst(args.base, models, seeds, args.tokens, args.temperature) + errors = sorted({r["error"] for r in solo + serial + together if r.get("error")}) + concurrent = [t["sha"] == s["sha"] and t["sha"] is not None for t, s in zip(together, solo)] + drafted = [s["sha"] == d["sha"] and s["sha"] is not None for s, d in zip(solo, serial)] + line = {"models": models, "concurrent_equals_solo": concurrent, "solo_equals_serial": drafted, + "errors": errors[:3]} + print(json.dumps(line), flush=True) + bad += (not all(concurrent)) + (not all(drafted)) + bool(errors) + return 1 if bad else 0 + + +if __name__ == "__main__": + sys.exit(main())