feat(atom): optimize atom pd performance

.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: |

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)

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"])
 

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
 

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" <image> \
+  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.