flash-attn-varlen paged expects k/v tensors of rank 4 (k: {k

Error message

flash-attn-varlen paged expects k/v tensors of rank 4 (k: {k_rank}, v: {v_rank})

What it means

In paged varlen flash-attention mode (a block_table was supplied), k and v must be rank-4 PagedCache tensors shaped (num_blocks, page_block_size, num_heads_k, head_dim), while q stays rank 3. The wrapper bails if either k or v is not rank 4 when paged attention is active.

Source

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

        let q_stride = q_l.stride();
        let k_stride = k_l.stride();
        let v_stride = v_l.stride();
        let o_stride = out_l.stride();

        let q_rank = q_stride.len();
        let k_rank = k_stride.len();
        let v_rank = v_stride.len();
        let o_rank = o_stride.len();

        let paged = block_table.is_some();
        if q_rank != 3 || (!paged && k_rank != 3) || (!paged && v_rank != 3) {
            candle::bail!(
                "flash-attn-varlen expects input tensors of rank 3 (q: {q_rank}, k: {k_rank}, v: {v_rank}"
            )
        }
        if paged && (k_rank != 4 || v_rank != 4) {
            candle::bail!(
                "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()? {

View on GitHub (pinned to d5fee525bf)

Solutions

Reshape/allocate k and v as rank-4 paged caches: (num_blocks, page_block_size, num_heads_k, head_dim)
If you do not intend paged attention, pass block_table = None so rank-3 k/v are accepted
Ensure the paged cache tensor is contiguous and strided as a proper 4D tensor (not a flattened view without shape metadata)

Example fix

// before (paged mode with rank-3 kv)
let out = fwd.forward(q, &k_cache, &v_cache, &seqlens_q, &seqlens_k, Some(&block_table))?;
// after
let k_cache = k_cache.reshape((num_blocks, page_block_size, h_kv, d))?;
let v_cache = v_cache.reshape((num_blocks, page_block_size, h_kv, d))?;
let out = fwd.forward(q, &k_cache, &v_cache, &seqlens_q, &seqlens_k, Some(&block_table))?;

Defensive patterns

Strategy: validation

Validate before calling

fn check_paged_kv(k: &candle_core::Tensor, v: &candle_core::Tensor) -> candle_core::Result<()> {
    for (name, t) in [("k", k), ("v", v)] {
        if t.dims().len() != 4 {
            candle_core::bail!("paged {name} cache must be rank 4 (blocks, page, heads, head_dim), got {:?}", t.dims());
        }
    }
    Ok(())
}
// call before forward when block_table.is_some()

Type guard

fn is_rank4(t: &candle_core::Tensor) -> bool { t.dims().len() == 4 }

Try / catch

match attn.forward(&q, &k, &v, &sq, &sk, block_table) {
    Ok(out) => out,
    Err(e) if e.to_string().contains("paged expects k/v tensors of rank 4") => {
        candle_core::bail!("KV cache not allocated as paged (blocks, page, heads, head_dim): {}", e)
    }
    Err(e) => return Err(e),
}

Prevention

Allocate K/V caches through a single PagedCache helper that enforces rank 4
Gate paged code paths with an explicit config flag and assert cache rank there
Never mix paged block_table with plain contiguous KV caches

When it happens

Trigger: Calling cuda_fwd_t with block_table = Some(...) but passing k or v with rank 3 (plain contiguous KV cache) or any other rank, instead of a 4D paged KV cache.

Common situations: Enabling paged attention (e.g. switching to a vLLM-style block KV cache) but forgetting to restructure the KV cache into blocks; mixing paged attention with non-paged cache tensors left over from the previous implementation.

Related errors

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