[Llama4] [multimodal] Fix misplaced dtype cast of cos_sin_cache in Llama4VisionRotaryEmbedding

[cjackal]

Purpose

Fix #25888.

Especially, this PR roll back the misplaced dtype cast introduced by #21126 which deteriorates the vision understanding performance of llama 4 family.

Note that the cos_sin_cache in Llama4VisionRotaryEmbedding is always complex (torch.complex64) but the query is always real (torch.bfloat16), so we cannot cast the trigonometric cache to query's dtype.

Test Plan

Check that Llama 4 model families are correctly recognize the "4 dice" image used in the vllm multimodal test CI pipeline.

import httpx
from openai import OpenAI

client = OpenAI(base_url="http://localhost:8080/v1", api_key="a")

# This is the 4 dice image used in the vllm multimodal test CI pipeline
image_url = "https://raw.githubusercontent.com/vllm-project/vllm/refs/heads/main/tests/multimodal/assets/rgba.png"

messages = [{"role": "user", "content": [{"type": "text", "text": "Explain the image."}, {"type": "image_url", "image_url": {"url": image_url}}]}]

for t in client.chat.completions.create(model="meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8", messages=messages, stream=True):
    if t.choices:
        d = t.choices[0].delta
        print(getattr(d, "reasoning_content", "") or d.content or "", end="", flush=True)

Test Result

The image depicts four dice in mid-air, each with a distinct color and orientation.

*   The dice are arranged in a staggered formation, with the red die at the center and the other three dice positioned around it.
    *   The red die is in sharp focus, while the other three dice are slightly blurred.
    *   The blue die is located in the upper left corner of the image, the green die is in the upper right corner, and the yellow die is below the red die.
*   The dice have a reflective surface, giving them a shiny appearance.
    *   The dice are translucent, allowing the background to be visible through them.
    *   The dice have white dots on their faces, which represent the numbers 1-6.
*   The background of the image is solid white, providing a clean and simple backdrop for the dice.
    *   There are no other objects or features in the image besides the dice.

Overall, the image presents a visually appealing and dynamic representation of four dice in motion, with a focus on their colors and reflective surfaces.

---

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

[gemini-code-assist]

Code Review

This pull request addresses a performance degradation in Llama-4 vision models by preventing an unintended data type cast of the cos_sin_cache. The change correctly preserves the precision of the cache. My review includes a suggestion to further optimize the implementation by avoiding an unconditional and potentially expensive device transfer operation on every forward pass, which could introduce a performance regression.

[gemini-code-assist]

This line unconditionally moves self.cos_sin_cache to the query tensor's device on every forward pass. This operation can be expensive and is unnecessary if the cache is already on the correct device. To optimize performance, it's better to add a check to perform the device transfer only when needed, similar to the logic in the _match_cos_sin_cache_dtype method that was previously used here.

if self.cos_sin_cache.device != query.device:
    self.cos_sin_cache = self.cos_sin_cache.to(query.device)

[cjackal]

Somewhat a valid comment, but as it was originally in this form before #21126, I'll wait for the vllm maintainer's opinition.

[cjackal]

@mgoin @huydhn @yeqcharlotte Ping to the author of relevant PR(#21126) and maintainers on meta side that I am aware of.

[mgoin]

My bad, thanks for finding this. Can you leave a comment noting this intentionality?

Two other notes:

• we are missing unit tests for these custom rope classes if this wasn't caught in CI
• since this custom rope is only overriding the native impl, I think it is possible that a hardware backend with a forward_cuda/hip/etc impl in the base class may not be running the right code. We should make sure to force this

[cjackal]

My bad, thanks for finding this. Can you leave a comment noting this intentionality?

Two other notes:

• we are missing unit tests for these custom rope classes if this wasn't caught in CI
• since this custom rope is only overriding the native impl, I think it is possible that a hardware backend with a forward_cuda/hip/etc impl in the base class may not be running the right code. We should make sure to force this

Comment added!

1. From the look of the test filetree, we indeed do not have a dedicated test for Llama4VisionRope. Do we want a full E2E test under test_maverick.py (which sounds too heavy I think) or a lightweight unit test like test_rope_module_cache in test_pos_encoding.py?
2. I'm not familiar with non-CUDA hardwares, do you want something like the following?

    def forward_hip(  # type: ignore[override]
        self,
        query: torch.Tensor,
        key: Optional[torch.Tensor] = None,
    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
        return self.forward_native(query, key)

[DarkLight1337]

Seems this was changed by @mgoin in #21126

Edit: Sorry I didn't see your comments earlier

[ywang96]

I think non-cuda devices are already using forward_native unless explicitly overidden, so this fix looks good to me.

@yeqcharlotte @houseroad Is it possible for meta side to add some unit tests on this specific rope implementation so we can prevent this from happening in the future?

[cjackal]

@mgoin It seems this PR hasn't been picked in v0.11.0rc6, would you mind double-checking it?