[ChengYao-amd]

Two changes came out of GLM-5.2 PD tuning on MI355X (decode TP4 + DCP4), where they were first measured as patches (AMD-AGI/Infera#209). The second, the DCP top-k exchange, has since landed on main in #2776 with the same code, so this PR now carries only the first, as one commit.

1. The API process no longer slows down with every request it serves

Two things in the API process grow with every completed request and are never freed:

• CoreManager.outputs_queue. The output thread puts every batch of finished sequences on it, but only LLMEngine.step() / generate() and asyncio mode read it; the OpenAI server and the atomesh standalone service deliver through the per-sequence stream callbacks. Online, every finished Sequence stays reachable for the life of the process (the same leak an earlier fix removed for stream_outputs_queue).
• atom:{prefill,decode}_request_context_tokens. One Gauge series per request, never removed. In multiprocess mode, every /metrics render reads and parses all of them.

Each gen-2 pass then walks a heap that grows with the requests served, /metrics renders get slower, and both freeze the event loop. Admission, streaming and the PD hand-back all wait on them.

• Whether outputs_queue has a reader is declared by CoreManager.consumes_outputs_queue, set in _init_shared_state. It defaults to True (offline LLMEngine, rollout and asyncio mode read the queue and are unchanged). The OpenAI server and AtomEngineService clear it when they wire up, and the output thread queues only while it is set. get_output() / get_output_async() raise if it is cleared instead of blocking on a queue nothing fills. _asyncio_mode is written only in _init_shared_state, and both managers derive async_output_queue from it.
• The two per-request gauges are opt-in with ATOM_METRICS_PER_REQUEST_CONTEXT=1 (envs.py, docs/environment_variables.md, docs/metrics_guide.md). pd_server_atom.sh sets it on both roles for aiperf_agentic, next to ATOM_ENABLE_METRICS_DEVICE_TIMER, so the agentic reports keep their request-context panels.

Behavior change: atom:prefill_request_context_tokens and atom:decode_request_context_tokens have no samples unless ATOM_METRICS_PER_REQUEST_CONTEXT=1 is set. Calling get_output() / get_output_async() (i.e. LLMEngine.step() / generate()) on an engine that the OpenAI server or the atomesh standalone service has wired up now raises RuntimeError instead of blocking forever. Other metrics are unchanged, and offline LLMEngine / asyncio mode still get every finished sequence.

2. DCP decode: the sparse-indexer top-k exchange uses the custom all-gather (now in main via #2776)

dcp_decode_candidate_exchange_fused gathers each rank's local top-k (score, id) candidates over the DCP group once per indexer layer, 24 times per GLM-5.2 decode step with MTP3. It called GroupCoordinator.all_gather, which is pynccl by default. dcp_all_gather, already used for the LSE and query-head gathers, takes aiter's IPC all-gather when the shape qualifies and pynccl otherwise. The float32 scores are gathered as they are. Both paths only copy bytes, so the result is bit-identical.

#2776 made the same change on main, so the commit was dropped when this branch was rebased onto ce46c7c8d. The "(2)" measurements below are kept for reference: they measured this change.

Validation

Rebased onto main c34e7d0b2 after review, and again onto ce46c7c8d on 10-30 (53 more commits; the only conflict was #2456's backpressure refill next to the guarded put_nowait, and both are kept). The GPU tables below were measured on the original base a1ebc608 and are not re-run; a short regression on c34e7d0b2 is listed under "After rebase".

Main a1ebc608 and this branch, each mounted as the ATOM source in rocm/atom-dev:nightly_202610221638 (aiter 665a43dc4) on one 8 × MI355X node. GLM-5.2-MXFP4 runs the standalone DCP setup from recipes/Agentic-GLM-5.2.md without LMCache (TP4 + DCP4, MTP3, block 64, FP4 index cache, --max-num-batched-tokens 16384, graphs up to 128). Main and this branch run at the same time on GPUs 0–3 and 4–7.

Unit tests (CPU, .github/scripts/run_unit_tests.sh, on ce46c7c8d): 10,160 passed, 1,085 skipped, 1 failed; main 10,152 passed, 1,085 skipped, 1 failed. The one failure, test_git_sha_of_checkout_with_different_owner, also fails on main with the local git 2.34. tests/test_engine_core_mgr_outputs_queue.py drives a real output thread: finished sequences are queued for LLMEngine and asyncio mode, not behind the OpenAI server or AtomEngineService, and get_output() / get_output_async() raise when nothing fills the queue. Those last three cases fail on the previous head. test_request_context_gauges_are_opt_in fails on main and passes here. Black and Ruff pass on the changed lines.

Accuracy (GSM8K with CI's lm_eval command: 3-shot, 65 concurrent, flexible-extract, real MTP acceptance): main 0.9158 ± 0.0076, this branch 0.9249 ± 0.0073.

(1) API process over a long run. 20,000 requests of 2,048 random input and 64 output tokens at concurrency 64 (29 min, real MTP acceptance). /metrics is scraped once a second, as a Prometheus job would, and the API process's gen-2 passes are timed with a gc.callbacks hook:

minutes	main: gen-2 share of wall clock	main: mean gen-2 pause	main: /metrics p50	main: API RSS	this branch: gen-2 share	this branch: /metrics p50	this branch: API RSS
0–5	1.9%	19.5 ms	94 ms	1,977 MB	0.0%	5.1 ms	1,737 MB
5–10	6.9%	45.5 ms	268 ms	2,199 MB	0.0%	5.1 ms	1,737 MB
10–15	14.5%	88.3 ms	504 ms	2,425 MB	0.0%	5.2 ms	1,738 MB
15–20	22.7%	135.0 ms	748 ms	2,654 MB	0.0%	5.3 ms	1,738 MB
20–25	31.3%	193.9 ms	1,009 ms	2,885 MB	0.0%	5.3 ms	1,738 MB
25–29	34.0%	239.1 ms	1,205 ms	3,022 MB	0.0%	5.3 ms	1,738 MB

• Over the run, main's API process ran 2,569 gen-2 passes taking 312.6 s (18.0% of the wall clock), and this branch's 30 passes took 0.1 s. Main's /metrics ended at 43,668 request-context series and a 6.5 MB body.
• TTFT mean / p99: main 661 / 1,757 ms, this branch 425 / 678 ms.
• Main completed this run 4.8% faster (12.34 vs 11.78 req/s). The two are equal in its first 5 minutes and on fresh servers (3,000 requests of the same load with forced acceptance, both GPU halves: main 15.05 / 15.07 req/s, this branch 15.08 / 15.13), so the gap comes with main's stalls. As they grow, fewer requests are in main's engine (58.6 vs 61.9 running on average in minutes 20–25) and its prompt throughput is higher (24.8K vs 23.2K tok/s). Most likely, new requests held back by a frozen API process reach the engine together and fill its prefill steps better. This load is prefill-bound (2K in, 64 out), where that trades latency for throughput; TTFT above shows the cost.

(2) Top-k gather, microbenchmark. [rows, 2048] int32 per rank on DCP4 (4 GPUs). GroupCoordinator.all_gather (main) is compared with dcp_all_gather (this branch). Both return every rank's exact bits, eagerly and on CUDA-graph replay. Time is per gather inside one graph of 24 gathers:

rows	MB per rank	main	this branch
16	0.12	18.0 µs	6.5 µs
64	0.50	22.8 µs	13.5 µs
160 (40 sequences × 4 tokens, MTP3)	1.25	37.1 µs	26.4 µs
256 (64 sequences × 4 tokens)	2.00	50.9 µs	39.3 µs
512	4.00	88.9 µs	76.0 µs

At 40 sequences that is −0.26 ms per decode step. Results are identical with AITER_REUSE_IDENTICAL_COMM_GROUPS=1.

(2) End-to-end decode. Fresh servers with forced MTP acceptance 2.99 and random tokens with --ignore-eos. Each pair runs twice, with the GPU halves swapped. One output chunk is one verify step, so the median inter-chunk latency (ITL p50) is the decode step:

load	GPUs (main / this branch)	ITL p50, main	ITL p50, this branch	Δ step	output tok/s, main → this branch
8K in / 1K out, C64	0–3 / 4–7	50.45 ms	50.02 ms	−0.43 ms	1,628.9 → 1,632.3
	4–7 / 0–3	50.48 ms	50.08 ms	−0.40 ms	1,649.4 → 1,660.0
32K in / 1K out, C32	0–3 / 4–7	38.10 ms	37.79 ms	−0.31 ms	588.8 → 586.6
	4–7 / 0–3	38.15 ms	37.81 ms	−0.34 ms	580.3 → 588.0
2K in / 64 out, C64, 3,000 requests	0–3 / 4–7	49.46 ms	49.02 ms	−0.44 ms	962.9 → 965.2
	4–7 / 0–3	49.50 ms	49.10 ms	−0.40 ms	964.3 → 968.5

The step is 0.31–0.44 ms shorter on both halves (−0.8% to −0.9%). Throughput moves −0.4% to +1.3%, within run-to-run noise.

In the GLM-5.2 PD deployment where these were first measured (ATOM 661ef766 with these changes as patches, AgentX at C64), both API processes had spent ~45% of the wall clock in gen-2 GC. The GC change raised throughput 1.9% and cut mean TTFT from 2.85 to 2.04 s; the top-k gather went from 54 to 31 µs in the decode trace.

After rebase (main c34e7d0b2 vs this branch, same setup, image rocm/atom-dev:nightly_202610251529):

• GSM8K: main 0.9174 ± 0.0076, this branch 0.9189 ± 0.0075.
• Top-k gather: bit-identical on every rank, eagerly and on graph replay, now gathering float32 directly; 160 rows 38.8 → 27.7 µs.
• 8K in / 1K out C64 ITL p50: 51.04 → 50.68 ms, and 51.17 → 50.68 ms with the GPU halves swapped.
• API process over 6,000 requests of 2K/64 at C64: main's gen-2 share rises from 0.7% to 7.7% per 2-minute window (18.6 s in total); this branch stays at 0.01% (0.05 s).

Not in this PR

Infera#209 also moves the DCP attention merge onto a custom all-to-all, and then into one aiter kernel together with the LSE combine and FP8 quant. Those need new aiter kernels (custom_all_to_all, custom_a2a_combine_quant) that are not in aiter yet. They will follow once the aiter side lands.

🤖 Generated with Claude Code

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 2680 --add-label <label>

[valarLip]

Thanks — I reviewed HEAD 20f55ac49 and found no regression in the paths I checked: the top-k exchange via dcp_all_gather gives the same shape and, since aiter's custom all-gather copies bytes, bit-identical results; the call site is inside the registered sparse_attn_indexer custom op (outside Dynamo tracing); and the side stream runs only the shared experts, so it doesn't contend with the indexer for the IPC input pool. I haven't run the unit or GPU tests.

The main gap is that the outputs_queue leak fix doesn't cover the atomesh standalone service, and the reader is inferred from the flush hook rather than declared (inline at engine_core_mgr.py:785). The rest are small cleanups, plus two items on lines outside the diff:

• atom/model_engine/engine_core_mgr.py:508: the worker-node early return still assigns self._asyncio_mode = False, although _init_shared_state now owns this field. Two writers for one field is the drift _init_shared_state's docstring warns about, and a worker node has no output threads, so this line can go.
• docs/metrics_guide.md:106 and :108: the table rows for atom:prefill_request_context_tokens / atom:decode_request_context_tokens don't say they now require ATOM_METRICS_PER_REQUEST_CONTEXT=1; only the later section does. Users reading the table will expect them by default.

[valarLip]

The atomesh standalone service still leaks. This guard uses "flush hook unset" as a stand-in for "someone reads outputs_queue", but AtomEngineService (entrypoints/atomesh/atom_standalone_service.py) also delivers through stream_callbacks (L394/439), calls core_mgr.add_request directly (L232/302), never sets _flush_stream_batch_fn (only api_server.py:3041 does), and never calls get_output/step(). So for atomesh _flush_stream_batch_fn stays None, every finished seq list still goes through outputs_queue.put_nowait(seqs) and is never consumed — the same gen-2 GC growth this PR fixes for the OpenAI server. The comment calling the API server "the one caller that wires the flush hook" already doesn't describe atomesh.

At the root, whether the queue has a reader should be declared by the readers, not inferred from the API server wiring a flush function: e.g. a consumes_outputs_queue flag that LLMEngine.step()/generate() and the asyncio handler opt into. The current inference is wrong in both directions — a callback-only entrypoint without the hook leaks (atomesh), and an entrypoint that sets the hook and later calls generate()/step() (warmup, rollout on the server) would block forever in outputs_queue.get().

[ChengYao-amd]

Agreed on both counts; fixed in 229f29f.

Whether the queue has a reader is now declared rather than inferred from the flush hook. CoreManager.consumes_outputs_queue is set in _init_shared_state, and the output thread checks only that flag before outputs_queue.put_nowait(seqs). The two callback-only entrypoints clear it as they wire up:

• the OpenAI server, in a new _deliver_through_stream_callbacks(core_mgr) that sets the flush hook and clears the flag together, so the two can't drift apart again;
• AtomEngineService.__init__ (atomesh standalone), which never sets the flush hook. This was the leak.

The flag defaults to True, so readers don't opt in; the two non-readers opt out. Reasons:

• The readers are LLMEngine.step()/generate()/generate_multimodal(), generate_hidden_states() in rollout, and asyncio mode. If they had to opt in, each would need to set the flag before its first add_request, or the first finished batch is dropped and get_output() blocks. With True as the default, these paths behave exactly as before.
• Only two entrypoints deliver through callbacks alone (api_server.py and AtomEngineService; nothing else calls core_mgr.add_request with stream callbacks), and each clears the flag at the single point where it takes over the engine.

The two failure directions you pointed out:

• Leak (callback-only entrypoint never reads): an entrypoint that forgets to clear the flag leaks the same way as before. That is why atomesh is now wired explicitly, and both wirings have tests.
• Hang (reader on a manager that doesn't queue): get_output() and get_output_async() now raise RuntimeError when the flag is False, so a generate()/step() after the server has wired up fails immediately instead of blocking in outputs_queue.get().

Tests (tests/test_engine_core_mgr_outputs_queue.py, which drives a real output thread):

• LLMEngine: queued, and LLMEngine.step() reads it;
• asyncio mode: queued;
• OpenAI server wiring: not queued;
• AtomEngineService: not queued;
• flag cleared: get_output() / get_output_async() raise.

The last three fail on the previous head and pass now. The comment at this line no longer names the API server as the only hook owner; it points to the flag.

[valarLip]

This comment ("Still moved as int32 so no float canonicalization can touch the bits") is no longer accurate: dcp_all_gather sends 4-byte payloads through _ag_custom_view_dtype → float32, and aiter's custom_all_gather also maps int32 to float32, so the custom path actually moves float32. It's bit-safe only because the aiter kernel is a pure memcpy; if that ever changes, the comment points at a protection that no longer exists.

[ChengYao-amd]

Right, the custom path moves float32, so the int32 view protected nothing. Rewritten in 4d3b997: the comment now says the result is bit-identical because both the custom all-gather and the pynccl fallback only copy bytes.

[valarLip]

Related simplification: send.view(torch.int32) followed by gathered_sc.view(torch.float32) is now a no-op round trip — both the custom path and the NCCL fallback copy bytes — so gathering the float32 send directly does the same job and drops the two views plus the comment justifying them.

[ChengYao-amd]

Done in 4d3b997. The float32 send is gathered directly, and both views and the comment defending them are gone. send is already contiguous float32, so dcp_all_gather hands it to device_comm.all_gather with no view at all. The views were metadata-only before, so there is no change to kernels or graph nodes. A 4-GPU check that gathers random 32-bit patterns (NaN payloads, signalling NaNs, denormals) as GroupCoordinator.all_gather on int32, dcp_all_gather on int32 and dcp_all_gather on float32 is bit-identical across all three, eagerly and on CUDA-graph replay, on DCP4 on MI355X (both GPU halves, and with AITER_REUSE_IDENTICAL_COMM_GROUPS=1). Timing inside a graph of 24 gathers is the same before and after this change: 160 rows (1.25 MB/rank) 27.7 µs for both, vs 38.8 µs for GroupCoordinator.all_gather; 16 rows 6.6 vs 18.2 µs.

[valarLip]

async_output_queue reads config.asyncio_mode with strict attribute access, while _asyncio_mode (set in _init_shared_state) uses getattr(config, 'asyncio_mode', False), so the same fact has two sources with different defaults. With a config lacking the attribute, _asyncio_mode is False but this line raises; if they ever diverge, the output thread queues for a handler that has no async_output_queue, or the reverse. Suggest deriving one from the other.

[ChengYao-amd]

Fixed in 229f29f. _init_shared_state is now the only writer of _asyncio_mode (still getattr(config, "asyncio_mode", False)), and both CoreManager.__init__ and DisaggCoreManager.__init__ derive async_output_queue from it. DisaggCoreManager used to hard-code None regardless of asyncio_mode, which was the same disagreement; with asyncio_mode=True it now gets the queue. No entrypoint in the repo sets asyncio_mode=True today.

[ChengYao-amd]

Thanks for the review. All four inline comments are addressed; replies are on each thread. The two items outside the diff:

• engine_core_mgr.py:508: removed. The worker-node early return no longer writes _asyncio_mode; _init_shared_state is its only writer.
• docs/metrics_guide.md:106/108: both table rows now start with "Opt-in: requires ATOM_METRICS_PER_REQUEST_CONTEXT=1, no samples by default."

Rebase. Rebased onto main c34e7d0b2 (13 new commits, including #2688 and #2644). The only conflict was in atom/utils/envs.py, next to ATOM_ENABLE_METRICS_DEVICE_TIMER. ATOM_METRICS_PER_REQUEST_CONTEXT now uses the same parenthesised lambda. docs/environment_variables.md and tests/test_envs.py merged cleanly. The review changes are folded into the two existing commits (229f29f GC, 4d3b997 DCP), so the history still has one commit per change.

Validation after the rebase.

• Unit tests (CPU, .github/scripts/run_unit_tests.sh): 9,798 passed, 960 skipped, 1 failed. Main c34e7d0b2 on the same machine: 9,790 passed, 960 skipped, 1 failed. The failure is test_git_sha_of_checkout_with_different_owner on both, as before. The 8 extra tests are this PR's, and no other test changes status. Black and Ruff pass on the changed lines.

GPU checks ran on one 8 × MI355X node, with main c34e7d0b2 and this branch side by side on GPUs 0–3 and 4–7, using the setup from the description. The image is rocm/atom-dev:nightly_202610251529; the description's nightly_202610221638 predates #2688 and fails torch.compile on the new base, on main as well.

• GSM8K (CI's lm_eval command, 3-shot, flexible-extract, real MTP acceptance): main 0.9174 ± 0.0076, this branch 0.9189 ± 0.0075.
• Top-k gather, DCP4: all three variants are bit-identical on every rank, eagerly and on graph replay, on both GPU halves. The variants are main (int32 view), the previous head (dcp_all_gather on int32) and this head (float32). Per gather inside a 24-gather graph: 160 rows 38.8 → 27.7 µs (previous head 27.7 µs); 16 rows 18.2 → 6.6 µs. tests/test_dcp_topk.py: 71 passed.
• Decode step, 8K in / 1K out, C64, forced AL 2.99, run twice with the GPU halves swapped. ITL p50 main → this branch:
  • GPUs 0–3 / 4–7: 51.04 → 50.68 ms (−0.36 ms); output 1,478.7 → 1,484.7 tok/s.
  • Swapped: 51.17 → 50.68 ms (−0.49 ms); output 1,488.6 → 1,499.7 tok/s.
• API process, short run: 6,000 requests of 2K in / 64 out at C64 (9 min), with /metrics scraped at 1 Hz. Main's gen-2 share of the wall clock grows per 2-minute window, 0.7 → 1.9 → 3.5 → 5.7 → 7.7%: 611 passes, 18.6 s in total, and the /metrics p50 rises from 51 to 280 ms. This branch has 9 passes (0.05 s) and a flat 4.6–4.8 ms /metrics p50.
• atomesh standalone: not run on GPU, because the published images don't ship atomesh_runner. The fix is covered by test_atomesh_engine_service_queues_nothing, which drives a real output thread behind AtomEngineService.

None of the review changes touch the paths measured in the description, so the full tables there are not re-run. The outputs_queue change leaves the OpenAI server exactly as before (nothing queued); it adds a boolean check per finished batch. The DCP change drops two metadata-only views. The GPU runs above check the rebased base, and match the description's numbers.

[yitingw1]

Please fix code conflicts.

[ChengYao-amd]

Rebased onto main ce46c7c8d; the conflicts are fixed.

• The only conflict was in atom/model_engine/engine_core_mgr.py. [AMD][DSR1] feat(dp): backpressure (lazy) DP dispatch to avoid forced-eager tail #2456 added the backpressure refill right after the outputs_queue.put_nowait(seqs) that this PR guards. Both are kept, and the refill still runs after every finished batch, whether or not the batch is queued.
• The DCP top-k commit is gone. refactor: shared DSA indexer forward and top-k schedule, sparse prefill decided once, DCP kernel cleanup #2776 made the same change on main (dcp_all_gather on the float32 scores), so after the rebase the commit was empty. This PR is now the GC commit only (5bfc096ad), and I've updated the title and description to match.
• Unit tests (CPU, .github/scripts/run_unit_tests.sh): 10,160 passed, 1,085 skipped, 1 failed. Main ce46c7c8d on the same machine: 10,152 passed, 1,085 skipped, 1 failed. The one failure is test_git_sha_of_checkout_with_different_owner on both. Black and Ruff pass on the changed lines.