{"record":{"id":"3174dc2121cc17b2","repo":"huggingface/candle","slug":"mm-prefix-ranges-shape-must-be-batch-size-max","errorCode":null,"errorMessage":"mm_prefix_ranges shape must be ({batch_size}, max_ranges, 2), got {:?}","messagePattern":"mm_prefix_ranges shape must be \\((.+?), max_ranges, 2\\), got (.+?)","errorType":"exception","errorClass":null,"httpStatus":null,"severity":"error","filePath":"candle-flash-attn/src/lib.rs","lineNumber":632,"sourceCode":"            candle::bail!(\"seqlens_q and seqlens_k should have the same number of elements {nseqlens_q} <> {nseqlens_k}\")\n        }\n\n        let batch_size = nseqlens_q - 1;\n        let mm_prefix_ranges = if let Some(mm_prefix_ranges) = self.mm_prefix_ranges.as_ref() {\n            let (storage, layout) = mm_prefix_ranges.storage_and_layout();\n            if mm_prefix_ranges.dtype() != DType::I32 {\n                candle::bail!(\n                    \"mm_prefix_ranges must be i32, got {:?}\",\n                    mm_prefix_ranges.dtype()\n                )\n            }\n            match &*storage {\n                candle::Storage::Cuda(_) => {}\n                _ => candle::bail!(\"mm_prefix_ranges must be a cuda tensor\"),\n            }\n            let (mm_batch, max_ranges, two) = layout.shape().dims3()?;\n            if mm_batch != batch_size || two != 2 {\n                candle::bail!(\n                    \"mm_prefix_ranges shape must be ({batch_size}, max_ranges, 2), got {:?}\",\n                    layout.shape()\n                )\n            }\n            if layout.stride().last().copied() != Some(1) {\n                candle::bail!(\"mm_prefix_ranges last dimension must be contiguous\")\n            }\n            Some((\n                storage,\n                layout.start_offset(),\n                layout.stride()[0],\n                max_ranges,\n            ))\n        } else {\n            None\n        };\n\n        let stream = dev.cuda_stream();","sourceCodeStart":614,"sourceCodeEnd":650,"githubUrl":"https://github.com/huggingface/candle/blob/d5fee525bfde3273eb7c9b75fd2bc4937be867ca/candle-flash-attn/src/lib.rs#L614-L650","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","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"],"exampleFix":"// before\nlet ranges = Tensor::from_vec(ranges_vec, (batch, max_ranges), &dev)?; // rank 2\n// after\nlet ranges = Tensor::from_vec(ranges_vec, (batch, max_ranges, 2), &dev)?;","handlingStrategy":"validation","validationCode":"fn check_prefix_ranges(r: &Tensor, batch_size: usize) -> candle::Result<()> {\n    let (b, _, two) = r.dims3()?;\n    if b != batch_size || two != 2 {\n        candle::bail!(\"mm_prefix_ranges must be ({batch_size}, max_ranges, 2), got {:?}\", r.shape());\n    }\n    Ok(())\n}","typeGuard":null,"tryCatchPattern":"match flash_attn_fwd(...) {\n    Err(e) if e.to_string().contains(\"mm_prefix_ranges shape\") => eprintln!(\"fix ranges shape to (B, max_ranges, 2)\"),\n    other => other?,\n}","preventionTips":["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"],"tags":["cuda","flash-attn","shape-mismatch","tensor-shape"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"d5fee525bfde3273eb7c9b75fd2bc4937be867ca","analyzedAt":"2026-09-02T00:15:47.023Z","contentChangedAt":null,"schemaVersion":2},"datasetVersion":"2026-09-09T06:17:21.866Z"}