feat(offload): enable LMCache offload via ATOM_KV_OFFLOAD env

atom/kv_transfer/offload/README.md

@@ -996,6 +996,22 @@ python -m atom.entrypoints.openai_server \
   --kv-transfer-config '{"kv_connector":"lmcache_offload","kv_role":"offload"}'
 ```
 
+`ATOM_KV_OFFLOAD` enables the same connector without taking
+`--kv-transfer-config`, so a launcher that owns that flag for P/D transfer can
+still add offload. `lmcache` selects `lmcache_offload` and `lmcache_mp` the
+standalone-server `lmcache_mp`, both with `kv_role: offload`.
+`ATOM_KV_OFFLOAD_EXTRA_CONFIG` (a JSON object) becomes the connector's
+`kv_connector_extra_config`, e.g. `lmcache.chunk_size` or `lmcache.mp.port`.
+With a P/D connector in `--kv-transfer-config`, both run behind a `multi`
+connector:
+
+```bash
+ATOM_KV_OFFLOAD=lmcache python -m atom.entrypoints.openai_server ...
+ATOM_KV_OFFLOAD=lmcache_mp ATOM_KV_OFFLOAD_EXTRA_CONFIG='{"lmcache.mp.port":5555}' \
+  python -m atom.entrypoints.openai_server ...
+# with P/D: --kv-transfer-config '{"kv_connector":"mooncake",...}' plus ATOM_KV_OFFLOAD=lmcache
+```
+
 Standalone DSV4 LMCache offload supports both FP8 and FP4 indexer layouts.
 FP4 PAGE objects include the packed data and separate e8m0 scale regions for
 every CSA layer. FP4 remains unsupported with PD connectors such as Mooncake

atom/kv_transfer/offload/__init__.py

@@ -5,12 +5,14 @@
 
 Registers the legacy in-process ``lmcache_offload`` backend and the standalone
 server ``lmcache_mp`` backend with the shared KV connector factory.
-Enable via ``--kv-transfer-config '{"kv_connector":"lmcache_offload","kv_role":"offload"}'``
+Enable via ``ATOM_KV_OFFLOAD=lmcache`` (or the equivalent
+``--kv-transfer-config '{"kv_connector":"lmcache_offload","kv_role":"offload"}'``)
 plus LMCache env (``LMCACHE_LOCAL_CPU=True``, ``LMCACHE_MAX_LOCAL_CPU_SIZE``,
 ``LMCACHE_CHUNK_SIZE=256``, optional ``LMCACHE_LOCAL_DISK`` for the NVMe L3 tier).
 
-For standalone MP mode, start ``lmcache server`` and select
-``{"kv_connector":"lmcache_mp","kv_role":"offload"}``; the active attention
+For standalone MP mode, start ``lmcache server`` and set
+``ATOM_KV_OFFLOAD=lmcache_mp`` (or select
+``{"kv_connector":"lmcache_mp","kv_role":"offload"}``); the active attention
 backend publishes its PAGE layout directly to the model-neutral connector.
 """
 

atom/model_engine/arg_utils.py

@@ -17,6 +17,7 @@
     SpeculativeConfig,
 )
 from atom.model_engine.engine_core_mgr import DP_LB_DEFAULT, DP_LB_STRATEGIES
+from atom.utils import envs
 
 logger = logging.getLogger("atom")
 
@@ -31,6 +32,78 @@ def parse_size_list(size_str: str) -> list[int]:
         raise ValueError(f"Error parsing size list: {size_str}") from e
 
 
+# Offload connector names and aliases (atom/kv_transfer/offload/__init__.py),
+# casefolded like KVConnectorFactory.canonical_name.
+_OFFLOAD_CONNECTORS = frozenset(
+    {
+        "lmcache_offload",
+        "lmcacheoffloadconnector",
+        "lmcacheconnectorv1",
+        "lmcache_mp",
+        "lmcachempconnector",
+    }
+)
+
+
+# ATOM_KV_OFFLOAD value -> the offload connector it selects.
+_KV_OFFLOAD_MODES = {"lmcache": "lmcache_offload", "lmcache_mp": "lmcache_mp"}
+
+
+def kv_offload_connector_config(mode: str, extra_config: str) -> dict | None:
+    """The offload connector ATOM_KV_OFFLOAD selects, or None when it is off.
+
+    ``lmcache`` is the in-process backend and ``lmcache_mp`` the standalone
+    ``lmcache server`` one. ATOM_KV_OFFLOAD_EXTRA_CONFIG, a JSON object, becomes
+    the connector's ``kv_connector_extra_config``, which both backends read.
+    """
+    mode = mode.strip().lower()
+    if mode in ("", "0", "off", "none"):
+        if extra_config.strip():
+            raise ValueError(
+                "ATOM_KV_OFFLOAD_EXTRA_CONFIG is set but ATOM_KV_OFFLOAD is off"
+            )
+        return None
+    if mode not in _KV_OFFLOAD_MODES:
+        raise ValueError(
+            f"ATOM_KV_OFFLOAD={mode!r}: expected one of {sorted(_KV_OFFLOAD_MODES)}"
+        )
+    offload = {"kv_connector": _KV_OFFLOAD_MODES[mode], "kv_role": "offload"}
+    if extra_config.strip():
+        extra = json.loads(extra_config)
+        if not isinstance(extra, dict):
+            raise TypeError("ATOM_KV_OFFLOAD_EXTRA_CONFIG must be a JSON object")
+        offload["kv_connector_extra_config"] = extra
+    return offload
+
+
+def compose_kv_offload_config(kv_transfer_config: str, offload: dict | None) -> str:
+    """Fold the ATOM_KV_OFFLOAD connector into the --kv-transfer-config JSON string.
+
+    Without a transfer connector the offload connector stands alone; with one,
+    both become subs of a ``multi`` connector (appended if it already is one).
+    """
+    if offload is None:
+        return kv_transfer_config
+    transfer = json.loads(kv_transfer_config or "{}")
+    if not transfer:
+        return json.dumps(offload)
+    is_multi = _connector_name(transfer) == "multi"
+    subs = transfer.get("connectors", []) if is_multi else [transfer]
+    if any(_connector_name(sub) in _OFFLOAD_CONNECTORS for sub in subs):
+        raise ValueError(
+            "ATOM_KV_OFFLOAD conflicts with an offload connector already set in "
+            "--kv-transfer-config; use only one of them"
+        )
+    if is_multi:
+        return json.dumps({**transfer, "connectors": [*subs, offload]})
+    return json.dumps({"kv_connector": "multi", "connectors": [transfer, offload]})
+
+
+def _connector_name(cfg: object) -> str:
+    name = cfg.get("kv_connector", "") if isinstance(cfg, dict) else ""
+    return str(name).strip().casefold()
+
+
 @dataclass
 class EngineArgs:
     """Arguments for configuring the LLM Engine."""
@@ -726,6 +799,19 @@ def _get_engine_kwargs(self) -> dict:
             "none": "none",
         }[all2all_backend]
 
+        # ATOM_KV_OFFLOAD → folded into kv_transfer_config, so offload does not
+        # compete with a P/D connector (or a launcher owning that flag) for it.
+        offload = kv_offload_connector_config(
+            envs.ATOM_KV_OFFLOAD, envs.ATOM_KV_OFFLOAD_EXTRA_CONFIG
+        )
+        kwargs["kv_transfer_config"] = compose_kv_offload_config(
+            kwargs["kv_transfer_config"], offload
+        )
+        if offload is not None:
+            logger.info(
+                "ATOM_KV_OFFLOAD: kv_transfer_config=%s", kwargs["kv_transfer_config"]
+            )
+
         # --dspark-config (JSON dict) → DSparkConfig object, passed through as
         # Config.dspark (no env vars).
         kwargs["dspark"] = DSparkConfig.from_dict(kwargs.pop("dspark_config", None))

atom/utils/envs.py

@@ -56,6 +56,14 @@ def _positive_float_env(name: str, default: str) -> float:
     "ATOM_PREFIX_CACHE_PROTECTED_RATIO": lambda: float(
         os.getenv("ATOM_PREFIX_CACHE_PROTECTED_RATIO", "0.5")
     ),
+    # LMCache KV offload without --kv-transfer-config: "lmcache" (in-process) or
+    # "lmcache_mp" (standalone `lmcache server`). Unset = off.
+    "ATOM_KV_OFFLOAD": lambda: os.getenv("ATOM_KV_OFFLOAD", ""),
+    # JSON object for that connector's kv_connector_extra_config, e.g.
+    # {"lmcache.chunk_size": 256} or {"lmcache.mp.port": 5556}.
+    "ATOM_KV_OFFLOAD_EXTRA_CONFIG": lambda: os.getenv(
+        "ATOM_KV_OFFLOAD_EXTRA_CONFIG", ""
+    ),
     # --- Data Parallelism ---
     "ATOM_DP_RANK": lambda: int(os.getenv("ATOM_DP_RANK", "0")),
     "ATOM_DP_RANK_LOCAL": lambda: int(os.getenv("ATOM_DP_RANK_LOCAL", "0")),

docs/environment_variables.md

@@ -288,6 +288,11 @@ flag below. Details in the state-checkpoint section of the
 
 ### LMCache offload tier
 
+| Variable | Type | Default | Description |
+|----------|------|---------|-------------|
+| **ATOM_KV_OFFLOAD** | str | "" (off) | Enables LMCache KV offload without `--kv-transfer-config`, so a launcher that owns that flag for P/D can still add offload. `lmcache` selects the in-process `lmcache_offload` connector, `lmcache_mp` the standalone-server `lmcache_mp` connector (start `lmcache server` first). With a P/D connector in `--kv-transfer-config`, both run behind a `multi` connector. Setting it while `--kv-transfer-config` already names an offload connector is an error. |
+| **ATOM_KV_OFFLOAD_EXTRA_CONFIG** | JSON object | "" | The offload connector's `kv_connector_extra_config`: `lmcache.<field>` overrides (e.g. `{"lmcache.chunk_size": 256}`) and, for `lmcache_mp`, `lmcache.mp.*` options (e.g. `{"lmcache.mp.port": 5556}`). Requires `ATOM_KV_OFFLOAD`. |
+
 The LMCache source-pin timeout and ATOM's pending-save bound are read directly
 via `os.environ` rather than through `atom.utils.envs`. Their behavior is
 defined in `atom/kv_transfer/offload/_offload_common.py` and documented in full

tests/test_arg_utils_spec.py

@@ -2,9 +2,12 @@
 # Regression tests for speculative-config validation in EngineArgs._get_engine_kwargs.
 
 import argparse
+import json
 import sys
 from unittest.mock import MagicMock, patch
 
+import pytest
+
 # conftest.py stubs atom.* and zmq before any atom imports are attempted,
 # but arg_utils.py imports LLMEngine from atom and CompilationConfig /
 # SpeculativeConfig from atom.config, which the minimal stub doesn't expose.
@@ -236,3 +239,116 @@ def test_index_cache_dtype_accepts_underscore_spelling(self):
         )
 
         assert args.index_cache_dtype == "fp8"
+
+
+class TestKVOffloadEnv:
+    """ATOM_KV_OFFLOAD folds into kv_transfer_config without owning that flag."""
+
+    @pytest.fixture(autouse=True)
+    def _clear_env(self, monkeypatch):
+        monkeypatch.delenv("ATOM_KV_OFFLOAD", raising=False)
+        monkeypatch.delenv("ATOM_KV_OFFLOAD_EXTRA_CONFIG", raising=False)
+
+    def _kv(self, argv=()):
+        parser = argparse.ArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        args = EngineArgs.from_cli_args(parser.parse_args(list(argv)))
+        return json.loads(args._get_engine_kwargs()["kv_transfer_config"])
+
+    def test_cli_flag_is_gone(self):
+        parser = argparse.ArgumentParser()
+        EngineArgs.add_cli_args(parser)
+        with pytest.raises(SystemExit):
+            parser.parse_args(["--kv-offload-config", "{}"])
+
+    @pytest.mark.parametrize("mode", ["", "0", "off", "none"])
+    def test_off_leaves_transfer_config_untouched(self, monkeypatch, mode):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", mode)
+        assert self._kv() == {}
+
+    @pytest.mark.parametrize(
+        "mode, connector",
+        [
+            ("lmcache", "lmcache_offload"),
+            ("LMCache", "lmcache_offload"),
+            ("lmcache_mp", "lmcache_mp"),
+        ],
+    )
+    def test_mode_selects_connector(self, monkeypatch, mode, connector):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", mode)
+        assert self._kv() == {"kv_connector": connector, "kv_role": "offload"}
+
+    def test_extra_config_becomes_connector_extra_config(self, monkeypatch):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", "lmcache_mp")
+        monkeypatch.setenv(
+            "ATOM_KV_OFFLOAD_EXTRA_CONFIG",
+            '{"lmcache.mp.port": 5556, "lmcache.chunk_size": 256}',
+        )
+        assert self._kv() == {
+            "kv_connector": "lmcache_mp",
+            "kv_role": "offload",
+            "kv_connector_extra_config": {
+                "lmcache.mp.port": 5556,
+                "lmcache.chunk_size": 256,
+            },
+        }
+
+    def test_pd_connector_and_offload_become_multi(self, monkeypatch):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", "lmcache")
+        pd = {"kv_connector": "mooncake", "kv_role": "kv_producer"}
+        cfg = self._kv(["--kv-transfer-config", json.dumps(pd)])
+        assert cfg == {
+            "kv_connector": "multi",
+            "connectors": [
+                pd,
+                {"kv_connector": "lmcache_offload", "kv_role": "offload"},
+            ],
+        }
+
+    def test_existing_multi_gets_offload_appended(self, monkeypatch):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", "lmcache_mp")
+        pd = {"kv_connector": "moriio", "kv_role": "kv_producer"}
+        multi = {"kv_connector": "multi", "connectors": [pd]}
+        cfg = self._kv(["--kv-transfer-config", json.dumps(multi)])
+        assert cfg["kv_connector"] == "multi"
+        assert [c["kv_connector"] for c in cfg["connectors"]] == [
+            "moriio",
+            "lmcache_mp",
+        ]
+
+    @pytest.mark.parametrize(
+        "transfer",
+        [
+            {"kv_connector": "lmcache_offload", "kv_role": "offload"},
+            {"kv_connector": "LMCacheMPConnector", "kv_role": "kv_both"},
+            {
+                "kv_connector": "multi",
+                "connectors": [
+                    {"kv_connector": "moriio"},
+                    {"kv_connector": "lmcache_offload"},
+                ],
+            },
+        ],
+    )
+    def test_duplicate_offload_is_rejected(self, monkeypatch, transfer):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", "lmcache")
+        with pytest.raises(ValueError, match="conflicts with an offload connector"):
+            self._kv(["--kv-transfer-config", json.dumps(transfer)])
+
+    def test_unknown_mode_is_rejected(self, monkeypatch):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", "1")
+        with pytest.raises(ValueError, match="expected one of"):
+            self._kv()
+
+    def test_non_object_extra_config_is_rejected(self, monkeypatch):
+        monkeypatch.setenv("ATOM_KV_OFFLOAD", "lmcache")
+        monkeypatch.setenv("ATOM_KV_OFFLOAD_EXTRA_CONFIG", "[]")
+        with pytest.raises(TypeError, match="must be a JSON object"):
+            self._kv()
+
+    def test_extra_config_without_mode_is_rejected(self, monkeypatch):
+        monkeypatch.setenv(
+            "ATOM_KV_OFFLOAD_EXTRA_CONFIG", '{"lmcache.chunk_size": 256}'
+        )
+        with pytest.raises(ValueError, match="ATOM_KV_OFFLOAD is off"):
+            self._kv()