What causes "shape mismatch k {:?} and v {:?}" in huggingface/candle?

Make v_cache match k_cache exactly in all four dims (num_blocks, page_block_size, num_heads_k, head_dim)

shape mismatch k {:?} and v {:?}

Error message

shape mismatch k {:?} and v {:?}

What it means

In paged varlen mode, the k and v paged caches must have identical shapes: (num_blocks, page_block_size, num_heads_k, head_dim). The wrapper compares k's dims4 with v's dims4 and bails on any mismatch.

Source

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

                "flash-attn-varlen paged expects k/v tensors of rank 4 (k: {k_rank}, v: {v_rank})"
            )
        }
        if q_stride[q_rank - 1] != 1 {
            candle::bail!("the last dim of q must be contiguous {q_stride:?}")
        }
        if k_stride[k_rank - 1] != 1 {
            candle::bail!("the last dim of k must be contiguous {k_stride:?}")
        }
        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})"
                )
            }

View on GitHub (pinned to d5fee525bf)

Solutions

Make v_cache match k_cache exactly in all four dims (num_blocks, page_block_size, num_heads_k, head_dim)
Re-allocate both caches from the same cache config at startup
Log/compare k.shape() and v.shape() before the call to find which dim differs

Example fix

// before
let k_cache = Tensor::zeros((n_blocks, 16, h_kv, d), ...)?;
let v_cache = Tensor::zeros((n_blocks, 32, h_kv, d), ...)?; // page size differs
// after
let k_cache = Tensor::zeros((n_blocks, 16, h_kv, d), ...)?;
let v_cache = Tensor::zeros((n_blocks, 16, h_kv, d), ...)?;

Defensive patterns

Strategy: validation

Validate before calling

fn check_kv_shapes_match(k: &candle_core::Tensor, v: &candle_core::Tensor) -> candle_core::Result<()> {
    if k.shape() != v.shape() {
        candle_core::bail!("kv cache shape mismatch: k {:?} vs v {:?}", k.shape(), v.shape());
    }
    Ok(())
}
// call before forward in paged mode

Type guard

fn kv_shapes_match(k: &candle_core::Tensor, v: &candle_core::Tensor) -> bool {
    k.shape() == v.shape()
}

Try / catch

match attn.forward(&q, &k, &v, &sq, &sk, Some(&bt)) {
    Ok(out) => out,
    Err(e) if e.to_string().contains("shape mismatch k") => {
        candle_core::bail!("misconfigured paged cache: {} (k={:?} v={:?})", e, k.shape(), v.shape())
    }
    Err(e) => return Err(e),
}

Prevention

Allocate K and V caches from one shared cache-config struct
Assert k.shape() == v.shape() whenever a cache is resized
Apply GQA/head-count changes to both K and V caches simultaneously

When it happens

Trigger: Calling forward with block_table set and k_cache of shape e.g. (N, P, H, D) but v_cache of a different number of blocks, page size, head count, or head dim.

Common situations: Allocating K and V caches with different page_block_size or block counts; a refactoring that resized one cache but not the other; GQA changes applied to K but not V.

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

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