2026-09-04 · written for the config-A executor · exact file:line anchors
Measure-first steps, anchored subtasks and GB10 verification for each item, plus the memory check that gated A3.
Context: gap study https://andrei.monkeycode.io/bench/gap/ (agent report + reviewer corrections). Measured on the GB10: llama.cpp 24–25 tok/s (42 ms/token), Atlas 15–17 tok/s (60 ms/token). CUDA graphs are worth ≈3% to llama.cpp on this model (on/off A/B on prod), so the gap is per-token host synchronisation and kernel efficiency, not graph capture. Work tree for all items: fork branch assist/flashnext-longform-fixes (head 7db64cca), spark clone ~/agentbench/atlas-work; executor = agent8 briefs (exact anchors, ≤3 subtasks).
Corrections to the agent report that these briefs already incorporate:
PLE at MODEL LAYER 1), so there is one prestage + one barrier per token, not 36. The PLE item is therefore small (≤1 ms/token); the QSA host round trips (12/token) and the MoE decode arm are the large terms.free -g: 111 GB used, 10 GB buff/cache). Atlas's row cache is slots × 160 B pinned: the derived default (~262K slots) is ~42 MB, and a 4M-slot cache is 640 MB. Per-sequence working set ≤ 16 rows/token.Why: crates/spark-model/src/layers/qsa.rs:386 copies the block scores device→host and copy_d2h_on_stream ends in cuStreamSynchronize (crates/spark-runtime/src/cuda_backend/gpu_copy.rs:92-137); the host then sorts all visible blocks (qsa.rs:391-401) and uploads the selection and its length with two pageable H2D copies (qsa.rs:417, :441). That is 12 hard stream barriers per token (12 attention layers). Because the selection is host-side, the QSA layer vetoes CUDA-graph capture for the whole model (crates/spark-model/src/layers/qwen3_attention/trait_impl.rs:110-112, consumed by the any() at crates/spark-model/src/model/trait_impl/decode_a.rs:250). llama.cpp keeps the indexer entirely on device (src/models/qwen4exp.cpp:402-512, ggml_top_k at :505).
Step 0 — measure first (no code): serve with ATLAS_QWEN4EXP_ATTN_PROF=1 and ATLAS_QWEN4EXP_VERIFY_PROF=1 (hooks documented at crates/spark-model/src/layers/qwen3_ssm/trait_layer.rs:134,224) on a 2K-token prompt decoding 500 tokens, and record per-stage µs; run nsys profile over 100 decode steps and count cudaStreamSynchronize and kernel launches per token. These numbers are the baseline every later step is judged against.
Subtask 1 — device-side selection. In qsa.rs decode_select (lines 300-450): replace the D2H copy, the host sort_by and the two H2D uploads with one device kernel that (a) selects the top-block_topk block scores over the complete blocks with torch's tie-break (lower index wins on equal score — this is what the host sort at :392-399 implements and what the reference model expects), (b) writes the selected block ids into self.sel_dev padded to a static block_topk budget, and (c) writes n_sel into self.seq_len_dev. Keep the host path behind ATLAS_QSA_HOST_TOPK=1 for A/B. Kernel home: kernels/gb10/qwen3.8-flash-next/nvfp4/ next to qsa_score / qsa_gather; a radix-select or a bitonic top-k over ≤ 512 blocks (65536 tokens / 128) is enough. qsa_gather (qsa.rs:418-432) then reads sel_dev unchanged.
Subtask 2 — lift the veto. Once Subtask 1 is in, make decode_graph_unsupported return false for QSA layers (qwen3_attention/trait_impl.rs:110-112). Re-read the PLE veto (crates/spark-model/src/layers/qwen3_ssm/trait_layer.rs:258-263): the single-decode path already prestages the PLE host work before replay (ple/verify.rs:113-116 says so explicitly), so the veto there should be limited to the batched path. Watch the other suppressors at decode_a.rs:212,246,255-267.
Subtask 3 — parity gate. Greedy-decode parity against the eager path over ≥200 prompts from bench/qwen4_exp/ (token-for-token), then tok/s. Confirm CUDA graph capture: appears in the log (decode_a.rs:338) and replay (:325) is the common path.
Expected: removing 12 syncs/token is the real gain (several ms); graph capture on top is a small extra (llama.cpp measures ≈3% for it). Risk: correctness-critical — a graph captured on the dense path "would silently replay WRONG attention once selection activates" (decode_a.rs:247-249), and the tie-break must match bit-for-bit. Keep the host path switchable until parity is proven.
Why: MoE is the largest GPU stage in Atlas's own profile (trait_layer.rs:136-142: ≈290 µs per layer call at K=2 → 48 layers ≈ 14 ms/token at verify width). The single-token decode arm is the "scalar FP32 FMA decode path" (crates/spark-model/src/layers/moe/forward.rs:88-92); the tensor-core grouped-GEMM arm exists but is reachable at M=1 only for DFlash capture layers behind ATLAS_FRANKENSTEIN_DECODE_VIA_PREFILL=1 (forward.rs:93-113, measured ≈250 µs/layer there), and the width-gated grouped decode arm only engages for wide batches (crates/spark-model/src/layers/moe_grouped_decode_tests.rs). llama.cpp runs the same 512×top-10 MoE as three mul_mat_id launches per token through mul_mat_vec_q (vectorised quantised GEMV, ggml-cuda.cu:1916-1941, mmvq.cu).
Step 0 — measure first (no code): at K=1 decode, per-layer MoE µs with ATLAS_QWEN4EXP_VERIFY_PROF=1; compute the bandwidth floor = bytes of the 10 active experts' gate/up/down NVFP4 weights per layer ÷ measured GB10 stream bandwidth (≈200–230 GB/s achievable); ncu --metrics dram__bytes_read.sum on one moe_expert_gate_up_shared and one moe_expert_silu_down_shared launch (forward.rs:502, :584). If the kernels already sit at the floor, A2 is done and the gap is elsewhere.
Subtask 1 — A/B the existing arms at M=1 (env only): extend the Frankenstein gate so ATLAS_FRANKENSTEIN_DECODE_VIA_PREFILL=all routes EVERY layer's single-token MoE through forward_prefill(M=1) (forward.rs:93-113), and measure tok/s and per-layer µs against the scalar arm.
Subtask 2 — vectorised NVFP4 GEMV for decode (code): if the scalar arm is well above the floor, write a w4a16 expert GEMV in the style of llama.cpp's mmvq (each warp handles 4–8 output rows, dequantise NVFP4 nibbles + FP8 group scales in registers, BF16 activations, top-10 experts and the shared expert in one launch per projection) and dispatch it from moe_expert_gate_up_shared / moe_expert_silu_down_shared. Keep the 2-launch structure (gate+up fused, silu+down fused). Kernel home: kernels/gb10/qwen3.8-flash-next/nvfp4/ (model-specific override; the generic ones live in kernels/gb10/*/nvfp4/). Keep the scalar arm behind ATLAS_MOE_DECODE_SCALAR=1.
Subtask 3 — verify: per-layer µs before/after; ncu bytes-per-token within 1.3× of the floor; logprob parity vs the scalar arm over 50 prompts (max |Δ logprob| on the argmax token < 1e-2); tok/s.
Expected: if the floor is ~5–6 ms/token for 48 layers and Atlas spends ~10–14 ms, this is the single largest item. Risk: numerics of the fused dequant; keep the fallback switch.
~/agentbench/briefs/a3.md)Memory check (the question asked before starting): unified memory is enough for what A3 actually needs, and NOT enough for a fully resident table — for either engine:
free -g reports 0 GB available after load). The 51.2 GB FP8 table cannot be resident; even a 4-bit copy (25.6 GB) cannot.free -g while prod runs: 111 GB used, 10 GB cache).ATLAS_PLE_CACHE_SLOTS=4194304) is 640 MB + 16 MB of scales and removes all capacity misses; the remaining misses are compulsory first-touch reads (4 KiB O_DIRECT preads, 32-deep), which are the same physics as llama.cpp's page faults.So A3 = (1) a bigger row cache by env (zero code, ~0.7 GB), (2) replace the one full stream barrier per token with an event recorded after the gather kernel, (3) stage the slot upload through pinned memory. Expected gain is small (≤1 ms/token) because PLE is one layer; it is cheap and it removes a barrier that would also block graph capture later. The executor brief is anchored at ple/aux_state.rs:153-165, ple/verify.rs:126,154-158,183-192, spark-runtime/src/gpu/backend.rs:334,360.
Benchmark after A3: driver_td3.sh ATLAS5 8192 with EXTRA="--attn-proj-dtype bf16 --content-loop-watchdog false" and ATLAS_PLE_CACHE_SLOTS=4194304, compared with ATLAS4 (same build minus A3): R1/R2 tok/s and the PLE gather resolve times in the serve log.
As handed to the executor (first pass).
Repo: Atlas (work tree = branch assist/flashnext-longform-fixes, current head 7db64cca). One commit. Model: Qwen3.8-Flash-Next (qwen4_exp). The PLE n-gram embedding sits on ONE model layer; every decoded token runs exactly one PleLayer::prestage on the host before the forward.
1. crates/spark-model/src/layers/ple/aux_state.rs:153-165 — release_prev_pins does a FULL
gpu.synchronize(stream) once per forward (line 158). Its own comment says: "An event recorded after gather_embed would avoid the full barrier; correctness first." The barrier exists only so that the pins taken by the previous gather outlive the previous gather KERNEL.
2. crates/spark-model/src/layers/ple/verify.rs:156-157 — the resolved slot list is built as a pageable
Vec<u8> and uploaded with gpu.copy_h2d_async(&bytes, self.slots_dev, stream). A pageable source makes the driver stage the copy synchronously on the host.
File: crates/spark-runtime/src/gpu/backend.rs (the GpuBackend trait; fn synchronize is at ~line 334, fn alloc_host_pinned at ~line 360). Add four trait methods with default implementations that fall back to synchronize semantics, so every existing backend/mock keeps compiling:
/// Create a timing-disabled event. Returns an opaque handle (0 = unsupported).
fn event_create(&self) -> Result<u64> { Ok(0) }
/// Record `event` on `stream`. No-op for handle 0.
fn event_record(&self, _event: u64, _stream: u64) -> Result<()> { Ok(()) }
/// Block the host until `event` has completed. For handle 0 this must fall back to
/// `self.synchronize(stream)` — callers pass the stream for exactly that reason.
fn event_synchronize(&self, _event: u64, stream: u64) -> Result<()> { self.synchronize(stream) }
fn event_destroy(&self, _event: u64) -> Result<()> { Ok(()) }Then implement them for the CUDA backend. Find it with search_code("fn synchronize", "crates/spark-runtime/src/cuda_backend") — implement with cuEventCreate(CU_EVENT_DISABLE_TIMING), cuEventRecord, cuEventSynchronize, cuEventDestroy through the same FFI module the file already uses for cuStreamSynchronize (look at how copy_d2h_on_stream_impl in crates/spark-runtime/src/cuda_backend/gpu_copy.rs:92-137 calls the driver and reuse that pattern / error mapping). If the FFI bindings lack the cuEvent* symbols, declare them next to the existing cuStreamSynchronize extern declaration.
Files: crates/spark-model/src/layers/ple/aux_state.rs, crates/spark-model/src/layers/ple/verify.rs, and the PleLayer struct (find it with search_code("pub struct PleLayer", "crates/spark-model/src/layers/ple")).
PleLayer a field gather_done: std::sync::Mutex<u64> (event handle, 0 until created) — initialise it to Mutex::new(0) wherever PleLayer is constructed (search_code("PleLayer {", "crates/spark-model/src")).gather_embed (verify.rs:183-192), right AFTER self.gather_embed_dispatch(...) returns Ok: lazily create the event if the handle is 0 (gpu.event_create()), then gpu.event_record(ev, stream).release_prev_pins (aux_state.rs:153-165): it currently takes (table, gpu, stream). Add a parameter event: u64 and replace gpu.synchronize(stream)? with gpu.event_synchronize(event, stream)? (handle 0 falls back to the full sync, so behaviour is unchanged on backends without events). Update the single call site at the top of gather_host (verify.rs:126) to pass the current handle.File: crates/spark-model/src/layers/ple/verify.rs:154-158 (the Cached arm of gather_host).
PleLayer a field slots_staging: std::sync::Mutex<(usize, usize)> = (pinned host pointer as usize, capacity in bytes), initialised to (0, 0).Cached arm, compute need = slots.len() * 4. If the staging capacity is < need, allocate a new pinned buffer with gpu.alloc_host_pinned(need.max(64 * 1024)) (never free the old one during a forward; it is fine to leak the rare regrow), and store it. Write the little-endian u32 slots into the pinned buffer with std::ptr::copy_nonoverlapping / a from_raw_parts_mut slice, then call gpu.copy_h2d_async(pinned_slice, self.slots_dev, stream) with a slice built from the pinned pointer (unsafe { std::slice::from_raw_parts(ptr, need) }). Document the safety argument in a // SAFETY: comment: the buffer is only rewritten at the top of the NEXT gather_host, which first waits on the event recorded after the previous gather kernel (Subtask 2), so the previous copy has completed.Bf16 arm may stay as it is.1. cargo fmt (then cargo fmt --check must be clean).
2. cargo clippy --workspace --tests must be clean (the workspace denies clippy::all).
3. cargo test -p spark-runtime --lib and cargo test -p spark-model --lib must pass. Add one unit test
in spark-runtime for the default trait methods (a mock backend: event_create returns 0 and event_synchronize(0, s) calls synchronize(s)).
4. Commit as one commit: ple: event-scoped pin release and pinned slot staging and finish.
Do NOT touch prestage, the n-gram hashing, the row cache (crates/spark-storage), or any other layer. Do not change what gather_host uploads (the slot ids) — only how the bytes are staged and how the previous gather is waited for.
Second pass: the three edits the first pass left out.
The previous pass added event_create / event_record / event_synchronize_on_stream / event_destroy to GpuBackend (crates/spark-runtime/src/gpu.rs, CUDA impl in cuda_backend/gpu_impl.rs), the fields gather_done: Mutex<u64> and slots_staging: Mutex<(usize, usize)> on PleLayer (crates/spark-model/src/layers/ple/layer.rs), and changed release_prev_pins to take an event: u64 (crates/spark-model/src/layers/ple/aux_state.rs:153-165). The tree does NOT compile yet because crates/spark-model/src/layers/ple/verify.rs was not updated. Make exactly these three edits in verify.rs and nothing else.
release_prev_pins (verify.rs, inside gather_host, currently line ~126)Current: Self::release_prev_pins(&mut table, gpu, stream)?; New:
let gather_done = *self
.gather_done
.lock()
.map_err(|_| anyhow::anyhow!("PLE gather_done mutex poisoned"))?;
Self::release_prev_pins(&mut table, gpu, stream, gather_done)?;gather_embed, currently lines ~183-192)Current body: self.gather_embed_dispatch(table_va, num_tokens, heads, gpu, stream). New body:
self.gather_embed_dispatch(table_va, num_tokens, heads, gpu, stream)?;
// The pins taken by this gather must outlive THIS kernel: record an event
// right behind it; `release_prev_pins` waits on it before freeing the slots.
let mut ev = self
.gather_done
.lock()
.map_err(|_| anyhow::anyhow!("PLE gather_done mutex poisoned"))?;
if *ev == 0 {
*ev = gpu.event_create()?;
}
if *ev != 0 {
gpu.event_record(*ev, stream)?;
}
Ok(())NgramTable::Cached arm, currently lines ~154-158)Current:
let bytes: Vec<u8> = slots.iter().flat_map(|s| s.to_le_bytes()).collect();
gpu.copy_h2d_async(&bytes, self.slots_dev, stream)?;New:
let need = slots.len() * 4;
let mut staging = self
.slots_staging
.lock()
.map_err(|_| anyhow::anyhow!("PLE slots_staging mutex poisoned"))?;
if staging.1 < need {
// Grow (never shrink). The old buffer is leaked on the rare regrow:
// a gather that may still read it could be in flight.
let cap = need.max(64 * 1024);
let ptr = gpu.alloc_host_pinned(cap)?;
*staging = (ptr as usize, cap);
}
let dst = staging.0 as *mut u8;
// SAFETY: `dst` is a live pinned allocation of at least `need` bytes owned by
// this layer; it is rewritten only here, and this point is reached only after
// `release_prev_pins` waited on the event recorded behind the previous gather
// kernel, so the previous H2D copy from this buffer has completed.
let staged: &[u8] = unsafe {
for (i, s) in slots.iter().enumerate() {
std::ptr::copy_nonoverlapping(s.to_le_bytes().as_ptr(), dst.add(i * 4), 4);
}
std::slice::from_raw_parts(dst as *const u8, need)
};
gpu.copy_h2d_async(staged, self.slots_dev, stream)?;Check the exact type of slots elements (cache.resolve(ids, &mut slots) fills it — u32); if it is not u32, cast each element with (*s as u32).to_le_bytes().
cargo fmt, then cargo clippy --workspace --tests (must be clean), then cargo test -p spark-runtime --lib and cargo test -p spark-model --lib (must pass). Finish when green.
Env-gated MoE decode arm selector.
Repo: Atlas (branch assist/flashnext-longform-fixes). One commit. Model: Qwen3.8-Flash-Next (qwen4_exp), 512 experts top-10 MoE on every one of 48 layers.
crates/spark-model/src/layers/moe/forward.rs:60-140 — MoeLayer::forward is the single-token decode entry. Lines 93-113 route a layer's M=1 MoE through self.forward_prefill(input, 1, ctx, stream) (the tensor-core grouped-GEMM arm, E2M1→E4M3 MMA) ONLY when self.is_dflash_capture_layer is true AND ATLAS_FRANKENSTEIN_DECODE_VIA_PREFILL=1; every other layer falls through to the scalar FP32-FMA decode arm (ops::moe_expert_gate_up_shared at forward.rs:502, ops::moe_expert_silu_down_shared at :584). The comment measured the grouped arm at ≈250 µs per layer; the scalar arm was measured at ≈290 µs per layer call in crates/spark-model/src/layers/qwen3_ssm/trait_layer.rs:136-142. We want to A/B the two arms for ALL 48 layers at M=1 on the GB10 without touching the default behaviour.
In crates/spark-model/src/layers/moe/forward.rs, directly above the existing if self.is_dflash_capture_layer && ... block (line ~96), add a cached process-wide selector:
/// `ATLAS_MOE_DECODE_ARM`: `scalar` (default — the fused FP32-FMA gemv arm) or
/// `grouped` (route every layer's single-token MoE through the tensor-core
/// grouped-GEMM arm via `forward_prefill(M=1)`). Read once per process.
fn moe_decode_arm_grouped() -> bool {
static ARM: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
*ARM.get_or_init(|| {
std::env::var("ATLAS_MOE_DECODE_ARM")
.map(|v| v.eq_ignore_ascii_case("grouped"))
.unwrap_or(false)
})
}(put the fn at module level, i.e. outside impl MoeLayer, next to the other free functions in that file — if the file has none, place it right before the impl block that contains forward.)
Then change the routing condition so the grouped arm is taken when EITHER the old Frankenstein gate holds OR the new selector says grouped:
if moe_decode_arm_grouped()
|| (self.is_dflash_capture_layer
&& std::env::var("ATLAS_FRANKENSTEIN_DECODE_VIA_PREFILL")
.ok()
.as_deref()
== Some("1"))
{Keep the body of the block as it is, but change the one-time log text to "MoE single-token decode: grouped-GEMM arm (forward_prefill M=1) for this layer (one-time log)".
In kernels/gb10/qwen3.8-flash-next/MODEL.toml, find the [behavior] section (search for use_sampling_preset_penalties_for_core) and add, right after that key, a comment block (comments only, no new keys):
# ATLAS_MOE_DECODE_ARM=grouped routes every layer's single-token MoE through the
# tensor-core grouped-GEMM arm (forward_prefill M=1) instead of the scalar FP32-FMA
# gemv arm. A/B lever for the 2026-09-04 llama.cpp gap study; default stays scalar.cargo fmt, cargo clippy --workspace --tests clean, cargo test -p spark-model --lib passes. Finish. Do NOT change any kernel, the scalar arm, forward_prefill, or anything outside the two files named.
Device-side indexer selection; the executor produced no change and the reviewer implemented it.
Repo: Atlas (branch assist/flashnext-longform-fixes). One commit. Model: Qwen3.8-Flash-Next (qwen4_exp).
crates/spark-model/src/layers/qsa.rs:383-441 (QsaIndexer::decode_select): after ops::qsa_score writes complete F32 block scores into self.scores_dev, the code copies them device→host (gpu.copy_d2h_on_stream, which ends in cuStreamSynchronize), sorts on the host by (score desc, index asc) — torch.topk's tie-break: ties go to the LOWER index — takes the first self.block_topk blocks, sorts those ascending, expands each block into ratio token ids (base + r), appends the partial-block tail complete*ratio .. visible, and uploads sel (i32) to self.sel_dev and n_sel to self.seq_len_dev. That is one hard host↔device round trip per attention layer per token (12 per token on this model). n_sel is deterministic: block_topk * ratio + (visible - complete*ratio) — the host does not need the scores to know it. Only the CONTENT of sel_dev depends on the scores.
Kernel conventions: the QSA kernels live in kernels/gb10/qwen3.8-flash-next/nvfp4/qsa_indexer.cu (qsa_score, qsa_gather, ...); their Rust launch wrappers are in crates/spark-model/src/layers/ops/qsa.rs (qsa_score at :83 — copy its shape: kernel handle, grid/block, argument packing); the handles are looked up in QsaIndexer::new (qsa.rs:150-200, search for k_score_k) with gpu.kernel("<module>", "<symbol>") — use the SAME module name as qsa_score.
qsa_select_topkAdd to qsa_indexer.cu a single-block kernel:
extern "C" __global__ void qsa_select_topk(const float* __restrict__ scores, int complete, int block_topk,
int ratio, int tail_start, int visible, int* __restrict__ sel)Semantics (must match the host code bit-for-bit in ORDER of sel):
1. Select the block_topk largest of scores[0..complete) with ties broken by lower index. A robust way
for complete ≤ 4096: each thread owns a strided subset of block ids; compute for every block b its rank = number of blocks j with (scores[j] > scores[b]) or (scores[j] == scores[b] and j < b); a block is selected iff rank < block_topk. This is O(complete²/threads) but complete is at most max_tokens/ratio (512 at 65536 tokens, ratio 128) — fine for one block of 1024 threads. NaN scores must never be selected (treat NaN as -inf).
2. Selected blocks are written in ASCENDING block order: block b with rank r_asc among the selected
(count of selected blocks with id < b) writes sel[r_asc*ratio + r] = b*ratio + r for r in 0..ratio.
3. Then the tail: for t in tail_start..visible, sel[block_topk*ratio + (t - tail_start)] = t.
Use shared memory for the selected flags / prefix counts; __syncthreads() between phases.
crates/spark-model/src/layers/ops/qsa.rs: add pub fn qsa_select_topk(gpu, k, scores_dev, complete: u32, block_topk: u32, ratio: u32, tail_start: u32, visible: u32, sel_dev, stream) -> Result<()> launching one block of 1024 threads (grid 1), packing args exactly like qsa_score.qsa.rs QsaIndexer: add a handle field k_select_k looked up in new next to k_score_k.decode_select, wrap the existing host block (from let mut raw = vec![0u8; complete * 4]; through gpu.copy_h2d_async(&sel_bytes, self.sel_dev, stream)?;) in if !qsa_device_topk() { ...existing... } else { ...new... } where the new arm computes let tail_start = complete * ratio; let n_sel = (self.block_topk as usize * ratio + (visible - tail_start)) as u32; and launches ops::qsa_select_topk(...) writing self.sel_dev directly — no D2H, no H2D of sel. Both arms must leave n_sel defined for the code that follows (qsa_gather, the identity table, the seq_len_dev upload). qsa_device_topk() is a OnceLock<bool> reading ATLAS_QSA_DEVICE_TOPK=1 (default false).Add ATLAS_QSA_TOPK_VERIFY=1 (another OnceLock<bool>): when set together with the device arm, ALSO run the host computation (the existing code, into a local Vec<i32>), then copy sel_dev back (gpu.copy_d2h_on_stream, n_sel*4 bytes) and compare element-wise; on the first mismatch tracing::error!("QSA device top-k mismatch at pos {pos}: index {i} host {h} device {d}") and return Err(anyhow!(...)); on success tracing::debug! once per 256 positions. This mode is for validation runs only.
cargo fmt; cargo clippy --workspace --tests clean; cargo test -p spark-model --lib passes. Add one unit test of the host reference ordering helper if you factor the host selection into a function (recommended: fn host_select(scores: &[f32], block_topk, ratio, visible) -> Vec<i32> used by both the host arm and the verify mode). Finish when green. Do not change qsa_gather, qsa_score, prefill paths, or the graph-veto flags.