[MM][Mistral] Add compile support for Pixtral vision encoders

[jzakrzew]

Purpose

Add compile_mm_encoder support for both Pixtral vision encoder implementations:

• PixtralForConditionalGeneration, including Mistral-format Mistral 3 models.
• The Hugging Face Pixtral encoder used by Mistral3ForConditionalGeneration.

The change applies support_torch_compile to the transformer blocks in both implementations and marks the token, rotary embedding, and packed-sequence dimensions as dynamic.

The feature remains gated behind compile_mm_encoder, which is disabled by default.

Performance results

Note: the results are only about GPU kernel time. Total E2E runtime improves much less, even with Cuda graphs from #56277, because apply_hf_processor is the main bottleneck now.

Profiling was performed on an NVIDIA GB200 system using the same 16-image BF16 capture before and after enabling compilation.
Total GPU kernel time decreased by 42%, from 1173.6 ms to 676.8 ms, corresponding to a 1.73x speedup.

The improvement came from eliminating or fusing much of execution overhead:

Kernel category	Before	After	Reduction
FlashAttention	273 ms	273 ms	0%
GEMM	272 ms	272 ms	0%
Elementwise operations	368 ms	11 ms	97%
SiLU	121 ms	44 ms	64%
Concat	96 ms	0.3 ms	>99%

CUDA-graph compatibility

This work is complementary to #56277, which adds Pixtral encoder CUDA-graph support. It does not duplicate or modify that implementation.

The CUDA-graph PR’s pixtral.py and mistral3.py changes apply cleanly over this change. Its packed encoder forwards continue to use the unchanged transformer interface. Encoder CUDA-graph capture also performs an eager warmup before capture,
allowing the transformer blocks to compile first.

A combined compile_mm_encoder=True plus cudagraph_mm_encoder=True integration test should be added once #56277 is merged.

Test Plan

Exercise both encoder implementations with compilation enabled and disabled:

• Mistral-format PixtralForConditionalGeneration
• Hugging Face Mistral3ForConditionalGeneration

The test uses dummy weights and verifies that the compilation counter increases only when compile_mm_encoder is enabled.

Test Result

  .venv/bin/python -m pytest \
  tests/compile/fullgraph/test_multimodal_compile.py \
  -k pixtral_compilation -vv -s

4 passed, 4 deselected

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[juliendenize]

LGTM thanks !

[github-actions]

✅ @jzakrzew, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[NickLucche]

/ci run

[NickLucche]

Thank you @jzakrzew !

[jzakrzew]

/ci run

[github-actions]

❌ This PR is 27 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[jzakrzew]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91555 for commit ea7cde312230.

[jzakrzew]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #91555.

[jzakrzew]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91862 for commit 8f51dc3188bb.

[jzakrzew]

/ci run

[github-actions]

✅ Triggered Buildkite CI #92093 for commit 92ee06225193.