metrale_model_arch/
dflash_head.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: DFlash block-diffusion draft head (arXiv 2602.06036),
4//! implementing [`DraftProposer`].
5//!
6//! A small Qwen3-style transformer drafter that emits a block of draft tokens
7//! in one forward pass with bidirectional in-block attention. It is
8//! conditioned on target hidden states captured at `target_layer_ids`,
9//! projected by one `fc` layer at the drafter's input.
10//!
11//! Owner: model-arch (DFlash drafter).
12//! Invariants:
13//! - `block_g()` stays in `2..=gamma.max(2)` (`set_block_g` clamps).
14
15use parking_lot::Mutex;
16use std::any::Any;
17
18use anyhow::Result;
19use metrale_cache::kv_cache::PagedKvCache;
20use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
21
22use metrale_model_layers::speculative::{DraftProposer, ProposerState};
23use metrale_model_layers::weight_map::{DenseWeight, QuantizedWeight};
24
25pub use kernels::DflashKernels;
26pub use state::{DflashProposerState, DflashScratch};
27
28/// 2026-09-25: Cross-sequence batch for one drafter forward: per sequence its
29/// last token, position, drafter block table (a device pointer) and filled
30/// ctx slot count.
31///
32/// Rows are sequence-major: sequence `i` owns rows `[i * w, (i + 1) * w)` of
33/// the scratch buffers, `w = block_g()`. Attention, the KV slot writes and
34/// the selector's chain seed are per sequence; the weight-bearing ops run
35/// once over all `n * w` rows.
36pub(super) struct DflashBatch<'a> {
37    pub last_tokens: &'a [u32],
38    pub positions: &'a [usize],
39    pub block_tables: Vec<DevicePtr>,
40    pub ctx_counts: Vec<u32>,
41}
42
43/// 2026-09-25: Drafter weight precision. `from_weights` builds the head as
44/// `Bf16` and switches it to `Fp8Weights` unless `METRALE_DFLASH_DRAFTER_FP8=0`
45/// or the `fp8_gemm_n128_row_scaled` kernels are missing.
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum DflashQuantization {
48    Bf16,
49    /// 2026-09-25: Weight-only FP8: the q/k/v/o/gate/up/down weights get FP8
50    /// E4M3 copies with per-row f32 scales at load, and the LM head is FP8
51    /// (`lm_head_shared_fp8`). The KV cache stays bf16.
52    Fp8Weights,
53}
54
55/// 2026-09-25: One drafter layer's weights: the bf16 norms (with per-head Q/K
56/// RMSNorm), projections and MLP; FP8 copies of the seven GEMM weights when
57/// the head is `Fp8Weights` (`None` otherwise); and the DFlash2 conv weights
58/// (`attention_conv.*`, `mlp_conv.*`), `None` when the checkpoint has none.
59#[allow(dead_code)]
60pub struct DflashLayer {
61    pub input_layernorm: DenseWeight,
62    pub post_attention_layernorm: DenseWeight,
63    pub q_proj: DenseWeight,
64    pub k_proj: DenseWeight,
65    pub v_proj: DenseWeight,
66    pub o_proj: DenseWeight,
67    pub q_norm: DenseWeight,
68    pub k_norm: DenseWeight,
69    pub gate_proj: DenseWeight,
70    pub up_proj: DenseWeight,
71    pub down_proj: DenseWeight,
72
73    pub q_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
74    pub k_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
75    pub v_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
76    pub o_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
77    pub gate_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
78    pub up_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
79    pub down_proj_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
80
81    pub attention_conv_base: Option<DenseWeight>,
82    pub attention_conv_proj: Option<DenseWeight>,
83    pub mlp_conv_base: Option<DenseWeight>,
84    pub mlp_conv_proj: Option<DenseWeight>,
85}
86
87#[allow(dead_code)]
88/// 2026-09-25: The propose graphs per block width (see
89/// `BlockDiffusionDraftHead::propose_graphs`). `warmup` counts the eager
90/// passes run at a width; its capture starts once the count reaches
91/// `DFlashLevers::propose_warmup_n` (`METRALE_DFLASH_PROPOSE_WARMUP_N`,
92/// default 2).
93#[derive(Default)]
94pub struct ProposeGraphs {
95    pub by_width: std::collections::HashMap<usize, Vec<metrale_gpu_runtime::gpu::GraphHandle>>,
96    pub warmup: std::collections::HashMap<usize, usize>,
97}
98
99/// 2026-09-25: Block-diffusion draft head; its public API is the
100/// [`DraftProposer`] trait. It uses the target's token embedding and LM head
101/// (`embed_tokens_shared`, `lm_head_shared` / `lm_head_nvfp4`) alongside the
102/// drafter checkpoint's own `fc`, `hidden_norm`, `norm` and layers. The
103/// geometry fields come from the drafter's config.
104pub struct BlockDiffusionDraftHead {
105    pub num_layers: usize,
106    pub hidden_size: usize,
107    pub intermediate_size: usize,
108    pub num_q_heads: usize,
109    pub num_kv_heads: usize,
110    pub head_dim: usize,
111    pub vocab_size: usize,
112    pub draft_vocab_size: usize,
113    /// 2026-09-25: The drafter's widest block (rows per sequence): the launch
114    /// `--dflash-gamma`, or `default_dflash_gamma` of the checkpoint's block
115    /// size. The gamma-sized buffers (scratch bands, the drafter KV pool) are
116    /// sized from it. The block a propose runs is [`Self::block_g`].
117    pub gamma: usize,
118    /// 2026-09-25: The block width of the propose in flight (rows per
119    /// sequence: the anchor plus `block_g - 1` mask rows), set by
120    /// `set_block_g` from the scheduler's `num_drafts` and read through
121    /// `block_g()`.
122    /// provenance-id: 526f6e616c6420522e205374657369616b
123    pub(super) block_gamma: std::sync::atomic::AtomicUsize,
124    /// 2026-09-25: Widest cross-sequence batch the scratch bands can hold.
125    pub(super) max_batch: usize,
126    pub mask_token_id: u32,
127    pub window_size: Option<usize>,
128    /// 2026-09-25: The target layers whose hidden states the drafter is
129    /// conditioned on, from the drafter config's `dflash_config.target_layer_ids`.
130    pub target_layer_ids: Vec<usize>,
131    /// 2026-09-25: The target's hidden size; `fc` takes
132    /// `target_layer_ids.len() * target_hidden_size` inputs.
133    pub target_hidden_size: usize,
134
135    /// 2026-09-25: The target's token embedding table (the drafter checkpoint
136    /// has none).
137    pub embed_tokens_shared: DevicePtr,
138    /// 2026-09-25: The target's bf16 LM head; not used when `lm_head_nvfp4` is
139    /// `Some`.
140    pub lm_head_shared: DevicePtr,
141    /// 2026-09-25: The target's NVFP4 LM head, when it ships one; the final
142    /// logits GEMM then uses it instead of `lm_head_shared`.
143    pub lm_head_nvfp4: Option<QuantizedWeight>,
144    /// 2026-09-25: The FP8 LM head for `Fp8Weights`: the target checkpoint's
145    /// native FP8 head when the loader passes one, else an FP8 copy of
146    /// `lm_head_shared`. `None` on the bf16 path.
147    pub lm_head_shared_fp8: Option<metrale_model_layers::weight_map::Fp8DenseWeight>,
148
149    // 2026-09-25: From the drafter checkpoint: `hidden_norm` (applied to the
150    // projected target context), the final `norm`, and `fc`, which projects
151    // the stacked captured target hiddens into the drafter's hidden size.
152    pub hidden_norm: DenseWeight,
153    pub norm: DenseWeight,
154    pub fc: DenseWeight,
155    /// 2026-09-25: A draft-vocab to target-vocab id remap; `from_weights` sets
156    /// `None`.
157    pub draft_id_to_target_id: Option<DevicePtr>,
158    pub layers: Vec<DflashLayer>,
159
160    /// 2026-09-25: Every drafter layer's K and V projection weights in one bf16
161    /// matrix `[L * 2 * kv_dim, h]`, rows `[K0; V0; K1; V1; ...]`, stitched from
162    /// `layers[i]` at construction, so `precompute_ctx_kv` computes every
163    /// layer's ctx K/V with one GEMM.
164    pub fused_kv_weight: Option<DevicePtr>,
165
166    /// 2026-09-25: The drafter's paged bf16 KV cache: one cache holding all
167    /// `num_layers` drafter layers, block-table-keyed like the target's.
168    pub kv_cache: Mutex<PagedKvCache>,
169
170    /// 2026-09-25: Per-step scratch buffers, allocated once at construction.
171    pub scratch: DflashScratch,
172
173    pub kernels: DflashKernels,
174
175    /// 2026-09-25: The model's `max_seq_len`; with `dflash_ctx_cap()` and the
176    /// request's budget it bounds each sequence's ctx accumulator.
177    pub max_seq_len: usize,
178
179    /// 2026-09-25: The RoPE inverse-frequency table, f32 `[rotary_dim / 2]` on
180    /// the device, YaRN-corrected when the drafter config has `rope_scaling`.
181    pub yarn_inv_freq: DevicePtr,
182
183    pub rope_theta: f32,
184
185    /// 2026-09-25: RoPE covers the whole head (`rotary_dim = head_dim`).
186    pub rotary_dim: usize,
187
188    /// 2026-09-25: RMSNorm epsilon; `from_weights` sets 1e-6.
189    pub rms_norm_eps: f32,
190
191    /// 2026-09-25: Most past target positions the drafter attends per step
192    /// (`METRALE_DFLASH_CTX_WINDOW`, default 4096).
193    pub ctx_window: usize,
194
195    /// 2026-09-25: The piecewise propose graphs, keyed by block width
196    /// (`block_g`): a capture bakes its row count, so each width gets its own,
197    /// captured once after the warm-up. Each width holds `2 * num_layers + 1`
198    /// handles, `[pre_0, post_0, ..., pre_{N-1}, post_{N-1}, tail]` (the tail
199    /// is the final norm, LM head and argmax). Attention runs eagerly between
200    /// a layer's halves and is never captured. `GraphHandle(0)` marks an empty
201    /// capture; that slot runs eagerly.
202    pub propose_graphs: Mutex<ProposeGraphs>,
203    /// 2026-09-25: When set, `forward_block` runs eagerly and captures nothing.
204    pub suppress_graphs: std::sync::atomic::AtomicBool,
205    /// 2026-09-25: Diagnostic and A/B levers, resolved from the environment
206    /// once when the head was built; `forward_block` and its per-layer
207    /// helpers read these instead of the environment. See
208    /// [`levers::DFlashLevers`].
209    pub levers: levers::DFlashLevers,
210    /// 2026-09-25: Not read; the warm-up is counted per width in
211    /// `ProposeGraphs::warmup`.
212    pub propose_warmup_count: std::sync::atomic::AtomicUsize,
213
214    pub quant: DflashQuantization,
215
216    /// 2026-09-25: DSpark Markov head rank, 0 when the drafter has no Markov
217    /// head (the `markov_w*` tensors absent).
218    pub markov_rank: usize,
219    pub markov_w1: Option<DenseWeight>,
220    pub markov_w2: Option<DenseWeight>,
221    pub confidence_proj: Option<DenseWeight>,
222    pub confidence_bias: Option<DenseWeight>,
223    /// 2026-09-25: The drafter config's `confidence_head_with_markov`.
224    pub confidence_with_markov: bool,
225    /// 2026-09-25: The shifted-row convention (drafter config
226    /// `dflash_config.projector_type == "dspark"`): row j's output is the
227    /// token at position j + 1, so the draft vector is rotated right by one.
228    /// `METRALE_DSPARK_SHIFT=0|1` overrides it.
229    pub shifted_rows: bool,
230
231    // 2026-09-25: DFlash2 geometry from the drafter's `dflash_config` (0 when
232    // absent) and its candidate-selector weights (`None` when absent).
233    pub conv_kernel_size: usize,
234    pub conv_group_size: usize,
235    pub selector_rank: usize,
236    pub selector_top_k: usize,
237    pub selector_pred: Option<DenseWeight>,
238    pub selector_succ: Option<DenseWeight>,
239    pub selector_hidden_proj: Option<DenseWeight>,
240}
241
242mod dflash2;
243/// 2026-09-25: Rows the batched ctx precompute staging holds per step.
244pub(super) const PRECOMPUTE_BATCH_ROWS: usize = 256;
245
246/// 2026-09-25: Whether the batched ctx precompute (one fc + fused-KV pass over
247/// every sequence's uncommitted ctx rows) runs on the batched propose path:
248/// yes unless `METRALE_DFLASH_NO_BATCHED_PRECOMPUTE=1`, which runs the
249/// per-sequence loop.
250pub(super) fn batched_precompute_enabled() -> bool {
251    std::env::var("METRALE_DFLASH_NO_BATCHED_PRECOMPUTE")
252        .ok()
253        .as_deref()
254        != Some("1")
255}
256
257/// 2026-09-25: The `METRALE_DFLASH_OPTION_B` predicate (`DFlashLevers::option_b`):
258/// on unless the value is `0`. Pure over the raw value, so `option_b_tests`
259/// tests it without setting the process environment.
260pub(super) fn option_b_from(v: Option<&str>) -> bool {
261    v != Some("0")
262}
263
264#[cfg(test)]
265mod option_b_tests;
266
267mod forward_block;
268mod forward_block_layer;
269mod forward_block_layer_paged;
270mod from_weights;
271mod kernels;
272pub mod levers;
273mod markov;
274mod precompute_ctx_kv;
275mod precompute_ctx_kv_batched;
276mod propose;
277mod proposer;
278mod state;
279
280/// 2026-09-25: Read once: false only when `METRALE_NO_DFLASH_FP8_RT=1`, so the
281/// kernel choice stays fixed across CUDA-graph captures.
282pub(crate) fn fp8_rt_enabled() -> bool {
283    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
284    *ON.get_or_init(|| std::env::var("METRALE_NO_DFLASH_FP8_RT").as_deref() != Ok("1"))
285}
286
287/// 2026-09-25: The DFlash context-window bound, in tokens: the most recent
288/// target positions the drafter may accumulate. Both the per-sequence ctx
289/// accumulator (`proposer.rs`) and the model's whole-prompt hidden capture
290/// (`metrale-model-engine`'s `impl_a1.rs`) are sized from it.
291///
292/// `METRALE_DFLASH_CTX_CAP=<tokens>` (default 16384); `0` disables the cap.
293pub fn dflash_ctx_cap() -> usize {
294    std::env::var("METRALE_DFLASH_CTX_CAP")
295        .ok()
296        .and_then(|v| v.parse::<usize>().ok())
297        .unwrap_or(16384)
298}
299
300impl BlockDiffusionDraftHead {
301    /// 2026-09-25: Block width (rows per sequence) of the propose in flight.
302    #[inline]
303    pub fn block_g(&self) -> usize {
304        self.block_gamma.load(std::sync::atomic::Ordering::Relaxed)
305    }
306
307    /// 2026-09-25: Arm the block width for the next propose from the
308    /// scheduler's draft count: `num_drafts + 1` rows (anchor + masks),
309    /// clamped to `2..=gamma.max(2)`, so a request above the head's sizing
310    /// gets the widest block.
311    #[inline]
312    pub(super) fn set_block_g(&self, num_drafts: usize) {
313        let g = (num_drafts + 1).clamp(2, self.gamma.max(2));
314        self.block_gamma
315            .store(g, std::sync::atomic::Ordering::Relaxed);
316    }
317}
metrale_model_arch/dflash_head.rs

metrale_model_arch/
dflash_head.rs
