metrale_model_arch/glm5next_layer/
profile.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Per-section decode timing for the GLM-5.3 layers, enabled by `METRALE_GLM_PROFILE`.
4//!
5//! Owner: model-arch (GLM-5.3).
6//! Invariants:
7//! - Unless `METRALE_GLM_PROFILE` is `1` or `2`, `start` and `start_hot` return `None` and no
8//!   span synchronises.
9//! - `end` and `end_us` synchronise the stream before reading the clock, so a profiled run is
10//!   serialised and its throughput is not a serving figure.
11
12use metrale_gpu_runtime::gpu::GpuBackend;
13use std::sync::atomic::{AtomicU64, Ordering::Relaxed};
14use std::time::Instant;
15
16pub const MHC: usize = 0;
17pub const NORM: usize = 1;
18pub const KDA: usize = 2;
19pub const DSA_PROJ: usize = 3;
20pub const DSA_INDEXER: usize = 4;
21pub const DSA_SELECT: usize = 5;
22pub const DSA_ATTEND: usize = 6;
23pub const REDUCE_ATTN: usize = 7;
24pub const MLP_DENSE: usize = 8;
25pub const MOE_ROUTER: usize = 9;
26pub const MOE_HOSTSYNC: usize = 10;
27pub const MOE_EXPERTS: usize = 11;
28pub const MOE_SHARED: usize = 12;
29pub const MOE_COMBINE: usize = 13;
30pub const REDUCE_MLP: usize = 14;
31/// 2026-09-25: Host time to enqueue the attention and MLP all-reduces, without a synchronise.
32/// `REDUCE_ATTN` and `REDUCE_MLP` then time only the `synchronize` that follows the enqueue.
33pub const REDUCE_ATTN_ENQ: usize = 15;
34pub const REDUCE_MLP_ENQ: usize = 16;
35/// 2026-09-25: A 2-byte all-reduce issued just before the real one while profiling
36/// (`Glm5NextLayer::reduce_probe`), so that the ranks' arrival skew is charged here and not to
37/// `REDUCE_ATTN` or `REDUCE_MLP`.
38pub const REDUCE_ATTN_BAR: usize = 17;
39pub const REDUCE_MLP_BAR: usize = 18;
40/// 2026-09-25: `hc_post`, plus `hc_head_mean` on the last layer; `MHC` holds `hc_expand` and
41/// `hc_pre`.
42pub const MHC_POST: usize = 19;
43/// 2026-09-25: Sub-spans of the KDA mixer: the projections before the recurrence, the
44/// recurrence, and the output side, timed in `glm5next_kda`'s decode and prefill.
45pub const KDA_FRONT: usize = 20;
46pub const KDA_RECUR: usize = 21;
47pub const KDA_BACK: usize = 22;
48const N: usize = 23;
49
50const NAMES: [&str; N] = [
51    "mhc",
52    "norm",
53    "kda_mixer",
54    "dsa_proj",
55    "dsa_indexer",
56    "dsa_select",
57    "dsa_attend",
58    "reduce_attn",
59    "mlp_dense",
60    "moe_router",
61    "moe_hostsync",
62    "moe_experts",
63    "moe_shared",
64    "moe_combine",
65    "reduce_mlp",
66    "reduce_attn_enq",
67    "reduce_mlp_enq",
68    "reduce_attn_bar",
69    "reduce_mlp_bar",
70    "mhc_post",
71    "kda_front",
72    "kda_recur",
73    "kda_back",
74];
75
76static NANOS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
77static CALLS: [AtomicU64; N] = [const { AtomicU64::new(0) }; N];
78static STEPS: AtomicU64 = AtomicU64::new(0);
79
80/// 2026-09-25: `METRALE_GLM_PROFILE`: `1` times every span, `2` only the collective spans
81/// (`start_hot`), anything else is off. Read once.
82fn level() -> u8 {
83    static L: std::sync::OnceLock<u8> = std::sync::OnceLock::new();
84    *L.get_or_init(|| match std::env::var("METRALE_GLM_PROFILE").as_deref() {
85        Ok("1") => 1,
86        Ok("2") => 2,
87        _ => 0,
88    })
89}
90
91pub fn on() -> bool {
92    level() != 0
93}
94
95/// 2026-09-25: True only at level 1, the per-kernel spans.
96pub fn full() -> bool {
97    level() == 1
98}
99
100/// 2026-09-25: Open a span at level 1; `None` otherwise, which makes [`end`] a no-op.
101pub fn start() -> Option<Instant> {
102    full().then(Instant::now)
103}
104
105/// 2026-09-25: Open a span at level 1 or 2, for the collectives and their rendezvous probe.
106pub fn start_hot() -> Option<Instant> {
107    on().then(Instant::now)
108}
109
110pub fn end(bucket: usize, t0: Option<Instant>, gpu: &dyn GpuBackend, stream: u64) {
111    let _ = end_us(bucket, t0, gpu, stream);
112}
113
114/// 2026-09-25: [`end`], returning the span in microseconds; 0.0 when `t0` is `None`.
115pub fn end_us(bucket: usize, t0: Option<Instant>, gpu: &dyn GpuBackend, stream: u64) -> f64 {
116    let Some(t0) = t0 else { return 0.0 };
117    let _ = gpu.synchronize(stream);
118    let ns = t0.elapsed().as_nanos() as u64;
119    NANOS[bucket].fetch_add(ns, Relaxed);
120    CALLS[bucket].fetch_add(1, Relaxed);
121    ns as f64 / 1e3
122}
123
124/// 2026-09-25: `METRALE_GLM_ROUTE_TRACE=1`: `trace_bar` logs one line per reduce site per layer
125/// with the rendezvous time and the expert ids `stash_route` last saved. Read once. It logs only
126/// while profiling is on too, because `trace_bar` is called from `reduce_probe`.
127pub fn trace_on() -> bool {
128    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
129    *ON.get_or_init(|| std::env::var("METRALE_GLM_ROUTE_TRACE").as_deref() == Ok("1"))
130}
131
132thread_local! {
133    /// 2026-09-25: The ids `forward_moe` last read back to the host, for the next trace line.
134    static ROUTE: std::cell::RefCell<Vec<i32>> = const { std::cell::RefCell::new(Vec::new()) };
135}
136
137/// 2026-09-25: Called from `forward_moe` after the ids are read back. No-op unless tracing.
138pub fn stash_route(ids: &[i32]) {
139    if !trace_on() {
140        return;
141    }
142    ROUTE.with(|r| {
143        let mut v = r.borrow_mut();
144        v.clear();
145        v.extend_from_slice(ids);
146    });
147}
148
149/// 2026-09-25: Log the trace line of one reduce site: `site` is `attn` or `mlp`, and `moe` says
150/// whether this layer's MLP is routed. No-op unless tracing.
151pub fn trace_bar(site: &str, layer: usize, moe: bool, us: f64) {
152    if !trace_on() {
153        return;
154    }
155    let step = STEPS.load(Relaxed);
156    ROUTE.with(|r| {
157        let v = r.borrow();
158        let ids = v
159            .iter()
160            .map(|x| x.to_string())
161            .collect::<Vec<_>>()
162            .join(",");
163        tracing::warn!(
164            "GLMTRACE step={step} site={site} L={layer} moe={} bar_us={us:.1} ids={ids}",
165            u8::from(moe)
166        );
167    });
168}
169
170/// 2026-09-25: 4-byte device buffer for the rendezvous probe, allocated on first use; 0 when the
171/// allocation failed, and `reduce_probe` then skips the probe.
172pub fn probe_buf(gpu: &dyn GpuBackend) -> u64 {
173    static P: std::sync::OnceLock<u64> = std::sync::OnceLock::new();
174    *P.get_or_init(|| gpu.alloc(4).map(|p| p.0).unwrap_or(0))
175}
176
177/// 2026-09-25: Record a span without synchronising: host wall time only.
178pub fn end_nosync(bucket: usize, t0: Option<Instant>) {
179    let Some(t0) = t0 else { return };
180    NANOS[bucket].fetch_add(t0.elapsed().as_nanos() as u64, Relaxed);
181    CALLS[bucket].fetch_add(1, Relaxed);
182}
183
184/// 2026-09-25: Count one token. Every 8th token, log the cumulative time per token and calls per
185/// token of each bucket.
186pub fn step() {
187    if !on() {
188        return;
189    }
190    let s = STEPS.fetch_add(1, Relaxed) + 1;
191    if !s.is_multiple_of(8) {
192        return;
193    }
194    let total: u64 = NANOS.iter().map(|n| n.load(Relaxed)).sum();
195    let mut rows: Vec<(usize, u64, u64)> = (0..N)
196        .map(|i| (i, NANOS[i].load(Relaxed), CALLS[i].load(Relaxed)))
197        .collect();
198    rows.sort_by_key(|r| std::cmp::Reverse(r.1));
199    let mut out = format!(
200        "GLM decode profile after {s} steps — {:.2} ms/token measured under profiling\n",
201        total as f64 / 1e6 / s as f64
202    );
203    for (i, ns, calls) in rows {
204        if calls == 0 {
205            continue;
206        }
207        out += &format!(
208            "  {:<13} {:>8.2} ms/tok  {:>6.1}%  {:>5} calls/tok  {:>7.1} us/call\n",
209            NAMES[i],
210            ns as f64 / 1e6 / s as f64,
211            100.0 * ns as f64 / total.max(1) as f64,
212            calls / s,
213            ns as f64 / 1e3 / calls.max(1) as f64,
214        );
215    }
216    tracing::warn!("{out}");
217}
metrale_model_arch/glm5next_layer/profile.rs

metrale_model_arch/glm5next_layer/
profile.rs
