Struct DflashDraft 

pub struct DflashDraft {
    pub cfg: DflashCfg,
    pub layers: Vec<DflashLayer>,
    pub fc: GpuTensor,
    pub hidden_norm: CudaSlice<f32>,
    pub norm: CudaSlice<f32>,
    pub markov: Option<MarkovHead>,
    pub confidence: Option<ConfidenceHead>,
    pub rope_yarn: Option<(CudaSlice<f32>, f32)>,
    pub dflash2: Option<Dflash2Head>,
}

Fields

cfg: DflashCfg

layers: Vec<DflashLayer>

fc: GpuTensor

hidden_norm: CudaSlice<f32>

norm: CudaSlice<f32>

markov: Option<MarkovHead>

DSpark semi-AR markov head (present in the repo-root checkpoint variant): draft logits at position k get + W2(W1[prev_realized_token]) — left-to-right within the block (the patch’s _markov_semiar_sample_block semantics, greedy). w1 = raw bf16 [V, rank] (row-gathered by device token id); w2 = q8_0 [rank->V].

confidence: Option<ConfidenceHead>

DSpark accept-rate head (trained with confidence loss). sglang’s DSPARK planner consumes it to SIZE VERIFY WINDOWS (cumprod survival — v0.5.16 headline; the earlier “reference serving loop never consumes it” note matched SpecForge’s legacy spec_generate only). memra schedules with it under MEMRA_DSPARK_VT=confidence (the H4 fix, DSPARK-POSTMORTEM-20260820.md: per-round verify window from cumprod survival, dspark_confidence_vt) and keeps it census+parity-only under the default ladder. Host-resident (5k floats).

rope_yarn: Option<(CudaSlice<f32>, f32)>

YaRN rope (q38 arm-a inherits the target’s rope_parameters: rope_type yarn, factor 32, original 8192, beta 32/1). ff = per-dim divisors for rope_neox_ff (effective inv_freq_j = base^(-2j/d)/ff[j] = the HF-yarn remapped frequency, verified vs Qwen3RotaryEmbedding to 1.6e-7), mscale = attention_scaling (0.1*ln(factor)+1) applied to q/k post-rope — cos/sin scaling distributes onto the rotated vector exactly. None = plain rope (gemma/z-lab drafters).

dflash2: Option<Dflash2Head>

DFlash2 head (z-lab DFlash2DraftModel, DFLASH2-EVAL-20260820.md): grouped dynamic causal convs around EVERY sublayer + the candidate path selector that replaces the markov chain. A DISTINCT semantic program from the DSpark head (no-generic-support law): present iff config architectures names DFlash2DraftModel, and then ALL 23 family tensors are REQUIRED — loading the 58 backbone tensors alone computes an untrained model (the census trap).

Implementations

impl DflashDraft

pub fn load(e: &Engine, dir: &Path) -> Result<Self, Box<dyn Error>>

Load the backbone-only checkpoint dir (config.json + model.safetensors, bf16). Config scalars ride a minimal extractor (no json dep in-tree — HfConfig precedent).

pub fn d2_conv_prepare( &self, e: &Engine, conv: &Dflash2Conv, xn: &CudaSlice<f32>, rows: usize, ) -> Result<(CudaSlice<f32>, CudaSlice<f32>), Box<dyn Error>>

DFlash2 conv prepare (reference GroupedDynamicCausalConv.prepare): projects the pre-conv rows to BOTH dynamic kernels, convolves the rows with base half 0

• dyn half 0, and returns (convolved rows, the dyn projection) — finish reuses the SAME projection’s half 1. Block-local causal shift (row 0 zero-pads).

pub fn d2_conv_finish( &self, e: &Engine, conv: &Dflash2Conv, y: &CudaSlice<f32>, dyn_: &CudaSlice<f32>, rows: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

DFlash2 conv finish: convolves the sublayer OUTPUT rows with base half 1 + dyn half 1 (dyn from the matching prepare).

pub fn dflash2_propose_greedy( &self, e: &Engine, dl: &CudaSlice<f32>, rows: &CudaSlice<f32>, nd: usize, n_vocab: usize, anchor: u32, ) -> Result<Vec<u32>, Box<dyn Error>>

DFlash2 proposal (reference DFlash2DraftModel.propose, greedy arm): device top-k over the draft logits + the rank-r hidden projection, ONE small dtoh (~nd*(2k+rank) floats — the same per-round sync slot the markov chain’s token readback occupies), then the host codebook walk. Returns the nd drafted tokens (mask-fill rows 1..b-1; the anchor row is not a draft).

pub fn ctx_features( &self, e: &Engine, taps: &CudaSlice<f32>, t: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

FIRST-LIGHT forward (oracle contract): full non-causal attention over [ctx_features ; block], NO draft KV cache, NO sliding window (the oracle bypasses the reference mask machinery the same way — window/caching land in the round arm).

target_hidden: [ctx, n_taps*hidden] (f32, device) — raw tapped states. noise_emb: [block, hidden] — target embed rows for [accepted, MASK x b-1]. pos: absolute positions for ctx rows THEN block rows (ctx+block i32). Returns final normed hidden [block, hidden] (feed target lm_head for draft logits). ctx features for t tapped rows: hidden_norm(fc(taps)) — the drafter’s context representation, cacheable across rounds (append-only in committed-token order).

pub fn forward( &self, e: &Engine, target_hidden: &CudaSlice<f32>, noise_emb: &CudaSlice<f32>, pos: &[i32], ctx: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

pub fn forward_block( &self, e: &Engine, ctx_f: &CudaSlice<f32>, noise_emb: &CudaSlice<f32>, pos: &[i32], ctx: usize, ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Block forward over PRECOMPUTED ctx features (the round arm’s entry: features are cached across rounds; only the block work repeats).

impl DflashDraft

pub fn ingest_ctx( &self, e: &Engine, kv: &mut DflashKv, feats: &CudaSlice<f32>, pos_new: &[i32], t: usize, ) -> Result<(), Box<dyn Error>>

Ingest t NEW ctx-feature rows (committed order, absolute positions pos_new) into the draft KV: per layer k/v projections + k head-norm + rope, appended at kv.len.

pub fn forward_round( &self, e: &Engine, kv: &mut DflashKv, noise_emb: &CudaSlice<f32>, pos_block: &[i32], ) -> Result<CudaSlice<f32>, Box<dyn Error>>

Block forward over the CACHED ctx KV: only the 16 block rows are projected per layer; block K/V land transiently at kv[len..len+b]. Bit-class-identical to forward_block (same kernels, same per-row programs; ONLY the ctx K/V recompute is cached).