From 18bf7d0db2f8375c2e1992a4c412b7cd8e3882fb Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 07:36:12 +0000 Subject: [PATCH 01/20] perf(atom): stop the API-process gen-2 GC stalls under AgentX load Both ATOM API processes spent ~45% of the wall clock in stop-the-world gen-2 GC passes that grew with every completed request (~0.7 s per TTFT after an hour of AgentX C64). - ATOM patch: finished sequences are no longer queued on CoreManager.outputs_queue when serving online (nothing reads it there), and the per-request context gauges are opt-in (ATOM_METRICS_PER_REQUEST_CONTEXT=1). - Router: the frontend /metrics no longer scrapes engines it reads nothing from (ATOM), halving the /metrics renders of each ATOM API. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../atom_nightly/patch_atom_api_gc_stalls.py | 78 +++++++++++++++++++ infera/common/engine_metrics.py | 19 +++++ infera/server/app.py | 11 ++- .../common/test_federate_engine_metrics.py | 14 +++- 4 files changed, 119 insertions(+), 3 deletions(-) create mode 100644 deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py diff --git a/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py b/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py new file mode 100644 index 000000000..3ca127684 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Keep the API process heap flat so its gen-2 GC pauses stay short. + +Under AgentX load both ATOM API processes spent ~45% of the wall clock in +stop-the-world gen-2 GC passes, nearly all triggered inside a /metrics render, +and each pass grew with the number of completed requests (D ~0.08 ms per +request: ~0.65 s per pass after an hour). Every pass freezes the event loop, +so request admission, the P result hand-back and SSE delivery all wait; it +added ~0.7 s to each TTFT. Two things grow with completed requests: + +1. ``CoreManager.outputs_queue``. The output thread put every batch of finished + sequences on it, but only ``LLMEngine.step()`` / ``generate()`` and asyncio + mode read it; the OpenAI server delivers through the per-sequence stream + callbacks instead. So every finished ``Sequence`` (token arrays, KV transfer + block-id lists) stayed reachable for the life of the process and each gen-2 + pass walked all of them -- the same leak ATOM already fixed for + ``stream_outputs_queue``. The API server is the only caller that wires + ``_flush_stream_batch_fn``; with it set and asyncio mode off, nothing will + ever read the queue, so the batch is not queued. +2. The per-request context gauges (``atom:{prefill,decode}_request_context_tokens``, + labelled by request id and never removed). The multiprocess collector + re-reads and json-parses every series on each render, so render time and the + objects it promotes grow with every request served. They are opt-in now: + ``ATOM_METRICS_PER_REQUEST_CONTEXT=1`` records them as before. + +Written against ATOM 661ef766 in rocm/atom-dev:nightly_202610131505. +Strict: see anchor_patch.py. +""" + +import sys + +from anchor_patch import apply + +PATCH = { + "model_engine/engine_core_mgr.py": ( + "# infera: nothing reads outputs_queue online", + [ + ( + " self.outputs_queue.put_nowait(seqs)\n", + " # infera: nothing reads outputs_queue online. The\n" + " # API server (the one caller that wires the flush\n" + " # hook) delivers through the stream callbacks; only\n" + " # LLMEngine.step()/generate() and asyncio mode read\n" + " # this queue. Queued anyway, every finished Sequence\n" + " # stayed reachable and each gen-2 GC pass grew.\n" + " if self._flush_stream_batch_fn is None or getattr(\n" + " self, \"_asyncio_mode\", False\n" + " ):\n" + " self.outputs_queue.put_nowait(seqs)\n", + ), + ], + ), + "metrics/scheduler.py": ( + "# infera: per-request context gauges are opt-in", + [ + ("import time\n", "import os\nimport time\n"), + ( + " def _record_request_context(self, metric, req_id, seq, tokens):\n", + " # infera: per-request context gauges are opt-in. Each request adds a\n" + " # series that is never removed, and every /metrics render of the API\n" + " # process re-reads all of them (render time and GC grow with requests).\n" + " _PER_REQUEST = os.environ.get(\"ATOM_METRICS_PER_REQUEST_CONTEXT\") == \"1\"\n" + "\n" + " def _record_request_context(self, metric, req_id, seq, tokens):\n" + " if not self._PER_REQUEST:\n" + " return\n", + ), + ], + ), +} + +if __name__ == "__main__": + sys.exit(apply("atom-patch api-gc-stalls", PATCH, package="atom")) diff --git a/infera/common/engine_metrics.py b/infera/common/engine_metrics.py index 408d61fdc..e384b7b77 100644 --- a/infera/common/engine_metrics.py +++ b/infera/common/engine_metrics.py @@ -251,6 +251,25 @@ def _inject_labels(sample_line: str, extra: dict[str, str]) -> str: return f"{name}{new_block} {rest}" +def frontend_reads(engine: str | EngineType) -> bool: + """Whether the frontend /metrics takes anything from ``engine``'s exposition. + + False for ATOM (and unknown engines): no gauge names, transfer queues or + federated families are known for it, so scraping it only costs the engine a + /metrics render. + """ + try: + engine = EngineType(engine) + except ValueError: + return False + return bool( + metric_names("kv_cache_usage", engine) + or metric_names("prefix_cache_hit_rate", engine) + or transfer_queue_names(engine) + or _FEDERATE_FAMILIES.get(engine) + ) + + def federate_engine_metrics( text: str, *, diff --git a/infera/server/app.py b/infera/server/app.py index 0a18ba03c..f2a3f6631 100644 --- a/infera/server/app.py +++ b/infera/server/app.py @@ -761,14 +761,21 @@ async def _scrape_engine_metrics() -> bytes: """ from infera.common.engine_metrics import ( federate_engine_metrics, + frontend_reads, merge_federated_exposition, ) if registry is None: return b"" # Scrape every active worker over its HTTP /metrics URL. NATS is only the - # request path; the engine still exposes Prometheus on w.url. - workers = [w for w in registry.list_all() if w.status == WorkerStatus.ACTIVE] + # request path; the engine still exposes Prometheus on w.url. Skip engines + # nothing here is read from (ATOM): each scrape costs the engine's API + # process a render, and under a 3/s client scrape that doubled its GC load. + workers = [ + w + for w in registry.list_all() + if w.status == WorkerStatus.ACTIVE and frontend_reads(w.engine) + ] metrics.engine_kv_cache_usage.clear() metrics.engine_prefix_cache_hit_rate.clear() metrics.engine_kv_transfer_queue_reqs.clear() diff --git a/tests/unit/common/test_federate_engine_metrics.py b/tests/unit/common/test_federate_engine_metrics.py index fc3e892ca..3e8484045 100644 --- a/tests/unit/common/test_federate_engine_metrics.py +++ b/tests/unit/common/test_federate_engine_metrics.py @@ -7,7 +7,11 @@ from __future__ import annotations -from infera.common.engine_metrics import federate_engine_metrics, merge_federated_exposition +from infera.common.engine_metrics import ( + federate_engine_metrics, + frontend_reads, + merge_federated_exposition, +) from infera.common.worker_pool import EngineType @@ -82,3 +86,11 @@ def test_merge_federated_dedupes_help_and_type(): assert out.count("# TYPE vllm:num_requests_running") == 1 assert 'worker_id="w1"' in out assert 'worker_id="w2"' in out + + +def test_frontend_reads_skips_engines_it_takes_nothing_from(): + assert frontend_reads(EngineType.VLLM) + assert frontend_reads("sglang") + # Nothing is mapped or federated for ATOM: scraping it only costs a render. + assert not frontend_reads(EngineType.ATOM) + assert not frontend_reads("unknown-engine") From 4957ac36d04a5328ed7b0063826a19e5ea059c97 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 07:42:10 +0000 Subject: [PATCH 02/20] perf(atom): DCP top-k exchange on the custom all-gather; MTP experts knob - ATOM patch: the sparse indexer's DCP top-k candidate exchange (24 gathers per decode step, ~1 ms) goes through dcp_all_gather, which picks aiter's IPC all-gather instead of pynccl. - Kit: MTP_EXPERTS=mxfp4 quantizes the MTP layer's routed experts online to MXFP4 (they are BF16 in the checkpoint). Default bf16 keeps the previous online quant config byte for byte. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../patch_atom_dcp_topk_custom_ag.py | 43 +++++++++++++++++++ examples/atom_4p4d_glm5.2/engine/common.sh | 2 +- examples/atom_4p4d_glm5.2/engine/worker.sh | 15 +++++-- examples/atom_4p4d_glm5.2/env.sh | 5 +++ 4 files changed, 61 insertions(+), 4 deletions(-) create mode 100644 deploy/docker/patches/atom_nightly/patch_atom_dcp_topk_custom_ag.py diff --git a/deploy/docker/patches/atom_nightly/patch_atom_dcp_topk_custom_ag.py b/deploy/docker/patches/atom_nightly/patch_atom_dcp_topk_custom_ag.py new file mode 100644 index 000000000..ce673d4ea --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_atom_dcp_topk_custom_ag.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Exchange the sparse-indexer DCP top-k candidates over aiter's custom all-gather. + +In DCP decode every indexer layer gathers each rank's local top-k (score, id) +candidates across the DCP group (21 indexer layers in the verify graph plus one +per draft for GLM-5.2: 24 gathers per step of 1.3 MB per rank at graph size 40). +That call used ``GroupCoordinator.all_gather``, whose default is pynccl, at +~100 GB/s bus bandwidth (~43 us each, ~1 ms per decode step), while ATOM's own +``dcp_all_gather`` -- already used for the LSE and query-head gathers -- picks +aiter's IPC all-gather when the shape qualifies and falls back to pynccl +otherwise. The payload is a bit-exact copy either way (moved as 4-byte words). + +Written against ATOM 661ef766 in rocm/atom-dev:nightly_202610131505. +Strict: see anchor_patch.py. +""" + +import sys + +from anchor_patch import apply + +PATCH = { + "model_ops/dcp_ops.py": ( + "# infera: custom all-gather for the top-k exchange", + [ + ( + " gathered_sc = (\n" + " get_dcp_group()\n" + " .all_gather(send.view(torch.int32), dim=0)\n", + " # infera: custom all-gather for the top-k exchange (pynccl otherwise).\n" + " gathered_sc = (\n" + " dcp_all_gather(get_dcp_group(), send.view(torch.int32), 0)\n", + ), + ], + ), +} + +if __name__ == "__main__": + sys.exit(apply("atom-patch dcp-topk-custom-ag", PATCH, package="atom")) diff --git a/examples/atom_4p4d_glm5.2/engine/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index 0ae4d9c23..3ed4e9a1b 100755 --- a/examples/atom_4p4d_glm5.2/engine/common.sh +++ b/examples/atom_4p4d_glm5.2/engine/common.sh @@ -34,7 +34,7 @@ CONTAINERS=() for kit_role in $KIT_ROLES; do CONTAINERS+=("$CTR-$kit_role"); done # Settings up.sh records per run, so later scripts report what actually runs. -RUN_VARS=(CONC MTP_K MTP_AL PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL +RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL PREFILL_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS DECODE_EXTRA_ARGS ROUTER_EXTRA_ARGS IMAGE MODEL) diff --git a/examples/atom_4p4d_glm5.2/engine/worker.sh b/examples/atom_4p4d_glm5.2/engine/worker.sh index 4861a3f51..84fc9f08d 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -15,9 +15,18 @@ case "$role" in prefill | decode) ;; *) die "usage: bash engine/worker.sh prefil case "$PREFILL_PP" in 1 | 4) ;; *) die "PREFILL_PP must be 1 or 4, not $PREFILL_PP" ;; esac resolve_model -# Experts of layers 0-77 stay MXFP4 and the MTP layer 78 stays BF16; the rest -# is quantized online to per-token FP8. -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.*"]}' +# Experts of layers 0-77 stay MXFP4 and the rest is quantized online to +# per-token FP8, except the MTP layer 78: BF16 as in the checkpoint, or with +# MTP_EXPERTS=mxfp4 its routed experts quantized online to MXFP4 at load. +case "$MTP_EXPERTS" in + bf16) mtp_quant="" mtp_exclude='"model.layers.78.*"' ;; + mxfp4) + mtp_quant=',"layer_quant_config":{"model.layers.78.mlp.experts*":"mxfp4"}' + mtp_exclude='"re:^model\\.layers\\.78\\.(?!mlp\\.experts)"' + ;; + *) die "MTP_EXPERTS must be bf16 or mxfp4, not $MTP_EXPERTS" ;; +esac +QUANT_CONFIG='{"global_quant_config":"ptpc_fp8"'"$mtp_quant"',"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*",'"$mtp_exclude"']}' env=( AITER_QUICK_REDUCE_QUANTIZATION=INT4 AITER_USE_FLYDSL_MOE_SORTING=1 AITER_LOG_LEVEL=WARNING diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index d17a13965..68ba8a31c 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -62,6 +62,11 @@ export MTP_K="${MTP_K:-$((CONC >= 48 ? 3 : 4))}" # Set it (K3 -> 2.99, K4 -> 3.33) only to reproduce the AgentX reference numbers. export MTP_AL="${MTP_AL:-}" +# Routed experts of the MTP layer (78), which the checkpoint keeps in BF16: +# bf16 as shipped, or mxfp4 to quantize them online at load like the experts +# of the other layers (the draft MoE then reads a quarter of the weight bytes). +export MTP_EXPERTS="${MTP_EXPERTS:-bf16}" + # The prefill memory, chunk, block size, layer split and LMCache load length # below follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md # in the image); the decode keeps this kit's settings. From 705b63eac2031bba061dfb21f8d34caa5f6073e4 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 07:43:48 +0000 Subject: [PATCH 03/20] feat(examples): DECODE_GRAPH_STEP sets the CUDA-graph size spacing above 8 Default 4 keeps the previous 34 sizes at CONC 64. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/engine/common.sh | 2 +- examples/atom_4p4d_glm5.2/engine/worker.sh | 7 ++++--- examples/atom_4p4d_glm5.2/env.sh | 5 +++++ 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/examples/atom_4p4d_glm5.2/engine/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index 3ed4e9a1b..1e3a34a91 100755 --- a/examples/atom_4p4d_glm5.2/engine/common.sh +++ b/examples/atom_4p4d_glm5.2/engine/common.sh @@ -34,7 +34,7 @@ CONTAINERS=() for kit_role in $KIT_ROLES; do CONTAINERS+=("$CTR-$kit_role"); done # Settings up.sh records per run, so later scripts report what actually runs. -RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL +RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS DECODE_GRAPH_STEP PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL PREFILL_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS DECODE_EXTRA_ARGS ROUTER_EXTRA_ARGS IMAGE MODEL) diff --git a/examples/atom_4p4d_glm5.2/engine/worker.sh b/examples/atom_4p4d_glm5.2/engine/worker.sh index 84fc9f08d..11d80b44b 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -41,10 +41,11 @@ env=( TRITON_CACHE_DIR=/root/.cache/triton TORCHINDUCTOR_CACHE_DIR=/root/.cache/inductor ) -# Graph capture bounded by the concurrency: 1, 2, 4, 8, then every 4 up to 2 x CONC. -# Left unset, a DCP prefill would capture up to batch 512; a PP prefill runs eager. +# Graph capture bounded by the concurrency: 1, 2, 4, 8, then every +# DECODE_GRAPH_STEP up to 2 x CONC. Left unset, a DCP prefill would capture up +# to batch 512; a PP prefill runs eager. cudagraph="[1,2,4,8" -for ((size = 12; size <= 2 * CONC; size += 4)); do cudagraph+=",$size"; done +for ((size = 8 + DECODE_GRAPH_STEP; size <= 2 * CONC; size += DECODE_GRAPH_STEP)); do cudagraph+=",$size"; done cudagraph+="]" # kv_mooncake KV_ROLE HANDSHAKE_PORT HTTP_PORT diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index 68ba8a31c..665e9c093 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -54,6 +54,11 @@ export DECODE_GPUS="${DECODE_GPUS:-4,5,6,7}" # 2 x CONC (also of a DCP prefill); restart the engines to change it. export CONC="${CONC:-48}" +# Spacing of the CUDA-graph batch sizes above 8. A batch runs in the next +# larger graph, so a finer step pads fewer rows per decode step at the cost of +# more graphs (capture time and graph memory). +export DECODE_GRAPH_STEP="${DECODE_GRAPH_STEP:-4}" + # MTP draft depth: K3 at CONC >= 48, K4 below, as in the reference runs. export MTP_K="${MTP_K:-$((CONC >= 48 ? 3 : 4))}" From 01a3fd49b71a2996a11cffa72b413e6c513de2b2 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 09:47:52 +0000 Subject: [PATCH 04/20] perf(atom): per-M Triton configs for the MLA absorbed batched GEMMs aiter's gfx950 configs for the MLA q/v up-projection batched GEMM use BLOCK_SIZE_M 64 from M_LEQ_64 up, so every decode M that is not a multiple of 64 (graph 36/40/44, drafts from 36) takes the masked path at 2-2.5x the time. One bucket per decode M with the sweep's fastest config (MI355X, output-checked): ~1.4 ms per decode step at graph 40. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../patch_aiter_mla_bmm_configs.py | 144 ++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py b/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py new file mode 100644 index 000000000..e90a791dc --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Per-M Triton configs for the MLA absorbed batched GEMMs of GLM-5.2 decode. + +ATOM's MLA decode runs aiter's batched_gemm_a8w8_a_per_token_group_prequant_w_ +per_batched_tensor_quant twice per layer over all 64 heads of the DCP group: +the q up-projection (N=512, K=192) and the v up-projection (N=256, K=512), with +M = query rows (4 x graph size in the MTP verify, the graph size in a draft). +Its gfx950 configs cover only the standard M buckets, and from M_LEQ_64 up they +use BLOCK_SIZE_M 64: any M that is not a multiple of 64 -- graph 36, 40, 44 +(M 144, 160, 176), every draft from 36 up -- takes the masked path at 2-2.5x +the time of a neighbouring multiple of 64 (q-up 19.1 us at M 160 vs 10.6 at +192). That is most of the irregular step-time staircase between graph sizes. + +This adds an M_BOUNDS list with one bucket per decode M (multiples of 4 up to +64, of 16 up to 512) and the fastest config for each from a sweep on MI355X +(72 candidates, CUDA-graph timing, output checked against the stock config): +at graph 40 q-up 19.1 -> 8.4 us and v-up 15.2 -> 8.1 us, ~1.4 ms per decode +step over 78 layers. The stock M_GEQ_*/any entries still serve M > 512. + +Written against aiter 9205e1059 in rocm/atom-dev:nightly_202610131505. Strict: +a file that no longer holds the stock buckets fails the build. +""" + +import importlib.util +import json +import os +import sys + +TAG = "aiter-patch mla-bmm-configs" +CFG_DIR = ( + "ops/triton/configs/gfx950/triton/gemm/" + "batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant" +) +NAME = "BATCHED_GEMM-A8W8-A_PER_TOKEN_GROUP_PREQUANT_W_PER_BATCHED_TENSOR_QUANT" +# The stock M_LEQ_256 entry of both files; checked before anything is written. +STOCK_LEQ_256 = { + "BLOCK_SIZE_M": 64, + "BLOCK_SIZE_N": 256, + "GROUP_SIZE_M": 1, + "num_warps": 8, + "num_stages": 2, + "waves_per_eu": 1, + "matrix_instr_nonkdim": 16, + "cache_modifier": ".cg", +} +# (N, K): [(M, BLOCK_SIZE_M, BLOCK_SIZE_N, num_warps, waves_per_eu, cache_modifier)] +TUNED = { + (512, 192): [ + (16, 16, 128, 4, 2, None), (20, 16, 128, 4, 2, None), (24, 16, 128, 4, 2, None), + (28, 16, 128, 4, 2, ".cg"), (32, 16, 128, 4, 2, ".cg"), + (36, 32, 128, 8, 1, ".cg"), (40, 32, 128, 8, 1, None), + (44, 32, 128, 8, 1, ".cg"), (48, 16, 256, 4, 2, None), + (52, 32, 128, 8, 1, None), (56, 32, 128, 8, 1, None), + (60, 32, 128, 8, 1, ".cg"), (64, 32, 128, 4, 2, None), + (80, 16, 128, 8, 1, None), (96, 32, 256, 8, 1, ".cg"), + (112, 16, 128, 8, 1, None), (128, 32, 256, 8, 1, None), + (144, 16, 128, 8, 1, None), (160, 32, 128, 8, 1, ".cg"), + (176, 16, 128, 8, 1, None), (192, 32, 128, 8, 1, ".cg"), + (208, 16, 128, 8, 1, None), (224, 32, 128, 8, 1, ".cg"), + (240, 16, 128, 8, 1, None), (256, 64, 256, 4, 2, None), + (272, 16, 128, 8, 1, None), (288, 32, 128, 8, 1, None), + (304, 16, 128, 8, 1, None), (320, 32, 128, 8, 1, ".cg"), + (336, 16, 128, 8, 1, ".cg"), (352, 32, 128, 8, 1, ".cg"), + (368, 16, 128, 8, 1, ".cg"), (384, 64, 256, 8, 1, ".cg"), + (400, 16, 128, 8, 1, None), (416, 32, 128, 8, 1, ".cg"), + (432, 16, 128, 8, 1, ".cg"), (448, 32, 128, 8, 1, None), + (464, 16, 128, 8, 1, ".cg"), (480, 32, 128, 8, 1, ".cg"), + (496, 16, 128, 8, 1, None), (512, 64, 256, 4, 2, None), + ], + (256, 512): [ + (16, 16, 64, 4, 2, None), (20, 16, 64, 4, 2, None), (24, 16, 64, 4, 2, None), + (28, 16, 64, 4, 2, None), (32, 16, 128, 4, 1, None), (36, 16, 128, 4, 2, None), + (40, 16, 128, 4, 2, None), (44, 16, 128, 4, 2, None), (48, 16, 128, 4, 2, None), + (52, 16, 128, 4, 2, None), (56, 16, 128, 4, 2, None), (60, 16, 128, 4, 2, None), + (64, 16, 128, 4, 2, None), (80, 16, 256, 4, 2, None), (96, 32, 256, 8, 1, None), + (112, 16, 256, 4, 2, None), (128, 32, 256, 8, 2, None), + (144, 16, 256, 4, 2, None), (160, 32, 128, 8, 1, None), + (176, 16, 256, 4, 2, None), (192, 32, 128, 8, 1, None), + (208, 16, 256, 4, 2, None), (224, 32, 256, 4, 2, None), + (240, 16, 256, 4, 2, None), (256, 32, 256, 4, 2, None), + (272, 16, 256, 4, 2, None), (288, 32, 256, 8, 1, None), + (304, 16, 256, 4, 2, None), (320, 32, 256, 8, 1, None), + (336, 16, 256, 4, 2, None), (352, 32, 256, 8, 1, None), + (368, 16, 256, 4, 2, None), (384, 32, 256, 8, 1, None), + (400, 16, 256, 4, 2, None), (416, 32, 256, 8, 1, None), + (432, 16, 256, 4, 2, None), (448, 64, 256, 8, 1, None), + (464, 16, 256, 4, 2, None), (480, 32, 256, 8, 1, None), + (496, 16, 256, 4, 2, None), (512, 64, 256, 8, 1, None), + ], +} +# Above 512 the stock M_GEQ_* / any entries keep serving (prefill sizes). +UPPER_BOUNDS = [1024, 2048, 4096, 8192] + + +def main() -> int: + spec = importlib.util.find_spec("aiter") + if spec is None or not spec.submodule_search_locations: + print(f"[{TAG}] FAILED: package 'aiter' not found") + return 1 + root = os.path.join(list(spec.submodule_search_locations)[0], CFG_DIR) + pending = {} + for (n, k), rows in TUNED.items(): + path = os.path.join(root, f"{NAME}-N={n}-K={k}.json") + if not os.path.isfile(path): + print(f"[{TAG}] FAILED: {path} not found") + return 1 + with open(path, encoding="utf-8") as f: + cfg = json.load(f) + bounds = [m for m, *_ in rows] + UPPER_BOUNDS + if cfg.get("M_BOUNDS") == bounds: + continue # already applied + if "M_BOUNDS" in cfg or cfg.get("M_LEQ_256") != STOCK_LEQ_256: + print(f"[{TAG}] FAILED: {path} is not the stock config") + return 1 + for key in [key for key in cfg if key.startswith("M_LEQ_")]: + del cfg[key] + cfg["M_BOUNDS"] = bounds + for m, bm, bn, warps, waves, cache in rows: + cfg[f"M_LEQ_{m}"] = { + "BLOCK_SIZE_M": bm, + "BLOCK_SIZE_N": bn, + "GROUP_SIZE_M": 1, + "num_warps": warps, + "num_stages": 2, + "waves_per_eu": waves, + "matrix_instr_nonkdim": 16, + "cache_modifier": cache, + } + pending[path] = cfg + for path, cfg in pending.items(): + with open(path, "w", encoding="utf-8") as f: + json.dump(cfg, f, indent=4) + f.write("\n") + print(f"[{TAG}] {'patched ' + ', '.join(pending) if pending else 'already applied'}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 2fcca8fce5797cb4045a6b7406ed10b78fb2955c Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 11:23:16 +0000 Subject: [PATCH 05/20] perf(atom): DCP attention merge over a custom IPC all-to-all The DCP decode merge runs one all_to_all_single per layer on RCCL (81 per GLM-5.2 step, ~3.3 ms), and no RCCL setting helps. - aiter patch: alltoall_vec next to allgather_vec on the custom all-reduce buffers, reached through dispatchAllGather's gather_dim (-1000) so the registered-buffer and CUDA-graph paths apply as they are; CustomAllreduce.custom_all_to_all; module rebuilt at build time. - ATOM patch: the DCP merge uses it when it applies, RCCL otherwise. 4 x MI355X, bit-identical to RCCL: 5.3 MB/rank 56 -> 34 us, 1.3 MB/rank 32 -> 13 us. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../patch_aiter_custom_alltoall.py | 199 ++++++++++++++++++ .../atom_nightly/patch_atom_dcp_custom_a2a.py | 55 +++++ 2 files changed, 254 insertions(+) create mode 100644 deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py create mode 100644 deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a.py diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py new file mode 100644 index 000000000..bb8a6edb1 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""An all-to-all on aiter's custom (IPC) all-reduce buffers, for the DCP merge. + +ATOM's DCP decode merges the partial attention of every rank with one +all_to_all_single per layer (81 per GLM-5.2 step, 5.3 MB per rank at graph +size 40), on RCCL: ~40 us each in the decode graph, ~3.3 ms per step, and no +RCCL setting helps (NCCL_PROTO / channel counts, measured). aiter's custom +all-gather moves the same bytes per rank (its 1.3 MB input x 4 ranks) in about +half the time, so this adds the all-to-all next to it: + +- ``alltoall_vec``: allgather_vec reading only the chunk addressed to this + rank: rank r receives chunk r of every rank's input, in rank order (the + all_to_all_single layout). Same start/end sync, same <= kMaxBlocks grid. +- ``dispatchAllGather`` runs it for gather_dim == -1000, so the existing + all_gather_reg/unreg bindings, the copy-in pool and the CUDA-graph buffer + registration all apply unchanged; no new pybind symbol. +- ``CustomAllreduce.custom_all_to_all(inp)``: the Python entry, None where the + path does not apply (gfx1250 bindings, a size the vector kernel cannot split). + +The patched module (module_custom_all_reduce) is rebuilt here, the way the +runtime JIT would; needs no GPU, the target archs come from GPU_ARCHS +(default PYTORCH_ROCM_ARCH). Written against aiter 9205e1059 in +rocm/atom-dev:nightly_202610131505. Strict: see anchor_patch.py. +""" + +import importlib.util +import os +import sys + +from anchor_patch import apply + +TAG = "aiter-patch custom-alltoall" + +KERNEL = r""" +// infera: all-to-all. size = 16-byte packs per destination chunk; rank r +// receives chunk r of every rank's registered input, in rank order. +constexpr int kAllToAllDim = -1000; + +template +__global__ void __launch_bounds__(512, 1) alltoall_vec( + RankData* _dp, RankSignals sg, Signal* self_sg, T* __restrict__ result, int rank, int size) +{ + constexpr int tnum_gpu = THREAD_NUM / ngpus; + constexpr int pack_size = 16 / sizeof(T); + using P = typename opus::vector_t; + int warp_id = threadIdx.x / tnum_gpu; + int lane_id = threadIdx.x % tnum_gpu; + int tid = blockIdx.x * tnum_gpu + lane_id; + int stride = gridDim.x * tnum_gpu; + const P* ptrs[ngpus]; + +#pragma unroll + for(int i = 0; i < ngpus; ++i) + { + ptrs[i] = (const P*)_dp->ptrs[i]; + } + start_sync(sg, self_sg, rank); + + for(int idx = tid; idx < size; idx += stride) + { + *(reinterpret_cast(&result[0]) + warp_id * size + idx) = ptrs[warp_id][rank * size + idx]; + } + end_sync(sg, self_sg, rank); +} + +""" + +DISPATCH = """ // infera: all-to-all, see alltoall_vec. + if(gather_dim == kAllToAllDim) + { + int chunk = size / d / world_size_; + int tnum_per_block = 512 / world_size_; + int block_num = (chunk + tnum_per_block - 1) / tnum_per_block; + dim3 grid(std::min(block_num, kMaxBlocks)); + switch(world_size_) + { + case 8: + alltoall_vec<<>>(ptrs, sg_, self_sg_, output, rank_, chunk); + break; + case 4: + alltoall_vec<<>>(ptrs, sg_, self_sg_, output, rank_, chunk); + break; + case 2: + alltoall_vec<<>>(ptrs, sg_, self_sg_, output, rank_, chunk); + break; + default: printf("alltoall world_size error\\n"); + } + return; + } +""" + +PY_METHOD = ''' # infera: dispatchAllGather runs the all-to-all kernel for this "dim". + _ALL_TO_ALL_DIM = -1000 + + def custom_all_to_all(self, inp: torch.Tensor) -> torch.Tensor | None: + """all_to_all_single over dim 0 of ``inp``: its world_size equal chunks + go to ranks 0..N-1 and the result holds every rank's chunk for this + rank, in rank order. Buffer and capture rules as custom_all_gather; + None when this path does not apply (the caller falls back to RCCL).""" + nbytes = inp.numel() * inp.element_size() + if ( + self._ops_all_gather is not None + or self._INT_TO_FP_VIEW.get(inp.dtype, inp.dtype) + not in (torch.float32, torch.float16, torch.bfloat16) + or nbytes % (16 * self.world_size) != 0 + or not self.should_custom_ag(inp) + ): + return None + view_dtype = self._INT_TO_FP_VIEW.get(inp.dtype) or inp.dtype + x = inp.view(view_dtype) + out = torch.empty_like(x) + capturing = self._IS_CAPTURING and torch.cuda.is_current_stream_capturing() + if capturing and self.enable_register_for_capturing: + ops.all_gather_reg(self._ptr, x, out, self._ALL_TO_ALL_DIM) + else: + ops.all_gather_unreg( + self._ptr, + x, + self._pool["input"].data_ptr, + out, + self._pool["input"].max_size, + self._ALL_TO_ALL_DIM, + ) + return out.view(inp.dtype) + + def fused_ar_rms( +''' + +PATCH_CSRC = { + "../csrc/include/custom_all_reduce.cuh": ( + "// infera: all-to-all.", + [ + ( + "// ========== reduce_scatter kernel start ==========\n", + KERNEL + "// ========== reduce_scatter kernel start ==========\n", + ), + ( + " // only support gather first dim and gather last dim\n", + DISPATCH + " // only support gather first dim and gather last dim\n", + ), + ], + ), + "dist/device_communicators/custom_all_reduce.py": ( + "# infera: dispatchAllGather runs the all-to-all kernel", + [(" def fused_ar_rms(\n", PY_METHOD)], + ), +} + +MODULE = "module_custom_all_reduce" + + +def rebuild() -> int: + """Rebuild the module from the patched sources, as the runtime JIT would.""" + root = importlib.util.find_spec("aiter").submodule_search_locations[0] + os.environ.setdefault("GPU_ARCHS", os.environ.get("PYTORCH_ROCM_ARCH", "gfx950")) + sys.path.insert(0, root) + from jit import core # noqa: E402 (import after GPU_ARCHS is set) + + so = f"{core.get_user_jit_dir()}/{MODULE}.so" + cuh = os.path.join(root, "../csrc/include/custom_all_reduce.cuh") + if os.path.exists(so): + if os.path.getmtime(so) > os.path.getmtime(cuh): + print(f"[{TAG}] {so} is newer than the patched header; not rebuilt") + return 0 + os.remove(so) + args = core.get_args_of_build(MODULE) + core.build_module( + MODULE, + args["srcs"], + args["flags_extra_cc"], + args["flags_extra_hip"], + args["blob_gen_cmd"], + args["extra_include"], + args["extra_ldflags"], + args["verbose"], + args["is_python_module"], + args["is_standalone"], + args["torch_exclude"], + args.get("third_party", []), + args.get("hipify", False), + flags_extra_hip_per_source=args.get("flags_extra_hip_per_source", {}), + ) + if not os.path.isfile(so): + print(f"[{TAG}] FAILED: AITER build left no {so}") + return 1 + print(f"[{TAG}] rebuilt {so}") + return 0 + + +if __name__ == "__main__": + rc = apply(TAG, PATCH_CSRC, package="aiter") + if rc == 0 and "--no-build" not in sys.argv: + rc = rebuild() + sys.exit(rc) diff --git a/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a.py b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a.py new file mode 100644 index 000000000..d977c9006 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""Merge DCP attention partials over aiter's custom all-to-all. + +Each DCP decode layer exchanges its per-head partial outputs and LSEs with one +all_to_all_single on RCCL (81 per GLM-5.2 step, ~3.3 ms). With +patch_aiter_custom_alltoall.py the DCP group's custom all-reduce communicator +offers custom_all_to_all; use it when it applies and keep RCCL otherwise. The +choice depends only on the shape, dtype and communicator, so every rank of the +group makes the same one. + +Written against ATOM 661ef766 in rocm/atom-dev:nightly_202610131505. +Strict: see anchor_patch.py. +""" + +import sys + +from anchor_patch import apply + +PATCH = { + "model_ops/dcp_ops.py": ( + "# infera: custom all-to-all for the DCP merge", + [ + ( + "class CPTritonContext:\n", + "def _dcp_custom_all_to_all(cp_group, send):\n" + ' """aiter\'s custom (IPC) all-to-all over the DCP group, or None for RCCL."""\n' + ' device_comm = getattr(cp_group, "device_communicator", None)\n' + ' ca = getattr(device_comm, "ca_comm", None)\n' + ' if ca is None or ca.disabled or not hasattr(ca, "custom_all_to_all"):\n' + " return None\n" + " return ca.custom_all_to_all(send)\n" + "\n" + "\n" + "class CPTritonContext:\n", + ), + ( + " recv = torch.empty_like(send)\n" + " torch.distributed.all_to_all_single(recv, send, group=cp_group.device_group)\n", + " # infera: custom all-to-all for the DCP merge, RCCL when it does not apply.\n" + " recv = _dcp_custom_all_to_all(cp_group, send)\n" + " if recv is None:\n" + " recv = torch.empty_like(send)\n" + " torch.distributed.all_to_all_single(recv, send, group=cp_group.device_group)\n", + ), + ], + ), +} + +if __name__ == "__main__": + sys.exit(apply("atom-patch dcp-custom-a2a", PATCH, package="atom")) From 3369354aa2937aa1bd60c5be349c5051a373e0b9 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 11:45:37 +0000 Subject: [PATCH 06/20] perf(atom): tuned FP8 GEMM rows for GLM-5.2 decode row counts above 128 aiter's GLM-5.2 a8w8_bpreshuffle table holds M = 64/128/256, so the MTP verify at graph >= 36 (M = 144 ...) runs padded or default FlyDSL kernels. Tuned M 144..240 for the nine GLM-5.2 shapes and kept the 27 rows at least 5% faster than the production path (e.g. qkv_a at M 160 14.9 -> 11.5 us), as an extra model_configs table. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../patch_aiter_glm52_decode_gemm_rows.py | 84 +++++++++++++++++++ 1 file changed, 84 insertions(+) create mode 100644 deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py b/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py new file mode 100644 index 000000000..f871feb8c --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""FP8 GEMM rows for the decode row counts of GLM-5.2 above 128. + +The per-token FP8 projections of the decode (qkv_a, q_b, o_proj, indexer ...) +look their FlyDSL kernel up in aiter's a8w8_bpreshuffle tuned table, which for +GLM-5.2 holds M = 64, 128, 256 (a few 96 / 192). The MTP verify runs +M = 4 x graph size, so every graph from 36 up (M 144 ...) takes a padded or +default kernel: the qkv_a GEMM costs 14.9 us at M 160 against 11.3 at 128. + +These rows come from aiter's tuner (gemm_a8w8_bpreshuffle_tune.py, FlyDSL +candidates, MI355X) for M 144..240; only the ones at least 5% faster than the +production path at the same shape (the tuner's --run_config, measured on the +same node) are kept, since for M not a multiple of 32 the tuned kernel is often +slower than the padded one. They go in as one more model_configs table, which +aiter merges with the others at startup (none of these keys exist there). + +Written against aiter 9205e1059 in rocm/atom-dev:nightly_202610131505. +""" + +import importlib.util +import os +import sys + +TAG = "aiter-patch glm52-decode-gemm-rows" +REL = "configs/model_configs/a8w8_bpreshuffle_tuned_gemm_glm5.2_decode.csv" +CSV = """\ +gfx,cu_num,M,N,K,q_dtype_w,libtype,kernelId,splitK,us,kernelName,tflops,bw,errRatio +gfx950,256,160,2048,2048,torch.float8_e4m3fn,flydsl,2371,0,5.6743,flydsl_bpreshuflle_32x32x512_F8_F8_B16_1x4x4x2_default,236.54,912.42,0.0 +gfx950,256,192,2048,2048,torch.float8_e4m3fn,flydsl,686,0,5.8112,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x2x4x2_default,277.16,924.76,0.0 +gfx950,256,224,2048,2048,torch.float8_e4m3fn,flydsl,914,0,5.8566,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x3x4x2_default,320.84,951.16,0.0 +gfx950,256,144,2688,6144,torch.float8_e4m3fn,flydsl,557,0,13.4479,flydsl_bpreshuflle_16x64x512_F8_F8_B16_0x2x4x2_default,353.69,1351.43,0.0 +gfx950,256,160,2688,6144,torch.float8_e4m3fn,flydsl,1117,0,11.5172,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x4x4x2_default,458.86,1593.99,0.0 +gfx950,256,160,3072,6144,torch.float8_e4m3fn,flydsl,1117,0,11.7838,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x4x4x2_default,512.55,1768.57,0.0 +gfx950,256,192,3072,6144,torch.float8_e4m3fn,flydsl,1117,0,14.2924,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x4x4x2_default,507.11,1485.66,0.0 +gfx950,256,224,3072,6144,torch.float8_e4m3fn,flydsl,72,0,14.8527,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x0x4x2_default,569.31,1456.09,0.0 +gfx950,256,144,3584,512,torch.float8_e4m3fn,flydsl,1012,0,3.7819,flydsl_bpreshuflle_16x128x512_F8_F8_B16_0x4x4x2_default,139.74,777.63,0.0 +gfx950,256,160,3584,512,torch.float8_e4m3fn,flydsl,860,0,4.1114,flydsl_bpreshuflle_32x128x256_F8_F8_B16_1x3x0x2_default,142.82,745.2,0.0 +gfx950,256,176,3584,512,torch.float8_e4m3fn,flydsl,2165,0,4.3114,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x4x4x1_default,149.82,739.13,0.0 +gfx950,256,192,3584,512,torch.float8_e4m3fn,flydsl,758,0,3.9518,flydsl_bpreshuflle_64x64x256_F8_F8_B16_0x3x0x2_default,178.31,837.48,0.0 +gfx950,256,208,3584,512,torch.float8_e4m3fn,flydsl,1009,0,4.5154,flydsl_bpreshuflle_16x64x256_F8_F8_B16_0x4x4x2_default,169.06,760.16,0.0 +gfx950,256,224,3584,512,torch.float8_e4m3fn,flydsl,443,0,4.3384,flydsl_bpreshuflle_32x64x256_F8_F8_B16_1x1x4x2_default,189.49,819.5,0.0 +gfx950,256,240,3584,512,torch.float8_e4m3fn,flydsl,29,0,4.3525,flydsl_bpreshuflle_80x64x256_F8_F8_B16_0x0x0x2_default,202.37,845.08,0.0 +gfx950,256,144,4096,2048,torch.float8_e4m3fn,flydsl,1075,0,7.6923,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x4x0x2_default,314.07,1282.21,0.0 +gfx950,256,160,4096,2048,torch.float8_e4m3fn,flydsl,686,0,7.2955,flydsl_bpreshuflle_32x64x512_F8_F8_B16_1x2x4x2_default,367.95,1374.41,0.0 +gfx950,256,192,4096,2048,torch.float8_e4m3fn,flydsl,566,0,7.4497,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x2x4x2_default,432.4,1389.95,0.0 +gfx950,256,224,4096,2048,torch.float8_e4m3fn,flydsl,802,0,8.0847,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x4x2_default,464.84,1321.31,0.0 +gfx950,256,144,6144,4096,torch.float8_e4m3fn,flydsl,451,0,13.1698,flydsl_bpreshuflle_48x64x256_F8_F8_B16_1x1x4x2_default,550.33,2090.02,0.0 +gfx950,256,160,6144,4096,torch.float8_e4m3fn,flydsl,802,0,13.7254,flydsl_bpreshuflle_32x64x512_F8_F8_B16_0x3x4x2_default,586.73,2024.51,0.0 +gfx950,256,192,6144,4096,torch.float8_e4m3fn,flydsl,452,0,14.0465,flydsl_bpreshuflle_48x128x256_F8_F8_B16_1x1x4x2_default,687.98,2015.56,0.0 +gfx950,256,144,7168,512,torch.float8_e4m3fn,flydsl,514,0,4.7557,flydsl_bpreshuflle_48x128x256_F8_F8_B16_0x2x0x2_default,222.25,1221.3,0.0 +gfx950,256,160,7168,512,torch.float8_e4m3fn,flydsl,29,0,4.7538,flydsl_bpreshuflle_80x64x256_F8_F8_B16_0x0x0x2_default,247.05,1271.76,0.0 +gfx950,256,192,7168,512,torch.float8_e4m3fn,flydsl,1438,0,4.7498,flydsl_bpreshuflle_96x64x256_F8_F8_B16_0x1x0x1_default,296.7,1372.86,0.0 +gfx950,256,224,7168,512,torch.float8_e4m3fn,flydsl,289,0,5.1789,flydsl_bpreshuflle_112x64x256_F8_F8_B16_0x1x0x2_default,317.47,1350.86,0.0 +gfx950,256,240,7168,512,torch.float8_e4m3fn,flydsl,586,0,5.4215,flydsl_bpreshuflle_80x128x256_F8_F8_B16_0x2x4x2_default,324.93,1334.23,0.0 +""" + + +def main() -> int: + spec = importlib.util.find_spec("aiter") + if spec is None or not spec.submodule_search_locations: + print(f"[{TAG}] FAILED: package 'aiter' not found") + return 1 + path = os.path.join(list(spec.submodule_search_locations)[0], REL) + if not os.path.isdir(os.path.dirname(path)): + print(f"[{TAG}] FAILED: {os.path.dirname(path)} not found") + return 1 + if os.path.exists(path): + with open(path, encoding="utf-8") as f: + same = f.read() == CSV + print(f"[{TAG}] {'already applied' if same else 'FAILED: a different ' + path + ' exists'}") + return 0 if same else 1 + with open(path, "w", encoding="utf-8") as f: + f.write(CSV) + print(f"[{TAG}] wrote {path}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 99cd422217b66516738e5792b17ba38170475fae Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 13:16:13 +0000 Subject: [PATCH 07/20] perf(router): ATOM PD decode reuses the prefill's prompt token ids Both ATOM legs rendered the chat template and tokenized the same prompt (~0.25 s each at 130K tokens). The prefill leg now asks for return_token_ids and the router passes them in the decode's kv_transfer_params['prompt_token_ids'], which ATOM takes as the prompt (and strips before the engine). A prefill that returns no ids leaves decode tokenizing as before. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../router/disagg_protocols/atom_mooncake.py | 14 +++++++++++++- .../unit/router/test_atom_mooncake_protocol.py | 18 ++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/infera/router/disagg_protocols/atom_mooncake.py b/infera/router/disagg_protocols/atom_mooncake.py index 34cc20865..f3aed100c 100644 --- a/infera/router/disagg_protocols/atom_mooncake.py +++ b/infera/router/disagg_protocols/atom_mooncake.py @@ -45,6 +45,7 @@ from typing import Literal from infera.common.worker_pool import WorkerInfo +from infera.router.token_ids import prompt_from_chunk def _params(w: WorkerInfo, side: str) -> dict: @@ -79,6 +80,9 @@ def annotate_prefill(self, body: dict, p: WorkerInfo, d: WorkerInfo, room_id: in if "max_completion_tokens" in out: out["max_completion_tokens"] = 1 out.pop("stream_options", None) + # Ask for the prompt ids, which extract_handoff passes on so decode + # skips rendering and tokenizing the same prompt again. + out["return_token_ids"] = True out["kv_transfer_params"] = { "remote_dp_size": _required(d_params, "dp_size", d, "decode"), "remote_tp_size": _required(d_params, "tp_size", d, "decode"), @@ -154,7 +158,15 @@ def extract_handoff(self, prefill_response_payload: dict) -> dict: # the decode leg reuses it verbatim (+ enrichment), exactly like ATOM's # own proxy / atomesh mesh. A shallow copy keeps the router from # mutating the parsed response. - return dict(kv) + handoff = dict(kv) + # ATOM decode takes kv_transfer_params['prompt_token_ids'] as the + # prompt (and drops it before the engine), skipping its own template + # rendering and tokenization: ~0.25 s per request at 130K tokens. A + # prefill that did not return them leaves decode tokenizing as before. + ids = prompt_from_chunk(prefill_response_payload) + if ids is not None: + handoff["prompt_token_ids"] = ids + return handoff def request_id_for(self, p: WorkerInfo, d: WorkerInfo, room_id: int) -> str | None: # ATOM correlates P/D via transfer_id in the body, not the request id. diff --git a/tests/unit/router/test_atom_mooncake_protocol.py b/tests/unit/router/test_atom_mooncake_protocol.py index 91b535203..2e5febab9 100644 --- a/tests/unit/router/test_atom_mooncake_protocol.py +++ b/tests/unit/router/test_atom_mooncake_protocol.py @@ -94,6 +94,24 @@ def test_extract_handoff_returns_full_kv(self): assert handoff == kv assert handoff is not kv + def test_prompt_ids_flow_from_prefill_to_decode(self): + body = {"model": "x", "messages": [], "max_tokens": 50, "stream": True} + pre = self.proto.annotate_prefill(body, self.p, self.d, room_id=1) + assert pre["return_token_ids"] is True + assert "return_token_ids" not in body + # ATOM puts them at the top level of a chat response, under the + # choice on completions; both reach decode's kv_transfer_params. + for payload in ( + {"kv_transfer_params": self._prefill_response_kv(), "prompt_token_ids": [5, 6, 7]}, + { + "kv_transfer_params": self._prefill_response_kv(), + "choices": [{"text": "", "prompt_token_ids": [5, 6, 7]}], + }, + ): + handoff = self.proto.extract_handoff(payload) + out = self.proto.annotate_decode(body, self.p, self.d, room_id=1, prefill_handoff=handoff) + assert out["kv_transfer_params"]["prompt_token_ids"] == [5, 6, 7] + def test_extract_handoff_missing_kv_raises(self): with pytest.raises(ValueError, match="kv_transfer_params"): self.proto.extract_handoff({"choices": []}) From 195f08e24d74c24d479808b42acf7a120ba38b72 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Mon, 19 Oct 2026 14:52:58 +0000 Subject: [PATCH 08/20] feat(examples): MTP layer experts in MXFP4 by default Validated on the 4P4D kit: ~1 ms less per decode step, GSM8K 0.990 and real MTP acceptance 0.780 (BF16: 0.975-0.985, 0.761). MTP_EXPERTS=bf16 keeps the checkpoint's BF16 experts. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/README.md | 2 +- examples/atom_4p4d_glm5.2/env.sh | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index 3b09d5c9c..529e73a8c 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -74,7 +74,7 @@ settings actually started); the other scripts read it and write `smoke/`, | max-num-batched-tokens | 8192 | 16384 | | mem util | 0.85 | 0.90 | | KV connector | `multi`: Mooncake `kv_producer` (4 send workers per stage) + `lmcache_offload` | Mooncake `kv_consumer` | -| shared | FP8 KV, block 16, prefix caching, MTP (`--method mtp`), online PTPC-FP8 with MXFP4 experts and BF16 MTP layer | | +| shared | FP8 KV, block 16, prefix caching, MTP (`--method mtp`), online PTPC-FP8 with MXFP4 experts (the MTP layer's quantized online, `MTP_EXPERTS`) and the rest of the MTP layer in BF16 | | The prefill settings are those of ATOM's GLM-5.2 AgentX recipe (`/app/ATOM/recipes/Agentic-GLM-5.2.md` in the image), apart from the send diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index 665e9c093..60035d6b8 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -68,9 +68,10 @@ export MTP_K="${MTP_K:-$((CONC >= 48 ? 3 : 4))}" export MTP_AL="${MTP_AL:-}" # Routed experts of the MTP layer (78), which the checkpoint keeps in BF16: -# bf16 as shipped, or mxfp4 to quantize them online at load like the experts -# of the other layers (the draft MoE then reads a quarter of the weight bytes). -export MTP_EXPERTS="${MTP_EXPERTS:-bf16}" +# mxfp4 quantizes them online at load like the experts of the other layers, so +# the draft MoE reads a quarter of the weight bytes (~1 ms per decode step; +# GSM8K and real MTP acceptance unchanged); bf16 keeps them as shipped. +export MTP_EXPERTS="${MTP_EXPERTS:-mxfp4}" # The prefill memory, chunk, block size, layer split and LMCache load length # below follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md From bf595a3ca1bed53f0a889b77e9104a73df6c116c Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Tue, 20 Oct 2026 06:34:02 +0000 Subject: [PATCH 09/20] perf(atom): DCP decode merge as one aiter kernel (all-to-all + LSE combine + FP8 quant) With the quantized combine (PBM + per-token FP8 o_proj input), the DCP decode merge ran pack -> custom all-to-all -> unpack/combine/quant. A new aiter kernel reads this rank's heads straight from the peers' registered buffers, merges them with the LSE weights and writes the FP8 rows and scales, so the received copy and the unpack launch go away. The pack writes a split layout (data, then fp32 LSE) for 16-byte loads. Matches the RCCL path to FP8 rounding (4.6e-6 of codes differ); merge microbenchmark at 160 rows 44.7 -> 38.1 us per layer; decode step -0.2..-0.6 ms (B16-64); GSM8K 0.980. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .../patch_aiter_custom_alltoall_combine.py | 323 ++++++++++++++++++ .../patch_atom_dcp_custom_a2a_combine.py | 124 +++++++ 2 files changed, 447 insertions(+) create mode 100644 deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py create mode 100644 deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py new file mode 100644 index 000000000..d7dfc1286 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py @@ -0,0 +1,323 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""The DCP attention merge as one kernel: all-to-all, LSE combine, FP8 quant. + +With patch_aiter_custom_alltoall.py, ATOM's DCP decode merge per layer is +pack -> custom all-to-all -> unpack + LSE combine + per-token FP8 quant (three +launches, ~6 + 31 + 8.5 us at graph size 40, 81 per GLM-5.2 step). The +receiving rank needs only its own heads of every rank's partials, so the +combine can read them straight from the peers' registered buffers, like the +all-to-all does, and write the quantized o_proj input directly: the received +copy and the unpack launch go away. + +- ``alltoall_combine_quant``: rank r reads chunk r of every rank's input + (data [N][rows][h_local][head_dim], then fp32 LSE [N][rows][h_local]), + merges each (row, head) over the N partials with the LSE weights, and stores + per-token FP8 rows plus one fp32 scale per row. The arithmetic is that of + ATOM's _dcp_a2a_unpack_combine_quant_kernel (non-finite LSE -> -inf, NaN + partials of empty ranks dropped, scale = row max * (1 / FP8 max)). +- ``dispatchAllGather`` runs it for gather_dim = -2000 - h_local, so the + registered-buffer and CUDA-graph paths of all_gather_reg / unreg apply as for + the all-to-all; no new pybind symbol. +- ``CustomAllreduce.supports_a2a_combine_quant`` / ``custom_a2a_combine_quant``: + the Python entry; the caller packs the split layout only when supported. + +Applies on top of patch_aiter_custom_alltoall.py and rebuilds +module_custom_all_reduce as that one does. Written against aiter 9205e1059 in +rocm/atom-dev:nightly_202610131505. Strict: see anchor_patch.py. +""" + +import importlib.util +import os +import shutil +import sys + +from anchor_patch import apply + +TAG = "aiter-patch custom-alltoall-combine" + +KERNEL = r""" +// infera: DCP merge in one kernel. Input of every rank: data +// [ngpus][rows][h_local][head_dim] (T), then fp32 LSE [ngpus][rows][h_local]. +// Rank r merges chunk r of all ranks: out = fp8 [rows][h_local * head_dim], +// then fp32 scale [rows]. Needs head_dim % pack == 0 and h_local * head_dim +// <= 512 * pack (one pack per thread and row). +constexpr int kA2ACombineDim = -2000; // gather_dim = kA2ACombineDim - h_local + +template +__global__ void __launch_bounds__(512, 1) alltoall_combine_quant(RankData* _dp, + RankSignals sg, + Signal* self_sg, + uint8_t* __restrict__ out, + int rank, + int rows, + int h_local, + int head_dim) +{ + constexpr int pack = 16 / sizeof(T); + constexpr int warps = 512 / 64; + using P = typename opus::vector_t; + __shared__ float warp_max[warps]; + const int row_len = h_local * head_dim; + const int64_t data_elems = (int64_t)ngpus * rows * row_len; + const float fp8_max = opus::cast(opus::numeric_limits::max()); + const int e0 = threadIdx.x * pack; + const bool active = e0 < row_len; + const int h = e0 / head_dim; + const T* src[ngpus]; +#pragma unroll + for(int i = 0; i < ngpus; ++i) + { + src[i] = (const T*)_dp->ptrs[i]; + } + start_sync(sg, self_sg, rank); + + for(int b = blockIdx.x; b < rows; b += gridDim.x) + { + const int64_t chunk_row = (int64_t)rank * rows + b; + float acc[pack]; +#pragma unroll + for(int j = 0; j < pack; ++j) + { + acc[j] = 0.f; + } + float amax = 0.f; + if(active) + { + P v[ngpus]; + float lse[ngpus]; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + v[s] = *reinterpret_cast(src[s] + chunk_row * row_len + e0); + lse[s] = reinterpret_cast(src[s] + data_elems)[chunk_row * h_local + h]; + } + // A rank without KV for the row reports lse -inf (or NaN) and o NaN. + float lse_max = -INFINITY; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + if(isnan(lse[s]) || lse[s] == INFINITY) + lse[s] = -INFINITY; + lse_max = fmaxf(lse_max, lse[s]); + } + if(lse_max == -INFINITY) + lse_max = 0.f; + float sum = 0.f; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + sum += expf(lse[s] - lse_max); + } + const float global_lse = logf(sum) + lse_max; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + const float factor = expf(lse[s] - global_lse); + if(isnan(factor) || factor == 0.f) + continue; // drops the NaN partials of empty ranks +#pragma unroll + for(int j = 0; j < pack; ++j) + { + acc[j] += upcast_s(v[s][j]) * factor; + } + } +#pragma unroll + for(int j = 0; j < pack; ++j) + { + amax = fmaxf(amax, fabsf(acc[j])); + } + } +#pragma unroll + for(int lane_mask = 32; lane_mask > 0; lane_mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor(amax, lane_mask, 64)); + } + if(threadIdx.x % 64 == 0) + warp_max[threadIdx.x / 64] = amax; + __syncthreads(); + float row_max = warp_max[0]; +#pragma unroll + for(int w = 1; w < warps; ++w) + { + row_max = fmaxf(row_max, warp_max[w]); + } + __syncthreads(); + // As the Triton combine: multiply by 1 / FP8 max, zero scale -> zero row. + const float scale = row_max * (1.0f / fp8_max); + const float inv = scale > 0.f ? 1.0f / scale : 0.f; + if(active) + { + opus::vector_t q; +#pragma unroll + for(int i = 0; i < pack / 4; ++i) + { + opus::fp32x4_t tmp; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + tmp[j] = fminf(fmaxf(acc[i * 4 + j] * inv, -fp8_max), fp8_max); + } + *(reinterpret_cast(&q) + i) = opus::cast(tmp); + } + *reinterpret_cast*>(out + (int64_t)b * row_len + e0) = q; + } + if(threadIdx.x == 0) + reinterpret_cast(out + (int64_t)rows * row_len)[b] = scale; + } + end_sync(sg, self_sg, rank); +} + +""" + +DISPATCH = """ // infera: DCP merge (all-to-all + LSE combine + FP8 quant), see alltoall_combine_quant. + if(gather_dim <= kA2ACombineDim) + { + int h_local = kA2ACombineDim - gather_dim; + int head_dim = last_dim_size - 2; // input viewed as [ngpus * rows * h_local, head_dim + 2] + int rows = size / (world_size_ * h_local * last_dim_size); + dim3 grid(std::min(rows, kMaxBlocks)); + auto* out = reinterpret_cast(output); + switch(world_size_) + { + case 8: + alltoall_combine_quant + <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); + break; + case 4: + alltoall_combine_quant + <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); + break; + case 2: + alltoall_combine_quant + <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); + break; + default: printf("alltoall_combine_quant world_size error\\n"); + } + return; + } +""" + +PY_METHOD = ''' # infera: dispatchAllGather runs the DCP merge kernel for this "dim" - h_local. + _A2A_COMBINE_DIM = -2000 + + def supports_a2a_combine_quant( + self, send: torch.Tensor, h_local: int, head_dim: int, quant_dtype + ) -> bool: + """Whether custom_a2a_combine_quant runs for this split-layout input + (see the method); otherwise the caller keeps its RCCL / all-to-all path.""" + return ( + self._ops_all_gather is None + and send.dtype in (torch.float16, torch.bfloat16) + and quant_dtype == fp8 + and head_dim % 8 == 0 + and 0 < h_local * head_dim <= 512 * 8 + and send.dim() == 2 + and send.shape[1] == head_dim + 2 + and send.shape[0] % (self.world_size * h_local) == 0 + and self.should_custom_ag(send) + ) + + def custom_a2a_combine_quant(self, send: torch.Tensor, h_local: int, head_dim: int): + """DCP attention merge of ``send`` = [world_size * rows * h_local, + head_dim + 2], holding data [world_size][rows][h_local][head_dim] and + then the fp32 LSE [world_size][rows][h_local] (chunk i is for rank i). + Returns (fp8 [rows, h_local, head_dim], fp32 scale [rows, 1]): the + LSE-weighted merge of this rank's chunks, quantized per row. Buffer and + capture rules as custom_all_gather; check supports_a2a_combine_quant first.""" + rows = send.shape[0] // (self.world_size * h_local) + row_bytes = h_local * head_dim + out = torch.empty(rows * (row_bytes + 4), dtype=torch.uint8, device=send.device) + dim = self._A2A_COMBINE_DIM - h_local + capturing = self._IS_CAPTURING and torch.cuda.is_current_stream_capturing() + if capturing and self.enable_register_for_capturing: + ops.all_gather_reg(self._ptr, send, out, dim) + else: + ops.all_gather_unreg( + self._ptr, + send, + self._pool["input"].data_ptr, + out, + self._pool["input"].max_size, + dim, + ) + q = out[: rows * row_bytes].view(fp8).view(rows, h_local, head_dim) + scale = out[rows * row_bytes :].view(torch.float32).view(rows, 1) + return q, scale + + def fused_ar_rms( +''' + +PATCH_CSRC = { + "../csrc/include/custom_all_reduce.cuh": ( + "// infera: DCP merge in one kernel.", + [ + ( + "// ========== reduce_scatter kernel start ==========\n", + KERNEL + "// ========== reduce_scatter kernel start ==========\n", + ), + ( + " // infera: all-to-all, see alltoall_vec.\n", + DISPATCH + " // infera: all-to-all, see alltoall_vec.\n", + ), + ], + ), + "dist/device_communicators/custom_all_reduce.py": ( + "# infera: dispatchAllGather runs the DCP merge kernel", + [(" def fused_ar_rms(\n", PY_METHOD)], + ), +} + +MODULE = "module_custom_all_reduce" + + +def rebuild() -> int: + """Rebuild the module from the patched sources, as the runtime JIT would.""" + root = importlib.util.find_spec("aiter").submodule_search_locations[0] + os.environ.setdefault("GPU_ARCHS", os.environ.get("PYTORCH_ROCM_ARCH", "gfx950")) + sys.path.insert(0, root) + from jit import core # noqa: E402 (import after GPU_ARCHS is set) + + so = f"{core.get_user_jit_dir()}/{MODULE}.so" + cuh = os.path.join(root, "../csrc/include/custom_all_reduce.cuh") + if os.path.exists(so): + if os.path.getmtime(so) > os.path.getmtime(cuh): + print(f"[{TAG}] {so} is newer than the patched header; not rebuilt") + return 0 + os.remove(so) + # The JIT compiles from a copy of the sources in its build dir; an image + # that already built the module would recompile the old copy. + shutil.rmtree(f"{core.get_user_jit_dir()}/build/{MODULE}", ignore_errors=True) + args = core.get_args_of_build(MODULE) + core.build_module( + MODULE, + args["srcs"], + args["flags_extra_cc"], + args["flags_extra_hip"], + args["blob_gen_cmd"], + args["extra_include"], + args["extra_ldflags"], + args["verbose"], + args["is_python_module"], + args["is_standalone"], + args["torch_exclude"], + args.get("third_party", []), + args.get("hipify", False), + flags_extra_hip_per_source=args.get("flags_extra_hip_per_source", {}), + ) + if not os.path.isfile(so): + print(f"[{TAG}] FAILED: AITER build left no {so}") + return 1 + print(f"[{TAG}] rebuilt {so}") + return 0 + + +if __name__ == "__main__": + rc = apply(TAG, PATCH_CSRC, package="aiter") + if rc == 0 and "--no-build" not in sys.argv: + rc = rebuild() + sys.exit(rc) diff --git a/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py new file mode 100644 index 000000000..2a4ea504b --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""DCP decode merge with the quantized combine in aiter's all-to-all kernel. + +When the DCP merge also quantizes for o_proj (PBM + per-token FP8, the decode +of the GLM-5.2 kit), cp_lse_a2a runs pack -> all-to-all -> unpack + combine + +quant. With patch_aiter_custom_alltoall_combine.py the last two are one aiter +kernel that reads each rank's partials straight from the peers: the pack writes +the split layout that kernel reads (data, then the fp32 LSE), and the unpack +launch and the received copy go away. Where the aiter kernel does not apply +(no custom communicator, other dtypes or sizes) the path is unchanged. + +Written against ATOM 661ef766 in rocm/atom-dev:nightly_202610131505, on top of +patch_atom_dcp_custom_a2a.py. Strict: see anchor_patch.py. +""" + +import sys + +from anchor_patch import apply + +HELPERS = '''# infera: one-kernel DCP merge (see _infera_dcp_merge_quant). +@triton.jit +def _infera_dcp_pack_split_kernel( + out_ptr, # [B, H, D] + lse_ptr, # [B, H] fp32 + owned_counts_ptr, # [B], or unused + send_ptr, # [N, B, H_LOCAL, D] + lse_send_ptr, # [N, B, H_LOCAL] fp32 + out_stride_b, + out_stride_h, + lse_stride_b, + lse_stride_h, + B, + H_LOCAL: tl.constexpr, + HEAD_DIM: tl.constexpr, + HAS_OWNED_COUNTS: tl.constexpr, +): + """_dcp_a2a_pack_kernel for the split layout of aiter's merge kernel: the + same records, with the LSE in its own fp32 array after the data.""" + b = tl.program_id(axis=0).to(tl.int64) + h = tl.program_id(axis=1).to(tl.int64) + row = ((h // H_LOCAL) * B + b) * H_LOCAL + h % H_LOCAL + d = tl.arange(0, HEAD_DIM) + src = tl.load(out_ptr + b * out_stride_b + h * out_stride_h + d) + tl.store(send_ptr + row * HEAD_DIM + d, src) + lse = tl.load(lse_ptr + b * lse_stride_b + h * lse_stride_h) + if HAS_OWNED_COUNTS: + lse = tl.where(tl.load(owned_counts_ptr + b) == 0, float("-inf"), lse) + tl.store(lse_send_ptr + row, lse) + + +def _infera_dcp_merge_quant( + cp_group, cp_attn_out, cp_attn_lse, owned_counts, has_owned_counts, h_local, quant_dtype +): + """Pack + aiter's one-kernel merge (all-to-all, LSE combine, per-token + quant) -> (out [B, H_local, D], scale [B, 1]); None where it does not apply.""" + device_comm = getattr(cp_group, "device_communicator", None) + ca = getattr(device_comm, "ca_comm", None) + if ca is None or ca.disabled or not hasattr(ca, "supports_a2a_combine_quant"): + return None + n_ranks = cp_group.world_size + b, h_total, head_dim = cp_attn_out.shape + send = torch.empty( + (n_ranks * b * h_local, head_dim + 2), + dtype=cp_attn_out.dtype, + device=cp_attn_out.device, + ) + if not ca.supports_a2a_combine_quant(send, h_local, head_dim, quant_dtype): + return None + flat = send.view(-1) + _infera_dcp_pack_split_kernel[(b, h_total)]( + cp_attn_out, + cp_attn_lse, + owned_counts, + flat, + flat[n_ranks * b * h_local * head_dim :].view(torch.float32), + cp_attn_out.stride(0), + cp_attn_out.stride(1), + cp_attn_lse.stride(0), + cp_attn_lse.stride(1), + b, + H_LOCAL=h_local, + HEAD_DIM=head_dim, + HAS_OWNED_COUNTS=has_owned_counts, + ) + return ca.custom_a2a_combine_quant(send, h_local, head_dim) + + +def cp_lse_a2a( +''' + +PATCH = { + "model_ops/dcp_ops.py": ( + "# infera: one-kernel DCP merge", + [ + ("def cp_lse_a2a(\n", HELPERS), + ( + " send = torch.empty((n_ranks, b, h_local, head_dim + pack), dtype=dtype, device=dev)\n", + " # infera: one-kernel DCP merge when the combine also quantizes.\n" + " if quant_dtype is not None:\n" + " merged = _infera_dcp_merge_quant(\n" + " cp_group,\n" + " cp_attn_out,\n" + " cp_attn_lse,\n" + " owned_counts,\n" + " has_owned_counts,\n" + " h_local,\n" + " quant_dtype,\n" + " )\n" + " if merged is not None:\n" + " return merged\n" + "\n" + " send = torch.empty((n_ranks, b, h_local, head_dim + pack), dtype=dtype, device=dev)\n", + ), + ], + ), +} + +if __name__ == "__main__": + sys.exit(apply("atom-patch dcp-custom-a2a-combine", PATCH, package="atom")) From 0dd2edb6a036a5ead25c5f2c4595462da56aaf7d Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Tue, 20 Oct 2026 06:34:02 +0000 Subject: [PATCH 10/20] feat(examples): DECODE_EXTRA_ENV passes extra environment to the decode container Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/engine/common.sh | 2 +- examples/atom_4p4d_glm5.2/engine/worker.sh | 2 ++ examples/atom_4p4d_glm5.2/env.sh | 2 ++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/examples/atom_4p4d_glm5.2/engine/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index 1e3a34a91..e7d4634d3 100755 --- a/examples/atom_4p4d_glm5.2/engine/common.sh +++ b/examples/atom_4p4d_glm5.2/engine/common.sh @@ -37,7 +37,7 @@ for kit_role in $KIT_ROLES; do CONTAINERS+=("$CTR-$kit_role"); done RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS DECODE_GRAPH_STEP PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL PREFILL_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS - DECODE_EXTRA_ARGS ROUTER_EXTRA_ARGS IMAGE MODEL) + DECODE_EXTRA_ARGS DECODE_EXTRA_ENV ROUTER_EXTRA_ARGS IMAGE MODEL) # resolve_model: check MODEL exists here and bind its resolved path; a symlink # across an NFS mount boundary binds an empty directory into the container. diff --git a/examples/atom_4p4d_glm5.2/engine/worker.sh b/examples/atom_4p4d_glm5.2/engine/worker.sh index 11d80b44b..5ba00960e 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -105,6 +105,8 @@ else --max-num-batched-tokens 16384 --max-num-seqs $((2 * CONC))) [[ -z "$MTP_AL" ]] || args+=(--spec-decode-acceptance-length "$MTP_AL") read -ra extra <<<"$DECODE_EXTRA_ARGS" + read -ra extra_env <<<"$DECODE_EXTRA_ENV" + env+=("${extra_env[@]}") fi env_args=() diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index 60035d6b8..926759fc1 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -120,6 +120,8 @@ export ROUTER_BACKEND="${ROUTER_BACKEND:-python}" # Extra flags, space-separated, appended last. export PREFILL_EXTRA_ARGS="${PREFILL_EXTRA_ARGS:-}" export DECODE_EXTRA_ARGS="${DECODE_EXTRA_ARGS:-}" +# Extra decode container environment, space-separated KEY=VALUE pairs. +export DECODE_EXTRA_ENV="${DECODE_EXTRA_ENV:-}" export ROUTER_EXTRA_ARGS="${ROUTER_EXTRA_ARGS:-}" # --------------------------------------------------------------------------- From ed8d0749daf058192ccfaa11d7071936976245ae Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Tue, 20 Oct 2026 09:30:43 +0000 Subject: [PATCH 11/20] perf(rust): ATOM PD decode reuses the prefill's prompt ids; /metrics skips ATOM The Rust router side of two Python router changes (f7e46219, 91910e57): - protocol.rs: the ATOM prefill leg asks for return_token_ids, and the hand-off passes the returned prompt ids to the decode in kv_transfer_params['prompt_token_ids'], so ATOM's decode skips rendering the chat template and tokenizing the same prompt (~0.25 s at 130K tokens). A reply without ids leaves the decode tokenizing as before. - handlers.rs / metrics.rs: the frontend /metrics scrapes only engines it reads something from (sglang, vllm). An ATOM scrape only cost its API process a /metrics render on top of the client's own scrapes. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- rust/router/src/handlers.rs | 11 ++++++-- rust/router/src/metrics.rs | 13 +++++++++ rust/router/src/protocol.rs | 53 +++++++++++++++++++++++++++++++++++-- 3 files changed, 73 insertions(+), 4 deletions(-) diff --git a/rust/router/src/handlers.rs b/rust/router/src/handlers.rs index 187755429..2a734c547 100644 --- a/rust/router/src/handlers.rs +++ b/rust/router/src/handlers.rs @@ -508,8 +508,15 @@ const ENGINE_METRICS_SCRAPE_TIMEOUT: Duration = Duration::from_secs(1); async fn scrape_engine_metrics(st: &AppState, snap: &crate::pool::Snapshot) -> String { crate::metrics::clear_engine_gauges(); // Scrape every active worker over its HTTP /metrics URL. NATS is only the - // request path; the engine still exposes Prometheus on w.url. - let workers: Vec<_> = snap.all.iter().filter(|w| w.is_active()).cloned().collect(); + // request path; the engine still exposes Prometheus on w.url. Skip engines + // nothing here is read from (ATOM): each scrape costs the engine's API + // process a render, and under a 3/s client scrape that doubled its GC load. + let workers: Vec<_> = snap + .all + .iter() + .filter(|w| w.is_active() && crate::metrics::frontend_reads(&w.engine)) + .cloned() + .collect(); if workers.is_empty() { return String::new(); } diff --git a/rust/router/src/metrics.rs b/rust/router/src/metrics.rs index 9c5cc8449..7381e8eff 100644 --- a/rust/router/src/metrics.rs +++ b/rust/router/src/metrics.rs @@ -534,6 +534,13 @@ const SGLANG_FEDERATE: &[&str] = &[ "sglang:num_used_tokens", ]; +/// Whether the frontend /metrics takes anything from `engine`'s exposition. +/// False for ATOM (and unknown engines): no gauge or federated family is known +/// for it, so a scrape only costs the engine's API process a /metrics render. +pub fn frontend_reads(engine: &str) -> bool { + matches!(engine.to_ascii_lowercase().as_str(), "sglang" | "vllm") +} + fn federate_families(engine: &str) -> &'static [&'static str] { match engine.to_ascii_lowercase().as_str() { "vllm" => VLLM_FEDERATE, @@ -923,6 +930,12 @@ mod tests { use super::*; + #[test] + fn frontend_scrapes_only_engines_it_reads() { + assert!(frontend_reads("vllm") && frontend_reads("SGLang")); + assert!(!frontend_reads("atom") && !frontend_reads("")); + } + #[test] fn federate_vllm_keeps_buckets_and_injects_labels() { let text = concat!( diff --git a/rust/router/src/protocol.rs b/rust/router/src/protocol.rs index 992ee6b9e..80523da7f 100644 --- a/rust/router/src/protocol.rs +++ b/rust/router/src/protocol.rs @@ -249,11 +249,19 @@ pub fn annotate_atom_prefill( atom_param(decode, "http_port")?.clone(), ); prefill_one_token(body, path); + // Ask for the prompt ids, which atom_handoff passes on so the decode skips + // rendering and tokenizing the same prompt again. + body.insert("return_token_ids".into(), Value::Bool(true)); body.insert("kv_transfer_params".into(), Value::Object(kv)); Ok(()) } -/// The hand-off the decode needs: the prefill response's `kv_transfer_params`. +/// The hand-off the decode needs: the prefill response's `kv_transfer_params`, +/// plus the prompt ids the prefill returned. ATOM's decode takes +/// `kv_transfer_params['prompt_token_ids']` as the prompt (and drops it before +/// the engine), skipping its own template rendering and tokenization: ~0.25 s +/// per request at 130K tokens. A prefill that returned no ids leaves the decode +/// tokenizing as before. pub fn atom_handoff(prefill_reply: &Value) -> anyhow::Result> { let kv = prefill_reply .get("kv_transfer_params") @@ -262,7 +270,26 @@ pub fn atom_handoff(prefill_reply: &Value) -> anyhow::Result> if !kv.contains_key("transfer_id") { anyhow::bail!("atom prefill kv_transfer_params has no transfer_id"); } - Ok(kv.clone()) + let mut handoff = kv.clone(); + if let Some(ids) = atom_prompt_ids(prefill_reply) { + handoff.insert("prompt_token_ids".into(), ids.clone()); + } + Ok(handoff) +} + +/// The prompt ids of a `return_token_ids` reply: top level for chat +/// completions, on the first choice for completions. Only a list of integers. +fn atom_prompt_ids(reply: &Value) -> Option<&Value> { + [ + reply.get("prompt_token_ids"), + reply.pointer("/choices/0/prompt_token_ids"), + ] + .into_iter() + .flatten() + .find(|ids| { + ids.as_array() + .is_some_and(|a| !a.is_empty() && a.iter().all(Value::is_u64)) + }) } /// Decode leg: the prefill's hand-off plus the producer's parallelism, which the @@ -504,6 +531,7 @@ mod tests { assert_eq!(body["max_tokens"], 1); assert_eq!(body["stream"], false); assert!(!body.contains_key("stream_options")); + assert_eq!(body["return_token_ids"], true); // What ATOM's own proxy sends, field for field. assert_eq!( body["kv_transfer_params"], @@ -543,6 +571,27 @@ mod tests { assert_eq!(kv["remote_tp_size"], 4); assert_eq!(kv["tp_size"], 4); assert_eq!(kv["remote_dp_rank"], 3); + assert!(!kv.as_object().unwrap().contains_key("prompt_token_ids")); + } + + #[test] + fn atom_handoff_carries_the_prompt_ids() { + let kv = json!({"transfer_id": "t"}); + // Chat completions: top level; completions: the first choice. + for reply in [ + json!({"kv_transfer_params": kv, "prompt_token_ids": [5, 6, 7]}), + json!({"kv_transfer_params": kv, "choices": [{"prompt_token_ids": [5, 6, 7]}]}), + ] { + let handoff = atom_handoff(&reply).unwrap(); + assert_eq!(handoff["prompt_token_ids"], json!([5, 6, 7])); + } + // Anything but a non-empty list of ids is left out. + for ids in [json!([]), json!(["a"]), json!([1, -2]), json!("1 2")] { + let reply = json!({"kv_transfer_params": kv, "prompt_token_ids": ids}); + assert!(!atom_handoff(&reply) + .unwrap() + .contains_key("prompt_token_ids")); + } } #[test] From 83e3a949e402ec6c9d2f2e4074dbb91ac3af3c7c Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Tue, 20 Oct 2026 09:34:48 +0000 Subject: [PATCH 12/20] feat(examples): GLM-5.2 4P4D kit on the ATOM recipe / InferenceX PR #4045 settings The kit's decode kept its own settings and its CPU tier was LMCache; this takes the rest of ATOM's GLM-5.2 AgentX recipe and InferenceX PR #4045's 1P1D: - decode: mem util 0.94, --max-num-batched-tokens 2048 (DECODE_MAX_BATCHED_TOKENS) and AITER_REUSE_IDENTICAL_COMM_GROUPS=1, ~196K KV blocks instead of ~156K, which moves the AgentX decode-prefix cliff from about C85 to C105. - MC_NUM_QP_PER_EP=1 on every Mooncake peer and ATOM's default 16 send workers (PREFILL_SEND_WORKERS empty; 4 stays the fallback of gotcha 13). - PREFILL_CPU_TIER=store (default): engine/store.sh starts a Mooncake Store master and one 768 GiB owner per NUMA node on the prefill's rails, backed by THP without root; the prefill offloads to it with one load worker. lmcache keeps the previous LMCache tier, off drops the CPU tier. - TOPK_FORCE_PATH=one and --timeout-keep-alive 900 on both legs. - ROUTER_BACKEND defaults to rust. Graph sizes, max-num-seqs and the MXFP4 MTP experts stay the kit's. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/README.md | 80 +++++++++++++++------- examples/atom_4p4d_glm5.2/engine/common.sh | 6 +- examples/atom_4p4d_glm5.2/engine/store.sh | 57 +++++++++++++++ examples/atom_4p4d_glm5.2/engine/up.sh | 8 ++- examples/atom_4p4d_glm5.2/engine/worker.sh | 25 ++++++- examples/atom_4p4d_glm5.2/env.sh | 64 +++++++++++------ 6 files changed, 189 insertions(+), 51 deletions(-) create mode 100755 examples/atom_4p4d_glm5.2/engine/store.sh diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index 529e73a8c..37729a58b 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -3,9 +3,10 @@ Runnable kit for **GLM-5.2-MXFP4** served the infera way on a single 8×MI355X (gfx950) node, PD-disaggregated: one **TP1 × PP4 prefill** on GPUs 0-3, set up as in ATOM's GLM-5.2 AgentX recipe, and one **TP4 + DCP4 decode** on GPUs 4-7, -KV moved by **Mooncake over the ionic RDMA rails**, a **1 TiB LMCache CPU tier** -behind the prefill, and the infera router (round-robin with -**`--session-affinity prefill`**) in front. `PREFILL_PP=1` runs a TP4 + DCP4 +KV moved by **Mooncake over the ionic RDMA rails**, a **1.5 TiB Mooncake Store +CPU tier** behind the prefill, and the infera Rust router (round-robin with +**`--session-affinity prefill`**) in front. Engine settings follow ATOM's GLM-5.2 +AgentX recipe and InferenceX PR #4045's 1P1D. `PREFILL_PP=1` runs a TP4 + DCP4 prefill instead ([below](#the-dcp-prefill)). One file is site-specific. Fill it in, then: @@ -51,7 +52,8 @@ docker build -f deploy/docker/Dockerfile.atom \ |---|---| | [`env.sh`](env.sh) | **the only file you edit** — node IP, weights, image, scratch dir, load shape, ports | | [`engine/worker.sh`](engine/worker.sh) | starts ONE engine container (`prefill` or `decode`); carries the tuned recipe | -| [`engine/up.sh`](engine/up.sh) | etcd → both workers → router; checks the router sees 1 prefill + 1 decode | +| [`engine/up.sh`](engine/up.sh) | etcd → Store → both workers → router; checks the router sees 1 prefill + 1 decode | +| [`engine/store.sh`](engine/store.sh) | the Mooncake Store CPU tier: a master and one owner per NUMA node | | [`engine/smoke.sh`](engine/smoke.sh), [`engine/smoke.py`](engine/smoke.py) | correctness checks through the router; need real MTP acceptance | | [`engine/eval.sh`](engine/eval.sh) | GSM8K (5-shot, 200 questions) with `lm_eval` from the image; needs real MTP acceptance | | [`engine/agentx/agentx.sh`](engine/agentx/agentx.sh), [`engine/agentx/aggregate.py`](engine/agentx/aggregate.py) | one InferenceX AgentX point, pinned at `918524ff`, with the single-node PD GPU count fixed | @@ -71,22 +73,29 @@ settings actually started); the other scripts read it and write `smoke/`, | parallelism | TP1 × PP4, layers per stage `20,20,20,18` | TP4 + DCP4, no DP-attention | | execution | `--enforce-eager` (ATOM runs PP without CUDA graphs) | CUDA graphs 1, 2, 4, 8, 12 … 2 × CONC | | max-num-seqs | 512 | 2 × CONC | -| max-num-batched-tokens | 8192 | 16384 | -| mem util | 0.85 | 0.90 | -| KV connector | `multi`: Mooncake `kv_producer` (4 send workers per stage) + `lmcache_offload` | Mooncake `kv_consumer` | +| max-num-batched-tokens | 8192 | 2048 | +| mem util | 0.85 | 0.94, communicator groups shared (`AITER_REUSE_IDENTICAL_COMM_GROUPS=1`) | +| KV connector | `multi`: Mooncake `kv_producer` (16 send workers per stage) + `mooncake_store` offload | Mooncake `kv_consumer` | | shared | FP8 KV, block 16, prefix caching, MTP (`--method mtp`), online PTPC-FP8 with MXFP4 experts (the MTP layer's quantized online, `MTP_EXPERTS`) and the rest of the MTP layer in BF16 | | -The prefill settings are those of ATOM's GLM-5.2 AgentX recipe -(`/app/ATOM/recipes/Agentic-GLM-5.2.md` in the image), apart from the send -workers (gotcha 13). Not taken from it: `OFFLOAD_PROFILE=1` (logging only), -`--timeout-keep-alive 900` (the infera disagg router does not reuse idle -connections) and the recipe's decode settings (mem util 0.94, -`--max-num-batched-tokens 2048`, `AITER_REUSE_IDENTICAL_COMM_GROUPS=1`, graphs -up to 256). - -MTP depth follows `CONC`: K3 at CONC ≥ 48, K4 below. LMCache holds 256 GiB per -prefill stage (`PREFILL_OFFLOAD_GB`), chunk 256, loads hits of any length, and -looks up all four stages because each holds the KV of its own layers. +The settings are those of ATOM's GLM-5.2 AgentX recipe +(`/app/ATOM/recipes/Agentic-GLM-5.2.md` in the image) and of InferenceX PR +#4045's 1P1D (`MC_NUM_QP_PER_EP=1` on every Mooncake peer, `TOPK_FORCE_PATH=one`, +`--timeout-keep-alive 900`, the Store). The decode's memory settings give it +~196K KV blocks instead of ~156K at 0.90 / 16384 with separate communicators: +AgentX lanes reuse their whole context every turn, and once C × mean prompt +length outgrows the decode's KV pool its prefix hits collapse, the prefill +re-sends whole prompts and TTFT explodes (about C85 before, C105 now). Not +taken from them: `OFFLOAD_PROFILE=1` (logging only), the decode's graphs (here +every 4 up to 2 × CONC, there every 8 up to 256) and BF16 MTP experts +(`MTP_EXPERTS`). + +MTP depth follows `CONC`: K3 at CONC ≥ 48, K4 below. The Store +(`PREFILL_CPU_TIER=store`, gotcha 14) is one pool of 2 × 768 GiB that the four +PP stages share, each through its own NIC; `PREFILL_CPU_TIER=lmcache` instead +holds 256 GiB of LMCache per stage (`PREFILL_OFFLOAD_GB`), chunk 256, and looks +up all four stages because each holds the KV of its own layers. Both load hits +of any length. KV moves from each PP stage (its own layers) to every decode DCP rank: the producer gathers the tokens a rank owns into a staging slot @@ -125,6 +134,9 @@ On shutdown each PP stage logs a few `TypeError` tracebacks from aiter ## Reference numbers +These are the previous defaults: Python router, decode at 0.90 / 16384 with +separate communicators, 4 send workers, LMCache. + AgentX agentic-coding replay, InferenceX `918524ff`, cold start per point, 10 warm-up requests per lane, 3600 s measured, forced MTP acceptance (`MTP_AL=2.99` with K3). Throughput is per GPU over all 8 GPUs and counts cache-hit input @@ -178,7 +190,7 @@ an operator decision, so the kit always passes it explicitly. drafts at a fixed rate whether they are right or not. Use it only for performance runs; the default (empty) is real acceptance. -**4. 1 TiB of pinned host memory, as GTT.** The four prefill stages (or DCP ranks) each pin +**4. LMCache: 1 TiB of pinned host memory, as GTT.** With `PREFILL_CPU_TIER=lmcache` the four prefill stages (or DCP ranks) each pin `PREFILL_OFFLOAD_GB` GiB. Lower it on smaller hosts, or set `0` to drop the CPU tier; the prefill then registers a plain Mooncake producer. Startup also waits a few minutes in an RCCL barrier after the pools are pinned. With the CPU tier @@ -201,9 +213,10 @@ socket can take the port before the listener binds. Mooncake handshakes use base + PP stage on the prefill and base + TP rank on the decode, so 21301-21304 and 21311-21314. -**7. Router backend and transports.** `ROUTER_BACKEND` picks the Python router -(default, used for the reference numbers) or the Rust one; both route ATOM PD -(`atom-mooncake`). ATOM speaks neither NATS requests nor NATS KV events, so +**7. Router backend and transports.** `ROUTER_BACKEND` picks the Rust router +(default) or the Python one (used for the reference numbers below); both route +ATOM PD (`atom-mooncake`) and hand the prefill's prompt token ids to the decode, +so the decode does not tokenize the prompt again. ATOM speaks neither NATS requests nor NATS KV events, so `up.sh` passes `--request-transport http --kv-event-transport zmq`. `kv-aware` would need engine KV events, so the policy is `round-robin`. @@ -250,9 +263,24 @@ CQ` and ionic `wqe is not posted` without end). Every later transfer to that rank failed, so every request did. It hit `ionic_4` and `ionic_5`, with forced and with real MTP acceptance. One send worker removed the hang, but a 443k-token prompt then missed ATOM's fixed 120 s transfer deadline (`PD_TRANSFER_DEADLINE_S`, -queueing included). Four workers (`PREFILL_SEND_WORKERS`, the default) ran C64 -and C96 without a transfer failure. The DCP prefill, whose ranks each write to -one decode rank without staging, never showed the hang. +queueing included). Four workers ran C64 and C96 without a transfer failure. +All of that was with Mooncake's default two QPs per endpoint. ATOM's offload +README reports the same symptom for concurrent reads at two QPs, and ATOM's CI +and InferenceX PR #4045 run the 16 workers with `MC_NUM_QP_PER_EP=1`, which the +kit now sets on every Mooncake peer; `PREFILL_SEND_WORKERS=4` is the fallback. +The DCP prefill, whose ranks each write to one decode rank without staging, +never showed the hang. + +**14. The Store owners need 2 MiB pages.** Each owner pins `STORE_OWNER_GB` +(768 GiB) of its NUMA node and registers it with the NICs in 64 GiB regions, +which the ionic NICs only accept when backed by huge pages. `store.sh` gets them +without root through glibc's THP malloc (`GLIBC_TUNABLES=glibc.malloc.hugetlb=1`, +THP `enabled` = `madvise` or `always`) and prints each owner's `AnonHugePages`; +a fragmented node can leave an owner on 4 KiB pages, and it then fails to mount. +Each NUMA node needs `STORE_OWNER_GB` free plus the engines' host memory; lower +it or use `PREFILL_CPU_TIER=lmcache` on smaller hosts. The prefill needs one +Store load worker (`OFFLOAD_LOAD_WORKERS=1`): more hung its pipeline at C96 in +ATOM's tests. ## What this kit exposes on the network @@ -263,6 +291,8 @@ one decode rank without staging, never showed the hang. | 19011 / 19012 | prefill / decode readiness probe | | 21301-21304 / 21311-21314 | Mooncake handshakes | | 23379 / 23380 | etcd client / peer (peer on 127.0.0.1) | +| 26051 / 26080 / 26090 | Store master RPC / metadata HTTP / metrics | +| 26052-26053 | Store owners | All bind on the host network without authentication. `--trust-remote-code` runs the checkpoint's Python; serve only weights you trust. diff --git a/examples/atom_4p4d_glm5.2/engine/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index e7d4634d3..951a6865b 100755 --- a/examples/atom_4p4d_glm5.2/engine/common.sh +++ b/examples/atom_4p4d_glm5.2/engine/common.sh @@ -24,7 +24,7 @@ ROUTER_URL="http://$NODE_IP:$ROUTER_PORT" # other nodes. PREFILL_NUM_WORKERS: prefills the router must list. ETCD_IP="${ETCD_IP:-$NODE_IP}" KIT_GPUS="${KIT_GPUS:-$PREFILL_GPUS,$DECODE_GPUS}" -KIT_ROLES="${KIT_ROLES:-agentx eval router decode prefill etcd}" +KIT_ROLES="${KIT_ROLES:-agentx eval router decode prefill store-owner0 store-owner1 store-master etcd}" REMOTE_PREFILLS="${REMOTE_PREFILLS:-}" PREFILL_NUM_WORKERS="${PREFILL_NUM_WORKERS:-1}" @@ -35,8 +35,8 @@ for kit_role in $KIT_ROLES; do CONTAINERS+=("$CTR-$kit_role"); done # Settings up.sh records per run, so later scripts report what actually runs. RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS DECODE_GRAPH_STEP PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL - PREFILL_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS - PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS + PREFILL_MAX_BATCHED_TOKENS DECODE_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS + PREFILL_CPU_TIER STORE_OWNER_GB STORE_OWNER_DEVICES PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS DECODE_EXTRA_ARGS DECODE_EXTRA_ENV ROUTER_EXTRA_ARGS IMAGE MODEL) # resolve_model: check MODEL exists here and bind its resolved path; a symlink diff --git a/examples/atom_4p4d_glm5.2/engine/store.sh b/examples/atom_4p4d_glm5.2/engine/store.sh new file mode 100755 index 000000000..569a4232d --- /dev/null +++ b/examples/atom_4p4d_glm5.2/engine/store.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# SPDX-License-Identifier: MIT +# +# Start the Mooncake Store the prefill offloads its KV to (PREFILL_CPU_TIER=store), +# laid out as InferenceX PR #4045's 1P1D: a master on NUMA node 1 and two owners, +# owner i holding STORE_OWNER_GB of NUMA node i's DRAM, both registered on the +# prefill's rails. Waits until both owners serve. Normally called by up.sh. +# +# Usage: bash engine/store.sh +KIT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +# shellcheck source=common.sh +source "$KIT/engine/common.sh" + +# Bind a process's CPUs and memory to one NUMA node; set_mempolicy needs SYS_NICE. +numa_exec=/app/ATOM/.github/scripts/atomesh/numa_exec.py +# No GPU: a Mooncake process that sees one probes every buffer through HIP. +common=(--network host --cap-add SYS_NICE -e HIP_VISIBLE_DEVICES=-1 -e CUDA_VISIBLE_DEVICES=-1 + -e MC_NUM_QP_PER_EP=1 -e "MC_GID_INDEX=$MC_GID_INDEX") + +docker run -d --name "$CTR-store-master" "${common[@]}" "$IMAGE" \ + python3 "$numa_exec" 1 mooncake_master --rpc_port="$STORE_MASTER_PORT" --rpc_thread_num=32 \ + --enable_http_metadata_server=true --http_metadata_server_host=0.0.0.0 \ + --http_metadata_server_port="$STORE_METADATA_PORT" --metrics_port="$STORE_METRICS_PORT" \ + --default_kv_lease_ttl=10000 --eviction_high_watermark_ratio=0.90 --eviction_ratio=0.05 \ + --allocation_strategy=random --memory_allocator=offset --client_ttl=10 >/dev/null +wait_health "http://$NODE_IP:$STORE_METRICS_PORT/metrics" "$CTR-store-master" + +mounts=() +if on_docker_host "test -e '$HOST_LIBIONIC'"; then + mounts+=(-v "$HOST_LIBIONIC:/host-libionic/libionic.so:ro") +fi +# glibc's THP malloc backs the segment with 2 MiB pages, which the NICs register +# in 64 GiB regions; 4 KiB pages they refuse (gotcha 14). +for i in 0 1; do + docker run -d --name "$CTR-store-owner$i" "${common[@]}" "${mounts[@]}" --device /dev/infiniband \ + --cap-add IPC_LOCK --ulimit memlock=-1:-1 --security-opt seccomp=unconfined \ + -e GLIBC_TUNABLES=glibc.malloc.hugetlb=1 -e MC_MAX_MR_SIZE=68719476736 -e "MC_TCP_BIND_ADDRESS=$NODE_IP" \ + "$IMAGE" python3 "$numa_exec" "$i" mooncake_client --host="$NODE_IP" --port=$((STORE_OWNER_PORT + i)) \ + --master_server_address="$NODE_IP:$STORE_MASTER_PORT" \ + --metadata_server="http://$NODE_IP:$STORE_METADATA_PORT/metadata" \ + --protocol=rdma --device_names="$STORE_OWNER_DEVICES" --global_segment_size=$((STORE_OWNER_GB << 30)) \ + --local_buffer_size=0 --threads=4 >/dev/null +done + +# An owner serves once its segment is faulted in and registered: minutes. +deadline=$((SECONDS + READY_TIMEOUT)) +for i in 0 1; do + c="$CTR-store-owner$i" + until docker logs "$c" 2>&1 | grep 'Starting real client service' >/dev/null; do + [[ "$(docker inspect --type container -f '{{.State.Running}}' "$c" 2>/dev/null)" == true ]] || + die "$c stopped while mounting its segment; see: docker logs $c" + ((SECONDS < deadline)) || die "timeout after ${READY_TIMEOUT}s waiting for $c" + sleep 10 + done + echo "$c: $STORE_OWNER_GB GiB on NUMA $i, $(docker exec "$c" grep AnonHugePages /proc/1/smaps_rollup | tr -s ' ')" +done diff --git a/examples/atom_4p4d_glm5.2/engine/up.sh b/examples/atom_4p4d_glm5.2/engine/up.sh index a65bac0a8..64ba313ec 100755 --- a/examples/atom_4p4d_glm5.2/engine/up.sh +++ b/examples/atom_4p4d_glm5.2/engine/up.sh @@ -3,7 +3,7 @@ # SPDX-License-Identifier: MIT # # Bring up the 4P4D deployment on this node: -# etcd -> prefill (GPUs 0-3) + decode (GPUs 4-7) -> router. +# etcd (+ Mooncake Store) -> prefill (GPUs 0-3) + decode (GPUs 4-7) -> router. # Each run gets a directory under $WORK_DIR/runs that records the settings used; # smoke.sh, eval.sh, agentx.sh and down.sh write into it. # @@ -20,6 +20,10 @@ case "$ROUTER_BACKEND" in python | rust) ;; *) die "ROUTER_BACKEND must be python or rust, not $ROUTER_BACKEND" ;; esac +case "$PREFILL_CPU_TIER" in + store | lmcache | off) ;; + *) die "PREFILL_CPU_TIER must be store, lmcache or off, not $PREFILL_CPU_TIER" ;; +esac for gpus in "$PREFILL_GPUS" "$DECODE_GPUS"; do [[ "$gpus" =~ ^[0-9]+(,[0-9]+){3}$ ]] || die "each leg needs exactly four GPUs, got $gpus" done @@ -47,6 +51,8 @@ docker run -d --name "$CTR-etcd" --network host --entrypoint /usr/local/bin/etcd --initial-advertise-peer-urls "http://127.0.0.1:$ETCD_PEER_PORT" \ --initial-cluster "default=http://127.0.0.1:$ETCD_PEER_PORT" >/dev/null wait_health "http://$NODE_IP:$ETCD_PORT/version" "$CTR-etcd" +# The prefill connects to the Store at startup. +[[ "$PREFILL_CPU_TIER" != store ]] || bash "$KIT/engine/store.sh" echo "===== 3. prefill and decode =====" bash "$KIT/engine/worker.sh" prefill diff --git a/examples/atom_4p4d_glm5.2/engine/worker.sh b/examples/atom_4p4d_glm5.2/engine/worker.sh index 5ba00960e..1a95308cf 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -32,7 +32,12 @@ env=( AITER_QUICK_REDUCE_QUANTIZATION=INT4 AITER_USE_FLYDSL_MOE_SORTING=1 AITER_LOG_LEVEL=WARNING PYTHONHASHSEED=0 ATOM_MLA_PAGE_SIZE=1 ATOM_ONLINE_QUANT_STREAMING=0 ATOM_SPARSE_INDEXER_LOGITS_BUDGET_MB=2047 ATOM_USE_TRITON_MLA=0 + # The indexer top-k always on aiter's one-block kernel, as ATOM's CI runs GLM-5.2. + TOPK_FORCE_PATH=one MC_ENABLE_DEST_DEVICE_AFFINITY=1 "MC_GID_INDEX=$MC_GID_INDEX" + # One QP per endpoint on every Mooncake peer (P, D and the Store): peers + # with different counts cannot connect, and the Store requires one. + MC_NUM_QP_PER_EP=1 "MOONCAKE_DISABLE_HIP_DMABUF=$MOONCAKE_DISABLE_HIP_DMABUF" # P and D see disjoint GPUs, which breaks the base Mooncake's same-host HIP # IPC path ("invalid device context"); send KV over the RDMA rails instead. @@ -63,7 +68,17 @@ if [[ "$role" == prefill ]]; then kv="$(kv_mooncake kv_producer "$PREFILL_HANDSHAKE_PORT" "$PREFILL_PORT")" # Append the send-worker count to that JSON object. [[ -z "$PREFILL_SEND_WORKERS" ]] || kv="${kv%\}},\"num_worker_threads\":$PREFILL_SEND_WORKERS}" - if ((PREFILL_OFFLOAD_GB > 0)); then + if [[ "$PREFILL_CPU_TIER" == store ]]; then + # Mooncake producer plus the Store of engine/store.sh, one pool shared by + # the four PP stages; each stage's client uses its GPU's NIC. + store="\"mooncake_store.protocol\":\"rdma\",\"mooncake_store.master\":\"$NODE_IP:$STORE_MASTER_PORT\"" + store+=",\"mooncake_store.metadata\":\"http://$NODE_IP:$STORE_METADATA_PORT/metadata\"" + store+=",\"mooncake_store.rdma_devices\":\"$STORE_OWNER_DEVICES\"" + kv="{\"kv_connector\":\"multi\",\"connectors\":[$kv,{\"kv_connector\":\"mooncake_store\",\"kv_role\":\"offload\",\"kv_connector_extra_config\":{$store}}]}" + # One load worker: more hung the PP prefill pipeline at C96 (ATOM #2459). + # The HBM transfer pool registers in 1 GiB regions. + env+=(OFFLOAD_LOAD_WORKERS=1 MC_MAX_MR_SIZE=1073741824 "OFFLOAD_MIN_LOAD_TOKENS=$OFFLOAD_MIN_LOAD_TOKENS") + elif [[ "$PREFILL_CPU_TIER" == lmcache ]] && ((PREFILL_OFFLOAD_GB > 0)); then # The four pools, one per TP rank or PP stage, are GTT, whose cap is # machine-wide (gotcha 4): check it now, not minutes into the cold # start. Every card reports the cap as its total; 64 GiB covers the @@ -101,14 +116,20 @@ else gpus="$DECODE_GPUS" port="$DECODE_PORT" util="$DECODE_MEM_UTIL" env+=("INFERA_READINESS_PORT=$DECODE_READINESS_PORT") kv="$(kv_mooncake kv_consumer "$DECODE_HANDSHAKE_PORT" "$DECODE_PORT")" + # DCP and EP groups reuse the TP group's communicators (same GPUs) instead + # of allocating their own buffers, which leaves that memory to the KV cache. + env+=(AITER_REUSE_IDENTICAL_COMM_GROUPS=1) args=(-tp 4 --decode-context-parallel-size 4 --cudagraph-capture-sizes "$cudagraph" - --max-num-batched-tokens 16384 --max-num-seqs $((2 * CONC))) + --max-num-batched-tokens "$DECODE_MAX_BATCHED_TOKENS" --max-num-seqs $((2 * CONC))) [[ -z "$MTP_AL" ]] || args+=(--spec-decode-acceptance-length "$MTP_AL") read -ra extra <<<"$DECODE_EXTRA_ARGS" read -ra extra_env <<<"$DECODE_EXTRA_ENV" env+=("${extra_env[@]}") fi +# The Rust router keeps idle connections pooled; ATOM's 5 s default closes them under it. +args+=(--timeout-keep-alive 900) + env_args=() for kv_pair in "${env[@]}"; do env_args+=(-e "$kv_pair"); done mounts=(-v "$MODEL:$MODEL:ro") diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index 926759fc1..a894b9abd 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -73,17 +73,21 @@ export MTP_AL="${MTP_AL:-}" # GSM8K and real MTP acceptance unchanged); bf16 keeps them as shipped. export MTP_EXPERTS="${MTP_EXPERTS:-mxfp4}" -# The prefill memory, chunk, block size, layer split and LMCache load length -# below follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md -# in the image); the decode keeps this kit's settings. +# Memory, chunk budgets, block size, layer split and offload load length below +# follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md in +# the image) and InferenceX PR #4045. The decode's 0.94 and 2048, with the +# shared communicator groups worker.sh sets, give it ~25% more KV blocks, which +# keeps the agents' prefixes on the decode up to about C100 (README "The recipe"). export PREFILL_MEM_UTIL="${PREFILL_MEM_UTIL:-0.85}" -export DECODE_MEM_UTIL="${DECODE_MEM_UTIL:-0.90}" +export DECODE_MEM_UTIL="${DECODE_MEM_UTIL:-0.94}" -# Prefill chunk budget (--max-num-batched-tokens); decode keeps 16384. ATOM's -# compile cache (/root/.cache/atom, kept in the JIT cache volumes) does not key -# on it: after a change, start with fresh volumes (another CTR), or the prefill -# fails at startup in assert_size_stride. +# Chunk budgets (--max-num-batched-tokens). ATOM's compile cache +# (/root/.cache/atom, kept in the JIT cache volumes) does not key on the +# prefill's: after a change, start with fresh volumes (another CTR), or the +# prefill fails at startup in assert_size_stride. MTP decode batches stay far +# below 2048 tokens. export PREFILL_MAX_BATCHED_TOKENS="${PREFILL_MAX_BATCHED_TOKENS:-8192}" +export DECODE_MAX_BATCHED_TOKENS="${DECODE_MAX_BATCHED_TOKENS:-2048}" # KV block size of both legs, which must match: the PP prefill lays the tokens # each decode DCP rank owns into that rank's pages. Block x DCP4 must divide @@ -98,24 +102,38 @@ export BLOCK_SIZE="${BLOCK_SIZE:-16}" export PREFILL_PP="${PREFILL_PP:-4}" export PREFILL_PP_PARTITION="${PREFILL_PP_PARTITION:-20,20,20,18}" -# Mooncake send workers per prefill rank (num_worker_threads). Set empty for -# ATOM's default, 16, which hangs an RDMA write of the PP prefill (README gotcha 13). -export PREFILL_SEND_WORKERS="${PREFILL_SEND_WORKERS-4}" - -# LMCache CPU KV tier per prefill PP stage, GiB (4 stages -> 1 TiB pinned host -# memory at the default, allocated as GTT: README gotcha 4). 0 turns it off and -# the prefill uses Mooncake only. +# Mooncake send workers per prefill rank (num_worker_threads). Empty = ATOM's +# default, 16, run with one QP per endpoint (MC_NUM_QP_PER_EP=1, worker.sh); +# with Mooncake's default two QPs 16 workers hung an RDMA write of the PP +# prefill (README gotcha 13). +export PREFILL_SEND_WORKERS="${PREFILL_SEND_WORKERS-}" + +# CPU KV tier behind the prefill: +# store Mooncake Store, as InferenceX PR #4045: one master and two owners +# of STORE_OWNER_GB each, one per NUMA node, one pool shared by the +# four PP stages, read and written over RDMA (engine/store.sh, +# README gotcha 14) +# lmcache LMCache, PREFILL_OFFLOAD_GB per PP stage, as GTT (gotcha 4) +# off Mooncake P->D only +export PREFILL_CPU_TIER="${PREFILL_CPU_TIER:-store}" +export STORE_OWNER_GB="${STORE_OWNER_GB:-768}" +# The prefill's rails in PREFILL_GPUS order: the owners register on all of +# them, and stage i's Store client uses the i-th, so each stage reaches the +# owners on its own rail (the fabric routes same-numbered rails only, gotcha +# 10). Explicit, because ATOM finds a GPU's NIC from the PCI tree, which a VM +# can flatten. +export STORE_OWNER_DEVICES="${STORE_OWNER_DEVICES:-ionic_0,ionic_1,ionic_2,ionic_3}" export PREFILL_OFFLOAD_GB="${PREFILL_OFFLOAD_GB:-256}" -# Shortest LMCache hit the prefill loads; shorter hits are recomputed. +# Shortest CPU-tier hit the prefill loads; shorter hits are recomputed. export OFFLOAD_MIN_LOAD_TOKENS="${OFFLOAD_MIN_LOAD_TOKENS:-0}" # Router --session-affinity on top of round-robin: prefill keeps an agent # session on one prefill, both also on one decode, off is stateless. export SESSION_AFFINITY="${SESSION_AFFINITY:-prefill}" -# Router data plane: python (the reference runs) or rust (the image's -# infera-router binary). Both route ATOM PD the same way. -export ROUTER_BACKEND="${ROUTER_BACKEND:-python}" +# Router data plane: rust (the image's infera-router binary) or python. Both +# route ATOM PD the same way. +export ROUTER_BACKEND="${ROUTER_BACKEND:-rust}" # Extra flags, space-separated, appended last. export PREFILL_EXTRA_ARGS="${PREFILL_EXTRA_ARGS:-}" @@ -141,9 +159,15 @@ export DECODE_READINESS_PORT="${DECODE_READINESS_PORT:-19012}" # Mooncake handshake: base + PP stage or TP rank, so 21301-21304 and 21311-21314. export PREFILL_HANDSHAKE_PORT="${PREFILL_HANDSHAKE_PORT:-21301}" export DECODE_HANDSHAKE_PORT="${DECODE_HANDSHAKE_PORT:-21311}" +# Mooncake Store: master RPC, metadata HTTP, metrics, owner i on base + i. +export STORE_MASTER_PORT="${STORE_MASTER_PORT:-26051}" +export STORE_METADATA_PORT="${STORE_METADATA_PORT:-26080}" +export STORE_METRICS_PORT="${STORE_METRICS_PORT:-26090}" +export STORE_OWNER_PORT="${STORE_OWNER_PORT:-26052}" export ETCD_IMAGE="${ETCD_IMAGE:-quay.io/coreos/etcd:v3.5.14}" -# Container name prefix: -etcd, -prefill, -decode, -router, -agentx, -eval. +# Container name prefix: -etcd, -prefill, -decode, -router, -agentx, -eval, +# -store-master, -store-owner0, -store-owner1. export CTR="${CTR:-glm52_4p4d}" # --------------------------------------------------------------------------- From 1dbe1ef11d5e5ffbb3b998001b4eb38d408f78d9 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Tue, 20 Oct 2026 10:11:52 +0000 Subject: [PATCH 13/20] docs(examples): the Store probe failure on hosts whose Store reads into HBM do not land Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/README.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index 37729a58b..04c36f473 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -278,7 +278,12 @@ without root through glibc's THP malloc (`GLIBC_TUNABLES=glibc.malloc.hugetlb=1` THP `enabled` = `madvise` or `always`) and prints each owner's `AnonHugePages`; a fragmented node can leave an owner on 4 KiB pages, and it then fails to mount. Each NUMA node needs `STORE_OWNER_GB` free plus the engines' host memory; lower -it or use `PREFILL_CPU_TIER=lmcache` on smaller hosts. The prefill needs one +it or use `PREFILL_CPU_TIER=lmcache` on smaller hosts. The prefill refuses to +start when its startup probe reads back other bytes than it wrote ("the chunk +read back differs"): on the crusoe MI355X VMs every Store read into HBM +completed without moving data (puts from HBM and transfers to host buffers were +fine, and so is the P→D path), with peer memory and DMA-BUF alike; use +`PREFILL_CPU_TIER=lmcache` there. The prefill needs one Store load worker (`OFFLOAD_LOAD_WORKERS=1`): more hung its pipeline at C96 in ATOM's tests. From d699ee05fb20d7a0e9bfd4e66ac396fa42400d5f Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Wed, 21 Oct 2026 07:37:25 +0000 Subject: [PATCH 14/20] fix(atom): review follow-ups for the GLM-5.2 ATOM PD patches and 4P4D kit Patches (deploy/docker/patches/atom_nightly/): - patch_aiter_custom_alltoall_combine.py is merged into patch_aiter_custom_alltoall.py: one set of edits and one rebuild of module_custom_all_reduce, which now always clears the JIT build dir. The combine patch anchored on text the other one inserted, so the build depended on the glob order (bash under en_US.UTF-8 sorts _combine first). The patched sources are as before, plus the next item. - Both Python entries refuse world sizes other than 2, 4 and 8, the only ones the kernels are instantiated for (should_custom_ag allows 6). - patch_aiter_glm52_decode_gemm_rows.py fails the build when a table aiter merges it with already holds one of its keys; aiter would raise at engine start instead. - check_dcp_merge.py: GPU check of both kernels against RCCL (bit-exact all-to-all, merge up to FP8 rounding, eager and from a CUDA graph), to run after a base move. - README.md and patch.upstream.status.md list the new patches with their contract, drop conditions and upstream status (none filed yet), and the per-request context gauges the GC fix turns off. - ruff format (0.13.0) on the PR's Python files; the bmm TUNED table stays compact inside fmt: off. Router: the Python ATOM hand-off leaves out an empty or negative prompt-id list, as the Rust one does; token_ids.py no longer says ATOM lacks return_token_ids. 4P4D kit: - DECODE_GRAPH_STEP is gone (no gain measured; 0 looped forever), and PREFILL_EXTRA_ENV joins DECODE_EXTRA_ENV, e.g. for the gauges. - The JIT cache volume name carries the leg's --max-num-batched-tokens, which ATOM's compile cache does not key on. - up.sh refuses a decode budget below (MTP_K + 1) x 2 x CONC; store.sh checks THP and each NUMA node's free memory before starting the owners. - README and the manual walkthrough describe the new defaults: validation and reference numbers with the Store, gotchas 8, 11, 13 and 14, the DCP prefill's validated command with PREFILL_CPU_TIER=lmcache, max-num-seqs among the settings not taken from PR #4045, no machine names. Checked: pytest tests/unit/{router,common,server,engine} 909 passed; ruff 0.13.0 format and check; shellcheck; all nightly patches applied to rocm/atom-dev:nightly_202610131505 in both glob orders (a second run is a no-op, the module rebuilds); check_dcp_merge.py passes on 4 MI300X GPUs and fails with a zero tolerance. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- deploy/docker/patch.upstream.status.md | 22 +- deploy/docker/patches/atom_nightly/README.md | 79 ++++- .../patches/atom_nightly/check_dcp_merge.py | 163 +++++++++ .../patch_aiter_custom_alltoall.py | 261 +++++++++++++- .../patch_aiter_custom_alltoall_combine.py | 323 ------------------ .../patch_aiter_glm52_decode_gemm_rows.py | 42 ++- .../patch_aiter_mla_bmm_configs.py | 6 +- .../atom_nightly/patch_atom_api_gc_stalls.py | 4 +- .../patch_atom_dcp_custom_a2a_combine.py | 11 +- examples/atom_4p4d_glm5.2/README.md | 166 +++++---- examples/atom_4p4d_glm5.2/engine/common.sh | 4 +- examples/atom_4p4d_glm5.2/engine/store.sh | 15 +- examples/atom_4p4d_glm5.2/engine/up.sh | 4 + examples/atom_4p4d_glm5.2/engine/worker.sh | 24 +- examples/atom_4p4d_glm5.2/env.sh | 23 +- .../router/disagg_protocols/atom_mooncake.py | 5 +- infera/router/token_ids.py | 4 +- manual/examples/atom_4p4d_glm5.2.md | 33 +- .../router/test_atom_mooncake_protocol.py | 8 +- 19 files changed, 719 insertions(+), 478 deletions(-) create mode 100644 deploy/docker/patches/atom_nightly/check_dcp_merge.py delete mode 100644 deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py diff --git a/deploy/docker/patch.upstream.status.md b/deploy/docker/patch.upstream.status.md index c5aba2913..0894ee97d 100644 --- a/deploy/docker/patch.upstream.status.md +++ b/deploy/docker/patch.upstream.status.md @@ -295,14 +295,24 @@ anchor is gone). ## ATOM nightly — `patches/atom_nightly/` (baked by `Dockerfile.atom`, `ATOM_NIGHTLY=1`) Strict anchor scripts against `rocm/atom-dev:nightly_202610131505` (ATOM -`661ef766`, LMCache 0.5.6.dev139+gf22dec28): drift fails the build. Rationale -is in each script's docstring and `atom_nightly/README.md`. LMCache `dev` and -its open PRs were checked on 2026-10-14. +`661ef766`, aiter `9205e1059`, LMCache 0.5.6.dev139+gf22dec28): drift fails the +build. Rationale is in each script's docstring and `atom_nightly/README.md`. +LMCache `dev` and its open PRs were checked on 2026-10-14. The ATOM and aiter +decode patches (the GLM-5.2 4P4D kit's) were written 2026-10-19/20 and none is +filed yet; aiter is public ([ROCm/aiter](https://github.com/ROCm/aiter)), so +its rows can go there. | patch | fixes | upstream issue | upstream PR | ours? | PR state | |---|---|---|---|---|---| | `atom_nightly/patch_lmcache_rocm_pinned_alloc.py` | LMCache's NUMA CPU pool is `hipHostRegister`ed (KFD USERPTR) on ROCm; a page migration (compaction, `move_pages`) evicts every GPU queue of the process → multi-second prefill stalls. With `HSA_USERPTR_FOR_PAGED_MEM=0` the pool comes from exact-size `hipHostMalloc` (GTT) instead | [LMCache#5361](https://github.com/LMCache/LMCache/issues/5361) (MP server pool) | none for the non-MP pool; [LMCache#5591](https://github.com/LMCache/LMCache/pull/5591) does the same for the MP server's L1 only; same approach as [sglang#42553](https://github.com/sgl-project/sglang/pull/42553) | no | #5591 draft | | `atom_nightly/patch_atom_dcp_sparse_pad_rows.py` | with sparse MLA + DCP the fused DCP candidate exchange writes `dcp_sparse_kv_indptr_buffer` for the scheduled rows only, so an eager decode step padded to a graph bucket (the P/D producer's q_len=1 decode with MTP on) reads stale offsets in its padded rows; a stale offset below the real end gives a negative-length row, a work with kv_end < kv_start, and the persistent ASM decode kernel reads past kv_indices → `Memory access fault` on the prefill node under long-prefix load. Padded rows now end where the last real row ends | not filed yet | none | — | — | +| `atom_nightly/patch_atom_api_gc_stalls.py` | ATOM's API process: the output thread queued every batch of finished sequences on `CoreManager.outputs_queue`, which nothing reads when serving online, and every request added a per-request context gauge series that was never removed; both kept growing with requests served, and so did every gen-2 GC pass (~45% of the wall clock under AgentX load, ~0.7 s on each TTFT). Online the batch is no longer queued, and the gauges are opt-in (`ATOM_METRICS_PER_REQUEST_CONTEXT=1`) | not filed yet | not filed yet | — | — | +| `atom_nightly/patch_atom_dcp_topk_custom_ag.py` | the sparse indexer's DCP top-k exchange used `GroupCoordinator.all_gather` (pynccl, ~43 us per gather, 24 per GLM-5.2 decode step) while ATOM's own `dcp_all_gather` picks aiter's IPC all-gather; now it does too | not filed yet | not filed yet | — | — | +| `atom_nightly/patch_aiter_custom_alltoall.py` | aiter has no IPC all-to-all: adds `alltoall_vec` and `alltoall_combine_quant` (all-to-all + LSE combine + per-token FP8 quant) to the custom all-reduce, dispatched through `all_gather_reg/unreg`, and rebuilds `module_custom_all_reduce` | not filed yet | not filed yet (ROCm/aiter) | — | — | +| `atom_nightly/patch_atom_dcp_custom_a2a.py` | the DCP attention merge's `all_to_all_single` runs on RCCL (~40 us per layer, ~3.3 ms per GLM-5.2 step); uses `custom_all_to_all` where it applies | not filed yet | not filed yet | — | — | +| `atom_nightly/patch_atom_dcp_custom_a2a_combine.py` | the quantized DCP merge (pack, all-to-all, unpack + combine + quant) as pack + one aiter kernel (`custom_a2a_combine_quant`) where it applies | not filed yet | not filed yet | — | — | +| `atom_nightly/patch_aiter_mla_bmm_configs.py` | aiter's gfx950 configs of the MLA absorbed bmm (N=512 K=192, N=256 K=512) use BLOCK_SIZE_M 64 from M 64 up, so any decode M that is not a multiple of 64 takes a masked path at 2-2.5x the time; adds per-M buckets up to 512 from a sweep on MI355X | not filed yet | not filed yet (ROCm/aiter) | — | — | +| `atom_nightly/patch_aiter_glm52_decode_gemm_rows.py` | aiter's GLM-5.2 a8w8_bpreshuffle table has no rows for the decode's M 144-240 (MTP verify at graph 36 and up), which then take a padded or default kernel; adds 27 tuned rows | not filed yet | not filed yet (ROCm/aiter) | — | — | | `atom_nightly/prebuild_aiter_mha_varlen.py` (a build step, not a source patch) | the base's AITER prebuilds five `mha_varlen_fwd` variants but not `..._mask_lse_ndropout_nskip_nqscale`, which chunked-MLA prefill calls once a prompt extends a cached prefix; its first call JIT-compiles it for ~30 s and stalls the prefill; built into the image instead. On gfx950 `nightly_202610131505` loads the ASM `fmha_v3` kernel instead; kept for gfx942 | not searched | not filed (add it to AITER `setup.py`'s prebuild list) | — | — | ## hipFile — `patches/hipfile_async/` (baked by `Dockerfile.vllm`) @@ -316,8 +326,10 @@ its open PRs were checked on 2026-10-14. Every file under `patches/` is covered above except these, which carry no fix of their own: `mooncake_cpp/apply_mooncake_cpp_patches.sh`, `sglang_dsa/README.md`, `sglang_disagg/README.md`, -`vllm-dsv4/legacy/README.md`, `atom_nightly/README.md` and -`atom_nightly/anchor_patch.py` (the strict anchor-edit helper). +`vllm-dsv4/legacy/README.md`, `atom_nightly/README.md`, +`atom_nightly/anchor_patch.py` (the strict anchor-edit helper) and +`atom_nightly/check_dcp_merge.py` (a GPU check of +`patch_aiter_custom_alltoall.py`, run by hand). ## Retired diff --git a/deploy/docker/patches/atom_nightly/README.md b/deploy/docker/patches/atom_nightly/README.md index c4ebeeea6..fd1b95a7e 100644 --- a/deploy/docker/patches/atom_nightly/README.md +++ b/deploy/docker/patches/atom_nightly/README.md @@ -5,8 +5,8 @@ an ATOM nightly rather than the default release base. With it the build skips the legacy `patches/atom/` loop (written for the release base; on the nightly its edits are duplicate or dead code), runs the scripts here, and pins `protobuf<7` (the nightly's OpenTelemetry needs it). Verified on -`rocm/atom-dev:nightly_202610131505` (ATOM `661ef766`, LMCache -`0.5.6.dev139+gf22dec28`). The build for this stack is: +`rocm/atom-dev:nightly_202610131505` (ATOM `661ef766`, aiter `9205e1059`, +LMCache `0.5.6.dev139+gf22dec28`). The build for this stack is: ```bash docker build -f deploy/docker/Dockerfile.atom \ @@ -27,7 +27,38 @@ Each script's docstring has the full rationale. |---|---|---| | `patch_lmcache_rocm_pinned_alloc.py` | `lmcache/v1/memory_management.py` | on ROCm with `HSA_USERPTR_FOR_PAGED_MEM=0`, backs LMCache's NUMA CPU pool with exact-size `hipHostMalloc` (GTT) instead of `hipHostRegister` (KFD USERPTR) | | `patch_atom_dcp_sparse_pad_rows.py` | `atom/model_ops/attention_mla.py` | in sparse DCP decode, gives the rows an eager step padded past the scheduled ones an empty KV range; their stale offsets made the decode kernel read garbage slot ids (prefill `Memory access fault`) | +| `patch_atom_api_gc_stalls.py` | `atom/model_engine/engine_core_mgr.py`, `atom/metrics/scheduler.py` | keeps the API process's heap flat, so its gen-2 GC passes stop growing with every request served: finished sequences no longer go on `CoreManager.outputs_queue` when serving online (nothing reads it), and the per-request context gauges are opt-in (`ATOM_METRICS_PER_REQUEST_CONTEXT=1`) | +| `patch_atom_dcp_topk_custom_ag.py` | `atom/model_ops/dcp_ops.py` | the sparse indexer's DCP top-k exchange goes through `dcp_all_gather` (aiter's IPC all-gather) instead of pynccl | +| `patch_aiter_custom_alltoall.py` | aiter `csrc/include/custom_all_reduce.cuh`, `dist/device_communicators/custom_all_reduce.py`; rebuilds `module_custom_all_reduce` | two kernels on aiter's custom all-reduce buffers: an all-to-all (`custom_all_to_all`) and the DCP merge in one kernel, all-to-all + LSE combine + per-token FP8 quant (`custom_a2a_combine_quant`) | +| `patch_atom_dcp_custom_a2a.py` | `atom/model_ops/dcp_ops.py` | the DCP attention merge exchanges its partials over `custom_all_to_all` where it applies, RCCL otherwise | +| `patch_atom_dcp_custom_a2a_combine.py` | `atom/model_ops/dcp_ops.py` | a DCP merge that also quantizes (the GLM-5.2 decode) packs the split layout and runs `custom_a2a_combine_quant` where it applies | +| `patch_aiter_mla_bmm_configs.py` | aiter's gfx950 Triton configs of `batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant`, N=512 K=192 and N=256 K=512 | one config bucket per decode M up to 512 for the MLA absorbed bmm; the stock buckets take a masked path at 2-2.5x the time whenever M is not a multiple of 64 | +| `patch_aiter_glm52_decode_gemm_rows.py` | new `configs/model_configs/a8w8_bpreshuffle_tuned_gemm_glm5.2_decode.csv` | 27 tuned FP8 GEMM rows for the GLM-5.2 decode's M 144-240 | | `prebuild_aiter_mha_varlen.py` | AITER JIT dir | builds an AITER kernel into the image (a build step, not a source patch) | +| `check_dcp_merge.py` | none | GPU check of the `patch_aiter_custom_alltoall.py` kernels, run by hand (below) | + +### ATOM per-request context gauges + +`patch_atom_api_gc_stalls.py` changes behavior for every model served from the +image: `atom:prefill_request_context_tokens` and +`atom:decode_request_context_tokens` are no longer recorded unless the engine +starts with `ATOM_METRICS_PER_REQUEST_CONTEXT=1`. Each request added a series +that was never removed, and every `/metrics` render of the API process re-read +all of them. Set it where those series are needed, such as profiling runs. + +### DCP merge kernel check + +The anchors only show that the source text did not drift, not that the two +aiter kernels still agree with RCCL once aiter's sync primitives or buffer +registration change. After a base move, run `check_dcp_merge.py` in the new +image on 2, 4 or 8 GPUs; it exits 1 on any mismatch: + +```bash +docker run --rm --network host --ipc host --device /dev/kfd --device /dev/dri \ + --group-add video --security-opt seccomp=unconfined -e HIP_VISIBLE_DEVICES=0,1,2,3 \ + -v "$PWD/deploy/docker/patches/atom_nightly:/check:ro" \ + torchrun --nproc-per-node 4 /check/check_dcp_merge.py +``` ### LMCache CPU pool on ROCm @@ -50,16 +81,31 @@ loads it; it is kept for other archs (gfx942). ## Contract -Patch scripts use `anchor_patch.py`. All files of a patch resolve under -one importable package root (`lmcache` or `atom`). Every anchor must occur exactly -once before any file is written; a missing or repeated anchor, or a missing -file, exits 1 and leaves the package unchanged. A file that already holds the -patch's marker is skipped, so a rerun is a no-op. Unlike `patches/atom/`, which -swallows failures, a base that drifted from the one above fails the build. +Source patches use `anchor_patch.py`. All files of a patch resolve under one +importable package root (`lmcache`, `atom` or `aiter`; aiter's C++ sources are +reached as `../csrc` from its package). Every anchor must occur exactly once +before any file is written; a missing or repeated anchor, or a missing file, +exits 1 and leaves the package unchanged. A file that already holds the +patch's marker is skipped, so a rerun is a no-op. + +The two aiter table patches edit data and check it themselves: +`patch_aiter_mla_bmm_configs.py` requires both files to hold the stock buckets, +and `patch_aiter_glm52_decode_gemm_rows.py` requires that its table is new or +identical and that none of its keys is in a table aiter merges it with (aiter +would raise at engine start). `patch_aiter_custom_alltoall.py` also rebuilds +`module_custom_all_reduce` from the patched sources, as the runtime JIT would; +that needs no GPU, and the archs come from `GPU_ARCHS` (default +`PYTORCH_ROCM_ARCH`). + +No script depends on another's edits, so the order the build globs them in +does not matter. Unlike `patches/atom/`, which swallows failures, a base that +drifted from the one above fails the build. ## Dropping them -Re-check each script whenever `ATOM_BASE_IMAGE` moves. +Re-check each script whenever `ATOM_BASE_IMAGE` moves, and run +`check_dcp_merge.py` in the new image. Upstream status: +[`patch.upstream.status.md`](../../patch.upstream.status.md). - `patch_lmcache_rocm_pinned_alloc.py`: drop once LMCache's non-MP CPU pool (`_resolve_pinned_alloc_free`) no longer uses `hipHostRegister` on ROCm. As @@ -69,5 +115,20 @@ Re-check each script whenever `ATOM_BASE_IMAGE` moves. exchange writes (or clamps) `dcp_sparse_kv_indptr_buffer` for the padded rows itself, or `_forward_decode` clamps them. Not in `661ef766`; not filed upstream yet. +- `patch_atom_api_gc_stalls.py`: drop once ATOM, serving online, no longer + queues finished sequences on `outputs_queue` and no longer keeps one gauge + series per request (removes it when the request finishes, or uses a + histogram). Not in `661ef766`; not filed upstream yet. +- `patch_atom_dcp_topk_custom_ag.py`: drop once ATOM's DCP top-k exchange uses + `dcp_all_gather` (or another custom all-gather) itself. +- `patch_aiter_custom_alltoall.py`, `patch_atom_dcp_custom_a2a.py`, + `patch_atom_dcp_custom_a2a_combine.py`: one change; drop the three together + once aiter offers an IPC all-to-all and a fused merge and ATOM's `cp_lse_a2a` + uses them. Not filed upstream yet. +- `patch_aiter_mla_bmm_configs.py`: drop once aiter's gfx950 configs for these + two shapes bucket M finely below 512 (a file with `M_BOUNDS` fails the + build); re-sweep if the stock entries change. +- `patch_aiter_glm52_decode_gemm_rows.py`: drop once aiter's own GLM-5.2 + a8w8_bpreshuffle table covers M 144-240 (a shared key fails the build). - `prebuild_aiter_mha_varlen.py`: drop once the base prebuilds that variant itself, or once ATOM no longer loads it. diff --git a/deploy/docker/patches/atom_nightly/check_dcp_merge.py b/deploy/docker/patches/atom_nightly/check_dcp_merge.py new file mode 100644 index 000000000..3d87032e3 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/check_dcp_merge.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""GPU check of the DCP merge kernels of patch_aiter_custom_alltoall.py. + +The strict anchors only prove that the patched source text did not drift. Run +this in the image whenever ATOM_BASE_IMAGE moves, on 2, 4 or 8 GPUs of a node +(see README.md): + + torchrun --nproc-per-node 4 check_dcp_merge.py + +It runs each path on aiter's custom all-reduce communicator eagerly and +replayed from a CUDA graph (the registered-buffer path of the decode), and +compares it with RCCL: + +- ``custom_all_to_all`` equals ``all_to_all_single`` bit for bit; +- ATOM's ``cp_lse_a2a`` with the per-token FP8 combine, at the GLM-5.2 decode + shape (64 heads, head_dim 256) with empty ranks (LSE -inf, NaN partials) and + rows whose owned count is 0, equals the RCCL path bit for bit over the custom + all-to-all, and up to FP8 rounding over the one-kernel merge + (``custom_a2a_combine_quant``; a few codes in a million differ). + +A case fails if it fails on any rank. Rank 0 prints one line per case, with its +own numbers; the exit code is 1 if any case failed. +""" + +import sys +from types import SimpleNamespace + +import torch +import torch.distributed as dist + +dist.init_process_group("nccl") +rank, world = dist.get_rank(), dist.get_world_size() +torch.cuda.set_device(rank) +dev = torch.device(f"cuda:{rank}") + +from aiter import dtypes # noqa: E402 +from aiter.dist.device_communicators.custom_all_reduce import CustomAllreduce # noqa: E402 +from atom.model_ops.dcp_ops import cp_lse_a2a # noqa: E402 + +ca = CustomAllreduce(group=dist.new_group(backend="gloo"), device=dev) +HEADS, HEAD_DIM = 64, 256 +# Merged rows (4 x graph size in the MTP verify): graph 40, 36, a draft, 128. +ROWS = (160, 144, 40, 512) +# One-kernel merge vs RCCL: FP8 codes that may differ, relative scale error. +MAX_CODE_DIFF, MAX_SCALE_REL = 1e-4, 1e-5 + + +def dcp_group(comm): + """The parts of ATOM's DCP group coordinator that cp_lse_a2a reads.""" + return SimpleNamespace( + world_size=world, + device_group=dist.group.WORLD, + device_communicator=SimpleNamespace(ca_comm=comm), + ) + + +PATHS = { + "rccl": dcp_group(None), + # Without supports_a2a_combine_quant: pack -> custom all-to-all -> unpack. + "a2a": dcp_group(SimpleNamespace(disabled=False, custom_all_to_all=ca.custom_all_to_all)), + "merge": dcp_group(ca), +} + + +def run(fn, graph): + """fn(), eagerly or captured in a CUDA graph and replayed once.""" + if not graph: + return fn() + g = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with ca.capture(), torch.cuda.stream(s): + fn() + with torch.cuda.graph(g, stream=s): + out = fn() + torch.cuda.current_stream().wait_stream(s) + torch.cuda.synchronize() + dist.barrier() + g.replay() + torch.cuda.synchronize() + return out + + +failed = [] + + +def report(case, ok, detail=""): + """Record a case that passed on every rank, or failed on any.""" + bad = torch.tensor([0 if ok else 1], device=dev) + dist.all_reduce(bad) + if bad.item(): + failed.append(case) + if rank == 0: + print( + f"[check_dcp_merge] {'FAIL' if bad.item() else 'ok'} {case} {detail}".rstrip(), + flush=True, + ) + + +def merge_inputs(rows): + g = torch.Generator(device=dev).manual_seed(rows * 100 + rank) + out = torch.randn(rows, HEADS, HEAD_DIM, generator=g, device=dev).to(torch.bfloat16) + lse = torch.rand(rows, HEADS, generator=g, device=dev) * 8 + # No KV on this rank for some rows: LSE -inf and NaN partials, as aiter returns. + empty = torch.rand(rows, generator=g, device=dev) < 0.05 + lse[empty] = float("-inf") + out[empty] = float("nan") + owned = torch.randint(1, 100, (rows,), generator=g, device=dev, dtype=torch.int32) + owned[torch.rand(rows, generator=g, device=dev) < 0.05] = 0 + return out, lse, owned + + +# The merge's all-to-all at graph 40 and at a draft: 5.3 / 1.3 MB per rank. +for mb in (5.28, 1.32): + n = int(mb * 1e6) // 2 // (8 * world) * (8 * world) + send = torch.randn(n, dtype=torch.bfloat16, device=dev) + ref = torch.empty_like(send) + dist.all_to_all_single(ref, send) + for graph in (False, True): + got = run(lambda send=send: ca.custom_all_to_all(send), graph) + ok = got is not None and torch.equal(got, ref) + report(f"custom_all_to_all {mb} MB/rank {'graph' if graph else 'eager'}", ok) + +h_local = HEADS // world +for rows in ROWS: + x = merge_inputs(rows) + probe = torch.empty((world * rows * h_local, HEAD_DIM + 2), dtype=torch.bfloat16, device=dev) + # Otherwise the "merge" path would fall back to RCCL and pass trivially. + report( + f"merge applies rows={rows}", + ca.supports_a2a_combine_quant(probe, h_local, HEAD_DIM, dtypes.fp8), + ) + ref_q, ref_s = cp_lse_a2a(x[0], x[1], PATHS["rccl"], owned_counts=x[2], quant_dtype=dtypes.fp8) + ref_codes = ref_q.view(torch.uint8) + for path in ("a2a", "merge"): + for graph in (False, True): + q, s = run( + lambda path=path, x=x: cp_lse_a2a( + x[0], x[1], PATHS[path], owned_counts=x[2], quant_dtype=dtypes.fp8 + ), + graph, + ) + code_diff = (q.view(torch.uint8) != ref_codes).float().mean().item() + scale_rel = ((s - ref_s).abs() / ref_s.abs().clamp_min(1e-30)).max().item() + finite = bool(torch.isfinite(q.float() * s.view(-1, 1, 1)).all()) + if path == "a2a": + ok = code_diff == 0 and torch.equal(s, ref_s) + else: + ok = finite and code_diff <= MAX_CODE_DIFF and scale_rel <= MAX_SCALE_REL + detail = f"code diff {code_diff:.2e}, scale rel {scale_rel:.2e}, finite {finite}" + report(f"cp_lse_a2a {path} rows={rows} {'graph' if graph else 'eager'}", ok, detail) + +# Free the IPC buffers now; left to interpreter exit, their __del__ fails noisily. +ca.close() +dist.destroy_process_group() +if rank == 0 and failed: + print(f"[check_dcp_merge] FAILED: {len(failed)} case(s)", flush=True) +sys.exit(1 if failed else 0) diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py index bb8a6edb1..8f0e076ad 100644 --- a/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py +++ b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py @@ -4,23 +4,40 @@ # # SPDX-License-Identifier: MIT ############################################################################### -"""An all-to-all on aiter's custom (IPC) all-reduce buffers, for the DCP merge. +"""Two kernels on aiter's custom (IPC) all-reduce buffers for the DCP merge. -ATOM's DCP decode merges the partial attention of every rank with one -all_to_all_single per layer (81 per GLM-5.2 step, 5.3 MB per rank at graph -size 40), on RCCL: ~40 us each in the decode graph, ~3.3 ms per step, and no -RCCL setting helps (NCCL_PROTO / channel counts, measured). aiter's custom +ATOM's DCP decode merges the partial attention of every rank once per layer +(81 per GLM-5.2 step): pack -> all_to_all_single -> unpack + LSE combine + +per-token FP8 quant. On RCCL the all-to-all alone (5.3 MB per rank at graph +size 40) takes ~40 us in the decode graph, ~3.3 ms per step, and no RCCL +setting helps (NCCL_PROTO / channel counts, measured). aiter's custom all-gather moves the same bytes per rank (its 1.3 MB input x 4 ranks) in about -half the time, so this adds the all-to-all next to it: +half the time, so this adds, next to it: - ``alltoall_vec``: allgather_vec reading only the chunk addressed to this rank: rank r receives chunk r of every rank's input, in rank order (the all_to_all_single layout). Same start/end sync, same <= kMaxBlocks grid. -- ``dispatchAllGather`` runs it for gather_dim == -1000, so the existing - all_gather_reg/unreg bindings, the copy-in pool and the CUDA-graph buffer - registration all apply unchanged; no new pybind symbol. -- ``CustomAllreduce.custom_all_to_all(inp)``: the Python entry, None where the - path does not apply (gfx1250 bindings, a size the vector kernel cannot split). + Python entry ``CustomAllreduce.custom_all_to_all(inp)``, None where the path + does not apply (gfx1250 bindings, a size the vector kernel cannot split); + patch_atom_dcp_custom_a2a.py uses it. +- ``alltoall_combine_quant``: the all-to-all, the LSE combine and the FP8 quant + in one kernel. With alltoall_vec the merge is still pack -> all-to-all -> + unpack + combine + quant (~6 + 31 + 8.5 us at graph size 40), but the + receiving rank needs only its own heads of every rank's partials, so this + reads them straight from the peers' registered buffers: rank r reads chunk r + of every rank's input (data [N][rows][h_local][head_dim], then fp32 LSE + [N][rows][h_local]), merges each (row, head) over the N partials with the LSE + weights, and stores per-token FP8 rows plus one fp32 scale per row; the + received copy and the unpack launch go away. The arithmetic is that of ATOM's + _dcp_a2a_unpack_combine_quant_kernel (non-finite LSE -> -inf, NaN partials of + empty ranks dropped, scale = row max * (1 / FP8 max)). Python entry + ``CustomAllreduce.supports_a2a_combine_quant`` / ``custom_a2a_combine_quant``; + patch_atom_dcp_custom_a2a_combine.py packs the split layout only when supported. + +``dispatchAllGather`` runs them for gather_dim == -1000 and -2000 - h_local, so +the existing all_gather_reg/unreg bindings, the copy-in pool and the CUDA-graph +buffer registration all apply unchanged; no new pybind symbol. Both are +instantiated for 2, 4 and 8 GPUs, and the Python entries refuse other sizes. The patched module (module_custom_all_reduce) is rebuilt here, the way the runtime JIT would; needs no GPU, the target archs come from GPU_ARCHS @@ -30,13 +47,14 @@ import importlib.util import os +import shutil import sys from anchor_patch import apply TAG = "aiter-patch custom-alltoall" -KERNEL = r""" +KERNELS = r""" // infera: all-to-all. size = 16-byte packs per destination chunk; rank r // receives chunk r of every rank's registered input, in rank order. constexpr int kAllToAllDim = -1000; @@ -68,9 +86,166 @@ end_sync(sg, self_sg, rank); } +// infera: DCP merge in one kernel. Input of every rank: data +// [ngpus][rows][h_local][head_dim] (T), then fp32 LSE [ngpus][rows][h_local]. +// Rank r merges chunk r of all ranks: out = fp8 [rows][h_local * head_dim], +// then fp32 scale [rows]. Needs head_dim % pack == 0 and h_local * head_dim +// <= 512 * pack (one pack per thread and row). +constexpr int kA2ACombineDim = -2000; // gather_dim = kA2ACombineDim - h_local + +template +__global__ void __launch_bounds__(512, 1) alltoall_combine_quant(RankData* _dp, + RankSignals sg, + Signal* self_sg, + uint8_t* __restrict__ out, + int rank, + int rows, + int h_local, + int head_dim) +{ + constexpr int pack = 16 / sizeof(T); + constexpr int warps = 512 / 64; + using P = typename opus::vector_t; + __shared__ float warp_max[warps]; + const int row_len = h_local * head_dim; + const int64_t data_elems = (int64_t)ngpus * rows * row_len; + const float fp8_max = opus::cast(opus::numeric_limits::max()); + const int e0 = threadIdx.x * pack; + const bool active = e0 < row_len; + const int h = e0 / head_dim; + const T* src[ngpus]; +#pragma unroll + for(int i = 0; i < ngpus; ++i) + { + src[i] = (const T*)_dp->ptrs[i]; + } + start_sync(sg, self_sg, rank); + + for(int b = blockIdx.x; b < rows; b += gridDim.x) + { + const int64_t chunk_row = (int64_t)rank * rows + b; + float acc[pack]; +#pragma unroll + for(int j = 0; j < pack; ++j) + { + acc[j] = 0.f; + } + float amax = 0.f; + if(active) + { + P v[ngpus]; + float lse[ngpus]; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + v[s] = *reinterpret_cast(src[s] + chunk_row * row_len + e0); + lse[s] = reinterpret_cast(src[s] + data_elems)[chunk_row * h_local + h]; + } + // A rank without KV for the row reports lse -inf (or NaN) and o NaN. + float lse_max = -INFINITY; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + if(isnan(lse[s]) || lse[s] == INFINITY) + lse[s] = -INFINITY; + lse_max = fmaxf(lse_max, lse[s]); + } + if(lse_max == -INFINITY) + lse_max = 0.f; + float sum = 0.f; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + sum += expf(lse[s] - lse_max); + } + const float global_lse = logf(sum) + lse_max; +#pragma unroll + for(int s = 0; s < ngpus; ++s) + { + const float factor = expf(lse[s] - global_lse); + if(isnan(factor) || factor == 0.f) + continue; // drops the NaN partials of empty ranks +#pragma unroll + for(int j = 0; j < pack; ++j) + { + acc[j] += upcast_s(v[s][j]) * factor; + } + } +#pragma unroll + for(int j = 0; j < pack; ++j) + { + amax = fmaxf(amax, fabsf(acc[j])); + } + } +#pragma unroll + for(int lane_mask = 32; lane_mask > 0; lane_mask >>= 1) + { + amax = fmaxf(amax, __shfl_xor(amax, lane_mask, 64)); + } + if(threadIdx.x % 64 == 0) + warp_max[threadIdx.x / 64] = amax; + __syncthreads(); + float row_max = warp_max[0]; +#pragma unroll + for(int w = 1; w < warps; ++w) + { + row_max = fmaxf(row_max, warp_max[w]); + } + __syncthreads(); + // As the Triton combine: multiply by 1 / FP8 max, zero scale -> zero row. + const float scale = row_max * (1.0f / fp8_max); + const float inv = scale > 0.f ? 1.0f / scale : 0.f; + if(active) + { + opus::vector_t q; +#pragma unroll + for(int i = 0; i < pack / 4; ++i) + { + opus::fp32x4_t tmp; +#pragma unroll + for(int j = 0; j < 4; ++j) + { + tmp[j] = fminf(fmaxf(acc[i * 4 + j] * inv, -fp8_max), fp8_max); + } + *(reinterpret_cast(&q) + i) = opus::cast(tmp); + } + *reinterpret_cast*>(out + (int64_t)b * row_len + e0) = q; + } + if(threadIdx.x == 0) + reinterpret_cast(out + (int64_t)rows * row_len)[b] = scale; + } + end_sync(sg, self_sg, rank); +} + """ -DISPATCH = """ // infera: all-to-all, see alltoall_vec. +DISPATCH = """ // infera: DCP merge (all-to-all + LSE combine + FP8 quant), see alltoall_combine_quant. + if(gather_dim <= kA2ACombineDim) + { + int h_local = kA2ACombineDim - gather_dim; + int head_dim = last_dim_size - 2; // input viewed as [ngpus * rows * h_local, head_dim + 2] + int rows = size / (world_size_ * h_local * last_dim_size); + dim3 grid(std::min(rows, kMaxBlocks)); + auto* out = reinterpret_cast(output); + switch(world_size_) + { + case 8: + alltoall_combine_quant + <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); + break; + case 4: + alltoall_combine_quant + <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); + break; + case 2: + alltoall_combine_quant + <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); + break; + default: printf("alltoall_combine_quant world_size error\\n"); + } + return; + } + // infera: all-to-all, see alltoall_vec. if(gather_dim == kAllToAllDim) { int chunk = size / d / world_size_; @@ -94,8 +269,10 @@ } """ -PY_METHOD = ''' # infera: dispatchAllGather runs the all-to-all kernel for this "dim". +PY_METHODS = ''' # infera: dispatchAllGather runs the all-to-all kernel for this "dim". _ALL_TO_ALL_DIM = -1000 + # Both infera kernels are instantiated for these world sizes only. + _A2A_WORLD_SIZES = (2, 4, 8) def custom_all_to_all(self, inp: torch.Tensor) -> torch.Tensor | None: """all_to_all_single over dim 0 of ``inp``: its world_size equal chunks @@ -105,6 +282,7 @@ def custom_all_to_all(self, inp: torch.Tensor) -> torch.Tensor | None: nbytes = inp.numel() * inp.element_size() if ( self._ops_all_gather is not None + or self.world_size not in self._A2A_WORLD_SIZES or self._INT_TO_FP_VIEW.get(inp.dtype, inp.dtype) not in (torch.float32, torch.float16, torch.bfloat16) or nbytes % (16 * self.world_size) != 0 @@ -128,6 +306,54 @@ def custom_all_to_all(self, inp: torch.Tensor) -> torch.Tensor | None: ) return out.view(inp.dtype) + # infera: dispatchAllGather runs the DCP merge kernel for this "dim" - h_local. + _A2A_COMBINE_DIM = -2000 + + def supports_a2a_combine_quant( + self, send: torch.Tensor, h_local: int, head_dim: int, quant_dtype + ) -> bool: + """Whether custom_a2a_combine_quant runs for this split-layout input + (see the method); otherwise the caller keeps its RCCL / all-to-all path.""" + return ( + self._ops_all_gather is None + and self.world_size in self._A2A_WORLD_SIZES + and send.dtype in (torch.float16, torch.bfloat16) + and quant_dtype == fp8 + and head_dim % 8 == 0 + and 0 < h_local * head_dim <= 512 * 8 + and send.dim() == 2 + and send.shape[1] == head_dim + 2 + and send.shape[0] % (self.world_size * h_local) == 0 + and self.should_custom_ag(send) + ) + + def custom_a2a_combine_quant(self, send: torch.Tensor, h_local: int, head_dim: int): + """DCP attention merge of ``send`` = [world_size * rows * h_local, + head_dim + 2], holding data [world_size][rows][h_local][head_dim] and + then the fp32 LSE [world_size][rows][h_local] (chunk i is for rank i). + Returns (fp8 [rows, h_local, head_dim], fp32 scale [rows, 1]): the + LSE-weighted merge of this rank's chunks, quantized per row. Buffer and + capture rules as custom_all_gather; check supports_a2a_combine_quant first.""" + rows = send.shape[0] // (self.world_size * h_local) + row_bytes = h_local * head_dim + out = torch.empty(rows * (row_bytes + 4), dtype=torch.uint8, device=send.device) + dim = self._A2A_COMBINE_DIM - h_local + capturing = self._IS_CAPTURING and torch.cuda.is_current_stream_capturing() + if capturing and self.enable_register_for_capturing: + ops.all_gather_reg(self._ptr, send, out, dim) + else: + ops.all_gather_unreg( + self._ptr, + send, + self._pool["input"].data_ptr, + out, + self._pool["input"].max_size, + dim, + ) + q = out[: rows * row_bytes].view(fp8).view(rows, h_local, head_dim) + scale = out[rows * row_bytes :].view(torch.float32).view(rows, 1) + return q, scale + def fused_ar_rms( ''' @@ -137,7 +363,7 @@ def fused_ar_rms( [ ( "// ========== reduce_scatter kernel start ==========\n", - KERNEL + "// ========== reduce_scatter kernel start ==========\n", + KERNELS + "// ========== reduce_scatter kernel start ==========\n", ), ( " // only support gather first dim and gather last dim\n", @@ -147,7 +373,7 @@ def fused_ar_rms( ), "dist/device_communicators/custom_all_reduce.py": ( "# infera: dispatchAllGather runs the all-to-all kernel", - [(" def fused_ar_rms(\n", PY_METHOD)], + [(" def fused_ar_rms(\n", PY_METHODS)], ), } @@ -168,6 +394,9 @@ def rebuild() -> int: print(f"[{TAG}] {so} is newer than the patched header; not rebuilt") return 0 os.remove(so) + # The JIT compiles from a copy of the sources in its build dir; an image + # that already built the module would recompile the old copy. + shutil.rmtree(f"{core.get_user_jit_dir()}/build/{MODULE}", ignore_errors=True) args = core.get_args_of_build(MODULE) core.build_module( MODULE, diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py deleted file mode 100644 index d7dfc1286..000000000 --- a/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall_combine.py +++ /dev/null @@ -1,323 +0,0 @@ -#!/usr/bin/env python3 -############################################################################### -# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. -# -# SPDX-License-Identifier: MIT -############################################################################### -"""The DCP attention merge as one kernel: all-to-all, LSE combine, FP8 quant. - -With patch_aiter_custom_alltoall.py, ATOM's DCP decode merge per layer is -pack -> custom all-to-all -> unpack + LSE combine + per-token FP8 quant (three -launches, ~6 + 31 + 8.5 us at graph size 40, 81 per GLM-5.2 step). The -receiving rank needs only its own heads of every rank's partials, so the -combine can read them straight from the peers' registered buffers, like the -all-to-all does, and write the quantized o_proj input directly: the received -copy and the unpack launch go away. - -- ``alltoall_combine_quant``: rank r reads chunk r of every rank's input - (data [N][rows][h_local][head_dim], then fp32 LSE [N][rows][h_local]), - merges each (row, head) over the N partials with the LSE weights, and stores - per-token FP8 rows plus one fp32 scale per row. The arithmetic is that of - ATOM's _dcp_a2a_unpack_combine_quant_kernel (non-finite LSE -> -inf, NaN - partials of empty ranks dropped, scale = row max * (1 / FP8 max)). -- ``dispatchAllGather`` runs it for gather_dim = -2000 - h_local, so the - registered-buffer and CUDA-graph paths of all_gather_reg / unreg apply as for - the all-to-all; no new pybind symbol. -- ``CustomAllreduce.supports_a2a_combine_quant`` / ``custom_a2a_combine_quant``: - the Python entry; the caller packs the split layout only when supported. - -Applies on top of patch_aiter_custom_alltoall.py and rebuilds -module_custom_all_reduce as that one does. Written against aiter 9205e1059 in -rocm/atom-dev:nightly_202610131505. Strict: see anchor_patch.py. -""" - -import importlib.util -import os -import shutil -import sys - -from anchor_patch import apply - -TAG = "aiter-patch custom-alltoall-combine" - -KERNEL = r""" -// infera: DCP merge in one kernel. Input of every rank: data -// [ngpus][rows][h_local][head_dim] (T), then fp32 LSE [ngpus][rows][h_local]. -// Rank r merges chunk r of all ranks: out = fp8 [rows][h_local * head_dim], -// then fp32 scale [rows]. Needs head_dim % pack == 0 and h_local * head_dim -// <= 512 * pack (one pack per thread and row). -constexpr int kA2ACombineDim = -2000; // gather_dim = kA2ACombineDim - h_local - -template -__global__ void __launch_bounds__(512, 1) alltoall_combine_quant(RankData* _dp, - RankSignals sg, - Signal* self_sg, - uint8_t* __restrict__ out, - int rank, - int rows, - int h_local, - int head_dim) -{ - constexpr int pack = 16 / sizeof(T); - constexpr int warps = 512 / 64; - using P = typename opus::vector_t; - __shared__ float warp_max[warps]; - const int row_len = h_local * head_dim; - const int64_t data_elems = (int64_t)ngpus * rows * row_len; - const float fp8_max = opus::cast(opus::numeric_limits::max()); - const int e0 = threadIdx.x * pack; - const bool active = e0 < row_len; - const int h = e0 / head_dim; - const T* src[ngpus]; -#pragma unroll - for(int i = 0; i < ngpus; ++i) - { - src[i] = (const T*)_dp->ptrs[i]; - } - start_sync(sg, self_sg, rank); - - for(int b = blockIdx.x; b < rows; b += gridDim.x) - { - const int64_t chunk_row = (int64_t)rank * rows + b; - float acc[pack]; -#pragma unroll - for(int j = 0; j < pack; ++j) - { - acc[j] = 0.f; - } - float amax = 0.f; - if(active) - { - P v[ngpus]; - float lse[ngpus]; -#pragma unroll - for(int s = 0; s < ngpus; ++s) - { - v[s] = *reinterpret_cast(src[s] + chunk_row * row_len + e0); - lse[s] = reinterpret_cast(src[s] + data_elems)[chunk_row * h_local + h]; - } - // A rank without KV for the row reports lse -inf (or NaN) and o NaN. - float lse_max = -INFINITY; -#pragma unroll - for(int s = 0; s < ngpus; ++s) - { - if(isnan(lse[s]) || lse[s] == INFINITY) - lse[s] = -INFINITY; - lse_max = fmaxf(lse_max, lse[s]); - } - if(lse_max == -INFINITY) - lse_max = 0.f; - float sum = 0.f; -#pragma unroll - for(int s = 0; s < ngpus; ++s) - { - sum += expf(lse[s] - lse_max); - } - const float global_lse = logf(sum) + lse_max; -#pragma unroll - for(int s = 0; s < ngpus; ++s) - { - const float factor = expf(lse[s] - global_lse); - if(isnan(factor) || factor == 0.f) - continue; // drops the NaN partials of empty ranks -#pragma unroll - for(int j = 0; j < pack; ++j) - { - acc[j] += upcast_s(v[s][j]) * factor; - } - } -#pragma unroll - for(int j = 0; j < pack; ++j) - { - amax = fmaxf(amax, fabsf(acc[j])); - } - } -#pragma unroll - for(int lane_mask = 32; lane_mask > 0; lane_mask >>= 1) - { - amax = fmaxf(amax, __shfl_xor(amax, lane_mask, 64)); - } - if(threadIdx.x % 64 == 0) - warp_max[threadIdx.x / 64] = amax; - __syncthreads(); - float row_max = warp_max[0]; -#pragma unroll - for(int w = 1; w < warps; ++w) - { - row_max = fmaxf(row_max, warp_max[w]); - } - __syncthreads(); - // As the Triton combine: multiply by 1 / FP8 max, zero scale -> zero row. - const float scale = row_max * (1.0f / fp8_max); - const float inv = scale > 0.f ? 1.0f / scale : 0.f; - if(active) - { - opus::vector_t q; -#pragma unroll - for(int i = 0; i < pack / 4; ++i) - { - opus::fp32x4_t tmp; -#pragma unroll - for(int j = 0; j < 4; ++j) - { - tmp[j] = fminf(fmaxf(acc[i * 4 + j] * inv, -fp8_max), fp8_max); - } - *(reinterpret_cast(&q) + i) = opus::cast(tmp); - } - *reinterpret_cast*>(out + (int64_t)b * row_len + e0) = q; - } - if(threadIdx.x == 0) - reinterpret_cast(out + (int64_t)rows * row_len)[b] = scale; - } - end_sync(sg, self_sg, rank); -} - -""" - -DISPATCH = """ // infera: DCP merge (all-to-all + LSE combine + FP8 quant), see alltoall_combine_quant. - if(gather_dim <= kA2ACombineDim) - { - int h_local = kA2ACombineDim - gather_dim; - int head_dim = last_dim_size - 2; // input viewed as [ngpus * rows * h_local, head_dim + 2] - int rows = size / (world_size_ * h_local * last_dim_size); - dim3 grid(std::min(rows, kMaxBlocks)); - auto* out = reinterpret_cast(output); - switch(world_size_) - { - case 8: - alltoall_combine_quant - <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); - break; - case 4: - alltoall_combine_quant - <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); - break; - case 2: - alltoall_combine_quant - <<>>(ptrs, sg_, self_sg_, out, rank_, rows, h_local, head_dim); - break; - default: printf("alltoall_combine_quant world_size error\\n"); - } - return; - } -""" - -PY_METHOD = ''' # infera: dispatchAllGather runs the DCP merge kernel for this "dim" - h_local. - _A2A_COMBINE_DIM = -2000 - - def supports_a2a_combine_quant( - self, send: torch.Tensor, h_local: int, head_dim: int, quant_dtype - ) -> bool: - """Whether custom_a2a_combine_quant runs for this split-layout input - (see the method); otherwise the caller keeps its RCCL / all-to-all path.""" - return ( - self._ops_all_gather is None - and send.dtype in (torch.float16, torch.bfloat16) - and quant_dtype == fp8 - and head_dim % 8 == 0 - and 0 < h_local * head_dim <= 512 * 8 - and send.dim() == 2 - and send.shape[1] == head_dim + 2 - and send.shape[0] % (self.world_size * h_local) == 0 - and self.should_custom_ag(send) - ) - - def custom_a2a_combine_quant(self, send: torch.Tensor, h_local: int, head_dim: int): - """DCP attention merge of ``send`` = [world_size * rows * h_local, - head_dim + 2], holding data [world_size][rows][h_local][head_dim] and - then the fp32 LSE [world_size][rows][h_local] (chunk i is for rank i). - Returns (fp8 [rows, h_local, head_dim], fp32 scale [rows, 1]): the - LSE-weighted merge of this rank's chunks, quantized per row. Buffer and - capture rules as custom_all_gather; check supports_a2a_combine_quant first.""" - rows = send.shape[0] // (self.world_size * h_local) - row_bytes = h_local * head_dim - out = torch.empty(rows * (row_bytes + 4), dtype=torch.uint8, device=send.device) - dim = self._A2A_COMBINE_DIM - h_local - capturing = self._IS_CAPTURING and torch.cuda.is_current_stream_capturing() - if capturing and self.enable_register_for_capturing: - ops.all_gather_reg(self._ptr, send, out, dim) - else: - ops.all_gather_unreg( - self._ptr, - send, - self._pool["input"].data_ptr, - out, - self._pool["input"].max_size, - dim, - ) - q = out[: rows * row_bytes].view(fp8).view(rows, h_local, head_dim) - scale = out[rows * row_bytes :].view(torch.float32).view(rows, 1) - return q, scale - - def fused_ar_rms( -''' - -PATCH_CSRC = { - "../csrc/include/custom_all_reduce.cuh": ( - "// infera: DCP merge in one kernel.", - [ - ( - "// ========== reduce_scatter kernel start ==========\n", - KERNEL + "// ========== reduce_scatter kernel start ==========\n", - ), - ( - " // infera: all-to-all, see alltoall_vec.\n", - DISPATCH + " // infera: all-to-all, see alltoall_vec.\n", - ), - ], - ), - "dist/device_communicators/custom_all_reduce.py": ( - "# infera: dispatchAllGather runs the DCP merge kernel", - [(" def fused_ar_rms(\n", PY_METHOD)], - ), -} - -MODULE = "module_custom_all_reduce" - - -def rebuild() -> int: - """Rebuild the module from the patched sources, as the runtime JIT would.""" - root = importlib.util.find_spec("aiter").submodule_search_locations[0] - os.environ.setdefault("GPU_ARCHS", os.environ.get("PYTORCH_ROCM_ARCH", "gfx950")) - sys.path.insert(0, root) - from jit import core # noqa: E402 (import after GPU_ARCHS is set) - - so = f"{core.get_user_jit_dir()}/{MODULE}.so" - cuh = os.path.join(root, "../csrc/include/custom_all_reduce.cuh") - if os.path.exists(so): - if os.path.getmtime(so) > os.path.getmtime(cuh): - print(f"[{TAG}] {so} is newer than the patched header; not rebuilt") - return 0 - os.remove(so) - # The JIT compiles from a copy of the sources in its build dir; an image - # that already built the module would recompile the old copy. - shutil.rmtree(f"{core.get_user_jit_dir()}/build/{MODULE}", ignore_errors=True) - args = core.get_args_of_build(MODULE) - core.build_module( - MODULE, - args["srcs"], - args["flags_extra_cc"], - args["flags_extra_hip"], - args["blob_gen_cmd"], - args["extra_include"], - args["extra_ldflags"], - args["verbose"], - args["is_python_module"], - args["is_standalone"], - args["torch_exclude"], - args.get("third_party", []), - args.get("hipify", False), - flags_extra_hip_per_source=args.get("flags_extra_hip_per_source", {}), - ) - if not os.path.isfile(so): - print(f"[{TAG}] FAILED: AITER build left no {so}") - return 1 - print(f"[{TAG}] rebuilt {so}") - return 0 - - -if __name__ == "__main__": - rc = apply(TAG, PATCH_CSRC, package="aiter") - if rc == 0 and "--no-build" not in sys.argv: - rc = rebuild() - sys.exit(rc) diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py b/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py index f871feb8c..b58468bf4 100644 --- a/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py +++ b/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py @@ -17,17 +17,25 @@ production path at the same shape (the tuner's --run_config, measured on the same node) are kept, since for M not a multiple of 32 the tuned kernel is often slower than the padded one. They go in as one more model_configs table, which -aiter merges with the others at startup (none of these keys exist there). +aiter merges with the others at startup. On a key that two tables share +(gfx, cu_num, M, N, K, q_dtype_w) aiter raises at engine start +(update_config_files), so the build fails instead when the base already holds +one of these keys: drop or re-cut the rows then. Written against aiter 9205e1059 in rocm/atom-dev:nightly_202610131505. """ +import csv +import glob import importlib.util +import io import os import sys TAG = "aiter-patch glm52-decode-gemm-rows" REL = "configs/model_configs/a8w8_bpreshuffle_tuned_gemm_glm5.2_decode.csv" +# The key aiter de-duplicates the merged a8w8_bpreshuffle tables on. +KEY = ("gfx", "cu_num", "M", "N", "K", "q_dtype_w") CSV = """\ gfx,cu_num,M,N,K,q_dtype_w,libtype,kernelId,splitK,us,kernelName,tflops,bw,errRatio gfx950,256,160,2048,2048,torch.float8_e4m3fn,flydsl,2371,0,5.6743,flydsl_bpreshuflle_32x32x512_F8_F8_B16_1x4x4x2_default,236.54,912.42,0.0 @@ -60,15 +68,45 @@ """ +def shared_keys(root: str) -> list[str]: + """The rows of CSV whose key a table aiter merges them with already has.""" + ours = [tuple(row[k] for k in KEY) for row in csv.DictReader(io.StringIO(CSV))] + configs = os.path.join(root, "configs") + # The tables aiter merges: the default one and this glob under model_configs/. + tables = glob.glob(os.path.join(configs, "a8w8_bpreshuffle_tuned_gemm.csv")) + tables += sorted( + glob.glob(os.path.join(configs, "model_configs", "*a8w8_bpreshuffle_tuned_gemm*.csv")) + ) + found = [] + for table in tables: + name = os.path.basename(table) + if "untuned" in name or table == os.path.join(root, REL): + continue + with open(table, encoding="utf-8", newline="") as f: + reader = csv.DictReader(f) + # aiter fills a missing gfx from cu_num; compare on the columns present. + cols = [i for i, k in enumerate(KEY) if k in (reader.fieldnames or ())] + have = {tuple(row[KEY[i]].strip() for i in cols) for row in reader} + found += [ + f"{name} M={k[2]} N={k[3]} K={k[4]}" for k in ours if tuple(k[i] for i in cols) in have + ] + return found + + def main() -> int: spec = importlib.util.find_spec("aiter") if spec is None or not spec.submodule_search_locations: print(f"[{TAG}] FAILED: package 'aiter' not found") return 1 - path = os.path.join(list(spec.submodule_search_locations)[0], REL) + root = list(spec.submodule_search_locations)[0] + path = os.path.join(root, REL) if not os.path.isdir(os.path.dirname(path)): print(f"[{TAG}] FAILED: {os.path.dirname(path)} not found") return 1 + shared = shared_keys(root) + if shared: + print(f"[{TAG}] FAILED: aiter already has these keys: {'; '.join(shared)}") + return 1 if os.path.exists(path): with open(path, encoding="utf-8") as f: same = f.read() == CSV diff --git a/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py b/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py index e90a791dc..fc217d04d 100644 --- a/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py +++ b/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py @@ -20,7 +20,9 @@ 64, of 16 up to 512) and the fastest config for each from a sweep on MI355X (72 candidates, CUDA-graph timing, output checked against the stock config): at graph 40 q-up 19.1 -> 8.4 us and v-up 15.2 -> 8.1 us, ~1.4 ms per decode -step over 78 layers. The stock M_GEQ_*/any entries still serve M > 512. +step over 78 layers. The stock M_GEQ_*/any entries still serve M > 512. The +sweep starts at M 16: smaller M (graphs 1-3, drafts below 16) takes the M 16 +row, BLOCK_SIZE_M 16 as the stock M_LEQ_16 it replaces. Written against aiter 9205e1059 in rocm/atom-dev:nightly_202610131505. Strict: a file that no longer holds the stock buckets fails the build. @@ -49,6 +51,7 @@ "cache_modifier": ".cg", } # (N, K): [(M, BLOCK_SIZE_M, BLOCK_SIZE_N, num_warps, waves_per_eu, cache_modifier)] +# fmt: off TUNED = { (512, 192): [ (16, 16, 128, 4, 2, None), (20, 16, 128, 4, 2, None), (24, 16, 128, 4, 2, None), @@ -93,6 +96,7 @@ (496, 16, 256, 4, 2, None), (512, 64, 256, 8, 1, None), ], } +# fmt: on # Above 512 the stock M_GEQ_* / any entries keep serving (prefill sizes). UPPER_BOUNDS = [1024, 2048, 4096, 8192] diff --git a/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py b/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py index 3ca127684..d803afb35 100644 --- a/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py +++ b/deploy/docker/patches/atom_nightly/patch_atom_api_gc_stalls.py @@ -49,7 +49,7 @@ " # this queue. Queued anyway, every finished Sequence\n" " # stayed reachable and each gen-2 GC pass grew.\n" " if self._flush_stream_batch_fn is None or getattr(\n" - " self, \"_asyncio_mode\", False\n" + ' self, "_asyncio_mode", False\n' " ):\n" " self.outputs_queue.put_nowait(seqs)\n", ), @@ -64,7 +64,7 @@ " # infera: per-request context gauges are opt-in. Each request adds a\n" " # series that is never removed, and every /metrics render of the API\n" " # process re-reads all of them (render time and GC grow with requests).\n" - " _PER_REQUEST = os.environ.get(\"ATOM_METRICS_PER_REQUEST_CONTEXT\") == \"1\"\n" + ' _PER_REQUEST = os.environ.get("ATOM_METRICS_PER_REQUEST_CONTEXT") == "1"\n' "\n" " def _record_request_context(self, metric, req_id, seq, tokens):\n" " if not self._PER_REQUEST:\n" diff --git a/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py index 2a4ea504b..217972eb3 100644 --- a/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py +++ b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py @@ -8,11 +8,12 @@ When the DCP merge also quantizes for o_proj (PBM + per-token FP8, the decode of the GLM-5.2 kit), cp_lse_a2a runs pack -> all-to-all -> unpack + combine + -quant. With patch_aiter_custom_alltoall_combine.py the last two are one aiter -kernel that reads each rank's partials straight from the peers: the pack writes -the split layout that kernel reads (data, then the fp32 LSE), and the unpack -launch and the received copy go away. Where the aiter kernel does not apply -(no custom communicator, other dtypes or sizes) the path is unchanged. +quant. With patch_aiter_custom_alltoall.py the last two are one aiter kernel +(alltoall_combine_quant) that reads each rank's partials straight from the +peers: the pack writes the split layout that kernel reads (data, then the fp32 +LSE), and the unpack launch and the received copy go away. Where the aiter +kernel does not apply (no custom communicator, other dtypes, sizes or world +sizes) the path is unchanged. Written against ATOM 661ef766 in rocm/atom-dev:nightly_202610131505, on top of patch_atom_dcp_custom_a2a.py. Strict: see anchor_patch.py. diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index 04c36f473..a82ec90de 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -21,8 +21,8 @@ bash engine/down.sh # save logs, remove containers, wait for VRAM to dr ## Image The engines need the ATOM nightly `661ef766` (`rocm/atom-dev:nightly_202610131505`) -plus an ATOM DCP fix and an LMCache fix, and the router needs Infera's -`--session-affinity`. Build the image from the repo root: +plus Infera's ATOM, aiter and LMCache patches for it, and the router needs +Infera's `--session-affinity`. Build the image from the repo root: ```bash docker build -f deploy/docker/Dockerfile.atom \ @@ -36,11 +36,16 @@ docker build -f deploy/docker/Dockerfile.atom \ [`deploy/docker/patches/atom_nightly/`](../../deploy/docker/patches/atom_nightly/) instead of the release-base `patches/atom/`: the ATOM sparse DCP padded-row fix (without it the prefill hits a `Memory access fault` under long-prefix - load), the LMCache ROCm pinned-pool fix (the CPU tier becomes GTT, gotcha 4), - the AITER kernel prebuild of gotcha 12 and the `protobuf<7` pin the nightly's - OpenTelemetry needs. This nightly already has the DCP MTP-draft and LMCache - lookup-pin fixes the previous image patched in. - A drifted anchor fails the build. + load), the LMCache ROCm pinned-pool fix (the LMCache tier becomes GTT, + gotcha 4), the decode patches (no gen-2 GC stalls in ATOM's API process, the + DCP exchanges on aiter's custom IPC kernels, tuned MLA bmm configs and FP8 + GEMM rows), the AITER kernel prebuild of gotcha 12 and the `protobuf<7` pin + the nightly's OpenTelemetry needs. A drifted anchor fails the build. +- The API-process GC fix leaves ATOM's per-request context gauges + (`atom:prefill_request_context_tokens`, `atom:decode_request_context_tokens`) + off: each request added a series that every `/metrics` render re-read. Set + `ATOM_METRICS_PER_REQUEST_CONTEXT=1` in `PREFILL_EXTRA_ENV` / + `DECODE_EXTRA_ENV` to record them, for example for profiling. - `BUILD_MOONCAKE=0` keeps the base's Mooncake 0.3.13. That build routes same-host targets by upstream `MC_DISABLE_HIP`, which this kit sets (gotcha 1). - The default `Dockerfile.atom` build (ATOM 0.1.4 base) is not validated with this @@ -82,12 +87,13 @@ The settings are those of ATOM's GLM-5.2 AgentX recipe (`/app/ATOM/recipes/Agentic-GLM-5.2.md` in the image) and of InferenceX PR #4045's 1P1D (`MC_NUM_QP_PER_EP=1` on every Mooncake peer, `TOPK_FORCE_PATH=one`, `--timeout-keep-alive 900`, the Store). The decode's memory settings give it -~196K KV blocks instead of ~156K at 0.90 / 16384 with separate communicators: -AgentX lanes reuse their whole context every turn, and once C × mean prompt -length outgrows the decode's KV pool its prefix hits collapse, the prefill -re-sends whole prompts and TTFT explodes (about C85 before, C105 now). Not -taken from them: `OFFLOAD_PROFILE=1` (logging only), the decode's graphs (here -every 4 up to 2 × CONC, there every 8 up to 256) and BF16 MTP experts +~203K KV blocks (~196K with BF16 MTP experts) instead of ~160K at +0.90 / 16384 with separate communicators: AgentX lanes reuse their whole +context every turn, and once C × mean prompt length outgrows the decode's KV +pool its prefix hits collapse, the prefill re-sends whole prompts and TTFT +explodes (about C85 before, C105 now). Not taken from them: `OFFLOAD_PROFILE=1` +(logging only), the decode's graphs (here every 4 up to 2 × CONC, there every 8 +up to 256), its max-num-seqs (here 2 × CONC, there 512) and BF16 MTP experts (`MTP_EXPERTS`). MTP depth follows `CONC`: K3 at CONC ≥ 48, K4 below. The Store @@ -108,63 +114,75 @@ throughput per GPU is over all 8 GPUs. `PREFILL_PP=1` starts the prefill as TP4 + DCP4 with CUDA graphs up to 2 × CONC, as the decode. It was validated with these settings, which are not -the defaults; set them with it: +the defaults, and with the LMCache tier; set them with it: ```bash PREFILL_PP=1 PREFILL_MEM_UTIL=0.90 PREFILL_MAX_BATCHED_TOKENS=16384 BLOCK_SIZE=64 \ - OFFLOAD_MIN_LOAD_TOKENS=8192 PREFILL_SEND_WORKERS= bash engine/up.sh + OFFLOAD_MIN_LOAD_TOKENS=8192 PREFILL_SEND_WORKERS= PREFILL_CPU_TIER=lmcache bash engine/up.sh ``` Each DCP rank then writes to one decode rank without staging, and the LMCache lookup covers the four DCP ranks, each holding a distinct KV shard. Under AgentX -load it is slower than the PP prefill (reference numbers below). +load it is slower than the PP prefill (reference numbers below). It has not +been run with the Store or with the current decode settings. ## Validation status -The image above, on one 8×MI355X node: +The image above with the defaults of `env.sh`, on 8×MI355X nodes: -| item | PP prefill (default) | DCP prefill | +| item | PP prefill, Store (default) | PP prefill, `PREFILL_CPU_TIER=lmcache` | |---|---|---| | correctness (`smoke.py`, real MTP acceptance) | **passed**, 6/6 | **passed**, 6/6 | -| GSM8K (200 questions) | **0.98** (16 send workers) | **0.975** | -| AgentX, 3600 s | **passed** at C64 and C96, no transfer failure | **passed** at C48, C64, C96, no `Memory access fault` | +| GSM8K (200 questions) | **0.970** | **0.985** | +| AgentX, 3600 s | **passed** at C56, C72, C96, C112, 0 errors | **passed** at C56, C72, C96, C112, 0 errors | + +The DCP prefill passed smoke 6/6, GSM8K 0.975 and AgentX at C48, C64 and C96 +with an image without the decode patches, the LMCache tier and the decode at +0.90 / 16384 with separate communicators; it has not been re-run since. On shutdown each PP stage logs a few `TypeError` tracebacks from aiter `shm_broadcast`; none appear while serving. ## Reference numbers -These are the previous defaults: Python router, decode at 0.90 / 16384 with -separate communicators, 4 send workers, LMCache. - -AgentX agentic-coding replay, InferenceX `918524ff`, cold start per point, 10 -warm-up requests per lane, 3600 s measured, forced MTP acceptance (`MTP_AL=2.99` -with K3). Throughput is per GPU over all 8 GPUs and counts cache-hit input -tokens; TTFT in s, ITL in ms. - -PP and DCP prefill on the same node (crsuse2-m2m-140), image -`infera-atom:nightly-202610131505-pr198fix`, one run per point: - -| CONC | prefill | input tok/s/GPU | output tok/s/GPU | TTFT p50 / p90 / p99 | ITL p50 / p90 | prefill hit GPU / CPU / total % | errors | +AgentX agentic-coding replay with the defaults above, one cold start per point, +warm-up, 3600 s measured, forced MTP acceptance on both legs (AL 2.99 with K3), +`SESSION_AFFINITY=off` and router `--breaker-failure-threshold 0`, as InferenceX +PR #4045's runs. The client was InferenceX CI's (InferenceX `288471b3`, +agentx-harness 1.0.11), not `agentx.sh`'s `918524ff`. Throughput is per GPU +over all 8 GPUs and counts cache-hit input tokens; interactivity is the inverse +of P90 inter-token latency. C56 and C96 ran on one node, C72 and C112 on +another. + +| CONC | input tok/s/GPU | output tok/s/GPU | P90 interactivity tok/s/user | TTFT mean / p50 / p90 s | ITL p50 / p90 ms | prefill hit GPU / CPU / total % | decode hit % | requests / errors | +|---:|---:|---:|---:|---|---|---|---:|---| +| 56 | 33578 | 249.0 | 66.1 | 1.12 / 0.84 / 2.00 | 13.62 / 15.12 | 86.0 / 10.2 / 96.2 | 92.6 | 7117 / 0 | +| 72 | 40844 | 327.4 | 53.6 | 1.51 / 1.05 / 2.68 | 16.11 / 18.64 | 76.4 / 18.5 / 94.9 | 85.5 | 9627 / 0 | +| 96 | 47444 | 393.4 | 46.8 | 2.42 / 1.50 / 4.72 | 19.25 / 21.35 | 65.6 / 28.2 / 93.8 | 60.4 | 11482 / 0 | +| 112 | **48859** | 416.6 | 43.9 | 4.42 / 2.86 / 9.19 | 20.52 / 22.76 | 59.9 / 33.1 / 93.0 | 32.1 | 12062 / 0 | + +The decode hit falls from C96 on, as C × prompt length outgrows the decode's +KV pool; the Store's RDMA loads do not block the prefill, so throughput still +rises to C112. With `PREFILL_CPU_TIER=lmcache` (other nodes) C56 and C72 match +(33392 / 40326), but C96 and C112 drop to 30690 / 23321: LMCache loads +synchronously and takes about half of each PP stage's time there. + +The previous defaults (Python router, decode at 0.90 / 16384 with separate +communicators, 4 send workers, LMCache, an image without the decode patches) +compared the two prefills on one node, with `agentx.sh`: + +| CONC | prefill | input tok/s/GPU | output tok/s/GPU | TTFT p50 / p90 / p99 s | ITL p50 / p90 ms | prefill hit GPU / CPU / total % | errors | |---:|---|---:|---:|---|---|---|---| | 64 | PP4 | **33389.60** (+8.6%) | 262.13 | 2.09 / 5.17 / 15.74 | 16.46 / 17.59 | 77.7 / 17.2 / 95.0 | 1 / 7682 | | 64 | DCP | 30739.05 | 242.78 | 3.23 / 17.81 / 49.41 | 16.03 / 17.30 | 79.8 / 16.7 / 96.5 | 1 / 7024 | | 96 | PP4 | **30115.94** (+44.1%) | 274.59 | 24.98 / 51.53 / 76.31 | 16.75 / 18.74 | 34.6 / 52.0 / 86.6 | 0 / 8142 | | 96 | DCP | 20902.07 | 193.03 | 51.39 / 103.55 / 161.52 | 13.75 / 16.48 | 16.3 / 77.6 / 93.9 | 0 / 5664 | -DCP prefill, image above, C48 and C64 on two different nodes: - -| CONC | MTP K / AL | total tok/s/GPU | output tok/s/GPU | input tok/s/GPU | P90 interactivity tok/s/user | TTFT p50 / p90 / p99 s | errors / tail cancels | -|---:|---|---:|---:|---:|---:|---|---| -| 48 | 3 / 2.99 | 26325.24 | 183.30 | 26141.94 | 66.52 | 1.759 / 7.662 / 24.695 | 2 / 4 | -| 64 | 3 / 2.99 | 30820.57 | 241.22 | 30579.35 | 56.83 | 3.114 / 17.256 / 47.358 | 1 / 27 | - -Interactivity is the inverse of P90 inter-token latency. - -To reproduce a point (cold start each time): +To run a point (cold start each time) with the settings of the first table: ```bash -CONC=64 MTP_AL=2.99 bash engine/up.sh # CONC < 48: MTP_AL=3.33 (K4) +CONC=96 MTP_AL=2.99 SESSION_AFFINITY=off ROUTER_EXTRA_ARGS="--breaker-failure-threshold 0" \ + PREFILL_EXTRA_ARGS="--spec-decode-acceptance-length 2.99" bash engine/up.sh # CONC < 48: 3.33 (K4) bash engine/agentx/agentx.sh # DURATION=3600 WARMUP_PER_LANE=10 by default bash engine/down.sh ``` @@ -214,7 +232,7 @@ base + PP stage on the prefill and base + TP rank on the decode, so 21301-21304 and 21311-21314. **7. Router backend and transports.** `ROUTER_BACKEND` picks the Rust router -(default) or the Python one (used for the reference numbers below); both route +(default) or the Python one (used for the previous-defaults numbers); both route ATOM PD (`atom-mooncake`) and hand the prefill's prompt token ids to the decode, so the decode does not tokenize the prompt again. ATOM speaks neither NATS requests nor NATS KV events, so `up.sh` passes `--request-transport http --kv-event-transport zmq`. `kv-aware` @@ -224,9 +242,9 @@ would need engine KV events, so the policy is `round-robin`. `prefill`) is passed as `--session-affinity`: requests with an `x-session-id` keep their session's prefill (`both`: and decode), others are round-robin. `agentx.sh` has AIPerf send it (`AIPERF_HTTP_X_SESSION_ID_FROM_CORRELATION_ID`). -The router also forwards `x-session-id` and `x-correlation-id` to both legs. -With a single prefill and decode it only carries the headers through; it -matters once more instances are added. +The Python router also forwards `x-session-id` and `x-correlation-id` to both +legs; the Rust router reads them for placement only. With a single prefill and +decode, affinity has nothing to choose; it matters once more instances are added. Tune the idle TTL (default 3600 s) through `ROUTER_EXTRA_ARGS` (`--session-affinity-ttl-seconds`). `SESSION_AFFINITY=off` is the stateless alternative. @@ -241,8 +259,11 @@ this ATOM; the PP prefill stages per decode DCP rank instead (see the recipe). **11. Cold start takes about 10 minutes**, mostly weights, AITER/FlyDSL JIT and graph capture. JIT caches persist in per-role, per-image Docker volumes -`-cache--` and `-aiter-jit--`; remove -them with `docker volume rm` when done with an image. +`-cache---mbt` and `-aiter-jit--`; +remove them with `docker volume rm` when done with an image. `` is the leg's +`--max-num-batched-tokens`: ATOM's compile cache does not key on it, and the +prefill fails at startup in `assert_size_stride` on artifacts compiled for +another value. **12. A JIT compile on live traffic stalls the prefill.** The base does not prebuild the AITER varlen attention that chunked-MLA prefill calls once a prompt @@ -254,20 +275,21 @@ archs (gfx942) or paths. AITER keeps such modules in `/app/aiter-test/aiter/jit` not `/root/.cache`, so `worker.sh` persists that directory too, for any other kernel compiled at run time. -**13. The PP prefill needs fewer Mooncake send workers.** With ATOM's default -16 send workers per stage, every AgentX run on MI355X (crsuse2-m2m-140) failed -within its first four requests: one RDMA write from a stage to one decode rank -never completed (`staged-mla` / `landed-mla RDMA chunk error -1 ... Mooncake's -sync timeout`, Mooncake `CQ completion timeout diagnostic`, then `Failed to poll -CQ` and ionic `wqe is not posted` without end). Every later transfer to that -rank failed, so every request did. It hit `ionic_4` and `ionic_5`, with forced -and with real MTP acceptance. One send worker removed the hang, but a 443k-token -prompt then missed ATOM's fixed 120 s transfer deadline (`PD_TRANSFER_DEADLINE_S`, -queueing included). Four workers ran C64 and C96 without a transfer failure. -All of that was with Mooncake's default two QPs per endpoint. ATOM's offload +**13. The PP prefill's 16 send workers need one QP per endpoint.** With +Mooncake's default two QPs per endpoint and ATOM's default 16 send workers per +stage, every AgentX run on one MI355X node failed within its first four +requests: one RDMA write from a stage to one decode rank never completed +(`staged-mla` / `landed-mla RDMA chunk error -1 ... Mooncake's sync timeout`, +Mooncake `CQ completion timeout diagnostic`, then `Failed to poll CQ` and ionic +`wqe is not posted` without end). Every later transfer to that rank failed, so +every request did. It hit `ionic_4` and `ionic_5`, with forced and with real MTP +acceptance. One send worker removed the hang, but a 443k-token prompt then +missed ATOM's fixed 120 s transfer deadline (`PD_TRANSFER_DEADLINE_S`, queueing +included); four ran C64 and C96 without a transfer failure. ATOM's offload README reports the same symptom for concurrent reads at two QPs, and ATOM's CI and InferenceX PR #4045 run the 16 workers with `MC_NUM_QP_PER_EP=1`, which the -kit now sets on every Mooncake peer; `PREFILL_SEND_WORKERS=4` is the fallback. +kit now sets on every Mooncake peer: 16 workers then ran 10 deployments (C56 to +C112) without a transfer failure. `PREFILL_SEND_WORKERS=4` remains the fallback. The DCP prefill, whose ranks each write to one decode rank without staging, never showed the hang. @@ -277,15 +299,19 @@ which the ionic NICs only accept when backed by huge pages. `store.sh` gets them without root through glibc's THP malloc (`GLIBC_TUNABLES=glibc.malloc.hugetlb=1`, THP `enabled` = `madvise` or `always`) and prints each owner's `AnonHugePages`; a fragmented node can leave an owner on 4 KiB pages, and it then fails to mount. -Each NUMA node needs `STORE_OWNER_GB` free plus the engines' host memory; lower -it or use `PREFILL_CPU_TIER=lmcache` on smaller hosts. The prefill refuses to -start when its startup probe reads back other bytes than it wrote ("the chunk -read back differs"): on the crusoe MI355X VMs every Store read into HBM -completed without moving data (puts from HBM and transfers to host buffers were -fine, and so is the P→D path), with peer memory and DMA-BUF alike; use -`PREFILL_CPU_TIER=lmcache` there. The prefill needs one -Store load worker (`OFFLOAD_LOAD_WORKERS=1`): more hung its pipeline at C96 in -ATOM's tests. +Each NUMA node needs `STORE_OWNER_GB` free plus the engines' host memory, and +`store.sh` refuses to start when THP is off or a NUMA node has less than +`STORE_OWNER_GB` free. Page cache does not count as free: on a node where it +holds most of the memory (a large weights read fills it), the owners reclaim it +slowly or end up on 4 KiB pages, so drop it first (as root, +`sync; echo 1 > /proc/sys/vm/drop_caches`, as InferenceX PR #4045's host setup +does). Otherwise lower `STORE_OWNER_GB` or use `PREFILL_CPU_TIER=lmcache`. The +prefill refuses to start when its startup probe reads back other bytes than it +wrote ("the chunk read back differs"): on some MI355X VM hosts every Store read +into HBM completed without moving data (puts from HBM and transfers to host +buffers were fine, and so is the P→D path), with peer memory and DMA-BUF alike; +use `PREFILL_CPU_TIER=lmcache` there. The prefill needs one Store load worker +(`OFFLOAD_LOAD_WORKERS=1`): more hung its pipeline at C96 in ATOM's tests. ## What this kit exposes on the network diff --git a/examples/atom_4p4d_glm5.2/engine/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index 951a6865b..8af413c7c 100755 --- a/examples/atom_4p4d_glm5.2/engine/common.sh +++ b/examples/atom_4p4d_glm5.2/engine/common.sh @@ -34,10 +34,10 @@ CONTAINERS=() for kit_role in $KIT_ROLES; do CONTAINERS+=("$CTR-$kit_role"); done # Settings up.sh records per run, so later scripts report what actually runs. -RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS DECODE_GRAPH_STEP PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL +RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL PREFILL_MAX_BATCHED_TOKENS DECODE_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS PREFILL_CPU_TIER STORE_OWNER_GB STORE_OWNER_DEVICES PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS - DECODE_EXTRA_ARGS DECODE_EXTRA_ENV ROUTER_EXTRA_ARGS IMAGE MODEL) + PREFILL_EXTRA_ENV DECODE_EXTRA_ARGS DECODE_EXTRA_ENV ROUTER_EXTRA_ARGS IMAGE MODEL) # resolve_model: check MODEL exists here and bind its resolved path; a symlink # across an NFS mount boundary binds an empty directory into the container. diff --git a/examples/atom_4p4d_glm5.2/engine/store.sh b/examples/atom_4p4d_glm5.2/engine/store.sh index 569a4232d..842dcad49 100755 --- a/examples/atom_4p4d_glm5.2/engine/store.sh +++ b/examples/atom_4p4d_glm5.2/engine/store.sh @@ -5,13 +5,26 @@ # Start the Mooncake Store the prefill offloads its KV to (PREFILL_CPU_TIER=store), # laid out as InferenceX PR #4045's 1P1D: a master on NUMA node 1 and two owners, # owner i holding STORE_OWNER_GB of NUMA node i's DRAM, both registered on the -# prefill's rails. Waits until both owners serve. Normally called by up.sh. +# prefill's rails. Checks THP and the free memory of both NUMA nodes first, and +# waits until both owners serve. Normally called by up.sh. # # Usage: bash engine/store.sh KIT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # shellcheck source=common.sh source "$KIT/engine/common.sh" +# Each owner faults in its segment as 2 MiB pages, which needs THP and that much +# free memory on its NUMA node (gotcha 14): fail now, not minutes into the mount. +thp="$(on_docker_host 'cat /sys/kernel/mm/transparent_hugepage/enabled')" || + die "cannot read the THP setting on the docker host" +[[ "$thp" != *"[never]"* ]] || die "THP is off on the docker host ($thp); the Store owners need madvise or always" +for i in 0 1; do + free_gb="$(on_docker_host "awk '/MemFree/ {print int(\$4 / 1048576)}' /sys/devices/system/node/node$i/meminfo")" || + die "cannot read the free memory of NUMA node $i on the docker host" + ((free_gb >= STORE_OWNER_GB)) || + die "NUMA node $i has $free_gb GiB free, its Store owner takes STORE_OWNER_GB=$STORE_OWNER_GB: drop the page cache or lower STORE_OWNER_GB (README gotcha 14)" +done + # Bind a process's CPUs and memory to one NUMA node; set_mempolicy needs SYS_NICE. numa_exec=/app/ATOM/.github/scripts/atomesh/numa_exec.py # No GPU: a Mooncake process that sees one probes every buffer through HIP. diff --git a/examples/atom_4p4d_glm5.2/engine/up.sh b/examples/atom_4p4d_glm5.2/engine/up.sh index 64ba313ec..d40c21270 100755 --- a/examples/atom_4p4d_glm5.2/engine/up.sh +++ b/examples/atom_4p4d_glm5.2/engine/up.sh @@ -24,6 +24,10 @@ case "$PREFILL_CPU_TIER" in store | lmcache | off) ;; *) die "PREFILL_CPU_TIER must be store, lmcache or off, not $PREFILL_CPU_TIER" ;; esac +# A full decode step, 2 x CONC sequences of MTP_K + 1 tokens, must fit one +# batch; otherwise every step is cut short of max-num-seqs. +(((MTP_K + 1) * 2 * CONC <= DECODE_MAX_BATCHED_TOKENS)) || + die "DECODE_MAX_BATCHED_TOKENS=$DECODE_MAX_BATCHED_TOKENS is below (MTP_K + 1) x 2 x CONC = $(((MTP_K + 1) * 2 * CONC))" for gpus in "$PREFILL_GPUS" "$DECODE_GPUS"; do [[ "$gpus" =~ ^[0-9]+(,[0-9]+){3}$ ]] || die "each leg needs exactly four GPUs, got $gpus" done diff --git a/examples/atom_4p4d_glm5.2/engine/worker.sh b/examples/atom_4p4d_glm5.2/engine/worker.sh index 1a95308cf..634be2246 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -46,11 +46,10 @@ env=( TRITON_CACHE_DIR=/root/.cache/triton TORCHINDUCTOR_CACHE_DIR=/root/.cache/inductor ) -# Graph capture bounded by the concurrency: 1, 2, 4, 8, then every -# DECODE_GRAPH_STEP up to 2 x CONC. Left unset, a DCP prefill would capture up -# to batch 512; a PP prefill runs eager. +# Graph capture bounded by the concurrency: 1, 2, 4, 8, then every 4 up to 2 x CONC. +# Left unset, a DCP prefill would capture up to batch 512; a PP prefill runs eager. cudagraph="[1,2,4,8" -for ((size = 8 + DECODE_GRAPH_STEP; size <= 2 * CONC; size += DECODE_GRAPH_STEP)); do cudagraph+=",$size"; done +for ((size = 12; size <= 2 * CONC; size += 4)); do cudagraph+=",$size"; done cudagraph+="]" # kv_mooncake KV_ROLE HANDSHAKE_PORT HTTP_PORT @@ -60,7 +59,7 @@ kv_mooncake() { } if [[ "$role" == prefill ]]; then - gpus="$PREFILL_GPUS" port="$PREFILL_PORT" util="$PREFILL_MEM_UTIL" + gpus="$PREFILL_GPUS" port="$PREFILL_PORT" util="$PREFILL_MEM_UTIL" mbt="$PREFILL_MAX_BATCHED_TOKENS" env+=("INFERA_READINESS_PORT=$PREFILL_READINESS_PORT" OFFLOAD_PROFILE=0 # The fabric routes same-numbered rails only: write on the rail of the # decode rank being served. @@ -110,22 +109,23 @@ if [[ "$role" == prefill ]]; then else args=(-tp 4 --decode-context-parallel-size 4 --cudagraph-capture-sizes "$cudagraph") fi - args+=(--max-num-batched-tokens "$PREFILL_MAX_BATCHED_TOKENS" --max-num-seqs 512) + args+=(--max-num-batched-tokens "$mbt" --max-num-seqs 512) read -ra extra <<<"$PREFILL_EXTRA_ARGS" + read -ra extra_env <<<"$PREFILL_EXTRA_ENV" else - gpus="$DECODE_GPUS" port="$DECODE_PORT" util="$DECODE_MEM_UTIL" + gpus="$DECODE_GPUS" port="$DECODE_PORT" util="$DECODE_MEM_UTIL" mbt="$DECODE_MAX_BATCHED_TOKENS" env+=("INFERA_READINESS_PORT=$DECODE_READINESS_PORT") kv="$(kv_mooncake kv_consumer "$DECODE_HANDSHAKE_PORT" "$DECODE_PORT")" # DCP and EP groups reuse the TP group's communicators (same GPUs) instead # of allocating their own buffers, which leaves that memory to the KV cache. env+=(AITER_REUSE_IDENTICAL_COMM_GROUPS=1) args=(-tp 4 --decode-context-parallel-size 4 --cudagraph-capture-sizes "$cudagraph" - --max-num-batched-tokens "$DECODE_MAX_BATCHED_TOKENS" --max-num-seqs $((2 * CONC))) + --max-num-batched-tokens "$mbt" --max-num-seqs $((2 * CONC))) [[ -z "$MTP_AL" ]] || args+=(--spec-decode-acceptance-length "$MTP_AL") read -ra extra <<<"$DECODE_EXTRA_ARGS" read -ra extra_env <<<"$DECODE_EXTRA_ENV" - env+=("${extra_env[@]}") fi +env+=("${extra_env[@]}") # The Rust router keeps idle connections pooled; ATOM's 5 s default closes them under it. args+=(--timeout-keep-alive 900) @@ -137,9 +137,11 @@ mounts=(-v "$MODEL:$MODEL:ro") if on_docker_host "test -e '$HOST_LIBIONIC'"; then mounts+=(-v "$HOST_LIBIONIC:/host-libionic/libionic.so:ro") fi -# JIT caches persist per role and image in a named volume. +# JIT caches persist per role and image in a named volume. ATOM's compile cache +# in it does not key on --max-num-batched-tokens (a stale one fails the prefill +# in assert_size_stride), so the volume name carries that too. image_id="$(docker image inspect -f '{{.Id}}' "$IMAGE")" -mounts+=(-v "$CTR-cache-$role-${image_id:7:12}:/root/.cache") +mounts+=(-v "$CTR-cache-$role-${image_id:7:12}-mbt$mbt:/root/.cache") # AITER builds its native modules inside its package, not under /root/.cache. # A new named volume starts as a copy of the image's directory. mounts+=(-v "$CTR-aiter-jit-$role-${image_id:7:12}:/app/aiter-test/aiter/jit") diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index a894b9abd..6cc42a800 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -54,11 +54,6 @@ export DECODE_GPUS="${DECODE_GPUS:-4,5,6,7}" # 2 x CONC (also of a DCP prefill); restart the engines to change it. export CONC="${CONC:-48}" -# Spacing of the CUDA-graph batch sizes above 8. A batch runs in the next -# larger graph, so a finer step pads fewer rows per decode step at the cost of -# more graphs (capture time and graph memory). -export DECODE_GRAPH_STEP="${DECODE_GRAPH_STEP:-4}" - # MTP draft depth: K3 at CONC >= 48, K4 below, as in the reference runs. export MTP_K="${MTP_K:-$((CONC >= 48 ? 3 : 4))}" @@ -76,16 +71,15 @@ export MTP_EXPERTS="${MTP_EXPERTS:-mxfp4}" # Memory, chunk budgets, block size, layer split and offload load length below # follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md in # the image) and InferenceX PR #4045. The decode's 0.94 and 2048, with the -# shared communicator groups worker.sh sets, give it ~25% more KV blocks, which -# keeps the agents' prefixes on the decode up to about C100 (README "The recipe"). +# shared communicator groups worker.sh sets, give it ~27% more KV blocks, which +# keeps the agents' prefixes on the decode up to about C105 (README "The recipe"). export PREFILL_MEM_UTIL="${PREFILL_MEM_UTIL:-0.85}" export DECODE_MEM_UTIL="${DECODE_MEM_UTIL:-0.94}" # Chunk budgets (--max-num-batched-tokens). ATOM's compile cache -# (/root/.cache/atom, kept in the JIT cache volumes) does not key on the -# prefill's: after a change, start with fresh volumes (another CTR), or the -# prefill fails at startup in assert_size_stride. MTP decode batches stay far -# below 2048 tokens. +# (/root/.cache/atom) does not key on them, so worker.sh keeps it in a volume per +# value (README gotcha 11). A full decode step takes (MTP_K + 1) x 2 x CONC +# tokens, 896 at CONC 112; up.sh refuses a decode budget below that. export PREFILL_MAX_BATCHED_TOKENS="${PREFILL_MAX_BATCHED_TOKENS:-8192}" export DECODE_MAX_BATCHED_TOKENS="${DECODE_MAX_BATCHED_TOKENS:-2048}" @@ -138,9 +132,12 @@ export ROUTER_BACKEND="${ROUTER_BACKEND:-rust}" # Extra flags, space-separated, appended last. export PREFILL_EXTRA_ARGS="${PREFILL_EXTRA_ARGS:-}" export DECODE_EXTRA_ARGS="${DECODE_EXTRA_ARGS:-}" -# Extra decode container environment, space-separated KEY=VALUE pairs. -export DECODE_EXTRA_ENV="${DECODE_EXTRA_ENV:-}" export ROUTER_EXTRA_ARGS="${ROUTER_EXTRA_ARGS:-}" +# Extra engine container environment, space-separated KEY=VALUE pairs, appended +# last; e.g. ATOM_METRICS_PER_REQUEST_CONTEXT=1 for ATOM's per-request context +# gauges, which the image leaves off (README "Image"). +export PREFILL_EXTRA_ENV="${PREFILL_EXTRA_ENV:-}" +export DECODE_EXTRA_ENV="${DECODE_EXTRA_ENV:-}" # --------------------------------------------------------------------------- # 3. Ports and names diff --git a/infera/router/disagg_protocols/atom_mooncake.py b/infera/router/disagg_protocols/atom_mooncake.py index f3aed100c..96bb5e2b7 100644 --- a/infera/router/disagg_protocols/atom_mooncake.py +++ b/infera/router/disagg_protocols/atom_mooncake.py @@ -162,9 +162,10 @@ def extract_handoff(self, prefill_response_payload: dict) -> dict: # ATOM decode takes kv_transfer_params['prompt_token_ids'] as the # prompt (and drops it before the engine), skipping its own template # rendering and tokenization: ~0.25 s per request at 130K tokens. A - # prefill that did not return them leaves decode tokenizing as before. + # prefill that did not return them leaves decode tokenizing as before; + # so does an empty list or a negative id, which decode would reject. ids = prompt_from_chunk(prefill_response_payload) - if ids is not None: + if ids and min(ids) >= 0: handoff["prompt_token_ids"] = ids return handoff diff --git a/infera/router/token_ids.py b/infera/router/token_ids.py index 83eacfc7f..4c10886b0 100644 --- a/infera/router/token_ids.py +++ b/infera/router/token_ids.py @@ -106,8 +106,8 @@ def supports_streaming_ids(engine: EngineType | None, path: str) -> bool: chat request into a 400. Its completions endpoint is fine, and vLLM supports both. - ATOM is left out: it is not known to accept the field, and an engine that - errors on an unknown parameter would fail requests that work today. + ATOM is left out: it rejects the field on a streaming request. Its PD + prefill leg asks for the ids without streaming (atom_mooncake). """ if engine == EngineType.VLLM: return True diff --git a/manual/examples/atom_4p4d_glm5.2.md b/manual/examples/atom_4p4d_glm5.2.md index 1befe8d4f..f2600187f 100644 --- a/manual/examples/atom_4p4d_glm5.2.md +++ b/manual/examples/atom_4p4d_glm5.2.md @@ -1,9 +1,9 @@ # Example: ATOM 4P4D — GLM-5.2-MXFP4 on one MI355X node A PD-disaggregated GLM-5.2-MXFP4 deployment on a single 8×MI355X node: one -**TP1 × PP4 prefill** (GPUs 0-3, ATOM's GLM-5.2 AgentX recipe) with an LMCache -CPU KV tier, one **TP4 + DCP4 decode** (GPUs 4-7), KV moved by Mooncake over -the ionic RDMA rails, and the Infera Python router (round-robin with +**TP1 × PP4 prefill** (GPUs 0-3, ATOM's GLM-5.2 AgentX recipe) with a Mooncake +Store CPU KV tier, one **TP4 + DCP4 decode** (GPUs 4-7), KV moved by Mooncake +over the ionic RDMA rails, and the Infera Rust router (round-robin with `--session-affinity prefill`). `PREFILL_PP=1` runs a TP4 + DCP4 prefill instead. The runnable kit is [`examples/atom_4p4d_glm5.2/`](https://github.com/AMD-AGI/Infera/tree/main/examples/atom_4p4d_glm5.2); @@ -13,9 +13,9 @@ For the concepts, read [PD disaggregation](../features/pd_disaggregation.md) fir ## 1. Build the image -The kit needs the ATOM nightly `661ef766` base plus an ATOM DCP fix and an -LMCache fix, enabled by one opt-in build argument of `Dockerfile.atom`. From the -repo root: +The kit needs the ATOM nightly `661ef766` base plus Infera's ATOM, aiter and +LMCache patches for it, enabled by one opt-in build argument of +`Dockerfile.atom`. From the repo root: ```bash docker build -f deploy/docker/Dockerfile.atom \ @@ -28,7 +28,9 @@ docker build -f deploy/docker/Dockerfile.atom \ `ATOM_NIGHTLY=1` applies `deploy/docker/patches/atom_nightly/`: the ATOM sparse DCP padded-row fix (without it the prefill hits a `Memory access fault` under long-prefix load), the LMCache ROCm pinned-pool fix (the pool becomes GTT, kit -README gotcha 4), an AITER kernel prebuild and a `protobuf<7` pin. +README gotcha 4), the decode patches (ATOM API-process GC, DCP exchanges on +aiter's custom IPC kernels, tuned MLA bmm configs and FP8 GEMM rows), an AITER +kernel prebuild and a `protobuf<7` pin. `BUILD_MOONCAKE=0` keeps the base's Mooncake 0.3.13; the kit sets its `MC_DISABLE_HIP=1` so same-host P/D on disjoint GPUs goes over RDMA. @@ -45,13 +47,16 @@ README gotcha 4), an AITER kernel prebuild and a `protobuf<7` pin. | `WORK_DIR` | host scratch directory for runs and caches | Check `MOONCAKE_DISABLE_HIP_DMABUF` for your driver (`1` = peer-memory) and -that the ports in section 3 are free. +that the ports in section 3 are free. The Store needs THP (`madvise` or +`always`) and `STORE_OWNER_GB` (768 GiB) free on each of the two NUMA nodes; +on smaller hosts, or hosts where its startup probe fails, set +`PREFILL_CPU_TIER=lmcache` (kit README gotcha 14). ## 3. Bring it up and check it ```bash cd examples/atom_4p4d_glm5.2 -bash engine/up.sh # etcd -> prefill + decode -> router; ~10 min cold start +bash engine/up.sh # etcd -> Store -> prefill + decode -> router; ~10 min cold start bash engine/smoke.sh # PD registration, completion, chat, stream, 66k-token prefix reuse bash engine/eval.sh # GSM8K, 5-shot, 200 questions ``` @@ -67,13 +72,15 @@ gives wrong text and is for performance runs only: ```bash bash engine/down.sh -CONC=64 MTP_AL=2.99 bash engine/up.sh +CONC=96 MTP_AL=2.99 SESSION_AFFINITY=off ROUTER_EXTRA_ARGS="--breaker-failure-threshold 0" \ + PREFILL_EXTRA_ARGS="--spec-decode-acceptance-length 2.99" bash engine/up.sh bash engine/agentx/agentx.sh # InferenceX 918524ff, 10 warm-up requests per lane, 3600 s ``` -With this image, CONC 64 reached 33390 input tokens/s/GPU and CONC 96 30116, -8.6% and 44.1% above the DCP prefill on the same node. The kit README lists -both, and the DCP prefill at C48 and C64. +With the kit's defaults, C56, C72, C96 and C112 reached 33578, 40844, 47444 and +48859 input tokens/s/GPU (InferenceX CI's AgentX client). The kit README lists +them, the LMCache tier, and the PP and DCP prefills under the previous +defaults. ## 5. Tear down diff --git a/tests/unit/router/test_atom_mooncake_protocol.py b/tests/unit/router/test_atom_mooncake_protocol.py index 2e5febab9..98b44eaa2 100644 --- a/tests/unit/router/test_atom_mooncake_protocol.py +++ b/tests/unit/router/test_atom_mooncake_protocol.py @@ -109,8 +109,14 @@ def test_prompt_ids_flow_from_prefill_to_decode(self): }, ): handoff = self.proto.extract_handoff(payload) - out = self.proto.annotate_decode(body, self.p, self.d, room_id=1, prefill_handoff=handoff) + out = self.proto.annotate_decode( + body, self.p, self.d, room_id=1, prefill_handoff=handoff + ) assert out["kv_transfer_params"]["prompt_token_ids"] == [5, 6, 7] + # Anything but a non-empty list of ids is left out, as the Rust router does. + for ids in ([], ["a"], [1, -2], "1 2"): + payload = {"kv_transfer_params": self._prefill_response_kv(), "prompt_token_ids": ids} + assert "prompt_token_ids" not in self.proto.extract_handoff(payload) def test_extract_handoff_missing_kv_raises(self): with pytest.raises(ValueError, match="kv_transfer_params"): From 17270a0c1319be3f926b974f6f2bc5d907e9c6c7 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Wed, 21 Oct 2026 07:52:32 +0000 Subject: [PATCH 15/20] revert(examples): drop the MXFP4 MTP experts knob InferenceX does not allow MXFP4 for the MTP draft, so the MTP layer 78 stays BF16 as in the checkpoint and MTP_EXPERTS goes away. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/README.md | 8 +++----- examples/atom_4p4d_glm5.2/engine/common.sh | 2 +- examples/atom_4p4d_glm5.2/engine/worker.sh | 15 +++------------ examples/atom_4p4d_glm5.2/env.sh | 8 +------- 4 files changed, 8 insertions(+), 25 deletions(-) diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index a82ec90de..2491fff41 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -81,20 +81,18 @@ settings actually started); the other scripts read it and write `smoke/`, | max-num-batched-tokens | 8192 | 2048 | | mem util | 0.85 | 0.94, communicator groups shared (`AITER_REUSE_IDENTICAL_COMM_GROUPS=1`) | | KV connector | `multi`: Mooncake `kv_producer` (16 send workers per stage) + `mooncake_store` offload | Mooncake `kv_consumer` | -| shared | FP8 KV, block 16, prefix caching, MTP (`--method mtp`), online PTPC-FP8 with MXFP4 experts (the MTP layer's quantized online, `MTP_EXPERTS`) and the rest of the MTP layer in BF16 | | +| shared | FP8 KV, block 16, prefix caching, MTP (`--method mtp`), online PTPC-FP8 with MXFP4 experts and BF16 MTP layer | | The settings are those of ATOM's GLM-5.2 AgentX recipe (`/app/ATOM/recipes/Agentic-GLM-5.2.md` in the image) and of InferenceX PR #4045's 1P1D (`MC_NUM_QP_PER_EP=1` on every Mooncake peer, `TOPK_FORCE_PATH=one`, `--timeout-keep-alive 900`, the Store). The decode's memory settings give it -~203K KV blocks (~196K with BF16 MTP experts) instead of ~160K at -0.90 / 16384 with separate communicators: AgentX lanes reuse their whole +~196K KV blocks instead of ~160K at 0.90 / 16384 with separate communicators: AgentX lanes reuse their whole context every turn, and once C × mean prompt length outgrows the decode's KV pool its prefix hits collapse, the prefill re-sends whole prompts and TTFT explodes (about C85 before, C105 now). Not taken from them: `OFFLOAD_PROFILE=1` (logging only), the decode's graphs (here every 4 up to 2 × CONC, there every 8 -up to 256), its max-num-seqs (here 2 × CONC, there 512) and BF16 MTP experts -(`MTP_EXPERTS`). +up to 256) and its max-num-seqs (here 2 × CONC, there 512). MTP depth follows `CONC`: K3 at CONC ≥ 48, K4 below. The Store (`PREFILL_CPU_TIER=store`, gotcha 14) is one pool of 2 × 768 GiB that the four diff --git a/examples/atom_4p4d_glm5.2/engine/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index 8af413c7c..73bf46d5b 100755 --- a/examples/atom_4p4d_glm5.2/engine/common.sh +++ b/examples/atom_4p4d_glm5.2/engine/common.sh @@ -34,7 +34,7 @@ CONTAINERS=() for kit_role in $KIT_ROLES; do CONTAINERS+=("$CTR-$kit_role"); done # Settings up.sh records per run, so later scripts report what actually runs. -RUN_VARS=(CONC MTP_K MTP_AL MTP_EXPERTS PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL +RUN_VARS=(CONC MTP_K MTP_AL PREFILL_GPUS DECODE_GPUS PREFILL_NUM_WORKERS PREFILL_MEM_UTIL DECODE_MEM_UTIL PREFILL_MAX_BATCHED_TOKENS DECODE_MAX_BATCHED_TOKENS BLOCK_SIZE PREFILL_PP PREFILL_PP_PARTITION PREFILL_SEND_WORKERS PREFILL_CPU_TIER STORE_OWNER_GB STORE_OWNER_DEVICES PREFILL_OFFLOAD_GB OFFLOAD_MIN_LOAD_TOKENS MOONCAKE_DISABLE_HIP_DMABUF SESSION_AFFINITY ROUTER_BACKEND PREFILL_EXTRA_ARGS PREFILL_EXTRA_ENV DECODE_EXTRA_ARGS DECODE_EXTRA_ENV ROUTER_EXTRA_ARGS IMAGE MODEL) diff --git a/examples/atom_4p4d_glm5.2/engine/worker.sh b/examples/atom_4p4d_glm5.2/engine/worker.sh index 634be2246..95f25b9e6 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -15,18 +15,9 @@ case "$role" in prefill | decode) ;; *) die "usage: bash engine/worker.sh prefil case "$PREFILL_PP" in 1 | 4) ;; *) die "PREFILL_PP must be 1 or 4, not $PREFILL_PP" ;; esac resolve_model -# Experts of layers 0-77 stay MXFP4 and the rest is quantized online to -# per-token FP8, except the MTP layer 78: BF16 as in the checkpoint, or with -# MTP_EXPERTS=mxfp4 its routed experts quantized online to MXFP4 at load. -case "$MTP_EXPERTS" in - bf16) mtp_quant="" mtp_exclude='"model.layers.78.*"' ;; - mxfp4) - mtp_quant=',"layer_quant_config":{"model.layers.78.mlp.experts*":"mxfp4"}' - mtp_exclude='"re:^model\\.layers\\.78\\.(?!mlp\\.experts)"' - ;; - *) die "MTP_EXPERTS must be bf16 or mxfp4, not $MTP_EXPERTS" ;; -esac -QUANT_CONFIG='{"global_quant_config":"ptpc_fp8"'"$mtp_quant"',"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*",'"$mtp_exclude"']}' +# Experts of layers 0-77 stay MXFP4 and the MTP layer 78 stays BF16; the rest +# is quantized online to per-token FP8. +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.*"]}' env=( AITER_QUICK_REDUCE_QUANTIZATION=INT4 AITER_USE_FLYDSL_MOE_SORTING=1 AITER_LOG_LEVEL=WARNING diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index 6cc42a800..dcb28b75a 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -62,16 +62,10 @@ export MTP_K="${MTP_K:-$((CONC >= 48 ? 3 : 4))}" # Set it (K3 -> 2.99, K4 -> 3.33) only to reproduce the AgentX reference numbers. export MTP_AL="${MTP_AL:-}" -# Routed experts of the MTP layer (78), which the checkpoint keeps in BF16: -# mxfp4 quantizes them online at load like the experts of the other layers, so -# the draft MoE reads a quarter of the weight bytes (~1 ms per decode step; -# GSM8K and real MTP acceptance unchanged); bf16 keeps them as shipped. -export MTP_EXPERTS="${MTP_EXPERTS:-mxfp4}" - # Memory, chunk budgets, block size, layer split and offload load length below # follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md in # the image) and InferenceX PR #4045. The decode's 0.94 and 2048, with the -# shared communicator groups worker.sh sets, give it ~27% more KV blocks, which +# shared communicator groups worker.sh sets, give it ~22% more KV blocks, which # keeps the agents' prefixes on the decode up to about C105 (README "The recipe"). export PREFILL_MEM_UTIL="${PREFILL_MEM_UTIL:-0.85}" export DECODE_MEM_UTIL="${DECODE_MEM_UTIL:-0.94}" From b239e797235d8108820ce21919c59a550ce3246b Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Wed, 21 Oct 2026 18:00:00 +0000 Subject: [PATCH 16/20] docs(examples): reference numbers with the BF16 MTP layer C56-C112 re-run on the Store after dropping the MXFP4 MTP experts: 1.4-2.5% below the MXFP4 runs (longer decode steps, ~198K instead of ~203K KV blocks) and still 5.0-23.8% above InferenceX PR #4045. GSM8K 0.975 with real acceptance. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- examples/atom_4p4d_glm5.2/README.md | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index 2491fff41..576b1d1fb 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -128,10 +128,10 @@ been run with the Store or with the current decode settings. The image above with the defaults of `env.sh`, on 8×MI355X nodes: -| item | PP prefill, Store (default) | PP prefill, `PREFILL_CPU_TIER=lmcache` | +| item | PP prefill, Store (default) | PP prefill, `PREFILL_CPU_TIER=lmcache`, MXFP4 MTP experts (earlier image) | |---|---|---| | correctness (`smoke.py`, real MTP acceptance) | **passed**, 6/6 | **passed**, 6/6 | -| GSM8K (200 questions) | **0.970** | **0.985** | +| GSM8K (200 questions) | **0.975** | **0.985** | | AgentX, 3600 s | **passed** at C56, C72, C96, C112, 0 errors | **passed** at C56, C72, C96, C112, 0 errors | The DCP prefill passed smoke 6/6, GSM8K 0.975 and AgentX at C48, C64 and C96 @@ -149,22 +149,27 @@ warm-up, 3600 s measured, forced MTP acceptance on both legs (AL 2.99 with K3), PR #4045's runs. The client was InferenceX CI's (InferenceX `288471b3`, agentx-harness 1.0.11), not `agentx.sh`'s `918524ff`. Throughput is per GPU over all 8 GPUs and counts cache-hit input tokens; interactivity is the inverse -of P90 inter-token latency. C56 and C96 ran on one node, C72 and C112 on -another. +of P90 inter-token latency. C72 and C112 ran on one node, C56 and C96 on two +others. | CONC | input tok/s/GPU | output tok/s/GPU | P90 interactivity tok/s/user | TTFT mean / p50 / p90 s | ITL p50 / p90 ms | prefill hit GPU / CPU / total % | decode hit % | requests / errors | |---:|---:|---:|---:|---|---|---|---:|---| -| 56 | 33578 | 249.0 | 66.1 | 1.12 / 0.84 / 2.00 | 13.62 / 15.12 | 86.0 / 10.2 / 96.2 | 92.6 | 7117 / 0 | -| 72 | 40844 | 327.4 | 53.6 | 1.51 / 1.05 / 2.68 | 16.11 / 18.64 | 76.4 / 18.5 / 94.9 | 85.5 | 9627 / 0 | -| 96 | 47444 | 393.4 | 46.8 | 2.42 / 1.50 / 4.72 | 19.25 / 21.35 | 65.6 / 28.2 / 93.8 | 60.4 | 11482 / 0 | -| 112 | **48859** | 416.6 | 43.9 | 4.42 / 2.86 / 9.19 | 20.52 / 22.76 | 59.9 / 33.1 / 93.0 | 32.1 | 12062 / 0 | +| 56 | 32925 | 246.2 | 64.0 | 1.12 / 0.84 / 1.96 | 14.17 / 15.62 | 84.1 / 11.9 / 96.0 | 92.3 | 7049 / 0 | +| 72 | 40075 | 321.9 | 52.1 | 1.50 / 1.05 / 2.56 | 16.75 / 19.21 | 72.7 / 21.7 / 94.4 | 84.1 | 9477 / 0 | +| 96 | 46234 | 387.3 | 45.6 | 2.68 / 1.66 / 5.47 | 19.88 / 21.94 | 62.4 / 31.2 / 93.5 | 52.9 | 11291 / 0 | +| 112 | **48154** | 406.0 | 42.1 | 5.05 / 3.64 / 9.63 | 21.21 / 23.73 | 55.5 / 36.6 / 92.1 | 25.1 | 12027 / 0 | The decode hit falls from C96 on, as C × prompt length outgrows the decode's KV pool; the Store's RDMA loads do not block the prefill, so throughput still -rises to C112. With `PREFILL_CPU_TIER=lmcache` (other nodes) C56 and C72 match -(33392 / 40326), but C96 and C112 drop to 30690 / 23321: LMCache loads +rises to C112. With `PREFILL_CPU_TIER=lmcache` (other nodes, and the MTP +layer's experts then quantized to MXFP4) C56 and C72 matched the Store +(33392 / 40326), but C96 and C112 dropped to 30690 / 23321: LMCache loads synchronously and takes about half of each PP stage's time there. +Those MXFP4 MTP experts gave the Store 1.4–2.5% more at each point +(33578 / 40844 / 47444 / 48859): shorter decode steps, and ~203K KV blocks +instead of ~198K. + The previous defaults (Python router, decode at 0.90 / 16384 with separate communicators, 4 send workers, LMCache, an image without the decode patches) compared the two prefills on one node, with `agentx.sh`: From 3424bf4da87bc3becb46b9fc87c36d243912aca1 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Thu, 22 Oct 2026 04:00:43 +0000 Subject: [PATCH 17/20] fix(examples): refuse the DCP prefill with the Store; label Store runs as Mooncake offload Second round of review follow-ups at 41bd647a. 4P4D kit: - up.sh refuses PREFILL_PP=1 with PREFILL_CPU_TIER=store. ATOM's Mooncake Store offload rejects decode context parallelism at startup (kv_transfer/offload/mooncake_store/config.py), so the DCP prefill died after store.sh had mounted 2 x 768 GiB. env.sh, the README and the manual now say the DCP prefill takes the LMCache tier. - agentx.sh writes the AgentX offload metadata from PREFILL_CPU_TIER. It only looked at PREFILL_OFFLOAD_GB, so Store runs (the default) were labelled LMCache and PREFILL_CPU_TIER=off runs DRAM offload. The Store is reported as KV_OFFLOAD_BACKEND=mooncake with mooncake_master's version, as in InferenceX PR #4045's results. - Numbers after dropping the MXFP4 MTP experts: ~198.5K KV blocks (+24%), the decode-prefix cliff at about C100, and the manual's C56-C112 throughput (32925 / 40075 / 46234 / 48154). The README no longer quotes the MXFP4 MTP runs; the LMCache comparison is labelled as an earlier image and kit. Patches: check_dcp_merge.py runs on 4 or 8 GPUs only. On 2 GPUs a rank merges 32 of GLM-5.2's 64 heads, more than the one-kernel merge takes, so its merge cases failed by construction. Router: comments say ATOM returns the prompt ids at the top level for completions too. Checked: up.sh dry run with a stub docker (PREFILL_PP=1 with the Store refused; with lmcache or off, and PREFILL_PP=4, past the check); the agentx.sh offload lines for every tier pass InferenceX 918524ff's aggregation checks; check_dcp_merge.py exits on 2 ranks before touching a GPU; pre-commit on the changed files, shellcheck, cargo test and clippy, pytest of the touched router tests. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- deploy/docker/patches/atom_nightly/README.md | 4 +- .../patches/atom_nightly/check_dcp_merge.py | 7 +++- examples/atom_4p4d_glm5.2/README.md | 37 +++++++++---------- .../atom_4p4d_glm5.2/engine/agentx/agentx.sh | 20 +++++++--- examples/atom_4p4d_glm5.2/engine/up.sh | 3 ++ examples/atom_4p4d_glm5.2/env.sh | 13 ++++--- manual/examples/atom_4p4d_glm5.2.md | 9 +++-- rust/router/src/protocol.rs | 7 ++-- .../router/test_atom_mooncake_protocol.py | 4 +- 9 files changed, 62 insertions(+), 42 deletions(-) diff --git a/deploy/docker/patches/atom_nightly/README.md b/deploy/docker/patches/atom_nightly/README.md index fd1b95a7e..fd962a3a3 100644 --- a/deploy/docker/patches/atom_nightly/README.md +++ b/deploy/docker/patches/atom_nightly/README.md @@ -51,7 +51,9 @@ all of them. Set it where those series are needed, such as profiling runs. The anchors only show that the source text did not drift, not that the two aiter kernels still agree with RCCL once aiter's sync primitives or buffer registration change. After a base move, run `check_dcp_merge.py` in the new -image on 2, 4 or 8 GPUs; it exits 1 on any mismatch: +image on 4 or 8 GPUs (its merge cases use GLM-5.2's 64 heads, which on 2 GPUs +leave each rank more heads than the merge kernel takes); it exits 1 on any +mismatch: ```bash docker run --rm --network host --ipc host --device /dev/kfd --device /dev/dri \ diff --git a/deploy/docker/patches/atom_nightly/check_dcp_merge.py b/deploy/docker/patches/atom_nightly/check_dcp_merge.py index 3d87032e3..af8497ac3 100644 --- a/deploy/docker/patches/atom_nightly/check_dcp_merge.py +++ b/deploy/docker/patches/atom_nightly/check_dcp_merge.py @@ -7,7 +7,7 @@ """GPU check of the DCP merge kernels of patch_aiter_custom_alltoall.py. The strict anchors only prove that the patched source text did not drift. Run -this in the image whenever ATOM_BASE_IMAGE moves, on 2, 4 or 8 GPUs of a node +this in the image whenever ATOM_BASE_IMAGE moves, on 4 or 8 GPUs of a node (see README.md): torchrun --nproc-per-node 4 check_dcp_merge.py @@ -35,6 +35,11 @@ dist.init_process_group("nccl") rank, world = dist.get_rank(), dist.get_world_size() +# On 2 GPUs a rank would merge 32 of the 64 heads, more than the one-kernel +# merge takes (h_local x head_dim <= 4096), so the merge cases could not pass. +if world not in (4, 8): + dist.destroy_process_group() + sys.exit(f"[check_dcp_merge] run on 4 or 8 GPUs, not {world}") torch.cuda.set_device(rank) dev = torch.device(f"cuda:{rank}") diff --git a/examples/atom_4p4d_glm5.2/README.md b/examples/atom_4p4d_glm5.2/README.md index 576b1d1fb..94fa00173 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -13,7 +13,7 @@ One file is site-specific. Fill it in, then: ```bash $EDITOR env.sh # the only file you edit -bash engine/up.sh # etcd -> prefill + decode -> router, waiting on health +bash engine/up.sh # etcd -> Store -> prefill + decode -> router, waiting on health bash engine/smoke.sh # PD registration, completion/chat/stream, 66k-token prefix reuse bash engine/down.sh # save logs, remove containers, wait for VRAM to drain ``` @@ -87,12 +87,13 @@ The settings are those of ATOM's GLM-5.2 AgentX recipe (`/app/ATOM/recipes/Agentic-GLM-5.2.md` in the image) and of InferenceX PR #4045's 1P1D (`MC_NUM_QP_PER_EP=1` on every Mooncake peer, `TOPK_FORCE_PATH=one`, `--timeout-keep-alive 900`, the Store). The decode's memory settings give it -~196K KV blocks instead of ~160K at 0.90 / 16384 with separate communicators: AgentX lanes reuse their whole -context every turn, and once C × mean prompt length outgrows the decode's KV -pool its prefix hits collapse, the prefill re-sends whole prompts and TTFT -explodes (about C85 before, C105 now). Not taken from them: `OFFLOAD_PROFILE=1` -(logging only), the decode's graphs (here every 4 up to 2 × CONC, there every 8 -up to 256) and its max-num-seqs (here 2 × CONC, there 512). +~198.5K KV blocks instead of ~160K at 0.90 / 16384 with separate communicators: +AgentX lanes reuse their whole context every turn, and once C × mean prompt +length outgrows the decode's KV pool its prefix hits collapse, the prefill +re-sends whole prompts and TTFT explodes (about C85 before, C100 now). Not +taken from them: `OFFLOAD_PROFILE=1` (logging only), the decode's graphs (here +every 4 up to 2 × CONC, there every 8 up to 256) and its max-num-seqs (here +2 × CONC, there 512). MTP depth follows `CONC`: K3 at CONC ≥ 48, K4 below. The Store (`PREFILL_CPU_TIER=store`, gotcha 14) is one pool of 2 × 768 GiB that the four @@ -111,8 +112,10 @@ throughput per GPU is over all 8 GPUs. ### The DCP prefill `PREFILL_PP=1` starts the prefill as TP4 + DCP4 with CUDA graphs up to -2 × CONC, as the decode. It was validated with these settings, which are not -the defaults, and with the LMCache tier; set them with it: +2 × CONC, as the decode. ATOM's Mooncake Store offload does not support DCP, so +this prefill takes `PREFILL_CPU_TIER=lmcache` (or `off`); `up.sh` refuses it +with the Store. It was validated with these settings, which are not the +defaults; set them with it: ```bash PREFILL_PP=1 PREFILL_MEM_UTIL=0.90 PREFILL_MAX_BATCHED_TOKENS=16384 BLOCK_SIZE=64 \ @@ -122,13 +125,13 @@ PREFILL_PP=1 PREFILL_MEM_UTIL=0.90 PREFILL_MAX_BATCHED_TOKENS=16384 BLOCK_SIZE=6 Each DCP rank then writes to one decode rank without staging, and the LMCache lookup covers the four DCP ranks, each holding a distinct KV shard. Under AgentX load it is slower than the PP prefill (reference numbers below). It has not -been run with the Store or with the current decode settings. +been run with the current decode settings. ## Validation status The image above with the defaults of `env.sh`, on 8×MI355X nodes: -| item | PP prefill, Store (default) | PP prefill, `PREFILL_CPU_TIER=lmcache`, MXFP4 MTP experts (earlier image) | +| item | PP prefill, Store (default) | PP prefill, `PREFILL_CPU_TIER=lmcache` (earlier image and kit) | |---|---|---| | correctness (`smoke.py`, real MTP acceptance) | **passed**, 6/6 | **passed**, 6/6 | | GSM8K (200 questions) | **0.975** | **0.985** | @@ -161,14 +164,10 @@ others. The decode hit falls from C96 on, as C × prompt length outgrows the decode's KV pool; the Store's RDMA loads do not block the prefill, so throughput still -rises to C112. With `PREFILL_CPU_TIER=lmcache` (other nodes, and the MTP -layer's experts then quantized to MXFP4) C56 and C72 matched the Store -(33392 / 40326), but C96 and C112 dropped to 30690 / 23321: LMCache loads -synchronously and takes about half of each PP stage's time there. - -Those MXFP4 MTP experts gave the Store 1.4–2.5% more at each point -(33578 / 40844 / 47444 / 48859): shorter decode steps, and ~203K KV blocks -instead of ~198K. +rises to C112. With `PREFILL_CPU_TIER=lmcache` (other nodes, an earlier image +and kit) C56 and C72 were on par with the Store (33392 / 40326), but C96 and +C112 dropped to 30690 / 23321: LMCache loads synchronously and takes about half +of each PP stage's time there. The previous defaults (Python router, decode at 0.90 / 16384 with separate communicators, 4 send workers, LMCache, an image without the decode patches) diff --git a/examples/atom_4p4d_glm5.2/engine/agentx/agentx.sh b/examples/atom_4p4d_glm5.2/engine/agentx/agentx.sh index 84092818a..2eb7d84ab 100755 --- a/examples/atom_4p4d_glm5.2/engine/agentx/agentx.sh +++ b/examples/atom_4p4d_glm5.2/engine/agentx/agentx.sh @@ -45,12 +45,20 @@ fi [[ "$(git -C "$ix" rev-parse HEAD)" == "$INFERENCEX_REF" ]] || die "InferenceX checkout $ix is not at $INFERENCEX_REF" -# AgentX metadata: the prefill's LMCache CPU tier is DRAM offload. -offload="KV_OFFLOADING=none" -if ((PREFILL_OFFLOAD_GB > 0)); then - lmcache_version="$(docker run --rm "$IMAGE" pip show lmcache | awk '/^Version:/ {print $2}')" - offload=$'KV_OFFLOADING=dram\nKV_OFFLOAD_BACKEND=lmcache\n' - offload+="KV_OFFLOAD_BACKEND_METADATA={\"name\":\"lmcache\",\"version\":\"${lmcache_version%%+*}\"}" +# AgentX metadata: the prefill's CPU tier (PREFILL_CPU_TIER) is DRAM offload, +# backed by the Mooncake Store or by LMCache. +offload="KV_OFFLOADING=none" backend="" version="" +if [[ "$PREFILL_CPU_TIER" == store ]]; then + backend=mooncake + version="$(docker run --rm "$IMAGE" mooncake_master --version | awk '/version/ {print $3}')" +elif [[ "$PREFILL_CPU_TIER" == lmcache ]] && ((PREFILL_OFFLOAD_GB > 0)); then + backend=lmcache + version="$(docker run --rm "$IMAGE" pip show lmcache | awk '/^Version:/ {print $2}')" +fi +if [[ -n "$backend" ]]; then + [[ -n "$version" ]] || die "cannot read the $backend version from $IMAGE" + offload=$'KV_OFFLOADING=dram\n'"KV_OFFLOAD_BACKEND=$backend"$'\n' + offload+="KV_OFFLOAD_BACKEND_METADATA={\"name\":\"$backend\",\"version\":\"${version%%+*}\"}" fi # The scraped /metrics come from the router's worker list, so workers on # other nodes are included. PREPARE_ONLY has no router to ask: this node's ports. diff --git a/examples/atom_4p4d_glm5.2/engine/up.sh b/examples/atom_4p4d_glm5.2/engine/up.sh index d40c21270..e2d0ff510 100755 --- a/examples/atom_4p4d_glm5.2/engine/up.sh +++ b/examples/atom_4p4d_glm5.2/engine/up.sh @@ -24,6 +24,9 @@ case "$PREFILL_CPU_TIER" in store | lmcache | off) ;; *) die "PREFILL_CPU_TIER must be store, lmcache or off, not $PREFILL_CPU_TIER" ;; esac +# ATOM's Mooncake Store offload refuses decode context parallelism at startup. +[[ "$PREFILL_PP" != 1 || "$PREFILL_CPU_TIER" != store ]] || + die "the DCP prefill (PREFILL_PP=1) cannot use the Store: set PREFILL_CPU_TIER=lmcache or off" # A full decode step, 2 x CONC sequences of MTP_K + 1 tokens, must fit one # batch; otherwise every step is cut short of max-num-seqs. (((MTP_K + 1) * 2 * CONC <= DECODE_MAX_BATCHED_TOKENS)) || diff --git a/examples/atom_4p4d_glm5.2/env.sh b/examples/atom_4p4d_glm5.2/env.sh index dcb28b75a..e41cdd0d0 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -65,8 +65,8 @@ export MTP_AL="${MTP_AL:-}" # Memory, chunk budgets, block size, layer split and offload load length below # follow ATOM's GLM-5.2 AgentX recipe (/app/ATOM/recipes/Agentic-GLM-5.2.md in # the image) and InferenceX PR #4045. The decode's 0.94 and 2048, with the -# shared communicator groups worker.sh sets, give it ~22% more KV blocks, which -# keeps the agents' prefixes on the decode up to about C105 (README "The recipe"). +# shared communicator groups worker.sh sets, give it ~24% more KV blocks, which +# keeps the agents' prefixes on the decode up to about C100 (README "The recipe"). export PREFILL_MEM_UTIL="${PREFILL_MEM_UTIL:-0.85}" export DECODE_MEM_UTIL="${DECODE_MEM_UTIL:-0.94}" @@ -84,9 +84,10 @@ export BLOCK_SIZE="${BLOCK_SIZE:-16}" # Prefill parallelism: 4 = TP1 x PP4, eager, with PREFILL_PP_PARTITION layers # per stage. 1 = TP4 + DCP4 with CUDA graphs, slower under AgentX load (README -# "Reference numbers"); its validated settings were PREFILL_MEM_UTIL=0.90, -# PREFILL_MAX_BATCHED_TOKENS=16384, BLOCK_SIZE=64, OFFLOAD_MIN_LOAD_TOKENS=8192 -# and PREFILL_SEND_WORKERS empty. +# "Reference numbers"); it cannot use the Store (ATOM's Store offload does not +# support DCP), and its validated settings were PREFILL_CPU_TIER=lmcache, +# PREFILL_MEM_UTIL=0.90, PREFILL_MAX_BATCHED_TOKENS=16384, BLOCK_SIZE=64, +# OFFLOAD_MIN_LOAD_TOKENS=8192 and PREFILL_SEND_WORKERS empty. export PREFILL_PP="${PREFILL_PP:-4}" export PREFILL_PP_PARTITION="${PREFILL_PP_PARTITION:-20,20,20,18}" @@ -100,7 +101,7 @@ export PREFILL_SEND_WORKERS="${PREFILL_SEND_WORKERS-}" # store Mooncake Store, as InferenceX PR #4045: one master and two owners # of STORE_OWNER_GB each, one per NUMA node, one pool shared by the # four PP stages, read and written over RDMA (engine/store.sh, -# README gotcha 14) +# README gotcha 14); not with PREFILL_PP=1 # lmcache LMCache, PREFILL_OFFLOAD_GB per PP stage, as GTT (gotcha 4) # off Mooncake P->D only export PREFILL_CPU_TIER="${PREFILL_CPU_TIER:-store}" diff --git a/manual/examples/atom_4p4d_glm5.2.md b/manual/examples/atom_4p4d_glm5.2.md index f2600187f..e927abfd5 100644 --- a/manual/examples/atom_4p4d_glm5.2.md +++ b/manual/examples/atom_4p4d_glm5.2.md @@ -4,8 +4,9 @@ A PD-disaggregated GLM-5.2-MXFP4 deployment on a single 8×MI355X node: one **TP1 × PP4 prefill** (GPUs 0-3, ATOM's GLM-5.2 AgentX recipe) with a Mooncake Store CPU KV tier, one **TP4 + DCP4 decode** (GPUs 4-7), KV moved by Mooncake over the ionic RDMA rails, and the Infera Rust router (round-robin with -`--session-affinity prefill`). `PREFILL_PP=1` runs a TP4 + DCP4 prefill instead. -The runnable kit is +`--session-affinity prefill`). `PREFILL_PP=1` runs a TP4 + DCP4 prefill instead, +with the LMCache tier: ATOM's Store offload does not support DCP. The runnable +kit is [`examples/atom_4p4d_glm5.2/`](https://github.com/AMD-AGI/Infera/tree/main/examples/atom_4p4d_glm5.2); its README carries the full recipe, reference numbers and gotchas. @@ -77,8 +78,8 @@ CONC=96 MTP_AL=2.99 SESSION_AFFINITY=off ROUTER_EXTRA_ARGS="--breaker-failure-th bash engine/agentx/agentx.sh # InferenceX 918524ff, 10 warm-up requests per lane, 3600 s ``` -With the kit's defaults, C56, C72, C96 and C112 reached 33578, 40844, 47444 and -48859 input tokens/s/GPU (InferenceX CI's AgentX client). The kit README lists +With the kit's defaults, C56, C72, C96 and C112 reached 32925, 40075, 46234 and +48154 input tokens/s/GPU (InferenceX CI's AgentX client). The kit README lists them, the LMCache tier, and the PP and DCP prefills under the previous defaults. diff --git a/rust/router/src/protocol.rs b/rust/router/src/protocol.rs index 80523da7f..7c90d6d23 100644 --- a/rust/router/src/protocol.rs +++ b/rust/router/src/protocol.rs @@ -277,8 +277,9 @@ pub fn atom_handoff(prefill_reply: &Value) -> anyhow::Result> Ok(handoff) } -/// The prompt ids of a `return_token_ids` reply: top level for chat -/// completions, on the first choice for completions. Only a list of integers. +/// The prompt ids of a `return_token_ids` reply: ATOM puts them at the top +/// level (chat and completions); the first choice is read too. Only a +/// non-empty list of non-negative integers. fn atom_prompt_ids(reply: &Value) -> Option<&Value> { [ reply.get("prompt_token_ids"), @@ -577,7 +578,7 @@ mod tests { #[test] fn atom_handoff_carries_the_prompt_ids() { let kv = json!({"transfer_id": "t"}); - // Chat completions: top level; completions: the first choice. + // At the top level, as ATOM returns them, or on the first choice. for reply in [ json!({"kv_transfer_params": kv, "prompt_token_ids": [5, 6, 7]}), json!({"kv_transfer_params": kv, "choices": [{"prompt_token_ids": [5, 6, 7]}]}), diff --git a/tests/unit/router/test_atom_mooncake_protocol.py b/tests/unit/router/test_atom_mooncake_protocol.py index 98b44eaa2..269b0c397 100644 --- a/tests/unit/router/test_atom_mooncake_protocol.py +++ b/tests/unit/router/test_atom_mooncake_protocol.py @@ -99,8 +99,8 @@ def test_prompt_ids_flow_from_prefill_to_decode(self): pre = self.proto.annotate_prefill(body, self.p, self.d, room_id=1) assert pre["return_token_ids"] is True assert "return_token_ids" not in body - # ATOM puts them at the top level of a chat response, under the - # choice on completions; both reach decode's kv_transfer_params. + # ATOM puts them at the top level (chat and completions); ids on the + # first choice are taken too. Both reach decode's kv_transfer_params. for payload in ( {"kv_transfer_params": self._prefill_response_kv(), "prompt_token_ids": [5, 6, 7]}, { From eb519a74a0e9415e49c716c10a437d3a73e5f07e Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Thu, 22 Oct 2026 10:35:10 +0000 Subject: [PATCH 18/20] Build libzmq with aligned queue allocation so the router tests stop crashing `cargo test` has died with SIGSEGV on every CI run since ubuntu-latest moved to Ubuntu 26.04. gcc 15 there defaults to -march=x86-64-v3, so libzmq's copy of its 64-byte-aligned command_t becomes a `vmovdqa ymm` store into a queue chunk that zeromq-src allocates with plain malloc (16-byte aligned, because it never defines HAVE_POSIX_MEMALIGN). It faults on the first zmq_connect, i.e. in the kv_event subscriber tests. Ubuntu 24.04's gcc 13 targets plain x86-64 and never emitted the aligned AVX store, which hid the bug. Define HAVE_POSIX_MEMALIGN (and the cache-line size it aligns to) for the x86_64 Linux C++ build through rust/.cargo/config.toml, which `cd rust && cargo build` in the Dockerfiles also picks up. Verified on the ubuntu-26.04 runner: fmt, clippy and the full `cargo test` pass 10 times out of 10 (10/10 crashed before); 24.04 is unchanged. Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- rust/.cargo/config.toml | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 rust/.cargo/config.toml diff --git a/rust/.cargo/config.toml b/rust/.cargo/config.toml new file mode 100644 index 000000000..7383aa812 --- /dev/null +++ b/rust/.cargo/config.toml @@ -0,0 +1,8 @@ +[env] +# libzmq (built from source by zeromq-src) declares command_t aligned(64) but, +# without HAVE_POSIX_MEMALIGN, allocates its queues with plain malloc (16-byte +# aligned). Any compiler that emits AVX for those copies -- gcc 15 on Ubuntu +# 26.04 defaults to -march=x86-64-v3 -- faults on the first zmq_connect. +# ZMQ_CACHELINE_SIZE is the alignment posix_memalign is then asked for; the +# cmake build that normally defines it is bypassed by zeromq-src. +CXXFLAGS_x86_64_unknown_linux_gnu = "-DHAVE_POSIX_MEMALIGN=1 -DZMQ_CACHELINE_SIZE=64" From 3b3b8310e011897c90a42c9bf472adec95153306 Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Thu, 22 Oct 2026 10:40:14 +0000 Subject: [PATCH 19/20] Pin Landlock in the env_mgr tests where bwrap is also installed The agent_sys job has failed 24 env_mgr tests on every CI run since ubuntu-latest moved to Ubuntu 26.04, whose runner image ships /usr/bin/bwrap (24.04 did not). `select` prefers bubblewrap whenever it is on PATH, and bubblewrap confines by exec: `apply` is a no-op in-process, and an ungranted path inside the bwrap'd child lands on bwrap's tmpfs instead of returning EACCES. Every denial those tests assert is Landlock's, so they ran against the wrong mechanism. The sandbox itself is not weakened. - conftest: an autouse fixture reports bwrap as absent to `select` when the kernel also has Landlock; `run_confined` reads the probe at call time so its forked child sees the same answer. `availability` still reports the real machine, and bubblewrap's branch keeps its fake-binary tests. - test_wrap_argv_refuses_when_the_binary_vanished: after deleting the fake bwrap, PATH is only the emptied directory, not the original PATH that may hold a real bwrap. Reproduced in an ubuntu:26.04 container with bubblewrap: the 24 CI failures all pass after this; without bwrap the full suite is unchanged (2436 passed). Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- agent_sys/tests/env_mgr/conftest.py | 28 +++++++++++++++++++++++++ agent_sys/tests/env_mgr/test_prepare.py | 4 +++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/agent_sys/tests/env_mgr/conftest.py b/agent_sys/tests/env_mgr/conftest.py index b8c5e7192..959823618 100644 --- a/agent_sys/tests/env_mgr/conftest.py +++ b/agent_sys/tests/env_mgr/conftest.py @@ -115,6 +115,10 @@ def run_confined(policy: Policy, body: Callable[[], Any], *, tier: Any = None) - outcome if the child dies on a signal. """ from env_mgr.isolation import apply as _apply + + # Imported here, not taken from this module's globals, so the child sees + # the probe `_landlock_over_bwrap` pinned. + from env_mgr.isolation.probe import probe from env_mgr.protocols import Tier # `pickle` here is this process talking to its own fork over an anonymous @@ -225,3 +229,27 @@ def _switch_off_the_kill_switch(monkeypatch: pytest.MonkeyPatch) -> None: applied after this and wins. """ monkeypatch.setenv("AGENT_SYS_NO_PERMISSIONS", "0") + + +@pytest.fixture(autouse=True) +def _landlock_over_bwrap(monkeypatch: pytest.MonkeyPatch) -> None: + """**This directory asserts Landlock, so it gets Landlock where bwrap is also present.** + + `select` prefers bubblewrap whenever it is on PATH, and bubblewrap confines + by exec: `apply` does nothing in-process, and inside a bwrap'd child an + ungranted path is an entry on bwrap's own tmpfs rather than an EACCES. Every + denial asserted here is Landlock's. Measured: GitHub's Ubuntu 26.04 runners + ship `bwrap` (24.04 did not), and 24 tests here failed there against the + wrong mechanism. bubblewrap's branch keeps its tests, which fake the binary + (`test_wrap_argv_builds_the_bwrap_command`, + `test_spawn_uses_bwrap_when_that_is_the_mechanism`). + + The session `availability` fixture still reports what the machine has. A + test that wants another mechanism patches `probe` itself, after this. + """ + av = probe() + if not (av.bwrap and av.landlock_abi): + return + pinned = av._replace(bwrap=None) + monkeypatch.setattr("env_mgr.isolation.probe.probe", lambda: pinned) + monkeypatch.setattr("env_mgr.prepare.probe", lambda: pinned) diff --git a/agent_sys/tests/env_mgr/test_prepare.py b/agent_sys/tests/env_mgr/test_prepare.py index a17ac9a65..97e3e7e17 100644 --- a/agent_sys/tests/env_mgr/test_prepare.py +++ b/agent_sys/tests/env_mgr/test_prepare.py @@ -577,7 +577,9 @@ def test_wrap_argv_refuses_when_the_binary_vanished(ctx, tmp_path: Path, monkeyp availability=Availability(bwrap=str(fake), landlock_abi=None), ) fake.unlink() - monkeypatch.setenv("PATH", os.environ["PATH"].split(os.pathsep, 1)[1]) + # Only the emptied directory: the rest of PATH may hold a real bwrap + # (GitHub's Ubuntu 26.04 runners ship one), which would not have vanished. + monkeypatch.setenv("PATH", str(fake.parent)) with pytest.raises(NoConfinement, match="does not start"): prepared.wrap_argv(["python3"]) From edd3e5955560c573bd1804587a0247a43727ae3d Mon Sep 17 00:00:00 2001 From: YaoCheng Date: Thu, 22 Oct 2026 10:53:58 +0000 Subject: [PATCH 20/20] Run the dsv4 parity check on a Python sglang publishes wheels for With `cargo test` no longer crashing, the rust job on Ubuntu 26.04 now reaches the dsv4 parity step and fails: "sglang is installed but has no dsv4 encoder". The step used the runner's system Python, which is 3.14 there. sglang 0.5.20+ ships only cp310-cp313 wheels, so `pip install sglang` falls back to 0.5.9, the last release with a pure-Python wheel or sdist, from before the dsv4 encoder existed. 24.04's Python 3.12 got the current release. Pin 3.12 with setup-python for this step (0.5.22 dropped cp310, so the 3.10 the other jobs use would not do either). Co-Authored-By: Claude Opus 5.5 Signed-off-by: YaoCheng --- .github/workflows/ci.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 12be3d7f4..253cec5ec 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -221,6 +221,12 @@ jobs: # a prefix the engine never produced and kv-aware settles at 0% hit rate -- # so diff the two here. sglang's module is pure stdlib and loads straight # off disk, hence --no-deps: no torch, no GPU, seconds not minutes. + # A pinned Python, because sglang only publishes wheels for the versions + # it supports: on the runner's own (3.14 on Ubuntu 26.04) pip quietly + # falls back to the last pure-Python release, 0.5.9, which predates dsv4. + - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0 + with: + python-version: "3.12" - name: dsv4 encoder parity with sglang working-directory: . run: |