feat: auto-size context to box VRAM + KV-cache quant + RTX PRO 6000 tier

- --context now defaults to 'auto': fill the rented box's spare VRAM
  instead of guessing or defaulting low. Working-VRAM headroom scales
  with GPU count (per-device), so it isn't biased toward big boxes.
- --kv-quant (q8_0/q4_0) quantizes the llama.cpp KV cache; sizing
  accounts for the saving so a longer context still finds a fitting box.
- Add the RTX PRO 6000 (96GB) GPU tier — previously invisible to the
  sizer, so big GGUF models skipped it for pricier 8x/H200 boxes.

Author: nathanmarlor

aiod/cli.py

@@ -22,7 +22,7 @@
 from .bootstrap import CONTAINER_PORT, ServerConfig
 from .config import Settings
 from .health import wait_until_ready
-from .sizing import QUANT_LABELS, size_any
+from .sizing import QUANT_LABELS, auto_context_for_offer, size_any
 from .vast import PricedOption, recommend_disk_gb
 
 app = typer.Typer(
@@ -317,7 +317,14 @@ def spin(
     idle: int = typer.Option(
         None, "--idle", help="Auto-shutdown after N idle minutes (starts a local watcher)"
     ),
-    context: int = typer.Option(None, "--context", help="Max model length to serve"),
+    context: str = typer.Option(
+        "auto", "--context",
+        help="Max model length: an int, or 'auto' to fill the rented box's spare VRAM",
+    ),
+    kv_quant: str = typer.Option(
+        None, "--kv-quant",
+        help="GGUF only: quantize the KV cache (q8_0/q4_0) to fit a longer context",
+    ),
     concurrency: int = typer.Option(None, "--concurrency", help="Concurrency for KV sizing"),
     yes: bool = typer.Option(False, "--yes", "-y", help="Skip the confirmation prompt"),
     no_ccr: bool = typer.Option(False, "--no-ccr", help="Don't touch the CCR config"),
@@ -343,7 +350,12 @@ def spin(
     quant = quant or (prof.quant if prof else None)
     provider = (provider or (prof.provider if prof else "vast")).lower()
     _require_provider_key(s, provider)
-    context = context if context is not None else (prof.context if prof else None)
+    # Context: explicit int wins; 'auto' (the default) defers to a pinned profile
+    # value, else fills the box's spare VRAM once the offer is known (GGUF only).
+    if context == "auto" and prof and prof.context is not None:
+        context = str(prof.context)
+    auto_context = context == "auto"
+    context = None if auto_context else int(context)
     concurrency = concurrency if concurrency is not None else (prof.concurrency if prof else 4)
     max_p = (
         max_price if max_price is not None
@@ -374,8 +386,12 @@ def spin(
         sizing = size_any(
             model, engine=engine, hf_token=s.hf_token,
             quants=[quant] if quant else None, context_len=context, concurrency=concurrency,
+            kv_quant=kv_quant,
         )
     eng = sizing.engine
+    if kv_quant:
+        # Quantize the llama.cpp KV cache (sizing already accounts for the saving).
+        extra_args += ["--cache-type-k", kv_quant, "--cache-type-v", kv_quant]
     m = sizing.model
 
     if quant is None:
@@ -406,6 +422,19 @@ def spin(
             raise typer.Exit(1)
 
         offer = best.offer
+
+        # Auto-context: now that a box is chosen, hand its spare VRAM to the context
+        # window. GGUF only — vLLM keeps deferring to the model's own max (None).
+        ctx_note = ""
+        if auto_context and eng == "llamacpp":
+            context = auto_context_for_offer(
+                offer.total_vram_gb, plan.weights_gb, kv_quant,
+                num_gpus=best.option.num_gpus,
+            )
+            ctx_note = f"\nContext: [bold]{context:,}[/] tokens (auto — fills spare VRAM)"
+        elif not auto_context:
+            ctx_note = f"\nContext: {context:,} tokens"
+
         console.print(
             Panel(
                 f"[bold]{m.repo_id}[/]  ·  {quant} ({QUANT_LABELS.get(quant, quant)})\n"
@@ -414,7 +443,8 @@ def spin(
                 f"Disk:  {disk} GB\n"
                 f"Price: [bold]${offer.dph_total:.2f}/hr[/]  (~${offer.dph_total * ttl_h:.2f} "
                 f"over a {ttl_h:g}h session)\n"
-                f"VRAM:  need ~{plan.required_vram_gb:.0f} GB / have {offer.total_vram_gb:.0f} GB",
+                f"VRAM:  need ~{plan.required_vram_gb:.0f} GB / have {offer.total_vram_gb:.0f} GB"
+                f"{ctx_note}",
                 title="Launch plan",
                 border_style="cyan",
                 expand=False,

aiod/sizing.py

@@ -69,6 +69,7 @@ class GpuTier:
     GpuTier("A100 80GB", 80, ["A100 SXM4 80GB", "A100 PCIE 80GB", "A100_SXM4_80GB", "A100"]),
     GpuTier("H100 80GB", 80, ["H100 SXM", "H100 PCIE", "H100_SXM", "H100"]),
     GpuTier("H100 NVL 94GB", 94, ["H100 NVL", "H100_NVL"]),
+    GpuTier("RTX PRO 6000 96GB", 96, ["RTX PRO 6000", "RTX 6000 Pro", "RTX PRO 6000 WS", "RTX6000Pro"]),
     GpuTier("H200 141GB", 141, ["H200"]),
     GpuTier("B200 180GB", 180, ["B200"]),
 ]
@@ -397,10 +398,63 @@ def detect_format(repo: str, hf_token: str | None = None, timeout: float = 20.0)
     return "unknown"
 
 
+# KV-cache size relative to f16, by llama.cpp --cache-type. Quantizing the KV
+# cache (q8_0/q4_0) is the main lever for fitting a longer context on the same box.
+KV_CACHE_FACTOR = {
+    None: 1.0, "f16": 1.0, "bf16": 1.0,
+    "q8_0": 0.5,
+    "q5_0": 0.34, "q5_1": 0.34,
+    "q4_0": 0.27, "q4_1": 0.27,
+}
+
+# llama.cpp needs working VRAM beyond weights+KV (compute graph, fragmentation).
+# This is per-device, so it scales with GPU count, not model size — held back when
+# auto-sizing context so a full-to-the-brim estimate doesn't OOM. Kept modest
+# because the KV estimate itself already over-counts (see size_gguf_model).
+GGUF_PER_GPU_BUFFER_GB = 2.0
+GGUF_MIN_BUFFER_GB = 3.0
+
+
+def context_for_vram(kv_budget_gb: float, kv_quant: str | None = None) -> int:
+    """Inverse of the GGUF KV model: how many tokens a VRAM budget buys.
+
+    `size_gguf_model` models f16 KV as ~8 GB per 8192 tokens; this reverses it,
+    scaled by the same `--cache-type` factor.
+    """
+    factor = KV_CACHE_FACTOR.get(kv_quant, 1.0)
+    if kv_budget_gb <= 0:
+        return 0
+    return int(kv_budget_gb / (8.0 * factor) * 8192)
+
+
+def auto_context_for_offer(
+    total_vram_gb: float,
+    weights_gb: float,
+    kv_quant: str | None = None,
+    *,
+    num_gpus: int = 1,
+    round_to: int = 4096,
+    floor: int = 8192,
+    cap: int = 262144,
+) -> int:
+    """Max context that fits a rented box's spare VRAM after weights + headroom.
+
+    Uses the offer's real VRAM (not the GPU-tier minimum), so a 98 GB card isn't
+    sized as 96. The working-VRAM headroom scales with `num_gpus` (it's a
+    per-device cost), so this isn't biased toward big multi-GPU boxes. Rounded
+    down to `round_to`, clamped to [`floor`, `cap`].
+    """
+    buffer_gb = max(GGUF_MIN_BUFFER_GB, GGUF_PER_GPU_BUFFER_GB * num_gpus)
+    kv_budget = total_vram_gb - weights_gb * 1.02 - buffer_gb
+    ctx = (context_for_vram(kv_budget, kv_quant) // round_to) * round_to
+    return max(floor, min(ctx, cap))
+
+
 def size_gguf_model(
     repo: str,
     hf_token: str | None = None,
     context_len: int | None = None,
+    kv_quant: str | None = None,
 ) -> SizingResult:
     """Size every GGUF quant in a repo by file size and produce a GPU plan each."""
     repo = parse_repo_id(repo)
@@ -410,7 +464,9 @@ def size_gguf_model(
 
     ctx = context_len or 8192
     # llama.cpp KV cache scales with context; rough budget (no per-model config here).
-    kv_gb = max(4.0, (ctx / 8192.0) * 8.0)
+    # A quantized KV cache (--cache-type-k/v) shrinks it by the factor above.
+    kv_factor = KV_CACHE_FACTOR.get(kv_quant, 1.0)
+    kv_gb = max(4.0, (ctx / 8192.0) * 8.0) * kv_factor
 
     plans: list[QuantPlan] = []
     for q in sorted(quants.values(), key=lambda x: x.size_gb):
@@ -451,6 +507,7 @@ def size_any(
     quants: list[str] | None = None,
     context_len: int | None = None,
     concurrency: int = 4,
+    kv_quant: str | None = None,
 ) -> SizingResult:
     """Engine-aware entry point: auto-detect GGUF vs safetensors (or force via
     `engine`) and return the matching SizingResult."""
@@ -459,7 +516,7 @@ def size_any(
     if engine == "auto":
         eng = "llamacpp" if detect_format(repo, hf_token=hf_token) == "gguf" else "vllm"
     if eng == "llamacpp":
-        return size_gguf_model(repo, hf_token=hf_token, context_len=context_len)
+        return size_gguf_model(repo, hf_token=hf_token, context_len=context_len, kv_quant=kv_quant)
     return size_model(
         repo, hf_token=hf_token, quants=quants, context_len=context_len, concurrency=concurrency
     )

tests/test_sizing.py

@@ -1,6 +1,8 @@
 from aiod.sizing import (
     ModelSpec,
     _params_from_name,
+    auto_context_for_offer,
+    context_for_vram,
     estimate_vram,
     parse_repo_id,
     plan_gpus,
@@ -68,3 +70,41 @@ def test_small_model_fits_one_gpu():
     _, _, req = estimate_vram(spec, "bf16", context_tokens=8192)
     opts = {o.tier.name: o for o in plan_gpus(req)}
     assert opts["A100 80GB"].num_gpus == 1
+
+
+def test_context_for_vram_scales_with_kv_quant():
+    # q4_0 KV is ~0.27x f16, so the same VRAM buys ~3.7x the tokens.
+    f16 = context_for_vram(16.0, None)
+    q4 = context_for_vram(16.0, "q4_0")
+    assert f16 == 16384  # 16 GB / 8 GB-per-8k * 8192
+    assert q4 > f16 * 3
+    assert context_for_vram(0.0) == 0
+
+
+def test_auto_context_fills_offer_vram():
+    # Tight box: weights 342.7 GB on a 392 GB box leaves little for f16 KV,
+    # but q4_0 stretches it well past the f16 result.
+    f16 = auto_context_for_offer(392.0, 342.7, None, num_gpus=4)
+    q4 = auto_context_for_offer(392.0, 342.7, "q4_0", num_gpus=4)
+    assert q4 > f16 > 0
+    assert f16 % 4096 == 0 and q4 % 4096 == 0  # rounded to multiple
+
+
+def test_auto_context_headroom_scales_with_gpu_count():
+    # The working-VRAM buffer is per-device, so more GPUs reserve more and
+    # leave less for context on an otherwise identical budget.
+    few = auto_context_for_offer(360.0, 300.0, None, num_gpus=1)
+    many = auto_context_for_offer(360.0, 300.0, None, num_gpus=8)
+    assert few > many
+
+
+def test_auto_context_small_model_not_starved_by_buffer():
+    # A 7B-ish Q4 (~4.6 GB) on a single 24 GB card should get a real window,
+    # not get floored because a fat flat buffer ate the spare VRAM.
+    assert auto_context_for_offer(24.0, 4.6, None, num_gpus=1) > 8192
+
+
+def test_auto_context_clamps_to_cap_and_floor():
+    # Huge box -> cap; box that can't even hold weights -> floor.
+    assert auto_context_for_offer(2000.0, 342.7, "q4_0", num_gpus=4) == 262144
+    assert auto_context_for_offer(350.0, 342.7, None, num_gpus=4) == 8192