Struct DflashScratch 

pub struct DflashScratch {

    pub stream_buf: DevicePtr,
    pub norm_buf: DevicePtr,
    pub q_buf: DevicePtr,
    pub k_buf: DevicePtr,
    pub v_buf: DevicePtr,
    pub attn_out: DevicePtr,
    pub mlp_intermediate: DevicePtr,
    pub mlp_up: DevicePtr,
    pub stream_acc: DevicePtr,
    pub fc_proj: DevicePtr,
    pub fused_kv_out: DevicePtr,
    pub slot_mapping_dev: DevicePtr,
    pub option_b_indirect_args_dev: DevicePtr,
    pub draft_tokens_host_pinned: AtomicPtr<u8>,
    pub draft_tokens_event: u64,
    pub logits: DevicePtr,
    pub draft_tokens_dev: DevicePtr,
    pub position_ids: DevicePtr,
    pub markov_embed: DevicePtr,
    pub markov_bias: DevicePtr,
    pub conf_out: DevicePtr,
    pub conv_dyn: DevicePtr,
    pub conv_tmp: DevicePtr,
    pub sel_vals: DevicePtr,
    pub sel_idx: DevicePtr,
    pub sel_hproj: DevicePtr,
}

Per-step scratch buffers for the γ-block forward.

Sized for n_attn_slots = ctx_window + γ rows, where ctx_window is the max number of past target positions the drafter attends to per step. The first ctx_window slots hold post-fc projected target context (K/V only — Q is zero-padded); the next γ slots hold the noise tokens.

At γ=16 and ctx_window=γ=16: 32 rows × 2048 BF16 × ~10 buffers = ~1.3 MB per head. lm_head logits buffer is the largest single alloc: 32 × 248320 × 2 = 15 MB.

Fields

stream_buf: DevicePtr

norm_buf: DevicePtr

q_buf: DevicePtr

k_buf: DevicePtr

v_buf: DevicePtr

attn_out: DevicePtr

mlp_intermediate: DevicePtr

mlp_up: DevicePtr

stream_acc: DevicePtr

fc_proj: DevicePtr

[ctx_window, draft_hidden] BF16 — fc-projected + hidden_norm’d ctx for the most recent ctx_window target positions.

fused_kv_out: DevicePtr

Phase 2 (Option B) scratch for precompute_ctx_kv: fused KV GEMM output, shape [max_new_ctx, L * 2 * kv_dim] BF16. max_new_ctx = ctx_window (worst case: first propose runs precompute over the entire prefix).

slot_mapping_dev: DevicePtr

Phase 2 scratch: i32 slot mapping for the per-layer reshape_and_cache calls. Sized [ctx_window].

option_b_indirect_args_dev: DevicePtr

Phase 5 (CUDA graph) scratch: 8 bytes ([u32 kv_len, u32 q_offset]) holding the per-call dynamic values that the indirect paged-attention kernel reads at entry. Host writes via copy_h2d BEFORE entering the captured region so the graph itself sees a stable device pointer.

draft_tokens_host_pinned: AtomicPtr<u8>

Phase E.2: pinned host buffer (γ × 4 bytes) for the per-propose draft-token D2H copy. Allocated once at construction via gpu.alloc_host_pinned; the async D2H lands here without touching the system pageable allocator each call.

Wrapped in AtomicPtr to keep DflashScratch: Send + Sync (the proposer is stored as Arc<dyn DraftProposer> which requires both auto-traits). Reads via Ordering::Relaxed are safe: the pointer itself never changes after construction; we only need atomic access for the Send/Sync bound, not for any actual concurrency.

draft_tokens_event: u64

Phase E.2: CUDA event recorded against the draft-tokens D2H so the host can block on completion just before reading the pinned buffer, without a full cuStreamSynchronize. Created once at construction.

logits: DevicePtr

draft_tokens_dev: DevicePtr

position_ids: DevicePtr

[ctx_window + γ] i32 positions. First ctx_window are historical target positions (decoded indices); last γ are the to-be-predicted noise positions.

markov_embed: DevicePtr

DSpark Markov scratch: [1, markov_rank] BF16 latent for the prev-token gather (markov_w1[prev]). DevicePtr(0) when the drafter has no Markov head.

markov_bias: DevicePtr

DSpark Markov scratch: [vocab] BF16 full-vocab bias (markov_w2 @ markov_embed), residual-added onto one logits row per sequential step. DevicePtr(0) when no Markov head.

conf_out: DevicePtr

DSpark confidence scratch: [γ] BF16 per-row acceptance logits (AcceptRatePredictor output). Read back host-side after the draft-token D2H to pick the confident prefix length. DevicePtr(0) when the drafter has no confidence head.

conv_dyn: DevicePtr

[γ, 2*kernel*groups] BF16 — dynamic conv kernels for one conv site (kernel_projection GEMM output at prepare; the finish application reads its slice after the sublayer). Reused sequentially by both conv sites of every layer.

conv_tmp: DevicePtr

[γ, hidden] BF16 — convolved-hidden staging (prepare writes here, the sublayer GEMMs read from here; finish stages here before the residual add).

sel_vals: DevicePtr

[γ, 16] f32 — selector top-16 unary logits per row.

sel_idx: DevicePtr

[γ, 16] u32 — selector top-16 candidate token ids per row.

sel_hproj: DevicePtr

[γ, selector_rank] BF16 — H(h_t) context-gate projections.