What causes "mm_prefix_ranges shape must be ({batch_size}, max_ranges, 2), got {:?}" in huggingface/candle?

Verify mm_prefix_ranges has exactly 3 dims (batch_size, max_ranges, 2) with dim0 equal to the q/k batch size

mm_prefix_ranges shape must be ({batch_size}, max_ranges, 2)

Error message

mm_prefix_ranges shape must be ({batch_size}, max_ranges, 2), got {:?}

What it means

flash-attn's varlen CUDA kernel requires the mm_prefix_ranges argument to be a CUDA tensor of shape (batch_size, max_ranges, 2) matching the attention batch size. The library checks this before extracting raw pointers and bails if the batch dimension or trailing dim of 2 does not match. This prevents passing misshapen per-request range data into the kernel.

Source

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

            candle::bail!("seqlens_q and seqlens_k should have the same number of elements {nseqlens_q} <> {nseqlens_k}")
        }

        let batch_size = nseqlens_q - 1;
        let mm_prefix_ranges = if let Some(mm_prefix_ranges) = self.mm_prefix_ranges.as_ref() {
            let (storage, layout) = mm_prefix_ranges.storage_and_layout();
            if mm_prefix_ranges.dtype() != DType::I32 {
                candle::bail!(
                    "mm_prefix_ranges must be i32, got {:?}",
                    mm_prefix_ranges.dtype()
                )
            }
            match &*storage {
                candle::Storage::Cuda(_) => {}
                _ => candle::bail!("mm_prefix_ranges must be a cuda tensor"),
            }
            let (mm_batch, max_ranges, two) = layout.shape().dims3()?;
            if mm_batch != batch_size || two != 2 {
                candle::bail!(
                    "mm_prefix_ranges shape must be ({batch_size}, max_ranges, 2), got {:?}",
                    layout.shape()
                )
            }
            if layout.stride().last().copied() != Some(1) {
                candle::bail!("mm_prefix_ranges last dimension must be contiguous")
            }
            Some((
                storage,
                layout.start_offset(),
                layout.stride()[0],
                max_ranges,
            ))
        } else {
            None
        };

        let stream = dev.cuda_stream();

View on GitHub (pinned to d5fee525bf)

Solutions

Verify mm_prefix_ranges has exactly 3 dims (batch_size, max_ranges, 2) with dim0 equal to the q/k batch size
Ensure each range entry is a (start, end) pair; unsqueeze the last dim if you passed a 2-D tensor
Recompute ranges for the current batch instead of caching tensors from another batch
Call .to_device(cuda) so the tensor is a CUDA storage

Example fix

// before
let ranges = Tensor::from_vec(ranges_vec, (batch, max_ranges), &dev)?; // rank 2
// after
let ranges = Tensor::from_vec(ranges_vec, (batch, max_ranges, 2), &dev)?;

Defensive patterns

Strategy: validation

Validate before calling

fn check_prefix_ranges(r: &Tensor, batch_size: usize) -> candle::Result<()> {
    let (b, _, two) = r.dims3()?;
    if b != batch_size || two != 2 {
        candle::bail!("mm_prefix_ranges must be ({batch_size}, max_ranges, 2), got {:?}", r.shape());
    }
    Ok(())
}

Try / catch

match flash_attn_fwd(...) {
    Err(e) if e.to_string().contains("mm_prefix_ranges shape") => eprintln!("fix ranges shape to (B, max_ranges, 2)"),
    other => other?,
}

Prevention

Always build mm_prefix_ranges as (batch, max_ranges, 2)
Assert the batch dim equals q's batch before every call
Keep a unit test that runs the forward with dummy ranges

When it happens

Trigger: Calling FlashAttention varlen forward (cuda_fwd_t) with an mm_prefix_ranges tensor whose dim0 differs from batch_size, or whose rank-3 last dimension is not 2 (e.g. forgot the (start,end) pair layout, or passed a (batch, max_ranges) 2-D tensor).

Common situations: Building paged/varlen attention batches by hand, reusing prefix-range tensors computed for a different batch size after dynamic batching, or reshaping errors when packing ranges.

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/3174dc2121cc17b2. Report an issue: GitHub.