metrale_model_layers/layers/ops/
ssm_gdn_tc_route.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The GDN prefill state spines' entry names, and the route lines
4//! built from them.
5//!
6//! # Why the name is a constant
7//!
8//! The tensor-core module holds two entries, `…_tcfuse` (one bf16 limb of
9//! `S_c`) and `…_tcfuse_x2` (two limbs), and the lever ships `_x2`. A route
10//! line that names only the family, such as
11//!
12//! ```text
13//! GDN state spine: gated_delta_rule_chunk_delta_h_tcfuse (METRALE_GDN_PREFILL_TC; …)
14//! ```
15//!
16//! reads as the 1-limb entry. So the name exists once, here:
17//! `qwen3_ssm::init_kernels` binds the handle with it and
18//! [`gdn_tc_spine_route_line`] prints it.
19//!
20//! Owner: model-layers ops (GDN).
21//! Invariants: none beyond the types.
22
23/// 2026-09-25: The spine entry the `gdn_prefill_tc` lever launches: two bf16
24/// limbs of `S_c`. Both entries take the same arguments, grid, block and
25/// shared memory, so nothing downstream depends on which one is bound.
26pub const GDN_TC_SPINE_ENTRY: &str = "gated_delta_rule_chunk_delta_h_tcfuse_x2";
27
28/// 2026-09-25: The module the entry lives in:
29/// `kernels/hopper/common/gated_delta_rule_chunk_tc.cu`.
30pub const GDN_TC_SPINE_MODULE: &str = "gated_delta_rule_chunk_tc";
31
32/// 2026-09-25: The scalar spine entries `init_kernels::fused_spine_kernel`
33/// binds, as [`gdn_scalar_spine`] picks them. The init route line and the
34/// handle are built from the same strings.
35/// 2026-09-28: The default: `_vfused`'s arithmetic with W, K and U double-buffered
36/// through cp.async (bit-identical to `_vfused`, 1.4x faster on GB10 at 1.6k-8k
37/// tokens per chunk).
38pub const GDN_SCALAR_SPINE_PIPE: &str = "gated_delta_rule_chunk_delta_h_pipe";
39/// 2026-09-25: SPLIT=4, 512 threads. Not the default; `fused_spine_kernel`'s
40/// doc says why.
41pub const GDN_SCALAR_SPINE_VTILE: &str = "gated_delta_rule_chunk_delta_h_vtile";
42/// 2026-09-25: SPLIT=2, 256 threads, single-buffered; `METRALE_GDN_PIPE=0`.
43pub const GDN_SCALAR_SPINE_VFUSED: &str = "gated_delta_rule_chunk_delta_h_vfused";
44
45/// 2026-09-28: The scalar spine a process binds (`init_kernels::fused_spine_kernel`) and
46/// launches (`ops::gdn_prefill_fla`, whose block size and shared memory depend on it). Both
47/// read it through [`gdn_scalar_spine`], so they cannot disagree.
48#[derive(Clone, Copy, Debug, PartialEq, Eq)]
49pub enum GdnScalarSpine {
50    Pipe,
51    Vtile,
52    Vfused,
53}
54
55impl GdnScalarSpine {
56    pub fn entry(self) -> &'static str {
57        match self {
58            Self::Pipe => GDN_SCALAR_SPINE_PIPE,
59            Self::Vtile => GDN_SCALAR_SPINE_VTILE,
60            Self::Vfused => GDN_SCALAR_SPINE_VFUSED,
61        }
62    }
63}
64
65/// 2026-09-28: From the values of `METRALE_GDN_PIPE` and `METRALE_GDN_VTILE`: `PIPE=1`
66/// picks the pipe spine (as before), else `VTILE=1` the vtile spine, else `PIPE=0` the
67/// vfused spine, else the pipe spine.
68pub fn gdn_scalar_spine_from(pipe: Option<&str>, vtile: Option<&str>) -> GdnScalarSpine {
69    match (pipe, vtile) {
70        (Some("1"), _) => GdnScalarSpine::Pipe,
71        (_, Some("1")) => GdnScalarSpine::Vtile,
72        (Some("0"), _) => GdnScalarSpine::Vfused,
73        _ => GdnScalarSpine::Pipe,
74    }
75}
76
77/// 2026-09-28: [`gdn_scalar_spine_from`] of this process's environment, read once.
78pub fn gdn_scalar_spine() -> GdnScalarSpine {
79    static SPINE: std::sync::OnceLock<GdnScalarSpine> = std::sync::OnceLock::new();
80    *SPINE.get_or_init(|| {
81        gdn_scalar_spine_from(
82            std::env::var("METRALE_GDN_PIPE").ok().as_deref(),
83            std::env::var("METRALE_GDN_VTILE").ok().as_deref(),
84        )
85    })
86}
87
88/// 2026-09-25: `GDN state spine: …`, the line `qwen3_ssm::init` prints for
89/// each layer as it binds the handles, before any prefill has run.
90///
91/// # Why it is not simply the scalar entry's name
92///
93/// With the tensor-core handle bound, the prefill launches
94/// [`GDN_TC_SPINE_ENTRY`], and the scalar entry stays bound only as the
95/// fallback the prefill's shape guards drop to. An init line that named the
96/// scalar entry would contradict the prefill's line:
97///
98/// ```text
99///    48  qwen3_ssm::init: GDN state spine: gated_delta_rule_chunk_delta_h_vfused
100/// 14400  GDN state spine: gated_delta_rule_chunk_delta_h_tcfuse_x2 (…)
101/// ```
102///
103/// So the line follows the handle the probe resolved: bound, it names
104/// [`GDN_TC_SPINE_ENTRY`]; unbound, `scalar_entry`.
105pub fn gdn_init_spine_line(tc_spine_bound: bool, scalar_entry: &str) -> String {
106    if tc_spine_bound {
107        format!(
108            "GDN state spine: {GDN_TC_SPINE_ENTRY} ([defaults] gdn_prefill_tc; the \
109             scalar spine stays bound as the fallback the prefill's shape guards \
110             drop to, and the prefill logs the entry it launches)"
111        )
112    } else {
113        format!("GDN state spine: {scalar_entry}")
114    }
115}
116
117/// 2026-09-25: `GDN state spine: …`, the line the prefill logs when the
118/// tensor-core spine runs, built from [`GDN_TC_SPINE_ENTRY`]. Pure and
119/// returning a `String`, so the tests below can check it.
120pub fn gdn_tc_spine_route_line(num_v_heads: u32, batch_size: u32, smem_bytes: u32) -> String {
121    format!(
122        "GDN state spine: {GDN_TC_SPINE_ENTRY} (METRALE_GDN_PREFILL_TC; bf16 mma.sync \
123         operands, f32 accumulator = the recurrent state, h stays f32) \
124         grid=[{num_v_heads},{batch_size}] block=256 smem={smem_bytes}B"
125    )
126}
127
128#[cfg(test)]
129mod tests {
130    use super::{
131        GDN_SCALAR_SPINE_PIPE, GDN_SCALAR_SPINE_VFUSED, GDN_SCALAR_SPINE_VTILE, GDN_TC_SPINE_ENTRY,
132        GdnScalarSpine, gdn_init_spine_line, gdn_scalar_spine_from, gdn_tc_spine_route_line,
133    };
134
135    /// 2026-09-28: The pipe spine by default and under `PIPE=1` (which also wins over
136    /// `VTILE=1`, as before), vtile under `VTILE=1`, vfused only under `PIPE=0`; any other
137    /// value is the default.
138    #[test]
139    fn scalar_spine_choice() {
140        use GdnScalarSpine::*;
141        for (pipe, vtile, want) in [
142            (None, None, Pipe),
143            (Some("1"), None, Pipe),
144            (Some("1"), Some("1"), Pipe),
145            (None, Some("1"), Vtile),
146            (Some("0"), Some("1"), Vtile),
147            (Some("0"), None, Vfused),
148            (Some("0"), Some("0"), Vfused),
149            (Some("yes"), None, Pipe),
150            (None, Some("0"), Pipe),
151        ] {
152            assert_eq!(
153                gdn_scalar_spine_from(pipe, vtile),
154                want,
155                "{pipe:?} {vtile:?}"
156            );
157        }
158        assert_eq!(Pipe.entry(), GDN_SCALAR_SPINE_PIPE);
159        assert_eq!(Vtile.entry(), GDN_SCALAR_SPINE_VTILE);
160        assert_eq!(Vfused.entry(), GDN_SCALAR_SPINE_VFUSED);
161    }
162
163    /// 2026-09-25: The line names the `_x2` entry, never the bare family name
164    /// (`…_tcfuse` followed by a space).
165    #[test]
166    fn the_route_line_names_the_entry_that_is_launched() {
167        let line = gdn_tc_spine_route_line(48, 1, 88_324);
168        assert!(line.contains(GDN_TC_SPINE_ENTRY), "{line}");
169        assert!(
170            !line.contains("gated_delta_rule_chunk_delta_h_tcfuse "),
171            "the line must not name the FAMILY where the binary launches the \
172             `_x2` member — round 12 stage 4b:\n{line}"
173        );
174    }
175
176    /// 2026-09-25: The line carries the launch geometry.
177    #[test]
178    fn the_route_line_carries_the_geometry() {
179        let line = gdn_tc_spine_route_line(48, 2, 88_324);
180        for field in [
181            "grid=[48,2]",
182            "block=256",
183            "smem=88324B",
184            "METRALE_GDN_PREFILL_TC",
185        ] {
186            assert!(line.contains(field), "missing `{field}` in:\n{line}");
187        }
188    }
189
190    /// 2026-09-25: With the tensor-core handle bound, the init line names that
191    /// entry, not the scalar spine bound behind it.
192    #[test]
193    fn the_init_line_names_the_tc_entry_when_its_handle_is_bound() {
194        let line = gdn_init_spine_line(true, GDN_SCALAR_SPINE_VFUSED);
195        assert!(line.contains(GDN_TC_SPINE_ENTRY), "{line}");
196        assert!(
197            !line.contains(GDN_SCALAR_SPINE_VFUSED),
198            "the init line must not NAME the scalar spine where the probe bound \
199             the tensor-core one:\n{line}"
200        );
201    }
202
203    /// 2026-09-25: The init line and the prefill's line name the same entry.
204    #[test]
205    fn the_two_route_lines_agree_on_the_entry() {
206        let init = gdn_init_spine_line(true, GDN_SCALAR_SPINE_VFUSED);
207        let dispatch = gdn_tc_spine_route_line(48, 1, 88_324);
208        for line in [&init, &dispatch] {
209            assert!(line.contains(GDN_TC_SPINE_ENTRY), "{line}");
210        }
211    }
212
213    /// 2026-09-25: With the probe off (by default, every target but
214    /// `kernels/hopper`), the line names the scalar entry, for each of the
215    /// three.
216    #[test]
217    fn the_init_line_names_the_scalar_entry_when_the_probe_is_off() {
218        for entry in [
219            GDN_SCALAR_SPINE_PIPE,
220            GDN_SCALAR_SPINE_VTILE,
221            GDN_SCALAR_SPINE_VFUSED,
222        ] {
223            assert_eq!(
224                gdn_init_spine_line(false, entry),
225                format!("GDN state spine: {entry}"),
226            );
227        }
228    }
229}
metrale_model_layers/layers/ops/ssm_gdn_tc_route.rs

metrale_model_layers/layers/ops/
ssm_gdn_tc_route.rs
