metrale_model_arch/glm5next_layer/
levers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GLM-5.3-Flash layer launch levers: the prefill sub-chunk width, cuBLASLt for
4//! wide projections, and the batched DSA indexer query.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants:
8//! - Each function reads its environment variable once per process and caches the result.
9
10/// 2026-09-25: Default tokens per batched prefill sub-chunk, the width `Glm5NextLayer::prefill`
11/// hands `Glm5NextLayer::forward_k` (overridable through [`prefill_rows`]).
12///
13/// 16 is `DENSE_GEMV_BATCHM_MAX_M`, the widest M that `ops::dense_mm_bf16` sends to the batched
14/// GEMV (`dense_gemv_bf16_batchm`), whose rows carry the same bits as the M = 1 GEMV. A wider
15/// sub-chunk moves the dense projections to cuBLASLt or the tile GEMM, which accumulate in a
16/// different order. Measured 2026-09-02: the 12 GLM prefill projection shapes cost 1.36-1.98x less
17/// per token at M = 16 than at M = 8 in 11 of 12 (N4096 K128: 0.77x).
18///
19/// The routed MoE does not follow the width: `glm5next_mlp::forward::forward_moe` splits a wider
20/// row group into even sub-groups of at most `row_batch_max()`, whose default and ceiling is
21/// `MOE_ROW_BATCH_MAX_ROWS` (8).
22pub(crate) const PREFILL_ROWS: usize = 16;
23
24/// 2026-09-25: Send GLM projections with M above `DENSE_GEMV_BATCHM_MAX_M` to cuBLASLt BF16
25/// instead of the tile GEMM. Read by the KDA, DSA and MLP blocks, always behind a
26/// `> DENSE_GEMV_BATCHM_MAX_M` row test, so the M = 1 GEMV and the batched GEMV are never
27/// replaced. On unless `METRALE_GLM_CUBLAS_PROJ=0`.
28pub(crate) fn cublas_wide_proj() -> bool {
29    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
30    *E.get_or_init(|| {
31        let on = std::env::var("METRALE_GLM_CUBLAS_PROJ").as_deref() != Ok("0");
32        tracing::warn!(
33            "METRALE_GLM_CUBLAS_PROJ: wide GLM projections (M > {}) use {}",
34            metrale_model_layers::layers::ops::DENSE_GEMV_BATCHM_MAX_M,
35            if on {
36                "cuBLASLt BF16"
37            } else {
38                "the scalar tile GEMM"
39            }
40        );
41        on
42    })
43}
44
45/// 2026-09-25: Compute the DSA indexer `wq_b` projection for all prefill rows in one cuBLASLt
46/// GEMM instead of one M = 1 GEMV per row. Off unless `METRALE_GLM_DSA_BATCH_QIDX=1`.
47///
48/// Prefill only: it is read in `select_rows_batched`, which runs only when
49/// `batch_select_enabled` holds (`is_prefill && k > 1`). Decode and the speculative verify keep
50/// the per-row GEMV. The GEMM accumulates in a different order from the GEMV.
51pub(crate) fn dsa_batch_qidx() -> bool {
52    static E: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
53    *E.get_or_init(|| {
54        let on = std::env::var("METRALE_GLM_DSA_BATCH_QIDX").as_deref() == Ok("1");
55        tracing::warn!(
56            "METRALE_GLM_DSA_BATCH_QIDX: prefill DSA indexer q uses {}",
57            if on {
58                "ONE batched cuBLASLt GEMM"
59            } else {
60                "one M=1 GEMV per row"
61            }
62        );
63        on
64    })
65}
66
67/// 2026-09-25: `PREFILL_ROWS`, overridable at launch with `METRALE_GLM_PREFILL_ROWS` (values
68/// below 1 or unparsable are ignored). `1` selects the per-token walk: `Glm5NextLayer::prefill`
69/// takes the batched sub-chunk path only when `rows > 1`. A width above
70/// `DENSE_GEMV_BATCHM_MAX_M` changes the projection numerics as well as the speed.
71pub fn prefill_rows() -> usize {
72    static ROWS: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
73    *ROWS.get_or_init(|| {
74        let r = std::env::var("METRALE_GLM_PREFILL_ROWS")
75            .ok()
76            .and_then(|v| v.parse::<usize>().ok())
77            .filter(|r| *r >= 1)
78            .unwrap_or(PREFILL_ROWS);
79        if r != PREFILL_ROWS {
80            tracing::warn!("GLM prefill sub-chunk overridden to {r} rows (default {PREFILL_ROWS})");
81        }
82        r
83    })
84}
metrale_model_arch/glm5next_layer/levers.rs

metrale_model_arch/glm5next_layer/
levers.rs
