spark_model/layers/glm5next_layer/
profile.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `AVAROK_GLM_PROFILE=1` — per-section decode timing for the GLM-5.3 stack.
4//!
5//! Off unless the variable is set. Every span ends in a `synchronize`, so enabling it
6//! SERIALISES the stream: read the split, not the total, and never quote a tok/s taken
7//! with it on.
8//!
9//! Sections are chosen to separate the three things that can each explain a 10x decode
10//! gap and look identical from the outside: weight bandwidth (the GEMM buckets), launch
11//! and host-sync latency (`moe_hostsync`, call counts), and collectives (`reduce_*`).
12
13use spark_runtime::gpu::GpuBackend;
14use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
15use std::time::Instant;
16
17pub const MHC: usize = 0;
18pub const NORM: usize = 1;
19pub const KDA: usize = 2;
20pub const DSA_PROJ: usize = 3;
21pub const DSA_INDEXER: usize = 4;
22pub const DSA_SELECT: usize = 5;
23pub const DSA_ATTEND: usize = 6;
24pub const REDUCE_ATTN: usize = 7;
25pub const MLP_DENSE: usize = 8;
26pub const MOE_ROUTER: usize = 9;
27pub const MOE_HOSTSYNC: usize = 10;
28pub const MOE_EXPERTS: usize = 11;
29pub const MOE_SHARED: usize = 12;
30pub const MOE_COMBINE: usize = 13;
31pub const REDUCE_MLP: usize = 14;
32/// Host-side ENQUEUE cost of the two collectives — the driver/NCCL calls only, no sync.
33/// Paired with [`REDUCE_ATTN`]/[`REDUCE_MLP`], which then time ONLY the `synchronize` that
34/// follows, i.e. the device + network + rank-skew wait. Splitting them is the difference
35/// between "the fabric is slow" and "we call it 90 times a token".
36pub const REDUCE_ATTN_ENQ: usize = 15;
37pub const REDUCE_MLP_ENQ: usize = 16;
38/// A 2-BYTE collective issued immediately before the real one, PROFILING ONLY. It is a
39/// rendezvous: neither rank leaves it until both have arrived, so it absorbs the per-call
40/// arrival jitter and charges the minimum-payload NCCL latency. The real 8 KB reduce that
41/// follows therefore starts with both ranks synchronised, which is what makes
42/// [`REDUCE_ATTN`]/[`REDUCE_MLP`] readable as network-and-kernel cost rather than "network
43/// plus whatever the other rank was still doing".
44///
45/// 🪤 Aggregate rank skew being ~0 does NOT mean per-call wait is ~0 — the two ranks trade the
46/// lead call by call, so the NET cancels while every individual call still pays |jitter|.
47/// That is exactly why this probe exists and why the both-rank profile diff was not enough.
48pub const REDUCE_ATTN_BAR: usize = 17;
49pub const REDUCE_MLP_BAR: usize = 18;
50/// mHC split by kernel: `hc_pre` (the mix + finish pair) vs `hc_post` (+ expand/head).
51pub const MHC_POST: usize = 19;
52/// `kda_mixer` split three ways. The parent bucket is 53.5% of prefill — by far the
53/// largest — but it spans three very different things: batched dense projections, a
54/// per-token recurrence that CANNOT be batched in the recurrent arm, and the output
55/// norm + o_proj. Chunking the recurrence measured ~4% SLOWER, which only makes sense
56/// if the recurrence is not the bulk; these say which part to actually attack.
57pub const KDA_FRONT: usize = 20;
58pub const KDA_RECUR: usize = 21;
59pub const KDA_BACK: usize = 22;
60const N: usize = 23;
61
62const NAMES: [&str; N] = [
63    "mhc",
64    "norm",
65    "kda_mixer",
66    "dsa_proj",
67    "dsa_indexer",
68    "dsa_select",
69    "dsa_attend",
70    "reduce_attn",
71    "mlp_dense",
72    "moe_router",
73    "moe_hostsync",
74    "moe_experts",
75    "moe_shared",
76    "moe_combine",
77    "reduce_mlp",
78    "reduce_attn_enq",
79    "reduce_mlp_enq",
80    "reduce_attn_bar",
81    "reduce_mlp_bar",
82    "mhc_post",
83    "kda_front",
84    "kda_recur",
85    "kda_back",
86];
87
88static NANOS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
89static CALLS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
90static STEPS: AtomicU64 = AtomicU64::new(0);
91
92/// `AVAROK_GLM_PROFILE=1` full · `=2` COLLECTIVES ONLY.
93///
94/// 🔴 Level 2 exists because level 1 cannot answer its own biggest question. Every span ends in
95/// a `synchronize`, ~15 of them per layer per rank, and any host-side scheduling difference
96/// between the two ranks accumulates between rendezvous points and is then charged to the next
97/// `reduce_*_bar`. Level 2 syncs ONLY the collective spans, so the bar it reports is arrival
98/// jitter the model actually has — not jitter the profiler manufactured.
99fn level() -> u8 {
100    static L: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
101    *L.get_or_init(|| match std::env::var("AVAROK_GLM_PROFILE").as_deref() {
102        Ok("1") => 1,
103        Ok("2") => 2,
104        _ => 0,
105    })
106}
107
108pub fn on() -> bool {
109    level() != 0
110}
111
112/// True only at level 1 — the per-kernel spans.
113pub fn full() -> bool {
114    level() == 1
115}
116
117/// Open a span. `None` when profiling is off, which makes [`end`] a no-op.
118pub fn start() -> Option<Instant> {
119    full().then(Instant::now)
120}
121
122/// A span that survives level 2: the collectives and their rendezvous probe.
123pub fn start_hot() -> Option<Instant> {
124    on().then(Instant::now)
125}
126
127pub fn end(bucket: usize, t0: Option<Instant>, gpu: &dyn GpuBackend, stream: u64) {
128    let _ = end_us(bucket, t0, gpu, stream);
129}
130
131/// Same, returning the measured microseconds (0.0 when profiling is off).
132pub fn end_us(bucket: usize, t0: Option<Instant>, gpu: &dyn GpuBackend, stream: u64) -> f64 {
133    let Some(t0) = t0 else { return 0.0 };
134    let _ = gpu.synchronize(stream);
135    let ns = t0.elapsed().as_nanos() as u64;
136    NANOS[bucket].fetch_add(ns, Relaxed);
137    CALLS[bucket].fetch_add(1, Relaxed);
138    ns as f64 / 1e3
139}
140
141/// `AVAROK_GLM_ROUTE_TRACE=1` — emit one line per reduce site per layer per token carrying the
142/// router's selected GLOBAL expert ids and the measured rendezvous (arrival-skew) time.
143///
144/// The router is REPLICATED and bit-identical on every rank, so rank 0's ids are the whole
145/// picture: any static ownership map can be scored offline from this one trace without a
146/// second run. Costly (one log line per MoE layer per token) — trace, then turn it off.
147pub fn trace_on() -> bool {
148    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
149    *ON.get_or_init(|| std::env::var("AVAROK_GLM_ROUTE_TRACE").as_deref() == Ok("1"))
150}
151
152thread_local! {
153    /// The ids `forward_moe` last read back to the host, for the trace line that follows.
154    static ROUTE: std::cell::RefCell<Vec<i32>> = const { std::cell::RefCell::new(Vec::new()) };
155}
156
157/// Called from `forward_moe` right after the routing D2H. No-op unless tracing.
158pub fn stash_route(ids: &[i32]) {
159    if !trace_on() {
160        return;
161    }
162    ROUTE.with(|r| {
163        let mut v = r.borrow_mut();
164        v.clear();
165        v.extend_from_slice(ids);
166    });
167}
168
169/// Emit the joined line. `site` is `attn` or `mlp`; `moe` says whether THIS layer's MLP is
170/// routed (the 3 dense layers are the natural control: no EP imbalance is possible there).
171pub fn trace_bar(site: &str, layer: usize, moe: bool, us: f64) {
172    if !trace_on() {
173        return;
174    }
175    let step = STEPS.load(Relaxed);
176    ROUTE.with(|r| {
177        let v = r.borrow();
178        let ids = v
179            .iter()
180            .map(|x| x.to_string())
181            .collect::<Vec<_>>()
182            .join(",");
183        tracing::warn!(
184            "GLMTRACE step={step} site={site} L={layer} moe={} bar_us={us:.1} ids={ids}",
185            u8::from(moe)
186        );
187    });
188}
189
190/// 4-byte device scratch for the rendezvous probe. Allocated once, PROFILING ONLY.
191pub fn probe_buf(gpu: &dyn GpuBackend) -> u64 {
192    static P: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
193    *P.get_or_init(|| gpu.alloc(4).map(|p| p.0).unwrap_or(0))
194}
195
196/// Record a span WITHOUT synchronising — host-side wall time only.
197pub fn end_nosync(bucket: usize, t0: Option<Instant>) {
198    let Some(t0) = t0 else { return };
199    NANOS[bucket].fetch_add(t0.elapsed().as_nanos() as u64, Relaxed);
200    CALLS[bucket].fetch_add(1, Relaxed);
201}
202
203/// Close one token. Dumps a cumulative per-token split every 8 steps, then keeps going —
204/// the totals are cumulative so a later dump is simply better averaged.
205pub fn step() {
206    if !on() {
207        return;
208    }
209    let s = STEPS.fetch_add(1, Relaxed) + 1;
210    if !s.is_multiple_of(8) {
211        return;
212    }
213    let total: u64 = NANOS.iter().map(|n| n.load(Relaxed)).sum();
214    let mut rows: Vec<(usize, u64, u64)> = (0..N)
215        .map(|i| (i, NANOS[i].load(Relaxed), CALLS[i].load(Relaxed)))
216        .collect();
217    rows.sort_by_key(|r| std::cmp::Reverse(r.1));
218    let mut out = format!(
219        "GLM decode profile after {s} steps — {:.2} ms/token measured under profiling\n",
220        total as f64 / 1e6 / s as f64
221    );
222    for (i, ns, calls) in rows {
223        if calls == 0 {
224            continue;
225        }
226        out += &format!(
227            "  {:<13} {:>8.2} ms/tok  {:>6.1}%  {:>5} calls/tok  {:>7.1} us/call\n",
228            NAMES[i],
229            ns as f64 / 1e6 / s as f64,
230            100.0 * ns as f64 / total.max(1) as f64,
231            calls / s,
232            ns as f64 / 1e3 / calls.max(1) as f64,
233        );
234    }
235    tracing::warn!("{out}");
236}
spark_model/layers/glm5next_layer/profile.rs

spark_model/layers/glm5next_layer/
profile.rs
