Skip to main content

attend_quantized

Function attend_quantized 

Source
pub fn attend_quantized(
    q: &[f32],
    layer: &QuantizedKvLayer,
    n_heads: usize,
    kv_heads: usize,
    head_dim: usize,
    scale: f32,
) -> Result<Vec<f32>>
Expand description

Single-query attention directly over a quantized K/V layer, dequantizing one row at a time so peak extra memory is O(kv_dim), not O(past_len × kv_dim) — the point of quantized-KV attention vs the dequantize-the-whole- cache path. GQA-aware: n_heads query heads share kv_heads K/V heads (group = n_heads / kv_heads).

  • q: the new token’s query, [n_heads × head_dim].
  • returns the attention output [n_heads × head_dim].

This is the host reference for a future first-class QuantizedAttention op; it lets long-context decode keep the cache small without ever materializing the full f32 history.