spark_model/layers/moe_v41/
init.rs

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

spark_model/layers/moe_v41/
init.rs
