flash-attn-v3-varlen expects input tensors of rank 3 (q: {q_

Error message

flash-attn-v3-varlen expects input tensors of rank 3 (q: {q_rank}, k: {k_rank}, v: {v_rank}

What it means

The varlen kernel operates on packed (total_tokens, num_heads, head_dim) tensors. The wrapper reads q/k/v stride lengths to get each rank and bails unless all three are exactly rank 3. Passing rank-2 (missing heads axis) or rank-4 (batched, non-varlen) tensors triggers this.

Source

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

        let q = q.as_cuda_slice::<T>()?;
        let k = k.as_cuda_slice::<T>()?;
        let v = v.as_cuda_slice::<T>()?;
        let q = q.slice(q_l.start_offset()..);
        let k = k.slice(k_l.start_offset()..);
        let v = v.slice(v_l.start_offset()..);

        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();

        if q_rank != 3 || k_rank != 3 || v_rank != 3 {
            candle::bail!(
                "flash-attn-v3-varlen expects input tensors of rank 3 (q: {q_rank}, 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 (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

Reshape inputs to (total_q, num_heads, head_dim) / (total_k, num_heads_k, head_dim): for (B,S,H,D) use q.reshape((b*s, h, d)).
If the tensor is (tokens, num_heads*head_dim), reshape the last dim into (heads, head_dim).
Check q.rank() == 3 (etc.) at the call site and reshape otherwise.

Example fix

// before
let q = q.reshape((b * s, h * d))?; // rank 2
// after
let q = q.reshape((b * s, h, d))?; // (total, heads, head_dim)

Defensive patterns

Strategy: validation

Validate before calling

if q.rank() != 3 || k.rank() != 3 || v.rank() != 3 {
    candle_core::bail!("varlen inputs must be (total, heads, head_dim); got q rank {}", q.rank());
}
// (B,S,H,D) -> (B*S,H,D)
let q = q.reshape((b * s, h, d))?;

Type guard

fn is_rank3(t: &candle_core::Tensor) -> bool { t.rank() == 3 }

Try / catch

match forward_varlen(&q, &k, &v, ...) {
    Err(e) if e.to_string().contains("rank 3") => {
        let q = flatten_batch_seq(&q)?;
        let k = flatten_batch_seq(&k)?;
        let v = flatten_batch_seq(&v)?;
        forward_varlen(&q, &k, &v, ...)
    }
    other => other,
}

Prevention

Remember varlen inputs are packed: merge batch and seq axes into a total-token axis, keeping (tokens, heads, head_dim).
Check rank == 3 in a shared wrapper around every flash-attn call.
Document the expected packed input shape at the attention module's public API.

When it happens

Trigger: Calling flash-attn-v3 varlen forward with q/k/v shaped (B, S, H, D) or (S, D) — e.g. reusing the batched (non-varlen) flash-attn input shapes or feeding projections without the heads axis.

Common situations: Migrating from the regular batched flash-attn API to the varlen API without reshaping; forgetting to merge batch and seq into a total-token axis; (tokens, heads*dim) tensors not split into two axes.

Related errors

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