feat(inference): configure KV offload per P/D role

docs/inference.md

@@ -219,7 +219,7 @@ Maximizing KV-Cache space is crucial to support high-concurrency workloads. You
 The `type` field selects the backend:
 
 - `native` — vLLM's built-in offloading. CPU-only uses `OffloadingConnector`; CPU+disk uses `TieringOffloadingSpec` (a CPU primary tier with a filesystem secondary tier). Fully self-contained — no extra processes.
-- `mooncake` — a [Mooncake](https://github.com/kvcache-ai/Mooncake) **shared distributed store** (SLURM only). One `mooncake_master` + metadata server runs on the head inference node; every inference node runs a `mooncake_client` that contributes its DRAM (and, with `disk`, SSD) segment to that *single* pool. Because blocks are keyed by model + parallel rank + content hash (no instance id), a prefix cached by one node/replica is reusable by all of them over RDMA — pooling every node's CPU RAM into one KV cache. Use `native` for local/single-process runs.
+- `mooncake` — a [Mooncake](https://github.com/kvcache-ai/Mooncake) **shared distributed store** (SLURM only). One `mooncake_master` + metadata server runs on the head inference node. Nodes configured with positive CPU capacity run a `mooncake_client` contributing DRAM (and, with `disk`, SSD) to that pool. Other Mooncake-connected nodes can reuse cached blocks over RDMA when their cache keys and layouts match. Use `native` for local/single-process runs.
 
 ```toml
 # Native CPU offload (reserves 128GB of CPU KV cache for this instance)
@@ -245,12 +245,67 @@ num_bytes = 128_000_000_000
 path = "/scratch/kv"
 ```
 
-For `native`, `cpu.num_bytes` is the aggregate CPU KV pool for the instance (vLLM shards it across workers). For `mooncake`, `cpu.num_bytes` is the DRAM each node contributes to the shared pool (so the total pool ≈ `num_bytes × #inference-nodes`); the store uses RDMA, so it requires an RDMA-capable fabric. Enabling offload automatically enables prefix caching.
+For `native`, `cpu.num_bytes` is the aggregate CPU KV pool for the instance (vLLM shards it across workers). For `mooncake`, `cpu.num_bytes` is the DRAM each node contributes to the shared pool; the store uses RDMA, so it requires an RDMA-capable fabric. Enabling offload automatically enables prefix caching.
+
+P/D deployments can override either backend separately with
+`inference.deployment.prefill_kv_cache_offload` and
+`inference.deployment.decode_kv_cache_offload`. An omitted override inherits
+`inference.kv_cache_offload`; `"None"` disables it for that role. For example,
+native offloading on prefill only:
+
+```toml
+[inference.deployment]
+type = "disaggregated"
+decode_kv_cache_offload = "None"
+
+[inference.deployment.prefill_kv_cache_offload]
+type = "native"
+[inference.deployment.prefill_kv_cache_offload.cpu]
+num_bytes = 128_000_000_000
+```
+
+Multi-node RL can also contribute trainer host RAM to the same Mooncake pool:
+
+```toml
+[deployment.trainer_mooncake]
+num_bytes = 536870912000 # 500 GiB per eligible trainer node
+# device_name = "mlx5_0" # optional; empty selects devices automatically
+```
+
+This requires SLURM and Mooncake offload on at least one inference role. It starts
+one DRAM-only storage client per trainer node, excluding the node hosting the
+orchestrator. With four trainer nodes and the default orchestrator placement,
+three nodes contribute 500 GiB each. If `deployment.orchestrator_on_inference`
+is enabled, all trainer nodes contribute. The size is independent of inference
+storage capacity; allow RAM for trainer activation and optimizer offloading when
+choosing it. Omit this section to disable trainer contributions. Clients share
+the inference master and its lease policy, stop with the Slurm job, and log to
+`<run_dir>/mooncake/trainer_<rank>/client.log`.
+
+For Mooncake, `cpu.num_bytes = 0` keeps a role connected to the shared pool
+without starting a storage client on its nodes. Staging buffers still consume
+RAM; `local_buffer_bytes` controls their per-worker size. Set positive capacity
+on the storage-owning role. Native offloading requires positive CPU capacity.
+The launcher starts one job-scoped Mooncake master and storage clients only on
+nodes contributing capacity. Both the standalone inference and RL launchers
+write separate resolved role configs. Mooncake disk tiers share the master's
+disk path across roles.
+
+An explicit `vllm.kv_transfer_config`, including one supplied through a role's
+`*_vllm_overrides`, takes precedence over automatic connector composition. Include
+all required connectors in that chain. The typed offload config still controls
+Mooncake service provisioning. This permits composing NIXL, Mooncake and
+HiSparse without replacing their connector-specific settings.
+
+For external-DP Slurm deployments, `deployment.local_rank_numa_nodes` and
+`deployment.local_rank_ucx_devices` optionally set NUMA binding and UCX NIC
+selection per local DP rank. Each list must have one entry per local rank;
+omit them to retain automatic placement.
 
 
 ### Optimized P/D disaggregation deployment
 
-For optimal P/D disaggregation deployment, we automatically set the decode `all2all_backend` to `deepep_low_latency` and the prefill `all2all_backend` to `deepep_high_throughput`. We currently don't support customizing all2all backends for P/D disaggragation out of the box. You can do this by overriding the slurm template only.
+P/D disaggregation defaults to `deepep_low_latency` for decode and `deepep_high_throughput` for prefill. Set `all2all_backend` in `deployment.prefill_vllm_overrides` or `deployment.decode_vllm_overrides` to select another backend for either role.
 
 For KV cache transfer, we utilize the NIXL connector. This is the default and only currently supported connector. We aim to support more advanced options, such as D->P transfer, or Mooncake Connector in the future.
 

examples/advanced/glm-5.3/infer/pd-nvfp4-b300-2p1d-disk.toml

@@ -0,0 +1,60 @@
+# Layer on top of pd-nvfp4-b300.toml. The disk tier uses shared Weka storage.
+[deployment.trainer_mooncake]
+num_bytes = 1000000000000 # 1 TB per trainer node, excluding the orchestrator host.
+device_name = "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_9,mlx5_12,mlx5_13"
+
+[trainer.log]
+level = "debug"
+
+[trainer.env_vars]
+TORCHINDUCTOR_COMPILE_THREADS = "1"
+
+[trainer.optim]
+type = "sign_sgd"
+
+[trainer.model]
+cp = 8
+
+[trainer.model.ac]
+mode = "full"
+
+[trainer.model.ac_offloading]
+max_inflight_activations = 1
+
+[inference.router]
+request_timeout_seconds = 3600
+log_level = "info"
+
+[inference.deployment]
+num_prefill_replicas = 2
+num_decode_replicas = 1
+
+[inference.deployment.prefill_kv_cache_offload]
+kv_lease_ttl_ms = 1200000
+
+[inference.deployment.decode_kv_cache_offload]
+kv_lease_ttl_ms = 1200000
+
+[inference.deployment.prefill_kv_cache_offload.disk]
+path = "/home/matej/prime-rl/outputs/glm53-rl-disk/job_${SLURM_JOB_ID}/node_${INFER_NODE_RANK}"
+
+[inference.deployment.decode_kv_cache_offload.disk]
+path = "/home/matej/prime-rl/outputs/glm53-rl-disk/job_${SLURM_JOB_ID}/node_${INFER_NODE_RANK}"
+
+[inference.env_vars]
+# Two prefill storage clients, each with a 2 TB logical disk budget.
+MOONCAKE_OFFLOAD_STORAGE_BACKEND_DESCRIPTOR = "bucket_storage_backend"
+MOONCAKE_OFFLOAD_TOTAL_SIZE_LIMIT_BYTES = "2000000000000"
+MOONCAKE_OFFLOAD_BUCKET_MAX_TOTAL_SIZE = "2000000000000"
+MOONCAKE_OFFLOAD_BUCKET_EVICTION_POLICY = "fifo"
+
+[orchestrator]
+batch_size = 128
+
+[orchestrator.log]
+level = "debug"
+
+[orchestrator.concurrency]
+initial_inflight = 2048
+min_inflight = 2048
+max_inflight = 2048

examples/advanced/glm-5.3/infer/pd-nvfp4-b300-4p2d-ram.toml

@@ -0,0 +1,68 @@
+# Layer on top of pd-nvfp4-b300.toml: 4P + 2D, DRAM-only Mooncake.
+[deployment.trainer_mooncake]
+num_bytes = 1000000000000 # 1 TB per trainer node, excluding the orchestrator host.
+device_name = "mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_9,mlx5_12,mlx5_13"
+
+[trainer.log]
+level = "debug"
+
+[trainer.env_vars]
+TORCHINDUCTOR_COMPILE_THREADS = "1"
+
+[trainer.optim]
+type = "sign_sgd"
+
+[trainer.model]
+cp = 8
+
+[trainer.model.ac]
+mode = "full"
+
+[trainer.model.ac_offloading]
+max_inflight_activations = 1
+
+[inference.router]
+request_timeout_seconds = 7200
+log_level = "info"
+
+[inference.deployment]
+num_prefill_replicas = 4
+num_decode_replicas = 2
+
+[inference.deployment.prefill_kv_cache_offload]
+kv_lease_ttl_ms = 1200000
+
+[inference.deployment.decode_kv_cache_offload]
+kv_lease_ttl_ms = 1200000
+
+[orchestrator]
+batch_size = 256
+max_off_policy_steps = 32
+
+[orchestrator.algo.length_penalty]
+type = "linear"
+num_output_tokens_weight = 0.2
+num_input_tokens_weight = 0.2
+num_turns_weight = 0.2
+
+[orchestrator.log]
+level = "debug"
+
+[orchestrator.concurrency]
+initial_inflight = 3600
+min_inflight = 3600
+max_inflight = 3600
+
+# 32 workers x 128 slots provide capacity for 3,600 fixed inflight rollouts.
+[[orchestrator.train.source]]
+name = "scaleswe-q0"
+env.taskset.id = "scaleswe"
+env.taskset.filter_fn = "lambda row: sum(row['instance_id'].encode()) % 4 == 0"
+env.agent.harness.id = "bash"
+env.agent.harness.compaction.summarize_at_tokens = "None"
+env.agent.runtime.labels = ["glm53-pd-train", "int4-syn-gen-bash"]
+env.agent.runtime.creates_per_min = 720
+env.agent.timeout.rollout = 21600
+serve.pool.type = "static"
+serve.pool.num_workers = 32
+serve.max_concurrent = 128

examples/advanced/glm-5.3/infer/pd-nvfp4-b300.md

@@ -0,0 +1,184 @@
+# GLM-5.3 NVFP4 Scale-SWE RL on B300
+
+From the repository root:
+
+```bash
+uv sync --frozen --all-extras --all-packages
+uv run --no-sync rl @ examples/advanced/glm-5.3/infer/pd-nvfp4-b300.toml --dry-run
+uv run --no-sync rl @ examples/advanced/glm-5.3/infer/pd-nvfp4-b300.toml
+```
+
+The dry run writes resolved trainer, orchestrator, prefill/decode and environment
+configs plus the Slurm script without allocating nodes. The launch requests six
+exclusive eight-GPU nodes: four trainer nodes, one prefill node and one decode
+node, with no job time limit. The router listens on port 8000 on the first
+inference node and accepts the alias `glm53` as well as the shared model path.
+
+The Slurm startup sync uses the frozen lockfile; component processes use
+`UV_NO_SYNC=1` after that sync. This avoids re-resolving remote wheel metadata
+on each process launch. The model checkpoint is already local; `HF_HOME` is
+not required to find its weights.
+
+Both the router and orchestrator allow 5,400 seconds for inference readiness.
+Cold FlashInfer compilation and autotuning can exceed the default 1,800-second
+orchestrator timeout. Compilation and autotune artifacts remain under the user's
+persistent cache directory across jobs.
+
+Training stops after 1,000 optimizer steps. `clean=true` deletes this named
+run's directory before each fresh launch, including its previous logs and
+checkpoints; explicit `--resume` instead follows resume semantics. The run
+directory is `outputs/glm53-nvfp4-rl/scaleswe-nvfp4-4over6`.
+
+Before a restart, preserve any needed logs outside that directory and wait for
+the previous Slurm allocation to finish cleanup. Check `nvidia-smi` and process
+start times on its hosts for leftover GPU allocations and job-owned Mooncake
+processes before launching the replacement. Clean up confirmed orphaned processes
+from the stopped job; keep the kernel compilation caches for reuse.
+
+Trainer and orchestrator log to the shared W&B project `glm53-nvfp4-rl`, with
+the run name inherited from `[run]`. The Slurm launcher sources the repository's
+git-ignored `.env`; `pre_run_command` exports `WANDB_API_KEY` to its child
+processes. Keep the key in `.env`, outside the TOML and generated configs.
+
+| Trainer setting | Value |
+|---|---|
+| Nodes / GPUs | 4 / 32 |
+| Expert parallelism | EP8, default torch dispatch with BF16 transport |
+| Context parallelism | CP4, ring |
+| Sequence length | 65,536 |
+| Optimizer | AdamW; default LR 1e-6, weight decay 0.01 |
+| CPU offload | Optimizer states; full CPU offload disabled |
+| Activation checkpointing | Selective, every layer, default retained operations |
+| Activation offloading | Enabled, one activation in flight |
+| Expert compute | NVFP4 4over6 for the first 85% of layers; BF16 elsewhere |
+| Index cache | Native checkpoint IndexShare schedule, matching inference |
+| Weight transfer | NCCL, BF16 weights quantized online by inference |
+
+Both sides start from the existing local `zai-org/GLM-5.3-BF16` download at
+revision `9d2398f478cab2de883137db3a36ad2c96205e24`. Trainer loading uses its
+converted Prime checkpoint. `apply_to="85%"` selects the first 66 of 78 layers
+(rounding down): routed experts in layers 3–65 use NVFP4, while the final 12
+layers remain BF16. Dense and shared-expert modules also remain BF16. Inference
+uses the same layer selection with `nvfp4_per_token` online quantization and
+4over6's 448 normalization / MAE scoring. NVFP4 trainer backward uses
+dequantized forward operands with BF16 GEMMs. Optimization and reduction dtypes
+retain their model defaults.
+
+| Inference setting | Prefill | Decode |
+|---|---|---|
+| Parallelism | DPEP8, TP1 | DPEP8, TP1 |
+| All-to-all backend | FlashInfer NVLink one-sided | FlashInfer NVLink one-sided |
+| Batched tokens | 32,768 | 320 |
+| Maximum sequences per GPU | 256 | 320 |
+| GPU memory utilization | 0.90 | 0.90 |
+| Graphs | None | Full + piecewise, breakable |
+| Mooncake store RAM per node | 2,560 GiB | 0 |
+| HiSparse host RAM per GPU | Disabled | 300 GiB |
+
+Both roles use FP8 KV, a 65,536-token context limit, the checkpoint's native
+IndexShare schedule, natural expert routing, and no EPLB, DBO or MTP. The
+schedule has 21 full-indexer layers and 57 layers that reuse indices. A trainer
+frequency override would replace this schedule and request weights absent from
+the checkpoint, so `trainer.model.index_cache` is omitted. The pinned wheel's
+DeepEP low-latency expert backend does not support online per-token NVFP4;
+both inference roles use FlashInfer one-sided. Routing is sticky least
+loaded with retries disabled. Each rank binds to its configured NUMA node and
+UCX NIC. These mappings and host-memory capacities target this B300 cluster;
+adjust them for a different machine topology.
+
+Use the pinned HiSparse/NVFP4 wheel with the standard Python frontend and native
+`NixlConnector`. Mooncake master/storage processes belong to the Slurm job, so
+cancellation also releases their memory. Both roles connect to the same store
+pool with separate role cache-key prefixes and 512 MiB staging buffers per
+connector. The prefixes identify the BF16-source online-quantized configuration.
+
+The training source uses Scale-SWE quarter-zero tasks selected by the instance-ID
+filter, bash harness, automatic compaction, 720 sandbox creates per minute, a
+six-hour rollout timeout, and no agentic judge. It retains 16 rollouts per task
+and `clear_thinking=false`. The static environment pool has 20 workers with 128
+concurrent slots each. Training sandboxes use `glm53-pd-train` and
+`int4-syn-gen-bash` labels.
+
+Rollout concurrency is fixed at 1,536 by setting initial, minimum and maximum
+in-flight episodes to the same value. Episodes doing tool work also occupy slots,
+so this does not guarantee 1,536 simultaneous inference requests.
+`orchestrator.batch_size=512` counts trainer-bound traces per optimizer step,
+independently of concurrency. The default `constant_trainer_batch_size=true`
+prunes zero-advantage RL samples before counting toward the batch and keeps
+collecting until 512 useful traces remain. The debug algorithm assigns advantage
+1 to every valid trainable rollout regardless of reward, so uniform-reward groups
+remain trainable. Errored rollouts are still excluded. This exercises the RL
+pipeline; the constant advantage is not a task-learning objective.
+
+Sampling uses temperature 1, top-p 0.95 and top-k 512, matching the prior sampling
+replay run. Top-p comes from the BF16 checkpoint's generation config; 512 is the
+bounded top-k used for replay (the checkpoint does not specify a top-k).
+Sampling-mask capture is enabled on inference, and the trainer replays those
+masks, including with CP. Inference returns processed logprobs so rollout and
+trainer probabilities use the same truncated distribution. Router replay is
+disabled. This config does not launch an additional eval or profiler.
+
+## Two prefill replicas with filesystem offload
+
+For two independent EP8 prefill replicas, one EP8 decode replica, and fixed
+2,048 rollout concurrency, add the disk-tier overlay:
+
+```bash
+uv run --no-sync rl @ examples/advanced/glm-5.3/infer/pd-nvfp4-b300.toml \
+  @ examples/advanced/glm-5.3/infer/pd-nvfp4-b300-2p1d-disk.toml
+```
+
+This uses seven nodes including the four trainer nodes. Each prefill node
+contributes 2,560 GiB DRAM and a 2 TB logical filesystem-offload budget, for
+4 TB total disk capacity. The filesystem tier uses shared Weka under
+`/home/matej/prime-rl/outputs/glm53-rl-disk/job_JOB_ID/node_RANK`; it does not
+measure local-NVMe performance. Separate directories isolate the storage clients.
+Mooncake's bucket limit controls logical occupancy, not a physical filesystem
+quota. Decode contributes no storage and retains access to both store tiers.
+The three trainer nodes without the orchestrator each contribute an additional
+1 TB (1,000,000,000,000 bytes) of DRAM to the same pool, configured separately
+under `deployment.trainer_mooncake.num_bytes`. Trainer storage is DRAM-only.
+Trainer storage explicitly uses the same eight 800 Gb/s NICs as inference;
+auto-discovery also selects incompatible interfaces on these hosts.
+Trainer and orchestrator logging use DEBUG; router logging uses INFO.
+The overlay uses CP=8, EP=8, full activation checkpointing with activation
+offloading, and batch size 128 (eight groups of 16). It uses stateless SignSGD,
+retaining the optimizer CPU-offload setting and disabling full CPU offload.
+The debug algorithm is inherited from the base config. The shared Mooncake master
+uses a 20-minute hard KV lease,
+and the router allows requests to run for 60 minutes.
+Trainer compilation uses `TORCHINDUCTOR_COMPILE_THREADS=1` to avoid a compiler
+pool shutdown deadlock in the installed PyTorch runtime. Compilation and caches
+remain enabled; inference keeps its own compiler settings.
+
+## Inference Grafana
+
+After Slurm starts the job, run the monitoring helper separately on the login node:
+
+```bash
+uv run --no-sync tools/inference_monitoring.py \
+  outputs/glm53-nvfp4-rl/scaleswe-nvfp4-4over6 JOB_ID
+```
+
+The helper uses the installed Prometheus/Grafana binaries under
+`~/.local/share/glm-monitoring` (override with `--binaries`). It discovers role
+hosts from the job's launch log, reads ports and GPU counts from the generated
+Slurm script, and scrapes every engine plus the router every 15 seconds. The
+dashboard includes per-role and per-rank throughput, latency, queues, KV usage,
+HiSparse reloads, Mooncake operations and NIXL transfers. Per-GPU rates use the
+allocated GPU count. This is independent of the disabled Prime-RL dashboard.
+
+On your laptop:
+
+```bash
+ssh -N -L 3000:127.0.0.1:3000 nebius
+```
+
+Open <http://localhost:3000/d/inference-live>. Grafana is a read-only view bound
+to loopback; Prometheus also binds to loopback. The helper stays in the foreground
+and stops its processes when the Slurm job ends or when interrupted. It does not
+install a system service or control the RL job. Its data and logs live under
+the run's `monitoring/job_JOB_ID/` directory; `clean=true` removes them on a fresh
+run. Retention is limited to seven days or 5 GB, whichever limit is reached first.
+Use a new job ID to monitor another deployment; `--grafana-port`,
+`--prometheus-port`, and `--router-port` allow multiple monitors concurrently.