From 4704d27ec23f20b57bc07a589731fa63a6eb6007 Mon Sep 17 00:00:00 2001 From: Jiong Gong Date: Mon, 21 Sep 2026 17:29:52 +0000 Subject: [PATCH 1/3] compass(runner): a model runner that constructs without device memory (RUNNER-1) The attachment point itself: a `ModelRunner` subclass, selected through the existing `runner_qualname` config field, that replaces the five methods owning weights, the KV tensors and the step. Nothing here runs a step. The behaviour sits in `overrides.py`, which imports nothing from the engine, because importing ATOM's model runner runs aiter's architecture probe and raises where there is no driver. `model_runner.py` binds it onto `ModelRunner` and is the only module here that needs one. Two things the base class does that decide the shape of this one: it warms the model from inside `__init__`, and warmup drives a forward -- so a runner with no weights constructs only because warmup is an override point. Both links are asserted over ATOM's own source. Co-Authored-By: Claude Opus 5 (1M context) --- atom/compass/runner/__init__.py | 27 +++ atom/compass/runner/model_runner.py | 25 ++ atom/compass/runner/overrides.py | 128 +++++++++++ tests/compass/test_runner_non_allocating.py | 239 ++++++++++++++++++++ 4 files changed, 419 insertions(+) create mode 100644 atom/compass/runner/__init__.py create mode 100644 atom/compass/runner/model_runner.py create mode 100644 atom/compass/runner/overrides.py create mode 100644 tests/compass/test_runner_non_allocating.py diff --git a/atom/compass/runner/__init__.py b/atom/compass/runner/__init__.py new file mode 100644 index 0000000000..876d21db70 --- /dev/null +++ b/atom/compass/runner/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""The model runner a simulated run installs in place of ATOM's own. + +Two modules, split by what each is allowed to import: + +- `overrides` holds the behaviour -- the methods that keep construction off the + device -- and imports nothing from the engine. That is not a stylistic + preference. Importing `atom.model_engine.model_runner` runs aiter's + architecture probe, which shells out to `rocminfo` and raises in a container + with no driver, so anything reachable from that import cannot be exercised by + a test that runs without one. +- `model_runner` binds those overrides onto `ModelRunner`, and is therefore the + only module here that needs a driver to import. + +This package stays empty at import time for the same reason, so naming it costs +nothing. + +`COMPASS_RUNNER_QUALNAME` is the string that selects the runner: a `Config` +field, read where the worker process instantiates its runner. Installing this +class needs no change to ATOM. +""" + +COMPASS_RUNNER_QUALNAME = "atom.compass.runner.model_runner.CompassModelRunner" + +__all__ = ["COMPASS_RUNNER_QUALNAME"] diff --git a/atom/compass/runner/model_runner.py b/atom/compass/runner/model_runner.py new file mode 100644 index 0000000000..ad560be8d0 --- /dev/null +++ b/atom/compass/runner/model_runner.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""The runner class that `runner_qualname` names.""" + +from __future__ import annotations + +from atom.compass.runner.overrides import NonAllocatingRunner +from atom.model_engine.model_runner import ModelRunner + + +class CompassModelRunner(NonAllocatingRunner, ModelRunner): + """A `ModelRunner` that constructs without owning any device memory. + + Everything the engine does around a step -- the scheduler, admission, the + block manager, the prefix index -- is ATOM's own and runs unchanged. What + this class removes is the weights, the KV tensors and the step itself, by + overriding the five methods that own them; see `overrides`, which holds the + bodies and says why each one does what it does. + + `NonAllocatingRunner` comes first so its methods win over the base's. There + is deliberately no `__init__`: the base runs all of its own before a + subclass body would get control, so there is no point at which state set + here would be visible to it. + """ diff --git a/atom/compass/runner/overrides.py b/atom/compass/runner/overrides.py new file mode 100644 index 0000000000..aab6ddde57 --- /dev/null +++ b/atom/compass/runner/overrides.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved. + +"""The model-runner methods that keep construction off the device. + +These are the memory-owning methods of ATOM's `ModelRunner`: the ones that read +a checkpoint, warm the model, size the KV pool, allocate it, and run a step. +They live here rather than beside the subclass because this module imports +nothing from the engine, and so can be executed on a machine with no driver. + +Two of them run during construction (`_build_and_load_model`, `_maybe_warmup`) +and decline to do their work. The rest run afterwards, over the worker's RPC +channel; `allocate_kv_cache` does the arithmetic and none of the allocation, +and the other two refuse by name because their answers are not this module's to +invent -- a wrong duration or a wrong block count would be indistinguishable +from a measured one. + +Order matters when mixing this in: `NonAllocatingRunner` must precede +`ModelRunner` in the bases so these definitions win. The class deliberately +defines no `__init__`. The base class runs the whole of its own `__init__` +before a subclass body would get control, so anything these methods read has to +come from `self.config`, which the base sets first. +""" + +from __future__ import annotations + +import logging +from typing import Any + +import torch + +logger = logging.getLogger(__name__) + + +class RunnerRefusal(RuntimeError): + """Raised where this runner has no answer and will not invent one.""" + + +class UnbuiltModel(torch.nn.Module): + """Stands in for the module tree a non-allocating runner never builds. + + It registers no parameter and no buffer, so it costs nothing on any device, + and calling it raises instead of returning something a caller could mistake + for a forward pass. + """ + + def __init__(self, model_class: Any) -> None: + super().__init__() + self.model_class_name = getattr(model_class, "__name__", repr(model_class)) + + def forward(self, *args: Any, **kwargs: Any) -> Any: + raise RunnerRefusal( + f"{self.model_class_name} was never built: this runner predicts what a " + "forward pass costs instead of running one, so there are no weights " + "to call." + ) + + +class NonAllocatingRunner: + """Overrides that construct a model runner without touching device memory.""" + + def _build_and_load_model(self, model_class: Any) -> None: + """Build nothing and read no checkpoint. + + The base class constructs the model with the default device set to this + rank's GPU and then fills it from disk. Both halves are skipped here, so + no weight byte reaches the device and no checkpoint is read. + """ + self.model = UnbuiltModel(model_class) + # Cleared on the way out, as both of ATOM's own implementations do: the + # caller set the default device to this rank's GPU before calling, and + # the code that runs next is written against a cleared default. + torch.set_default_device(None) + logger.info( + "%s not built and no checkpoint read; no weight bytes on the device.", + self.model.model_class_name, + ) + + def _maybe_warmup(self) -> None: + """Skip warmup. + + Warmup runs a forward over a dummy batch against real weights. There are + none here, so there is nothing to warm. Skipping is also what lets + construction finish at all: the base class warms the model from inside + its own `__init__`, and the forward that warmup drives is this class's, + which refuses. + """ + return + + def get_num_blocks(self) -> dict[str, object]: + """Refuse to size the KV pool. + + The base sizes it from what a real device reports free after the weights + are resident. Neither figure exists here, and the substitute belongs to + the memory model rather than to the runner, so answering would mean + inventing a block count that the scheduler would then treat as measured. + """ + raise RunnerRefusal( + "a non-allocating runner cannot size the KV pool from a device it " + "never allocated on; the block count has to come from a memory model " + "this runner has not been given." + ) + + def allocate_kv_cache(self, num_kvcache_blocks: int) -> bool: + """Record the block count and allocate nothing. + + The block accounting above this -- the pool, the prefix index, eviction, + preemption -- is integer arithmetic and runs unmodified against the count + recorded here. Only the tensors behind the blocks are absent. + """ + self.config.num_kvcache_blocks = num_kvcache_blocks + logger.info( + "kv cache: %d blocks accounted, 0 bytes allocated.", + num_kvcache_blocks, + ) + return True + + def forward(self, batch: Any) -> Any: + """Refuse to run a step. + + This class is the attachment point and not the replacement for a step. A + step's output has to reproduce what the scheduler reads back from a real + one, and a plausible-looking stub that does not is worse than no answer. + """ + raise RunnerRefusal( + "this runner has no cost model and no step semantics yet, so it " + "cannot say what a step produced or how long it took." + ) diff --git a/tests/compass/test_runner_non_allocating.py b/tests/compass/test_runner_non_allocating.py new file mode 100644 index 0000000000..7a39570034 --- /dev/null +++ b/tests/compass/test_runner_non_allocating.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: MIT +"""The runner subclass, and the construction that owns no device memory. + +These run without a driver, which is the reason the behaviour lives in +`atom.compass.runner.overrides` rather than beside the subclass: importing +`atom.model_engine.model_runner` runs aiter's architecture probe, which shells +out to `rocminfo` and raises where there is no GPU. So the bodies are exercised +directly and the binding to ATOM's class is read from the source instead. + +The two facts worth naming, because both were surprises: + +* The base class warms the model from inside its own `__init__`, and warmup + drives a forward. A runner with no weights therefore cannot construct unless + it skips warmup -- the forward that warmup would drive is its own. That chain + is asserted over ATOM's source below, so this stops being true loudly rather + than quietly. +* Allocation is checked by counting dispatched operators, not by reading a CUDA + allocator, so the check is meaningful on a machine with no CUDA allocator to + read. A control asserts the counter sees an allocation when there is one; + without it a broken counter and a clean runner look identical. +""" + +import ast +import dataclasses +import pathlib +from types import SimpleNamespace + +import pytest +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from atom.compass.runner import COMPASS_RUNNER_QUALNAME +from atom.compass.runner.overrides import ( + NonAllocatingRunner, + RunnerRefusal, + UnbuiltModel, +) + +REPO = pathlib.Path(__file__).resolve().parents[2] +PACKAGE = REPO / "atom" / "compass" / "runner" +ENGINE = REPO / "atom" / "model_engine" +ATOM_RUNNER = ENGINE / "model_runner.py" + +# The methods that own memory or run a step, and so are the ones replaced. +OVERRIDDEN = { + "_build_and_load_model", + "_maybe_warmup", + "get_num_blocks", + "allocate_kv_cache", + "forward", +} + + +def _classes(path): + tree = ast.parse(path.read_text()) + return {n.name: n for n in ast.walk(tree) if isinstance(n, ast.ClassDef)} + + +def _methods(node): + return {n.name for n in node.body if isinstance(n, ast.FunctionDef)} + + +def _self_calls(node): + """Names of `self.x(...)` calls anywhere inside a function definition.""" + return { + n.func.attr + for n in ast.walk(node) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "self" + } + + +class Recorder(TorchDispatchMode): + """Every operator torch dispatches while this is entered.""" + + def __init__(self): + self.ops = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + self.ops.append(str(func)) + return func(*args, **(kwargs or {})) + + +class Runner(NonAllocatingRunner): + """The overrides over a base that supplies only what they read.""" + + def __init__(self): + self.config = SimpleNamespace(num_kvcache_blocks=None) + + +class SomeModel: + """Stands in for the model class the engine hands the runner.""" + + +# --- what is replaced, and how that differs from the runner already in tree --- + + +def test_the_replaced_methods_are_the_ones_that_own_memory_or_run_a_step(): + assert _methods(_classes(PACKAGE / "overrides.py")["NonAllocatingRunner"]) == ( + OVERRIDDEN + ) + + +def test_every_replaced_method_exists_on_the_class_being_replaced(): + """A rename upstream turns an override into a new method, silently.""" + assert OVERRIDDEN <= _methods(_classes(ATOM_RUNNER)["ModelRunner"]) + + +def test_the_difference_from_the_in_tree_non_allocating_runner_is_two_methods(): + """`RapidServeModelRunner` is the working non-allocating runner in the tree. + + It overrides two things this one does not. `__init__`: it has to bind a + method before the base runs, and this class has nothing to bind, so leaving + `__init__` alone is what keeps every read lazy. `_kv_budget_extra_reserve`: + it holds bytes back because a second process shares its GPU, and a runner + that allocates nothing has no tenant to hold anything back from. + + Its `_init_weight_params_on_meta` is not in the difference because it is not + an override -- it is a helper the base does not have. + """ + base = _methods(_classes(ATOM_RUNNER)["ModelRunner"]) + template = _methods(_classes(ATOM_RUNNER)["RapidServeModelRunner"]) & base + assert template - OVERRIDDEN == {"__init__", "_kv_budget_extra_reserve"} + assert OVERRIDDEN - template == set() + + +# --- construction allocates nothing ----------------------------------------- + + +def test_construction_dispatches_no_operator_at_all(): + """The measured result: zero, against a base that reads a checkpoint.""" + runner = Runner() + with Recorder() as recorded: + runner._build_and_load_model(SomeModel) + runner._maybe_warmup() + runner.allocate_kv_cache(2048) + assert recorded.ops == [] + + +def test_the_recorder_sees_an_allocation_when_there_is_one(): + """Otherwise a dead recorder reads exactly like a clean runner.""" + with Recorder() as recorded: + torch.empty(4) + assert recorded.ops + + +def test_the_model_that_replaces_the_weights_owns_nothing(): + runner = Runner() + runner._build_and_load_model(SomeModel) + assert isinstance(runner.model, UnbuiltModel) + assert list(runner.model.parameters()) == [] + assert list(runner.model.buffers()) == [] + + +def test_calling_that_model_refuses_and_names_the_class_it_stands_for(): + with pytest.raises(RunnerRefusal, match="SomeModel"): + UnbuiltModel(SomeModel)(torch.zeros(1)) + + +# --- warmup drives a forward, which is why it is skipped --------------------- + + +def test_the_base_class_warms_the_model_from_init_and_warmup_runs_a_forward(): + """The chain that decides whether a weightless runner can construct.""" + runner = _classes(ATOM_RUNNER)["ModelRunner"] + bodies = {n.name: n for n in runner.body if isinstance(n, ast.FunctionDef)} + assert "_maybe_warmup" in _self_calls(bodies["__init__"]) + assert "warmup_model" in _self_calls(bodies["_maybe_warmup"]) + assert "forward" in _self_calls(bodies["warmup_model"]) + + +def test_skipping_warmup_is_what_lets_construction_finish(): + """`forward` refuses, so a warmup that ran one would raise out of `__init__`.""" + runner = Runner() + assert runner._maybe_warmup() is None + with pytest.raises(RunnerRefusal): + runner.forward(object()) + + +# --- the two answers that are refused, and the one that is arithmetic -------- + + +def test_sizing_the_kv_pool_refuses_rather_than_inventing_a_block_count(): + with pytest.raises(RunnerRefusal, match="memory model"): + Runner().get_num_blocks() + + +def test_allocating_the_kv_cache_records_the_count_and_no_bytes(): + runner = Runner() + assert runner.allocate_kv_cache(2048) is True + assert runner.config.num_kvcache_blocks == 2048 + + +# --- the seam: reachable without changing ATOM ------------------------------- + + +def test_the_qualname_names_a_class_this_package_defines(): + module, _, name = COMPASS_RUNNER_QUALNAME.rpartition(".") + path = REPO.joinpath(*module.split(".")).with_suffix(".py") + assert name in _classes(path) + + +def test_selecting_a_runner_is_a_config_field_that_already_exists(): + from atom.config import Config + + field = {f.name: f for f in dataclasses.fields(Config)}["runner_qualname"] + assert field.default == "atom.model_engine.model_runner.ModelRunner" + + +def test_the_worker_process_instantiates_whatever_that_field_names(): + assert "config.runner_qualname" in (ENGINE / "engine_core.py").read_text() + assert ( + "resolve_obj_by_qualname(runner_qualname)" + in (ENGINE / "async_proc.py").read_text() + ) + + +@pytest.mark.parametrize( + "path", + sorted(PACKAGE.rglob("*.py")), + ids=lambda p: p.name, +) +def test_only_the_binding_module_reaches_the_engine(path): + """Everything else stays runnable where the engine cannot be imported.""" + tree = ast.parse(path.read_text()) + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported |= {a.name for a in node.names} + elif isinstance(node, ast.ImportFrom): + imported.add("." * node.level + (node.module or "")) + engine = {m for m in imported if m.split(".")[0] == "atom"} - { + m for m in imported if m.startswith("atom.compass.runner") + } + assert engine == ( + {"atom.model_engine.model_runner"} if path.name == "model_runner.py" else set() + ) From fda25e63daffa9f5d07cdc70c8f58b11dce75789 Mon Sep 17 00:00:00 2001 From: Jiong Gong Date: Mon, 21 Sep 2026 18:03:10 +0000 Subject: [PATCH 2/3] compass(runner): qualify the class docstring to what was measured (RUNNER-1 r2) "constructs without owning any device memory" was unqualified and is false by 16.9 MiB. End-to-end construction on the device leaves 2,168,320 bytes at a 1024-token / 4-sequence budget and 17,668,096 at 8192 / 256, of which `allocate_forward_vars` is 99.1% and 99.4%; named tensors on the runner are 0 bytes in both. The docstring now says weights, KV tensors and the step, and names the forward-vars ring as the remainder -- sized by the batch budget, not by the model. Docstring only; no behaviour changes. Co-Authored-By: Claude Opus 5 (1M context) --- atom/compass/runner/model_runner.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/atom/compass/runner/model_runner.py b/atom/compass/runner/model_runner.py index ad560be8d0..c7e9dacd0d 100644 --- a/atom/compass/runner/model_runner.py +++ b/atom/compass/runner/model_runner.py @@ -10,7 +10,7 @@ class CompassModelRunner(NonAllocatingRunner, ModelRunner): - """A `ModelRunner` that constructs without owning any device memory. + """A `ModelRunner` that constructs without weights, KV tensors or a step. Everything the engine does around a step -- the scheduler, admission, the block manager, the prefix index -- is ATOM's own and runs unchanged. What @@ -18,6 +18,15 @@ class CompassModelRunner(NonAllocatingRunner, ModelRunner): overriding the five methods that own them; see `overrides`, which holds the bodies and says why each one does what it does. + Construction is not free of device memory, and what remains is the base's + forward-vars ring. Measured at TP1 on Qwen3-0.6B with `enforce_eager`, the + whole of `__init__` leaves 2,168,320 bytes allocated at a 1024-token, + 4-sequence budget and 17,668,096 at 8192 / 256; `allocate_forward_vars` + accounts for 99.1% and 99.4% of those, and named tensors on the runner for + 0 bytes both times. That remainder is sized by the batch budget and not by + the model: 16.9 MiB against the 1.40 GiB of weights the base makes resident + for Qwen3-0.6B, and against 51.7 GiB for Qwen3.8-27B. + `NonAllocatingRunner` comes first so its methods win over the base's. There is deliberately no `__init__`: the base runs all of its own before a subclass body would get control, so there is no point at which state set From 1a30e6e7ffbd8b9469f2e65f47c2d01a0cd2c5d8 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 21 Sep 2026 18:43:55 +0000 Subject: [PATCH 3/3] compass(runner): state the shape of the construction residue, not its bytes The class docstring's residue paragraph carried two measured byte figures and set one of them against two models. The dominant term of the forward-vars ring is `torch.empty(max_num_batched_tokens, hidden_size, ...)` in `ModelRunner.allocate_forward_vars` (`model_runner.py:1288`, `:1308-1313`), so the ring scales with the model's hidden size; "sized by the batch budget and not by the model" was wide, and 16.9 MiB is the value at one hidden size only. Replace the figures with the shape: the remainder is the base's forward-vars ring, whose dominant term is a `max_num_batched_tokens` by `hidden_size` output buffer, sized by the batch budget and the model's hidden size rather than by its weights, and no named tensor on the runner holds any of it. True at any model and any budget, with no conditions to go stale. The measurements, their conditions and their attribution move to the PR body, which is the only place that can carry them: a docstring cannot carry a measurement's instrument, date and attribution, and nothing in CI reads this one. Docstring only; no behaviour change, and no test reads it. Co-Authored-By: Claude Opus 5 (1M context) --- atom/compass/runner/model_runner.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/atom/compass/runner/model_runner.py b/atom/compass/runner/model_runner.py index c7e9dacd0d..cd2792f723 100644 --- a/atom/compass/runner/model_runner.py +++ b/atom/compass/runner/model_runner.py @@ -18,14 +18,11 @@ class CompassModelRunner(NonAllocatingRunner, ModelRunner): overriding the five methods that own them; see `overrides`, which holds the bodies and says why each one does what it does. - Construction is not free of device memory, and what remains is the base's - forward-vars ring. Measured at TP1 on Qwen3-0.6B with `enforce_eager`, the - whole of `__init__` leaves 2,168,320 bytes allocated at a 1024-token, - 4-sequence budget and 17,668,096 at 8192 / 256; `allocate_forward_vars` - accounts for 99.1% and 99.4% of those, and named tensors on the runner for - 0 bytes both times. That remainder is sized by the batch budget and not by - the model: 16.9 MiB against the 1.40 GiB of weights the base makes resident - for Qwen3-0.6B, and against 51.7 GiB for Qwen3.8-27B. + Construction is not free of device memory. What remains is the base's + forward-vars ring from `allocate_forward_vars`, whose dominant term is a + `max_num_batched_tokens` by `hidden_size` output buffer. It is sized by the + batch budget and the model's hidden size, not by the model's weights, and + no named tensor on the runner holds any of it. `NonAllocatingRunner` comes first so its methods win over the base's. There is deliberately no `__init__`: the base runs all of its own before a