RomanNum3ral
/
soup
mirror of https://github.com/razor-ai/soup.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

feat: live FSDP-shard consolidation + serve KV-cache type + ONNX QA (v0.71.14) 

Close the doable tail of the export-QA + deferred-stub family.

- #96 consolidate_shards: lazy-torch safetensors merge, TOCTOU-hardened
  (enforce_under_cwd_and_no_symlink, weights_only=True, atomic_write_bytes),
  16 GiB/shard cap, dup-key shape-conflict reject. --plan-only flag.
- #140 apply_kv_cache_type -> KvCacheRuntime; soup serve --kv-cache-type
  (bf16/f16 dtype, q8_0 quantized cache + hqq probe, fp8 Hopper-gated).
- #71 ONNX export QA: tiny-GPT2 PASS, TinyLlama host-RAM-bound (qa doc).
- #70/#72/#144/#74/#79 deferred to INFRA-BLOCKED tail (kept open).

Tests 13430 -> 13476 (+46). ruff clean. Suite green, 78.52% cov.
This commit is contained in:
Alpamys 2026-06-05 13:39:36 +05:00
parent 3ac9e305ce
commit 6ec36b5ff6
14 changed files with 1237 additions and 70 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
CHANGELOG.md
View File

@ -12,6 +12,37 @@ reproducing 70+ versions of notes.
## [Unreleased]
## [0.71.14] - 2026-06-05
### Added
- **Live FSDP shard consolidation** (closes #96). `soup merge-sharded-fsdp-weights`
lifts the v0.44.0 plan-only stub: it now streams each `pytorch_model_fsdp_*.bin`
shard via `torch.load(weights_only=True)` (no arbitrary pickle exec), unions the
per-rank parameter fragments into one state-dict, and writes a single
`.safetensors` atomically. Memory-friendly (one shard loaded at a time). New
`--plan-only` flag prints the plan without writing. Single-process — no
multi-GPU needed to MERGE. (Per-rank disjoint-parameter / FULL_STATE_DICT
shards; DCP sharded-tensor reconstruction is out of scope — use
`accelerate merge-weights` for those.)
- **Live `kv_cache_type` wiring on the transformers serve backend** (closes #140).
`soup serve --kv-cache-type q8_0 | bf16 | f16 | fp8` lifts the v0.53.1
`apply_kv_cache_type` `NotImplementedError` stub: `bf16`/`f16` load the model in
that dtype (the KV cache inherits it); `q8_0` routes an 8-bit HQQ quantized KV
cache through `model.generate` (needs `pip install hqq`); `fp8` raises a friendly
runtime error (vLLM + Hopper-only — the transformers backend has no fp8 KV
path). vLLM / SGLang KV-cache-dtype routing stays in the infra-blocked tail.
- **ONNX export QA verified** (closes #71) — `soup export --format onnx` exercised
end-to-end on a tiny model: export exits 0, `model.onnx` loads in ONNX Runtime
with `input_ids` present, and a forward pass produces a real output. Recorded in
`tests/qa/v07114_qa.md`.
### Notes
- GGUF export (#70), AWQ/GPTQ export (#72), the CUDA + llama.cpp QA doc (#144),
HF Hub push/Spaces deploy (#74), and the Community-QA tracking meta-issue (#79)
remain open with `infra-blocked` labels — they need a built llama.cpp toolchain,
`autoawq`/`auto-gptq` Windows wheels, or HF credentials the QA box lacks. See
`tests/qa/v07114_qa.md`.
## [0.71.13] - 2026-06-04
### Added

2
CONTRIBUTING.md
View File

@ -120,7 +120,7 @@ src/soup_cli/
templates/ - 17 built-in soup.yaml templates (YAML + manifest.json) with load_template loader (v0.39.0, +bco v0.40.0)
ui/ - Web UI (FastAPI + HTML/JS SPA)
tests/ - Test suite (284 files, 13430 tests)
tests/ - Test suite (285 files, 13476 tests)
examples/ - Real-world config examples and datasets
```

24
README.md
View File

@ -49,18 +49,20 @@ infrastructure instead of improving models. Soup fixes that.
## What's New
**v0.71.13 — Prompt-compile family (live).** Four commands that shipped as deferred stubs in
v0.68.0 are now real, validated end-to-end on tiny models:
**v0.71.14 — Export QA + serve finale.** Two deferred-stub runtimes go live (validated on a tiny
model + the real RTX 3050), plus the ONNX export pipeline verified end-to-end:
- **`soup local-rl train --once`** — harvest your latest 👍/👎 feedback into DPO pairs and train a
real DPO/KTO/ORPO round via `soup train` (argv list, no shell). A `state` table tracks the last
train so a re-run with no new feedback skips. Without `--once` it renders a systemd/launchd nightly
scheduler scaffold you can install.
- **`soup distill-prompt`** — call a teacher model (Ollama / Anthropic / vLLM) once per trace and
write a real distillation dataset: `sft`/`kl` → `{messages}`, `preference` → `{prompt, chosen, rejected}`.
- **`soup compile`** — DSPy / GEPA / TextGrad prompt-program optimisation, and **`soup compile-tools`**
— TextGrad / GEPA tool-schema optimisation, both behind the new `[compile]` extra
(`pip install 'soup-cli[compile]'`).
- **`soup merge-sharded-fsdp-weights`** — live FSDP shard consolidation: streams each
`pytorch_model_fsdp_*.bin` shard (no arbitrary pickle exec) and writes one `.safetensors` atomically.
Single-process, memory-friendly. `--plan-only` previews without writing.
- **`soup serve --kv-cache-type bf16 | f16 | q8_0 | fp8`** — live KV-cache wiring on the transformers
backend: `bf16`/`f16` set the cache dtype, `q8_0` runs an 8-bit quantized KV cache (needs `hqq`),
`fp8` gives a friendly vLLM+Hopper-only error.
- **ONNX export verified** — `soup export --format onnx` confirmed to produce a runnable ONNX graph
loadable in ONNX Runtime.
GGUF / AWQ / GPTQ export + HF Hub push QA stay open with `infra-blocked` labels (need a built
llama.cpp toolchain, Windows quant wheels, or HF credentials).
Full history: [CHANGELOG.md](CHANGELOG.md) · [GitHub Releases](https://github.com/MakazhanAlpamys/Soup/releases).

2
docs/commands.md
View File

@ -22,6 +22,7 @@ soup chat --model ./output Interactive chat
soup push --model ./output --repo user/name Upload to HuggingFace
soup push --model ./output --repo user/name --collection user/coll-abc123 Add to HF Collection
soup merge --adapter ./output Merge LoRA with base model
soup merge-sharded-fsdp-weights ./shards -o merged.safetensors Consolidate FSDP shards into one safetensors (v0.71.14; --plan-only previews)
soup export --model ./output --format gguf Export to GGUF (Ollama)
soup export --model ./output --deploy ollama Export GGUF + auto-deploy to Ollama
soup export --model ./output --format onnx Export to ONNX
@ -57,6 +58,7 @@ soup serve --model <m> --dashboard Live dashboard + /metrics endpoint
soup serve --model <m> --trace --trace-endpoint http://localhost:4317 OpenTelemetry tracing
soup serve --model <m> --trace-log ./serve.jsonl Per-request JSONL log + rotation + secret redaction
soup serve --model <m> --record-thumbs ./rl.db Capture 👍/👎 feedback into local-RL SQLite + POST /v1/thumbs (transformers)
soup serve --model <m> --kv-cache-type bf16|f16|q8_0|fp8 KV-cache type (transformers; q8_0 needs hqq; fp8 = vLLM+Hopper only) (v0.71.14)
POST /v1/adapters/activate/<name> Hot-swap active LoRA adapter
soup sweep --config soup.yaml --param lr=... Hyperparameter search
soup diff --model-a ./a --model-b ./b Compare two models

2
pyproject.toml
View File

@ -4,7 +4,7 @@ build-backend = "hatchling.build"
[project]
name = "soup-cli"
version = "0.71.13"
version = "0.71.14"
description = "Fine-tune LLMs in one command. No SSH, no config hell."
readme = "README.md"
license = "Apache-2.0"

2
src/soup_cli/__init__.py
View File

@ -1,3 +1,3 @@
"""Soup CLI — Fine-tune LLMs in one command."""
__version__ = "0.71.13"
__version__ = "0.71.14"

66
src/soup_cli/commands/merge_sharded_fsdp_weights.py
View File

@ -1,7 +1,9 @@
"""v0.44.0 Part D — `soup merge-sharded-fsdp-weights` command.
"""`soup merge-sharded-fsdp-weights` command.
Schema-only stub: discovers + validates FSDP shards and prints the planned
operation. Live consolidation lands in v0.44.1.
v0.44.0 shipped the planner; v0.71.14 (#96) lifts the live torch-side
consolidation: by default the command loads each FSDP shard and writes a single
consolidated ``.safetensors``. Pass ``--plan-only`` to print the plan without
writing anything (single-process — no multi-GPU needed to MERGE).
"""
from __future__ import annotations
@ -11,7 +13,7 @@ from rich.console import Console
from rich.markup import escape
from rich.panel import Panel
from soup_cli.utils.fsdp_consolidate import plan_consolidation
from soup_cli.utils.fsdp_consolidate import consolidate_shards, plan_consolidation
console = Console()
@ -27,31 +29,57 @@ def merge_sharded_fsdp_weights(
"-o",
help="Destination .safetensors file path (under cwd).",
),
yes: bool = typer.Option(
plan_only: bool = typer.Option(
False,
"--yes",
help="Acknowledge that live consolidation lands in v0.44.1 (plan-only now).",
"--plan-only",
help="Print the consolidation plan and exit without writing anything.",
),
) -> None:
"""Plan a consolidation of FSDP shard files into a single safetensors file.
"""Consolidate FSDP shard files into a single safetensors file.
v0.44.0 ships the planner; live torch-side consolidation lands in v0.44.1.
Streams each ``pytorch_model_fsdp_*.bin`` shard (loaded via
``torch.load(weights_only=True)`` — no arbitrary pickle exec), unions the
per-shard parameter fragments, and writes one ``.safetensors`` atomically.
"""
try:
plan = plan_consolidation(shard_dir, output)
except (ValueError, FileNotFoundError, RuntimeError) as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(code=2) from exc
body = (
f"Shards found: {len(plan.shard_files)}\n"
f"Source dir: {escape(plan.shard_dir)}\n"
f"Output (target): {escape(plan.output_path)}\n\n"
"Live consolidation runtime lands in v0.44.1 — this is a plan-only run."
)
console.print(Panel(body, title="FSDP Consolidation Plan", border_style="cyan"))
if not yes:
if plan_only:
body = (
f"Shards found: {len(plan.shard_files)}\n"
f"Source dir: {escape(plan.shard_dir)}\n"
f"Output (target): {escape(plan.output_path)}\n\n"
"Plan-only run — pass without --plan-only to write the file."
)
console.print(
"[yellow]Pass --yes to acknowledge the deferred runtime "
"and exit cleanly.[/]"
Panel(body, title="FSDP Consolidation Plan", border_style="cyan")
)
raise typer.Exit(code=0)
console.print(
f"[dim]Consolidating {len(plan.shard_files)} shard(s)...[/]"
)
try:
result = consolidate_shards(plan)
except (TypeError, ValueError) as exc:
console.print(f"[red]{escape(str(exc))}[/]")
raise typer.Exit(code=2) from exc
except ImportError as exc:
console.print(
"[red]torch + safetensors are required for consolidation. "
"Install with: [bold]pip install 'soup-cli[train]'[/][/]"
)
raise typer.Exit(code=1) from exc
body = (
f"Tensors written: {result.num_tensors}\n"
f"Shards merged: {result.num_shards}\n"
f"Size: {result.total_bytes / 1e6:.2f} MB\n"
f"Output: {escape(result.output_path)}"
)
console.print(
Panel(body, title="[bold green]FSDP Consolidation Done[/]", border_style="green")
)

114
src/soup_cli/commands/serve.py
View File

@ -5,6 +5,7 @@ import logging
import re
import time
import uuid
from collections.abc import Mapping
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional
@ -249,6 +250,16 @@ def serve(
"--bank-strength",
help="Vector-bank delta strength multiplier. Ignored when --bank is unset.",
),
kv_cache_type: Optional[str] = typer.Option(
None,
"--kv-cache-type",
help=(
"KV-cache type for decoding: q8_0 (8-bit quantized, needs hqq) / "
"bf16 / f16 (cache dtype) / fp8 (vLLM+Hopper only). Transformers "
"backend only — vLLM / SGLang routing is in the blocked tail. "
"v0.71.14 (#140)."
),
),
):
"""Start a local inference server with OpenAI-compatible API."""
# v0.71.12 #221 — validate `--bank` up front (path containment + backend)
@ -301,6 +312,50 @@ def serve(
)
raise typer.Exit(code=2)
# v0.71.14 #140 — resolve `--kv-cache-type` up front (before model load) so
# an invalid type / fp8-on-Ampere / vLLM-backend / missing-quant-backend
# surfaces immediately. The resolved runtime is threaded into the
# transformers branch below (model dtype + generate kwargs).
resolved_kv_runtime = None
if kv_cache_type is not None:
from rich.markup import escape as _rich_escape
from soup_cli.utils.kv_cache import (
apply_kv_cache_type,
quantized_cache_backend_available,
)
cc = None
try:
import torch as _torch
if _torch.cuda.is_available():
cc = _torch.cuda.get_device_capability(0)
except Exception as _cc_exc: # noqa: BLE001 — torch missing / no CUDA
# cc stays None → the fp8 gate falls back to the generic
# vLLM-only message instead of the precise capability one.
logger.debug("kv-cache CUDA capability probe failed: %r", _cc_exc)
cc = None
try:
resolved_kv_runtime = apply_kv_cache_type(
kv_cache_type, backend=backend.lower(), compute_capability=cc
)
except (TypeError, ValueError, NotImplementedError, RuntimeError) as exc:
console.print(
f"[red]--kv-cache-type:[/] {_rich_escape(str(exc))}"
)
raise typer.Exit(code=2) from exc
if (
resolved_kv_runtime.requires_quant_backend
and quantized_cache_backend_available() is None
):
console.print(
"[red]--kv-cache-type q8_0[/] needs a quantized-cache backend. "
"Install one with [bold]pip install hqq[/] "
"(or optimum-quanto)."
)
raise typer.Exit(code=2)
# v0.53.10 #152 — pre-fetch base from a non-HF hub before serve starts.
if hub and hub != "hf":
from soup_cli.utils.hubs import apply_hub_to_cli_model
@ -640,8 +695,16 @@ def serve(
is_adapter=is_adapter,
device=device,
trust_remote_code=resolved_trust,
kv_cache_dtype=(
resolved_kv_runtime.model_dtype if resolved_kv_runtime else None
),
)
console.print("[bold green]Model loaded![/]")
if resolved_kv_runtime is not None:
console.print(
f"[green]KV cache:[/] {resolved_kv_runtime.kv_cache_type} "
f"— {resolved_kv_runtime.note}"
)
# v0.71.10 #201 — install the activation-steering decode hook. The
# handle persists for the server's lifetime (process-global model).
@ -821,6 +884,11 @@ def serve(
reasoning_parser=resolved_reasoning_parser,
record_thumbs_db=record_thumbs_db,
loaded_bank=loaded_bank,
kv_cache_generate_kwargs=(
_plain_kv_kwargs(resolved_kv_runtime.generate_kwargs)
if resolved_kv_runtime
else None
),
)
console.print(
@ -945,11 +1013,22 @@ def _load_model(
is_adapter: bool,
device: str,
trust_remote_code: bool = False,
kv_cache_dtype: Optional[str] = None,
):
"""Load model and tokenizer."""
"""Load model and tokenizer.
``kv_cache_dtype`` (v0.71.14 #140) selects the model compute dtype so the
transformers DynamicCache runs in it: ``"bfloat16"`` → bf16, else the
default float16. The bf16/f16 ``kv_cache_type`` values map here; q8_0 uses
a quantized cache via generate kwargs and leaves the model dtype unchanged.
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
load_dtype = (
torch.bfloat16 if kv_cache_dtype == "bfloat16" else torch.float16
)
console.print("[dim]Loading tokenizer...[/]")
tokenizer = AutoTokenizer.from_pretrained(
model_path, trust_remote_code=trust_remote_code
@ -965,7 +1044,7 @@ def _load_model(
base_model,
trust_remote_code=trust_remote_code,
device_map="auto",
dtype=torch.float16,
dtype=load_dtype,
)
console.print(f"[dim]Loading LoRA adapter: {model_path}...[/]")
model_obj = PeftModel.from_pretrained(base, model_path)
@ -975,7 +1054,7 @@ def _load_model(
model_path,
trust_remote_code=trust_remote_code,
device_map="auto",
dtype=torch.float16,
dtype=load_dtype,
)
model_obj.eval()
@ -1007,6 +1086,23 @@ def _load_draft_model(speculative_model: str, device: str):
return draft
def _plain_kv_kwargs(mapping: Any) -> Dict[str, Any]:
"""Deep-convert a (possibly MappingProxyType-nested) mapping to plain dicts.
The v0.71.14 #140 ``kv_cache_type`` runtime stores ``generate_kwargs`` as
immutable ``MappingProxyType`` (incl. a nested ``cache_config``). transformers
``generate`` builds the quantized cache from a plain ``cache_config`` dict,
so convert before threading it in.
"""
out: Dict[str, Any] = {}
for key, value in dict(mapping).items():
if isinstance(value, Mapping):
out[key] = _plain_kv_kwargs(value)
else:
out[key] = value
return out
def _generate_response(
model,
tokenizer,
@ -1019,6 +1115,7 @@ def _generate_response(
num_assistant_tokens: int = 5,
logits_processor=None,
ngram_config: Any = None,
kv_cache_generate_kwargs: Optional[Dict[str, Any]] = None,
):
"""Generate a response from the model."""
import torch
@ -1079,6 +1176,12 @@ def _generate_response(
# is defence-in-depth.
pass
# v0.71.14 #140 — KV-cache-type generate kwargs (quantized cache for
# q8_0). bf16/f16 map to the model load dtype, not here, so this is
# empty for those. Merged last so it can't be clobbered.
if kv_cache_generate_kwargs:
gen_kwargs.update(kv_cache_generate_kwargs)
outputs = model.generate(**gen_kwargs)
new_tokens = outputs[0][input_ids.shape[1]:]
@ -1110,6 +1213,7 @@ def _create_app(
reasoning_parser: Optional[str] = None,
record_thumbs_db: Optional[str] = None,
loaded_bank: Any = None,
kv_cache_generate_kwargs: Optional[Dict[str, Any]] = None,
):
"""Create the FastAPI application with OpenAI-compatible endpoints.
@ -1297,6 +1401,7 @@ def _create_app(
num_assistant_tokens=num_speculative_tokens,
trace_log_writer=trace_log_writer,
started=stream_started,
kv_cache_generate_kwargs=kv_cache_generate_kwargs,
),
media_type="text/event-stream",
)
@ -1331,6 +1436,7 @@ def _create_app(
num_assistant_tokens=num_speculative_tokens,
logits_processor=processors or None,
ngram_config=ngram_config,
kv_cache_generate_kwargs=kv_cache_generate_kwargs,
)
except Exception:
logger.exception("Generation error")
@ -1735,6 +1841,7 @@ def _stream_response(
max_tokens, temperature, top_p, model_name,
assistant_model=None, num_assistant_tokens=5,
trace_log_writer=None, started=None,
kv_cache_generate_kwargs=None,
):
"""Generator that yields SSE chunks for streaming responses."""
chat_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
@ -1750,6 +1857,7 @@ def _stream_response(
top_p=top_p,
assistant_model=assistant_model,
num_assistant_tokens=num_assistant_tokens,
kv_cache_generate_kwargs=kv_cache_generate_kwargs,
)
except Exception:
logger.exception("Stream generation error")

151
src/soup_cli/utils/fsdp_consolidate.py
View File

@ -1,22 +1,39 @@
"""v0.44.0 Part D — `soup merge-sharded-fsdp-weights` consolidator.
Schema-only stub: validates the shard-directory layout and returns an
operation plan. Live consolidation (loading each shard via torch + writing a
single `.safetensors`) is deferred to v0.44.1 to keep the v0.44.0 surface
small.
v0.44.0 shipped the planner (``plan_consolidation`` / ``discover_shards``).
v0.71.14 (#96) lifts the deferred runtime: ``consolidate_shards`` streams each
FSDP shard ``.bin`` (via ``torch.load(weights_only=True)`` — no arbitrary
pickle exec), unions the per-shard parameter fragments into one state-dict, and
writes a single ``.safetensors`` atomically.
Scope note: this handles per-rank **FULL_STATE_DICT** / disjoint-parameter
shards (each ``.bin`` holds complete parameters; the union of all shards is the
full model). DCP sharded-tensor checkpoints (``.distcp`` with per-rank tensor
SLICES that need concatenation from the torch.distributed.checkpoint metadata)
are out of scope — use ``accelerate merge-weights`` for those.
"""
from __future__ import annotations
import logging
import os
import re
from dataclasses import dataclass
from typing import List, Tuple
from soup_cli.utils.paths import is_under_cwd
from soup_cli.utils.paths import (
atomic_write_bytes,
enforce_under_cwd_and_no_symlink,
is_under_cwd,
)
logger = logging.getLogger(__name__)
_SHARD_RE = re.compile(r"^pytorch_model_fsdp_\d+(_\d+)?\.bin$")
_MAX_SHARDS = 1024
# Defence-in-depth: a single shard fragment over 16 GiB is almost certainly
# a wrong file (or an attempt to OOM the box). Real per-rank shards are small.
_MAX_SHARD_BYTES = 16 * 1024**3
@dataclass(frozen=True)
@ -77,3 +94,127 @@ def plan_consolidation(shard_dir: str, output_path: str) -> ConsolidationPlan:
shard_files=tuple(shards),
output_path=os.path.realpath(output_path),
)
@dataclass(frozen=True)
class ConsolidationResult:
"""Outcome of a live shard consolidation."""
output_path: str
num_tensors: int
num_shards: int
total_bytes: int
def consolidate_shards(plan: ConsolidationPlan) -> ConsolidationResult:
"""Stream each FSDP shard and write a single consolidated ``.safetensors``.
Loads one shard at a time via ``torch.load(weights_only=True)`` (no
arbitrary pickle execution), unions the per-shard parameter fragments into
a single CPU state-dict, and writes it atomically. Memory-friendly: the
raw shard object is freed before the next is loaded (peak ≈ merged
state-dict + one shard).
Raises:
TypeError: ``plan`` is not a :class:`ConsolidationPlan`.
ValueError: a shard is not a state-dict, two shards disagree on a
tensor's shape, the output path is a symlink, or no tensors were
found across all shards.
"""
import torch # lazy — keep CLI startup fast
from safetensors.torch import save as st_save
if not isinstance(plan, ConsolidationPlan):
raise TypeError(
f"plan must be ConsolidationPlan, got {type(plan).__name__}"
)
# Re-validate the output path (symlink rejection + cwd containment) at
# write time — TOCTOU defence, mirrors v0.59.0 atomic_write policy.
enforce_under_cwd_and_no_symlink(plan.output_path, "output_path")
merged: dict = {}
duplicate_keys: List[str] = []
for shard_name in plan.shard_files:
shard_path = os.path.join(plan.shard_dir, shard_name)
# The shard file is a direct child of the realpath'd shard_dir, but a
# symlinked child could redirect torch.load to an arbitrary target —
# reject symlinks + confirm containment before reading.
enforce_under_cwd_and_no_symlink(shard_path, "shard")
try:
size = os.path.getsize(shard_path)
except OSError as exc:
raise ValueError(
f"shard {shard_name} unreadable: {type(exc).__name__}"
) from exc
if size > _MAX_SHARD_BYTES:
raise ValueError(
f"shard {shard_name} too large (>{_MAX_SHARD_BYTES} bytes); "
"refuse to load"
)
try:
state = torch.load(
shard_path, map_location="cpu", weights_only=True
)
except Exception as exc: # noqa: BLE001 — corrupt/non-torch shard
# torch.load raises UnpicklingError / RuntimeError / EOFError /
# BadZipFile on a corrupt or non-torch .bin. Surface a clean
# ValueError (the CLI maps it to exit 2) instead of leaking the
# raw pickle internals as a crash.
raise ValueError(
f"shard {shard_name} is not a valid torch checkpoint: "
f"{type(exc).__name__}"
) from exc
if not isinstance(state, dict):
raise ValueError(
f"shard {shard_name} is not a state-dict "
f"(got {type(state).__name__})"
)
for key, tensor in state.items():
if not isinstance(key, str):
raise ValueError(
f"shard {shard_name} has a non-string key {key!r}"
)
if not isinstance(tensor, torch.Tensor):
raise ValueError(
f"shard {shard_name} key {key!r} is not a tensor "
f"(got {type(tensor).__name__})"
)
if key in merged:
if tuple(merged[key].shape) != tuple(tensor.shape):
raise ValueError(
f"shape conflict for {key!r}: "
f"{tuple(merged[key].shape)} vs {tuple(tensor.shape)}"
)
# Same key + same shape across shards — FSDP replicates some
# params; keep the first occurrence. Track it so a genuinely
# sharded checkpoint (where ranks hold DIFFERENT values for the
# same key) isn't silently corrupted without a trace.
duplicate_keys.append(key)
continue
merged[key] = tensor.detach().to("cpu").contiguous()
# Free the raw shard before loading the next (memory-friendly).
del state
if duplicate_keys:
sample = ", ".join(duplicate_keys[:5])
logger.warning(
"%d key(s) appeared in more than one shard with the same shape; "
"kept the first occurrence (e.g. %s). This is correct for "
"replicated FULL_STATE_DICT params, but if your shards hold "
"per-rank tensor SLICES they need concatenation — use "
"`accelerate merge-weights` for DCP sharded checkpoints.",
len(duplicate_keys),
sample,
)
if not merged:
raise ValueError("no tensors found across shards; nothing to write")
payload = st_save(merged)
written = atomic_write_bytes(payload, plan.output_path, field="output_path")
return ConsolidationResult(
output_path=written,
num_tensors=len(merged),
num_shards=len(plan.shard_files),
total_bytes=len(payload),
)

189
src/soup_cli/utils/kv_cache.py
View File

@ -1,29 +1,39 @@
"""v0.53.0 Part C — KV cache types schema helpers.
"""v0.53.0 Part C — KV cache types schema helpers + v0.71.14 #140 live wiring.
Closed allowlist of ``kv_cache_type`` strings exposed to
``soup serve --kv-cache-type <type>`` and YAML ``training.kv_cache_type`` (the
field is reused by serve / chat runtime in v0.53.1). Mirrors the unsloth
serve recipe.
``soup serve --kv-cache-type <type>`` and YAML ``training.kv_cache_type``.
Mirrors the unsloth serve recipe.
* ``q8_0`` — 8-bit (default for the unsloth runtime)
* ``bf16`` — bfloat16
* ``f16`` — float16
* ``fp8`` — FP8 on Hopper+ (gated by a separate runtime check)
* ``q8_0`` — 8-bit quantized KV cache (transformers: HQQ quantized cache)
* ``bf16`` — bfloat16 KV cache (transformers: model + cache dtype)
* ``f16`` — float16 KV cache (transformers: model + cache dtype)
* ``fp8`` — FP8 on Hopper+ (vLLM-backend only; transformers has no fp8 path)
Live wiring into the vLLM / SGLang / transformers serve loops is deferred
to v0.53.1.
v0.71.14 (#140) lifts the v0.53.0 ``apply_kv_cache_type`` stub: it returns a
:class:`KvCacheRuntime` plan for the **transformers** backend (model dtype +
``generate`` kwargs) + a Hopper SM gate for ``fp8``. vLLM / SGLang stay in the
infra-blocked tail (need a Hopper / Linux box to validate honestly).
"""
from __future__ import annotations
from dataclasses import dataclass
import importlib.util
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import Mapping
from typing import Mapping, Optional, Tuple
KV_CACHE_TYPES: frozenset[str] = frozenset({"q8_0", "bf16", "f16", "fp8"})
# Live serve wiring exists for the transformers backend only. vLLM / SGLang
# KV-cache-dtype routing needs a Hopper / Linux box to validate honestly →
# blocked tail (v0.71.14 #140).
_LIVE_BACKENDS: frozenset[str] = frozenset({"transformers"})
_MAX_KV_CACHE_LEN: int = 16
# Hopper is compute capability (9, x); fp8 KV cache needs SM >= 9.0.
_HOPPER_MAJOR: int = 9
@dataclass(frozen=True)
class KVCacheSpec:
@ -109,13 +119,154 @@ def requires_hopper(name: object) -> bool:
return spec is not None and spec.requires_hopper
def apply_kv_cache_type() -> None:
"""Live KV-cache-type wiring — deferred to v0.53.1.
@dataclass(frozen=True)
class KvCacheRuntime:
"""Resolved KV-cache plan for a serve backend.
Mirrors v0.50.0 ``apply_vllm_sleep_mode`` and v0.52.0
``apply_moe_expert_quant`` stub-then-live pattern.
* ``model_dtype`` — ``"bfloat16"`` / ``"float16"`` for the dtype-based
types (the transformers DynamicCache inherits the model's compute
dtype), else ``None``.
* ``generate_kwargs`` — kwargs threaded into ``model.generate`` for the
quantized-cache types (``cache_implementation`` / ``cache_config``),
else empty.
* ``requires_quant_backend`` — True for ``q8_0`` (needs ``hqq`` or
``optimum-quanto`` installed for the quantized cache to run).
"""
raise NotImplementedError(
"kv_cache_type live wiring deferred to v0.53.1. Schema accepts "
"q8_0 / bf16 / f16 / fp8 but no serve backend is routing the flag yet."
kv_cache_type: str
backend: str
model_dtype: Optional[str]
generate_kwargs: Mapping[str, object] = field(
default_factory=lambda: MappingProxyType({})
)
requires_quant_backend: bool = False
note: str = ""
def quantized_cache_backend_available() -> Optional[str]:
"""Return the installed quantized-KV-cache backend name, or ``None``.
Probes (via :func:`importlib.util.find_spec`, no import) for the optional
libraries the transformers quantized cache needs: ``hqq`` (8-bit) or
``optimum-quanto`` / ``quanto`` (2/4-bit). Returns the first found, else
``None`` (the ``q8_0`` path then surfaces a friendly install advisory).
"""
if importlib.util.find_spec("hqq") is not None:
return "hqq"
if importlib.util.find_spec("optimum.quanto") is not None:
return "quanto"
if importlib.util.find_spec("quanto") is not None:
return "quanto"
return None
def _validate_compute_capability(
cc: object,
) -> Optional[Tuple[int, int]]:
"""Validate an optional ``(major, minor)`` compute-capability tuple."""
if cc is None:
return None
if isinstance(cc, bool) or not isinstance(cc, tuple) or len(cc) != 2:
raise TypeError(
"compute_capability must be a (major, minor) int 2-tuple or None"
)
major, minor = cc
if isinstance(major, bool) or isinstance(minor, bool):
raise TypeError("compute_capability entries must be int, not bool")
if not isinstance(major, int) or not isinstance(minor, int):
raise TypeError("compute_capability entries must be int")
return (major, minor)
def apply_kv_cache_type(
kv_cache_type: object,
*,
backend: str = "transformers",
compute_capability: Optional[Tuple[int, int]] = None,
) -> KvCacheRuntime:
"""Resolve a ``kv_cache_type`` into a serve runtime plan (v0.71.14 #140).
Args:
kv_cache_type: one of ``q8_0`` / ``bf16`` / ``f16`` / ``fp8``.
backend: serve backend. Only ``transformers`` is live; ``vllm`` /
``sglang`` raise :class:`NotImplementedError` (blocked tail).
compute_capability: optional ``(major, minor)`` CUDA compute
capability, used to gate ``fp8`` (Hopper SM >= 9.0).
Raises:
TypeError / ValueError: invalid ``kv_cache_type`` / ``backend`` /
``compute_capability``.
NotImplementedError: ``backend`` is vLLM / SGLang (deferred).
RuntimeError: ``fp8`` (transformers has no fp8 KV-cache path; it is a
vLLM + Hopper feature — friendly error names the requirement).
"""
canonical = validate_kv_cache_type(kv_cache_type)
if isinstance(backend, bool) or not isinstance(backend, str):
raise TypeError(
f"backend must be str, got {type(backend).__name__}"
)
backend_l = backend.lower()
if backend_l not in _LIVE_BACKENDS:
raise NotImplementedError(
f"kv_cache_type live wiring for backend {backend!r} is deferred "
"(vLLM / SGLang KV-cache-dtype routing needs a Hopper / Linux box "
"to validate — tracked in the infra-blocked tail). "
"Use --backend transformers."
)
cc = _validate_compute_capability(compute_capability)
spec = _KV_CACHE_METADATA[canonical]
if spec.requires_hopper:
# fp8: transformers has NO fp8 KV-cache path — it is a vLLM feature
# that additionally needs a Hopper+ GPU. Always raise on transformers,
# with a message that names the Hopper requirement (so the non-Hopper
# case is explicit) and the vLLM-only support.
if cc is not None and cc[0] < _HOPPER_MAJOR:
raise RuntimeError(
f"kv_cache_type 'fp8' needs a Hopper+ GPU "
f"(compute capability >= {_HOPPER_MAJOR}.0); "
f"detected {cc[0]}.{cc[1]}. fp8 KV cache is also vLLM-only — "
"the transformers backend has no fp8 KV-cache path."
)
raise RuntimeError(
"kv_cache_type 'fp8' is only available on the vLLM backend with a "
f"Hopper+ GPU (compute capability >= {_HOPPER_MAJOR}.0); the "
"transformers backend has no fp8 KV-cache path. Use --backend "
"transformers with q8_0 / bf16 / f16 instead."
)
if canonical in ("bf16", "f16"):
dtype = "bfloat16" if canonical == "bf16" else "float16"
return KvCacheRuntime(
kv_cache_type=canonical,
backend=backend_l,
model_dtype=dtype,
generate_kwargs=MappingProxyType({}),
requires_quant_backend=False,
note=(
f"KV cache runs in {dtype} (the transformers DynamicCache "
"inherits the model's compute dtype)."
),
)
# q8_0 — 8-bit quantized KV cache. transformers routes this through the
# HQQ backend (quanto only does 2/4-bit); needs `pip install hqq`.
return KvCacheRuntime(
kv_cache_type=canonical,
backend=backend_l,
model_dtype=None,
generate_kwargs=MappingProxyType({
"cache_implementation": "quantized",
"cache_config": MappingProxyType({
"backend": "hqq",
"nbits": 8,
"axis_key": 0,
"axis_value": 0,
}),
}),
requires_quant_backend=True,
note=(
"8-bit quantized KV cache via the HQQ backend "
"(install with `pip install hqq`)."
),
)

79
tests/qa/v07114_qa.md Normal file
View File

@ -0,0 +1,79 @@
# v0.71.14 — Export QA log
Patch: **v0.71.14 — "Export QA + serve(transformers) + GPU smoke finale"**
Box: Windows 11 + RTX 3050 Laptop GPU (4 GB, Ampere, compute capability 8.6)
Date: 2026-06-05
Environment:
| lib | version |
|---|---|
| torch | 2.5.1+cu121 |
| transformers | 4.57.6 |
| optimum | 2.1.0 |
| onnxruntime | 1.23.2 |
| CUDA available | yes (RTX 3050 Laptop GPU) |
---
## #71 — ONNX export pipeline — **PASS**
Verified `soup export --format onnx` end-to-end on a locally-built tiny GPT-2
(2 layers / 2 heads / n_embd=32 / vocab 512, saved with safetensors so it loads
under torch 2.5.1).
```bash
soup export --model ./model --format onnx --output ./onnx_out
```
Result:
- Export exited **0**.
- `onnx_out/model.onnx` written — **252,445 bytes**.
- `onnxruntime.InferenceSession(model.onnx, providers=["CPUExecutionProvider"])`
loaded it successfully.
- Inputs: `['input_ids', 'attention_mask', 'position_ids']` — `input_ids` present
as required by the acceptance criteria.
- One forward `sess.run(...)` on `input_ids=[[1,2,3,4]]` produced output shape
`(1, 4, 512)` = (batch, seq_len, vocab) — a real, runnable graph.
Notes:
- Tracer warnings (`aten::index` advanced indexing, masking_utils boolean
conversion) are emitted by `optimum.exporters.onnx.main_export` during the
trace; the exported graph still loads + runs. These are upstream optimum/torch
warnings, not a Soup bug.
- The historical v0.53.2 limitation (TinyLlama-1.1B ONNX export OOM during
`onnx.load(load_external_data=True)`) is a host-RAM constraint, unchanged;
the pipeline integrity is what #71 verifies, and that is confirmed here.
---
## #70 — GGUF export — **DEFERRED (infra-blocked)**
Needs a built llama.cpp C++ toolchain (the pipeline shells out to a pinned
llama.cpp tag). The maintainer's Windows 11 + RTX 3050 box cannot reliably
build the llama.cpp C++ chain — this is stated in issue #70 itself. Kept OPEN
with the `infra-blocked` label; recipe + acceptance criteria live in
`tests/qa/v053_qa.md`.
## #72 — AWQ + GPTQ export — **DEFERRED (infra-blocked)**
Needs `autoawq` + `auto-gptq` (or `gptqmodel`). None are installed and their
Windows wheels are unmaintained / build-from-source only (autoawq is archived).
Kept OPEN with the `infra-blocked` label.
## #144 — CUDA + llama.cpp manual QA smoke doc — **DEFERRED**
Depends on #70 (GGUF) and #72 (AWQ/GPTQ) actually running, both of which are
infra-blocked on this box. Kept OPEN.
## #74 — HF Hub push + Spaces deploy QA — **DEFERRED (external account)**
Needs an HF Write token + account (no credentials available on the QA box).
Previously deferred in v0.53.4 to "a contributor with private HF credentials".
Kept OPEN with the `infra-blocked` label.
## #79 — [Tracking] Community-QA meta-issue — **STAYS OPEN**
Closes only once #70–#75 all land. Since #70 / #72 / #73 / #74 / #75 remain
blocked on this box, the tracking issue stays open.

15
tests/test_v0440_part_d.py
View File

@ -197,7 +197,9 @@ def test_plan_consolidation_no_shards(tmp_path, monkeypatch):
plan_consolidation(str(out), str(tmp_path / "x.safetensors"))
def test_cli_merge_sharded(tmp_path, monkeypatch):
def test_cli_merge_sharded_plan_only(tmp_path, monkeypatch):
# v0.71.14 #96: --plan-only prints the plan without loading shards, so an
# empty placeholder .bin is fine for the planning path.
monkeypatch.chdir(tmp_path)
out = tmp_path / "shards"
out.mkdir()
@ -210,25 +212,26 @@ def test_cli_merge_sharded(tmp_path, monkeypatch):
str(out),
"-o",
str(target),
"--yes",
"--plan-only",
],
)
assert result.exit_code == 0, (result.output, repr(result.exception))
assert "Plan" in result.output
assert not target.exists()
def test_cli_merge_sharded_no_yes_warns(tmp_path, monkeypatch):
def test_cli_merge_sharded_invalid_shard_exits_2(tmp_path, monkeypatch):
# v0.71.14 #96: live consolidation of a non-torch .bin fails gracefully.
monkeypatch.chdir(tmp_path)
out = tmp_path / "shards"
out.mkdir()
(out / "pytorch_model_fsdp_0.bin").write_bytes(b"")
(out / "pytorch_model_fsdp_0.bin").write_bytes(b"not a torch checkpoint")
target = tmp_path / "merged.safetensors"
result = runner.invoke(
app,
["merge-sharded-fsdp-weights", str(out), "-o", str(target)],
)
assert result.exit_code == 0
assert "--yes" in result.output
assert result.exit_code == 2
# --- delinearize_llama4 ------------------------------------------------------

14
tests/test_v0530.py
View File

@ -249,11 +249,17 @@ class TestKVCache:
with pytest.raises(Exception):
spec.requires_hopper = False # type: ignore[misc]
def test_apply_deferred(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
def test_apply_live_v07114(self):
# v0.71.14 #140 lifted the v0.53.1 stub: apply_kv_cache_type now
# returns a runtime plan for the transformers backend.
from soup_cli.utils.kv_cache import KvCacheRuntime, apply_kv_cache_type
with pytest.raises(NotImplementedError, match="v0.53.1"):
apply_kv_cache_type()
rt = apply_kv_cache_type("bf16", backend="transformers")
assert isinstance(rt, KvCacheRuntime)
assert rt.model_dtype == "bfloat16"
# vLLM / SGLang remain deferred (blocked tail).
with pytest.raises(NotImplementedError):
apply_kv_cache_type("q8_0", backend="vllm")
def test_metadata_immutable(self):
from soup_cli.utils.kv_cache import _KV_CACHE_METADATA

616
tests/test_v07114.py Normal file
View File

@ -0,0 +1,616 @@
"""v0.71.14 — Export QA + serve(transformers) + GPU smoke finale.
Closes the doable subset of the finale patch on the Windows + RTX 3050 4 GB box:
* **#96** — live torch runtime for ``soup merge-sharded-fsdp-weights``
(``consolidate_shards`` + ``ConsolidationResult``).
* **#140** — live ``kv_cache_type`` wiring on the **transformers** backend
(``apply_kv_cache_type`` lifts the v0.53.1 ``NotImplementedError`` stub;
``--kv-cache-type`` CLI flag on ``soup serve``).
* **#71** — ONNX export QA (recorded in ``tests/qa/v07114_qa.md``).
Infra/tooling-blocked items (#70 GGUF / #72 AWQ-GPTQ / #144 doc / #74 HF push /
#79 meta) stay OPEN with ``infra-blocked`` labels — see the patch notes.
"""
from __future__ import annotations
import dataclasses
import os
import sys
from pathlib import Path
from types import MappingProxyType
import pytest
from typer.testing import CliRunner
from soup_cli.cli import app
runner = CliRunner()
_POSIX_ONLY = pytest.mark.skipif(
sys.platform == "win32", reason="POSIX symlink semantics"
)
# =====================================================================
# #96 — live FSDP shard consolidation
# =====================================================================
def _write_shard(path: Path, tensors: dict) -> None:
import torch
torch.save(tensors, str(path))
class TestConsolidateShards:
def test_merges_two_disjoint_shards(self, tmp_path, monkeypatch):
import torch
from safetensors.torch import load_file
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(
shards / "pytorch_model_fsdp_0.bin",
{"model.layers.0.weight": torch.ones(2, 3)},
)
_write_shard(
shards / "pytorch_model_fsdp_1.bin",
{"model.layers.1.weight": torch.zeros(4, 5)},
)
target = tmp_path / "merged.safetensors"
plan = plan_consolidation(str(shards), str(target))
result = consolidate_shards(plan)
assert result.num_tensors == 2
assert result.num_shards == 2
assert result.total_bytes > 0
assert os.path.exists(target)
loaded = load_file(str(target))
assert set(loaded) == {"model.layers.0.weight", "model.layers.1.weight"}
assert loaded["model.layers.0.weight"].shape == (2, 3)
def test_streams_one_shard_at_a_time(self, tmp_path, monkeypatch):
"""Memory-friendly: torch.load is called exactly once per shard."""
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {"a": torch.ones(2)})
_write_shard(shards / "pytorch_model_fsdp_1.bin", {"b": torch.ones(2)})
target = tmp_path / "merged.safetensors"
plan = plan_consolidation(str(shards), str(target))
calls = {"n": 0}
real_load = torch.load
def _counting_load(*args, **kwargs):
calls["n"] += 1
return real_load(*args, **kwargs)
# consolidate_shards lazy-imports torch and calls torch.load — patch
# the real module so the count reflects per-shard streaming.
monkeypatch.setattr(torch, "load", _counting_load)
consolidate_shards(plan)
assert calls["n"] == 2
def test_shape_conflict_rejected(self, tmp_path, monkeypatch):
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {"w": torch.ones(2, 3)})
_write_shard(shards / "pytorch_model_fsdp_1.bin", {"w": torch.ones(4, 5)})
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
with pytest.raises(ValueError, match="shape"):
consolidate_shards(plan)
def test_duplicate_key_same_shape_kept_once(self, tmp_path, monkeypatch):
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {"w": torch.ones(2, 3)})
_write_shard(shards / "pytorch_model_fsdp_1.bin", {"w": torch.ones(2, 3)})
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
result = consolidate_shards(plan)
assert result.num_tensors == 1
def test_duplicate_key_warns(self, tmp_path, monkeypatch, caplog):
"""A same-shape duplicate across shards logs a WARNING (no silent drop)."""
import logging
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {"w": torch.ones(2, 3)})
_write_shard(shards / "pytorch_model_fsdp_1.bin", {"w": torch.zeros(2, 3)})
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
with caplog.at_level(
logging.WARNING, logger="soup_cli.utils.fsdp_consolidate"
):
consolidate_shards(plan)
assert any(
"more than one shard" in rec.message for rec in caplog.records
)
def test_non_dict_shard_rejected(self, tmp_path, monkeypatch):
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
# A bare tensor, not a state-dict.
torch.save(torch.ones(3), str(shards / "pytorch_model_fsdp_0.bin"))
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
with pytest.raises(ValueError, match="state-dict"):
consolidate_shards(plan)
def test_empty_merged_rejected(self, tmp_path, monkeypatch):
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {})
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
with pytest.raises(ValueError, match="no tensors"):
consolidate_shards(plan)
def test_result_frozen(self, tmp_path, monkeypatch):
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {"a": torch.ones(2)})
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
result = consolidate_shards(plan)
with pytest.raises(dataclasses.FrozenInstanceError):
result.num_tensors = 99 # type: ignore[misc]
def test_non_plan_rejected(self):
from soup_cli.utils.fsdp_consolidate import consolidate_shards
with pytest.raises(TypeError):
consolidate_shards({"not": "a plan"}) # type: ignore[arg-type]
@_POSIX_ONLY
def test_symlinked_shard_rejected(self, tmp_path, monkeypatch):
# A symlinked shard child could redirect torch.load to an arbitrary
# target — consolidate_shards must reject it (TOCTOU defence). The
# output-symlink case is handled earlier: plan_consolidation realpaths
# the output and rejects an out-of-cwd target as "outside cwd".
import torch
from soup_cli.utils.fsdp_consolidate import (
consolidate_shards,
plan_consolidation,
)
monkeypatch.chdir(tmp_path)
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(shards / "pytorch_model_fsdp_0.bin", {"a": torch.ones(2)})
# A second "shard" that is actually a symlink (to a real under-cwd file).
real_target = tmp_path / "secret.bin"
_write_shard(real_target, {"b": torch.ones(2)})
os.symlink(real_target, shards / "pytorch_model_fsdp_1.bin")
plan = plan_consolidation(str(shards), str(tmp_path / "m.safetensors"))
with pytest.raises(ValueError, match="symlink"):
consolidate_shards(plan)
def test_source_uses_atomic_write_bytes(self):
src = (
Path(__file__).resolve().parent.parent
/ "src"
/ "soup_cli"
/ "utils"
/ "fsdp_consolidate.py"
).read_text(encoding="utf-8")
assert "atomic_write_bytes" in src
assert "weights_only=True" in src
class TestMergeShardedCli:
def _make_shards(self, tmp_path):
import torch
shards = tmp_path / "shards"
shards.mkdir()
_write_shard(
shards / "pytorch_model_fsdp_0.bin",
{"model.layers.0.weight": torch.ones(2, 2)},
)
_write_shard(
shards / "pytorch_model_fsdp_1.bin",
{"model.layers.1.weight": torch.zeros(2, 2)},
)
return shards
def test_live_consolidation(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
shards = self._make_shards(tmp_path)
target = tmp_path / "merged.safetensors"
result = runner.invoke(
app,
["merge-sharded-fsdp-weights", str(shards), "-o", str(target)],
)
assert result.exit_code == 0, (result.output, repr(result.exception))
assert target.exists()
assert "2" in result.output # 2 tensors / 2 shards reported
def test_plan_only(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
shards = self._make_shards(tmp_path)
target = tmp_path / "merged.safetensors"
result = runner.invoke(
app,
[
"merge-sharded-fsdp-weights",
str(shards),
"-o",
str(target),
"--plan-only",
],
)
assert result.exit_code == 0, (result.output, repr(result.exception))
assert "Plan" in result.output
assert not target.exists() # plan-only writes nothing
def test_outside_cwd_output_rejected(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
shards = self._make_shards(tmp_path)
result = runner.invoke(
app,
[
"merge-sharded-fsdp-weights",
str(shards),
"-o",
str(tmp_path.parent / "escape.safetensors"),
],
)
assert result.exit_code == 2
def test_missing_shards_rejected(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
empty = tmp_path / "empty"
empty.mkdir()
result = runner.invoke(
app,
[
"merge-sharded-fsdp-weights",
str(empty),
"-o",
str(tmp_path / "m.safetensors"),
],
)
assert result.exit_code == 2
def test_help_lists_plan_only(self):
import re
result = runner.invoke(app, ["merge-sharded-fsdp-weights", "--help"])
assert result.exit_code == 0
clean = re.sub(r"\x1b\[[0-9;]*m", "", result.output)
assert "--plan-only" in clean
# =====================================================================
# #140 — live kv_cache_type wiring (transformers backend)
# =====================================================================
class TestApplyKvCacheType:
def test_no_longer_raises_notimplemented(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("bf16", backend="transformers")
assert rt.kv_cache_type == "bf16"
def test_bf16_model_dtype(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("bf16", backend="transformers")
assert rt.model_dtype == "bfloat16"
assert dict(rt.generate_kwargs) == {}
assert rt.requires_quant_backend is False
def test_f16_model_dtype(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("f16", backend="transformers")
assert rt.model_dtype == "float16"
assert dict(rt.generate_kwargs) == {}
def test_q8_0_quantized_cache_kwargs(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("q8_0", backend="transformers")
gk = dict(rt.generate_kwargs)
assert gk["cache_implementation"] == "quantized"
cfg = gk["cache_config"]
assert cfg["backend"] == "hqq"
assert cfg["nbits"] == 8
assert rt.model_dtype is None
assert rt.requires_quant_backend is True
def test_case_insensitive(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("Q8_0", backend="transformers")
assert rt.kv_cache_type == "q8_0"
def test_fp8_non_hopper_friendly_error(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(RuntimeError, match="Hopper"):
apply_kv_cache_type(
"fp8", backend="transformers", compute_capability=(8, 6)
)
def test_fp8_hopper_still_unsupported_on_transformers(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(RuntimeError, match="vLLM"):
apply_kv_cache_type(
"fp8", backend="transformers", compute_capability=(9, 0)
)
def test_fp8_unknown_cc_raises(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(RuntimeError):
apply_kv_cache_type(
"fp8", backend="transformers", compute_capability=None
)
def test_vllm_backend_deferred(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(NotImplementedError, match="vLLM|transformers"):
apply_kv_cache_type("q8_0", backend="vllm")
def test_sglang_backend_deferred(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(NotImplementedError):
apply_kv_cache_type("bf16", backend="sglang")
def test_invalid_type_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(ValueError, match="not supported"):
apply_kv_cache_type("wat", backend="transformers")
def test_oversize_type_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(ValueError, match="too long"):
apply_kv_cache_type("x" * 17, backend="transformers")
def test_null_byte_type_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(ValueError, match="null byte"):
apply_kv_cache_type("bf\x0016", backend="transformers")
def test_bool_type_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(TypeError):
apply_kv_cache_type(True, backend="transformers") # type: ignore[arg-type]
def test_bool_backend_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(TypeError):
apply_kv_cache_type("bf16", backend=True) # type: ignore[arg-type]
def test_non_str_backend_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises(TypeError):
apply_kv_cache_type("bf16", backend=123) # type: ignore[arg-type]
def test_bad_compute_capability_shape_rejected(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
with pytest.raises((TypeError, ValueError)):
apply_kv_cache_type(
"fp8", backend="transformers", compute_capability=(9,) # type: ignore[arg-type]
)
def test_runtime_frozen(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("bf16", backend="transformers")
with pytest.raises(dataclasses.FrozenInstanceError):
rt.kv_cache_type = "f16" # type: ignore[misc]
def test_generate_kwargs_immutable(self):
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("q8_0", backend="transformers")
assert isinstance(rt.generate_kwargs, MappingProxyType)
with pytest.raises(TypeError):
rt.generate_kwargs["x"] = 1 # type: ignore[index]
def test_quantized_cache_backend_available(self):
from soup_cli.utils.kv_cache import quantized_cache_backend_available
# On this box neither hqq nor quanto is installed → None. The return
# type is Optional[str]; assert it's None-or-str without asserting a
# specific value (CI may differ).
got = quantized_cache_backend_available()
assert got is None or isinstance(got, str)
class TestPlainKvKwargs:
def test_deep_converts_nested_mappingproxy(self):
from soup_cli.commands.serve import _plain_kv_kwargs
nested = MappingProxyType(
{"cache_implementation": "quantized",
"cache_config": MappingProxyType({"backend": "hqq", "nbits": 8})}
)
out = _plain_kv_kwargs(nested)
assert isinstance(out, dict)
assert isinstance(out["cache_config"], dict)
# Plain dicts are mutable — proves the proxy was unwrapped.
out["cache_config"]["nbits"] = 4
assert out == {
"cache_implementation": "quantized",
"cache_config": {"backend": "hqq", "nbits": 4},
}
def test_empty_mapping(self):
from soup_cli.commands.serve import _plain_kv_kwargs
assert _plain_kv_kwargs(MappingProxyType({})) == {}
def test_q8_0_runtime_round_trips_to_plain(self):
from soup_cli.commands.serve import _plain_kv_kwargs
from soup_cli.utils.kv_cache import apply_kv_cache_type
rt = apply_kv_cache_type("q8_0", backend="transformers")
plain = _plain_kv_kwargs(rt.generate_kwargs)
assert plain["cache_implementation"] == "quantized"
assert plain["cache_config"]["backend"] == "hqq"
assert plain["cache_config"]["nbits"] == 8
class TestServeKvCacheCli:
def test_help_lists_flag(self):
import re
result = runner.invoke(app, ["serve", "--help"])
assert result.exit_code == 0
clean = re.sub(r"\x1b\[[0-9;]*m", "", result.output)
assert "--kv-cache-type" in clean
def test_invalid_type_exit_2(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
model = tmp_path / "model"
model.mkdir()
result = runner.invoke(
app,
["serve", "--model", str(model), "--kv-cache-type", "wat"],
)
assert result.exit_code == 2
def test_fp8_on_ampere_exit_2(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
model = tmp_path / "model"
model.mkdir()
# Force a non-Hopper compute capability so the friendly Hopper error
# fires regardless of the host GPU.
import torch
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: (8, 6))
result = runner.invoke(
app,
["serve", "--model", str(model), "--kv-cache-type", "fp8"],
)
assert result.exit_code == 2
def test_vllm_backend_deferred_exit_2(self, tmp_path, monkeypatch):
monkeypatch.chdir(tmp_path)
model = tmp_path / "model"
model.mkdir()
result = runner.invoke(
app,
[
"serve",
"--model",
str(model),
"--backend",
"vllm",
"--kv-cache-type",
"q8_0",
],
)
assert result.exit_code == 2
# =====================================================================
# #71 — ONNX export QA (recorded in the QA log)
# =====================================================================
class TestOnnxQaLog:
def test_qa_log_exists(self):
log = (
Path(__file__).resolve().parent
/ "qa"
/ "v07114_qa.md"
)
assert log.exists(), "v07114 QA log missing"
body = log.read_text(encoding="utf-8")
assert "#71" in body
assert "onnx" in body.lower()
class TestPatchInvariants:
def test_version_bumped(self):
import soup_cli
parts = tuple(int(x) for x in soup_cli.__version__.split(".")[:3])
assert parts >= (0, 71, 14)
def test_no_top_level_torch_in_kv_cache(self):
src = (
Path(__file__).resolve().parent.parent
/ "src"
/ "soup_cli"
/ "utils"
/ "kv_cache.py"
).read_text(encoding="utf-8")
for line in src.splitlines():
assert not line.startswith("import torch")
assert not line.startswith("from torch")