{"record":{"id":"ee0ff84521b3836a","repo":"huggingface/transformers","slug":"attention-sinks-cannot-be-run-on-cpu-with-flex-att","errorCode":null,"errorMessage":"Attention sinks cannot be run on CPU with flex attention. Please switch to a different device, e.g. CUDA","messagePattern":"Attention sinks cannot be run on CPU with flex attention\\. Please switch to a different device, e\\.g\\. CUDA","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"src/transformers/integrations/flex_attention.py","lineNumber":316,"sourceCode":"        # because it requires operating on the full attention matrix before softmax.\n        # ==> this is done after flex attention\n        return score\n\n    enable_gqa = True\n    num_local_query_heads = query.shape[1]\n\n    # When running TP this helps:\n    if (num_local_query_heads & (num_local_query_heads - 1)) != 0:\n        key = repeat_kv(key, query.shape[1] // key.shape[1])\n        value = repeat_kv(value, query.shape[1] // value.shape[1])\n        enable_gqa = False\n\n    kernel_options = kwargs.get(\"kernel_options\")\n    # On CPU we must skip returning LSE due to a runtime issue; elsewhere, follow PyTorch API and return it\n    return_lse = query.device.type != \"cpu\"\n\n    if not return_lse and s_aux is not None:\n        raise ValueError(\n            \"Attention sinks cannot be run on CPU with flex attention. Please switch to a different device, e.g. CUDA\"\n        )\n\n    flex_attention_output = compile_friendly_flex_attention(\n        query,\n        key,\n        value,\n        score_mod=score_mod,\n        block_mask=block_mask,\n        enable_gqa=enable_gqa,\n        scale=scaling,\n        kernel_options=kernel_options,\n        # Last time checked on PyTorch == 2.5.1: Flex Attention always computes the lse regardless.\n        # For simplification, we thus always return it as no additional computations are introduced.\n        training=module.training,\n        # inject the lse args\n        **get_flex_attention_lse_kwargs(return_lse),\n    )","sourceCodeStart":298,"sourceCodeEnd":334,"githubUrl":"https://github.com/huggingface/transformers/blob/a597f974857b3d92939971296bc0deb93d33d780/src/transformers/integrations/flex_attention.py#L298-L334","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Move the model to a CUDA device: model.to(\"cuda\")","On CPU, use a different attention implementation (sdpa/eager) — omit attn_implementation=\"flex_attention\""],"exampleFix":"# before\nmodel = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=\"flex_attention\")\nmodel(**inputs)  # on CPU -> ValueError\n\n# after\nmodel = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=\"flex_attention\").to(\"cuda\")\nmodel(**{k: v.to(\"cuda\") for k, v in inputs.items()})","handlingStrategy":"validation","validationCode":"import torch\nattn_impl = \"flex_attention\" if torch.cuda.is_available() else \"sdpa\"\nmodel = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation=attn_impl)","typeGuard":"def flex_attention_ok(config, device) -> bool:\n    \"\"\"Sink-attention models need LSE, which flex_attention cannot return on CPU.\"\"\"\n    has_sinks = getattr(config, \"attention_sink\", None) is not None or getattr(config, \"use_attention_sinks\", False)\n    return (not has_sinks) or device.type != \"cpu\"","tryCatchPattern":null,"preventionTips":["Select flex_attention only when the model runs on CUDA","Make CPU test fixtures use sdpa rather than forcing flex_attention"],"tags":["flex-attention","attention-sinks","cpu","device","lse"],"backgroundTag":null,"analyzedSha":"a597f974857b3d92939971296bc0deb93d33d780","analyzedAt":"2026-08-14T18:24:08.354Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}