causal::run_causal_attn_cpu is B=1 only (got B={b}). Multi-b

Error message

causal::run_causal_attn_cpu is B=1 only (got B={b}). Multi-batch should be routed through the varlen path.

What it means

The CPU causal attention kernel run_causal_attn_cpu only supports a batch size of 1; it squeezes the batch dim immediately after the check. Multi-batch inputs must go through the variable-length (varlen) attention path instead.

Source

Thrown at candle-nn/src/attention/cpu_flash/causal.rs:68

///
/// Squeezes batch dim, extracts contiguous slices, dispatches to
/// f32 or generic kernel. The inner kernels operate on raw slices only.
#[allow(clippy::too_many_arguments)]
pub fn run_causal_attn_cpu<T>(
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
    softmax_scale: f32,
    kv_offset: usize,
    max_bias: Option<f32>,
    softcap: Option<f32>,
) -> Result<Tensor>
where
    T: WithDType + num_traits::Float,
{
    let b = q.dims()[0];
    if b != 1 {
        candle::bail!(
            "causal::run_causal_attn_cpu is B=1 only (got B={b}). \
             Multi-batch should be routed through the varlen path."
        );
    }

    let q = q.squeeze(0)?.contiguous()?;
    let k = k.squeeze(0)?.contiguous()?;
    let v = v.squeeze(0)?.contiguous()?;

    let (s_q, h_q, d) = q.dims3()?;
    let (s_kv, h_kv, _) = k.dims3()?;
    let (_, h_v, _) = v.dims3()?;

    let max_bias = max_bias.unwrap_or(0.0);
    let softcap = softcap.unwrap_or(0.0);

    if q.dtype() == DType::F32 {
        let (q_g, q_l) = q.storage_and_layout();

View on GitHub (pinned to d5fee525bf)

Solutions

Route batched inputs through the varlen attention path instead
Split the batch and call run_causal_attn_cpu per element with b=1
Squeeze to a single sequence if the input truly has one example (fix upstream shape)

Example fix

// before
let y = causal::run_causal_attn_cpu(&q, &k, &v, &mask)?; // q dims (4, s, h, d)
// after
let y = varlen::run_varlen_attn_cpu(&q, &k, &v, &cu_seqlens)?; // multi-batch path

Defensive patterns

Strategy: validation

Validate before calling

if q.dim(0)? != 1 {
    return Err(anyhow!("use the varlen path for batch>1"));
}

Type guard

fn is_single_batch(q: &Tensor) -> bool { q.dim(0).map(|b| b == 1).unwrap_or(false) }

Try / catch

match result {
    Err(e) if e.to_string().contains("B=1 only") => {
        // route to varlen::run_varlen_attn_cpu instead
    }
    other => other?,
}

Prevention

Dispatch to the varlen path whenever batch > 1 at the call site
Keep batch handling in one dispatcher function
Add a batch-size assert in tests covering the CPU attention path

When it happens

Trigger: Calling run_causal_attn_cpu (or the CPU causal flash-attention wrapper) with a q tensor whose dims()[0] > 1.

Common situations: Batched inference on CPU with batch>1 and the code path selecting the B=1 causal kernel; refactoring code that previously ran single sequences; tests with stacked examples.

Related errors

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