diff --git a/benchmarks/benchmark_lib.sh b/benchmarks/benchmark_lib.sh index 4819f50599..e84caa0b33 100644 --- a/benchmarks/benchmark_lib.sh +++ b/benchmarks/benchmark_lib.sh @@ -49,19 +49,19 @@ require_agentic_kv_offload_backend() { fi return 1 ;; - dram) + dram|nvme|dram+nvme) if [[ "${KV_OFFLOAD_BACKEND:-}" != "$expected_backend" ]]; then - echo "Error: expected KV_OFFLOAD_BACKEND=$expected_backend when KV_OFFLOADING=dram, got '${KV_OFFLOAD_BACKEND:-}'" >&2 + echo "Error: expected KV_OFFLOAD_BACKEND=$expected_backend when KV_OFFLOADING=$KV_OFFLOADING, got '${KV_OFFLOAD_BACKEND:-}'" >&2 exit 1 fi - if [[ ! "${TOTAL_CPU_DRAM_GB:-}" =~ ^[1-9][0-9]*$ ]]; then - echo "Error: DRAM KV offloading requires a positive TOTAL_CPU_DRAM_GB capacity" >&2 + if [[ "$KV_OFFLOADING" != "nvme" && ! "${TOTAL_CPU_DRAM_GB:-}" =~ ^[1-9][0-9]*$ ]]; then + echo "Error: $KV_OFFLOADING KV offloading requires a positive TOTAL_CPU_DRAM_GB capacity" >&2 exit 1 fi return 0 ;; *) - echo "Error: unsupported KV_OFFLOADING value '$KV_OFFLOADING' (expected one of: none, dram)" >&2 + echo "Error: unsupported KV_OFFLOADING value '$KV_OFFLOADING' (expected one of: none, dram, nvme, dram+nvme)" >&2 exit 1 ;; esac @@ -86,18 +86,18 @@ if [[ "$_benchmark_caller" == */agentic/* || exit 1 fi ;; - dram) + dram|nvme|dram+nvme) if [[ -z "${KV_OFFLOAD_BACKEND:-}" || "${KV_OFFLOAD_BACKEND:-}" == "none" ]]; then - echo "Error: KV_OFFLOAD_BACKEND is required when KV_OFFLOADING=dram" >&2 + echo "Error: KV_OFFLOAD_BACKEND is required when KV_OFFLOADING=$KV_OFFLOADING" >&2 exit 1 fi - if [[ ! "${TOTAL_CPU_DRAM_GB:-}" =~ ^[1-9][0-9]*$ ]]; then - echo "Error: DRAM KV offloading requires a positive configured TOTAL_CPU_DRAM_GB capacity" >&2 + if [[ "$KV_OFFLOADING" != "nvme" && ! "${TOTAL_CPU_DRAM_GB:-}" =~ ^[1-9][0-9]*$ ]]; then + echo "Error: $KV_OFFLOADING KV offloading requires a positive configured TOTAL_CPU_DRAM_GB capacity" >&2 exit 1 fi ;; *) - echo "Error: unsupported KV_OFFLOADING value '$KV_OFFLOADING' (expected one of: none, dram)" >&2 + echo "Error: unsupported KV_OFFLOADING value '$KV_OFFLOADING' (expected one of: none, dram, nvme, dram+nvme)" >&2 exit 1 ;; esac diff --git a/benchmarks/single_node/agentic/minimaxm3_fp8_h100_mtp.sh b/benchmarks/single_node/agentic/minimaxm3_fp8_h100_mtp.sh index 9e5f8fc8d6..04a25bf10c 100755 --- a/benchmarks/single_node/agentic/minimaxm3_fp8_h100_mtp.sh +++ b/benchmarks/single_node/agentic/minimaxm3_fp8_h100_mtp.sh @@ -2,7 +2,7 @@ set -eo pipefail set -x -# H100 MiniMax-M3 MXFP8 AgentX with EAGLE3 and optional Mooncake DRAM KV offload. +# H100 MiniMax-M3 MXFP8 AgentX with EAGLE3 and optional DRAM or NVMe KV offload. source "$(dirname "$0")/../../benchmark_lib.sh" @@ -97,6 +97,30 @@ EOF --kv-transfer-config '{"kv_connector":"MooncakeStoreConnector","kv_role":"kv_both","kv_connector_extra_config":{"load_async":true}}' ) +elif [ "$KV_OFFLOADING" = "nvme" ]; then + require_agentic_kv_offload_backend vllm-simple + : "${NVME_OFFLOAD_DIR:?NVME_OFFLOAD_DIR must be mounted by the H100 launcher}" + NVME_OFFLOAD_TOTAL_BYTES=8000000000000 + NVME_OFFLOAD_PER_RANK_BYTES=$((NVME_OFFLOAD_TOTAL_BYTES / TP)) + # vLLM appends .rank_ to give each TP rank its own file. + OFFLOAD_ARGS=( + --kv-transfer-config + "{\"kv_connector\":\"SimpleCPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"kv_offload_backend\":\"disk\",\"disk_path\":\"$NVME_OFFLOAD_DIR/cache.bin\",\"disk_capacity_bytes\":$NVME_OFFLOAD_PER_RANK_BYTES,\"disk_buffer_slots\":4,\"lazy_offload\":false}}" + ) +elif [ "$KV_OFFLOADING" = "dram+nvme" ]; then + require_agentic_kv_offload_backend vllm-native + : "${NVME_OFFLOAD_DIR:?NVME_OFFLOAD_DIR must be mounted by the H100 launcher}" + TOTAL_CPU_DRAM_GIB=$((TOTAL_CPU_DRAM_GB * 1000000000 / 1073741824)) + PER_RANK_GIB=$(((TOTAL_CPU_DRAM_GIB - MODEL_CHECKPOINT_PAGE_CACHE_GIB) / TP - MODEL_CPU_OFFLOAD_GB - MOONCAKE_LOCAL_BUFFER_GIB)) + if (( PER_RANK_GIB <= 0 )); then + echo "Error: CPU DRAM budget is too small for checkpoint cache, model, and DRAM+NVMe KV offload" >&2 + exit 1 + fi + CPU_OFFLOAD_TOTAL_BYTES=$((PER_RANK_GIB * TP * 1073741824)) + OFFLOAD_ARGS=( + --kv-transfer-config + "{\"kv_connector\":\"OffloadingConnector\",\"kv_role\":\"kv_both\",\"kv_connector_extra_config\":{\"spec_name\":\"TieringOffloadingSpec\",\"cpu_bytes_to_use\":$CPU_OFFLOAD_TOTAL_BYTES,\"eviction_policy\":\"lru\",\"secondary_tiers\":[{\"type\":\"fs\",\"root_dir\":\"$NVME_OFFLOAD_DIR\",\"n_read_threads\":32,\"n_write_threads\":16,\"locality\":\"LOCAL\"}]}}" + ) else echo "Error: unsupported KV_OFFLOADING='$KV_OFFLOADING'" >&2 exit 1 diff --git a/configs/nvidia-master.yaml b/configs/nvidia-master.yaml index 72adfa2fa7..8fcb7765a1 100644 --- a/configs/nvidia-master.yaml +++ b/configs/nvidia-master.yaml @@ -7341,7 +7341,7 @@ qwen3.5-fp4-b200-trt-mtp: - { tp: 8, ep: 8, dp-attn: true, spec-decoding: "mtp", conc-list: [128, 256, 1024] } minimaxm3-fp8-h100-vllm-agentic-mtp: - image: vllm/vllm-openai:v0.27.1 + image: ttl.sh/cquil11-vllm-tier-d67b417bca-pr53087-20260901:24h model: MiniMaxAI/MiniMax-M3-MXFP8 model-prefix: minimaxm3 runner: cluster:h100-dgxc @@ -7350,11 +7350,12 @@ minimaxm3-fp8-h100-vllm-agentic-mtp: multinode: false scenarios: agentic-coding: - # The fast sweep places the resident HBM cliff between c5 and c6. + # Narrow validation sweep for physical cache-source attribution. - dram-utilization: 0.80 search-space: - - { tp: 8, spec-decoding: mtp, kv-offloading: none, conc-list: [1, 2, 3, 4, 5] } - - { tp: 8, spec-decoding: mtp, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [6, 8] } + - { tp: 8, spec-decoding: mtp, kv-offloading: dram, kv-offload-backend: { name: mooncake, version: "0.3.11.post1" }, conc-list: [8] } + - { tp: 8, spec-decoding: mtp, kv-offloading: nvme, kv-offload-backend: { name: vllm-simple }, conc-list: [30] } + - { tp: 8, spec-decoding: mtp, kv-offloading: [dram, nvme], kv-offload-backend: { name: vllm-native }, conc-list: [20] } minimaxm3-fp8-h200-vllm-agentic-mtp: image: vllm/vllm-openai:v0.27.1 diff --git a/perf-changelog.yaml b/perf-changelog.yaml index a619ab7c7d..b4c76c64ec 100644 --- a/perf-changelog.yaml +++ b/perf-changelog.yaml @@ -1881,7 +1881,7 @@ description: - "Update GPTOSS-120B FP4 MI355X Atom benchmark (rocm/atom:rocm7.2.2_ubuntu24.04_py3.12_pytorch_release_2.10.0_atom0.1.2.post)" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/1195 - + - config-keys: - dsv4-fp4-b300-vllm @@ -1989,7 +1989,7 @@ - "Add conc=8192 recipe for 1k1k: deepep mega_moe backend with cuda-graph-max-bs 1088, max-running-requests 8192, mem-fraction-static 0.80, swa-full-tokens-ratio 0.3, tokenizer-worker-num 16" - "conc=8192 enables SGLANG_OPT_USE_ONLINE_COMPRESS=1 and --stream-interval 30" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/1209 - + - config-keys: - dsv4-fp4-b300-vllm @@ -2514,7 +2514,7 @@ - "Improves tput/GPU by up to +31% at low concurrency (tp=4, isl=1024, c=4-16)" - "Ref ATOM upstream benchmark run https://github.com/ROCm/ATOM/actions/runs/25686894636" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/1387 - + - config-keys: - dsv4-fp4-b300-vllm-mtp @@ -6577,6 +6577,18 @@ - "Runner: launch_gb300-nv.sh bumped from NVIDIA/srt-slurm@v1.0.29 to v1.0.72 for the dynamo-trt+qwen3.5+fp4 path." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2730 +- config-keys: + - minimaxm3-fp8-h100-vllm-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Run a narrow H100 MiniMax-M3 AgentX validation on the patched vLLM image that exports the bounded cached-token sources device, cpu, disk, p2p, and external and includes vLLM PR #53087's bounded fallback for stalled tier-primary writes." + - "Validate one Mooncake DRAM point at TP8 concurrency 8 and one NVMe-only point at TP8 concurrency 30 using SimpleCPUOffloadConnector's disk backend with 8 TB aggregate capacity." + - "Validate one declarative kv-offloading [dram, nvme] point at TP8 concurrency 20, mapped by the vLLM recipe to OffloadingConnector's TieringOffloadingSpec with an LRU DRAM primary tier and node-local filesystem secondary tier." + - "Mount a job-scoped directory from the H100 node's native NVMe filesystem into the Pyxis container." + - "Collect the vLLM Prometheus endpoint through AIPerf so artifacts include vllm:prompt_tokens_cached_by_source alongside the native KV-offload tiering counters." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2775 + - config-keys: - dsv4-fp4-mi355x-atom-agentic-mtp scenario-type: @@ -6629,7 +6641,7 @@ - "Drop the --hf-overrides that forced text_config use_index_cache true and index_topk_freq 4. The pinned nightly supplies the sparse-PA index-cache and top-k refresh behavior itself, so the override is redundant; because it rewrites the served model config, removing it is architecture-affecting rather than a measurement no-op." - "Drop the TP4 LMCache DRAM KV-offloading sweep points at concurrency 32 and 40 (kv-offload-backend lmcache 0.5.3). The agentic-coding search space is now GPU-resident only, TP4 at concurrency 1 through 32 and TP2 at 1, 2, and 5, so this config no longer reports any DRAM-offload points." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2781 - + - config-keys: - dsv4-fp4-mi355x-sglang-agentic-mtp @@ -6653,7 +6665,7 @@ - "Halve the prefill budget, --max-prefill-tokens and --chunked-prefill-size from 32768 to 16384, matching the B200 sibling recipe qwen3.5_fp4_b200_sglang_mtp.sh." - "Capture the decode CUDA graph to min(2*CONC, 128) instead of min(CONC, 64). The replay keeps --max-running-requests 2*CONC in flight, so a CONC-sized graph dropped every decode batch above CONC onto the eager path. The 128 cap follows the sibling MI355X AgentX recipe dsv4_fp4_mi355x_sglang_mtp.sh." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2737 - + - config-keys: - qwen3.5-fp4-mi355x-sglang - qwen3.5-fp4-mi355x-sglang-mtp @@ -6664,7 +6676,7 @@ - "Serve both arms with an fp8_e4m3 KV cache." - "Route multi-GPU collectives through INT8-quantized ROCm quick all-reduce (ROCM_QUICK_REDUCE_QUANTIZATION=INT8) and drop --enable-aiter-allreduce-fusion, aligning both arms with the published SGLang cookbook recipe for MXFP4 on MI355X (https://docs.sglang.io/cookbook/autoregressive/Qwen/Qwen3.5). The two paths are mutually exclusive in SGLang: the AITER fused AR+RMSNorm path is gated on --enable-aiter-allreduce-fusion, and only with it off do collectives fall back to custom all-reduce where the quick-reduce regime applies. This supersedes the aiter allreduce fusion enabled for these two config keys in #1680, so throughput on this arm is not a like-for-like continuation of the previous MXFP4 series." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2640 - + - config-keys: - qwen3.5-fp8-mi355x-sglang - qwen3.5-fp8-mi355x-sglang-mtp @@ -6673,4 +6685,3 @@ description: - "Bump image from lmsysorg/sglang-rocm:v0.5.16-rocm720-mi35x-20260726 to lmsysorg/sglang-rocm:v0.5.18-rocm720-mi35x-20260828" pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/2763 - diff --git a/runners/launch_h100-dgxc-slurm.sh b/runners/launch_h100-dgxc-slurm.sh index 1334c95542..a8f86e411f 100644 --- a/runners/launch_h100-dgxc-slurm.sh +++ b/runners/launch_h100-dgxc-slurm.sh @@ -289,13 +289,36 @@ else export GPU_COUNT="${GPU_COUNT:-${TP:?TP must be set}}" - salloc --partition=$SLURM_PARTITION --account=$SLURM_ACCOUNT --gres=gpu:$GPU_COUNT --exclusive --time=180 --no-shell --job-name="$RUNNER_NAME" + SALLOC_TIME_LIMIT="${SALLOC_TIME_LIMIT:-300}" + salloc --partition="$SLURM_PARTITION" --account="$SLURM_ACCOUNT" \ + --gres="gpu:$GPU_COUNT" --exclusive --time="$SALLOC_TIME_LIMIT" \ + --no-shell --job-name="$RUNNER_NAME" JOB_ID=$(squeue --name="$RUNNER_NAME" -u "$USER" -h -o %A | head -n1) if [[ -z "$JOB_ID" ]]; then echo "ERROR: failed to resolve H100 Slurm allocation" >&2 exit 1 fi - trap 'rc=$?; scancel "$JOB_ID" 2>/dev/null || true; exit "$rc"' EXIT + cleanup_allocation() { + local rc=$? + trap - EXIT INT TERM + scancel "$JOB_ID" 2>/dev/null || true + exit "$rc" + } + trap cleanup_allocation EXIT INT TERM + + NVME_CONTAINER_MOUNT="" + if [[ "${KV_OFFLOADING:-none}" == "nvme" || "${KV_OFFLOADING:-none}" == "dram+nvme" ]]; then + NVME_HOST_ROOT="/mnt/numa0/enroot/cache/group-$(id -g)" + NVME_HOST_DIR="$NVME_HOST_ROOT/inferencex-kv-$JOB_ID" + srun --jobid="$JOB_ID" bash -c " + set -e + test -w '$NVME_HOST_ROOT' + mkdir -m 700 '$NVME_HOST_DIR' + findmnt -T '$NVME_HOST_DIR' + " + NVME_CONTAINER_MOUNT=",$NVME_HOST_DIR:/kv-offload" + export NVME_OFFLOAD_DIR=/kv-offload + fi # flock-serialize the enroot import so concurrent sweep jobs on the same # shared NFS path don't race each other into 'File already exists' (race @@ -317,12 +340,10 @@ else srun --jobid=$JOB_ID \ --container-image=$SQUASH_FILE \ - --container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE,$AIPERF_MMAP_CACHE_HOST_PATH:/aiperf_mmap_cache \ + --container-mounts=$GITHUB_WORKSPACE:/workspace/,$HF_HUB_CACHE_MOUNT:$HF_HUB_CACHE,$AIPERF_MMAP_CACHE_HOST_PATH:/aiperf_mmap_cache$NVME_CONTAINER_MOUNT \ --no-container-mount-home \ --container-workdir=/workspace/ \ --no-container-entrypoint --export=ALL,PORT=8888,AIPERF_DATASET_MMAP_CACHE_DIR=/aiperf_mmap_cache \ bash benchmarks/single_node/${SCENARIO_SUBDIR}${EXP_NAME%%_*}_${PRECISION}_h100${SPEC_SUFFIX}.sh - scancel $JOB_ID - fi diff --git a/utils/matrix_logic/generate_sweep_configs.py b/utils/matrix_logic/generate_sweep_configs.py index e3872c3d77..48ba2d96d2 100644 --- a/utils/matrix_logic/generate_sweep_configs.py +++ b/utils/matrix_logic/generate_sweep_configs.py @@ -322,7 +322,7 @@ def agentic_dram_offload_gb( budgeted separately if it ever gains its own pool). """ kv_offloading = benchmark.get(Fields.KV_OFFLOADING.value, "none") - if kv_offloading != "dram": + if kv_offloading != "dram" and kv_offloading != ["dram", "nvme"]: return 0 available_mib = min( @@ -356,13 +356,21 @@ def agentic_dram_offload_gb( def agentic_kv_offload_suffix( - kv_offloading: str, + kv_offloading: str | list[str], kv_offload_backend: dict | None, ) -> str: """Return a compact exp-name suffix for agentic KV offload settings.""" if kv_offloading == "none": return "kvnone" - return f"kv{kv_offloading}-{kv_offload_backend['name']}" + mode = "+".join(kv_offloading) if isinstance(kv_offloading, list) else kv_offloading + return f"kv{mode}-{kv_offload_backend['name']}" + + +def agentic_kv_offload_runtime_value(kv_offloading: str | list[str]) -> str: + """Convert declarative tier lists into the workflow's string input.""" + if isinstance(kv_offloading, list): + return "+".join(kv_offloading) + return kv_offloading def multinode_agentic_exp_name( @@ -1029,7 +1037,7 @@ def generate_full_sweep(args, all_config_data, runner_data): Fields.PREFILL.value: prefill, Fields.DECODE.value: decode, Fields.CONC.value: conc_batch, - Fields.KV_OFFLOADING.value: kv_offloading, + Fields.KV_OFFLOADING.value: agentic_kv_offload_runtime_value(kv_offloading), Fields.TOTAL_CPU_DRAM_GB.value: total_cpu_dram_gb, Fields.DURATION.value: duration, Fields.EXP_NAME.value: multinode_agentic_exp_name( @@ -1066,7 +1074,7 @@ def generate_full_sweep(args, all_config_data, runner_data): Fields.DP_ATTN.value: dp_attn if dp_attn is not None else False, Fields.SPEC_DECODING.value: spec_decoding, Fields.CONC.value: conc, - Fields.KV_OFFLOADING.value: kv_offloading, + Fields.KV_OFFLOADING.value: agentic_kv_offload_runtime_value(kv_offloading), Fields.TOTAL_CPU_DRAM_GB.value: total_cpu_dram_gb, Fields.DURATION.value: duration, Fields.EXP_NAME.value: ( @@ -1328,7 +1336,7 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): Fields.PREFILL.value: prefill, Fields.DECODE.value: decode, Fields.CONC.value: conc_batch, - Fields.KV_OFFLOADING.value: kv_offloading, + Fields.KV_OFFLOADING.value: agentic_kv_offload_runtime_value(kv_offloading), Fields.TOTAL_CPU_DRAM_GB.value: total_cpu_dram_gb, Fields.DURATION.value: duration, Fields.EXP_NAME.value: multinode_agentic_exp_name( @@ -1364,7 +1372,7 @@ def generate_test_config_sweep(args, all_config_data, runner_data=None): Fields.DP_ATTN.value: dp_attn if dp_attn is not None else False, Fields.SPEC_DECODING.value: spec_decoding, Fields.CONC.value: conc, - Fields.KV_OFFLOADING.value: kv_offloading, + Fields.KV_OFFLOADING.value: agentic_kv_offload_runtime_value(kv_offloading), Fields.TOTAL_CPU_DRAM_GB.value: total_cpu_dram_gb, Fields.DURATION.value: duration, Fields.EXP_NAME.value: ( diff --git a/utils/matrix_logic/test_generate_sweep_configs.py b/utils/matrix_logic/test_generate_sweep_configs.py index cbde1afc4a..6cf79d1b89 100644 --- a/utils/matrix_logic/test_generate_sweep_configs.py +++ b/utils/matrix_logic/test_generate_sweep_configs.py @@ -2385,6 +2385,56 @@ def test_agentic_node_dram_uses_explicit_gpu_count(self, sample_runner_config): } assert all(entry["duration"] == 3600 for entry in result) + def test_multi_tier_agentic_uses_dram_budget_and_distinct_name( + self, sample_runner_config + ): + config = { + "dsv4-b300-agentic": { + "image": "vllm/vllm-openai:v0.23.0", + "model": "deepseek-ai/DeepSeek-V4-Pro", + "model-prefix": "dsv4", + "precision": "fp4", + "framework": "vllm", + "runner": "cluster:b300-nv", + "multinode": False, + "scenarios": { + "agentic-coding": [{ + "dram-utilization": 0.80, + "search-space": [ + { + "tp": 8, + "kv-offloading": "nvme", + "kv-offload-backend": {"name": "vllm-simple"}, + "conc-list": [7], + }, + { + "tp": 8, + "kv-offloading": ["dram", "nvme"], + "kv-offload-backend": {"name": "vllm-native"}, + "conc-list": [7], + }, + ], + }], + }, + }, + } + args = argparse.Namespace( + config_keys=["dsv4-b300-agentic"], + seq_lens=None, + conc=None, + scenario_type=["agentic-coding"], + runner_node_filter=None, + ) + + result = generate_test_config_sweep(args, config, sample_runner_config) + + assert [entry["kv-offloading"] for entry in result] == ["nvme", "dram+nvme"] + assert [entry["total-cpu-dram-gb"] for entry in result] == [0, 2399] + assert [entry["exp-name"] for entry in result] == [ + "dsv4_tp8_conc7_kvnvme-vllm-simple", + "dsv4_tp8_conc7_kvdram+nvme-vllm-native", + ] + def test_agentic_node_dram_rejects_tp_above_runner_gpus(self, sample_runner_config): config = { "dsv4-b300-agentic": { diff --git a/utils/matrix_logic/test_validation.py b/utils/matrix_logic/test_validation.py index f20c36b3de..be34d6f393 100644 --- a/utils/matrix_logic/test_validation.py +++ b/utils/matrix_logic/test_validation.py @@ -530,6 +530,49 @@ def test_dram_kv_offload_requires_dram_utilization(self): }], }) + def test_multi_tier_kv_offload_requires_dram_utilization(self): + with pytest.raises(Exception, match="dram-utilization"): + AgenticCodingConfig(**{ + "search-space": [{ + "tp": 8, + "kv-offloading": ["dram", "nvme"], + "kv-offload-backend": {"name": "vllm-native"}, + "conc-list": [7, 8], + }], + }) + + def test_multi_tier_kv_offload_accepts_dram_capacity_config(self): + config = AgenticCodingConfig(**{ + "dram-utilization": 0.99, + "search-space": [{ + "tp": 8, + "kv-offloading": ["dram", "nvme"], + "kv-offload-backend": {"name": "vllm-native"}, + "conc-list": [7, 8], + }], + }) + + assert config.search_space[0].kv_offloading == ["dram", "nvme"] + assert config.dram_utilization == 0.99 + + @pytest.mark.parametrize("kv_offloading", ["nvme", ["dram", "nvme"]]) + def test_nvme_kv_offload_rejects_multinode_agentic_entries( + self, kv_offloading + ): + with pytest.raises(Exception, match="only for single-node agentic"): + AgenticCodingSearchSpaceEntry(**{ + "worker": { + "tp": 8, + "pp": 1, + "ep": 1, + "dp-attn": False, + }, + "num-nodes": 2, + "kv-offloading": kv_offloading, + "kv-offload-backend": {"name": "vllm-native"}, + "conc-list": [7], + }) + def test_agentic_search_space_rejects_total_cpu_dram_gb(self): with pytest.raises(Exception, match="total-cpu-dram-gb"): AgenticCodingSearchSpaceEntry(**{ diff --git a/utils/matrix_logic/validation.py b/utils/matrix_logic/validation.py index 436ea97b39..201c69c738 100644 --- a/utils/matrix_logic/validation.py +++ b/utils/matrix_logic/validation.py @@ -14,6 +14,11 @@ CLUSTER_LABEL_PREFIX = "cluster:" DEFAULT_AGENTIC_DURATION_SECONDS = 3600 +KVOffloadingTier = Literal["dram", "nvme"] +KVOffloadingConfig = Union[ + Literal["none", "dram", "nvme"], + List[KVOffloadingTier], +] """ The below class defines the field names expected to be present in the JSON entries @@ -318,7 +323,7 @@ class SingleNodeAgenticMatrixEntry(BaseModel): default="none", alias=Fields.SPEC_DECODING.value ) conc: int - kv_offloading: Literal["none", "dram"] = Field( + kv_offloading: Literal["none", "dram", "nvme", "dram+nvme"] = Field( alias=Fields.KV_OFFLOADING.value ) kv_offload_backend: Optional[KVOffloadBackendMetadata] = Field( @@ -513,7 +518,13 @@ def _validate_kv_offload_fields(self): f"{Fields.KV_OFFLOADING.value}" ) return self - if self.kv_offloading == "none": + if isinstance(self.kv_offloading, list): + if self.kv_offloading != ["dram", "nvme"]: + raise ValueError( + f"The only supported multi-tier {Fields.KV_OFFLOADING.value} " + "configuration is ['dram', 'nvme']" + ) + elif self.kv_offloading == "none": if backend is not None: raise ValueError( f"{Fields.KV_OFFLOAD_BACKEND.value} can only be set when " @@ -644,7 +655,7 @@ class AgenticCodingSearchSpaceEntry(BaseModel): decode: Optional[WorkerConfig] = None num_nodes: Optional[int] = Field( default=None, alias=Fields.NUM_NODES.value, gt=0, strict=True) - kv_offloading: Optional[Literal["none", "dram"]] = Field( + kv_offloading: Optional[KVOffloadingConfig] = Field( default=None, alias=Fields.KV_OFFLOADING.value ) kv_offload_backend: Optional[KVOffloadBackendMetadata] = Field( @@ -692,6 +703,11 @@ def validate_topology_fields(self): ) _validate_tp_context_topology(self) if has_aggregate_worker or has_complete_multinode: + if self.kv_offloading in ("nvme", ["dram", "nvme"]): + raise ValueError( + f"{Fields.KV_OFFLOADING.value}={self.kv_offloading!r} is " + "currently supported only for single-node agentic entries" + ) explicitly_single_node_fields = { "pp", "dcp_size", @@ -726,7 +742,7 @@ class AgenticCodingConfig(BaseModel): @model_validator(mode='after') def validate_dram_offload_capacity(self): for entry in self.search_space: - if entry.kv_offloading != "dram": + if entry.kv_offloading != "dram" and entry.kv_offloading != ["dram", "nvme"]: continue if self.dram_utilization is None: raise ValueError(