1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
//! Gate for door `MEMRA_GLM5_Q8_FUSE` (lane/b200-q8-fuse-20260902).
//!
//! Compares the SHIPPED two-launch chain (`rms_norm` then `quantize_q8_1`) against the fused
//! one-launch producer (`rms_norm_zq8_f32`) on synthetic inputs, per fused site shape:
//!
//! - glm5_next trunk width (n_embd = 4096) — the mHC FFN-input norm this lane wired
//! (`hyper_range_decode` / `hyper_range_decode_ws_body`'s post_attn_norm -> MoE zq8).
//! - the glm5_next KDA qkv width (64 heads * 128 head_dim = 8192) and MoE expert-ff width
//! (n_ff_exp = 1536, the glm5_next pack's expert_ff_length) — carried as candidate widths
//! for the SAME kernel shape (ncols a multiple of 32); this lane's actual dispatch wiring
//! touches only the n_embd site (see research/b200-q8-fuse-20260902/LANE.md "open items"
//! for the KDA/MoE-activation producer sites still on the unfused chain).
//!
//! Asserts BYTE-IDENTICAL bytes (f32 z, int8 qs, f32 per-32 scale) between the two arms and
//! prints N=5 per-launch us for each half of the chain and for the fused kernel. No model
//! checkpoint needed — this is a pure kernel-shape gate, run on any CUDA device (no CPU-time
//! quota concern: five short launches per shape).
use memra_engine::Engine;
/// Deterministic, dependency-free pseudo-random f32 generator (xorshift64) — avoids adding a
/// `rand` dependency for a gate binary. Seeded per-shape so the two arms compare on IDENTICAL
/// generated bytes and re-runs reproduce the same failure if one ever fires.
struct Xorshift64(u64);
impl Xorshift64 {
fn new(seed: u64) -> Self {
Xorshift64(seed | 1)
}
fn next_u64(&mut self) -> u64 {
let mut x = self.0;
x ^= x << 13;
x ^= x >> 7;
x ^= x << 17;
self.0 = x;
x
}
/// f32 in roughly [-4, 4), never exactly 0 (keeps amax > 0 on every 32-block so the
/// chain's `id = 1/d` branch and the fused kernel's take the SAME branch).
fn next_f32(&mut self) -> f32 {
let bits = (self.next_u64() >> 40) as u32; // 24 bits
let unit = (bits as f32) / (1u32 << 24) as f32; // [0,1)
(unit - 0.5) * 8.0
}
}
fn gen_vec(seed: u64, n: usize) -> Vec<f32> {
let mut rng = Xorshift64::new(seed);
(0..n).map(|_| rng.next_f32()).collect()
}
/// Run the chain and the fused arm on the same (x, w) and assert byte-identical outputs.
/// Returns (chain_launch_us_pass1_plus_pass2, fused_launch_us) medians over N iterations.
fn run_shape(
e: &Engine,
label: &str,
ncols: usize,
n: usize,
) -> Result<(), Box<dyn std::error::Error>> {
assert!(
ncols.is_multiple_of(32),
"{label}: ncols must be a multiple of 32"
);
let eps = 1e-6f32;
let x = gen_vec(0x9E3779B97F4A7C15 ^ ncols as u64, ncols);
let w = gen_vec(0xD1B54A32D192ED03 ^ ncols as u64, ncols);
let x_d = e.htod(&x)?;
let w_d = e.htod(&w)?;
// ---- correctness: one comparison pair, byte-identical required ----
let mut z_chain = e.uninit(ncols)?;
e.rms_norm(&x_d, &w_d, &mut z_chain, ncols, 1, eps)?;
let (q_chain, d_chain) = e.quantize_q8_1(&z_chain, 1, ncols)?;
let mut z_fused = e.uninit(ncols)?;
let (q_fused, d_fused) = e.rms_norm_zq8_f32(&x_d, &w_d, &mut z_fused, ncols, 1, eps)?;
let z_chain_h = e.dtoh(&z_chain)?;
let z_fused_h = e.dtoh(&z_fused)?;
// CudaSlice<i8> has no `Engine::dtoh` twin (only f32/u32/i32/u8 are wrapped) — the
// `clone_dtoh` raw-stream form is kernel_check.rs's precedent for reading a q8_1 `qs`
// buffer back (its rms_norm_q8_1 gate, same shape as this one).
let q_chain_h: Vec<i8> = e.stream().clone_dtoh(&q_chain)?;
let q_fused_h: Vec<i8> = e.stream().clone_dtoh(&q_fused)?;
e.stream().synchronize()?;
let d_chain_h = e.dtoh(&d_chain)?;
let d_fused_h = e.dtoh(&d_fused)?;
let z_match = z_chain_h == z_fused_h;
let q_match = q_chain_h == q_fused_h;
let d_match = d_chain_h == d_fused_h;
let ok = z_match && q_match && d_match;
println!(
"[q8-fuse-gate] {label} ncols={ncols} z_bytes_match={z_match} q_bytes_match={q_match} \
d_bytes_match={d_match} -> {}",
if ok { "PASS" } else { "FAIL" }
);
if !ok {
// Report the first mismatch index of whichever array diverged, for triage.
if !z_match {
for i in 0..z_chain_h.len() {
if z_chain_h[i].to_bits() != z_fused_h[i].to_bits() {
println!(
" first z mismatch at i={i}: chain={:.9e} fused={:.9e}",
z_chain_h[i], z_fused_h[i]
);
break;
}
}
}
if !q_match {
for i in 0..q_chain_h.len() {
if q_chain_h[i] != q_fused_h[i] {
println!(
" first q mismatch at i={i}: chain={} fused={}",
q_chain_h[i], q_fused_h[i]
);
break;
}
}
}
if !d_match {
for i in 0..d_chain_h.len() {
if d_chain_h[i].to_bits() != d_fused_h[i].to_bits() {
println!(
" first d mismatch at i={i}: chain={:.9e} fused={:.9e}",
d_chain_h[i], d_fused_h[i]
);
break;
}
}
}
return Err(format!("{label}: fused arm diverged from the shipped chain").into());
}
// ---- timing: N launches per arm, host-clock wall around a synchronize (rig-safe: five
// short launches per shape, no sustained GPU load) ----
let s = e.stream();
s.synchronize()?;
let mut chain_us = Vec::with_capacity(n);
for _ in 0..n {
let t0 = std::time::Instant::now();
let mut z = e.uninit(ncols)?;
e.rms_norm(&x_d, &w_d, &mut z, ncols, 1, eps)?;
let _ = e.quantize_q8_1(&z, 1, ncols)?;
s.synchronize()?;
chain_us.push(t0.elapsed().as_micros());
}
let mut fused_us = Vec::with_capacity(n);
for _ in 0..n {
let t0 = std::time::Instant::now();
let mut z = e.uninit(ncols)?;
let _ = e.rms_norm_zq8_f32(&x_d, &w_d, &mut z, ncols, 1, eps)?;
s.synchronize()?;
fused_us.push(t0.elapsed().as_micros());
}
println!(
"[q8-fuse-gate] {label} ncols={ncols} N={n} chain_us={chain_us:?} fused_us={fused_us:?}"
);
Ok(())
}
fn main() -> Result<(), Box<dyn std::error::Error>> {
let n: usize = std::env::args()
.nth(1)
.and_then(|s| s.parse().ok())
.unwrap_or(5);
let e = Engine::new(0)?;
// glm5_next trunk width — the site this lane actually wired (hyper_range_decode's
// post_attn_norm -> moe_ffn_il_zq8).
run_shape(&e, "glm5_next-n_embd", 4096, n)?;
// glm5_next KDA qkv width (64 heads * 128) — candidate shape for a future producer fusion
// (the KDA gated-norm site remains on the unfused chain; see LANE.md open items).
run_shape(&e, "glm5_next-kda_qkv", 8192, n)?;
// glm5_next MoE expert-ff width (expert_ff_length) — candidate shape for the activation
// (silu/gate*up -> q8_1) fusion site, also still unfused; see LANE.md open items.
run_shape(&e, "glm5_next-moe_ff_exp", 1536, n)?;
println!("[q8-fuse-gate] ALL SHAPES PASS (byte-identical fused vs chain)");
Ok(())
}