tracel-ai/burn · error

attention: unsupported dtype {:?}

Error message

attention: unsupported dtype {:?}

What it means

The flex attention kernel macro handles F32, F64, F16 and BF16 (with half-precision paths casting mask/bias to f32), and panics with "attention: unsupported dtype" for any other query/key dtype. Attention math is inherently floating-point, so integer/bool inputs are rejected deliberately.

Source

Thrown at crates/burn-flex/src/ops/attention.rs:114

                    mask,
                    attn_bias.map(|b| cast_to_f32(b, f16::to_f32)),
                    options,
                );
                cast_from_f32(r, f16::from_f32)
            }
            DType::BF16 => {
                use burn_std::bf16;
                let r = $impl_fn::<f32>(
                    cast_to_f32(query, bf16::to_f32),
                    cast_to_f32(key, bf16::to_f32),
                    cast_to_f32(value, bf16::to_f32),
                    mask,
                    attn_bias.map(|b| cast_to_f32(b, bf16::to_f32)),
                    options,
                );
                cast_from_f32(r, bf16::from_f32)
            }
            dtype => panic!("attention: unsupported dtype {:?}", dtype),
        }
    }};
}

/// Contiguous mask/bias tensor plus the per-batch and per-head element offsets the
/// inner loop should use to locate the `[seq_q, seq_kv]` tile for each `(batch, head)`
/// pair. When a leading dim (batch or heads) is `1` in the source, its step is `0`, so
/// the inner loop re-reads the same tile for every pair without allocating an expanded
/// copy. The tile length itself is always `seq_q * seq_kv` and is computed at the call
/// site, so it is not stored here.
struct BroadcastMaskBias {
    tensor: FlexTensor,
    batch_step: usize,
    head_step: usize,
}

/// Prepare an attention mask or bias for the inner loop, accepting ONNX Attention-23
/// broadcast shapes.

View on GitHub (pinned to d16f7ba2ed)

Solutions

1. Cast q/k/v to F32 (or F16/BF16) before the attention call: q.cast(DType::F32) etc.
2. Verify tokens go through the embedding layer so attention receives float tensors, not token indices.
3. If a hook (new_with_hook / dtype_usage) advertises an unsupported dtype, restrict supported_dtype to F32/F64/F16/BF16 on the backend.

Example fix

// before
let attn = attention(q_ids, k_ids, v, mask, options); // integer
// after
let attn = attention(q_ids.cast(DType::F32), k_ids.cast(DType::F32), v.cast(DType::F32), mask, options);

for t in [&q, &k, &v] {
    if !matches!(t.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16) {
        panic!("attention inputs must be float, got {:?}", t.dtype());
    }
}

fn is_float_tensor(t: &FlexTensor) -> bool {
    matches!(t.dtype(), DType::F32 | DType::F64 | DType::F16 | DType::BF16)
}

let out = std::panic::catch_unwind(|| attention(q.clone(), k.clone(), v.clone(), mask, options))
    .unwrap_or_else(|_| attention(q.cast(DType::F32), k.cast(DType::F32), v.cast(DType::F32), mask, options));

When it happens

Trigger: Invoking the flex attention op with query/key/value tensors of an integer or bool dtype. Note the DECLARED-AS hint: a dtype value flowing through burn-autodiff's supports_dtype/dtype_usage hook surface (crates/burn-autodiff/src/backend.rs:100) that was accepted upstream can still reach this kernel unhandled if the inner backend's dtype filtering and the kernel's match disagree.

Common situations: Token ids (integer embeddings indices) passed straight to attention instead of after the embedding lookup; a custom dtype-usage hook allowing a dtype the flex kernel lacks; backend version drift where autodiff forwards dtypes the flex impl never added.

Related errors

• softmax: unsupported dtype {:?}
• burn_flex::layer_norm: unsupported dtype {:?}
• binary_op: unsupported dtype {:?}
• scalar_op: unsupported dtype {:?}
• int_binary_op: unsupported dtype {:?}

AI-assisted analysis of tracel-ai/burn@d16f7ba2ed (2026-09-05). Data as JSON: /api/errors/263d2e004a73e9f7. Report an issue: GitHub.