Struct Dflash2Head 

pub struct Dflash2Head {
    pub attn_conv: Vec<Dflash2Conv>,
    pub mlp_conv: Vec<Dflash2Conv>,
    pub hidden_proj: GpuTensor,
    pub pred_codebook: Vec<u8>,
    pub succ_codebook: Vec<u8>,
    pub rank: usize,
    pub top_k: usize,
    pub conv_k: usize,
    pub group_size: usize,
    pub vocab: usize,
}

Fields

attn_conv: Vec<Dflash2Conv>

mlp_conv: Vec<Dflash2Conv>

hidden_proj: GpuTensor

candidate_selector.hidden_projection.weight [rank, hidden].

pred_codebook: Vec<u8>

Codebooks [V, rank] raw bf16, HOST-resident: the walk gathers ~1+16 rows per draft slot (~70KB/round) — host math beside the round’s existing chain dtoh, no device residency for 2x127MB tables. Checkpoint quirk: stored WITHOUT the .weight suffix (reference from_pretrained installs a key_mapping).

succ_codebook: Vec<u8>

rank: usize

top_k: usize

conv_k: usize

group_size: usize

vocab: usize

Implementations

impl Dflash2Head

pub fn walk_greedy( &self, unary: &[f32], cand: &[u32], hproj: &[f32], anchor: u32, nd: usize, ) -> Vec<u32>

Greedy selector walk over this head’s codebooks — see dflash2_walk_greedy.

pub fn walk_sampled( &self, unary: &[f32], cand: &[u32], hproj: &[f32], anchor: u32, nd: usize, temp: f32, uniforms: &mut dyn FnMut() -> f32, ) -> (Vec<u32>, Vec<f32>, Vec<f32>)

Sampled (T>0) selector walk — see dflash2_walk_sampled.