inferencelayer 0.2.13

//! Prefill attention-mask policy — the per-query KEY SPAN a decoder column may read.
//!
//! The engine's attention kernels ([`crate::forward`]'s `ATTN_K` family) never materialize a 2-D
//! mask. Each query column carries its own contiguous key span and the kernel walks it, which is
//! why a mask policy here is a SPAN RULE and not a new kernel family: see `ATTN_K`'s
//! `t = cmeta[..] + 1` (span end) and `s0 = t - window` (span start).
//!
//! Today every decoder column is causal — span end is the column's own position. That is correct
//! for text and for Qwen3-VL, and WRONG for Gemma 3, whose decoder attends bidirectionally within
//! each contiguous run of image tokens. See `bench/gemma3_mask_fixture/` for the HF-derived
//! reference this module is gated against, and `docs/LITERT_LM_PARITY_PLAN.md` item V1.

/// The inclusive key span `[lo, hi]` a query column attends.
///
/// Both ends are absolute positions in the sequence, so `hi` is the LAST key read (not a
/// one-past-the-end bound) — `ATTN_K` consumes it as `t = hi + 1`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct KeySpan {
    pub lo: u32,
    pub hi: u32,
}

/// Which mask a prefill column uses.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum MaskPolicy {
    /// Query `q` attends `[causal_start, q]`. Text models, and Qwen3-VL including its image runs.
    #[default]
    Causal,
    /// [`MaskPolicy::Causal`], but an image token additionally reads FORWARD to the end of its own
    /// contiguous image run — Gemma 3 / 3n.
    ///
    /// Scope, all three parts load-bearing (verified against `bench/gemma3_mask_fixture/`):
    /// - **its OWN run only.** An earlier run must not reach a later one; a later run reaches an
    ///   earlier one only through the ordinary causal floor.
    /// - **forward-unbounded even on sliding layers.** HF composes the window as
    ///   `AND(k > q - window, OR(causal, blockwise))`, so the window clips the span FLOOR and
    ///   never its ceiling.
    /// - **prefill only.** HF drops `token_type_ids` after the first iteration
    ///   (`modeling_gemma3.py`: *"Don't pass to not apply bidirectional mask on top"*), so a
    ///   decode step is plain causal and must keep using [`MaskPolicy::Causal`].
    VisionBidirectional,
}

/// Per-layer key-range floor: `None` for a global (full-attention) layer, `Some(window)` for a
/// sliding-window layer, where the layer may read keys `k > q - window`.
///
/// This is the global-vs-local split: the same policy produces different spans per layer type,
/// and a model can be causal-global / bidirectional-local.
pub type Window = Option<u32>;

/// The span start for query `q` under `window` — identical to `ATTN_K`'s `s0`.
///
/// The floor is a function of the QUERY position, never of the span end: HF composes the window
/// as `AND(k > q - window, ...)`, so widening the top of a span must not drag its floor up with
/// it. Getting this wrong is invisible until an image run ends beyond `window` tokens in.
fn span_lo(q: u32, window: Window) -> u32 {
    match window {
        // k > q - window  <=>  k >= q - window + 1
        Some(w) if w > 0 => q.saturating_sub(w - 1),
        _ => 0,
    }
}

/// The key spans for a prefill chunk of `n` columns, one per query column.
///
/// `is_image[q]` marks a column belonging to a bidirectionally-attended block. It is deliberately
/// NOT a token-id test (no hard-coded linguistic data) and NOT "the embedding came from the vision
/// arena" — Gemma 3 is exactly the case that separates the two, attending bidirectionally over
/// image runs while having no M-RoPE and therefore no arena. See `BatchCol::is_image`.
pub fn spans(policy: MaskPolicy, is_image: &[bool], window: Window) -> Vec<KeySpan> {
    let n = is_image.len();
    // hi[q]: the last key q may read. Causal everywhere, then — under the vision policy — each
    // contiguous image run is walked once and its members all raised to the run's last column.
    let mut hi: Vec<u32> = (0..n as u32).collect();
    if policy == MaskPolicy::VisionBidirectional {
        let mut q = 0;
        while q < n {
            if !is_image[q] {
                q += 1;
                continue;
            }
            let start = q;
            while q + 1 < n && is_image[q + 1] {
                q += 1;
            }
            // `q` is now the run's last column. A run that reaches the end of the chunk keeps the
            // chunk's own end as its ceiling — see `spans` caller contract on chunk splitting.
            hi[start..=q].fill(q as u32);
            q += 1;
        }
    }
    (0..n)
        .map(|q| KeySpan {
            lo: span_lo(q as u32, window),
            hi: hi[q],
        })
        .collect()
}

/// One column of a batch step, as far as the mask is concerned.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MaskCol {
    /// Linear position in its own sequence — the KV slot and the causal span end.
    pub pos: u32,
    /// Block-table row: the SEQUENCE identity. Columns of different sequences never share a run.
    pub btrow: u32,
    /// True when this column belongs to a bidirectionally-attended block — `BatchCol::is_image`,
    /// which is a MASK question, not a "where did the embedding come from" question.
    pub is_image: bool,
}

/// Per-column span END (`hi`) for a whole batch step — the value the attention meta carries.
///
/// A batch step interleaves columns from DIFFERENT sequences (continuous batching) in arbitrary
/// order, so a run is a maximal set of columns that share a `btrow`, are all image, and occupy
/// CONSECUTIVE positions. Two image columns that merely sit next to each other in the batch are
/// not a run — that would let one sequence's image widen into another's, which is why this takes
/// `MaskCol` rather than the flat `is_image` slice [`spans`] uses.
///
/// A run's ceiling is the largest position of that run PRESENT IN THIS STEP. Keys beyond the step
/// are not in the KV cache yet, so they cannot be read regardless of policy — see
/// [`chunk_splits_image_run`], which is how a caller detects that it has cut a run short.
pub fn batch_span_ends(policy: MaskPolicy, cols: &[MaskCol]) -> Vec<u32> {
    let mut hi: Vec<u32> = cols.iter().map(|c| c.pos).collect();
    if policy != MaskPolicy::VisionBidirectional {
        return hi;
    }
    for (i, c) in cols.iter().enumerate() {
        if !c.is_image {
            continue;
        }
        // Walk this column's run by POSITION within its own sequence, not by batch order.
        let mut end = c.pos;
        loop {
            let next = cols
                .iter()
                .any(|o| o.is_image && o.btrow == c.btrow && o.pos == end + 1);
            if !next {
                break;
            }
            end += 1;
        }
        hi[i] = end;
    }
    hi
}

/// Does this chunk split an image run at its END boundary?
///
/// A column may only read keys already written to the KV cache. Under
/// [`MaskPolicy::VisionBidirectional`] an image column reads FORWARD, so if a prefill chunk ends
/// mid-image-run the keys past the boundary do not exist yet and the run's spans would be
/// silently truncated — the failure mode V1 exists to prevent, one step removed. Callers that
/// chunk prefill must extend the chunk to the run's end (or refuse) when this returns true.
#[must_use]
pub fn chunk_splits_image_run(is_image: &[bool], chunk_end: usize) -> bool {
    chunk_end > 0 && chunk_end < is_image.len() && is_image[chunk_end - 1] && is_image[chunk_end]
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn causal_spans_are_the_engines_current_rule() {
        let is_image = vec![false; 6];
        let g = spans(MaskPolicy::Causal, &is_image, None);
        assert_eq!(g[0], KeySpan { lo: 0, hi: 0 });
        assert_eq!(g[5], KeySpan { lo: 0, hi: 5 });
        // window = 3 ⇒ q reads k > q-3, i.e. three keys ending at q.
        let l = spans(MaskPolicy::Causal, &is_image, Some(3));
        assert_eq!(l[5], KeySpan { lo: 3, hi: 5 });
        assert_eq!(l[1], KeySpan { lo: 0, hi: 1 }, "floor clamps at 0");
    }
}