metrale_model_layers/layers/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The layer types (`FfnComponent` and the modules below) and the
4//! kernel-handle resolvers several layers share.
5//!
6//! Owner: model-layers.
7//! Invariants: none beyond the types.
8
9pub mod dense_ffn;
10pub mod ep_dispatch;
11pub mod fp8_calibration;
12mod gemv_tier;
13pub mod glm_vit;
14pub mod moe;
15pub mod mtp_head;
16pub mod mtp_meta;
17pub mod mtp_multi;
18pub mod ngram_embed;
19pub mod ops;
20pub mod ple;
21pub mod qsa;
22pub mod qwen3_attention;
23pub mod qwen3_ssm;
24pub mod vision_encoder;
25pub mod vision_tower;
26pub mod w4a16_gemv_tiers;
27
28/// 2026-09-25: Smallest K at which the tile-GEMM sites (`dense_ffn.rs`,
29/// `qwen3_attention/trait_impl/multi_seq/qkv.rs`, `qwen3_ssm/kernel_select.rs`)
30/// take the deep-K `w4a16_gemm_t_k64` handle.
31///
32/// Resolved once per process: `METRALE_W4A16_K64_MIN_K` (a parsed `u32`) if
33/// set, else 8192 when `METRALE_NO_W4A16_K64=1`, else 6144.
34pub(crate) fn w4a16_k64_min_k() -> u32 {
35    static MIN_K: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
36    *MIN_K.get_or_init(|| {
37        if let Some(n) = std::env::var("METRALE_W4A16_K64_MIN_K")
38            .ok()
39            .and_then(|v| v.parse::<u32>().ok())
40        {
41            return n;
42        }
43        if std::env::var("METRALE_NO_W4A16_K64").ok().as_deref() == Some("1") {
44            8192
45        } else {
46            6144
47        }
48    })
49}
50
51pub use dense_ffn::{DenseFfnLayer, DenseFfnWeights, FfnActivation};
52
53pub use glm_vit::{GlmVit, GlmVitBlock, GlmVitMerger};
54
55pub use moe::MoeLayer;
56// 2026-09-28: The MoE experts' decode activation cell (`moe/fp8_grouped_tc_w8a8.rs`).
57pub use moe::{moe_expert_fp8_act, set_moe_expert_fp8_act};
58pub use mtp_head::{
59    MtpHead, MtpQuantization, mtp_drafter_prefill_enabled, set_mtp_experts_nvfp4_from_cli,
60};
61
62pub use qwen3_attention::Qwen3AttentionLayer;
63pub use qwen3_ssm::Qwen3SsmLayer;
64pub use vision_encoder::{MergerLayer, ViTBlock, VisionEncoder};
65pub use vision_tower::VisionTower;
66
67use crate::layer::ForwardContext;
68use anyhow::Result;
69use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
70
71/// 2026-09-26: The `w4a16_gemm_t_m128_v2` handle for the layer builders in
72/// `dense_ffn_init.rs`, `qwen3_attention/init_prefill_kernels.rs` and
73/// `qwen3_ssm/init.rs`.
74///
75/// `KernelHandle(0)` unless `METRALE_W4A16_VARIANT` is `v2` or `v3`, with no
76/// lookup. For `v2`/`v3` it panics when the target lacks the kernel.
77#[track_caller]
78pub fn w4a16_v2_kernel(gpu: &dyn GpuBackend) -> KernelHandle {
79    let variant = std::env::var("METRALE_W4A16_VARIANT").ok();
80    // 2026-09-25: Off unless requested. Measured 2026-07-30 with
81    // `w4a16_bf16_v2_bench` on the 27B FFN shapes: v2 ran at 0.78-0.82x of v1.
82    if !matches!(variant.as_deref(), Some("v2") | Some("v3")) {
83        if variant.as_deref() == Some("v1") {
84            tracing::debug!("METRALE_W4A16_VARIANT=v1: w4a16 m128 v2 suppressed (explicit)");
85        }
86        return KernelHandle(0);
87    }
88    let h = try_kernel(gpu, "w4a16_v2", "w4a16_gemm_t_m128_v2");
89    if h.0 == 0 {
90        panic!(
91            "METRALE_W4A16_VARIANT={} requested but w4a16_v2::w4a16_gemm_t_m128_v2 is not in this \
92             target's kernel set — refusing to start with a silently-degraded config",
93            variant.unwrap()
94        );
95    }
96    tracing::debug!(
97        handle = h.0,
98        "w4a16_gemm_t_m128_v2 resolution (explicit opt-in)"
99    );
100    h
101}
102
103/// 2026-09-25: The `w4a16_gemm_t_m128_v3` handle. `KernelHandle(0)`, with no
104/// lookup and so no boot-audit row, unless `METRALE_W4A16_VARIANT=v3`; then it
105/// panics when the target lacks the kernel. Its user
106/// (`qwen3_attention/prefill_weights.rs`) launches it only for `v == 3` and a
107/// nonzero handle.
108#[track_caller]
109pub fn w4a16_v3_kernel(gpu: &dyn GpuBackend) -> KernelHandle {
110    if std::env::var("METRALE_W4A16_VARIANT").as_deref() != Ok("v3") {
111        return KernelHandle(0);
112    }
113    let h = try_kernel(gpu, "w4a16_v3", "w4a16_gemm_t_m128_v3");
114    if h.0 == 0 {
115        panic!(
116            "METRALE_W4A16_VARIANT=v3 requested but w4a16_v3::w4a16_gemm_t_m128_v3 is not in this \
117             target's kernel set — refusing to start with a silently-degraded config"
118        );
119    }
120    h
121}
122
123/// 2026-09-25: The N128/M64 tile GEMM. Prefers `w4a16_gemm_t_p3`; takes
124/// `w4a16_gemm_t` when `METRALE_NO_TGEMM_PIPELINE3` is set to any value
125/// (`0` included) or the target lacks `_p3`.
126#[track_caller]
127pub fn tgemm_kernel(gpu: &dyn GpuBackend) -> KernelHandle {
128    if std::env::var("METRALE_NO_TGEMM_PIPELINE3").is_err() {
129        let h = try_kernel(gpu, "w4a16", "w4a16_gemm_t_p3");
130        if h.0 != 0 {
131            return h;
132        }
133    }
134    try_kernel(gpu, "w4a16", "w4a16_gemm_t")
135}
136
137/// 2026-09-25: False for `kimi_k3` and `kimi_linear`. For those the model
138/// builder (`model-engine` `model/impl_a1.rs`) skips the `tgemm_kernel` lookup
139/// and keeps a zero handle.
140pub fn tgemm_probe_ok(model_type: &str) -> bool {
141    !matches!(model_type, "kimi_k3" | "kimi_linear")
142}
143
144/// 2026-09-25: The k64 deep-K tile GEMM. Prefers `w4a16_gemm_t_k64_p3`; takes
145/// `w4a16_gemm_t_k64` when `METRALE_NO_K64_PIPELINE3` is set to any value
146/// (`0` included) or the target lacks `_p3`, and errors when that is missing too.
147#[track_caller]
148pub fn k64_kernel(gpu: &dyn GpuBackend) -> Result<KernelHandle> {
149    let want_p3 = std::env::var("METRALE_NO_K64_PIPELINE3").is_err();
150    if want_p3 {
151        let h = try_kernel(gpu, "w4a16", "w4a16_gemm_t_k64_p3");
152        if h.0 != 0 {
153            return Ok(h);
154        }
155    }
156    gpu.kernel("w4a16", "w4a16_gemm_t_k64")
157}
158
159/// 2026-09-25: The N_TILE=64 deep-K twin `w4a16_gemm_t_k64_n64_p3`.
160/// `KernelHandle(0)` when the target lacks it, or, with no lookup, when
161/// `METRALE_NO_K64_N64` is set to any value (`0` included).
162#[track_caller]
163pub fn k64_n64_kernel(gpu: &dyn GpuBackend) -> KernelHandle {
164    if std::env::var("METRALE_NO_K64_N64").is_ok() {
165        return KernelHandle(0);
166    }
167    try_kernel(gpu, "w4a16", "w4a16_gemm_t_k64_n64_p3")
168}
169
170/// 2026-09-25: Largest wide-tile grid at which the N_TILE=64 twin serves instead.
171///
172/// The wide tile launches `ceil(n/128) * ceil(m/64)` CTAs of 128 threads
173/// (`ops::w4a16_gemm_n128_ldb`); the twin launches `ceil(n/64) * ceil(m/64)`
174/// (`ops::w4a16_gemm`), about twice the grid.
175const K64_N64_MAX_WIDE_CTAS: u32 = 64;
176
177/// 2026-09-25: Whether the N_TILE=64 twin serves an `m × n` launch
178/// (`K64_N64_MAX_WIDE_CTAS`).
179pub fn k64_n64_wins(m: u32, n: u32) -> bool {
180    n.div_ceil(128) * m.div_ceil(64) <= K64_N64_MAX_WIDE_CTAS
181}
182
183mod moe_grouped_decode;
184pub use moe_grouped_decode::*;
185
186mod expert_quantization;
187mod moe_expert_tables;
188pub use expert_quantization::{
189    ExpertQuantization, expert_quantization, set_expert_quantization_from_cli,
190};
191pub use moe_expert_tables::{MoeExpertTables, moe_expert_tables, set_moe_expert_tables_from_plan};
192
193mod activation_quantization;
194pub use activation_quantization::{
195    act_route, activation_quantization, any_fixed, family_fixed, fixed_act,
196    set_activation_quantization_from_cli,
197};
198mod weight_quantization;
199pub use weight_quantization::{kernel_caps, set_weight_quantization_from_cli, weight_quantization};
200
201mod row_tiers;
202pub use row_tiers::{
203    RowTiers, prefill_row_invariant, publish_row_tiers, resolve_row_tiers, row_invariant,
204    row_tiers, row_tiers_from,
205};
206
207mod kernel_probe;
208pub use kernel_probe::{try_kernel, try_target_kernel};
209
210mod w8a8_layer;
211pub use w8a8_layer::{W8a8Ctx, W8a8Ffn, W8a8Mixer};
212
213/// 2026-09-25: A layer's FFN: MoE, dense, or none.
214#[allow(clippy::large_enum_variant)]
215pub enum FfnComponent {
216    Moe(MoeLayer),
217    Dense(DenseFfnLayer),
218    /// 2026-09-25: No FFN: `forward` returns its input and the other passes launch nothing.
219    None,
220}
221
222impl FfnComponent {
223    /// 2026-09-28: Add the FFN's weights to a circuit binding; a MoE or absent FFN is reported
224    /// as unmodelled (the executor binds dense FFNs so far).
225    pub(crate) fn circuit_bind(
226        &self,
227        levers: &ops::ModelLevers,
228        weights: &mut std::collections::BTreeMap<
229            crate::circuit_exec::WeightSlot,
230            crate::circuit_exec::BoundWeight,
231        >,
232        unmodelled: &mut Vec<String>,
233    ) {
234        match self {
235            Self::Dense(d) => d.circuit_bind(levers, weights, unmodelled),
236            Self::Moe(_) => unmodelled.push("a MoE FFN (not bound yet)".to_string()),
237            Self::None => unmodelled.push("no FFN".to_string()),
238        }
239    }
240
241    /// 2026-09-28: Build what a dense FFN's binding hands out (`DenseFfnLayer::circuit_prepare`);
242    /// nothing for a MoE or an absent FFN, which bind as unmodelled.
243    pub(crate) fn circuit_prepare(
244        &self,
245        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
246        config: &metrale_config::ModelConfig,
247        levers: &ops::ModelLevers,
248        stream: u64,
249    ) -> Result<()> {
250        match self {
251            Self::Dense(d) => d.circuit_prepare(gpu, config, levers, stream),
252            Self::Moe(_) | Self::None => Ok(()),
253        }
254    }
255
256    pub fn is_none(&self) -> bool {
257        matches!(self, Self::None)
258    }
259
260    pub fn is_dense(&self) -> bool {
261        matches!(self, Self::Dense(_))
262    }
263
264    /// 2026-09-25: Whether this is a MoE whose `MoeLayer::grouped_decode_ok` holds,
265    /// so multi-sequence decode may route it through `forward_prefill`
266    /// (`qwen3_attention/trait_impl/multi_seq/ffn.rs`). False for dense and none.
267    pub fn moe_grouped_decode_ok(&self) -> bool {
268        match self {
269            Self::Moe(m) => m.grouped_decode_ok(),
270            _ => false,
271        }
272    }
273
274    /// 2026-09-25: `MoeLayer::fp32_routing_active` for a MoE; false otherwise.
275    pub fn fp32_routing_active(&self, levers: &ops::ModelLevers) -> bool {
276        match self {
277            Self::Moe(m) => m.fp32_routing_active(levers),
278            _ => false,
279        }
280    }
281
282    pub fn forward(
283        &self,
284        input: DevicePtr,
285        ctx: &ForwardContext,
286        stream: u64,
287    ) -> Result<DevicePtr> {
288        match self {
289            Self::Moe(m) => m.forward(input, ctx, stream),
290            Self::Dense(d) => d.forward(input, ctx, stream),
291            Self::None => Ok(input),
292        }
293    }
294
295    pub fn forward_k2(&self, input: DevicePtr, ctx: &ForwardContext, stream: u64) -> Result<()> {
296        match self {
297            Self::Moe(m) => m.forward_k2(input, ctx, stream),
298            Self::Dense(d) => d.forward_k2(input, ctx, stream),
299            Self::None => Ok(()),
300        }
301    }
302
303    pub fn forward_k3(&self, input: DevicePtr, ctx: &ForwardContext, stream: u64) -> Result<()> {
304        match self {
305            Self::Moe(m) => m.forward_k3(input, ctx, stream),
306            Self::Dense(d) => d.forward_k3(input, ctx, stream),
307            Self::None => Ok(()),
308        }
309    }
310
311    /// 2026-09-30: Whether this is a dense FFN that runs `m` rows under a fixed
312    /// `--activation-quantization` (`DenseFfnLayer::fixed_ok`). False for MoE and none.
313    pub fn dense_fixed_ok(&self, m: usize, ctx: &ForwardContext) -> bool {
314        matches!(self, Self::Dense(d) if d.fixed_ok(m, ctx))
315    }
316
317    /// 2026-09-30: `DenseFfnLayer::forward_fixed` over `m` rows into `moe_output`; call only
318    /// when `dense_fixed_ok(m)` holds.
319    pub fn forward_dense_fixed(
320        &self,
321        input: DevicePtr,
322        m: usize,
323        ctx: &ForwardContext,
324        stream: u64,
325    ) -> Result<()> {
326        match self {
327            Self::Dense(d) => d.forward_fixed(input, m, ctx, stream),
328            _ => anyhow::bail!("forward_dense_fixed on a non-dense FFN"),
329        }
330    }
331
332    /// 2026-09-25: Whether this is a dense FFN whose
333    /// `DenseFfnLayer::can_forward_km(m)` holds: a batched-GEMV tier serves `m`
334    /// rows and NVFP4 or FP8 weights are loaded. False for MoE and none.
335    /// Multi-sequence attention decode (`qwen3_attention/trait_impl/multi_seq/ffn.rs`)
336    /// asks before computing the pre-FFN norm.
337    pub fn can_forward_km(&self, m: u32) -> bool {
338        matches!(self, Self::Dense(d) if d.can_forward_km(m))
339    }
340
341    /// 2026-09-28: Row edge of a dense FFN's narrow decode arms (`DenseFfnLayer::narrow_rows`:
342    /// its W4A4 edge under the weight-quantization tier, capped at 32, else the W4A16 edge).
343    /// MoE and none answer the W4A16 edge; `can_forward_km` is false for them anyway.
344    pub fn narrow_rows(&self) -> u32 {
345        match self {
346            Self::Dense(d) => d.narrow_rows(),
347            _ => ops::gemv_tc::narrow_gemv_max_rows().min(ops::w4a4_proj::W4A4_MAX_M),
348        }
349    }
350
351    /// 2026-09-25: `DenseFfnLayer::forward_km` over `m` rows. Returns `Ok(false)`
352    /// without launching when `can_forward_km(m)` is false.
353    pub fn try_forward_km(
354        &self,
355        input: DevicePtr,
356        m: u32,
357        ctx: &ForwardContext,
358        stream: u64,
359    ) -> Result<bool> {
360        match self {
361            Self::Dense(d) if d.can_forward_km(m) => {
362                d.forward_km(input, m, ctx, stream)?;
363                Ok(true)
364            }
365            _ => Ok(false),
366        }
367    }
368
369    pub fn forward_prefill(
370        &self,
371        input: DevicePtr,
372        num_tokens: usize,
373        ctx: &ForwardContext,
374        stream: u64,
375    ) -> Result<()> {
376        match self {
377            Self::Moe(m) => m.forward_prefill(input, num_tokens, ctx, stream),
378            Self::Dense(d) => d.forward_prefill(input, num_tokens, ctx, stream),
379            Self::None => {
380                let _ = (input, num_tokens);
381                Ok(())
382            }
383        }
384    }
385
386    pub fn forward_batched(
387        &self,
388        input: DevicePtr,
389        num_tokens: usize,
390        ctx: &ForwardContext,
391        stream: u64,
392    ) -> Result<()> {
393        match self {
394            Self::Moe(m) => m.forward_batched(input, num_tokens, ctx, stream),
395            Self::Dense(d) => d.forward_batched(input, num_tokens, ctx, stream),
396            Self::None => {
397                let _ = (input, num_tokens);
398                Ok(())
399            }
400        }
401    }
402
403    pub fn forward_token_major_decode(
404        &self,
405        input: DevicePtr,
406        num_tokens: usize,
407        ctx: &ForwardContext,
408        stream: u64,
409    ) -> Result<()> {
410        match self {
411            Self::Moe(m) => m.forward_token_major_decode(input, num_tokens, ctx, stream),
412            Self::Dense(d) => d.forward_batched(input, num_tokens, ctx, stream),
413            Self::None => {
414                let _ = (input, num_tokens);
415                Ok(())
416            }
417        }
418    }
419
420    pub fn forward_atomic_c4_decode(
421        &self,
422        input: DevicePtr,
423        num_tokens: usize,
424        ctx: &ForwardContext,
425        stream: u64,
426    ) -> Result<()> {
427        match self {
428            Self::Moe(m) => m.forward_atomic_c4_decode(input, num_tokens, ctx, stream),
429            Self::Dense(d) => d.forward_batched(input, num_tokens, ctx, stream),
430            Self::None => {
431                let _ = (input, num_tokens);
432                Ok(())
433            }
434        }
435    }
436
437    /// 2026-09-25: Whether the target model's cross-row grouped FP8 MoE decode
438    /// serves `m` rows: `ModelLevers::moe_fp8_grouped_decode_target`
439    /// (`METRALE_FP8_MOE_GROUPED_DECODE`) is on and this is a MoE whose
440    /// `MoeLayer::fp8_grouped_decode_ok(m)` holds. The MTP drafter calls
441    /// `MoeLayer::fp8_grouped_decode_ok` directly, without the lever.
442    pub fn fp8_grouped_decode_ok(&self, m: usize, ctx: &ForwardContext) -> bool {
443        ctx.levers.moe_fp8_grouped_decode_target
444            && matches!(self, Self::Moe(moe) if moe.fp8_grouped_decode_ok(m, ctx))
445    }
446
447    /// 2026-09-26: Whether this is a MoE whose grouped FP8 decode serves `m`
448    /// rows with `routing` (`MoeLayer::fp8_grouped_routing_ok`). Unlike
449    /// [`Self::fp8_grouped_decode_ok`] it does not read the
450    /// `moe_fp8_grouped_decode_target` lever: the exact routings reproduce the
451    /// path they replace byte for byte, so they need no opt-in.
452    pub fn fp8_grouped_routing_ok(
453        &self,
454        m: usize,
455        routing: moe::GroupedRouting,
456        ctx: &ForwardContext,
457    ) -> bool {
458        matches!(self, Self::Moe(moe) if moe.fp8_grouped_routing_ok(m, routing, ctx))
459    }
460
461    /// 2026-09-26: [`Self::forward_fp8_grouped_decode`] with `routing`. Errors
462    /// for dense and none; callers gate on `fp8_grouped_routing_ok` first.
463    pub fn forward_fp8_grouped_decode_routed(
464        &self,
465        input: DevicePtr,
466        m: usize,
467        routing: moe::GroupedRouting,
468        ctx: &ForwardContext,
469        stream: u64,
470    ) -> Result<()> {
471        match self {
472            Self::Moe(moe) => moe.forward_fp8_grouped_decode_routed(input, m, routing, ctx, stream),
473            _ => anyhow::bail!("forward_fp8_grouped_decode_routed is MoE-only (m={m})"),
474        }
475    }
476
477    /// 2026-09-25: Cross-row grouped FP8 MoE decode over `[m, H]` rows into
478    /// `moe_output()`. Errors for dense and none; callers gate on
479    /// `fp8_grouped_decode_ok` first.
480    pub fn forward_fp8_grouped_decode(
481        &self,
482        input: DevicePtr,
483        m: usize,
484        ctx: &ForwardContext,
485        stream: u64,
486    ) -> Result<()> {
487        match self {
488            Self::Moe(moe) => moe.forward_fp8_grouped_decode(input, m, ctx, stream),
489            _ => anyhow::bail!("forward_fp8_grouped_decode is MoE-only (m={m})"),
490        }
491    }
492}
493
494pub(crate) use gemv_tier::batch8_kernel;
metrale_model_layers/layers/mod.rs

metrale_model_layers/layers/
mod.rs
