hf2q 0.1.1

//! DFlash target hidden-state capture buffer (ADR-030 Phase 4).
//!
//! Defines the buffer layout + helpers used by the spec-decode
//! orchestrator to consume per-layer hidden states from the TARGET
//! model's forward pass at `cfg.target_layer_ids` positions.
//!
//! ## Capture buffer layout
//!
//! The capture buffer is a single flat `[f32]` slice of length
//! `num_capture_layers * seq_len * hidden_size`. Layout is row-major
//! by (capture_layer_idx, token_pos, dim):
//!
//! ```text
//!   captured[(capture_layer_idx * seq_len + token_pos) * hidden_size + dim]
//! ```
//!
//! `capture_layer_idx` indexes into `target_layer_ids` (0..N), NOT
//! into the target's full layer count (0..30 for gemma-4).
//!
//! ## Concat semantics for drafter input
//!
//! `dispatch_dflash_fc` expects `target_hidden_concat` shape
//! `[ctx_seq_len, num_capture_layers * hidden_size]` — i.e., for each
//! token position, the N captured hidden vectors are concatenated
//! along the last axis. The capture buffer layout above gives the
//! TRANSPOSE of this. The orchestrator must permute
//! `[capture_layer, token, dim]` → `[token, capture_layer, dim]`
//! before feeding to `dispatch_dflash_fc`.
//!
//! The permute is a one-time per-step cost; for our shapes
//! (num_capture_layers=6, seq_len ≤ block_size + ctx_chunk ≈ 16,
//! hidden_size=2816) the buffer is small (~270KB).
//!
//! ## Why a flat slice (not a Vec<Vec<f32>>)
//!
//! - Single contiguous allocation = one upload to GPU
//! - Caller controls allocation; the capture can write into a
//!   pre-allocated pool buffer to avoid per-call alloc
//! - Easy to pass to forward_prefill_batched via &mut [f32]

use super::config::DFlashConfig;
use anyhow::{anyhow, Result};

/// Strict shape descriptor + flat buffer for target hidden capture.
///
/// Lifetime parameter: the inner buffers borrow from caller-supplied
/// storage; `PrefillCapture` itself is a typed view.
pub struct PrefillCapture<'a> {
    /// Layer indices (in target's 0..num_target_layers numbering) at
    /// which to capture pf_hidden.
    pub target_layer_ids: &'a [usize],
    /// Captured hidden states: flat F32, length
    /// `target_layer_ids.len() * seq_len * hidden_size`.
    /// Layout (row-major): `[(layer_idx_in_capture, token_pos, dim)]`.
    pub hidden_output: &'a mut [f32],
    /// Per-position argmaxes from the target's LM head. Length =
    /// seq_len. Pre-allocated by caller. When None, the legacy
    /// single-row argmax return is preserved.
    pub per_position_argmaxes: Option<&'a mut [u32]>,
}

/// Trim a `DFlashCaptureSession`'s hidden_output down to the first
/// `new_seq_len` positions per captured layer.
///
/// Used by the multi-round orchestrator after `step_round_from_argmaxes`
/// reports `accept_count + 1` accepted tokens. Per Python
/// `model_mlx.py:567` `hidden = hidden[:, :accepted + 1, :]`, the
/// drafter's ctx K/V for the next round must reflect ONLY the
/// target-accepted positions. Without this trim, the drafter would
/// see all K+1 verify positions as ctx, including rejected ones —
/// breaks Phase 4 coherence.
///
/// This is an in-place buffer rewrite (allocates a smaller Vec,
/// copies slabs, swaps). Cost is proportional to
/// `target_layer_ids.len() * new_seq_len * hidden_size` floats.
pub fn trim_capture_to(
    session: &mut super::hidden_capture::DFlashCaptureSession,
    new_seq_len: usize,
) {
    if new_seq_len >= session.seq_len {
        return;
    }
    let n_layers = session.target_layer_ids.len();
    let hs = session.hidden_size;
    let old_seq_len = session.seq_len;
    let mut new_buf = vec![0.0f32; n_layers * new_seq_len * hs];
    for l in 0..n_layers {
        for t in 0..new_seq_len {
            let src = (l * old_seq_len + t) * hs;
            let dst = (l * new_seq_len + t) * hs;
            new_buf[dst..dst + hs].copy_from_slice(&session.hidden_output[src..src + hs]);
        }
    }
    session.hidden_output = new_buf;
    session.seq_len = new_seq_len;
    if let Some(pa) = session.per_position_argmaxes.as_mut() {
        pa.truncate(new_seq_len);
    }
}

impl<'a> PrefillCapture<'a> {
    /// Validate that the buffer sizes match the declared seq_len +
    /// hidden_size + num_capture_layers. Cheap check at call boundary
    /// to avoid silent buffer overruns deep in the layer loop.
    pub fn validate(&self, seq_len: usize, hidden_size: usize) -> Result<()> {
        let expected = self.target_layer_ids.len() * seq_len * hidden_size;
        if self.hidden_output.len() != expected {
            return Err(anyhow!(
                "PrefillCapture: hidden_output len {} != target_layer_ids({}) * seq_len({}) * hidden_size({}) = {}",
                self.hidden_output.len(),
                self.target_layer_ids.len(),
                seq_len,
                hidden_size,
                expected
            ));
        }
        if let Some(ref pa) = self.per_position_argmaxes {
            if pa.len() != seq_len {
                return Err(anyhow!(
                    "PrefillCapture: per_position_argmaxes len {} != seq_len {}",
                    pa.len(),
                    seq_len
                ));
            }
        }
        // Validate target_layer_ids: strictly increasing, all in bounds
        // (caller's responsibility — we can't know num_target_layers here).
        for w in self.target_layer_ids.windows(2) {
            if w[0] >= w[1] {
                return Err(anyhow!(
                    "PrefillCapture: target_layer_ids must be strictly increasing; got {:?}",
                    self.target_layer_ids
                ));
            }
        }
        Ok(())
    }

    /// Compute the byte offset (in `hidden_output`) where the captured
    /// row for `(capture_layer_idx, token_pos)` begins. `hidden_size`
    /// elements follow contiguously.
    pub fn offset_for(
        capture_layer_idx: usize,
        token_pos: usize,
        seq_len: usize,
        hidden_size: usize,
    ) -> usize {
        (capture_layer_idx * seq_len + token_pos) * hidden_size
    }

    /// Write one (capture_layer_idx, all token positions) slab into
    /// `hidden_output`. Used by the target's layer-loop end hook:
    /// pf_hidden contains the layer's `[seq_len, hidden_size]` output;
    /// this copies all rows in one memcpy if `pf_hidden_data` is
    /// row-major contiguous.
    pub fn write_layer_slab(
        &mut self,
        capture_layer_idx: usize,
        pf_hidden_data: &[f32],
        seq_len: usize,
        hidden_size: usize,
    ) -> Result<()> {
        let slab_len = seq_len * hidden_size;
        if pf_hidden_data.len() != slab_len {
            return Err(anyhow!(
                "PrefillCapture::write_layer_slab: pf_hidden_data len {} != seq_len({}) * hidden_size({}) = {}",
                pf_hidden_data.len(), seq_len, hidden_size, slab_len
            ));
        }
        let start = Self::offset_for(capture_layer_idx, 0, seq_len, hidden_size);
        let end = start + slab_len;
        if end > self.hidden_output.len() {
            return Err(anyhow!(
                "PrefillCapture::write_layer_slab: write out of bounds (start {}, end {}, buf {})",
                start,
                end,
                self.hidden_output.len()
            ));
        }
        self.hidden_output[start..end].copy_from_slice(pf_hidden_data);
        Ok(())
    }

    /// Permute `[capture_layer, token, dim]` → `[token, capture_layer, dim]`
    /// into a freshly-allocated Vec<f32>. Used by the orchestrator to
    /// build the `target_hidden_concat` argument for the drafter's
    /// `dispatch_dflash_fc` (which expects `[seq_len, num_layers * hidden]`
    /// flat = `[seq_len, num_layers, hidden]` with last 2 dims merged).
    pub fn permute_to_concat(&self, seq_len: usize, hidden_size: usize) -> Vec<f32> {
        let n_layers = self.target_layer_ids.len();
        let mut out = vec![0.0f32; seq_len * n_layers * hidden_size];
        for layer_idx in 0..n_layers {
            for t in 0..seq_len {
                let src = (layer_idx * seq_len + t) * hidden_size;
                let dst = (t * n_layers + layer_idx) * hidden_size;
                out[dst..dst + hidden_size]
                    .copy_from_slice(&self.hidden_output[src..src + hidden_size]);
            }
        }
        out
    }
}

/// Compute total flat buffer length for the hidden capture.
pub fn capture_buffer_len(cfg: &DFlashConfig, seq_len: usize) -> usize {
    cfg.target_layer_ids.len() * seq_len * cfg.hidden_size
}

/// Extract drafter's `target_layer_ids` slabs from a COMBINED capture
/// buffer (which may include additional layers like the final layer for
/// argmax purposes) and produce the drafter-input concat layout
/// `[seq_len, drafter_n_layers * hidden_size]`.
///
/// ## When this is used
///
/// The multi-round orchestrator wrapper installs a SINGLE capture
/// session per target forward call, with `target_layer_ids` = drafter's
/// `[1, 6, 11, 17, 22, 27]` ∪ `[final_layer_idx]`. After the forward,
/// it needs to produce TWO things from the same `hidden_output`:
/// - drafter's input: `[seq_len, drafter_n_layers * hs]` (this helper)
/// - target's argmax input: `[seq_len, hs]` (final layer slab only)
///
/// ## Layout
///
/// Input: `hidden_output` layout `[combined_capture_layer, seq_len, hs]`
/// row-major where `combined_capture_layer_ids` is the merged set.
///
/// Output: `[seq_len, drafter_n_layers, hs]` flat row-major (= same as
/// `[seq_len, drafter_n_layers * hs]` after final-axis merge), matching
/// what `dispatch_dflash_fc` expects as `target_hidden_concat`.
pub fn extract_drafter_concat(
    hidden_output: &[f32],
    combined_capture_layer_ids: &[usize],
    drafter_target_layer_ids: &[usize],
    seq_len: usize,
    hidden_size: usize,
) -> Result<Vec<f32>> {
    let drafter_n = drafter_target_layer_ids.len();
    let total_in = combined_capture_layer_ids.len() * seq_len * hidden_size;
    if hidden_output.len() != total_in {
        return Err(anyhow!(
            "extract_drafter_concat: hidden_output len {} != combined_layers({}) * seq_len({}) * hs({}) = {}",
            hidden_output.len(), combined_capture_layer_ids.len(), seq_len, hidden_size, total_in
        ));
    }

    // Map each drafter layer id to its index in combined_capture_layer_ids
    let mut drafter_to_combined: Vec<usize> = Vec::with_capacity(drafter_n);
    for &dl in drafter_target_layer_ids {
        match combined_capture_layer_ids.iter().position(|&c| c == dl) {
            Some(idx) => drafter_to_combined.push(idx),
            None => {
                return Err(anyhow!(
                    "extract_drafter_concat: drafter target_layer_id {} not in combined capture set {:?}",
                    dl, combined_capture_layer_ids
                ));
            }
        }
    }

    // Output layout: for each (t, drafter_l, d) put hidden_output at
    // (combined_idx, t, d) into out at ((t * drafter_n + drafter_l), d).
    let mut out = vec![0.0f32; seq_len * drafter_n * hidden_size];
    for (drafter_l, &combined_idx) in drafter_to_combined.iter().enumerate() {
        for t in 0..seq_len {
            let src = (combined_idx * seq_len + t) * hidden_size;
            let dst = (t * drafter_n + drafter_l) * hidden_size;
            out[dst..dst + hidden_size].copy_from_slice(&hidden_output[src..src + hidden_size]);
        }
    }
    Ok(out)
}

/// Extract the final-layer slab `[seq_len, hidden_size]` from a
/// combined capture buffer. Used by the multi-round orchestrator's
/// per-round argmax step.
pub fn extract_final_layer_slab(
    hidden_output: &[f32],
    combined_capture_layer_ids: &[usize],
    final_layer_idx: usize,
    seq_len: usize,
    hidden_size: usize,
) -> Result<Vec<f32>> {
    let final_combined_idx = combined_capture_layer_ids
        .iter()
        .position(|&c| c == final_layer_idx)
        .ok_or_else(|| {
            anyhow!(
                "extract_final_layer_slab: final_layer_idx {} not in combined capture set {:?}",
                final_layer_idx,
                combined_capture_layer_ids
            )
        })?;
    let start = final_combined_idx * seq_len * hidden_size;
    let end = start + seq_len * hidden_size;
    if end > hidden_output.len() {
        return Err(anyhow!(
            "extract_final_layer_slab: end offset {} > buffer len {}",
            end,
            hidden_output.len()
        ));
    }
    Ok(hidden_output[start..end].to_vec())
}

/// Owned variant of `PrefillCapture` for installation as an
/// `MlxModelWeights` field. Avoids the lifetime parameter (Vec-backed)
/// so the capture session can live across the forward call without
/// borrow-checker gymnastics.
///
/// Set `seq_len` and `hidden_size` at install time so the layer-loop
/// hook can compute offsets without re-deriving them. The
/// `target_layer_ids` are checked once at install via `validate()`.
#[derive(Debug, Default, Clone)]
pub struct DFlashCaptureSession {
    pub target_layer_ids: Vec<usize>,
    pub hidden_output: Vec<f32>,
    pub per_position_argmaxes: Option<Vec<u32>>,
    /// Cached so the hook doesn't need to re-derive.
    pub seq_len: usize,
    pub hidden_size: usize,
}

impl DFlashCaptureSession {
    /// Allocate a capture session with the given layout. Buffers are
    /// pre-sized to `target_layer_ids.len() * seq_len * hidden_size`
    /// (and `seq_len` for argmaxes if requested).
    pub fn new(
        target_layer_ids: Vec<usize>,
        seq_len: usize,
        hidden_size: usize,
        with_argmaxes: bool,
    ) -> Self {
        let hidden_output = vec![0.0f32; target_layer_ids.len() * seq_len * hidden_size];
        let per_position_argmaxes = if with_argmaxes {
            Some(vec![0u32; seq_len])
        } else {
            None
        };
        Self {
            target_layer_ids,
            hidden_output,
            per_position_argmaxes,
            seq_len,
            hidden_size,
        }
    }

    /// Find the capture index (0..num_capture_layers) for a given
    /// target layer index, or None if this layer is not captured.
    pub fn capture_index_for(&self, target_layer_idx: usize) -> Option<usize> {
        self.target_layer_ids
            .iter()
            .position(|&i| i == target_layer_idx)
    }

    /// Write one layer's full `[seq_len, hidden_size]` row-major slab
    /// from a target hidden buffer into the capture. Used by the
    /// layer-loop hook in `forward_prefill_batched`.
    pub fn write_layer_slab(
        &mut self,
        capture_layer_idx: usize,
        pf_hidden_data: &[f32],
    ) -> Result<()> {
        let slab_len = self.seq_len * self.hidden_size;
        if pf_hidden_data.len() != slab_len {
            return Err(anyhow!(
                "DFlashCaptureSession::write_layer_slab: pf_hidden_data len {} != seq_len({}) * hidden_size({}) = {}",
                pf_hidden_data.len(), self.seq_len, self.hidden_size, slab_len
            ));
        }
        let start =
            PrefillCapture::offset_for(capture_layer_idx, 0, self.seq_len, self.hidden_size);
        let end = start + slab_len;
        if end > self.hidden_output.len() {
            return Err(anyhow!(
                "DFlashCaptureSession::write_layer_slab: write out of bounds (start {}, end {}, buf {})",
                start, end, self.hidden_output.len()
            ));
        }
        self.hidden_output[start..end].copy_from_slice(pf_hidden_data);
        Ok(())
    }

    /// Convenience: borrow as a `PrefillCapture` view for consumers
    /// that need the borrowed-buffer API.
    pub fn as_view(&mut self) -> PrefillCapture<'_> {
        PrefillCapture {
            target_layer_ids: &self.target_layer_ids,
            hidden_output: &mut self.hidden_output,
            per_position_argmaxes: self.per_position_argmaxes.as_deref_mut(),
        }
    }
}

/// ADR-030 iter-76 — append `n_committed` accepted positions from
/// `verify_captured` onto `prior_captured`, returning a fresh
/// `DFlashCaptureSession` whose `seq_len = prior.seq_len + n_committed`.
///
/// Used by the Option A orchestrator path: each round's verify pass
/// captures K+1 positions starting at `start_pos = output.len() - 1`.
/// After accept-prefix, the first `n_committed` of those positions
/// are accepted into the committed prefix and need to be merged into
/// the persistent prior_captured slab for the next round's drafter
/// input.
///
/// Layout: both sessions store `hidden_output` as flat
/// `[num_target_layers, seq_len, hidden_size]` row-major.  The
/// returned session's seq_len = prior.seq_len + n_committed; for each
/// capture layer, positions `[0..prior.seq_len)` come from `prior`
/// and positions `[prior.seq_len..prior.seq_len + n_committed)` come
/// from `verify_captured`'s first `n_committed` rows.
///
/// Both sessions must have IDENTICAL `target_layer_ids` and
/// `hidden_size`.  Panics in debug mode otherwise.
pub fn append_capture_positions(
    prior: &DFlashCaptureSession,
    verify_captured: &DFlashCaptureSession,
    n_committed: usize,
) -> Result<DFlashCaptureSession> {
    debug_assert_eq!(prior.target_layer_ids, verify_captured.target_layer_ids);
    debug_assert_eq!(prior.hidden_size, verify_captured.hidden_size);
    if n_committed > verify_captured.seq_len {
        return Err(anyhow!(
            "append_capture_positions: n_committed ({}) > verify_captured.seq_len ({})",
            n_committed,
            verify_captured.seq_len
        ));
    }
    let hs = prior.hidden_size;
    let num_layers = prior.target_layer_ids.len();
    let prior_seq = prior.seq_len;
    let new_seq = prior_seq + n_committed;
    let mut new_hidden = vec![0.0f32; num_layers * new_seq * hs];
    for layer in 0..num_layers {
        let new_layer_base = layer * new_seq * hs;
        // (a) Copy prior's positions [0..prior_seq) for this layer.
        let prior_layer_base = layer * prior_seq * hs;
        new_hidden[new_layer_base..new_layer_base + prior_seq * hs].copy_from_slice(
            &prior.hidden_output[prior_layer_base..prior_layer_base + prior_seq * hs],
        );
        // (b) Append verify_captured's positions [0..n_committed) for this layer.
        let verify_layer_base = layer * verify_captured.seq_len * hs;
        let new_extended = new_layer_base + prior_seq * hs;
        new_hidden[new_extended..new_extended + n_committed * hs].copy_from_slice(
            &verify_captured.hidden_output[verify_layer_base..verify_layer_base + n_committed * hs],
        );
    }
    Ok(DFlashCaptureSession {
        target_layer_ids: prior.target_layer_ids.clone(),
        hidden_output: new_hidden,
        per_position_argmaxes: None,
        seq_len: new_seq,
        hidden_size: hs,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::inference::spec_decode::dflash::config::DFlashConfig;

    fn gemma4_cfg() -> DFlashConfig {
        DFlashConfig::from_json_str(super::super::config::tests::GEMMA4_26B_A4B_DFLASH_CONFIG)
            .expect("config parse")
    }

    #[test]
    fn extract_drafter_concat_picks_right_slabs() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Combined capture has 7 layers: drafter's [1,6,11,17,22,27] + final 29
        let combined = vec![1, 6, 11, 17, 22, 27, 29];
        let drafter_ids = vec![1, 6, 11, 17, 22, 27];
        let seq_len = 3usize;
        let hs = 4usize;
        // hidden_output layout: [combined_layer=7][t=3][dim=4]
        // Fill with values = combined_layer_idx * 100 + t * 10 + dim
        let mut hidden = vec![0.0f32; 7 * 3 * 4];
        for cli in 0..7 {
            for t in 0..3 {
                for d in 0..4 {
                    hidden[(cli * 3 + t) * 4 + d] = (cli * 100 + t * 10 + d) as f32;
                }
            }
        }
        let out = extract_drafter_concat(&hidden, &combined, &drafter_ids, seq_len, hs).unwrap();
        // Expected layout: [t=3][drafter_l=6][dim=4]
        // out[(t * 6 + drafter_l) * 4 + d] = hidden[(combined_idx * 3 + t) * 4 + d]
        // For drafter_l=0 (id=1, combined_idx=0): values from cli=0
        // For drafter_l=5 (id=27, combined_idx=5): values from cli=5
        // Combined idx 6 (id=29) is the FINAL layer; should NOT appear in out.
        for t in 0..seq_len {
            for drafter_l in 0..drafter_ids.len() {
                let combined_idx = drafter_l; // since drafter ids are first 6 of combined
                for d in 0..hs {
                    let expected = (combined_idx * 100 + t * 10 + d) as f32;
                    let actual = out[(t * drafter_ids.len() + drafter_l) * hs + d];
                    assert_eq!(actual, expected, "t={t} drafter_l={drafter_l} d={d}");
                }
            }
        }
    }

    #[test]
    fn extract_final_layer_slab_picks_correct_layer() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Combined: [1, 6, 11, 17, 22, 27, 29]
        // Final layer = 29 → combined_idx = 6 (last)
        let combined = vec![1, 6, 11, 17, 22, 27, 29];
        let seq_len = 3usize;
        let hs = 4usize;
        let mut hidden = vec![0.0f32; 7 * 3 * 4];
        // Make final-layer slab (combined_idx=6) all 9.0; rest all 1.0.
        for cli in 0..7 {
            for t in 0..3 {
                for d in 0..4 {
                    hidden[(cli * 3 + t) * 4 + d] = if cli == 6 { 9.0 } else { 1.0 };
                }
            }
        }
        let slab = extract_final_layer_slab(&hidden, &combined, 29, seq_len, hs).unwrap();
        assert_eq!(slab.len(), seq_len * hs);
        for v in slab.iter() {
            assert_eq!(*v, 9.0);
        }
    }

    #[test]
    fn trim_capture_to_compacts_buffer_and_keeps_first_positions() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // 3 layers × 5 positions × 4 dim. Trim to 2 positions.
        let mut sess = DFlashCaptureSession::new(vec![1, 6, 11], 5, 4, true);
        // Fill: hidden[(l * 5 + t) * 4 + d] = l * 100 + t * 10 + d
        for l in 0..3 {
            for t in 0..5 {
                for d in 0..4 {
                    sess.hidden_output[(l * 5 + t) * 4 + d] = (l * 100 + t * 10 + d) as f32;
                }
            }
        }
        if let Some(pa) = sess.per_position_argmaxes.as_mut() {
            for (i, v) in pa.iter_mut().enumerate() {
                *v = i as u32;
            }
        }
        trim_capture_to(&mut sess, 2);
        assert_eq!(sess.seq_len, 2);
        assert_eq!(sess.hidden_output.len(), 3 * 2 * 4);
        // Verify positions 0 + 1 preserved per layer
        for l in 0..3 {
            for t in 0..2 {
                for d in 0..4 {
                    let expected = (l * 100 + t * 10 + d) as f32;
                    let actual = sess.hidden_output[(l * 2 + t) * 4 + d];
                    assert_eq!(actual, expected, "l={l} t={t} d={d}");
                }
            }
        }
        // argmaxes truncated
        assert_eq!(
            sess.per_position_argmaxes.as_ref().map(|v| v.len()),
            Some(2)
        );
    }

    #[test]
    fn trim_capture_to_no_op_when_new_geq_old() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let mut sess = DFlashCaptureSession::new(vec![1, 6], 3, 4, false);
        let orig_len = sess.hidden_output.len();
        for v in sess.hidden_output.iter_mut() {
            *v = 7.0;
        }
        trim_capture_to(&mut sess, 5); // 5 > 3, no-op
        assert_eq!(sess.seq_len, 3);
        assert_eq!(sess.hidden_output.len(), orig_len);
        assert!(sess.hidden_output.iter().all(|&v| v == 7.0));
    }

    #[test]
    fn extract_drafter_concat_errors_on_missing_layer() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // Drafter wants layer 100 which isn't in combined
        let combined = vec![1, 6, 29];
        let drafter_ids = vec![1, 100];
        let mut hidden = vec![0.0f32; 3 * 3 * 4];
        for v in hidden.iter_mut() {
            *v = 1.0;
        }
        let err = extract_drafter_concat(&hidden, &combined, &drafter_ids, 3, 4).unwrap_err();
        assert!(format!("{err}").contains("not in combined"));
    }

    #[test]
    fn capture_buffer_len_matches_drafter_fc_input() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let cfg = gemma4_cfg();
        // Per drafter config: 6 target_layer_ids × hidden_size 2816.
        // For seq_len=4: 6 × 4 × 2816 = 67584 floats.
        assert_eq!(capture_buffer_len(&cfg, 4), 6 * 4 * 2816);
    }

    #[test]
    fn validate_catches_wrong_buffer_size() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let target_layer_ids = vec![1, 6, 11, 17, 22, 27];
        let mut hidden = vec![0.0f32; 100]; // way too small
        let cap = PrefillCapture {
            target_layer_ids: &target_layer_ids,
            hidden_output: &mut hidden,
            per_position_argmaxes: None,
        };
        let err = cap.validate(4, 2816).unwrap_err();
        assert!(format!("{err}").contains("hidden_output len"));
    }

    #[test]
    fn validate_catches_non_monotonic_layer_ids() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let target_layer_ids = vec![1, 11, 6, 17, 22, 27]; // 11 > 6
        let mut hidden = vec![0.0f32; 6 * 4 * 2816];
        let cap = PrefillCapture {
            target_layer_ids: &target_layer_ids,
            hidden_output: &mut hidden,
            per_position_argmaxes: None,
        };
        let err = cap.validate(4, 2816).unwrap_err();
        assert!(format!("{err}").contains("strictly increasing"));
    }

    #[test]
    fn validate_catches_argmax_size_mismatch() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let target_layer_ids = vec![1, 6, 11, 17, 22, 27];
        let mut hidden = vec![0.0f32; 6 * 4 * 2816];
        let mut argmaxes = vec![0u32; 99]; // wrong: should be seq_len=4
        let cap = PrefillCapture {
            target_layer_ids: &target_layer_ids,
            hidden_output: &mut hidden,
            per_position_argmaxes: Some(&mut argmaxes),
        };
        let err = cap.validate(4, 2816).unwrap_err();
        assert!(format!("{err}").contains("per_position_argmaxes"));
    }

    #[test]
    fn offset_for_matches_row_major_layout() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        // (layer=0, token=0): offset 0
        assert_eq!(PrefillCapture::offset_for(0, 0, 4, 2816), 0);
        // (layer=0, token=1): one row in
        assert_eq!(PrefillCapture::offset_for(0, 1, 4, 2816), 2816);
        // (layer=1, token=0): one (seq_len * hidden) slab in
        assert_eq!(PrefillCapture::offset_for(1, 0, 4, 2816), 4 * 2816);
        // (layer=5, token=3): last position in 6×4×2816 buffer
        assert_eq!(
            PrefillCapture::offset_for(5, 3, 4, 2816),
            (5 * 4 + 3) * 2816
        );
    }

    #[test]
    fn write_layer_slab_places_data_correctly() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let target_layer_ids = vec![1, 6, 11, 17, 22, 27];
        let mut hidden = vec![0.0f32; 6 * 4 * 2816];
        let mut cap = PrefillCapture {
            target_layer_ids: &target_layer_ids,
            hidden_output: &mut hidden,
            per_position_argmaxes: None,
        };
        // Write layer 2's slab as ones
        let slab = vec![1.0f32; 4 * 2816];
        cap.write_layer_slab(2, &slab, 4, 2816)
            .expect("write_layer_slab");
        // Verify: layer 2's slab is ones; layer 0/1/3/4/5 still zero
        let layer2_start = 2 * 4 * 2816;
        let layer2_end = 3 * 4 * 2816;
        for i in 0..hidden.len() {
            let expected = if i >= layer2_start && i < layer2_end {
                1.0
            } else {
                0.0
            };
            assert_eq!(hidden[i], expected, "at index {i}");
        }
    }

    /// DFlashCaptureSession: new() pre-allocates correct buffer sizes.
    #[test]
    fn session_new_allocates_correct_sizes() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let sess = DFlashCaptureSession::new(vec![1, 6, 11, 17, 22, 27], 4, 2816, true);
        assert_eq!(sess.hidden_output.len(), 6 * 4 * 2816);
        assert_eq!(
            sess.per_position_argmaxes.as_ref().map(|v| v.len()),
            Some(4)
        );
        assert_eq!(sess.seq_len, 4);
        assert_eq!(sess.hidden_size, 2816);

        // No-argmaxes path
        let sess2 = DFlashCaptureSession::new(vec![1, 6], 8, 128, false);
        assert_eq!(sess2.hidden_output.len(), 2 * 8 * 128);
        assert!(sess2.per_position_argmaxes.is_none());
    }

    #[test]
    fn session_capture_index_for_finds_target_layers() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let sess = DFlashCaptureSession::new(vec![1, 6, 11, 17, 22, 27], 4, 2816, false);
        assert_eq!(sess.capture_index_for(1), Some(0));
        assert_eq!(sess.capture_index_for(11), Some(2));
        assert_eq!(sess.capture_index_for(27), Some(5));
        assert_eq!(sess.capture_index_for(0), None);
        assert_eq!(sess.capture_index_for(5), None);
        assert_eq!(sess.capture_index_for(28), None);
    }

    #[test]
    fn session_write_layer_slab_places_data_at_correct_offset() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let mut sess = DFlashCaptureSession::new(vec![1, 6, 11, 17, 22, 27], 3, 4, false);
        // Layer-2 slab: capture_layer_idx=2; offset = 2 * 3 * 4 = 24
        let slab = vec![5.0f32; 3 * 4];
        sess.write_layer_slab(2, &slab).expect("write");
        for i in 0..sess.hidden_output.len() {
            let expected = if (24..36).contains(&i) { 5.0 } else { 0.0 };
            assert_eq!(sess.hidden_output[i], expected, "i={i}");
        }
    }

    #[test]
    fn session_as_view_borrows_buffers_consistently() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let mut sess = DFlashCaptureSession::new(vec![1, 6], 2, 4, true);
        // Write something via session
        let slab = vec![3.0f32; 2 * 4];
        sess.write_layer_slab(0, &slab).expect("write");
        // Borrow as view — should see the same data
        let view = sess.as_view();
        for i in 0..8 {
            assert_eq!(view.hidden_output[i], 3.0);
        }
        assert!(view.per_position_argmaxes.is_some());
    }

    /// ADR-030 iter-105 — CPU-only validation of the Option A round-2
    /// drafter-input pipeline.  Targets the iter-87 "all-pad [0]*7"
    /// hypothesis (degenerate h_final on long Option A prompts) by
    /// asserting that the round-1 → round-2 data plumbing for
    /// `prior_captured` correctly delivers the NEWLY committed token's
    /// hidden state into round-2's `drafter_concat_new` slice.
    ///
    /// Scenario:
    /// - Initial prefill capture seeds prior_captured with prompt_len rows
    ///   (synthetic values: layer_idx * 1000 + pos * 10 + dim).
    /// - Round-1 verify produces a separate capture with block_size rows
    ///   (synthetic values: 100000 + layer_idx * 1000 + pos * 10 + dim).
    /// - n_committed = 1 (worst-case 0% accept).
    /// - append_capture_positions stacks them → prior_captured grows by 1.
    /// - extract_drafter_concat permutes → round-2 takes the last 1 row.
    /// - That row's contents at each drafter layer MUST equal the
    ///   round-1 verify capture at position 0 for that layer
    ///   (= the hidden state for `last_token` at the verify start_pos).
    ///
    /// If this assertion fails, the iter-87 all-pad pattern is plausibly
    /// caused by data-plumbing corruption.  If it passes (expected),
    /// the bug lies downstream in the GPU drafter forward.
    #[test]
    fn option_a_round2_prior_captured_delivers_correct_new_row() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let drafter_target_layer_ids: Vec<usize> = vec![1, 6, 11, 17, 22, 27];
        let final_layer_idx = 29usize;
        let mut combined_ids: Vec<usize> = drafter_target_layer_ids.clone();
        combined_ids.push(final_layer_idx);
        combined_ids.sort_unstable();
        combined_ids.dedup();
        let prompt_len = 10usize;
        let block_size = 8usize;
        let hs = 16usize; // small for fast test

        // Initial prefill capture (prompt only).
        let mut prior_captured =
            DFlashCaptureSession::new(combined_ids.clone(), prompt_len, hs, false);
        for (cli, &_clid) in combined_ids.iter().enumerate() {
            for t in 0..prompt_len {
                for d in 0..hs {
                    let off = (cli * prompt_len + t) * hs + d;
                    prior_captured.hidden_output[off] = (cli * 1000 + t * 10 + d) as f32;
                }
            }
        }

        // Round-1 verify capture (block_size rows at start_pos=prompt_len).
        let mut verify_captured =
            DFlashCaptureSession::new(combined_ids.clone(), block_size, hs, false);
        for (cli, &_clid) in combined_ids.iter().enumerate() {
            for t in 0..block_size {
                for d in 0..hs {
                    let off = (cli * block_size + t) * hs + d;
                    verify_captured.hidden_output[off] =
                        100_000.0 + (cli * 1000 + t * 10 + d) as f32;
                }
            }
        }

        // Round-1 accept-prefix = 0 → n_committed = 1 (target's bonus token).
        let n_committed = 1usize;
        let next_prior = append_capture_positions(&prior_captured, &verify_captured, n_committed)
            .expect("append_capture_positions");
        assert_eq!(next_prior.seq_len, prompt_len + n_committed);

        // Round-2 drafter extraction.
        let drafter_concat = extract_drafter_concat(
            &next_prior.hidden_output,
            &combined_ids,
            &drafter_target_layer_ids,
            next_prior.seq_len,
            hs,
        )
        .expect("extract_drafter_concat");

        // Round-2 takes the LAST n_committed = 1 row from the extraction
        // (mirrors orchestrator.rs:668-669 `new_rows_start =
        // drafter_cached_seq_len * row_stride`).
        let drafter_n = drafter_target_layer_ids.len();
        let row_stride = drafter_n * hs;
        let drafter_cached_seq_len = prompt_len;
        let new_rows_start = drafter_cached_seq_len * row_stride;
        let drafter_concat_new: &[f32] = &drafter_concat[new_rows_start..];
        assert_eq!(drafter_concat_new.len(), n_committed * row_stride);

        // The single new row's layout (per extract_drafter_concat):
        //   for drafter_l in 0..drafter_n:
        //     drafter_concat_new[drafter_l * hs .. (drafter_l+1) * hs]
        //     == verify_captured[combined_idx_of(target_layer_ids[drafter_l]), 0, :]
        for (drafter_l, &dl_id) in drafter_target_layer_ids.iter().enumerate() {
            let combined_idx = combined_ids
                .iter()
                .position(|&c| c == dl_id)
                .expect("drafter layer in combined");
            let verify_layer_base = combined_idx * block_size * hs;
            let verify_pos0_base = verify_layer_base; // pos 0 → +0
            let expected = &verify_captured.hidden_output[verify_pos0_base..verify_pos0_base + hs];
            let actual = &drafter_concat_new[drafter_l * hs..(drafter_l + 1) * hs];
            assert_eq!(
                actual, expected,
                "drafter_l={} (target_layer_id={}, combined_idx={}): \
                 drafter_concat_new row 0 must equal verify_captured[combined={}, pos=0, :]",
                drafter_l, dl_id, combined_idx, combined_idx,
            );
        }
    }

    /// ADR-030 iter-109 — covers the n_committed > 1 case (drafter
    /// acceptance > 0).  iter-105 tested the 0%-accept worst case
    /// (n_committed = 1).  This test exercises the multi-token-accept
    /// path: n_committed = 4 (4 drafts accepted + 1 target bonus from
    /// position 4).  All 5 new rows MUST come from verify_captured
    /// positions [0..5) at each drafter target layer.
    #[test]
    fn option_a_round2_prior_captured_multi_accept_plumbing() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let drafter_target_layer_ids: Vec<usize> = vec![1, 6, 11, 17, 22, 27];
        let final_layer_idx = 29usize;
        let mut combined_ids: Vec<usize> = drafter_target_layer_ids.clone();
        combined_ids.push(final_layer_idx);
        combined_ids.sort_unstable();
        combined_ids.dedup();
        let prompt_len = 8usize;
        let block_size = 8usize;
        let hs = 16usize;

        let mut prior_captured =
            DFlashCaptureSession::new(combined_ids.clone(), prompt_len, hs, false);
        for cli in 0..combined_ids.len() {
            for t in 0..prompt_len {
                for d in 0..hs {
                    let off = (cli * prompt_len + t) * hs + d;
                    prior_captured.hidden_output[off] = (cli * 1000 + t * 10 + d) as f32;
                }
            }
        }

        let mut verify_captured =
            DFlashCaptureSession::new(combined_ids.clone(), block_size, hs, false);
        for cli in 0..combined_ids.len() {
            for t in 0..block_size {
                for d in 0..hs {
                    let off = (cli * block_size + t) * hs + d;
                    verify_captured.hidden_output[off] =
                        100_000.0 + (cli * 1000 + t * 10 + d) as f32;
                }
            }
        }

        let n_committed = 5usize; // 4 accepts + 1 target bonus
        let next_prior = append_capture_positions(&prior_captured, &verify_captured, n_committed)
            .expect("append");
        assert_eq!(next_prior.seq_len, prompt_len + n_committed);

        let drafter_concat = extract_drafter_concat(
            &next_prior.hidden_output,
            &combined_ids,
            &drafter_target_layer_ids,
            next_prior.seq_len,
            hs,
        )
        .expect("extract");

        let drafter_n = drafter_target_layer_ids.len();
        let row_stride = drafter_n * hs;
        let drafter_cached_seq_len = prompt_len;
        let new_rows_start = drafter_cached_seq_len * row_stride;
        let drafter_concat_new: &[f32] = &drafter_concat[new_rows_start..];
        assert_eq!(drafter_concat_new.len(), n_committed * row_stride);

        // For each of the 5 new rows, verify it matches
        // verify_captured[combined_idx, t, :] for t in [0..5).
        for t in 0..n_committed {
            for (drafter_l, &dl_id) in drafter_target_layer_ids.iter().enumerate() {
                let combined_idx = combined_ids
                    .iter()
                    .position(|&c| c == dl_id)
                    .expect("drafter layer in combined");
                let verify_pos_t_base = (combined_idx * block_size + t) * hs;
                let expected =
                    &verify_captured.hidden_output[verify_pos_t_base..verify_pos_t_base + hs];
                let actual_row_base = (t * drafter_n + drafter_l) * hs;
                let actual = &drafter_concat_new[actual_row_base..actual_row_base + hs];
                assert_eq!(
                    actual, expected,
                    "t={t} drafter_l={drafter_l} (target_layer_id={dl_id}, combined_idx={combined_idx}): \
                     drafter_concat_new row {t} must equal verify_captured[combined={combined_idx}, pos={t}, :]",
                );
            }
        }
    }

    /// End-to-end integration: simulates the post-target-verify state
    /// (capture buffer populated by hypothetical layer-loop hook) and
    /// drives the drafter forward on the permuted concat. Validates
    /// the Phase 4 data path from capture → permute → drafter forward
    /// → finite output.
    #[test]
    #[ignore = "requires Metal device + drafter HF cache"]
    fn smoke_capture_to_drafter_forward_pipeline() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        use crate::inference::spec_decode::dflash::{
            forward::dispatch_dflash_model_forward,
            kv_cache::DFlashKvCache,
            tensors::DFlashModelTensors,
            weights::{DFlashWeights, DFlashWeightsFile},
        };
        use mlx_native::{DType, KernelRegistry, MlxDevice};

        let cfg = gemma4_cfg();
        let device = MlxDevice::new().expect("Metal device available on M5 Max");
        let mut registry = KernelRegistry::new();

        // Load drafter
        let home = std::env::var("HOME").expect("HOME set");
        let path = format!(
            "{home}/.cache/huggingface/hub/models--z-lab--gemma-4-26B-A4B-it-DFlash/snapshots/77d4202772dfe50b2396ec7bac9cfffc7b9e7057/model.safetensors"
        );
        let file = DFlashWeightsFile::open(&path).expect("file open");
        let weights = DFlashWeights::load(file.bytes(), &cfg).expect("validated load");
        let tensors = DFlashModelTensors::upload(&device, &cfg, &weights).expect("GPU upload");
        let mut cache = DFlashKvCache::new(&device, &cfg, 128).expect("cache");

        // Allocate capture buffer + populate as if target forward had
        // run and hooks had been written.
        let ctx_chunk = 4usize;
        let hidden = cfg.hidden_size;
        let block_size = 8u32;

        let mut hidden_buf = vec![0.0f32; capture_buffer_len(&cfg, ctx_chunk)];
        let target_layer_ids: Vec<usize> = cfg.target_layer_ids.clone();

        // Populate with synthetic values per (capture_layer_idx, token,
        // dim) so we can verify permute placement
        for cli in 0..target_layer_ids.len() {
            for t in 0..ctx_chunk {
                for d in 0..hidden {
                    let off = (cli * ctx_chunk + t) * hidden + d;
                    hidden_buf[off] = 0.1 + ((off % 31) as f32) / 310.0;
                }
            }
        }

        // Build PrefillCapture view + validate.
        // per_position_argmaxes is for the VERIFY forward pass, not the
        // ctx capture; we leave it None here since this test only
        // exercises the ctx-capture → drafter-forward path.
        let cap = PrefillCapture {
            target_layer_ids: &target_layer_ids,
            hidden_output: &mut hidden_buf,
            per_position_argmaxes: None,
        };
        cap.validate(ctx_chunk, hidden).expect("validate");

        // Permute capture buffer to drafter's expected concat layout
        // [seq_len, num_capture_layers * hidden_size]
        let concat = cap.permute_to_concat(ctx_chunk, hidden);
        assert_eq!(concat.len(), ctx_chunk * target_layer_ids.len() * hidden);

        // Upload concat to GPU as target_hidden_concat for the drafter
        let mut target_hidden = device
            .alloc_buffer(
                concat.len() * 4,
                DType::F32,
                vec![ctx_chunk, target_layer_ids.len() * hidden],
            )
            .expect("alloc target_hidden");
        target_hidden
            .as_mut_slice::<f32>()
            .expect("target_hidden slice")
            .copy_from_slice(&concat);

        // Allocate input h (simulated embed_tokens output)
        let h_elem = (block_size as usize) * hidden;
        let mut h = device
            .alloc_buffer(h_elem * 4, DType::F32, vec![block_size as usize, hidden])
            .expect("alloc h");
        {
            let s = h.as_mut_slice::<f32>().expect("h slice");
            for v in s.iter_mut() {
                *v = 1.0;
            }
        }

        // Run drafter forward end-to-end on the permuted captured hidden
        let h_final = dispatch_dflash_model_forward(
            &mut registry,
            &device,
            &h,
            &target_hidden,
            &tensors,
            &mut cache,
            &cfg,
            block_size,
            ctx_chunk as u32,
        )
        .expect("drafter forward on permuted capture");

        // Validate output [L, hidden] all finite
        assert_eq!(h_final.element_count(), (block_size as usize) * hidden);
        let host: &[f32] = h_final.as_slice::<f32>().expect("h_final slice");
        let n_finite = host.iter().filter(|v| v.is_finite()).count();
        assert_eq!(
            n_finite,
            host.len(),
            "drafter output on permuted capture must be all finite"
        );
    }

    #[test]
    fn permute_to_concat_round_trips_layout() {
        let _gpu = crate::inference::hf2q_gpu_test_lock();
        let target_layer_ids = vec![1, 6];
        let n_layers = target_layer_ids.len();
        let seq_len = 3;
        let hidden_size = 4;
        // Build a recognizable input: value = layer_idx * 10 + token * 1 + dim * 0.1
        let mut hidden = vec![0.0f32; n_layers * seq_len * hidden_size];
        for layer_idx in 0..n_layers {
            for t in 0..seq_len {
                for d in 0..hidden_size {
                    let src = (layer_idx * seq_len + t) * hidden_size + d;
                    hidden[src] = (layer_idx * 10) as f32 + t as f32 + (d as f32) * 0.1;
                }
            }
        }
        let cap = PrefillCapture {
            target_layer_ids: &target_layer_ids,
            hidden_output: &mut hidden,
            per_position_argmaxes: None,
        };
        let concat = cap.permute_to_concat(seq_len, hidden_size);
        assert_eq!(concat.len(), seq_len * n_layers * hidden_size);
        // Verify: concat[(t * n_layers + layer_idx) * hidden_size + d]
        //       == hidden[(layer_idx * seq_len + t) * hidden_size + d]
        for layer_idx in 0..n_layers {
            for t in 0..seq_len {
                for d in 0..hidden_size {
                    let dst = (t * n_layers + layer_idx) * hidden_size + d;
                    let expected = (layer_idx * 10) as f32 + t as f32 + (d as f32) * 0.1;
                    assert_eq!(
                        concat[dst], expected,
                        "concat mismatch t={t} layer={layer_idx} d={d}"
                    );
                }
            }
        }
    }
}