diff --git a/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/glm5.2/atom/mi355x-fp4/agentx/disagg-variants.yaml b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/glm5.2/atom/mi355x-fp4/agentx/disagg-variants.yaml new file mode 100644 index 0000000000..21c171b2b7 --- /dev/null +++ b/inferencex-e2e/benchmarks/multi_node/srt-slurm-recipes/glm5.2/atom/mi355x-fp4/agentx/disagg-variants.yaml @@ -0,0 +1,360 @@ +base: + schema: 2 + name: mi355x-glm52-atom-agentx-disagg-mtp + slurm: + time_limit: "08:00:00" + resources: + gpu_type: mi355x + gpus_per_node: 8 + frontend: + type: atomesh + enable_multiple_frontends: false + args: + log-level: info + disable-health-check: true + disable-circuit-breaker: true + prometheus-port: 29100 + prefill-policy: round_robin + decode-policy: round_robin + atom-pd-rank-mapping-policy: none + engine: atom + roles: + prefill: + nodes: 1 + workers: 1 + gpus: 4 + env: + <<: &environment + PYTHONUNBUFFERED: "1" + PYTHONHASHSEED: "0" + # srtctl's pre-download sees only HF_HOME, so the node-local cache is mounted at $HF_HOME/hub. + HF_HOME: /mnt/hf_home + HF_HUB_CACHE: /mnt/hf_home/hub + AITER_LOG_LEVEL: WARNING + AITER_QUICK_REDUCE_QUANTIZATION: INT4 + AITER_USE_FLYDSL_MOE_SORTING: "1" + ATOM_MLA_PAGE_SIZE: "1" + ATOM_ONLINE_QUANT_STREAMING: "0" + ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB: "2047" + ATOM_USE_TRITON_MLA: "0" + TOPK_FORCE_PATH: one + MAX_JOBS: "16" + HSA_ENABLE_IPC_MODE_LEGACY: "1" + MC_DISABLE_HIP: "1" + # Rail-isolated RDMA: each transfer uses the same rail on both ends. + MC_GID_INDEX: "@fabric.rdma-gid-index" + ATOM_MOONCAKE_MATCHED_RAILS: auto + NCCL_IB_DISABLE: "1" + VLLM_PP_LAYER_PARTITION: "20,20,20,18" + OFFLOAD_PROFILE: "1" + OFFLOAD_MIN_LOAD_TOKENS: "0" + args: + <<: &server + served-model-name: amd/GLM-5.2-MXFP4 + trust-remote-code: true + kv_cache_dtype: fp8 + block-size: 16 + # model.layers.78 is the MTP head, shipped in BF16; excluding it keeps + # online ptpc_fp8 quantization off the draft while experts stay MXFP4. + online_quant_config: '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*.mlp.gate","model.layers.[0-9].mlp.*expert*","model.layers.[1-6][0-9].mlp.*expert*","model.layers.7[0-7].mlp.*expert*","model.layers.78.*"]}' + level: 3 + cudagraph-mode: FULL + method: mtp + timeout-keep-alive: 900 + pipeline-parallel-size: 4 + enforce-eager: true + gpu-memory-utilization: 0.85 + max-num-seqs: 512 + max-num-batched-tokens: 8192 + decode: + nodes: 1 + workers: 1 + gpus: 4 + env: *environment + args: *server + sbatch_directives: + cpus-per-task: "128" + mem: "0" + srun_options: + mem: "0" + container-writable: "" + container-remap-root: "" + extra_mount: + # The node-local hub cache the single-node GLM-5.2 points already fill. + - /var/lib/hf-hub-cache:/mnt/hf_home/hub + health_check: + max_attempts: 1080 + interval_seconds: 5 + benchmark: + env: + AIPERF_REQUIRED_SERVER_METRIC_PREFIX: "atom:" + AIPERF_APPLY_CHAT_TEMPLATE: "true" + TOKENIZERS_PARALLELISM: "false" + TRANSFORMERS_VERBOSITY: error + +override_1p1d_mcstore: + host_setup: + nodes: all + timeout_seconds: 1800 + commands: + # A recipe host_setup replaces the cluster default, which also drops the page cache under DRAM KV offloading. + - IBDEVICES=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/setup.sh" + # Owners take 786432 2 MiB pages, half per NUMA node; reserve 2% more above the pages host services hold. + # Only the prefill node holds owners; srtctl starts the prefill on the head node, which launches this srun. + - >- + if hostname -I | tr ' ' '\n' | grep -qxF "${SLURM_LAUNCH_NODE_IPADDR:?srun did not set SLURM_LAUNCH_NODE_IPADDR}"; then + echo "[hugepage] $(hostname -s) is the head node; reserving the owners' pages"; + held=$(awk '/^HugePages_Total:/ {total = $2} /^HugePages_Free:/ {free = $2} /^HugePages_Rsvd:/ {rsvd = $2} END {print total - free + rsvd}' /proc/meminfo); + HUGEPAGES_TARGET=$(( held + 786432 * 102 / 100 )) HUGEPAGE_GROW_TOLERANCE_PCT=1 + bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/hugepage_sanity.sh" || exit 1; + for node in 0 1; do + free=$(cat /sys/devices/system/node/node$node/hugepages/hugepages-2048kB/free_hugepages); + echo "[hugepage] NUMA$node: $free free 2 MiB pages, owners need 393216"; + [[ $free -ge 393216 ]] || exit 1; + done; fi + teardown: + - HUGEPAGES_TARGET=0 bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/hugepage_sanity.sh" + services: + - &master + name: mc-master + type: generic + placement: + node: prefill + start: before_workers + critical: true + inherit_discovery_env: false + # numa_exec.py binds the process's CPUs and memory to that NUMA node. + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_master, --rpc_port=26051, --rpc_thread_num=32, --enable_http_metadata_server=true, --http_metadata_server_host=0.0.0.0, --http_metadata_server_port=26080, --metrics_port=26090, --default_kv_lease_ttl=10000, --eviction_high_watermark_ratio=0.90, --eviction_ratio=0.05, --allocation_strategy=random, --memory_allocator=offset, --client_ttl=10] + env: &store_env + # A Mooncake process that sees a GPU probes every buffer through HIP. + HIP_VISIBLE_DEVICES: "-1" + CUDA_VISIBLE_DEVICES: "-1" + # Peers with different QP counts cannot connect. + MC_NUM_QP_PER_EP: "1" + MC_GID_INDEX: "@fabric.rdma-gid-index" + readiness: + port: 26051 + timeout_seconds: 300 + - &owner + <<: *master + name: mc-owner-numa0 + # Each owner registers on rdma0-3, so every stage reaches it on its own rail. + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "0", mooncake_client, "--host={node_ip}", --port=26052, --master_server_address=127.0.0.1:26051, --metadata_server=http://127.0.0.1:26080/metadata, --protocol=rdma, "--device_names=rdma0,rdma1,rdma2,rdma3", --global_segment_size=824633720832, --local_buffer_size=0, --threads=4] + env: + <<: *store_env + # Back the segment with the reserved 2 MiB pages; the NIC cannot register 64 GiB regions of 4 KiB pages. + MC_STORE_USE_HUGEPAGE: "1" + MC_MAX_MR_SIZE: "68719476736" + MC_TCP_BIND_ADDRESS: "{node_ip}" + # The remapped container root lacks host CAP_IPC_LOCK, so registering the segment counts against RLIMIT_MEMLOCK. + preamble: ulimit -l unlimited + readiness: + log: + pattern: Starting real client service + timeout_seconds: 1800 + interval_seconds: 5 + - <<: *owner + name: mc-owner-numa1 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_client, "--host={node_ip}", --port=26053, --master_server_address=127.0.0.1:26051, --metadata_server=http://127.0.0.1:26080/metadata, --protocol=rdma, "--device_names=rdma0,rdma1,rdma2,rdma3", --global_segment_size=824633720832, --local_buffer_size=0, --threads=4] + roles: + prefill: + env: + <<: &store_prefill_env + # More than one load worker hung the PP prefill pipeline at C96. + OFFLOAD_LOAD_WORKERS: "1" + MC_NUM_QP_PER_EP: "1" + # Registers the worker's HBM transfer pool in 1 GiB regions. + MC_MAX_MR_SIZE: "1073741824" + # Every owner sits on rdma0-3 and this fabric is rail-isolated, so each stage must reach an owner through the NIC of its own name. + MC_ENABLE_DEST_DEVICE_AFFINITY: "1" + args: + num-speculative-tokens: 3 + # One pool shared by all four stages, as in ATOM's CI; each client still uses its GPU's NIC. + extra-kv-connectors: + - kv_connector: mooncake_store + kv_role: offload + kv_connector_extra_config: + mooncake_store.protocol: rdma + mooncake_store.master: "127.0.0.1:26051" + mooncake_store.metadata: "http://127.0.0.1:26080/metadata" + decode: + nodes: colocate + env: + AITER_REUSE_IDENTICAL_COMM_GROUPS: "1" + # The P/D transfer must match the prefill's QP count, one with the Store. + MC_NUM_QP_PER_EP: "1" + args: + <<: &dcp_decode + decode-context-parallel-size: 4 + gpu-memory-utilization: 0.94 + max-num-seqs: 512 + max-num-batched-tokens: 2048 + cudagraph-capture-sizes: '[1,2,4,8,16,24,32,40,48,56,64,72,80,88,96,104,112,120,128,136,144,152,160,168,176,184,192,200,208,216,224,232,240,248,256]' + num-speculative-tokens: 3 + +override_1p2d_mcstore_c56: + resources: &spread + # A node per decode; srtctl otherwise packs two TP4 decodes onto one node. + spread_workers: true + frontend: + args: &cache_aware + policy: cache_aware + prefill-policy: cache_aware + decode-policy: cache_aware + cache-threshold: 0.8 + balance-abs-threshold: 2 + balance-rel-threshold: 2.0 + eviction-interval: 300 + host_setup: &prefill_budget_host_setup + nodes: all + timeout_seconds: 1800 + commands: + # A recipe host_setup replaces the cluster default, which also drops the page cache under DRAM KV offloading. + - IBDEVICES=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/setup.sh" + # Owners take 714752 2 MiB pages, half per NUMA node; reserve 2% more above the pages host services hold. + # Only the prefill node holds owners; srtctl starts the prefill on the head node, which launches this srun. + - >- + if hostname -I | tr ' ' '\n' | grep -qxF "${SLURM_LAUNCH_NODE_IPADDR:?srun did not set SLURM_LAUNCH_NODE_IPADDR}"; then + echo "[hugepage] $(hostname -s) is the head node; reserving the owners' pages"; + held=$(awk '/^HugePages_Total:/ {total = $2} /^HugePages_Free:/ {free = $2} /^HugePages_Rsvd:/ {rsvd = $2} END {print total - free + rsvd}' /proc/meminfo); + HUGEPAGES_TARGET=$(( held + 714752 * 102 / 100 )) HUGEPAGE_GROW_TOLERANCE_PCT=1 + bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/hugepage_sanity.sh" || exit 1; + for node in 0 1; do + free=$(cat /sys/devices/system/node/node$node/hugepages/hugepages-2048kB/free_hugepages); + echo "[hugepage] NUMA$node: $free free 2 MiB pages, owners need 357376"; + [[ $free -ge 357376 ]] || exit 1; + done; fi + teardown: + - HUGEPAGES_TARGET=0 bash "${GITHUB_WORKSPACE}/runners/srt-slurm/hooks/mi355x-amds/hugepage_sanity.sh" + services: &prefill_budget_services + - <<: *master + name: mc-master-rdma0 + - <<: *master + name: mc-master-rdma1 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_master, --rpc_port=26151, --rpc_thread_num=32, --enable_http_metadata_server=true, --http_metadata_server_host=0.0.0.0, --http_metadata_server_port=26180, --metrics_port=26190, --default_kv_lease_ttl=10000, --eviction_high_watermark_ratio=0.90, --eviction_ratio=0.05, --allocation_strategy=random, --memory_allocator=offset, --client_ttl=10] + readiness: + port: 26151 + timeout_seconds: 300 + - <<: *master + name: mc-master-rdma2 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_master, --rpc_port=26251, --rpc_thread_num=32, --enable_http_metadata_server=true, --http_metadata_server_host=0.0.0.0, --http_metadata_server_port=26280, --metrics_port=26290, --default_kv_lease_ttl=10000, --eviction_high_watermark_ratio=0.90, --eviction_ratio=0.05, --allocation_strategy=random, --memory_allocator=offset, --client_ttl=10] + readiness: + port: 26251 + timeout_seconds: 300 + - <<: *master + name: mc-master-rdma3 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_master, --rpc_port=26351, --rpc_thread_num=32, --enable_http_metadata_server=true, --http_metadata_server_host=0.0.0.0, --http_metadata_server_port=26380, --metrics_port=26390, --default_kv_lease_ttl=10000, --eviction_high_watermark_ratio=0.90, --eviction_ratio=0.05, --allocation_strategy=random, --memory_allocator=offset, --client_ttl=10] + readiness: + port: 26351 + timeout_seconds: 300 + - <<: *owner + name: mc-owner-rdma0 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "0", mooncake_client, "--host={node_ip}", --port=26052, --master_server_address=127.0.0.1:26051, --metadata_server=http://127.0.0.1:26080/metadata, --protocol=rdma, --device_names=rdma0, --global_segment_size=374735896576, --local_buffer_size=0, --threads=4] + - <<: *owner + name: mc-owner-rdma1 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "0", mooncake_client, "--host={node_ip}", --port=26053, --master_server_address=127.0.0.1:26151, --metadata_server=http://127.0.0.1:26180/metadata, --protocol=rdma, --device_names=rdma1, --global_segment_size=374735896576, --local_buffer_size=0, --threads=4] + - <<: *owner + name: mc-owner-rdma2 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_client, "--host={node_ip}", --port=26054, --master_server_address=127.0.0.1:26251, --metadata_server=http://127.0.0.1:26280/metadata, --protocol=rdma, --device_names=rdma2, --global_segment_size=374735896576, --local_buffer_size=0, --threads=4] + - <<: *owner + name: mc-owner-rdma3 + command: [python3, /app/ATOM/.github/scripts/atomesh/numa_exec.py, "1", mooncake_client, "--host={node_ip}", --port=26055, --master_server_address=127.0.0.1:26351, --metadata_server=http://127.0.0.1:26380/metadata, --protocol=rdma, --device_names=rdma3, --global_segment_size=374735896576, --local_buffer_size=0, --threads=4] + roles: + prefill: + env: *store_prefill_env + args: + num-speculative-tokens: 4 + # Each stage's Store client uses its GPU's NIC and that NIC's pool. + extra-kv-connectors: &per_nic_store_connector + - kv_connector: mooncake_store + kv_role: offload + kv_connector_extra_config: + mooncake_store.protocol: rdma + mooncake_store.pools: + rdma0: {master: "127.0.0.1:26051", metadata: "http://127.0.0.1:26080/metadata"} + rdma1: {master: "127.0.0.1:26151", metadata: "http://127.0.0.1:26180/metadata"} + rdma2: {master: "127.0.0.1:26251", metadata: "http://127.0.0.1:26280/metadata"} + rdma3: {master: "127.0.0.1:26351", metadata: "http://127.0.0.1:26380/metadata"} + decode: + nodes: 2 + workers: 2 + env: + AITER_REUSE_IDENTICAL_COMM_GROUPS: "1" + MC_NUM_QP_PER_EP: "1" + args: + <<: *dcp_decode + num-speculative-tokens: 4 + +override_1p3d_mcstore_c12: + resources: *spread + frontend: + args: *cache_aware + host_setup: *prefill_budget_host_setup + services: *prefill_budget_services + roles: + prefill: + env: *store_prefill_env + args: + num-speculative-tokens: 5 + extra-kv-connectors: *per_nic_store_connector + decode: + nodes: 3 + workers: 3 + env: &tp4_decode_env + PYTHONNOUSERSITE: "1" + MC_NUM_QP_PER_EP: "1" + args: + num-speculative-tokens: 5 + gpu-memory-utilization: 0.95 + max-num-seqs: 24 + max-num-batched-tokens: 16384 + cudagraph-capture-sizes: '[1,2,4,8,12,16,20,24]' + disable-uvicorn-access-log: true + +override_1p3d_mcstore_c16: + resources: *spread + frontend: + args: *cache_aware + host_setup: *prefill_budget_host_setup + services: *prefill_budget_services + roles: + prefill: + env: *store_prefill_env + args: + num-speculative-tokens: 5 + extra-kv-connectors: *per_nic_store_connector + decode: + nodes: 3 + workers: 3 + env: *tp4_decode_env + args: + num-speculative-tokens: 5 + gpu-memory-utilization: 0.95 + max-num-seqs: 32 + max-num-batched-tokens: 16384 + cudagraph-capture-sizes: '[1,2,4,8,12,16,20,24,28,32]' + disable-uvicorn-access-log: true + +override_1p3d_mcstore_c24: + resources: *spread + frontend: + args: *cache_aware + host_setup: *prefill_budget_host_setup + services: *prefill_budget_services + roles: + prefill: + env: *store_prefill_env + args: + num-speculative-tokens: 5 + extra-kv-connectors: *per_nic_store_connector + decode: + nodes: 3 + workers: 3 + env: *tp4_decode_env + args: + num-speculative-tokens: 5 + gpu-memory-utilization: 0.95 + max-num-seqs: 48 + max-num-batched-tokens: 16384 + cudagraph-capture-sizes: '[1,2,4,8,12,16,20,24,28,32,36,40,44,48]' + disable-uvicorn-access-log: true diff --git a/inferencex-e2e/configs/amd-master.yaml b/inferencex-e2e/configs/amd-master.yaml index 6bba904af0..167d1e7392 100644 --- a/inferencex-e2e/configs/amd-master.yaml +++ b/inferencex-e2e/configs/amd-master.yaml @@ -1486,6 +1486,59 @@ glm5.2-fp4-mi355x-atom-agentic-mtp: - { tp: 4, kv-offloading: none, conc-list: [2, 4, 8, 10], spec-decoding: mtp, srt-recipe: agentic.yaml } - { tp: 8, kv-offloading: none, conc-list: [1, 2, 4], spec-decoding: mtp, srt-recipe: agentic.yaml } +glm5.2-fp4-mi355x-atom-disagg-agentic-mtp: + image: rocm/atom-dev:nightly_202610131249@sha256:d5b3fcd99bd44e3489740d1b026037f49f2e3f6160dc70b70c3fd23a125a4905 + model: amd/GLM-5.2-MXFP4 + model-prefix: glm5.2 + runner: cluster:mi355x-amds + precision: fp4 + framework: atom-disagg + router: { name: atomesh, version: "6fd1bb1be97c79b571eb86674786d0d888353b73" } + kv-p2p-transfer: mooncake + multinode: true + disagg: true + srt-recipe-dir: glm5.2/atom/mi355x-fp4/agentx + scenarios: + agentic-coding: + # 1.0 gives each 4-GPU server 1499 GB: 1P1D's 2 x 768 GiB owners hold 825 GB per NUMA node of its prefill's and colocated decode's shares; 1P2D and 1P3D owners hold the prefill's 1499 GB. + - dram-utilization: 1.0 + search-space: + - spec-decoding: "mtp" + conc-list: [ 56, 72, 96, 112, 128 ] + kv-offloading: dram + kv-offload-backend: { name: mooncake, version: "0.3.13" } + srt-recipe: disagg-variants.yaml:override_1p1d_mcstore + prefill: { num-worker: 1, tp: 1, pp: 4, ep: 1, dp-attn: false } + decode: { num-worker: 1, tp: 4, dcp-size: 4, ep: 1, dp-attn: false } + - spec-decoding: "mtp" + conc-list: [ 56 ] + kv-offloading: dram + kv-offload-backend: { name: mooncake, version: "0.3.13" } + srt-recipe: disagg-variants.yaml:override_1p2d_mcstore_c56 + prefill: { num-worker: 1, tp: 1, pp: 4, ep: 1, dp-attn: false } + decode: { num-worker: 2, tp: 4, dcp-size: 4, ep: 1, dp-attn: false } + - spec-decoding: "mtp" + conc-list: [ 12 ] + kv-offloading: dram + kv-offload-backend: { name: mooncake, version: "0.3.13" } + srt-recipe: disagg-variants.yaml:override_1p3d_mcstore_c12 + prefill: { num-worker: 1, tp: 1, pp: 4, ep: 1, dp-attn: false } + decode: { num-worker: 3, tp: 4, ep: 1, dp-attn: false } + - spec-decoding: "mtp" + conc-list: [ 16 ] + kv-offloading: dram + kv-offload-backend: { name: mooncake, version: "0.3.13" } + srt-recipe: disagg-variants.yaml:override_1p3d_mcstore_c16 + prefill: { num-worker: 1, tp: 1, pp: 4, ep: 1, dp-attn: false } + decode: { num-worker: 3, tp: 4, ep: 1, dp-attn: false } + - spec-decoding: "mtp" + conc-list: [ 24 ] + kv-offloading: dram + kv-offload-backend: { name: mooncake, version: "0.3.13" } + srt-recipe: disagg-variants.yaml:override_1p3d_mcstore_c24 + prefill: { num-worker: 1, tp: 1, pp: 4, ep: 1, dp-attn: false } + decode: { num-worker: 3, tp: 4, ep: 1, dp-attn: false } + dsv4-fp4-mi355x-sglang-agentic-mtp: image: lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019 model: deepseek-ai/DeepSeek-V4-Pro-0813 diff --git a/inferencex-e2e/infx/results/agentic/__init__.py b/inferencex-e2e/infx/results/agentic/__init__.py index 67ada94c91..638ab24a84 100644 --- a/inferencex-e2e/infx/results/agentic/__init__.py +++ b/inferencex-e2e/infx/results/agentic/__init__.py @@ -231,6 +231,7 @@ def build_result( framework=framework, records=records, server_logs=server_logs, + env=env, ) if "total_tput_tps" in request_flat and num_gpus > 0: diff --git a/inferencex-e2e/infx/results/agentic/backends/atom.py b/inferencex-e2e/infx/results/agentic/backends/atom.py index 9e5c661de2..32543a005b 100644 --- a/inferencex-e2e/infx/results/agentic/backends/atom.py +++ b/inferencex-e2e/infx/results/agentic/backends/atom.py @@ -2,12 +2,24 @@ from __future__ import annotations +import logging +import math +from collections.abc import Mapping from typing import Any -from infx.results.agentic.common import gauge_stat, normalize_fraction, rate, sum_stat +from infx.results.agentic.common import ( + gauge_stat, + metric_series, + normalize_fraction, + rate, + series_stat, + sum_stat, +) from .base import ServerMetricsBackend, counter_int +logger = logging.getLogger(__name__) + def _atom_names(suffix: str) -> list[str]: """Accept direct worker metrics and Atomesh's colon-normalized form.""" @@ -19,6 +31,98 @@ def _atom_counter_names(stem: str) -> list[str]: return [*_atom_names(stem), *_atom_names(f"{stem}_total")] +def _counter_values( + metrics: dict[str, dict[str, Any]], stem: str, endpoint: str +) -> dict[tuple, float] | None: + names = _atom_counter_names(f"prefix_cache_{stem}_tokens") + values = {} + for series in metric_series(metrics, names): + if series.get("endpoint_url") != endpoint: + continue + labels = series.get("labels") + labels = {} if labels is None else labels + value = series_stat(series, preferred_keys=("total",)) + if ( + not isinstance(labels, dict) + or any(not isinstance(k, str) or not isinstance(v, str) for k, v in labels.items()) + or value is None + or not math.isfinite(value) + or value < 0 + ): + return None + key = tuple(sorted(labels.items())) + if key in values: + logger.warning("ATOM %s cache counter duplicated at %s", stem, endpoint) + return None + values[key] = value + return values + + +def _cache_counts( + metrics: dict[str, dict[str, Any]], + endpoint_roles: Mapping[str, str] | None, + disaggregated: bool | None, + cpu_offload_disabled: bool = False, +) -> tuple[float | None, float | None, float | None]: + names = [ + name + for stem in ("full", "cached", "offload") + for name in _atom_counter_names(f"prefix_cache_{stem}_tokens") + ] + endpoints = [s.get("endpoint_url") for s in metric_series(metrics, names)] + if any(not isinstance(endpoint, str) or not endpoint for endpoint in endpoints): + return None, None, None + observed = set(endpoints) + if endpoint_roles is None: + accounted = observed if disaggregated is False and len(observed) == 1 else set() + else: + # Decode counts the prompt again without serving its prefill cache reuse. + accounted = {url for url, role in endpoint_roles.items() if role in {"prefill", "unified"}} + if observed - endpoint_roles.keys(): + accounted = set() + if not accounted: + logger.warning("ATOM cache endpoint roles missing or ambiguous") + return None, None, None + full_total = cached_total = offload_total = 0.0 + cpu_complete = True + for endpoint in accounted: + full, cached, offload = ( + _counter_values(metrics, stem, endpoint) for stem in ("full", "cached", "offload") + ) + if ( + not full + or cached is None + or full.keys() != cached.keys() + or any(cached[key] > full[key] for key in full) + ): + return None, None, None + if cpu_offload_disabled: + offload = ( + dict.fromkeys(full, 0.0) + if offload is not None + and offload.keys() <= full.keys() + and not any(offload.values()) + else None + ) + full_total += sum(full.values()) + cached_total += sum(cached.values()) + if ( + offload is None + or offload.keys() != full.keys() + or any(cached[key] + offload[key] > full[key] for key in full) + ): + cpu_complete = False + else: + offload_total += sum(offload.values()) + if full_total <= 0 or not all(math.isfinite(v) for v in (full_total, cached_total)): + return None, None, None + return ( + full_total, + cached_total, + offload_total if cpu_complete and math.isfinite(offload_total) else None, + ) + + class AtomBackend(ServerMetricsBackend): name = "atom" @@ -34,12 +138,32 @@ def populate( metrics: dict[str, dict[str, Any]], flat: dict[str, Any], nested: dict[str, Any], + *, + endpoint_roles: Mapping[str, str] | None = None, + disaggregated: bool | None = None, + cpu_offload_disabled: bool = False, ) -> None: """Fill token totals, prefix and LMCache hit rates and KV usage from ATOM ``metrics``.""" + prompt_names = _atom_counter_names("prompt_tokens") + prompt_endpoints = {s.get("endpoint_url") for s in metric_series(metrics, prompt_names)} + if endpoint_roles is None: + accounted = ( + prompt_endpoints if disaggregated is False and len(prompt_endpoints) == 1 else set() + ) + else: + accounted = { + url for url, role in endpoint_roles.items() if role in {"prefill", "unified"} + } + if prompt_endpoints - endpoint_roles.keys(): + accounted = set() + accounted = {url for url in accounted if isinstance(url, str) and url} + if accounted - prompt_endpoints: + accounted = set() prompt_total = sum_stat( metrics, - _atom_counter_names("prompt_tokens"), - preferred_keys=("total", "sum", "max", "avg"), + prompt_names, + preferred_keys=("total",), + series_filter=lambda s: s.get("endpoint_url") in accounted, ) generation_total = sum_stat( metrics, @@ -49,24 +173,17 @@ def populate( flat["total_prompt_tokens"] = counter_int(prompt_total) flat["total_generation_tokens"] = counter_int(generation_total) - cached_tokens = sum_stat( - metrics, - _atom_counter_names("prefix_cache_cached_tokens"), - preferred_keys=("total", "sum", "max", "avg"), - ) - full_tokens = sum_stat( - metrics, - _atom_counter_names("prefix_cache_full_tokens"), - preferred_keys=("total", "sum", "max", "avg"), + full_tokens, cached_tokens, external_tokens = _cache_counts( + metrics, endpoint_roles, disaggregated, cpu_offload_disabled ) cache_hit_rate = rate(cached_tokens, full_tokens) - external_tokens = sum_stat( - metrics, - _atom_counter_names("lmcache_loaded_tokens"), - preferred_keys=("total", "sum", "max", "avg"), - ) + # Offload is CPU reuse beyond the GPU prefix; loaded tokens measure transfer volume. external_hit_rate = rate(external_tokens, full_tokens) - overall_hit_rate = cache_hit_rate + overall_hit_rate = ( + rate(cached_tokens + external_tokens, full_tokens) + if cached_tokens is not None and external_tokens is not None + else None + ) flat["server_gpu_cache_hit_rate"] = cache_hit_rate flat["server_cpu_cache_hit_rate"] = external_hit_rate @@ -90,7 +207,9 @@ def populate( "prefix_cache_hits": cached_tokens, "prefix_cache_queries": full_tokens, "external_prefix_cache_hits": external_tokens, - "external_prefix_cache_queries": full_tokens, + "external_prefix_cache_queries": ( + full_tokens if external_tokens is not None else None + ), } ) nested["kv_cache"]["gpu_usage_pct"] = flat["gpu_kv_cache_usage_pct"] diff --git a/inferencex-e2e/infx/results/agentic/server_metrics.py b/inferencex-e2e/infx/results/agentic/server_metrics.py index 0ba5ffdb55..5d176339fe 100644 --- a/inferencex-e2e/infx/results/agentic/server_metrics.py +++ b/inferencex-e2e/infx/results/agentic/server_metrics.py @@ -2,10 +2,11 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from typing import Any from .backends import detect_backend +from .backends.atom import AtomBackend from .backends.base import ( SERVER_CACHE_FLAT_FIELDS, apply_profile_totals, @@ -14,12 +15,36 @@ from .common import index_server_metrics +def endpoint_roles_from_env(env: Mapping[str, str]) -> dict[str, str] | None: + if env.get("SRTCTL_FRONTEND_TYPE") == "dynamo": + return {} + if env.get("AIPERF_SERVER_METRICS_URLS"): + return None + groups = ( + [("unified", env["SRT_AGG_ENDPOINTS"])] + if env.get("SRT_AGG_ENDPOINTS") + else [ + ("prefill", env.get("SRT_PREFILL_ENDPOINTS", "")), + ("decode", env.get("SRT_DECODE_ENDPOINTS", "")), + ] + ) + roles: dict[str, str] = {} + for role, endpoints in groups: + for endpoint in endpoints.split(","): + if endpoint: + url = f"http://{endpoint}/metrics" + if roles.setdefault(url, role) != role: + return {} + return roles or None + + def compute_server_metrics( server_metrics: dict[str, Any], *, framework: str, records: list[dict[str, Any]], server_logs: Iterable[str | None] = (), + env: Mapping[str, str] | None = None, ) -> tuple[dict[str, Any], dict[str, Any], list[str]]: """Return flat and nested server metrics plus warnings from aiperf's ``server_metrics`` export. @@ -50,7 +75,22 @@ def compute_server_metrics( f"framework={framework.lower()!r}, metric_names=[{preview}]" ) nested["adapter"] = adapter.name - adapter.populate(metrics, flat, nested) + if isinstance(adapter, AtomBackend): + source_env = env or {} + adapter.populate( + metrics, + flat, + nested, + endpoint_roles=endpoint_roles_from_env(source_env), + disaggregated={"true": True, "false": False}.get(source_env.get("DISAGG", "").lower()), + cpu_offload_disabled=(source_env.get("KV_OFFLOADING", "").strip().lower() == "none") + and not ( + source_env.get("KV_OFFLOAD_BACKEND") + or source_env.get("KV_OFFLOAD_BACKEND_METADATA") + ), + ) + else: + adapter.populate(metrics, flat, nested) if framework.lower() == "dynamo-sglang" and adapter.name == "sglang": nested["kv_cache"]["gpu_total_tokens"] = None diff --git a/inferencex-e2e/infx/tests/results/agentic/test_atom_cache_counters.py b/inferencex-e2e/infx/tests/results/agentic/test_atom_cache_counters.py new file mode 100644 index 0000000000..8356e5bd39 --- /dev/null +++ b/inferencex-e2e/infx/tests/results/agentic/test_atom_cache_counters.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import math + +import pytest + +from infx.results.agentic.server_metrics import compute_server_metrics + + +def _export(full: dict, cached: dict, offload: dict) -> dict: + return { + "metrics": { + f"atom:prefix_cache_{stem}_tokens": { + "series": [ + {"endpoint_url": f"http://{endpoint}/metrics", "stats": {"total": value}} + for endpoint, value in values.items() + ] + } + for stem, values in (("full", full), ("cached", cached), ("offload", offload)) + } + } + + +def _compute(export: dict, env: dict) -> tuple[dict, dict]: + flat, nested, _ = compute_server_metrics(export, framework="mooncake-atom", records=[], env=env) + return flat, nested + + +def _assert_rates( + flat: dict, nested: dict, gpu: float | None, cpu: float | None, overall: float | None +): + for name, expected in (("gpu", gpu), ("cpu", cpu), ("external", cpu), ("overall", overall)): + actual = flat[f"server_{name}_cache_hit_rate"] + assert actual == (pytest.approx(expected) if expected is not None else None) + assert nested["cache"][f"{name}_cache_hit_rate"] == actual + + +@pytest.mark.parametrize("decode_cached", [800, math.nan]) +def test_prefill_counters_exclude_decode_and_transfer_volume(decode_cached: float) -> None: + export = _export({"p": 1000, "d": 1000}, {"p": 600, "d": decode_cached}, {"p": 100, "d": -1}) + export["metrics"]["atom:lmcache_loaded_tokens"] = { + "series": [{"endpoint_url": "http://p/metrics", "stats": {"total": 900}}] + } + export["metrics"]["atom:prompt_tokens"] = { + "series": [ + {"endpoint_url": f"http://{endpoint}/metrics", "stats": {"total": 1000}} + for endpoint in ("p", "d") + ] + } + flat, nested = _compute( + export, {"SRT_PREFILL_ENDPOINTS": "p", "SRT_DECODE_ENDPOINTS": "d", "DISAGG": "true"} + ) + assert flat["total_prompt_tokens"] == nested["tokens"]["prompt_total"] == 1000 + _assert_rates(flat, nested, 0.6, 0.1, 0.7) + assert nested["cache"]["prefix_cache_queries"] == 1000 + assert nested["cache"]["external_prefix_cache_hits"] == 100 + assert nested["cache"]["external_prefix_cache_queries"] == 1000 + + +@pytest.mark.parametrize( + ("env", "offload", "expected"), + [ + ({"SRT_AGG_ENDPOINTS": "worker", "SRT_PREFILL_ENDPOINTS": "unused"}, 50, 0.5), + ({"DISAGG": "false"}, 0, 0), + ({"DISAGG": "false", "AIPERF_SERVER_METRICS_URLS": "http://worker/metrics"}, 50, 0.5), + ], +) +def test_unified_zero_gpu_and_offload_counter_alias( + env: dict, offload: int, expected: float +) -> None: + export = _export({"worker": 100}, {"worker": 0}, {"worker": offload}) + export["metrics"] = { + name.replace("atom:", "atom_") + "_total": value + for name, value in export["metrics"].items() + } + flat, nested = _compute(export, env) + _assert_rates(flat, nested, 0, expected, expected) + + +def test_multiple_prefills_and_labelled_series_are_token_weighted() -> None: + export = _export( + {"p": 100, "q": 900, "idle": 0}, + {"p": 80, "q": 20, "idle": 0}, + {"p": 10, "q": 180, "idle": 0}, + ) + for stem, values in (("full", (400, 500)), ("cached", (5, 15)), ("offload", (80, 100))): + series = export["metrics"][f"atom:prefix_cache_{stem}_tokens"]["series"] + series[:] = [s for s in series if s["endpoint_url"] != "http://q/metrics"] + series.extend( + { + "endpoint_url": "http://q/metrics", + "labels": {"engine": str(i)}, + "stats": {"total": v}, + } + for i, v in enumerate(values) + ) + flat, nested = _compute(export, {"SRT_PREFILL_ENDPOINTS": "p,q,idle"}) + _assert_rates(flat, nested, 0.1, 0.19, 0.29) + + +@pytest.mark.parametrize( + "env", + [ + {}, + {"DISAGG": "true"}, + {"SRT_PREFILL_ENDPOINTS": "p", "SRT_DECODE_ENDPOINTS": "p,d"}, + {"SRT_PREFILL_ENDPOINTS": "p,missing", "SRT_DECODE_ENDPOINTS": "d"}, + {"SRT_PREFILL_ENDPOINTS": "p"}, + { + "SRT_PREFILL_ENDPOINTS": "p", + "SRT_DECODE_ENDPOINTS": "d", + "AIPERF_SERVER_METRICS_URLS": "custom", + }, + { + "SRT_PREFILL_ENDPOINTS": "p", + "SRT_DECODE_ENDPOINTS": "d", + "SRTCTL_FRONTEND_TYPE": "dynamo", + }, + ], +) +def test_unknown_conflicting_or_incomplete_roles_do_not_mix_endpoints(env: dict) -> None: + flat, nested = _compute( + _export({"p": 100, "d": 100}, {"p": 60, "d": 0}, {"p": 10, "d": 0}), env + ) + _assert_rates(flat, nested, None, None, None) + + +def test_single_endpoint_without_explicit_unified_topology_stays_unknown() -> None: + flat, nested = _compute(_export({"worker": 100}, {"worker": 60}, {"worker": 10}), {}) + _assert_rates(flat, nested, None, None, None) + + +@pytest.mark.parametrize( + ("offload", "cpu", "overall"), + [({"p": 0, "q": 0}, 0, 0.2), ({}, None, None), ({"p": 10}, None, None)], +) +def test_missing_or_partial_offload_is_distinct_from_zero( + offload: dict, cpu: float | None, overall: float | None +) -> None: + flat, nested = _compute( + _export({"p": 100, "q": 100}, {"p": 20, "q": 20}, offload), + {"SRT_PREFILL_ENDPOINTS": "p,q"}, + ) + _assert_rates(flat, nested, 0.2, cpu, overall) + assert nested["cache"]["external_prefix_cache_queries"] == (200 if cpu is not None else None) + + +@pytest.mark.parametrize( + ("stem", "value"), + [ + ("full", 0), + ("cached", -1), + ("cached", 101), + ("cached", math.nan), + ("cached", math.inf), + ("cached", None), + ], +) +def test_invalid_counted_counter_does_not_produce_rates(stem: str, value: object) -> None: + export = _export({"worker": 100}, {"worker": 20}, {"worker": 10}) + export["metrics"][f"atom:prefix_cache_{stem}_tokens"]["series"][0]["stats"]["total"] = value + flat, nested = _compute(export, {"DISAGG": "false"}) + _assert_rates(flat, nested, None, None, None) + + +def test_unreadable_series_does_not_get_silently_skipped() -> None: + export = _export({"worker": 100}, {"worker": 20}, {"worker": 10}) + export["metrics"]["atom:prefix_cache_cached_tokens"]["series"].append( + { + "endpoint_url": "http://worker/metrics", + "labels": {"engine": "other"}, + "stats": {"max": 5}, + } + ) + flat, nested = _compute(export, {"DISAGG": "false"}) + _assert_rates(flat, nested, None, None, None) + + +@pytest.mark.parametrize("offload", [30, math.nan]) +@pytest.mark.parametrize("labelled", [False, True]) +def test_invalid_worker_cpu_count_cannot_hide_in_another_workers_denominator( + offload: float, labelled: bool +) -> None: + export = _export({"p": 100, "q": 900}, {"p": 80, "q": 20}, {"p": offload, "q": 0}) + if labelled: + for metric in export["metrics"].values(): + for series in metric["series"]: + series["labels"] = {"engine": series["endpoint_url"]} + series["endpoint_url"] = "http://worker/metrics" + flat, nested = _compute(export, {"SRT_PREFILL_ENDPOINTS": "worker" if labelled else "p,q"}) + _assert_rates(flat, nested, 0.1, None, None) + + +def test_unknown_cache_roles_preserve_other_atom_telemetry() -> None: + export = _export({"p": 100, "d": 100}, {"p": 60, "d": 0}, {"p": 10, "d": 0}) + export["metrics"].update( + { + "atom:prompt_tokens": {"series": [{"stats": {"total": 200}}]}, + "atom:generation_tokens": {"series": [{"stats": {"total": 30}}]}, + "atom:kv_cache_usage_ratio": {"series": [{"stats": {"max": 0.75}}]}, + } + ) + flat, nested, _ = compute_server_metrics( + export, + framework="mooncake-atom", + records=[{"metrics": {"input_sequence_length": {"value": 123}}}], + env={}, + ) + _assert_rates(flat, nested, None, None, None) + assert flat["total_prompt_tokens"] == nested["tokens"]["prompt_total"] == 123 + assert flat["total_generation_tokens"] == nested["tokens"]["generation_total"] == 30 + assert flat["gpu_kv_cache_usage_pct"] == nested["kv_cache"]["gpu_usage_pct"] == 0.75 + + +@pytest.mark.parametrize( + ("env", "cpu", "overall"), + [ + ({"KV_OFFLOADING": " NONE "}, 0, 0.6), + ({"KV_OFFLOADING": "dram"}, None, None), + ({"KV_OFFLOADING": "none", "KV_OFFLOAD_BACKEND": "mooncake"}, None, None), + ], +) +@pytest.mark.parametrize("offload", [{}, {"p": 0}]) +def test_missing_offload_is_zero_only_with_confirmed_disabled_tier( + env: dict, cpu: float | None, overall: float | None, offload: dict +) -> None: + flat, nested = _compute( + _export({"p": 100, "q": 100}, {"p": 60, "q": 60}, offload), + {"SRT_PREFILL_ENDPOINTS": "p,q", **env}, + ) + _assert_rates(flat, nested, 0.6, cpu, overall) + + +@pytest.mark.parametrize("offload", [10, math.nan, None]) +def test_disabled_cpu_tier_does_not_hide_invalid_present_offload(offload: object) -> None: + flat, nested = _compute( + _export({"p": 100, "q": 100}, {"p": 60, "q": 60}, {"p": offload}), + {"SRT_PREFILL_ENDPOINTS": "p,q", "KV_OFFLOADING": "none"}, + ) + _assert_rates(flat, nested, 0.6, None, None) + + +def test_disabled_cpu_tier_does_not_fill_an_unmatched_label_domain() -> None: + export = _export({"worker": 100}, {"worker": 60}, {"worker": 0}) + export["metrics"]["atom:prefix_cache_offload_tokens"]["series"][0]["labels"] = { + "engine": "other" + } + flat, nested = _compute(export, {"DISAGG": "false", "KV_OFFLOADING": "none"}) + _assert_rates(flat, nested, 0.6, None, None) diff --git a/inferencex-e2e/infx/tests/results/agentic/test_process_agentic_result.py b/inferencex-e2e/infx/tests/results/agentic/test_process_agentic_result.py index 363ce7ba1a..c736a237db 100644 --- a/inferencex-e2e/infx/tests/results/agentic/test_process_agentic_result.py +++ b/inferencex-e2e/infx/tests/results/agentic/test_process_agentic_result.py @@ -365,6 +365,55 @@ def test_processor_records_run_attempt(tmp_path: Path): assert "run_attempt" not in local +def test_processor_uses_launcher_roles_for_atom_cache_counts(tmp_path: Path) -> None: + result_dir = _write_fixture(tmp_path) + metrics = { + f"atom:prefix_cache_{stem}_tokens": { + "series": [ + {"endpoint_url": f"http://{endpoint}/metrics", "stats": {"total": value}} + for endpoint, value in values.items() + ] + } + for stem, values in ( + ("full", {"prefill:8000": 1000, "decode:8000": 1000}), + ("cached", {"prefill:8000": 600, "decode:8000": 0}), + ("offload", {"prefill:8000": 100, "decode:8000": 0}), + ) + } + (result_dir / "aiperf_artifacts" / "server_metrics_export.json").write_text( + json.dumps({"metrics": metrics}) + ) + agg = _run_processor( + result_dir, + tmp_path / "out", + { + "FRAMEWORK": "mooncake-atom", + "DISAGG": "true", + "IS_MULTINODE": "true", + "PREFILL_NUM_WORKERS": "1", + "PREFILL_TP": "4", + "PREFILL_EP": "1", + "DECODE_NUM_WORKERS": "1", + "DECODE_TP": "4", + "DECODE_EP": "1", + "SRT_PREFILL_ENDPOINTS": "prefill:8000", + "SRT_DECODE_ENDPOINTS": "decode:8000", + "SRT_AGG_ENDPOINTS": "", + "AIPERF_SERVER_METRICS_URLS": "", + "SRTCTL_FRONTEND_TYPE": "", + "KV_OFFLOADING": "dram", + "KV_OFFLOAD_BACKEND": "lmcache", + "KV_OFFLOAD_BACKEND_METADATA": '{"name":"lmcache","version":"0.4.5"}', + }, + ) + cache = agg["server_metrics"]["cache"] + assert cache["gpu_cache_hit_rate"] == 0.6 + assert cache["cpu_cache_hit_rate"] == cache["external_cache_hit_rate"] == 0.1 + assert cache["overall_cache_hit_rate"] == 0.7 + assert cache["prefix_cache_queries"] == cache["external_prefix_cache_queries"] == 1000 + assert agg["num_gpus"] == 8 + + def test_processor_preserves_dataset_provenance(tmp_path: Path): result_dir = _write_fixture(tmp_path) output_dir = tmp_path / "out" diff --git a/inferencex-e2e/perf-changelog.yaml b/inferencex-e2e/perf-changelog.yaml index 518bc9b1ad..0738f9cd16 100644 --- a/inferencex-e2e/perf-changelog.yaml +++ b/inferencex-e2e/perf-changelog.yaml @@ -11210,3 +11210,15 @@ - "Add TP2 points with DRAM KV offload, and adjust the memory settings on the DRAM KV offload points." pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/XXX scenario-type: [agentic-coding] + +- config-keys: + - glm5.2-fp4-mi355x-atom-disagg-agentic-mtp + scenario-type: + - agentic-coding + description: + - "Add GLM-5.2 MXFP4 on MI355X with ATOM prefill/decode disaggregation behind AToMesh, Mooncake RDMA from prefill to decode and MTP, on rocm/atom-dev:nightly_202610131249 (digest-pinned). The prefill is one CPP4 worker (PP4 x TP1 on four GPUs)." + - "Every point keeps its prefill's prefix tier in a Mooncake Store (0.3.13) on the prefill node, with owners on 2 MiB HugeTLB pages. 1P1D, as in ATOM's CI, has one pool shared by all four PP stages: one master and two 768 GiB owners, one per NUMA node (1649 GB, within the prefill's and the colocated decode's 1499 GB DRAM shares), each registered on rdma0-3, with MC_ENABLE_DEST_DEVICE_AFFINITY so each stage reaches them through the NIC of its own name on this rail-isolated fabric. 1P2D and 1P3D keep a pool per prefill NIC, each one master and a 349 GiB owner, the prefill's 1499 GB. After the cluster host setup drops the page cache, host_setup on the head node, where srtctl starts the prefill, reserves the owners' pages plus 2% above the hugepages host services hold and fails unless each NUMA node has its owners' pages free; teardown releases them." + - "1P1D on one node at concurrency 56, 72, 96, 112 and 128 with a DCP4 decode on GPUs 4-7 and three draft tokens; 1P2D with two DCP4 decodes at concurrency 56 (four draft tokens); and 1P3D with three TP4 decodes at concurrency 12, 16 and 24 (five draft tokens). In 1P2D and 1P3D each decode has a node of its own." + - "Workers read the checkpoint from each node's local Hugging Face hub cache, as the single-node GLM-5.2 points do; srtctl fills the first node's cache before they start." + - "Patch srt-slurm so its ATOM backend passes -tp as the worker's GPUs over its pipeline and prefill-context stages (NVIDIA/srt-slurm#604); unpatched, the PP4 prefill launches with -tp 4." + pr-link: https://github.com/SemiAnalysisAI/InferenceX/pull/4045 diff --git a/inferencex-e2e/runners/srt-slurm/patches/604-atom-tp-from-pp.patch b/inferencex-e2e/runners/srt-slurm/patches/604-atom-tp-from-pp.patch new file mode 100644 index 0000000000..33a6019d1f --- /dev/null +++ b/inferencex-e2e/runners/srt-slurm/patches/604-atom-tp-from-pp.patch @@ -0,0 +1,60 @@ +From e23ca1bed9aee9a59905fb30b27f09b3ef9e9224 Mon Sep 17 00:00:00 2001 +From: Jasen2201 +Date: Wed, 14 Oct 2026 04:15:36 +0000 +Subject: [PATCH] fix(atom): size -tp from the worker's GPUs over its PP and + PCP stages + +ATOM runs a worker on tp x pp x pcp GPUs, but srtctl always passed +`-tp `. A 4-GPU prefill worker with +`pipeline-parallel-size: 4` therefore launched as +`-tp 4 --pipeline-parallel-size 4`, and ATOM sized its world to 16 GPUs. + +`-tp` is now the worker's GPU count divided by `pipeline-parallel-size` +and `prefill-context-parallel-size`. Decode context parallelism reuses the +TP GPUs and does not change it. A split that does not divide the worker's +GPUs fails before launch. + +Co-Authored-By: Claude Opus 5.5 +Signed-off-by: Jasen2201 +--- + src/srtctl/backends/atom.py | 20 +++++++++++++++++++- + 1 file changed, 19 insertions(+), 1 deletion(-) + +diff --git a/src/srtctl/backends/atom.py b/src/srtctl/backends/atom.py +index edf9a9b7..4e5b836e 100644 +--- a/src/srtctl/backends/atom.py ++++ b/src/srtctl/backends/atom.py +@@ -155,7 +155,7 @@ class AtomBackend(Backend): + "--server-port", + str(process.http_port), + "-tp", +- str(len(process.gpu_indices)), ++ str(_tensor_parallel_size(len(process.gpu_indices), config)), + ] + ) + kv_transfer_config = self._kv_transfer_config(process, worker_ip, extra_connectors) +@@ -183,3 +183,21 @@ def _config_to_cli_args(config: dict[str, Any]) -> list[str]: + + def _canonical_arg_key(key: str) -> str: + return key.lstrip("-").replace("_", "-") ++ ++ ++# ATOM sizes a worker as tp x pp x pcp GPUs; decode context parallelism reuses the TP GPUs. ++_GPU_SPLITTING_ARGS = ("pipeline-parallel-size", "prefill-context-parallel-size") ++ ++ ++def _tensor_parallel_size(gpu_count: int, config: Mapping[str, Any]) -> int: ++ """The worker's GPUs left for tensor parallelism after its pipeline and prefill-context stages.""" ++ stages = 1 ++ for key, value in config.items(): ++ if _canonical_arg_key(key) in _GPU_SPLITTING_ARGS: ++ stages *= int(value) ++ if stages < 1 or gpu_count % stages: ++ raise ValueError( ++ f"ATOM worker has {gpu_count} GPU(s), which its pipeline x prefill-context parallelism of {stages} " ++ "does not divide" ++ ) ++ return gpu_count // stages +-- +2.43.0 + diff --git a/inferencex-e2e/runners/srt-slurm/patches/README.md b/inferencex-e2e/runners/srt-slurm/patches/README.md index d9052d1aa8..a6b80e31c3 100644 --- a/inferencex-e2e/runners/srt-slurm/patches/README.md +++ b/inferencex-e2e/runners/srt-slurm/patches/README.md @@ -9,3 +9,4 @@ Each patch backports an upstream fix that the shared submodule pin does not incl | Patch | Upstream PR | Fix | |-------|-------------|-----| | `573-participating-gpus.patch` | [NVIDIA/srt-slurm#573](https://github.com/NVIDIA/srt-slurm/pull/573) | Record only the GPUs workers occupy on a worker node; an exclusive node's exporter (DCGM or AMD) reports every GPU, which the collector otherwise flags as `unexpected_device`. The pin includes #573, but upstream merged a different design: it keeps every row and excuses idle worker-node GPUs only after a DCGM worker GPU UUID probe. [`infx/results/power/multinode.py`](../../../infx/results/power/multinode.py) does not apply that excusal yet, so keep this filter until it does. | +| `604-atom-tp-from-pp.patch` | [NVIDIA/srt-slurm#604](https://github.com/NVIDIA/srt-slurm/pull/604) | ATOM's `-tp` is the worker's GPUs divided by its pipeline and prefill-context stages, so a 4-GPU PP4 prefill launches with `-tp 1` instead of `-tp 4` |