perf(glm52): launch-ahead decode — hide the ~0.7ms cuGraphLaunch under the previous step's execution

docs/models/glm52/whole-step-decode-graph.md

@@ -1,6 +1,6 @@
 # GLM5.2 whole-step decode CUDA graph (PR5c)
 
-> **TL;DR:** Execution record of PR5c: the whole per-rank decode step (embed → 78 layers → lm_head → device argmax) is captured into one CUDA graph and replayed every step. Measured on jz-38 8×H200, single request: **200 → 37.5 ms/step from the graph alone** (byte-identical to the PR5b record), **→ 31.3** after switching every m=1 projection to the weight-only fp8 GEMV (activation quant removed — re-gated via the #499 oracle in the new `--precision gemv` mode), **→ 25.3** after grid-striding the capacity-sized MoE quant/SiLU launches (block *scheduling*, not arithmetic, was their cost), **→ 23.4** after packing gate|up into one GEMV, overlapping the shared expert with the MoE collectives on a second stream (fork/join events inside the graph), and staging the relayout's expert ranges in shared memory, **→ 22.9 single** after fusing each layer's closing add with the next layer's input norm (ping-ponged attn buffers, bit-identical `_round` kernel) and running the DSA indexer concurrently with the MLA front's q_b/kv_a, **→ 22.6 single / 22.3 at 8-way (~346 tok/s aggregate)** after two-tier attention graphs (short-context FlashMLA topk 256 — while `seq_len <= 256` the DSA top-256 IS the full token set, so the short-tier graph attends the same tokens at 1/8 the padded index walk; tier-crossing and mixed-tier concurrency e2e-gated, short tier oracle-gated), **→ 22.5 single** after swapping the device argmax to the shared two-stage split kernel (bit-identical; the single-CTA scan was the last serial kernel in the step). A step-timing probe puts the inter-step host gap at ~0.05 ms — the whole step lives inside the graph; **vLLM GLM5.2 DP8/EP8 measured on the same node/workload: steady-state 20.0 ms/step (TPOT 19.8)** — the remaining 2.6 ms gap sits mostly in the expert GEMM (their ~10.8 µs/instance vs our 64-row-M-tile TRTLLM grouped at 22.7 µs) and the collectives' wait structure (#542; fp8 dispatch payload is measured perf-NEUTRAL — dispatch is rank-arrival-wait-bound, not byte-bound). **→ 19.6 single (2026-07-05, below the vLLM reference)** after replacing the TRTLLM grouped expert GEMM with the DeepGEMM masked grouped GEMM (the "swapAB" attribution was wrong — retraction + A/B record in the masked-GEMM section; c64 diverse 1113 → 1475 tok/s, solo span-8 32.2 → 28.2 ms). Indexer oracle reference drift is #541 (pre-existing on main).
+> **TL;DR:** Execution record of PR5c: the whole per-rank decode step (embed → 78 layers → lm_head → device argmax) is captured into one CUDA graph and replayed every step. Measured on jz-38 8×H200, single request: **200 → 37.5 ms/step from the graph alone** (byte-identical to the PR5b record), **→ 31.3** after switching every m=1 projection to the weight-only fp8 GEMV (activation quant removed — re-gated via the #499 oracle in the new `--precision gemv` mode), **→ 25.3** after grid-striding the capacity-sized MoE quant/SiLU launches (block *scheduling*, not arithmetic, was their cost), **→ 23.4** after packing gate|up into one GEMV, overlapping the shared expert with the MoE collectives on a second stream (fork/join events inside the graph), and staging the relayout's expert ranges in shared memory, **→ 22.9 single** after fusing each layer's closing add with the next layer's input norm (ping-ponged attn buffers, bit-identical `_round` kernel) and running the DSA indexer concurrently with the MLA front's q_b/kv_a, **→ 22.6 single / 22.3 at 8-way (~346 tok/s aggregate)** after two-tier attention graphs (short-context FlashMLA topk 256 — while `seq_len <= 256` the DSA top-256 IS the full token set, so the short-tier graph attends the same tokens at 1/8 the padded index walk; tier-crossing and mixed-tier concurrency e2e-gated, short tier oracle-gated), **→ 22.5 single** after swapping the device argmax to the shared two-stage split kernel (bit-identical; the single-CTA scan was the last serial kernel in the step). A step-timing probe puts the inter-step host gap at ~0.05 ms — the whole step lives inside the graph; **vLLM GLM5.2 DP8/EP8 measured on the same node/workload: steady-state 20.0 ms/step (TPOT 19.8)** — the remaining 2.6 ms gap sits mostly in the expert GEMM (their ~10.8 µs/instance vs our 64-row-M-tile TRTLLM grouped at 22.7 µs) and the collectives' wait structure (#542; fp8 dispatch payload is measured perf-NEUTRAL — dispatch is rank-arrival-wait-bound, not byte-bound). **→ 19.6 single (2026-07-05, below the vLLM reference)** after replacing the TRTLLM grouped expert GEMM with the DeepGEMM masked grouped GEMM (the "swapAB" attribution was wrong — retraction + A/B record in the masked-GEMM section; c64 diverse 1113 → 1475 tok/s, solo span-8 32.2 → 28.2 ms). **Launch-ahead decode (D8, 2026-07-05): the ~0.7 ms host-side cuGraphLaunch + step-boundary idle move off the critical path (inter-step device idle 641–804 → 43 µs; solo 19.5 ms/step, zero variance) — feed kernel + pinned D2H + coordinator-global speculation lease, byte-parity gated.** Indexer oracle reference drift is #541 (pre-existing on main).
 >
 > **Last touched:** 2026-07
 
@@ -341,10 +341,61 @@ spec-8 verify step. Second lever, independent: mask pad-row routing
 (`topk_idx = -1`) in partial buckets to kill the artificial combine
 straggler and 7/8 of solo collective payload.
 
+## Launch-ahead decode (D8, jz-38, 2026-07-05)
+
+The stagger investigation (docs PR #568) found `cuGraphLaunch` costing
+~700 µs of host-side driver work per rank per step **with the GPU idle
+through it**, plus a ~364 µs cross-rank launch spread — together the
+641–810 µs step-boundary idle. The double-replay probe proved the launch
+hides completely under the previous replay's execution (5.4 µs intra-step
+gap). This branch turns that into the serving path:
+
+- a **feed kernel** advances `token_ids`/`positions`/`slot_mapping`/
+  `seq_lens` in place from the step's device argmax; rope rows re-gather on
+  device; the argmax D2H lands in per-bucket **pinned** buffers so the copy
+  is truly asynchronous — the next step's replay is enqueued before the
+  host blocks on this step's result. Zero graph changes: everything is
+  eager launches stream-ordered behind the executing replay, so the leased
+  path is byte-identical by construction.
+- **the speculation is all-ranks-or-none**: a speculative replay is a full
+  set of collectives. The coordinator decides both flags on global data —
+  `lease` (next step repeats this shape; includes model-length headroom,
+  and pad rows can never outrun active rows) and `consume` (this step IS
+  the speculation; needs unchanged shapes AND no slot handoff, because a
+  finish+admission can reuse a slot id under an identical shape). Rank-side
+  checks are hard `ensure!`s, not silent fallbacks — the first gate run
+  desynced exactly there (the 1/8 near-tie variant hit EOS early, that rank
+  fell back while others consumed, and the collective pairing slipped into
+  the ~100 s DeepEP device-timeout trap).
+- **every bucket × tier graph is pre-captured at startup** (8 warm pad
+  steps while the ranks are trivially in lock-step): captured-ness must be
+  uniform for the same reason, and the old mid-serving capture stall goes
+  away with it.
+
+Measured (jz-38, same-day baselines):
+
+| metric | main | launch-ahead |
+|---|---|---|
+| inter-step device idle (own-gap, graph trace) | 641–804 µs | **43.0 µs** (p90 44.8) |
+| cross-rank step-start spread | 350–364 µs (max 965–1097) | **32.6 µs** (max 42.5) |
+| `cuGraphLaunch` p50 | 700 µs on the critical path | 660 µs, fully overlapped |
+| solo ms/step (5 runs) | 19.6–20.4, run-to-run wander | **19.5, zero variance** |
+
+Gates: determinism + byte-parity vs the main record PASS (short + 128-tok
+head), 8-way identical (7/8 + the known deterministic near-tie variant),
+8-way 128-tok concurrent 391 tok/s aggregate, staggered-admission churn +
+post-churn determinism PASS, post-disconnect + teardown PASS.
+
+Pitfalls recorded: cudarc's `memcpy_dtoh` copies the DESTINATION's byte
+count (an oversized pinned buffer reads past the device allocation —
+CUDA_ERROR_INVALID_VALUE); and a speculative-consume decision made from
+per-rank data (inputs, positions, captured-ness) desyncs the collectives —
+every such condition must come from the coordinator or be an invariant.
+
 ## Next action
 
-Build the device-fed decode loop (launch-ahead validated above — ~1 ms/step);
-then pad-row masked routing. Remaining from #542: per-layer combine straggler
-under real diverse load = expert imbalance (EPLB territory, deferred by user
-2026-07-05). MegaMoE fusion measured OUT (section above). #559 residual
-(combine +1.9, MLA +1.2 at b8), #541 indexer reference re-pin.
+Follow-ups: #542 wait-structure — launch-ahead lands the boundary half
+(section above); the per-layer combine straggler under diverse load is
+expert imbalance (EPLB, deferred). Pad-row masked routing (`topk_idx = -1`)
+is the next small lever. #559 residual (combine +1.9, MLA +1.2 at b8),
+#541 indexer reference re-pin.

openinfer-core/src/cuda_graph.rs

@@ -65,6 +65,27 @@ impl CudaGraphState {
         }
     }
 
+    /// Whether a graph has been captured and instantiated — i.e. whether
+    /// [`Self::launch_captured`] can replay without running the kernel closure.
+    pub fn is_captured(&self) -> bool {
+        !self.exec.is_null()
+    }
+
+    /// Replay the captured graph. Errors if nothing has been captured yet —
+    /// callers that may hit an uncaptured shape go through
+    /// [`Self::run_or_capture`] instead.
+    pub fn launch_captured(&mut self, ctx: &DeviceContext) -> Result<()> {
+        anyhow::ensure!(
+            self.is_captured(),
+            "CUDA graph replay requested before any capture"
+        );
+        let stream = active_cu_stream(ctx);
+        check(
+            unsafe { sys::cuGraphLaunch(self.exec, stream) },
+            "cuGraphLaunch",
+        )
+    }
+
     /// Run kernel closure directly, or capture into a graph and replay.
     ///
     /// `kernels` must be a pure GPU kernel sequence — no CPU-GPU sync, no allocation.

openinfer-glm52/src/model/launch_ahead.rs

@@ -0,0 +1,173 @@
+//! Launch-ahead speculation state machine: the speculative next-step replay
+//! (feed kernel + rope re-gather + graph launch enqueued while the current
+//! step executes) and its harvest. Split from `model/mod.rs`; see
+//! `decode_step` there for the consume/lease protocol.
+
+use anyhow::{Result, ensure};
+use openinfer_kernels::ops::{embedding_rows_into, glm52_decode_feed_launch};
+use openinfer_kernels::tensor::DeviceContext;
+
+use super::{
+    GLM52_MAX_BATCH_PER_RANK, GLM52_MAX_MODEL_LEN, GLM52_MLA_TOPK_SHORT, GLM52_VOCAB,
+    Glm52RankModel, Glm52StepShape, TIER_FULL, TIER_SHORT,
+};
+
+/// A speculative next-step whole-graph replay already enqueued on the decode
+/// stream: the feed kernel advanced the device input buffers in place and
+/// the bucket graph was launched ahead, hiding the ~0.7 ms host-side
+/// `cuGraphLaunch` under the current step's execution (measured on jz-38:
+/// 5.4 µs intra-step replay gap vs 810 µs at an unhidden step boundary).
+/// Consumed only when the coordinator's next step matches `expect` exactly;
+/// on any mismatch the full host prologue overwrites the advanced inputs and
+/// the stale replay degrades to a recompute whose rows never influence the
+/// following real replay (per-row math is row-independent).
+pub(super) struct Glm52SpeculatedStep {
+    pub(super) bucket: usize,
+    pub(super) active_rows: usize,
+    pub(super) slots: [u8; GLM52_MAX_BATCH_PER_RANK],
+    /// The coordinator inputs this speculation assumed: active rows carry
+    /// (this step's argmax, position + 1); padding rows echo this step's
+    /// padding input verbatim (their device rows keep self-feeding, which
+    /// is harmless — pad outputs are never read and rows are isolated).
+    pub(super) expect: [(u32, usize); GLM52_MAX_BATCH_PER_RANK],
+}
+
+impl Glm52RankModel {
+    /// Post-replay tail shared by both step paths: enqueue the argmax D2H,
+    /// then — with the copy still in flight — optionally speculate the next
+    /// step (feed kernel + rope-row gathers + launch-ahead replay), and only
+    /// then block for this step's result. The speculative `cuGraphLaunch`
+    /// thereby overlaps this step's execution instead of idling the GPU at
+    /// the next step boundary.
+    pub(super) fn decode_step_harvest(
+        &mut self,
+        ctx: &DeviceContext,
+        inputs: &[(u32, usize); GLM52_MAX_BATCH_PER_RANK],
+        shape: Glm52StepShape,
+        lease: bool,
+    ) -> Result<[u32; GLM52_MAX_BATCH_PER_RANK]> {
+        let batch = shape.bucket;
+        let bucket = self
+            .buckets
+            .iter_mut()
+            .find(|bucket| bucket.rows == batch)
+            .expect("decode_step validated the bucket");
+        ctx.stream.memcpy_dtoh(
+            &bucket.scratch.argmax_values,
+            &mut bucket.argmax_values_host,
+        )?;
+        ctx.stream.memcpy_dtoh(
+            &bucket.scratch.argmax_indices,
+            &mut bucket.argmax_indices_host,
+        )?;
+
+        // Speculation must be all-ranks-or-none: the speculative replay is a
+        // full set of collectives, so every precondition here is a GLOBAL
+        // invariant the coordinator's lease already guarantees (single-token
+        // rows from the shape, position headroom from its slot bookkeeping —
+        // pad rows never outrun active rows — and captured graphs from the
+        // startup pre-capture). A rank that silently skipped would desync
+        // the collective pairing; crash early instead.
+        let mut speculated = false;
+        if lease {
+            let mut seen = [false; GLM52_MAX_BATCH_PER_RANK];
+            for &slot in &shape.slots[..batch] {
+                let slot = slot as usize;
+                ensure!(
+                    slot < GLM52_MAX_BATCH_PER_RANK && !std::mem::replace(&mut seen[slot], true),
+                    "GLM5.2 launch-ahead lease granted for a step with repeated slot {slot} \
+                     (span rows are never leased)"
+                );
+            }
+            for &position in &self.device_positions[..batch] {
+                ensure!(
+                    position + 1 < GLM52_MAX_MODEL_LEN,
+                    "GLM5.2 launch-ahead lease granted with a row at position {position} — the \
+                     advanced step would breach the model-length cap"
+                );
+            }
+            let longest_next = self.device_positions[..batch]
+                .iter()
+                .map(|&position| position + 2)
+                .max()
+                .expect("decode buckets forward at least one row");
+            let tier = if longest_next <= GLM52_MLA_TOPK_SHORT {
+                TIER_SHORT
+            } else {
+                TIER_FULL
+            };
+            ensure!(
+                bucket.graphs[tier].is_captured(),
+                "GLM5.2 launch-ahead lease granted before the (bucket {batch}, tier {tier}) graph \
+                 was captured — the startup pre-capture must cover every shape"
+            );
+            // From the first enqueue below onward, a host-side failure
+            // (launch error) leaves the other ranks' speculative collectives
+            // half-paired — teardown then eats the ~100 s device timeout.
+            // Every fallible CHECK is above for exactly that reason; do not
+            // insert anything that can fail between these launches.
+            glm52_decode_feed_launch(
+                ctx,
+                &bucket.scratch.argmax_indices,
+                &mut self.token_ids,
+                &mut self.positions,
+                &mut self.slot_mapping,
+                &mut self.seq_lens,
+                batch,
+            )?;
+            embedding_rows_into(ctx, &self.cos_table, &self.positions, batch, &mut self.cos)?;
+            embedding_rows_into(ctx, &self.sin_table, &self.positions, batch, &mut self.sin)?;
+            bucket.graphs[tier].launch_captured(ctx)?;
+            for position in &mut self.device_positions[..batch] {
+                *position += 1;
+            }
+            speculated = true;
+        }
+
+        // Block on the D2H (the pinned slices synchronize on their own copy
+        // events) and unpack — identical semantics to the old blocking
+        // readback.
+        let top_values = bucket
+            .argmax_values_host
+            .as_slice()
+            .map_err(|err| anyhow::anyhow!("GLM5.2 argmax values D2H sync failed: {err}"))?;
+        let top_indices = bucket
+            .argmax_indices_host
+            .as_slice()
+            .map_err(|err| anyhow::anyhow!("GLM5.2 argmax indices D2H sync failed: {err}"))?;
+        let mut outputs = [0u32; GLM52_MAX_BATCH_PER_RANK];
+        for row in 0..batch {
+            outputs[row] = top_indices[row].max(0) as u32;
+        }
+        // Validate ACTIVE rows only: padding rows self-feed their own argmax
+        // through lease streaks, and nobody consumes their outputs — an
+        // engine-fatal assertion has no business reading them.
+        for row in 0..shape.active_rows {
+            let slot = shape.slots[row] as usize;
+            let top_value = top_values[row].to_f32();
+            let top_index = top_indices[row];
+            ensure!(
+                top_value.is_finite(),
+                "GLM5.2 slot {slot} greedy argmax found no finite logit (top = {top_value})"
+            );
+            ensure!(
+                (0..GLM52_VOCAB as i32).contains(&top_index),
+                "GLM5.2 slot {slot} greedy argmax index {top_index} outside the vocab"
+            );
+        }
+
+        if speculated {
+            let mut expect = *inputs;
+            for row in 0..shape.active_rows {
+                expect[row] = (outputs[row], inputs[row].1 + 1);
+            }
+            self.speculated = Some(Glm52SpeculatedStep {
+                bucket: batch,
+                active_rows: shape.active_rows,
+                slots: shape.slots,
+                expect,
+            });
+        }
+        Ok(outputs)
+    }
+}