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<Option<Vec<GraphHandle>>>,
    pub suppress_graphs: AtomicBool,
    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 */
}

Block-diffusion draft head. Public API is the DraftProposer trait.

The drafter shares embed_tokens and lm_head with the target — these are NOT in the drafter’s safetensors checkpoint (verified against z-lab/Qwen3.6-35B-A3B-DFlash commit 42d3b34). The constructor takes the target’s embed_tokens_shared and lm_head_shared device pointers at build time and slots them in alongside the drafter’s own fc, hidden_norm, norm, and per-layer weights.

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

mask_token_id: u32

window_size: Option<usize>

target_layer_ids: Vec<usize>

target_layer_ids. Same data as TransformerModel::dflash_capture_layers, repeated here so the loader is the single source of truth; the model reads these to size its capture buffer.

target_hidden_size: usize

Target-side hidden_size (used for the fc projection input width: target_layer_ids.len() * target_hidden_size).

embed_tokens_shared: DevicePtr

Target’s embed_tokens GPU pointer. The drafter’s checkpoint has no own embeddings — both vocab and embedding dim must match the target (Qwen3.6-35B-A3B-DFlash: vocab=248320, hidden=2048 — same as target).

lm_head_shared: DevicePtr

Target’s lm_head GPU pointer. Used for the drafter’s per-position argmax over [γ, vocab] logits. Valid only when the target lm_head is BF16; when lm_head_nvfp4 is Some, the NVFP4 path is used instead.

lm_head_nvfp4: Option<QuantizedWeight>

Target’s NVFP4 lm_head (packed + scales), shared with the drafter for the final logits GEMM. Some when the target ships an NVFP4 lm_head (e.g. Holo) — required because a BF16 dense_gemm on the NVFP4 buffer reads garbage and OOB. None → use the BF16 lm_head_shared.

lm_head_shared_fp8: Option<Fp8DenseWeight>

Phase G — optional FP8 mirror of the shared lm_head weight, [vocab_size, hidden_size] FP8 E4M3 + per-row f32 scales. Built at model load when ATLAS_DFLASH_DRAFTER_FP8=1. Owned by the drafter (separate allocation from the shared BF16 ptr) since it must not mutate the target model’s lm_head. None on the BF16 path.

hidden_norm: DenseWeight

Hidden-norm applied to the projected target context before mixing with the embedded tokens (Qwen3-DFlash convention; see vLLM DFlashQwen3Model.hidden_norm).

norm: DenseWeight

Final RMSNorm before LM head.

fc: DenseWeight

fc projection — [draft_hidden, target_layer_ids.len() * target_hidden_size] BF16. Maps the stack of captured target hiddens to drafter’s input space once at model entry. Replaces the earlier (incorrect) “per-layer KV injection” design.

draft_id_to_target_id: Option<DevicePtr>

Optional draft-vocab-id → target-vocab-id remap. None when the drafter shares vocab with the target (Qwen3.6-35B-A3B-DFlash case: vocab_size == draft_vocab_size == 248320).

layers: Vec<DflashLayer>

Drafter transformer layers (8 for Qwen3.6-35B-A3B-DFlash).

fused_kv_weight: Option<DevicePtr>

Phase 2 (Option B) fused K/V projection across all L drafter layers. Shape: [L × 2 × kv_dim, h] BF16 — concatenated [K0; V0; K1; V1; …] (per-layer K then V interleaved). Built once at construction by copy_d2d-stitching the per-layer k_proj.weight and v_proj.weight pointers from layers[i]. Lets precompute_ctx_kv derive every drafter layer’s ctx K/V via a single dense_gemm of shape [new_ctx_count, h] × [h, L·2·kv_dim] instead of 2·L per-layer GEMMs.

None until Phase 2 lands the build (stage 1: kernel/dispatcher scaffolding; stage 2: this allocation + the precompute_ctx_kv module; stage 3: pyref bit-exact diff). Layout (K then V per layer) chosen to match vLLM’s _fused_kv_weight in qwen3_dflash.py:381-389.

kv_cache: Mutex<PagedKvCache>

Paged FP8 KV cache. One cache holding all num_layers drafter layers, laid out the same way the target’s KV cache is — block-table-keyed, num_layers × num_kv_heads × head_dim per slot. Allocating a single multi-layer cache (vs. one per drafter layer) matches Atlas’s existing PagedKvCache ABI and lets us reuse the existing reshape_and_cache kernel without per-layer dispatch overhead.

scratch: DflashScratch

Per-step scratch buffers (allocated once at construction, reused).

kernels: DflashKernels

All kernel handles needed by propose() and the eventual prefill projection (precompute_and_store_context_kv).

max_seq_len: usize

Per-sequence ctx accumulator capacity (mirrors model’s max_seq_len). Used by alloc_state to size each new sequence’s ctx_hidden_acc.

yarn_inv_freq: DevicePtr

Pre-computed yarn inv_freq table ([head_dim/2] f32 on GPU). Drafter rope_scaling: factor=64, beta_fast=32, beta_slow=1, original_max_position_embeddings=4096 (per drafter config.json).

rope_theta: f32

rope_theta (10000000 for Qwen3.6-DFlash). Stored to pass into the rope_yarn kernel each step.

rotary_dim: usize

rotary_dim. Drafter uses full-rotation (rotary_dim = head_dim = 128).

rms_norm_eps: f32

RMSNorm epsilon (drafter inherits Qwen3 default 1e-6).

ctx_window: usize

Max number of past target positions injected into the drafter’s K/V per step. Default γ — drafter sees at most γ ctx + γ noise = 2γ attention positions per step. ctx_window=0 disables ctx conditioning (degraded quality, ablation only).

propose_graphs: Mutex<Option<Vec<GraphHandle>>>

Per-subgraph captured handles. None until warm-up completes and the first capture pass lands; on the capture pass we fill this Vec with 2 × num_layers + 1 handles laid out as [pre_0, post_0, pre_1, post_1, ..., pre_{N-1}, post_{N-1}, tail]. Slot index = layer_idx * 2 + half for the layer halves (half = 0 for pre_attn, 1 for post_attn) and num_layers * 2 for the tail (final norm + lm_head + argmax). GraphHandle(0) is the “empty capture” sentinel and means that slot replays eager.

Phase F.2 (2026-05-28): replaces the single full-region capture with one capture per subgraph. Attention is NEVER captured — it’s the natural sync barrier between captured subgraphs (vLLM piecewise convention). See design doc §15.

suppress_graphs: AtomicBool

When set, all forward_block calls run eagerly. Mirrors target-model TransformerModel::suppress_graphs so external code can disable graphs at runtime (e.g. while calibrating FP8 KV).

propose_warmup_count: AtomicUsize

How many eager warm-up calls we’ve executed against the graph path. Default warmup target is 2 (override via ATLAS_DFLASH_PROPOSE_WARMUP_N). Two eager passes warm the PTX→SASS cache, ramp GB10 clocks to steady state, and bring hot weight tiles into L2 before the capture freezes SASS variants the driver picks. Shared across all subgraphs — every subgraph captures on the same propose call after the warmup target is hit.

quant: DflashQuantization

markov_rank: usize

Markov head rank (0 when the drafter has no Markov head). RadixArk Qwen3.8-27B-DSpark: 256.

markov_w1: Option<DenseWeight>

markov_w1: [vocab, rank] BF16 prev-token embedding table.

markov_w2: Option<DenseWeight>

markov_w2: [vocab, rank] BF16 latent→vocab projection (Linear(rank, vocab, bias=False).weight, [N, K] GEMV layout).

confidence_proj: Option<DenseWeight>

Confidence head (AcceptRatePredictor) weight [1, hidden(+rank)]. Loaded for the dynamic-K phase; not consumed by the Markov fixup.

confidence_bias: Option<DenseWeight>

Confidence head bias [1].

confidence_with_markov: bool

Whether the confidence input is [hidden ‖ markov_embed] (true) or hidden only (false). Mirrors confidence_head_with_markov.

shifted_rows: bool

SpecForge shifted row convention (drafter config dflash_config.projector_type == "dspark"): row j’s output is the token at position j+1, so the returned draft vector is rotated right by one to line up with Atlas’s z-lab-convention verify indexing. Overridable for A/B via ATLAS_DSPARK_SHIFT=0|1.

conv_kernel_size: usize

Conv kernel size (2) — taps per conv application.

conv_group_size: usize

Channels per conv group (16).

selector_rank: usize

Selector codebook rank (256).

selector_top_k: usize

Candidates per position for the selector walk (16; the kernels are specialized to 16 — other values refuse to arm).

selector_pred: Option<DenseWeight>

candidate_selector.predecessor_codebook [vocab, rank].

selector_succ: Option<DenseWeight>

candidate_selector.successor_codebook [vocab, rank].

selector_hidden_proj: Option<DenseWeight>

candidate_selector.hidden_projection.weight [rank, hidden].

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<()>

Borrow-validate the drafter dimensions against the target’s hidden_size at construction time. Mismatch is a hard error — the fc projection width is baked from target_hidden_size and a runtime mismatch would produce silent garbage (vLLM’s loader hits this same check).

Trait Implementations

impl DraftProposer for BlockDiffusionDraftHead

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

Widest batch one drafter forward can carry. Bounded by the scratch bands (max_batch); 1 means the batched path cannot run and the caller stays on propose.

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>>>>

Cross-sequence batched propose: ONE drafter forward over n * gamma rows instead of n forwards.

Per-sequence preparation (ctx append, Option-B block growth, the incremental ctx precompute) still runs per sequence — it is cheap, touching only the uncommitted ctx tail — and it reuses propose_drafts’ own prep through the collect_prep sink so the two paths cannot drift. The expensive part, the drafter layers plus an lm_head against a 248k vocab, runs ONCE for the whole batch. That is the entire win.

Returns Ok(None) to decline, and the caller falls back to the per-sequence loop — never a wrong answer.

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

The proposer’s trained block size γ, when it is a block-diffusion drafter (DFlash/DFlash2). The serve layer derives num_drafts from this — the head resolved it from the drafter checkpoint and is the SSOT. None = not a block drafter.

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

Allocate per-sequence proposer state.

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

Self::alloc_state with the sequence’s KNOWN token budget (prompt_len + max_tokens), so a proposer whose per-sequence state scales with context can size to what this request can actually reach instead of the global --max-seq-len ceiling. That distinction is what OOMs a high-concurrency long-context serve: the ceiling is per-sequence and paid n times, while a typical request needs a fraction of it.

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>>

Propose up to num_drafts tokens autoregressively.

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

Called after target verification to trim proposer state.

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

Free per-sequence proposer state (KV cache blocks, device buffers, etc.).

fn last_confidence(&self) -> Option<f32>

Chain confidence of the most recent propose (min top-1 softmax prob across its drafts), when the proposer computes it (draft_conf_tau > 0). None = not computed; callers must not gate on it then.

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

Rows this proposer can actually consume from mtp_prefill_hidden, given the served --max-seq-len.

fn needs_comm(&self) -> bool

True when this proposer’s block is SHARDED across ranks and its forward therefore needs the communicator (a routed-MoE all-reduce and a row-parallel o_proj reduce), like any target layer.

fn prefill_uses_shared_buffers(&self) -> bool

True when this proposer’s context prefill uses the SHARED forward scratch (ctx.buffers), so it must not run from the end-of-prefill eager hook — only from the first propose, where the target owns nothing.

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

Current drafter KV length (rows), for the catch-up append point. 0 = unknown / not applicable (catch-up is skipped).

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

Sequence-space pair key of the newest drafter row (None = untracked; catch-up is skipped). The drafter row space is compacted, so rows cannot locate the drafter in the sequence — this can.

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

ATLAS_MTP_CARRY_DRAFTER: move this sequence’s drafter KV blocks OUT of its proposer state, so free_state releases nothing and the model can hold them for the next turn. Returns (blocks, rows, last_pair_key); None = unsupported or nothing to carry. After this call the state must behave as if freshly allocated.

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

Inverse of Self::take_drafter_kv: install carried blocks into a fresh proposer state. Returns false when unsupported (caller must then free the blocks itself).

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

Release drafter KV blocks that no proposer state owns (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>

Append drafter rows at KV slots row_base .. with RoPE positions pos_base .. from (tokens, hiddens) pairs — the catch-up feed. Returns rows written (0 = unsupported/no-op).

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

Prefill the drafter’s own context (KV cache) over the prompt, before the first propose() of a sequence (ATLAS_MTP_DRAFTER_PREFILL).

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

Read the draft token ID stored on GPU by the last propose() call that used draft_embed_target = Some(...). Returns 0 if not supported.