Struct BlockDiffusionDraftHead 

pub struct BlockDiffusionDraftHead {

    pub num_layers: usize,
    pub hidden_size: usize,
    pub intermediate_size: usize,
    pub num_q_heads: usize,
    pub num_kv_heads: usize,
    pub head_dim: usize,
    pub vocab_size: usize,
    pub draft_vocab_size: usize,
    pub gamma: usize,
    pub mask_token_id: u32,
    pub window_size: Option<usize>,
    pub target_layer_ids: Vec<usize>,
    pub target_hidden_size: usize,
    pub embed_tokens_shared: DevicePtr,
    pub lm_head_shared: DevicePtr,
    pub lm_head_nvfp4: Option<QuantizedWeight>,
    pub lm_head_shared_fp8: Option<Fp8DenseWeight>,
    pub hidden_norm: DenseWeight,
    pub norm: DenseWeight,
    pub fc: DenseWeight,
    pub draft_id_to_target_id: Option<DevicePtr>,
    pub layers: Vec<DflashLayer>,
    pub fused_kv_weight: Option<DevicePtr>,
    pub kv_cache: Mutex<PagedKvCache>,
    pub scratch: DflashScratch,
    pub kernels: DflashKernels,
    pub max_seq_len: usize,
    pub yarn_inv_freq: DevicePtr,
    pub rope_theta: f32,
    pub rotary_dim: usize,
    pub rms_norm_eps: f32,
    pub ctx_window: usize,
    pub propose_graphs: Mutex<ProposeGraphs>,
    pub suppress_graphs: AtomicBool,
    pub levers: DFlashLevers,
    pub propose_warmup_count: AtomicUsize,
    pub quant: DflashQuantization,
    pub markov_rank: usize,
    pub markov_w1: Option<DenseWeight>,
    pub markov_w2: Option<DenseWeight>,
    pub confidence_proj: Option<DenseWeight>,
    pub confidence_bias: Option<DenseWeight>,
    pub confidence_with_markov: bool,
    pub shifted_rows: bool,
    pub conv_kernel_size: usize,
    pub conv_group_size: usize,
    pub selector_rank: usize,
    pub selector_top_k: usize,
    pub selector_pred: Option<DenseWeight>,
    pub selector_succ: Option<DenseWeight>,
    pub selector_hidden_proj: Option<DenseWeight>,
    /* private fields */
}

2026-09-25: Block-diffusion draft head; its public API is the DraftProposer trait. It uses the target’s token embedding and LM head (embed_tokens_shared, lm_head_shared / lm_head_nvfp4) alongside the drafter checkpoint’s own fc, hidden_norm, norm and layers. The geometry fields come from the drafter’s config.

Fields

num_layers: usize

hidden_size: usize

intermediate_size: usize

num_q_heads: usize

num_kv_heads: usize

head_dim: usize

vocab_size: usize

draft_vocab_size: usize

gamma: usize

2026-09-25: The drafter’s widest block (rows per sequence): the launch --dflash-gamma, or default_dflash_gamma of the checkpoint’s block size. The gamma-sized buffers (scratch bands, the drafter KV pool) are sized from it. The block a propose runs is Self::block_g.

mask_token_id: u32

window_size: Option<usize>

target_layer_ids: Vec<usize>

2026-09-25: The target layers whose hidden states the drafter is conditioned on, from the drafter config’s dflash_config.target_layer_ids.

target_hidden_size: usize

2026-09-25: The target’s hidden size; fc takes target_layer_ids.len() * target_hidden_size inputs.

embed_tokens_shared: DevicePtr

2026-09-25: The target’s token embedding table (the drafter checkpoint has none).

lm_head_shared: DevicePtr

2026-09-25: The target’s bf16 LM head; not used when lm_head_nvfp4 is Some.

lm_head_nvfp4: Option<QuantizedWeight>

2026-09-25: The target’s NVFP4 LM head, when it ships one; the final logits GEMM then uses it instead of lm_head_shared.

lm_head_shared_fp8: Option<Fp8DenseWeight>

2026-09-25: The FP8 LM head for Fp8Weights: the target checkpoint’s native FP8 head when the loader passes one, else an FP8 copy of lm_head_shared. None on the bf16 path.

hidden_norm: DenseWeight

norm: DenseWeight

fc: DenseWeight

draft_id_to_target_id: Option<DevicePtr>

2026-09-25: A draft-vocab to target-vocab id remap; from_weights sets None.

layers: Vec<DflashLayer>

fused_kv_weight: Option<DevicePtr>

2026-09-25: Every drafter layer’s K and V projection weights in one bf16 matrix [L * 2 * kv_dim, h], rows [K0; V0; K1; V1; ...], stitched from layers[i] at construction, so precompute_ctx_kv computes every layer’s ctx K/V with one GEMM.

kv_cache: Mutex<PagedKvCache>

2026-09-25: The drafter’s paged bf16 KV cache: one cache holding all num_layers drafter layers, block-table-keyed like the target’s.

scratch: DflashScratch

2026-09-25: Per-step scratch buffers, allocated once at construction.

kernels: DflashKernels

max_seq_len: usize

2026-09-25: The model’s max_seq_len; with dflash_ctx_cap() and the request’s budget it bounds each sequence’s ctx accumulator.

yarn_inv_freq: DevicePtr

2026-09-25: The RoPE inverse-frequency table, f32 [rotary_dim / 2] on the device, YaRN-corrected when the drafter config has rope_scaling.

rope_theta: f32

rotary_dim: usize

2026-09-25: RoPE covers the whole head (rotary_dim = head_dim).

rms_norm_eps: f32

2026-09-25: RMSNorm epsilon; from_weights sets 1e-6.

ctx_window: usize

2026-09-25: Most past target positions the drafter attends per step (METRALE_DFLASH_CTX_WINDOW, default 4096).

propose_graphs: Mutex<ProposeGraphs>

2026-09-25: The piecewise propose graphs, keyed by block width (block_g): a capture bakes its row count, so each width gets its own, captured once after the warm-up. Each width holds 2 * num_layers + 1 handles, [pre_0, post_0, ..., pre_{N-1}, post_{N-1}, tail] (the tail is the final norm, LM head and argmax). Attention runs eagerly between a layer’s halves and is never captured. GraphHandle(0) marks an empty capture; that slot runs eagerly.

suppress_graphs: AtomicBool

2026-09-25: When set, forward_block runs eagerly and captures nothing.

levers: DFlashLevers

2026-09-25: Diagnostic and A/B levers, resolved from the environment once when the head was built; forward_block and its per-layer helpers read these instead of the environment. See levers::DFlashLevers.

propose_warmup_count: AtomicUsize

2026-09-25: Not read; the warm-up is counted per width in ProposeGraphs::warmup.

quant: DflashQuantization

markov_rank: usize

2026-09-25: DSpark Markov head rank, 0 when the drafter has no Markov head (the markov_w* tensors absent).

markov_w1: Option<DenseWeight>

markov_w2: Option<DenseWeight>

confidence_proj: Option<DenseWeight>

confidence_bias: Option<DenseWeight>

confidence_with_markov: bool

2026-09-25: The drafter config’s confidence_head_with_markov.

shifted_rows: bool

2026-09-25: The shifted-row convention (drafter config dflash_config.projector_type == "dspark"): row j’s output is the token at position j + 1, so the draft vector is rotated right by one. METRALE_DSPARK_SHIFT=0|1 overrides it.

conv_kernel_size: usize

conv_group_size: usize

selector_rank: usize

selector_top_k: usize

selector_pred: Option<DenseWeight>

selector_succ: Option<DenseWeight>

selector_hidden_proj: Option<DenseWeight>

Implementations

impl BlockDiffusionDraftHead

pub fn from_weights( weights: DflashWeights, embed_tokens_shared: DevicePtr, lm_head_shared: DevicePtr, lm_head_nvfp4: Option<QuantizedWeight>, lm_head_native_fp8: Option<(Fp8DenseWeight, usize)>, target_hidden_size: usize, gamma: Option<usize>, window_size: Option<usize>, gpu: &dyn GpuBackend, max_seq_len: usize, max_batch_size: usize, ) -> Result<Self>

pub fn validate_against_target(&self, target_hidden_size: usize) -> Result<()>

2026-09-25: Errors when target_hidden_size differs from the value the head was built with, which sets the fc input width (target_layer_ids.len() * target_hidden_size). Nothing calls it.

impl BlockDiffusionDraftHead

pub fn block_g(&self) -> usize

2026-09-25: Block width (rows per sequence) of the propose in flight.

Trait Implementations

impl DraftProposer for BlockDiffusionDraftHead

fn propose_batch_max( &self, _buffers: &BufferArena, _config: &ModelConfig, ) -> usize

2026-09-25: 1 without the DFlash2 path or when METRALE_DFLASH_BATCH_PROPOSE is below 2; otherwise that width, capped at the scratch bands (max_batch).

fn propose_batch( &self, last_tokens: &[u32], _target_hiddens: &[DevicePtr], positions: &[usize], num_drafts: usize, states: &mut [&mut dyn ProposerState], ctx: &ForwardContext<'_>, stream: u64, _out_conf: Option<&mut Vec<Vec<f32>>>, ) -> Result<Option<Vec<Vec<u32>>>>

2026-09-25: One drafter forward over n * block_g() rows instead of n forwards. The per-sequence prep runs through propose_drafts with a collect_prep sink; the batched ctx precompute and the forward then run once for all n sequences.

Returns Ok(None) to decline: fewer than 2 or more than max_batch sequences, no DFlash2 path, a failed prep, precompute or forward, or a forward that returned fewer than n * block_g() rows.

fn block_gamma(&self) -> Option<usize>

2026-09-25: The block size γ of a block-diffusion (DFlash) drafter; None for other proposers. With --dflash, serve sets the draft count to γ - 1, at least 1.

fn alloc_state(&self, gpu: &dyn GpuBackend) -> Result<Box<dyn ProposerState>>

2026-09-25: Allocate per-sequence proposer state.

fn alloc_state_for( &self, gpu: &dyn GpuBackend, budget_tokens: usize, ) -> Result<Box<dyn ProposerState>>

2026-09-25: Self::alloc_state given the sequence’s token budget (prompt length plus max_tokens), so state that scales with context can be sized to the request instead of --max-seq-len. usize::MAX means unknown. The default ignores the budget.

fn propose( &self, last_token: u32, target_hidden: DevicePtr, position: usize, num_drafts: usize, state: &mut dyn ProposerState, ctx: &ForwardContext<'_>, stream: u64, draft_embed_target: Option<DevicePtr>, grammar_bitmask: Option<&[i32]>, target_hidden_stack: Option<DevicePtr>, ) -> Result<Vec<u32>>

2026-09-25: Propose up to num_drafts tokens autoregressively.

fn after_verify( &self, num_accepted: usize, state: &mut dyn ProposerState, _stream: u64, ) -> Result<()>

2026-09-25: Trim the proposer’s state after a verify that accepted num_accepted drafts.

fn free_state( &self, gpu: &dyn GpuBackend, state: &mut dyn ProposerState, ) -> Result<()>

2026-09-25: Free per-sequence proposer state when a sequence finishes. DevicePtr has no Drop, so device memory alloc_state allocated leaks unless it is freed here; the default frees nothing.

fn prefill_hidden_rows(&self, max_seq_len: usize) -> usize

2026-09-25: Rows of mtp_prefill_hidden this proposer can use, given the served --max-seq-len; the model sizes that [rows, hidden] BF16 buffer from it. The default is max_seq_len.

fn needs_comm(&self) -> bool

2026-09-25: Whether this proposer’s block is sharded across ranks, so its forward is given the communicator. The default is false; the GLM-5.3 MTP head returns mtp_ep_propose_enabled. When it is true on a multi-rank serve, the model sends the worker EP_CMD_MTP_PROPOSE before each propose, so both ranks run the same forward and collectives.

fn prefill_uses_shared_buffers(&self) -> bool

2026-09-25: Whether this proposer’s context prefill uses the shared forward scratch (ctx.buffers). If so, the model skips the prefill at the end of the target’s prefill, while the target still owns those buffers, and it runs at the first propose instead. Measured 2026-08-29 on GLM-5.3 (2x GB10): running it at the end of prefill changed the target’s output on 2 of 6 probes.

fn drafter_rows(&self, _state: &mut dyn ProposerState) -> usize

2026-09-25: Drafter KV rows held in state; 0 when unknown.

fn last_pair_key(&self, _state: &mut dyn ProposerState) -> Option<usize>

2026-09-25: Sequence-space pair key of the newest drafter row; None when untracked. Rows are compacted, so the row count does not give the sequence position.

fn take_drafter_kv( &self, _state: &mut dyn ProposerState, ) -> Option<(Vec<u32>, usize, Option<usize>)>

2026-09-25: Move this sequence’s drafter KV blocks out of its state, so free_state does not release them and the model can carry them to the next turn. Returns (blocks, rows, last_pair_key), or None when unsupported or empty. Afterwards the state must behave as freshly allocated.

fn install_drafter_kv( &self, _state: &mut dyn ProposerState, _blocks: Vec<u32>, _rows: usize, _last_pair_key: Option<usize>, ) -> bool

2026-09-25: Install carried blocks into a fresh state, the inverse of Self::take_drafter_kv. Returns false when unsupported; the caller then still owns the blocks and must free them.

fn free_drafter_kv(&self, _blocks: &[u32])

2026-09-25: Release drafter KV blocks that no proposer state owns, such as a carried entry being replaced or dropped.

fn catchup_drafter( &self, _tokens: &[u32], _hiddens: DevicePtr, _row_base: usize, _pos_base: usize, _state: &mut dyn ProposerState, _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<usize, Error>

2026-09-25: Append drafter rows at KV slots row_base .. with RoPE positions pos_base .. from (tokens, hiddens) pairs. Returns the rows written; the default writes none.

fn catchup_batch( &self, _tokens: &[Vec<u32>], _hiddens: &[Vec<DevicePtr>], _first_pos: &[usize], _states: &mut [&mut dyn ProposerState], _ctx: &ForwardContext<'_>, _stream: u64, ) -> Result<usize, Error>

2026-09-25: Append one drafter row per accepted draft of the last verify, for several sequences; the model calls it only with ModelLevers::mtp_kv_exact on. tokens[i] are sequence i’s accepted drafts, hiddens[i][k] the target hidden paired with tokens[i][k], and first_pos[i] the RoPE position of tokens[i][0]. Returns the rows written; the default writes none.

fn prefill_drafter( &self, prompt_tokens: &[u32], hiddens: DevicePtr, state: &mut dyn ProposerState, ctx: &ForwardContext<'_>, stream: u64, ) -> Result<usize, Error>

2026-09-25: Prefill the drafter’s KV over the prompt before the sequence’s first propose. hiddens is [P, hidden_size] BF16, row i the target’s hidden for prompt_tokens[i] before the final norm. Returns the drafter rows written; the default writes none.

fn read_deferred_draft_token(&self, gpu: &dyn GpuBackend) -> Result<u32, Error>

2026-09-25: The draft token id the last propose with draft_embed_target left on the device; the default returns 0.

fn circuit_draft( &self, _config: &ModelConfig, _levers: &ModelLevers, ) -> Option<DraftBinding>

2026-09-29: This proposer’s draft head as the circuit executor binds it; None for a proposer the circuit does not model.

fn set_circuit_draft(&self, _runner: Option<Arc<dyn DraftRunner>>)

2026-09-29: Run runner in place of the draft head’s own forward from now on, or its own forward again with None. A proposer without a circuit binding ignores it.