spark_model/layers/glm5next_layer/
profile.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `ATLAS_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;
52const N: usize = 20;
53
54const NAMES: [&str; N] = [
55    "mhc",
56    "norm",
57    "kda_mixer",
58    "dsa_proj",
59    "dsa_indexer",
60    "dsa_select",
61    "dsa_attend",
62    "reduce_attn",
63    "mlp_dense",
64    "moe_router",
65    "moe_hostsync",
66    "moe_experts",
67    "moe_shared",
68    "moe_combine",
69    "reduce_mlp",
70    "reduce_attn_enq",
71    "reduce_mlp_enq",
72    "reduce_attn_bar",
73    "reduce_mlp_bar",
74    "mhc_post",
75];
76
77static NANOS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
78static CALLS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
79static STEPS: AtomicU64 = AtomicU64::new(0);
80
81/// `ATLAS_GLM_PROFILE=1` full · `=2` COLLECTIVES ONLY.
82///
83/// 🔴 Level 2 exists because level 1 cannot answer its own biggest question. Every span ends in
84/// a `synchronize`, ~15 of them per layer per rank, and any host-side scheduling difference
85/// between the two ranks accumulates between rendezvous points and is then charged to the next
86/// `reduce_*_bar`. Level 2 syncs ONLY the collective spans, so the bar it reports is arrival
87/// jitter the model actually has — not jitter the profiler manufactured.
88fn level() -> u8 {
89    static L: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
90    *L.get_or_init(|| match std::env::var("ATLAS_GLM_PROFILE").as_deref() {
91        Ok("1") => 1,
92        Ok("2") => 2,
93        _ => 0,
94    })
95}
96
97pub fn on() -> bool {
98    level() != 0
99}
100
101/// True only at level 1 — the per-kernel spans.
102pub fn full() -> bool {
103    level() == 1
104}
105
106/// Open a span. `None` when profiling is off, which makes [`end`] a no-op.
107pub fn start() -> Option<Instant> {
108    full().then(Instant::now)
109}
110
111/// A span that survives level 2: the collectives and their rendezvous probe.
112pub fn start_hot() -> Option<Instant> {
113    on().then(Instant::now)
114}
115
116pub fn end(bucket: usize, t0: Option<Instant>, gpu: &dyn GpuBackend, stream: u64) {
117    let _ = end_us(bucket, t0, gpu, stream);
118}
119
120/// Same, returning the measured microseconds (0.0 when profiling is off).
121pub fn end_us(bucket: usize, t0: Option<Instant>, gpu: &dyn GpuBackend, stream: u64) -> f64 {
122    let Some(t0) = t0 else { return 0.0 };
123    let _ = gpu.synchronize(stream);
124    let ns = t0.elapsed().as_nanos() as u64;
125    NANOS[bucket].fetch_add(ns, Relaxed);
126    CALLS[bucket].fetch_add(1, Relaxed);
127    ns as f64 / 1e3
128}
129
130/// `ATLAS_GLM_ROUTE_TRACE=1` — emit one line per reduce site per layer per token carrying the
131/// router's selected GLOBAL expert ids and the measured rendezvous (arrival-skew) time.
132///
133/// The router is REPLICATED and bit-identical on every rank, so rank 0's ids are the whole
134/// picture: any static ownership map can be scored offline from this one trace without a
135/// second run. Costly (one log line per MoE layer per token) — trace, then turn it off.
136pub fn trace_on() -> bool {
137    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
138    *ON.get_or_init(|| std::env::var("ATLAS_GLM_ROUTE_TRACE").as_deref() == Ok("1"))
139}
140
141thread_local! {
142    /// The ids `forward_moe` last read back to the host, for the trace line that follows.
143    static ROUTE: std::cell::RefCell<Vec<i32>> = const { std::cell::RefCell::new(Vec::new()) };
144}
145
146/// Called from `forward_moe` right after the routing D2H. No-op unless tracing.
147pub fn stash_route(ids: &[i32]) {
148    if !trace_on() {
149        return;
150    }
151    ROUTE.with(|r| {
152        let mut v = r.borrow_mut();
153        v.clear();
154        v.extend_from_slice(ids);
155    });
156}
157
158/// Emit the joined line. `site` is `attn` or `mlp`; `moe` says whether THIS layer's MLP is
159/// routed (the 3 dense layers are the natural control: no EP imbalance is possible there).
160pub fn trace_bar(site: &str, layer: usize, moe: bool, us: f64) {
161    if !trace_on() {
162        return;
163    }
164    let step = STEPS.load(Relaxed);
165    ROUTE.with(|r| {
166        let v = r.borrow();
167        let ids = v
168            .iter()
169            .map(|x| x.to_string())
170            .collect::<Vec<_>>()
171            .join(",");
172        tracing::warn!(
173            "GLMTRACE step={step} site={site} L={layer} moe={} bar_us={us:.1} ids={ids}",
174            u8::from(moe)
175        );
176    });
177}
178
179/// 4-byte device scratch for the rendezvous probe. Allocated once, PROFILING ONLY.
180pub fn probe_buf(gpu: &dyn GpuBackend) -> u64 {
181    static P: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
182    *P.get_or_init(|| gpu.alloc(4).map(|p| p.0).unwrap_or(0))
183}
184
185/// Record a span WITHOUT synchronising — host-side wall time only.
186pub fn end_nosync(bucket: usize, t0: Option<Instant>) {
187    let Some(t0) = t0 else { return };
188    NANOS[bucket].fetch_add(t0.elapsed().as_nanos() as u64, Relaxed);
189    CALLS[bucket].fetch_add(1, Relaxed);
190}
191
192/// Close one token. Dumps a cumulative per-token split every 8 steps, then keeps going —
193/// the totals are cumulative so a later dump is simply better averaged.
194pub fn step() {
195    if !on() {
196        return;
197    }
198    let s = STEPS.fetch_add(1, Relaxed) + 1;
199    if !s.is_multiple_of(8) {
200        return;
201    }
202    let total: u64 = NANOS.iter().map(|n| n.load(Relaxed)).sum();
203    let mut rows: Vec<(usize, u64, u64)> = (0..N)
204        .map(|i| (i, NANOS[i].load(Relaxed), CALLS[i].load(Relaxed)))
205        .collect();
206    rows.sort_by_key(|r| std::cmp::Reverse(r.1));
207    let mut out = format!(
208        "GLM decode profile after {s} steps — {:.2} ms/token measured under profiling\n",
209        total as f64 / 1e6 / s as f64
210    );
211    for (i, ns, calls) in rows {
212        if calls == 0 {
213            continue;
214        }
215        out += &format!(
216            "  {:<13} {:>8.2} ms/tok  {:>6.1}%  {:>5} calls/tok  {:>7.1} us/call\n",
217            NAMES[i],
218            ns as f64 / 1e6 / s as f64,
219            100.0 * ns as f64 / total.max(1) as f64,
220            calls / s,
221            ns as f64 / 1e3 / calls.max(1) as f64,
222        );
223    }
224    tracing::warn!("{out}");
225}
spark_model/layers/glm5next_layer/profile.rs

spark_model/layers/glm5next_layer/
profile.rs
