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: | 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"]) 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..fd962a3a3 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,40 @@ 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 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 \ + --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 +83,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 +117,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..af8497ac3 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/check_dcp_merge.py @@ -0,0 +1,168 @@ +#!/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 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() +# 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}") + +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 new file mode 100644 index 000000000..8f0e076ad --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_custom_alltoall.py @@ -0,0 +1,428 @@ +#!/usr/bin/env python3 +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# +# SPDX-License-Identifier: MIT +############################################################################### +"""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 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, 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. + 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 +(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 shutil +import sys + +from anchor_patch import apply + +TAG = "aiter-patch custom-alltoall" + +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; + +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); +} + +// 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; + } + // 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_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 + 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.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 + 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) + + # 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( +''' + +PATCH_CSRC = { + "../csrc/include/custom_all_reduce.cuh": ( + "// infera: all-to-all.", + [ + ( + "// ========== reduce_scatter kernel start ==========\n", + KERNELS + "// ========== 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_METHODS)], + ), +} + +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 new file mode 100644 index 000000000..b58468bf4 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_glm52_decode_gemm_rows.py @@ -0,0 +1,122 @@ +#!/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. 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 +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 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 + 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 + 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()) 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..fc217d04d --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_aiter_mla_bmm_configs.py @@ -0,0 +1,148 @@ +#!/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. 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. +""" + +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)] +# fmt: off +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), + ], +} +# fmt: on +# 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()) 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..d803afb35 --- /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/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")) 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..217972eb3 --- /dev/null +++ b/deploy/docker/patches/atom_nightly/patch_atom_dcp_custom_a2a_combine.py @@ -0,0 +1,125 @@ +#!/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.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. +""" + +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")) 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/README.md b/examples/atom_4p4d_glm5.2/README.md index 3b09d5c9c..94fa00173 100644 --- a/examples/atom_4p4d_glm5.2/README.md +++ b/examples/atom_4p4d_glm5.2/README.md @@ -3,16 +3,17 @@ 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: ```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 ``` @@ -20,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 \ @@ -35,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 @@ -51,7 +57,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 +78,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 and BF16 MTP layer | | -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 +~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 +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 @@ -98,61 +112,79 @@ 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; 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 \ - 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 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` (earlier image and kit) | |---|---|---| | 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.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 +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 -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. 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 | 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, 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) +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 ``` @@ -178,7 +210,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 +233,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 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` would need engine KV events, so the policy is `round-robin`. @@ -211,9 +244,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. @@ -228,8 +261,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 @@ -241,18 +277,43 @@ 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 (`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. +**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: 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. + +**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, 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 @@ -263,6 +324,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/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/common.sh b/examples/atom_4p4d_glm5.2/engine/common.sh index 0ae4d9c23..73bf46d5b 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,9 +35,9 @@ 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 - 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) + 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) # 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 new file mode 100755 index 000000000..842dcad49 --- /dev/null +++ b/examples/atom_4p4d_glm5.2/engine/store.sh @@ -0,0 +1,70 @@ +#!/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. 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. +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..e2d0ff510 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,17 @@ 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 +# 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)) || + 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 @@ -47,6 +58,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 4861a3f51..95f25b9e6 100755 --- a/examples/atom_4p4d_glm5.2/engine/worker.sh +++ b/examples/atom_4p4d_glm5.2/engine/worker.sh @@ -23,7 +23,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. @@ -45,7 +50,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. @@ -53,7 +58,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 @@ -85,17 +100,26 @@ 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 16384 --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" 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) env_args=() for kv_pair in "${env[@]}"; do env_args+=(-e "$kv_pair"); done @@ -104,9 +128,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 d17a13965..e41cdd0d0 100755 --- a/examples/atom_4p4d_glm5.2/env.sh +++ b/examples/atom_4p4d_glm5.2/env.sh @@ -62,17 +62,20 @@ 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:-}" -# 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 ~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.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) 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}" # 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 @@ -81,35 +84,55 @@ 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}" -# 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); 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}" +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:-}" export DECODE_EXTRA_ARGS="${DECODE_EXTRA_ARGS:-}" 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 @@ -128,9 +151,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}" # --------------------------------------------------------------------------- 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/router/disagg_protocols/atom_mooncake.py b/infera/router/disagg_protocols/atom_mooncake.py index 34cc20865..96bb5e2b7 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,16 @@ 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; + # so does an empty list or a negative id, which decode would reject. + ids = prompt_from_chunk(prefill_response_payload) + if ids and min(ids) >= 0: + 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/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/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/manual/examples/atom_4p4d_glm5.2.md b/manual/examples/atom_4p4d_glm5.2.md index 1befe8d4f..e927abfd5 100644 --- a/manual/examples/atom_4p4d_glm5.2.md +++ b/manual/examples/atom_4p4d_glm5.2.md @@ -1,11 +1,12 @@ # 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 -`--session-affinity prefill`). `PREFILL_PP=1` runs a TP4 + DCP4 prefill instead. -The runnable kit is +**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, +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. @@ -13,9 +14,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 +29,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 +48,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 +73,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 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. ## 5. Tear down 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" 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..7c90d6d23 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,27 @@ 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: 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"), + 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 +532,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 +572,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"}); + // 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]}]}), + ] { + 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] 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") diff --git a/tests/unit/router/test_atom_mooncake_protocol.py b/tests/unit/router/test_atom_mooncake_protocol.py index 91b535203..269b0c397 100644 --- a/tests/unit/router/test_atom_mooncake_protocol.py +++ b/tests/unit/router/test_atom_mooncake_protocol.py @@ -94,6 +94,30 @@ 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 (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]}, + { + "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] + # 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"): self.proto.extract_handoff({"choices": []})