| |
| """K1: write offload/launch/configs/D-kv6-1m-k1mem.env, D-kv6-pack.env and D-kv6-pack-1m.env from D-cal1-kv6.env (remote Engram |
| + remote DSpark capacity, KV 6 GiB/rank) plus the K1 overlay, and the same three on the async production candidate |
| D-kv6-async-process-nap-lazy.env (S8R2 FIX-2: D-kv6-async-1m-k1mem.env, D-kv6-async-pack.env, D-kv6-async-pack-1m.env; the |
| GB10 plan runs these, the sync ones stay as the control). Every K1 item is its own variable with a default that a |
| start can override (empty = off), e.g. K1_TF_ALIAS= offload/launch/start_exp.sh offload/launch/configs/D-kv6-async-pack-1m.env |
| python3 offload/kvpack/make_configs.py # write |
| python3 offload/kvpack/make_configs.py --check # the committed configs equal a fresh generation |
| """ |
| import re |
| import subprocess |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| CFG = ROOT / "offload/launch/configs" |
| SYNC_BASE = "D-cal1-kv6.env" |
| ASYNC_BASE = "D-kv6-async-process-nap-lazy.env" |
| NODROP_SPEC = ('{"method":"dspark","num_speculative_tokens":5,"num_speculative_tokens_per_batch_size":[[1,2,5],[3,4,4]],' |
| '"disable_eagle_block_drop":true}') |
|
|
|
|
| def base_parts(base: str) -> tuple[str, str, str]: |
| text = (CFG / base).read_text() |
| get = lambda k: re.search(rf'^{k}="(.*)"$', text, re.M).group(1) |
| env = get("EXP_DOCKER_ENV") |
| if "VLLM_INDEXER_PREFILL_BUFFER_FACTOR" in env: |
| raise SystemExit("base config already sets the indexer factor") |
| return get("OVERLAYS"), get("EXTRA_DOCKER_ARGS"), env |
|
|
|
|
| def k1_overlays() -> str: |
| return subprocess.run([str(ROOT / "offload/kvpack/stage_overlay.sh"), "overlays"], check=True, |
| capture_output=True, text=True).stdout.strip() |
|
|
|
|
| def render(label: str, header: str, max_len: int | None, kv_pair: str, base: str) -> str: |
| overlays, extra, env = base_parts(base) |
| lines = [ |
| f"# {header}", |
| f"# Generated by offload/kvpack/make_configs.py from {base} + the K1 overlay (offload/kvpack, MANIFEST.json).", |
| "# Before the first start on a node: offload/kvpack/stage_overlay.sh check spark sparkworker (image base == overlay base).", |
| "# K1 switches (override at start; empty = off). Without the overlay mounted none of them exists in the image.", |
| f"K1_KV_PAIR=${{K1_KV_PAIR-{kv_pair}}} # KV page packing: 3 = 0060 (three 64-state pages), 1 = 0057; empty = v5 layout", |
| "K1_DROP_MTP=${K1_DROP_MTP-1} # 0076b: no [4096, 20480] bf16 MTP residual buffer (160 MiB/rank, never read by DSpark)", |
| "K1_TF_TRANSIENT=${K1_TF_TRANSIENT-1} # 0076d: TensorFold prompt scratch per MoE call instead of 1,014 MiB/rank held", |
| "K1_TF_ALIAS=${K1_TF_ALIAS-1} # 0076e (with 0076d): down partials in the rotated rows' storage, call peak 1,014 -> 750 MiB", |
| "K1_INDEXER_FACTOR=${K1_INDEXER_FACTOR-4} # indexer prefill K-gather workspace factor (image default 8); empty = 8", |
| "K1_MEMPROBE=${K1_MEMPROBE-} # 0076a: per-step memory JSONL prefix /root/.cache/dsv41-memprobe-<run>; empty = off", |
| "# memprobe opens <prefix>_r<rank>.jsonl and creates no directory: the prefix must sit directly in /root/.cache", |
| "# (= ~/.cache/dsv41-serve on each node, mounted by the launcher), or the worker dies at its first dummy step.", |
| 'case "$K1_MEMPROBE" in "") ;; /root/.cache/*/*|/root/.cache/) echo "K1_MEMPROBE=$K1_MEMPROBE: use /root/.cache/<name> (no subdirectory)" >&2; return 1 ;; /root/.cache/*) ;; *) echo "K1_MEMPROBE=$K1_MEMPROBE: must start with /root/.cache/" >&2; return 1 ;; esac', |
| "K1_NODROP=${K1_NODROP-} # disable_eagle_block_drop: NO effect in our serve (S8 N-4: with swa_bounded_replay the", |
| "# flagged SWA group is not prefix-cacheable, so the drop never applies); a switch only", |
| f"LABEL={label}", |
| "KV_BYTES_OVERRIDE=${K1_KV_BYTES-6442450944} # 6 GiB/rank; K1_KV_BYTES=1610612736 (1.5 GiB) = pool-pressure run R1b", |
| ] |
| if max_len: |
| lines.append(f"MAX_LEN_OVERRIDE={max_len}") |
| lines += [ |
| f"if [ -n \"$K1_NODROP\" ]; then SPEC_OVERRIDE='{NODROP_SPEC}'; fi", |
| f'OVERLAYS="{overlays} {k1_overlays()}"', |
| f'EXTRA_DOCKER_ARGS="{extra}"', |
| f'EXP_DOCKER_ENV="{env}' |
| '${K1_KV_PAIR:+ VLLM_DSV41_KV_PAIR=$K1_KV_PAIR}' |
| '${K1_DROP_MTP:+ VLLM_DSV41_DROP_MTP_BUFFER=$K1_DROP_MTP}' |
| '${K1_TF_TRANSIENT:+ VLLM_EXL3_TF_PROMPT_TRANSIENT=$K1_TF_TRANSIENT}' |
| '${K1_TF_ALIAS:+ VLLM_EXL3_TF_PROMPT_ALIAS=$K1_TF_ALIAS}' |
| '${K1_INDEXER_FACTOR:+ VLLM_INDEXER_PREFILL_BUFFER_FACTOR=$K1_INDEXER_FACTOR}' |
| '${K1_MEMPROBE:+ VLLM_DSV41_MEMPROBE=$K1_MEMPROBE VLLM_DSV41_MEMPROBE_LAYERS=2,14,20}"', |
| ] |
| return "\n".join(lines) + "\n" |
|
|
|
|
| CONFIGS = { |
| "D-kv6-1m-k1mem.env": ("D-kv6-1m-k1mem", "D-kv6 with max-model-len 1,048,576 (the model's max_position_embeddings; YaRN x16 from 64K), " |
| "no KV packing, 0076b/d/e + indexer factor 4 (upstream: what makes 1M fit on GB10). Needs the x86host " |
| "worker started with --max-model-len 1048576. (D-kv6-1m.env is the coordinator's plain 1M run; not generated here.)", 1_048_576, "", |
| SYNC_BASE), |
| "D-kv6-pack.env": ("D-kv6-pack", "D-kv6 (327,680) with KV page packing 0060 (VLLM_DSV41_KV_PAIR=3) + 0076b/d/e + factor 4; " |
| "predicted pool 2,502,600 tokens (v5: 1,604,926).", None, "3", SYNC_BASE), |
| "D-kv6-pack-1m.env": ("D-kv6-pack-1m", "D-kv6-1m-k1mem with KV page packing 0060 + 0076b/d/e + factor 4; predicted pool " |
| "3,154,934 tokens = 3.01 requests of 1,048,576 (v5: 1,850,243). Needs the 1M x86host worker.", |
| 1_048_576, "3", SYNC_BASE), |
| } |
| |
| |
| |
| CONFIGS |= {name.replace("D-kv6-", "D-kv6-async-", 1): (label.replace("D-kv6-", "D-kv6-async-", 1), header, max_len, kv_pair, |
| ASYNC_BASE) |
| for name, (label, header, max_len, kv_pair, _) in list(CONFIGS.items())} |
|
|
|
|
| def main() -> None: |
| check = "--check" in sys.argv[1:] |
| bad = [] |
| for name, (label, header, max_len, kv_pair, base) in CONFIGS.items(): |
| text = render(label, header, max_len, kv_pair, base) |
| path = CFG / name |
| if check: |
| if not path.exists() or path.read_text() != text: |
| bad.append(name) |
| else: |
| path.write_text(text) |
| print("wrote", path.relative_to(ROOT)) |
| if bad: |
| raise SystemExit("configs differ from a fresh generation: " + ", ".join(bad)) |
| if check: |
| print(f"verified {len(CONFIGS)} configs") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
