metrale_model_arch/dflash_head/
proposer.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `DraftProposer` for `BlockDiffusionDraftHead`, and the allocation of
4//! its per-sequence state.
5//!
6//! Owner: model-arch (DFlash drafter).
7//! Invariants:
8//! - `propose_batch` declines with `Ok(None)`, never an error, so the caller
9//!   falls back to the per-sequence loop.
10
11use super::*;
12
13impl DraftProposer for BlockDiffusionDraftHead {
14    fn block_gamma(&self) -> Option<usize> {
15        Some(self.gamma)
16    }
17
18    fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>> {
19        self.alloc_state_windowed(gpu, usize::MAX)
20    }
21
22    fn alloc_state_for(
23        &self,
24        gpu: &dyn GpuBackend,
25        budget_tokens: usize,
26    ) -> Result<Box<dyn ProposerState>> {
27        self.alloc_state_windowed(gpu, budget_tokens)
28    }
29
30    fn propose(
31        &self,
32        last_token: u32,
33        target_hidden: metrale_gpu_runtime::gpu::DevicePtr,
34        position: usize,
35        num_drafts: usize,
36        state: &mut dyn ProposerState,
37        ctx: &metrale_model_layers::layer::ForwardContext,
38        stream: u64,
39        draft_embed_target: Option<metrale_gpu_runtime::gpu::DevicePtr>,
40        grammar_bitmask: Option<&[i32]>,
41        target_hidden_stack: Option<metrale_gpu_runtime::gpu::DevicePtr>,
42    ) -> Result<Vec<u32>> {
43        self.propose_drafts(
44            last_token,
45            target_hidden,
46            position,
47            num_drafts,
48            state,
49            ctx,
50            stream,
51            draft_embed_target,
52            grammar_bitmask,
53            target_hidden_stack,
54            None,
55        )
56    }
57
58    /// 2026-09-25: 1 without the DFlash2 path or when
59    /// `METRALE_DFLASH_BATCH_PROPOSE` is below 2; otherwise that width, capped
60    /// at the scratch bands (`max_batch`).
61    fn propose_batch_max(
62        &self,
63        _buffers: &metrale_gpu_runtime::buffers::BufferArena,
64        _config: &metrale_config::ModelConfig,
65    ) -> usize {
66        if !self.dflash2_active() {
67            return 1;
68        }
69        let want = self.levers.batch_propose_width;
70        if want < 2 {
71            return 1;
72        }
73        want.min(self.max_batch.max(1))
74    }
75
76    /// 2026-09-25: One drafter forward over `n * block_g()` rows instead of n
77    /// forwards. The per-sequence prep runs through `propose_drafts` with a
78    /// `collect_prep` sink; the batched ctx precompute and the forward then run
79    /// once for all n sequences.
80    ///
81    /// Returns `Ok(None)` to decline: fewer than 2 or more than `max_batch`
82    /// sequences, no DFlash2 path, a failed prep, precompute or forward, or a
83    /// forward that returned fewer than `n * block_g()` rows.
84    fn propose_batch(
85        &self,
86        last_tokens: &[u32],
87        _target_hiddens: &[metrale_gpu_runtime::gpu::DevicePtr],
88        positions: &[usize],
89        num_drafts: usize,
90        states: &mut [&mut dyn ProposerState],
91        ctx: &metrale_model_layers::layer::ForwardContext,
92        stream: u64,
93        _out_conf: Option<&mut Vec<Vec<f32>>>,
94    ) -> Result<Option<Vec<Vec<u32>>>> {
95        let n = last_tokens.len();
96        if n < 2
97            || n > self.max_batch
98            || positions.len() != n
99            || states.len() != n
100            || !self.dflash2_active()
101        {
102            return Ok(None);
103        }
104
105        // 2026-09-25: Per-sequence prep, collecting each sequence's paged
106        // descriptor. One failed prep declines the whole batch.
107        self.set_block_g(num_drafts);
108        let mut prep: Vec<(metrale_gpu_runtime::gpu::DevicePtr, u32)> = Vec::with_capacity(n);
109        for (i, st) in states.iter_mut().enumerate() {
110            let before = prep.len();
111            match self.propose_drafts(
112                last_tokens[i],
113                metrale_gpu_runtime::gpu::DevicePtr::NULL,
114                positions[i],
115                num_drafts,
116                *st,
117                ctx,
118                stream,
119                None,
120                None,
121                None,
122                Some(&mut prep),
123            ) {
124                Ok(_) if prep.len() == before + 1 => {}
125                Ok(_) => return Ok(None),
126                Err(e) => {
127                    tracing::warn!("DFlash batched propose prep (seq {i}): {e:#} — per-seq path");
128                    return Ok(None);
129                }
130            }
131        }
132
133        // 2026-09-25: One ctx precompute over every sequence's uncommitted
134        // rows, which the prep deferred. A failure leaves `ctx_committed`
135        // unchanged for every tail not yet written, so the per-sequence path
136        // recomputes those.
137        if batched_precompute_enabled()
138            && let Err(e) = self.precompute_ctx_kv_batched(states, ctx, stream)
139        {
140            tracing::warn!("DFlash batched ctx precompute: {e:#} — per-seq path");
141            return Ok(None);
142        }
143
144        let batch = DflashBatch {
145            last_tokens,
146            positions,
147            block_tables: prep.iter().map(|p| p.0).collect(),
148            ctx_counts: prep.iter().map(|p| p.1).collect(),
149        };
150        let all = match self.forward_block(
151            last_tokens[0],
152            positions[0],
153            ctx,
154            stream,
155            None,
156            Some(prep[0]),
157            Some(&batch),
158        ) {
159            Ok(v) => v,
160            Err(e) => {
161                tracing::warn!("DFlash batched forward_block: {e:#} — falling back to per-seq");
162                return Ok(None);
163            }
164        };
165        let g = self.block_g();
166        if all.len() < n * g {
167            tracing::warn!(
168                "DFlash batched forward returned {} rows, expected {} — per-seq path",
169                all.len(),
170                n * g
171            );
172            return Ok(None);
173        }
174
175        // 2026-09-25: Split the bands. As in `propose_drafts`, a drafter with a
176        // mask token drops each band's row 0 (the anchor).
177        let cap = self.levers.draft_cap.unwrap_or(g);
178        let mut out: Vec<Vec<u32>> = Vec::with_capacity(n);
179        for (i, st) in states.iter_mut().enumerate() {
180            let band = &all[i * g..(i + 1) * g];
181            let drafts: Vec<u32> = if self.mask_token_id != 0 {
182                band.iter().skip(1).copied().take(cap).collect()
183            } else {
184                band.iter().copied().take(cap).collect()
185            };
186            if let Some(d) = st.as_any_mut().downcast_mut::<DflashProposerState>() {
187                d.last_num_drafted = drafts.len();
188            }
189            out.push(drafts);
190        }
191        Ok(Some(out))
192    }
193
194    fn after_verify(
195        &self,
196        num_accepted: usize,
197        state: &mut dyn ProposerState,
198        _stream: u64,
199    ) -> Result<()> {
200        let dstate = state
201            .as_any_mut()
202            .downcast_mut::<DflashProposerState>()
203            .ok_or_else(|| anyhow::anyhow!("Invalid DFlash proposer state"))?;
204        // 2026-09-25: `ctx_committed` counts the ctx rows whose K/V is in the
205        // paged cache. The paths that shrink `ctx_len`, the window slides in
206        // model-engine's `commit_ctx` and `dflash_serial_ctx_append`, reset it
207        // to 0; `propose_drafts` also clamps it to `ctx_len`.
208        let _ = num_accepted;
209        dstate.last_num_drafted = 0;
210        Ok(())
211    }
212
213    fn free_state(&self, gpu: &dyn GpuBackend, state: &mut dyn ProposerState) -> Result<()> {
214        // 2026-09-25: Return the drafter's paged blocks to the shared pool.
215        let dstate = match state.as_any_mut().downcast_mut::<DflashProposerState>() {
216            Some(s) => s,
217            None => return Ok(()),
218        };
219        if !dstate.block_table.is_empty() {
220            self.kv_cache.lock().free_blocks(&dstate.block_table);
221            dstate.block_table.clear();
222        }
223        // 2026-09-25: Free the ctx accumulator (`max_ctx_len` rows of
224        // `ctx_slot_bytes`); `DevicePtr` is `Copy` and frees nothing on drop.
225        // The null check and reset make a second `free_state` skip it.
226        if dstate.ctx_hidden_acc.0 != 0 {
227            gpu.free(dstate.ctx_hidden_acc)?;
228            dstate.ctx_hidden_acc = DevicePtr(0);
229        }
230        if let Some(bt) = dstate.block_table_dev.take() {
231            gpu.free(bt)?;
232        }
233        // 2026-09-25: Counters and flags back to their `alloc_state` values.
234        dstate.max_ctx_count_drafter = 0;
235        dstate.ctx_count_drafter = 0;
236        dstate.ctx_committed = 0;
237        dstate.ctx_positions.clear();
238        dstate.seq_len = 0;
239        dstate.ctx_len = 0;
240        dstate.prefill_done = false;
241        dstate.last_num_drafted = 0;
242        dstate.last_num_accepted = 0;
243        dstate.skip_next_decode_append = false;
244        Ok(())
245    }
246}
247
248impl BlockDiffusionDraftHead {
249    /// 2026-09-25: Allocate proposer state with a ctx accumulator of `window`
250    /// rows: the smallest of `budget_tokens + gamma + 1`,
251    /// `METRALE_DFLASH_CTX_CAP` (unless `0`) and the head's `max_seq_len`.
252    fn alloc_state_windowed(
253        &self,
254        gpu: &dyn GpuBackend,
255        budget_tokens: usize,
256    ) -> Result<Box<dyn ProposerState>> {
257        let bf16 = 2usize;
258        let ctx_slot_bytes = self.target_layer_ids.len() * self.target_hidden_size * bf16;
259        // 2026-09-25: The accumulator is per sequence, `window *
260        // ctx_slot_bytes` bytes. `METRALE_DFLASH_CTX_CAP` (default 16384, `0`
261        // uncapped) bounds it below the head's `max_seq_len`. When it fills,
262        // model-engine's `commit_ctx` and `dflash_serial_ctx_append` slide it,
263        // keeping the newest rows and their positions.
264        let cap = dflash_ctx_cap();
265        let ceiling = if cap == 0 {
266            self.max_seq_len
267        } else {
268            self.max_seq_len.min(cap)
269        };
270        // 2026-09-25: `budget_tokens` is the prompt plus `max_tokens` when the
271        // caller knows it, `usize::MAX` otherwise.
272        let window = ceiling.min(budget_tokens.saturating_add(self.gamma + 1));
273        if ceiling < self.max_seq_len {
274            static LOGGED: std::sync::atomic::AtomicBool =
275                std::sync::atomic::AtomicBool::new(false);
276            if !LOGGED.swap(true, std::sync::atomic::Ordering::Relaxed) {
277                tracing::info!(
278                    "DFlash ctx window capped to {} of --max-seq-len {} ({} MB/seq instead of \
279                     {} MB): the accumulator is PER SEQUENCE, so the uncapped size is what \
280                     OOMs a high-concurrency long-context serve. Override with \
281                     METRALE_DFLASH_CTX_CAP=<tokens> (0 = uncapped).",
282                    ceiling,
283                    self.max_seq_len,
284                    ceiling * ctx_slot_bytes / (1024 * 1024),
285                    self.max_seq_len * ctx_slot_bytes / (1024 * 1024),
286                );
287            }
288        }
289        let total = window * ctx_slot_bytes;
290        let ctx_hidden_acc = gpu.alloc(total)?;
291        // 2026-09-25: Zeroed, so no earlier allocation's data is visible.
292        gpu.memset(ctx_hidden_acc, 0, total)?;
293        Ok(Box::new(DflashProposerState {
294            block_table: Vec::with_capacity(64),
295            seq_len: 0,
296            last_num_drafted: 0,
297            prefill_done: false,
298            ctx_hidden_acc,
299            ctx_len: 0,
300            last_num_accepted: 0,
301            skip_next_decode_append: false,
302            max_ctx_len: window,
303            ctx_slot_bytes,
304            // 2026-09-25: Allocated by the first propose with the paged cache on.
305            block_table_dev: None,
306            ctx_count_drafter: 0,
307            max_ctx_count_drafter: 0,
308            ctx_committed: 0,
309            ctx_positions: Vec::new(),
310        }))
311    }
312}
metrale_model_arch/dflash_head/proposer.rs

metrale_model_arch/dflash_head/
proposer.rs
