metrale_model_arch/moe_v41/
init.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: `MoeV41` construction and teardown: the kernel handles and the
5//! `max_tokens` workspace.
6//!
7//! Owner: model-arch (DeepSeek-V4.1 MoE).
8//! Invariants:
9//! - `new` refuses a config whose `dim` or `inter` is not a multiple of 256.
10//! - The per-expert buffers (`gate_out`, `up_out`, `h`, `h_q8`, `down_out`)
11//!   hold at least `topk` rows whatever `max_tokens` is.
12
13use anyhow::{Result, ensure};
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
15
16use super::{GEMM_MODULE, Kernels, MODULE, MoeV41, MoeV41Cfg, MoeV41Timing, SLOT_TABLE_LAYERS};
17use metrale_model_layers::layers::ops::{
18    KQUANT_MODULE, kquant_mmq_act_bytes, kquant_q8_1_rows_bytes,
19};
20
21impl MoeV41 {
22    pub fn new(gpu: &dyn GpuBackend, cfg: MoeV41Cfg) -> Result<Self> {
23        ensure!(
24            cfg.dim.is_multiple_of(256) && cfg.inter.is_multiple_of(256),
25            "K-quant experts need dim and inter to be multiples of 256 (got {} / {})",
26            cfg.dim,
27            cfg.inter
28        );
29        let m = cfg.max_tokens;
30        // 2026-09-25: The single-token arm (`routed_m1`) runs every selected
31        // expert as a row of the per-expert buffers, so they hold at least
32        // `topk` rows however small `max_tokens` is; the token-indexed buffers
33        // keep `m`.
34        let me = m.max(cfg.topk);
35        let alloc = |bytes: usize| gpu.alloc(bytes.max(16));
36        // 2026-09-25: Warps a block of the single-token expert GEMV batch,
37        // `METRALE_DS41_EXPERT_WARPS` in {2, 4, 8}: 8 (the `_w8` entries) when
38        // unset or any other value, 2 selects `_w2`, 4 the unsuffixed `_w`.
39        // Measured 2026-09-19 (hot decode): 8 was the fastest of the three,
40        // 18.9 tok/s.
41        let experts_warps: u32 = std::env::var("METRALE_DS41_EXPERT_WARPS")
42            .ok()
43            .and_then(|v| v.parse().ok())
44            .filter(|w| matches!(w, 2 | 4 | 8))
45            .unwrap_or(8);
46        let sfx = match experts_warps {
47            2 => "2",
48            8 => "8",
49            _ => "",
50        };
51        Ok(MoeV41 {
52            last: std::cell::Cell::new(MoeV41Timing::default()),
53            timing_sync: std::env::var("METRALE_DS41_DIAG").is_ok_and(|v| v == "1"),
54            k: Kernels {
55                gemm: gpu.kernel(GEMM_MODULE, "dense_gemm_bf16")?,
56                gemv: gpu.kernel("gemv", "dense_gemv_bf16")?,
57                gemm_f32out: gpu.kernel(GEMM_MODULE, "dense_gemm_bf16_f32out")?,
58                router_gemv: gpu.kernel(MODULE, "moe_v41_router_gemv_f32out")?,
59                router_gemv_staged: gpu.kernel(MODULE, "moe_v41_router_gemv_f32out_products")?,
60                q8_rows: gpu.kernel(KQUANT_MODULE, "kquant_q8_1_rows_bf16")?,
61                mmvq_q2k: gpu.kernel(KQUANT_MODULE, "kquant_mmvq_q2_k_w")?,
62                mmvq_q3k: gpu.kernel(KQUANT_MODULE, "kquant_mmvq_q3_k_w")?,
63                mmvq_q2k_experts: gpu
64                    .kernel(KQUANT_MODULE, &format!("kquant_mmvq_q2_k_experts_w{sfx}"))?,
65                mmvq_q3k_experts: gpu
66                    .kernel(KQUANT_MODULE, &format!("kquant_mmvq_q3_k_experts_w{sfx}"))?,
67                experts_warps,
68                swiglu: gpu.kernel(MODULE, "moe_v41_swiglu")?,
69                swiglu_q8: gpu.kernel(KQUANT_MODULE, "kquant_swiglu_q8_1_rows_bf16")?,
70                accumulate: gpu.kernel(MODULE, "moe_v41_accumulate")?,
71                finish: gpu.kernel(MODULE, "moe_v41_finish")?,
72                gather: gpu.kernel(MODULE, "moe_v41_gather_rows")?,
73                scatter_add: gpu.kernel(MODULE, "moe_v41_scatter_add")?,
74                sum_rows: gpu.kernel(MODULE, "moe_v41_sum_rows")?,
75                quant_d2s6: gpu.kernel(KQUANT_MODULE, "metrale_q8_1_quantize_d2s6_bf16")?,
76                quant_d4: gpu.kernel(KQUANT_MODULE, "metrale_q8_1_quantize_d4_bf16")?,
77                mmq_q2k_nc: gpu.kernel(KQUANT_MODULE, "metrale_q2_k_mmq128_nc")?,
78                mmq_q2k_wc: gpu.kernel(KQUANT_MODULE, "metrale_q2_k_mmq128_wc")?,
79                mmq_q3k_nc: gpu.kernel(KQUANT_MODULE, "metrale_q3_k_mmq128_nc")?,
80                mmq_q3k_wc: gpu.kernel(KQUANT_MODULE, "metrale_q3_k_mmq128_wc")?,
81                route_select: gpu.kernel(MODULE, "moe_v41_route_select")?,
82                slot_table_set: gpu.kernel(MODULE, "moe_v41_slot_table_set")?,
83            },
84            logits: alloc(m * cfg.n_routed * 4)?,
85            pred_logits: alloc(cfg.n_routed * 4)?,
86            a_rows: alloc(m * cfg.dim * 2)?,
87            a_q8: alloc(
88                kquant_mmq_act_bytes(m as u32, cfg.dim as u32)
89                    .max(kquant_q8_1_rows_bytes(m as u32, cfg.dim as u32)),
90            )?,
91            // 2026-09-25: `[2 * me, inter]`: the single-token path runs gate
92            // and up as one 2 * ne expert batch and reads up at
93            // gate_out + ne * inter.
94            gate_out: alloc(2 * me * cfg.inter * 2)?,
95            up_out: alloc(me * cfg.inter * 2)?,
96            h: alloc(me * cfg.inter * 2)?,
97            h_q8: alloc(
98                kquant_mmq_act_bytes(me as u32, cfg.inter as u32)
99                    .max(kquant_q8_1_rows_bytes(me as u32, cfg.inter as u32)),
100            )?,
101            down_out: alloc(me * cfg.dim * 2)?,
102            rows_dev: alloc(m * cfg.topk * 4)?,
103            weight_dev: alloc(m * cfg.topk * 4)?,
104            route_hdr: alloc(SLOT_TABLE_LAYERS * (1 + 3 * cfg.topk) * 4)?,
105            slot_table: {
106                let t = alloc(SLOT_TABLE_LAYERS * cfg.n_routed * 4)?;
107                gpu.memset(t, 0xFF, SLOT_TABLE_LAYERS * cfg.n_routed * 4)?;
108                t
109            },
110            ptrs_dev: alloc(3 * cfg.topk * 8)?,
111            sg: alloc(m * cfg.inter * 2)?,
112            su: alloc(m * cfg.inter * 2)?,
113            sh: alloc(m * cfg.inter * 2)?,
114            sd: alloc(m * cfg.dim * 2)?,
115            acc: alloc(m * cfg.dim * 4)?,
116            out: alloc(m * cfg.dim * 2)?,
117            sh_q8: alloc(kquant_q8_1_rows_bytes(1, cfg.inter as u32))?,
118            side: gpu.create_stream()?,
119            ev_in: gpu.create_event()?,
120            ev_out: gpu.create_event()?,
121            cfg,
122        })
123    }
124
125    pub fn out_ptr(&self) -> DevicePtr {
126        self.out
127    }
128
129    pub fn free(self, gpu: &dyn GpuBackend) -> Result<()> {
130        for p in [
131            self.logits,
132            self.pred_logits,
133            self.a_rows,
134            self.a_q8,
135            self.rows_dev,
136            self.route_hdr,
137            self.slot_table,
138            self.gate_out,
139            self.up_out,
140            self.h,
141            self.h_q8,
142            self.down_out,
143            self.weight_dev,
144            self.sg,
145            self.su,
146            self.sh,
147            self.sd,
148            self.acc,
149            self.out,
150            self.sh_q8,
151        ] {
152            gpu.free(p)?;
153        }
154        gpu.destroy_event(self.ev_in)?;
155        gpu.destroy_event(self.ev_out)?;
156        Ok(())
157    }
158}
metrale_model_arch/moe_v41/init.rs

metrale_model_arch/moe_v41/
init.rs
