metrale_model_arch/dflash_head/
state.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `DflashScratch`, a drafter head's device scratch, and
4//! `DflashProposerState`, its per-sequence state.
5//!
6//! Owner: model-arch (DFlash drafter).
7//! Invariants: none beyond the types.
8
9use super::*;
10
11/// 2026-09-25: Device scratch for the γ-block forward, allocated once per head
12/// in `from_weights.rs`. `nb` below is the head's `max_batch`.
13pub struct DflashScratch {
14    pub stream_buf: DevicePtr,
15    pub norm_buf: DevicePtr,
16    pub q_buf: DevicePtr,
17    pub k_buf: DevicePtr,
18    pub v_buf: DevicePtr,
19    pub attn_out: DevicePtr,
20    pub mlp_intermediate: DevicePtr,
21    pub mlp_up: DevicePtr,
22    pub stream_acc: DevicePtr,
23    /// 2026-09-25: `[ctx_window, hidden]` BF16: ctx rows after `fc` and
24    /// `hidden_norm`.
25    pub fc_proj: DevicePtr,
26    /// 2026-09-25: `[ctx_window, L * 2 * kv_dim]` BF16: the ctx precompute's KV
27    /// GEMM output.
28    pub fused_kv_out: DevicePtr,
29    /// 2026-09-25: `i64[ctx_window]` paged-cache slot indices for
30    /// `reshape_and_cache`.
31    pub slot_mapping_dev: DevicePtr,
32    /// 2026-09-25: `[PRECOMPUTE_BATCH_ROWS, L_t * h_t]` BF16 staging for the
33    /// batched ctx precompute: the uncommitted ctx rows of several sequences,
34    /// contiguous.
35    pub precompute_in: DevicePtr,
36    /// 2026-09-25: `nb` triples `[u32 kv_len, u32 q_offset, u32 q_rope_pos]`,
37    /// 12 bytes each, one per band, read by the indirect paged-attention kernel
38    /// at entry. The host writes them before the captured region, so a captured
39    /// graph sees a fixed pointer.
40    pub option_b_indirect_args_dev: DevicePtr,
41    /// 2026-09-25: Pinned host buffer, `nb * gamma * 4` bytes, for the
42    /// draft-token D2H copy; allocated once in `from_weights.rs`. `AtomicPtr`
43    /// keeps `DflashScratch` `Send + Sync`. Nothing stores to it after
44    /// construction, so `Relaxed` loads are enough.
45    pub draft_tokens_host_pinned: std::sync::atomic::AtomicPtr<u8>,
46    /// 2026-09-25: Event recorded after the draft-token D2H copy; the host waits
47    /// on it before reading the pinned buffer. Created once in `from_weights.rs`.
48    pub draft_tokens_event: u64,
49    pub logits: DevicePtr,
50    pub draft_tokens_dev: DevicePtr,
51    /// 2026-09-25: i32 RoPE positions: the `eff_ctx` ctx rows, then each
52    /// sequence's block rows from that sequence's position.
53    pub position_ids: DevicePtr,
54    /// 2026-09-25: DSpark Markov scratch: `[1, markov_rank]` BF16 latent for the
55    /// prev-token gather (`markov_w1[prev]`). `DevicePtr(0)` when the
56    /// drafter has no Markov head.
57    pub markov_embed: DevicePtr,
58    /// 2026-09-25: DSpark Markov scratch: `[vocab]` BF16 bias
59    /// (`markov_w2 @ markov_embed`), residual-added onto one logits row
60    /// per sequential step. `DevicePtr(0)` when no Markov head.
61    pub markov_bias: DevicePtr,
62    /// 2026-09-25: DSpark confidence scratch: `[γ]` BF16 per-row acceptance
63    /// logits, read back on the host after the draft-token D2H to pick how many
64    /// drafts to keep. `DevicePtr(0)` when the drafter has no confidence head.
65    pub conf_out: DevicePtr,
66
67    // 2026-09-25: DFlash2 scratch, `DevicePtr(0)` when the checkpoint has no
68    // selector.
69    /// 2026-09-25: `[nb * γ, 2 * conv_kernel_size * (hidden / conv_group_size)]`
70    /// BF16 dynamic conv kernels, reused by each conv site in turn.
71    pub conv_dyn: DevicePtr,
72    /// 2026-09-25: `[nb * γ, hidden]` BF16 staging for the convolved hidden.
73    pub conv_tmp: DevicePtr,
74    /// 2026-09-25: `[nb * γ, 16]` f32: the selector's top-16 logits per row.
75    pub sel_vals: DevicePtr,
76    /// 2026-09-25: `[nb * γ, 16]` u32: the selector's top-16 token ids per row.
77    pub sel_idx: DevicePtr,
78    /// 2026-09-25: `[nb * γ, selector_rank]` BF16: the selector's hidden
79    /// projection.
80    pub sel_hproj: DevicePtr,
81}
82
83/// 2026-09-25: Per-sequence DFlash drafter state. One block table serves every
84/// drafter layer: the paged cache has a pool per layer, addressed by the same
85/// slots.
86pub struct DflashProposerState {
87    /// 2026-09-25: Host copy of the drafter block table.
88    pub block_table: Vec<u32>,
89    pub seq_len: usize,
90    /// 2026-09-25: Drafts the last propose returned; `after_verify` and
91    /// `free_state` reset it to 0.
92    pub last_num_drafted: usize,
93    pub prefill_done: bool,
94    /// 2026-09-25: `max_ctx_len` rows of `ctx_slot_bytes`: captured target
95    /// hiddens, one row per ctx position, the first `ctx_len` valid. Model-engine
96    /// appends rows and slides the window when it is full; `propose_drafts`
97    /// appends while there is room.
98    pub ctx_hidden_acc: DevicePtr,
99    /// 2026-09-25: Valid rows in `ctx_hidden_acc`, at most `max_ctx_len`.
100    pub ctx_len: usize,
101    pub last_num_accepted: usize,
102    /// 2026-09-25: Model-engine has already appended the latest capture; the
103    /// next `propose_drafts` skips its own append and clears the flag.
104    pub skip_next_decode_append: bool,
105    /// 2026-09-25: Rows `ctx_hidden_acc` holds (the `window` of
106    /// `alloc_state_windowed`).
107    pub max_ctx_len: usize,
108    /// 2026-09-25: Bytes per `ctx_hidden_acc` row:
109    /// `target_layer_ids.len() * target_hidden_size * 2`.
110    pub ctx_slot_bytes: usize,
111
112    /// 2026-09-25: Device copy of `block_table`, made once when the first
113    /// propose allocates the table; `None` before that and after `free_state`.
114    pub block_table_dev: Option<DevicePtr>,
115    /// 2026-09-25: Ctx rows of the paged cache the γ block attends to;
116    /// `propose_drafts` sets it to `ctx_len`.
117    pub ctx_count_drafter: usize,
118    /// 2026-09-25: Slots the allocated block table covers
119    /// (`block_table.len() * BLOCK_SIZE`).
120    pub max_ctx_count_drafter: usize,
121    /// 2026-09-25: Ctx rows `[0, ctx_committed)` whose K/V is already in the
122    /// paged cache; the next precompute starts here. Model-engine's window
123    /// slides reset it to 0.
124    pub ctx_committed: usize,
125    /// 2026-09-25: One RoPE position per ctx row (`len() == ctx_len`), stamped
126    /// when the row is appended and kept when the window slides. Model-engine's
127    /// `update_dflash_ctx_len_after_prefill` seeds prefill row i with i.
128    pub ctx_positions: Vec<i32>,
129}
130
131impl ProposerState for DflashProposerState {
132    fn as_any(&self) -> &dyn Any {
133        self
134    }
135    fn as_any_mut(&mut self) -> &mut dyn Any {
136        self
137    }
138}
metrale_model_arch/dflash_head/state.rs

metrale_model_arch/dflash_head/
state.rs
