Make the compiler memory headroom a preparation setting

Forward compilation and kernel tuning size their worker pools against host
memory and always leave a fixed 8 GiB free. On a unified-memory GPU such as
DGX Spark (GB10), the model's weights already occupy most of that memory: a
DeepSeek V4.1 Flash TP3 rank has about 13 GiB left after loading, and with
the fixed floor the sizing rule (4 GiB per compiler, 512 MiB per Inductor
thread) allows one forward compiler with two Inductor threads per rank, so
"Compile model forwards" runs mostly serially. The right floor depends on
the host and the model, so it should be a setting rather than a constant.
Lowering it keeps the existing safety rule: on unified-memory GPUs, sizing
still leaves the greater of this floor and the preparation headroom that
protects KV capacity and later processes.

Before this change the floor was the private constant `_HEADROOM_BYTES` in
`ignition/preparation/host_memory.py`. It is now `COMPILE_HEADROOM_BYTES`
(still 8 GiB) used as the default of a `compile_headroom_bytes` argument in
`worker_limit`, `compile_threads`, `tuning_limits`,
`candidate_scratch_limit` and `CompilePool`. The value comes from the new
`[preparation] compile_headroom` setting (a size string such as "6G",
default "8G") or the advanced `--compile-headroom SIZE` flag of `export`,
`debug-serve` and `bench`, which overrides the file. It is carried as
`Plan.compile_headroom_bytes`, passed through every recipe's
`PreparationConfig` to the tuning session, `ForwardCache` and `CompilePool`,
and sent from rank zero to multi-host peers, as `compile_workers` is. The
plan dictionary, the "Select GPUs." line, the tuning report and the bench
result record it. Negative, empty and unparsable values are usage errors.
The low-memory threshold in `release_unused` keeps 8 GiB. When a pool has no
room for one worker, its MemoryError now also reports the headroom and
per-worker bytes. docs/cli-design.md and docs/model-preparation.md document
the setting and flag. With the default, pool sizing is unchanged.

Measured result: with a cold forward cache and restored tuning on three DGX
Spark GB10 hosts (DeepSeek V4.1 Flash TP3), "Compile model forwards" took
507.7 s with `--compile-headroom 6G` and 511.2 s with `[preparation]
compile_headroom = "6G"`, against 905.7 s with the default. By the sizing
rule, at 13 GiB available a 6 GiB floor gives each rank one compiler with
six Inductor threads instead of two; the export logs do not record the
thread count. This is a component result for one model and host type,
measured once per setting with this change applied to main at `ca1f960e`,
184 commits before `d2b5c5f3`; those commits change how `ForwardCache`
releases memory, so the timing may differ on `d2b5c5f3`. Kernel tuning was
restored from cache, so tuning pools and candidate scratch sized with a
lower headroom were not exercised on hardware. A host with less memory left
after loading can exhaust host memory with a low value.

Evidence: the plans of those three exports recorded 8 GiB, 6 GiB and 6 GiB.
With 6 GiB the lowest MemAvailable on any host during the phase was 9 GiB,
with no OOM. The 6 GiB package served with decode speed and 262k-token
prefill time in line with two earlier exports of the same checkpoint and
recipe with the default headroom, and answered a 262k-token prompt that asks
for 16 values placed through the context with all 16 correct. On `d2b5c5f3`
plus this commit, the CPU suites tests/cli, tests/preparation,
tests/recipes/test_declarations.py, tests/recipes/test_minimax_m3.py and
tests/recipes/muse_spark (CUDA hidden, environment from uv.lock) gave 2063
passed, 65 failed and 222 skipped; `d2b5c5f3` alone fails the same 65 tests
in the same environment (2016 passed, 222 skipped). Those failures come from
a locally built extension without `validate_prepared_contract`,
chat-template fixtures that were not prepared, a missing Kimi K3 checkpoint,
and CLI and multi-host routing expectations that also fail on `d2b5c5f3`.
New tests cover setting and flag parsing and validation, the plan and
startup line, every recipe's `PreparationConfig` (Kimi K3.1 Pro reuses the
Kimi K3 adapter), `worker_limit`, `compile_threads`, `tuning_limits`,
`candidate_scratch_limit`, `CompilePool` sizing and the tuning report. The
export measurement was not repeated on `d2b5c5f3`, and no GPU tests were run
on it.

Design-Signed-Off-By: Claude Opus 5.5; design-review=3ad6
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Javier López <59578422+javieratomic@users.noreply.github.com>

Author: javieratomic

docs/cli-design.md

@@ -419,6 +419,7 @@ explicit flags take precedence.
 | `--collective-tuning static\|runtime` | Static autotunes collectives during preparation; runtime uses the required ladder file for final compiled geometry and captures runtime submode/active-CTA selection for eligible BF16 PCIe reductions. Overrides `[preparation] collective_tuning`, then the recipe declaration. | Static |
 | `--lock-gpu-clocks` | Lock each tuning GPU at its maximum supported graphics frequency during b12x and Inductor autotuning, then reset driver clock defaults. Frozen warmup, capture, and serving remain unlocked. Requires root or noninteractive sudo access to `nvidia-smi`. Command line only; inactive with `--autotune-effort off`. | off |
 | `--compile-workers N` | CPU compiler process budget per host, divided among local preparation ranks and bounded by available RAM and device memory. Applies to b12x kernel compilation and complete-forward compilation. Overrides `[preparation] compile_workers`. | Automatic |
+| `--compile-headroom SIZE` | System memory that compiler processes and kernel-tuning races leave free on each host, such as `6G`; binary size suffixes as in `[memory] reserve`. On unified-memory GPUs, a larger preparation headroom from the plan and `[memory] reserve` still applies. Lowering it permits more compiler processes and Inductor threads on unified-memory GPUs such as DGX Spark, at the risk of exhausting host memory. Recorded in the plan, the tuning report and the `bench` result. Overrides `[preparation] compile_headroom`. | `8G` |
 | `--tuning-gpus LIST` | Local kernel-tuning GPUs. Start with the serving GPU list, then append complete replicas of that TP group on matching hardware. Serving TP size stays fixed. Command line only; cannot be combined with multi-host `--ranks`. | Serving GPUs |
 | `--no-compile` | Development only: capture graphs without compiling forwards. A recipe never sets it. | off |
 | `--force-moe-decode-a16` / `--no-force-moe-decode-a16` | Prepare Qwen3.8 Flash NVFP4 target routed experts as one packed weight owner: decode and speculative verification use W4A16 at every token-row count, and separate prefill graphs use calibrated A4. Requires SM120/121, checkpoint activation scales, and rank-local expert intermediate width divisible by 64 (TP1/2). Mixed prefill and other recipes/precisions reject the flag. Shared experts and stored A16 MTP experts retain their precision. Overrides `[preparation] force_moe_decode_a16`; the resolved value is recorded in the plan and package rebuild metadata. | off |
@@ -1192,6 +1193,7 @@ max_concurrency = 4
 [preparation]
 # force_moe_decode_a16 = false    # Qwen Flash NVFP4: A16 decode/verify, calibrated A4 prefill.
 # compile_workers = 24           # Optional CPU process budget per host.
+# compile_headroom = "8G"        # System memory compiler and tuning pools leave free per host.
 # autotune_effort = "default"     # off, default, audit.
 # collective_tuning = "static"    # static, runtime; omitted uses the recipe default.
 
@@ -1280,6 +1282,7 @@ Configuration assigns each decision to one owner:
 | Queue, frontend workers, event capacity, body limit | Concurrency and recipe image limits |
 | Compiler workers | Automatic memory-bounded sizing; `[preparation] compile_workers` or `--compile-workers` overrides |
 | Collective selection | `--collective-tuning`, then `[preparation] collective_tuning`, otherwise the recipe default (`static`) |
+| Compiler memory headroom | 8 GiB of system memory per host; `[preparation] compile_headroom` or `--compile-headroom` overrides |
 | Kernel search effort | `--autotune-effort`, then `[preparation] autotune_effort`, otherwise `default` |
 | Local tuning replicas | `--tuning-gpus`; independent of serving TP size |
 | CUDA toolkit | `[paths] cuda`, otherwise native build libraries discover the system toolkit |

docs/model-preparation.md

@@ -245,6 +245,15 @@ device memory can reduce it; local preparation ranks share the budget. Set
 overrides that setting. Omitting both retains automatic sizing, including
 shared-memory limits on DGX Spark.
 
+Automatic sizing leaves 8 GiB of system memory free beside compiler processes
+and tuning races. On unified-memory GPUs, the plan's preparation headroom
+applies instead when it is larger. On DGX Spark, the memory left after loading
+weights bounds both the compiler count and each compiler's Inductor threads.
+`--compile-headroom SIZE`, or `[preparation] compile_headroom = "SIZE"`,
+selects another value, such as `6G`; the flag overrides the setting. A smaller
+headroom compiles faster when memory is tight but leaves less for the rest of
+the host. The plan and tuning report record the value.
+
 On integrated GPUs, kernel tuning is divided into memory-bounded jobs. The
 terminal dashboard shows the active job number and total job count, with a
 separate live display for each job. Redirecting output records plain kernel

ignition/cli/__init__.py

@@ -232,6 +232,9 @@ def _preparation_arguments(running, *, model_default="defaults.model"):
     advanced.add_argument("--compile-workers", type=int, metavar="N",
                           help="CPU compiler process budget per host, shared by local preparation ranks;"
                                " bounded by available RAM and device memory (default: automatic)")
+    advanced.add_argument("--compile-headroom", metavar="SIZE",
+                          help="system memory compiler and tuning pools leave free on each host, such as 6G;"
+                               " a larger preparation headroom still applies (default: 8G)")
     advanced.add_argument("--tuning-gpus", metavar="LIST",
                           help="local GPUs for kernel tuning, starting with --gpus and followed by complete"
                                " TP replicas; serving TP size is unchanged (default: serving GPUs)")

ignition/cli/bench.py

@@ -149,7 +149,7 @@ def run(args):
     plan = derive(configuration, model, Host.discover(), gpus=args.gpus, concurrency=concurrency,
         speculation=speculation(args, model, configuration), prefill_chunk=args.prefill_chunk_size,
         context_parallel=context_parallel(args, model), stream_overlap=stream_overlap(args, model),
-        compile_workers=args.compile_workers, tuning_gpus=args.tuning_gpus,
+        compile_workers=args.compile_workers, compile_headroom=args.compile_headroom, tuning_gpus=args.tuning_gpus,
         autotune_effort=args.autotune_effort, collective_tuning=args.collective_tuning, collective_ladder=args.collective_ladder,
         lock_gpu_clocks=args.lock_gpu_clocks, native_dense_activations=args.native_dense_activations,
         force_moe_decode_a16=args.force_moe_decode_a16,
@@ -225,6 +225,7 @@ def _measure(model, plan, configuration, settings, progress, *, compile_forwards
         "preparation": {key: preparation[key] for key in (
             "tensor_parallel", "prefill_chunk", "stream_overlap", "context_parallel",
             "native_dense_activations", "force_moe_decode_a16", "autotune_effort", "collective_tuning", "collective_ladder", "compile_workers",
+            "compile_headroom_bytes",
             "speculative_token_capacity", "speculation", "expert_scale_compression", "mla_compute_mode",
             "online_mxfp8", "online_mxfp8_scope", "language_model_only")}}
 

ignition/cli/config.py

@@ -23,7 +23,7 @@
     "defaults": {"model": str, "max_concurrency": int, "image_concurrency": int,
                  "speculative_token_capacity": int, "speculative_graph_budget": int, "speculative_method": str},
     "memory": {"reserve": str},
-    "preparation": {"compile_workers": int, "autotune_effort": str, "collective_tuning": str, "checkpoint_token_file": str, "force_moe_decode_a16": bool},
+    "preparation": {"compile_workers": int, "compile_headroom": str, "autotune_effort": str, "collective_tuning": str, "checkpoint_token_file": str, "force_moe_decode_a16": bool},
     "prefix_cache": {"memory": str, "disk_limit": str, "drain_on_shutdown": str, "max_age": str},
     "paths": {"hub_cache": str, "kernel_cache": str, "package_cache": str, "export_staging": str,
               "prefix_cache": str, "cuda": str},
@@ -37,6 +37,7 @@
     "memory.reserve": "0",
     "preparation.force_moe_decode_a16": False,
     "preparation.compile_workers": 0,
+    "preparation.compile_headroom": "8G",
     "preparation.autotune_effort": "default",
     "preparation.collective_tuning": None,
     "preparation.checkpoint_token_file": "~/.config/flare/agent.token",
@@ -92,6 +93,7 @@ def load(cls, path: Path | None = None) -> "Configuration":
                 raise UsageError(f"{resolved}: unknown section [{section}]")
         configuration = cls(resolved, values, True)
         parse_size(configuration.get("memory.reserve").value)
+        parse_size(configuration.get("preparation.compile_headroom").value, "compile_headroom in [preparation]")
         for key in ("prefix_cache.memory", "prefix_cache.disk_limit"):
             if configuration.get(key).value:
                 parse_size(configuration.get(key).value, key.replace(".", " "))
@@ -182,6 +184,8 @@ def write_initial(self, model: str, max_concurrency: int) -> None:
             '# speculative_graph_budget = 64  # packed round graphs; excludes width-one and ordinary graphs',
             "", "[preparation]",
             '# compile_workers = 24           # CPU processes per host; default: automatic, bounded by memory',
+            f'# compile_headroom = "{self.get("preparation.compile_headroom").value}"'
+            '        # system memory compiler and tuning pools leave free per host',
             '# force_moe_decode_a16 = false   # Qwen Flash NVFP4 decode/verify A16, calibrated A4 prefill',
             '# autotune_effort = "default"    # off, default, audit',
             '# collective_tuning = "static"  # static, runtime; omitted uses recipe default',

ignition/cli/preparation.py

@@ -12,7 +12,7 @@
 from ignition.recipes import ResolvedModel, resolve
 from ignition.recipes.declaration import Speculation
 
-from .config import AUTOTUNE_EFFORTS, Configuration, Setting
+from .config import AUTOTUNE_EFFORTS, Configuration, Setting, parse_size
 
 
 @dataclass(frozen=True)
@@ -48,6 +48,7 @@ def derive(configuration: Configuration, model: ResolvedModel, host: Host, *, gp
            mixed_prefill: bool = False, prompt_logprobs: bool = False, expert_scale_compression: str | None = None,
            online_mxfp8: bool = False, online_mxfp8_scope: str | None = None, language_model_only: bool = False,
            compile_workers: int | None = None, tuning_gpus: str | None = None, autotune_effort: str | None = None, collective_tuning: str | None = None, collective_ladder=None,
+           compile_headroom: str | None = None,
            image_concurrency: int | None = None, mla_compute_mode: str | None = None,
            lock_gpu_clocks: bool = False, native_dense_activations: bool = False,
            joint_tuning: bool = False, shard_hyperconnections: bool = False,
@@ -66,6 +67,7 @@ def derive(configuration: Configuration, model: ResolvedModel, host: Host, *, gp
         image_concurrency=image_capacity_setting(configuration, model, image_concurrency),
         online_mxfp8=online_mxfp8, online_mxfp8_scope=online_mxfp8_scope, language_model_only=language_model_only, mla_compute_mode=mla_compute_mode)
     result = replace(result, compile_workers=workers, hub_cache=configuration.directory("hub_cache"),
+                     compile_headroom_bytes=compiler_headroom(configuration, compile_headroom),
                      autotune_effort=tuning_effort(configuration, autotune_effort),
                      collective_tuning=mode, collective_ladder=ladder,
                      force_moe_decode_a16=configuration.get("preparation.force_moe_decode_a16", force_moe_decode_a16).value,
@@ -102,6 +104,13 @@ def compiler_workers(configuration, flag=None):
     return value or None
 
 
+def compiler_headroom(configuration, flag=None):
+    """Resolve the system memory compiler pools and tuning races leave free, in bytes."""
+    setting = configuration.get("preparation.compile_headroom", flag)
+    return parse_size(setting.value, "--compile-headroom" if setting.source == "flag"
+                      else "compile_headroom in [preparation]")
+
+
 def tuning_devices(plan, host, flag):
     if plan.compile_workers is not None and plan.compile_workers < plan.world_size:
         raise UsageError("--compile-workers must provide at least one process per local serving rank")

ignition/cli/serve.py

@@ -17,7 +17,7 @@
 
 from . import display, environment
 
-from .preparation import (Host, collective_tuning_mode, compiler_workers, context_parallel, derive, fetch, image_capacity_setting, resolve_model, speculation,
+from .preparation import (Host, collective_tuning_mode, compiler_headroom, compiler_workers, context_parallel, derive, fetch, image_capacity_setting, resolve_model, speculation,
                           stream_overlap, tuning_devices, tuning_effort)
 from .config import Configuration
 
@@ -202,6 +202,7 @@ def _export_hosts(args):
                          mla_compute_mode=args.mla_compute_mode)
     collective_mode = collective_tuning_mode(configuration, model, args.collective_tuning)
     plan = replace(plan, compile_workers=compiler_workers(configuration, args.compile_workers), hub_cache=hub_cache,
+                   compile_headroom_bytes=compiler_headroom(configuration, args.compile_headroom),
                    autotune_effort=tuning_effort(configuration, args.autotune_effort),
                    collective_tuning=collective_mode,
                    collective_ladder=resolve_collective_ladder(collective_mode, args.collective_ladder),
@@ -270,7 +271,7 @@ def __call__(self, message):
                                                     args.max_concurrency)
         host = Host.discover()
         plan = derive(configuration, model, host, gpus=args.gpus, concurrency=concurrency,
-                      compile_workers=args.compile_workers, tuning_gpus=args.tuning_gpus, autotune_effort=args.autotune_effort,
+                      compile_workers=args.compile_workers, compile_headroom=args.compile_headroom, tuning_gpus=args.tuning_gpus, autotune_effort=args.autotune_effort,
                       collective_tuning=args.collective_tuning, collective_ladder=args.collective_ladder,
                       lock_gpu_clocks=args.lock_gpu_clocks, native_dense_activations=args.native_dense_activations,
                       force_moe_decode_a16=configuration.get("preparation.force_moe_decode_a16", args.force_moe_decode_a16).value,
@@ -304,6 +305,7 @@ def __call__(self, message):
                      f"{(' · online MXFP8 ' + plan.online_mxfp8_scope) if plan.online_mxfp8 else ''}"
                      f"{' · language model only' if plan.language_model_only else ''}"
                      f" · compile workers {plan.compile_workers or 'auto'} per host{tuning}"
+                     f" · compile headroom {plan.compile_headroom_bytes / (1 << 30):g} GiB"
                      f" · autotune effort {plan.autotune_effort} · collective tuning {plan.collective_tuning}{ladder_name}"
                      f"{' · native dense activations' if plan.native_dense_activations else ''}"
                      f"{' · MoE decode A16 / prefill A4' if plan.force_moe_decode_a16 else ''}"
@@ -464,6 +466,7 @@ def debug_serve(args) -> int:
                          language_model_only=args.language_model_only, mla_compute_mode=args.mla_compute_mode)
     collective_mode = collective_tuning_mode(configuration, model, args.collective_tuning)
     plan = replace(plan, compile_workers=compiler_workers(configuration, args.compile_workers), hub_cache=hub_cache,
+                   compile_headroom_bytes=compiler_headroom(configuration, args.compile_headroom),
                    autotune_effort=tuning_effort(configuration, args.autotune_effort),
                    collective_tuning=collective_mode,
                    collective_ladder=resolve_collective_ladder(collective_mode, args.collective_ladder),

ignition/planning.py

@@ -15,6 +15,7 @@
 from pathlib import Path
 
 from ignition.errors import CapacityError, UsageError
+from ignition.preparation.host_memory import COMPILE_HEADROOM_BYTES
 from ignition.preparation.tuning_replicas import TuningReplica
 from ignition.preparation.selection import BenchmarkSelection
 from ignition.recipes.declaration import Checkpoint, Recipe, Speculation
@@ -185,6 +186,8 @@ class Plan:
     roce_ring: bool = False
     roce_chain: bool = False
     compile_workers: int | None = None
+    # System memory compiler pools and tuning races leave free on each host.
+    compile_headroom_bytes: int = COMPILE_HEADROOM_BYTES
     autotune_effort: str = "default"
     collective_ladder: dict | None = None
     collective_tuning: str | None = None
@@ -369,6 +372,7 @@ def to_dict(self) -> dict:
             "expert_scale_compression": self.expert_scale_compression,
             "mla_compute_mode": self.mla_compute_mode,
             "compile_workers": self.compile_workers,
+            "compile_headroom_bytes": self.compile_headroom_bytes,
             "autotune_effort": self.autotune_effort,
             "collective_tuning": self.collective_tuning,
             "collective_ladder": self.collective_ladder,

ignition/preparation/aot.py

@@ -34,7 +34,7 @@
 from .compilation import (DEFAULT_CACHE_ROOT, _INVOCATIONS, _OPTIONS, _Invocation, _artifact_kernel_table,
                           _byte_storage, _storage_region, _lower_storage, _compiler_recursion)
 from .device_binding import bind_device, relocate_inductor, restore_device, restore_prepared_device
-from .host_memory import release_unused
+from .host_memory import COMPILE_HEADROOM_BYTES, release_unused
 
 _GUARD_LOCK = RLock()
 _GUARD_TENSOR_MODE = ContextVar("flare_guard_tensor_mode", default=None)
@@ -566,7 +566,7 @@ class ForwardCache:
     """Load complete compiled forwards before tracing the serving catalog."""
 
     def __init__(self, *, namespace, device, root=None, headroom_bytes=0, options=None, compile_workers=None,
-                 lock_gpu_clocks=False):
+                 lock_gpu_clocks=False, compile_headroom_bytes=COMPILE_HEADROOM_BYTES):
         runtime_patches.apply()
         import flashinfer
         from torch._inductor.codecache import torch_key
@@ -575,6 +575,7 @@ def __init__(self, *, namespace, device, root=None, headroom_bytes=0, options=No
         self.headroom_bytes = headroom_bytes
         self.compile_workers = compile_workers
         self.lock_gpu_clocks = lock_gpu_clocks
+        self.compile_headroom_bytes = compile_headroom_bytes
         self.options = MappingProxyType(dict(_OPTIONS if options is None else options))
         properties = torch.cuda.get_device_properties(device)
         options = {name: value for name, value in self.options.items() if name != 'compile_threads'}
@@ -710,7 +711,8 @@ def _prepare_many(self, functions, signatures, *, progress, coordinator, materia
         selections = InductorSelections(self.device, lock_gpu_clocks=self.lock_gpu_clocks)
         compiled, paths, pending = {}, {}, []
         with phase(progress, "Compile model forwards", total=len(functions)) as stage, CompilePool(self.device,
-                headroom_bytes=self.headroom_bytes, options=self.options, compile_workers=self.compile_workers) as compiler:
+                headroom_bytes=self.headroom_bytes, options=self.options, compile_workers=self.compile_workers,
+                compile_headroom_bytes=self.compile_headroom_bytes) as compiler:
             def finish():
                 name, bound, path = _completed(pending)
                 try:
@@ -799,7 +801,8 @@ def _prepare_ranks(self, functions, signatures, coordinator, *, progress, materi
             try:
                 compiler = compiler_scope.enter_context(CompilePool(
                     self.device, headroom_bytes=self.headroom_bytes, pools=coordinator.local_ranks,
-                    options=self.options, compile_workers=self.compile_workers))
+                    options=self.options, compile_workers=self.compile_workers,
+                    compile_headroom_bytes=self.compile_headroom_bytes))
             except MemoryError as failure:
                 # A rank whose GPU has no room for a compiler lends its forwards to the others on this host.
                 unavailable = f"{type(failure).__name__}: {failure}"