add SP support for `_flash_3_varlen_hub` backend by zhtmike · Pull Request #13809 · huggingface/diffusers

zhtmike · 2026-05-26T06:18:49Z

What does this PR do?

A follow up work for #13479. I have added _flash_3_varlen_hub support for SP forward & backward.

Tested with QwenImage pipeline, the result image is expected.
Tested with QwenImage training with SP, there is no error.
The UTs for Flux and QwenImage are passed.

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline?
Did you read our philosophy doc (important for complex PRs)?
Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

@sayakpaul

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

sayakpaul · 2026-05-29T11:38:17Z

- if attn_mask is not None:
- attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
-
- (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
- _prepare_for_flash_attn_or_sage_varlen(
- batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
- )
- )
-
- key_valid, value_valid = [], []
- for b in range(batch_size):
- valid_len = seqlens_k[b]
- key_valid.append(key[b, :valid_len])
- value_valid.append(value[b, :valid_len])


This seems like should come under if _parallel_config is None and attn_mask is not None:?

There is a little refactor in this PR, and basically it is doing the same thing as
#13479 (comment)

sayakpaul · 2026-05-29T11:39:14Z

+ max_seqlen_k=max_seqlen_k,
+ softmax_scale=scale,
+ causal=is_causal,
+ return_attn_probs=return_lse,


This seems like an extra argument?

yes, we following the other implementation's style.
And actually return_attn_probs by default is False (see `https://github.com/huggingface/kernels-community/blob/main/flash-attn3/torch-ext/flash_attn3/flash_attn_interface.py#L648)
So it will only return single tensor (see https://github.com/huggingface/kernels-community/blob/main/flash-attn3/torch-ext/flash_attn3/flash_attn_interface.py#L691),
the previous code should not be runnable.

sayakpaul · 2026-05-29T11:39:32Z

+ return_attn_probs=return_lse,
+ )
+ if return_lse:
+ out, lse, *_ = out


Why do we need to initialize lse = None above?

It is a guard similar to _native_cudnn_attention, _sage_attention, etc. To make the code in same style

sayakpaul · 2026-05-29T11:41:23Z

@askserge could you do a review?

github-actions

🤗 Serge says:

This PR adds sequence parallel (context parallel) support for the _flash_3_varlen_hub attention backend, following the same pattern established by the existing _flash_varlen_hub (flash-attn2) and _flash_3_hub implementations. The code is well-structured and closely mirrors the existing patterns.

Correctness

Potential bug: indices_k used before assignment on the no-mask path. In _flash_attention_3_varlen_hub_forward_op, when attn_mask is None, the variable indices_k is never assigned, but at line 1721 ctx.indices_k = indices_k if attn_mask is not None else None — this is actually fine because the conditional guards it. However, if _save_ctx is False, indices_k is never referenced at all on the no-mask path, so there's no issue. This matches the flash-attn2 varlen pattern exactly.
Positional argument fragility in wrapped_forward_fn call. The non-varlen _flash_attention_3_hub_forward_op uses keyword arguments for causal, window_size_left, etc., but the new varlen forward op passes everything positionally (lines 1677–1712). This makes the code harder to read and more fragile if the upstream _flash_attn_forward signature changes. Consider using keyword arguments for at least the trailing parameters, consistent with the non-varlen version.
return_lse handling change in the non-SP path. The original code always unpacked out, lse, *_ from the function call. The new code passes return_attn_probs=return_lse and conditionally unpacks. This is a behavioral change for the non-SP path — if return_lse=False, the function now returns a single tensor instead of a tuple. This should be verified to work correctly with the flash_attn_varlen_func API. The flash-attn2 varlen hub uses return_attn_probs=return_lse similarly, so this is likely correct.

Tests

The new backend is properly added to ContextParallelAttentionBackendsTesterMixin and the ring_degree skip logic.
The _FLASH_3_VARLEN_HUB is added to the hub kernels set in tests/models/testing_utils/utils.py.

Minor Issues

Bug in existing test code (pre-existing). Line 413 (deleted): attention_backend in ("flash_varlen_hub") — using in with a parenthesized string (not a tuple) means this is just attention_backend in "flash_varlen_hub", which checks character membership, not string membership. The fix on line 417 adds a proper tuple ("flash_varlen_hub", "_flash_3_varlen_hub"), which is correct. However, the original single-element check was buggy. Good that it's fixed now.

Overall the implementation follows established patterns well and looks correct.

17 LLM turns · 18 tool calls · 92.5s · 309778 in / 4124 out tokens

github-actions · 2026-05-29T11:43:24Z

+ max_seqlen_k=max_seqlen_k,
+ softmax_scale=scale,
+ causal=is_causal,
+ return_attn_probs=return_lse,


Note: The original code always unpacked out, lse, *_ = func(...). Now with return_attn_probs=return_lse, when return_lse=False the return value may be different (single tensor vs tuple). Make sure flash_attn_varlen_func from flash-attn3 returns a single tensor (not a tuple) when return_attn_probs=False. The flash-attn2 varlen hub uses the same pattern, so this is likely fine, but worth verifying.

Verified with the source code
https://github.com/huggingface/kernels-community/blob/main/flash-attn3/torch-ext/flash_attn3/flash_attn_interface.py#L648
https://github.com/huggingface/kernels-community/blob/main/flash-attn3/torch-ext/flash_attn3/flash_attn_interface.py#L691

zhtmike · 2026-06-03T03:45:34Z

Hi @sayakpaul sorry for late update. The review should be addressed, can you take a look?

sayakpaul · 2026-06-04T02:07:35Z

Waiting for PR #13756.

zhtmike · 2026-06-12T06:46:39Z

Waiting for PR #13756.

Per the discussion in PR #13756, I have added a precision test in ContextParallelAttentionBackendsTesterMixin.

I have tested with:

tests/models/transformers/test_models_transformer_qwenimage.py::TestQwenImageTransformerContextParallelAttnBackends
tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerContextParallelAttnBackends

and the tests passed.

Meanwhile, there is a breaking change in huggingface/kernels-community@1471f94 that causes FA2 to fail to run with recent pre‑built kernels (missing _wrapped_flash_attn_forward), so I have also fixed that in the PR. Please take a look!

zhtmike · 2026-06-29T02:48:41Z

@sayakpaul please take a look if you have time~

sayakpaul · 2026-07-03T03:21:49Z

Could you please resolve the conflicts?

sergereview

🤗 Serge says:

Clean, well-scoped port of the existing FLASH_VARLEN_HUB context-parallel implementation to the _FLASH_3_VARLEN_HUB backend. The new forward/backward ops mirror the established FA3 (_flash_attention_3_hub_*) and FA2-varlen (_flash_varlen_attention_hub_*) patterns closely, and the wiring all cross-checks:

The FA3 _flash_attn_forward/_flash_attn_backward positional call signatures match those already used by _flash_attention_3_hub_forward_op/_backward_op, with the varlen-specific cu_seqlens/max_seqlen args inserted in the right positions.
ctx.indices_k = indices_k if attn_mask is not None else None is safe — the ternary short-circuits before indices_k (only defined in the masked branch) is evaluated.
The non-CP path change to gate lse via return_attn_probs=return_lse and unpack conditionally is correct and avoids requesting lse unnecessarily.
ring_degree > 1 is explicitly rejected in both the dispatch function and the test skip list; ulysses is the supported path.
functools is imported; the registry wrapped_*_attr entries point at the FA3 flash_attn_interface._flash_attn_{forward,backward} which matches the long-signature calls.

Tests

The _context_parallel_worker already accepts the new state_dict positional arg, so the added argument at the call site is consistent. Capturing weights from the (bf16-cast) reference model and reloading them into a fresh model in the worker before its own cast ensures both sides use identical weights — this upgrades the test from a shape-only check to a real numerical assert_close, which is a genuine improvement.

No correctness or security concerns found. The description's claims (SP forward/backward support, Flux/QwenImage tests) are consistent with the diff.

serge v0.1.0 · model: claude-opus-4-8 · 14 LLM turns · 14 tool calls · 89.5s · 439489 in / 5510 out tokens

# Conflicts: # src/diffusers/models/attention_dispatch.py

zhtmike · 2026-07-03T04:11:23Z

Could you please resolve the conflicts?

Done!

sayakpaul · 2026-07-03T04:17:01Z

+ wrapped_forward_attr="flash_attn_interface._flash_attn_forward",
+ wrapped_backward_attr="flash_attn_interface._flash_attn_backward",


Can you explain these changes?

there is a breaking change in huggingface/kernels-community@1471f94 in PR huggingface/kernels-community#877 that causes FA2 to fail to run with recent pre‑built kernels (missing _wrapped_flash_attn_forward), so I have to fix it here to pass the tests.

Cc: @drbh. Does this need a version bump?

Just run with the latest kernels version; it seems to be fixed. I’ve reverted the changes.

sayakpaul · 2026-07-03T04:17:06Z

- wrapped_forward_attr="flash_attn_interface._wrapped_flash_attn_varlen_forward",
- wrapped_backward_attr="flash_attn_interface._wrapped_flash_attn_varlen_backward",
+ wrapped_forward_attr="flash_attn_interface._flash_attn_varlen_forward",
+ wrapped_backward_attr="flash_attn_interface._flash_attn_varlen_backward",


sayakpaul · 2026-07-03T04:18:02Z

 AttentionBackendName._FLASH_3_VARLEN_HUB,
 constraints=[_check_device, _check_qkv_dtype_bf16_or_fp16, _check_shape],
- supports_context_parallel=False,
+ supports_context_parallel=True,


sayakpaul · 2026-07-03T04:19:06Z

 batch_size, seq_len_q, _, _ = query.shape
 _, seq_len_kv, _, _ = key.shape

- if attn_mask is not None:


Sorry I have lost the context here since it's been a while. Can you explain why do we need these changes? I don't think the if _parallel_config is None: branch is doing what we had previously.

For if _parallel_config is None: branch, I major do a fix for non-contiguous mark handling for models like Qwen-Image, following the pattern of https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_dispatch.py#L2783.
Both implementations are equivalent for contiguous masks, but the previous implementation was buggy for non‑contiguous masks.

Can that come in a separate PR before we merge this PR? I think the two are slightly unrelated?

Please check #14115

zhtmike · 2026-07-08T03:06:29Z

Updated the PR with latest main, tested with tests/models/transformers/test_models_transformer_qwenimage.py::TestQwenImageTransformerContextParallelAttnBackends and tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerContextParallelAttnBackends, they are passing

sayakpaul · 2026-07-08T03:07:27Z

tests/models/transformers/test_models_transformer_qwenimage.py::TestQwenImageTransformerContextParallelAttnBackends

For this I think only varlen applies since Qwen has masks?

zhtmike · 2026-07-08T03:08:08Z

tests/models/transformers/test_models_transformer_qwenimage.py::TestQwenImageTransformerContextParallelAttnBackends
For this I think only varlen applies since Qwen has masks?

yeah right!

HuggingFaceDocBuilderDev · 2026-07-08T03:14:37Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

sayakpaul

Ran pytest tests/models/transformers/ -k "test_context_parallel_attn_backend_inference" on 2 H100s and all were green! Let's go.

add _flash_3_varlen_hub support for SP
40618ff

github-actionsBot added size/L PR with diff > 200 LOC models tests and removed size/L PR with diff > 200 LOC labels May 26, 2026

sayakpaul reviewed May 29, 2026
View reviewed changes

github-actionsBot reviewed May 29, 2026
View reviewed changes

Merge branch 'main' into fa3
330eb13

github-actionsBot added the size/L PR with diff > 200 LOC label Jun 3, 2026

zhtmike added 2 commits June 3, 2026 03:31

make signiture clear
a7a3b9c

update
83ef0c7

zhtmike added 6 commits June 11, 2026 14:08

Merge branch 'main' into fa3
e21870c

add accuracy test
363f8fd

fix dtype
2a3c2b5

loosen threshold because of bfloat16
2a0dfb1

fix the break change ofhuggingface/kernels-community@1471f94
aceb39d

fix stale error message
300427f

sergereviewBot reviewed Jul 3, 2026
View reviewed changes

Merge remote-tracking branch 'official/main' into fa3
6dbc061
# Conflicts: # src/diffusers/models/attention_dispatch.py

github-actionsBot added lora examples labels Jul 3, 2026

zhtmike force-pushed the fa3 branch from 22c98b9 to 4d2d7aaCompare July 3, 2026 04:03

github-actionsBot removed lora examples labels Jul 3, 2026

fix merge
ac34358

zhtmike force-pushed the fa3 branch from 4d2d7aa to ac34358Compare July 3, 2026 04:08

sayakpaul reviewed Jul 3, 2026
View reviewed changes

zhtmike mentioned this pull request Jul 3, 2026
fix _flash_3_varlen_hub mask handling #14115
Merged
9 tasks

zhtmike added 2 commits July 8, 2026 01:54

Merge branch 'main' into fa3
f0f36af

revert change
f53378c

qflen mentioned this pull request Jul 8, 2026
Non-hub varlen attention backends mispack keys for non-contiguous masks (missed in #14114/#14115) #14139
Open

Merge branch 'main' into fa3
1a81f29

sayakpaul approved these changes Jul 8, 2026
View reviewed changes

sayakpaul merged commit 59d4e5a into huggingface:mainJul 8, 2026
13 of 15 checks passed

zhtmike deleted the fa3 branch July 9, 2026 02:48

zhtmike mentioned this pull request Aug 25, 2026
[RFC][tracking] Upstream patch debt: audited ledger of temporary compatibility patches verl-project/verl-omni#445
Open
37 tasks

		wrapped_forward_attr="flash_attn_interface._flash_attn_forward",
		wrapped_backward_attr="flash_attn_interface._flash_attn_backward",

Uh oh!

Conversation

zhtmike commented May 26, 2026

What does this PR do?

Before submitting

Who can review?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhtmikeJun 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhtmikeJun 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented May 29, 2026

Uh oh!

github-actionsBot left a comment

Choose a reason for hiding this comment

Correctness

Tests

Minor Issues

Uh oh!

Uh oh!

Uh oh!

github-actionsBotMay 29, 2026

Choose a reason for hiding this comment

Uh oh!

zhtmikeJun 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhtmike commented Jun 3, 2026

Uh oh!

sayakpaul commented Jun 4, 2026

Uh oh!

zhtmike commented Jun 12, 2026

Uh oh!

zhtmike commented Jun 29, 2026

Uh oh!

sayakpaul commented Jul 3, 2026

Uh oh!

sergereviewBot left a comment

Choose a reason for hiding this comment

Uh oh!

zhtmike commented Jul 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhtmikeJul 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zhtmikeJul 3, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

zhtmikeJun 3, 2026•
edited
Loading

zhtmikeJun 3, 2026•
edited
Loading

zhtmikeJun 3, 2026•
edited
Loading

zhtmikeJul 3, 2026•
edited
Loading

zhtmikeJul 3, 2026•
edited
Loading