System Info

[Bug] Flash Attention crashes with illegal memory access on Qwen3.5 due to 3D position_ids being misinterpreted as packed sequence

We fixed it in https://github.com/ouroborosscr/transformers/tree/fix/qwen35-flash-attn-3d-position-ids

Description

When using attn_implementation="flash_attention_2" with Qwen3.5 models, all forward passes crash with CUDA error: an illegal memory access was encountered. This affects both training and inference.

Root cause: Qwen3.5 uses a hybrid architecture (GatedDeltaNet linear attention + standard attention) and passes 3D position_ids with shape [3, batch_size, seq_len] (for multi-dimensional rotary embedding). The function _is_packed_sequence() in modeling_flash_attention_utils.py misinterprets this 3D tensor as a packed sequence indicator, causing cu_seqlens to be constructed with 3× the actual token count. Flash Attention then reads beyond the q/k/v tensor boundaries, resulting in an illegal memory access.

Reproduction

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3.5-9B",  # or any Qwen3.5 variant
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    ),
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    attn_implementation="flash_attention_2",
)

# This crashes immediately
input_ids = torch.randint(100, 5000, (1, 256), device="cuda")
with torch.no_grad():
    out = model(input_ids=input_ids, use_cache=False)

Error:

torch.AcceleratorError: CUDA error: an illegal memory access was encountered

Traceback (abbreviated):

File "transformers/modeling_flash_attention_utils.py", line 692, in _flash_attention_forward
    out = flash_varlen_fn(
File "flash_attn/flash_attn_interface.py", line 1443, in flash_attn_varlen_func
    return FlashAttnVarlenFunc.apply(
File "flash_attn/flash_attn_interface.py", line 165, in _flash_attn_varlen_forward
    out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.varlen_fwd(
torch.AcceleratorError: CUDA error: an illegal memory access was encountered

Root Cause Analysis

Qwen3.5 hybrid architecture

Qwen3.5 uses a mixed attention architecture: 24 layers of Qwen3_5GatedDeltaNet (linear attention) and 8 layers of Qwen3_5Attention (standard attention, at layers 3, 7, 11, 15, 19, 23, 27, 31). Only the standard attention layers use flash attention.

Qwen3.5 passes 3D position_ids with shape [3, batch_size, seq_len] for its multi-dimensional rotary embedding (3 sets of position indices).

The bug

In modeling_flash_attention_utils.py, the function _is_packed_sequence() (line 444) does not handle tensors with more than 2 dimensions:

def _is_packed_sequence(position_ids, batch_size):
    if position_ids is None:
        return False
    increasing_position_sequences = (
        torch.arange(position_ids.shape[1], device=position_ids.device) + position_ids.min()
    )
    return batch_size == 1 and (increasing_position_sequences - position_ids).abs().sum().bool()

When position_ids has shape [3, 1, 256]:

• position_ids.shape[1] returns 1 (not 256 as expected for a 2D [batch, seq_len] tensor)
• The function returns True, misidentifying this as a packed sequence

This triggers the packed-sequence code path at line 677:

elif is_fa_with_varlen_kwargs or is_fa_with_position_ids:
    if cu_seq_lens_q is None or cu_seq_lens_k is None:
        q, k, v, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _prepare_from_posids(
            query_states, key_states, value_states, position_ids
        )

Inside prepare_fa_kwargs_from_position_ids() (line 362):

position_ids = position_ids.reshape(-1)  # [3, 1, 256] → [768]
indices_q = (position_ids == 0).nonzero().view(-1)  # Finds 3 zero positions

This constructs cu_seqlens = [0, 256, 512, 768], claiming 3 sequences with 768 total tokens. But the actual q/k/v tensors only contain 256 tokens. Flash Attention reads up to index 768, causing the illegal memory access.

Intercepted parameters confirming the mismatch

🔍 varlen_fwd parameters:
  q: torch.Size([256, 16, 256])           ← 256 tokens
  cu_seqlens_q: tensor([0, 256, 512, 768]) ← claims 768 tokens
  q total=256 vs cu_seqlens_q[-1]=768      ← MISMATCH → crash

Fix

Add a dimensionality check in _is_packed_sequence() to reject tensors with more than 2 dimensions, since packed sequences always use 2D position_ids [batch, seq_len]:

def _is_packed_sequence(position_ids, batch_size):
    if position_ids is None:
        return False
    if position_ids.dim() > 2:
        return False
    increasing_position_sequences = (
        torch.arange(position_ids.shape[1], device=position_ids.device) + position_ids.min()
    )
    return batch_size == 1 and (increasing_position_sequences - position_ids).abs().sum().bool()

This fix has been validated: all 8 standard attention layers in Qwen3.5-9B pass flash attention forward successfully after applying the patch.

Environment

• Model: Qwen3.5-9B (hybrid GatedDeltaNet + standard attention)
• GPU: NVIDIA A100-SXM4-80GB
• PyTorch: 2.9.0 / 2.10.0 (both affected)
• Transformers: 5.3.0
• flash-attn: 2.8.3
• CUDA: 12.8

Impact

• Affects all Qwen3.5 variants (and potentially any future model using >2D position_ids)
• Blocks both training and inference when using flash_attention_2
• No workaround other than falling back to sdpa or eager attention implementations

Who can help?

@vasqu @ArthurZucker (attention)

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

GRPO reinforcement learning training with Qwen3.5-9B using TRL GRPOTrainer When using attn_implementation="flash_attention_2" with Qwen3.5, all forward passes crash with CUDA illegal memory access. Minimal reproduction:

import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3.5-9B",
    quantization_config=BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
    ),
    torch_dtype=torch.bfloat16,
    device_map={"": 0},
    attn_implementation="flash_attention_2",
)

input_ids = torch.randint(100, 5000, (1, 256), device="cuda")
with torch.no_grad():
    out = model(input_ids=input_ids, use_cache=False)  # crashes here

Error:

torch.AcceleratorError: CUDA error: an illegal memory access was encountered

Root cause:

Qwen3.5 is a hybrid architecture (24 GatedDeltaNet layers + 8 standard attention layers). It uses 3D position_ids with shape [3, batch_size, seq_len] for multi-dimensional rotary embedding.

_is_packed_sequence() in modeling_flash_attention_utils.py (line 444) does not handle >2D tensors:

def _is_packed_sequence(position_ids, batch_size):
    if position_ids is None:
        return False
    increasing_position_sequences = (
        torch.arange(position_ids.shape[1], device=position_ids.device) + position_ids.min()
    )
    return batch_size == 1 and (increasing_position_sequences - position_ids).abs().sum().bool()

When position_ids has shape [3, 1, 256], position_ids.shape[1] returns 1 instead of the sequence length, and the function returns True, misidentifying this as a packed sequence.

This triggers prepare_fa_kwargs_from_position_ids() which does:

position_ids = position_ids.reshape(-1)  # [3, 1, 256] → [768]
indices_q = (position_ids == 0).nonzero()  # finds 3 zero positions
# constructs cu_seqlens = [0, 256, 512, 768] — claims 768 tokens

But q/k/v only contain 256 tokens. Flash attention reads up to index 768, causing the illegal memory access.

Intercepted evidence:

q: torch.Size([256, 16, 256])           ← 256 tokens
cu_seqlens_q: tensor([0, 256, 512, 768]) ← claims 768 tokens
q total=256 vs cu_seqlens_q[-1]=768      ← MISMATCH → crash

Environment:

• Model: Qwen3.5-9B
• GPU: NVIDIA A100-SXM4-80GB
• PyTorch: 2.9.0 and 2.10.0 (both affected)
• transformers: 5.3.0
• flash-attn: 2.8.3
• CUDA: 12.8

Fix: Add if position_ids.dim() > 2: return False in _is_packed_sequence(). Packed sequences always use 2D [batch, seq_len] position_ids.

Expected behavior

Model forward pass with attn_implementation="flash_attention_2" should complete successfully without CUDA errors. After the fix (adding dimensionality check for >2D position_ids), all 8 standard attention layers in Qwen3.5-9B pass flash attention forward correctly.