metrale_model_arch/glm5next_layer/steps/
drafter.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The MTP drafter's entry points into its `Glm5NextLayer`: one decode token, and
4//! one context row's DSA cache writes.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11impl Glm5NextLayer {
12    /// 2026-09-25: One token through this block for the MTP drafter; errors on a block with a
13    /// hyper-connection. `hidden` is both input and residual and is updated in place. No disk
14    /// tiers: the disk block lists passed down are empty.
15    #[allow(clippy::too_many_arguments)]
16    pub fn decode_one_for_drafter(
17        &self,
18        hidden: DevicePtr,
19        state: &mut dyn LayerState,
20        kv_cache: &mut PagedKvCache,
21        seq_len: usize,
22        block_table: &mut Vec<u32>,
23        ctx: &ForwardContext,
24        stream: u64,
25    ) -> Result<()> {
26        if self.mhc.is_some() {
27            bail!(
28                "GLM layer {}: decode_one_for_drafter is the MTP block's path; this layer has a \
29                 hyper-connection",
30                self.layer_idx
31            );
32        }
33        let (mut disk_a, mut disk_b) = (Vec::new(), Vec::new());
34        self.forward_one_plain(
35            hidden,
36            hidden,
37            state,
38            kv_cache,
39            seq_len,
40            block_table,
41            &mut disk_a,
42            &mut disk_b,
43            ctx,
44            stream,
45        )
46    }
47
48    /// 2026-09-25: One drafter context row: `input_norm`, then `Glm5NextDsaLayer::write_kv_row`,
49    /// which writes the row's KV latent and indexer entry without computing the block output.
50    /// Errors on a block whose mixer is not DSA.
51    #[allow(clippy::too_many_arguments)]
52    pub fn drafter_write_kv_row(
53        &self,
54        x: DevicePtr,
55        state: &mut dyn LayerState,
56        kv_cache: &mut PagedKvCache,
57        seq_len: usize,
58        block_table: &mut Vec<u32>,
59        ctx: &ForwardContext,
60        stream: u64,
61    ) -> Result<()> {
62        let layer = match &self.mixer {
63            Glm5NextMixer::Dsa(l) => l,
64            _ => bail!(
65                "GLM layer {}: drafter_write_kv_row is the MTP block's path; this layer is not \
66                 a DSA layer",
67                self.layer_idx
68            ),
69        };
70        let normed = ctx.buffers.norm_output();
71        self.norm(ctx.gpu, x, self.input_norm, normed, 1, stream)?;
72        layer.write_kv_row(normed, state, kv_cache, seq_len, block_table, ctx, stream)
73    }
74}
metrale_model_arch/glm5next_layer/steps/drafter.rs

metrale_model_arch/glm5next_layer/steps/
drafter.rs
