Fix static-cache generate() with host-side inputs and under TP by dacorvo · Pull Request #49317 · huggingface/transformers

dacorvo · 2026-10-05T12:17:21Z

What does this PR do?

Fixes #49261, two generate() bugs with a static cache:

With inputs kept on another device, tensors nested in a model input (e.g. the mask dict keyed by layer type) stayed behind. They are now moved at any depth.
Under TP, the early static cache divided the KV heads by tp_size even when the key projection isn't sharded. It now checks the TP plan.

Note for #47811: colwise_units should count as sharded in get_kv_heads_per_rank. cc @kaixuanliu

Tests

test_generate_with_inputs_on_cpu_moves_tensors_in_dict_inputs (new)
test_get_kv_heads_per_rank (new)
test_chunked_prefill_initializes_static_cache_eagerly (extended: empty plan, gathered k_proj)

Also checked on Neuron, including Qwen3-VL-8B at TP=2.

Who can review?

@ArthurZucker @vasqu @Cyrilvallez @3outeille

With inputs kept on another device (#46590), `prepare_inputs_for_generation` moves only the top-level tensors of `model_inputs` onto the model device, so a tensor nested in a dict stayed behind (example: causal masks for some models). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…ithout TP Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…olds Under tensor parallelism `_get_static_cache_init_shape` divided the KV head count by `tp_size` whenever `tp_size > 1`, whether or not the plan shards attention. With no plan (Qwen3-VL), a gathered key projection (`colwise_gather_output`) or an expert-parallel plan, every rank keeps all heads, so the chunked-prefill cache is allocated with too few and the first `index_copy_` fails: "Incompatible shapes for broadcasting: (1, 8, 512, 128) and requested shape (1, 4, 512, 128)" for Qwen3-VL-8B-Instruct at tp=2. Only divide when the key projection resolves to `colwise`, the style that shards heads; skip early init when no key projection is found. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

HuggingFaceDocBuilderDev · 2026-10-05T12:28:52Z

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

Some initial comments, definitely see the overall issues but we can keep the PR description a bit shorter especially since the issue is already detailed

- Move tensors nested at any depth in dicts, lists and tuples to the device. - Read the KV heads per rank from the TP plan instead of scanning parameters. - Use a tiny random checkpoint in the nested inputs test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

vasqu

Smaller comments overall but cc @3outeille whenever you have time to sanity check as well

…jections into the TP util Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

kaixuanliu · 2026-10-08T01:31:23Z

@dacorvo Thx for the heads up! Not sure #47811 will be adopted, will keep this in mind and make related update if so.

dacorvo · 2026-10-12T07:05:42Z

gently pinging @3outeille

Cyrilvallez · 2026-10-13T13:27:28Z

+    def _move_to_device(self: "GenerativePreTrainedModel", value):
+        """Moves the tensors in `value`, including those nested in dicts, lists and tuples, onto the model device."""
+        if isinstance(value, torch.Tensor):
+            return value.to(self.device)
+        if isinstance(value, dict):
+            return {k: self._move_to_device(v) for k, v in value.items()}
+        if isinstance(value, (list, tuple)):
+            return type(value)(self._move_to_device(v) for v in value)
+        return value


I don't think it should be a method of GenerationMixin, better inside pytorch_utils or something as a separate function

Lol just saw @vasqu said the opposite before on a resolved comment 😆 IMO it should rather be a general util, let me know your thoughts and sorry for the back and forth 😅

I dont mind either way, I see the point of it being more general as well

OK, just tell me which you prefer, and i will lazily ask claude to move it 😅

Lets move it to something more general, e.g. torch utils or similar (not sure about the exact file name again)

Moved to pytorch_utils.py

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dacorvo · 2026-10-14T08:00:22Z

@vasqu @Cyrilvallez any blockers ?

github-actions · 2026-10-14T09:08:24Z

CI recap

Dashboard: View test results in Grafana
Latest run: 38531945519:1
Result: success | Jobs: 15 | Tests: 200,517 | Failures: 0 | Duration: 15h 44m

vasqu

Should be good, sorry for the delays

dacorvo and others added 3 commits October 5, 2026 09:33

refactor(generation): return early from the static cache init shape w…

aebd7ea

…ithout TP Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dacorvo mentioned this pull request Oct 5, 2026

Fix chunked prefill with multi-axis position ids (Qwen3-VL) #49319

Merged

vasqu reviewed Oct 5, 2026

View reviewed changes

Comment thread src/transformers/generation/utils.py Resolved

Comment thread src/transformers/generation/utils.py Resolved

Comment thread src/transformers/generation/utils.py Resolved

Comment thread tests/generation/test_utils.py Resolved

dacorvo and others added 3 commits October 5, 2026 16:53

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

5f4fdb6

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

1f4df27

dacorvo requested a review from vasqu October 6, 2026 06:47

vasqu reviewed Oct 6, 2026

View reviewed changes

Comment thread src/transformers/generation/utils.py Resolved

Comment thread src/transformers/distributed/tensor_parallel.py Resolved

dacorvo and others added 2 commits October 6, 2026 09:33

refactor: move the device helper into GenerationMixin and the key pro…

5fa702f

…jections into the TP util Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

5f847c3

dacorvo requested a review from vasqu October 6, 2026 09:52

dacorvo added 3 commits October 6, 2026 12:24

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

478c9ed

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

2318b93

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

0417709

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

ce90dec

vasqu requested a review from ArthurZucker October 8, 2026 13:43

dacorvo requested a review from 3outeille October 12, 2026 07:05

Cyrilvallez reviewed Oct 13, 2026

View reviewed changes

dacorvo and others added 2 commits October 13, 2026 15:52

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

5c7aea9

review: move the device helper to pytorch_utils

a51dd47

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

dacorvo requested a review from Cyrilvallez October 13, 2026 14:38

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

e874135

Merge branch 'main' into fix-generate-dict-masks-and-tp-cache-heads

f0ee84b

vasqu approved these changes Oct 14, 2026

View reviewed changes

dacorvo added this pull request to the merge queue Oct 14, 2026

Merged via the queue into main with commit 245167e Oct 14, 2026
116 checks passed

dacorvo deleted the fix-generate-dict-masks-and-tp-cache-heads branch October 14, 2026 13:36

digit50 mentioned this pull request Oct 17, 2026

generate() fails with KeyError/AttributeError when input_ids are off the model device #49699

Closed

Cyrilvallez mentioned this pull request Oct 19, 2026

Fix inputs dict casting in generate #49731

Merged

stevhliu mentioned this pull request Oct 23, 2026

[docs] Translation workflow #49889

Closed

Conversation

dacorvo commented Oct 5, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Tests

Who can review?

Uh oh!

HuggingFaceDocBuilderDev commented Oct 5, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

kaixuanliu commented Oct 8, 2026

Uh oh!

dacorvo commented Oct 12, 2026

Uh oh!

Cyrilvallez Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

Cyrilvallez Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

dacorvo Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

vasqu Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

dacorvo Oct 13, 2026

Choose a reason for hiding this comment

Uh oh!

dacorvo commented Oct 14, 2026

Uh oh!

github-actions Bot commented Oct 14, 2026

CI recap

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

dacorvo commented Oct 5, 2026 •

edited by github-actions Bot

Loading