metrale_model_layers/layer/transformer_layer/
weight_setup.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `LayerWeightSetup`, the `&mut self` hooks that change a layer's weights after
4//! construction: the downcast for weight overlays and the MoE expert transposes for
5//! prefill.
6//!
7//! Owner: model-layers.
8//! Invariants: none beyond the types.
9
10use anyhow::Result;
11use metrale_config::ModelConfig;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13
14/// 2026-09-26: A supertrait of `TransformerLayer`; see the module header.
15pub trait LayerWeightSetup {
16    /// 2026-09-25: `&mut dyn Any` downcast hook for weight overlays applied after
17    /// construction, such as a LoRA install. Default `None`.
18    fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
19        None
20    }
21
22    /// 2026-09-25: Build persistent transposed gate/up/down MoE expert weights for the
23    /// prefill GEMMs. The factory's MoE transpose pass (`factory/m2_setup.rs`) calls it,
24    /// outside the hybrid and unified layouts, when free memory covers the full transpose
25    /// plus a safety margin. Default: nothing.
26    fn transpose_moe_for_prefill(
27        &mut self,
28        _gpu: &dyn GpuBackend,
29        _config: &ModelConfig,
30    ) -> Result<()> {
31        Ok(())
32    }
33
34    /// 2026-09-25: Like `transpose_moe_for_prefill`, for gate and up only. Outside the
35    /// hybrid and unified layouts, the transpose pass calls it when the full transpose does
36    /// not fit but gate and up do. Default: nothing.
37    fn transpose_moe_gate_up_for_prefill(
38        &mut self,
39        _gpu: &dyn GpuBackend,
40        _config: &ModelConfig,
41    ) -> Result<()> {
42        Ok(())
43    }
44
45    /// 2026-09-25: Give this layer's MoE block the per-prefill `down_proj` transpose
46    /// scratch that every MoE layer shares, after `transpose_moe_gate_up_for_prefill`.
47    /// Default: nothing.
48    fn set_moe_down_transpose_scratch(
49        &mut self,
50        _scratch_packed: DevicePtr,
51        _scratch_scale: DevicePtr,
52        _packed_ptrs_t: DevicePtr,
53        _scale_ptrs_t: DevicePtr,
54    ) {
55    }
56
57    /// 2026-09-25: Unified layout (`METRALE_UNIFIED_MOE_LAYOUT=1`): build transposed
58    /// gate/up/down and free the untransposed copies. Decode reads the transposed copies
59    /// when `MoeLayer::use_t_layout_for_decode` holds. Default: nothing.
60    fn transpose_moe_for_prefill_unified(
61        &mut self,
62        _gpu: &dyn GpuBackend,
63        _config: &ModelConfig,
64    ) -> Result<()> {
65        Ok(())
66    }
67
68    /// 2026-09-25: Hybrid layout (`METRALE_HYBRID_MOE_LAYOUT=1`): build transposed
69    /// gate/up/down beside the originals and free nothing, so decode keeps the originals
70    /// and prefill uses the transposed copies. The transpose pass calls it only when twice
71    /// the full transpose fits. Default: nothing.
72    fn transpose_moe_for_prefill_hybrid(
73        &mut self,
74        _gpu: &dyn GpuBackend,
75        _config: &ModelConfig,
76    ) -> Result<()> {
77        Ok(())
78    }
79}
metrale_model_layers/layer/transformer_layer/weight_setup.rs

metrale_model_layers/layer/transformer_layer/
weight_setup.rs
