Codestin Search App

dxqb · 2025-12-26T18:46:55Z

brings back attention selection, that was removed in 482333f when xformers wasn't used anymore

Can be used to select flash attention, which is faster on windows because torch SDP doesn't support flash internally on Windows

It can be selected, but it must be manually installed from here: https://github.com/zzlol63/flash-attention-prebuild-wheels/releases

thank you to @zzlol63 for the investigation: #1090

Uses huggingface/diffusers#12892 to raise an error if flash attention cannot be used because of attention masks.

dxqb · 2026-01-05T15:25:35Z

pre-built flash-attn can also be downloaded from here https://github.com/mjun0812/flash-attention-prebuild-wheels, which might be the better source because they have automated the build process, so more likely will keep it up to date.

the (automated, I guess) versioning there is a bit strange though. It doesn't always update all platforms. This version seems to be the latest release for windows:
https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/tag/v0.4.19

@Calamdor FYI

dxqb · 2026-01-07T16:19:00Z

diffusers have merged the PR used by this PR

Koratahiu · 2026-02-27T21:41:12Z

Notes after testing (Windows 10, Torch 2.9):

It is incompatible with FP32 LoRA/OFT; I encountered an error stating that Flash is only compatible with FP16/BF16. Setting the LoRA weight to BF16 resolved this.
The speed-up is noticeable, and it worked for Klein.

dxqb · 2026-03-01T07:29:42Z

It is incompatible with FP32 LoRA/OFT; I encountered an error stating that Flash is only compatible with FP16/BF16. Setting the LoRA weight to BF16 resolved this.

that's a bug. can you show how to reproduce?

Koratahiu · 2026-03-01T09:06:39Z

that's a bug. can you show how to reproduce?

Setting LoRA weight to FP32:

Traceback (most recent call last):
  File "\OneTrainer\modules\ui\TrainUI.py", line 716, in __training_thread_function
    trainer.train()
  File "\OneTrainer\modules\trainer\GenericTrainer.py", line 744, in train
    model_output_data = self.model_setup.predict(self.model, batch, self.config, train_progress)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\modules\modelSetup\BaseChromaSetup.py", line 241, in predict
    packed_predicted_flow = model.transformer(
                            ^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1773, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "D:\expermential guff\venv\src\diffusers\src\diffusers\models\transformers\transformer_chroma.py", line 577, in forward
    encoder_hidden_states, hidden_states = block(
                                           ^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1771, in _wrapped_call_impl
    return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 736, in compile_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\nn\modules\module.py", line 1784, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\modules\util\checkpointing_util.py", line 90, in forward
    def forward(self, *args, **kwargs):
  File "\OneTrainer\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 929, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\aot_autograd.py", line 1241, in forward
    return compiled_fn(full_args)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 370, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\utils.py", line 126, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
                            ^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\utils.py", line 100, in g
    return f(*args)
           ^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\autograd\function.py", line 576, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 2074, in forward
    fw_outs = call_func_at_runtime_with_args(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\utils.py", line 126, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
                            ^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 556, in wrapper
    return compiled_fn(runtime_args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_functorch\_aot_autograd\runtime_wrappers.py", line 750, in inner_fn
    outs = compiled_fn(args)
           ^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_inductor\output_code.py", line 584, in __call__
    return self.current_callable(inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_inductor\utils.py", line 2716, in run
    out = model(new_inputs)
          ^^^^^^^^^^^^^^^^^
  File "\AppData\Local\Temp\torchinductor\bq\cbqrfpkp4z55klkpln3fescrhmpvo7357m3j6qqgugdzy6ldyfze.py", line 3114, in call
    buf137 = torch.ops.flash_attn._flash_attn_forward.default(buf135, buf136, buf134, 0.0, 0.08838834764831845, False, -1, -1, 0.0, None, False)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_ops.py", line 829, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_library\autograd.py", line 111, in autograd_impl
    result = forward_no_grad(*args, Metadata(keyset, keyword_only_args))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_library\autograd.py", line 40, in forward_no_grad
    result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_ops.py", line 836, in redispatch
    return self._handle.redispatch_boxed(keyset, *args, **kwargs)  # type: ignore[return-value]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_library\custom_ops.py", line 344, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_compile.py", line 53, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_dynamo\eval_frame.py", line 929, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\torch\_library\custom_ops.py", line 377, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "\OneTrainer\venv\Lib\site-packages\flash_attn\flash_attn_interface.py", line 91, in _flash_attn_forward
    out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.fwd(
                                           ^^^^^^^^^^^^^^^^^^^
RuntimeError: FlashAttention only support fp16 and bf16 data type
Exception raised from mha_fwd at C:\Users\junya\actions-runner\_work\flash-attention-prebuild-wheels\flash-attention-prebuild-wheels\flash-attention\csrc\flash_attn\flash_api.cpp:374 (most recent call first):
00007FFC8050199400007FFC805018F0 c10.dll!c10::Error::Error [<unknown file> @ <unknown line number>]
00007FFC8050040A00007FFC805003B0 c10.dll!c10::detail::torchCheckFail [<unknown file> @ <unknown line number>]
00007FFAF25B42FD00007FFAF25A5EA0 flash_attn_2_cuda.cp312-win_amd64.pyd!c10::ivalue::Object::operator= [<unknown file> @ <unknown line number>]
00007FFAF25C3C9B00007FFAF25BD050 flash_attn_2_cuda.cp312-win_amd64.pyd!PyInit_flash_attn_2_cuda [<unknown file> @ <unknown line number>]
00007FFAF25C3DB400007FFAF25BD050 flash_attn_2_cuda.cp312-win_amd64.pyd!PyInit_flash_attn_2_cuda [<unknown file> @ <unknown line number>]
00007FFAF25AF2BC00007FFAF25A5EA0 flash_attn_2_cuda.cp312-win_amd64.pyd!c10::ivalue::Object::operator= [<unknown file> @ <unknown line number>]
00007FFC780B79E400007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [<unknown file> @ <unknown line number>]
00007FFC7810201800007FFC78101990 python312.dll!PyObject_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC781019C500007FFC78101990 python312.dll!PyObject_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC78102EA500007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC60AF7BBA00007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [<unknown file> @ <unknown line number>]
00007FFC60AEDAC800007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [<unknown file> @ <unknown line number>]
00007FFC60AF15EB00007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [<unknown file> @ <unknown line number>]
00007FFC60AD17D600007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [<unknown file> @ <unknown line number>]
00007FFC601E25F600007FFC601E1730 torch_python.dll!c10::ivalue::Future::devices [<unknown file> @ <unknown line number>]
00007FFC780B79E400007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [<unknown file> @ <unknown line number>]
00007FFC780E560E00007FFC780E556C python312.dll!PyObject_MakeTpCall [<unknown file> @ <unknown line number>]
00007FFC7824673800007FFC78223794 python312.dll!PyThread_tss_is_created [<unknown file> @ <unknown line number>]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC780FBD8E00007FFC780FB9BC python312.dll!PyArg_CheckPositional [<unknown file> @ <unknown line number>]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC7813275F00007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC60AF7B2F00007FFC60ACFF50 torch_python.dll!torch::FunctionParameter::type_name [<unknown file> @ <unknown line number>]
00007FFB5B08924100007FFB5B089080 torch_cpu.dll!c10::Dispatcher::callBoxed [<unknown file> @ <unknown line number>]
00007FFC608BF02500007FFC608BEF60 torch_python.dll!torch::jit::invokeOperatorFromPython [<unknown file> @ <unknown line number>]
00007FFC608BC08D00007FFC608BBF50 torch_python.dll!torch::jit::_get_operation_for_overload_or_packet [<unknown file> @ <unknown line number>]
00007FFC6081AFCA00007FFC608153F0 torch_python.dll!c10d::PythonOnCompletionHook::PythonOnCompletionHook [<unknown file> @ <unknown line number>]
00007FFC6082F6C100007FFC608153F0 torch_python.dll!c10d::PythonOnCompletionHook::PythonOnCompletionHook [<unknown file> @ <unknown line number>]
00007FFC607C820600007FFC6074CED0 torch_python.dll!THPPointer<_frame>::dup [<unknown file> @ <unknown line number>]
00007FFC601E25F600007FFC601E1730 torch_python.dll!c10::ivalue::Future::devices [<unknown file> @ <unknown line number>]
00007FFC780B79E400007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [<unknown file> @ <unknown line number>]
00007FFC781327A600007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC78133E8B00007FFC78133E18 python312.dll!PyObject_FastCallDictTstate [<unknown file> @ <unknown line number>]
00007FFC781D319B00007FFC781D311C python312.dll!PyObject_Call_Prepend [<unknown file> @ <unknown line number>]
00007FFC781D30C600007FFC781D2DF0 python312.dll!PyDictProxy_New [<unknown file> @ <unknown line number>]
00007FFC7810201800007FFC78101990 python312.dll!PyObject_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC781019C500007FFC78101990 python312.dll!PyObject_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC78102EA500007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC78133E8B00007FFC78133E18 python312.dll!PyObject_FastCallDictTstate [<unknown file> @ <unknown line number>]
00007FFC781D319B00007FFC781D311C python312.dll!PyObject_Call_Prepend [<unknown file> @ <unknown line number>]
00007FFC781D30C600007FFC781D2DF0 python312.dll!PyDictProxy_New [<unknown file> @ <unknown line number>]
00007FFC7810201800007FFC78101990 python312.dll!PyObject_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC781019C500007FFC78101990 python312.dll!PyObject_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC78102EA500007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]
00007FFC7810103C00007FFC78100EC0 python312.dll!PyFunction_Vectorcall [<unknown file> @ <unknown line number>]
00007FFC7813281500007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC780893B700007FFC7808931C python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC6069275F00007FFC6068D170 torch_python.dll!torch::autograd::registerFunctionPreHook [<unknown file> @ <unknown line number>]
00007FFC780B7A1600007FFC780B73F4 python312.dll!PyThread_acquire_lock_timed [<unknown file> @ <unknown line number>]
00007FFC781327A600007FFC781326F0 python312.dll!PyObject_Call [<unknown file> @ <unknown line number>]
00007FFC781066EE00007FFC78102610 python312.dll!PyEval_EvalFrameDefault [<unknown file> @ <unknown line number>]

…nto Nerogar:merge Adds AttentionMechanism (SDP/FLASH) selection, restoring the explicit attention-backend choice removed when xformers was dropped. Rebased from PR Nerogar#1227 (originally against master) onto Nerogar:merge; the UI addition was re-targeted at BaseTrainingTabView (the ctk/Qt6 split from PR Nerogar#1566) since the old TrainingTab.py no longer exists on this base. Co-Authored-By: Claude Sonnet 5 <[email protected]>

…verage - Add AttentionMechanism.CUDNN (diffusers "_native_cudnn" backend). - Make _set_attention_backend's mask argument required (no default) so every call site has to make an explicit decision; drop the unused varlen parameter (FLASH hard-raises on any mask, it doesn't "maybe fail"). - Roll out to the models the original PR missed: Ernie (mask=True, its transformer builds a real attention mask internally from text_lens), Flux2 (mask=False), Wuerstchen (mask=False, StableCascadeUNet.forward has no attention_mask parameter at all). - Switch the UI dropdown to options_kv (same pattern as LoraTabController.get_peft_types()) with human-readable labels sourced from TrainingTabController.get_attention_mechanisms(), keeping the stored/serialized enum identity unchanged. Co-Authored-By: Claude Sonnet 5 <[email protected]>

dxqb · 2026-07-01T19:34:03Z

RuntimeError: FlashAttention only support fp16 and bf16 data type

reported upstream: huggingface/diffusers#14104

dxqb · 2026-07-01T19:34:24Z

updated and rebased because Krea 2 needs cuDNN attention

Removed an extra newline before setting the attention backend.

dxqb linked an issue Dec 26, 2025 that may be closed by this pull request

[Feat]: Attention backend selection for Diffusers #1090

Open

dxqb mentioned this pull request Jan 5, 2026

Add flash-attn support for Windows #1107

Closed

dxqb marked this pull request as ready for review January 7, 2026 18:22

dxqb commented Feb 11, 2026

View reviewed changes

Comment thread modules/modelSetup/BaseModelSetup.py

dxqb marked this pull request as draft March 1, 2026 07:46

dxqb changed the title ~~Attention backend selection~~ Attention backend selection / flash attention on Windows Mar 9, 2026

dxqb force-pushed the attn_selection branch from cd2bd85 to 58d996b Compare July 1, 2026 18:21

dxqb changed the base branch from master to merge July 1, 2026 18:22

dxqb marked this pull request as ready for review July 1, 2026 19:34

dxqb mentioned this pull request Jul 1, 2026

Flux-family attention: flash-attn and other backends fail in an autocast context huggingface/diffusers#14104

Open

dxqb added 2 commits July 1, 2026 21:42

Remove unnecessary blank line in BaseHunyuanVideoSetup.py

1efc82f

Remove extra newline in BaseWuerstchenSetup.py

0c05969

Removed an extra newline before setting the attention backend.

dxqb merged commit 58955db into Nerogar:merge Jul 1, 2026
1 check passed

dxqb deleted the attn_selection branch July 1, 2026 19:57

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Attention backend selection / flash attention on Windows#1227

Attention backend selection / flash attention on Windows#1227
dxqb merged 4 commits into
Nerogar:mergefrom
dxqb:attn_selection

dxqb commented Dec 26, 2025 •

edited

Loading

Uh oh!

dxqb commented Jan 5, 2026

Uh oh!

dxqb commented Jan 7, 2026

Uh oh!

Uh oh!

Koratahiu commented Feb 27, 2026

Uh oh!

dxqb commented Mar 1, 2026

Uh oh!

Koratahiu commented Mar 1, 2026

Uh oh!

dxqb commented Jul 1, 2026

Uh oh!

dxqb commented Jul 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

dxqb commented Dec 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

dxqb commented Jan 5, 2026

Uh oh!

dxqb commented Jan 7, 2026

Uh oh!

Uh oh!

Koratahiu commented Feb 27, 2026

Uh oh!

dxqb commented Mar 1, 2026

Uh oh!

Koratahiu commented Mar 1, 2026

Uh oh!

dxqb commented Jul 1, 2026

Uh oh!

dxqb commented Jul 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

dxqb commented Dec 26, 2025 •

edited

Loading