feat(offload): enable LMCache offload via ATOM_KV_OFFLOAD env

[yhl-amd]

Motivation

LMCache KV offload is enabled only through --kv-transfer-config. Launchers that own that flag cannot add offload. For example, srtctl (NVIDIA/srt-slurm) generates --kv-transfer-config for P/D workers and rejects it in recipe args. As a result, single-node ATOM AgentX recipes on srt-slurm (MiniMax-M3, Kimi-K3, GLM-5.2 on MI355X) had to drop all of their LMCache offload points.

srtctl passes each role's env to the worker unchanged. It also serializes recipe args with str(), so a JSON-valued flag is fragile there. An environment variable avoids both problems.

Change

Add two environment variables, registered in atom/utils/envs.py. EngineArgs._get_engine_kwargs folds them into kv_transfer_config, so nothing changes after argument parsing:

ATOM_KV_OFFLOAD	Resulting kv_transfer_config
unset / 0 / off / none	unchanged
lmcache	{"kv_connector":"lmcache_offload","kv_role":"offload"} (in-process)
lmcache_mp	{"kv_connector":"lmcache_mp","kv_role":"offload"} (standalone lmcache server)
any mode, plus a P/D connector in --kv-transfer-config	{"kv_connector":"multi","connectors":[<pd>, <offload>]}
any mode, with --kv-transfer-config already multi	offload appended to connectors
any mode, with --kv-transfer-config already naming an offload connector (name or alias)	ValueError

ATOM_KV_OFFLOAD_EXTRA_CONFIG takes a JSON object and becomes the offload connector's kv_connector_extra_config. Both backends read it: lmcache.<field> overrides and, for lmcache_mp, lmcache.mp.* options such as lmcache.mp.port. An unknown mode, a non-object extra config, or an extra config without a mode raises an error. When offload is enabled, the resulting kv_transfer_config is logged at startup.

The first revision of this PR used a --kv-offload-config CLI flag. That flag is removed; the environment variables replace it. Without them, behavior is unchanged. docs/environment_variables.md, the offload README and the module docstring document the variables.

Test

• tests/test_arg_utils_spec.py: 17 environment-variable cases covering every row above, case-insensitive modes, extra-config pass-through, invalid input, and confirming --kv-offload-config is gone. With tests/test_envs.py and tests/entrypoints/test_metrics.py, 112 tests pass locally (CPU torch).
• black is clean on the changed files. ruff check reports only the two E402 imports that tests/test_arg_utils_spec.py already had.
• GPU (MI355X, rocm/atom-dev:nightly_202609231534 with this branch mounted), Llama-3.1-8B-Instruct TP1. The test sends a ~23k-token prompt, sends enough distinct prompts to evict it from the GPU prefix cache, then sends it again:
  • ATOM_KV_OFFLOAD=lmcache: startup logs the composed kv_transfer_config and KV transfer config loaded: {'kv_connector': 'lmcache_offload', ...}. The resend loads 22,784 tokens from LMCache (atom:lmcache_loaded_tokens, atom:prefix_cache_offload_tokens) with zero GPU prefix hits. Output matches the first response, and there are no errors.
  • ATOM_KV_OFFLOAD unset: no offload connector, and every atom:lmcache_* counter stays at 0.
  • ATOM_KV_OFFLOAD=lmcache_mp selects lmcache_mp and passes ATOM_KV_OFFLOAD_EXTRA_CONFIG through (lmcache.mp.port). On DeepSeek-R1-0528 TP8 it registers all ranks and saves to the lmcache server. MP data-plane reuse is out of scope for this PR: lookups there fail in the existing TP-collapse path (No GPU context found ... world size 1), and lmcache_mp requires an MLA backend. The equivalent --kv-transfer-config behaves the same way.

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2414 --add-label <label>