[sudhu2k]

Motivation

This PR adds Rollout Router Replay (R3) support to the ATOM runtime. Inspired by
vllm-project/vllm#28284, it records the expert routing
decision for every token at every MoE layer during inference, so RL
post-training can replay the exact routing the rollout used.

vLLM exposes this as enable_return_routed_experts →
CompletionOutput.routed_experts (int16 [seq_len-1, num_layers, top_k]).
ATOM LLM.generate() did not, so Lumen-RL and similar trainers could not
hard-assign ATOM routes on the actor forward.

Goals:

• Match the vLLM field contract on native ATOM generate(), including
  CUDA-graph decode and prefix-cache hits.
• Refuse at startup any configuration where the routes could be incomplete or
  wrong, instead of returning bad routes.
• No route work at all when the flag is off.
• Leave OpenAI serving and vLLM/SGLang plugin mode unchanged.

Routes are returned as routed_experts in each request's generate() output
dict, read from the finished Sequence. RequestOutput (the per-step stream
callback) does not carry them: no in-repo consumer reads them there, and the
OpenAI server never returned them.

Technical Details

Flag (default False): Config.enable_return_routed_experts /
--enable-return-routed-experts. ModelRunner and Scheduler read it once
into _return_routed_experts; when it is off, none of the paths below run.

Capture (GPU, inside the forward)

atom/model_ops/fused_moe/routed_experts_capturer.py owns one GPU buffer,
int32 [num_slots + 1, num_layers, top_k], indexed by physical KV slot and
initialized to -1 (UNCAPTURED).

• allocate_kv_cache binds each target FusedMoE its layer's view,
  buffer[:, layer] (routed_experts_view), on TP rank 0 only. Sleep clears it
  with the KV views and wake rebinds it, so it shares the KV cache's lifecycle;
  there is no process-global capturer.
• _select_logical, which every quant method's select_experts returns
  through, calls maybe_capture_routed_experts(layer, topk_ids). The write
  happens on device, so CUDA-graph replay still records it.
• The buffer's extra row is its last row. A CUDA-graph pad token has
  slot_mapping == PAD_SLOT_ID == -1, which indexes that row, so padding never
  overwrites slot 0. slot_mapping indexes the buffer directly; there is no
  separate destination computation.
• k < top_k padding is -1, so an uncaptured cell is never mistaken for
  expert 0.
• Only the target model's MoE layers get a view, so an MTP drafter's MoE
  cannot write into the target model's rows.
• The buffer's bytes are charged to the PAGE pool budget (per KV block, plus
  the pad row) before the expected-vs-actual KV check.

Export (runner → scheduler)

Per step, on rank 0:

1. store_step gathers buffer[slot_mapping[:total_tokens]] into int16 rows on
   the compute stream. It has to run there, before the next forward rewrites
   the slot_mapping staging buffer in place.
2. The rows are copied to the host without blocking on the existing
   async_copy_stream, with their own event, and kept on the batch as
   PendingRoutes.
3. When that batch's output is built one step later, PendingRoutes.patch()
   waits on the event and attaches one
   RoutedExpertsPatch(req_ids, starts, offsets, rows). Each batch is attached
   at most once; a pure middle prefill chunk, which builds no output, holds its
   routes until the next output that does.

Output is always deferred under R3. The only non-deferred configurations (PP,
and a P/D producer with recurrent state) are already refused, so there is one
export path.

The IPC payload is the step's new rows plus three small index arrays,
independent of context length.

Accumulation (scheduler)

• Scheduler._apply_routed_experts_patch writes each request's segment into
  Sequence's growable int16 buffer (doubling capacity) at its absolute start
  position. A re-forwarded range overwrites the old rows. A patch that would
  leave a gap raises instead of handing the trainer uninitialized rows. Rows
  for a request that has left running are dropped: a finished request's last
  in-flight step is past its num_tokens - 1 rows, and a preempted one asks
  for them again on resume.
• Preemption: routed_forwarded_end counts rows planned, not delivered, so
  preempt rewinds it to routed_expert_rows. A range whose patch was lost is
  re-requested on resume (from the cache on a hit, otherwise recomputed).
• Prefix-cache hits: _plan_routed_prefixes records, at schedule time, any
  prefill that starts past what the request has forwarded. The runner turns
  that range into slots with prefix_slots (the same slot formula the forward
  uses) and exports those rows once, as an extra segment of the same copy.
• Finish: every finish path, ABORTED included, calls
  finish_routed_experts(num_tokens - 1), which trims to the live rows and
  freezes them read-only. Rejected speculative tokens are already excluded from
  num_tokens. A non-aborted finish with fewer than num_tokens - 1 rows
  raises. An aborted request may be short, because its last step's patch can
  still be in flight; it returns the prefix it has. Later patches for a
  finished request are ignored.

Sleep / wake

release_memory(tags=["kv_cache"]) clears each MoE layer's
routed_experts_view with the KV views and drops the runner's capturer, so the
buffer is freed with the KV cache and wake's get_num_blocks() /
allocate_kv_cache() don't count it as overhead and allocate a second one.
Wake recreates and rebinds it in allocate_kv_cache. Sleep also drops the
pending routes (_held_routes / _deferred_routes) next to
tokenID_processor.clean(), so a stale step's rows can't attach to the first
batch after wake.

Refusals

Where	Refused	Why
check_return_routed_experts (Config, Scheduler, RapidServeModelRunner)	DCP or PCP ≠ 1	Routes would be a rank-local shard
	PP ≠ 1	Layers split across ranks
	KV transfer / offload	Route rows are not moved with KV blocks
	RapidServe P/D	Decode never allocates KV or the capture buffer
ModelRunner	No FusedMoE in the model	Dense models have no routes
	Attention builder without supplies_token_slots (e.g. DeepSeek-V4 / V4.1)	The metadata has no per-token slots, so routes would never be written
Mono decode (routed_experts_refusal)	MiniMax-M3, DeepSeek-V4.1, GLM-5.2, Kimi-K3, Qwen3.8 mono (not the shared DSA indexer, which routes no experts)	Fused kernels route in-kernel and never write the buffer; the model falls back to the per-layer path
Qwen4-Exp	bs = 1 single-token experts	Routes outside FusedMoE; disabled under R3

Downstream: Lumen-RL's Ray ATOM rollout sets this flag when moe.r3.enabled
and forwards routed_experts from ATOMRayServer.

Test Plan

On MI355X (TP = 2):

pytest tests/test_routed_experts_capturer.py tests/test_sequence.py \
  tests/test_scheduler.py tests/test_rollout_memory_manager.py \
  tests/test_postprocess_width.py tests/test_page_unit_checkpoint.py \
  tests/test_output_transfer_runner.py tests/test_model_runner_staging_fence.py \
  tests/models/glm52/test_mono_config.py \
  tests/models/test_kimi_k3_mono_config_refusal.py \
  tests/test_deepseek_v41_mono_dispatch.py tests/test_minimax_m3_mono_dispatch.py -q
ATOM_RUN_ROUTED_EXPERTS_E2E=1 pytest tests/test_routed_experts_generate_e2e.py -v

Covers batch-order segments, prefix segments and slot-formula parity, graph pad
tail, -1 for uncaptured rows and a short slot_mapping, attach-once across a
deferred step and a middle chunk, the async copy snapshot before staging reuse,
the gap and short-finish errors, finish trim/freeze, the preempt rewind, the
disabled path, drafter MoE exclusion, every refusal (each mono model's included),
and sleep unbinding the views. The e2e (TP = 2) covers CUDA-graph decode with a
prefix hit, concurrent requests, chunked prefill followed by a prefix hit, and
MTP 1 / 3 verify, checking each against a prefill of the same tokens.

Test Result

• Nov 4, on 1105f69f1: the unit tests above are 585 passed, and the e2e on
  2 × MI355X is 5 passed. The full tests/ suite has no failures that are not
  also on the merge base with main: every file that failed was rerun on both
  trees, and the sets match.
• GPU runs on tiny dummy-weight models (Qwen3-Next GDN + MTP under graphs,
  MiniMax-M3 TP 4, Qwen4-Exp bs 1, DeepSeek-V3 MLA prefix hit): routes have the
  right shape, contain no -1, and a prefix hit returns identical cached rows.

Host cost per decode step (MI355X, ctx 4096)

Geometry	bs	R3 off	R3 on
Qwen3-30B-A3B (L = 48, K = 8)	256	0.3 µs	272 µs
DeepSeek-R1 (L = 61, K = 8)	256	0.3 µs	275 µs

(2.33 ms per step before the batch-order export.)

The per-layer capture views cost the same as the old global scatter (±0.5%;
62 layers × top-k 8 at 256 tokens: 194.3 vs 193.9 µs).

Serving, R3 off vs on (MI355X, Nov 4)

Qwen3-30B-A3B-Base, TP 1, max_num_seqs 256, random 1024 in / 1024 out,
concurrency 256, 2560 requests. Runs alternate off/on with the compile cache
cleared before each start.

run	R3	tok/s	mean TTFT (ms)	mean TPOT (ms)	P99 TPOT (ms)
2	off	21,452	222.7	23.05	26.95
3	on	21,324	232.2	23.18	25.61
4	on	21,481	228.8	23.00	25.27

Run 1 (off, 20,177 tok/s) overlapped the unit-test suite on the host and is
left out. Throughput and TPOT are within noise; mean TTFT is about 3% higher
with R3 on. An earlier round on the previous revision measured 2.6% lower
throughput with R3 on.

Lumen-RL 1-step DAPO smoke (KL)

Image zhangdanyangamd/lumen-rl:dapo-gfx950-rocm7.2.3-260910, 8× gfx950, seed
10086, Qwen3-30B-A3B-Base, Megatron native (TP = 1, PP = 1, EP = 8, a2a),
ATOM Ray rollout (TP = 1, CUDA graphs, sleep_level=2), 64 sequences, 256-token
responses, moe.r3.enabled + replay_mode=hard_assignment.

Same-weight r3_verify before the optimizer step:

rollout	r3	tokens	kl	abs_diff	frac_abs_gt_0.1
ATOM	off	16087	0.00196	0.0256	0.087
ATOM	on	15530	0.000765	0.0142	0.043
vLLM (reference)	off	15623	0.00222	0.0240	0.081
vLLM (reference)	on	15533	0.000835	0.0135	0.039

Native ATOM R3 cuts the large-mismatch tail ~2×, matching vLLM. The residual is
train/infer kernel mismatch, not missing routes.

Generate performance

Qwen3-30B-A3B serving, max_num_seqs=256 (measured before the A3 fixes,
which only remove host syncs):

	R3 off	R3 on	Change
Total tok/s	23,431	22,827	−2.6%
Median TPOT	21.74 ms	22.32 ms	+2.7%
P99 TPOT	23.4 ms	24.1 ms	+3.2%
Mean TTFT	210 ms	228 ms	+9%

MiniMax-M3-MXFP8, TP=4, 8192 in / 1024 out: −0.8% tok/s at concurrency 16,
−1.4% at concurrency 64; TTFT unchanged.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[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 2274 --add-label <label>

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

🟡 Changes recommended

GPU-buffer lifecycle and PAGE-backed checkpoint accounting issues remain unresolved, with additional avoidable overhead in the default path.

Get a fresh assessment by requesting another Copilot review.

Review details

• Files reviewed: 12/12 changed files
• Comments generated: 3
• Review effort level: Lite

[sudhu2k]

The geometry concern is right, though DeepSeek-V4 can't hit it.

DeepSeek-V4 / V4.1: R3 refuses these builders before any surcharge is added. Both inherit CommonAttentionBuilder, where supplies_token_slots = False, so _routed_experts_geometry() raises at startup.

Where it does cause an issue: Kimi's MLA+GDN builders (kimi_linear, glm5_next_text). They supply token slots, use PAGE checkpoints, and don't override checkpoint_page_unit_bytes(). R3 refuses pipeline parallelism and RapidServe, which are the only cases where Kimi turns checkpoints off, so with R3 on, page_unit_bytes took the merged "kv" charge including the route rows. PageUnitGeometryMixin._page_unit_regions then failed the equality check. Qwen GDN was unaffected because its override already narrows the unit to its MHA regions.

Fix: I kept the route rows in "kv". They grow with the block count, so charging them per block is what keeps the capture buffer inside --gpu-memory-utilization. What changed is where the checkpoint unit comes from: the new checkpoint_page_unit_bytes(builder) in page_unit_checkpoint.py sums only the target builder's own PAGE specs, and a backend override can still narrow that further. Sizing still uses the merged charge.

Why not refuse R3 when state_transfer().copies is on: that would also rule out R3 on Qwen3-Next and Qwen3.5-MoE, which work.

[Copilot]

🟡 Changes recommended

Unresolved critical routing and compilation issues, plus multiple memory-accounting defects, block approval.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (4)

atom/model_engine/model_runner.py:1459

• The route buffer is allocated and charged to the PAGE pool on every TP worker, but _store_routed_experts_step and _attach_routed_experts explicitly skip every rank != 0 worker. Those workers therefore hold a full GPU buffer and CPU twin that are never read, multiplying the memory surcharge by TP size and reducing KV capacity unnecessarily. Either make initialization/pool pricing rank-0-only or consume the per-rank buffers in an explicit aggregation path.

        geo = self._routed_experts_geometry()
        if geo is not None:
            from atom.model_ops.fused_moe.routed_experts_capturer import (
                capture_bytes_per_kv_block,
            )

atom/model_engine/model_runner.py:2006

• The pad row is subtracted from available_for_kv, but it is not represented in pool_plan.total_reserved_bytes. allocate_kv_cache() then allocates that extra row, so the later expected-vs-actual comparison is always short by capture_pad_row_bytes(*geo) (and can exceed the 3% threshold for small pools). Include the pad allocation in the same accounting used by the mismatch check rather than only reducing the budget.

        # Capture buffer is part of the PAGE budget (per-slot int32 routes).
        # Allocate it before the expected-vs-actual check so the measurement
        # includes those bytes rather than warning on a false mismatch.
        self._maybe_init_routed_experts_capturer(num_kvcache_blocks)

atom/model_engine/model_runner.py:2087

• num_layers is derived from the last MoE layer rather than the model's layer count. Qwen3-Next-style configs can place MoE layers sparsely (decoder_sparse_step/mlp_only_layers), so a dense suffix makes routed_experts shorter than the documented [seq_len - 1, num_layers, top_k] contract. Size this dimension from the target HF config and leave dense-layer rows zero-filled.

        self._routed_experts_geo = (
            max(layer_ids) + 1,
            max(int(m.top_k) for m in moes),
        )

atom/model_engine/model_runner.py:1465

• The R3 surcharge is merged into the PAGE entry size, but the route tensor is allocated separately by _maybe_init_routed_experts_capturer and is not part of the builder's paged pool. For builders with state_transfer.copies (for example Kimi MLA/GDN), get_num_blocks() passes this inflated PAGE size as PagedStateCheckpointSpec.page_unit_bytes, while the checkpoint geometry measures only the actual KV/index regions and raises a page-unit mismatch. Keep the surcharge out of checkpoint page-unit geometry or explicitly reject page-backed state transfer when R3 is enabled.

                page_pool(
                    capture_bytes_per_kv_block(self.block_size, num_layers, top_k)
                )

• Files reviewed: 14/14 changed files
• Comments generated: 2
• Review effort level: Lite

[Copilot]

🟡 Changes recommended

A moderate active issue remains because DP-attention batches can return incomplete per-sequence routes instead of failing closed for a normal Aiter/GDN topology.

12 open issues

• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• Reject DP attention before exporting incomplete routes · New
• Size capture geometry using the full decoder layer count
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

1 resolved since last review

• Prefix rows are double-counted when padding captured routes

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

An active nit remains because the contract comment contradicts the implementation and tests about whether routed-expert capture supports DP attention, potentially misleading users about supported topologies.

13 open issues

• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Correct the routed-expert DP-attention contract comment · New
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[valar-cop]

LGTM

ci Run non-GPU unit tests failed — not sure yet (low): The Pre Checkin job failed but its log (.valar/ci/job_126820971975.log) is 0 bytes, so the cause cannot be classified; the same job passes on main and other PR runs show only cancellations.

next: Re-run the job (or fetch a non-empty log) before merging; if it reproduces, inspect the non-GPU unit test output. https://github.com/ROCm/ATOM/actions/runs/41900249814/job/126820971975

[Copilot]

🔵 Needs a closer look

Human assessment is needed because the change has multiple hidden correctness issues spanning deferred route handling, mixed prefix-cache batches, and KV sleep/wake state cleanup.

13 open issues

• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Correct the routed-expert DP-attention contract comment
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[valarLip]

Reviewed head c9e710c75. The capture/export design has converged a lot since the last round, but there is one crash on a supported topology, and several earlier asks are still open (folded in below where they overlap, listed at the end otherwise). Memory and KV-capacity figures below are estimates from the stated formulas, not measurements; the one timing figure is from a CPU microbenchmark of this PR's own apply_routed_expert_patch.

Must change

1. Pure middle chunks read the drafter's slot_mapping (model_runner.py:3884). In the middle-chunk branch, propose_draft_token_ids(..., align_only=True) runs first, and only then _store_routed_experts_step(batch, stream=...) with slot_mapping=None, so store_step reads current_slot_mapping() after the drafter replaced it (eagle_proposer.py:509 sets attn_metadata.slot_mapping = var["slot_mapping"].gpu[:running_bs]; dspark_proposer.py:905 does the same). With DP>1 + MTP k>=2 (or DSpark) + a prompt long enough to produce a pure middle chunk + R3 -- none of which is refused -- store_step raises the step forwarded N tokens but its slot_mapping has running_bs slots and the engine core dies; if the lengths ever fit, the rows are gathered from draft slots. postprocess already snapshots slots = current_slot_mapping() before propose; do the same here and pass it explicitly. No test covers the align path.

Worth addressing

2. One capture point instead of five (moe.py:5448). The hook is pasted after select_experts at moe.py:5448/5524, comm_fused_moe.py:370 and twice in qwen4_exp.py. Both quant methods' select_experts already go through _select_logical (moe.py:506), which has the layer and returns logical ids before place_experts. Capturing there covers FusedMoE and comm-fused at once, leaving only Qwen4Exp's in-kernel branch. Since the buffer is filled with -1 once and never cleared on slot reuse, a future forward path that misses the hook returns the previous occupant's routes rather than -1, with nothing to catch it.
3. The prefix-hit gather is unbudgeted and unbounded (routed_experts_capturer.py:340). self.buffer[rows_at].to(int16) materializes an int32 gather plus an int16 copy sized to this step's tokens plus the whole cached prefix, and the profile run never allocates it (_store_routed_experts_step returns on is_dummy_run). Estimate, tokens x L x K x (4 + 2) B: DeepSeek-R1 (L=61, K=8) is 2.9 KB/token, ~290 MB for a 100k-token prefix hit, taken from HBM the planner already gave to KV. The int16 rows (~98 MB for 100k tokens) are also queued on the copy stream ahead of the token D2H, so the next step's copy_done wait covers them. Bound it per step (chunk routed_prefix_ranges) or read the prefix once at finish.
4. Capture buffer sizing (routed_experts_capturer.py:30), also A4. CAPTURE_DTYPE = torch.int32 while every reader casts to int16 and both the ids and -1 fit int16, so the PAGE surcharge is twice what is needed; the layer axis also charges dense layers, whose rows are always -1. Estimate, surcharge L x K x 4 B per token against per-rank KV bytes per token: R1 fp8 MLA 1952 B vs 35,136 B, ~5.3% of KV blocks (int16: ~2.7%); Qwen3-235B TP8 fp8 (1 KV head/rank) 3008 B vs 24,064 B, ~11%. The reply to A4 deferred int16 as "a separate follow-up"; it is a one-constant change here and halves the gather as well.
5. Only rank 0 reads the routes (model_runner.py:2455), also A4. Every TP rank still allocates the full torch.full buffer and does an index_put per MoE layer per forward. Keeping the PAGE charge on all ranks keeps block counts identical, as the reply notes, but the buffer and the scatter are only needed on rank 0. Estimate, (num_slots + 1) x L x K x 4 B: ~1.9 GB per non-zero rank for ~1M slots on R1, unread.
6. RequestOutput.routed_experts has no reader (scheduler.py:4651). generate() reads Sequence.routed_experts (llm_engine.py:1077); nothing in atom/ reads the RequestOutput field. Meanwhile the final RequestOutput goes out as ("STREAM", ...) and finished_seqs is pickled with __dict__ including _routed_expert_buf (engine_core.py:575/581), so each finished request's routes cross processes twice (estimate: 32k tokens x 61 x 8 x 2 B ~ 31 MB per copy on R1). Please drop the field and its kwarg.
7. The drop warning fires on every normal finish (scheduler.py:4153). With deferred output a request whose EOS was sampled in step N is already in batch N+1. Postprocess finishes it and removes it from self.running; batch N+1's routes arrive with the next output and running_by_id has no entry, so every finished request logs "not in the batch it was scheduled in" -- which is also not what happened. Skip ids that are already finished (or log at debug) and keep the warning for the case it was added for (earlier thread).
8. Finish truncates but never checks completeness (scheduler.py:4775, also :3080). On abort/reject the previous step's patch can still be in flight; finish_routed_experts freezes the buffer and the late patch is dropped, so the returned array has fewer than seq_len - 1 rows without -1 fill and without an error, while RequestOutput documents [seq_len - 1, ...]. Either assert rows == keep for non-aborted finishes or document that aborted outputs are short.
9. Sleep leaves stale route state (memory_manager.py:335). release_memory runs tokenID_processor.clean() (drops prev_batch) and _release_kv_cache resets the capturer, but _deferred_routes / _held_routes are only cleared in exit. After wake the first deferred step attaches the pre-sleep PendingRoutes to an output whose tokens belong to a new batch. This is one more symptom of 10/11.
10. A second deferred pipeline beside the token one (model_runner.py:3662). Routes get their own after_recv callback threaded into prepare_sampled_ids, a _deferred_routes mirror of prev_batch, a _held_routes list, held_first / take_held ordering rules, and a per-step record_stream + event. Logprobs already ride the token D2H entry and its copy_done, and output_transfer.D2HQueue exists for exactly this. Make routes another payload of the token entry and the ordering, clean() and sleep cases come for free. The non-deferred branch (rows.cpu()) is reachable only with PP or a P/D producer, both refused; please delete it.
11. Process-global singleton (routed_experts_capturer.py:25), A7 still open. _INSTANCE, the weakref owner, the runner's _routed_experts_capturer plus the per-step get() is not capturer check, and teardown written twice (model_runner.py:1210, memory_manager.py:335) -- and item 9 shows the copies already drifting. Binding each target FusedMoE a buffer[:, layer] view in the KV bind loop makes drafter exclusion and sleep/wake follow the KV cache's lifecycle. ATOM runs one runner per process, so the multi-owner machinery protects a configuration that isn't supported.
12. _routed_experts_geometry is not a getter (model_runner.py:2368). It stamps moe_capture_layer_id on modules, enforces refusals and caches via hasattr, and runs as a side effect of _sub_pool_specs / get_num_blocks. Please make it an explicit _bind_routed_experts_layers() called once after load. Its enable_dp_attention and builder_cls.supports_token_graph branch (:2395) can never fire: the only supports_token_graph builder (V4.1) is already refused one check earlier for supplies_token_slots = False. Delete it, and the matching sentence in the description.
13. Mono refusal is in the wrong place (deployment.py:147). common_refusals is shared with atom/mono/dsa_indexer/dispatch.py, which routes no experts, so on GLM-5 with R3 on the DSA indexer mono turns off silently for no reason. Put the refusal on the MoE mono installers that route in-kernel. (The new docstring line is also well past the surrounding width.)
14. Per-request host work in the scheduler (sequence.py:647). _apply_routed_expert_patch does one Python iteration plus numpy calls per request per step: ~1.05 µs/request/step in a CPU microbenchmark of this PR's code (R1 geometry), i.e. ~270 µs/step at bs=256 on the engine-core thread. The doubling buffer copies the whole history on each growth and the frozen view pins the doubled base (estimate: ~125 MB live / up to 250 MB capacity for a 128k-token R1 request). The final length is known at admission (prompt + max_tokens - 1), so allocate once. _routed_experts_final + read-only flag + the copy-while-running property mode can go too: both readers run after finish, and seq.is_finished already says that.
15. Docs (earlier, section 4). docs/configuration_guide.md (Config and EngineArgs tables) still lacks enable_return_routed_experts / --enable-return-routed-experts; pool_entries there still says entries come only from the attention builders' sub_pool_specs(); docs/scheduling_kv_cache_guide.md lacks ScheduledBatchOutput.routed_experts and ScheduledBatch.routed_prefix_ranges; gdn_attn.py:1689's docstring still says the runner takes the unit from the paged entry_bytes. The config comment and CLI help list only the config-time refusals: dense models, builders without token slots (V4/V4.1) and "mono turns off" are only discovered after weight load.

Smaller

• checkpoint_page_unit_bytes(builder) (module function, always resolves) and builder.checkpoint_page_unit_bytes() (optional override, may return None) share a name; rename the resolver (e.g. resolve_checkpoint_page_unit_bytes). It also adds an engine -> model_ops import into page_unit_checkpoint.py; this arithmetic sits more naturally next to plan_pools.
• That change also fixes a startup RuntimeError for a copy-checkpoint backend without its own unit (Kimi MLA+GDN) plus an MHA draft that owns a pool. Worth saying so in the description.
• scheduler.py:4043: the preempt fix-up of routed_forwarded_end is the only R3 hook not gated on self._return_routed_experts.
• is_fused_moe_module's docstring attributes the wrapper to plugin mode, but the decorator is applied unconditionally (moe.py:3452).
• moe_decode_single_token(..., return_topk_ids=False) -> Tensor | tuple[...]: two callers, so always returning (out, topk_ids) gives one return type.
• PendingRoutes and RoutedExpertsPatch differ only in tensor-vs-ndarray and the event; with 10 one record type is enough and merge_routed_patches can become a list applied in order.

Still open from earlier rounds

Not covered above (A4 is in 4/5, A7 in 11, the docs ask in 15):

• A2: per-step trimming (scheduler.py:4626). Kept because the gap check needs the live length. That check could compare a patch's start against seq.num_tokens - 1, which postprocess already maintains, instead of keeping a second counter in step with it every step.
• A3: have the scheduler send the computed prefix slots. The upload is pinned now, but the runner still rebuilds slots from batch.block_tables by batch position with a fresh pin_memory() per prefix-hit step.
• --enable-return-routed-experts is exposed on the server, where nothing reads routed_experts (section 3). A server with the flag pays the capture and KV cost and returns nothing; either refuse it in openai_server or document it as offline-only.
• test_rollout_memory_manager.py installs the global capturer without fixture teardown (section 4). If the test fails before _release_kv_cache, every later MoE test in the process captures.

[valarLip]

[1] Must change. propose_draft_token_ids(align_only=True) just above has already replaced attn_metadata.slot_mapping (eagle_proposer.py:509, dspark_proposer.py:905), so store_step reads current_slot_mapping() and gets running_bs draft slots. DP>1 + MTP k>=2 / DSpark + a pure middle chunk raises forwarded N tokens but its slot_mapping has running_bs slots. Snapshot current_slot_mapping() before the propose and pass it, as postprocess does.

[sudhu2k]

Fixed. A snapshot alone wouldn't be enough here: the drafter's decode metadata is a view of the same forward_vars slot_mapping buffer, and eagle_proposer.py:888/892 writes it in place. A saved reference would read draft slots in its first running_bs entries, so the crash becomes silently wrong routes. postprocess is only safe because its gather is queued inside prepare_sampled_ids, ahead of the propose. So the middle-chunk branch now stores its routes before the aligning propose, and stream order covers the rest. Added test_output_less_forward_stores_routes_before_the_aligning_propose to test_model_runner_staging_fence.py, which checks that order in forward.

[valarLip]

[2] Both quant methods' select_experts go through _select_logical (L506), which has the layer and returns logical ids before place_experts. One capture there replaces this, L5524 and comm_fused_moe.py:370. A path that misses a pasted hook returns the slot's previous occupant's routes, not -1.

[sudhu2k]

Done. Capture now lives in _select_logical, which every quant method's select_experts returns through.

[valarLip]

[3] This gather is sized to the step's tokens plus the whole cached prefix, int32 then an int16 copy, and the profile run never allocates it. Estimate tokens x L x K x 6 B: ~290 MB for a 100k-token prefix hit on R1, from HBM the planner already gave to KV. Bound it per step or read the prefix at finish.

[sudhu2k]

Agreed, the prefix-hit gather isn't bounded, and the profile run never sees it. It only happens on the first step after a prefix-cache hit, and the memory is freed once the D2H completes. Splitting routed_prefix_ranges across steps would trip the gap check, and reading at finish needs a new runner round-trip, so I'd like to do this as a follow-up rather than grow this PR.

[valarLip]

[4] (also A4) int16 holds every id and -1; int32 doubles the PAGE surcharge (estimate: R1 fp8 ~5.3% of KV blocks vs ~2.7%, Qwen3-235B TP8 fp8 ~11%) and adds a cast kernel per gather. The layer axis also charges dense layers that are always -1.

[sudhu2k]

We kept int32 on purpose. topk_ids come out of routing as int32, so capture scatters them with no conversion. An int16 buffer adds a cast kernel per MoE layer per forward: on MI355X under a CUDA graph the performance gets worse. Switching the buffer to int16 while keeping today's index_put scatter makes capture slower: about 190 µs → 355 µs per forward across R1's 61 MoE layers. That's +165 µs, from one extra cast kernel per layer, and it adds to every forward, decode steps included. The memory saving is real (R1 about 5.3% to 2.7% of KV blocks) but only matters when serving is KV-bound. If we want both, the way to get them is a fused scatter that casts in-register, which we can do as a follow-up.

[valarLip]

[5] (also A4) Only rank 0 reads, but every rank allocates the buffer and scatters every MoE layer. Keep the PAGE charge on all ranks so block counts match, and create the capturer on rank 0 only; the hook is already a no-op without one.

[sudhu2k]

Done. The capturer is now created on TP rank 0 only, at startup and on wake, so the other ranks' MoE hook does nothing and they skip the buffer and the per-layer scatter. Every rank still charges the capture rows to its KV pool, so block counts match across ranks. store_step's separate rank check is gone, since "no capturer" already covers it. Added test_only_tp_rank0_allocates_the_capture_buffer.

[valarLip]

[14] The final length is known at admission (prompt + max_tokens - 1); allocate once instead of doubling and copying the history. _routed_experts_final and the copy-while-running property mode are only exercised after finish, where seq.is_finished already says so.

[sudhu2k]

Dropped _routed_experts_final and the copy-on-read mode. routed_experts is now a view of the live rows, and the late-patch guard uses seq.is_finished. finish_routed_experts still trims and freezes, so the pickle carries only the live rows.

I'd keep the doubling instead of allocating at admission. max_tokens is a cap, not the length: RL rollouts set it to the response budget, and most responses stop at EOS well short of it. At about 1 KB per row (61 layers × top-8 × int16), 256 concurrent requests with max_tokens=32k would reserve around 8 GB of host memory per engine, mostly for rows that are never written. Doubling copies each row at most about once more over the request's life, and the spare capacity is trimmed at finish.

[valarLip]

[15] Not in docs/configuration_guide.md (Config / EngineArgs tables) or docs/scheduling_kv_cache_guide.md. This comment and the CLI help list only the config-time refusals; dense models, V4/V4.1 builders and "mono turns off" are found only after weight load.

[sudhu2k]

Added docs.

[valarLip]

A2 still open. The gap check could compare a patch's start against seq.num_tokens - 1, which postprocess already maintains, rather than truncating a second counter every step.

[sudhu2k]

Keeping it. Checking against num_tokens - 1 would stop the check from catching a lost patch. postprocess advances num_tokens from the sampled tokens whether or not that step's routes arrived, so the next patch always starts within num_tokens - 1 and passes. Its rows would land after a hole of np.empty rows: uninitialized garbage, not -1, which the trainer would replay as real expert ids.

[valarLip]

nit: same name as the optional builder hook checkpoint_page_unit_bytes() it reads, with a different meaning (resolved vs override). A resolve_ prefix would separate them.

[sudhu2k]

Renamed to resolve_checkpoint_page_unit_bytes, so it no longer shares a name with the builder hook it reads. I also fixed the hook's docstring in gdn_attn.py, which still described the old fallback.

[valarLip]

nit: the only R3 hook not gated on self._return_routed_experts.

[sudhu2k]

Gated on self._return_routed_experts like the other hooks.

[Copilot]

🔵 Needs a closer look

Human assessment is needed because the change contains a moderate decode-loop blocking risk and a critical lifecycle issue where stale deferred route state can survive KV capture reset.

12 open issues

• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

1 resolved since last review

• Correct the routed-expert DP-attention contract comment

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

Approval should be deferred because the change spans correctness-critical CUDA-graph, async D2H, TP/DP, and MTP paths without available GPU workloads, CI results, or a runnable PR head to validate the advertised behavior.

12 open issues

• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

A critical active finding remains: DP attention with prefill token graphs can overwrite real slot-0 routing data with dummy execution routes, so the capture path is unsafe until this topology is rejected or handled.

13 open issues

• Dummy prefill capture can overwrite cached routes in slot 0 · New
• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

A moderate active finding shows the DeepSeek-V4 mono path raises an unrelated AttributeError instead of cleanly refusing the unsupported routed-expert topology, so the startup check still requires correction.

14 open issues

• Dummy prefill capture can overwrite cached routes in slot 0
• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• V4 mono startup check crashes on missing class attribute · New
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[Copilot]

🟡 Changes recommended

Critical unresolved issues remain in IQ2R route capture and memory-manager teardown, including a missing quantization path and a potential asynchronous write race, while the documentation also needs reconciliation.

17 open issues

• Capture routed experts in the IQ2R MoE path · New
• Drain pending route copies before clearing held routes · New
• Dummy prefill capture can overwrite cached routes in slot 0
• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• V4 mono startup check crashes on missing class attribute
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Reconcile Qwen4-Exp R3 support docs with implementation · New
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

Two critical active issues remain: IQ2R routing bypasses capture and can yield missing rows, while releasing pending routes without synchronizing async copies can cause buffer reuse races.

17 open issues

• Drain pending route copies before clearing held routes
• Capture routed experts in the IQ2R MoE path
• Dummy prefill capture can overwrite cached routes in slot 0
• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• V4 mono startup check crashes on missing class attribute
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Reconcile Qwen4-Exp R3 support docs with implementation
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

A critical active finding shows the capture hook is bypassed on the IQ2R quantization path, causing routed-expert capture to fail closed for those layers and requiring a code fix plus regression test.

17 open issues

• Drain pending route copies before clearing held routes
• Capture routed experts in the IQ2R MoE path
• Dummy prefill capture can overwrite cached routes in slot 0
• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• V4 mono startup check crashes on missing class attribute
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Reconcile Qwen4-Exp R3 support docs with implementation
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}

[Copilot]

🔵 Needs a closer look

An active moderate issue leaves IQ2R/GLM-5.3 replay data incomplete because its routing bypasses the capture hook, with no coverage test confirming the corrected path.

17 open issues

• Drain pending route copies before clearing held routes
• Capture routed experts in the IQ2R MoE path
• Dummy prefill capture can overwrite cached routes in slot 0
• Clear reused routing slots before forward execution
• Capture routes in DeepSeek V4.1 mono decode
• Exclude torch.nn.Module from fused MoE module detection
• Reject MoE decoders that bypass route capture hooks
• Disable route capture during draft model execution
• Handle backends without slot_mapping during route capture
• Initialize routed-expert capture before level-3 compilation
• Adding the capture surcharge to the shared PAGED_CLASS changes plan.entry_bytes["kv"], which is…
• V4 mono startup check crashes on missing class attribute
• Reject DP attention before exporting incomplete routes
• Size capture geometry using the full decoder layer count
• Reconcile Qwen4-Exp R3 support docs with implementation
• Test KV-cache refusal through runner entry points
• This capture-only slot lookup runs on every forward even when enable_return_routed_experts is…

_{Review effort: Lite}