如果令牌数量过大,则会发生 RuntimeError
Traceback (most recent call last): File "/home/t-ykong/remesh_datamaker/1.py", line 22, in feat = flash_attn.flash_attn_varlen_qkvpacked_func( File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 1267, in flash_attn_varlen_qkvpacked_func return FlashAttnVarlenQKVPackedFunc.apply( File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/autograd/function.py", line 575, in apply return super().apply(*args, **kwargs) # type: ignore[misc] File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 553, in forward out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_varlen_forward( File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_ops.py", line 1158, in call return self._op(*args, **(kwargs or {})) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_library/autograd.py", line 113, in autograd_impl result = forward_no_grad(*args, Metadata(keyset, keyword_only_args)) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_library/autograd.py", line 40, in forward_no_grad result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_ops.py", line 761, in redispatch return self._handle.redispatch_boxed(keyset, *args, **kwargs) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_library/custom_ops.py", line 335, in backend_impl result = self._backend_fns[device_type](*args, **kwargs) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_compile.py", line 51, in inner return disable_fn(*args, **kwargs) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 838, in _fn return fn(*args, **kwargs) File "/home/t-ykong/anaconda3/envs/py310/lib/python3.10/site-packages/torch/_library/custom_ops.py", line 367, in wrapped_fn return fn(*args, **kwargs) File
内容来源: Dao-AILab/flash-attention