[dacorvo]

What does this PR do?

With a static cache, q_offset is a tensor. Under torch.compile (torch ≥ 2.14), the condition in _ignore_causal_mask_sdpa tests its value, which breaks the graph (Unsupported: Data-dependent branching with fullgraph=True). A traced q_offset now materializes the mask instead.

In generate(), this affects chunked prefill, the only path that runs multi-token inputs through the compiled forward. For the first chunk of an unpadded prompt, prepare_inputs_for_generation skips the mask because the cache is empty, so the compiled forward rebuilds it and hits that condition.

Tests

• test_static_cache_prefill_mask_under_compile (new): fails on main, passes with the fix (torch 2.14.1).
• Compiled chunked prefill (static cache, fullgraph=True) fails on main with the error above; with the fix, it matches a dynamic-cache generate().

Who can review?

@vasqu @Cyrilvallez

[HuggingFaceDocBuilderDev]

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]

Hmm so this happens since forever tho right? Why is it only torch 2.14? Is this rather a regression on torch or something we could fix on cache side?

[vasqu]

cc @Cyrilvallez would love your take here

[dacorvo]

@vasqu good question, that made me wonder why I had not seen it before.
Turns out this only happens with torch ≥ 2.14 (now the default version shipped with the Neuron SDK).
I asked Claude to dig in, and it's a regression from #48975, revealed by the torch bump.

Before #48975: is_tracing(padding_mask) returned early under compile on all versions, so q_offset == 0 was never reached.

After #48975: the guard now uses is_torchdynamo_exporting(), which dynamo hard-codes to True under compile before torch 2.14, so older versions still return early. On ≥ 2.14 it only returns True during export and False under compile. This is why we don't return early and run into the data-dependent graph issue.

[vasqu]

Oh ok this is now revealed as torch 2.14 can properly use it during compile. That was a super recent PR

2 things tho

1. Let's move this close where we already check for padding mask tracing (and add a comment that this only happens under (static) caches that need a tensor object (otherwise its another graph break))
2. Let's just use is_tracing no? I think it indirectly checks for this to be a tensor like object no?

[dacorvo]

Let's move this close where we already check for padding mask tracing (and add a comment that this only happens under (static) caches that need a tensor object (otherwise its another graph break))

Done

Let's just use is_tracing no? I think it indirectly checks for this to be a tensor like object no?

I double checked: is_tracing(x) returns True as soon as is_torchdynamo_compiling() is True, without looking at x. So without the isinstance additional check, an int q_offset (and not only the Tensor q_offset we have when using a static cache) would make _ignore_causal_mask_sdpa return False and force the causal mask to be materialized, hurting performance even for non static caches.

[vasqu]

Just some smaller comments, this is indeed something that only got revealed through a recent PR

[vasqu]

Oh ok this is now revealed as torch 2.14 can properly use it during compile. That was a super recent PR

2 things tho

1. Let's move this close where we already check for padding mask tracing (and add a comment that this only happens under (static) caches that need a tensor object (otherwise its another graph break))
2. Let's just use is_tracing no? I think it indirectly checks for this to be a tensor like object no?

[github-actions]

---

CI recap

Dashboard: View test results in Grafana
Latest run: 37777671026:1
Result: success | Jobs: 16 | Tests: 198,274 | Failures: 0 | Duration: 12h 35m