2026-09-04 · config A as investigator · 129 tool calls · Atlas tree only

Atlas decode anatomy — the first throughput investigation, verbatim

Report

Where ~60 ms per decoded token was thought to go before anything was measured; several of its estimates were later overruled by the nsys profile.

Atlas Flash-Next decode throughput investigation

Scope note on evidence quality: everything below is traced in the tree at main @ 4df6e609 with path:line. Where I could not confirm a number from code (per-kernel times, the bench operator's tok/s split, the five open questions' measured values), I say so explicitly and mark the claim as inferred. The repo itself is unusually well instrumented, and several of its own comments are the strongest available measurements — I cite those as reported-by-comment, not as my own measurement.

1. Per-token decode anatomy (qwen4_exp, single stream, batch 1)

Entry point: TransformerModel::decode in crates/spark-model/src/model/trait_impl/decode_a.rs. The step is split into an explicit Phase 1 (outside any graph) and Phase 2/3 (graph replay or capture) — and the central structural fact is that for this model Phase 2/3 never happens: CUDA graphs are vetoed.

Graph gating — the decisive branch

decode_a.rs:250 computes layer_veto = self.layers.iter().any(|l| l.decode_graph_unsupported()), and use_graphs is false whenever any layer vetoes (decode_a.rs:266-275). Two layer families veto unconditionally on presence:

QSA attention layers: crates/spark-model/src/layers/qwen3_attention/trait_impl.rs:110-112 → self.qsa.is_some(), with the comment "QSA selection does a host top-k per step — never capturable".
PLE/SSM layers: crates/spark-model/src/layers/qwen3_ssm/trait_layer.rs:258-263 → self.ple.is_some(), "PLE's per-seq host hash … is capture-illegal (pageable reads); the single-decode path prestages around it, the batched path does not".

So on Qwen3.8-Flash-Next (12 QSA layers and 36 PLE-bearing SSM layers) use_graphs == false for every token. decode_forward_body therefore runs eagerly, and the ~1,500–2,000 launches per token are all host-driven cuLaunchKernels. The code even acknowledges the eager launch cost it would have removed: decode_a.rs:224-240 describes the GDN decode-graph hatch as removing "the per-token host launch cost that dominates" — the same cost we are looking for.

Additional suppressors that would bite even without the veto: seq64_boundary defers capture every 64th step (decode_a.rs:255-265, a real 716/700 misaligned-address fault on first replay), hss_engaged (:212), self.profile (:267), lora_eager (:246).

Ordered stage list for one decoded token

Phase 1 — eager host work, before any replay/capture (decode_a.rs:90-176)

1. Optional arena zero — skipped for this model: only when kv_lora_rank > 0 && o_lora_rank == 0 (decode_a.rs:100-102). Good: avoids a ~1.7 GB memset/step.

2. embed_ctx(&seq.tokens, token, …) (decode_a.rs:107) — token embedding; seq.tokens is the pre-append history, which is the n-gram contract.

3. ensure_blocks_through_decode (decode_a.rs:112-120) — KV block allocation + block-table maintenance.

4. Five separate pageable H2D uploads on the compute stream: position (:126-127), global KV slot (:133-134), actual_seq_len (:136-138), the whole block table as i32 (:140-145), and the decode token id (:150-151). Each copy_h2d_async is a small memcpy; the block-table upload is re-sent in full every token.

5. upload_seq_slot_uniform (decode_a.rs:174-175) — 6th H2D.

6. l.decode_prestage(...) for all 48 layers (decode_a.rs:157-164) → qwen3_ssm/trait_layer.rs:301-310 → PleLayer::prestage (crates/spark-model/src/layers/ple/aux_state.rs:99-127). Per PLE layer this does, on the host:

st.history_ckpt = st.history.clone() (aux_state.rs:113) — a per-layer, per-token Vec clone;
ple_ngram_ids(&self.dims, &window) (aux_state.rs:117) — the n-gram hash over the context window for all ngram_heads;
gather_host(&flat, …) (aux_state.rs:120 → ple/verify.rs:115-179), which
takes a global mutex on the shared n-gram table (verify.rs:121-124),
gpu.synchronize(stream)? — a full stream barrier (aux_state.rs:153-165, release_prev_pins; the comment states plainly: "Costs one stream sync per forward. An event recorded after gather_embed would avoid the full barrier; correctness first."),
cache.resolve(ids, &mut slots) — row→slot lookup plus NVMe fault-in of missing rows (verify.rs:137),
a pageable H2D of the slot list into slots_dev (verify.rs:156-157).
NVMe cost model: rows are ~90 B read as 4 KiB O_DIRECT blocks, one block per row plus a second when a row straddles a boundary, plus a second 4 KiB read for the per-row FP8 scale when a scale file exists (crates/spark-storage/src/ngram_cache_fault.rs:99-133). Fault parallelism is ATLAS_PLE_FAULT_THREADS, default 32 (ngram_cache_fault.rs:29-34), measured there at 9.55× over serial for a 22,462-miss prefill (~74 µs per issued serial read).
Inferred but important: the sync at step (b) is per PLE layer, and it is a barrier on the compute stream. On an eager step the stream is nearly drained by the time prestage runs, so the barrier is cheap; the serialization cost is that all 36 layers' host-hash + NVMe fault-ins run strictly sequentially with respect to each other and to the GPU forward, under one table mutex.

7. ForwardContext construction (decode_a.rs:278-300), including host_token_ids: Some(slice) so PLE never reads token ids back from the device (decode_a.rs:294-296, ple/layer.rs:278-313).

Phase 2 — graph replay: decode_a.rs:309-329. Unreachable for this model (use_graphs == false ⇒ graph_cache is None).

Phase 3 — eager body: decode_forward_body (crates/spark-model/src/model/trait_impl/decode_a3.rs:30-118), a plain loop over 48 layers (:41-53), then final RMS norm (:103-113) and self.lm_head(normed, stream) (:116).

Per SSM/GDN layer (36×), via qwen3_ssm/trait_layer.rs:312+ → trait_decode_hc.rs (stage names at trait_decode_hc.rs:119): hc_pre_attn → GDN block → hc_post → hc_pre_ffn → MoE → hc_post (+ hc_head on the last layer).

Per attention layer (12×), qwen3_attention/trait_impl/decode_inner.rs:

hc_pre_attn (:678-692 region for FFN; attn equivalent earlier), attention forward (:576), optional post-attn norm (:596-608), hc_post (:644-656), hc_pre_ffn (:678-692), rms_norm or a copy_d2d for the block-input norm (:718-732 — on Qwen the low-rank variant takes the copy_d2d branch, i.e. one extra memcpy per layer, :731), FFN/MoE (:734), optional post-FFN norm (:737-749), hc_post (:755-767), hc_head on the last layer (:785-798).

QSA indexer per attention layer (crates/spark-model/src/layers/qsa.rs:330-450) — this is the worst per-token offender:

cublas_bf16_proj_dense qk projection (:331-340), a copy_d2d_async of the raw key (:341-346), pool_new_blocks (:348), qsa_qprep (:357-370), qsa_score (:371-381);
gpu.copy_d2h_on_stream(self.scores_dev, &mut raw, stream) (:386) — and copy_d2h_on_stream_impl ends in cuStreamSynchronize(stream) (crates/spark-runtime/src/cuda_backend/gpu_copy.rs:92-137, sync at :128). That is a hard host↔device round trip inside the layer loop, 12 times per token, each one draining everything enqueued so far.
host sort_by over complete block scores (:391-401) — an O(n log n) host sort over the whole visible context in 128-token blocks;
two more pageable H2D uploads (:417 selection list, :441 n_sel), plus qsa_gather (:418-432) which materialises the selected K/V into a scratch "paged cache" that the dense attention kernel then reads.

MoE: 512 experts, top-10, NVFP4. In the arms I could read, MoE is the single largest GPU stage: qwen3_ssm/trait_layer.rs:136-142 reports ATLAS_QWEN4EXP_VERIFY_PROF means of moe 287.6 µs/layer-call at K=2 vs hc_pre_attn 184.9 µs, hc_pre_ffn 183.8 µs, gdn_block 16.8 µs.

HC (hyper-connections): the low-rank (rank 320) family, dispatched in crates/spark-model/src/layers/ops/hyper_connection_dispatch.rs:75-124 / 134-175 / 183-225. At decode (num_tokens <= 64) hc_pre_lowrank and hc_head_lowrank deliberately take a 3-launch split path (crates/spark-model/src/layers/ops/hyper_connection_lowrank.rs:64-79 and :141-156) because the fused grid=[T] kernel is one block on one SM streaming ~13 MB of weights — "measured 2.0 ms; the whole token was 96× that" (hyper_connection_lowrank.rs:59-63), i.e. the historical eager token was ~192 ms. The split path trades that for 3 launches × 96 sites.

Head: decode_a3.rs:116 calls self.lm_head(normed, stream). The batched-GEMV ladder (crates/spark-model/src/model/trait_impl/lm_head_batched.rs:31-39, ON unless ATLAS_NO_LM_HEAD_BATCH_GEMV=1) is used by the mixed head (decode_b2.rs:58-68) and by decode_a2.rs:579, but the single-token body still calls the plain lm_head. For a [248320, 5120] NVFP4 head that is ~254 MB streamed per token (decode_b2.rs:60-63), which at GB10-class bandwidth is a floor of a few ms — real but not the 60 ms.

Launch/sync/host-IO tally per token (counts are code-derived; the per-launch cost is not knowable from source):

Stage	Count/token	Graph?	Host syncs	NVMe
Phase 1 uploads	6 H2D	eager	0	0
PLE prestage	36 × (hash + resolve + 1 H2D)	eager	36 (`aux_state.rs:158`)	up to 36 × 16 rows × 4 KiB (×2 with scale file)
HC pre/head	96 sites × 3 launches = 288	eager (uncaptured)	0	0
HC post	96	eager	0	0
QSA indexer	12 × (1 cuBLAS + 2 kernels + 1 gather + 3 copies) ≈ 84	eager	12 (`qsa.rs:386`)	0
GDN block / attention / norms / d2d	~10–20/layer ≈ 500–700	eager	0	0
MoE	48 × (router + 10 experts) ≈ 500+	eager	0	0
final norm + lm_head	2	eager	0	0
Total	~1,500–2,000 launches	0 captured	~48	≤576 4 KiB reads

2. Answers to the five open questions

I must be upfront: read_findings was not re-readable inside this budget-exhausted session, so I answer the five questions as the tree frames them (the questions the code itself poses). If the operator's wording differs, the mapping is: (a) where the time goes, (b) is it launch-bound, (c) is it PLE/NVMe, (d) is it QSA, (e) what to change first.

(a) Where does the ~60 ms go? Three terms, in order of confidence:

1. Eager launch + host-serialization overhead across ~1,500–2,000 kernels, because graphs are vetoed (decode_a.rs:250,266-275; vetoes at qwen3_attention/trait_impl.rs:110 and qwen3_ssm/trait_layer.rs:261). This is the only term that is proportional to the whole model rather than to one stage, and it is the term the codebase itself names as dominant in the analogous GDN case (decode_a.rs:236-238).

2. 48 hard stream barriers inside the step: 12 QSA copy_d2h_on_stream (qsa.rs:386 → gpu_copy.rs:128) and 36 PLE release_prev_pins (aux_state.rs:158). Each barrier converts GPU-parallel work into a host round trip and prevents any cross-stage overlap.

3. MoE and HC weight streaming: MoE ≈ 287.6 µs/layer-call and HC pre ≈ 185 µs/site at K=2 (qwen3_ssm/trait_layer.rs:136-142). Even reading those as verify-width numbers, they are the two largest GPU stages, and 96 HC sites/token is the multiplier.

Not the head: trait_layer.rs:197-201 reports a probe giving "norm 9 us / lm_head 1508 us / argmax 289 us", and explicitly rules the lm_head/argmax tail out. Not the SSM recurrence: gdn_block 16.8 µs (:141).

(b) Is it launch-bound? Almost certainly yes, and it is provably graph-free by construction rather than by misconfiguration — the veto is a any() over layers, so a single QSA or PLE layer disables capture model-wide (decode_a.rs:250). I could not measure launch cost from source; the decisive experiment is ATLAS_NO_DECODE_GRAPHS=1 vs the current state (decode_a.rs:253) which will show no difference here, plus a nsys launch-count/duty-cycle trace.

(c) Is PLE/NVMe on the critical path at decode? Partly, and mostly as serialization, not bandwidth. Decode asks for 16 ids/layer (verify.rs:141: "Decode-scale (16 ids) stays at debug"), the row cache is O_DIRECT 4 KiB reads with QD=32 (ngram_cache_fault.rs:29-34), and misses are compulsory first-touch ("65536 → 1048576 slots changed the miss count by zero", :13-15). A warm steady-state decode should be near-all-hits, in which case the cost is the 36 mutex acquisitions + 36 stream syncs + 36 pageable H2Ds, not the disk. I could not confirm the decode-time hit rate from code — it is logged at debug (verify.rs:150-155) and must be read off a live run.

(d) Is QSA helping or hurting? As implemented it is a net loss at decode: it (i) vetoes graphs for the whole model (qwen3_attention/trait_impl.rs:110-112), (ii) inserts a host round trip per layer per token (qsa.rs:386), (iii) does a host sort over all visible blocks (qsa.rs:391-401), and (iv) then pays a qsa_gather into a scratch paged view (qsa.rs:418-432) before attention runs. It only pays off once complete > block_topk — qsa.rs:350-354 returns None and takes the exact dense path while the context is short. So at short-to-moderate context, QSA costs the graph and buys nothing.

(e) What is the repo's own unresolved mystery? Worth stating so nobody re-derives it: the K=3-vs-K=2 verify cost has +26.3 ms "unattributed" in the non-SSM half of the forward, with four prior explanations explicitly refuted in qwen3_ssm/trait_layer.rs:130-221. That is a verify-path number, not the single-token decode number, but it is direct evidence that ~26 ms of a step in this model is currently unexplained by any stage the authors instrumented — consistent with my (b): per-kernel launch/host overhead, which their stage timers (all bracketed by synchronize, e.g. decode_inner.rs:561-575) attribute to whichever bucket surrounds them.

3. Ranked optimization plan (smallest-safe first)

Config / serve-flag changes (no code)

1. Measure before touching anything: get the real anatomy.

Change: run the serve line (serve_qwen4exp_tui.sh, which already raises ATLAS_PLE_MAX_TOKENS, QWEN4_EXP_PORT_LOG.md:339) with ATLAS_QWEN4EXP_ATTN_PROF=1 (decode_inner.rs:561) and ATLAS_QWEN4EXP_VERIFY_PROF=1 (trait_layer.rs:134), plus RUST_LOG=debug to capture the PLE line "PLE gather: 16 ids, H hits / M misses, resolve Uus" (verify.rs:150-155), and one nsys profile of 100 decode steps.
Gain: 0 directly; it discriminates the three hypotheses (launch-bound / barrier-bound / bandwidth-bound) for ~free.
Risk: the profilers insert their own synchronizes (decode_inner.rs:562,570) and ATLAS_NO_DECODE_GRAPHS is force-enabled when verify tracing is on (decode_a.rs:253-254) — read them as relative stage costs only.
Verify: decode tok/s from the server log; kernel-launch count and GPU duty cycle from nsys; PLE decode hit-rate.

2. ATLAS_D2H_SPIN_SYNC=1.

Change: one env var. It replaces the blocking cuStreamSynchronize after every on-stream D2H with a cuStreamQuery spin loop (gpu_copy.rs:108-129).
Why: the driver comment documents a measured pathology — "a blocked thread appears to leave the submission ring unflushed until a ~130 ms driver housekeeping tick doorbells the GPU; each query forces a flush" (gpu_copy.rs:110-113). We hit that sync 48 times per token (12 QSA + 36 PLE). If even a small fraction of those ticks land, this is free throughput.
Risk: burns a core spinning; can increase latency if the GPU genuinely has work left. Env-gated, instant revert.
Verify: A/B decode tok/s over ≥500 tokens; also watch that tok/s variance drops (the signature of a periodic stall).

3. ATLAS_PLE_FAULT_THREADS sweep (32 → 64/128) and confirm the row cache is warm.

Change: env only (ngram_cache_fault.rs:29-34).
Why: the file's own table shows monotone gains with depth (1→32 = 9.55×) and states "the depth is the whole lever" (:9-26). The table was measured on prefill; decode has 36 concurrent per-layer resolves under one mutex, so deeper QD may or may not help — cheap to test.
Risk: I/O contention with the KV/prefix paths on a shared NVMe device.
Verify: the resolve {us} field in the PLE gather log at decode, and TTFT + tok/s.

4. Do not expect anything from graph flags. ATLAS_GDN_DECODE_GRAPH=1 / ATLAS_EP_GRAPHS=1 (decode_a.rs:223,240) are OR'd into the use_graphs precondition but && !layer_veto (:274) still kills it. Anyone "testing CUDA graphs" on this model without removing the veto is testing nothing. This is a negative config finding worth recording.

Code changes

5. Make QSA's host top-k device-side, then lift the QSA graph veto. (highest expected gain; largest work)

Files: layers/qsa.rs:350-450 (replace copy_d2h_on_stream + host sort_by + 2 pageable H2Ds with a device top-k writing sel_dev directly), qwen3_attention/trait_impl.rs:110-112 (return false once the path is capture-safe), decode_a.rs:250.
Why: removes 12 host round trips and unlocks capture of the whole 48-layer body — the ~1,500–2,000 launches collapse into one replay. This is the only change that attacks term (1) and (2) simultaneously.
Risk: correctness-critical. The comment at decode_a.rs:247-249 is explicit that a graph captured on the dense path "would silently replay WRONG attention once selection activates". Device top-k must reproduce torch's tie-break "by LOWER index" (qsa.rs:392-399) bit-for-bit or outputs diverge. Also the sel length is data-dependent (:404-414) — it must be padded to a static budget for capture. The seq64_boundary 716-fault (:255-264) shows this codebase's capture path already has one latent alignment bug; assume another.
Verify: greedy-decode parity vs eager over ≥200 prompts (bench/qwen4_exp/ has golden harnesses); then tok/s; then confirm CUDA graph capture: appears in the log (decode_a.rs:338) and that replay (:325) is the common path.

6. Replace PLE's per-layer full barrier with a recorded event.

Files: ple/aux_state.rs:153-165 (release_prev_pins) — record an event after gather_embed (ple/verify.rs:183-192) and cuEventSynchronize/cuStreamWaitEvent on it instead of gpu.synchronize(stream); keep ATLAS_PLE_VERIFY_SNAPSHOTS-style hatch (ple/layer.rs:65-69) for bisection.
Why: 36 full-stream barriers/token → 1 event. The code already says the barrier is a deliberate deferral: "An event recorded after gather_embed would avoid the full barrier; correctness first" (aux_state.rs:151-152). The pin must still cover the kernel (aux_state.rs:129-149 — the historical garbage-embedding bug), so the event must be recorded after the gather launch, not before.
Risk: moderate — a too-early release reintroduces the documented "user has simply sent FDFDFDFFDFD" corruption (aux_state.rs:138-142). Guard with gather_guard.rs.
Verify: tok/s A/B; and a long-cache stress run (many sequences, cache full so evictions happen — the bug "grew worse as the cache filled", :142) checking output sanity.

7. Cut HC launches from 3 to 1 per site at decode by splitting along rank/hidden, not along the pipeline.

Files: ops/hyper_connection_lowrank.rs:64-79 and :141-156 (hc_pre_split), kernels in kernels/gb10/qwen3.8-flash-next/nvfp4/hyper_connection.cu:249-540 (hc_pre_stage/hc_pre_down/hc_pre_finish, plus the hc_pre_stage_bf16/hc_pre_mix pair at :555-667).
Why: 96 sites × 3 launches = 288 launches/token, and the sites are the second-largest GPU stage (~185 µs reported, trait_layer.rs:139-140). The fused kernel's problem was grid=[T]=1 on one SM (hyper_connection_lowrank.rs:59-63), which a grid=[T, k] split-K variant fixes without the 3-launch round trip.
Risk: numerics — the GEMM variant already "rounds normed to BF16 before the projections" (:25-28), and HC is on the residual highway, so drift is global. Keep ATLAS_QWEN4EXP_NO_HC_GEMM=1 (:30-33) as the kill switch.
Verify: VERIFY_PROF per-site µs before/after; greedy parity; tok/s.

8. Batch the Phase-1 metadata uploads into one pinned-memory H2D.

Files: decode_a.rs:122-151 (position, slot, seq_len, block table, token id) + :174-175.
Why: 6 pageable H2Ds/token, one of them the full block table, on the critical path before any compute. One cuMemcpyHtoDAsync from a pinned staging buffer into the same meta_base layout (+0/+8/+16/+128/+256, :177-188) is a drop-in, layout-preserving change.
Risk: low — the layout is already documented as graph-safe fixed addressing (:166-173). Must keep using a pinned ring so the host buffer isn't freed before the copy completes (gpu_copy.rs:59-67 synchronizes today precisely because memory is pageable).
Verify: launch/copy count in nsys; small but consistent tok/s gain.

Explicitly not recommended yet

lm_head work. The head is ~1.5 ms and already ruled out by measurement (trait_layer.rs:197-201). Note the inconsistency worth a one-line fix later: the single-token body calls self.lm_head (decode_a3.rs:116) while every other decode head uses the batched ladder (decode_a2.rs:579, decode_b2.rs:68) — worth unifying for numerics parity, not for speed.
Chasing the K=3 "+26.3 ms unattributed" (trait_layer.rs:191-215). That is the verify path, and the file's own guidance is to first build a bucket "verified to fire" before trusting any new number.

What I could not confirm from the code

Any absolute per-kernel or per-stage decode latency. All µs figures cited are reported-by-comment and mostly from verify (K=2/K=3) widths, not single-token decode.
The decode-time PLE row-cache hit rate (logged at debug, verify.rs:150-155) and therefore whether NVMe is on the decode critical path at all.
Whether the bench operator's 16 tok/s run had MTP/speculative decode enabled; the K=2 numbers in trait_layer.rs:186 (80.01 ms/step @ 1.715 tok/step ≈ 21 tok/s) are close to the reported 16 tok/s and would change the whole anatomy if that is the same configuration.
The exact MoE decode kernel and whether the 10 experts are one grouped launch or 10 GEMV launches — search_code for gemv/w4a16 sites surfaced the dense/lm_head paths but I did not reach the MoE decode dispatch inside budget. Given MoE is the largest reported stage, this is the highest-value unread file.
Whether embed_ctx (decode_a.rs:107) itself touches NVMe, or whether all n-gram I/O is confined to PLE prestage as the structure suggests.