@tohtana - the HPU is down currently, so I'll remove this test for now.

When will support for flashattn 2.8.0 be available?

Hi @hijkzzz,

It is more about how PyTorch supports flash-attention. With the packing option enabled in OpenRLHF, transformers tried to use flash_attn_varlen_forward. But PyTorch compiler can't compile it. I tried flash-attn v2.8 but the matrix was the same.

Here is the error I got from flash-attention+packing+compile. This error happens without DeepSpeed.

Dynamo failed to run FX node with fake tensors: call_function flash_attn._flash_attn_varlen_forward(*(FakeTensor(..., device='cuda:0', size=(2328, 32, 128), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(2328, 8, 128), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(2328, 8, 128), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(3,), dtype=torch.int32), FakeTensor(..., device='cuda:0', size=(3,), dtype=torch.int32), FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64), FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64), 0.0, 0.08838834764831845), **{'causal': True, 'window_size_left': -1, 'window_size_right': -1, 'softcap': 0.0, 'alibi_slopes': None, 'return_softmax': False, 'block_table': None}): got RuntimeError("flash_attn::_flash_attn_varlen_forward() Expected a value of type 'int' for argument 'max_seqlen_q' but instead found type 'FakeTensor'.\nPosition: 5\nValue: FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64)\nDeclaration: flash_attn::_flash_attn_varlen_forward(Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, SymInt max_seqlen_q, SymInt max_seqlen_k, float dropout_p, float softmax_scale, bool causal, SymInt window_size_left=-1, SymInt window_size_right=-1, float softcap=0., Tensor? alibi_slopes=None, bool return_softmax=False, Tensor? block_table=None, Tensor? leftpad_k=None, Tensor? seqused_k=None, bool zero_tensors=False) -> (Tensor, Tensor, Tensor, Tensor)\nCast error details: Unable to cast Python instance of type <class 'torch._subclasses.fake_tensor.FakeTensor'> to C++ type '?' (#define PYBIND11_DETAILED_ERROR_MESSAGES or compile in debug mode for details)")

from user code:
   File "/home/mtanaka/.local/lib/python3.10/site-packages/transformers/modeling_flash_attention_utils.py", line 279, in torch_dynamo_resume_in__flash_attention_forward_at_272
    attn_output = flash_attn_varlen_func(
  File "/home/mtanaka/.local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 1443, in flash_attn_varlen_func
    return FlashAttnVarlenFunc.apply(
  File "/home/mtanaka/.local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 925, in forward
    out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward(

Hi @hijkzzz,

It is more about how PyTorch supports flash-attention. With the packing option enabled in OpenRLHF, transformers tried to use flash_attn_varlen_forward. But PyTorch compiler can't compile it. I tried flash-attn v2.8 but the matrix was the same.

Here is the error I got from flash-attention+packing+compile. This error happens without DeepSpeed.

Dynamo failed to run FX node with fake tensors: call_function flash_attn._flash_attn_varlen_forward(*(FakeTensor(..., device='cuda:0', size=(2328, 32, 128), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(2328, 8, 128), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(2328, 8, 128), dtype=torch.bfloat16), FakeTensor(..., device='cuda:0', size=(3,), dtype=torch.int32), FakeTensor(..., device='cuda:0', size=(3,), dtype=torch.int32), FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64), FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64), 0.0, 0.08838834764831845), **{'causal': True, 'window_size_left': -1, 'window_size_right': -1, 'softcap': 0.0, 'alibi_slopes': None, 'return_softmax': False, 'block_table': None}): got RuntimeError("flash_attn::_flash_attn_varlen_forward() Expected a value of type 'int' for argument 'max_seqlen_q' but instead found type 'FakeTensor'.\nPosition: 5\nValue: FakeTensor(..., device='cuda:0', size=(), dtype=torch.int64)\nDeclaration: flash_attn::_flash_attn_varlen_forward(Tensor q, Tensor k, Tensor v, Tensor cu_seqlens_q, Tensor cu_seqlens_k, SymInt max_seqlen_q, SymInt max_seqlen_k, float dropout_p, float softmax_scale, bool causal, SymInt window_size_left=-1, SymInt window_size_right=-1, float softcap=0., Tensor? alibi_slopes=None, bool return_softmax=False, Tensor? block_table=None, Tensor? leftpad_k=None, Tensor? seqused_k=None, bool zero_tensors=False) -> (Tensor, Tensor, Tensor, Tensor)\nCast error details: Unable to cast Python instance of type <class 'torch._subclasses.fake_tensor.FakeTensor'> to C++ type '?' (#define PYBIND11_DETAILED_ERROR_MESSAGES or compile in debug mode for details)")

from user code:
   File "/home/mtanaka/.local/lib/python3.10/site-packages/transformers/modeling_flash_attention_utils.py", line 279, in torch_dynamo_resume_in__flash_attention_forward_at_272
    attn_output = flash_attn_varlen_func(
  File "/home/mtanaka/.local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 1443, in flash_attn_varlen_func
    return FlashAttnVarlenFunc.apply(
  File "/home/mtanaka/.local/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 925, in forward
    out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward(

@tohtana Flash-attn is currently a strong dependency for frameworks like OpenRLHF and VerL Alibaba ROLL. Is there any way to work around this issue?

@hijkzzz Yeah, I understand the eager implementation is not a choice. What do you think about SDPF or disabling "packing sample" option? In my environment, DeepCompile+SDPF showed better performance with OpenRLHF than using flash attention.

@hijkzzz Yeah, I understand the eager implementation is not a choice. What do you think about SDPF or disabling "packing sample" option? In my environment, DeepCompile+SDPF showed better performance with OpenRLHF than using flash attention.

@tohtana flash-attn and packing_samples are extremely important for RLHF, and even for SFT... they're not something we can afford to drop for now.

@hijkzzz Sure, I understand the importance. I'm just wondering how much the performance difference with SDPA is (Sorry for the typo in my last message). From my understanding, SDPA calls the flash-attention kernels (or almost equivalent) inside, and you can enable it just by settinging sdpa to attn_implementation if you use HF models. At least, I didn't see much performance difference when I tried it.

@hijkzzz Sure, I understand the importance. I'm just wondering how much the performance difference with SDPA is (Sorry for the typo in my last message). From my understanding, SDPA calls the flash-attention kernels (or almost equivalent) inside, and you can enable it just by settinging sdpa to attn_implementation if you use HF models. At least, I didn't see much performance difference when I tried it.

The packing samples / ring attn directly depends on the API interface of flash attn lib, so we are temporarily unable to do it this way.