[Enhancement] Get rid of "paged|" implementations by remi-or · Pull Request #48931 · huggingface/transformers

remi-or · 2026-09-18T10:52:01Z

Draft

HuggingFaceDocBuilderDev · 2026-09-18T11:04:06Z

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.

vasqu

Ok long review but trying to summarize my key points here

The cache handling imo should happen at the same update where we call the base cache. It's a bit weird to hide this (with a different cache name). It would be nice if we really could just hide it within the update
Similar for the masking, I'd rather we hide this in the masking utils. We already pass the cache there so it should be unifiable
There are a lot of unrelated changes to refactor the code. In some cases it might be necessary but in others it's just to make the code refactored. In my eyes as reviewer it's hard to judge what is truly necessary now. Additionally, we risk potentially being very breaking, a lot of other libraries rely on our existing functions so renames or removals will be a high risk. Better to isolate what is truly necessary

vasqu · 2026-09-22T15:40:20Z

 _use_top_left_mask = flash_attn_supports_top_left_mask()


-def get_target_dtype(query: torch.Tensor, module: torch.nn.Module) -> torch.dtype:


removing this is breaking, any reason we don't keep the same order and just pass the extracted target dtype? (as before)

I just found the design weird TBH: target_dtype is computed here, but it is not used before entering the modeling_flash_attention_utils file, which itself calls a function to do the cast, as soon as _flash_attention_forward is entered. It seems to make sense just to cast the tensors before entering the _flash_attention_forward and to fuse get_target_dtype and the casting fn, both of which are only used once.
But I agree I went a little too fast and this is breaking, so reverting, but this could simplify the code a bit.

vasqu · 2026-09-22T15:40:56Z

-    # This might slowdown training & inference so it is recommended to not cast the LayerNorms
-    # in fp32. (usually our RMSNorm modules handle it correctly)
-    target_dtype = get_target_dtype(query, module)
+    # _flash_attention_forward needs non-transposed inputs, with shape [batch_size, seq_len, num_heads, head_dim]


Suggested change

# _flash_attention_forward needs non-transposed inputs, with shape [batch_size, seq_len, num_heads, head_dim]

# FA uses non-transposed inputs

let's keep it simple no?

If it's ok for you I prefer to leave some shape hints, it's always one of the most confusing part when reading the code IMO. Simplified the beginning tho.

vasqu · 2026-09-22T15:53:38Z

+@torch.compiler.disable
+def _uncompiled_flash_attention_forward(**kwargs) -> torch.Tensor:
+    return _flash_attention_forward(**kwargs)


I feel like im missing context here: why are we intentionally introducing a graph break here? Ideally the modeling flash attn utils should mark those portions only no? i assume it's due to paged FA doing something compile incompatible

Exactly. I think before I marked the whole function because I already had a graph break at the update, so adding a second graph break just around the attention actually made things worst. But now that the update is outside the second flash attention function, it makes less sense.

vasqu · 2026-09-22T15:54:37Z

-# Compile is disabled because the cache update mutates in place aliased views of the cache tensor, which compile's
-# functionalization handles by making a copy of the full cache for every layer.


ok here it is then, but I feel like we tackle this too early and modeling flash attn utils should differentiate between paged call or not and mark it only there. this feels to early / too preemptive

I see what you mean. Will run some benchmarks to see where it's best to put it.

vasqu · 2026-09-22T16:26:18Z

+    if attention_mask is not None:
+        padded_out = torch.zeros((batch_size * query_length, *out.shape[1:]), device=out.device, dtype=out.dtype)
+        padded_out[indices_q] = out
+        return padded_out.view(batch_size, query_length, *out.shape[1:])


this is weird as well was this _pad before?

Yes but it is called only once, and it is 3 lines of code, might as well inline it no?

vasqu · 2026-09-22T16:32:10Z

-    query_length: int,
-    is_causal: bool,
-    dropout: float = 0.0,
    position_ids: torch.Tensor | None = None,
-    softmax_scale: float | None = None,
-    sliding_window: int | None = None,
-    use_top_left_mask: bool = False,
-    softcap: float | None = None,
-    deterministic: bool | None = None,
    cu_seq_lens_q: torch.LongTensor | None = None,
    cu_seq_lens_k: torch.LongTensor | None = None,
    max_length_q: int | None = None,
    max_length_k: int | None = None,
-    target_dtype: torch.dtype | None = None,


A lot of breaking stuff happening here, changing the signature and all, also not sure if it's completely necessary. As mentioned before let's keep the diff to a minimum?

A lot of those were deleted because they are always passed as a kwarg, and only used in the processing kwarg. So it's exactly the same behavior, unless someone call this with only positionnal arguments. WDYT ?

vasqu · 2026-09-22T16:34:16Z

+    #         It is safe to use `flash_varlen_fn` knowing we already have all necessary the kwargs.
+
+    is_fa_with_varlen_kwargs = None not in (cu_seq_lens_q, cu_seq_lens_k, max_length_q, max_length_k)
+    is_fa_with_block_table = None not in (k_cache, v_cache, cache_seqlens, block_table)


Imo the only portion we need is a block after varlen and before the base fa where we check that these kwargs indeed exist

vasqu · 2026-09-22T16:34:55Z

    )

-    # Extract the flash attention kwargs that have been requested (and are supported by the implementation)
-    flash_kwargs = partial(


we delay this call either way so doesn't hurt to keep?

Yes but it is not used for the paged version. Although I could unify the two.

vasqu · 2026-09-22T16:35:42Z

We probably do need to freeze this tho to keep some deprecation cycle in (so not immediate delete)

github-actions · 2026-09-24T04:13:22Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: bark, falcon, gpt_neo, gptj

github-actions · 2026-09-24T04:18:48Z

CI recap

Dashboard: View test results in Grafana
Latest run: 35839207559:1
Result: failure | Jobs: 2 | Tests: 23 | Failures: 1 | Duration: 3m 3s

Code quality check failed: test jobs were skipped. Fix the code quality issues and push again to run tests.

remi-or · 2026-09-24T15:28:58Z

Not discarding this PR, but closing in favor of a stack of PR starting with #49069

remi-or added 5 commits September 17, 2026 22:24

Start of refacto

418e734

Nits

a572627

Remove ref to fa2 unpad bert fn

02f8616

Update FA

92670ed

Nits

d2912e1

remi-or added 12 commits September 21, 2026 11:21

Nit

4023ce2

Merge branch 'main' into merge-paged

f41cc8c

Everything but the processing

1d31315

Compile stuff

e492b6a

Remove flash paged attn

05b43ae

Review 1/n

184e6b1

Renames

911edf8

Flash should be good but bloated

9fdb0f1

SDPA TOO WOW

a2abbbe

Cleanup

aed5644

GTG

bf64e83

Simplification

d5656fe

vasqu reviewed Sep 22, 2026

View reviewed changes

remi-or and others added 2 commits September 23, 2026 21:12

Review

42eecd6

Merge branch 'main' into merge-paged

7105ef2

remi-or closed this Sep 24, 2026

		_use_top_left_mask = flash_attn_supports_top_left_mask()


		def get_target_dtype(query: torch.Tensor, module: torch.nn.Module) -> torch.dtype:

	# _flash_attention_forward needs non-transposed inputs, with shape [batch_size, seq_len, num_heads, head_dim]
	# FA uses non-transposed inputs

		# Compile is disabled because the cache update mutates in place aliased views of the cache tensor, which compile's
		# functionalization handles by making a copy of the full cache for every layer.

Conversation

remi-or commented Sep 18, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Sep 18, 2026

Uh oh!

vasqu left a comment

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!

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!

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!

github-actions Bot commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026

CI recap

Uh oh!

remi-or commented Sep 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

remi-or commented Sep 18, 2026 •

edited by github-actions Bot

Loading