metrale_model_layers/layers/
mtp_meta.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The one definition of an MTP proposer's attention-metadata
4//! slab, packed by `MtpHead::forward_one`, `forward_batch_position` and
5//! `DeepseekV4MtpHead::propose`.
6//!
7//! Owner: model-layers (MTP head).
8//! Invariants:
9//! - `pack_mtp_attn_meta` returns a slab only when it fits the caller's
10//!   `region_bytes`; otherwise it returns an error and builds nothing.
11
12use anyhow::{Result, ensure};
13
14/// 2026-09-25: Bytes of header before the block table. The three scalars sit
15/// at offsets 0, 8 and 16 and the rest is zero.
16pub(crate) const MTP_META_HEADER_BYTES: usize = 256;
17
18/// 2026-09-25: `scratch` byte offset of the single-sequence MTP metadata
19/// slab, used by `MtpHead::forward_one` and `DeepseekV4MtpHead::propose`.
20/// The target's batched decode metadata starts at `scratch + 32768`. The
21/// batched propose writes its own `propose_meta` allocation instead.
22pub const MTP_META_OFFSET: usize = 49152;
23
24fn mtp_meta_len(block_entries: usize) -> Result<usize> {
25    block_entries
26        .checked_mul(4)
27        .and_then(|bt_bytes| MTP_META_HEADER_BYTES.checked_add(bt_bytes))
28        .ok_or_else(|| {
29            anyhow::anyhow!(
30                "MTP attention metadata exceeds meta stride: byte size overflows for a {block_entries}-entry block table"
31            )
32        })
33}
34
35/// 2026-09-25: Pack one MTP attention-metadata slab, or return an error when
36/// it would not fit `region_bytes`.
37///
38/// Layout: `position` u32 at 0, `global_slot` i64 at 8, `seq_len` i32 at 16,
39/// then the block table from `MTP_META_HEADER_BYTES`, four little-endian
40/// bytes per entry. `region_bytes` is what the caller owns at the
41/// destination: the per-sequence stride for the batched propose, or
42/// `scratch_bytes - MTP_META_OFFSET` for the single-sequence callers.
43/// 2026-09-29: The attention metadata [`pack_mtp_attn_meta`] lays out at `base` (position @0,
44/// slot @8, sequence length @16, block table @`MTP_META_HEADER_BYTES`), for one row.
45pub fn mtp_attn_meta_dev(
46    base: metrale_gpu_runtime::gpu::DevicePtr,
47    max_blocks_per_seq: u32,
48) -> crate::layer::AttnMetadataDev {
49    crate::layer::AttnMetadataDev {
50        positions: base,
51        positions_h: base,
52        positions_w: base,
53        slot: base.offset(8),
54        seq_len: base.offset(16),
55        block_table: base.offset(MTP_META_HEADER_BYTES),
56        max_blocks_per_seq,
57        num_seqs: 1,
58        seq_slot: metrale_gpu_runtime::gpu::DevicePtr::NULL,
59        moe_row_adapter: metrale_gpu_runtime::gpu::DevicePtr::NULL,
60    }
61}
62
63pub fn pack_mtp_attn_meta(
64    position: u32,
65    global_slot: i64,
66    seq_len: i32,
67    block_table: &[u32],
68    region_bytes: usize,
69) -> Result<Vec<u8>> {
70    let need = mtp_meta_len(block_table.len())?;
71    // 2026-09-25: `scheduler/mtp_bootstrap_step.rs` matches the phrase
72    // "exceeds meta stride" to log this error at debug instead of error.
73    // Rewording it breaks that match; the test below pins the phrase.
74    ensure!(
75        need <= region_bytes,
76        "MTP attention metadata exceeds meta stride: needs {need} B for a {}-entry \
77         block table, have {region_bytes} B",
78        block_table.len()
79    );
80
81    let mut buf = vec![0u8; need];
82    buf[0..4].copy_from_slice(&position.to_le_bytes());
83    buf[8..16].copy_from_slice(&global_slot.to_le_bytes());
84    buf[16..20].copy_from_slice(&seq_len.to_le_bytes());
85    for (i, &block) in block_table.iter().enumerate() {
86        let at = MTP_META_HEADER_BYTES + i * 4;
87        buf[at..at + 4].copy_from_slice(&block.to_le_bytes());
88    }
89    Ok(buf)
90}
91
92/// 2026-09-29: Byte offsets of the batched propose's attention metadata for `n` rows: `u32`
93/// positions, `i64` KV slots, `i32` sequence lengths, then `n` block tables of `max_blocks`
94/// entries each, row-major.
95#[derive(Clone, Copy, Debug, PartialEq, Eq)]
96pub(crate) struct MtpBatchMetaLayout {
97    pub positions: usize,
98    pub slots: usize,
99    pub seq_lens: usize,
100    pub block_tables: usize,
101    pub len: usize,
102}
103
104impl MtpBatchMetaLayout {
105    pub(crate) fn new(n: usize, max_blocks: usize) -> Self {
106        let slots = (4 * n).next_multiple_of(8);
107        let seq_lens = slots + 8 * n;
108        let block_tables = (seq_lens + 4 * n).next_multiple_of(16);
109        Self {
110            positions: 0,
111            slots,
112            seq_lens,
113            block_tables,
114            len: block_tables + 4 * n * max_blocks,
115        }
116    }
117}
118
119/// 2026-09-30: The attention metadata `pack_mtp_attn_meta_batch` lays out at `base` for `n`
120/// rows (`MtpBatchMetaLayout`); `max_blocks_per_seq` is the step's, set by its reader.
121pub fn mtp_attn_meta_batch_dev(
122    base: metrale_gpu_runtime::gpu::DevicePtr,
123    n: usize,
124) -> crate::layer::AttnMetadataDev {
125    let l = MtpBatchMetaLayout::new(n, 0);
126    crate::layer::AttnMetadataDev {
127        positions: base.offset(l.positions),
128        positions_h: base.offset(l.positions),
129        positions_w: base.offset(l.positions),
130        slot: base.offset(l.slots),
131        seq_len: base.offset(l.seq_lens),
132        block_table: base.offset(l.block_tables),
133        max_blocks_per_seq: 0,
134        num_seqs: n as u32,
135        seq_slot: metrale_gpu_runtime::gpu::DevicePtr::NULL,
136        moe_row_adapter: metrale_gpu_runtime::gpu::DevicePtr::NULL,
137    }
138}
139
140/// 2026-09-29: Pack the attention metadata of `n` propose rows in [`MtpBatchMetaLayout`], for
141/// one upload and one launch per attention kernel. Row `i` has `positions[i]`, `slots[i]`,
142/// `seq_lens[i]` and `tables[i]`, zero-padded to `max_blocks` entries. Refused, with the phrase
143/// `mtp_bootstrap_step.rs` matches, when it would not fit `region_bytes` or a table is longer
144/// than `max_blocks`.
145pub(crate) fn pack_mtp_attn_meta_batch(
146    positions: &[u32],
147    slots: &[i64],
148    seq_lens: &[i32],
149    tables: &[&[u32]],
150    max_blocks: usize,
151    region_bytes: usize,
152) -> Result<(Vec<u8>, MtpBatchMetaLayout)> {
153    let n = positions.len();
154    ensure!(
155        slots.len() == n && seq_lens.len() == n && tables.len() == n,
156        "MTP batch metadata: row counts differ"
157    );
158    let layout = MtpBatchMetaLayout::new(n, max_blocks);
159    ensure!(
160        layout.len <= region_bytes && tables.iter().all(|t| t.len() <= max_blocks),
161        "MTP attention metadata exceeds meta stride: needs {} B for {n} rows of {max_blocks} \
162         block entries, have {region_bytes} B",
163        layout.len
164    );
165    let mut buf = vec![0u8; layout.len];
166    for i in 0..n {
167        let p = layout.positions + 4 * i;
168        buf[p..p + 4].copy_from_slice(&positions[i].to_le_bytes());
169        let s = layout.slots + 8 * i;
170        buf[s..s + 8].copy_from_slice(&slots[i].to_le_bytes());
171        let l = layout.seq_lens + 4 * i;
172        buf[l..l + 4].copy_from_slice(&seq_lens[i].to_le_bytes());
173        for (j, &block) in tables[i].iter().enumerate() {
174            let at = layout.block_tables + 4 * (i * max_blocks + j);
175            buf[at..at + 4].copy_from_slice(&block.to_le_bytes());
176        }
177    }
178    Ok((buf, layout))
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    /// 2026-09-29: Two rows: every field lands at its layout offset, a short table is
186    /// zero-padded, and a table longer than `max_blocks` or a region too small is refused with
187    /// the matched phrase.
188    #[test]
189    fn packs_the_batched_layout() {
190        let (buf, l) =
191            pack_mtp_attn_meta_batch(&[7, 9], &[100, 2000], &[8, 10], &[&[3], &[4, 5]], 2, 4096)
192                .unwrap();
193        assert_eq!(
194            l,
195            MtpBatchMetaLayout {
196                positions: 0,
197                slots: 8,
198                seq_lens: 24,
199                block_tables: 32,
200                len: 48
201            }
202        );
203        let u32_at = |o: usize| u32::from_le_bytes(buf[o..o + 4].try_into().unwrap());
204        let i64_at = |o: usize| i64::from_le_bytes(buf[o..o + 8].try_into().unwrap());
205        assert_eq!((u32_at(0), u32_at(4)), (7, 9));
206        assert_eq!((i64_at(8), i64_at(16)), (100, 2000));
207        assert_eq!((u32_at(24), u32_at(28)), (8, 10));
208        assert_eq!(
209            [u32_at(32), u32_at(36), u32_at(40), u32_at(44)],
210            [3, 0, 4, 5]
211        );
212        let long = pack_mtp_attn_meta_batch(&[1], &[1], &[1], &[&[1, 2, 3]], 2, 4096);
213        assert!(
214            long.unwrap_err()
215                .to_string()
216                .contains("exceeds meta stride")
217        );
218        let small = pack_mtp_attn_meta_batch(&[1], &[1], &[1], &[&[1]], 1, 16);
219        assert!(
220            small
221                .unwrap_err()
222                .to_string()
223                .contains("exceeds meta stride")
224        );
225    }
226
227    #[test]
228    fn packs_the_documented_layout() {
229        let buf = pack_mtp_attn_meta(7, 0x1234_5678_9abc, 42, &[3, 9], 4096).unwrap();
230        assert_eq!(buf.len(), MTP_META_HEADER_BYTES + 8);
231        assert_eq!(u32::from_le_bytes(buf[0..4].try_into().unwrap()), 7);
232        assert_eq!(
233            i64::from_le_bytes(buf[8..16].try_into().unwrap()),
234            0x1234_5678_9abc
235        );
236        assert_eq!(i32::from_le_bytes(buf[16..20].try_into().unwrap()), 42);
237        assert_eq!(u32::from_le_bytes(buf[256..260].try_into().unwrap()), 3);
238        assert_eq!(u32::from_le_bytes(buf[260..264].try_into().unwrap()), 9);
239        // 2026-09-25: Every byte the layout does not define is zero.
240        assert!(buf[20..256].iter().all(|&b| b == 0));
241    }
242
243    /// 2026-09-30: The device view of the batched metadata reads every field where the packer
244    /// wrote it, at every width (the circuit's n-row draft reads it through this view).
245    #[test]
246    fn the_batched_device_view_is_the_packed_layout() {
247        let base = metrale_gpu_runtime::gpu::DevicePtr(0x1000);
248        for n in [1usize, 2, 3, 5, 16, 32, 128] {
249            let pos = vec![0u32; n];
250            let slots = vec![0i64; n];
251            let lens = vec![0i32; n];
252            let t = [0u32; 2];
253            let tables = vec![&t[..]; n];
254            let (_, l) =
255                pack_mtp_attn_meta_batch(&pos, &slots, &lens, &tables, 2, 1 << 20).unwrap();
256            let m = mtp_attn_meta_batch_dev(base, n);
257            assert_eq!(m.positions, base.offset(l.positions), "n={n}");
258            assert_eq!(m.slot, base.offset(l.slots), "n={n}");
259            assert_eq!(m.seq_len, base.offset(l.seq_lens), "n={n}");
260            assert_eq!(m.block_table, base.offset(l.block_tables), "n={n}");
261            assert_eq!(m.num_seqs as usize, n);
262        }
263    }
264
265    /// 2026-09-25: A block table that would overrun the region is refused.
266    #[test]
267    fn refuses_a_block_table_that_would_overrun_the_region() {
268        // 2026-09-25: A 2048-byte region with the 256-byte header holds 448
269        // entries.
270        let fits = vec![0u32; 448];
271        assert_eq!(
272            pack_mtp_attn_meta(0, 0, 1, &fits, 2048).unwrap().len(),
273            2048
274        );
275        let over = vec![0u32; 449];
276        let e = pack_mtp_attn_meta(0, 0, 1, &over, 2048).unwrap_err();
277        // 2026-09-25: Pins the phrase `scheduler/mtp_bootstrap_step.rs`
278        // matches.
279        assert!(
280            e.to_string().contains("exceeds meta stride"),
281            "the demotion phrase mtp_bootstrap_step matches must survive: {e}"
282        );
283    }
284
285    /// 2026-09-25: A region smaller than the header is refused.
286    #[test]
287    fn refuses_a_region_too_small_for_the_header() {
288        assert!(pack_mtp_attn_meta(0, 0, 1, &[], MTP_META_HEADER_BYTES - 1).is_err());
289        assert!(pack_mtp_attn_meta(0, 0, 1, &[], MTP_META_HEADER_BYTES).is_ok());
290    }
291
292    #[test]
293    fn refuses_an_unrepresentable_block_table_length() {
294        let err = mtp_meta_len(usize::MAX).unwrap_err();
295        assert!(err.to_string().contains("exceeds meta stride"), "{err}");
296    }
297}
metrale_model_layers/layers/mtp_meta.rs

metrale_model_layers/layers/
mtp_meta.rs
