huggingface/candle · error

page_block_size {page_block_size_arg} does not match k shape

Error message

page_block_size {page_block_size_arg} does not match k shape {:?}

What it means

When paged attention is active, the page_block_size passed to the builder must equal the page_block_size dimension (dim 1) of the 4D k cache. The wrapper compares them and bails on disagreement, since the kernel indexes the block table assuming that layout.

Source

Thrown at candle-flash-attn/src/lib.rs:572

        if v_stride[v_rank - 1] != 1 {
            candle::bail!("the last dim of v must be contiguous {v_stride:?}")
        }

        let (total_q, num_heads, head_size_og) = q_l.shape().dims3()?;
        let (num_heads_k, page_block_size) = if paged {
            let (_, page_block_size, num_heads_k, k_head_size) = k_l.shape().dims4()?;
            let expected_v = k_l.shape().dims4()?;
            if expected_v != v_l.shape().dims4()? {
                candle::bail!("shape mismatch k {:?} and v {:?}", k_l.shape(), v_l.shape())
            }
            if k_head_size != head_size_og {
                candle::bail!("shape mismatch q {:?} and k {:?}", q_l.shape(), k_l.shape())
            }
            let Some(page_block_size_arg) = self.page_block_size else {
                candle::bail!("paged flash-attn requires page_block_size")
            };
            if page_block_size_arg != page_block_size {
                candle::bail!(
                    "page_block_size {page_block_size_arg} does not match k shape {:?}",
                    k_l.shape()
                )
            }
            if page_block_size % 32 != 0 {
                candle::bail!(
                    "paged flash-attn requires page_block_size to be a multiple of 32 (got {page_block_size})"
                )
            }
            if head_size_og > 512 {
                candle::bail!("paged flash-attn supports head sizes up to 512 (got {head_size_og})")
            }
            (num_heads_k, page_block_size_arg)
        } else {
            let (total_k, num_heads_k, _head_size_og) = k_l.shape().dims3()?;
            let expected_kv = (total_k, num_heads_k, head_size_og);
            if expected_kv != k_l.shape().dims3()? {
                candle::bail!("shape mismatch q {:?} and k {:?}", q_l.shape(), k_l.shape())

View on GitHub (pinned to d5fee525bf)

Solutions

1. Use a single block-size constant for both the builder option and the cache allocation
2. Re-allocate the k/v caches with page_block_size matching the builder value
3. Read the actual page size from the cache shape and pass it to the builder instead of hardcoding

Example fix

// before
let attn = FlashAttentionVarLen::new(s, wl, wr, Some(16))?;
let k_cache = Tensor::zeros((n, 32, h_kv, d), ...)?; // 32 != 16
// after
let page_block_size = 16;
let attn = FlashAttentionVarLen::new(s, wl, wr, Some(page_block_size))?;
let k_cache = Tensor::zeros((n, page_block_size, h_kv, d), ...)?;

fn check_page_block_size(builder_pbs: u32, k: &candle_core::Tensor) -> candle_core::Result<()> {
    let k_dims = k.dims4()?;
    if builder_pbs as usize != k_dims[1] {
        candle_core::bail!(
            "page_block_size {} != k cache page dim {}",
            builder_pbs, k_dims[1]
        );
    }
    Ok(())
}

fn page_block_size_matches(pbs: u32, k: &candle_core::Tensor) -> bool {
    k.dims().get(1) == Some(&(pbs as usize))
}

match attn.forward(&q, &k, &v, &sq, &sk, Some(&bt)) {
    Ok(out) => out,
    Err(e) if e.to_string().contains("does not match k shape") => {
        candle_core::bail!("cache page size vs builder page_block_size mismatch: {}", e)
    }
    Err(e) => return Err(e),
}

When it happens

Trigger: Calling forward with block_table set, builder page_block_size = P1, but a k cache whose shape is (blocks, P2, Hk, D) with P1 != P2.

Common situations: Changing the KV cache block size in the runtime config without updating the attention builder; reusing a preallocated cache across models with different page sizes; copy-pasted configs where two block-size constants diverge.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• shape mismatch k {:?} and v {:?}
• shape mismatch alibi_slopes {:?}, expected {:?}
• block_table must be a cuda tensor
• block_table last dimension must be contiguous
• flash-attn-varlen paged expects k/v tensors of rank 4 (k: {k

AI-assisted analysis of huggingface/candle@d5fee525bf (2026-09-02). Data as JSON: /api/errors/0f0a80db0b182fd2. Report an issue: GitHub.