metrale_model_layers/circuit_exec/
fixed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The device addresses a program is compiled against: the model's buffers and
4//! metadata uploads ([`Fixed`]) and the MTP draft head's ([`DraftFixed`]). Split from
5//! `compile.rs`.
6//!
7//! Owner: model-layers circuit executor.
8//! Invariants:
9//! - Every address here is fixed after boot, so a captured program stays valid.
10
11use metrale_gpu_runtime::gpu::DevicePtr;
12
13use super::bindings::DraftRows;
14use crate::layer::AttnMetadataDev;
15
16/// 2026-09-28: Device addresses that never move after boot.
17#[derive(Clone)]
18pub struct Fixed {
19    /// 2026-09-28: The residual stream (`BufferArena::hidden_states`).
20    pub hidden: DevicePtr,
21    /// 2026-09-28: The copy of the stream the norms write (`BufferArena::residual`).
22    pub residual: DevicePtr,
23    /// 2026-09-28: The logits buffer the step returns.
24    pub logits: DevicePtr,
25    /// 2026-09-29: Where a verify's argmax writes each row's token (`i32` per row): the
26    /// scratch buffer's start, which the verify reads back.
27    pub tokens: DevicePtr,
28    /// 2026-09-28: The single-sequence step's attention metadata, at its fixed upload
29    /// address. `max_blocks_per_seq` is ignored: a step supplies it.
30    pub meta: AttnMetadataDev,
31    /// 2026-09-28: The multi-sequence step's attention metadata (one row per sequence), at its
32    /// fixed upload address; `max_blocks_per_seq` is ignored.
33    pub batch_meta: AttnMetadataDev,
34    /// 2026-09-29: The MTP verify step's attention metadata (one row per verified token), at
35    /// its fixed upload address; `max_blocks_per_seq` is ignored.
36    pub verify_meta: AttnMetadataDev,
37    /// 2026-09-29: The MTP draft head's buffers; `None` without a bound draft head.
38    pub draft: Option<DraftFixed>,
39    /// 2026-09-30: The batched MTP verify's attention metadata (one row per verified row,
40    /// `stage_verify_metadata` at scratch + 32768); `max_blocks_per_seq` is ignored.
41    pub verify_batch_meta: AttnMetadataDev,
42    /// 2026-09-30: The batched verify's per-GDN-layer WY pointer tables
43    /// (`upload_verify_wy_tables`, `VERIFY_WY_LAYER_STRIDE_BYTES` apart); NULL without MTP.
44    pub verify_wy_tables: DevicePtr,
45    /// 2026-09-30: Where the batched verify's argmax writes each row's token: the mapped host
46    /// blob's device alias, or scratch when there is none (`verify_rows_argmax`).
47    pub verify_batch_tokens: DevicePtr,
48    /// 2026-09-28: The quantized-activation scratch the NVFP4 MMQ GEMMs read
49    /// (`BufferArena::ffn_act_q8`), sized for the widest batch.
50    pub ffn_act_q8: DevicePtr,
51    /// 2026-09-28: K pool per attention layer (`PagedKvCache::k_pool_ptr`).
52    pub k_pools: Vec<DevicePtr>,
53    /// 2026-09-28: V pool per attention layer.
54    pub v_pools: Vec<DevicePtr>,
55    /// 2026-09-28: Tokens per KV block.
56    pub block_size: u32,
57    /// 2026-09-28: `PagedKvCache::cache_stride`.
58    pub cache_stride: u64,
59}
60
61/// 2026-09-29: The MTP draft head's fixed buffers.
62#[derive(Clone)]
63pub struct DraftFixed {
64    /// 2026-09-29: Where the host puts the token's embedding row (`MtpHead::forward_one`'s
65    /// `ssm_qkvz`).
66    pub embed: DevicePtr,
67    /// 2026-09-29: The draft step's attention metadata (`mtp_meta::mtp_attn_meta_dev`);
68    /// `max_blocks_per_seq` is ignored.
69    pub meta: AttnMetadataDev,
70    /// 2026-09-29: The draft cache's pools and geometry.
71    pub k_pool: DevicePtr,
72    pub v_pool: DevicePtr,
73    pub block_size: u32,
74    pub cache_stride: u64,
75    /// 2026-09-29: The vocabulary rows the draft lm_head scores (`DraftBinding::vocab`).
76    pub vocab: u32,
77    /// 2026-09-30: The head's n-row draft (the batched propose) facts (`DraftBinding::rows`);
78    /// `None` when the head drafts one row at a time.
79    pub rows: Option<DraftRows>,
80}
81
82impl DraftFixed {
83    /// 2026-09-30: The attention metadata an `n`-row draft step uploads: the batched
84    /// propose's, at its allocation (`mtp_meta::mtp_attn_meta_batch_dev`); one row reads the
85    /// single-row step's.
86    pub fn meta_rows(&self, n: u64) -> anyhow::Result<AttnMetadataDev> {
87        if n == 1 {
88            return Ok(self.meta);
89        }
90        let rows = self.rows.as_ref().ok_or_else(|| {
91            anyhow::anyhow!("a {n}-row draft plan for a head outside the batched propose")
92        })?;
93        Ok(crate::layers::mtp_meta::mtp_attn_meta_batch_dev(
94            rows.meta,
95            usize::try_from(n)?,
96        ))
97    }
98}
metrale_model_layers/circuit_exec/fixed.rs

metrale_model_layers/circuit_exec/
fixed.rs
