diff --git a/docs/models/glm52/whole-step-decode-graph.md b/docs/models/glm52/whole-step-decode-graph.md index d48095f1..9548f54e 100644 --- a/docs/models/glm52/whole-step-decode-graph.md +++ b/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. diff --git a/openinfer-core/src/cuda_graph.rs b/openinfer-core/src/cuda_graph.rs index 26ade04c..e248f1df 100644 --- a/openinfer-core/src/cuda_graph.rs +++ b/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. diff --git a/openinfer-glm52/src/model/launch_ahead.rs b/openinfer-glm52/src/model/launch_ahead.rs new file mode 100644 index 00000000..544ade48 --- /dev/null +++ b/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) + } +} diff --git a/openinfer-glm52/src/model.rs b/openinfer-glm52/src/model/mod.rs similarity index 89% rename from openinfer-glm52/src/model.rs rename to openinfer-glm52/src/model/mod.rs index abcaa9fb..7fde0860 100644 --- a/openinfer-glm52/src/model.rs +++ b/openinfer-glm52/src/model/mod.rs @@ -15,7 +15,7 @@ //! shapes identical within a bucket (the whole-step CUDA graphs' contract). use anyhow::{Context as _, Result, ensure}; -use cudarc::driver::CudaSlice; +use cudarc::driver::{CudaSlice, PinnedHostSlice}; use half::bf16; use openinfer_core::cuda_graph::CudaGraphState; use openinfer_kernels::ops::{ @@ -44,6 +44,9 @@ use crate::moe_ep8::{Glm52MoeEp8LayerWeights, Glm52MoeEp8State, glm52_moe_ep8_ro use crate::scratch::Glm52DecodeScratch; use crate::weights::{Glm52RankGpuWeights, retype_owned}; +mod launch_ahead; +use launch_ahead::Glm52SpeculatedStep; + /// Per-request context cap: `prompt + max_tokens - 1 <= GLM52_MAX_MODEL_LEN`. /// Sizes each slot's region of the per-layer MLA and index-K caches at build /// time. @@ -72,6 +75,10 @@ const _: () = assert!( <= openinfer_kernels::ops::GLM52_DEEPGEMM_MASKED_CAP ); +// The decode feed kernel runs one 32-thread block (`glm52_decode_feed.cu`); +// a batch-cap bump past it must widen the kernel, not silently truncate. +const _: () = assert!(GLM52_MAX_BATCH_PER_RANK <= 32); + /// The step's forward shape, agreed globally by the coordinator: `bucket` /// rows per rank (a member of [`GLM52_DECODE_BUCKETS`] — the MoE collectives /// require every rank to enter with the same global row count), with @@ -89,6 +96,11 @@ const _: () = assert!( pub(crate) struct Glm52StepShape { pub(crate) bucket: usize, pub(crate) slots: [u8; GLM52_MAX_BATCH_PER_RANK], + /// Rows `0..active_rows` carry real requests; `active_rows..bucket` are + /// padding. Carried explicitly because a padding input is NOT + /// value-distinguishable from an active one (a single-token prompt `[0]` + /// legally feeds `(token 0, position 0)`). + pub(crate) active_rows: usize, } /// Short-context attention tier: while `seq_len <= 256` the DSA top-256 IS @@ -284,6 +296,12 @@ pub(crate) struct Glm52RankModel { cos: CudaSlice, sin: CudaSlice, token_ids: CudaSlice, + /// In-flight speculative next-step replay, if any (see `decode_step`). + speculated: Option, + /// What the per-row `positions` device buffer currently holds (padding + /// rows included): the feed kernel advances it without host readback, + /// and a speculation must keep every row under the model-length cap. + device_positions: [usize; GLM52_MAX_BATCH_PER_RANK], } /// Everything one decode bucket owns: batch-`rows` FlashMLA plans (per @@ -303,6 +321,12 @@ struct Glm52BucketState { scratch: Glm52DecodeScratch, graphs: [CudaGraphState; 2], block_table: CudaSlice, + /// Pinned landing buffers for this bucket's argmax D2H, sized exactly + /// `rows` (`memcpy_dtoh` copies the DESTINATION's byte count). Pinned + /// memory keeps the copy asynchronous so the next step's replay can be + /// enqueued launch-ahead before the host blocks on the result. + argmax_values_host: PinnedHostSlice, + argmax_indices_host: PinnedHostSlice, } /// Tier index into the per-tier arrays: every consumer selects with the same @@ -467,6 +491,10 @@ impl Glm52RankModel { scratch: Glm52DecodeScratch::new(ctx, &contract_rows, mqa_shape)?, graphs: [CudaGraphState::new(), CudaGraphState::new()], block_table: bucket_table, + // Read only after a D2H lands in them (the write-combined + // pages start uninitialized). + argmax_values_host: unsafe { ctx.ctx.alloc_pinned::(rows)? }, + argmax_indices_host: unsafe { ctx.ctx.alloc_pinned::(rows)? }, }); } let buckets: [Glm52BucketState; GLM52_DECODE_BUCKETS.len()] = buckets @@ -542,6 +570,8 @@ impl Glm52RankModel { cos: ctx.stream.alloc_zeros::(batch * ROPE_HALF)?, sin: ctx.stream.alloc_zeros::(batch * ROPE_HALF)?, token_ids: ctx.stream.alloc_zeros::(batch)?, + speculated: None, + device_positions: [0; GLM52_MAX_BATCH_PER_RANK], }) } @@ -575,7 +605,55 @@ impl Glm52RankModel { ep8: &mut Glm52MoeEp8State, inputs: &[(u32, usize); GLM52_MAX_BATCH_PER_RANK], shape: Glm52StepShape, + flags: crate::runner::Glm52StepFlags, ) -> Result<[u32; GLM52_MAX_BATCH_PER_RANK]> { + let batch = shape.bucket; + if flags.consume { + // Launch-ahead fast path: the coordinator says this step IS the + // replay every rank speculatively enqueued last step. That claim + // is global — a speculative replay is a full set of collectives, + // so ranks must consume together or not at all. Any mismatch is + // a protocol bug; failing the step beats a silent fallback that + // would desync the collective pairing (measured as the ~100 s + // DeepEP device-timeout trap). + let speculated = self.speculated.take().context( + "GLM5.2 launch-ahead desync: the coordinator consumed a speculation this rank \ + never enqueued", + )?; + ensure!( + speculated.bucket == batch + && speculated.active_rows == shape.active_rows + && speculated.slots[..batch] == shape.slots[..batch] + && speculated.expect[..batch] == inputs[..batch], + "GLM5.2 launch-ahead desync: consumed speculation (bucket {}, slots {:?}, expect \ + {:?}) does not match the step (bucket {batch}, slots {:?}, inputs {:?})", + speculated.bucket, + &speculated.slots[..speculated.bucket], + &speculated.expect[..speculated.bucket], + &shape.slots[..batch], + &inputs[..batch], + ); + } else { + // Any stale speculation was enqueued by EVERY rank (the lease is + // a global grant), so the stale replay's collectives pair up and + // it degrades to a harmless recompute the prologue overwrites. + self.speculated = None; + self.decode_step_prologue_and_replay(ctx, aux, ep8, inputs, shape)?; + } + self.decode_step_harvest(ctx, inputs, shape, flags.lease) + } + + /// The non-leased step path: validate the shape, rewrite every per-step + /// device input buffer from the coordinator's `inputs`, and run (or lazily + /// capture) the whole-step graph for the step's bucket × tier. + fn decode_step_prologue_and_replay( + &mut self, + ctx: &DeviceContext, + aux: &DeviceContext, + ep8: &mut Glm52MoeEp8State, + inputs: &[(u32, usize); GLM52_MAX_BATCH_PER_RANK], + shape: Glm52StepShape, + ) -> Result<()> { // The bucket state's `rows` is the lookup key — an unknown bucket is // a coordinator bug and fails the step before touching the GPU. let bucket = self @@ -640,6 +718,9 @@ impl Glm52RankModel { ctx.stream .memcpy_htod(&slots_host, &mut self.slot_mapping)?; ctx.stream.memcpy_htod(&seq_lens_host, &mut self.seq_lens)?; + for row in 0..batch { + self.device_positions[row] = inputs[row].1; + } // Gather each row's rotary table row (a bit-exact row copy). 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)?; @@ -710,27 +791,7 @@ impl Glm52RankModel { ) }); bucket.graphs[tier] = graph; - result?; - - let s = &bucket.scratch; - let top_values = ctx.stream.clone_dtoh(&s.argmax_values)?; - let top_indices = ctx.stream.clone_dtoh(&s.argmax_indices)?; - let mut outputs = [0u32; GLM52_MAX_BATCH_PER_RANK]; - for row in 0..batch { - 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" - ); - outputs[row] = top_index as u32; - } - Ok(outputs) + result } } diff --git a/openinfer-glm52/src/runner.rs b/openinfer-glm52/src/runner.rs index f612f4ed..96a3c985 100644 --- a/openinfer-glm52/src/runner.rs +++ b/openinfer-glm52/src/runner.rs @@ -44,6 +44,29 @@ pub(crate) struct Glm52RankWeightLoadReport { pub(crate) loaded_to_gpu: bool, } +/// The coordinator's launch-ahead directives for one step — both are GLOBAL +/// claims (a speculative replay is a full set of collectives, so ranks must +/// act on them together or not at all). +#[derive(Clone, Copy, Debug)] +pub(crate) struct Glm52StepFlags { + /// This step IS the speculative replay every rank enqueued last step. + pub(crate) consume: bool, + /// The next step is guaranteed to repeat this shape with each row + /// advanced by its own argmax — every rank MUST enqueue that next + /// replay launch-ahead (see `Glm52RankModel::decode_step`). + pub(crate) lease: bool, +} + +impl Glm52StepFlags { + /// No speculation in either direction (warm pre-capture, tests). + pub(crate) fn plain() -> Self { + Self { + consume: false, + lease: false, + } + } +} + enum Glm52RankCommand { LoadWeights { model_path: PathBuf, @@ -71,6 +94,7 @@ enum Glm52RankCommand { Step { inputs: Box<[(u32, usize); GLM52_MAX_BATCH_PER_RANK]>, shape: Glm52StepShape, + flags: Glm52StepFlags, resp: Sender>, }, /// Non-collective: load the DSpark draft model onto this rank. Issued to @@ -174,12 +198,14 @@ impl Glm52RankWorker { &self, inputs: [(u32, usize); GLM52_MAX_BATCH_PER_RANK], shape: Glm52StepShape, + flags: Glm52StepFlags, ) -> Result>> { let (resp_tx, resp_rx) = bounded(1); self.tx .send(Glm52RankCommand::Step { inputs: Box::new(inputs), shape, + flags, resp: resp_tx, }) .map_err(|_| anyhow::anyhow!("GLM5.2 rank worker channel closed"))?; @@ -456,6 +482,7 @@ impl Glm52RankThreadState { &mut self, inputs: &[(u32, usize); GLM52_MAX_BATCH_PER_RANK], shape: Glm52StepShape, + flags: Glm52StepFlags, ) -> Result<[u32; GLM52_MAX_BATCH_PER_RANK]> { let dev_ctx = self.ctx.device_context()?; let runtime = self @@ -468,7 +495,7 @@ impl Glm52RankThreadState { .context("GLM5.2 step before setup_comm")?; runtime .model - .decode_step(&dev_ctx, &runtime.aux_ctx, ep8, inputs, shape) + .decode_step(&dev_ctx, &runtime.aux_ctx, ep8, inputs, shape, flags) } } @@ -487,9 +514,10 @@ fn rank_worker_loop(rx: Receiver, mut state: Glm52RankThreadSt Glm52RankCommand::Step { inputs, shape, + flags, resp, } => { - let _ = resp.send(state.step(&inputs, shape)); + let _ = resp.send(state.step(&inputs, shape, flags)); } Glm52RankCommand::LoadDspark { path, resp } => { let _ = resp.send(state.load_dspark(&path)); diff --git a/openinfer-glm52/src/scheduler.rs b/openinfer-glm52/src/scheduler.rs index a62baeaa..33bc5ae6 100644 --- a/openinfer-glm52/src/scheduler.rs +++ b/openinfer-glm52/src/scheduler.rs @@ -28,9 +28,10 @@ use tokio::sync::mpsc; use crate::dspark::{GLM52_DSPARK_DRAFTS, accept_greedy}; use crate::model::{ - GLM52_DECODE_BUCKETS, GLM52_MAX_BATCH_PER_RANK, GLM52_MAX_MODEL_LEN, Glm52StepShape, + GLM52_DECODE_BUCKETS, GLM52_MAX_BATCH_PER_RANK, GLM52_MAX_MODEL_LEN, GLM52_MLA_TOPK_SHORT, + Glm52StepShape, }; -use crate::runner::Glm52RankWorker; +use crate::runner::{Glm52RankWorker, Glm52StepFlags}; /// What a rank forwards this step. Idle ranks feed the padding input; its /// KV/index-cache writes land in the idle rank's own dead cache slots and are @@ -390,16 +391,50 @@ fn plan_step_shapes(wants: &[[usize; GLM52_MAX_BATCH_PER_RANK]]) -> Vec= actives and bucket <= MAX, so bucket - used <= frees. + let active_rows = dst; let mut frees = (0..GLM52_MAX_BATCH_PER_RANK).filter(|&slot| row[slot] == 0); while dst < bucket { slots[dst] = frees.next().expect("bucket - used <= free slots") as u8; dst += 1; } - Glm52StepShape { bucket, slots } + Glm52StepShape { + bucket, + slots, + active_rows, + } }) .collect() } +/// The launch-ahead flag decision — pure so the desync rules are testable. +/// `consume`: this step IS the speculation every rank enqueued (same shapes +/// AND no slot changed hands — a finish + admission can reuse a slot id +/// under an identical-looking shape). `lease`: every rank must enqueue the +/// next step speculatively — pure single-token decode everywhere with +/// model-length headroom, nothing queued, no draft round. Both are global +/// claims: a speculative replay is a full set of collectives, so per-rank +/// discretion would desync the pairing. +fn launch_ahead_flags( + shapes: &[Glm52StepShape], + leased_shapes: Option<&[Glm52StepShape]>, + slots_changed: bool, + pending_empty: bool, + dspark_enabled: bool, + slots: &[RankSlots], +) -> Glm52StepFlags { + let consume = !slots_changed && leased_shapes == Some(shapes); + let lease = pending_empty + && !dspark_enabled + && slots + .iter() + .flat_map(|rank_slots| rank_slots.iter().flatten()) + .all(|active| { + active.state.feed_want() == 1 + && active.state.next_input_at(0).position + 1 < GLM52_MAX_MODEL_LEN + }); + Glm52StepFlags { consume, lease } +} + fn feed_wants(slots: &[RankSlots]) -> Vec<[usize; GLM52_MAX_BATCH_PER_RANK]> { slots .iter() @@ -447,6 +482,67 @@ pub(crate) fn run_dp8_coordinator( .all(|rank_slots| rank_slots.iter().all(Option::is_none)) }; + // Pre-capture every whole-step graph (bucket × attention tier) while the + // ranks are idle and trivially in lock-step. Launch-ahead speculation + // requires captured-ness to be UNIFORM across ranks — a lazily capturing + // rank would skip the speculative replay the others enqueued and desync + // the collectives — and pre-capturing also removes the old mid-serving + // capture stall. Row 0 at position GLM52_MLA_TOPK_SHORT lifts the step + // into the full tier; every row is a padding write into a free slot's + // dead cache region. + for &bucket in GLM52_DECODE_BUCKETS.iter() { + for full_tier in [false, true] { + let mut shape = Glm52StepShape { + bucket, + slots: [0; GLM52_MAX_BATCH_PER_RANK], + active_rows: 0, + }; + for (slot, dst) in shape.slots.iter_mut().enumerate().take(bucket) { + *dst = slot as u8; + } + let mut inputs = + [(GLM52_PADDING_STEP.token, GLM52_PADDING_STEP.position); GLM52_MAX_BATCH_PER_RANK]; + if full_tier { + inputs[0] = (GLM52_PADDING_STEP.token, GLM52_MLA_TOPK_SHORT); + } + let responses = match workers + .iter() + .map(|worker| worker.step_async(inputs, shape, Glm52StepFlags::plain())) + .collect::>>() + { + Ok(responses) => responses, + Err(err) => { + log::error!("GLM5.2 graph pre-capture failed to submit: {err:#}"); + return; + } + }; + for (rank, resp) in responses.into_iter().enumerate() { + let result = resp + .recv() + .map_err(|_| anyhow::anyhow!("rank dropped its pre-capture response")) + .and_then(|r| r); + if let Err(err) = result { + log::error!( + "GLM5.2 graph pre-capture (bucket {bucket}, full_tier {full_tier}) \ + failed on rank {rank}: {err:#}" + ); + return; + } + } + } + } + log::info!( + "GLM5.2 whole-step graphs pre-captured: {} buckets x 2 tiers", + GLM52_DECODE_BUCKETS.len() + ); + + // The step shapes that carried the last launch-ahead lease, and whether + // any slot changed hands since it was granted — together they decide + // whether this step consumes the speculation (must be all-ranks-or-none, + // so the decision lives here, on global data, not on the ranks). + let mut leased_shapes: Option> = None; + let mut slots_changed = false; + 'serve: loop { // Intake: block when fully idle, otherwise drain what's queued. if channel_open && all_idle(&slots) && pending.is_empty() { @@ -495,6 +591,7 @@ pub(crate) fn run_dp8_coordinator( pending_resets[rank].push(slot); } slots[rank][slot] = Some(ActiveRequest { req, state }); + slots_changed = true; } if all_idle(&slots) { continue; @@ -505,6 +602,16 @@ pub(crate) fn run_dp8_coordinator( // free slots — and all responses are joined before any output is // interpreted. let shapes = plan_step_shapes(&feed_wants(&slots)); + let flags = launch_ahead_flags( + &shapes, + leased_shapes.as_deref(), + slots_changed, + pending.is_empty(), + dspark_enabled, + &slots, + ); + leased_shapes = flags.lease.then(|| shapes.clone()); + slots_changed = false; let responses = slots .iter() .zip(&workers) @@ -523,7 +630,7 @@ pub(crate) fn run_dp8_coordinator( *input = (step.token, step.position); } } - worker.step_async(inputs, *shape) + worker.step_async(inputs, *shape, flags) }) .collect::>>(); let responses = match responses { @@ -635,6 +742,7 @@ pub(crate) fn run_dp8_coordinator( pending_resets[rank].push(slot_id); } *slot = None; + slots_changed = true; } else if dspark_enabled { // Committed rows' captured hidden feeds the draft // context; then re-propose from the new anchor. @@ -770,6 +878,56 @@ fn fail_step(slots: &mut [RankSlots], err: &anyhow::Error) { } } +#[cfg(test)] +mod launch_ahead_flag_tests { + use super::*; + + fn shape(bucket: usize, active_rows: usize) -> Glm52StepShape { + let mut slots = [0u8; GLM52_MAX_BATCH_PER_RANK]; + for (slot, dst) in slots.iter_mut().enumerate().take(bucket) { + *dst = slot as u8; + } + Glm52StepShape { + bucket, + slots, + active_rows, + } + } + + #[test] + fn consume_requires_unchanged_shapes_and_untouched_slots() { + let shapes = vec![shape(1, 1)]; + let flags = launch_ahead_flags(&shapes, Some(&shapes), false, true, false, &[]); + assert!(flags.consume); + } + + #[test] + fn slot_handoff_blocks_consume_even_under_identical_shapes() { + // A finish + admission can reuse a slot id without changing the + // shape — the desync class the first gate run hit. + let shapes = vec![shape(1, 1)]; + let flags = launch_ahead_flags(&shapes, Some(&shapes), true, true, false, &[]); + assert!(!flags.consume); + } + + #[test] + fn active_row_count_is_part_of_shape_equality() { + // Same bucket/slots but a row flipped active <-> pad must not consume: + // a padding input is not value-distinguishable from an active one. + let leased = vec![shape(1, 1)]; + let shapes = vec![shape(1, 0)]; + let flags = launch_ahead_flags(&shapes, Some(&leased), false, true, false, &[]); + assert!(!flags.consume); + } + + #[test] + fn no_lease_without_an_empty_queue() { + let shapes = vec![shape(1, 1)]; + let flags = launch_ahead_flags(&shapes, None, false, false, false, &[]); + assert!(!flags.lease && !flags.consume); + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/openinfer-kernels/csrc/glm52/glm52_decode_feed.cu b/openinfer-kernels/csrc/glm52/glm52_decode_feed.cu new file mode 100644 index 00000000..67163ca7 --- /dev/null +++ b/openinfer-kernels/csrc/glm52/glm52_decode_feed.cu @@ -0,0 +1,38 @@ +// Device-side decode-step input feed: advance the per-row step inputs in +// place so the next whole-step graph replay can be enqueued without a host +// round-trip. Runs as an eager launch between two replays on the decode +// stream; on a non-leased step the host prologue rewrites all four buffers, +// so an unused advance is harmless by construction. +#include +#include + +__global__ void glm52_decode_feed_kernel(const int32_t* __restrict__ argmax_indices, + uint32_t* __restrict__ token_ids, + uint32_t* __restrict__ positions, + int64_t* __restrict__ slot_mapping, + int32_t* __restrict__ seq_lens, + int rows) { + const int row = static_cast(threadIdx.x); + if (row >= rows) { + return; + } + token_ids[row] = static_cast(argmax_indices[row]); + positions[row] += 1u; + slot_mapping[row] += 1; + seq_lens[row] += 1; +} + +extern "C" cudaError_t glm52_decode_feed_launch_cuda(const int32_t* argmax_indices, + uint32_t* token_ids, + uint32_t* positions, + int64_t* slot_mapping, + int32_t* seq_lens, + int rows, + cudaStream_t stream) { + if (rows <= 0 || rows > 32) { + return cudaErrorInvalidValue; + } + glm52_decode_feed_kernel<<<1, 32, 0, stream>>>(argmax_indices, token_ids, positions, + slot_mapping, seq_lens, rows); + return cudaGetLastError(); +} diff --git a/openinfer-kernels/src/ffi/glm52.rs b/openinfer-kernels/src/ffi/glm52.rs index 0a6fc307..c4630134 100644 --- a/openinfer-kernels/src/ffi/glm52.rs +++ b/openinfer-kernels/src/ffi/glm52.rs @@ -15,6 +15,16 @@ mod deepgemm_mqa; pub use deepgemm_mqa::*; unsafe extern "C" { + pub fn glm52_decode_feed_launch_cuda( + argmax_indices: *const i32, + token_ids: *mut u32, + positions: *mut u32, + slot_mapping: *mut i64, + seq_lens: *mut i32, + rows: i32, + stream: CUstream, + ) -> CUresult; + pub fn glm52_deepgemm_mn_major_tma_aligned_f32_cuda( input: *const f32, output: *mut f32, diff --git a/openinfer-kernels/src/ops/glm52.rs b/openinfer-kernels/src/ops/glm52.rs index 649f379e..a7437a6f 100644 --- a/openinfer-kernels/src/ops/glm52.rs +++ b/openinfer-kernels/src/ops/glm52.rs @@ -1,3 +1,4 @@ +mod decode_feed; mod deepgemm_grouped; mod deepgemm_layout; mod deepgemm_mqa; @@ -14,6 +15,7 @@ mod topk; mod trtllm_grouped; mod trtllm_linear; +pub use decode_feed::*; pub use deepgemm_grouped::*; pub use deepgemm_layout::*; pub use deepgemm_mqa::*; diff --git a/openinfer-kernels/src/ops/glm52/decode_feed.rs b/openinfer-kernels/src/ops/glm52/decode_feed.rs new file mode 100644 index 00000000..f0346cc4 --- /dev/null +++ b/openinfer-kernels/src/ops/glm52/decode_feed.rs @@ -0,0 +1,55 @@ +use anyhow::{Result, anyhow, ensure}; +use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut}; + +use crate::ffi; +use crate::tensor::DeviceContext; + +/// Advance the decode-step device inputs in place for a self-fed next step: +/// `token_ids[row] = argmax_indices[row]`, and `positions`/`slot_mapping`/ +/// `seq_lens` each move one position forward. Enqueued between two whole-step +/// graph replays; a host prologue that follows instead simply overwrites all +/// four buffers. +pub fn glm52_decode_feed_launch( + ctx: &DeviceContext, + argmax_indices: &CudaSlice, + token_ids: &mut CudaSlice, + positions: &mut CudaSlice, + slot_mapping: &mut CudaSlice, + seq_lens: &mut CudaSlice, + rows: usize, +) -> Result<()> { + ensure!(rows > 0, "GLM5.2 decode feed needs at least one row"); + for (len, name) in [ + (argmax_indices.len(), "argmax_indices"), + (token_ids.len(), "token_ids"), + (positions.len(), "positions"), + (slot_mapping.len(), "slot_mapping"), + (seq_lens.len(), "seq_lens"), + ] { + ensure!( + len >= rows, + "GLM5.2 decode feed buffer {name} holds {len} rows, step needs {rows}" + ); + } + + let (argmax_ptr, _argmax_guard) = argmax_indices.device_ptr(&ctx.stream); + let (token_ptr, _token_guard) = token_ids.device_ptr_mut(&ctx.stream); + let (pos_ptr, _pos_guard) = positions.device_ptr_mut(&ctx.stream); + let (slot_ptr, _slot_guard) = slot_mapping.device_ptr_mut(&ctx.stream); + let (seq_ptr, _seq_guard) = seq_lens.device_ptr_mut(&ctx.stream); + + let result = unsafe { + ffi::glm52_decode_feed_launch_cuda( + argmax_ptr as *const i32, + token_ptr as *mut u32, + pos_ptr as *mut u32, + slot_ptr as *mut i64, + seq_ptr as *mut i32, + rows as i32, + ctx.stream.cu_stream(), + ) + }; + result + .result() + .map_err(|err| anyhow!("GLM5.2 decode feed CUDA launch failed: {err}")) +}