huggingface/transformers · error · ValueError

Attention sinks cannot be run on CPU with flex attention. Pl

Error message

Attention sinks cannot be run on CPU with flex attention. Please switch to a different device, e.g. CUDA

What it means

With attention sinks (the s_aux stream in Gemma-3/Nemotron-style sink attention), the flex_attention path must return the LSE (log-sum-exp) tensor for the sink computation. On CPU the kernel options disable returning LSE due to a PyTorch runtime issue (return_lse = query.device.type != "cpu"), so a model with attention sinks cannot run under flex_attention on CPU — the code detects this combination (not return_lse and s_aux is not None) and raises ValueError directing you to CUDA.

Source

Thrown at src/transformers/integrations/flex_attention.py:316

        # because it requires operating on the full attention matrix before softmax.
        # ==> this is done after flex attention
        return score

    enable_gqa = True
    num_local_query_heads = query.shape[1]

    # When running TP this helps:
    if (num_local_query_heads & (num_local_query_heads - 1)) != 0:
        key = repeat_kv(key, query.shape[1] // key.shape[1])
        value = repeat_kv(value, query.shape[1] // value.shape[1])
        enable_gqa = False

    kernel_options = kwargs.get("kernel_options")
    # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it
    return_lse = query.device.type != "cpu"

    if not return_lse and s_aux is not None:
        raise ValueError(
            "Attention sinks cannot be run on CPU with flex attention. Please switch to a different device, e.g. CUDA"
        )

    flex_attention_output = compile_friendly_flex_attention(
        query,
        key,
        value,
        score_mod=score_mod,
        block_mask=block_mask,
        enable_gqa=enable_gqa,
        scale=scaling,
        kernel_options=kernel_options,
        # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.
        # For simplification, we thus always return it as no additional computations are introduced.
        training=module.training,
        # inject the lse args
        **get_flex_attention_lse_kwargs(return_lse),
    )

View on GitHub (pinned to a597f97485)

Solutions

1. Move the model to a CUDA device: model.to("cuda")
2. On CPU, use a different attention implementation (sdpa/eager) — omit attn_implementation="flex_attention"

Example fix

# before
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="flex_attention")
model(**inputs)  # on CPU -> ValueError

# after
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="flex_attention").to("cuda")
model(**{k: v.to("cuda") for k, v in inputs.items()})

import torch
attn_impl = "flex_attention" if torch.cuda.is_available() else "sdpa"
model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=attn_impl)

def flex_attention_ok(config, device) -> bool:
    """Sink-attention models need LSE, which flex_attention cannot return on CPU."""
    has_sinks = getattr(config, "attention_sink", None) is not None or getattr(config, "use_attention_sinks", False)
    return (not has_sinks) or device.type != "cpu"

When it happens

Trigger: Loading a sink-attention model (e.g. Gemma 3, certain Nemotron configs) with attn_implementation="flex_attention" and running it on CPU: the forward passes s_aux, but return_lse is False on CPU, triggering the guard.

Common situations: Testing locally on a laptop/CPU-only box with flex_attention (e.g. for block-sparse mask prototyping) on a sink-attention architecture; CPU unit tests that force the flex implementation.

Related errors

• Attempting to cast a BatchFeature to type {str(arg)}. This i
• `flex_attention` does not support `dropout`. Please use it w

AI-assisted analysis of huggingface/transformers@a597f97485 (2026-08-14). Data as JSON: /api/errors/ee0ff84521b3836a. Report an issue: GitHub.