metrale_model_arch/glm5next_mlp/
forward.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The GLM-5.3 MLP forward over `rows` rows: the dense SwiGLU FFN (also used for the
4//! shared expert) and the routed MoE site.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants:
8//! - `forward_dense` and `forward_moe` leave a partial sum in `out` when the site is split over
9//!   TP or EP ranks; the caller all-reduces.
10//! - `forward_moe` zeroes `expert_out` before any expert writes it, so an expert another EP rank
11//!   owns adds zero.
12
13use anyhow::{Result, bail};
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::KernelLaunch;
16
17use super::forward_prefill_gemm;
18use super::weights::Glm5NextMoeWeights;
19use super::{Glm5NextMlpConfig, Glm5NextMlpKernels};
20
21mod dense;
22mod launch;
23mod moe_experts;
24mod workspace;
25
26pub use dense::forward_dense;
27use launch::{gemm, swiglu};
28pub use workspace::{mlp_ws_bytes, mlp_ws_total_bytes};
29
30const ACT_BLOCK: u32 = 256;
31
32/// 2026-09-25: `out = silu(min(gate, limit)) * clamp(up, -limit, limit)` over `n` elements
33/// (`glm5next_swiglu_clamp`), for the grouped-prefill module.
34pub(super) fn swiglu_rows(
35    gpu: &dyn GpuBackend,
36    k: KernelHandle,
37    gate: DevicePtr,
38    up: DevicePtr,
39    out: DevicePtr,
40    n: usize,
41    limit: f32,
42    stream: u64,
43) -> Result<()> {
44    swiglu(gpu, k, gate, up, out, n, limit, stream)
45}
46
47/// 2026-09-25: MLP scratch for up to `max_rows` rows, allocated once at load. The buffer sizes
48/// are listed in `mlp_ws_bytes`.
49pub struct Glm5NextMlpWorkspace {
50    /// 2026-09-25: Gate, up and activated values, BF16, each
51    /// `max(rows * max_inter, rows * top_k * moe_intermediate)` elements.
52    a_gate: DevicePtr,
53    a_up: DevicePtr,
54    a_act: DevicePtr,
55    /// 2026-09-25: `[rows, num_experts]` F32 router logits.
56    logits: DevicePtr,
57    /// 2026-09-25: `[rows, top_k]` selected ids (I32); `wts` holds their F32 weights.
58    ids: DevicePtr,
59    wts: DevicePtr,
60    /// 2026-09-25: `[rows, top_k, hidden]` BF16 routed outputs; `forward_moe` zeroes it before
61    /// the experts run.
62    expert_out: DevicePtr,
63    /// 2026-09-25: `[rows, hidden]` BF16 shared-expert output.
64    shared_out: DevicePtr,
65    /// 2026-09-25: `[rows * top_k]` I32 union expert ids, `-1` for an unused entry. Row-batched
66    /// path only.
67    u_eid: DevicePtr,
68    /// 2026-09-25: `[rows * top_k, rows]` I32: the slot row `r` gave union entry `u`, `-1` when
69    /// row `r` did not select it.
70    u_slot: DevicePtr,
71    /// 2026-09-25: `[rows * top_k]` I32: sorted row to original token. Grouped GEMM only.
72    sorted_token_ids: DevicePtr,
73    /// 2026-09-25: `[rows * top_k]` I32: sorted row to expert id. Grouped GEMM only.
74    sorted_expert_ids: DevicePtr,
75    /// 2026-09-25: `[num_experts + 1]` I32 prefix sum over the expert-sorted rows.
76    expert_offsets: DevicePtr,
77    /// 2026-09-25: `[rows, top_k]` I32: a slot's row in the expert-sorted output, read by
78    /// `glm5next_moe_combine_indexed`.
79    token_to_perm: DevicePtr,
80    max_inter: usize,
81    /// 2026-09-25: Widest row group this scratch serves, at least 1.
82    max_rows: usize,
83    /// 2026-09-25: `max_rows * top_k`, the routed-slot count the grouped-path buffers hold.
84    max_total_expanded: usize,
85}
86
87/// 2026-09-25: Whether the text layers share one MLP workspace: yes unless
88/// `METRALE_GLM_MLP_WS_SHARED=0`. Read once.
89pub fn mlp_ws_shared() -> bool {
90    static S: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
91    *S.get_or_init(|| {
92        let shared = std::env::var("METRALE_GLM_MLP_WS_SHARED").as_deref() != Ok("0");
93        if !shared {
94            tracing::warn!(
95                "GLM MLP workspace: PER-LAYER (METRALE_GLM_MLP_WS_SHARED=0) — the pre-P2 \
96                 allocation, one scratch per layer"
97            );
98        }
99        shared
100    })
101}
102
103impl Glm5NextMlpWorkspace {
104    /// 2026-09-25: Widest row group this scratch serves.
105    pub(super) fn max_rows(&self) -> usize {
106        self.max_rows
107    }
108    /// 2026-09-25: `max_rows * top_k`, the routed-slot count the grouped buffers hold.
109    pub(super) fn max_total_expanded(&self) -> usize {
110        self.max_total_expanded
111    }
112    pub(super) fn ids(&self) -> DevicePtr {
113        self.ids
114    }
115    pub(super) fn a_gate(&self) -> DevicePtr {
116        self.a_gate
117    }
118    pub(super) fn a_up(&self) -> DevicePtr {
119        self.a_up
120    }
121    pub(super) fn a_act(&self) -> DevicePtr {
122        self.a_act
123    }
124    pub(super) fn expert_out(&self) -> DevicePtr {
125        self.expert_out
126    }
127    pub(super) fn sorted_token_ids(&self) -> DevicePtr {
128        self.sorted_token_ids
129    }
130    pub(super) fn sorted_expert_ids(&self) -> DevicePtr {
131        self.sorted_expert_ids
132    }
133    pub(super) fn expert_offsets(&self) -> DevicePtr {
134        self.expert_offsets
135    }
136    pub(super) fn token_to_perm(&self) -> DevicePtr {
137        self.token_to_perm
138    }
139}
140
141/// 2026-09-25: `METRALE_NO_GLM_MOE_ROW_BATCH=1` turns the row-batched routed path off; each row
142/// then runs its own expert launches. Read once.
143fn row_batch_disabled() -> bool {
144    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
145    *F.get_or_init(|| std::env::var("METRALE_NO_GLM_MOE_ROW_BATCH").as_deref() == Ok("1"))
146}
147
148/// 2026-09-25: Widest sub-group of the row-batched routed path: `METRALE_GLM_MOE_ROW_BATCH_MAX`,
149/// default `MOE_ROW_BATCH_MAX_ROWS`, clamped to `1..=MOE_ROW_BATCH_MAX_ROWS`. Read once.
150pub fn row_batch_max() -> usize {
151    static M: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
152    *M.get_or_init(|| {
153        let m = std::env::var("METRALE_GLM_MOE_ROW_BATCH_MAX")
154            .ok()
155            .and_then(|v| v.parse::<usize>().ok())
156            .unwrap_or(MOE_ROW_BATCH_MAX_ROWS)
157            .clamp(1, MOE_ROW_BATCH_MAX_ROWS);
158        if m != MOE_ROW_BATCH_MAX_ROWS {
159            tracing::warn!(
160                "GLM MoE row-batch width capped at {m} (default {MOE_ROW_BATCH_MAX_ROWS})"
161            );
162        }
163        m
164    })
165}
166
167/// 2026-09-25: Widest compiled `w4a16_gemv_sw_moe_batchm_mR` tier: the
168/// `METRALE_MOE_BATCHM_ENTRY(2..=8)` instances in `kernels/gb10/common/w4a16_gemv.cu`, held at
169/// index `R - 2` of `Glm5NextMlpKernels::w4a16_gemv_sw_moe_batchm`. `forward_moe` splits a wider
170/// row group into sub-groups no wider than this.
171pub const MOE_ROW_BATCH_MAX_ROWS: usize = 8;
172
173/// 2026-09-25: Split `rows` into consecutive `(start, width)` sub-groups of at most `cap`, as
174/// even as the count allows. There is no width-1 tier; at `cap = MOE_ROW_BATCH_MAX_ROWS` no
175/// group of a split of 2 or more rows is one row wide. At a smaller `cap` one can be (3 rows at
176/// cap 2), and `forward_moe` then does not take the row-batched path.
177///
178/// The split does not change a row's result: its routed output is the sum over its own top-k
179/// slots, and each slot's GEMV accumulates in per-row registers, whatever other rows share the
180/// sweep.
181fn moe_row_groups(rows: usize, cap: usize) -> Vec<(usize, usize)> {
182    let n = rows.div_ceil(cap.max(1)).max(1);
183    let mut out = Vec::with_capacity(n);
184    let mut start = 0usize;
185    for i in 0..n {
186        let w = (rows - start).div_ceil(n - i);
187        out.push((start, w));
188        start += w;
189    }
190    out
191}
192
193/// 2026-09-25: Most ids `glm5next_moe_row_union` resolves: it runs as one block of
194/// `rows * top_k` threads, and `forward_moe` does not take the row-batched path above this.
195pub const MOE_ROW_UNION_MAX_IDS: usize = 64;
196
197/// 2026-09-25: Log once per row count (counts from 15 up share one bit) whether the routed
198/// experts took the row-batched path.
199fn announce_row_batch(batched: bool, rows: usize) {
200    use std::sync::atomic::{AtomicU16, Ordering};
201    static SEEN: AtomicU16 = AtomicU16::new(0);
202    let bit = 1u16 << rows.min(15);
203    if SEEN.fetch_or(bit, Ordering::Relaxed) & bit != 0 {
204        return;
205    }
206    if batched {
207        tracing::info!("GLM MoE: row-batched expert union ({rows} rows, one sweep each)");
208    } else {
209        tracing::info!("GLM MoE: per-row expert sweeps ({rows} rows)");
210    }
211}
212
213/// 2026-09-25: Log once per row count (counts from 31 up share one bit) whether a row group
214/// wider than `MOE_ROW_BATCH_MAX_ROWS` took the grouped GEMM; narrower groups log nothing.
215fn announce_grouped_prefill(on: bool, rows: usize) {
216    use std::sync::atomic::{AtomicU32, Ordering};
217    static SEEN: AtomicU32 = AtomicU32::new(0);
218    if rows <= MOE_ROW_BATCH_MAX_ROWS {
219        return;
220    }
221    let bit = 1u32 << (rows.min(31));
222    if SEEN.fetch_or(bit, Ordering::Relaxed) & bit != 0 {
223        return;
224    }
225    if on {
226        let t = super::forward_prefill_gemm::gemm_tile();
227        tracing::info!(
228            "GLM MoE prefill: grouped tensor-core W4A16 GEMM, {rows} rows in ONE launch \
229             per projection, tile `{}` (M_TILE {}, N_TILE {}, {} threads) \
230             (METRALE_GLM_MOE_PREFILL_GEMM=0 to restore the GEMV path; \
231             METRALE_GLM_MOE_GEMM_TILE=base for the prior tile)",
232            t.name,
233            t.m_tile,
234            t.n_tile,
235            t.threads
236        );
237    } else {
238        tracing::info!(
239            "GLM MoE prefill: row-batched GEMV, {rows} rows split into \
240             {MOE_ROW_BATCH_MAX_ROWS}-row sweeps"
241        );
242    }
243}
244
245/// 2026-09-25: `METRALE_GLM_MOE_HOST_DISPATCH=1` forces the host-dispatch expert loop (ids read
246/// back to the host, one GEMV per local expert) and turns off every device-dispatched routed
247/// path. Read once.
248fn host_dispatch_forced() -> bool {
249    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
250    *F.get_or_init(|| std::env::var("METRALE_GLM_MOE_HOST_DISPATCH").as_deref() == Ok("1"))
251}
252
253/// 2026-09-25: Log, once per process, whether the per-row expert path is device-dispatched
254/// (`w4a16_gemv_sw_moe`) or host-dispatched.
255fn announce_dispatch(grouped: bool) {
256    static ONCE: std::sync::Once = std::sync::Once::new();
257    ONCE.call_once(|| {
258        if grouped {
259            tracing::info!("GLM MoE: grouped device dispatch (no per-layer D2H)");
260        } else {
261            tracing::info!("GLM MoE: host dispatch (per-layer stream sync + D2H)");
262        }
263    });
264}
265
266/// 2026-09-25: One routed MoE site over `rows` rows of `x` into `out`: router, top-k, routed
267/// experts, shared expert, and the weighted combine.
268///
269/// The routed experts take one of three paths:
270/// - the grouped GEMM (`forward_prefill_gemm`), for `rows > MOE_ROW_BATCH_MAX_ROWS` and
271///   `rows >= prefill_gemm_min_rows()`;
272/// - the row-batched union GEMV over the sub-groups of `moe_row_groups`, for `rows >= 2`;
273/// - per row otherwise, device-dispatched (`w4a16_gemv_sw_moe`) or host-dispatched.
274///
275/// Route tracing (`METRALE_GLM_ROUTE_TRACE=1`) forces the per-row path, the one that reads `ids`
276/// back. Errors when `rows` does not fit the workspace or the bound expert count is not this
277/// rank's share.
278#[allow(clippy::too_many_arguments)]
279pub fn forward_moe(
280    gpu: &dyn GpuBackend,
281    k: &Glm5NextMlpKernels,
282    cfg: &Glm5NextMlpConfig,
283    w: &Glm5NextMoeWeights,
284    x: DevicePtr,
285    out: DevicePtr,
286    rows: usize,
287    ws: &Glm5NextMlpWorkspace,
288    stream: u64,
289) -> Result<()> {
290    if rows == 0 || rows > ws.max_rows {
291        bail!(
292            "GLM MoE: {rows} rows do not fit a workspace built for {}",
293            ws.max_rows
294        );
295    }
296    if w.experts.len() != cfg.local_experts {
297        bail!(
298            "GLM MoE: {} bound experts but this rank owns {} of {}",
299            w.experts.len(),
300            cfg.local_experts,
301            cfg.num_experts
302        );
303    }
304
305    use crate::glm5next_layer::profile;
306
307    let groups = moe_row_groups(rows, row_batch_max());
308    let grouped_prefill = rows > MOE_ROW_BATCH_MAX_ROWS
309        && rows >= forward_prefill_gemm::prefill_gemm_min_rows()
310        && forward_prefill_gemm::prefill_gemm_enabled()
311        && !host_dispatch_forced()
312        && !profile::trace_on()
313        && k.moe_sort_by_expert.0 != 0
314        && k.moe_grouped_gemm.0 != 0
315        && k.combine_indexed.0 != 0
316        && rows * cfg.top_k <= ws.max_total_expanded();
317    let batched = !grouped_prefill
318        && rows >= 2
319        && !host_dispatch_forced()
320        && !row_batch_disabled()
321        && !profile::trace_on()
322        && k.moe_row_union.0 != 0
323        && groups.iter().all(|&(_, w)| {
324            w >= 2
325                && w * cfg.top_k <= MOE_ROW_UNION_MAX_IDS
326                && k.w4a16_gemv_sw_moe_batchm[w - 2].0 != 0
327        });
328    announce_row_batch(batched, rows);
329    announce_grouped_prefill(grouped_prefill, rows);
330
331    let t = profile::start();
332    if rows > metrale_model_layers::layers::ops::DENSE_GEMV_BATCHM_MAX_M as usize
333        && crate::glm5next_layer::cublas_wide_proj()
334    {
335        metrale_model_layers::layers::ops::cublas_bf16_proj_dense_f32_out(
336            x,
337            w.router,
338            ws.logits,
339            rows as u32,
340            cfg.num_experts as u32,
341            cfg.hidden as u32,
342            stream,
343        )?;
344    } else {
345        for r in 0..rows {
346            gemm(
347                gpu,
348                k.gemm_f32,
349                k.gemv_f32,
350                KernelHandle(0),
351                x.offset(r * cfg.hidden * 2),
352                w.router,
353                ws.logits.offset(r * cfg.num_experts * 4),
354                1,
355                cfg.num_experts,
356                cfg.hidden,
357                stream,
358            )?;
359        }
360    }
361    // 2026-09-25: One top-k launch for all rows: `glm5next_router_topk` handles row `blockIdx.x`.
362    KernelLaunch::new(gpu, k.router)
363        .grid([rows as u32, 1, 1])
364        .block([ACT_BLOCK, 1, 1])
365        .arg_ptr(ws.logits)
366        .arg_ptr(w.router_bias)
367        .arg_ptr(ws.ids)
368        .arg_ptr(ws.wts)
369        .arg_u32(cfg.num_experts as u32)
370        .arg_u32(cfg.top_k as u32)
371        .arg_u32(1)
372        .arg_f32(cfg.routed_scale)
373        .arg_u32(u32::from(cfg.renormalize))
374        .arg_u32(u32::from(cfg.router_bf16_ladder))
375        .launch(stream)?;
376    profile::end(profile::MOE_ROUTER, t, gpu, stream);
377
378    // 2026-09-25: Zero every routed output first: an expert this rank does not own writes
379    // nothing, and its slots must add zero to the sum.
380    gpu.memset_async(ws.expert_out, 0, rows * cfg.top_k * cfg.hidden * 2, stream)?;
381
382    let site = moe_experts::MoeSite {
383        gpu,
384        k,
385        cfg,
386        w,
387        x,
388        rows,
389        ws,
390        stream,
391    };
392    moe_experts::per_row_experts(&site, batched, grouped_prefill)?;
393
394    if grouped_prefill {
395        // 2026-09-25: Leaves the routed outputs in expert-sorted order; the combine below reads
396        // them through `token_to_perm`.
397        let t = profile::start();
398        forward_prefill_gemm::forward_moe_grouped_prefill(gpu, k, cfg, w, x, rows, ws, stream)?;
399        profile::end(profile::MOE_EXPERTS, t, gpu, stream);
400    }
401
402    if batched {
403        moe_experts::row_batched_experts(&site, &groups)?;
404    }
405
406    let t = profile::start();
407    forward_dense(
408        gpu,
409        k,
410        cfg,
411        &w.shared,
412        cfg.local_shared_intermediate,
413        x,
414        ws.shared_out,
415        rows,
416        ws,
417        stream,
418    )?;
419
420    // 2026-09-25: The shared expert joins the routed sum in the combine, before the caller's
421    // all-reduce, so one collective reduces both partial sums.
422    profile::end(profile::MOE_SHARED, t, gpu, stream);
423    let t = profile::start();
424    // 2026-09-25: One combine launch for all rows (row = `blockIdx.x`). After the grouped GEMM
425    // the routed rows are in expert-sorted order, so that path uses
426    // `glm5next_moe_combine_indexed`, which finds each slot's row through `token_to_perm`.
427    if grouped_prefill {
428        KernelLaunch::new(gpu, k.combine_indexed)
429            .grid([rows as u32, 1, 1])
430            .block([ACT_BLOCK, 1, 1])
431            .arg_ptr(ws.expert_out)
432            .arg_ptr(ws.token_to_perm)
433            .arg_ptr(ws.wts)
434            .arg_ptr(ws.shared_out)
435            .arg_ptr(out)
436            .arg_u32(cfg.hidden as u32)
437            .arg_u32(cfg.top_k as u32)
438            .launch(stream)?;
439    } else {
440        KernelLaunch::new(gpu, k.combine)
441            .grid([rows as u32, 1, 1])
442            .block([ACT_BLOCK, 1, 1])
443            .arg_ptr(ws.expert_out)
444            .arg_ptr(ws.wts)
445            .arg_ptr(ws.shared_out)
446            .arg_ptr(out)
447            .arg_u32(cfg.hidden as u32)
448            .arg_u32(cfg.top_k as u32)
449            .launch(stream)?;
450    }
451    profile::end(profile::MOE_COMBINE, t, gpu, stream);
452    Ok(())
453}
454
455#[cfg(test)]
456mod tests;
metrale_model_arch/glm5next_mlp/forward.rs

metrale_model_arch/glm5next_mlp/
forward.rs
